diff --git a/assets/schemas/trial.json b/assets/schemas/trial.json index c8ec3d9..5c83d90 100644 --- a/assets/schemas/trial.json +++ b/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" diff --git a/src/agents/installed/claude_code.rs b/src/agents/installed/claude_code.rs index fd01c97..1d82ad8 100644 --- a/src/agents/installed/claude_code.rs +++ b/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 { self.version.clone() } diff --git a/src/agents/installed/codex.rs b/src/agents/installed/codex.rs index d1b4805..13b358b 100644 --- a/src/agents/installed/codex.rs +++ b/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 { self.version.clone() } diff --git a/src/agents/installed/deepseek_harness.rs b/src/agents/installed/deepseek_harness.rs index 8adf85f..41b0d09 100644 --- a/src/agents/installed/deepseek_harness.rs +++ b/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 { 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; } diff --git a/src/agents/installed/opencode.rs b/src/agents/installed/opencode.rs index e2b0a19..2fa1bc1 100644 --- a/src/agents/installed/opencode.rs +++ b/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 { self.version.clone() } diff --git a/src/agents/mod.rs b/src/agents/mod.rs index 5c4e1aa..f6e4da1 100644 --- a/src/agents/mod.rs +++ b/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 { 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 { serde_json::from_value(v).unwrap() }; diff --git a/src/config.rs b/src/config.rs index 16fab4a..5fd0861 100644 --- a/src/config.rs +++ b/src/config.rs @@ -278,6 +278,10 @@ pub struct EnvConfig { /// Docker kwargs; `keep_containers: true` stops instead of removing containers. pub kwargs: BTreeMap, pub extra_allowed_hosts: Vec, + /// 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(k: &Option, s: S) -> Result { @@ -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] diff --git a/src/env.rs b/src/env.rs index f81367a..8cd29b5 100644 --- a/src/env.rs +++ b/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, } pub struct DockerEnv { @@ -445,6 +450,27 @@ pub struct DockerEnv { /// The services behind the sidecar. egress_services: Vec, log: Arc, + setup_cache: Option, +} + +/// 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, + hit: bool, +} + +/// `hbench-setup-cache:` 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> { + 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 = 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() }; diff --git a/src/main.rs b/src/main.rs index a7694dd..38fb166 100644 --- a/src/main.rs +++ b/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 { 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,6 +1828,7 @@ 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(); @@ -1830,7 +1836,7 @@ mod tests { 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())]); } diff --git a/src/tools.rs b/src/tools.rs index 667da0f..55abf5d 100644 --- a/src/tools.rs +++ b/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"), } diff --git a/src/trial.rs b/src/trial.rs index bbbe9db..54c22da 100644 --- a/src/trial.rs +++ b/src/trial.rs @@ -434,9 +434,44 @@ impl Trial { keep_containers: e.keep_containers(), delete: e.delete, log_dirs: if agent_env { LOG_DIRS.to_vec() } else { vec!["/logs/verifier"] }, + setup_cache_key: None, }) } + /// What, besides the main image and the task's environment dir, this trial's agent setup image depends on: + /// None when the cache is off, the agent's setup is not cacheable, or setup depends on more than the agent and + /// its options. + fn setup_cache_key(&self, step: Option<&StepConfig>) -> Option { + let a = &self.cfg.agent; + if !self.cfg.environment.cache_agent_setup { + return None; + } + let task_env = &self.task.config.environment; + let unshareable = [ + (!self.agent.as_ref().is_some_and(|agent| agent.setup_is_cacheable()), "this agent's setup is not cacheable"), + (!a.skills.is_empty() || task_env.skills_dir.is_some(), "skills"), + (!task_env.mcp_servers.is_empty(), "the task's MCP servers"), + (self.cfg.user_agent.is_some(), "a user agent"), + (self.cfg.environment.force_build, "--force-build"), + (a.kwargs.get("version").and_then(Value::as_str) == Some("latest"), "version=latest"), + ]; + if let Some((_, why)) = unshareable.iter().find(|(applies, _)| *applies) { + self.log.line(&format!("agent setup cache: not used ({why})")); + return None; + } + let key = serde_json::json!({ + "hbench": crate::build_info::VERSION, + "agent": a.name(), + "model": a.model_name, + "kwargs": a.kwargs, + "mcp_servers": a.mcp_servers, + "env": a.env.keys().collect::>(), + "user": self.agent_user(step), + "scrub": crate::integrity::scrub_command(&crate::integrity::leaks(&self.task.dir)), + }); + Some(key.to_string()) + } + fn record(&mut self, e: &anyhow::Error) { self.log.line(&format!("exception: {e:?}")); if self.result.exception_info.is_some() { @@ -619,6 +654,7 @@ impl Trial { if self.user_agent.is_some() { spec.log_dirs.push("/logs/user-agent"); } + spec.setup_cache_key = self.setup_cache_key(steps[0]); let env = DockerEnv::new(spec, self.log.clone()).await?; self.env = Some(env); @@ -726,6 +762,7 @@ impl Trial { self.result.agent_info.version = a.version().unwrap_or_else(|| "unknown".into()); self.log.line(&format!("agent {} setup done (version {})", a.name(), self.result.agent_info.version)); r?; + self.env().commit_setup_cache().await; if self.cfg.install_only { return Ok(()); }