Cache agent setup in an image, behind --cache-agent-setup (experiment)

assets/schemas/trial.json

@@ -308,6 +308,12 @@
           "title": "Stream",
           "type": "boolean"
         },
+        "cache_agent_setup": {
+          "default": false,
+          "description": "Start the main container from an image of it as agent setup left it, committed by the first trial with the same image, agent, options and hbench version.",
+          "title": "Cache Agent Setup",
+          "type": "boolean"
+        },
         "cpu_enforcement_policy": {
           "$ref": "#/$defs/ResourceMode",
           "default": "auto"

src/agents/installed/claude_code.rs

@@ -320,6 +320,9 @@ impl Agent for ClaudeCode {
     fn name(&self) -> &str {
         "claude-code"
     }
+    fn setup_is_cacheable(&self) -> bool {
+        true
+    }
     fn version(&self) -> Option<String> {
         self.version.clone()
     }

src/agents/installed/codex.rs

@@ -645,6 +645,9 @@ impl Agent for Codex {
     fn name(&self) -> &str {
         "codex"
     }
+    fn setup_is_cacheable(&self) -> bool {
+        true
+    }
     fn version(&self) -> Option<String> {
         self.version.clone()
     }

src/agents/installed/deepseek_harness.rs

@@ -323,6 +323,9 @@ impl Agent for DeepSeekHarness {
     fn name(&self) -> &str {
         NAME
     }
+    fn setup_is_cacheable(&self) -> bool {
+        true
+    }
     fn version(&self) -> Option<String> {
         self.version.clone()
     }
@@ -332,14 +335,18 @@ impl Agent for DeepSeekHarness {
 
     async fn setup(&mut self, cx: &AgentCtx<'_>) -> Result<()> {
         installed_dir(cx).await?;
-        ensure_system_dependencies(cx, &["curl", "bash", "coreutils"]).await?;
-        let spec = pins::DEEPSEEK_HARNESS.npm_spec(self.version.as_deref());
-        exec_agent(
-            cx,
-            &format!("set -euo pipefail; {} && npm i -g {spec} && dsh --version", nvm_node_install_snippet(22)),
-            &BTreeMap::new(),
-        )
-        .await?;
+        let wanted = pins::DEEPSEEK_HARNESS.or(self.version.as_deref()).to_string();
+        let installed = detect_version(cx, &format!("{NVM} dsh --version"), |s| s.trim().to_string()).await;
+        if wanted == "latest" || installed.as_deref() != Some(wanted.as_str()) {
+            ensure_system_dependencies(cx, &["curl", "bash", "coreutils"]).await?;
+            let spec = pins::DEEPSEEK_HARNESS.npm_spec(self.version.as_deref());
+            exec_agent(
+                cx,
+                &format!("set -euo pipefail; {} && npm i -g {spec} && dsh --version", nvm_node_install_snippet(22)),
+                &BTreeMap::new(),
+            )
+            .await?;
+        }
         if self.version.is_none() {
             self.version = detect_version(cx, &format!("{NVM} dsh --version"), |s| s.trim().to_string()).await;
         }

src/agents/installed/opencode.rs

@@ -225,6 +225,9 @@ impl Agent for OpenCode {
     fn name(&self) -> &str {
         "opencode"
     }
+    fn setup_is_cacheable(&self) -> bool {
+        true
+    }
     fn version(&self) -> Option<String> {
         self.version.clone()
     }

src/agents/mod.rs

@@ -101,6 +101,11 @@ pub trait Agent: Send + Sync {
     fn supports_resume(&self) -> bool {
         false
     }
+    /// Setup writes only what hbench's version and the agent's options decide (no credentials, no host
+    /// files), so `--cache-agent-setup` may commit the container it leaves for later trials.
+    fn setup_is_cacheable(&self) -> bool {
+        false
+    }
     /// Can seed a session from a native (`atif=false`) or ATIF (`atif=true`) trajectory.
     fn supports_load(&self, _atif: bool) -> bool {
         false
@@ -294,6 +299,20 @@ pub fn model_endpoint(cfg: &AgentConfig, task: &Task) -> Option<ModelEndpoint> {
 mod tests {
     use super::*;
 
+    #[test]
+    fn only_agents_whose_setup_writes_no_credentials_or_host_files_are_cacheable() {
+        let task = Arc::new(Task::load(std::path::Path::new("examples/tasks/hello-world")).unwrap());
+        let cacheable = |name: &str| {
+            let cfg: AgentConfig =
+                serde_json::from_value(serde_json::json!({"name": name, "model_name": "anthropic/claude-sonnet-4-5"})).unwrap();
+            create(&cfg, std::env::temp_dir(), task.clone(), None).is_ok_and(|agent| agent.setup_is_cacheable())
+        };
+        let expected = ["claude-code", "codex", "deepseek-harness", "opencode"];
+        for name in agent_names() {
+            assert_eq!(cacheable(name), expected.contains(&name), "{name}");
+        }
+    }
+
     #[test]
     fn options_come_from_the_schema() {
         let kw = |v: Value| -> BTreeMap<String, Value> { serde_json::from_value(v).unwrap() };

src/config.rs

@@ -278,6 +278,10 @@ pub struct EnvConfig {
     /// Docker kwargs; `keep_containers: true` stops instead of removing containers.
     pub kwargs: BTreeMap<String, Value>,
     pub extra_allowed_hosts: Vec<String>,
+    /// Start the main container from an image of it as agent setup left it, committed by the first trial with
+    /// the same image, agent, options and hbench version (`--cache-agent-setup`).
+    #[serde(skip_serializing_if = "std::ops::Not::not")]
+    pub cache_agent_setup: bool,
 }
 
 fn ser_env_type<S: serde::Serializer>(k: &Option<String>, s: S) -> Result<S::Ok, S::Error> {
@@ -310,6 +314,7 @@ impl Default for EnvConfig {
             env: BTreeMap::new(),
             kwargs: BTreeMap::new(),
             extra_allowed_hosts: vec![],
+            cache_agent_setup: false,
         }
     }
 }
@@ -1124,6 +1129,11 @@ mod tests {
             j2.config_json(),
             serde_json::json!({"job_name": "j", "n_attempts": 3, "agents": [{"name": "claude-code", "model_name": "m"}]})
         );
+        let cached: JobConfig =
+            serde_json::from_value(serde_json::json!({"job_name": "j", "environment": {"cache_agent_setup": true}})).unwrap();
+        assert!(cached.environment.cache_agent_setup);
+        assert_eq!(cached.config_json()["environment"]["cache_agent_setup"], true);
+        assert!(serde_json::to_value(JobConfig::default()).unwrap()["environment"].get("cache_agent_setup").is_none());
     }
 
     #[test]

src/env.rs

@@ -273,6 +273,8 @@ const CLEAR_DIRS_TIMEOUT: Duration = Duration::from_secs(120);
 const EGRESS_SETUP_TIMEOUT: Duration = Duration::from_secs(120);
 /// Building the egress sidecar image, which pulls its base image the first time.
 const SIDECAR_BUILD_TIMEOUT: Duration = Duration::from_secs(900);
+/// Committing a main container to an agent setup image: its writable layer can hold a large install.
+const SETUP_COMMIT_TIMEOUT: Duration = Duration::from_secs(600);
 /// `docker version` before a run: an ssh context may take a few seconds to connect.
 const DAEMON_CHECK_TIMEOUT: Duration = Duration::from_secs(30);
 
@@ -434,6 +436,9 @@ pub struct EnvSpec {
     pub delete: bool,
     /// Log dirs to create (agent env: all three; separate verifier: /logs/verifier).
     pub log_dirs: Vec<&'static str>,
+    /// What, besides the main image, an agent setup image depends on: set when this trial may start from
+    /// one, or commit one after agent setup.
+    pub setup_cache_key: Option<String>,
 }
 
 pub struct DockerEnv {
@@ -445,6 +450,27 @@ pub struct DockerEnv {
     /// The services behind the sidecar.
     egress_services: Vec<String>,
     log: Arc<TrialLog>,
+    setup_cache: Option<SetupCache>,
+}
+
+/// What the agent setup image depends on besides the main image; the image, unless the main image is pulled by
+/// this trial and its id is known only once the container runs; and whether this trial started from it.
+struct SetupCache {
+    key: String,
+    image: Option<String>,
+    hit: bool,
+}
+
+/// `hbench-setup-cache:<hash>` of the main image and `key`: the name of an image hbench builds, which hashes its
+/// build context and goes away with the trial, else the id of the pulled image, whose tag may move.
+pub fn setup_cache_image(key: &str, main_image: &str) -> String {
+    let hash = sha2::Digest::finalize(sha2::Digest::chain_update(sha2::Sha256::default(), format!("{main_image}\n{key}")));
+    format!("hbench-setup-cache:{}", &crate::task::hex(&hash)[..16])
+}
+
+async fn image_id(image: &str) -> Result<Option<String>> {
+    let out = docker_output(&["image", "inspect", "--format", "{{.Id}}", image], QUICK_COMMAND_TIMEOUT).await?;
+    Ok(Some(String::from_utf8_lossy(&out.stdout).trim().to_string()).filter(|id| out.status.success() && !id.is_empty()))
 }
 
 /// Docker image/project names: lowercase, alnum first, restricted charset.
@@ -746,6 +772,40 @@ impl DockerEnv {
             std::fs::write(&p, serde_json::to_string_pretty(&v)?)?;
             Ok(p)
         };
+        let plain = egress.is_none()
+            && ec.env.is_empty()
+            && spec.persistent_env.is_empty()
+            && spec.mounts.is_empty()
+            && spec.extra_compose.is_empty()
+            && !task_compose.exists();
+        let mut build_image = build_image;
+        let setup_cache = match (&spec.setup_cache_key, main.get("image").and_then(Value::as_str)) {
+            (Some(key), Some(base)) if plain => {
+                let dir = env_dir.clone();
+                let env_dir_hash = tokio::task::spawn_blocking(move || crate::task::hash_dir(&dir)).await??;
+                let key = format!("{key}\n{env_dir_hash}");
+                let image = match &build_image {
+                    Some(built) => Some(setup_cache_image(&key, built)),
+                    None => image_id(base).await?.map(|id| setup_cache_image(&key, &id)),
+                };
+                let hit = match &image {
+                    Some(image) => image_id(image).await?.is_some(),
+                    None => false,
+                };
+                if let Some(image) = image.as_ref().filter(|_| hit) {
+                    log.line(&format!("agent setup cache: starting from {image}"));
+                    main.insert("image".into(), json!(image));
+                    main.remove("build");
+                    build_image = None;
+                }
+                Some(SetupCache { key, image, hit })
+            }
+            (Some(_), _) => {
+                log.line("agent setup cache: not used (the main container sets env, egress control, mounts or its own compose)");
+                None
+            }
+            (None, _) => None,
+        };
         let mut files = vec![write("docker-compose-base.json", json!({"services": {"main": main}}))?];
         let mut user_files = vec![];
         if task_compose.exists() {
@@ -816,13 +876,51 @@ impl DockerEnv {
             std::env::var("DOCKER_CONFIG").unwrap_or_else(|_| cache_home().join(".docker").display().to_string()),
         );
         crate::redact::register_env(&compose_env);
-        Ok(DockerEnv { spec, files, compose_env, build_image, egress, egress_services, log })
+        Ok(DockerEnv { spec, files, compose_env, build_image, egress, egress_services, log, setup_cache })
     }
 
     pub fn workdir(&self) -> Option<&str> {
         self.spec.config.workdir.as_deref()
     }
 
+    /// Commit the main container, as agent setup left it, to the setup cache image; a failure is logged,
+    /// not raised, since the trial does not need the image.
+    pub async fn commit_setup_cache(&self) {
+        let Some(cache) = self.setup_cache.as_ref().filter(|c| !c.hit) else { return };
+        let committed = async {
+            let ps = self.run_ok(to_owned_args(&["ps", "--quiet", MAIN]), Some(QUICK_COMMAND_TIMEOUT)).await?;
+            let id = ps.stdout.trim();
+            if id.is_empty() {
+                bail!("no main container");
+            }
+            let inspected =
+                docker_output(&["inspect", "--format", "{{.Image}} {{json .Config.Labels}}", id], QUICK_COMMAND_TIMEOUT).await?;
+            let inspected = String::from_utf8_lossy(&inspected.stdout);
+            let (main_image_id, labels) = inspected.trim().split_once(' ').ok_or_else(|| anyhow!("cannot inspect the main container"))?;
+            let image = cache.image.clone().unwrap_or_else(|| setup_cache_image(&cache.key, main_image_id));
+            if image_id(&image).await?.is_some() {
+                return anyhow::Ok(format!("{image} is already there (another trial committed it)"));
+            }
+            let labels: BTreeMap<String, String> = serde_json::from_str(labels).unwrap_or_default();
+            let mut changes = vec!["LABEL dev.hbench.setup-cache=1".to_string()];
+            changes.extend(labels.keys().filter(|k| k.starts_with("com.docker.compose.")).map(|k| format!("LABEL {k}=")));
+            let mut args: Vec<&str> = vec!["commit"];
+            for c in &changes {
+                args.extend(["--change", c]);
+            }
+            args.extend([id, &image]);
+            let out = docker_output(&args, SETUP_COMMIT_TIMEOUT).await?;
+            if !out.status.success() {
+                bail!("{}", String::from_utf8_lossy(&out.stderr).trim());
+            }
+            anyhow::Ok(format!("committed {image}"))
+        };
+        match committed.await {
+            Ok(done) => self.log.line(&format!("agent setup cache: {done}")),
+            Err(e) => self.log.line(&format!("agent setup cache: commit failed: {e:#}")),
+        }
+    }
+
     fn compose(&self, args: &[String]) -> Command {
         let mut c = Command::new("docker");
         c.arg("compose").arg("--project-name").arg(&self.spec.project);
@@ -1417,6 +1515,15 @@ pub async fn preflight() -> Result<()> {
 
 #[cfg(test)]
 mod tests {
+    #[test]
+    fn a_setup_cache_image_is_named_by_its_main_image_and_key() {
+        let image = setup_cache_image("{\"agent\":\"codex\"}", "sha256:aa");
+        assert!(image.starts_with("hbench-setup-cache:") && image.len() == "hbench-setup-cache:".len() + 16, "{image}");
+        assert_eq!(image, setup_cache_image("{\"agent\":\"codex\"}", "sha256:aa"));
+        assert_ne!(image, setup_cache_image("{\"agent\":\"codex\"}", "sha256:bb"));
+        assert_ne!(image, setup_cache_image("{\"agent\":\"claude-code\"}", "sha256:aa"));
+    }
+
     #[test]
     fn ssh_not_connected_only_before_a_session() {
         let r = |code, stderr: &str| ExecResult { code, stdout: String::new(), stderr: stderr.into() };

src/main.rs

@@ -749,6 +749,10 @@ struct RunArgs {
     #[arg(long, help_heading = "Agent")]
     skip_preflight: bool,
 
+    /// Start each trial's container from an image of it as agent setup left it, committed by the first trial
+    /// with the same image, agent, options and hbench version (experimental; see `hbench cache clean`).
+    #[arg(long, help_heading = "Environment")]
+    cache_agent_setup: bool,
     /// Environment (container) env var KEY=VALUE, repeatable.
     #[arg(long = "ee", alias = "environment-env", value_parser = parse_kv, help_heading = "Environment")]
     env_env: Vec<(String, String)>,
@@ -1003,6 +1007,7 @@ fn build_job_config(a: &RunArgs) -> Result<JobConfig> {
     cfg.verifier.env.extend(a.verifier_env.iter().cloned());
     cfg.verifier.disable |= a.disable_verification;
     cfg.install_only |= a.install_only;
+    cfg.environment.cache_agent_setup |= a.cache_agent_setup;
     cfg.quiet |= a.quiet;
     cfg.artifacts.extend(a.artifact.iter().cloned().map(Artifact::Path));
     cfg.extra_instructions.extend(a.extra_instruction.iter().cloned());
@@ -1823,14 +1828,15 @@ mod tests {
             "--no-delete",
             "--artifact",
             "/app/x",
+            "--cache-agent-setup",
         ]);
         let Cmd::Run(a) = cli.cmd else { panic!() };
         let c = build_job_config(&a).unwrap();
         assert_eq!(c.agents.len(), 2);
         assert_eq!(c.agents[1].model_name.as_deref(), Some("anthropic/b"));
         assert_eq!(c.agents[0].kwargs["max_turns"], 5);
         assert_eq!(c.agents[0].env["X"], "1");
-        assert_eq!((c.n_attempts, c.environment.delete), (3, false));
+        assert_eq!((c.n_attempts, c.environment.delete, c.environment.cache_agent_setup), (3, false, true));
         assert_eq!(c.datasets[0].task_names.as_ref().unwrap(), &["t*"]);
         assert_eq!(c.artifacts, vec![Artifact::Path("/app/x".into())]);
     }

src/tools.rs

@@ -39,6 +39,7 @@ pub async fn start_env(path: &Path, all: bool, interactive: bool, agent: Option<
         keep_containers: false,
         delete: true,
         log_dirs: LOG_DIRS.to_vec(),
+        setup_cache_key: None,
     };
     let env = DockerEnv::new(spec, log).await?;
     let r = async {
@@ -81,7 +82,8 @@ pub async fn start_env(path: &Path, all: bool, interactive: bool, agent: Option<
     r
 }
 
-/// Remove the hbench cache dir and hbench-built images (`hb__*`, `hbench-egress-sidecar`).
+/// Remove the hbench cache dir and hbench-built images (`hb__*`, `hbench-egress-sidecar`, `hbench-setup-cache`, and
+/// setup cache images a later commit of the same tag untagged).
 pub async fn cache_clean(force: bool, dry: bool, no_docker: bool, no_cache_dir: bool) -> Result<()> {
     let dir = cache_root();
     let size: u64 = if dir.exists() && !no_cache_dir { walk_size(&dir) } else { 0 };
@@ -91,9 +93,14 @@ pub async fn cache_clean(force: bool, dry: bool, no_docker: bool, no_cache_dir:
             Ok(o) if o.status.success() => {
                 images = String::from_utf8_lossy(&o.stdout)
                     .lines()
-                    .filter(|i| i.starts_with("hb__") || i.starts_with("hbench-egress-sidecar"))
+                    .filter(|i| i.starts_with("hb__") || i.starts_with("hbench-egress-sidecar") || i.starts_with("hbench-setup-cache"))
                     .map(String::from)
                     .collect();
+                let untagged_setup_caches = tokio::process::Command::new("docker")
+                    .args(["images", "--quiet", "--filter", "dangling=true", "--filter", "label=dev.hbench.setup-cache"])
+                    .output()
+                    .await?;
+                images.extend(String::from_utf8_lossy(&untagged_setup_caches.stdout).lines().map(String::from));
             }
             _ => eprintln!("could not list Docker images; skipping image cleanup"),
         }