From 219188bfcf23a6a7e0d774a41ab2ba1e34da73eb Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Tue, 11 Aug 2026 21:32:15 +0000 Subject: [PATCH 1/3] air/convert-to-dabs: carry a schedule onto the bundle job MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A run YAML may now set a `schedule` block (quartz_cron_expression, timezone_id, and optional pause_status). `air run` submits a one-time run and can't honor it, so it detects the schedule and offers to convert the config to a bundle instead — where `bundle deploy` creates the persistent, scheduled job. convert-to-dabs maps the block 1:1 onto the bundle job's schedule. Co-authored-by: Isaac --- .../air/convert-to-dabs/output.txt | 55 +++++++++++++++++++ .../air/convert-to-dabs/schedule.yaml | 13 +++++ .../experimental/air/convert-to-dabs/script | 4 ++ experimental/air/cmd/convert_to_dabs.go | 13 +++++ experimental/air/cmd/convert_to_dabs_test.go | 42 ++++++++++++++ experimental/air/cmd/run.go | 39 +++++++++++++ experimental/air/cmd/runconfig.go | 34 ++++++++++++ experimental/air/cmd/runconfig_test.go | 12 ++++ 8 files changed, 212 insertions(+) create mode 100644 acceptance/experimental/air/convert-to-dabs/schedule.yaml diff --git a/acceptance/experimental/air/convert-to-dabs/output.txt b/acceptance/experimental/air/convert-to-dabs/output.txt index 8b94bf55a7c..a508140eb6c 100644 --- a/acceptance/experimental/air/convert-to-dabs/output.txt +++ b/acceptance/experimental/air/convert-to-dabs/output.txt @@ -97,3 +97,58 @@ job and its uploaded files with: Error: environment.docker_image is not yet supported by convert-to-dabs Exit code: 1 + +=== a schedule maps onto the bundle job (air run can't schedule; the bundle can) +>>> [CLI] experimental air convert-to-dabs schedule.yaml --force +Wrote a Databricks Asset Bundle to .: + databricks.yml + generated_artifacts/training_config.yaml + generated_artifacts/command.sh + +To deploy and run this workload as a bundle: + 1. [CLI] bundle validate + 2. [CLI] bundle deploy + 3. [CLI] bundle run torchrun-a10-smoke-test + +bundle deploy uploads the code source and launch scripts automatically. +To see what it deployed and where: [CLI] bundle summary + +Unlike `air run` (which submits an ephemeral run), bundle deploy creates a +persistent job that is not garbage-collected. When you are done, remove the +job and its uploaded files with: + [CLI] bundle destroy + +>>> cat databricks.yml +bundle: + name: torchrun-a10-smoke-test +sync: + paths: + - generated_artifacts +targets: + dev: + mode: development + default: true +resources: + jobs: + torchrun-a10-smoke-test: + name: torchrun-a10-smoke-test + tasks: + - task_key: torchrun-a10-smoke-test + environment_key: default + max_retries: 3 + ai_runtime_task: + experiment: torchrun-a10-smoke-test + deployments: + - command_path: ./generated_artifacts/command.sh + compute: + accelerator_type: GPU_1xA10 + accelerator_count: 1 + code_source_path: ./src + environments: + - environment_key: default + spec: + environment_version: "4" + schedule: + quartz_cron_expression: 0 0 9 * * ? + timezone_id: America/Los_Angeles + pause_status: PAUSED diff --git a/acceptance/experimental/air/convert-to-dabs/schedule.yaml b/acceptance/experimental/air/convert-to-dabs/schedule.yaml new file mode 100644 index 00000000000..54d2e80c25b --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/schedule.yaml @@ -0,0 +1,13 @@ +experiment_name: torchrun-a10-smoke-test +command: torchrun --nproc_per_node=1 train.py +compute: + accelerator_type: GPU_1xA10 + num_accelerators: 1 +code_source: + type: snapshot + snapshot: + root_path: ./src +schedule: + quartz_cron_expression: "0 0 9 * * ?" + timezone_id: America/Los_Angeles + pause_status: PAUSED diff --git a/acceptance/experimental/air/convert-to-dabs/script b/acceptance/experimental/air/convert-to-dabs/script index 77b6e06938a..cb7eabeb081 100644 --- a/acceptance/experimental/air/convert-to-dabs/script +++ b/acceptance/experimental/air/convert-to-dabs/script @@ -18,3 +18,7 @@ trace $CLI experimental air convert-to-dabs train.yaml --force title "docker_image is not supported yet" errcode trace $CLI experimental air convert-to-dabs docker.yaml --output-dir generated-docker + +title "a schedule maps onto the bundle job (air run can't schedule; the bundle can)" +trace $CLI experimental air convert-to-dabs schedule.yaml --force +trace cat databricks.yml diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index f5563136176..d0e032ddb2a 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -301,6 +301,19 @@ func buildBundleValue(ctx context.Context, cfg *runConfig, configPath, codeSourc if perms := buildPermissionsValue(cfg.Permissions); perms.Kind() != dyn.KindInvalid { job["permissions"] = nv(perms.MustSequence(), 5) } + // schedule maps 1:1 onto the job's CronSchedule. `air run` can't honor it (it + // submits a one-time run); it exists so a converted bundle deploys as a recurring + // job. pause_status is omitted when unset so the Jobs default (UNPAUSED) applies. + if cfg.Schedule != nil { + sched := map[string]dyn.Value{ + "quartz_cron_expression": nv(cfg.Schedule.QuartzCronExpression, 1), + "timezone_id": nv(cfg.Schedule.TimezoneID, 2), + } + if cfg.Schedule.PauseStatus != "" { + sched["pause_status"] = nv(cfg.Schedule.PauseStatus, 3) + } + job["schedule"] = nv(sched, 6) + } rootValue := map[string]dyn.Value{ "bundle": nv(map[string]dyn.Value{ diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index 0350f010367..ddcc87514ba 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -357,6 +357,48 @@ func TestConvertToDabsMapsUsagePolicyID(t *testing.T) { assert.Equal(t, "12345678-90ab-cdef-1234-567890abcdef", get(t, root, "resources.jobs."+loaded.ExperimentName+".budget_policy_id").MustString()) } +// A schedule block maps 1:1 onto the bundle job's schedule (the field air run +// can't honor but convert-to-dabs carries so bundle deploy schedules the job). +func TestConvertToDabsMapsSchedule(t *testing.T) { + cfg := minimalConfig + ` +schedule: + quartz_cron_expression: "0 0 9 * * ?" + timezone_id: America/Los_Angeles + pause_status: PAUSED +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) + require.NoError(t, err) + + sched := "resources.jobs." + loaded.ExperimentName + ".schedule" + assert.Equal(t, "0 0 9 * * ?", get(t, root, sched+".quartz_cron_expression").MustString()) + assert.Equal(t, "America/Los_Angeles", get(t, root, sched+".timezone_id").MustString()) + assert.Equal(t, "PAUSED", get(t, root, sched+".pause_status").MustString()) +} + +// pause_status is optional; when omitted it is left off the emitted schedule so +// the Jobs default (UNPAUSED) applies. +func TestConvertToDabsScheduleOmitsEmptyPauseStatus(t *testing.T) { + cfg := minimalConfig + ` +schedule: + quartz_cron_expression: "0 0 9 * * ?" + timezone_id: UTC +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) + require.NoError(t, err) + + sched := "resources.jobs." + loaded.ExperimentName + ".schedule" + assert.True(t, has(root, sched+".quartz_cron_expression")) + assert.False(t, has(root, sched+".pause_status"), "empty pause_status must be omitted") +} + func TestConvertToDabsMapsPermissions(t *testing.T) { cfg := minimalConfig + ` permissions: diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index ea00368679f..5931452e77c 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -2,7 +2,9 @@ package aircmd import ( "context" + "errors" "fmt" + "path/filepath" "strconv" "github.com/databricks/cli/cmd/root" @@ -71,6 +73,13 @@ The workload is described by a YAML config file (see --file).`, return renderEnvelope(ctx, runResult{Status: "DRY_RUN_OK", DryRun: true}) } + // A schedule needs a persistent job; `air run` only submits a one-time run, so + // it can't honor one. Offer to convert to a bundle (where deploy schedules it) + // rather than silently submitting an unscheduled run. + if cfg.Schedule != nil { + return handleScheduledRun(ctx, cfg, file) + } + w := cmdctx.WorkspaceClient(ctx) runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey) if err != nil { @@ -128,6 +137,36 @@ The workload is described by a YAML config file (see --file).`, return cmd } +// handleScheduledRun responds to a `schedule` in the run config. `air run` submits a +// one-time run and can't schedule it; scheduling needs a persistent job, which a +// Databricks Asset Bundle provides. It offers to convert the config to a bundle +// (writing it next to the YAML) and prints the deploy step; it never submits, since a +// scheduled run can't go through the submit path. Declining, or a non-interactive +// session, exits non-zero with the manual command. +func handleScheduledRun(ctx context.Context, cfg *runConfig, configPath string) error { + self := cliInvocation() + cmdio.LogString(ctx, "'schedule' is set, but 'air run' submits a one-time run and can't schedule it.") + cmdio.LogString(ctx, "Scheduling needs a persistent job, which a Databricks Asset Bundle provides.") + + convertCmd := fmt.Sprintf("%s experimental air convert-to-dabs %s", self, configPath) + manual := fmt.Sprintf("Nothing submitted. To schedule this job:\n %s\n databricks bundle deploy", convertCmd) + + convert, err := cmdio.AskYesOrNo(ctx, "Convert this config to a bundle now?") + if err != nil || !convert { + // A non-interactive session (no TTY) surfaces as an error here; treat it the + // same as declining and point at the manual command. + return errors.New(manual) + } + + dir := filepath.Dir(configPath) + written, err := writeBundle(ctx, cfg, configPath, dir, false) + if err != nil { + return err + } + printConvertNextSteps(ctx, dir, written, bundleResourceKey(cfg.ExperimentName)) + return nil +} + // watchTerminalStatus resolves a watched run's final display state for the // closing --watch envelope. The run is terminal once streaming returns; if the // status can't be re-fetched, "UNKNOWN" is reported rather than guessing. diff --git a/experimental/air/cmd/runconfig.go b/experimental/air/cmd/runconfig.go index 96dcf96d5a1..857731be42d 100644 --- a/experimental/air/cmd/runconfig.go +++ b/experimental/air/cmd/runconfig.go @@ -57,6 +57,10 @@ type runConfig struct { Permissions []permission `yaml:"permissions"` UsagePolicyName *string `yaml:"usage_policy_name"` UsagePolicyID *string `yaml:"usage_policy_id"` + // Schedule turns the workload into a recurring job. `air run` submits a one-time + // run and cannot honor it (see run.go); it is carried through convert-to-dabs, + // which emits it onto the bundle job where `bundle deploy` schedules it. + Schedule *scheduleConfig `yaml:"schedule"` } // validate runs structural validation over the whole config, returning the first @@ -179,6 +183,12 @@ func (c *runConfig) validate() error { } } + if c.Schedule != nil { + if err := c.Schedule.validate(); err != nil { + return err + } + } + return nil } @@ -410,6 +420,30 @@ func (s *snapshotSourceConfig) validate() error { return nil } +// scheduleConfig mirrors the Jobs CronSchedule proto so it maps 1:1 onto the +// bundle job's schedule block (see convert_to_dabs.go). pause_status is optional +// and defaults to UNPAUSED, matching the Jobs default. +type scheduleConfig struct { + QuartzCronExpression string `yaml:"quartz_cron_expression"` + TimezoneID string `yaml:"timezone_id"` + PauseStatus string `yaml:"pause_status"` +} + +func (s *scheduleConfig) validate() error { + if strings.TrimSpace(s.QuartzCronExpression) == "" { + return errors.New("schedule.quartz_cron_expression is required") + } + if strings.TrimSpace(s.TimezoneID) == "" { + return errors.New("schedule.timezone_id is required (for example, 'America/Los_Angeles' or 'UTC')") + } + switch s.PauseStatus { + case "", "PAUSED", "UNPAUSED": + default: + return fmt.Errorf("schedule.pause_status must be PAUSED or UNPAUSED, got %q", s.PauseStatus) + } + return nil +} + // gitRef pins a snapshot to a specific git ref. branch and commit are mutually // exclusive; remote is only meaningful with branch. type gitRef struct { diff --git a/experimental/air/cmd/runconfig_test.go b/experimental/air/cmd/runconfig_test.go index cd7804cd4af..20681c80dbc 100644 --- a/experimental/air/cmd/runconfig_test.go +++ b/experimental/air/cmd/runconfig_test.go @@ -247,6 +247,18 @@ func TestRunConfigValidate_FieldRules(t *testing.T) { // A name pasted into the id field gets pointed at the right field. {"policy name in usage_policy_id", func(c *runConfig) { c.UsagePolicyID = str("team-a") }, "use usage_policy_name"}, {"uuid usage_policy_id alone is ok", func(c *runConfig) { c.UsagePolicyID = str("12345678-90ab-cdef-1234-567890abcdef") }, ""}, + {"schedule missing cron", func(c *runConfig) { + c.Schedule = &scheduleConfig{TimezoneID: "UTC"} + }, "quartz_cron_expression is required"}, + {"schedule missing timezone", func(c *runConfig) { + c.Schedule = &scheduleConfig{QuartzCronExpression: "0 0 9 * * ?"} + }, "timezone_id is required"}, + {"schedule bad pause_status", func(c *runConfig) { + c.Schedule = &scheduleConfig{QuartzCronExpression: "0 0 9 * * ?", TimezoneID: "UTC", PauseStatus: "MAYBE"} + }, "pause_status must be PAUSED or UNPAUSED"}, + {"schedule ok", func(c *runConfig) { + c.Schedule = &scheduleConfig{QuartzCronExpression: "0 0 9 * * ?", TimezoneID: "UTC"} + }, ""}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { From ec35629abf081641af72d3d8d15fcbaba23ddb32 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Tue, 11 Aug 2026 22:26:31 +0000 Subject: [PATCH 2/3] air/run: add schedule to the --override available-fields golden Co-authored-by: Isaac --- acceptance/experimental/air/run/output.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/acceptance/experimental/air/run/output.txt b/acceptance/experimental/air/run/output.txt index 13c9b360ece..bad81ede14a 100644 --- a/acceptance/experimental/air/run/output.txt +++ b/acceptance/experimental/air/run/output.txt @@ -22,7 +22,7 @@ Dry run: configuration for "smoke-test" is valid; not submitting. === override of an unknown field is rejected >>> [CLI] experimental air run -f valid.yaml --dry-run --override bogus=1 -Error: invalid --override "bogus": "bogus" is not a known field; available fields are: code_source, command, compute, env_variables, environment, experiment_name, idempotency_token, max_retries, mlflow_experiment_directory, mlflow_run_name, parameters, permissions, secrets, timeout_minutes, usage_policy_id, usage_policy_name +Error: invalid --override "bogus": "bogus" is not a known field; available fields are: code_source, command, compute, env_variables, environment, experiment_name, idempotency_token, max_retries, mlflow_experiment_directory, mlflow_run_name, parameters, permissions, schedule, secrets, timeout_minutes, usage_policy_id, usage_policy_name Exit code: 1 From 5aa03c813b1098c9567d2f36c5bf80971198b21b Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Tue, 11 Aug 2026 23:36:31 +0000 Subject: [PATCH 3/3] air/run: create a persistent scheduled job natively (no DABs) Pivot the schedule handling: instead of offering to convert to a bundle, a run config with a `schedule` block now makes `air run` create (or update) a persistent, scheduled Databricks job directly via jobs/create. `air run` already uploads the code snapshot + launch scripts to a durable workspace path and builds the ai_runtime_task; this reuses that prep and swaps the ephemeral jobs/runs/submit for a jobs/create carrying a CronSchedule. Since the CLI keeps no state, the job is upserted by name: an existing job with the same experiment_name is reset in place (so re-running doesn't pile up duplicates), and an ambiguous name (>1 match) errors rather than guessing. convert-to-dabs still emits the schedule block, so converting a scheduled config to a bundle keeps the schedule. Co-authored-by: Isaac --- experimental/air/cmd/run.go | 61 +++++------ experimental/air/cmd/runschedule.go | 128 +++++++++++++++++++++++ experimental/air/cmd/runschedule_test.go | 113 ++++++++++++++++++++ experimental/air/cmd/runsubmit.go | 109 +++++++++++++------ 4 files changed, 347 insertions(+), 64 deletions(-) create mode 100644 experimental/air/cmd/runschedule.go create mode 100644 experimental/air/cmd/runschedule_test.go diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index 5931452e77c..15622462501 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -2,9 +2,7 @@ package aircmd import ( "context" - "errors" "fmt" - "path/filepath" "strconv" "github.com/databricks/cli/cmd/root" @@ -20,6 +18,7 @@ type runResult struct { Status string `json:"status"` DryRun bool `json:"dry_run,omitempty"` RunID string `json:"run_id,omitempty"` + JobID string `json:"job_id,omitempty"` DashboardURL string `json:"dashboard_url,omitempty"` } @@ -73,11 +72,11 @@ The workload is described by a YAML config file (see --file).`, return renderEnvelope(ctx, runResult{Status: "DRY_RUN_OK", DryRun: true}) } - // A schedule needs a persistent job; `air run` only submits a one-time run, so - // it can't honor one. Offer to convert to a bundle (where deploy schedules it) - // rather than silently submitting an unscheduled run. + // A schedule turns the workload into a persistent, scheduled job instead of a + // one-time run: create (or update) the job and return, since there is no + // immediate run to submit or stream. if cfg.Schedule != nil { - return handleScheduledRun(ctx, cfg, file) + return runScheduled(ctx, cmd, cfg, file) } w := cmdctx.WorkspaceClient(ctx) @@ -137,33 +136,35 @@ The workload is described by a YAML config file (see --file).`, return cmd } -// handleScheduledRun responds to a `schedule` in the run config. `air run` submits a -// one-time run and can't schedule it; scheduling needs a persistent job, which a -// Databricks Asset Bundle provides. It offers to convert the config to a bundle -// (writing it next to the YAML) and prints the deploy step; it never submits, since a -// scheduled run can't go through the submit path. Declining, or a non-interactive -// session, exits non-zero with the manual command. -func handleScheduledRun(ctx context.Context, cfg *runConfig, configPath string) error { - self := cliInvocation() - cmdio.LogString(ctx, "'schedule' is set, but 'air run' submits a one-time run and can't schedule it.") - cmdio.LogString(ctx, "Scheduling needs a persistent job, which a Databricks Asset Bundle provides.") - - convertCmd := fmt.Sprintf("%s experimental air convert-to-dabs %s", self, configPath) - manual := fmt.Sprintf("Nothing submitted. To schedule this job:\n %s\n databricks bundle deploy", convertCmd) - - convert, err := cmdio.AskYesOrNo(ctx, "Convert this config to a bundle now?") - if err != nil || !convert { - // A non-interactive session (no TTY) surfaces as an error here; treat it the - // same as declining and point at the manual command. - return errors.New(manual) - } - - dir := filepath.Dir(configPath) - written, err := writeBundle(ctx, cfg, configPath, dir, false) +// runScheduled creates (or updates) a persistent, scheduled job for a workload +// whose config carries a `schedule`. Unlike a submit, there is no immediate run +// to stream, so --watch does not apply here. +func runScheduled(ctx context.Context, cmd *cobra.Command, cfg *runConfig, configPath string) error { + w := cmdctx.WorkspaceClient(ctx) + jobID, jobURL, created, err := createScheduledJob(ctx, w, cfg, configPath) if err != nil { return err } - printConvertNextSteps(ctx, dir, written, bundleResourceKey(cfg.ExperimentName)) + jobIDStr := strconv.FormatInt(jobID, 10) + + if root.OutputType(cmd) == flags.OutputJSON { + status := "SCHEDULED_UPDATED" + if created { + status = "SCHEDULED_CREATED" + } + return renderEnvelope(ctx, runResult{Status: status, JobID: jobIDStr, DashboardURL: jobURL}) + } + + verb := "Updated" + if created { + verb = "Created" + } + cmdio.LogString(ctx, fmt.Sprintf("%s scheduled job %s", verb, jobIDStr)) + cmdio.LogString(ctx, "View at: "+jobURL) + cmdio.LogString(ctx, fmt.Sprintf("Runs on schedule: %s (%s)", cfg.Schedule.QuartzCronExpression, cfg.Schedule.TimezoneID)) + if cfg.Schedule.PauseStatus == "PAUSED" { + cmdio.LogString(ctx, "The schedule is PAUSED; set pause_status: UNPAUSED (or unpause it in the Jobs UI) to activate it.") + } return nil } diff --git a/experimental/air/cmd/runschedule.go b/experimental/air/cmd/runschedule.go new file mode 100644 index 00000000000..be166573f1e --- /dev/null +++ b/experimental/air/cmd/runschedule.go @@ -0,0 +1,128 @@ +package aircmd + +import ( + "context" + "fmt" + "strconv" + "strings" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +// buildCronSchedule maps the run YAML's schedule block onto the Jobs +// CronSchedule. pause_status is optional; an empty value lets the Jobs default +// (UNPAUSED) apply. +func buildCronSchedule(s *scheduleConfig) *jobs.CronSchedule { + if s == nil { + return nil + } + c := &jobs.CronSchedule{ + QuartzCronExpression: s.QuartzCronExpression, + TimezoneId: s.TimezoneID, + } + if s.PauseStatus != "" { + c.PauseStatus = jobs.PauseStatus(s.PauseStatus) + } + return c +} + +// buildJobSettings assembles the persistent-job settings for a scheduled +// workload: the same ai_runtime_task and environment as an ephemeral submit, +// carried on a Task (not SubmitTask) with the cron schedule attached. +func buildJobSettings(cfg *runConfig, commandPath, dlImage, usagePolicyID string, snap snapshotResult, deps []string) jobs.JobSettings { + task := buildAiRuntimeTask(cfg, commandPath, snap) + + maxRetries := cfg.maxRetries() + t := jobs.Task{ + TaskKey: cfg.ExperimentName, + RunIf: jobs.RunIfAllSuccess, + AiRuntimeTask: &task, + EnvironmentKey: aiRuntimeEnvironmentKey, + MaxRetries: maxRetries, + // retry_on_timeout only makes sense when retries are allowed (matches the + // ephemeral submit path). + RetryOnTimeout: maxRetries > 0, + ForceSendFields: []string{"MaxRetries"}, + } + + return jobs.JobSettings{ + Name: cfg.ExperimentName, + BudgetPolicyId: usagePolicyID, + TimeoutSeconds: cfg.timeoutSeconds(), + Tasks: []jobs.Task{t}, + Environments: buildAiRuntimeEnvironments(dlImage, deps), + Schedule: buildCronSchedule(cfg.Schedule), + } +} + +// findJobByName resolves an existing job to update. The CLI keeps no state, so a +// scheduled workload is keyed on its (unique-by-convention) name: zero matches +// means create a new job, one means update it in place, and more than one is +// ambiguous — the caller can't tell which to overwrite. +func findJobByName(ctx context.Context, w *databricks.WorkspaceClient, name string) (int64, error) { + it := w.Jobs.List(ctx, jobs.ListJobsRequest{Name: name, Limit: 100}) + var matches []int64 + for it.HasNext(ctx) { + job, err := it.Next(ctx) + if err != nil { + return 0, err + } + // The list Name filter is a case-insensitive match; require an exact name + // so "run" doesn't collide with "Run". + if job.Settings != nil && job.Settings.Name == name { + matches = append(matches, job.JobId) + } + } + switch len(matches) { + case 0: + return 0, nil + case 1: + return matches[0], nil + default: + return 0, fmt.Errorf("found %d jobs named %q; the name is not unique, so air run cannot tell which to update — rename or delete the duplicates, or use a unique experiment_name", len(matches), name) + } +} + +// createScheduledJob turns a workload with a schedule into a persistent, +// scheduled Databricks job. It uploads the same launch artifacts as an ephemeral +// run, then upserts by name: an existing job with the same name is updated in +// place (so re-running doesn't pile up duplicates), otherwise a new one is +// created. It returns the job id, its URL, and whether the job was created (vs +// updated). +func createScheduledJob(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath string) (jobID int64, url string, created bool, err error) { + prep, err := prepareWorkload(ctx, w, cfg, configPath) + if err != nil { + return 0, "", false, err + } + settings := buildJobSettings(cfg, prep.commandPath, prep.dlImage, prep.usagePolicyID, prep.snap, prep.deps) + + existingID, err := findJobByName(ctx, w, cfg.ExperimentName) + if err != nil { + return 0, "", false, err + } + + if existingID != 0 { + if err := w.Jobs.Reset(ctx, jobs.ResetJob{JobId: existingID, NewSettings: settings}); err != nil { + return 0, "", false, err + } + jobID = existingID + } else { + resp, err := w.Jobs.Create(ctx, jobs.CreateJob{ + Name: settings.Name, + BudgetPolicyId: settings.BudgetPolicyId, + TimeoutSeconds: settings.TimeoutSeconds, + Tasks: settings.Tasks, + Environments: settings.Environments, + Schedule: settings.Schedule, + }) + if err != nil { + return 0, "", false, err + } + jobID = resp.JobId + created = true + } + + url = strings.TrimRight(w.Config.Host, "/") + "/jobs/" + strconv.FormatInt(jobID, 10) + return jobID, url, created, nil +} diff --git a/experimental/air/cmd/runschedule_test.go b/experimental/air/cmd/runschedule_test.go new file mode 100644 index 00000000000..c513b8136bb --- /dev/null +++ b/experimental/air/cmd/runschedule_test.go @@ -0,0 +1,113 @@ +package aircmd + +import ( + "strconv" + "testing" + + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const scheduledConfig = minimalConfig + ` +schedule: + quartz_cron_expression: "0 0 9 * * ?" + timezone_id: America/Los_Angeles + pause_status: PAUSED +` + +// buildJobSettings carries the same ai_runtime_task/environment as a submit, on a +// persistent Task with the cron schedule attached. +func TestBuildJobSettings(t *testing.T) { + cfg, err := loadRunConfig(writeConfigFile(t, "run.yaml", scheduledConfig)) + require.NoError(t, err) + + js := buildJobSettings(cfg, "/d/command.sh", "5", "policy-1", snapshotResult{}, []string{"numpy"}) + assert.Equal(t, "my-run", js.Name) + assert.Equal(t, "policy-1", js.BudgetPolicyId) + + require.Len(t, js.Tasks, 1) + tk := js.Tasks[0] + assert.Equal(t, "my-run", tk.TaskKey) + assert.Equal(t, aiRuntimeEnvironmentKey, tk.EnvironmentKey) + require.NotNil(t, tk.AiRuntimeTask) + assert.Equal(t, "my-run", tk.AiRuntimeTask.Experiment) + assert.Equal(t, "/d/command.sh", tk.AiRuntimeTask.Deployments[0].CommandPath) + + require.Len(t, js.Environments, 1) + require.NotNil(t, js.Schedule) + assert.Equal(t, "0 0 9 * * ?", js.Schedule.QuartzCronExpression) + assert.Equal(t, "America/Los_Angeles", js.Schedule.TimezoneId) + assert.Equal(t, jobs.PauseStatusPaused, js.Schedule.PauseStatus) +} + +func TestBuildCronSchedule(t *testing.T) { + assert.Nil(t, buildCronSchedule(nil)) + + // An empty pause_status is left off so the Jobs default (UNPAUSED) applies. + c := buildCronSchedule(&scheduleConfig{QuartzCronExpression: "* * * * * ?", TimezoneID: "UTC"}) + require.NotNil(t, c) + assert.Equal(t, jobs.PauseStatus(""), c.PauseStatus) +} + +func loadScheduledConfig(t *testing.T) (*databricks.WorkspaceClient, *runConfig, string) { + t.Helper() + server := testserver.New(t) + t.Cleanup(server.Close) + testserver.AddDefaultHandlers(server) + + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + cfgPath := writeConfigFile(t, "run.yaml", scheduledConfig) + cfg, err := loadRunConfig(cfgPath) + require.NoError(t, err) + return w, cfg, cfgPath +} + +// A scheduled run creates a persistent job the first time and updates it in place +// on a re-run (upsert by name), so re-running never piles up duplicate jobs. +func TestCreateScheduledJobCreatesThenUpdates(t *testing.T) { + w, cfg, cfgPath := loadScheduledConfig(t) + + id1, url1, created, err := createScheduledJob(t.Context(), w, cfg, cfgPath) + require.NoError(t, err) + assert.True(t, created) + assert.NotZero(t, id1) + assert.Contains(t, url1, "/jobs/"+strconv.FormatInt(id1, 10)) + + id2, _, created2, err := createScheduledJob(t.Context(), w, cfg, cfgPath) + require.NoError(t, err) + assert.False(t, created2, "a re-run with the same name updates in place") + assert.Equal(t, id1, id2) + + // Exactly one job carries the name, and it kept the schedule through the reset. + var found []jobs.BaseJob + it := w.Jobs.List(t.Context(), jobs.ListJobsRequest{Name: cfg.ExperimentName}) + for it.HasNext(t.Context()) { + j, err := it.Next(t.Context()) + require.NoError(t, err) + if j.Settings != nil && j.Settings.Name == cfg.ExperimentName { + found = append(found, j) + } + } + require.Len(t, found, 1) + require.NotNil(t, found[0].Settings.Schedule) + assert.Equal(t, "0 0 9 * * ?", found[0].Settings.Schedule.QuartzCronExpression) +} + +// When two jobs share the experiment name, the CLI can't tell which to update, so +// it errors rather than guessing. +func TestCreateScheduledJobAmbiguousName(t *testing.T) { + w, cfg, cfgPath := loadScheduledConfig(t) + + for range 2 { + _, err := w.Jobs.Create(t.Context(), jobs.CreateJob{Name: cfg.ExperimentName}) + require.NoError(t, err) + } + + _, _, _, err := createScheduledJob(t.Context(), w, cfg, cfgPath) + require.ErrorContains(t, err, "not unique") +} diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 54189fba956..8360711fcce 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -69,7 +69,9 @@ func environmentDependencies(cfg *runConfig, configPath string) (deps []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, usagePolicyID string, snap snapshotResult, deps []string) jobs.SubmitRun { +// buildAiRuntimeTask assembles the ai_runtime_task shared by the ephemeral +// (SubmitTask) and persistent (Task) payloads. +func buildAiRuntimeTask(cfg *runConfig, commandPath string, snap snapshotResult) jobs.AiRuntimeTask { task := jobs.AiRuntimeTask{ Experiment: cfg.ExperimentName, Deployments: []jobs.DeploymentSpec{{ @@ -92,6 +94,25 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID stri if cfg.MLflowExperimentDirectory != nil { task.MlflowExperimentDirectory = *cfg.MLflowExperimentDirectory } + return task +} + +// buildAiRuntimeEnvironments carries the user's declared deps inline on +// spec.dependencies; the AI Runtime backend installs them via --deps-config. The +// SDK marshaler drops nil and empty slices, so a no-deps run omits the key. +func buildAiRuntimeEnvironments(dlImage string, deps []string) []jobs.JobEnvironment { + envSpec := &compute.Environment{EnvironmentVersion: dlImage} + if len(deps) > 0 { + envSpec.Dependencies = deps + } + return []jobs.JobEnvironment{{ + EnvironmentKey: aiRuntimeEnvironmentKey, + Spec: envSpec, + }} +} + +func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID string, snap snapshotResult, deps []string) jobs.SubmitRun { + task := buildAiRuntimeTask(cfg, commandPath, snap) maxRetries := cfg.maxRetries() st := jobs.SubmitTask{ @@ -107,14 +128,6 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID stri ForceSendFields: []string{"MaxRetries"}, } - // Carry the user's declared deps inline on spec.dependencies; the AI Runtime - // backend installs them via --deps-config. The SDK marshaler drops nil and empty - // slices, so a no-deps run omits the key. - envSpec := &compute.Environment{EnvironmentVersion: dlImage} - if len(deps) > 0 { - envSpec.Dependencies = deps - } - return jobs.SubmitRun{ RunName: cfg.ExperimentName, // budget_policy_id matches what the Python CLI and `ssh connect` send; @@ -122,10 +135,7 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID stri BudgetPolicyId: usagePolicyID, TimeoutSeconds: cfg.timeoutSeconds(), Tasks: []jobs.SubmitTask{st}, - Environments: []jobs.JobEnvironment{{ - EnvironmentKey: aiRuntimeEnvironmentKey, - Spec: envSpec, - }}, + Environments: buildAiRuntimeEnvironments(dlImage, deps), } } @@ -146,18 +156,23 @@ func submitToken(flag string, cfg *runConfig) (string, error) { return token, nil } -// submitWorkload runs the submit happy path: ensure the experiment directory, -// upload the launch artifacts, assemble the Jobs payload, and submit it. It -// returns the new run_id and its dashboard URL. -func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string) (int64, string, error) { - // Resolve the idempotency token first so a bad key fails before any upload, - // and before the policy lookup below spends a round trip on it. - token, err := submitToken(idempotencyKey, cfg) - if err != nil { - return 0, "", err - } +// preparedWorkload holds everything resolved and uploaded for a workload, ready +// to become either an ephemeral SubmitRun (submitWorkload) or a persistent, +// scheduled CreateJob (createScheduledJob). +type preparedWorkload struct { + commandPath string + dlImage string + usagePolicyID string + snap snapshotResult + deps []string +} - // Resolve the usage policy to its id next, so a bad name fails fast with a +// prepareWorkload runs the shared pre-submit work: resolve the usage policy and +// dependencies, ensure the experiment directory, prepare any docker image, and +// upload the launch artifacts + code snapshot. Ordering fails cheap checks before +// any upload so a bad config leaves no orphaned artifacts in the workspace. +func prepareWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath string) (preparedWorkload, error) { + // Resolve the usage policy to its id first, so a bad name fails fast with a // clear (caller-fixable) message before we upload any artifacts. Validation // guarantees name and id are mutually exclusive: a literal id is used as-is, a // name is resolved against the workspace. @@ -166,9 +181,10 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run usagePolicyID = strings.TrimSpace(*cfg.UsagePolicyID) } if cfg.UsagePolicyName != nil { + var err error usagePolicyID, err = resolveUsagePolicyIDByName(ctx, w, *cfg.UsagePolicyName) if err != nil { - return 0, "", err + return preparedWorkload{}, err } } @@ -176,7 +192,7 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run // fast without leaving orphaned artifacts in the workspace. deps, fileVersion, err := environmentDependencies(cfg, configPath) if err != nil { - return 0, "", err + return preparedWorkload{}, err } experimentDir := "" @@ -184,19 +200,19 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run experimentDir = *cfg.MLflowExperimentDirectory } if err := ensureExperimentDirectory(ctx, w, experimentDir); err != nil { - return 0, "", err + return preparedWorkload{}, err } base, err := userWorkspaceDir(ctx, w) if err != nil { - return 0, "", err + return preparedWorkload{}, err } // After the cheap workspace checks (a tag_policy=latest refresh can block for // minutes) but before any upload, so a bad image wastes no artifact work. if img := cfg.dockerImage(); img != nil { if err := prepareDockerImage(ctx, w, img); err != nil { - return 0, "", err + return preparedWorkload{}, err } } @@ -208,14 +224,14 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run fc, err := filer.NewWorkspaceFilesClient(w, funcDir) if err != nil { - return 0, "", err + return preparedWorkload{}, err } items, err := buildArtifacts(cfg, configPath) if err != nil { - return 0, "", err + return preparedWorkload{}, err } if err := uploadArtifacts(ctx, fc, items); err != nil { - return 0, "", err + return preparedWorkload{}, err } // Package and upload the code snapshot, if any, via DABs' artifact-upload @@ -226,7 +242,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. snap, err = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir) if err != nil { - return 0, "", err + return preparedWorkload{}, err } } @@ -236,7 +252,32 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run if !ok { runtimeVersion = fileVersion } - payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion), usagePolicyID, snap, deps) + + return preparedWorkload{ + commandPath: path.Join(funcDir, commandScriptName), + dlImage: dlRuntimeImage(ctx, runtimeVersion), + usagePolicyID: usagePolicyID, + snap: snap, + deps: deps, + }, nil +} + +// submitWorkload runs the submit happy path: prepare the workload, assemble the +// ephemeral Jobs payload, and submit it. It returns the new run_id and its +// dashboard URL. +func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string) (int64, string, error) { + // Resolve the idempotency token first so a bad key fails before any upload. + token, err := submitToken(idempotencyKey, cfg) + if err != nil { + return 0, "", err + } + + prep, err := prepareWorkload(ctx, w, cfg, configPath) + if err != nil { + return 0, "", err + } + + payload := buildSubmitPayload(cfg, prep.commandPath, prep.dlImage, prep.usagePolicyID, prep.snap, prep.deps) payload.IdempotencyToken = token // Submit returns as soon as the run is created; we don't wait for it to finish.