[erarpitagarwal]

Changes

Add a hidden databricks air deploy command. It registers a model in Unity Catalog and serves it on a Model Serving endpoint, from one YAML file.

• model.type picks the source:

  • checkpoint: weights in a volume, plus the command that serves them.
  • code: a Python file, with the files it imports from code_source, sent with the request.
  • saved_model_dir: a model directory saved by MLflow, in a volume.
  • registered: an existing model version.

• To register, it makes three calls:

  1. POST /api/2.0/serving-endpoints:registerModelPackage with an idempotency key, once. It logs the model, creates a pending version and starts copying a checkpoint's weights into it.
  2. The read-only serving-endpoints:checkModelDeploymentStagingStatus, polled while the weights copy.
  3. serving-endpoints:finalizeModelPackage, which finalizes the version once the server sees the weights copied.

  A version an earlier run finalized comes back ready from the registration, and isn't finalized again. Transient server errors are retried per request with libs/retry, including each poll of the endpoint, so failures spread over a long wait don't end it. On Ctrl-C or --timeout, it prints the command that continues. --retry-failed-staging copies the weights again after a failure.

• It then creates the endpoint, or updates it to serve the new version, and waits until it is ready. It only replaces a single version of the same model, and it checks this before registering, so a wrong endpoint name fails before any weights are copied. An endpoint that already serves the version is left alone. An update that an interrupted run started is waited for, not repeated, and so is anyone else's update.

• It reuses the air run YAML loader (--override, -h config.<field>, rejecting unknown fields), the code_source file selection and root path resolution, and the idempotency token helper. The shared helpers now take the schema, or the token, as a parameter. air run behaves as before, except that one error now reads "config must be a YAML mapping" instead of "run config must be a YAML mapping".

Why

Serving a checkpoint or an agent today means calling mlflow.<flavor>.log_model from Python: pickling a model, copying weights into it, and packing an environment. With air deploy, the user hands over text (YAML and .py files) and the files a training job already wrote to a volume. The server writes the MLmodel and builds the environment from the declared runtime and dependencies. The YAML keeps the air run names for the fields the two share: environment, code_source, env_variables, secrets and idempotency_token.

The command is hidden, with no changelog entry, until the registration API is on in every workspace. The API is off by default while it rolls out. Some decisions:

• The registration key is the only state the CLI keeps. Repeating the same registration request is how a run resumes, so the CLI builds the request once and a later run re-sends it unchanged. Code files are read in git ls-files order, so a later run sends the same bytes.
• A model from code is stored in the model version, where anyone who can read the model can read it. So code_source.include_paths is required and must name files and directories literally (no patterns), and within it files git ignores, hidden files and directories, and __pycache__ aren't sent; the server can't place names that start with . anyway. A file git tracks although .gitignore matches it is refused. --dry-run lists every file it would send. This also fixes the shared code_source listing outside git, which listed paths relative to the working directory instead of root_path (it affected air run too).
• The registration and finalize APIs share one server flag, and the staging status API has its own. All are off by default while they roll out, and must be on together for air deploy to work.
• The idempotency key is printed to stderr before the first request in every output mode, so a killed run can always be continued. In JSON mode, errors are JSON error envelopes with the server's error code, or INVALID_ARGS, INTERRUPTED, TIMEOUT, STAGING_FAILED or INTERNAL_ERROR.
• The CLI accepts code and saved_model_dir now, but the server registers only checkpoints until its follow-up lands, and until then answers the other two with "isn't supported yet". The CLI's checks follow the API's documented limits: checkpoint tasks, workload types (from the SDK), and the path, command, file and size caps.

Tests

• go test ./cmd/air/ passes, and so does -race on the new tests. golangci-lint reports 0 issues on cmd/air. ./task checks passes, and GOOS=windows go build ./cmd/air/ builds. New unit tests cover:
  • the YAML validation and code selection
  • registering, status polling and finalizing, including retries, staging failure, a version that is already ready, a refused finalize, timeouts and cancellation
  • the endpoint create, update and refusal paths, and waits through transient failures
  • resume commands and JSON error codes
• End to end with the register, poll and finalize flow, against a staging workspace with the APIs on, deploying a Qwen3-1.7B checkpoint on GPU_MEDIUM:
  • One registration request, 5 status polls 15 s apart while the weights copied (77 s), one finalize request, then the endpoint was created. It became READY about 8.5 minutes later and answered a chat request.
  • Running the same command again took 2 s: the registration came back ready, with no poll, no finalize and no endpoint update. The endpoint was deleted afterwards.
  • One test-environment step on the server side was done by hand between the first two runs; it doesn't involve the CLI.
• Earlier end-to-end runs, before the finalize call, with the same model and workspace and the same hand step, covered the rest:
  • --dry-run validated the config without credentials.
  • --timeout 45s stopped while the weights were copying and printed the resume command. Running it registered the version.
  • Ctrl-C during the endpoint start stopped with the resume command. Running it waited on the create already in progress instead of repeating it. The endpoint became READY and answered a chat request.
  • Running the same command again, and model.type: registered with that version, changed nothing and took 1 s. A changed env var updated the endpoint to config version 2, which became READY.
  • With another model's name, air deploy refused the endpoint before registering anything.
  • model.type: code sent its files, and the server answered "isn't supported yet", as expected until its follow-up.
• There are no acceptance tests: air has none yet, the command is hidden, and the test server doesn't model these APIs.

This pull request and its description were written by Isaac.

[github-actions]

Thanks for the contribution! This pull request was opened from a fork.

If you're a member of the Databricks organization, please request access to the cli-contrib team and open your pull request from a branch in this repository instead of a fork — fork PRs don't run integration tests automatically and have reduced CI access.

If you're an external contributor, no action is needed; a maintainer will review your pull request.

[github-actions]

An authorized user can trigger integration tests manually by following the instructions below:

Trigger:
go/deco-tests-run/cli

Inputs:

• PR number: 7642
• Commit SHA: 261b2028fb2eab0149eb41b31b90dfb1ec6daa68

Checks will be approved automatically on success.