databricks-air 0.1.0__py3-none-any.whl

cli/__init__.py

@@ -0,0 +1,5 @@
+"""Databricks AIR CLI package."""
+
+# NOTE: version and git sha are determined at release time; the versions listed here are placeholders
+__version__ = "0.1.0"
+__git_sha__ = "4a95901bfafec7edf7c9f9f6919b19b1b2749338"

cli/ai_training_client.py

@@ -0,0 +1,51 @@
+"""Client for AiTrainingService — the per-user training-workflow surface on AICM.
+
+Authenticated as the calling user via the Databricks SDK's [[WorkspaceClient]], matching how
+[[cli.image_registration.image_client.ImageClient]] talks to AICM today.
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+from cli.utils import get_workspace_client
+from databricks.sdk import WorkspaceClient
+
+log = logging.getLogger(__name__)
+
+
+class AiTrainingClient:
+    """REST client for /api/2.0/ai-training/workflows."""
+
+    API_PATH = "/api/2.0/ai-training/workflows"
+
+    def __init__(self, workspace_host: Optional[str] = None):
+        self.workspace_host = workspace_host
+        self._workspace_client: Optional[WorkspaceClient] = None
+
+    def _get_workspace_client(self) -> WorkspaceClient:
+        if self._workspace_client is None:
+            self._workspace_client = get_workspace_client(workspace_host=self.workspace_host)
+        return self._workspace_client
+
+    def _api_request(
+        self,
+        method: str,
+        endpoint: str,
+        data: Optional[Dict[str, Any]] = None,
+        params: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        w = self._get_workspace_client()
+        path = f"{self.API_PATH}{endpoint}"
+        try:
+            response = w.api_client.do(method, path, body=data, query=params)
+            return response if response else {}
+        except Exception as e:
+            raise RuntimeError(f"AiTrainingService request failed: {e}") from e
+
+    def cancel_workflow(self, job_run_id: str) -> None:
+        """Cancel a training workflow by its Jobs job_run_id.
+
+        Maps to CancelTrainingWorkflow with a job_run_id ref. No-op server-side if the run is
+        already in a terminal state.
+        """
+        self._api_request(method="POST", endpoint=f"/by-run-id/{job_run_id}/cancel", data={})

cli/base_client.py

@@ -0,0 +1,101 @@
+"""Base client interface and types for workload management.
+This module defines the abstract interface that all workload clients must implement,
+ensuring consistent behavior across different backend APIs (Jobs API, CMv3, etc.).
+It also defines the client type enumeration.
+"""
+
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Dict, Any, Optional
+from cli.compute import GPUType
+
+
+class WorkloadClientType(Enum):
+    """Enum for different workload client types."""
+
+    JOBS_API = "Jobs API"
+    AI_COMPUTE_MANAGER = "AI Compute Manager (CMv3)"
+
+    def __str__(self) -> str:
+        """Return the string representation of the client type."""
+        return self.value
+
+
+class WorkloadClient(ABC):
+    """Abstract base class for workload management clients.
+    This interface ensures that all workload clients (Jobs API, CMv3, etc.)
+    provide the same core functionality with consistent method signatures.
+    """
+
+    @abstractmethod
+    def create_workload(
+        self,
+        gpus: int,
+        gpus_per_node: int,
+        python_file: str,
+        func_dir: str,
+        experiment_name: str,
+        timeout_seconds: int,
+        param_to_args: dict,
+        env_sync_timeout: int,
+        gpu_type: GPUType,
+        max_retries: int = 3,
+    ) -> str:
+        """Create a workload and return workload/run ID.
+        Args:
+            gpus: Number of GPUs required
+            gpus_per_node: GPUs per node (for distributed workloads)
+            python_file: Path to the Python script to execute
+            func_dir: Directory containing the function
+            experiment_name: Name of the function to execute
+            timeout_seconds: Timeout for the workload
+            param_to_args: Parameters to pass to the function
+            env_sync_timeout: Environment synchronization timeout
+            gpu_type: Type of GPU required
+            max_retries: Maximum number of retries for failed tasks (default: 3)
+        Returns:
+            Workload/run ID as string
+        Raises:
+            Exception: If workload creation fails
+        """
+        pass
+
+    @abstractmethod
+    def get_workload_status(self, workload_id: str) -> Dict[str, Any]:
+        """Get workload status and metadata.
+        Args:
+            workload_id: Unique workload identifier
+        Returns:
+            Dictionary containing workload status and metadata
+        """
+        pass
+
+    @abstractmethod
+    def terminate_workload(self, workload_id: str, reason: Optional[str] = None) -> bool:
+        """Terminate a running workload.
+        Args:
+            workload_id: Unique workload identifier
+            reason: Optional reason for termination
+        Returns:
+            True if termination was successful, False otherwise
+        """
+        pass
+
+    @abstractmethod
+    def list_workloads(self, **kwargs) -> Dict[str, Any]:
+        """List workloads.
+        Args:
+            **kwargs: Client-specific parameters for listing (pagination, filters, etc.)
+        Returns:
+            Dictionary containing workloads and pagination info
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def client_type(self) -> WorkloadClientType:
+        """Get the client type for logging and debugging.
+        Returns:
+            WorkloadClientType enum identifying the client type
+        """
+        pass

cli/changelog.md

@@ -0,0 +1,57 @@
+# Start of CLI Changelog
+
+Version 0.1.0
+[CNXT-2129] air run --dry-run now implies -v unless verbosity is already set.
+[CNXT-2130] air cancel: print a friendly "Run X not found" message instead of leaking the Jobs API URL.
+[CNXT-2155] air now emits a "watch" telemetry event capturing submission-to-first-streamed-log latency on run --watch and the logs command.
+[CNXT-2164] Fix `air run` crash on native Windows by skipping the POSIX-only SIGALRM upload timeout
+Fix --json error envelope classification: derive retryable from a closed ErrorKind set (only TRANSIENT retries), classify uniformly by SDK type / HTTP status / filesystem error across all commands, and emit error{NOT_FOUND} for subscribe on a missing run
+air --json now reports usage/early-exit failures (unknown subcommand, missing or invalid arguments, mutually-exclusive misuse, no command, and register/logs validation) as a structured USER error envelope on stdout with a non-zero exit instead of plain-text argparse usage or human help, and no longer double-prints the profile-not-found warning around the envelope
+Launch artifacts now upload under .air/cli_launch/<experiment>/<run_name>_<uuid> with the command script named command.sh, and --help/--version help text is capitalized.
+Deprecation warnings for renamed YAML keys now route through the `air` logger (shown on stderr, including in `--json` mode); the config schema and GPU types are now defined in the internal `air` SDK.
+Show [Beta] badges in help output: `air --help` (AIR CLI is in Beta) and Beta GPU types (e.g. GPU_1xH100) in `air -h config.compute`
+Stamp the installed CLI wheel version on BYOT submissions (client_version field) so the training service can correlate runs to a CLI release.
+`air cancel --all` now lists target runs and workspace before prompting for y/yes confirmation; skip with -y/--yes.
+Route `air cancel` through AiTrainingService instead of the Jobs API (no user-visible change).
+Forward `credentials-for-read` response headers to the pre-signed URL fetch in `mlflow_rest_client.download_artifact`. Fixes log/artifact downloads on Azure-backed workspaces where SAS URIs require headers like `x-ms-version` to authenticate.
+Route air diagnostics ([INFO]/[WARNING] messages) to stderr so `air logs | tee /dev/tty | pbcopy` captures only training log lines
+Fix cryptic "'str' object has no attribute 'decode'" error on submit when no Databricks profile is configured; surface the real authentication error instead
+Fix `sgcli --json run --watch` returning immediately with `status=PENDING` instead of blocking until terminal. The JSON path now emits a `SUBMITTED` event with run_id, streams `STATUS`/`LOG`/`ALERT` JSONL events through the watch, then emits a final success envelope whose `status` reflects the actual terminal state (`SUCCESS`/`FAILED`/`TIMEDOUT`/`CANCELED`).
+Fix post-submit MLflow run-name update clobbering the user-supplied `mlflow_run_name`. The CLI now passes the user's run_name (when set) to the post-submit `/api/2.0/mlflow/runs/update` call, falling back to `job_run_name` (== `experiment_name`) only when the user did not set one. Previously the CLI always passed `job_run_name`, which overwrote the run name the ai-training server had just set from `gen_ai_compute_task.mlflow_run_name`.
+Fix the "Waiting for run to start" spinner smearing into the first lines of log output during `air logs` and `air run --watch`. The Rich Live render thread was still drawing the spinner while raw bytes were being written to stdout from inside `print_new_logs`; `live.stop()` only ran in the caller after the write returned. The fix threads an `on_before_print` hook through `print_new_logs` that the streaming caller uses to tear down the spinner strictly before any byte hits stdout, and only when there is actually new content to emit (a no-op poll keeps the spinner running).
+Fix spinner artifacts in redirected output and JSON mode by detecting TTY status before creating Rich console objects
+[air-cli] Align `air get run` display with `air list runs`: rename "Task Run ID" -> "Run ID", adopt list-runs color palette, add User and Accelerators rows.
+`air get run`: started_at and duration_seconds now describe the reported (latest) attempt instead of the whole run, so a retried run no longer reports time spent on earlier failed attempts.
+validate GPU_* accelerator availability via AICM ListComputeOptions
+[air-cli] Help/error polish: drop duplicate `--download_to` alias in `air logs -h`; add description lines to subcommand help; `air` (no args) and `air config[.field]` now print help; remove UNCOMMITTED-CHANGES table from `air -h config.code_source`; long YAML help (`air -h config[.field]`) now opens in a pager unless `--json` / piped; pull_wheel.sh gains `--download-only` (downloads to /tmp and prints the `uv tool install` command); `Failed to get run output: 404` during list-runs no longer logs at ERROR; YAML validation deduplicates the "Available fields are: …" list when multiple unknown fields share a parent.
+Added support for inline dependencies in the workload YAML: environment.dependencies now accepts a list of packages alongside an environment.version field, as an alternative to pointing at a separate requirements.yaml file.
+Remove `no_interpolation` field and `--no-interpolation` flag; OmegaConf variable interpolation is now disabled unconditionally. Literal `${VAR}` strings in YAML are preserved as-is. Bash `$VAR` shell expansion at runtime on the worker node is unaffected.
+[air-cli] Plug --json envelope coverage gaps so machine consumers always get a structured envelope: `--version --json` returns `{version}` (instead of the ASCII banner); `--json changelog` returns `{version, changelog}`; `--json get pools` returns `{pools[], workspace_url, workspace_id}`; `--json run --dry-run` returns `{status: DRY_RUN_OK, dry_run: true}`; `--json register image` returns success envelopes (cached + new) and a classified error envelope (USER/INTERNAL_ERROR) on failure.
+PuPr filter consolidation: replace `get runs` flags `--user`, `--experiment`, and the deprecated `--all` with a single repeatable `--filter KEY=VALUE` flag (supported keys: user, experiment). `--all-users` and the other listing flags are unchanged.
+Pass user run_name through gen_ai_compute_task.mlflow_run_name on the TS path.
+sgcli logs: follow multi-chunk MLflow log streams; default to last 10000 lines for completed runs
+Group code_source.snapshot.{git_branch,git_commit,use_remote_head,remote_alias} under a nested `git:` object; consolidate remote behavior into bool-or-string `git.remote`.
+Allow --override to add new YAML fields and auto-create missing nested blocks; typos still surface from Pydantic schema validation.
+--override with an unknown field now reports a clean error naming the invalid key instead of a raw pydantic validation dump
+Without git_branch/git_commit, package working tree as plain tar; upload git state (base/tip/dirty + diff) as WSFS sidecars for the backend to apply MLflow tags.
+PuPr command renames: monitor->subscribe, get runs->list runs, get status->get run, get logs->logs; cancel now takes one or more run IDs or --all; removed --no-interpolation, get-run --watch, logs --local-rank; hid get pools, register image, logs --review; deprecated aliases (get runs, get status, get logs) still parse with a warning.
+PuPr cleanup: `environment` block is now optional in YAML, and `--host` is removed from `sgcli list runs` / `sgcli get pools` (workspace comes from `-p/--profile`).
+Expand `air register image -h` to walk through the three credential methods (docker login -- recommended, `--interactive-authenticate`, `--scope`/`--key`) with concrete examples.
+[air-cli] Reject workload configs whose YAML exceeds 1 MB before submission, with a clear error pointing at oversized parameters/command fields.
+Remove the `bash_script` YAML field. Use the top-level `command:` field instead. Configs that still specify `bash_script` now fail validation with a hint pointing at the replacement.
+Removed legacy YAML field and command aliases. Unrecognized fields and commands are now rejected as errors. Older client image versions are accepted again.
+Remove deprecated YAML fields: 'workspace' (use -p/--profile) and 'code_source.snapshot.allow_uncommitted' (commit changes or pin with git_commit / git_branch + use_remote_head). Raise 'command' cap from 500 to 1000 lines.
+Remove --local-rank from `sgcli logs` / `sgcli subscribe` and the `local_rank` parameter from `sdk.get_logs()` — per-rank logs are not produced by the platform; every rank funnels into the consolidated per-node stream
+Remove the deprecated nested `environment.env_variables` and `environment.env_variables_secrets` YAML fields. Use the top-level `env_variables:` and `secrets:` fields instead. Configs that still set the nested forms now fail validation with a hint pointing at the replacement.
+[BREAKING] Rename CLI from `sgcli` to `air`; wheel from `databricks-serverless-gpu-cli` to `databricks-air`. The `SGCLI_DISABLE_TELEMETRY` env var is now `AIR_DISABLE_TELEMETRY`. No deprecation alias — update scripts that invoke `sgcli`.
+Fix path expansion to support ~ (tilde) home directory shorthand in YAML config file paths
+Revert task_run_id as the user-facing identifier; the CLI again uses job_run_id as the canonical handle (#1929533 undone).
+[air-cli] Validation hardening: reject experiment_name > 100 chars client-side (Jobs API task_key limit); reject empty/whitespace docker_image_url in `air register image`; fix `air get logs --review` (without `--lines`) rendering "Last None lines per node" by coercing the None default to 200.
+Faster snapshot submission: detect uncommitted changes once and scope the check to include_paths, avoiding a redundant full-repo `git status` (slow on large monorepos).
+[air-cli] `air run` now shows self-clearing spinners while uploading the YAML configuration files and while packaging/uploading the code snapshot.
+Surface the server's error and run-termination messages (e.g. version upgrade notices) so job submission failures and runs that fail before producing logs report the actual reason instead of a raw HTTP error or bare status.
+Telemetry: emit job_run_id (Jobs API run_id) on sgcli run events to enable joins against AI scheduler activity logs.
+Assign a usage policy by name via usage_policy_name; removed usage_policy_id/budget_policy_id fields and the DATABRICKS_USAGE_POLICY_ID env var
+Install requirements.yaml deps with `uv pip install` layered on the base environment (no forced `-U`, so the image's torch and other packages are preserved). `--trusted-host` is ignored with a warning (uv configures trust per index URL).
+`air --version` now prints an ASCII art banner alongside the version string.
+Rename user-facing YAML fields to align with cross-product Databricks naming. Old names are temporarily accepted with a deprecation warning and will be removed in an upcoming release. Renames: env_variables_secrets -> secrets, run_name -> mlflow_run_name, experiment_directory -> mlflow_experiment_directory, budget_policy_id -> usage_policy_id, code_source.snapshot.repo_path -> code_source.snapshot.root_path, compute.gpus -> compute.num_accelerators, compute.gpu_type -> compute.accelerator_type, compute.gpu_node_pool_id -> compute.node_pool_id, compute.gpu_pool_name -> compute.pool_name.