diff --git a/api/endpoint.go b/api/endpoint.go index 1a34160b..25b4a7d8 100644 --- a/api/endpoint.go +++ b/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"` diff --git a/cmd/project/functions.go b/cmd/project/functions.go index c8cffc50..5de1d300 100644 --- a/cmd/project/functions.go +++ b/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") diff --git a/cmd/project/tomlBuilder.go b/cmd/project/tomlBuilder.go index 8fca466b..0662ab58 100644 --- a/cmd/project/tomlBuilder.go +++ b/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. # diff --git a/cmd/serverless/update.go b/cmd/serverless/update.go index 3cd6d05a..b9bab9f2 100644 --- a/cmd/serverless/update.go +++ b/cmd/serverless/update.go @@ -14,19 +14,34 @@ import ( var updateCmd = &cobra.Command{ Use: "update ", 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 --name my-endpoint + + # set model references (replaces existing) + runpodctl serverless update --model-reference https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct:main + + # attach multiple models + runpodctl serverless update --model-reference --model-reference + + # clear all model references + runpodctl serverless update --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) diff --git a/cmd/serverless/update_test.go b/cmd/serverless/update_test.go index 85a27d6a..d1ae4fda 100644 --- a/cmd/serverless/update_test.go +++ b/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"]) + } +} diff --git a/internal/api/endpoints.go b/internal/api/endpoints.go index 045bfbcb..a0c0d186 100644 --- a/internal/api/endpoints.go +++ b/internal/api/endpoints.go @@ -188,6 +188,106 @@ func (c *Client) DeleteEndpoint(endpointID string) error { return err } +// UpdateEndpointModels sets the model references on an existing endpoint via +// saveEndpoint. The full current config is round-tripped so that only +// modelReferences changes — saveEndpoint is a full replace and omitting fields +// would reset them to server defaults. Pass nil or an empty slice to clear all +// model references. +func (c *Client) UpdateEndpointModels(endpointID string, modelRefs []string) (*Endpoint, error) { + endpoint, err := c.GetEndpoint(endpointID, false, false) + if err != nil { + return nil, fmt.Errorf("failed to fetch endpoint: %w", err) + } + + if modelRefs == nil { + modelRefs = []string{} + } + + // saveEndpoint expects networkVolumeIds as [{networkVolumeId}] objects; the + // REST read shape uses bare id strings which UnmarshalJSON normalises into + // EndpointNetworkVolume — convert to the GraphQL write shape here. + nvIDs := make([]NetworkVolumeIDInput, len(endpoint.NetworkVolumeIDs)) + for i, nv := range endpoint.NetworkVolumeIDs { + nvIDs[i] = NetworkVolumeIDInput{NetworkVolumeID: nv.NetworkVolumeID} + } + + query := ` + mutation SaveEndpoint($input: EndpointInput!) { + saveEndpoint(input: $input) { + id + name + templateId + gpuIds + gpuCount + instanceIds + workersMin + workersMax + locations + networkVolumeId + networkVolumeIds { networkVolumeId } + idleTimeout + scalerType + scalerValue + executionTimeoutMs + minCudaVersion + flashBootType + modelReferences + } + } + ` + + variables := map[string]interface{}{ + "input": map[string]interface{}{ + "id": endpointID, + "name": endpoint.Name, + "templateId": endpoint.TemplateID, + "gpuIds": endpoint.GpuIDs, + "gpuCount": endpoint.GpuCount, + "instanceIds": endpoint.InstanceIDs, + "workersMin": endpoint.WorkersMin, + "workersMax": endpoint.WorkersMax, + "locations": endpoint.Locations, + "networkVolumeId": endpoint.NetworkVolumeID, + "networkVolumeIds": nvIDs, + "idleTimeout": endpoint.IdleTimeout, + "scalerType": endpoint.ScalerType, + "scalerValue": endpoint.ScalerValue, + "executionTimeoutMs": endpoint.ExecutionTimeoutMs, + "minCudaVersion": endpoint.MinCudaVersion, + "flashBootType": endpoint.FlashBootType, + "modelReferences": modelRefs, + }, + } + + data, err := c.graphqlRequest(query, variables) + if err != nil { + return nil, err + } + + var resp struct { + Data struct { + SaveEndpoint *Endpoint `json:"saveEndpoint"` + } `json:"data"` + Errors []struct { + Message string `json:"message"` + } `json:"errors"` + } + + if err := json.Unmarshal(data, &resp); err != nil { + return nil, fmt.Errorf("failed to parse response: %w", err) + } + + if len(resp.Errors) > 0 { + return nil, fmt.Errorf("graphql error: %s", resp.Errors[0].Message) + } + + if resp.Data.SaveEndpoint == nil { + return nil, fmt.Errorf("update returned nil response") + } + + return resp.Data.SaveEndpoint, nil +} + // NetworkVolumeIDInput is a single multi-region network volume entry for the // graphql saveEndpoint mutation (rest uses a flat []string instead). type NetworkVolumeIDInput struct { diff --git a/internal/api/endpoints_test.go b/internal/api/endpoints_test.go index b5f17065..e0906b3a 100644 --- a/internal/api/endpoints_test.go +++ b/internal/api/endpoints_test.go @@ -291,6 +291,99 @@ func TestUpdateEndpointTemplate_GraphQLError(t *testing.T) { } } +func TestUpdateEndpointModels_RoundTripsConfig(t *testing.T) { + oldAPIURL := viper.GetString("apiUrl") + t.Cleanup(func() { viper.Set("apiUrl", oldAPIURL) }) + + var gqlInput map[string]interface{} + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/endpoints/"): + // REST read: return wire shape with custom config and a bare-string networkVolumeId. + w.Write([]byte(`{ + "id": "ep-abc", + "name": "my-ep", + "templateId": "tpl-1", + "gpuIds": "ADA_24", + "workersMin": 1, + "workersMax": 5, + "idleTimeout": 42, + "scalerType": "REQUEST_COUNT", + "scalerValue": 9, + "networkVolumeId": "vol-9", + "networkVolumeIds": ["vol-9"] + }`)) + case r.Method == http.MethodPost: + // GraphQL saveEndpoint call. + var body struct { + Variables struct { + Input map[string]interface{} `json:"input"` + } `json:"variables"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatalf("decode gql body: %v", err) + } + gqlInput = body.Variables.Input + json.NewEncoder(w).Encode(map[string]interface{}{ + "data": map[string]interface{}{ + "saveEndpoint": map[string]interface{}{ + "id": "ep-abc", + "name": "my-ep", + }, + }, + }) + default: + t.Fatalf("unexpected request: %s %s", r.Method, r.URL.Path) + } + })) + defer server.Close() + + t.Setenv("RUNPOD_API_KEY", "test-key") + viper.Set("apiUrl", server.URL) + + client, _ := NewClient() + client.baseURL = server.URL + + _, err := client.UpdateEndpointModels("ep-abc", []string{"https://huggingface.co/org/model:main"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // All existing config fields must be present in the mutation input. + checks := map[string]interface{}{ + "id": "ep-abc", + "name": "my-ep", + "templateId": "tpl-1", + "gpuIds": "ADA_24", + "workersMax": float64(5), + "idleTimeout": float64(42), + "scalerType": "REQUEST_COUNT", + "scalerValue": float64(9), + } + for field, want := range checks { + if got := gqlInput[field]; got != want { + t.Errorf("input.%s = %v, want %v", field, got, want) + } + } + + // modelReferences must carry the new value, not the old one. + refs, _ := gqlInput["modelReferences"].([]interface{}) + if len(refs) != 1 || refs[0] != "https://huggingface.co/org/model:main" { + t.Errorf("unexpected modelReferences: %v", gqlInput["modelReferences"]) + } + + // networkVolumeIds must be passed as objects, not bare strings. + nvids, _ := gqlInput["networkVolumeIds"].([]interface{}) + if len(nvids) != 1 { + t.Fatalf("expected 1 networkVolumeId, got %v", gqlInput["networkVolumeIds"]) + } + nvobj, _ := nvids[0].(map[string]interface{}) + if nvobj["networkVolumeId"] != "vol-9" { + t.Errorf("expected networkVolumeId vol-9, got %v", nvobj) + } +} + func TestDeleteEndpoint(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodDelete {