[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 calls POST /api/2.0/serving-endpoints:registerModelPackage with an idempotency key. It repeats the same request until the version is ready, because the server finalizes a checkpoint's version once its weights are copied. 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 request is how a run resumes, so the CLI builds the request once and 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 files git ignores aren't sent, and nor are hidden files and directories or __pycache__; the server can't place names that start with . anyway. Outside a git repository, code_source.include_paths is required. --dry-run lists every file it would send.
• 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

Running now: go test ./cmd/air/, ./task lint, and an end-to-end run of the built binary against a staging workspace that has the registration API. Results follow here.

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: bcdf2ddd195d29051543d5b6ee174fe96242060a

Checks will be approved automatically on success.