Edit Instance Count Enabled

horizon/internal/externalcall/gcs_client.go

@@ -9,14 +9,15 @@ import (
 	"os"
 	"path"
 	"path/filepath"
-	"regexp"
 	"strings"
 	"sync"
 	"time"
 
 	"cloud.google.com/go/storage"
+	"github.com/Meesho/BharatMLStack/horizon/internal/predator/proto/protogen"
 	"github.com/rs/zerolog/log"
 	"google.golang.org/api/iterator"
+	"google.golang.org/protobuf/encoding/prototext"
 )
 
 type GCSClientInterface interface {
@@ -268,7 +269,10 @@ func (g *GCSClient) transferSingleConfigFile(objAttrs storage.ObjectAttrs, srcBu
 
 	// Replace model name
 	log.Info().Msgf("Processing config.pbtxt file: %s -> %s", objAttrs.Name, destObjectPath)
-	modified := replaceModelNameInConfig(content, destModelName)
+	modified, err := ReplaceModelNameInConfig(content, destModelName)
+	if err != nil {
+		return fmt.Errorf("failed to replace model name: %w", err)
+	}
 
 	// Upload modified content
 	destWriter := g.client.Bucket(destBucket).Object(destObjectPath).NewWriter(g.ctx)
@@ -464,49 +468,36 @@ func (g *GCSClient) TransferAndDeleteFolder(srcBucket, srcPath, srcModelName, de
 
 // replaceModelNameInConfig modifies only the top-level `name:` field in config.pbtxt content
 // It replaces only the first occurrence to avoid modifying nested names in inputs/outputs/instance_groups
-func replaceModelNameInConfig(data []byte, destModelName string) []byte {
-	content := string(data)
-	lines := strings.Split(content, "\n")
-
-	for i, line := range lines {
-		trimmed := strings.TrimSpace(line)
-		// Match top-level "name:" field - should be at the start of line (or minimal indentation)
-		// Skip nested names which are typically indented with 2+ spaces
-		if strings.HasPrefix(trimmed, "name:") {
-			// Check indentation: top-level fields have minimal/no indentation
-			leadingWhitespace := len(line) - len(strings.TrimLeft(line, " \t"))
-			// Skip if heavily indented (nested field)
-			if leadingWhitespace >= 2 {
-				continue
-			}
+func ReplaceModelNameInConfig(data []byte, destModelName string) ([]byte, error) {
+	if destModelName == "" {
+		return nil, fmt.Errorf("destination model name cannot be empty")
+	}
 
-			// Match the first occurrence of name: "value" pattern
-			namePattern := regexp.MustCompile(`name\s*:\s*"([^"]+)"`)
-			matches := namePattern.FindStringSubmatch(line)
-			if len(matches) > 1 {
-				oldModelName := matches[1]
-				// Replace only the FIRST occurrence to avoid replacing nested names
-				loc := namePattern.FindStringIndex(line)
-				if loc != nil {
-					// Replace only the matched portion (first occurrence)
-					before := line[:loc[0]]
-					matched := line[loc[0]:loc[1]]
-					after := line[loc[1]:]
-					// Replace the value inside quotes while preserving the "name:" format
-					valuePattern := regexp.MustCompile(`"([^"]+)"`)
-					valueReplaced := valuePattern.ReplaceAllString(matched, fmt.Sprintf(`"%s"`, destModelName))
-					lines[i] = before + valueReplaced + after
-				} else {
-					// Fallback: replace all (shouldn't happen with valid input)
-					lines[i] = namePattern.ReplaceAllString(line, fmt.Sprintf(`name: "%s"`, destModelName))
-				}
-				log.Info().Msgf("Replacing top-level model name in config.pbtxt: '%s' -> '%s'", oldModelName, destModelName)
-				break
-			}
-		}
+	var modelConfig protogen.ModelConfig
+
+	if err := prototext.Unmarshal(data, &modelConfig); err != nil {
+		return nil, fmt.Errorf("failed to parse config.pbtxt: %w", err)
 	}
 
-	return []byte(strings.Join(lines, "\n"))
+	oldModelName := modelConfig.Name
+	modelConfig.Name = destModelName
+
+	opts := prototext.MarshalOptions{
+		Multiline: true,
+		Indent:    "  ",
+	}
+
+	out, err := opts.Marshal(&modelConfig)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal updated config.pbtxt: %w", err)
+	}
+
+	log.Info().
+		Str("old_model_name", oldModelName).
+		Str("new_model_name", destModelName).
+		Msg("replaced top-level model name in config.pbtxt")
+
+	return out, nil
 }
 
 func (g *GCSClient) ListFolders(bucket, prefix string) ([]string, error) {

horizon/internal/externalcall/gcs_client_test.go

@@ -10,50 +10,60 @@ import (
 
 func TestReplaceModelNameInConfig(t *testing.T) {
 	tests := []struct {
-		name           string
-		data           []byte
-		destModelName  string
-		expectContains string
+		name          string
+		data          []byte
+		destModelName string
+		wantTopLevel  string
+		wantNested    string
+		expectError   bool
 	}{
 		{
 			name: "replaces top-level name only",
-			data: []byte(`name: "old_model"
+			data: []byte(`
+name: "old_model"
 instance_group {
-  name: "old_model"
+  name: "nested_model"
 }
 `),
-			destModelName:  "new_model",
-			expectContains: `name: "new_model"`,
+			destModelName: "new_model",
+			wantTopLevel:  `"new_model"`,
+			wantNested:    "nested_model",
 		},
 		{
-			name: "preserves nested name with indentation",
-			data: []byte(`name: "top_level"
-  instance_group {
-    name: "nested_name"
-  }
-`),
-			destModelName:  "replaced",
-			expectContains: `name: "replaced"`,
+			name:          "single line config",
+			data:          []byte(`name: "single_model"` + "\n"),
+			destModelName: "replaced_model",
+			wantTopLevel:  `"replaced_model"`,
 		},
 		{
-			name: "single line config",
-			data: []byte(`name: "single_model"` + "\n"),
-			destModelName:  "replaced_model",
-			expectContains: `name: "replaced_model"`,
+			name:          "empty dest model name returns error",
+			data:          []byte(`name: "some_model"`),
+			destModelName: "",
+			expectError:   true,
 		},
 		{
-			name: "no name field returns unchanged",
-			data: []byte(`platform: "tensorflow"
-version: 1
-`),
+			name:          "malformed pbtxt returns error",
+			data:          []byte(`name: "unclosed`),
 			destModelName: "any",
-			expectContains: `platform: "tensorflow"`,
+			expectError:   true,
 		},
 	}
+
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got := replaceModelNameInConfig(tt.data, tt.destModelName)
-			assert.Contains(t, string(got), tt.expectContains)
+			got, err := ReplaceModelNameInConfig(tt.data, tt.destModelName)
+
+			if tt.expectError {
+				require.Error(t, err)
+				return
+			}
+
+			require.NoError(t, err)
+			assert.Contains(t, string(got), tt.wantTopLevel)
+
+			if tt.wantNested != "" {
+				assert.Contains(t, string(got), tt.wantNested)
+			}
 		})
 	}
 }

horizon/internal/predator/handler/model.go

@@ -3,6 +3,8 @@ package handler
 import (
 	"encoding/json"
 	"time"
+
+	"github.com/Meesho/BharatMLStack/horizon/internal/predator/proto/protogen"
 )
 
 type Payload struct {
@@ -169,7 +171,7 @@ type ModelParamsResponse struct {
 	Backend                string           `json:"backend"`
 	DynamicBatchingEnabled bool             `json:"dynamic_batching_enabled"`
 	Platform               string           `json:"platform"`
-	EnsembleScheduling     *ModelEnsembling `json:"ensemble_scheduling,omitempty"`
+	EnsembleScheduling     *protogen.ModelEnsembling `json:"ensemble_scheduling,omitempty"`
 }
 
 type RequestGenerationRequest struct {

horizon/internal/predator/handler/predator.go

@@ -36,6 +36,7 @@ import (
 	"github.com/Meesho/BharatMLStack/horizon/pkg/random"
 	"github.com/Meesho/BharatMLStack/horizon/pkg/serializer"
 	"github.com/rs/zerolog/log"
+	"github.com/Meesho/BharatMLStack/horizon/internal/predator/proto/protogen"
 )
 
 type Predator struct {
@@ -353,7 +354,7 @@ func (p *Predator) FetchModelConfig(req FetchModelConfigRequest) (ModelParamsRes
 		}
 	}
 
-	var modelConfig ModelConfig
+	var modelConfig protogen.ModelConfig
 	if err := prototext.Unmarshal(configData, &modelConfig); err != nil {
 		return ModelParamsResponse{}, http.StatusInternalServerError, fmt.Errorf(errUnmarshalProtoFormat, err)
 	}
@@ -391,7 +392,7 @@ func parseModelPath(modelPath string) (bucket, objectPath string) {
 	return parts[0], parts[1]
 }
 
-func validateModelConfig(cfg *ModelConfig) error {
+func validateModelConfig(cfg *protogen.ModelConfig) error {
 	switch {
 	case cfg.Name == constant.EmptyString:
 		return errors.New(errModelNameMissing)
@@ -423,7 +424,7 @@ func convertFields(name string, dims []int64, dataType string) (IO, bool) {
 	}, true
 }
 
-func convertInputWithFeatures(fields []*ModelInput, featureMap map[string][]string) []IO {
+func convertInputWithFeatures(fields []*protogen.ModelInput, featureMap map[string][]string) []IO {
 	ios := make([]IO, 0, len(fields))
 	for _, f := range fields {
 		if io, ok := convertFields(f.Name, f.Dims, f.DataType.String()); ok {
@@ -437,7 +438,7 @@ func convertInputWithFeatures(fields []*ModelInput, featureMap map[string][]stri
 	return ios
 }
 
-func convertOutput(fields []*ModelOutput) []IO {
+func convertOutput(fields []*protogen.ModelOutput) []IO {
 	ios := make([]IO, 0, len(fields))
 	for _, f := range fields {
 		if io, ok := convertFields(f.Name, f.Dims, f.DataType.String()); ok {
@@ -447,7 +448,7 @@ func convertOutput(fields []*ModelOutput) []IO {
 	return ios
 }
 
-func createModelParamsResponse(modelConfig *ModelConfig, objectPath string, inputs, outputs []IO) ModelParamsResponse {
+func createModelParamsResponse(modelConfig *protogen.ModelConfig, objectPath string, inputs, outputs []IO) ModelParamsResponse {
 	var resp ModelParamsResponse
 
 	if len(modelConfig.InstanceGroup) > 0 {

horizon/internal/predator/handler/predator_approval.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"path"
 	"strings"
 	"time"
 
@@ -13,7 +14,9 @@ import (
 	"github.com/Meesho/BharatMLStack/horizon/internal/repositories/sql/predatorconfig"
 	"github.com/Meesho/BharatMLStack/horizon/internal/repositories/sql/predatorrequest"
 	"github.com/rs/zerolog/log"
+	"google.golang.org/protobuf/encoding/prototext"
 	"gorm.io/gorm"
+	"github.com/Meesho/BharatMLStack/horizon/internal/predator/proto/protogen"
 )
 
 func (p *Predator) processRequest(requestIdPayloadMap map[uint]*Payload, predatorRequestList []predatorrequest.PredatorRequest, req ApproveRequest) {
@@ -191,6 +194,12 @@ func (p *Predator) processEditGCSCopyStage(requestIdPayloadMap map[uint]*Payload
 		} else {
 			configBucket := pred.GcsConfigBucket
 			configPath := pred.GcsConfigBasePath
+			if configBucket != "" && configPath != "" && payload.MetaData.InstanceCount > 0 {
+				if err := p.updateInstanceCountInConfigSource(configBucket, configPath, modelName, payload.MetaData.InstanceCount); err != nil {
+					log.Error().Err(err).Msgf("Failed to update instance count in config-source for model %s", modelName)
+					return transferredGcsModelData, err
+				}
+			}
 			if err := p.GcsClient.TransferFolderWithSplitSources(
 				sourceBucket, sourceBasePath, configBucket, configPath,
 				sourceModelName, targetBucket, targetPath, modelName,
@@ -212,6 +221,53 @@ func (p *Predator) processEditGCSCopyStage(requestIdPayloadMap map[uint]*Payload
 	return transferredGcsModelData, nil
 }
 
+func (p *Predator) updateInstanceCountInConfigSource(bucket, basePath, modelName string, instanceCount int) error {
+	if modelName == "" {
+		return fmt.Errorf("model name is empty, required to update instance count in config-source")
+	}
+
+	configPath := path.Join(basePath, modelName, configFile)
+	configData, err := p.GcsClient.ReadFile(bucket, configPath)
+	if err != nil {
+		return fmt.Errorf("failed to read config.pbtxt from config-source for model %s: %w", modelName, err)
+	}
+
+	var modelConfig protogen.ModelConfig
+	if err := prototext.Unmarshal(configData, &modelConfig); err != nil {
+		return fmt.Errorf("failed to parse config.pbtxt from config-source for model %s: %w", modelName, err)
+	}
+	if len(modelConfig.InstanceGroup) == 0 {
+		return fmt.Errorf("%s (model %s)", errNoInstanceGroup, modelName)
+	}
+
+	currentCount := modelConfig.InstanceGroup[0].Count
+	if currentCount == int32(instanceCount) {
+		log.Info().
+			Str("model", modelName).
+			Int("instance_count", instanceCount).
+			Msg("instance_count unchanged, skipping config update")
+		return nil
+	}
+
+	modelConfig.InstanceGroup[0].Count = int32(instanceCount)
+
+	opts := prototext.MarshalOptions{
+		Multiline: true,
+		Indent:    "  ",
+	}
+
+	newConfigData, err := opts.Marshal(&modelConfig)
+	if err != nil {
+		return fmt.Errorf("failed to marshal config.pbtxt for model %s: %w", modelName, err)
+	}
+	if err := p.GcsClient.UploadFile(bucket, configPath, newConfigData); err != nil {
+		return fmt.Errorf("failed to upload config.pbtxt to config-source for model %s: %w", modelName, err)
+	}
+
+	log.Info().Msgf("Updated instance_count to %d in config-source for model %s", instanceCount, modelName)
+	return nil
+}
+
 // processEditDBUpdateStage updates predator config for edit approval
 // This updates the existing predator config with new config.pbtxt and metadata.json changes
 func (p *Predator) processEditDBUpdateStage(requestIdPayloadMap map[uint]*Payload, predatorRequestList []predatorrequest.PredatorRequest, approvedBy string) error {

horizon/internal/predator/handler/predator_upload.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/Meesho/BharatMLStack/horizon/internal/externalcall"
 	pred "github.com/Meesho/BharatMLStack/horizon/internal/predator"
+	"github.com/Meesho/BharatMLStack/horizon/internal/predator/proto/protogen"
 	"github.com/rs/zerolog/log"
 	"google.golang.org/protobuf/encoding/prototext"
 )
@@ -88,8 +89,10 @@ func (p *Predator) copyConfigToProdConfigSource(gcsPath, modelName string) error
 		return fmt.Errorf("failed to read config.pbtxt from source: %w", err)
 	}
 
-	updatedConfigData := p.replaceModelNameInConfigPreservingFormat(configData, modelName)
-
+	updatedConfigData, err := externalcall.ReplaceModelNameInConfig(configData, modelName)
+	if err != nil {
+		return fmt.Errorf("failed to replace model name in config: %w", err)
+	}
 	destConfigPath := path.Join(pred.GcsConfigBasePath, modelName, configFile)
 	if err := p.GcsClient.UploadFile(pred.GcsConfigBucket, destConfigPath, updatedConfigData); err != nil {
 		return fmt.Errorf("failed to upload config.pbtxt to config source: %w", err)
@@ -437,7 +440,7 @@ func (p *Predator) validateModelConfiguration(gcsPath string) error {
 		return fmt.Errorf("failed to read config.pbtxt from %s/%s: %w", srcBucket, configPath, err)
 	}
 
-	var modelConfig ModelConfig
+	var modelConfig protogen.ModelConfig
 	if err := prototext.Unmarshal(configData, &modelConfig); err != nil {
 		return fmt.Errorf("failed to parse config.pbtxt as proto: %w", err)
 	}
@@ -481,8 +484,10 @@ func (p *Predator) copyConfigToNewNameInConfigSource(oldModelName, newModelName
 		return fmt.Errorf("failed to read config.pbtxt from %s: %w", srcConfigPath, err)
 	}
 
-	updatedConfigData := p.replaceModelNameInConfigPreservingFormat(configData, newModelName)
-
+	updatedConfigData, err := externalcall.ReplaceModelNameInConfig(configData, newModelName)
+	if err != nil {
+		return fmt.Errorf("failed to replace model name in config: %w", err)
+	}
 	if err := p.GcsClient.UploadFile(pred.GcsConfigBucket, destConfigPath, updatedConfigData); err != nil {
 		return fmt.Errorf("failed to upload config.pbtxt to %s: %w", destConfigPath, err)
 	}