From 723e0632038dc69a47ca13f4f6f2b1e2d5502ce0 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Thu, 27 Aug 2026 22:37:20 +0000 Subject: [PATCH] Add AIR submission and YAML parity --- .nextchanges/cli/air-submission-parity.md | 1 + .../experimental/air/config-help/output.txt | 30 +++++- .../experimental/air/config-help/script | 5 + .../air/run-submit-deps/output.txt | 60 +++++++++++- .../air/run-submit-deps/run-file.yaml | 2 + .../experimental/air/run-submit-deps/script | 10 +- .../air/run-submit-deps/test.toml | 4 + experimental/air/cmd/compute.go | 19 +++- experimental/air/cmd/compute_test.go | 4 + experimental/air/cmd/convert_to_dabs.go | 4 +- experimental/air/cmd/runconfig.go | 54 +++++++++-- experimental/air/cmd/runconfig_launch.go | 18 +++- experimental/air/cmd/runconfig_load.go | 96 +++++++++++++++++++ experimental/air/cmd/runconfig_test.go | 88 +++++++++++++++-- experimental/air/cmd/runsubmit.go | 81 +++++++++++++++- experimental/air/cmd/runsubmit_test.go | 36 +++++++ experimental/air/cmd/runupload.go | 9 ++ experimental/air/cmd/runupload_test.go | 19 ++++ experimental/air/cmd/validateconfig.go | 5 + 19 files changed, 509 insertions(+), 36 deletions(-) create mode 100644 .nextchanges/cli/air-submission-parity.md diff --git a/.nextchanges/cli/air-submission-parity.md b/.nextchanges/cli/air-submission-parity.md new file mode 100644 index 00000000000..aedb55d2cf0 --- /dev/null +++ b/.nextchanges/cli/air-submission-parity.md @@ -0,0 +1 @@ +Restored AIR requirements YAML support and added Docker image, MLflow artifact location, and provisioned capacity submission fields. diff --git a/acceptance/experimental/air/config-help/output.txt b/acceptance/experimental/air/config-help/output.txt index 2d7dffc135d..498af7d180b 100644 --- a/acceptance/experimental/air/config-help/output.txt +++ b/acceptance/experimental/air/config-help/output.txt @@ -50,6 +50,7 @@ config 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. + mlflow_artifact_location DBFS location where MLflow artifacts are written. 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. @@ -62,11 +63,34 @@ 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. + num_accelerators Total number of GPUs to allocate. + accelerator_type Which accelerator to run on, e.g. GPU_1xA10. + provisioned_capacity_id Pre-provisioned AIR capacity reservation id. Use "-h config.compute." for details on a field. +=== new submission fields are documented +>>> [CLI] experimental air run -h config.mlflow_artifact_location +config.mlflow_artifact_location + DBFS location where MLflow artifacts are written. A /Volumes path is normalized to dbfs:/Volumes/... . + + Type: string + Required: no + +>>> [CLI] experimental air run -h config.compute.provisioned_capacity_id +config.compute.provisioned_capacity_id + Pre-provisioned AIR capacity reservation id. Must be 1-255 characters. + + Type: string + Required: no + +>>> [CLI] experimental air run -h config.environment.dependencies +config.environment.dependencies + Inline package list or path to a requirements YAML file. Not allowed alongside docker_image. + + Type: list of strings + Required: no + === leaf field >>> [CLI] experimental air run -h config.compute.accelerator_type config.compute.accelerator_type @@ -103,7 +127,7 @@ config.compute.num_accelerators >>> [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 +fields under "config.compute" are: accelerator_type, num_accelerators, provisioned_capacity_id === free-form map keys are not schema fields >>> [CLI] experimental air run -h config.parameters.learning_rate diff --git a/acceptance/experimental/air/config-help/script b/acceptance/experimental/air/config-help/script index d1c25be1100..21e6c529a00 100644 --- a/acceptance/experimental/air/config-help/script +++ b/acceptance/experimental/air/config-help/script @@ -12,6 +12,11 @@ trace $CLI experimental air run -h config title "nested object lists its fields" trace $CLI experimental air run -h config.compute +title "new submission fields are documented" +trace $CLI experimental air run -h config.mlflow_artifact_location +trace $CLI experimental air run -h config.compute.provisioned_capacity_id +trace $CLI experimental air run -h config.environment.dependencies + title "leaf field" trace $CLI experimental air run -h config.compute.accelerator_type diff --git a/acceptance/experimental/air/run-submit-deps/output.txt b/acceptance/experimental/air/run-submit-deps/output.txt index 3c7280bc5ba..026f9c0da00 100644 --- a/acceptance/experimental/air/run-submit-deps/output.txt +++ b/acceptance/experimental/air/run-submit-deps/output.txt @@ -56,6 +56,62 @@ Tip: use --watch to stream logs until the run completes. } } -=== a requirements.yaml file path is rejected; deps must be inline +=== submit with file-backed dependencies >>> [CLI] experimental air run -f run-file.yaml -Error: invalid config run-file.yaml: environment.dependencies must be a list of packages or reference a requirements.txt (see https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference). A direct file reference is not supported +Submitting experiment: deps-file-smoke +Submitted workload with Job Run ID: 555 +View job run at: [DATABRICKS_URL]/jobs/runs/555 + +Tip: use --watch to stream logs until the run completes. + +=== requirements.yaml is uploaded beside the launch config +>>> print_requests.py //api/2.0/workspace-files/import-file --oneline --sort --unique --keep +{"method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/cli_launch/deps-file-smoke/deps-file-smoke_[RUN_ID]/command.sh", "q": {"overwrite": "true"}, "raw_body": "python train.py"} +{"method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/cli_launch/deps-file-smoke/deps-file-smoke_[RUN_ID]/requirements.yaml", "q": {"overwrite": "true"}, "raw_body": "version: 5\ndependencies:\n - numpy\n - torch==2.3.0\n"} +{"method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/cli_launch/deps-file-smoke/deps-file-smoke_[RUN_ID]/training_config.yaml", "q": {"overwrite": "true"}, "raw_body": "experiment_name: deps-file-smoke\ncommand: python train.py\ncompute:\n accelerator_type: GPU_1xH100\n num_accelerators: 1\n provisioned_capacity_id: capacity-123\nenvironment:\n dependencies: ./reqs.yaml\nmlflow_artifact_location: /Volumes/main/default/air-artifacts\n"} + +=== file deps, artifact location, and capacity id reach the submit payload +>>> print_requests.py //api/2.2/jobs/runs/submit +{ + "method": "POST", + "path": "/api/2.2/jobs/runs/submit", + "body": { + "environments": [ + { + "environment_key": "default", + "spec": { + "dependencies": [ + "numpy", + "torch==2.3.0" + ], + "environment_version": "5" + } + } + ], + "idempotency_token": "[UUID]", + "run_name": "deps-file-smoke", + "tasks": [ + { + "ai_runtime_task": { + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/deps-file-smoke/deps-file-smoke_[RUN_ID]/command.sh", + "compute": { + "accelerator_count": 1, + "accelerator_type": "GPU_1xH100", + "provisioned_capacity_id": "capacity-123" + } + } + ], + "experiment": "deps-file-smoke", + "mlflow_artifact_location": "dbfs:/Volumes/main/default/air-artifacts" + }, + "environment_key": "default", + "max_retries": 3, + "retry_on_timeout": true, + "run_if": "ALL_SUCCESS", + "task_key": "deps-file-smoke" + } + ] + } +} diff --git a/acceptance/experimental/air/run-submit-deps/run-file.yaml b/acceptance/experimental/air/run-submit-deps/run-file.yaml index 3e5fe54b7b0..122233f07d6 100644 --- a/acceptance/experimental/air/run-submit-deps/run-file.yaml +++ b/acceptance/experimental/air/run-submit-deps/run-file.yaml @@ -3,5 +3,7 @@ command: python train.py compute: accelerator_type: GPU_1xH100 num_accelerators: 1 + provisioned_capacity_id: capacity-123 environment: dependencies: ./reqs.yaml +mlflow_artifact_location: /Volumes/main/default/air-artifacts diff --git a/acceptance/experimental/air/run-submit-deps/script b/acceptance/experimental/air/run-submit-deps/script index 12dc884740d..ae45af188e1 100644 --- a/acceptance/experimental/air/run-submit-deps/script +++ b/acceptance/experimental/air/run-submit-deps/script @@ -7,5 +7,11 @@ trace print_requests.py //api/2.0/workspace-files/import-file --oneline --sort - title "declared deps ride on environments[].spec.dependencies" trace print_requests.py //api/2.2/jobs/runs/submit -title "a requirements.yaml file path is rejected; deps must be inline" -musterr trace $CLI experimental air run -f run-file.yaml +title "submit with file-backed dependencies" +trace $CLI experimental air run -f run-file.yaml + +title "requirements.yaml is uploaded beside the launch config" +trace print_requests.py //api/2.0/workspace-files/import-file --oneline --sort --unique --keep + +title "file deps, artifact location, and capacity id reach the submit payload" +trace print_requests.py //api/2.2/jobs/runs/submit diff --git a/acceptance/experimental/air/run-submit-deps/test.toml b/acceptance/experimental/air/run-submit-deps/test.toml index abe4f7f5ebf..097eb39d8c9 100644 --- a/acceptance/experimental/air/run-submit-deps/test.toml +++ b/acceptance/experimental/air/run-submit-deps/test.toml @@ -21,3 +21,7 @@ Response.Body = ''' [[Repls]] Old = 'deps-smoke_[0-9a-f]{16}' New = 'deps-smoke_[RUN_ID]' + +[[Repls]] +Old = 'deps-file-smoke_[0-9a-f]{16}' +New = 'deps-file-smoke_[RUN_ID]' diff --git a/experimental/air/cmd/compute.go b/experimental/air/cmd/compute.go index e7b69d470be..658f7291d51 100644 --- a/experimental/air/cmd/compute.go +++ b/experimental/air/cmd/compute.go @@ -1,6 +1,7 @@ package aircmd import ( + "errors" "fmt" "strings" ) @@ -54,12 +55,13 @@ func gpusPerNode(g gpuType) (int, error) { // computeConfig is the `compute` block of the run YAML: which accelerators to // use and how many. type computeConfig struct { - NumAccelerators int `yaml:"num_accelerators" help:"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."` - AcceleratorType string `yaml:"accelerator_type" help:"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."` + NumAccelerators int `yaml:"num_accelerators" help:"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."` + AcceleratorType string `yaml:"accelerator_type" help:"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."` + ProvisionedCapacityID *string `yaml:"provisioned_capacity_id" help:"Pre-provisioned AIR capacity reservation id. Must be 1-255 characters."` } // validate checks the compute block against the backend's constraints. -func (c computeConfig) validate() error { +func (c *computeConfig) validate() error { g, err := parseGPUType(c.AcceleratorType) if err != nil { return fmt.Errorf("compute.accelerator_type: %w", err) @@ -77,5 +79,16 @@ func (c computeConfig) validate() error { return fmt.Errorf("compute.num_accelerators for %s must be a multiple of %d, got %d", c.AcceleratorType, perNode, c.NumAccelerators) } + if c.ProvisionedCapacityID != nil { + v := strings.TrimSpace(*c.ProvisionedCapacityID) + if v == "" { + return errors.New("compute.provisioned_capacity_id cannot be empty") + } + if len(v) > 255 { + return fmt.Errorf("compute.provisioned_capacity_id must be 255 characters or less, got %d", len(v)) + } + *c.ProvisionedCapacityID = v + } + return nil } diff --git a/experimental/air/cmd/compute_test.go b/experimental/air/cmd/compute_test.go index 3464afbe9ea..de6d42becd6 100644 --- a/experimental/air/cmd/compute_test.go +++ b/experimental/air/cmd/compute_test.go @@ -1,6 +1,7 @@ package aircmd import ( + "strings" "testing" "github.com/stretchr/testify/assert" @@ -67,6 +68,9 @@ func TestComputeConfigValidate(t *testing.T) { {"single node", computeConfig{NumAccelerators: 8, AcceleratorType: "GPU_8xH100"}, ""}, {"multiple nodes", computeConfig{NumAccelerators: 16, AcceleratorType: "GPU_8xH100"}, ""}, {"single-gpu partitions", computeConfig{NumAccelerators: 3, AcceleratorType: "GPU_1xH100"}, ""}, + {"capacity id", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new(" capacity ")}, ""}, + {"empty capacity id", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new(" ")}, "cannot be empty"}, + {"long capacity id", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new(strings.Repeat("a", 256))}, "255 characters or less"}, {"unknown type", computeConfig{NumAccelerators: 8, AcceleratorType: "b200"}, "accelerator_type"}, {"legacy type rejected", computeConfig{NumAccelerators: 8, AcceleratorType: "h100_80gb"}, "accelerator_type"}, {"non-positive count", computeConfig{NumAccelerators: 0, AcceleratorType: "GPU_1xH100"}, "must be positive"}, diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index 4f22233667d..55461389ba0 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -329,8 +329,8 @@ func buildBundleValue(ctx context.Context, cfg *runConfig, configPath, codeSourc // bundleEnvironmentDeps resolves the runtime version and the inline dependency // list to emit in the bundle's environments[] spec. The aicode mutator synthesizes // requirements.yaml from that spec at deploy, so the whole set must be here. -// Dependencies are inline-only (a requirements-file path is rejected at config -// load), so an unset list yields no dependencies. +// File-backed dependencies are resolved during config load, so an unset list +// yields no dependencies here. func bundleEnvironmentDeps(ctx context.Context, cfg *runConfig) (version string, deps []string) { cfgVersion, _ := cfg.runtimeVersion() version = dlRuntimeImage(ctx, cfgVersion) diff --git a/experimental/air/cmd/runconfig.go b/experimental/air/cmd/runconfig.go index 86fd9d2b4e4..195a8af733c 100644 --- a/experimental/air/cmd/runconfig.go +++ b/experimental/air/cmd/runconfig.go @@ -56,6 +56,7 @@ type runConfig struct { Parameters map[string]any `yaml:"parameters" help:"Free-form values passed through to the workload. Any nested structure is allowed."` MLflowRunName *string `yaml:"mlflow_run_name" help:"Name for the MLflow run. Max 100 characters, alphanumerics, hyphens, and underscores only."` MLflowExperimentDirectory *string `yaml:"mlflow_experiment_directory" help:"Workspace directory holding the MLflow experiment. Must start with /Workspace."` + MLflowArtifactLocation *string `yaml:"mlflow_artifact_location" help:"DBFS location where MLflow artifacts are written. A /Volumes path is normalized to dbfs:/Volumes/... ."` Permissions []permission `yaml:"permissions" help:"Who may view or manage the run, as a list of principal plus level grants."` UsagePolicyName *string `yaml:"usage_policy_name" help:"Usage policy to bill the run to, by name. Max 127 characters. Mutually exclusive with usage_policy_id."` UsagePolicyID *string `yaml:"usage_policy_id" help:"Usage policy to bill the run to, by id. Mutually exclusive with usage_policy_name."` @@ -150,6 +151,20 @@ func (c *runConfig) validate() error { } } + if c.MLflowArtifactLocation != nil { + v := strings.TrimSpace(*c.MLflowArtifactLocation) + if v == "" { + return errors.New("mlflow_artifact_location cannot be empty") + } + if strings.HasPrefix(v, "/Volumes/") { + v = "dbfs:" + v + } + if !strings.HasPrefix(v, "dbfs:/") { + return fmt.Errorf("mlflow_artifact_location must be a dbfs: URI, got: %s", v) + } + *c.MLflowArtifactLocation = v + } + for i := range c.Permissions { if err := c.Permissions[i].validate(); err != nil { return err @@ -229,7 +244,7 @@ func validateSecretRefs(secrets map[string]string) error { // environmentConfig is the `environment` block: dependencies and/or a custom // docker image. type environmentConfig struct { - Dependencies dependencies `yaml:"dependencies" help:"Inline list of packages to install. Not allowed alongside docker_image."` + Dependencies dependencies `yaml:"dependencies" help:"Inline package list or path to a requirements YAML file. Not allowed alongside docker_image."` Version stringOrInt `yaml:"version" help:"Client image version to pin. Only valid alongside inline dependencies."` DockerImage *dockerImageConfig `yaml:"docker_image" help:"Custom image supplying the whole runtime. Not allowed alongside dependencies or version."` } @@ -256,24 +271,43 @@ func (e *environmentConfig) validate() error { if e.Version.set && !e.Dependencies.set { return errors.New("'environment.version' requires inline 'dependencies' (a list of packages)") } + if e.Version.set { + version, err := validateRuntimeVersion(e.Version.raw, "environment.version") + if err != nil { + return err + } + e.Version.raw = version + } return nil } -// dependencies is environment.dependencies: an inline list of packages. A scalar -// (e.g. a path to a requirements file) is rejected — the list may itself reference -// a requirements.txt, but dependencies must be given as a list. +// dependencies is environment.dependencies: either an inline list of packages or +// a path to a requirements YAML file resolved relative to the run config. type dependencies struct { - set bool - list []string + set bool + list []string + path string + resolvedPath string + version string } func (d *dependencies) UnmarshalYAML(node *yaml.Node) error { - if node.Kind != yaml.SequenceNode { - return errors.New("environment.dependencies must be a list of packages or reference a requirements.txt (see https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference). A direct file reference is not supported") - } d.set = true - return node.Decode(&d.list) + switch node.Kind { + case yaml.SequenceNode: + return node.Decode(&d.list) + case yaml.ScalarNode: + if err := node.Decode(&d.path); err != nil { + return err + } + if strings.TrimSpace(d.path) == "" { + return errors.New("environment.dependencies requirements YAML path cannot be empty") + } + return nil + default: + return errors.New("environment.dependencies must be a list of packages or a requirements YAML path") + } } // stringOrInt holds a scalar that may be a string or an integer in YAML diff --git a/experimental/air/cmd/runconfig_launch.go b/experimental/air/cmd/runconfig_launch.go index 6a658c89b29..b4c76773ee7 100644 --- a/experimental/air/cmd/runconfig_launch.go +++ b/experimental/air/cmd/runconfig_launch.go @@ -26,9 +26,6 @@ func (c *runConfig) maxRetries() int { } // dockerImageURL returns the custom docker image URL, or "" when none is set. -// -// TODO: not wired into submission yet — the native ai_runtime_task carries no -// docker field, and full support needs image registration (pending the DCS work). func (c *runConfig) dockerImageURL() string { if c.Environment != nil && c.Environment.DockerImage != nil { return c.Environment.DockerImage.URL @@ -53,10 +50,23 @@ func (c *runConfig) inlineDependencies() ([]string, bool) { return c.Environment.Dependencies.list, true } +func (c *runConfig) requirementsPath() string { + if c.Environment == nil { + return "" + } + return c.Environment.Dependencies.resolvedPath +} + // runtimeVersion returns the client image version from environment.version when // set. func (c *runConfig) runtimeVersion() (string, bool) { - if c.Environment == nil || !c.Environment.Version.set { + if c.Environment == nil { + return "", false + } + if c.Environment.Dependencies.version != "" { + return c.Environment.Dependencies.version, true + } + if !c.Environment.Version.set { return "", false } return c.Environment.Version.raw, true diff --git a/experimental/air/cmd/runconfig_load.go b/experimental/air/cmd/runconfig_load.go index 7f6ad8b5e17..05343e280e7 100644 --- a/experimental/air/cmd/runconfig_load.go +++ b/experimental/air/cmd/runconfig_load.go @@ -7,6 +7,10 @@ import ( "fmt" "io" "os" + "path/filepath" + "regexp" + "strconv" + "strings" "go.yaml.in/yaml/v3" ) @@ -57,6 +61,9 @@ func loadRunConfig(path string) (*runConfig, error) { if err := validateRunConfig(cfg); err != nil { return nil, err } + if err := resolveRequirementsFile(cfg, path); err != nil { + return nil, err + } return cfg, nil } @@ -107,5 +114,94 @@ func loadRunConfigWithOverrides(ctx context.Context, path string, overrides []st if err := validateRunConfig(cfg); err != nil { return nil, err } + if err := resolveRequirementsFile(cfg, path); err != nil { + return nil, err + } return cfg, nil } + +var runtimeVersionRe = regexp.MustCompile(`^[0-9]+$`) + +const databricksAIPrefix = "databricks_ai_v" + +func validateRuntimeVersion(version, source string) (string, error) { + normalized := strings.ToLower(version) + numeric := normalized + usesDatabricksAI := strings.HasPrefix(normalized, databricksAIPrefix) + if usesDatabricksAI { + numeric = strings.TrimPrefix(normalized, databricksAIPrefix) + } + if !runtimeVersionRe.MatchString(numeric) { + return "", fmt.Errorf("unsupported client image version %q in %s: version must be an integer, optionally prefixed with databricks_ai_v", version, source) + } + if !usesDatabricksAI { + return numeric, nil + } + major, err := strconv.Atoi(numeric) + if err != nil { + return "", fmt.Errorf("failed to parse client image version %q in %s: %w", version, source, err) + } + if major < 5 { + return "", fmt.Errorf("databricks_ai_v in %s requires AI Runtime version 5 or higher, got %q", source, version) + } + return databricksAIPrefix + numeric, nil +} + +type requirementsConfig struct { + Version stringOrInt `yaml:"version"` + Dependencies []string `yaml:"dependencies"` +} + +func resolveRequirementsFile(cfg *runConfig, configPath string) error { + if cfg.Environment == nil || cfg.Environment.Dependencies.path == "" { + return nil + } + if cfg.Environment.Version.set { + return errors.New("'environment.version' is only supported with inline 'dependencies'") + } + + dep := &cfg.Environment.Dependencies + resolved := dep.path + if !filepath.IsAbs(resolved) { + resolved = filepath.Join(filepath.Dir(configPath), resolved) + } + info, err := os.Stat(resolved) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("environment.dependencies: requirements YAML not found at %q", resolved) + } + return fmt.Errorf("failed to inspect requirements YAML %q: %w", resolved, err) + } + if !info.Mode().IsRegular() { + return fmt.Errorf("environment.dependencies: requirements YAML %q is not a file", resolved) + } + + f, err := os.Open(resolved) + if err != nil { + return fmt.Errorf("failed to open requirements YAML %q: %w", resolved, err) + } + defer f.Close() + dec := yaml.NewDecoder(f) + var req requirementsConfig + if err := dec.Decode(&req); err != nil { + return fmt.Errorf("failed to parse requirements YAML %q: %w", resolved, err) + } + version := "4" + if req.Version.set { + version = req.Version.raw + } + version, err = validateRuntimeVersion(version, fmt.Sprintf("requirements YAML %q", resolved)) + if err != nil { + return err + } + for _, item := range req.Dependencies { + fields := strings.Fields(item) + if len(fields) > 0 && (fields[0] == "-r" || fields[0] == "--requirement") { + return fmt.Errorf("requirements YAML dependency %q uses unsupported -r/--requirement include", item) + } + } + dep.resolvedPath = resolved + dep.version = version + dep.list = req.Dependencies + return nil +} diff --git a/experimental/air/cmd/runconfig_test.go b/experimental/air/cmd/runconfig_test.go index f99bb344e79..38b613e5561 100644 --- a/experimental/air/cmd/runconfig_test.go +++ b/experimental/air/cmd/runconfig_test.go @@ -95,16 +95,63 @@ permissions: } // TestLoadRunConfig_PolymorphicFields exercises the str|int and bool|str unions -// decoded by custom UnmarshalYAML, plus the rejection of the removed -// dependencies string form. +// decoded by custom UnmarshalYAML, including file-backed dependencies. func TestLoadRunConfig_PolymorphicFields(t *testing.T) { - t.Run("dependencies as string path is rejected", func(t *testing.T) { - _, err := loadRunConfig(writeConfig(t, minimalConfig+` + t.Run("dependencies as string path is resolved", func(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "requirements.yaml"), []byte("version: 5\ndependencies:\n - torch\n"), 0o600)) + configPath := filepath.Join(dir, "run.yaml") + require.NoError(t, os.WriteFile(configPath, []byte(minimalConfig+` environment: dependencies: requirements.yaml -`)) - require.Error(t, err) - assert.Contains(t, err.Error(), "must be a list of packages") +`), 0o600)) + cfg, err := loadRunConfig(configPath) + require.NoError(t, err) + assert.Equal(t, filepath.Join(dir, "requirements.yaml"), cfg.requirementsPath()) + assert.Equal(t, []string{"torch"}, cfg.Environment.Dependencies.list) + version, ok := cfg.runtimeVersion() + assert.True(t, ok) + assert.Equal(t, "5", version) + }) + + t.Run("requirements defaults version and normalizes prefix", func(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "requirements.yaml"), []byte("version: DATABRICKS_AI_V5\ndependencies:\n - torch\n"), 0o600)) + configPath := filepath.Join(dir, "run.yaml") + require.NoError(t, os.WriteFile(configPath, []byte(minimalConfig+"environment:\n dependencies: requirements.yaml\n"), 0o600)) + cfg, err := loadRunConfig(configPath) + require.NoError(t, err) + version, ok := cfg.runtimeVersion() + assert.True(t, ok) + assert.Equal(t, "databricks_ai_v5", version) + }) + + t.Run("requirements permits worker-side fields", func(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "requirements.yaml"), []byte("version: 5\ndependencies:\n - torch\npip_options:\n index_url: https://packages.example/simple\n"), 0o600)) + configPath := filepath.Join(dir, "run.yaml") + require.NoError(t, os.WriteFile(configPath, []byte(minimalConfig+"environment:\n dependencies: requirements.yaml\n"), 0o600)) + cfg, err := loadRunConfig(configPath) + require.NoError(t, err) + assert.Equal(t, []string{"torch"}, cfg.Environment.Dependencies.list) + }) + + t.Run("requirements rejects old databricks ai version", func(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "requirements.yaml"), []byte("version: databricks_ai_v4\n"), 0o600)) + configPath := filepath.Join(dir, "run.yaml") + require.NoError(t, os.WriteFile(configPath, []byte(minimalConfig+"environment:\n dependencies: requirements.yaml\n"), 0o600)) + _, err := loadRunConfig(configPath) + require.ErrorContains(t, err, "requires AI Runtime version 5 or higher") + }) + + t.Run("requirements rejects includes", func(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "requirements.yaml"), []byte("dependencies:\n - -r other.txt\n"), 0o600)) + configPath := filepath.Join(dir, "run.yaml") + require.NoError(t, os.WriteFile(configPath, []byte(minimalConfig+"environment:\n dependencies: requirements.yaml\n"), 0o600)) + _, err := loadRunConfig(configPath) + require.ErrorContains(t, err, "unsupported -r/--requirement include") }) t.Run("git remote as bool true is rejected", func(t *testing.T) { @@ -229,6 +276,9 @@ func TestRunConfigValidate_FieldRules(t *testing.T) { {"long idempotency", func(c *runConfig) { c.IdempotencyToken = str(string(make([]byte, 65))) }, "64 characters or less"}, {"bad mlflow_run_name", func(c *runConfig) { c.MLflowRunName = str("bad name") }, "invalid mlflow_run_name"}, {"bad experiment dir", func(c *runConfig) { c.MLflowExperimentDirectory = str("/Users/me") }, "must start with '/Workspace'"}, + {"artifact volume path normalizes", func(c *runConfig) { c.MLflowArtifactLocation = str(" /Volumes/main/default/artifacts ") }, ""}, + {"empty artifact location", func(c *runConfig) { c.MLflowArtifactLocation = str(" ") }, "mlflow_artifact_location cannot be empty"}, + {"non-dbfs artifact location", func(c *runConfig) { c.MLflowArtifactLocation = str("s3://bucket/path") }, "must be a dbfs: URI"}, {"empty usage policy", func(c *runConfig) { c.UsagePolicyName = str(" ") }, "usage_policy_name must not be empty"}, {"bad secret ref", func(c *runConfig) { c.Secrets = map[string]string{"T": "noslash"} }, "expected format 'scope/key'"}, {"empty secret scope", func(c *runConfig) { c.Secrets = map[string]string{"T": "/key"} }, "scope and key cannot be empty"}, @@ -254,6 +304,9 @@ func TestRunConfigValidate_FieldRules(t *testing.T) { err := c.validate() if tt.errFrag == "" { assert.NoError(t, err) + if tt.name == "artifact volume path normalizes" { + assert.Equal(t, "dbfs:/Volumes/main/default/artifacts", *c.MLflowArtifactLocation) + } return } require.Error(t, err) @@ -299,12 +352,31 @@ func TestEnvironmentConfigValidate(t *testing.T) { }, "", }, + { + "version prefix normalizes", + environmentConfig{ + Version: stringOrInt{set: true, raw: "DATABRICKS_AI_V5"}, + Dependencies: dependencies{set: true, list: []string{"torch"}}, + }, + "", + }, + { + "old databricks ai version rejected", + environmentConfig{ + Version: stringOrInt{set: true, raw: "databricks_ai_v4"}, + Dependencies: dependencies{set: true, list: []string{"torch"}}, + }, + "requires AI Runtime version 5 or higher", + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { err := tt.env.validate() if tt.errFrag == "" { assert.NoError(t, err) + if tt.name == "version prefix normalizes" { + assert.Equal(t, "databricks_ai_v5", tt.env.Version.raw) + } return } require.Error(t, err) @@ -474,7 +546,7 @@ func TestResolveConfigField_Containers(t *testing.T) { compute, err := resolveConfigField("config.compute") require.NoError(t, err) assert.Equal(t, "object", compute.typeName) - require.Len(t, compute.children, 2) + require.Len(t, compute.children, 3) } func TestResolveConfigField_Errors(t *testing.T) { diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 3be0e455f46..2456ef3eca4 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -1,16 +1,22 @@ package aircmd import ( + "bytes" "context" + "encoding/json" + "errors" "fmt" + "net/http" "path" "strconv" "strings" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/filer" "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/service/compute" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/google/uuid" @@ -67,6 +73,10 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID stri if cfg.MLflowExperimentDirectory != nil { task.MlflowExperimentDirectory = *cfg.MLflowExperimentDirectory } + if cfg.MLflowArtifactLocation != nil { + task.MlflowArtifactLocation = *cfg.MLflowArtifactLocation + } + task.DockerImageUrl = cfg.dockerImageURL() maxRetries := cfg.maxRetries() st := jobs.SubmitTask{ @@ -104,6 +114,70 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID stri } } +func submitRun(ctx context.Context, w *databricks.WorkspaceClient, payload jobs.SubmitRun, provisionedCapacityID string) (int64, error) { + if provisionedCapacityID == "" { + wait, err := w.Jobs.Submit(ctx, payload) + if err != nil { + return 0, err + } + return wait.RunId, nil + } + + raw, err := json.Marshal(payload) + if err != nil { + return 0, fmt.Errorf("failed to marshal AIR submit payload: %w", err) + } + var body map[string]any + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + if err := decoder.Decode(&body); err != nil { + return 0, fmt.Errorf("failed to decode AIR submit payload: %w", err) + } + if err := injectProvisionedCapacityID(body, provisionedCapacityID); err != nil { + return 0, err + } + + apiClient, err := client.New(w.Config) + if err != nil { + return 0, fmt.Errorf("failed to create API client: %w", err) + } + var response jobs.SubmitRunResponse + err = apiClient.Do(ctx, http.MethodPost, "/api/2.2/jobs/runs/submit", auth.WorkspaceIDHeaders(w.Config), nil, body, &response) + if err != nil { + return 0, err + } + return response.RunId, nil +} + +func injectProvisionedCapacityID(body map[string]any, provisionedCapacityID string) error { + tasks, ok := body["tasks"].([]any) + if !ok || len(tasks) != 1 { + return errors.New("AIR submit payload must contain exactly one task") + } + task, ok := tasks[0].(map[string]any) + if !ok { + return errors.New("AIR submit payload task has an invalid shape") + } + aiRuntimeTask, ok := task["ai_runtime_task"].(map[string]any) + if !ok { + return errors.New("AIR submit payload is missing ai_runtime_task") + } + deployments, ok := aiRuntimeTask["deployments"].([]any) + if !ok || len(deployments) != 1 { + return errors.New("AIR submit payload must contain exactly one deployment") + } + deployment, ok := deployments[0].(map[string]any) + if !ok { + return errors.New("AIR submit payload deployment has an invalid shape") + } + computeSpec, ok := deployment["compute"].(map[string]any) + if !ok { + return errors.New("AIR submit payload is missing deployment compute") + } + computeSpec["provisioned_capacity_id"] = provisionedCapacityID + return nil +} + // submitToken resolves the idempotency token: the --idempotency-key flag wins, // then the config's token, else a generated one. Over-long tokens error rather // than truncate, since truncation could make two distinct tokens collide. @@ -235,12 +309,15 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run payload := buildSubmitPayload(cfg, commandPath, dlRuntimeImage(ctx, runtimeVersion), usagePolicyID, snap, deps) payload.IdempotencyToken = token + provisionedCapacityID := "" + if cfg.Compute.ProvisionedCapacityID != nil { + provisionedCapacityID = *cfg.Compute.ProvisionedCapacityID + } // Submit returns as soon as the run is created; we don't wait for it to finish. - wait, err := w.Jobs.Submit(ctx, payload) + runID, err := submitRun(ctx, w, payload, provisionedCapacityID) if err != nil { return 0, "", err } - runID := wait.RunId dashboardURL := strings.TrimRight(w.Config.Host, "/") + "/jobs/runs/" + strconv.FormatInt(runID, 10) return runID, dashboardURL, nil diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index dd8471bd3f5..f033801408d 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -47,6 +47,10 @@ func TestBuildSubmitPayload(t *testing.T) { TimeoutMinutes: new(30), MLflowRunName: new("run-v2"), MLflowExperimentDirectory: new("/Workspace/Users/me/exp"), + MLflowArtifactLocation: new("dbfs:/Volumes/main/default/artifacts"), + Environment: &environmentConfig{DockerImage: &dockerImageConfig{ + URL: "registry.example.com/team/image:tag", + }}, } p := buildSubmitPayload(cfg, "/d/command.sh", "5", "", snapshotResult{}, nil) @@ -73,11 +77,43 @@ func TestBuildSubmitPayload(t *testing.T) { assert.Equal(t, "exp", at.Experiment) assert.Equal(t, "run-v2", at.MlflowRun) assert.Equal(t, "/Workspace/Users/me/exp", at.MlflowExperimentDirectory) + assert.Equal(t, "dbfs:/Volumes/main/default/artifacts", at.MlflowArtifactLocation) + assert.Equal(t, "registry.example.com/team/image:tag", at.DockerImageUrl) require.Len(t, at.Deployments, 1) assert.Equal(t, "/d/command.sh", at.Deployments[0].CommandPath) assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu8xH100, AcceleratorCount: 16}, at.Deployments[0].Compute) } +func TestSubmitRunInjectsProvisionedCapacityID(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + assert.Equal(t, "123", req.Headers.Get("X-Databricks-Workspace-Id")) + var body map[string]any + require.NoError(t, json.Unmarshal(req.Body, &body)) + tasks := body["tasks"].([]any) + task := tasks[0].(map[string]any) + airTask := task["ai_runtime_task"].(map[string]any) + deployments := airTask["deployments"].([]any) + deployment := deployments[0].(map[string]any) + compute := deployment["compute"].(map[string]any) + assert.Equal(t, "capacity-1", compute["provisioned_capacity_id"]) + return jobs.SubmitRunResponse{RunId: 42} + }) + + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token", WorkspaceID: "123"}) + require.NoError(t, err) + payload := buildSubmitPayload(&runConfig{ + ExperimentName: "exp", + Command: new("x"), + Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, + }, "/command.sh", "4", "", snapshotResult{}, nil) + + runID, err := submitRun(t.Context(), w, payload, "capacity-1") + require.NoError(t, err) + assert.Equal(t, int64(42), runID) +} + func TestBuildSubmitPayloadDefaultRetries(t *testing.T) { // max_retries unset defaults to 3 (matching the Python native path), so both // retry fields are sent. diff --git a/experimental/air/cmd/runupload.go b/experimental/air/cmd/runupload.go index ca0ffd9c480..a19a3744eea 100644 --- a/experimental/air/cmd/runupload.go +++ b/experimental/air/cmd/runupload.go @@ -24,6 +24,7 @@ const ( hyperparametersName = "hyperparameters.yaml" envVarsName = "env_vars.json" secretEnvVarsName = "secret_env_vars.json" + requirementsName = "requirements.yaml" ) // maxConfigYAMLBytes caps training_config.yaml. It is referenced by the Jobs @@ -74,6 +75,14 @@ func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) { items = append(items, uploadItem{hyperparametersName, data}) } + if requirementsPath := cfg.requirementsPath(); requirementsPath != "" { + data, err := os.ReadFile(requirementsPath) + if err != nil { + return nil, fmt.Errorf("failed to read requirements YAML %s: %w", requirementsPath, err) + } + items = append(items, uploadItem{requirementsName, data}) + } + // The ai_runtime_task proto carries no inline env vars or secrets; stage them // as JSON files co-located with command.sh for the server-side launcher. if len(cfg.EnvVariables) > 0 { diff --git a/experimental/air/cmd/runupload_test.go b/experimental/air/cmd/runupload_test.go index 1b3e04ee00d..d7c90de929e 100644 --- a/experimental/air/cmd/runupload_test.go +++ b/experimental/air/cmd/runupload_test.go @@ -74,6 +74,25 @@ func TestBuildArtifacts_ParametersButNoRequirements(t *testing.T) { assert.Equal(t, []string{trainingConfigName, commandScriptName, hyperparametersName}, itemNames(items)) } +func TestBuildArtifacts_FileRequirements(t *testing.T) { + dir := t.TempDir() + configPath := filepath.Join(dir, "run.yaml") + requirementsPath := filepath.Join(dir, "deps.yaml") + require.NoError(t, os.WriteFile(configPath, []byte("x: y\n"), 0o600)) + require.NoError(t, os.WriteFile(requirementsPath, []byte("version: 5\ndependencies:\n - torch\n"), 0o600)) + cfg := &runConfig{ + Command: new("echo hi"), + Environment: &environmentConfig{Dependencies: dependencies{ + set: true, resolvedPath: requirementsPath, + }}, + } + + items, err := buildArtifacts(cfg, configPath) + require.NoError(t, err) + assert.Equal(t, []string{trainingConfigName, commandScriptName, requirementsName}, itemNames(items)) + assert.Equal(t, "version: 5\ndependencies:\n - torch\n", string(items[2].data)) +} + func TestBuildArtifacts_EnvVarsAndSecrets(t *testing.T) { path := writeConfigFile(t, "run.yaml", "x: y\n") cfg := &runConfig{ diff --git a/experimental/air/cmd/validateconfig.go b/experimental/air/cmd/validateconfig.go index d65339878f0..5df25ce862d 100644 --- a/experimental/air/cmd/validateconfig.go +++ b/experimental/air/cmd/validateconfig.go @@ -66,6 +66,7 @@ func validateConfigRequest(cfg *runConfig, commandPath string) map[string]any { if cfg.Compute != nil { compute["accelerator_type"] = cfg.Compute.AcceleratorType compute["accelerator_count"] = cfg.Compute.NumAccelerators + putOpt(compute, "provisioned_capacity_id", cfg.Compute.ProvisionedCapacityID) } task := map[string]any{ "experiment": cfg.ExperimentName, @@ -73,6 +74,10 @@ func validateConfigRequest(cfg *runConfig, commandPath string) map[string]any { } putOpt(task, "mlflow_run", cfg.MLflowRunName) putOpt(task, "mlflow_experiment_directory", cfg.MLflowExperimentDirectory) + putOpt(task, "mlflow_artifact_location", cfg.MLflowArtifactLocation) + if dockerImageURL := cfg.dockerImageURL(); dockerImageURL != "" { + task["docker_image_url"] = dockerImageURL + } req := map[string]any{"task": task} if runOptions := validateConfigRunOptions(cfg); len(runOptions) > 0 {