From 5aca38756ef903cc88e2db42440ff8e3579195ab Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Fri, 31 Jul 2026 19:03:33 +0000 Subject: [PATCH] air run: send dependencies via environment spec, not requirements.yaml Dependencies declared under environment.dependencies (inline list or a requirements.yaml file path) were uploaded as a co-located requirements.yaml artifact in the run's launch directory. Send them instead on the Jobs serverless environment spec (environments[].spec.dependencies), the modern delivery mechanism. The server installs them identically, so this drops the requirements.yaml round-trip. The file form of environment.dependencies is now read at submit time and parsed into the dependency list (its version field is still resolved separately via environment.version). --- experimental/air/cmd/runsubmit.go | 12 +++-- experimental/air/cmd/runsubmit_test.go | 42 +++++++++++++++-- experimental/air/cmd/runupload.go | 65 +++++++++++++++----------- experimental/air/cmd/runupload_test.go | 47 ++++++++++++------- 4 files changed, 114 insertions(+), 52 deletions(-) diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index d29bea427ba..b80a903b244 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -47,7 +47,7 @@ func dlRuntimeImage(ctx context.Context, runtimeVersion string) string { // omitempty so the wire form matches the Python CLI (which never emits a bare // "false"). Jobs performs the retries — each attempt is a fresh AI Runtime // workload. -func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapshotResult) jobs.SubmitRun { +func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, deps []string, snap snapshotResult) jobs.SubmitRun { task := jobs.AiRuntimeTask{ Experiment: cfg.ExperimentName, Deployments: []jobs.DeploymentSpec{{ @@ -86,7 +86,9 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapsh Tasks: []jobs.SubmitTask{st}, Environments: []jobs.JobEnvironment{{ EnvironmentKey: aiRuntimeEnvironmentKey, - Spec: &compute.Environment{EnvironmentVersion: dlImage}, + // Dependencies ride the serverless environment spec (the modern Jobs mechanism); + // the server installs them the same way it did the co-located requirements.yaml. + Spec: &compute.Environment{EnvironmentVersion: dlImage, Dependencies: deps}, }}, } } @@ -150,6 +152,10 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run if err != nil { return 0, "", err } + deps, err := resolveDependencies(cfg, configPath) + if err != nil { + return 0, "", err + } if err := uploadArtifacts(ctx, fc, items); err != nil { return 0, "", err } @@ -167,7 +173,7 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run } runtimeVersion, _ := cfg.runtimeVersion() - payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion), snap) + payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion), deps, snap) payload.IdempotencyToken = token // Submit returns as soon as the run is created; we don't wait for it to finish. diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index 6418eed5ffe..ae6236bdafd 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -41,7 +41,7 @@ func TestBuildSubmitPayload(t *testing.T) { MLflowExperimentDirectory: new("/Workspace/Users/me/exp"), } - p := buildSubmitPayload(cfg, "/d/command.sh", "5", snapshotResult{}) + p := buildSubmitPayload(cfg, "/d/command.sh", "5", []string{"torch==2.4.0", "numpy"}, snapshotResult{}) assert.Equal(t, "exp", p.RunName) assert.Equal(t, 1800, p.TimeoutSeconds) @@ -49,6 +49,8 @@ func TestBuildSubmitPayload(t *testing.T) { assert.Equal(t, aiRuntimeEnvironmentKey, p.Environments[0].EnvironmentKey) require.NotNil(t, p.Environments[0].Spec) assert.Equal(t, "5", p.Environments[0].Spec.EnvironmentVersion) + // Dependencies ride the environment spec rather than a co-located requirements.yaml. + assert.Equal(t, []string{"torch==2.4.0", "numpy"}, p.Environments[0].Spec.Dependencies) require.Len(t, p.Tasks, 1) task := p.Tasks[0] @@ -76,7 +78,7 @@ func TestBuildSubmitPayloadDefaultRetries(t *testing.T) { Command: new("x"), Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, } - task := buildSubmitPayload(cfg, "/d/command.sh", "4", snapshotResult{}).Tasks[0] + task := buildSubmitPayload(cfg, "/d/command.sh", "4", nil, snapshotResult{}).Tasks[0] assert.Equal(t, defaultMaxRetries, task.MaxRetries) assert.True(t, task.RetryOnTimeout) } @@ -91,7 +93,7 @@ func TestBuildSubmitPayloadNoRetries(t *testing.T) { Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, MaxRetries: new(0), } - task := buildSubmitPayload(cfg, "/d/command.sh", "4", snapshotResult{}).Tasks[0] + task := buildSubmitPayload(cfg, "/d/command.sh", "4", nil, snapshotResult{}).Tasks[0] assert.Equal(t, 0, task.MaxRetries) assert.False(t, task.RetryOnTimeout) @@ -160,6 +162,40 @@ func TestSubmitWorkload(t *testing.T) { assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu1xH100, AcceleratorCount: 1}, d.Compute) } +// TestSubmitWorkloadSendsDependenciesOnEnvironment proves inline +// environment.dependencies reach the runs/submit payload on the environment spec +// (the modern Jobs mechanism) rather than as a co-located requirements.yaml artifact. +func TestSubmitWorkloadSendsDependenciesOnEnvironment(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + var got jobs.SubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return jobs.SubmitRunResponse{RunId: 777} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + cfg := minimalConfig + ` +environment: + dependencies: + - torch==2.4.0 + - numpy +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + _, _, err = submitWorkload(t.Context(), w, loaded, cfgPath, "idem-key") + require.NoError(t, err) + + require.Len(t, got.Environments, 1) + require.NotNil(t, got.Environments[0].Spec) + assert.Equal(t, []string{"torch==2.4.0", "numpy"}, got.Environments[0].Spec.Dependencies) +} + // TestSubmitWorkloadHonorsOverride proves a --override reaches the actual // runs/submit payload on a real submit, not just dry-run validation: the config // pins num_accelerators=1, the override bumps it to 4, and the recorded request diff --git a/experimental/air/cmd/runupload.go b/experimental/air/cmd/runupload.go index fb9ca00b987..d81ae6d6c39 100644 --- a/experimental/air/cmd/runupload.go +++ b/experimental/air/cmd/runupload.go @@ -17,12 +17,11 @@ import ( ) // Launch artifact basenames, uploaded into the run's cli_launch directory. The -// server-side launcher derives requirements.yaml / hyperparameters.yaml from the -// same directory, so these names are part of the contract. +// server-side launcher derives hyperparameters.yaml from the same directory, so +// these names are part of the contract. const ( trainingConfigName = "training_config.yaml" commandScriptName = "command.sh" - requirementsName = "requirements.yaml" hyperparametersName = "hyperparameters.yaml" envVarsName = "env_vars.json" secretEnvVarsName = "secret_env_vars.json" @@ -45,16 +44,44 @@ type fileWriter interface { Write(ctx context.Context, name string, reader io.Reader, mode ...filer.WriteMode) error } -// requirementsDoc mirrors the on-disk requirements.yaml format so the worker -// parses synthesized inline dependencies identically to a user-provided file. +// requirementsDoc mirrors the on-disk requirements.yaml format used by the file form of +// environment.dependencies. type requirementsDoc struct { Version string `yaml:"version,omitempty"` Dependencies []string `yaml:"dependencies"` } +// resolveDependencies returns the pip dependency list to send on the serverless environment +// (environments[].spec.dependencies). environment.dependencies is either an inline list or a path +// to a requirements.yaml file; for the file form the list is read here (its version field is +// resolved separately via environment.version). Returns nil when no dependencies are declared. +func resolveDependencies(cfg *runConfig, configPath string) ([]string, error) { + if deps, ok := cfg.inlineDependencies(); ok { + return deps, nil + } + reqPath, ok := cfg.requirementsFile() + if !ok { + return nil, nil + } + // Resolve a relative requirements path against the config's directory. + if !filepath.IsAbs(reqPath) { + reqPath = filepath.Join(filepath.Dir(configPath), reqPath) + } + data, err := os.ReadFile(reqPath) + if err != nil { + return nil, fmt.Errorf("failed to read requirements file %s: %w", reqPath, err) + } + var doc requirementsDoc + if err := yaml.Unmarshal(data, &doc); err != nil { + return nil, fmt.Errorf("failed to parse requirements file %s: %w", reqPath, err) + } + return doc.Dependencies, nil +} + // buildArtifacts assembles the files to upload for a run: the merged config, the -// inline command as a script, requirements (from a file or synthesized from -// inline dependencies), and hyperparameters. configPath is the local YAML path. +// inline command as a script, hyperparameters, and env var / secret sidecars. +// Dependencies are not uploaded here — they ride the environment spec (see +// resolveDependencies / buildSubmitPayload). configPath is the local YAML path. func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) { // TODO(DABs): with no _bases_/overrides ported yet, the merged config is the // file as-is; once those land, upload the re-serialized merged YAML instead. @@ -72,27 +99,9 @@ func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) { {commandScriptName, []byte(*cfg.Command)}, } - switch reqPath, ok := cfg.requirementsFile(); { - case ok: - // Resolve a relative requirements path against the config's directory. - if !filepath.IsAbs(reqPath) { - reqPath = filepath.Join(filepath.Dir(configPath), reqPath) - } - data, err := os.ReadFile(reqPath) - if err != nil { - return nil, fmt.Errorf("failed to read requirements file %s: %w", reqPath, err) - } - items = append(items, uploadItem{requirementsName, data}) - default: - if deps, ok := cfg.inlineDependencies(); ok { - version, _ := cfg.runtimeVersion() - data, err := yaml.Marshal(requirementsDoc{Version: version, Dependencies: deps}) - if err != nil { - return nil, fmt.Errorf("failed to synthesize requirements.yaml: %w", err) - } - items = append(items, uploadItem{requirementsName, data}) - } - } + // Dependencies are sent on the Jobs serverless environment (environments[].spec.dependencies, + // see buildSubmitPayload), not as a co-located requirements.yaml. The server installs them + // identically, so we no longer upload a requirements.yaml artifact. if len(cfg.Parameters) > 0 { data, err := yaml.Marshal(cfg.Parameters) diff --git a/experimental/air/cmd/runupload_test.go b/experimental/air/cmd/runupload_test.go index 0c87524735d..969d54a4d36 100644 --- a/experimental/air/cmd/runupload_test.go +++ b/experimental/air/cmd/runupload_test.go @@ -57,7 +57,7 @@ func TestBuildArtifacts_CommandAndConfig(t *testing.T) { assert.Equal(t, "python train.py", string(items[1].data)) } -func TestBuildArtifacts_InlineRequirementsAndParameters(t *testing.T) { +func TestBuildArtifacts_ParametersNoRequirements(t *testing.T) { path := writeConfigFile(t, "run.yaml", "x: y\n") cfg := &runConfig{ Command: new("echo hi"), @@ -68,19 +68,30 @@ func TestBuildArtifacts_InlineRequirementsAndParameters(t *testing.T) { Parameters: map[string]any{"lr": 0.1}, } + // Dependencies no longer produce a requirements.yaml artifact; they ride the + // environment spec (see resolveDependencies / buildSubmitPayload). items, err := buildArtifacts(cfg, path) require.NoError(t, err) - assert.Equal(t, []string{trainingConfigName, commandScriptName, requirementsName, hyperparametersName}, itemNames(items)) + assert.Equal(t, []string{trainingConfigName, commandScriptName, hyperparametersName}, itemNames(items)) +} - var reqIdx int - for i, it := range items { - if it.name == requirementsName { - reqIdx = i - } +func TestResolveDependencies_Inline(t *testing.T) { + path := writeConfigFile(t, "run.yaml", "x: y\n") + cfg := &runConfig{ + Environment: &environmentConfig{ + Dependencies: dependencies{set: true, isList: true, list: []string{"torch", "numpy"}}, + }, } - req := string(items[reqIdx].data) - assert.Contains(t, req, "version: \"5\"") - assert.Contains(t, req, "- torch") + deps, err := resolveDependencies(cfg, path) + require.NoError(t, err) + assert.Equal(t, []string{"torch", "numpy"}, deps) +} + +func TestResolveDependencies_None(t *testing.T) { + path := writeConfigFile(t, "run.yaml", "x: y\n") + deps, err := resolveDependencies(&runConfig{}, path) + require.NoError(t, err) + assert.Nil(t, deps) } func TestBuildArtifacts_EnvVarsAndSecrets(t *testing.T) { @@ -103,18 +114,19 @@ func TestBuildArtifacts_EnvVarsAndSecrets(t *testing.T) { assert.JSONEq(t, `[{"name":"HF_TOKEN","secret_scope":"myscope","secret_key":"hf"}]`, string(byName[secretEnvVarsName])) } -func TestBuildArtifacts_RequirementsFile(t *testing.T) { +func TestResolveDependencies_RequirementsFile(t *testing.T) { dir := t.TempDir() require.NoError(t, os.WriteFile(filepath.Join(dir, "run.yaml"), []byte("x: y\n"), 0o600)) - require.NoError(t, os.WriteFile(filepath.Join(dir, "reqs.yaml"), []byte("version: 4\n"), 0o600)) + // The file form of environment.dependencies carries the pip list under `dependencies`; + // the `version` field is resolved separately (environment.version), so it is ignored here. + require.NoError(t, os.WriteFile(filepath.Join(dir, "reqs.yaml"), []byte("version: 4\ndependencies:\n - torch\n - numpy\n"), 0o600)) cfg := &runConfig{ - Command: new("echo hi"), Environment: &environmentConfig{Dependencies: dependencies{set: true, isList: false, path: "reqs.yaml"}}, } - items, err := buildArtifacts(cfg, filepath.Join(dir, "run.yaml")) + deps, err := resolveDependencies(cfg, filepath.Join(dir, "run.yaml")) require.NoError(t, err) - assert.Contains(t, itemNames(items), requirementsName) + assert.Equal(t, []string{"torch", "numpy"}, deps) } func TestBuildArtifacts_OversizeConfigRejected(t *testing.T) { @@ -144,12 +156,11 @@ func TestUploadArtifacts_WriteError(t *testing.T) { require.ErrorContains(t, err, "failed to upload "+trainingConfigName) } -func TestBuildArtifacts_MissingRequirementsFile(t *testing.T) { +func TestResolveDependencies_MissingRequirementsFile(t *testing.T) { cfgPath := writeConfigFile(t, "run.yaml", "x: y\n") cfg := &runConfig{ - Command: new("echo hi"), Environment: &environmentConfig{Dependencies: dependencies{set: true, isList: false, path: "nope.yaml"}}, } - _, err := buildArtifacts(cfg, cfgPath) + _, err := resolveDependencies(cfg, cfgPath) require.ErrorContains(t, err, "failed to read requirements file") }