fix: add modelReferences to endpoint create/update and document in runpod.toml template

api/endpoint.go

@@ -33,12 +33,13 @@ type CreateEndpointInput struct {
 	WorkersMin      int    `json:"workersMin"`
 	WorkersMax      int    `json:"workersMax"`
 	FlashBootType   string `json:"flashBootType"`
+	ModelReferences []string `json:"modelReferences"`
 }
 
 // there are many more fields in the result of the query but I just care about these for CLI port
 type Endpoint struct {
 	Name string `json:"name"`
-	Id   string
+	Id   string `json:"id"`
 }
 type EndpointOut struct {
 	Data   *EndpointData   `json:"data"`

cmd/project/functions.go

@@ -581,6 +581,7 @@ func deployProject(networkVolumeId string) (endpointId string, err error) {
 	flashBootType := "FLASHBOOT"
 	idleTimeout := 5
 	endpointConfig, ok := config.Get("endpoint").(*toml.Tree)
+	var modelRefs []string
 	if ok {
 		if min, ok := endpointConfig.Get("active_workers").(int64); ok {
 			minWorkers = int(min)
@@ -597,6 +598,13 @@ func deployProject(networkVolumeId string) (endpointId string, err error) {
 		if idle, ok := endpointConfig.Get("idle_timeout").(int64); ok {
 			idleTimeout = int(idle)
 		}
+		if refs, ok := endpointConfig.Get("model_refs").([]interface{}); ok {
+			for _, r := range refs {
+				if str, ok := r.(string); ok {
+					modelRefs = append(modelRefs, str)
+				}
+			}
+		}
 	}
 	if err != nil {
 		deployedEndpointId, err = api.CreateEndpoint(&api.CreateEndpointInput{
@@ -610,6 +618,7 @@ func deployProject(networkVolumeId string) (endpointId string, err error) {
 			WorkersMin:      minWorkers,
 			WorkersMax:      maxWorkers,
 			FlashBootType:   flashBootType,
+			ModelReferences: modelRefs,
 		})
 		if err != nil {
 			fmt.Println("error making endpoint")

cmd/project/tomlBuilder.go

@@ -77,6 +77,11 @@ active_workers = 0
 max_workers = 3
 flashboot = true
 
+# model_refs - List of model references to cache on the endpoint workers.
+#              Format: "owner/model-name" or "owner/model-name:branch".
+#              Example: ["runpod/stable-diffusion-v1-5", "meta-llama/Llama-2-7b-chat-hf"]
+# model_refs = []
+
 [runtime]
 # python_version    - Python version to use for the project.
 #

cmd/serverless/update.go

@@ -14,19 +14,34 @@ import (
 var updateCmd = &cobra.Command{
 	Use:   "update <endpoint-id>",
 	Short: "update an endpoint",
-	Long:  "update an existing serverless endpoint",
-	Args:  cobra.ExactArgs(1),
-	RunE:  runUpdate,
+	Long: `update an existing serverless endpoint.
+
+examples:
+  # rename an endpoint
+  runpodctl serverless update <id> --name my-endpoint
+
+  # set model references (replaces existing)
+  runpodctl serverless update <id> --model-reference https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct:main
+
+  # attach multiple models
+  runpodctl serverless update <id> --model-reference <ref-a> --model-reference <ref-b>
+
+  # clear all model references
+  runpodctl serverless update <id> --clear-models`,
+	Args: cobra.ExactArgs(1),
+	RunE: runUpdate,
 }
 
 var (
-	updateName        string
-	updateTemplateID  string
-	updateWorkersMin  int
-	updateWorkersMax  int
-	updateIdleTimeout int
-	updateScaleBy     string
+	updateName           string
+	updateTemplateID     string
+	updateWorkersMin     int
+	updateWorkersMax     int
+	updateIdleTimeout    int
+	updateScaleBy        string
 	updateScaleThreshold int
+	updateModelRefs      []string
+	updateClearModels    bool
 )
 
 func init() {
@@ -37,11 +52,17 @@ func init() {
 	updateCmd.Flags().IntVar(&updateIdleTimeout, "idle-timeout", -1, "new idle timeout in seconds")
 	updateCmd.Flags().StringVar(&updateScaleBy, "scale-by", "", "autoscale strategy: delay (seconds of queue wait) or requests (pending request count)")
 	updateCmd.Flags().IntVar(&updateScaleThreshold, "scale-threshold", -1, "trigger point for autoscaler (delay: seconds, requests: count)")
+	updateCmd.Flags().StringArrayVar(&updateModelRefs, "model-reference", nil, "model reference to cache on the endpoint (repeatable); replaces existing model references")
+	updateCmd.Flags().BoolVar(&updateClearModels, "clear-models", false, "remove all model references from the endpoint")
 }
 
 func runUpdate(cmd *cobra.Command, args []string) error {
 	endpointID := args[0]
 
+	if updateClearModels && len(updateModelRefs) > 0 {
+		return fmt.Errorf("--clear-models and --model-reference are mutually exclusive")
+	}
+
 	client, err := api.NewClient()
 	if err != nil {
 		return err
@@ -97,6 +118,16 @@ func runUpdate(cmd *cobra.Command, args []string) error {
 		}
 	}
 
+	if len(updateModelRefs) > 0 || updateClearModels {
+		var refs []string
+		if !updateClearModels {
+			refs = updateModelRefs
+		}
+		if _, err := client.UpdateEndpointModels(endpointID, refs); err != nil {
+			return fmt.Errorf("failed to update model references: %w", err)
+		}
+	}
+
 	endpoint, err := client.GetEndpoint(endpointID, false, false)
 	if err != nil {
 		return fmt.Errorf("failed to get updated endpoint: %w", err)

cmd/serverless/update_test.go

@@ -20,6 +20,18 @@ func TestUpdateCmd_HasTemplateIDFlag(t *testing.T) {
 	}
 }
 
+func TestUpdateCmd_HasModelReferenceFlag(t *testing.T) {
+	if flag := updateCmd.Flags().Lookup("model-reference"); flag == nil {
+		t.Fatal("expected model-reference flag")
+	}
+}
+
+func TestUpdateCmd_HasClearModelsFlag(t *testing.T) {
+	if flag := updateCmd.Flags().Lookup("clear-models"); flag == nil {
+		t.Fatal("expected clear-models flag")
+	}
+}
+
 func captureStderr(t *testing.T, fn func()) string {
 	t.Helper()
 
@@ -43,23 +55,32 @@ func captureStderr(t *testing.T, fn func()) string {
 	return string(data)
 }
 
-func TestRunUpdate_WarnsWhenTemplateSwapFailsAfterRESTUpdate(t *testing.T) {
+func resetUpdateVars(t *testing.T) {
+	t.Helper()
 	origName := updateName
 	origTemplateID := updateTemplateID
 	origWorkersMin := updateWorkersMin
 	origWorkersMax := updateWorkersMax
 	origIdleTimeout := updateIdleTimeout
-	origScalerType := updateScaleBy
-	origScalerValue := updateScaleThreshold
+	origScaleBy := updateScaleBy
+	origScaleThreshold := updateScaleThreshold
+	origModelRefs := updateModelRefs
+	origClearModels := updateClearModels
 	t.Cleanup(func() {
 		updateName = origName
 		updateTemplateID = origTemplateID
 		updateWorkersMin = origWorkersMin
 		updateWorkersMax = origWorkersMax
 		updateIdleTimeout = origIdleTimeout
-		updateScaleBy = origScalerType
-		updateScaleThreshold = origScalerValue
+		updateScaleBy = origScaleBy
+		updateScaleThreshold = origScaleThreshold
+		updateModelRefs = origModelRefs
+		updateClearModels = origClearModels
 	})
+}
+
+func TestRunUpdate_WarnsWhenTemplateSwapFailsAfterRESTUpdate(t *testing.T) {
+	resetUpdateVars(t)
 
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch {
@@ -102,6 +123,8 @@ func TestRunUpdate_WarnsWhenTemplateSwapFailsAfterRESTUpdate(t *testing.T) {
 	updateIdleTimeout = -1
 	updateScaleBy = ""
 	updateScaleThreshold = -1
+	updateModelRefs = nil
+	updateClearModels = false
 
 	cmd := &cobra.Command{}
 	cmd.Flags().String("output", "json", "")
@@ -124,3 +147,105 @@ func TestRunUpdate_WarnsWhenTemplateSwapFailsAfterRESTUpdate(t *testing.T) {
 		t.Fatalf("expected no json error output, got %q", stderr)
 	}
 }
+
+func TestRunUpdate_ClearModelsAndModelReferenceMutuallyExclusive(t *testing.T) {
+	resetUpdateVars(t)
+
+	updateModelRefs = []string{"https://huggingface.co/org/model:main"}
+	updateClearModels = true
+	updateWorkersMin = -1
+	updateWorkersMax = -1
+	updateIdleTimeout = -1
+	updateScaleThreshold = -1
+
+	cmd := &cobra.Command{}
+	cmd.Flags().String("output", "json", "")
+
+	err := runUpdate(cmd, []string{"ep-123"})
+	if err == nil {
+		t.Fatal("expected error for mutually exclusive flags")
+	}
+	if !strings.Contains(err.Error(), "mutually exclusive") {
+		t.Fatalf("unexpected error: %v", err)
+	}
+}
+
+func TestRunUpdate_ModelReferences(t *testing.T) {
+	resetUpdateVars(t)
+
+	var gqlBody map[string]interface{}
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch {
+		case r.Method == http.MethodGet && r.URL.Path == "/endpoints/ep-123":
+			// serve raw rest wire shape with non-default config to catch round-trip regressions.
+			_, _ = w.Write([]byte(`{
+				"id":          "ep-123",
+				"name":        "my-endpoint",
+				"idleTimeout": 42,
+				"scalerType":  "REQUEST_COUNT",
+				"scalerValue": 9,
+				"workersMax":  5
+			}`))
+		case r.Method == http.MethodPost && r.URL.Path == "/":
+			if err := json.NewDecoder(r.Body).Decode(&gqlBody); err != nil {
+				t.Fatalf("decode gql request: %v", err)
+			}
+			_ = json.NewEncoder(w).Encode(map[string]interface{}{
+				"data": map[string]interface{}{
+					"saveEndpoint": map[string]interface{}{
+						"id":              "ep-123",
+						"name":            "my-endpoint",
+						"modelReferences": []string{"https://huggingface.co/org/model:main"},
+					},
+				},
+			})
+		default:
+			t.Fatalf("unexpected request: %s %s", r.Method, r.URL.Path)
+		}
+	}))
+	defer server.Close()
+
+	t.Setenv("RUNPOD_API_KEY", "test-key")
+	viper.Set("restApiUrl", server.URL)
+	viper.Set("apiUrl", server.URL)
+	t.Cleanup(func() {
+		viper.Set("restApiUrl", "")
+		viper.Set("apiUrl", "")
+	})
+
+	updateModelRefs = []string{"https://huggingface.co/org/model:main"}
+	updateClearModels = false
+	updateWorkersMin = -1
+	updateWorkersMax = -1
+	updateIdleTimeout = -1
+	updateScaleThreshold = -1
+
+	cmd := &cobra.Command{}
+	cmd.Flags().String("output", "json", "")
+
+	err := runUpdate(cmd, []string{"ep-123"})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	vars, _ := gqlBody["variables"].(map[string]interface{})
+	input, _ := vars["input"].(map[string]interface{})
+
+	// model references must carry the new value.
+	refs, _ := input["modelReferences"].([]interface{})
+	if len(refs) != 1 || refs[0] != "https://huggingface.co/org/model:main" {
+		t.Fatalf("expected modelReferences to contain the provided ref, got %#v", refs)
+	}
+
+	// existing config must be round-tripped, not reset to defaults.
+	if input["idleTimeout"] != float64(42) {
+		t.Errorf("idleTimeout not round-tripped: got %v", input["idleTimeout"])
+	}
+	if input["scalerValue"] != float64(9) {
+		t.Errorf("scalerValue not round-tripped: got %v", input["scalerValue"])
+	}
+	if input["workersMax"] != float64(5) {
+		t.Errorf("workersMax not round-tripped: got %v", input["workersMax"])
+	}
+}