From bcdf2ddd195d29051543d5b6ee174fe96242060a Mon Sep 17 00:00:00 2001 From: Arpit Agarwal Date: Sun, 15 Nov 2026 05:30:24 +0000 Subject: [PATCH 1/4] air: add a hidden `air deploy` command that registers and serves a model `databricks air deploy -f deploy.yaml` registers a Unity Catalog model version through the registerModelPackage API (a checkpoint in a volume, a model from code, a model directory saved by MLflow) and serves it on a Model Serving endpoint, or serves an existing version. Registration is idempotent per key; the CLI repeats the same request until the version is ready, and prints the command that continues after Ctrl-C or --timeout. The command is hidden until the API is on everywhere. The run YAML loader, --override, -h config., code_source file selection and root path resolution are shared with `air run`. Co-authored-by: Isaac --- cmd/air/air.go | 2 +- cmd/air/air_test.go | 2 +- cmd/air/deploy.go | 382 ++++++++++++++++++++++++++++ cmd/air/deploy_test.go | 366 +++++++++++++++++++++++++++ cmd/air/deploycode.go | 155 ++++++++++++ cmd/air/deploycode_test.go | 183 ++++++++++++++ cmd/air/deployconfig.go | 387 +++++++++++++++++++++++++++++ cmd/air/deployconfig_test.go | 306 +++++++++++++++++++++++ cmd/air/deployendpoint.go | 186 ++++++++++++++ cmd/air/deployendpoint_test.go | 280 +++++++++++++++++++++ cmd/air/deployregister.go | 172 +++++++++++++ cmd/air/deployregister_test.go | 165 ++++++++++++ cmd/air/run.go | 4 +- cmd/air/runconfig.go | 47 ++-- cmd/air/runconfig_load.go | 33 ++- cmd/air/runconfig_override.go | 7 +- cmd/air/runconfig_override_test.go | 2 +- cmd/air/runconfig_test.go | 6 +- cmd/air/runsubmit.go | 8 +- cmd/air/runsubmit_test.go | 6 +- 20 files changed, 2658 insertions(+), 41 deletions(-) create mode 100644 cmd/air/deploy.go create mode 100644 cmd/air/deploy_test.go create mode 100644 cmd/air/deploycode.go create mode 100644 cmd/air/deploycode_test.go create mode 100644 cmd/air/deployconfig.go create mode 100644 cmd/air/deployconfig_test.go create mode 100644 cmd/air/deployendpoint.go create mode 100644 cmd/air/deployendpoint_test.go create mode 100644 cmd/air/deployregister.go create mode 100644 cmd/air/deployregister_test.go diff --git a/cmd/air/air.go b/cmd/air/air.go index 4154bab891c..584f4165c58 100644 --- a/cmd/air/air.go +++ b/cmd/air/air.go @@ -47,7 +47,7 @@ active runs with "databricks air list".`, wrapRunErrorWithDebugTip(runCommand) addToGroup(cmd, runGroupID, runCommand) addToGroup(cmd, manageGroupID, newGetCommand(), newListCommand(), newLogsCommand(), newCancelCommand()) - addToGroup(cmd, packageGroupID, newConvertToDabsCommand(), newImagesCommand()) + addToGroup(cmd, packageGroupID, newDeployCommand(), newConvertToDabsCommand(), newImagesCommand()) addToGroup(cmd, capacityGroupID, newPoolsCommand()) // Subcommands inherit this help func, so the footer is printed only for `air` itself. diff --git a/cmd/air/air_test.go b/cmd/air/air_test.go index 11e4546d21d..1c877160270 100644 --- a/cmd/air/air_test.go +++ b/cmd/air/air_test.go @@ -18,7 +18,7 @@ func TestNewRegistersAllSubcommands(t *testing.T) { registered[c.Name()] = true } - want := []string{"run", "get", "list", "pools", "logs", "cancel", "convert-to-dabs", "images"} + want := []string{"run", "get", "list", "pools", "logs", "cancel", "deploy", "convert-to-dabs", "images"} for _, name := range want { assert.True(t, registered[name], "subcommand %q is not registered", name) } diff --git a/cmd/air/deploy.go b/cmd/air/deploy.go new file mode 100644 index 00000000000..1b8f0932e92 --- /dev/null +++ b/cmd/air/deploy.go @@ -0,0 +1,382 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "net/http" + "path/filepath" + "strings" + "time" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/shellquote" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/retries" + "github.com/databricks/databricks-sdk-go/service/serving" + "github.com/spf13/cobra" +) + +// defaultDeployTimeout bounds the whole wait: staging a checkpoint's weights and +// starting the endpoint. +const defaultDeployTimeout = time.Hour + +// Errors air deploy wraps when it stops waiting. The server keeps working, and +// the error says how to continue. +var ( + errDeployInterrupted = errors.New("air deploy was interrupted") + errDeployTimedOut = errors.New("air deploy timed out") +) + +// deployResult is the JSON payload for `air deploy`. +type deployResult struct { + Status string `json:"status"` + DryRun bool `json:"dry_run,omitempty"` + // CodeFiles lists, on a dry run of a model from code, the files it would send. + CodeFiles []string `json:"code_files,omitempty"` + RegisteredModel string `json:"registered_model,omitempty"` + ModelVersion string `json:"model_version,omitempty"` + ModelID string `json:"model_id,omitempty"` + IdempotencyKey string `json:"idempotency_key,omitempty"` + // EnvironmentStagingFailed is set when the version's serving environment + // couldn't be prepared ahead of time, so the endpoint can be slower to start. + EnvironmentStagingFailed bool `json:"environment_staging_failed,omitempty"` + Endpoint string `json:"endpoint,omitempty"` + EndpointURL string `json:"endpoint_url,omitempty"` +} + +// deployOptions are the `air deploy` flags. +type deployOptions struct { + file string + overrides []string + dryRun bool + idempotencyKey string + retryFailedStaging bool + timeout time.Duration +} + +func newDeployCommand() *cobra.Command { + return newDeployCommandWithPollInterval(registrationPollInterval) +} + +func newDeployCommandWithPollInterval(pollInterval time.Duration) *cobra.Command { + var opts deployOptions + + cmd := &cobra.Command{ + Use: "deploy", + Args: root.NoArgs, + Short: "Register a model in Unity Catalog and serve it on an endpoint", + Long: `Register a model in Unity Catalog and serve it on a Model Serving endpoint. + +The deployment is described by a YAML file (see --file). The model is a +checkpoint (weights in a Unity Catalog volume and the command that serves +them), a model from code, a model directory saved by MLflow, or an existing +model version. air deploy registers a new version, waits until it is ready, +then creates or updates the endpoint to serve it and waits until the endpoint +is ready. + +Registration is idempotent: run air deploy again with the idempotency key it +prints to continue an interrupted deployment. To look up a YAML field, pass +its path to -h as a separate argument: + + databricks air deploy -h config all fields + databricks air deploy -h config.model one section + databricks air deploy -h config.model.type one field`, + Example: ` # Register a checkpoint and serve it + databricks air deploy -f deploy.yaml + + # Validate the config, and list the code it would send, without deploying + databricks air deploy -f deploy.yaml --dry-run + + # Continue a deployment that was interrupted + databricks air deploy -f deploy.yaml --idempotency-key 3f6c0b52-4c7e-4d0b-9b8e-0f7ad6a4c2d1`, + // Hidden until the registration API is available in every workspace. + Hidden: true, + } + + // cobra passes -h's positional args to the help func before Args/required-flag + // validation, so a config path documents a field without needing -f. + cmd.SetHelpFunc(func(c *cobra.Command, args []string) { + fields := c.Flags().Args() + if len(fields) == 0 { + // Parent() is nil for a detached command (unit tests). + if parent := c.Parent(); parent != nil { + parent.HelpFunc()(c, args) + return + } + _ = c.Usage() + return + } + field, err := resolveSchemaField(deploySchema(), fields[0]) + if err != nil { + c.PrintErrln("Error:", err) + return + } + renderConfigField(c.OutOrStdout(), deployCommandPath, field, nil) + }) + + // List flags in declaration order (most used first) rather than alphabetically. + cmd.Flags().SortFlags = false + cmd.Flags().StringVarP(&opts.file, "file", "f", "", "Path to the deployment YAML config") + cmd.Flags().StringArrayVar(&opts.overrides, "override", nil, "Override a YAML field, e.g. compute.workload_size=Medium (repeatable)") + cmd.Flags().BoolVar(&opts.dryRun, "dry-run", false, "Validate the config, and list the code it would send, without deploying") + cmd.Flags().StringVar(&opts.idempotencyKey, "idempotency-key", "", "Continue the registration that used this key, instead of registering another version") + cmd.Flags().BoolVar(&opts.retryFailedStaging, "retry-failed-staging", false, "Copy a checkpoint's weights again after copying them failed") + cmd.Flags().DurationVar(&opts.timeout, "timeout", defaultDeployTimeout, "How long to wait for the version and the endpoint to be ready") + _ = cmd.MarkFlagRequired("file") + + // --dry-run never authenticates, so it reports config mistakes without + // credentials. + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + if opts.dryRun { + return nil + } + return root.MustWorkspaceClient(cmd, args) + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + err := runDeploy(cmd, opts, pollInterval) + if err == nil { + return nil + } + code, retryable := deployErrorCode(err) + kind := "PERMANENT" + if retryable { + kind = "TRANSIENT" + } + return renderError(cmd.Context(), cmd, code, kind, retryable, err) + } + return cmd +} + +func runDeploy(cmd *cobra.Command, opts deployOptions, pollInterval time.Duration) error { + ctx := cmd.Context() + cfg, err := loadDeployConfig(ctx, opts.file, opts.overrides) + if err != nil { + return err + } + if err := validateDeployFlags(cfg, opts); err != nil { + return err + } + + var code *modelPackageCode + if cfg.Model.Type == modelTypeCode { + code, err = collectModelCode(ctx, cfg, opts.file) + if err != nil { + return err + } + } + + jsonOut := root.OutputType(cmd) == flags.OutputJSON + if opts.dryRun { + return renderDeployDryRun(ctx, cfg, code, jsonOut) + } + + w := cmdctx.WorkspaceClient(ctx) + ctx, cancel := context.WithTimeout(ctx, opts.timeout) + defer cancel() + ctx, stop := notifyInterrupt(ctx) + defer stop() + + result := deployResult{RegisteredModel: cfg.RegisteredModel, Endpoint: cfg.Endpoint} + if cfg.registers() { + // Fail on an endpoint air deploy won't change before copying any weights. + if err := checkEndpoint(ctx, w, cfg.Endpoint, cfg.RegisteredModel); err != nil { + return err + } + result.IdempotencyKey = effectiveIdempotencyToken(opts.idempotencyKey, cfg.IdempotencyToken) + resp, err := registerDeployment(ctx, w, cfg, opts, code, result.IdempotencyKey, pollInterval, !jsonOut) + if err != nil { + return err + } + result.ModelVersion, result.ModelID = resp.ModelVersion, resp.ModelID + result.EnvironmentStagingFailed = resp.environmentState() == stagingStateFailed + if !jsonOut { + cmdio.LogString(ctx, fmt.Sprintf("Version %s of %s is ready.", resp.ModelVersion, cfg.RegisteredModel)) + if result.EnvironmentStagingFailed { + cmdio.LogString(ctx, "Warning: the version's serving environment couldn't be prepared ahead of time, so the endpoint can be slower to start, or fail to start.") + } + } + } else { + result.ModelVersion = *cfg.Model.Version + } + + entity := cfg.servedEntity(result.ModelVersion) + endpoint, err := serveVersion(ctx, w, cfg.Endpoint, entity, !jsonOut) + if err != nil { + resume := deployResumeCommand(w.Config.Profile, opts, result.IdempotencyKey) + if stopped := deployStopped(ctx, err); stopped != nil { + return fmt.Errorf("%w while waiting for endpoint %s, which keeps updating. To continue, run:\n %s", stopped, cfg.Endpoint, resume) + } + return fmt.Errorf("serving version %s of %s on endpoint %s failed: %w\n\nTo try again without registering another version, run:\n %s", result.ModelVersion, cfg.RegisteredModel, cfg.Endpoint, err, resume) + } + result.Status = "READY" + result.EndpointURL = endpointInvocationsURL(w.Config.Host, endpoint.Name) + if jsonOut { + return renderEnvelope(ctx, result) + } + out := cmd.OutOrStdout() + fmt.Fprintf(out, "Endpoint %s serves version %s of %s.\n", endpoint.Name, result.ModelVersion, cfg.RegisteredModel) + fmt.Fprintln(out, "Query it at: "+result.EndpointURL) + return nil +} + +// renderDeployDryRun reports a valid config and, for a model from code, the +// files a deployment would send. +func renderDeployDryRun(ctx context.Context, cfg *deployConfig, code *modelPackageCode, jsonOut bool) error { + var files []string + if code != nil { + files = codeFilePaths(code) + } + if jsonOut { + return renderEnvelope(ctx, deployResult{Status: "DRY_RUN_OK", DryRun: true, CodeFiles: files}) + } + if code != nil { + cmdio.LogString(ctx, fmt.Sprintf("Dry run: the model from code would include %d file(s):", len(files))) + for _, f := range files { + cmdio.LogString(ctx, " "+f) + } + } + cmdio.LogString(ctx, fmt.Sprintf("Dry run: configuration for endpoint %q is valid; not deploying.", cfg.Endpoint)) + return nil +} + +// registerDeployment registers the version and waits until it is ready, showing +// progress on a spinner when show is true. Its errors say how to continue. +func registerDeployment(ctx context.Context, w *databricks.WorkspaceClient, cfg *deployConfig, opts deployOptions, code *modelPackageCode, key string, pollInterval time.Duration, show bool) (*registerModelPackageResponse, error) { + if show { + cmdio.LogString(ctx, fmt.Sprintf("Registering a version of %s with idempotency key %s", cfg.RegisteredModel, key)) + } + req := cfg.registrationRequest(key, code) + req.RetryFailedStaging = opts.retryFailedStaging + + progress := func(*registerModelPackageResponse) {} + if show { + sp := cmdio.NewSpinner(ctx) + defer sp.Close() + progress = func(resp *registerModelPackageResponse) { + sp.Update(registrationProgress(cfg, resp)) + } + } + resp, err := registerUntilReady(ctx, w, req, pollInterval, progress) + if err == nil { + return resp, nil + } + resume := deployResumeCommand(w.Config.Profile, opts, key) + if errors.Is(err, errStagingFailed) { + source := "the checkpoint's weights" + if cfg.Model.Path != nil { + source = "the weights in " + *cfg.Model.Path + } + return nil, fmt.Errorf("copying %s into version %s of %s failed. Check that you can read every file, then run:\n %s --retry-failed-staging", source, resp.ModelVersion, cfg.RegisteredModel, resume) + } + if stopped := deployStopped(ctx, err); stopped != nil { + return nil, fmt.Errorf("%w while registering a version of %s. Copying a checkpoint's weights continues, and the version is finalized when you continue. To continue, run:\n %s", stopped, cfg.RegisteredModel, resume) + } + return nil, fmt.Errorf("registering a version of %s with idempotency key %s failed: %w", cfg.RegisteredModel, key, err) +} + +// deployStopped returns why air deploy stopped waiting, or nil if it didn't. The +// SDK's endpoint poller reports a canceled context as its own timeout, so the +// context decides between Ctrl-C and --timeout. +func deployStopped(ctx context.Context, err error) error { + switch { + case errors.Is(ctx.Err(), context.Canceled): + return errDeployInterrupted + case errors.Is(ctx.Err(), context.DeadlineExceeded): + return errDeployTimedOut + } + if _, ok := errors.AsType[*retries.ErrTimedOut](err); ok { + return errDeployTimedOut + } + return nil +} + +// deployErrorCode classifies a failed deployment for the JSON error envelope. +func deployErrorCode(err error) (code string, retryable bool) { + switch { + case errors.Is(err, errDeployInterrupted): + return "INTERRUPTED", true + case errors.Is(err, errDeployTimedOut): + return "TIMEOUT", true + } + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok && apiErr.ErrorCode != "" { + return apiErr.ErrorCode, apiErr.StatusCode == http.StatusTooManyRequests || apiErr.StatusCode >= http.StatusInternalServerError + } + return "INVALID_ARGS", false +} + +// validateDeployFlags rejects flags the config's model source can't use. +func validateDeployFlags(cfg *deployConfig, opts deployOptions) error { + if opts.timeout <= 0 { + return fmt.Errorf("--timeout must be positive, got %s", opts.timeout) + } + if opts.idempotencyKey != "" { + if !cfg.registers() { + return errors.New("--idempotency-key applies only when air deploy registers a version; model.type is registered") + } + if err := validateDeployIdempotencyKey(opts.idempotencyKey, "--idempotency-key"); err != nil { + return err + } + } + if opts.retryFailedStaging && cfg.Model.Type != modelTypeCheckpoint { + return fmt.Errorf("--retry-failed-staging applies only to a checkpoint; model.type is %s", cfg.Model.Type) + } + return nil +} + +// registrationProgress describes a registration response on the spinner. +func registrationProgress(cfg *deployConfig, resp *registerModelPackageResponse) string { + state := resp.weightsState() + if state == "" { + return fmt.Sprintf("Registering version %s of %s…", resp.ModelVersion, cfg.RegisteredModel) + } + return fmt.Sprintf("Copying the weights into version %s of %s (%s)…", resp.ModelVersion, cfg.RegisteredModel, state) +} + +// serveVersion serves entity on the endpoint, showing the endpoint's state on a +// spinner when show is true. +func serveVersion(ctx context.Context, w *databricks.WorkspaceClient, name string, entity serving.ServedEntityInput, show bool) (*serving.ServingEndpointDetailed, error) { + if !show { + return serveOnEndpoint(ctx, w, name, entity, nil) + } + cmdio.LogString(ctx, fmt.Sprintf("Serving version %s of %s on endpoint %s", entity.EntityVersion, entity.EntityName, name)) + sp := cmdio.NewSpinner(ctx) + defer sp.Close() + return serveOnEndpoint(ctx, w, name, entity, func(endpoint *serving.ServingEndpointDetailed) { + if endpoint.State != nil { + sp.Update(fmt.Sprintf("Waiting for endpoint %s (%s, %s)…", name, endpoint.State.Ready, endpoint.State.ConfigUpdate)) + } + }) +} + +// deployResumeCommand is the command that continues this deployment: the same +// config (by absolute path, so it runs from any directory), overrides and +// profile, with the registration's key when it has one. The overrides are part +// of the request the key was first used with, so they must be repeated. +func deployResumeCommand(profile string, opts deployOptions, key string) string { + file := opts.file + if abs, err := filepath.Abs(file); err == nil { + file = abs + } + args := []string{deployCommandPath, "-f", shellquote.BashArg(filepath.ToSlash(file))} + for _, override := range opts.overrides { + args = append(args, "--override", shellquote.BashArg(override)) + } + if key != "" { + args = append(args, "--idempotency-key", shellquote.BashArg(key)) + } + if profile != "" { + args = append(args, "-p", shellquote.BashArg(profile)) + } + return strings.Join(args, " ") +} + +// endpointInvocationsURL is the REST URL that queries the endpoint. +func endpointInvocationsURL(host, name string) string { + return strings.TrimSuffix(host, "/") + "/serving-endpoints/" + name + "/invocations" +} diff --git a/cmd/air/deploy_test.go b/cmd/air/deploy_test.go new file mode 100644 index 00000000000..ce232b46ff4 --- /dev/null +++ b/cmd/air/deploy_test.go @@ -0,0 +1,366 @@ +package aircmd + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http/httptest" + "path/filepath" + "strconv" + "strings" + "testing" + "time" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/shellquote" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/retries" + "github.com/databricks/databricks-sdk-go/service/serving" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// withDeadline returns the test's context with a deadline, which air deploy +// always sets from --timeout. +func withDeadline(t *testing.T) context.Context { + ctx, cancel := context.WithTimeout(t.Context(), time.Minute) + t.Cleanup(cancel) + return ctx +} + +// deployServer serves registerModelPackage with responses and the serving +// endpoints API from endpoints. +func deployServer(t *testing.T, endpoints *fakeEndpoints, responses ...string) (*httptest.Server, *fakeModelServing) { + t.Helper() + fake := newFakeModelServing(t, responses...) + fake.endpoints = endpoints.ServeHTTP + return fake.start(), fake +} + +// runDeployCmd runs `air deploy` on config with the given flags. An empty srvURL +// runs it without a workspace client, as --dry-run does. +func runDeployCmd(t *testing.T, out flags.Output, buf *bytes.Buffer, srvURL, config string, flagValues ...string) error { + t.Helper() + cmd := withOutput(newDeployCommandWithPollInterval(time.Millisecond), out) + require.NoError(t, cmd.Flags().Set("file", writeConfigFile(t, "deploy.yaml", config))) + for i := 0; i < len(flagValues); i += 2 { + require.NoError(t, cmd.Flags().Set(flagValues[i], flagValues[i+1])) + } + + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), out, nil, buf, buf, "", "")) + if srvURL != "" { + ctx = cmdctx.SetWorkspaceClient(ctx, newTestWorkspaceClient(t, srvURL)) + } + cmd.SetContext(ctx) + cmd.SetOut(buf) + return cmd.RunE(cmd, nil) +} + +func TestDeployCheckpointJSON(t *testing.T) { + endpoints := newFakeEndpoints(t, nil) + srv, fake := deployServer(t, endpoints, stagingInProgress, stagingReady) + + var buf bytes.Buffer + err := runDeployCmd(t, flags.OutputJSON, &buf, srv.URL, checkpointDeployConfig, "idempotency-key", "key-1") + require.NoError(t, err) + + var got struct { + Data deployResult `json:"data"` + } + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, deployResult{ + Status: "READY", + RegisteredModel: "main.models.qwen", + ModelVersion: "3", + ModelID: "m-1", + IdempotencyKey: "key-1", + Endpoint: "qwen-chat", + EndpointURL: srv.URL + "/serving-endpoints/qwen-chat/invocations", + }, got.Data) + + bodies := fake.requests() + require.Len(t, bodies, 2) + assert.Equal(t, "key-1", bodies[0].IdempotencyKey) + assert.Equal(t, bodies[0], bodies[1]) + assert.Equal(t, []string{"create"}, endpoints.writes) + served := endpoints.endpoint.Config.ServedEntities[0] + assert.Equal(t, "3", served.EntityVersion) + assert.Equal(t, map[string]string{"LOG_LEVEL": "info", "HF_TOKEN": "{{secrets/llm/hf_token}}"}, served.EnvironmentVars) +} + +func TestDeployCheckpointText(t *testing.T) { + srv, _ := deployServer(t, newFakeEndpoints(t, nil), stagingReady) + + var buf bytes.Buffer + require.NoError(t, runDeployCmd(t, flags.OutputText, &buf, srv.URL, checkpointDeployConfig)) + out := buf.String() + assert.Contains(t, out, "Registering a version of main.models.qwen with idempotency key ") + assert.Contains(t, out, "Version 3 of main.models.qwen is ready.") + assert.Contains(t, out, "Serving version 3 of main.models.qwen on endpoint qwen-chat") + assert.Contains(t, out, "Endpoint qwen-chat serves version 3 of main.models.qwen.\n") + assert.Contains(t, out, "Query it at: "+srv.URL+"/serving-endpoints/qwen-chat/invocations\n") + assert.NotContains(t, out, "Warning") +} + +func TestDeployWarnsWhenTheEnvironmentFailedToStage(t *testing.T) { + environmentFailed := strings.Replace(stagingReady, `"environment_artifact": {"state": "SUCCEEDED"}`, `"environment_artifact": {"state": "FAILED"}`, 1) + srv, _ := deployServer(t, newFakeEndpoints(t, nil), environmentFailed) + + var buf bytes.Buffer + require.NoError(t, runDeployCmd(t, flags.OutputText, &buf, srv.URL, checkpointDeployConfig)) + assert.Contains(t, buf.String(), "Warning: the version's serving environment couldn't be prepared ahead of time") +} + +func TestDeployUsesTheConfigTokenThenANewKey(t *testing.T) { + srv, fake := deployServer(t, newFakeEndpoints(t, nil), stagingReady) + + var buf bytes.Buffer + require.NoError(t, runDeployCmd(t, flags.OutputJSON, &buf, srv.URL, checkpointDeployConfig+"idempotency_token: from-config\n")) + require.NoError(t, runDeployCmd(t, flags.OutputJSON, &buf, srv.URL, checkpointDeployConfig+"idempotency_token: from-config\n", "idempotency-key", "from-flag")) + require.NoError(t, runDeployCmd(t, flags.OutputJSON, &buf, srv.URL, checkpointDeployConfig)) + + bodies := fake.requests() + require.Len(t, bodies, 3) + assert.Equal(t, "from-config", bodies[0].IdempotencyKey) + assert.Equal(t, "from-flag", bodies[1].IdempotencyKey) + assert.Regexp(t, deployIdempotencyKeyRe(), bodies[2].IdempotencyKey) + assert.NotContains(t, []string{"from-config", "from-flag"}, bodies[2].IdempotencyKey) +} + +func TestDeployRetriesFailedStagingOnce(t *testing.T) { + srv, fake := deployServer(t, newFakeEndpoints(t, nil), stagingInProgress, stagingReady) + + var buf bytes.Buffer + require.NoError(t, runDeployCmd(t, flags.OutputJSON, &buf, srv.URL, checkpointDeployConfig, "retry-failed-staging", "true")) + bodies := fake.requests() + require.Len(t, bodies, 2) + assert.True(t, bodies[0].RetryFailedStaging) + assert.False(t, bodies[1].RetryFailedStaging) +} + +func TestDeployFailedStagingSaysHowToRetry(t *testing.T) { + srv, _ := deployServer(t, newFakeEndpoints(t, nil), stagingFailed) + + var buf bytes.Buffer + err := runDeployCmd(t, flags.OutputText, &buf, srv.URL, checkpointDeployConfig, "idempotency-key", "key-1") + require.Error(t, err) + assert.Contains(t, err.Error(), "copying the weights in /Volumes/main/models/weights/qwen3_8b into version 3 of main.models.qwen failed") + assert.Contains(t, err.Error(), " --idempotency-key key-1 --retry-failed-staging") +} + +func TestDeployReportsRegistrationErrors(t *testing.T) { + srv, _ := deployServer(t, newFakeEndpoints(t, nil), "") + + var buf bytes.Buffer + err := runDeployCmd(t, flags.OutputText, &buf, srv.URL, checkpointDeployConfig, "idempotency-key", "key-1") + assert.ErrorContains(t, err, "registering a version of main.models.qwen with idempotency key key-1 failed: The idempotency key can't be used for this registration") +} + +func TestDeployReportsErrorsAsJSON(t *testing.T) { + srv, _ := deployServer(t, newFakeEndpoints(t, nil), "") + + var buf bytes.Buffer + err := runDeployCmd(t, flags.OutputJSON, &buf, srv.URL, checkpointDeployConfig) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + var got errorEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, "RESOURCE_CONFLICT", got.Error.Code) + assert.Equal(t, "PERMANENT", got.Error.Kind) + assert.False(t, got.Error.Retryable) + + buf.Reset() + err = runDeployCmd(t, flags.OutputJSON, &buf, "", "registered_model: x\n", "dry-run", "true") + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, "INVALID_ARGS", got.Error.Code) + assert.Contains(t, got.Error.Message, "registered_model must be a Unity Catalog model name") +} + +func TestDeployChecksTheEndpointBeforeRegistering(t *testing.T) { + other := testEntity("5") + other.EntityName = "main.models.llama" + srv, fake := deployServer(t, newFakeEndpoints(t, endpointServing(other, readyState)), stagingReady) + + var buf bytes.Buffer + err := runDeployCmd(t, flags.OutputText, &buf, srv.URL, checkpointDeployConfig) + assert.ErrorContains(t, err, "endpoint qwen-chat serves version 5 of main.models.llama") + assert.Empty(t, fake.requests()) +} + +func TestDeployEndpointFailureSaysHowToRetry(t *testing.T) { + endpoints := newFakeEndpoints(t, nil) + endpoints.finalState = serving.EndpointState{Ready: serving.EndpointStateReadyNotReady, ConfigUpdate: serving.EndpointStateConfigUpdateUpdateFailed} + srv, _ := deployServer(t, endpoints, stagingReady) + + var buf bytes.Buffer + err := runDeployCmd(t, flags.OutputText, &buf, srv.URL, checkpointDeployConfig, "idempotency-key", "key-1") + require.Error(t, err) + assert.Contains(t, err.Error(), "serving version 3 of main.models.qwen on endpoint qwen-chat failed: ") + assert.Contains(t, err.Error(), "To try again without registering another version, run:\n databricks air deploy -f ") + assert.Contains(t, err.Error(), " --idempotency-key key-1") +} + +func TestDeployRegisteredVersionSkipsRegistration(t *testing.T) { + endpoints := newFakeEndpoints(t, nil) + srv, fake := deployServer(t, endpoints) + + var buf bytes.Buffer + err := runDeployCmd(t, flags.OutputJSON, &buf, srv.URL, `registered_model: main.models.qwen +model: {type: registered, version: 7} +endpoint: qwen-chat +compute: {workload_type: GPU_MEDIUM, workload_size: Small} +`) + require.NoError(t, err) + assert.Empty(t, fake.requests()) + assert.Equal(t, "7", endpoints.endpoint.Config.ServedEntities[0].EntityVersion) + assert.NotContains(t, buf.String(), "idempotency_key") +} + +func TestDeployCodeSendsTheFiles(t *testing.T) { + cfg, _ := writeCodeTree(t, map[string]string{"agent.py": "import tools\n", "tools/__init__.py": "x = 1\n"}) + codeReady := `{"model_id": "m-2", "registered_model": "main.agents.support", "model_version": "1", "ready": true}` + srv, fake := deployServer(t, newFakeEndpoints(t, nil), codeReady) + + var buf bytes.Buffer + require.NoError(t, runDeployCmd(t, flags.OutputJSON, &buf, srv.URL, codeConfigAt(cfg.CodeSource.RootPath))) + bodies := fake.requests() + require.Len(t, bodies, 1) + assert.Equal(t, &modelPackageCode{ + EntryFile: "agent.py", + EntrySource: "import tools\n", + Files: []modelPackageFile{{Path: "tools/__init__.py", Content: "x = 1\n"}}, + Task: "agent/v1/responses", + Environment: modelPackageEnvironment{Runtime: "client.5", Dependencies: []string{"mlflow==3.5.0"}}, + }, bodies[0].Code) + assert.Nil(t, bodies[0].Checkpoint) +} + +func TestDeployDryRunNeedsNoWorkspace(t *testing.T) { + cmd := newDeployCommand() + require.NoError(t, cmd.Flags().Set("dry-run", "true")) + assert.NoError(t, cmd.PreRunE(cmd, nil)) + + var buf bytes.Buffer + require.NoError(t, runDeployCmd(t, flags.OutputText, &buf, "", checkpointDeployConfig, "dry-run", "true")) + assert.Contains(t, buf.String(), `Dry run: configuration for endpoint "qwen-chat" is valid; not deploying.`) + + buf.Reset() + require.NoError(t, runDeployCmd(t, flags.OutputJSON, &buf, "", checkpointDeployConfig, "dry-run", "true")) + var got struct { + Data deployResult `json:"data"` + } + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, deployResult{Status: "DRY_RUN_OK", DryRun: true}, got.Data) +} + +func TestDeployDryRunListsCodeFiles(t *testing.T) { + cfg, _ := writeCodeTree(t, map[string]string{"agent.py": "x\n", "tools/a.py": "a\n"}) + + var buf bytes.Buffer + require.NoError(t, runDeployCmd(t, flags.OutputText, &buf, "", codeConfigAt(cfg.CodeSource.RootPath), "dry-run", "true")) + assert.Contains(t, buf.String(), "Dry run: the model from code would include 2 file(s):\n agent.py\n tools/a.py\n") + + buf.Reset() + require.NoError(t, runDeployCmd(t, flags.OutputJSON, &buf, "", codeConfigAt(cfg.CodeSource.RootPath), "dry-run", "true")) + var got struct { + Data deployResult `json:"data"` + } + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, []string{"agent.py", "tools/a.py"}, got.Data.CodeFiles) +} + +// codeConfigAt is codeDeployConfig with its code at root. +func codeConfigAt(root string) string { + return strings.Replace(codeDeployConfig, "root_path: .", "root_path: "+strconv.Quote(filepath.ToSlash(root)), 1) +} + +func TestDeployFlagValidation(t *testing.T) { + registered := "registered_model: a.b.c\nmodel: {type: registered, version: 1}\nendpoint: e\ncompute: {workload_type: CPU, workload_size: Small}\n" + tests := []struct { + name string + config string + flagValues []string + wantErr string + }{ + {"key for a registered version", registered, []string{"idempotency-key", "k"}, "--idempotency-key applies only when air deploy registers a version"}, + {"bad key", checkpointDeployConfig, []string{"idempotency-key", "a/b"}, "--idempotency-key must be 1 to 128 letters"}, + {"retry for a registered version", registered, []string{"retry-failed-staging", "true"}, "--retry-failed-staging applies only to a checkpoint"}, + {"zero timeout", checkpointDeployConfig, []string{"timeout", "0s"}, "--timeout must be positive"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var buf bytes.Buffer + err := runDeployCmd(t, flags.OutputText, &buf, "", tt.config, append(tt.flagValues, "dry-run", "true")...) + assert.ErrorContains(t, err, tt.wantErr) + }) + } +} + +func TestDeployResumeCommand(t *testing.T) { + abs, err := filepath.Abs("my deploy.yaml") + require.NoError(t, err) + opts := deployOptions{file: "my deploy.yaml", overrides: []string{"compute.workload_size=Medium", "task=llm/v1/chat"}} + assert.Equal(t, + "databricks air deploy -f "+shellquote.BashArg(filepath.ToSlash(abs))+" --override 'compute.workload_size=Medium' --override 'task=llm/v1/chat' --idempotency-key k-1 -p dev", + deployResumeCommand("dev", opts, "k-1")) + + abs, err = filepath.Abs("d.yaml") + require.NoError(t, err) + assert.Equal(t, "databricks air deploy -f "+shellquote.BashArg(filepath.ToSlash(abs)), deployResumeCommand("", deployOptions{file: "d.yaml"}, "")) +} + +func TestDeployStopped(t *testing.T) { + assert.NoError(t, deployStopped(t.Context(), errors.New("boom"))) + + canceled, cancel := context.WithCancel(t.Context()) + cancel() + // The SDK's endpoint poller reports Ctrl-C as its own timeout. + assert.ErrorIs(t, deployStopped(canceled, &retries.ErrTimedOut{}), errDeployInterrupted) + assert.ErrorIs(t, deployStopped(canceled, context.Canceled), errDeployInterrupted) + + expired, cancelExpired := context.WithTimeout(t.Context(), -time.Second) + defer cancelExpired() + assert.ErrorIs(t, deployStopped(expired, context.DeadlineExceeded), errDeployTimedOut) + assert.ErrorIs(t, deployStopped(t.Context(), &retries.ErrTimedOut{}), errDeployTimedOut) +} + +func TestDeployErrorCode(t *testing.T) { + for _, tt := range []struct { + err error + code string + retryable bool + }{ + {fmt.Errorf("x: %w", errDeployInterrupted), "INTERRUPTED", true}, + {fmt.Errorf("x: %w", errDeployTimedOut), "TIMEOUT", true}, + {fmt.Errorf("x: %w", &apierr.APIError{ErrorCode: "TEMPORARILY_UNAVAILABLE", StatusCode: 503}), "TEMPORARILY_UNAVAILABLE", true}, + {fmt.Errorf("x: %w", &apierr.APIError{ErrorCode: "INVALID_PARAMETER_VALUE", StatusCode: 400}), "INVALID_PARAMETER_VALUE", false}, + {errors.New("model: section is required"), "INVALID_ARGS", false}, + } { + code, retryable := deployErrorCode(tt.err) + assert.Equal(t, tt.code, code, tt.err.Error()) + assert.Equal(t, tt.retryable, retryable, tt.err.Error()) + } +} + +func TestDeployHelpDocumentsAField(t *testing.T) { + var out strings.Builder + cmd := newDeployCommand() + cmd.SetOut(&out) + cmd.SetArgs([]string{"-h", "config.model.type"}) + require.NoError(t, cmd.Execute()) + assert.Contains(t, out.String(), "config.model.type\n One of checkpoint") + assert.Contains(t, out.String(), "Required: yes") + + out.Reset() + cmd = newDeployCommand() + cmd.SetOut(&out) + cmd.SetArgs([]string{"-h", "config"}) + require.NoError(t, cmd.Execute()) + // The coding agent skill writes run YAML, so the deploy help leaves it out. + assert.NotContains(t, out.String(), "databricks-ai-runtime") +} diff --git a/cmd/air/deploycode.go b/cmd/air/deploycode.go new file mode 100644 index 00000000000..eb2c1ae1ee0 --- /dev/null +++ b/cmd/air/deploycode.go @@ -0,0 +1,155 @@ +package aircmd + +import ( + "context" + "fmt" + "os" + "path/filepath" + "regexp" + "slices" + "strings" + "sync" + "unicode/utf8" +) + +// Limits of a model from code, as the server enforces them. +const ( + // maxCodePackageBytes caps the code files plus the dependency lines. The server + // also counts the small environment file it writes, so this check rejects only + // packages the server would reject too. + maxCodePackageBytes = 4 * 1024 * 1024 + // maxCodeFiles caps the files besides the entry file. + maxCodeFiles = 1000 + maxCodePathDepth = 16 +) + +// pythonCacheDir holds bytecode Python writes next to imported sources. +const pythonCacheDir = "__pycache__" + +// codePathSegmentRe matches a segment of a code file path the server can place. +var codePathSegmentRe = sync.OnceValue(func() *regexp.Regexp { return regexp.MustCompile(`^[A-Za-z0-9_][A-Za-z0-9_.-]*$`) }) + +// collectModelCode reads a model from code from code_source: model.entry and the +// other files a code_source snapshot would include, as UTF-8 text. Files git +// ignores aren't sent, and nor are hidden files and directories or __pycache__ +// directories, whose names the server can't place or which hold no source. +// configPath locates a relative code_source.root_path. +// +// The files are stored in the model version, where anyone who can read the model +// can read them. Outside a git repository no .gitignore filters them, so the user +// must list them in code_source.include_paths. +func collectModelCode(ctx context.Context, cfg *deployConfig, configPath string) (*modelPackageCode, error) { + root, err := resolveRootPath(ctx, cfg.CodeSource.RootPath, filepath.Dir(configPath)) + if err != nil { + return nil, err + } + isGitRepo := newGitRepo(root).isRepository(ctx) + if !isGitRepo && len(cfg.CodeSource.IncludePaths) == 0 { + return nil, fmt.Errorf("code_source.include_paths is required because %s isn't in a git repository: air deploy stores the files in the model version, so list the ones to send, e.g. include_paths: [%s]", filepath.ToSlash(root), *cfg.Model.Entry) + } + files, err := snapshotFiles(ctx, root, cfg.CodeSource.IncludePaths, isGitRepo) + if err != nil { + return nil, err + } + files = slices.DeleteFunc(files, skippedCodeFile) + + code := &modelPackageCode{EntryFile: *cfg.Model.Entry, Environment: cfg.servingEnvironment()} + if cfg.Task != nil { + code.Task = *cfg.Task + } + // Check the listing before reading anything, so a root_path that holds large + // data fails fast instead of being read into memory. + if err := checkCodeFiles(files, code.Environment.Dependencies); err != nil { + return nil, err + } + + foundEntry := false + for _, f := range files { + rel := filepath.ToSlash(f.rel) + content, err := os.ReadFile(filepath.Join(root, f.rel)) + if err != nil { + return nil, fmt.Errorf("failed to read code_source file %s: %w", rel, err) + } + if !utf8.Valid(content) { + return nil, fmt.Errorf("code_source file %s isn't UTF-8 text; air deploy sends source files only, so leave it out with code_source.include_paths", rel) + } + if rel == code.EntryFile { + code.EntrySource = string(content) + foundEntry = true + continue + } + code.Files = append(code.Files, modelPackageFile{Path: rel, Content: string(content)}) + } + if !foundEntry { + return nil, fmt.Errorf("model.entry %s isn't among the files code_source sends from %s; check that it exists, isn't ignored by git and is under code_source.include_paths", code.EntryFile, filepath.ToSlash(root)) + } + return code, nil +} + +// checkCodeFiles checks the files to send against the server's limits, from +// their listing alone. +func checkCodeFiles(files []snapshotFile, dependencies []string) error { + // One of the files is the entry file, which the limit doesn't count. + if len(files) > maxCodeFiles+1 { + return fmt.Errorf("code_source sends %d files, but a model from code can include at most %d besides model.entry; narrow it with code_source.include_paths", len(files), maxCodeFiles) + } + var size int64 + for _, dependency := range dependencies { + size += int64(len(dependency)) + 1 + } + for _, f := range files { + rel := filepath.ToSlash(f.rel) + if !os.FileMode(f.mode).IsRegular() { + return fmt.Errorf("code_source path %s isn't a regular file; air deploy sends the contents of files only, so replace it with a file or leave it out with code_source.include_paths", rel) + } + if err := validateCodePath(rel); err != nil { + return err + } + size += f.size + if size > maxCodePackageBytes { + return fmt.Errorf("code_source files and environment.dependencies total more than %d bytes, the most a model from code can hold; narrow code_source.include_paths", maxCodePackageBytes) + } + } + return nil +} + +// codeFilePaths lists the files of code, entry file first, as dry runs show them. +func codeFilePaths(code *modelPackageCode) []string { + paths := []string{code.EntryFile} + for _, f := range code.Files { + paths = append(paths, f.Path) + } + return paths +} + +// sentCodePath reports whether a code_source file at rel (slash-separated) is +// sent: hidden files and directories and __pycache__ directories are skipped. +func sentCodePath(rel string) bool { + segments := strings.Split(rel, "/") + directories := segments[:len(segments)-1] + return !slices.ContainsFunc(segments, isHiddenName) && !slices.Contains(directories, pythonCacheDir) +} + +// skippedCodeFile reports whether a code_source file isn't sent. +func skippedCodeFile(f snapshotFile) bool { + return !sentCodePath(filepath.ToSlash(f.rel)) +} + +// isHiddenName reports whether a file or directory name is hidden by convention. +func isHiddenName(name string) bool { + return strings.HasPrefix(name, ".") +} + +// validateCodePath checks that the server can place a code file at rel. +func validateCodePath(rel string) error { + segments := strings.Split(rel, "/") + if len(segments) <= maxCodePathDepth && !slices.ContainsFunc(segments, invalidCodePathSegment) { + return nil + } + return fmt.Errorf("code_source file %s can't be sent: a path must have at most %d parts of letters, digits, '_', '-' and '.', each starting with a letter, digit or '_'; rename it or leave it out with code_source.include_paths", rel, maxCodePathDepth) +} + +// invalidCodePathSegment reports whether the server can't place a path segment. +func invalidCodePathSegment(segment string) bool { + return !codePathSegmentRe().MatchString(segment) +} diff --git a/cmd/air/deploycode_test.go b/cmd/air/deploycode_test.go new file mode 100644 index 00000000000..c5d50fb88ce --- /dev/null +++ b/cmd/air/deploycode_test.go @@ -0,0 +1,183 @@ +package aircmd + +import ( + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// writeCodeTree writes files (slash-separated path to content) into a new git +// repository and returns a deployment config whose code_source is its root. +func writeCodeTree(t *testing.T, files map[string]string) (*deployConfig, string) { + t.Helper() + return codeTreeConfig(t, newTestRepo(t), files) +} + +// codeTreeConfig writes files under dir and returns a deployment config whose +// code_source is dir. +func codeTreeConfig(t *testing.T, dir string, files map[string]string) (*deployConfig, string) { + t.Helper() + for rel, content := range files { + writeRepoFile(t, dir, filepath.FromSlash(rel), content) + } + entry, task := "agent.py", "agent/v1/responses" + cfg := &deployConfig{ + Model: &deployModelConfig{Type: modelTypeCode, Entry: &entry}, + Task: &task, + CodeSource: &deployCodeSourceConfig{RootPath: dir}, + Environment: &deployEnvironmentConfig{Dependencies: []string{"mlflow==3.5.0"}}, + Compute: &deployComputeConfig{WorkloadType: "CPU", WorkloadSize: "Small"}, + } + return cfg, filepath.Join(dir, "deploy.yaml") +} + +func TestCollectModelCode(t *testing.T) { + cfg, configPath := writeCodeTree(t, map[string]string{ + "agent.py": "import tools.search\n", + "tools/__init__.py": "", + "tools/search.py": "def search(): ...\n", + "tools/__pycache__/search.pyc": "\x00\xff", + ".env": "TOKEN=x\n", + ".venv/lib/site.py": "x = 1\n", + "notes/.draft.py": "y = 2\n", + "deploy.yaml": "endpoint: e\n", + "configs/prompts/system.md": "You are helpful.\n", + "configs/prompts/__init__.py.md": "not hidden\n", + }) + + code, err := collectModelCode(t.Context(), cfg, configPath) + require.NoError(t, err) + assert.Equal(t, "agent.py", code.EntryFile) + assert.Equal(t, "import tools.search\n", code.EntrySource) + assert.Equal(t, "agent/v1/responses", code.Task) + assert.Equal(t, modelPackageEnvironment{Runtime: "client.6", Dependencies: []string{"mlflow==3.5.0"}}, code.Environment) + + var paths []string + for _, f := range code.Files { + paths = append(paths, f.Path) + } + assert.ElementsMatch(t, []string{ + "configs/prompts/__init__.py.md", + "configs/prompts/system.md", + "deploy.yaml", + "tools/__init__.py", + "tools/search.py", + }, paths) +} + +func TestCollectModelCodeIsDeterministic(t *testing.T) { + cfg, configPath := writeCodeTree(t, map[string]string{ + "agent.py": "x = 1\n", "b.py": "b\n", "a/c.py": "c\n", "a.py": "a\n", + }) + first, err := collectModelCode(t.Context(), cfg, configPath) + require.NoError(t, err) + second, err := collectModelCode(t.Context(), cfg, configPath) + require.NoError(t, err) + assert.Equal(t, first, second) +} + +func TestCollectModelCodeIncludePaths(t *testing.T) { + cfg, configPath := writeCodeTree(t, map[string]string{ + "agent.py": "x = 1\n", "tools/a.py": "a\n", "data/big.csv": "1,2\n", + }) + cfg.CodeSource.IncludePaths = []string{"agent.py", "tools"} + code, err := collectModelCode(t.Context(), cfg, configPath) + require.NoError(t, err) + require.Len(t, code.Files, 1) + assert.Equal(t, "tools/a.py", code.Files[0].Path) +} + +func TestCollectModelCodeErrors(t *testing.T) { + tests := []struct { + name string + files map[string]string + include []string + wantErr string + }{ + {"missing entry", map[string]string{"main.py": "x\n"}, nil, "model.entry agent.py isn't among the files code_source sends"}, + {"entry left out", map[string]string{"agent.py": "x\n", "tools/a.py": "a\n"}, []string{"tools"}, "model.entry agent.py isn't among the files code_source sends"}, + {"binary file", map[string]string{"agent.py": "x\n", "model.bin": "\xff\xfe\x00"}, nil, "code_source file model.bin isn't UTF-8 text"}, + {"unplaceable name", map[string]string{"agent.py": "x\n", "my tools/a.py": "a\n"}, nil, "code_source file my tools/a.py can't be sent"}, + {"too deep", map[string]string{"agent.py": "x\n", strings.Repeat("d/", maxCodePathDepth) + "a.py": "a\n"}, nil, "can't be sent: a path must have at most 16 parts"}, + {"too large", map[string]string{"agent.py": "x\n", "blob.txt": strings.Repeat("a", maxCodePackageBytes)}, nil, "total more than 4194304 bytes"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg, configPath := writeCodeTree(t, tt.files) + cfg.CodeSource.IncludePaths = tt.include + _, err := collectModelCode(t.Context(), cfg, configPath) + assert.ErrorContains(t, err, tt.wantErr) + }) + } +} + +func TestCollectModelCodeLeavesOutIgnoredFiles(t *testing.T) { + cfg, configPath := writeCodeTree(t, map[string]string{ + "agent.py": "x\n", "credentials.json": "{}\n", "tools/a.py": "a\n", + }) + require.NoError(t, os.WriteFile(filepath.Join(cfg.CodeSource.RootPath, ".gitignore"), []byte("credentials.json\n"), 0o600)) + code, err := collectModelCode(t.Context(), cfg, configPath) + require.NoError(t, err) + assert.Equal(t, []string{"agent.py", "tools/a.py"}, codeFilePaths(code)) +} + +func TestCollectModelCodeOutsideGitNeedsIncludePaths(t *testing.T) { + cfg, configPath := codeTreeConfig(t, t.TempDir(), map[string]string{"agent.py": "x\n", "credentials.json": "{}\n"}) + _, err := collectModelCode(t.Context(), cfg, configPath) + assert.ErrorContains(t, err, "code_source.include_paths is required because") + assert.ErrorContains(t, err, "isn't in a git repository") + + cfg.CodeSource.IncludePaths = []string{"agent.py"} + code, err := collectModelCode(t.Context(), cfg, configPath) + require.NoError(t, err) + assert.Equal(t, []string{"agent.py"}, codeFilePaths(code)) +} + +func TestCollectModelCodeChecksSizesBeforeReading(t *testing.T) { + // Read in full, this file would fail as binary; the size check must reject it + // from its listing first. + cfg, configPath := writeCodeTree(t, map[string]string{ + "agent.py": "x\n", "data.bin": strings.Repeat("\xff", maxCodePackageBytes+1), + }) + _, err := collectModelCode(t.Context(), cfg, configPath) + assert.ErrorContains(t, err, "total more than 4194304 bytes") +} + +func TestCollectModelCodeRejectsSymlink(t *testing.T) { + cfg, configPath := writeCodeTree(t, map[string]string{"agent.py": "x\n", "real.py": "r\n"}) + dir := filepath.Dir(configPath) + if err := os.Symlink("real.py", filepath.Join(dir, "link.py")); err != nil { + t.Skipf("symlinks unavailable: %v", err) + } + _, err := collectModelCode(t.Context(), cfg, configPath) + assert.ErrorContains(t, err, "code_source path link.py isn't a regular file") +} + +func TestCollectModelCodeTooManyFiles(t *testing.T) { + files := map[string]string{"agent.py": "x\n"} + for i := range maxCodeFiles + 1 { + files[fmt.Sprintf("pkg/m%d.py", i)] = "" + } + cfg, configPath := writeCodeTree(t, files) + _, err := collectModelCode(t.Context(), cfg, configPath) + assert.ErrorContains(t, err, "code_source sends 1002 files, but a model from code can include at most 1000 besides model.entry") +} + +func TestSentCodePath(t *testing.T) { + for rel, want := range map[string]bool{ + "agent.py": true, + "tools/search.py": true, + ".env": false, + "tools/.cache/x.py": false, + "__pycache__/a.pyc": false, + "tools/__pycache__/a.p": false, + "__pycache__": true, + } { + assert.Equal(t, want, sentCodePath(rel), rel) + } +} diff --git a/cmd/air/deployconfig.go b/cmd/air/deployconfig.go new file mode 100644 index 00000000000..4b035cd616e --- /dev/null +++ b/cmd/air/deployconfig.go @@ -0,0 +1,387 @@ +package aircmd + +import ( + "bytes" + "context" + "errors" + "fmt" + "maps" + "reflect" + "regexp" + "slices" + "strconv" + "strings" + "sync" + "unicode/utf8" + + "github.com/databricks/databricks-sdk-go/service/serving" + "go.yaml.in/yaml/v3" +) + +// This file defines the `air deploy` YAML schema and its structural validation. +// Fields shared with `air run` (environment, code_source, env_variables, secrets, +// idempotency_token) keep their meaning; the server validates what needs the +// workspace, such as task names, runtimes and dependencies. + +// Model sources, as model.type. +const ( + modelTypeCheckpoint = "checkpoint" + modelTypeCode = "code" + modelTypeSavedModelDir = "saved_model_dir" + modelTypeRegistered = "registered" +) + +// modelTypes lists model.type values in the order the help and errors name them. +var modelTypes = []string{modelTypeCheckpoint, modelTypeCode, modelTypeSavedModelDir, modelTypeRegistered} + +// checkpointTasks are the tasks a checkpoint can serve, as the registration API +// documents them; models from code serve agent tasks too. +var checkpointTasks = []string{"llm/v1/chat", "llm/v1/completions", "llm/v1/embeddings"} + +// deployCommandPath is the command whose -h documents deployConfig. +const deployCommandPath = "databricks air deploy" + +// defaultServingEnvironmentVersion is the serverless environment version a model +// is served on when environment.version is unset, as for `air run`. +const defaultServingEnvironmentVersion = 6 + +// Limits the server enforces on a registration, checked here to fail early. +const ( + maxServingCommandChars = 32 * 1024 + maxVolumePathChars = 1024 + maxDependencies = 1000 +) + +// cpuWorkloadTypePrefix starts the name of every CPU serving workload type (CPU, +// CPU_MEDIUM, CPU_LARGE); the others are GPU workload types. +const cpuWorkloadTypePrefix = "CPU" + +// registeredModelRe matches .. with the characters the +// server accepts in each part. +var registeredModelRe = sync.OnceValue(func() *regexp.Regexp { return regexp.MustCompile("^[^\\s./`]+\\.[^\\s./`]+\\.[^\\s./`]+$") }) + +// deployIdempotencyKeyRe matches the registration keys the server records. +var deployIdempotencyKeyRe = sync.OnceValue(func() *regexp.Regexp { return regexp.MustCompile(`^[A-Za-z0-9._-]{1,128}$`) }) + +// modelVersionRe matches a Unity Catalog model version number. +var modelVersionRe = sync.OnceValue(func() *regexp.Regexp { return regexp.MustCompile(`^[1-9][0-9]*$`) }) + +// entryFileRe matches the name of a model from code's entry file, which the server +// places at the model's root. +var entryFileRe = sync.OnceValue(func() *regexp.Regexp { return regexp.MustCompile(`^[A-Za-z0-9_][A-Za-z0-9_.-]*\.py$`) }) + +// deployConfig is the `air deploy` YAML schema: a model to register in Unity +// Catalog and the serving endpoint that serves it. +type deployConfig struct { + RegisteredModel string `yaml:"registered_model" help:"Unity Catalog model, as ... air deploy registers a new version of it, creating the model if it doesn't exist, unless model.type is registered." required:"yes" group:"Model"` + Model *deployModelConfig `yaml:"model" help:"What to serve: a checkpoint, a model from code, a model directory saved by MLflow, or an existing version." required:"yes" group:"Model"` + Task *string `yaml:"task" help:"The model's task as MLflow names it: llm/v1/chat, llm/v1/completions or llm/v1/embeddings for a checkpoint, or e.g. agent/v1/responses for a model from code. A model from code without a task has no signature." required:"when model.type is checkpoint" group:"Model"` + Command *string `yaml:"command" help:"Shell command that starts the model server of a checkpoint, at most 32,768 characters. It runs from the directory where model.path appears under its own name." required:"when model.type is checkpoint" group:"Model"` + Environment *deployEnvironmentConfig `yaml:"environment" help:"Serverless environment version to serve on and the packages to install, for a checkpoint or a model from code." group:"Model"` + CodeSource *deployCodeSourceConfig `yaml:"code_source" help:"Local code of a model from code, sent with the registration and stored in the model version, where anyone who can read the model can read it. Hidden files and directories (names starting with '.') and __pycache__ directories aren't sent; run with --dry-run to list the files that are." required:"when model.type is code" group:"Model"` + + Endpoint string `yaml:"endpoint" help:"Serving endpoint to serve the model on. It is created if it doesn't exist. An endpoint that serves one version of registered_model is updated to serve this version instead; air deploy doesn't change an endpoint that serves anything else." required:"yes" group:"Serving"` + Compute *deployComputeConfig `yaml:"compute" help:"Serving compute and how the endpoint scales." required:"yes" group:"Serving"` + EnvVariables map[string]string `yaml:"env_variables" help:"Plain environment variables of the served model, as NAME: value. A name here cannot also appear in secrets." group:"Serving"` + Secrets map[string]string `yaml:"secrets" help:"Environment variables of the served model sourced from secrets, as NAME: scope/key." group:"Serving"` + + IdempotencyToken *string `yaml:"idempotency_token" help:"Key of the registration: 1 to 128 letters, digits, '.', '_' or '-'. Running air deploy again with the same key and unchanged inputs continues that registration instead of registering another version. Defaults to a new key, which air deploy prints." group:"Registration"` + MLflowExperimentPath *string `yaml:"mlflow_experiment_path" help:"Workspace path of the MLflow experiment that records the registered model (default: /Users//air-deployments)." group:"Registration"` +} + +// deployModelConfig is the `model` block: where the served model comes from. +type deployModelConfig struct { + Type string `yaml:"type" help:"One of checkpoint (weights in a volume, started by command), code (a Python file from code_source), saved_model_dir (a model directory saved by MLflow in a volume) or registered (an existing version of registered_model)." required:"yes"` + Path *string `yaml:"path" help:"Unity Catalog volume directory, as /Volumes////, holding a checkpoint's weights or a model saved by MLflow. Model Serving copies it into the version." required:"when type is checkpoint or saved_model_dir"` + Entry *string `yaml:"entry" help:"Python file directly under code_source.root_path that builds the model and calls mlflow.models.set_model, e.g. agent.py." required:"when type is code"` + Version *string `yaml:"version" help:"Version of registered_model to serve, e.g. 3." required:"when type is registered"` +} + +// deployEnvironmentConfig is the `environment` block of a deployment. +type deployEnvironmentConfig struct { + Version *int `yaml:"version" help:"Serverless environment version to serve on, e.g. 6 (default: 6). The model uses its Python version."` + Dependencies []string `yaml:"dependencies" help:"pip requirement lines to install, e.g. einops==0.8.0, at most 1,000. A model from code must list mlflow or mlflow-skinny."` +} + +// deployCodeSourceConfig is the `code_source` block of a model from code. +type deployCodeSourceConfig struct { + RootPath string `yaml:"root_path" help:"Directory holding the model's code. A relative path resolves against the YAML file's directory." required:"yes"` + IncludePaths []string `yaml:"include_paths" help:"Send only these paths, relative to root_path and without \"..\". Omit to send every file that git doesn't ignore; required when root_path isn't in a git repository."` +} + +// deployComputeConfig is the `compute` block: the served entity's workload. +type deployComputeConfig struct { + WorkloadType string `yaml:"workload_type" help:"Serving compute, e.g. CPU, GPU_SMALL, GPU_MEDIUM or GPU_LARGE. A checkpoint is served on a GPU workload type." required:"yes"` + WorkloadSize string `yaml:"workload_size" help:"How far the endpoint scales: Small, Medium or Large." required:"yes"` + ScaleToZero bool `yaml:"scale_to_zero" help:"Whether the endpoint scales to zero when idle (default: false)."` +} + +// sourceFields lists, per model.type, the optional fields that apply to it, in +// addition to the required ones in requiredSourceFields. Every other source field +// is rejected, so a field the source can't use is never silently dropped. +var sourceFields = map[string][]string{ + modelTypeCheckpoint: {"environment", "idempotency_token", "mlflow_experiment_path"}, + modelTypeCode: {"task", "environment", "idempotency_token", "mlflow_experiment_path"}, + modelTypeSavedModelDir: {"idempotency_token", "mlflow_experiment_path"}, + modelTypeRegistered: {}, +} + +// requiredSourceFields lists the fields each model.type requires. +var requiredSourceFields = map[string][]string{ + modelTypeCheckpoint: {"model.path", "task", "command"}, + modelTypeCode: {"model.entry", "code_source"}, + modelTypeSavedModelDir: {"model.path"}, + modelTypeRegistered: {"model.version"}, +} + +// validate runs structural validation over the whole config, returning the +// first failure. Fields are checked in declaration order to keep errors stable. +func (c *deployConfig) validate() error { + if !registeredModelRe().MatchString(c.RegisteredModel) { + return fmt.Errorf("registered_model must be a Unity Catalog model name as .., got %q", c.RegisteredModel) + } + if c.Model == nil { + return errors.New("model: section is required") + } + if err := c.validateSource(); err != nil { + return err + } + if strings.TrimSpace(c.Endpoint) == "" { + return errors.New("endpoint is required") + } + if c.Compute == nil { + return errors.New("compute: section is required") + } + if workloadTypes := servingWorkloadTypes(); !slices.Contains(workloadTypes, c.Compute.WorkloadType) { + return fmt.Errorf("compute.workload_type must be one of %s, got %q", strings.Join(workloadTypes, ", "), c.Compute.WorkloadType) + } + if strings.TrimSpace(c.Compute.WorkloadSize) == "" { + return errors.New("compute.workload_size is required") + } + if c.Model.Type == modelTypeCheckpoint && !c.servesOnGPU() { + return fmt.Errorf("a checkpoint is served on GPU; set compute.workload_type to a GPU workload type such as GPU_MEDIUM, got %q", c.Compute.WorkloadType) + } + if err := validateEnvironmentVariables(c.EnvVariables, c.Secrets); err != nil { + return err + } + if c.IdempotencyToken != nil { + if err := validateDeployIdempotencyKey(*c.IdempotencyToken, "idempotency_token"); err != nil { + return err + } + } + if c.MLflowExperimentPath != nil && !strings.HasPrefix(*c.MLflowExperimentPath, "/") { + return fmt.Errorf("mlflow_experiment_path must be an absolute workspace path such as /Users//air-deployments, got %q", *c.MLflowExperimentPath) + } + return nil +} + +// validateSource checks model.type and the fields that describe the source. +func (c *deployConfig) validateSource() error { + m := c.Model + if !slices.Contains(modelTypes, m.Type) { + return fmt.Errorf("model.type must be one of %s, got %q", strings.Join(modelTypes, ", "), m.Type) + } + + set := map[string]bool{ + "model.path": m.Path != nil, + "model.entry": m.Entry != nil, + "model.version": m.Version != nil, + "task": c.Task != nil, + "command": c.Command != nil, + "environment": c.Environment != nil, + "code_source": c.CodeSource != nil, + "idempotency_token": c.IdempotencyToken != nil, + "mlflow_experiment_path": c.MLflowExperimentPath != nil, + } + required := requiredSourceFields[m.Type] + for _, name := range required { + if !set[name] { + return fmt.Errorf("%s is required when model.type is %s", name, m.Type) + } + } + for _, name := range slices.Sorted(maps.Keys(set)) { + if set[name] && !slices.Contains(required, name) && !slices.Contains(sourceFields[m.Type], name) { + return fmt.Errorf("%s can't be set when model.type is %s", name, m.Type) + } + } + + if m.Path != nil { + v := *m.Path + if !strings.HasPrefix(v, "/Volumes/") || utf8.RuneCountInString(v) > maxVolumePathChars { + return fmt.Errorf("model.path must be a Unity Catalog volume directory as /Volumes////, of at most %d characters, got %q", maxVolumePathChars, v) + } + } + if m.Entry != nil && !entryFileRe().MatchString(*m.Entry) { + return fmt.Errorf("model.entry must name a .py file directly under code_source.root_path, e.g. agent.py, got %q", *m.Entry) + } + if m.Version != nil && !modelVersionRe().MatchString(*m.Version) { + return fmt.Errorf("model.version must be a model version number such as 3, got %q", *m.Version) + } + if c.Task != nil && strings.TrimSpace(*c.Task) == "" { + return errors.New("task cannot be empty") + } + if m.Type == modelTypeCheckpoint && !slices.Contains(checkpointTasks, *c.Task) { + return fmt.Errorf("task of a checkpoint must be one of %s, got %q", strings.Join(checkpointTasks, ", "), *c.Task) + } + if c.Command != nil { + if err := validateCommand(*c.Command); err != nil { + return err + } + if n := utf8.RuneCountInString(*c.Command); n > maxServingCommandChars { + return fmt.Errorf("command must be at most %d characters, got %d", maxServingCommandChars, n) + } + } + if c.Environment != nil { + if err := c.Environment.validate(); err != nil { + return err + } + } + if c.CodeSource != nil { + snapshot := snapshotSourceConfig{RootPath: c.CodeSource.RootPath, IncludePaths: c.CodeSource.IncludePaths} + if err := snapshot.validate("code_source"); err != nil { + return err + } + } + return nil +} + +// validate checks the environment version and dependency lines. +func (e *deployEnvironmentConfig) validate() error { + if e.Version != nil && *e.Version < 1 { + return fmt.Errorf("environment.version must be a serverless environment version such as 6, got %d", *e.Version) + } + if len(e.Dependencies) > maxDependencies { + return fmt.Errorf("environment.dependencies can list at most %d packages, got %d", maxDependencies, len(e.Dependencies)) + } + for i, dependency := range e.Dependencies { + if strings.TrimSpace(dependency) == "" { + return fmt.Errorf("environment.dependencies[%d] cannot be empty", i) + } + } + return nil +} + +// validateDeployIdempotencyKey checks a registration key from source (the YAML +// field or the flag) against the server's rule. +func validateDeployIdempotencyKey(key, source string) error { + if !deployIdempotencyKeyRe().MatchString(key) { + return fmt.Errorf("%s must be 1 to 128 letters, digits, '.', '_' or '-', got %q", source, key) + } + return nil +} + +// servingWorkloadTypes lists the served entity workload types, sorted. +func servingWorkloadTypes() []string { + var workloadType serving.ServingModelWorkloadType + var types []string + for _, t := range workloadType.Values() { + types = append(types, string(t)) + } + slices.Sort(types) + return types +} + +// registers reports whether air deploy registers a new version, rather than +// serving an existing one. +func (c *deployConfig) registers() bool { + return c.Model.Type != modelTypeRegistered +} + +// servesOnGPU reports whether compute.workload_type is a GPU workload type. +func (c *deployConfig) servesOnGPU() bool { + return !strings.HasPrefix(c.Compute.WorkloadType, cpuWorkloadTypePrefix) +} + +// servingRuntime is the serverless runtime the model is served on, as the server +// names it: client., with a -gpu suffix on GPU compute. +func (c *deployConfig) servingRuntime() string { + version := defaultServingEnvironmentVersion + if c.Environment != nil && c.Environment.Version != nil { + version = *c.Environment.Version + } + runtime := "client." + strconv.Itoa(version) + if c.servesOnGPU() { + runtime += "-gpu" + } + return runtime +} + +// servingEnvironment is the registration's runtime and dependencies. +func (c *deployConfig) servingEnvironment() modelPackageEnvironment { + environment := modelPackageEnvironment{Runtime: c.servingRuntime()} + if c.Environment != nil { + environment.Dependencies = c.Environment.Dependencies + } + return environment +} + +// servedEnvironmentVars merges env_variables and secrets into the served +// entity's environment variables, with each secret as {{secrets/scope/key}}. +func (c *deployConfig) servedEnvironmentVars() map[string]string { + if len(c.EnvVariables) == 0 && len(c.Secrets) == 0 { + return nil + } + vars := maps.Clone(c.EnvVariables) + if vars == nil { + vars = map[string]string{} + } + for name, ref := range c.Secrets { + vars[name] = "{{secrets/" + ref + "}}" + } + return vars +} + +// deploySchema is the reflection walk over deployConfig, shared by +// `air deploy -h config.` and --override path validation. +func deploySchema() configField { + return configField{ + path: configHelpRoot, + help: `Deployment YAML schema. For details on a field, pass its path to -h, for example: + + databricks air deploy -h config.model.type + + Minimal example, a checkpoint served by vLLM: + registered_model: main.models.qwen + model: + type: checkpoint + path: /Volumes/main/models/weights/qwen3_8b + task: llm/v1/chat + command: python -m vllm.entrypoints.openai.api_server --model qwen3_8b --served-model-name base --host 0.0.0.0 --port 8080 + endpoint: qwen-chat + compute: + workload_type: GPU_MEDIUM + workload_size: Small`, + children: describeStruct(reflect.TypeFor[deployConfig](), configHelpRoot), + } +} + +// loadDeployConfig reads a deployment YAML, applies --override entries, decodes +// it with unknown keys rejected and validates it. ctx is used only to log applied +// overrides. +func loadDeployConfig(ctx context.Context, path string, overrides []string) (*deployConfig, error) { + var ( + document *yaml.Node + validationYAML []byte + err error + ) + if len(overrides) == 0 { + document, validationYAML, err = readRunConfigDocument(path) + } else { + document, validationYAML, err = readConfigWithOverrides(ctx, path, overrides, deploySchema()) + } + if err != nil { + return nil, err + } + if len(document.Content) == 0 { + return nil, fmt.Errorf("config %s is empty", path) + } + + cfg, err := decodeConfigReader[deployConfig](bytes.NewReader(validationYAML), path) + if err != nil { + return nil, err + } + if err := cfg.validate(); err != nil { + return nil, err + } + return cfg, nil +} diff --git a/cmd/air/deployconfig_test.go b/cmd/air/deployconfig_test.go new file mode 100644 index 00000000000..80f64e76012 --- /dev/null +++ b/cmd/air/deployconfig_test.go @@ -0,0 +1,306 @@ +package aircmd + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const checkpointDeployConfig = `registered_model: main.models.qwen +model: + type: checkpoint + path: /Volumes/main/models/weights/qwen3_8b +task: llm/v1/chat +command: python -m vllm.entrypoints.openai.api_server --model qwen3_8b --port 8080 +environment: + dependencies: [einops==0.8.0] +endpoint: qwen-chat +compute: + workload_type: GPU_MEDIUM + workload_size: Small +secrets: + HF_TOKEN: llm/hf_token +env_variables: + LOG_LEVEL: info +` + +const codeDeployConfig = `registered_model: main.agents.support +model: + type: code + entry: agent.py +task: agent/v1/responses +code_source: + root_path: . +environment: + version: 5 + dependencies: [mlflow==3.5.0] +endpoint: support-agent +compute: + workload_type: CPU + workload_size: Small + scale_to_zero: true +` + +func loadDeployConfigText(t *testing.T, text string, overrides ...string) (*deployConfig, error) { + t.Helper() + return loadDeployConfig(t.Context(), writeConfigFile(t, "deploy.yaml", text), overrides) +} + +func TestLoadDeployConfigCheckpoint(t *testing.T) { + cfg, err := loadDeployConfigText(t, checkpointDeployConfig) + require.NoError(t, err) + assert.True(t, cfg.registers()) + assert.Equal(t, "client.6-gpu", cfg.servingRuntime()) + assert.Equal(t, modelPackageEnvironment{Runtime: "client.6-gpu", Dependencies: []string{"einops==0.8.0"}}, cfg.servingEnvironment()) + assert.Equal(t, map[string]string{"LOG_LEVEL": "info", "HF_TOKEN": "{{secrets/llm/hf_token}}"}, cfg.servedEnvironmentVars()) + + req := cfg.registrationRequest("key-1", nil) + assert.Equal(t, registerModelPackageRequest{ + RegisteredModel: "main.models.qwen", + IdempotencyKey: "key-1", + Checkpoint: &modelPackageCheckpoint{ + VolumePath: "/Volumes/main/models/weights/qwen3_8b", + ServingCommand: "python -m vllm.entrypoints.openai.api_server --model qwen3_8b --port 8080", + Task: "llm/v1/chat", + Environment: modelPackageEnvironment{Runtime: "client.6-gpu", Dependencies: []string{"einops==0.8.0"}}, + }, + }, req) +} + +func TestLoadDeployConfigCodeOnCPU(t *testing.T) { + cfg, err := loadDeployConfigText(t, codeDeployConfig) + require.NoError(t, err) + assert.Equal(t, "client.5", cfg.servingRuntime()) + assert.Nil(t, cfg.servedEnvironmentVars()) + + entity := cfg.servedEntity("4") + assert.Equal(t, "main.agents.support", entity.EntityName) + assert.Equal(t, "4", entity.EntityVersion) + assert.True(t, entity.ScaleToZeroEnabled) + assert.Contains(t, entity.ForceSendFields, "ScaleToZeroEnabled") +} + +func TestLoadDeployConfigRegisteredVersion(t *testing.T) { + cfg, err := loadDeployConfigText(t, `registered_model: main.ml.churn +model: {type: registered, version: 7} +endpoint: churn +compute: {workload_type: CPU, workload_size: Small} +`) + require.NoError(t, err) + assert.False(t, cfg.registers()) + assert.Equal(t, "7", *cfg.Model.Version) +} + +func TestLoadDeployConfigSavedModelDir(t *testing.T) { + cfg, err := loadDeployConfigText(t, `registered_model: main.ml.churn +model: {type: saved_model_dir, path: /Volumes/main/ml/runs/model} +mlflow_experiment_path: /Shared/churn +endpoint: churn +compute: {workload_type: CPU, workload_size: Small} +`) + require.NoError(t, err) + assert.Equal(t, registerModelPackageRequest{ + RegisteredModel: "main.ml.churn", + IdempotencyKey: "k", + SavedModel: &modelPackageSavedModel{VolumePath: "/Volumes/main/ml/runs/model"}, + ExperimentPath: "/Shared/churn", + }, cfg.registrationRequest("k", nil)) +} + +func TestLoadDeployConfigOverride(t *testing.T) { + cfg, err := loadDeployConfigText(t, checkpointDeployConfig, "compute.workload_size=Medium", "environment.version=5") + require.NoError(t, err) + assert.Equal(t, "Medium", cfg.Compute.WorkloadSize) + assert.Equal(t, "client.5-gpu", cfg.servingRuntime()) +} + +func TestLoadDeployConfigRejectsUnknownOverride(t *testing.T) { + _, err := loadDeployConfigText(t, checkpointDeployConfig, "compute.accelerator_type=GPU_1xH100") + assert.ErrorContains(t, err, `invalid --override "compute.accelerator_type": "accelerator_type" is not a known field; available fields are: scale_to_zero, workload_size, workload_type`) +} + +func TestLoadDeployConfigRejectsUnknownField(t *testing.T) { + _, err := loadDeployConfigText(t, checkpointDeployConfig+"experiment_name: x\n") + assert.ErrorContains(t, err, "field experiment_name not found") +} + +func TestLoadDeployConfigRejectsEmptyFile(t *testing.T) { + _, err := loadDeployConfigText(t, "") + assert.ErrorContains(t, err, "is empty") +} + +func TestDeployConfigValidation(t *testing.T) { + tests := []struct { + name string + edit func(string) string + wantErr string + }{ + {"bad model name", replaceFirst("main.models.qwen", "qwen"), "registered_model must be a Unity Catalog model name"}, + {"backticked model name", replaceFirst("main.models.qwen", "main.models.`qwen`"), "registered_model must be a Unity Catalog model name"}, + {"unknown type", replaceFirst("type: checkpoint", "type: pickle"), `model.type must be one of checkpoint, code, saved_model_dir, registered, got "pickle"`}, + {"checkpoint without task", replaceFirst("task: llm/v1/chat\n", ""), "task is required when model.type is checkpoint"}, + {"checkpoint without command", dropLine("command:"), "command is required when model.type is checkpoint"}, + {"checkpoint with entry", replaceFirst(" type: checkpoint\n", " type: checkpoint\n entry: a.py\n"), "model.entry can't be set when model.type is checkpoint"}, + {"checkpoint with code_source", appendText("code_source: {root_path: .}\n"), "code_source can't be set when model.type is checkpoint"}, + {"path outside volumes", replaceFirst("/Volumes/main/models/weights/qwen3_8b", "/dbfs/weights"), "model.path must be a Unity Catalog volume directory"}, + {"path too long", replaceFirst("/Volumes/main/models/weights/qwen3_8b", "/Volumes/"+strings.Repeat("a", 1100)), "model.path must be a Unity Catalog volume directory"}, + {"empty command", replaceFirst("command: python -m vllm.entrypoints.openai.api_server --model qwen3_8b --port 8080", `command: " "`), "command cannot be empty"}, + {"command too long", replaceFirst("command: python", "command: python "+strings.Repeat("x", maxServingCommandChars)), "command must be at most 32768 characters"}, + {"path too long in characters", replaceFirst("/Volumes/main/models/weights/qwen3_8b", "/Volumes/"+strings.Repeat("é", 1020)), "model.path must be a Unity Catalog volume directory"}, + {"checkpoint on CPU", replaceFirst("GPU_MEDIUM", "CPU"), "a checkpoint is served on GPU"}, + {"unknown workload type", replaceFirst("GPU_MEDIUM", "GPU_MEDUIM"), `compute.workload_type must be one of CPU, CPU_LARGE, CPU_MEDIUM, GPU_LARGE, GPU_MEDIUM, GPU_SMALL, GPU_XLARGE, GPU_XLARGE_8, MULTIGPU_MEDIUM, got "GPU_MEDUIM"`}, + {"lowercase workload type", replaceFirst("GPU_MEDIUM", "gpu_medium"), "compute.workload_type must be one of"}, + {"agent task on a checkpoint", replaceFirst("task: llm/v1/chat", "task: agent/v1/responses"), `task of a checkpoint must be one of llm/v1/chat, llm/v1/completions, llm/v1/embeddings, got "agent/v1/responses"`}, + {"zero environment version", replaceFirst(" dependencies: [einops==0.8.0]", " version: 0"), "environment.version must be a serverless environment version"}, + {"empty dependency", replaceFirst("[einops==0.8.0]", `[""]`), "environment.dependencies[0] cannot be empty"}, + {"no endpoint", dropLine("endpoint:"), "endpoint is required"}, + {"no compute", dropBlock("compute:"), "compute: section is required"}, + {"no workload size", dropLine(" workload_size:"), "compute.workload_size is required"}, + {"secret without key", replaceFirst("llm/hf_token", "llm"), "expected format 'scope/key'"}, + {"env var and secret", replaceFirst("LOG_LEVEL: info", "HF_TOKEN: x"), `"HF_TOKEN" is set in both env_variables and secrets`}, + {"bad idempotency token", appendText("idempotency_token: a b\n"), "idempotency_token must be 1 to 128 letters"}, + {"long idempotency token", appendText("idempotency_token: " + strings.Repeat("k", 129) + "\n"), "idempotency_token must be 1 to 128 letters"}, + {"relative experiment", appendText("mlflow_experiment_path: Users/me/x\n"), "mlflow_experiment_path must be an absolute workspace path"}, + {"no model", dropBlock("model:"), "model: section is required"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := loadDeployConfigText(t, tt.edit(checkpointDeployConfig)) + assert.ErrorContains(t, err, tt.wantErr) + }) + } +} + +func TestDeployConfigSourceFields(t *testing.T) { + base := "endpoint: e\ncompute: {workload_type: CPU, workload_size: Small}\nregistered_model: a.b.c\n" + tests := []struct { + name string + source string + wantErr string + }{ + {"code without entry", "model: {type: code}\ncode_source: {root_path: .}\n", "model.entry is required when model.type is code"}, + {"code without code_source", "model: {type: code, entry: agent.py}\n", "code_source is required when model.type is code"}, + {"code with command", "model: {type: code, entry: agent.py}\ncode_source: {root_path: .}\ncommand: x\n", "command can't be set when model.type is code"}, + {"code entry in a directory", "model: {type: code, entry: src/agent.py}\ncode_source: {root_path: .}\n", "model.entry must name a .py file directly under code_source.root_path"}, + {"code entry not python", "model: {type: code, entry: agent.ipynb}\ncode_source: {root_path: .}\n", "model.entry must name a .py file"}, + {"code include traversal", "model: {type: code, entry: agent.py}\ncode_source: {root_path: ., include_paths: [../x]}\n", "code_source.include_paths cannot contain '..' traversal"}, + {"saved model with task", "model: {type: saved_model_dir, path: /Volumes/a/b/c/m}\ntask: llm/v1/chat\n", "task can't be set when model.type is saved_model_dir"}, + {"saved model with environment", "model: {type: saved_model_dir, path: /Volumes/a/b/c/m}\nenvironment: {version: 6}\n", "environment can't be set when model.type is saved_model_dir"}, + {"registered without version", "model: {type: registered}\n", "model.version is required when model.type is registered"}, + {"registered with path", "model: {type: registered, version: 3, path: /Volumes/a/b/c/m}\n", "model.path can't be set when model.type is registered"}, + {"registered with token", "model: {type: registered, version: 3}\nidempotency_token: k\n", "idempotency_token can't be set when model.type is registered"}, + {"registered bad version", "model: {type: registered, version: v3}\n", "model.version must be a model version number"}, + {"registered version zero", "model: {type: registered, version: 0}\n", "model.version must be a model version number"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := loadDeployConfigText(t, base+tt.source) + assert.ErrorContains(t, err, tt.wantErr) + }) + } +} + +func TestDeployConfigCountsCharactersNotBytes(t *testing.T) { + // 1,000 two-byte characters fit the 1,024-character cap on a volume path. + _, err := loadDeployConfigText(t, replaceFirst("/Volumes/main/models/weights/qwen3_8b", "/Volumes/"+strings.Repeat("é", 1000))(checkpointDeployConfig)) + assert.NoError(t, err) +} + +func TestDeployConfigServingRuntimeByWorkload(t *testing.T) { + for workloadType, want := range map[string]string{ + "CPU": "client.6", + "CPU_LARGE": "client.6", + "GPU_SMALL": "client.6-gpu", + "MULTIGPU_MEDIUM": "client.6-gpu", + } { + cfg := deployConfig{Compute: &deployComputeConfig{WorkloadType: workloadType}} + assert.Equal(t, want, cfg.servingRuntime(), workloadType) + } +} + +func TestDeploySchemaHelp(t *testing.T) { + root, err := resolveSchemaField(deploySchema(), "config") + require.NoError(t, err) + var out strings.Builder + renderConfigField(&out, deployCommandPath, root, nil) + help := out.String() + assert.Contains(t, help, " databricks air deploy -h config.model.type\n") + assert.Contains(t, help, " Model:\n") + assert.Contains(t, help, " Serving:\n") + assert.Contains(t, help, " Registration:\n") + assert.Contains(t, help, `For details on a field, run "databricks air deploy -h config.".`) + + field, err := resolveSchemaField(deploySchema(), "model.type") + require.NoError(t, err) + assert.Equal(t, "yes", field.required) + + _, err = resolveSchemaField(deploySchema(), "model.typ") + assert.ErrorContains(t, err, `did you mean "type"?`) + + _, err = resolveSchemaField(deploySchema(), "secrets.HF_TOKEN") + assert.ErrorContains(t, err, "holds user-defined keys") +} + +// Every source field named in sourceFields and requiredSourceFields must be a +// schema path, so the per-source checks can't drift from the YAML. +func TestDeploySourceFieldsAreSchemaPaths(t *testing.T) { + for _, table := range []map[string][]string{sourceFields, requiredSourceFields} { + for _, fields := range table { + for _, f := range fields { + _, err := resolveSchemaField(deploySchema(), f) + assert.NoError(t, err, f) + } + } + } + for _, modelType := range modelTypes { + assert.Contains(t, requiredSourceFields, modelType) + assert.Contains(t, sourceFields, modelType) + } +} + +func replaceFirst(old, with string) func(string) string { + return func(s string) string { + if !strings.Contains(s, old) { + panic("fixture does not contain " + old) + } + return strings.Replace(s, old, with, 1) + } +} + +func appendText(extra string) func(string) string { + return func(s string) string { return s + extra } +} + +// dropLine removes the first line starting with prefix. +func dropLine(prefix string) func(string) string { + return func(s string) string { + lines := strings.SplitAfter(s, "\n") + for i, line := range lines { + if strings.HasPrefix(line, prefix) { + return strings.Join(append(lines[:i:i], lines[i+1:]...), "") + } + } + panic("fixture has no line starting with " + prefix) + } +} + +// dropBlock removes the top-level key starting with prefix and its indented lines. +func dropBlock(prefix string) func(string) string { + return func(s string) string { + var kept []string + dropping := false + for _, line := range strings.SplitAfter(s, "\n") { + switch { + case strings.HasPrefix(line, prefix): + dropping = true + case dropping && strings.HasPrefix(line, " "): + default: + dropping = false + kept = append(kept, line) + } + } + return strings.Join(kept, "") + } +} diff --git a/cmd/air/deployendpoint.go b/cmd/air/deployendpoint.go new file mode 100644 index 00000000000..7a5d4d5f90a --- /dev/null +++ b/cmd/air/deployendpoint.go @@ -0,0 +1,186 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "maps" + "strings" + "time" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/serving" +) + +// servedEntity is the one served entity air deploy puts on the endpoint. +func (c *deployConfig) servedEntity(version string) serving.ServedEntityInput { + return serving.ServedEntityInput{ + EntityName: c.RegisteredModel, + EntityVersion: version, + WorkloadType: serving.ServingModelWorkloadType(c.Compute.WorkloadType), + WorkloadSize: c.Compute.WorkloadSize, + ScaleToZeroEnabled: c.Compute.ScaleToZero, + EnvironmentVars: c.servedEnvironmentVars(), + // scale_to_zero: false is the user's choice, so send it rather than leave it + // to the server's default. + ForceSendFields: []string{"ScaleToZeroEnabled"}, + } +} + +// getEndpoint returns the endpoint named name, or nil if it doesn't exist. +func getEndpoint(ctx context.Context, w *databricks.WorkspaceClient, name string) (*serving.ServingEndpointDetailed, error) { + endpoint, err := w.ServingEndpoints.GetByName(ctx, name) + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return nil, nil + } + return endpoint, err +} + +// checkEndpoint fails when the endpoint named name, if it exists, serves +// something air deploy won't replace with a version of model. air deploy checks +// it before registering, so a wrong endpoint name fails before any weights are +// copied. +func checkEndpoint(ctx context.Context, w *databricks.WorkspaceClient, name, model string) error { + endpoint, err := getEndpoint(ctx, w, name) + if err != nil || endpoint == nil { + return err + } + return checkReplaceable(endpoint, model) +} + +// checkReplaceable fails unless the endpoint serves at most one entity, a version +// of model, in its current and pending configs. Replacing anything else would +// remove another model or a traffic split the user didn't ask to change. +func checkReplaceable(endpoint *serving.ServingEndpointDetailed, model string) error { + for _, served := range [][]serving.ServedEntityOutput{currentEntities(endpoint), pendingEntities(endpoint)} { + if len(served) > 1 || (len(served) == 1 && !strings.EqualFold(served[0].EntityName, model)) { + return fmt.Errorf("endpoint %s serves %s, and air deploy replaces only a single version of %s; set endpoint to another name, or change this endpoint with databricks serving-endpoints update-config", endpoint.Name, describeEntities(served), model) + } + } + return nil +} + +// serveOnEndpoint makes the endpoint named name serve only entity and waits until +// it is ready, calling progress on every poll. It creates the endpoint if it +// doesn't exist. An endpoint that already serves entity isn't updated, and an +// update to entity that an interrupted air deploy started is waited for. +func serveOnEndpoint(ctx context.Context, w *databricks.WorkspaceClient, name string, entity serving.ServedEntityInput, progress func(*serving.ServingEndpointDetailed)) (*serving.ServingEndpointDetailed, error) { + endpoint, err := getEndpoint(ctx, w, name) + if err != nil { + return nil, err + } + if endpoint == nil { + _, err := w.ServingEndpoints.Create(ctx, serving.CreateServingEndpoint{ + Name: name, + Config: &serving.EndpointCoreConfigInput{ServedEntities: []serving.ServedEntityInput{entity}}, + }) + if err != nil { + return nil, err + } + return waitServes(ctx, w, name, entity, progress) + } + if updating(endpoint) { + if servesOnly(pendingEntities(endpoint), entity) { + return waitServes(ctx, w, name, entity, progress) + } + // The endpoint takes no other update until this one finishes. + endpoint, err = w.ServingEndpoints.WaitGetServingEndpointNotUpdating(ctx, name, untilDeadline(ctx), progress) + if err != nil { + return nil, err + } + } + if err := checkReplaceable(endpoint, entity.EntityName); err != nil { + return nil, err + } + if servesOnly(currentEntities(endpoint), entity) { + return checkServes(endpoint, entity) + } + _, err = w.ServingEndpoints.UpdateConfig(ctx, serving.EndpointCoreConfigInput{ + Name: name, + ServedEntities: []serving.ServedEntityInput{entity}, + }) + if err != nil { + return nil, err + } + return waitServes(ctx, w, name, entity, progress) +} + +// waitServes waits until the endpoint's update finishes, then checks that it +// serves entity. +func waitServes(ctx context.Context, w *databricks.WorkspaceClient, name string, entity serving.ServedEntityInput, progress func(*serving.ServingEndpointDetailed)) (*serving.ServingEndpointDetailed, error) { + endpoint, err := w.ServingEndpoints.WaitGetServingEndpointNotUpdating(ctx, name, untilDeadline(ctx), progress) + if err != nil { + return nil, err + } + return checkServes(endpoint, entity) +} + +// checkServes returns endpoint if it serves only entity and is ready to serve. +func checkServes(endpoint *serving.ServingEndpointDetailed, entity serving.ServedEntityInput) (*serving.ServingEndpointDetailed, error) { + if !servesOnly(currentEntities(endpoint), entity) { + return nil, fmt.Errorf("endpoint %s was changed to serve %s while air deploy waited for it", endpoint.Name, describeEntities(currentEntities(endpoint))) + } + if endpoint.State == nil || endpoint.State.Ready != serving.EndpointStateReadyReady { + return nil, fmt.Errorf("endpoint %s serves version %s of %s but isn't ready; check its events in the Serving UI", endpoint.Name, entity.EntityVersion, entity.EntityName) + } + return endpoint, nil +} + +func updating(endpoint *serving.ServingEndpointDetailed) bool { + return endpoint.State != nil && endpoint.State.ConfigUpdate == serving.EndpointStateConfigUpdateInProgress +} + +func currentEntities(endpoint *serving.ServingEndpointDetailed) []serving.ServedEntityOutput { + if endpoint.Config == nil { + return nil + } + return endpoint.Config.ServedEntities +} + +func pendingEntities(endpoint *serving.ServingEndpointDetailed) []serving.ServedEntityOutput { + if endpoint.PendingConfig == nil { + return nil + } + return endpoint.PendingConfig.ServedEntities +} + +// describeEntities names served entities for an error, e.g. "version 2 of a.b.c". +func describeEntities(served []serving.ServedEntityOutput) string { + if len(served) == 0 { + return "nothing" + } + names := make([]string, 0, len(served)) + for _, s := range served { + if s.EntityVersion == "" { + names = append(names, s.EntityName) + continue + } + names = append(names, fmt.Sprintf("version %s of %s", s.EntityVersion, s.EntityName)) + } + return strings.Join(names, " and ") +} + +// servesOnly reports whether served is exactly entity, with the same workload and +// environment variables. +func servesOnly(served []serving.ServedEntityOutput, entity serving.ServedEntityInput) bool { + if len(served) != 1 { + return false + } + s := served[0] + // Unity Catalog names are case-insensitive, and the endpoint reports them as + // Unity Catalog stores them. + return strings.EqualFold(s.EntityName, entity.EntityName) && + s.EntityVersion == entity.EntityVersion && + s.WorkloadType == entity.WorkloadType && + s.WorkloadSize == entity.WorkloadSize && + s.ScaleToZeroEnabled == entity.ScaleToZeroEnabled && + maps.Equal(s.EnvironmentVars, entity.EnvironmentVars) +} + +// untilDeadline is the time left before ctx's deadline, which air deploy always +// sets from --timeout. +func untilDeadline(ctx context.Context) time.Duration { + deadline, _ := ctx.Deadline() + return time.Until(deadline) +} diff --git a/cmd/air/deployendpoint_test.go b/cmd/air/deployendpoint_test.go new file mode 100644 index 00000000000..3733a6a4ae9 --- /dev/null +++ b/cmd/air/deployendpoint_test.go @@ -0,0 +1,280 @@ +package aircmd + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + "github.com/databricks/databricks-sdk-go/service/serving" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const servingEndpointsPath = "/api/2.0/serving-endpoints" + +// fakeEndpoints is a stateful fake of the serving endpoints API for one endpoint. +// A create or config update is applied at once, and the endpoint then reports +// finalState (READY and NOT_UPDATING unless a test sets otherwise). +type fakeEndpoints struct { + t *testing.T + mu sync.Mutex + endpoint *serving.ServingEndpointDetailed + finalState serving.EndpointState + writes []string + // rawWrites holds the body of every create and config update, as sent. + rawWrites []string +} + +func newFakeEndpoints(t *testing.T, existing *serving.ServingEndpointDetailed) *fakeEndpoints { + return &fakeEndpoints{ + t: t, + endpoint: existing, + finalState: serving.EndpointState{Ready: serving.EndpointStateReadyReady, ConfigUpdate: serving.EndpointStateConfigUpdateNotUpdating}, + } +} + +func (f *fakeEndpoints) ServeHTTP(w http.ResponseWriter, r *http.Request) { + f.mu.Lock() + defer f.mu.Unlock() + switch { + case r.Method == http.MethodPost && r.URL.Path == servingEndpointsPath: + var req serving.CreateServingEndpoint + if !f.decode(w, r, &req) { + return + } + f.writes = append(f.writes, "create") + f.endpoint = &serving.ServingEndpointDetailed{Name: req.Name, Config: f.output(req.Config.ServedEntities)} + f.endpoint.State = &f.finalState + case r.Method == http.MethodPut && strings.HasSuffix(r.URL.Path, "/config"): + var req serving.EndpointCoreConfigInput + if !f.decode(w, r, &req) { + return + } + if updating(f.endpoint) { + w.WriteHeader(http.StatusConflict) + _, _ = w.Write([]byte(`{"error_code": "RESOURCE_CONFLICT", "message": "The endpoint is being updated."}`)) + return + } + f.writes = append(f.writes, "update") + f.endpoint.Config = f.output(req.ServedEntities) + f.endpoint.PendingConfig = nil + f.endpoint.State = &f.finalState + case r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, servingEndpointsPath+"/"): + if f.endpoint == nil { + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "Endpoint not found"}`)) + return + } + default: + _, _ = w.Write([]byte(`{}`)) + return + } + assert.NoError(f.t, json.NewEncoder(w).Encode(f.endpoint)) +} + +// decode reads a request body into v, failing the request (not the handler +// goroutine) when it can't. +func (f *fakeEndpoints) decode(w http.ResponseWriter, r *http.Request, v any) bool { + raw, err := io.ReadAll(r.Body) + if !assert.NoError(f.t, err) || !assert.NoError(f.t, json.Unmarshal(raw, v)) { + w.WriteHeader(http.StatusInternalServerError) + return false + } + f.rawWrites = append(f.rawWrites, string(raw)) + return true +} + +func (f *fakeEndpoints) output(entities []serving.ServedEntityInput) *serving.EndpointCoreConfigOutput { + out := &serving.EndpointCoreConfigOutput{ConfigVersion: 1} + for _, e := range entities { + out.ServedEntities = append(out.ServedEntities, servedOutput(e)) + } + return out +} + +func servedOutput(e serving.ServedEntityInput) serving.ServedEntityOutput { + return serving.ServedEntityOutput{ + EntityName: e.EntityName, + EntityVersion: e.EntityVersion, + WorkloadType: e.WorkloadType, + WorkloadSize: e.WorkloadSize, + ScaleToZeroEnabled: e.ScaleToZeroEnabled, + EnvironmentVars: e.EnvironmentVars, + } +} + +func (f *fakeEndpoints) start() *httptest.Server { + srv := httptest.NewServer(f) + f.t.Cleanup(srv.Close) + return srv +} + +func testEntity(version string) serving.ServedEntityInput { + return serving.ServedEntityInput{ + EntityName: "main.models.qwen", + EntityVersion: version, + WorkloadType: serving.ServingModelWorkloadTypeGpuMedium, + WorkloadSize: "Small", + EnvironmentVars: map[string]string{"HF_TOKEN": "{{secrets/llm/hf_token}}"}, + } +} + +func endpointServing(entity serving.ServedEntityInput, state serving.EndpointState) *serving.ServingEndpointDetailed { + return &serving.ServingEndpointDetailed{ + Name: "qwen-chat", + Config: &serving.EndpointCoreConfigOutput{ServedEntities: []serving.ServedEntityOutput{servedOutput(entity)}}, + State: &state, + } +} + +var readyState = serving.EndpointState{Ready: serving.EndpointStateReadyReady, ConfigUpdate: serving.EndpointStateConfigUpdateNotUpdating} + +func TestServeOnEndpointCreatesAMissingEndpoint(t *testing.T) { + fake := newFakeEndpoints(t, nil) + srv := fake.start() + + endpoint, err := serveOnEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", testEntity("3"), nil) + require.NoError(t, err) + assert.Equal(t, []string{"create"}, fake.writes) + assert.Equal(t, "3", endpoint.Config.ServedEntities[0].EntityVersion) + // scale_to_zero: false is sent, not left to the server's default. + require.Len(t, fake.rawWrites, 1) + assert.Contains(t, fake.rawWrites[0], `"scale_to_zero_enabled":false`) +} + +func TestServeOnEndpointUpdatesAnotherVersion(t *testing.T) { + fake := newFakeEndpoints(t, endpointServing(testEntity("2"), readyState)) + srv := fake.start() + + _, err := serveOnEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", testEntity("3"), nil) + require.NoError(t, err) + assert.Equal(t, []string{"update"}, fake.writes) +} + +func TestServeOnEndpointUpdatesChangedCompute(t *testing.T) { + fake := newFakeEndpoints(t, endpointServing(testEntity("3"), readyState)) + srv := fake.start() + entity := testEntity("3") + entity.WorkloadSize = "Medium" + + _, err := serveOnEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", entity, nil) + require.NoError(t, err) + assert.Equal(t, []string{"update"}, fake.writes) +} + +func TestServeOnEndpointLeavesAServedVersion(t *testing.T) { + fake := newFakeEndpoints(t, endpointServing(testEntity("3"), readyState)) + srv := fake.start() + // Unity Catalog reports names lowercased. + entity := testEntity("3") + entity.EntityName = "Main.Models.Qwen" + + _, err := serveOnEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", entity, nil) + require.NoError(t, err) + assert.Empty(t, fake.writes) +} + +func TestServeOnEndpointWaitsForItsOwnPendingUpdate(t *testing.T) { + existing := endpointServing(testEntity("2"), serving.EndpointState{Ready: serving.EndpointStateReadyReady, ConfigUpdate: serving.EndpointStateConfigUpdateInProgress}) + existing.PendingConfig = &serving.EndpointPendingConfig{ServedEntities: []serving.ServedEntityOutput{servedOutput(testEntity("3"))}} + fake := newFakeEndpoints(t, existing) + srv := fake.start() + + polls := 0 + _, err := serveOnEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", testEntity("3"), func(*serving.ServingEndpointDetailed) { + // The pending update finishes after the first poll. + polls++ + fake.mu.Lock() + defer fake.mu.Unlock() + fake.endpoint = endpointServing(testEntity("3"), readyState) + }) + require.NoError(t, err) + assert.Empty(t, fake.writes) + assert.Positive(t, polls) +} + +func TestServeOnEndpointWaitsForAnotherUpdateBeforeUpdating(t *testing.T) { + existing := endpointServing(testEntity("1"), serving.EndpointState{Ready: serving.EndpointStateReadyReady, ConfigUpdate: serving.EndpointStateConfigUpdateInProgress}) + existing.PendingConfig = &serving.EndpointPendingConfig{ServedEntities: []serving.ServedEntityOutput{servedOutput(testEntity("2"))}} + fake := newFakeEndpoints(t, existing) + srv := fake.start() + + _, err := serveOnEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", testEntity("3"), func(endpoint *serving.ServingEndpointDetailed) { + // The other update finishes after the first poll. + if updating(endpoint) { + fake.mu.Lock() + defer fake.mu.Unlock() + fake.endpoint = endpointServing(testEntity("2"), readyState) + } + }) + require.NoError(t, err) + assert.Equal(t, []string{"update"}, fake.writes) + assert.Equal(t, "3", fake.endpoint.Config.ServedEntities[0].EntityVersion) +} + +func TestServeOnEndpointRefusesToReplaceOtherEntities(t *testing.T) { + other := testEntity("5") + other.EntityName = "main.models.llama" + split := endpointServing(testEntity("2"), readyState) + split.Config.ServedEntities = append(split.Config.ServedEntities, servedOutput(testEntity("1"))) + tests := []struct { + name string + existing *serving.ServingEndpointDetailed + wantErr string + }{ + {"another model", endpointServing(other, readyState), "endpoint qwen-chat serves version 5 of main.models.llama, and air deploy replaces only a single version of main.models.qwen"}, + {"a traffic split", split, "endpoint qwen-chat serves version 2 of main.models.qwen and version 1 of main.models.qwen"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + fake := newFakeEndpoints(t, tt.existing) + w := newTestWorkspaceClient(t, fake.start().URL) + + assert.ErrorContains(t, checkEndpoint(withDeadline(t), w, "qwen-chat", "main.models.qwen"), tt.wantErr) + _, err := serveOnEndpoint(withDeadline(t), w, "qwen-chat", testEntity("3"), nil) + assert.ErrorContains(t, err, tt.wantErr) + assert.Empty(t, fake.writes) + }) + } +} + +func TestCheckEndpointAllowsAMissingEndpoint(t *testing.T) { + srv := newFakeEndpoints(t, nil).start() + assert.NoError(t, checkEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", "main.models.qwen")) +} + +func TestServeOnEndpointReportsAFailedUpdate(t *testing.T) { + fake := newFakeEndpoints(t, nil) + fake.finalState = serving.EndpointState{Ready: serving.EndpointStateReadyNotReady, ConfigUpdate: serving.EndpointStateConfigUpdateUpdateFailed} + srv := fake.start() + + _, err := serveOnEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", testEntity("3"), nil) + assert.ErrorContains(t, err, "UPDATE_FAILED") +} + +func TestServeOnEndpointReportsAnEndpointThatIsNotReady(t *testing.T) { + fake := newFakeEndpoints(t, endpointServing(testEntity("3"), serving.EndpointState{Ready: serving.EndpointStateReadyNotReady, ConfigUpdate: serving.EndpointStateConfigUpdateNotUpdating})) + srv := fake.start() + + _, err := serveOnEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", testEntity("3"), nil) + assert.ErrorContains(t, err, "endpoint qwen-chat serves version 3 of main.models.qwen but isn't ready") +} + +func TestServesOnly(t *testing.T) { + entity := testEntity("3") + assert.True(t, servesOnly([]serving.ServedEntityOutput{servedOutput(entity)}, entity)) + assert.False(t, servesOnly(nil, entity)) + assert.False(t, servesOnly([]serving.ServedEntityOutput{servedOutput(entity), servedOutput(testEntity("2"))}, entity)) + + changed := entity + changed.EnvironmentVars = map[string]string{"HF_TOKEN": "{{secrets/llm/other}}"} + assert.False(t, servesOnly([]serving.ServedEntityOutput{servedOutput(entity)}, changed)) + + scaled := entity + scaled.ScaleToZeroEnabled = true + assert.False(t, servesOnly([]serving.ServedEntityOutput{servedOutput(entity)}, scaled)) +} diff --git a/cmd/air/deployregister.go b/cmd/air/deployregister.go new file mode 100644 index 00000000000..f09867aa04f --- /dev/null +++ b/cmd/air/deployregister.go @@ -0,0 +1,172 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "net/http" + "time" + + "github.com/databricks/cli/libs/auth" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" +) + +// registerModelPackagePath registers a model whose files Model Serving writes. It +// runs as the caller and is idempotent per key: a repeat request with the same key +// and inputs continues the registration, and finalizes a checkpoint's version once +// its weights are staged. +const registerModelPackagePath = "/api/2.0/serving-endpoints:registerModelPackage" + +// registrationPollInterval is how often air deploy repeats the registration while +// a checkpoint's weights are staged. +const registrationPollInterval = 15 * time.Second + +// stagingStateFailed is the ModelDeploymentStagingStatus state of an artifact that +// failed to stage and won't be retried automatically. +const stagingStateFailed = "FAILED" + +// registerModelPackageRequest is the RegisterModelPackageRequest body. Exactly one +// of Checkpoint, Code and SavedModel is set. +type registerModelPackageRequest struct { + RegisteredModel string `json:"registered_model"` + IdempotencyKey string `json:"idempotency_key"` + Checkpoint *modelPackageCheckpoint `json:"checkpoint,omitempty"` + Code *modelPackageCode `json:"code,omitempty"` + SavedModel *modelPackageSavedModel `json:"saved_model,omitempty"` + ExperimentPath string `json:"experiment_path,omitempty"` + RetryFailedStaging bool `json:"retry_failed_staging,omitempty"` +} + +type modelPackageCheckpoint struct { + VolumePath string `json:"volume_path"` + ServingCommand string `json:"serving_command"` + Task string `json:"task"` + Environment modelPackageEnvironment `json:"environment"` +} + +type modelPackageCode struct { + EntryFile string `json:"entry_file"` + EntrySource string `json:"entry_source"` + Files []modelPackageFile `json:"files,omitempty"` + Task string `json:"task,omitempty"` + Environment modelPackageEnvironment `json:"environment"` +} + +type modelPackageFile struct { + Path string `json:"path"` + Content string `json:"content"` +} + +type modelPackageSavedModel struct { + VolumePath string `json:"volume_path"` +} + +type modelPackageEnvironment struct { + Runtime string `json:"runtime"` + Dependencies []string `json:"dependencies,omitempty"` +} + +// registerModelPackageResponse is the RegisterModelPackageResponse body. +type registerModelPackageResponse struct { + ModelID string `json:"model_id"` + RegisteredModel string `json:"registered_model"` + ModelVersion string `json:"model_version"` + Ready bool `json:"ready"` + StagingStatus *modelStagingStatus `json:"staging_status,omitempty"` +} + +type modelStagingStatus struct { + State string `json:"state,omitempty"` + ModelArtifact *artifactStagingStatus `json:"model_artifact,omitempty"` + EnvironmentArtifact *artifactStagingStatus `json:"environment_artifact,omitempty"` +} + +type artifactStagingStatus struct { + State string `json:"state,omitempty"` +} + +// weightsState is the staging state of a checkpoint's weights, or "" when the +// response reports none. +func (r *registerModelPackageResponse) weightsState() string { + if r.StagingStatus == nil || r.StagingStatus.ModelArtifact == nil { + return "" + } + return r.StagingStatus.ModelArtifact.State +} + +// environmentState is the staging state of the serving environment, or "" when the +// response reports none. +func (r *registerModelPackageResponse) environmentState() string { + if r.StagingStatus == nil || r.StagingStatus.EnvironmentArtifact == nil { + return "" + } + return r.StagingStatus.EnvironmentArtifact.State +} + +// errStagingFailed reports that Model Serving couldn't copy a checkpoint's weights +// into the version. Rerunning with --retry-failed-staging stages them again. +var errStagingFailed = errors.New("staging failed") + +// registrationRequest builds the request that registers cfg's model under key. +// code is the collected model from code, and is nil for the other sources. +func (c *deployConfig) registrationRequest(key string, code *modelPackageCode) registerModelPackageRequest { + req := registerModelPackageRequest{ + RegisteredModel: c.RegisteredModel, + IdempotencyKey: key, + } + if c.MLflowExperimentPath != nil { + req.ExperimentPath = *c.MLflowExperimentPath + } + switch c.Model.Type { + case modelTypeCheckpoint: + req.Checkpoint = &modelPackageCheckpoint{ + VolumePath: *c.Model.Path, + ServingCommand: *c.Command, + Task: *c.Task, + Environment: c.servingEnvironment(), + } + case modelTypeCode: + req.Code = code + case modelTypeSavedModelDir: + req.SavedModel = &modelPackageSavedModel{VolumePath: *c.Model.Path} + } + return req +} + +// registerUntilReady sends req until the version it registers is ready to serve, +// calling progress with every response. A checkpoint's version becomes ready once +// its weights are staged; the other sources are ready after the first request. +// It returns errStagingFailed, with the last response, when staging the weights +// failed. +func registerUntilReady(ctx context.Context, w *databricks.WorkspaceClient, req registerModelPackageRequest, interval time.Duration, progress func(*registerModelPackageResponse)) (*registerModelPackageResponse, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + headers := auth.WorkspaceIDHeaders(w.Config) + for { + var resp registerModelPackageResponse + if err := apiClient.Do(ctx, http.MethodPost, registerModelPackagePath, headers, nil, req, &resp); err != nil { + return nil, err + } + progress(&resp) + if resp.Ready { + return &resp, nil + } + if resp.weightsState() == stagingStateFailed { + return &resp, errStagingFailed + } + // Only the first request restages a failed checkpoint: repeating it would + // restart a staging that keeps failing instead of reporting the failure. + req.RetryFailedStaging = false + + timer := time.NewTimer(interval) + select { + case <-ctx.Done(): + timer.Stop() + return &resp, ctx.Err() + case <-timer.C: + } + } +} diff --git a/cmd/air/deployregister_test.go b/cmd/air/deployregister_test.go new file mode 100644 index 00000000000..90732e14df5 --- /dev/null +++ b/cmd/air/deployregister_test.go @@ -0,0 +1,165 @@ +package aircmd + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// fakeModelServing serves registerModelPackage with the given responses in turn, +// repeating the last, and records every registration body. Other paths are served +// by endpoints when set, else get an empty object (the SDK's config probe). +type fakeModelServing struct { + t *testing.T + mu sync.Mutex + responses []string + bodies []registerModelPackageRequest + endpoints http.HandlerFunc +} + +func newFakeModelServing(t *testing.T, responses ...string) *fakeModelServing { + return &fakeModelServing{t: t, responses: responses} +} + +func (f *fakeModelServing) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != registerModelPackagePath { + if f.endpoints != nil { + f.endpoints(w, r) + return + } + _, _ = w.Write([]byte(`{}`)) + return + } + raw, err := io.ReadAll(r.Body) + var body registerModelPackageRequest + if !assert.NoError(f.t, err) || !assert.NoError(f.t, json.Unmarshal(raw, &body)) { + w.WriteHeader(http.StatusInternalServerError) + return + } + + f.mu.Lock() + defer f.mu.Unlock() + f.bodies = append(f.bodies, body) + i := min(len(f.bodies), len(f.responses)) - 1 + if f.responses[i] == "" { + w.WriteHeader(http.StatusConflict) + _, _ = w.Write([]byte(`{"error_code": "RESOURCE_CONFLICT", "message": "The idempotency key can't be used for this registration. Use a new key."}`)) + return + } + _, _ = w.Write([]byte(f.responses[i])) +} + +func (f *fakeModelServing) start() *httptest.Server { + srv := httptest.NewServer(f) + f.t.Cleanup(srv.Close) + return srv +} + +func (f *fakeModelServing) requests() []registerModelPackageRequest { + f.mu.Lock() + defer f.mu.Unlock() + return append([]registerModelPackageRequest(nil), f.bodies...) +} + +const ( + stagingInProgress = `{"model_id": "m-1", "registered_model": "main.models.qwen", "model_version": "3", "ready": false, + "staging_status": {"state": "IN_PROGRESS", "model_artifact": {"state": "IN_PROGRESS"}, "environment_artifact": {"state": "SUCCEEDED"}}}` + stagingReady = `{"model_id": "m-1", "registered_model": "main.models.qwen", "model_version": "3", "ready": true, + "staging_status": {"state": "SUCCEEDED", "model_artifact": {"state": "SUCCEEDED"}, "environment_artifact": {"state": "SUCCEEDED"}}}` + stagingFailed = `{"model_id": "m-1", "registered_model": "main.models.qwen", "model_version": "3", "ready": false, + "staging_status": {"state": "FAILED", "model_artifact": {"state": "FAILED"}, "environment_artifact": {"state": "SUCCEEDED"}}}` +) + +func checkpointRequest() registerModelPackageRequest { + return registerModelPackageRequest{ + RegisteredModel: "main.models.qwen", + IdempotencyKey: "key-1", + Checkpoint: &modelPackageCheckpoint{ + VolumePath: "/Volumes/main/models/weights/qwen3_8b", + ServingCommand: "python -m vllm.entrypoints.openai.api_server --model qwen3_8b", + Task: "llm/v1/chat", + Environment: modelPackageEnvironment{Runtime: "client.6-gpu"}, + }, + } +} + +func TestRegisterUntilReadyRepeatsTheSameRequest(t *testing.T) { + fake := newFakeModelServing(t, stagingInProgress, stagingInProgress, stagingReady) + srv := fake.start() + + var states []string + req := checkpointRequest() + req.RetryFailedStaging = true + resp, err := registerUntilReady(t.Context(), newTestWorkspaceClient(t, srv.URL), req, time.Millisecond, func(r *registerModelPackageResponse) { + states = append(states, r.weightsState()) + }) + require.NoError(t, err) + assert.True(t, resp.Ready) + assert.Equal(t, "3", resp.ModelVersion) + assert.Equal(t, []string{"IN_PROGRESS", "IN_PROGRESS", "SUCCEEDED"}, states) + + bodies := fake.requests() + require.Len(t, bodies, 3) + assert.Equal(t, req, bodies[0]) + // Only the first request restages; the repeats continue the same registration. + req.RetryFailedStaging = false + assert.Equal(t, req, bodies[1]) + assert.Equal(t, req, bodies[2]) +} + +func TestRegisterUntilReadyReportsFailedStaging(t *testing.T) { + fake := newFakeModelServing(t, stagingInProgress, stagingFailed) + srv := fake.start() + + resp, err := registerUntilReady(t.Context(), newTestWorkspaceClient(t, srv.URL), checkpointRequest(), time.Millisecond, func(*registerModelPackageResponse) {}) + require.ErrorIs(t, err, errStagingFailed) + assert.Equal(t, "3", resp.ModelVersion) + assert.Len(t, fake.requests(), 2) +} + +func TestRegisterUntilReadyReturnsServerErrors(t *testing.T) { + srv := newFakeModelServing(t, "").start() + + _, err := registerUntilReady(t.Context(), newTestWorkspaceClient(t, srv.URL), checkpointRequest(), time.Millisecond, func(*registerModelPackageResponse) {}) + assert.ErrorContains(t, err, "The idempotency key can't be used for this registration") +} + +func TestRegisterUntilReadyStopsWhenCanceled(t *testing.T) { + srv := newFakeModelServing(t, stagingInProgress).start() + ctx, cancel := context.WithCancel(t.Context()) + + _, err := registerUntilReady(ctx, newTestWorkspaceClient(t, srv.URL), checkpointRequest(), time.Hour, func(*registerModelPackageResponse) { cancel() }) + assert.ErrorIs(t, err, context.Canceled) +} + +func TestRegisterModelPackageRequestJSON(t *testing.T) { + raw, err := json.Marshal(registerModelPackageRequest{ + RegisteredModel: "a.b.c", + IdempotencyKey: "k", + Code: &modelPackageCode{ + EntryFile: "agent.py", + EntrySource: "x = 1\n", + Files: []modelPackageFile{{Path: "tools/a.py", Content: ""}}, + Environment: modelPackageEnvironment{Runtime: "client.6", Dependencies: []string{"mlflow"}}, + }, + }) + require.NoError(t, err) + assert.JSONEq(t, `{ + "registered_model": "a.b.c", + "idempotency_key": "k", + "code": { + "entry_file": "agent.py", + "entry_source": "x = 1\n", + "files": [{"path": "tools/a.py", "content": ""}], + "environment": {"runtime": "client.6", "dependencies": ["mlflow"]} + } + }`, string(raw)) +} diff --git a/cmd/air/run.go b/cmd/air/run.go index 20a7683461f..6492764d3c0 100644 --- a/cmd/air/run.go +++ b/cmd/air/run.go @@ -131,7 +131,7 @@ field, pass its path to -h as a separate argument: if field.path == "config.compute" || field.path == "config.compute.accelerator_type" { computeOptions = loadComputeOptionsForHelp(c, args) } - renderConfigField(c.OutOrStdout(), field, computeOptions) + renderConfigField(c.OutOrStdout(), runCommandPath, field, computeOptions) }) // List flags in declaration order (most used first) rather than alphabetically. @@ -170,7 +170,7 @@ field, pass its path to -h as a separate argument: } snapshotFileCount = &count } - idempotencyToken := effectiveIdempotencyToken(idempotencyKey, cfg) + idempotencyToken := effectiveIdempotencyToken(idempotencyKey, cfg.IdempotencyToken) ctx, validationErr := validateDryRunInWorkspace(ctx, cmd, args, cfg, idempotencyToken, validationTimeout) if err := validateIdempotencyToken(idempotencyToken); err != nil { return err diff --git a/cmd/air/runconfig.go b/cmd/air/runconfig.go index 78554b02183..7306c1ab3c6 100644 --- a/cmd/air/runconfig.go +++ b/cmd/air/runconfig.go @@ -119,18 +119,10 @@ func (c *runConfig) validate() error { } } - if err := validateSecretRefs(c.Secrets); err != nil { + if err := validateEnvironmentVariables(c.EnvVariables, c.Secrets); err != nil { return err } - // A name can't be both a plain env var and a secret: the precedence would be - // ambiguous and could leak the secret. Sorted for a stable error. - for _, name := range slices.Sorted(maps.Keys(c.EnvVariables)) { - if _, ok := c.Secrets[name]; ok { - return fmt.Errorf("%q is set in both env_variables and secrets; remove it from one", name) - } - } - if c.CodeSource != nil { if err := c.CodeSource.validate(); err != nil { return err @@ -335,6 +327,23 @@ func validateCommand(v string) error { return nil } +// validateEnvironmentVariables checks the env_variables and secrets of a config: +// secret references use the "scope/key" format, and no name is in both. +func validateEnvironmentVariables(envVariables, secrets map[string]string) error { + if err := validateSecretRefs(secrets); err != nil { + return err + } + + // A name can't be both a plain env var and a secret: the precedence would be + // ambiguous and could leak the secret. Sorted for a stable error. + for _, name := range slices.Sorted(maps.Keys(envVariables)) { + if _, ok := secrets[name]; ok { + return fmt.Errorf("%q is set in both env_variables and secrets; remove it from one", name) + } + } + return nil +} + // validateSecretRefs checks that secret references use the "scope/key" format. func validateSecretRefs(secrets map[string]string) error { for varName, ref := range secrets { @@ -608,13 +617,16 @@ func (p *permission) validate() error { // or bare `compute`). const configHelpRoot = "config" +// runCommandPath is the command whose -h documents runConfig. +const runCommandPath = "databricks air run" + // writeConfigFieldHelp resolves a dotted config path and writes its docs. func writeConfigFieldHelp(w io.Writer, path string) error { field, err := resolveConfigField(path) if err != nil { return err } - renderConfigField(w, field, nil) + renderConfigField(w, runCommandPath, field, nil) return nil } @@ -675,8 +687,13 @@ func configSchema() configField { // resolveConfigField resolves a dotted YAML path against the run config schema. // The leading "config." is optional; an empty path describes the whole schema. func resolveConfigField(path string) (configField, error) { + return resolveSchemaField(configSchema(), path) +} + +// resolveSchemaField resolves a dotted YAML path against a config schema rooted +// at root, as resolveConfigField does for the run config. +func resolveSchemaField(root configField, path string) (configField, error) { trimmed := strings.TrimPrefix(strings.TrimPrefix(path, configHelpRoot), ".") - root := configSchema() if trimmed == "" { return root, nil } @@ -867,7 +884,8 @@ func underlyingConfigStruct(t reflect.Type) reflect.Type { // renderConfigField writes a resolved field's documentation. An object lists its // immediate children; a leaf gets its type, required-ness, and description. -func renderConfigField(w io.Writer, f configField, computeOptions []computeOption) { +// command is the command whose -h documents the schema. +func renderConfigField(w io.Writer, command string, f configField, computeOptions []computeOption) { fmt.Fprintf(w, "%s\n", f.path) if f.help != "" { fmt.Fprintf(w, " %s\n", f.help) @@ -916,8 +934,9 @@ func renderConfigField(w io.Writer, f configField, computeOptions []computeOptio fmt.Fprintf(w, " %-*s %s\n", width, configLeafName(c.path), configFieldSummary(c)) } } - fmt.Fprintf(w, "\nFor details on a field, run \"databricks air run -h %s.\".\n", f.path) - if f.path == configHelpRoot { + fmt.Fprintf(w, "\nFor details on a field, run \"%s -h %s.\".\n", command, f.path) + // The skill writes run YAML, so only `air run -h config` points at it. + if f.path == configHelpRoot && command == runCommandPath { fmt.Fprint(w, agentSkillHelp) } } diff --git a/cmd/air/runconfig_load.go b/cmd/air/runconfig_load.go index 2e05e1b50e6..e272bc8f096 100644 --- a/cmd/air/runconfig_load.go +++ b/cmd/air/runconfig_load.go @@ -18,10 +18,16 @@ import ( // decodeRunConfigReader decodes and unknown-key-checks a run YAML from r. path is // used only for error messages. func decodeRunConfigReader(r io.Reader, path string) (*runConfig, error) { + return decodeConfigReader[runConfig](r, path) +} + +// decodeConfigReader decodes and unknown-key-checks a YAML config from r into a +// T. path is used only for error messages. +func decodeConfigReader[T any](r io.Reader, path string) (*T, error) { dec := yaml.NewDecoder(r) dec.KnownFields(true) - var cfg runConfig + var cfg T if err := dec.Decode(&cfg); err != nil { if errors.Is(err, io.EOF) { return nil, fmt.Errorf("config %s is empty", path) @@ -59,30 +65,41 @@ func loadRunConfigWithOverrides(ctx context.Context, path string, overrides []st if len(overrides) == 0 { return loadRunConfig(path) } + document, validationYAML, err := readConfigWithOverrides(ctx, path, overrides, configSchema()) + if err != nil { + return nil, err + } + return finishRunConfigLoad(path, document, validationYAML) +} +// readConfigWithOverrides reads a YAML config and applies --override KEY=VALUE +// entries, checked against schema, to its ordered YAML tree. It returns the tree +// and its serialization, which the caller decodes with unknown keys rejected. +// ctx is used only to log applied overrides. +func readConfigWithOverrides(ctx context.Context, path string, overrides []string, schema configField) (*yaml.Node, []byte, error) { entries, err := parseOverrides(overrides) if err != nil { - return nil, err + return nil, nil, err } - if err := validateOverridePaths(entries); err != nil { - return nil, err + if err := validateOverridePaths(entries, schema); err != nil { + return nil, nil, err } document, _, err := readRunConfigDocument(path) if err != nil { - return nil, err + return nil, nil, err } if len(document.Content) == 0 { document.Kind = yaml.DocumentNode document.Content = []*yaml.Node{newMappingNode()} } if err := applyOverrides(ctx, document, entries); err != nil { - return nil, err + return nil, nil, err } validationYAML, err := yaml.Marshal(document) if err != nil { - return nil, fmt.Errorf("failed to serialize config %s: %w", path, err) + return nil, nil, fmt.Errorf("failed to serialize config %s: %w", path, err) } - return finishRunConfigLoad(path, document, validationYAML) + return document, validationYAML, nil } func readRunConfigDocument(path string) (*yaml.Node, []byte, error) { diff --git a/cmd/air/runconfig_override.go b/cmd/air/runconfig_override.go index 3fbccd076cb..d7585c8214a 100644 --- a/cmd/air/runconfig_override.go +++ b/cmd/air/runconfig_override.go @@ -45,11 +45,10 @@ type overrideEntry struct { raw string } -// validateOverridePaths checks every dotted path against the runConfig schema +// validateOverridePaths checks every dotted path against a config schema // before mutation, so an error names the exact --override key rather than the // re-decode's Go-type language. -func validateOverridePaths(entries []overrideEntry) error { - schema := configSchema() +func validateOverridePaths(entries []overrideEntry, schema configField) error { for _, e := range entries { if err := checkOverridePath(strings.Split(e.path, "."), schema, e.path); err != nil { return err @@ -100,7 +99,7 @@ func applyOverrides(ctx context.Context, document *yaml.Node, entries []override root = document.Content[0] } if root.Kind != yaml.MappingNode { - return errors.New("run config must be a YAML mapping") + return errors.New("config must be a YAML mapping") } for _, e := range entries { diff --git a/cmd/air/runconfig_override_test.go b/cmd/air/runconfig_override_test.go index c0dc4937138..982f9ad74a0 100644 --- a/cmd/air/runconfig_override_test.go +++ b/cmd/air/runconfig_override_test.go @@ -90,7 +90,7 @@ func TestValidateOverridePaths(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - err := validateOverridePaths([]overrideEntry{{path: tt.path, raw: "x"}}) + err := validateOverridePaths([]overrideEntry{{path: tt.path, raw: "x"}}, configSchema()) if tt.wantErr != "" { require.ErrorContains(t, err, tt.wantErr) return diff --git a/cmd/air/runconfig_test.go b/cmd/air/runconfig_test.go index 74b4c187628..64413f5d789 100644 --- a/cmd/air/runconfig_test.go +++ b/cmd/air/runconfig_test.go @@ -1057,14 +1057,14 @@ func TestConfigSchemaSharedByHelpAndOverride(t *testing.T) { for _, p := range []string{"code_source.type", "code_source.snapshot.root_path"} { t.Run(p+" compatibility override", func(t *testing.T) { - require.NoError(t, validateOverridePaths([]overrideEntry{{path: p, raw: "x"}})) + require.NoError(t, validateOverridePaths([]overrideEntry{{path: p, raw: "x"}}, configSchema())) _, err := resolveConfigField(p) require.Error(t, err) }) } for _, p := range paths { t.Run(p, func(t *testing.T) { - require.NoError(t, validateOverridePaths([]overrideEntry{{path: p, raw: "x"}})) + require.NoError(t, validateOverridePaths([]overrideEntry{{path: p, raw: "x"}}, configSchema())) // The help path accepts the same field (free-form sub-paths resolve to // the map itself, which is the schema's most specific node). _, err := resolveConfigField(p) @@ -1075,7 +1075,7 @@ func TestConfigSchemaSharedByHelpAndOverride(t *testing.T) { } // A field unknown to one is unknown to the other. - require.Error(t, validateOverridePaths([]overrideEntry{{path: "compute.bogus", raw: "x"}})) + require.Error(t, validateOverridePaths([]overrideEntry{{path: "compute.bogus", raw: "x"}}, configSchema())) _, err := resolveConfigField("compute.bogus") require.Error(t, err) } diff --git a/cmd/air/runsubmit.go b/cmd/air/runsubmit.go index c06165dd2b4..8ee0beca466 100644 --- a/cmd/air/runsubmit.go +++ b/cmd/air/runsubmit.go @@ -270,10 +270,10 @@ func aiRuntimeTaskFromSubmitBody(body map[string]any) (map[string]any, error) { // effectiveIdempotencyToken resolves the token sent to validation and submit: // the --idempotency-key flag wins, then the config's token, else a generated one. -func effectiveIdempotencyToken(flag string, cfg *runConfig) string { +func effectiveIdempotencyToken(flag string, configured *string) string { token := flag - if token == "" && cfg.IdempotencyToken != nil { - token = *cfg.IdempotencyToken + if token == "" && configured != nil { + token = *configured } if token == "" { token = uuid.New().String() @@ -345,7 +345,7 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run return 0, "", err } containers := submittedContainers(cfg, funcDir) - token := effectiveIdempotencyToken(idempotencyKey, cfg) + token := effectiveIdempotencyToken(idempotencyKey, cfg.IdempotencyToken) deps, _ := cfg.inlineDependencies() diff --git a/cmd/air/runsubmit_test.go b/cmd/air/runsubmit_test.go index e101bbc28af..daa87788e8b 100644 --- a/cmd/air/runsubmit_test.go +++ b/cmd/air/runsubmit_test.go @@ -33,9 +33,9 @@ func stubValidateConfig(server *testserver.Server) { func TestEffectiveIdempotencyToken(t *testing.T) { cfg := &runConfig{IdempotencyToken: new("from-config")} - assert.Equal(t, "from-flag", effectiveIdempotencyToken("from-flag", cfg)) - assert.Equal(t, "from-config", effectiveIdempotencyToken("", cfg)) - assert.NotEmpty(t, effectiveIdempotencyToken("", &runConfig{})) + assert.Equal(t, "from-flag", effectiveIdempotencyToken("from-flag", cfg.IdempotencyToken)) + assert.Equal(t, "from-config", effectiveIdempotencyToken("", cfg.IdempotencyToken)) + assert.NotEmpty(t, effectiveIdempotencyToken("", nil)) } func TestDlRuntimeImage(t *testing.T) { From 11000234dab303d74666b05149897198b4f06287 Mon Sep 17 00:00:00 2001 From: Arpit Agarwal Date: Sun, 15 Nov 2026 06:27:26 +0000 Subject: [PATCH 2/4] air deploy: poll staging status, guard endpoints and code, type JSON errors - Register once, poll the read-only staging status API while a checkpoint's weights copy, then register again to finalize. - Retry transient failures of the polled requests via libs/retry. - Check the endpoint before registering: replace only a single version of the same model, and wait out another update first. - Require code_source.include_paths, refuse pathspec magic and files git tracks although .gitignore matches them, and list the files on --dry-run. Fix the shared non-git listing, which was relative to the working directory instead of root_path. - Print the idempotency key before the first request in every mode; reject an empty --idempotency-key; give JSON errors typed codes. - Validate endpoint names, workload types and sizes, checkpoint tasks, and count length limits as the server does. Co-authored-by: Isaac --- cmd/air/deploy.go | 112 ++++++++++++++++++----- cmd/air/deploy_test.go | 103 ++++++++++++++++++--- cmd/air/deploycode.go | 56 +++++++++--- cmd/air/deploycode_test.go | 19 ++-- cmd/air/deployconfig.go | 49 +++++++--- cmd/air/deployconfig_test.go | 9 +- cmd/air/deployendpoint.go | 39 ++++++-- cmd/air/deployendpoint_test.go | 23 ++++- cmd/air/deployregister.go | 101 ++++++++++++++++----- cmd/air/deployregister_test.go | 148 ++++++++++++++++++++++++------- cmd/air/snapshot_package.go | 4 +- cmd/air/snapshot_package_test.go | 18 ++++ 12 files changed, 550 insertions(+), 131 deletions(-) diff --git a/cmd/air/deploy.go b/cmd/air/deploy.go index 1b8f0932e92..8ac4dd903ba 100644 --- a/cmd/air/deploy.go +++ b/cmd/air/deploy.go @@ -6,7 +6,9 @@ import ( "fmt" "net/http" "path/filepath" + "regexp" "strings" + "sync" "time" "github.com/databricks/cli/cmd/root" @@ -32,6 +34,38 @@ var ( errDeployTimedOut = errors.New("air deploy timed out") ) +// JSON error codes of air deploy's own failures; API errors keep the server's. +const ( + codeInvalidArgs = "INVALID_ARGS" + codeInterrupted = "INTERRUPTED" + codeTimeout = "TIMEOUT" + codeStagingFailed = "STAGING_FAILED" + // codeEndpointNotReady: the endpoint serves the version but isn't ready, which + // running air deploy again doesn't change. + codeEndpointNotReady = "ENDPOINT_NOT_READY" + codeInternalError = "INTERNAL_ERROR" +) + +// apiErrorCodeRe matches an error code the server sent. The SDK fills ErrorCode +// with "" or "" when the response body has none. +var apiErrorCodeRe = sync.OnceValue(func() *regexp.Regexp { return regexp.MustCompile(`^[A-Z][A-Z_]*$`) }) + +// deployError carries the JSON error code and retryability of a failure whose +// message is err's. +type deployError struct { + code string + retryable bool + err error +} + +func (e *deployError) Error() string { return e.err.Error() } + +func (e *deployError) Unwrap() error { return e.err } + +func withCode(code string, retryable bool, err error) error { + return &deployError{code: code, retryable: retryable, err: err} +} + // deployResult is the JSON payload for `air deploy`. type deployResult struct { Status string `json:"status"` @@ -51,9 +85,12 @@ type deployResult struct { // deployOptions are the `air deploy` flags. type deployOptions struct { - file string - overrides []string - dryRun bool + file string + overrides []string + dryRun bool + // idempotencyKeySet records an explicit --idempotency-key, so an empty value + // (say, from an unset shell variable) fails instead of registering anew. + idempotencyKeySet bool idempotencyKey string retryFailedStaging bool timeout time.Duration @@ -139,6 +176,7 @@ its path to -h as a separate argument: } cmd.RunE = func(cmd *cobra.Command, args []string) error { + opts.idempotencyKeySet = cmd.Flags().Changed("idempotency-key") err := runDeploy(cmd, opts, pollInterval) if err == nil { return nil @@ -157,17 +195,17 @@ func runDeploy(cmd *cobra.Command, opts deployOptions, pollInterval time.Duratio ctx := cmd.Context() cfg, err := loadDeployConfig(ctx, opts.file, opts.overrides) if err != nil { - return err + return withCode(codeInvalidArgs, false, err) } if err := validateDeployFlags(cfg, opts); err != nil { - return err + return withCode(codeInvalidArgs, false, err) } var code *modelPackageCode if cfg.Model.Type == modelTypeCode { code, err = collectModelCode(ctx, cfg, opts.file) if err != nil { - return err + return withCode(codeInvalidArgs, false, err) } } @@ -186,6 +224,9 @@ func runDeploy(cmd *cobra.Command, opts deployOptions, pollInterval time.Duratio if cfg.registers() { // Fail on an endpoint air deploy won't change before copying any weights. if err := checkEndpoint(ctx, w, cfg.Endpoint, cfg.RegisteredModel); err != nil { + if stopped := deployStopped(ctx, err); stopped != nil { + return fmt.Errorf("%w while checking endpoint %s, before registering anything", stopped, cfg.Endpoint) + } return err } result.IdempotencyKey = effectiveIdempotencyToken(opts.idempotencyKey, cfg.IdempotencyToken) @@ -196,7 +237,7 @@ func runDeploy(cmd *cobra.Command, opts deployOptions, pollInterval time.Duratio result.ModelVersion, result.ModelID = resp.ModelVersion, resp.ModelID result.EnvironmentStagingFailed = resp.environmentState() == stagingStateFailed if !jsonOut { - cmdio.LogString(ctx, fmt.Sprintf("Version %s of %s is ready.", resp.ModelVersion, cfg.RegisteredModel)) + cmdio.LogString(ctx, fmt.Sprintf("Version %s of %s is ready.", stripControl(resp.ModelVersion), cfg.RegisteredModel)) if result.EnvironmentStagingFailed { cmdio.LogString(ctx, "Warning: the version's serving environment couldn't be prepared ahead of time, so the endpoint can be slower to start, or fail to start.") } @@ -206,21 +247,24 @@ func runDeploy(cmd *cobra.Command, opts deployOptions, pollInterval time.Duratio } entity := cfg.servedEntity(result.ModelVersion) - endpoint, err := serveVersion(ctx, w, cfg.Endpoint, entity, !jsonOut) + _, err = serveVersion(ctx, w, cfg.Endpoint, entity, !jsonOut) if err != nil { resume := deployResumeCommand(w.Config.Profile, opts, result.IdempotencyKey) if stopped := deployStopped(ctx, err); stopped != nil { return fmt.Errorf("%w while waiting for endpoint %s, which keeps updating. To continue, run:\n %s", stopped, cfg.Endpoint, resume) } + if errors.Is(err, errEndpointNotReady) { + return err + } return fmt.Errorf("serving version %s of %s on endpoint %s failed: %w\n\nTo try again without registering another version, run:\n %s", result.ModelVersion, cfg.RegisteredModel, cfg.Endpoint, err, resume) } result.Status = "READY" - result.EndpointURL = endpointInvocationsURL(w.Config.Host, endpoint.Name) + result.EndpointURL = endpointInvocationsURL(w.Config.Host, cfg.Endpoint) if jsonOut { return renderEnvelope(ctx, result) } out := cmd.OutOrStdout() - fmt.Fprintf(out, "Endpoint %s serves version %s of %s.\n", endpoint.Name, result.ModelVersion, cfg.RegisteredModel) + fmt.Fprintf(out, "Endpoint %s serves version %s of %s.\n", cfg.Endpoint, result.ModelVersion, cfg.RegisteredModel) fmt.Fprintln(out, "Query it at: "+result.EndpointURL) return nil } @@ -248,9 +292,9 @@ func renderDeployDryRun(ctx context.Context, cfg *deployConfig, code *modelPacka // registerDeployment registers the version and waits until it is ready, showing // progress on a spinner when show is true. Its errors say how to continue. func registerDeployment(ctx context.Context, w *databricks.WorkspaceClient, cfg *deployConfig, opts deployOptions, code *modelPackageCode, key string, pollInterval time.Duration, show bool) (*registerModelPackageResponse, error) { - if show { - cmdio.LogString(ctx, fmt.Sprintf("Registering a version of %s with idempotency key %s", cfg.RegisteredModel, key)) - } + // Printed to stderr in every output mode, before the first request, so a run + // that is killed can still be continued with this key. + cmdio.LogString(ctx, fmt.Sprintf("Registering a version of %s with idempotency key %s", cfg.RegisteredModel, key)) req := cfg.registrationRequest(key, code) req.RetryFailedStaging = opts.retryFailedStaging @@ -272,12 +316,22 @@ func registerDeployment(ctx context.Context, w *databricks.WorkspaceClient, cfg if cfg.Model.Path != nil { source = "the weights in " + *cfg.Model.Path } - return nil, fmt.Errorf("copying %s into version %s of %s failed. Check that you can read every file, then run:\n %s --retry-failed-staging", source, resp.ModelVersion, cfg.RegisteredModel, resume) + return nil, withCode(codeStagingFailed, true, fmt.Errorf("copying %s into version %s of %s failed. Check that you can read every file, then run:\n %s --retry-failed-staging", source, stripControl(resp.ModelVersion), cfg.RegisteredModel, resume)) } if stopped := deployStopped(ctx, err); stopped != nil { - return nil, fmt.Errorf("%w while registering a version of %s. Copying a checkpoint's weights continues, and the version is finalized when you continue. To continue, run:\n %s", stopped, cfg.RegisteredModel, resume) + continuing := "" + if resp != nil && resp.weightsState() != "" { + continuing = " Copying the weights continues, and the version is finalized when you continue." + } + return nil, fmt.Errorf("%w while registering a version of %s.%s To continue, run:\n %s", stopped, cfg.RegisteredModel, continuing, resume) + } + failed := fmt.Errorf("registering a version of %s with idempotency key %s failed: %w", cfg.RegisteredModel, key, err) + // A retryable failure is retried with the same key, so no other version is + // registered. + if _, retryable := deployErrorCode(err); retryable { + return nil, fmt.Errorf("%w\n\nTo try again, run:\n %s", failed, resume) } - return nil, fmt.Errorf("registering a version of %s with idempotency key %s failed: %w", cfg.RegisteredModel, key, err) + return nil, failed } // deployStopped returns why air deploy stopped waiting, or nil if it didn't. The @@ -297,17 +351,27 @@ func deployStopped(ctx context.Context, err error) error { } // deployErrorCode classifies a failed deployment for the JSON error envelope. +// Failures that aren't the user's input or the server's answer, such as a lost +// connection or a failed endpoint update, are worth retrying with the printed +// command. func deployErrorCode(err error) (code string, retryable bool) { switch { case errors.Is(err, errDeployInterrupted): - return "INTERRUPTED", true + return codeInterrupted, true case errors.Is(err, errDeployTimedOut): - return "TIMEOUT", true + return codeTimeout, true } - if apiErr, ok := errors.AsType[*apierr.APIError](err); ok && apiErr.ErrorCode != "" { - return apiErr.ErrorCode, apiErr.StatusCode == http.StatusTooManyRequests || apiErr.StatusCode >= http.StatusInternalServerError + if coded, ok := errors.AsType[*deployError](err); ok { + return coded.code, coded.retryable + } + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { + retryable = apiErr.StatusCode == http.StatusTooManyRequests || apiErr.StatusCode >= http.StatusInternalServerError + if !apiErrorCodeRe().MatchString(apiErr.ErrorCode) { + return codeInternalError, retryable + } + return apiErr.ErrorCode, retryable } - return "INVALID_ARGS", false + return codeInternalError, true } // validateDeployFlags rejects flags the config's model source can't use. @@ -315,7 +379,7 @@ func validateDeployFlags(cfg *deployConfig, opts deployOptions) error { if opts.timeout <= 0 { return fmt.Errorf("--timeout must be positive, got %s", opts.timeout) } - if opts.idempotencyKey != "" { + if opts.idempotencyKeySet { if !cfg.registers() { return errors.New("--idempotency-key applies only when air deploy registers a version; model.type is registered") } @@ -333,9 +397,9 @@ func validateDeployFlags(cfg *deployConfig, opts deployOptions) error { func registrationProgress(cfg *deployConfig, resp *registerModelPackageResponse) string { state := resp.weightsState() if state == "" { - return fmt.Sprintf("Registering version %s of %s…", resp.ModelVersion, cfg.RegisteredModel) + return fmt.Sprintf("Registering version %s of %s…", stripControl(resp.ModelVersion), cfg.RegisteredModel) } - return fmt.Sprintf("Copying the weights into version %s of %s (%s)…", resp.ModelVersion, cfg.RegisteredModel, state) + return fmt.Sprintf("Copying the weights into version %s of %s (%s)…", stripControl(resp.ModelVersion), cfg.RegisteredModel, stripControl(state)) } // serveVersion serves entity on the endpoint, showing the endpoint's state on a diff --git a/cmd/air/deploy_test.go b/cmd/air/deploy_test.go index ce232b46ff4..2ae6ac0917c 100644 --- a/cmd/air/deploy_test.go +++ b/cmd/air/deploy_test.go @@ -6,6 +6,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "net/http/httptest" "path/filepath" "strconv" @@ -17,6 +18,7 @@ import ( "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/retry" "github.com/databricks/cli/libs/shellquote" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/retries" @@ -28,7 +30,7 @@ import ( // withDeadline returns the test's context with a deadline, which air deploy // always sets from --timeout. func withDeadline(t *testing.T) context.Context { - ctx, cancel := context.WithTimeout(t.Context(), time.Minute) + ctx, cancel := context.WithTimeout(fastRetries(t.Context()), time.Minute) t.Cleanup(cancel) return ctx } @@ -52,7 +54,12 @@ func runDeployCmd(t *testing.T, out flags.Output, buf *bytes.Buffer, srvURL, con require.NoError(t, cmd.Flags().Set(flagValues[i], flagValues[i+1])) } - ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), out, nil, buf, buf, "", "")) + // JSON tests parse buf, so stderr's progress lines go elsewhere. + var stderr io.Writer = buf + if out == flags.OutputJSON { + stderr = io.Discard + } + ctx := cmdio.InContext(fastRetries(t.Context()), cmdio.NewIO(t.Context(), out, nil, buf, stderr, "", "")) if srvURL != "" { ctx = cmdctx.SetWorkspaceClient(ctx, newTestWorkspaceClient(t, srvURL)) } @@ -91,6 +98,9 @@ func TestDeployCheckpointJSON(t *testing.T) { served := endpoints.endpoint.Config.ServedEntities[0] assert.Equal(t, "3", served.EntityVersion) assert.Equal(t, map[string]string{"LOG_LEVEL": "info", "HF_TOKEN": "{{secrets/llm/hf_token}}"}, served.EnvironmentVars) + // scale_to_zero: false is sent, not left to the server's default. + require.Len(t, endpoints.rawWrites, 1) + assert.Contains(t, endpoints.rawWrites[0], `"scale_to_zero_enabled":false`) } func TestDeployCheckpointText(t *testing.T) { @@ -132,6 +142,20 @@ func TestDeployUsesTheConfigTokenThenANewKey(t *testing.T) { assert.NotContains(t, []string{"from-config", "from-flag"}, bodies[2].IdempotencyKey) } +func TestDeployPrintsTheKeyBeforeRegisteringInJSONMode(t *testing.T) { + srv, _ := deployServer(t, newFakeEndpoints(t, nil), conflictResponse) + + var stdout, stderr bytes.Buffer + cmd := withOutput(newDeployCommandWithPollInterval(time.Millisecond), flags.OutputJSON) + require.NoError(t, cmd.Flags().Set("file", writeConfigFile(t, "deploy.yaml", checkpointDeployConfig))) + require.NoError(t, cmd.Flags().Set("idempotency-key", "key-1")) + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputJSON, nil, &stdout, &stderr, "", "")) + cmd.SetContext(cmdctx.SetWorkspaceClient(ctx, newTestWorkspaceClient(t, srv.URL))) + require.ErrorIs(t, cmd.RunE(cmd, nil), root.ErrAlreadyPrinted) + assert.Equal(t, "Registering a version of main.models.qwen with idempotency key key-1\n", stderr.String()) + assert.Contains(t, stdout.String(), `"code": "RESOURCE_CONFLICT"`) +} + func TestDeployRetriesFailedStagingOnce(t *testing.T) { srv, fake := deployServer(t, newFakeEndpoints(t, nil), stagingInProgress, stagingReady) @@ -154,7 +178,7 @@ func TestDeployFailedStagingSaysHowToRetry(t *testing.T) { } func TestDeployReportsRegistrationErrors(t *testing.T) { - srv, _ := deployServer(t, newFakeEndpoints(t, nil), "") + srv, _ := deployServer(t, newFakeEndpoints(t, nil), conflictResponse) var buf bytes.Buffer err := runDeployCmd(t, flags.OutputText, &buf, srv.URL, checkpointDeployConfig, "idempotency-key", "key-1") @@ -162,7 +186,7 @@ func TestDeployReportsRegistrationErrors(t *testing.T) { } func TestDeployReportsErrorsAsJSON(t *testing.T) { - srv, _ := deployServer(t, newFakeEndpoints(t, nil), "") + srv, _ := deployServer(t, newFakeEndpoints(t, nil), conflictResponse) var buf bytes.Buffer err := runDeployCmd(t, flags.OutputJSON, &buf, srv.URL, checkpointDeployConfig) @@ -192,6 +216,56 @@ func TestDeployChecksTheEndpointBeforeRegistering(t *testing.T) { assert.Empty(t, fake.requests()) } +func TestDeployReportsAnEndpointThatIsNotReady(t *testing.T) { + endpoints := newFakeEndpoints(t, nil) + endpoints.finalState = serving.EndpointState{Ready: serving.EndpointStateReadyNotReady, ConfigUpdate: serving.EndpointStateConfigUpdateNotUpdating} + srv, _ := deployServer(t, endpoints, stagingReady) + + var buf bytes.Buffer + err := runDeployCmd(t, flags.OutputJSON, &buf, srv.URL, checkpointDeployConfig, "idempotency-key", "key-1") + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + var got errorEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + // Running air deploy again doesn't make the endpoint ready, so nothing says to. + assert.Equal(t, "ENDPOINT_NOT_READY", got.Error.Code) + assert.False(t, got.Error.Retryable) + assert.NotContains(t, got.Error.Message, "run:") +} + +func TestDeployTimesOutWhileCopyingTheWeights(t *testing.T) { + srv, fake := deployServer(t, newFakeEndpoints(t, nil), stagingInProgress) + fake.statuses = []string{statusInProgress} + + // Polls are spaced as the SDK's rate limiter allows, so the deadline, not the + // limiter, stops the wait. + var buf bytes.Buffer + cmd := withOutput(newDeployCommandWithPollInterval(100*time.Millisecond), flags.OutputJSON) + require.NoError(t, cmd.Flags().Set("file", writeConfigFile(t, "deploy.yaml", checkpointDeployConfig))) + require.NoError(t, cmd.Flags().Set("idempotency-key", "key-1")) + require.NoError(t, cmd.Flags().Set("timeout", "500ms")) + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputJSON, nil, &buf, io.Discard, "", "")) + cmd.SetContext(cmdctx.SetWorkspaceClient(ctx, newTestWorkspaceClient(t, srv.URL))) + require.ErrorIs(t, cmd.RunE(cmd, nil), root.ErrAlreadyPrinted) + var got errorEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, "TIMEOUT", got.Error.Code) + assert.True(t, got.Error.Retryable) + assert.Contains(t, got.Error.Message, "air deploy timed out while registering a version of main.models.qwen. Copying the weights continues") + assert.Contains(t, got.Error.Message, " --idempotency-key key-1") + assert.Len(t, fake.requests(), 1) +} + +func TestDeployRetryableRegistrationFailureSaysHowToRetry(t *testing.T) { + srv, fake := deployServer(t, newFakeEndpoints(t, nil), serverErrorResponse) + + var buf bytes.Buffer + err := runDeployCmd(t, flags.OutputText, &buf, srv.URL, checkpointDeployConfig, "idempotency-key", "key-1") + require.Error(t, err) + assert.Contains(t, err.Error(), "failed: Something went wrong.\n\nTo try again, run:\n databricks air deploy -f ") + assert.Contains(t, err.Error(), " --idempotency-key key-1") + assert.Len(t, fake.requests(), retry.MaxRetries+1) +} + func TestDeployEndpointFailureSaysHowToRetry(t *testing.T) { endpoints := newFakeEndpoints(t, nil) endpoints.finalState = serving.EndpointState{Ready: serving.EndpointStateReadyNotReady, ConfigUpdate: serving.EndpointStateConfigUpdateUpdateFailed} @@ -291,6 +365,7 @@ func TestDeployFlagValidation(t *testing.T) { {"bad key", checkpointDeployConfig, []string{"idempotency-key", "a/b"}, "--idempotency-key must be 1 to 128 letters"}, {"retry for a registered version", registered, []string{"retry-failed-staging", "true"}, "--retry-failed-staging applies only to a checkpoint"}, {"zero timeout", checkpointDeployConfig, []string{"timeout", "0s"}, "--timeout must be positive"}, + {"empty key", checkpointDeployConfig, []string{"idempotency-key", ""}, `--idempotency-key must be 1 to 128 letters, digits, '.', '_' or '-', got ""`}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -331,19 +406,25 @@ func TestDeployStopped(t *testing.T) { func TestDeployErrorCode(t *testing.T) { for _, tt := range []struct { + name string err error code string retryable bool }{ - {fmt.Errorf("x: %w", errDeployInterrupted), "INTERRUPTED", true}, - {fmt.Errorf("x: %w", errDeployTimedOut), "TIMEOUT", true}, - {fmt.Errorf("x: %w", &apierr.APIError{ErrorCode: "TEMPORARILY_UNAVAILABLE", StatusCode: 503}), "TEMPORARILY_UNAVAILABLE", true}, - {fmt.Errorf("x: %w", &apierr.APIError{ErrorCode: "INVALID_PARAMETER_VALUE", StatusCode: 400}), "INVALID_PARAMETER_VALUE", false}, - {errors.New("model: section is required"), "INVALID_ARGS", false}, + {"interrupted", fmt.Errorf("x: %w", errDeployInterrupted), "INTERRUPTED", true}, + {"timed out", fmt.Errorf("x: %w", errDeployTimedOut), "TIMEOUT", true}, + {"invalid input", withCode(codeInvalidArgs, false, errors.New("model: section is required")), "INVALID_ARGS", false}, + {"staging failed", withCode(codeStagingFailed, true, errors.New("copying failed")), "STAGING_FAILED", true}, + {"server unavailable", fmt.Errorf("x: %w", &apierr.APIError{ErrorCode: "TEMPORARILY_UNAVAILABLE", StatusCode: 503}), "TEMPORARILY_UNAVAILABLE", true}, + {"bad request", fmt.Errorf("x: %w", &apierr.APIError{ErrorCode: "INVALID_PARAMETER_VALUE", StatusCode: 400}), "INVALID_PARAMETER_VALUE", false}, + {"5xx without a body", &apierr.APIError{StatusCode: 502}, "INTERNAL_ERROR", true}, + {"error body without a code", &apierr.APIError{ErrorCode: "", StatusCode: 400}, "INTERNAL_ERROR", false}, + {"lost connection", errors.New("connection reset by peer"), "INTERNAL_ERROR", true}, + {"endpoint not ready", withCode(codeEndpointNotReady, false, errEndpointNotReady), "ENDPOINT_NOT_READY", false}, } { code, retryable := deployErrorCode(tt.err) - assert.Equal(t, tt.code, code, tt.err.Error()) - assert.Equal(t, tt.retryable, retryable, tt.err.Error()) + assert.Equal(t, tt.code, code, tt.name) + assert.Equal(t, tt.retryable, retryable, tt.name) } } diff --git a/cmd/air/deploycode.go b/cmd/air/deploycode.go index eb2c1ae1ee0..433941d77af 100644 --- a/cmd/air/deploycode.go +++ b/cmd/air/deploycode.go @@ -1,9 +1,12 @@ package aircmd import ( + "bytes" "context" + "errors" "fmt" "os" + "os/exec" "path/filepath" "regexp" "slices" @@ -30,28 +33,31 @@ const pythonCacheDir = "__pycache__" var codePathSegmentRe = sync.OnceValue(func() *regexp.Regexp { return regexp.MustCompile(`^[A-Za-z0-9_][A-Za-z0-9_.-]*$`) }) // collectModelCode reads a model from code from code_source: model.entry and the -// other files a code_source snapshot would include, as UTF-8 text. Files git -// ignores aren't sent, and nor are hidden files and directories or __pycache__ -// directories, whose names the server can't place or which hold no source. -// configPath locates a relative code_source.root_path. -// -// The files are stored in the model version, where anyone who can read the model -// can read them. Outside a git repository no .gitignore filters them, so the user -// must list them in code_source.include_paths. +// other files under code_source.include_paths, as UTF-8 text, selected as a +// code_source snapshot selects them. Files git ignores aren't sent, and nor are +// hidden files and directories or __pycache__ directories, whose names the +// server can't place or which hold no source. configPath locates a relative +// code_source.root_path. func collectModelCode(ctx context.Context, cfg *deployConfig, configPath string) (*modelPackageCode, error) { root, err := resolveRootPath(ctx, cfg.CodeSource.RootPath, filepath.Dir(configPath)) if err != nil { return nil, err } isGitRepo := newGitRepo(root).isRepository(ctx) - if !isGitRepo && len(cfg.CodeSource.IncludePaths) == 0 { - return nil, fmt.Errorf("code_source.include_paths is required because %s isn't in a git repository: air deploy stores the files in the model version, so list the ones to send, e.g. include_paths: [%s]", filepath.ToSlash(root), *cfg.Model.Entry) - } files, err := snapshotFiles(ctx, root, cfg.CodeSource.IncludePaths, isGitRepo) if err != nil { return nil, err } files = slices.DeleteFunc(files, skippedCodeFile) + if isGitRepo { + ignored, err := ignoredTrackedFiles(ctx, root, files) + if err != nil { + return nil, err + } + if len(ignored) > 0 { + return nil, fmt.Errorf("code_source file %s matches a .gitignore rule but is tracked by git; air deploy doesn't send ignored files, so leave it out of code_source.include_paths or stop tracking it", ignored[0]) + } + } code := &modelPackageCode{EntryFile: *cfg.Model.Entry, Environment: cfg.servingEnvironment()} if cfg.Task != nil { @@ -86,6 +92,34 @@ func collectModelCode(ctx context.Context, cfg *deployConfig, configPath string) return code, nil } +// ignoredTrackedFiles returns the files that a .gitignore rule matches although +// git tracks them: ls-files lists tracked files whatever the rules say, and the +// rules are the user's statement of what not to share. +func ignoredTrackedFiles(ctx context.Context, root string, files []snapshotFile) ([]string, error) { + var paths bytes.Buffer + for _, f := range files { + paths.WriteString(filepath.ToSlash(f.rel)) + paths.WriteByte(0) + } + cmd := exec.CommandContext(ctx, "git", "-C", root, "check-ignore", "--no-index", "-z", "--stdin") + cmd.Stdin = &paths + out, err := cmd.Output() + // check-ignore exits 1 when no path is ignored. + if exitErr, ok := errors.AsType[*exec.ExitError](err); ok && exitErr.ExitCode() == 1 { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("failed to evaluate git ignore rules: %w", err) + } + var ignored []string + for raw := range bytes.SplitSeq(out, []byte{0}) { + if len(raw) > 0 { + ignored = append(ignored, string(raw)) + } + } + return ignored, nil +} + // checkCodeFiles checks the files to send against the server's limits, from // their listing alone. func checkCodeFiles(files []snapshotFile, dependencies []string) error { diff --git a/cmd/air/deploycode_test.go b/cmd/air/deploycode_test.go index c5d50fb88ce..8fa2fa7e0ab 100644 --- a/cmd/air/deploycode_test.go +++ b/cmd/air/deploycode_test.go @@ -126,16 +126,23 @@ func TestCollectModelCodeLeavesOutIgnoredFiles(t *testing.T) { assert.Equal(t, []string{"agent.py", "tools/a.py"}, codeFilePaths(code)) } -func TestCollectModelCodeOutsideGitNeedsIncludePaths(t *testing.T) { - cfg, configPath := codeTreeConfig(t, t.TempDir(), map[string]string{"agent.py": "x\n", "credentials.json": "{}\n"}) +func TestCollectModelCodeRefusesTrackedIgnoredFiles(t *testing.T) { + cfg, configPath := writeCodeTree(t, map[string]string{"agent.py": "x\n", "creds.json": "{}\n"}) + repo := cfg.CodeSource.RootPath + commitAll(t, repo, "init") + writeRepoFile(t, repo, ".gitignore", "creds.json\n") + cfg.CodeSource.IncludePaths = []string{"agent.py", "creds.json"} + _, err := collectModelCode(t.Context(), cfg, configPath) - assert.ErrorContains(t, err, "code_source.include_paths is required because") - assert.ErrorContains(t, err, "isn't in a git repository") + assert.ErrorContains(t, err, "code_source file creds.json matches a .gitignore rule but is tracked by git") +} - cfg.CodeSource.IncludePaths = []string{"agent.py"} +func TestCollectModelCodeOutsideGit(t *testing.T) { + cfg, configPath := codeTreeConfig(t, t.TempDir(), map[string]string{"agent.py": "x\n", "tools/a.py": "a\n", "credentials.json": "{}\n"}) + cfg.CodeSource.IncludePaths = []string{"agent.py", "tools"} code, err := collectModelCode(t.Context(), cfg, configPath) require.NoError(t, err) - assert.Equal(t, []string{"agent.py"}, codeFilePaths(code)) + assert.Equal(t, []string{"agent.py", "tools/a.py"}, codeFilePaths(code)) } func TestCollectModelCodeChecksSizesBeforeReading(t *testing.T) { diff --git a/cmd/air/deployconfig.go b/cmd/air/deployconfig.go index 4b035cd616e..29619ef9db7 100644 --- a/cmd/air/deployconfig.go +++ b/cmd/air/deployconfig.go @@ -12,7 +12,7 @@ import ( "strconv" "strings" "sync" - "unicode/utf8" + "unicode/utf16" "github.com/databricks/databricks-sdk-go/service/serving" "go.yaml.in/yaml/v3" @@ -66,6 +66,13 @@ var deployIdempotencyKeyRe = sync.OnceValue(func() *regexp.Regexp { return regex // modelVersionRe matches a Unity Catalog model version number. var modelVersionRe = sync.OnceValue(func() *regexp.Regexp { return regexp.MustCompile(`^[1-9][0-9]*$`) }) +// endpointNameRe matches a serving endpoint name: letters, digits, '-' and '_'. +// The name is part of the API's URL paths, so nothing else may reach a request. +var endpointNameRe = sync.OnceValue(func() *regexp.Regexp { return regexp.MustCompile(`^[A-Za-z0-9_-]+$`) }) + +// workloadSizes are the concurrency tiers a served entity scales to. +var workloadSizes = []string{"Small", "Medium", "Large"} + // entryFileRe matches the name of a model from code's entry file, which the server // places at the model's root. var entryFileRe = sync.OnceValue(func() *regexp.Regexp { return regexp.MustCompile(`^[A-Za-z0-9_][A-Za-z0-9_.-]*\.py$`) }) @@ -78,9 +85,9 @@ type deployConfig struct { Task *string `yaml:"task" help:"The model's task as MLflow names it: llm/v1/chat, llm/v1/completions or llm/v1/embeddings for a checkpoint, or e.g. agent/v1/responses for a model from code. A model from code without a task has no signature." required:"when model.type is checkpoint" group:"Model"` Command *string `yaml:"command" help:"Shell command that starts the model server of a checkpoint, at most 32,768 characters. It runs from the directory where model.path appears under its own name." required:"when model.type is checkpoint" group:"Model"` Environment *deployEnvironmentConfig `yaml:"environment" help:"Serverless environment version to serve on and the packages to install, for a checkpoint or a model from code." group:"Model"` - CodeSource *deployCodeSourceConfig `yaml:"code_source" help:"Local code of a model from code, sent with the registration and stored in the model version, where anyone who can read the model can read it. Hidden files and directories (names starting with '.') and __pycache__ directories aren't sent; run with --dry-run to list the files that are." required:"when model.type is code" group:"Model"` + CodeSource *deployCodeSourceConfig `yaml:"code_source" help:"Local code of a model from code, sent with the registration and stored in the model version, where anyone who can read the model can read it. Only the files under include_paths are sent, leaving out files git ignores, hidden files and directories (names starting with '.') and __pycache__ directories; run with --dry-run to list them." required:"when model.type is code" group:"Model"` - Endpoint string `yaml:"endpoint" help:"Serving endpoint to serve the model on. It is created if it doesn't exist. An endpoint that serves one version of registered_model is updated to serve this version instead; air deploy doesn't change an endpoint that serves anything else." required:"yes" group:"Serving"` + Endpoint string `yaml:"endpoint" help:"Serving endpoint to serve the model on. It is created if it doesn't exist. An endpoint that serves one version of registered_model is updated to serve this version instead, with the compute, environment variables and secrets in this file replacing the ones it has; air deploy doesn't change an endpoint that serves anything else." required:"yes" group:"Serving"` Compute *deployComputeConfig `yaml:"compute" help:"Serving compute and how the endpoint scales." required:"yes" group:"Serving"` EnvVariables map[string]string `yaml:"env_variables" help:"Plain environment variables of the served model, as NAME: value. A name here cannot also appear in secrets." group:"Serving"` Secrets map[string]string `yaml:"secrets" help:"Environment variables of the served model sourced from secrets, as NAME: scope/key." group:"Serving"` @@ -106,7 +113,7 @@ type deployEnvironmentConfig struct { // deployCodeSourceConfig is the `code_source` block of a model from code. type deployCodeSourceConfig struct { RootPath string `yaml:"root_path" help:"Directory holding the model's code. A relative path resolves against the YAML file's directory." required:"yes"` - IncludePaths []string `yaml:"include_paths" help:"Send only these paths, relative to root_path and without \"..\". Omit to send every file that git doesn't ignore; required when root_path isn't in a git repository."` + IncludePaths []string `yaml:"include_paths" help:"Files and directories to send, relative to root_path and without \"..\", e.g. [agent.py, tools]." required:"yes"` } // deployComputeConfig is the `compute` block: the served entity's workload. @@ -146,8 +153,8 @@ func (c *deployConfig) validate() error { if err := c.validateSource(); err != nil { return err } - if strings.TrimSpace(c.Endpoint) == "" { - return errors.New("endpoint is required") + if !endpointNameRe().MatchString(c.Endpoint) { + return fmt.Errorf("endpoint must be a serving endpoint name of letters, digits, '-' and '_', got %q", c.Endpoint) } if c.Compute == nil { return errors.New("compute: section is required") @@ -155,8 +162,8 @@ func (c *deployConfig) validate() error { if workloadTypes := servingWorkloadTypes(); !slices.Contains(workloadTypes, c.Compute.WorkloadType) { return fmt.Errorf("compute.workload_type must be one of %s, got %q", strings.Join(workloadTypes, ", "), c.Compute.WorkloadType) } - if strings.TrimSpace(c.Compute.WorkloadSize) == "" { - return errors.New("compute.workload_size is required") + if !slices.Contains(workloadSizes, c.Compute.WorkloadSize) { + return fmt.Errorf("compute.workload_size must be one of %s, got %q", strings.Join(workloadSizes, ", "), c.Compute.WorkloadSize) } if c.Model.Type == modelTypeCheckpoint && !c.servesOnGPU() { return fmt.Errorf("a checkpoint is served on GPU; set compute.workload_type to a GPU workload type such as GPU_MEDIUM, got %q", c.Compute.WorkloadType) @@ -207,7 +214,7 @@ func (c *deployConfig) validateSource() error { if m.Path != nil { v := *m.Path - if !strings.HasPrefix(v, "/Volumes/") || utf8.RuneCountInString(v) > maxVolumePathChars { + if !strings.HasPrefix(v, "/Volumes/") || serverLength(v) > maxVolumePathChars { return fmt.Errorf("model.path must be a Unity Catalog volume directory as /Volumes////, of at most %d characters, got %q", maxVolumePathChars, v) } } @@ -227,7 +234,7 @@ func (c *deployConfig) validateSource() error { if err := validateCommand(*c.Command); err != nil { return err } - if n := utf8.RuneCountInString(*c.Command); n > maxServingCommandChars { + if n := serverLength(*c.Command); n > maxServingCommandChars { return fmt.Errorf("command must be at most %d characters, got %d", maxServingCommandChars, n) } } @@ -237,6 +244,18 @@ func (c *deployConfig) validateSource() error { } } if c.CodeSource != nil { + // The files are stored in the model version for anyone who can read the + // model, so the user names them rather than sending a whole directory. + if c.CodeSource.IncludePaths == nil { + return errors.New("code_source.include_paths is required: list the files and directories to send, e.g. include_paths: [agent.py, tools]") + } + // git reads a leading ':' as pathspec magic, such as ":!x" for everything + // but x, which would send files the list doesn't name. + for _, p := range c.CodeSource.IncludePaths { + if strings.HasPrefix(strings.TrimSpace(p), ":") { + return fmt.Errorf("code_source.include_paths must name files and directories, got %q", p) + } + } snapshot := snapshotSourceConfig{RootPath: c.CodeSource.RootPath, IncludePaths: c.CodeSource.IncludePaths} if err := snapshot.validate("code_source"); err != nil { return err @@ -270,6 +289,16 @@ func validateDeployIdempotencyKey(key, source string) error { return nil } +// serverLength is the length of s as the server's limits count it, in UTF-16 +// code units. +func serverLength(s string) int { + n := 0 + for _, r := range s { + n += utf16.RuneLen(r) + } + return n +} + // servingWorkloadTypes lists the served entity workload types, sorted. func servingWorkloadTypes() []string { var workloadType serving.ServingModelWorkloadType diff --git a/cmd/air/deployconfig_test.go b/cmd/air/deployconfig_test.go index 80f64e76012..6d5efc47302 100644 --- a/cmd/air/deployconfig_test.go +++ b/cmd/air/deployconfig_test.go @@ -33,6 +33,7 @@ model: task: agent/v1/responses code_source: root_path: . + include_paths: [agent.py, tools] environment: version: 5 dependencies: [mlflow==3.5.0] @@ -155,9 +156,11 @@ func TestDeployConfigValidation(t *testing.T) { {"agent task on a checkpoint", replaceFirst("task: llm/v1/chat", "task: agent/v1/responses"), `task of a checkpoint must be one of llm/v1/chat, llm/v1/completions, llm/v1/embeddings, got "agent/v1/responses"`}, {"zero environment version", replaceFirst(" dependencies: [einops==0.8.0]", " version: 0"), "environment.version must be a serverless environment version"}, {"empty dependency", replaceFirst("[einops==0.8.0]", `[""]`), "environment.dependencies[0] cannot be empty"}, - {"no endpoint", dropLine("endpoint:"), "endpoint is required"}, + {"no endpoint", dropLine("endpoint:"), `endpoint must be a serving endpoint name of letters, digits, '-' and '_', got ""`}, + {"endpoint with a path", replaceFirst("endpoint: qwen-chat", "endpoint: x/../../jobs"), "endpoint must be a serving endpoint name"}, {"no compute", dropBlock("compute:"), "compute: section is required"}, - {"no workload size", dropLine(" workload_size:"), "compute.workload_size is required"}, + {"no workload size", dropLine(" workload_size:"), `compute.workload_size must be one of Small, Medium, Large, got ""`}, + {"lowercase workload size", replaceFirst("workload_size: Small", "workload_size: medium"), "compute.workload_size must be one of Small, Medium, Large"}, {"secret without key", replaceFirst("llm/hf_token", "llm"), "expected format 'scope/key'"}, {"env var and secret", replaceFirst("LOG_LEVEL: info", "HF_TOKEN: x"), `"HF_TOKEN" is set in both env_variables and secrets`}, {"bad idempotency token", appendText("idempotency_token: a b\n"), "idempotency_token must be 1 to 128 letters"}, @@ -185,6 +188,8 @@ func TestDeployConfigSourceFields(t *testing.T) { {"code with command", "model: {type: code, entry: agent.py}\ncode_source: {root_path: .}\ncommand: x\n", "command can't be set when model.type is code"}, {"code entry in a directory", "model: {type: code, entry: src/agent.py}\ncode_source: {root_path: .}\n", "model.entry must name a .py file directly under code_source.root_path"}, {"code entry not python", "model: {type: code, entry: agent.ipynb}\ncode_source: {root_path: .}\n", "model.entry must name a .py file"}, + {"code without include_paths", "model: {type: code, entry: agent.py}\ncode_source: {root_path: .}\n", "code_source.include_paths is required: list the files and directories to send"}, + {"code include pathspec magic", "model: {type: code, entry: agent.py}\ncode_source: {root_path: ., include_paths: [\":!creds.json\"]}\n", `code_source.include_paths must name files and directories, got ":!creds.json"`}, {"code include traversal", "model: {type: code, entry: agent.py}\ncode_source: {root_path: ., include_paths: [../x]}\n", "code_source.include_paths cannot contain '..' traversal"}, {"saved model with task", "model: {type: saved_model_dir, path: /Volumes/a/b/c/m}\ntask: llm/v1/chat\n", "task can't be set when model.type is saved_model_dir"}, {"saved model with environment", "model: {type: saved_model_dir, path: /Volumes/a/b/c/m}\nenvironment: {version: 6}\n", "environment can't be set when model.type is saved_model_dir"}, diff --git a/cmd/air/deployendpoint.go b/cmd/air/deployendpoint.go index 7a5d4d5f90a..e240dbc1f0b 100644 --- a/cmd/air/deployendpoint.go +++ b/cmd/air/deployendpoint.go @@ -8,6 +8,7 @@ import ( "strings" "time" + "github.com/databricks/cli/libs/retry" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/service/serving" @@ -30,13 +31,25 @@ func (c *deployConfig) servedEntity(version string) serving.ServedEntityInput { // getEndpoint returns the endpoint named name, or nil if it doesn't exist. func getEndpoint(ctx context.Context, w *databricks.WorkspaceClient, name string) (*serving.ServingEndpointDetailed, error) { - endpoint, err := w.ServingEndpoints.GetByName(ctx, name) + endpoint, err := retry.With(ctx, retry.IsTransientRead, func() (*serving.ServingEndpointDetailed, error) { + return w.ServingEndpoints.GetByName(ctx, name) + }) if errors.Is(err, apierr.ErrResourceDoesNotExist) { return nil, nil } return endpoint, err } +// waitNotUpdating waits until the endpoint has no update in progress. The SDK's +// wait stops at the first failed poll, so a transient failure starts it again: +// an endpoint can take many minutes to start. Creates and config updates aren't +// repeated this way, since a repeat after an applied write would fail. +func waitNotUpdating(ctx context.Context, w *databricks.WorkspaceClient, name string, progress func(*serving.ServingEndpointDetailed)) (*serving.ServingEndpointDetailed, error) { + return retry.With(ctx, retry.IsTransientRead, func() (*serving.ServingEndpointDetailed, error) { + return w.ServingEndpoints.WaitGetServingEndpointNotUpdating(ctx, name, untilDeadline(ctx), progress) + }) +} + // checkEndpoint fails when the endpoint named name, if it exists, serves // something air deploy won't replace with a version of model. air deploy checks // it before registering, so a wrong endpoint name fails before any weights are @@ -51,11 +64,13 @@ func checkEndpoint(ctx context.Context, w *databricks.WorkspaceClient, name, mod // checkReplaceable fails unless the endpoint serves at most one entity, a version // of model, in its current and pending configs. Replacing anything else would -// remove another model or a traffic split the user didn't ask to change. +// remove another model or a traffic split the user didn't ask to change. The +// endpoint lists every served model under served_entities too (served_models +// mirrors it), so served_entities covers both. func checkReplaceable(endpoint *serving.ServingEndpointDetailed, model string) error { for _, served := range [][]serving.ServedEntityOutput{currentEntities(endpoint), pendingEntities(endpoint)} { if len(served) > 1 || (len(served) == 1 && !strings.EqualFold(served[0].EntityName, model)) { - return fmt.Errorf("endpoint %s serves %s, and air deploy replaces only a single version of %s; set endpoint to another name, or change this endpoint with databricks serving-endpoints update-config", endpoint.Name, describeEntities(served), model) + return withCode(codeInvalidArgs, false, fmt.Errorf("endpoint %s serves %s, and air deploy replaces only a single version of %s; set endpoint to another name, or change this endpoint with databricks serving-endpoints update-config", stripControl(endpoint.Name), describeEntities(served), model)) } } return nil @@ -85,7 +100,7 @@ func serveOnEndpoint(ctx context.Context, w *databricks.WorkspaceClient, name st return waitServes(ctx, w, name, entity, progress) } // The endpoint takes no other update until this one finishes. - endpoint, err = w.ServingEndpoints.WaitGetServingEndpointNotUpdating(ctx, name, untilDeadline(ctx), progress) + endpoint, err = waitNotUpdating(ctx, w, name, progress) if err != nil { return nil, err } @@ -109,20 +124,25 @@ func serveOnEndpoint(ctx context.Context, w *databricks.WorkspaceClient, name st // waitServes waits until the endpoint's update finishes, then checks that it // serves entity. func waitServes(ctx context.Context, w *databricks.WorkspaceClient, name string, entity serving.ServedEntityInput, progress func(*serving.ServingEndpointDetailed)) (*serving.ServingEndpointDetailed, error) { - endpoint, err := w.ServingEndpoints.WaitGetServingEndpointNotUpdating(ctx, name, untilDeadline(ctx), progress) + endpoint, err := waitNotUpdating(ctx, w, name, progress) if err != nil { return nil, err } return checkServes(endpoint, entity) } +// errEndpointNotReady reports an endpoint that serves the version but isn't ready, +// which running air deploy again doesn't change. +var errEndpointNotReady = errors.New("endpoint isn't ready") + // checkServes returns endpoint if it serves only entity and is ready to serve. func checkServes(endpoint *serving.ServingEndpointDetailed, entity serving.ServedEntityInput) (*serving.ServingEndpointDetailed, error) { + name := stripControl(endpoint.Name) if !servesOnly(currentEntities(endpoint), entity) { - return nil, fmt.Errorf("endpoint %s was changed to serve %s while air deploy waited for it", endpoint.Name, describeEntities(currentEntities(endpoint))) + return nil, fmt.Errorf("endpoint %s was changed to serve %s while air deploy waited for it", name, describeEntities(currentEntities(endpoint))) } if endpoint.State == nil || endpoint.State.Ready != serving.EndpointStateReadyReady { - return nil, fmt.Errorf("endpoint %s serves version %s of %s but isn't ready; check its events in the Serving UI", endpoint.Name, entity.EntityVersion, entity.EntityName) + return nil, withCode(codeEndpointNotReady, false, fmt.Errorf("%w: endpoint %s serves version %s of %s but isn't ready; check its events and logs in the Serving UI", errEndpointNotReady, name, entity.EntityVersion, entity.EntityName)) } return endpoint, nil } @@ -152,11 +172,12 @@ func describeEntities(served []serving.ServedEntityOutput) string { } names := make([]string, 0, len(served)) for _, s := range served { + name := stripControl(s.EntityName) if s.EntityVersion == "" { - names = append(names, s.EntityName) + names = append(names, name) continue } - names = append(names, fmt.Sprintf("version %s of %s", s.EntityVersion, s.EntityName)) + names = append(names, fmt.Sprintf("version %s of %s", stripControl(s.EntityVersion), name)) } return strings.Join(names, " and ") } diff --git a/cmd/air/deployendpoint_test.go b/cmd/air/deployendpoint_test.go index 3733a6a4ae9..e267a1bdaf4 100644 --- a/cmd/air/deployendpoint_test.go +++ b/cmd/air/deployendpoint_test.go @@ -27,6 +27,8 @@ type fakeEndpoints struct { writes []string // rawWrites holds the body of every create and config update, as sent. rawWrites []string + // failGets is how many GETs answer 500 before the endpoint is served again. + failGets int } func newFakeEndpoints(t *testing.T, existing *serving.ServingEndpointDetailed) *fakeEndpoints { @@ -64,6 +66,12 @@ func (f *fakeEndpoints) ServeHTTP(w http.ResponseWriter, r *http.Request) { f.endpoint.PendingConfig = nil f.endpoint.State = &f.finalState case r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, servingEndpointsPath+"/"): + if f.failGets > 0 && f.endpoint != nil { + f.failGets-- + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL_ERROR", "message": "Something went wrong."}`)) + return + } if f.endpoint == nil { w.WriteHeader(http.StatusNotFound) _, _ = w.Write([]byte(`{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "Endpoint not found"}`)) @@ -141,9 +149,6 @@ func TestServeOnEndpointCreatesAMissingEndpoint(t *testing.T) { require.NoError(t, err) assert.Equal(t, []string{"create"}, fake.writes) assert.Equal(t, "3", endpoint.Config.ServedEntities[0].EntityVersion) - // scale_to_zero: false is sent, not left to the server's default. - require.Len(t, fake.rawWrites, 1) - assert.Contains(t, fake.rawWrites[0], `"scale_to_zero_enabled":false`) } func TestServeOnEndpointUpdatesAnotherVersion(t *testing.T) { @@ -247,6 +252,18 @@ func TestCheckEndpointAllowsAMissingEndpoint(t *testing.T) { assert.NoError(t, checkEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", "main.models.qwen")) } +func TestServeOnEndpointWaitsThroughATransientFailure(t *testing.T) { + fake := newFakeEndpoints(t, nil) + // The create succeeds; the first poll of the wait then fails. + fake.failGets = 1 + srv := fake.start() + + _, err := serveOnEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", testEntity("3"), nil) + require.NoError(t, err) + assert.Equal(t, []string{"create"}, fake.writes) + assert.Zero(t, fake.failGets) +} + func TestServeOnEndpointReportsAFailedUpdate(t *testing.T) { fake := newFakeEndpoints(t, nil) fake.finalState = serving.EndpointState{Ready: serving.EndpointStateReadyNotReady, ConfigUpdate: serving.EndpointStateConfigUpdateUpdateFailed} diff --git a/cmd/air/deployregister.go b/cmd/air/deployregister.go index f09867aa04f..a5ae7277553 100644 --- a/cmd/air/deployregister.go +++ b/cmd/air/deployregister.go @@ -8,6 +8,7 @@ import ( "time" "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/retry" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/client" ) @@ -15,16 +16,24 @@ import ( // registerModelPackagePath registers a model whose files Model Serving writes. It // runs as the caller and is idempotent per key: a repeat request with the same key // and inputs continues the registration, and finalizes a checkpoint's version once -// its weights are staged. +// its weights are staged. Staging never finalizes the version by itself. const registerModelPackagePath = "/api/2.0/serving-endpoints:registerModelPackage" -// registrationPollInterval is how often air deploy repeats the registration while -// a checkpoint's weights are staged. +// stagingStatusPath reports the recorded staging progress of a model version, a +// read. +const stagingStatusPath = "/api/2.0/serving-endpoints:checkModelDeploymentStagingStatus" + +// registrationPollInterval is how often air deploy checks on a checkpoint's +// weights while they are staged. const registrationPollInterval = 15 * time.Second -// stagingStateFailed is the ModelDeploymentStagingStatus state of an artifact that -// failed to stage and won't be retried automatically. -const stagingStateFailed = "FAILED" +// ModelDeploymentStagingStatus states air deploy acts on. +const ( + // stagingStateInProgress: an upload is queued, running or waiting to be retried. + stagingStateInProgress = "IN_PROGRESS" + // stagingStateFailed: the upload failed and won't be retried automatically. + stagingStateFailed = "FAILED" +) // registerModelPackageRequest is the RegisterModelPackageRequest body. Exactly one // of Checkpoint, Code and SavedModel is set. @@ -86,6 +95,17 @@ type artifactStagingStatus struct { State string `json:"state,omitempty"` } +// stagingStatusRequest is the CheckModelDeploymentStagingStatusRequest body. +type stagingStatusRequest struct { + Model string `json:"model"` + ModelVersion string `json:"model_version"` +} + +// stagingStatusResponse is the CheckModelDeploymentStagingStatusResponse body. +type stagingStatusResponse struct { + Status *modelStagingStatus `json:"status,omitempty"` +} + // weightsState is the staging state of a checkpoint's weights, or "" when the // response reports none. func (r *registerModelPackageResponse) weightsState() string { @@ -134,22 +154,26 @@ func (c *deployConfig) registrationRequest(key string, code *modelPackageCode) r return req } -// registerUntilReady sends req until the version it registers is ready to serve, -// calling progress with every response. A checkpoint's version becomes ready once -// its weights are staged; the other sources are ready after the first request. -// It returns errStagingFailed, with the last response, when staging the weights -// failed. +// registerUntilReady registers req's version and waits until it is ready to +// serve, calling progress with every registration response and staging update. +// The other sources are ready after the first request. A checkpoint's version +// becomes ready once its weights are staged: while they copy, it polls the staging +// status, a read, then repeats the registration, which finalizes the version. On +// error it returns the last response it got, nil before the first one, and +// errStagingFailed when staging failed. func registerUntilReady(ctx context.Context, w *databricks.WorkspaceClient, req registerModelPackageRequest, interval time.Duration, progress func(*registerModelPackageResponse)) (*registerModelPackageResponse, error) { apiClient, err := client.New(w.Config) if err != nil { return nil, fmt.Errorf("failed to create API client: %w", err) } - headers := auth.WorkspaceIDHeaders(w.Config) + p := modelServingPoller{api: apiClient, headers: auth.WorkspaceIDHeaders(w.Config), interval: interval} + var last *registerModelPackageResponse for { var resp registerModelPackageResponse - if err := apiClient.Do(ctx, http.MethodPost, registerModelPackagePath, headers, nil, req, &resp); err != nil { - return nil, err + if err := p.post(ctx, registerModelPackagePath, req, &resp); err != nil { + return last, err } + last = &resp progress(&resp) if resp.Ready { return &resp, nil @@ -160,13 +184,50 @@ func registerUntilReady(ctx context.Context, w *databricks.WorkspaceClient, req // Only the first request restages a failed checkpoint: repeating it would // restart a staging that keeps failing instead of reporting the failure. req.RetryFailedStaging = false + if err := p.waitForStaging(ctx, &resp, progress); err != nil { + return &resp, err + } + } +} - timer := time.NewTimer(interval) - select { - case <-ctx.Done(): - timer.Stop() - return &resp, ctx.Err() - case <-timer.C: +// modelServingPoller sends the requests of a registration that can take many +// minutes. +type modelServingPoller struct { + api *client.DatabricksClient + headers map[string]string + interval time.Duration +} + +// waitForStaging polls the staging status of resp's version, updating resp, until +// its weights are no longer being copied. +func (p modelServingPoller) waitForStaging(ctx context.Context, resp *registerModelPackageResponse, progress func(*registerModelPackageResponse)) error { + statusReq := stagingStatusRequest{Model: resp.RegisteredModel, ModelVersion: resp.ModelVersion} + for { + if err := sleepOrCancel(ctx, p.interval); err != nil { + return err + } + if resp.weightsState() != stagingStateInProgress { + return nil + } + var status stagingStatusResponse + if err := p.post(ctx, stagingStatusPath, statusReq, &status); err != nil { + return err + } + resp.StagingStatus = status.Status + progress(resp) + if resp.weightsState() != stagingStateInProgress { + return nil } } } + +// post sends a request, repeating it after a transient failure the SDK doesn't +// retry: a wait of many minutes outlasts the SDK's own retries. Both requests air +// deploy polls with are safe to repeat, the registration because it is +// idempotent per key. +func (p modelServingPoller) post(ctx context.Context, path string, body, out any) error { + _, err := retry.With(ctx, retry.IsTransientRead, func() (struct{}, error) { + return struct{}{}, p.api.Do(ctx, http.MethodPost, path, p.headers, nil, body, out) + }) + return err +} diff --git a/cmd/air/deployregister_test.go b/cmd/air/deployregister_test.go index 90732e14df5..f548faf2da2 100644 --- a/cmd/air/deployregister_test.go +++ b/cmd/air/deployregister_test.go @@ -10,19 +10,30 @@ import ( "testing" "time" + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/retry" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -// fakeModelServing serves registerModelPackage with the given responses in turn, -// repeating the last, and records every registration body. Other paths are served -// by endpoints when set, else get an empty object (the SDK's config probe). +// Canned fake responses that stand for an error rather than a body. +const ( + conflictResponse = "conflict" + serverErrorResponse = "server error" +) + +// fakeModelServing serves registerModelPackage with responses, and the staging +// status API with statuses, each in turn and then repeating the last; statuses +// default to SUCCEEDED. It records every request. Other paths are served by +// endpoints when set, else get an empty object (the SDK's config probe). type fakeModelServing struct { - t *testing.T - mu sync.Mutex - responses []string - bodies []registerModelPackageRequest - endpoints http.HandlerFunc + t *testing.T + mu sync.Mutex + responses []string + statuses []string + bodies []registerModelPackageRequest + statusRequests []stagingStatusRequest + endpoints http.HandlerFunc } func newFakeModelServing(t *testing.T, responses ...string) *fakeModelServing { @@ -30,31 +41,61 @@ func newFakeModelServing(t *testing.T, responses ...string) *fakeModelServing { } func (f *fakeModelServing) ServeHTTP(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != registerModelPackagePath { + switch r.URL.Path { + case registerModelPackagePath: + var body registerModelPackageRequest + if !f.decode(w, r, &body) { + return + } + f.mu.Lock() + defer f.mu.Unlock() + f.bodies = append(f.bodies, body) + f.respond(w, f.responses, len(f.bodies)) + case stagingStatusPath: + var body stagingStatusRequest + if !f.decode(w, r, &body) { + return + } + f.mu.Lock() + defer f.mu.Unlock() + f.statusRequests = append(f.statusRequests, body) + statuses := f.statuses + if len(statuses) == 0 { + statuses = []string{statusSucceeded} + } + f.respond(w, statuses, len(f.statusRequests)) + default: if f.endpoints != nil { f.endpoints(w, r) return } _, _ = w.Write([]byte(`{}`)) - return } +} + +// decode reads a request body into v, failing the request (not the handler +// goroutine) when it can't. +func (f *fakeModelServing) decode(w http.ResponseWriter, r *http.Request, v any) bool { raw, err := io.ReadAll(r.Body) - var body registerModelPackageRequest - if !assert.NoError(f.t, err) || !assert.NoError(f.t, json.Unmarshal(raw, &body)) { + if !assert.NoError(f.t, err) || !assert.NoError(f.t, json.Unmarshal(raw, v)) { w.WriteHeader(http.StatusInternalServerError) - return + return false } + return true +} - f.mu.Lock() - defer f.mu.Unlock() - f.bodies = append(f.bodies, body) - i := min(len(f.bodies), len(f.responses)) - 1 - if f.responses[i] == "" { +// respond writes the nth (1-based) of responses, repeating the last. +func (f *fakeModelServing) respond(w http.ResponseWriter, responses []string, n int) { + switch body := responses[min(n, len(responses))-1]; body { + case conflictResponse: w.WriteHeader(http.StatusConflict) _, _ = w.Write([]byte(`{"error_code": "RESOURCE_CONFLICT", "message": "The idempotency key can't be used for this registration. Use a new key."}`)) - return + case serverErrorResponse: + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL_ERROR", "message": "Something went wrong."}`)) + default: + _, _ = w.Write([]byte(body)) } - _, _ = w.Write([]byte(f.responses[i])) } func (f *fakeModelServing) start() *httptest.Server { @@ -76,8 +117,17 @@ const ( "staging_status": {"state": "SUCCEEDED", "model_artifact": {"state": "SUCCEEDED"}, "environment_artifact": {"state": "SUCCEEDED"}}}` stagingFailed = `{"model_id": "m-1", "registered_model": "main.models.qwen", "model_version": "3", "ready": false, "staging_status": {"state": "FAILED", "model_artifact": {"state": "FAILED"}, "environment_artifact": {"state": "SUCCEEDED"}}}` + + statusInProgress = `{"status": {"state": "IN_PROGRESS", "model_artifact": {"state": "IN_PROGRESS"}, "environment_artifact": {"state": "SUCCEEDED"}}}` + statusSucceeded = `{"status": {"state": "SUCCEEDED", "model_artifact": {"state": "SUCCEEDED"}, "environment_artifact": {"state": "SUCCEEDED"}}}` + statusFailed = `{"status": {"state": "FAILED", "model_artifact": {"state": "FAILED"}, "environment_artifact": {"state": "SUCCEEDED"}}}` ) +// fastRetries shortens the pause before air deploy repeats a failed request. +func fastRetries(ctx context.Context) context.Context { + return env.Set(ctx, retry.IntervalMsVariable, "1") +} + func checkpointRequest() registerModelPackageRequest { return registerModelPackageRequest{ RegisteredModel: "main.models.qwen", @@ -91,45 +141,75 @@ func checkpointRequest() registerModelPackageRequest { } } -func TestRegisterUntilReadyRepeatsTheSameRequest(t *testing.T) { - fake := newFakeModelServing(t, stagingInProgress, stagingInProgress, stagingReady) +func registerTestCheckpoint(t *testing.T, fake *fakeModelServing, req registerModelPackageRequest) (*registerModelPackageResponse, []string, error) { + t.Helper() srv := fake.start() - var states []string - req := checkpointRequest() - req.RetryFailedStaging = true - resp, err := registerUntilReady(t.Context(), newTestWorkspaceClient(t, srv.URL), req, time.Millisecond, func(r *registerModelPackageResponse) { + resp, err := registerUntilReady(fastRetries(t.Context()), newTestWorkspaceClient(t, srv.URL), req, time.Millisecond, func(r *registerModelPackageResponse) { states = append(states, r.weightsState()) }) + return resp, states, err +} + +func TestRegisterUntilReadyPollsStagingThenFinalizes(t *testing.T) { + fake := newFakeModelServing(t, stagingInProgress, stagingReady) + fake.statuses = []string{statusInProgress, statusSucceeded} + req := checkpointRequest() + req.RetryFailedStaging = true + + resp, states, err := registerTestCheckpoint(t, fake, req) require.NoError(t, err) assert.True(t, resp.Ready) assert.Equal(t, "3", resp.ModelVersion) - assert.Equal(t, []string{"IN_PROGRESS", "IN_PROGRESS", "SUCCEEDED"}, states) + assert.Equal(t, []string{"IN_PROGRESS", "IN_PROGRESS", "SUCCEEDED", "SUCCEEDED"}, states) + // The staging status is polled for the registered version; the registration is + // repeated once the weights are staged, which finalizes the version. + assert.Equal(t, []stagingStatusRequest{{Model: "main.models.qwen", ModelVersion: "3"}, {Model: "main.models.qwen", ModelVersion: "3"}}, fake.statusRequests) bodies := fake.requests() - require.Len(t, bodies, 3) + require.Len(t, bodies, 2) assert.Equal(t, req, bodies[0]) - // Only the first request restages; the repeats continue the same registration. + // Only the first request restages; the repeat continues the same registration. req.RetryFailedStaging = false assert.Equal(t, req, bodies[1]) - assert.Equal(t, req, bodies[2]) } func TestRegisterUntilReadyReportsFailedStaging(t *testing.T) { fake := newFakeModelServing(t, stagingInProgress, stagingFailed) - srv := fake.start() + fake.statuses = []string{statusFailed} - resp, err := registerUntilReady(t.Context(), newTestWorkspaceClient(t, srv.URL), checkpointRequest(), time.Millisecond, func(*registerModelPackageResponse) {}) + resp, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) require.ErrorIs(t, err, errStagingFailed) assert.Equal(t, "3", resp.ModelVersion) + // The failure the status reports is confirmed by the registration itself. assert.Len(t, fake.requests(), 2) } +func TestRegisterUntilReadyRetriesTransientErrors(t *testing.T) { + fake := newFakeModelServing(t, serverErrorResponse, stagingInProgress, stagingReady) + fake.statuses = []string{serverErrorResponse, statusSucceeded} + + resp, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) + require.NoError(t, err) + assert.True(t, resp.Ready) + assert.Len(t, fake.requests(), 3) + assert.Len(t, fake.statusRequests, 2) +} + +func TestRegisterUntilReadyGivesUpOnPersistentErrors(t *testing.T) { + fake := newFakeModelServing(t, serverErrorResponse) + + _, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) + assert.ErrorContains(t, err, "Something went wrong.") + assert.Len(t, fake.requests(), retry.MaxRetries+1) +} + func TestRegisterUntilReadyReturnsServerErrors(t *testing.T) { - srv := newFakeModelServing(t, "").start() + fake := newFakeModelServing(t, conflictResponse) - _, err := registerUntilReady(t.Context(), newTestWorkspaceClient(t, srv.URL), checkpointRequest(), time.Millisecond, func(*registerModelPackageResponse) {}) + _, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) assert.ErrorContains(t, err, "The idempotency key can't be used for this registration") + assert.Len(t, fake.requests(), 1) } func TestRegisterUntilReadyStopsWhenCanceled(t *testing.T) { diff --git a/cmd/air/snapshot_package.go b/cmd/air/snapshot_package.go index c1dc9c5f89d..3a4f6ff6ddd 100644 --- a/cmd/air/snapshot_package.go +++ b/cmd/air/snapshot_package.go @@ -146,7 +146,9 @@ func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, if output, err := cmd.CombinedOutput(); err != nil { return nil, fmt.Errorf("failed to initialize temporary git metadata: %w: %s", err, strings.TrimSpace(string(output))) } - args = []string{"--git-dir", gitDir, "--work-tree", repoPath, "ls-files", "-z", "--cached", "--others", "--exclude-standard"} + // -C makes ls-files list paths relative to repoPath; without it they are + // relative to the process's working directory, which may be a subdirectory. + args = []string{"-C", repoPath, "--git-dir", gitDir, "--work-tree", repoPath, "ls-files", "-z", "--cached", "--others", "--exclude-standard"} } if len(includePaths) > 0 { diff --git a/cmd/air/snapshot_package_test.go b/cmd/air/snapshot_package_test.go index 068997e2671..37c7bdbfec3 100644 --- a/cmd/air/snapshot_package_test.go +++ b/cmd/air/snapshot_package_test.go @@ -324,3 +324,21 @@ func TestSnapshotFiles_RejectsIncludePathsMatchingNothing(t *testing.T) { assert.Contains(t, err.Error(), "contains no files") assert.Contains(t, err.Error(), "include_paths") } + +func TestSnapshotFilesOutsideGitFromASubdirectory(t *testing.T) { + dir := t.TempDir() + writeRepoFile(t, dir, ".gitignore", "/creds.json\n") + writeRepoFile(t, dir, "creds.json", "{}") + writeRepoFile(t, dir, "agent.py", "x") + writeRepoFile(t, dir, "tools/a.py", "a") + // The listing must be relative to the root, not to the working directory. + t.Chdir(filepath.Join(dir, "tools")) + + files, err := snapshotFiles(t.Context(), dir, nil, false) + require.NoError(t, err) + var rels []string + for _, f := range files { + rels = append(rels, filepath.ToSlash(f.rel)) + } + assert.ElementsMatch(t, []string{".gitignore", "agent.py", "tools/a.py"}, rels) +} From 8d2d9f9272d5cfdb04a786b95028f24c31e42e66 Mon Sep 17 00:00:00 2001 From: Arpit Agarwal Date: Sun, 15 Nov 2026 06:38:05 +0000 Subject: [PATCH 3/4] air deploy: retry each endpoint poll, stop on a not-ready code model - Poll the endpoint per GET with libs/retry, so transient failures spread over a long start don't end the wait. - Don't resend a model from code that isn't ready after registering. - Keep git's stderr in check-ignore failures; reject include_paths patterns and padding; don't mark a local client error retryable. Co-authored-by: Isaac --- cmd/air/deploycode.go | 4 +++- cmd/air/deployconfig.go | 7 +++--- cmd/air/deployconfig_test.go | 4 +++- cmd/air/deployendpoint.go | 39 +++++++++++++++++++++++++++------- cmd/air/deployendpoint_test.go | 33 ++++++++++++++++++---------- cmd/air/deployregister.go | 7 +++++- cmd/air/deployregister_test.go | 9 ++++++++ 7 files changed, 78 insertions(+), 25 deletions(-) diff --git a/cmd/air/deploycode.go b/cmd/air/deploycode.go index 433941d77af..721fadd7825 100644 --- a/cmd/air/deploycode.go +++ b/cmd/air/deploycode.go @@ -101,15 +101,17 @@ func ignoredTrackedFiles(ctx context.Context, root string, files []snapshotFile) paths.WriteString(filepath.ToSlash(f.rel)) paths.WriteByte(0) } + var stderr bytes.Buffer cmd := exec.CommandContext(ctx, "git", "-C", root, "check-ignore", "--no-index", "-z", "--stdin") cmd.Stdin = &paths + cmd.Stderr = &stderr out, err := cmd.Output() // check-ignore exits 1 when no path is ignored. if exitErr, ok := errors.AsType[*exec.ExitError](err); ok && exitErr.ExitCode() == 1 { return nil, nil } if err != nil { - return nil, fmt.Errorf("failed to evaluate git ignore rules: %w", err) + return nil, fmt.Errorf("failed to evaluate git ignore rules: %w: %s", err, strings.TrimSpace(stderr.String())) } var ignored []string for raw := range bytes.SplitSeq(out, []byte{0}) { diff --git a/cmd/air/deployconfig.go b/cmd/air/deployconfig.go index 29619ef9db7..9c3bdac4cf0 100644 --- a/cmd/air/deployconfig.go +++ b/cmd/air/deployconfig.go @@ -250,10 +250,11 @@ func (c *deployConfig) validateSource() error { return errors.New("code_source.include_paths is required: list the files and directories to send, e.g. include_paths: [agent.py, tools]") } // git reads a leading ':' as pathspec magic, such as ":!x" for everything - // but x, which would send files the list doesn't name. + // but x, and expands globs, either of which would send files the list + // doesn't name; padding would match nothing. for _, p := range c.CodeSource.IncludePaths { - if strings.HasPrefix(strings.TrimSpace(p), ":") { - return fmt.Errorf("code_source.include_paths must name files and directories, got %q", p) + if strings.HasPrefix(p, ":") || strings.ContainsAny(p, "*?[") || p != strings.TrimSpace(p) { + return fmt.Errorf("code_source.include_paths must name files and directories, without patterns or surrounding spaces, got %q", p) } } snapshot := snapshotSourceConfig{RootPath: c.CodeSource.RootPath, IncludePaths: c.CodeSource.IncludePaths} diff --git a/cmd/air/deployconfig_test.go b/cmd/air/deployconfig_test.go index 6d5efc47302..722d1acb350 100644 --- a/cmd/air/deployconfig_test.go +++ b/cmd/air/deployconfig_test.go @@ -189,7 +189,9 @@ func TestDeployConfigSourceFields(t *testing.T) { {"code entry in a directory", "model: {type: code, entry: src/agent.py}\ncode_source: {root_path: .}\n", "model.entry must name a .py file directly under code_source.root_path"}, {"code entry not python", "model: {type: code, entry: agent.ipynb}\ncode_source: {root_path: .}\n", "model.entry must name a .py file"}, {"code without include_paths", "model: {type: code, entry: agent.py}\ncode_source: {root_path: .}\n", "code_source.include_paths is required: list the files and directories to send"}, - {"code include pathspec magic", "model: {type: code, entry: agent.py}\ncode_source: {root_path: ., include_paths: [\":!creds.json\"]}\n", `code_source.include_paths must name files and directories, got ":!creds.json"`}, + {"code include pathspec magic", "model: {type: code, entry: agent.py}\ncode_source: {root_path: ., include_paths: [\":!creds.json\"]}\n", `code_source.include_paths must name files and directories, without patterns or surrounding spaces, got ":!creds.json"`}, + {"code include glob", "model: {type: code, entry: agent.py}\ncode_source: {root_path: ., include_paths: [\"*\"]}\n", `code_source.include_paths must name files and directories, without patterns or surrounding spaces, got "*"`}, + {"code include padded", "model: {type: code, entry: agent.py}\ncode_source: {root_path: ., include_paths: [\" agent.py\"]}\n", `got " agent.py"`}, {"code include traversal", "model: {type: code, entry: agent.py}\ncode_source: {root_path: ., include_paths: [../x]}\n", "code_source.include_paths cannot contain '..' traversal"}, {"saved model with task", "model: {type: saved_model_dir, path: /Volumes/a/b/c/m}\ntask: llm/v1/chat\n", "task can't be set when model.type is saved_model_dir"}, {"saved model with environment", "model: {type: saved_model_dir, path: /Volumes/a/b/c/m}\nenvironment: {version: 6}\n", "environment can't be set when model.type is saved_model_dir"}, diff --git a/cmd/air/deployendpoint.go b/cmd/air/deployendpoint.go index e240dbc1f0b..1effc2a8303 100644 --- a/cmd/air/deployendpoint.go +++ b/cmd/air/deployendpoint.go @@ -11,6 +11,7 @@ import ( "github.com/databricks/cli/libs/retry" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/retries" "github.com/databricks/databricks-sdk-go/service/serving" ) @@ -40,13 +41,34 @@ func getEndpoint(ctx context.Context, w *databricks.WorkspaceClient, name string return endpoint, err } -// waitNotUpdating waits until the endpoint has no update in progress. The SDK's -// wait stops at the first failed poll, so a transient failure starts it again: -// an endpoint can take many minutes to start. Creates and config updates aren't -// repeated this way, since a repeat after an applied write would fail. +// waitNotUpdating waits until the endpoint has no update in progress, calling +// progress on every poll. It polls as the SDK's WaitGetServingEndpointNotUpdating +// does, except that each poll repeats a transient failure, which ends the SDK's +// wait: an endpoint can take many minutes to start. Creates and config updates +// aren't repeated this way, since a repeat after an applied write would fail. func waitNotUpdating(ctx context.Context, w *databricks.WorkspaceClient, name string, progress func(*serving.ServingEndpointDetailed)) (*serving.ServingEndpointDetailed, error) { - return retry.With(ctx, retry.IsTransientRead, func() (*serving.ServingEndpointDetailed, error) { - return w.ServingEndpoints.WaitGetServingEndpointNotUpdating(ctx, name, untilDeadline(ctx), progress) + return retries.Poll(ctx, untilDeadline(ctx), func() (*serving.ServingEndpointDetailed, *retries.Err) { + endpoint, err := retry.With(ctx, retry.IsTransientRead, func() (*serving.ServingEndpointDetailed, error) { + return w.ServingEndpoints.GetByName(ctx, name) + }) + if err != nil { + return nil, retries.Halt(err) + } + if progress != nil { + progress(endpoint) + } + var state serving.EndpointStateConfigUpdate + if endpoint.State != nil { + state = endpoint.State.ConfigUpdate + } + switch state { + case serving.EndpointStateConfigUpdateNotUpdating: + return endpoint, nil + case serving.EndpointStateConfigUpdateUpdateFailed, serving.EndpointStateConfigUpdateUpdateCanceled: + return nil, retries.Halt(fmt.Errorf("the endpoint's update ended in %s; check its events and logs in the Serving UI", state)) + default: + return nil, retries.Continues(string(state)) + } }) } @@ -131,8 +153,9 @@ func waitServes(ctx context.Context, w *databricks.WorkspaceClient, name string, return checkServes(endpoint, entity) } -// errEndpointNotReady reports an endpoint that serves the version but isn't ready, -// which running air deploy again doesn't change. +// errEndpointNotReady reports an endpoint that serves the version but isn't ready. +// Running air deploy again doesn't make it ready; it reports success once the +// endpoint recovers, so the error isn't marked retryable. var errEndpointNotReady = errors.New("endpoint isn't ready") // checkServes returns endpoint if it serves only entity and is ready to serve. diff --git a/cmd/air/deployendpoint_test.go b/cmd/air/deployendpoint_test.go index e267a1bdaf4..0a4dc57e019 100644 --- a/cmd/air/deployendpoint_test.go +++ b/cmd/air/deployendpoint_test.go @@ -27,8 +27,11 @@ type fakeEndpoints struct { writes []string // rawWrites holds the body of every create and config update, as sent. rawWrites []string - // failGets is how many GETs answer 500 before the endpoint is served again. - failGets int + // gets counts the GETs of an existing endpoint. Those numbered in failOnGets + // answer 500, and from readyAfterGets on (when set) the endpoint is ready. + gets int + failOnGets map[int]bool + readyAfterGets int } func newFakeEndpoints(t *testing.T, existing *serving.ServingEndpointDetailed) *fakeEndpoints { @@ -66,11 +69,16 @@ func (f *fakeEndpoints) ServeHTTP(w http.ResponseWriter, r *http.Request) { f.endpoint.PendingConfig = nil f.endpoint.State = &f.finalState case r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, servingEndpointsPath+"/"): - if f.failGets > 0 && f.endpoint != nil { - f.failGets-- - w.WriteHeader(http.StatusInternalServerError) - _, _ = w.Write([]byte(`{"error_code": "INTERNAL_ERROR", "message": "Something went wrong."}`)) - return + if f.endpoint != nil { + f.gets++ + if f.failOnGets[f.gets] { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL_ERROR", "message": "Something went wrong."}`)) + return + } + if f.readyAfterGets > 0 && f.gets >= f.readyAfterGets { + f.endpoint.State = &readyState + } } if f.endpoint == nil { w.WriteHeader(http.StatusNotFound) @@ -252,16 +260,19 @@ func TestCheckEndpointAllowsAMissingEndpoint(t *testing.T) { assert.NoError(t, checkEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", "main.models.qwen")) } -func TestServeOnEndpointWaitsThroughATransientFailure(t *testing.T) { +func TestServeOnEndpointWaitsThroughTransientFailures(t *testing.T) { fake := newFakeEndpoints(t, nil) - // The create succeeds; the first poll of the wait then fails. - fake.failGets = 1 + // The endpoint starts over three polls, and the first GET of each fails: more + // failures than one request may retry, spread over the wait. + fake.finalState = serving.EndpointState{Ready: serving.EndpointStateReadyNotReady, ConfigUpdate: serving.EndpointStateConfigUpdateInProgress} + fake.failOnGets = map[int]bool{1: true, 3: true, 5: true} + fake.readyAfterGets = 6 srv := fake.start() _, err := serveOnEndpoint(withDeadline(t), newTestWorkspaceClient(t, srv.URL), "qwen-chat", testEntity("3"), nil) require.NoError(t, err) assert.Equal(t, []string{"create"}, fake.writes) - assert.Zero(t, fake.failGets) + assert.Equal(t, 6, fake.gets) } func TestServeOnEndpointReportsAFailedUpdate(t *testing.T) { diff --git a/cmd/air/deployregister.go b/cmd/air/deployregister.go index a5ae7277553..cbc128fadb8 100644 --- a/cmd/air/deployregister.go +++ b/cmd/air/deployregister.go @@ -164,7 +164,7 @@ func (c *deployConfig) registrationRequest(key string, code *modelPackageCode) r func registerUntilReady(ctx context.Context, w *databricks.WorkspaceClient, req registerModelPackageRequest, interval time.Duration, progress func(*registerModelPackageResponse)) (*registerModelPackageResponse, error) { apiClient, err := client.New(w.Config) if err != nil { - return nil, fmt.Errorf("failed to create API client: %w", err) + return nil, withCode(codeInternalError, false, fmt.Errorf("failed to create API client: %w", err)) } p := modelServingPoller{api: apiClient, headers: auth.WorkspaceIDHeaders(w.Config), interval: interval} var last *registerModelPackageResponse @@ -181,6 +181,11 @@ func registerUntilReady(ctx context.Context, w *databricks.WorkspaceClient, req if resp.weightsState() == stagingStateFailed { return &resp, errStagingFailed } + // Only a checkpoint waits on staging; repeating any other registration + // would resend its files to no effect. + if req.Checkpoint == nil { + return &resp, fmt.Errorf("version %s of %s isn't ready after registering", stripControl(resp.ModelVersion), req.RegisteredModel) + } // Only the first request restages a failed checkpoint: repeating it would // restart a staging that keeps failing instead of reporting the failure. req.RetryFailedStaging = false diff --git a/cmd/air/deployregister_test.go b/cmd/air/deployregister_test.go index f548faf2da2..22e4f6b567b 100644 --- a/cmd/air/deployregister_test.go +++ b/cmd/air/deployregister_test.go @@ -196,6 +196,15 @@ func TestRegisterUntilReadyRetriesTransientErrors(t *testing.T) { assert.Len(t, fake.statusRequests, 2) } +func TestRegisterUntilReadyDoesNotRepeatANotReadyModelFromCode(t *testing.T) { + fake := newFakeModelServing(t, `{"model_id": "m-2", "registered_model": "a.b.c", "model_version": "1", "ready": false}`) + req := registerModelPackageRequest{RegisteredModel: "a.b.c", IdempotencyKey: "k", Code: &modelPackageCode{EntryFile: "agent.py"}} + + _, _, err := registerTestCheckpoint(t, fake, req) + assert.ErrorContains(t, err, "version 1 of a.b.c isn't ready after registering") + assert.Len(t, fake.requests(), 1) +} + func TestRegisterUntilReadyGivesUpOnPersistentErrors(t *testing.T) { fake := newFakeModelServing(t, serverErrorResponse) From 261b2028fb2eab0149eb41b31b90dfb1ec6daa68 Mon Sep 17 00:00:00 2001 From: Arpit Agarwal Date: Sun, 15 Nov 2026 16:43:22 +0000 Subject: [PATCH 4/4] air deploy: register once, poll staging, then finalize the version The server now leaves a registered version pending: air deploy registers it once, polls the read-only staging status until the weights are copied or fail, then calls finalizeModelPackage. A status without a state is waited out, a finalize failure prints the command that continues, and the finalize response must name the same version. Co-authored-by: Isaac --- cmd/air/deploy.go | 20 +++- cmd/air/deploy_test.go | 40 +++++-- cmd/air/deployregister.go | 150 +++++++++++++++++-------- cmd/air/deployregister_test.go | 200 ++++++++++++++++++++++++++------- 4 files changed, 305 insertions(+), 105 deletions(-) diff --git a/cmd/air/deploy.go b/cmd/air/deploy.go index 8ac4dd903ba..4d2b0d876f5 100644 --- a/cmd/air/deploy.go +++ b/cmd/air/deploy.go @@ -289,7 +289,7 @@ func renderDeployDryRun(ctx context.Context, cfg *deployConfig, code *modelPacka return nil } -// registerDeployment registers the version and waits until it is ready, showing +// registerDeployment registers the version and makes it ready to serve, showing // progress on a spinner when show is true. Its errors say how to continue. func registerDeployment(ctx context.Context, w *databricks.WorkspaceClient, cfg *deployConfig, opts deployOptions, code *modelPackageCode, key string, pollInterval time.Duration, show bool) (*registerModelPackageResponse, error) { // Printed to stderr in every output mode, before the first request, so a run @@ -306,7 +306,7 @@ func registerDeployment(ctx context.Context, w *databricks.WorkspaceClient, cfg sp.Update(registrationProgress(cfg, resp)) } } - resp, err := registerUntilReady(ctx, w, req, pollInterval, progress) + resp, err := registerAndFinalize(ctx, w, req, pollInterval, progress) if err == nil { return resp, nil } @@ -320,11 +320,18 @@ func registerDeployment(ctx context.Context, w *databricks.WorkspaceClient, cfg } if stopped := deployStopped(ctx, err); stopped != nil { continuing := "" - if resp != nil && resp.weightsState() != "" { - continuing = " Copying the weights continues, and the version is finalized when you continue." + switch { + case resp == nil: + case resp.weightsState() == stagingStateSucceeded: + continuing = " The weights are copied, so continuing finalizes the version." + case resp.weightsState() != "": + continuing = " Copying the weights continues meanwhile." } return nil, fmt.Errorf("%w while registering a version of %s.%s To continue, run:\n %s", stopped, cfg.RegisteredModel, continuing, resume) } + if _, ok := errors.AsType[*finalizeError](err); ok { + return nil, fmt.Errorf("finalizing version %s of %s failed: %w\n\nTo try again, run:\n %s", stripControl(resp.ModelVersion), cfg.RegisteredModel, err, resume) + } failed := fmt.Errorf("registering a version of %s with idempotency key %s failed: %w", cfg.RegisteredModel, key, err) // A retryable failure is retried with the same key, so no other version is // registered. @@ -396,8 +403,11 @@ func validateDeployFlags(cfg *deployConfig, opts deployOptions) error { // registrationProgress describes a registration response on the spinner. func registrationProgress(cfg *deployConfig, resp *registerModelPackageResponse) string { state := resp.weightsState() - if state == "" { + switch state { + case "": return fmt.Sprintf("Registering version %s of %s…", stripControl(resp.ModelVersion), cfg.RegisteredModel) + case stagingStateSucceeded: + return fmt.Sprintf("Finalizing version %s of %s…", stripControl(resp.ModelVersion), cfg.RegisteredModel) } return fmt.Sprintf("Copying the weights into version %s of %s (%s)…", stripControl(resp.ModelVersion), cfg.RegisteredModel, stripControl(state)) } diff --git a/cmd/air/deploy_test.go b/cmd/air/deploy_test.go index 2ae6ac0917c..05cbdab8d6b 100644 --- a/cmd/air/deploy_test.go +++ b/cmd/air/deploy_test.go @@ -35,8 +35,9 @@ func withDeadline(t *testing.T) context.Context { return ctx } -// deployServer serves registerModelPackage with responses and the serving -// endpoints API from endpoints. +// deployServer serves registerModelPackage with responses, the staging status +// and finalizeModelPackage with their defaults, and the serving endpoints API +// from endpoints. func deployServer(t *testing.T, endpoints *fakeEndpoints, responses ...string) (*httptest.Server, *fakeModelServing) { t.Helper() fake := newFakeModelServing(t, responses...) @@ -70,7 +71,7 @@ func runDeployCmd(t *testing.T, out flags.Output, buf *bytes.Buffer, srvURL, con func TestDeployCheckpointJSON(t *testing.T) { endpoints := newFakeEndpoints(t, nil) - srv, fake := deployServer(t, endpoints, stagingInProgress, stagingReady) + srv, fake := deployServer(t, endpoints, stagingInProgress) var buf bytes.Buffer err := runDeployCmd(t, flags.OutputJSON, &buf, srv.URL, checkpointDeployConfig, "idempotency-key", "key-1") @@ -91,9 +92,9 @@ func TestDeployCheckpointJSON(t *testing.T) { }, got.Data) bodies := fake.requests() - require.Len(t, bodies, 2) + require.Len(t, bodies, 1) assert.Equal(t, "key-1", bodies[0].IdempotencyKey) - assert.Equal(t, bodies[0], bodies[1]) + assert.Equal(t, []finalizeModelPackageRequest{{RegisteredModel: "main.models.qwen", ModelVersion: "3"}}, fake.finalizeRequests) assert.Equal(t, []string{"create"}, endpoints.writes) served := endpoints.endpoint.Config.ServedEntities[0] assert.Equal(t, "3", served.EntityVersion) @@ -156,15 +157,14 @@ func TestDeployPrintsTheKeyBeforeRegisteringInJSONMode(t *testing.T) { assert.Contains(t, stdout.String(), `"code": "RESOURCE_CONFLICT"`) } -func TestDeployRetriesFailedStagingOnce(t *testing.T) { - srv, fake := deployServer(t, newFakeEndpoints(t, nil), stagingInProgress, stagingReady) +func TestDeploySendsRetryFailedStaging(t *testing.T) { + srv, fake := deployServer(t, newFakeEndpoints(t, nil), stagingInProgress) var buf bytes.Buffer require.NoError(t, runDeployCmd(t, flags.OutputJSON, &buf, srv.URL, checkpointDeployConfig, "retry-failed-staging", "true")) bodies := fake.requests() - require.Len(t, bodies, 2) + require.Len(t, bodies, 1) assert.True(t, bodies[0].RetryFailedStaging) - assert.False(t, bodies[1].RetryFailedStaging) } func TestDeployFailedStagingSaysHowToRetry(t *testing.T) { @@ -253,6 +253,28 @@ func TestDeployTimesOutWhileCopyingTheWeights(t *testing.T) { assert.Contains(t, got.Error.Message, "air deploy timed out while registering a version of main.models.qwen. Copying the weights continues") assert.Contains(t, got.Error.Message, " --idempotency-key key-1") assert.Len(t, fake.requests(), 1) + assert.Empty(t, fake.finalizeRequests) +} + +func TestDeployFinalizeFailureSaysHowToRetry(t *testing.T) { + srv, fake := deployServer(t, newFakeEndpoints(t, nil), stagingSucceeded) + fake.finalizations = []string{notFinalizableResponse} + + var buf bytes.Buffer + err := runDeployCmd(t, flags.OutputText, &buf, srv.URL, checkpointDeployConfig, "idempotency-key", "key-1") + require.Error(t, err) + assert.Contains(t, err.Error(), "finalizing version 3 of main.models.qwen failed: Version 3 of model main.models.qwen can't be finalized yet") + assert.Contains(t, err.Error(), "To try again, run:\n databricks air deploy -f ") + assert.Contains(t, err.Error(), " --idempotency-key key-1") +} + +func TestDeployWarnsWhenTheFinalizedEnvironmentFailedToStage(t *testing.T) { + srv, fake := deployServer(t, newFakeEndpoints(t, nil), stagingInProgress) + fake.finalizations = []string{strings.Replace(finalizedResponse, `"environment_artifact": {"state": "IN_PROGRESS"}`, `"environment_artifact": {"state": "FAILED"}`, 1)} + + var buf bytes.Buffer + require.NoError(t, runDeployCmd(t, flags.OutputText, &buf, srv.URL, checkpointDeployConfig)) + assert.Contains(t, buf.String(), "Warning: the version's serving environment couldn't be prepared ahead of time") } func TestDeployRetryableRegistrationFailureSaysHowToRetry(t *testing.T) { diff --git a/cmd/air/deployregister.go b/cmd/air/deployregister.go index cbc128fadb8..f3f24c99693 100644 --- a/cmd/air/deployregister.go +++ b/cmd/air/deployregister.go @@ -13,24 +13,33 @@ import ( "github.com/databricks/databricks-sdk-go/client" ) -// registerModelPackagePath registers a model whose files Model Serving writes. It -// runs as the caller and is idempotent per key: a repeat request with the same key -// and inputs continues the registration, and finalizes a checkpoint's version once -// its weights are staged. Staging never finalizes the version by itself. +// registerModelPackagePath registers a model whose files Model Serving writes and +// starts copying a checkpoint's weights into the version. It runs as the caller and +// is idempotent per key: a repeat request with the same key and inputs continues +// the registration. It leaves the version pending. const registerModelPackagePath = "/api/2.0/serving-endpoints:registerModelPackage" // stagingStatusPath reports the recorded staging progress of a model version, a // read. const stagingStatusPath = "/api/2.0/serving-endpoints:checkModelDeploymentStagingStatus" +// finalizeModelPackagePath finalizes a registered version once its weights are +// copied, so it can be served. It is idempotent: a finalized version is returned +// as it is. +const finalizeModelPackagePath = "/api/2.0/serving-endpoints:finalizeModelPackage" + // registrationPollInterval is how often air deploy checks on a checkpoint's // weights while they are staged. const registrationPollInterval = 15 * time.Second // ModelDeploymentStagingStatus states air deploy acts on. const ( + // stagingStateNotStaged: no staging record exists for the version. + stagingStateNotStaged = "NOT_STAGED" // stagingStateInProgress: an upload is queued, running or waiting to be retried. stagingStateInProgress = "IN_PROGRESS" + // stagingStateSucceeded: the upload is recorded as done. + stagingStateSucceeded = "SUCCEEDED" // stagingStateFailed: the upload failed and won't be retried automatically. stagingStateFailed = "FAILED" ) @@ -95,6 +104,19 @@ type artifactStagingStatus struct { State string `json:"state,omitempty"` } +// finalizeModelPackageRequest is the FinalizeModelPackageRequest body. +type finalizeModelPackageRequest struct { + RegisteredModel string `json:"registered_model"` + ModelVersion string `json:"model_version"` +} + +// finalizeModelPackageResponse is the FinalizeModelPackageResponse body, the +// fields air deploy reads. +type finalizeModelPackageResponse struct { + ModelVersion string `json:"model_version"` + StagingStatus *modelStagingStatus `json:"staging_status,omitempty"` +} + // stagingStatusRequest is the CheckModelDeploymentStagingStatusRequest body. type stagingStatusRequest struct { Model string `json:"model"` @@ -128,6 +150,14 @@ func (r *registerModelPackageResponse) environmentState() string { // into the version. Rerunning with --retry-failed-staging stages them again. var errStagingFailed = errors.New("staging failed") +// finalizeError is a failed request to finalize the registered version. +// Rerunning with the same key registers the same version again and retries it. +type finalizeError struct{ err error } + +func (e *finalizeError) Error() string { return e.err.Error() } + +func (e *finalizeError) Unwrap() error { return e.err } + // registrationRequest builds the request that registers cfg's model under key. // code is the collected model from code, and is nil for the other sources. func (c *deployConfig) registrationRequest(key string, code *modelPackageCode) registerModelPackageRequest { @@ -154,45 +184,69 @@ func (c *deployConfig) registrationRequest(key string, code *modelPackageCode) r return req } -// registerUntilReady registers req's version and waits until it is ready to -// serve, calling progress with every registration response and staging update. -// The other sources are ready after the first request. A checkpoint's version -// becomes ready once its weights are staged: while they copy, it polls the staging -// status, a read, then repeats the registration, which finalizes the version. On -// error it returns the last response it got, nil before the first one, and -// errStagingFailed when staging failed. -func registerUntilReady(ctx context.Context, w *databricks.WorkspaceClient, req registerModelPackageRequest, interval time.Duration, progress func(*registerModelPackageResponse)) (*registerModelPackageResponse, error) { +// registerAndFinalize registers req's version and makes it ready to serve, +// calling progress with the registration response and every staging update. A +// checkpoint's weights are copied into the version after registering: while they +// are, it polls the staging status, a read. It then finalizes the version, for +// every source; the server checks that the version is ready to finalize. A +// version that is already ready, such as one an earlier run finalized, is +// returned as it is. On error it returns the last response it got, nil before +// the first one; errStagingFailed when staging failed, and a *finalizeError when +// finalizing did. +func registerAndFinalize(ctx context.Context, w *databricks.WorkspaceClient, req registerModelPackageRequest, interval time.Duration, progress func(*registerModelPackageResponse)) (*registerModelPackageResponse, error) { apiClient, err := client.New(w.Config) if err != nil { return nil, withCode(codeInternalError, false, fmt.Errorf("failed to create API client: %w", err)) } p := modelServingPoller{api: apiClient, headers: auth.WorkspaceIDHeaders(w.Config), interval: interval} - var last *registerModelPackageResponse - for { - var resp registerModelPackageResponse - if err := p.post(ctx, registerModelPackagePath, req, &resp); err != nil { - return last, err - } - last = &resp - progress(&resp) - if resp.Ready { - return &resp, nil - } - if resp.weightsState() == stagingStateFailed { - return &resp, errStagingFailed + var resp registerModelPackageResponse + if err := p.post(ctx, registerModelPackagePath, req, &resp); err != nil { + return nil, err + } + if resp.ModelVersion == "" { + return nil, withCode(codeInternalError, true, fmt.Errorf("registering a version of %s returned no version", req.RegisteredModel)) + } + progress(&resp) + if resp.Ready { + return &resp, nil + } + if req.Checkpoint != nil { + if err := p.waitForWeights(ctx, req.RegisteredModel, &resp, progress); err != nil { + return &resp, err } - // Only a checkpoint waits on staging; repeating any other registration - // would resend its files to no effect. - if req.Checkpoint == nil { - return &resp, fmt.Errorf("version %s of %s isn't ready after registering", stripControl(resp.ModelVersion), req.RegisteredModel) + } + if resp.weightsState() == stagingStateFailed { + return &resp, errStagingFailed + } + var finalized finalizeModelPackageResponse + finalizeReq := finalizeModelPackageRequest{RegisteredModel: req.RegisteredModel, ModelVersion: resp.ModelVersion} + if err := p.post(ctx, finalizeModelPackagePath, finalizeReq, &finalized); err != nil { + return &resp, &finalizeError{err} + } + if finalized.ModelVersion != resp.ModelVersion { + return &resp, &finalizeError{withCode(codeInternalError, true, fmt.Errorf("the server finalized version %q instead", stripControl(finalized.ModelVersion)))} + } + resp.Ready = true + resp.StagingStatus = mergeStagingStatus(resp.StagingStatus, finalized.StagingStatus) + return &resp, nil +} + +// mergeStagingStatus is newer, with each artifact it doesn't report kept from +// older. +func mergeStagingStatus(older, newer *modelStagingStatus) *modelStagingStatus { + if newer == nil { + return older + } + merged := *newer + if older != nil { + if merged.ModelArtifact == nil { + merged.ModelArtifact = older.ModelArtifact } - // Only the first request restages a failed checkpoint: repeating it would - // restart a staging that keeps failing instead of reporting the failure. - req.RetryFailedStaging = false - if err := p.waitForStaging(ctx, &resp, progress); err != nil { - return &resp, err + if merged.EnvironmentArtifact == nil { + merged.EnvironmentArtifact = older.EnvironmentArtifact } } + return &merged } // modelServingPoller sends the requests of a registration that can take many @@ -203,33 +257,33 @@ type modelServingPoller struct { interval time.Duration } -// waitForStaging polls the staging status of resp's version, updating resp, until -// its weights are no longer being copied. -func (p modelServingPoller) waitForStaging(ctx context.Context, resp *registerModelPackageResponse, progress func(*registerModelPackageResponse)) error { - statusReq := stagingStatusRequest{Model: resp.RegisteredModel, ModelVersion: resp.ModelVersion} +// waitForWeights polls the staging status of resp's version of model, updating +// resp, until its weights are copied, failed to copy, or have no staging record. +// A state it doesn't know, or a status without one, is waited out like a copy in +// progress, until the context ends. +func (p modelServingPoller) waitForWeights(ctx context.Context, model string, resp *registerModelPackageResponse, progress func(*registerModelPackageResponse)) error { + statusReq := stagingStatusRequest{Model: model, ModelVersion: resp.ModelVersion} for { + switch resp.weightsState() { + case stagingStateSucceeded, stagingStateFailed, stagingStateNotStaged: + return nil + } if err := sleepOrCancel(ctx, p.interval); err != nil { return err } - if resp.weightsState() != stagingStateInProgress { - return nil - } var status stagingStatusResponse if err := p.post(ctx, stagingStatusPath, statusReq, &status); err != nil { return err } - resp.StagingStatus = status.Status + resp.StagingStatus = mergeStagingStatus(resp.StagingStatus, status.Status) progress(resp) - if resp.weightsState() != stagingStateInProgress { - return nil - } } } // post sends a request, repeating it after a transient failure the SDK doesn't -// retry: a wait of many minutes outlasts the SDK's own retries. Both requests air -// deploy polls with are safe to repeat, the registration because it is -// idempotent per key. +// retry: a wait of many minutes outlasts the SDK's own retries. Every request air +// deploy sends here is safe to repeat: the registration is idempotent per key, +// the status is a read, and finalizing a finalized version returns it. func (p modelServingPoller) post(ctx context.Context, path string, body, out any) error { _, err := retry.With(ctx, retry.IsTransientRead, func() (struct{}, error) { return struct{}{}, p.api.Do(ctx, http.MethodPost, path, p.headers, nil, body, out) diff --git a/cmd/air/deployregister_test.go b/cmd/air/deployregister_test.go index 22e4f6b567b..2efeb1a9028 100644 --- a/cmd/air/deployregister_test.go +++ b/cmd/air/deployregister_test.go @@ -3,9 +3,11 @@ package aircmd import ( "context" "encoding/json" + "errors" "io" "net/http" "net/http/httptest" + "strings" "sync" "testing" "time" @@ -18,22 +20,26 @@ import ( // Canned fake responses that stand for an error rather than a body. const ( - conflictResponse = "conflict" - serverErrorResponse = "server error" + conflictResponse = "conflict" + serverErrorResponse = "server error" + notFinalizableResponse = "not finalizable" ) -// fakeModelServing serves registerModelPackage with responses, and the staging -// status API with statuses, each in turn and then repeating the last; statuses -// default to SUCCEEDED. It records every request. Other paths are served by -// endpoints when set, else get an empty object (the SDK's config probe). +// fakeModelServing serves registerModelPackage with responses, the staging status +// API with statuses, and finalizeModelPackage with finalizations, each in turn and +// then repeating the last; statuses default to SUCCEEDED, and finalizations to the +// finalized version. It records every request. Other paths are served by endpoints +// when set, else get an empty object (the SDK's config probe). type fakeModelServing struct { - t *testing.T - mu sync.Mutex - responses []string - statuses []string - bodies []registerModelPackageRequest - statusRequests []stagingStatusRequest - endpoints http.HandlerFunc + t *testing.T + mu sync.Mutex + responses []string + statuses []string + finalizations []string + bodies []registerModelPackageRequest + statusRequests []stagingStatusRequest + finalizeRequests []finalizeModelPackageRequest + endpoints http.HandlerFunc } func newFakeModelServing(t *testing.T, responses ...string) *fakeModelServing { @@ -64,6 +70,19 @@ func (f *fakeModelServing) ServeHTTP(w http.ResponseWriter, r *http.Request) { statuses = []string{statusSucceeded} } f.respond(w, statuses, len(f.statusRequests)) + case finalizeModelPackagePath: + var body finalizeModelPackageRequest + if !f.decode(w, r, &body) { + return + } + f.mu.Lock() + defer f.mu.Unlock() + f.finalizeRequests = append(f.finalizeRequests, body) + finalizations := f.finalizations + if len(finalizations) == 0 { + finalizations = []string{finalizedResponse} + } + f.respond(w, finalizations, len(f.finalizeRequests)) default: if f.endpoints != nil { f.endpoints(w, r) @@ -93,6 +112,9 @@ func (f *fakeModelServing) respond(w http.ResponseWriter, responses []string, n case serverErrorResponse: w.WriteHeader(http.StatusInternalServerError) _, _ = w.Write([]byte(`{"error_code": "INTERNAL_ERROR", "message": "Something went wrong."}`)) + case notFinalizableResponse: + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error_code": "BAD_REQUEST", "message": "Version 3 of model main.models.qwen can't be finalized yet: the checkpoint's weights haven't been copied into it."}`)) default: _, _ = w.Write([]byte(body)) } @@ -113,6 +135,8 @@ func (f *fakeModelServing) requests() []registerModelPackageRequest { const ( stagingInProgress = `{"model_id": "m-1", "registered_model": "main.models.qwen", "model_version": "3", "ready": false, "staging_status": {"state": "IN_PROGRESS", "model_artifact": {"state": "IN_PROGRESS"}, "environment_artifact": {"state": "SUCCEEDED"}}}` + stagingSucceeded = `{"model_id": "m-1", "registered_model": "main.models.qwen", "model_version": "3", "ready": false, + "staging_status": {"state": "SUCCEEDED", "model_artifact": {"state": "SUCCEEDED"}, "environment_artifact": {"state": "SUCCEEDED"}}}` stagingReady = `{"model_id": "m-1", "registered_model": "main.models.qwen", "model_version": "3", "ready": true, "staging_status": {"state": "SUCCEEDED", "model_artifact": {"state": "SUCCEEDED"}, "environment_artifact": {"state": "SUCCEEDED"}}}` stagingFailed = `{"model_id": "m-1", "registered_model": "main.models.qwen", "model_version": "3", "ready": false, @@ -121,6 +145,9 @@ const ( statusInProgress = `{"status": {"state": "IN_PROGRESS", "model_artifact": {"state": "IN_PROGRESS"}, "environment_artifact": {"state": "SUCCEEDED"}}}` statusSucceeded = `{"status": {"state": "SUCCEEDED", "model_artifact": {"state": "SUCCEEDED"}, "environment_artifact": {"state": "SUCCEEDED"}}}` statusFailed = `{"status": {"state": "FAILED", "model_artifact": {"state": "FAILED"}, "environment_artifact": {"state": "SUCCEEDED"}}}` + + finalizedResponse = `{"model_id": "m-1", "registered_model": "main.models.qwen", "model_version": "3", + "staging_status": {"state": "IN_PROGRESS", "model_artifact": {"state": "SUCCEEDED"}, "environment_artifact": {"state": "IN_PROGRESS"}}}` ) // fastRetries shortens the pause before air deploy repeats a failed request. @@ -145,14 +172,14 @@ func registerTestCheckpoint(t *testing.T, fake *fakeModelServing, req registerMo t.Helper() srv := fake.start() var states []string - resp, err := registerUntilReady(fastRetries(t.Context()), newTestWorkspaceClient(t, srv.URL), req, time.Millisecond, func(r *registerModelPackageResponse) { + resp, err := registerAndFinalize(fastRetries(t.Context()), newTestWorkspaceClient(t, srv.URL), req, time.Millisecond, func(r *registerModelPackageResponse) { states = append(states, r.weightsState()) }) return resp, states, err } -func TestRegisterUntilReadyPollsStagingThenFinalizes(t *testing.T) { - fake := newFakeModelServing(t, stagingInProgress, stagingReady) +func TestRegisterAndFinalizePollsStagingThenFinalizes(t *testing.T) { + fake := newFakeModelServing(t, stagingInProgress) fake.statuses = []string{statusInProgress, statusSucceeded} req := checkpointRequest() req.RetryFailedStaging = true @@ -161,51 +188,136 @@ func TestRegisterUntilReadyPollsStagingThenFinalizes(t *testing.T) { require.NoError(t, err) assert.True(t, resp.Ready) assert.Equal(t, "3", resp.ModelVersion) - assert.Equal(t, []string{"IN_PROGRESS", "IN_PROGRESS", "SUCCEEDED", "SUCCEEDED"}, states) + assert.Equal(t, []string{"IN_PROGRESS", "IN_PROGRESS", "SUCCEEDED"}, states) - // The staging status is polled for the registered version; the registration is - // repeated once the weights are staged, which finalizes the version. + // One registration; the status is polled for the registered version until the + // weights are copied, and the version is then finalized. + assert.Equal(t, []registerModelPackageRequest{req}, fake.requests()) assert.Equal(t, []stagingStatusRequest{{Model: "main.models.qwen", ModelVersion: "3"}, {Model: "main.models.qwen", ModelVersion: "3"}}, fake.statusRequests) - bodies := fake.requests() - require.Len(t, bodies, 2) - assert.Equal(t, req, bodies[0]) - // Only the first request restages; the repeat continues the same registration. - req.RetryFailedStaging = false - assert.Equal(t, req, bodies[1]) + assert.Equal(t, []finalizeModelPackageRequest{{RegisteredModel: "main.models.qwen", ModelVersion: "3"}}, fake.finalizeRequests) } -func TestRegisterUntilReadyReportsFailedStaging(t *testing.T) { - fake := newFakeModelServing(t, stagingInProgress, stagingFailed) +func TestRegisterAndFinalizeReportsTheFinalizedStagingStatus(t *testing.T) { + fake := newFakeModelServing(t, stagingSucceeded) + + resp, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) + require.NoError(t, err) + // The environment is still being staged when the version is finalized. + assert.Equal(t, stagingStateInProgress, resp.environmentState()) + assert.Empty(t, fake.statusRequests) +} + +func TestRegisterAndFinalizeReturnsAReadyVersion(t *testing.T) { + fake := newFakeModelServing(t, stagingReady) + + resp, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) + require.NoError(t, err) + assert.True(t, resp.Ready) + assert.Empty(t, fake.statusRequests) + assert.Empty(t, fake.finalizeRequests) +} + +func TestRegisterAndFinalizeReportsFailedStaging(t *testing.T) { + fake := newFakeModelServing(t, stagingInProgress) fake.statuses = []string{statusFailed} resp, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) require.ErrorIs(t, err, errStagingFailed) assert.Equal(t, "3", resp.ModelVersion) - // The failure the status reports is confirmed by the registration itself. - assert.Len(t, fake.requests(), 2) + assert.Len(t, fake.requests(), 1) + assert.Empty(t, fake.finalizeRequests) } -func TestRegisterUntilReadyRetriesTransientErrors(t *testing.T) { - fake := newFakeModelServing(t, serverErrorResponse, stagingInProgress, stagingReady) - fake.statuses = []string{serverErrorResponse, statusSucceeded} +func TestRegisterAndFinalizeReportsAFailureTheRegistrationReports(t *testing.T) { + fake := newFakeModelServing(t, stagingFailed) + + _, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) + require.ErrorIs(t, err, errStagingFailed) + assert.Empty(t, fake.statusRequests) + assert.Empty(t, fake.finalizeRequests) +} + +func TestRegisterAndFinalizeReturnsWhyTheServerWontFinalize(t *testing.T) { + notStaged := strings.ReplaceAll(stagingInProgress, "IN_PROGRESS", "NOT_STAGED") + fake := newFakeModelServing(t, notStaged) + fake.finalizations = []string{notFinalizableResponse} resp, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) + _, ok := errors.AsType[*finalizeError](err) + assert.True(t, ok) + assert.ErrorContains(t, err, "haven't been copied into it") + assert.False(t, resp.Ready) + // A version with no staging record isn't polled; the server says what to do. + assert.Empty(t, fake.statusRequests) + assert.Len(t, fake.finalizeRequests, 1) +} + +func TestRegisterAndFinalizeWaitsOutAStatusWithoutAState(t *testing.T) { + fake := newFakeModelServing(t, stagingInProgress) + fake.statuses = []string{`{}`, `{"status": {"state": "IN_PROGRESS"}}`, statusSucceeded} + + resp, states, err := registerTestCheckpoint(t, fake, checkpointRequest()) require.NoError(t, err) assert.True(t, resp.Ready) - assert.Len(t, fake.requests(), 3) - assert.Len(t, fake.statusRequests, 2) + assert.Equal(t, []string{"IN_PROGRESS", "IN_PROGRESS", "IN_PROGRESS", "SUCCEEDED"}, states) + assert.Len(t, fake.statusRequests, 3) + assert.Len(t, fake.finalizeRequests, 1) +} + +func TestRegisterAndFinalizeKeepsAnEnvironmentStateTheFinalizeOmits(t *testing.T) { + environmentFailed := strings.Replace(stagingSucceeded, `"environment_artifact": {"state": "SUCCEEDED"}`, `"environment_artifact": {"state": "FAILED"}`, 1) + fake := newFakeModelServing(t, environmentFailed) + fake.finalizations = []string{`{"model_version": "3", "staging_status": {"model_artifact": {"state": "SUCCEEDED"}}}`} + + resp, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) + require.NoError(t, err) + assert.Equal(t, stagingStateFailed, resp.environmentState()) +} + +func TestRegisterAndFinalizeRejectsAnotherFinalizedVersion(t *testing.T) { + fake := newFakeModelServing(t, stagingSucceeded) + fake.finalizations = []string{`{"model_version": "4"}`} + + _, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) + _, ok := errors.AsType[*finalizeError](err) + assert.True(t, ok) + assert.ErrorContains(t, err, `the server finalized version "4" instead`) +} + +func TestRegisterAndFinalizeRejectsARegistrationWithoutAVersion(t *testing.T) { + fake := newFakeModelServing(t, `{"model_id": "m-1", "registered_model": "main.models.qwen", "ready": false}`) + + _, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) + assert.ErrorContains(t, err, "registering a version of main.models.qwen returned no version") + assert.Empty(t, fake.finalizeRequests) } -func TestRegisterUntilReadyDoesNotRepeatANotReadyModelFromCode(t *testing.T) { +func TestRegisterAndFinalizeFinalizesAModelFromCodeWithoutPolling(t *testing.T) { fake := newFakeModelServing(t, `{"model_id": "m-2", "registered_model": "a.b.c", "model_version": "1", "ready": false}`) + fake.finalizations = []string{`{"model_id": "m-2", "registered_model": "a.b.c", "model_version": "1"}`} req := registerModelPackageRequest{RegisteredModel: "a.b.c", IdempotencyKey: "k", Code: &modelPackageCode{EntryFile: "agent.py"}} - _, _, err := registerTestCheckpoint(t, fake, req) - assert.ErrorContains(t, err, "version 1 of a.b.c isn't ready after registering") - assert.Len(t, fake.requests(), 1) + resp, _, err := registerTestCheckpoint(t, fake, req) + require.NoError(t, err) + assert.True(t, resp.Ready) + assert.Empty(t, fake.statusRequests) + assert.Equal(t, []finalizeModelPackageRequest{{RegisteredModel: "a.b.c", ModelVersion: "1"}}, fake.finalizeRequests) +} + +func TestRegisterAndFinalizeRetriesTransientErrors(t *testing.T) { + fake := newFakeModelServing(t, serverErrorResponse, stagingInProgress) + fake.statuses = []string{serverErrorResponse, statusSucceeded} + fake.finalizations = []string{serverErrorResponse, finalizedResponse} + + resp, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) + require.NoError(t, err) + assert.True(t, resp.Ready) + assert.Len(t, fake.requests(), 2) + assert.Len(t, fake.statusRequests, 2) + assert.Len(t, fake.finalizeRequests, 2) } -func TestRegisterUntilReadyGivesUpOnPersistentErrors(t *testing.T) { +func TestRegisterAndFinalizeGivesUpOnPersistentErrors(t *testing.T) { fake := newFakeModelServing(t, serverErrorResponse) _, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) @@ -213,7 +325,7 @@ func TestRegisterUntilReadyGivesUpOnPersistentErrors(t *testing.T) { assert.Len(t, fake.requests(), retry.MaxRetries+1) } -func TestRegisterUntilReadyReturnsServerErrors(t *testing.T) { +func TestRegisterAndFinalizeReturnsServerErrors(t *testing.T) { fake := newFakeModelServing(t, conflictResponse) _, _, err := registerTestCheckpoint(t, fake, checkpointRequest()) @@ -221,12 +333,14 @@ func TestRegisterUntilReadyReturnsServerErrors(t *testing.T) { assert.Len(t, fake.requests(), 1) } -func TestRegisterUntilReadyStopsWhenCanceled(t *testing.T) { - srv := newFakeModelServing(t, stagingInProgress).start() +func TestRegisterAndFinalizeStopsWhenCanceled(t *testing.T) { + fake := newFakeModelServing(t, stagingInProgress) + srv := fake.start() ctx, cancel := context.WithCancel(t.Context()) - _, err := registerUntilReady(ctx, newTestWorkspaceClient(t, srv.URL), checkpointRequest(), time.Hour, func(*registerModelPackageResponse) { cancel() }) + _, err := registerAndFinalize(ctx, newTestWorkspaceClient(t, srv.URL), checkpointRequest(), time.Hour, func(*registerModelPackageResponse) { cancel() }) assert.ErrorIs(t, err, context.Canceled) + assert.Empty(t, fake.finalizeRequests) } func TestRegisterModelPackageRequestJSON(t *testing.T) {