Skip to content

experimental/air: warm snapshot cache for the plain_tar path #6572

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ben-hansen-db wants to merge 4 commits into
base: air-plain-tar-pgzip
Choose a base branch
from
Open

experimental/air: warm snapshot cache for the plain_tar path #6572

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4be3730
experimental/air: warm snapshot cache for the plain_tar path
ben-hansen-db Sep 8, 2026
85981d5
experimental/air: use default gzip level in the warm cache too
ben-hansen-db Sep 8, 2026
8bfd8da
experimental/air: update config-help golden for --no-cache flag
ben-hansen-db Sep 8, 2026
ed64774
experimental/air: make the warm cache crash- and concurrency-safe
ben-hansen-db Sep 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions acceptance/experimental/air/config-help/output.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ Flags:
-f, --file string Path to the workload YAML config
-h, --help help for run
--idempotency-key string Return the existing run if this key was already used
--no-cache Bypass the local snapshot cache and re-pack the code tarball from scratch
--override stringArray Override a YAML field, e.g. compute.num_accelerators=8 (repeatable)
--watch Stream logs until the run completes

Expand Down
4 changes: 3 additions & 1 deletion experimental/air/cmd/run.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ func newRunCommand() *cobra.Command {
overrides []string
dryRun bool
idempotencyKey string
noCache bool
)

cmd := &cobra.Command{
Expand Down Expand Up @@ -78,6 +79,7 @@ The path must be a separate argument: cobra reserves -h as a boolean, so
cmd.Flags().StringArrayVar(&overrides, "override", nil, "Override a YAML field, e.g. compute.num_accelerators=8 (repeatable)")
cmd.Flags().BoolVar(&dryRun, "dry-run", false, "Validate the config without submitting")
cmd.Flags().StringVar(&idempotencyKey, "idempotency-key", "", "Return the existing run if this key was already used")
cmd.Flags().BoolVar(&noCache, "no-cache", false, "Bypass the local snapshot cache and re-pack the code tarball from scratch")
_ = cmd.MarkFlagRequired("file")

// --dry-run only validates the config locally, so it needs no workspace.
Expand Down Expand Up @@ -114,7 +116,7 @@ The path must be a separate argument: cobra reserves -h as a boolean, so
}

w := cmdctx.WorkspaceClient(ctx)
runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey, !jsonOut)
runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey, !jsonOut, noCache)
if err != nil {
return err
}
Expand Down
4 changes: 2 additions & 2 deletions experimental/air/cmd/runsubmit.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -235,7 +235,7 @@ func withSpinner(ctx context.Context, show bool, msg string, fn func() error) er
// upload the launch artifacts, assemble the Jobs payload, and submit it. It
// returns the new run_id and its dashboard URL. showProgress enables the
// stderr upload/packaging spinners (text mode only).
func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress bool) (int64, string, error) {
func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress, noCache bool) (int64, string, error) {
// Compute the launch dir and command_path up front — a read-only workspace lookup plus a
// local path build, no writes yet — so the pre-flight validates the real command_path. The
// same path is reused for the upload and submit below, so the validated path is the submitted
Expand Down Expand Up @@ -320,7 +320,7 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run
// Sidecars land in the run's launch dir (funcDir) via fc, next to command.sh.
err = withSpinner(ctx, showProgress, "Packaging code snapshot…", func() error {
var e error
snap, e = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir)
snap, e = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir, noCache)
return e
})
if err != nil {
Expand Down
24 changes: 12 additions & 12 deletions experimental/air/cmd/runsubmit_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,7 +242,7 @@ func TestSubmitWorkload(t *testing.T) {
cfg, err := loadRunConfig(cfgPath)
require.NoError(t, err)

runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false)
runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false, false)
require.NoError(t, err)
assert.Equal(t, int64(777), runID)
assert.Contains(t, dashboardURL, "/jobs/runs/777")
Expand Down Expand Up @@ -286,7 +286,7 @@ func TestSubmitWorkloadHonorsOverride(t *testing.T) {
cfg, err := loadRunConfigWithOverrides(t.Context(), cfgPath, []string{"compute.num_accelerators=4"})
require.NoError(t, err)

_, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false)
_, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false, false)
require.NoError(t, err)

require.Len(t, got.Tasks, 1)
Expand Down Expand Up @@ -329,7 +329,7 @@ code_source:

// The DABs upload path logs via cmdio; the real `air run` context carries it.
ctx := cmdio.MockDiscard(t.Context())
_, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false)
_, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false, false)
require.NoError(t, err)

at := got.Tasks[0].AiRuntimeTask
Expand Down Expand Up @@ -373,7 +373,7 @@ code_source:
require.NoError(t, err)

ctx := cmdio.MockDiscard(t.Context())
_, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false)
_, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false, false)
require.NoError(t, err)

at := got.Tasks[0].AiRuntimeTask
Expand Down Expand Up @@ -423,7 +423,7 @@ code_source:
// The uploaded name carries a discriminator (timestamp), not the bare dir name.
ctx := cmdio.MockDiscard(t.Context())
sidecarStore, sidecarBase := testSidecarStore(t, w)
snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase)
snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase, false)
require.NoError(t, err)
base := path.Base(snap.CodeSourcePath)
assert.NotEqual(t, "src.tar.gz", base, "plain-tar name must be unique, not the bare dir name")
Expand Down Expand Up @@ -475,9 +475,9 @@ code_source:

ctx := cmdio.MockDiscard(t.Context())
sidecarStore, sidecarBase := testSidecarStore(t, w)
first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase)
first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase, false)
require.NoError(t, err)
second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase)
second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase, false)
require.NoError(t, err)

// Same pinned commit → identical content-addressed remote path, uploaded once
Expand Down Expand Up @@ -519,7 +519,7 @@ code_source:

ctx := cmdio.MockDiscard(t.Context())
sidecarStore, sidecarBase := testSidecarStore(t, w)
snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase)
snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase, false)
require.NoError(t, err)

assert.Empty(t, snap.GitStatePath)
Expand Down Expand Up @@ -568,7 +568,7 @@ code_source:
require.NoError(t, err)

ctx := cmdio.MockDiscard(t.Context())
_, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false)
_, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false, false)
require.NoError(t, err)

at := got.Tasks[0].AiRuntimeTask
Expand Down Expand Up @@ -603,7 +603,7 @@ func TestSubmitWorkloadGuards(t *testing.T) {

cfg := *base
cfg.UsagePolicyName = new("nope")
_, _, err = submitWorkload(t.Context(), pw, &cfg, cfgPath, "", false)
_, _, err = submitWorkload(t.Context(), pw, &cfg, cfgPath, "", false, false)
require.ErrorContains(t, err, `no usage policy named "nope"`)
for _, p := range paths {
assert.NotContains(t, p, "/workspace/", "no workspace write may precede policy resolution")
Expand Down Expand Up @@ -640,7 +640,7 @@ func TestSubmitWorkloadSendsUsagePolicy(t *testing.T) {
cfg, err := loadRunConfig(cfgPath)
require.NoError(t, err)

_, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false)
_, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false, false)
require.NoError(t, err)
assert.Equal(t, policyID, got.BudgetPolicyId)
})
Expand All @@ -651,7 +651,7 @@ func TestSubmitWorkloadSendsUsagePolicy(t *testing.T) {
cfg, err := loadRunConfig(cfgPath)
require.NoError(t, err)

_, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false)
_, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false, false)
require.NoError(t, err)
assert.Equal(t, policyID, got.BudgetPolicyId)
})
Expand Down
Loading
Loading