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..4d2b0d876f5 --- /dev/null +++ b/cmd/air/deploy.go @@ -0,0 +1,456 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "net/http" + "path/filepath" + "regexp" + "strings" + "sync" + "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") +) + +// 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"` + 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 + // 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 +} + +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 { + opts.idempotencyKeySet = cmd.Flags().Changed("idempotency-key") + 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 withCode(codeInvalidArgs, false, err) + } + if err := validateDeployFlags(cfg, opts); err != nil { + return withCode(codeInvalidArgs, false, err) + } + + var code *modelPackageCode + if cfg.Model.Type == modelTypeCode { + code, err = collectModelCode(ctx, cfg, opts.file) + if err != nil { + return withCode(codeInvalidArgs, false, 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 { + 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) + 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.", 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.") + } + } + } else { + result.ModelVersion = *cfg.Model.Version + } + + entity := cfg.servedEntity(result.ModelVersion) + _, 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, cfg.Endpoint) + if jsonOut { + return renderEnvelope(ctx, result) + } + out := cmd.OutOrStdout() + 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 +} + +// 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 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 + // 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 + + progress := func(*registerModelPackageResponse) {} + if show { + sp := cmdio.NewSpinner(ctx) + defer sp.Close() + progress = func(resp *registerModelPackageResponse) { + sp.Update(registrationProgress(cfg, resp)) + } + } + resp, err := registerAndFinalize(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, 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 { + continuing := "" + 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. + if _, retryable := deployErrorCode(err); retryable { + return nil, fmt.Errorf("%w\n\nTo try again, run:\n %s", failed, resume) + } + return nil, failed +} + +// 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. +// 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 codeInterrupted, true + case errors.Is(err, errDeployTimedOut): + return codeTimeout, true + } + 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 codeInternalError, true +} + +// 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.idempotencyKeySet { + 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() + 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)) +} + +// 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..05cbdab8d6b --- /dev/null +++ b/cmd/air/deploy_test.go @@ -0,0 +1,469 @@ +package aircmd + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "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/retry" + "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(fastRetries(t.Context()), time.Minute) + t.Cleanup(cancel) + return ctx +} + +// 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...) + 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])) + } + + // 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)) + } + 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) + + 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, 1) + assert.Equal(t, "key-1", bodies[0].IdempotencyKey) + 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) + 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) { + 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 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 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, 1) + assert.True(t, bodies[0].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), conflictResponse) + + 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), conflictResponse) + + 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 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) + 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) { + 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} + 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"}, + {"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) { + 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 { + name string + err error + code string + retryable bool + }{ + {"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.name) + assert.Equal(t, tt.retryable, retryable, tt.name) + } +} + +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..721fadd7825 --- /dev/null +++ b/cmd/air/deploycode.go @@ -0,0 +1,191 @@ +package aircmd + +import ( + "bytes" + "context" + "errors" + "fmt" + "os" + "os/exec" + "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 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) + 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 { + 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 +} + +// 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) + } + 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: %s", err, strings.TrimSpace(stderr.String())) + } + 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 { + // 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..8fa2fa7e0ab --- /dev/null +++ b/cmd/air/deploycode_test.go @@ -0,0 +1,190 @@ +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 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 file creds.json matches a .gitignore rule but is tracked by git") +} + +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", "tools/a.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..9c3bdac4cf0 --- /dev/null +++ b/cmd/air/deployconfig.go @@ -0,0 +1,417 @@ +package aircmd + +import ( + "bytes" + "context" + "errors" + "fmt" + "maps" + "reflect" + "regexp" + "slices" + "strconv" + "strings" + "sync" + "unicode/utf16" + + "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]*$`) }) + +// 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$`) }) + +// 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. 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, 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"` + + 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:"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. +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 !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") + } + 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 !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) + } + 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/") || 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) + } + } + 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 := serverLength(*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 { + // 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, 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(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} + 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 +} + +// 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 + 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..722d1acb350 --- /dev/null +++ b/cmd/air/deployconfig_test.go @@ -0,0 +1,313 @@ +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: . + include_paths: [agent.py, tools] +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 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 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"}, + {"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 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, 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"}, + {"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..1effc2a8303 --- /dev/null +++ b/cmd/air/deployendpoint.go @@ -0,0 +1,230 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "maps" + "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/retries" + "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 := 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, 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 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)) + } + }) +} + +// 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. 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 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 +} + +// 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 = waitNotUpdating(ctx, w, name, 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 := 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. +// 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. +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", name, describeEntities(currentEntities(endpoint))) + } + if endpoint.State == nil || endpoint.State.Ready != serving.EndpointStateReadyReady { + 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 +} + +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 { + name := stripControl(s.EntityName) + if s.EntityVersion == "" { + names = append(names, name) + continue + } + names = append(names, fmt.Sprintf("version %s of %s", stripControl(s.EntityVersion), name)) + } + 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..0a4dc57e019 --- /dev/null +++ b/cmd/air/deployendpoint_test.go @@ -0,0 +1,308 @@ +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 + // 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 { + 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 { + 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) + _, _ = 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) +} + +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 TestServeOnEndpointWaitsThroughTransientFailures(t *testing.T) { + fake := newFakeEndpoints(t, nil) + // 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.Equal(t, 6, fake.gets) +} + +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..f3f24c99693 --- /dev/null +++ b/cmd/air/deployregister.go @@ -0,0 +1,292 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "net/http" + "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" +) + +// 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" +) + +// 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"` +} + +// 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"` + 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 { + 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") + +// 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 { + 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 +} + +// 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 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 + } + } + 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 + } + if merged.EnvironmentArtifact == nil { + merged.EnvironmentArtifact = older.EnvironmentArtifact + } + } + return &merged +} + +// 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 +} + +// 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 + } + var status stagingStatusResponse + if err := p.post(ctx, stagingStatusPath, statusReq, &status); err != nil { + return err + } + resp.StagingStatus = mergeStagingStatus(resp.StagingStatus, status.Status) + progress(resp) + } +} + +// 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. 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) + }) + return err +} diff --git a/cmd/air/deployregister_test.go b/cmd/air/deployregister_test.go new file mode 100644 index 00000000000..2efeb1a9028 --- /dev/null +++ b/cmd/air/deployregister_test.go @@ -0,0 +1,368 @@ +package aircmd + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/retry" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// Canned fake responses that stand for an error rather than a body. +const ( + conflictResponse = "conflict" + serverErrorResponse = "server error" + notFinalizableResponse = "not finalizable" +) + +// 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 + finalizations []string + bodies []registerModelPackageRequest + statusRequests []stagingStatusRequest + finalizeRequests []finalizeModelPackageRequest + 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) { + 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)) + 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) + return + } + _, _ = w.Write([]byte(`{}`)) + } +} + +// 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) + if !assert.NoError(f.t, err) || !assert.NoError(f.t, json.Unmarshal(raw, v)) { + w.WriteHeader(http.StatusInternalServerError) + return false + } + return true +} + +// 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."}`)) + 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)) + } +} + +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"}}}` + 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, + "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"}}}` + + 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. +func fastRetries(ctx context.Context) context.Context { + return env.Set(ctx, retry.IntervalMsVariable, "1") +} + +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 registerTestCheckpoint(t *testing.T, fake *fakeModelServing, req registerModelPackageRequest) (*registerModelPackageResponse, []string, error) { + t.Helper() + srv := fake.start() + var states []string + 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 TestRegisterAndFinalizePollsStagingThenFinalizes(t *testing.T) { + fake := newFakeModelServing(t, stagingInProgress) + 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) + + // 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) + assert.Equal(t, []finalizeModelPackageRequest{{RegisteredModel: "main.models.qwen", ModelVersion: "3"}}, fake.finalizeRequests) +} + +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) + assert.Len(t, fake.requests(), 1) + assert.Empty(t, fake.finalizeRequests) +} + +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.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 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"}} + + 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 TestRegisterAndFinalizeGivesUpOnPersistentErrors(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 TestRegisterAndFinalizeReturnsServerErrors(t *testing.T) { + fake := newFakeModelServing(t, conflictResponse) + + _, _, 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 TestRegisterAndFinalizeStopsWhenCanceled(t *testing.T) { + fake := newFakeModelServing(t, stagingInProgress) + srv := fake.start() + ctx, cancel := context.WithCancel(t.Context()) + + _, 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) { + 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) { 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) +}