[air] Deliver env vars only via the common Jobs env-var API (no sidecar)

Match the Python CLI (#2237609): env vars and secrets ride the job-level
environment_variables profile the converter emits, not the legacy
env_vars.json / secret_env_vars.json workspace side-channel. Drop that
sidecar staging from buildArtifacts and delete the now-dead entry
encoders and filename constants.

The other launch artifacts buildArtifacts stages (command.sh,
training_config.yaml, requirements.yaml, hyperparameters.yaml) are
unchanged — those fixed the "requirements.yaml not found" harness error
and are unrelated to env vars.

Known gap (verified on staging): the vendored databricks-sdk-go
jobs.JobSettings/Task has no environment_variables field, so the DABs
bundle schema strips it client-side ("unknown field") before the wire.
Env vars therefore aren't delivered on the DABs path until the SDK
carries the common Jobs env-var API. The Jobs server already accepts the
field (gated by allowEnvironmentVariables); this is purely a
DABs/SDK-schema lag.

Co-authored-by: Isaac

Author: vinchenzo-db

experimental/air/cmd/rundabs.go

@@ -107,23 +107,28 @@ func renderBundle(cfg *runConfig, configPath string) (string, error) {
 	return string(body) + bundleTargetsBlock(), nil
 }
 
-// writeBundleProject renders databricks.yml (via renderBundle) plus command.sh and,
-// for a code_source snapshot, the user's code tree, into a temp bundle root. Deploy
-// uploads the whole root as an immutable-folder snapshot. Returns the root and a
-// cleanup func.
+// writeBundleProject renders databricks.yml plus the launch artifacts the AI Runtime
+// harness reads (command.sh, training_config.yaml, requirements.yaml, and — when
+// present — hyperparameters.yaml and env-var sidecars) and, for a code_source
+// snapshot, the user's code tree, into a temp bundle root. Deploy uploads the whole
+// root as an immutable-folder snapshot. Returns the root and a cleanup func.
 func writeBundleProject(ctx context.Context, cfg *runConfig, configPath string) (string, func(), error) {
 	body, err := renderBundle(cfg, configPath)
 	if err != nil {
 		return "", func() {}, err
 	}
+	artifacts, err := buildArtifacts(cfg, configPath)
+	if err != nil {
+		return "", func() {}, err
+	}
 	bundleRoot, err := os.MkdirTemp("", "air-dabs-*")
 	if err != nil {
 		return "", func() {}, err
 	}
 	cleanup := func() { _ = os.RemoveAll(bundleRoot) }
 
-	// Copy the code_source working tree into the bundle root first, so a stray
-	// command.sh / databricks.yml in the user's tree can't shadow ours below.
+	// Copy the code_source working tree first, so a stray command.sh / databricks.yml
+	// in the user's tree can't shadow the generated files written below.
 	if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil {
 		if err := stageCodeSource(ctx, cfg.CodeSource.Snapshot, configPath, bundleRoot); err != nil {
 			cleanup()
@@ -136,9 +141,10 @@ func writeBundleProject(ctx context.Context, cfg *runConfig, configPath string)
 		return "", func() {}, err
 	}
 
-	// command.sh: the entrypoint the task's command_path points at.
-	if cfg.Command != nil {
-		if err := os.WriteFile(filepath.Join(bundleRoot, bundleCommandScript), []byte(*cfg.Command), 0o600); err != nil {
+	// The launch artifacts the harness expects co-located with command.sh, the same
+	// set the retired ephemeral path uploaded (buildArtifacts).
+	for _, it := range artifacts {
+		if err := os.WriteFile(filepath.Join(bundleRoot, it.name), it.data, 0o600); err != nil {
 			cleanup()
 			return "", func() {}, err
 		}

experimental/air/cmd/rundabs_test.go

@@ -15,13 +15,11 @@ import (
 )
 
 func TestWriteBundleProject(t *testing.T) {
-	cfg := &runConfig{
-		ExperimentName: "exp",
-		Command:        new("echo hi"),
-		Compute:        &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1},
-	}
+	configPath := writeConfigFile(t, "train.yaml", minimalConfig)
+	cfg, err := loadRunConfig(configPath)
+	require.NoError(t, err)
 
-	root, cleanup, err := writeBundleProject(t.Context(), cfg, "train.yaml")
+	root, cleanup, err := writeBundleProject(t.Context(), cfg, configPath)
 	require.NoError(t, err)
 	defer cleanup()
 
@@ -32,10 +30,12 @@ func TestWriteBundleProject(t *testing.T) {
 	assert.Contains(t, string(body), "immutable_folder: true")
 	assert.Contains(t, string(body), "mode: development")
 
-	// command.sh holds the run's command verbatim.
+	// The launch artifacts the AI Runtime harness reads are staged next to command.sh.
+	assert.FileExists(t, filepath.Join(root, bundleCommandScript))
+	assert.FileExists(t, filepath.Join(root, trainingConfigName))
 	script, err := os.ReadFile(filepath.Join(root, bundleCommandScript))
 	require.NoError(t, err)
-	assert.Equal(t, "echo hi", string(script))
+	assert.Equal(t, "python train.py", string(script))
 
 	// cleanup removes the temp root.
 	cleanup()
@@ -49,14 +49,16 @@ func TestWriteBundleProjectStagesCodeSource(t *testing.T) {
 	require.NoError(t, os.MkdirAll(filepath.Join(src, "pkg"), 0o755))
 	require.NoError(t, os.WriteFile(filepath.Join(src, "pkg", "mod.py"), []byte("x=1"), 0o644))
 
-	cfg := &runConfig{
-		ExperimentName: "exp",
-		Command:        new("python train.py"),
-		Compute:        &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1},
-		CodeSource:     &codeSourceConfig{Type: "snapshot", Snapshot: &snapshotSourceConfig{RootPath: src}},
-	}
+	configPath := writeConfigFile(t, "train.yaml", minimalConfig+`
+code_source:
+  type: snapshot
+  snapshot:
+    root_path: `+src+`
+`)
+	cfg, err := loadRunConfig(configPath)
+	require.NoError(t, err)
 
-	root, cleanup, err := writeBundleProject(t.Context(), cfg, "train.yaml")
+	root, cleanup, err := writeBundleProject(t.Context(), cfg, configPath)
 	require.NoError(t, err)
 	defer cleanup()
 
@@ -69,35 +71,42 @@ func TestWriteBundleProjectStagesCodeSource(t *testing.T) {
 
 func TestWriteBundleProjectCommandShadowProtection(t *testing.T) {
 	// A command.sh in the user's tree must not shadow the one we generate from the
-	// run's command: our write happens after the tree copy.
+	// run's command: the artifact writes happen after the tree copy.
 	src := t.TempDir()
 	require.NoError(t, os.WriteFile(filepath.Join(src, bundleCommandScript), []byte("STALE"), 0o644))
 
-	cfg := &runConfig{
-		ExperimentName: "exp",
-		Command:        new("FRESH"),
-		Compute:        &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1},
-		CodeSource:     &codeSourceConfig{Type: "snapshot", Snapshot: &snapshotSourceConfig{RootPath: src}},
-	}
+	configPath := writeConfigFile(t, "train.yaml", minimalConfig+`
+code_source:
+  type: snapshot
+  snapshot:
+    root_path: `+src+`
+`)
+	cfg, err := loadRunConfig(configPath)
+	require.NoError(t, err)
 
-	root, cleanup, err := writeBundleProject(t.Context(), cfg, "train.yaml")
+	root, cleanup, err := writeBundleProject(t.Context(), cfg, configPath)
 	require.NoError(t, err)
 	defer cleanup()
 
+	// minimalConfig's command is "python train.py", not the STALE tree copy.
 	script, err := os.ReadFile(filepath.Join(root, bundleCommandScript))
 	require.NoError(t, err)
-	assert.Equal(t, "FRESH", string(script))
+	assert.Equal(t, "python train.py", string(script))
 }
 
 func TestWriteBundleProjectRejectsUnconvertible(t *testing.T) {
 	// A gate failure (here: a $CODE_SOURCE_PATH command) surfaces before any temp
 	// directory is created.
-	cfg := &runConfig{
-		ExperimentName: "exp",
-		Command:        new("cd $CODE_SOURCE_PATH && python train.py"),
-		Compute:        &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1},
-	}
-	_, _, err := writeBundleProject(t.Context(), cfg, "train.yaml")
+	configPath := writeConfigFile(t, "train.yaml", `
+experiment_name: exp
+command: cd $CODE_SOURCE_PATH && python train.py
+compute:
+  accelerator_type: GPU_1xA10
+  num_accelerators: 1
+`)
+	cfg, err := loadRunConfig(configPath)
+	require.NoError(t, err)
+	_, _, err = writeBundleProject(t.Context(), cfg, configPath)
 	require.Error(t, err)
 }
 

experimental/air/cmd/runupload.go

@@ -3,14 +3,10 @@ package aircmd
 import (
 	"bytes"
 	"context"
-	"encoding/json"
 	"fmt"
 	"io"
-	"maps"
 	"os"
 	"path/filepath"
-	"slices"
-	"strings"
 
 	"github.com/databricks/cli/libs/filer"
 	"go.yaml.in/yaml/v3"
@@ -24,8 +20,6 @@ const (
 	commandScriptName   = "command.sh"
 	requirementsName    = "requirements.yaml"
 	hyperparametersName = "hyperparameters.yaml"
-	envVarsName         = "env_vars.json"
-	secretEnvVarsName   = "secret_env_vars.json"
 )
 
 // maxConfigYAMLBytes caps training_config.yaml. It is referenced by the Jobs
@@ -102,59 +96,14 @@ func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) {
 		items = append(items, uploadItem{hyperparametersName, 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 {
-		data, err := json.Marshal(envVarEntries(cfg.EnvVariables))
-		if err != nil {
-			return nil, fmt.Errorf("failed to serialize env_variables: %w", err)
-		}
-		items = append(items, uploadItem{envVarsName, data})
-	}
-	if len(cfg.Secrets) > 0 {
-		data, err := json.Marshal(secretEnvVarEntries(cfg.Secrets))
-		if err != nil {
-			return nil, fmt.Errorf("failed to serialize secrets: %w", err)
-		}
-		items = append(items, uploadItem{secretEnvVarsName, data})
-	}
+	// Env vars and secrets are NOT staged as sidecar files: they ride the common
+	// Jobs env-var API (the environment_variables profile the converter emits, see
+	// exportbundle.go), matching the Python CLI which dropped the env_vars.json /
+	// secret_env_vars.json side-channel.
 
 	return items, nil
 }
 
-// envVarEntry is one entry in env_vars.json.
-type envVarEntry struct {
-	Name  string `json:"name"`
-	Value string `json:"value"`
-}
-
-// secretEnvVarEntry is one entry in secret_env_vars.json. The YAML side is
-// {ENV_VAR: "scope/key"}; the launcher wants the split form.
-type secretEnvVarEntry struct {
-	Name        string `json:"name"`
-	SecretScope string `json:"secret_scope"`
-	SecretKey   string `json:"secret_key"`
-}
-
-// envVarEntries renders env_variables sorted by name for deterministic output.
-func envVarEntries(vars map[string]string) []envVarEntry {
-	out := make([]envVarEntry, 0, len(vars))
-	for _, name := range slices.Sorted(maps.Keys(vars)) {
-		out = append(out, envVarEntry{Name: name, Value: vars[name]})
-	}
-	return out
-}
-
-// secretEnvVarEntries renders secrets sorted by name for deterministic output.
-func secretEnvVarEntries(secrets map[string]string) []secretEnvVarEntry {
-	out := make([]secretEnvVarEntry, 0, len(secrets))
-	for _, name := range slices.Sorted(maps.Keys(secrets)) {
-		scope, key, _ := strings.Cut(secrets[name], "/")
-		out = append(out, secretEnvVarEntry{Name: name, SecretScope: scope, SecretKey: key})
-	}
-	return out
-}
-
 // uploadArtifacts writes each artifact into the launch directory, overwriting and
 // creating parents as needed.
 //

experimental/air/cmd/runupload_test.go

@@ -83,7 +83,10 @@ func TestBuildArtifacts_InlineRequirementsAndParameters(t *testing.T) {
 	assert.Contains(t, req, "- torch")
 }
 
-func TestBuildArtifacts_EnvVarsAndSecrets(t *testing.T) {
+func TestBuildArtifacts_EnvVarsNotStagedAsSidecars(t *testing.T) {
+	// Env vars and secrets ride the common Jobs env-var API (the converter's
+	// environment_variables profile), not sidecar files — matching the Python CLI,
+	// which dropped the env_vars.json / secret_env_vars.json side-channel.
 	path := writeConfigFile(t, "run.yaml", "x: y\n")
 	cfg := &runConfig{
 		Command:      new("echo hi"),
@@ -93,14 +96,8 @@ func TestBuildArtifacts_EnvVarsAndSecrets(t *testing.T) {
 
 	items, err := buildArtifacts(cfg, path)
 	require.NoError(t, err)
-	assert.Subset(t, itemNames(items), []string{envVarsName, secretEnvVarsName})
-
-	byName := map[string][]byte{}
-	for _, it := range items {
-		byName[it.name] = it.data
-	}
-	assert.JSONEq(t, `[{"name":"WANDB","value":"demo"}]`, string(byName[envVarsName]))
-	assert.JSONEq(t, `[{"name":"HF_TOKEN","secret_scope":"myscope","secret_key":"hf"}]`, string(byName[secretEnvVarsName]))
+	assert.NotContains(t, itemNames(items), "env_vars.json")
+	assert.NotContains(t, itemNames(items), "secret_env_vars.json")
 }
 
 func TestBuildArtifacts_RequirementsFile(t *testing.T) {