AIR CLI Integration: Batch Merge Phase 3

acceptance/experimental/air/cancel/test.toml

@@ -1,6 +1,3 @@
-# This command does not deploy a bundle, so no engine matrix is needed.
-EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
-
 # The SDK occasionally probes host reachability with a HEAD request; stub it so
 # the test is deterministic.
 [[Server]]

acceptance/experimental/air/config-help/output.txt

@@ -0,0 +1,110 @@
+
+=== command help
+>>> [CLI] experimental air run --help
+Submit a training workload to Databricks serverless GPU compute.
+
+The workload is described by a YAML config file (see --file).
+
+To look up a config field, pass its path to -h:
+
+  databricks experimental air run -h config
+  databricks experimental air run -h config.compute
+  databricks experimental air run -h config.compute.accelerator_type
+
+The path must be a separate argument: cobra reserves -h as a boolean, so
+-h=config.compute and -hconfig.compute are not accepted.
+
+Usage:
+  databricks experimental air run [flags]
+
+Flags:
+      --dry-run                  Validate the config without submitting
+  -f, --file string              Path to the workload YAML config
+  -h, --help                     help for run
+      --idempotency-key string   Return the existing run if this key was already used
+      --override stringArray     Override a YAML field, e.g. compute.num_accelerators=8 (repeatable)
+      --watch                    Stream logs until the run completes
+
+Global Flags:
+      --debug            enable debug logging
+  -o, --output type      output type: text or json (default text)
+  -p, --profile string   ~/.databrickscfg profile
+  -t, --target string    bundle target to use (if applicable)
+
+=== schema overview
+>>> [CLI] experimental air run -h config
+config
+  The run YAML schema. Pass a field path for details, e.g. config.compute.accelerator_type.
+
+  Fields:
+    experiment_name              (required) Name of the experiment.
+    compute                      (required) Which accelerators to run on and how many.
+    environment                  Python dependencies, or a custom Docker image, for the run's runtime.
+    command                      (required) Shell command that starts the workload.
+    env_variables                Plain environment variables, as NAME: value.
+    secrets                      Environment variables sourced from secrets, as NAME: scope/key.
+    code_source                  Local code to upload and make available to the run.
+    max_retries                  How many times to retry a failed run.
+    timeout_minutes              Wall-clock limit for the run in minutes.
+    idempotency_token            Reuse token: a repeat submission with the same token returns the existing run instead of starting another.
+    parameters                   Free-form values passed through to the workload.
+    mlflow_run_name              Name for the MLflow run.
+    mlflow_experiment_directory  Workspace directory holding the MLflow experiment.
+    permissions                  Who may view or manage the run, as a list of principal plus level grants.
+    usage_policy_name            Usage policy to bill the run to, by name.
+    usage_policy_id              Usage policy to bill the run to, by id.
+
+Use "-h config.<field>" for details on a field.
+
+=== nested object lists its fields
+>>> [CLI] experimental air run -h config.compute
+config.compute
+  Which accelerators to run on and how many.
+
+  Fields:
+    num_accelerators  Total number of GPUs to allocate.
+    accelerator_type  Which accelerator to run on, e.g. GPU_1xA10.
+
+Use "-h config.compute.<field>" for details on a field.
+
+=== leaf field
+>>> [CLI] experimental air run -h config.compute.accelerator_type
+config.compute.accelerator_type
+  Which accelerator to run on, e.g. GPU_1xA10. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for the current list of supported GPU types. Matched case-sensitively.
+
+  Type:     string
+  Required: no
+
+=== conditionally required leaf
+>>> [CLI] experimental air run -h config.environment.docker_image.url
+config.environment.docker_image.url
+  Fully qualified image URL, e.g. myregistry.io/team/train:v3.
+
+  Type:     string
+  Required: when environment.docker_image is set
+
+=== union field reports both accepted shapes
+>>> [CLI] experimental air run -h config.code_source.snapshot.git.remote
+config.code_source.snapshot.git.remote
+  No longer supported: the snapshot archives your local copy. Only false is accepted; use commit to pin a revision.
+
+  Type:     bool or string
+  Required: no
+
+=== the config. prefix is optional
+>>> [CLI] experimental air run -h compute.num_accelerators
+config.compute.num_accelerators
+  Total number of GPUs to allocate. Must be a positive multiple of the accelerator type's per-node GPU count. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for supported GPU types.
+
+  Type:     int
+  Required: no
+
+=== unknown field suggests the near match
+>>> [CLI] experimental air run -h config.compute.acclerator_type
+Error: unknown config field "config.compute.acclerator_type"; did you mean "accelerator_type"?
+
+fields under "config.compute" are: accelerator_type, num_accelerators
+
+=== free-form map keys are not schema fields
+>>> [CLI] experimental air run -h config.parameters.learning_rate
+Error: "config.parameters" holds user-defined keys, so "learning_rate" is not part of the schema; see "config.parameters" instead

acceptance/experimental/air/config-help/script

@@ -0,0 +1,31 @@
+# Pin the rendered config field help. Any change to a `help:` struct tag on the
+# run YAML schema shows up as a diff here.
+
+# Plain -h must keep documenting the command itself, and advertise the config
+# path syntax so the feature is discoverable.
+title "command help"
+trace $CLI experimental air run --help
+
+title "schema overview"
+trace $CLI experimental air run -h config
+
+title "nested object lists its fields"
+trace $CLI experimental air run -h config.compute
+
+title "leaf field"
+trace $CLI experimental air run -h config.compute.accelerator_type
+
+title "conditionally required leaf"
+trace $CLI experimental air run -h config.environment.docker_image.url
+
+title "union field reports both accepted shapes"
+trace $CLI experimental air run -h config.code_source.snapshot.git.remote
+
+title "the config. prefix is optional"
+trace $CLI experimental air run -h compute.num_accelerators
+
+title "unknown field suggests the near match"
+trace $CLI experimental air run -h config.compute.acclerator_type
+
+title "free-form map keys are not schema fields"
+trace $CLI experimental air run -h config.parameters.learning_rate

acceptance/experimental/air/get-ai-runtime/output.txt

@@ -27,7 +27,7 @@
 │  MLflow Run    my-run                                          │
 │  User          user@example.com                                │
 │  Accelerators  1x A10                                          │
-│  Environment   N/A                                             │
+│  Environment   4                                               │
 │                                                                │
 ╰────────────────────────────────────────────────────────────────╯
 

acceptance/experimental/air/get-ai-runtime/test.toml

@@ -1,6 +1,3 @@
-# This command does not deploy a bundle, so no engine matrix is needed.
-EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
-
 # On Windows, Git Bash rewrites the leading-/ workspace paths passed to
 # `workspace mkdirs`/`import` into C:/... paths; disable that conversion.
 [Env]
@@ -12,8 +9,10 @@ MSYS_NO_PATHCONV = "1"
 Pattern = "HEAD /"
 Response.Body = ''
 
-# The typed SDK GetRun response: an ai_runtime_task run has no gen_ai_compute_task,
-# so the task comes back empty (the SDK has no field for ai_runtime_task).
+# The GetRun response for an ai_runtime_task run. It has no gen_ai_compute_task,
+# so the Environment cell instead comes from the run's environments[].spec (keyed
+# by the task's environment_key), which the typed SDK Run drops — `air get` reads
+# it with a raw request, so this route is hit twice (typed GetRun + that request).
 [[Server]]
 Pattern = "GET /api/2.2/jobs/runs/get"
 Response.Body = '''
@@ -24,12 +23,16 @@ Response.Body = '''
   "start_time": 1700000000000,
   "end_time": 1700000012000,
   "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"},
+  "environments": [
+    {"environment_key": "default", "spec": {"environment_version": "4"}}
+  ],
   "tasks": [
     {
       "task_key": "train",
       "run_id": 456,
       "attempt_number": 0,
       "max_retries": 3,
+      "environment_key": "default",
       "ai_runtime_task": {
         "experiment": "my-exp",
         "deployments": [

acceptance/experimental/air/get/output.txt

@@ -27,7 +27,7 @@
 │  MLflow Run    my-run                                          │
 │  User          user@example.com                                │
 │  Accelerators  1x A10                                          │
-│  Environment   ml-runtime-gpu:1.0                              │
+│  Environment   N/A                                             │
 │                                                                │
 ╰────────────────────────────────────────────────────────────────╯
 

acceptance/experimental/air/get/test.toml

@@ -1,6 +1,3 @@
-# This command does not deploy a bundle, so no engine matrix is needed.
-EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
-
 # The SDK occasionally probes host reachability with a HEAD request; stub it so
 # the test is deterministic.
 [[Server]]

acceptance/experimental/air/help/output.txt

@@ -47,3 +47,24 @@ Global Flags:
   -o, --output type      output type: text or json (default text)
   -p, --profile string   ~/.databrickscfg profile
   -t, --target string    bundle target to use (if applicable)
+
+=== logs help
+>>> [CLI] experimental air logs --help
+Stream logs from an active run, or fetch logs from a completed run.
+
+Usage:
+  databricks experimental air logs JOB_RUN_ID [flags]
+
+Flags:
+      --download-to string   Download all logs to this directory instead of printing
+  -h, --help                 help for logs
+      --lines int            For completed runs, print the last N lines (default 10000)
+      --minutes int          Fetch only logs from the last N minutes
+      --node int             Fetch logs from this node
+      --retry int            View logs from a specific retry attempt; -1 means latest (default -1)
+
+Global Flags:
+      --debug            enable debug logging
+  -o, --output type      output type: text or json (default text)
+  -p, --profile string   ~/.databrickscfg profile
+  -t, --target string    bundle target to use (if applicable)

acceptance/experimental/air/help/script

@@ -6,3 +6,6 @@ trace $CLI experimental air --help
 
 title "list help"
 trace $CLI experimental air list --help
+
+title "logs help"
+trace $CLI experimental air logs --help

acceptance/experimental/air/list/output.txt

@@ -1,8 +1,8 @@
 
 === list (text)
 >>> [CLI] experimental air list
-   Run ID           Experiment  Status     Started              Duration  MLflow              User                   Accelerators
-   [NUMID]  qwen-train  ● SUCCESS  [TIMESTAMP]       12s  …/runs/run1         [USERNAME]  8x H100
+   Run ID           Experiment  Status     Started              Duration  MLflow                  User                   Accelerators
+   [NUMID]  qwen-train  ● SUCCESS  [TIMESTAMP]       12s  qwen-train-001          [USERNAME]  8x H100
 
 === list (json)
 >>> [CLI] experimental air list -o json
@@ -25,8 +25,8 @@
 
 === list --all-status (text, via AiTrainingService index)
 >>> [CLI] experimental air list --all-status
-   Run ID           Experiment  Status     Started              Duration  MLflow              User                   Accelerators
-   [NUMID]  qwen-train  ● SUCCESS  [TIMESTAMP]       12s  …/runs/run1         [USERNAME]  8x H100
+   Run ID           Experiment  Status     Started              Duration  MLflow                  User                   Accelerators
+   [NUMID]  qwen-train  ● SUCCESS  [TIMESTAMP]       12s  qwen-train-001          [USERNAME]  8x H100
 
 === list --all-status (json)
 >>> [CLI] experimental air list --all-status -o json

acceptance/experimental/air/list/test.toml

@@ -1,6 +1,3 @@
-# This command does not deploy a bundle, so no engine matrix is needed.
-EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
-
 # Disable the on-disk run cache so --all-status output is deterministic across runs.
 [Env]
 DATABRICKS_CACHE_ENABLED = "false"
@@ -64,6 +61,19 @@ Response.Body = '''
 {"training_workflows": [{"job_run_id": "334747067049496", "submit_time": "2024-06-05T17:32:39Z"}]}
 '''
 
+# MLflow run names (for the MLflow column label) are fetched per AIR run (text mode).
+[[Server]]
+Pattern = "GET /api/2.0/mlflow/runs/get"
+Response.Body = '''
+{
+  "run": {
+    "info": {
+      "run_name": "qwen-train-001"
+    }
+  }
+}
+'''
+
 # runs/get hydrates one index id into the same shape as a runs/list element.
 [[Server]]
 Pattern = "GET /api/2.2/jobs/runs/get"

acceptance/experimental/air/logs-download/output.txt

@@ -0,0 +1,16 @@
+
+=== download-to reports no logs when none are available
+>>> [CLI] experimental air logs 123 --download-to dl-logs
+No logs available for run 123. Run terminated in state SUCCESS
+
+=== download-to with an out-of-range node is rejected
+>>> [CLI] experimental air logs 123 --download-to dl-logs --node 5
+Error: invalid --node 5: run has 2 node(s), indexed 0 to 1
+
+Exit code: 1
+
+=== download-to cannot be combined with --lines
+>>> [CLI] experimental air logs 123 --download-to dl-logs --lines 50
+Error: --download-to writes complete logs, so it cannot be combined with --lines or --minutes
+
+Exit code: 1

acceptance/experimental/air/logs-download/script

@@ -0,0 +1,12 @@
+# --download-to resolves the run's node count, then downloads each node's logs.
+# This run resolves no MLflow run id, so it reports no logs (the full byte
+# download is covered by unit tests, since the pre-signed URL host is dynamic).
+
+title "download-to reports no logs when none are available"
+errcode trace $CLI experimental air logs 123 --download-to dl-logs
+
+title "download-to with an out-of-range node is rejected"
+errcode trace $CLI experimental air logs 123 --download-to dl-logs --node 5
+
+title "download-to cannot be combined with --lines"
+errcode trace $CLI experimental air logs 123 --download-to dl-logs --lines 50

acceptance/experimental/air/logs-download/test.toml

@@ -0,0 +1,40 @@
+# The command creates this download directory; don't treat it as test output.
+Ignore = ["dl-logs"]
+
+# The SDK occasionally probes host reachability with a HEAD request; stub it so
+# the test is deterministic.
+[[Server]]
+Pattern = "HEAD /"
+Response.Body = ''
+
+# A completed 2-node run (GPU_1xA10 x 2 = 2 nodes).
+[[Server]]
+Pattern = "GET /api/2.2/jobs/runs/get"
+Response.Body = '''
+{
+  "run_id": 123,
+  "start_time": 1700000000000,
+  "end_time": 1700000012000,
+  "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"},
+  "tasks": [
+    {
+      "task_key": "train",
+      "run_id": 456,
+      "attempt_number": 0,
+      "ai_runtime_task": {
+        "experiment": "dl-exp",
+        "deployments": [
+          {"command_path": "/x/command.sh", "compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 2}}
+        ]
+      }
+    }
+  ]
+}
+'''
+
+# No MLflow run id resolvable, so the download reports no logs rather than
+# attempting a (host-dynamic) pre-signed artifact fetch. The full download path
+# is covered by unit tests.
+[[Server]]
+Pattern = "GET /api/2.2/jobs/runs/get-output"
+Response.Body = '{}'

acceptance/experimental/air/logs-mlflow-fallback/output.txt

@@ -0,0 +1,8 @@
+
+=== logs falls back to mlflow (no logs)
+>>> [CLI] experimental air logs 123
+No logs available for run 123. Run terminated in state SUCCESS
+
+=== logs falls back to mlflow (json)
+>>> [CLI] experimental air logs 123 -o json
+{"type":"ERROR","ts":"[TIMESTAMP]","node":0,"line":"No logs available for run 123. Run terminated in state SUCCESS"}

acceptance/experimental/air/logs-mlflow-fallback/script

@@ -0,0 +1,9 @@
+# Bricklens is gated off (FEATURE_DISABLED), so the command falls back to the
+# MLflow log path. With no MLflow run id resolvable, the fallback reports no
+# logs and exits non-zero — proving the try/catch routes to MLflow.
+
+title "logs falls back to mlflow (no logs)"
+errcode trace $CLI experimental air logs 123
+
+title "logs falls back to mlflow (json)"
+errcode trace $CLI experimental air logs 123 -o json