Air cli drop requirements yaml

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.

experimental/air/cmd/runsubmit_test.go

@@ -41,14 +41,16 @@ 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)
 	require.Len(t, p.Environments, 1)
 	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

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)

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")
 }