Add --override-image-id CLI flag for the Modal provider

.beads/issues.jsonl

@@ -147,6 +147,7 @@
 {"id":"code-b9d4c2","title":"Git module -- src/git.rs","description":"Create new src/git.rs module (and register in src/lib.rs). All git interactions via spawn_blocking with std::process::Command. Async public API. Contains: NOTES_REF constant, ImageEntry struct, NoteContents type alias, and functions: head_sha, repo_root, compute_build_inputs_hash, read_note, write_note, push_notes, fetch_notes, configure_notes_fetch, commit_touches_paths, ancestors, canonicalize_config_path. push_notes force-pushes unconditionally (no fetch). See plans/checkpoint-images.plan.md Commit 2 for full details including all tests.","status":"closed","priority":0,"issue_type":"task","created_at":"2026-04-12T21:37:33.154961Z","created_by":"danver","updated_at":"2026-04-12T21:57:46.045945Z","closed_at":"2026-04-12T21:57:46.045945Z","source_repo":".","compaction_level":0,"original_size":0}
 {"id":"code-c6e8f3","title":"Checkpoint resolution logic -- src/checkpoint.rs","description":"Create new src/checkpoint.rs module (and register in src/lib.rs). Contains ResolvedCheckpoint struct (checkpoint_sha, image_id -- NO sandbox_project_root), resolve_checkpoint async fn, resolve_cached_image async fn. Uses SandboxProvider trait (not BuildableProvider). See plans/checkpoint-images.plan.md Commit 3 for logic and tests.","status":"closed","priority":0,"issue_type":"task","created_at":"2026-04-12T21:37:33.202922Z","created_by":"danver","updated_at":"2026-04-12T22:04:59.052080Z","closed_at":"2026-04-12T22:04:59.052080Z","source_repo":".","compaction_level":0,"original_size":0}
 {"id":"code-d2a5b7","title":"Provider checkpoint support","description":"Add CheckpointContext struct to providers. Add checkpoint field and with_checkpoint builder to ModalProvider and DefaultProvider. Modify build_prepare_command: when checkpoint is Some, append --from-checkpoint, --checkpoint-sha, --sandbox-project-root flags and omit --include-cwd, --copy-dir, --sandbox-init-cmd. Remove --cached flag logic from both providers. Tests: test_prepare_command_with_checkpoint, test_prepare_command_without_checkpoint. See plans/checkpoint-images.plan.md Commit 4.","status":"closed","priority":0,"issue_type":"task","created_at":"2026-04-12T21:37:33.251745Z","created_by":"danver","updated_at":"2026-04-12T22:15:27.690901Z","closed_at":"2026-04-12T22:15:27.690901Z","source_repo":".","compaction_level":0,"original_size":0}
+{"id":"code-d9y3","title":"Add --override-image-id CLI flag to bypass Modal image build","description":"## Goal\n\nAdd a per-invocation CLI escape hatch: `offload run --override-image-id <ID>`.\nWhen passed, Offload runs tests against the given pre-built Modal image exactly\nas-is and skips its entire image-setup pipeline (Dockerfile build, thin-diff\npatch of local changes, checkpoint cache). For users who have already built\ntheir own Modal image and do not want Offload to build one for them.\n\nThis is a jury-rigged, invocation-scoped override. It is NOT a config option.\n\n## Hard constraints (do not violate)\n\n- Do NOT add any field to `ModalProviderConfig` or anything else in\n  `src/config/schema.rs`. Do NOT add config validation in `src/config.rs`.\n  The override is a CLI flag ONLY — nothing in the TOML schema.\n- The branch has been reset to a clean base. You are implementing fresh: there\n  is NO pre-existing `override_image_id` config field anywhere. Do not look for\n  one or try to reuse one.\n\n## Approach\n\n1. CLI flag: add `override_image_id: Option<String>` to the `Commands::Run`\n   struct in `src/main.rs` with `#[arg(long)]` (clap renders it as\n   `--override-image-id`). Give it a clear `help` string: an escape hatch that\n   makes Offload run tests against the given Modal image ID directly, bypassing\n   all image build/patch/cache setup.\n\n2. Thread it to `prepare()` via `PrepareContext` (`src/provider.rs`), mirroring\n   exactly how the existing `no_cache` CLI flag is threaded. Add an\n   `override_image_id` field to `PrepareContext` and populate it in\n   `run_prepare()` in `src/main.rs` from the CLI value.\n\n3. `ModalProvider::prepare()` (`src/provider/modal.rs`): when the\n   `PrepareContext` override is `Some`, short-circuit — skip\n   `prepare_with_prewarm` entirely, set `self.image_id` to the override value,\n   return `Ok(Some(value))`. Emit ONE loud `tracing::warn!` stating that the\n   configured image build/patch/cache is bypassed for this run and the image is\n   used as-is.\n\n4. Non-Modal guard: `--override-image-id` only makes sense for the `modal`\n   provider. If it is passed while `config.provider` is not Modal, fail early in\n   `run_tests()` with a clear error (e.g. \"--override-image-id is only\n   supported with the modal provider\").\n\n## Tests (TDD — failing test first)\n\n- CLI parsing: `offload run --override-image-id im-x ...` parses and the value\n  lands on the `Run` command struct.\n- `prepare()` short-circuit: with the `PrepareContext` override set to\n  `Some(...)`, `ModalProvider::prepare()` returns the override ID and does NOT\n  run a build. Mirror the existing modal.rs prepare-test style.\n- Non-Modal guard returns an error when the flag is set with a non-modal\n  provider, if reasonably testable at that seam.\n\n## Docs\n\n- The clap `help` string is the primary documentation.\n- If README enumerates `offload run` flags, add one line for\n  `--override-image-id` matching that list's style. Do NOT touch the\n  offload-onboard skill — that skill is about config setup; this is not config.\n\n## Out of scope\n\n- No config/schema changes. No base-commit incremental mode. No registry images.\n- No drive-by refactors.\n\n## Definition of Done\n\n- Read README.md, AGENTS.md, STYLE_GUIDE.md, TESTING.md before changing code.\n- `cargo fmt --check`, `cargo clippy` (no warnings), `cargo nextest run`, and\n  `ratchets check` all pass.\n- Prefer a small number of coherent atomic commits (the user explicitly wants a\n  clean commit history) — e.g. one feature+tests commit, optionally one docs\n  commit. Each commit must typecheck and pass checks.\n- Commit messages include the `Co-authored-by: Sculptor <sculptor@imbue.com>`\n  trailer.\n- Do NOT commit `.beads/` changes — the Coordinator manages bead state.\n- Do NOT close the bead — report back to the Coordinator.","status":"closed","priority":1,"issue_type":"task","owner":"danver","created_at":"2026-05-20T12:17:58.863828Z","created_by":"danver","updated_at":"2026-05-20T12:28:23.873452Z","closed_at":"2026-05-20T12:28:23.873292Z","close_reason":"done","source_repo":".","compaction_level":0,"original_size":0}
 {"id":"code-de3b01","title":"Remove 10s MAX_BATCH_DURATION cap from LPT scheduling","description":"Remove the hardcoded MAX_BATCH_DURATION constant (10s) in orchestrator.rs and pass None to Scheduler::new. Also remove the max_batch_duration parameter from Scheduler::new and Batch::would_fit since it is no longer used in production. Remove associated dead-code tests and the test constant MAX_BATCH_DURATION.","status":"closed","priority":1,"issue_type":"task","created_at":"2026-04-15T00:15:37.328773Z","created_by":"Jacob Kirmayer","updated_at":"2026-04-15T00:19:58.452539Z","closed_at":"2026-04-15T00:19:58.452539Z","source_repo":".","compaction_level":0,"original_size":0}
 {"id":"code-e8f1d4","title":"Python prepare modifications","description":"Modify scripts/modal_sandbox.py: Remove old caching (CACHE_FILE, read_cached_image_id, write_cached_image_id, clear_image_cache, --cached flag). Add --from-checkpoint, --checkpoint-sha, --sandbox-project-root options to prepare command. Implement thin diff application: git diff, untracked file collection, add_local_file, git apply. On failure, exit non-zero (no fallback in Python). See plans/checkpoint-images.plan.md Commit 5.","status":"closed","priority":0,"issue_type":"task","created_at":"2026-04-12T21:37:33.299476Z","created_by":"danver","updated_at":"2026-04-12T22:20:57.471990Z","closed_at":"2026-04-12T22:20:57.471990Z","source_repo":".","compaction_level":0,"original_size":0}
 {"id":"code-ear","title":"Run post_patch_cmd in full_build_fallback path","description":"When thin-diff fails and Offload falls back to full_build_fallback, post_patch_cmd is skipped. The fallback path calls build_full (which runs sandbox_init_cmd) but never runs post_patch_cmd. Fix: after build_full succeeds in full_build_fallback, if post_patch_cmd is configured, call build_incremental(base=new_image_id, patch=None, post_patch_cmd=Some(...)) to route through _derive_image_from_base. Discovered via sculptor end-to-end testing.","status":"closed","priority":0,"issue_type":"bug","created_at":"2026-05-07T05:32:26.823328Z","created_by":"danver","updated_at":"2026-05-07T05:36:49.157160Z","closed_at":"2026-05-07T05:36:49.157160Z","source_repo":".","compaction_level":0,"original_size":0}

.beads/last-touched

@@ -1 +1 @@
-code-2fzk
+code-d9y3

README.md

@@ -183,6 +183,7 @@ Run tests in parallel.
 | `--copy-dir LOCAL:REMOTE` | Copy a directory into each sandbox (repeatable) |
 | `--env KEY=VALUE` | Set an environment variable in sandboxes (repeatable) |
 | `--no-cache` | Skip cached image lookup during prepare (forces fresh build) |
+| `--override-image-id ID` | Escape hatch: run tests against the given pre-built Modal image ID as-is, bypassing all image build, thin-diff patch, and cache setup. Only valid with the `modal` provider |
 | `--trace` | Emit a Perfetto trace to `{output_dir}/trace.json` |
 | `--fail-fast` | Stop on first test failure. Passes a framework-level stop flag (`-x` for pytest, `--fail-fast` for nextest, `--bail` for vitest) and cancels remaining batches at the orchestrator level |
 | `--show-estimated-cost` | Show estimated sandbox cost after run (client-side estimate, may not reflect actual billing) |

src/main.rs

@@ -73,6 +73,12 @@ enum Commands {
         #[arg(long)]
         no_cache: bool,
 
+        /// Escape hatch: run tests against this pre-built Modal image ID
+        /// directly, bypassing all image build, thin-diff patch, and cache
+        /// setup. Only valid with the modal provider.
+        #[arg(long, value_name = "IMAGE_ID")]
+        override_image_id: Option<String>,
+
         /// Emit a Perfetto trace to {output_dir}/trace.json
         #[arg(long)]
         trace: bool,
@@ -224,6 +230,7 @@ async fn main() -> Result<()> {
             copy_dir,
             env_vars,
             no_cache,
+            override_image_id,
             trace,
             show_estimated_cost,
             fail_fast,
@@ -238,6 +245,7 @@ async fn main() -> Result<()> {
                 copy_dir,
                 env_vars,
                 no_cache,
+                override_image_id,
                 cli.verbose,
                 trace,
                 show_estimated_cost,
@@ -464,6 +472,7 @@ async fn run_prepare<P: SandboxProvider>(
     config_path: &Path,
     copy_dirs: &[(PathBuf, PathBuf)],
     no_cache: bool,
+    override_image_id: Option<&str>,
     tracer: &offload::trace::Tracer,
     discovery_done: &AtomicBool,
 ) -> Result<Option<String>> {
@@ -475,6 +484,7 @@ async fn run_prepare<P: SandboxProvider>(
         config,
         config_path,
         no_cache,
+        override_image_id,
         tracer,
         discovery_done,
     };
@@ -498,6 +508,7 @@ async fn run_remote_provider<P: SandboxProvider>(
     run_id: &str,
     copy_dirs: &[CopyDir],
     no_cache: bool,
+    override_image_id: Option<&str>,
     verbose: bool,
     tracer: &offload::trace::Tracer,
     show_estimated_cost: bool,
@@ -517,6 +528,7 @@ async fn run_remote_provider<P: SandboxProvider>(
             config_path,
             copy_dir_tuples,
             no_cache,
+            override_image_id,
             tracer,
             &discovery_done,
         ),
@@ -552,6 +564,7 @@ async fn run_tests(
     copy_dir_args: Vec<String>,
     env_vars: Vec<String>,
     no_cache: bool,
+    override_image_id: Option<String>,
     verbose: bool,
     trace: bool,
     show_estimated_cost: bool,
@@ -587,6 +600,11 @@ async fn run_tests(
         anyhow::bail!("--record-history requires a [history] section in the config file");
     }
 
+    // Validate --override-image-id flag: only the modal provider builds images.
+    if override_image_id.is_some() && !matches!(config.provider, ProviderConfig::Modal(_)) {
+        anyhow::bail!("--override-image-id is only supported with the modal provider");
+    }
+
     // Generate run ID for history recording
     let run_id = offload::generate_run_id();
     let config_filename = config_path
@@ -719,6 +737,8 @@ async fn run_tests(
                 &run_id,
                 &copy_dirs,
                 no_cache,
+                // --override-image-id is rejected above for non-modal providers.
+                None,
                 verbose,
                 &tracer,
                 show_estimated_cost,
@@ -744,6 +764,7 @@ async fn run_tests(
                 &run_id,
                 &copy_dirs,
                 no_cache,
+                override_image_id.as_deref(),
                 verbose,
                 &tracer,
                 show_estimated_cost,
@@ -879,6 +900,7 @@ async fn build_image(config_path: &Path, no_cache: bool) -> Result<()> {
                 config_path,
                 &copy_dir_tuples,
                 no_cache,
+                None,
                 &tracer,
                 &discovery_done,
             )
@@ -893,6 +915,7 @@ async fn build_image(config_path: &Path, no_cache: bool) -> Result<()> {
                 config_path,
                 &copy_dir_tuples,
                 no_cache,
+                None,
                 &tracer,
                 &discovery_done,
             )
@@ -1513,4 +1536,29 @@ deleted file mode 100644
 
         Ok(())
     }
+
+    /// Parses CLI args and extracts `override_image_id` from a `Run` command.
+    ///
+    /// Returns the outer `Some` when the parsed command is `Run`; the inner
+    /// `Option` is the flag value. Avoids panic-family macros in tests.
+    fn parse_run_override_image_id<const N: usize>(args: [&str; N]) -> Option<Option<String>> {
+        match Cli::parse_from(args).command {
+            Commands::Run {
+                override_image_id, ..
+            } => Some(override_image_id),
+            _ => None,
+        }
+    }
+
+    #[test]
+    fn run_parses_override_image_id_flag() {
+        let parsed = parse_run_override_image_id(["offload", "run", "--override-image-id", "im-x"]);
+        assert_eq!(parsed, Some(Some("im-x".to_string())));
+    }
+
+    #[test]
+    fn run_override_image_id_defaults_to_none() {
+        let parsed = parse_run_override_image_id(["offload", "run"]);
+        assert_eq!(parsed, Some(None));
+    }
 }

src/provider.rs

@@ -31,6 +31,12 @@ pub struct PrepareContext<'a> {
     pub config_path: &'a Path,
     /// Skip cached image lookup.
     pub no_cache: bool,
+    /// Per-invocation escape hatch: use this pre-built image ID as-is.
+    ///
+    /// When `Some`, image-building providers skip their entire setup pipeline
+    /// (build, thin-diff patch, checkpoint cache). Sourced from the
+    /// `--override-image-id` CLI flag; never from config.
+    pub override_image_id: Option<&'a str>,
     /// Tracer for performance tracing.
     pub tracer: &'a crate::trace::Tracer,
     /// Signal flag set when test discovery is complete.
@@ -573,11 +579,13 @@ mod tests {
             config: &config,
             config_path: Path::new("offload.toml"),
             no_cache: false,
+            override_image_id: None,
             tracer: &tracer,
             discovery_done: &discovery_done,
         };
         assert_eq!(ctx.repo, Path::new("/tmp"));
         assert!(!ctx.no_cache);
+        assert!(ctx.override_image_id.is_none());
         assert!(ctx.sandbox_init_cmd.is_none());
         assert!(ctx.copy_dirs.is_empty());
     }

src/provider/modal.rs

@@ -220,6 +220,17 @@ impl SandboxProvider for ModalProvider {
     type Sandbox = DefaultSandbox;
 
     async fn prepare(&mut self, ctx: &PrepareContext<'_>) -> ProviderResult<Option<String>> {
+        // Escape hatch: when --override-image-id is set, run against the given
+        // pre-built image as-is and skip the entire setup pipeline.
+        if let Some(image_id) = ctx.override_image_id {
+            tracing::warn!(
+                "--override-image-id={image_id}: bypassing all Modal image \
+                 build/patch/cache setup; using the image as-is for this run",
+            );
+            self.image_id = Some(image_id.to_string());
+            return Ok(Some(image_id.to_string()));
+        }
+
         let result = prepare_with_prewarm(self, ctx).await?;
         self.image_id = result.clone();
         Ok(result)
@@ -468,4 +479,64 @@ mod tests {
         assert!(!cmd.contains("--patch-file"));
         assert!(cmd.contains("--post-patch-cmd=scripts/regen-clients.sh"));
     }
+
+    // -- override_image_id (--override-image-id escape hatch) tests --
+
+    /// Builds a minimal config suitable for exercising `prepare()`.
+    fn minimal_config() -> crate::config::Config {
+        crate::config::Config {
+            offload: crate::config::schema::OffloadConfig {
+                max_parallel: 1,
+                test_timeout_secs: 60,
+                working_dir: None,
+                sandbox_project_root: None,
+                sandbox_repo_root: None,
+                sandbox_init_cmd: None,
+                post_patch_cmd: None,
+            },
+            provider: crate::config::schema::ProviderConfig::Modal(ModalProviderConfig::default()),
+            framework: crate::config::schema::FrameworkConfig::Default(
+                crate::config::schema::DefaultFrameworkConfig {
+                    discover_command: "echo test".to_string(),
+                    run_command: "run {tests}".to_string(),
+                    result_file: None,
+                    working_dir: None,
+                    test_id_format: "{name}".to_string(),
+                },
+            ),
+            groups: Default::default(),
+            report: Default::default(),
+            checkpoint: None,
+            history: None,
+        }
+    }
+
+    /// With the override set, `prepare()` returns the override ID directly and
+    /// never spawns a build (a real build would shell out to `modal_sandbox.py`).
+    #[tokio::test]
+    async fn test_prepare_short_circuits_on_override_image_id() -> ProviderResult<()> {
+        let mut p = provider(ModalProviderConfig::default());
+        let config = minimal_config();
+        let discovery_done = AtomicBool::new(true);
+        let tracer = crate::trace::Tracer::noop();
+        let ctx = PrepareContext {
+            copy_dirs: &[],
+            sandbox_init_cmd: None,
+            post_patch_cmd: None,
+            repo: Path::new("/tmp"),
+            config: &config,
+            config_path: Path::new("offload.toml"),
+            no_cache: false,
+            override_image_id: Some("im-override-xyz"),
+            tracer: &tracer,
+            discovery_done: &discovery_done,
+        };
+
+        let result = p.prepare(&ctx).await?;
+
+        assert_eq!(result.as_deref(), Some("im-override-xyz"));
+        // The override ID is stored so create_sandbox() uses it.
+        assert_eq!(p.image_id.as_deref(), Some("im-override-xyz"));
+        Ok(())
+    }
 }