From 434a28332100442383a1258c6ad2c1e20261b275 Mon Sep 17 00:00:00 2001 From: Myles Anderson Date: Thu, 22 Oct 2026 14:39:35 -0700 Subject: [PATCH 01/10] Configurable remote storage root for SSH and Slurm (OR-510) Co-Authored-By: Claude Opus 5.5 --- .../orx-compute/references/configuration.md | 2 + agent-skills/orx-compute/references/slurm.md | 5 + agent-skills/orx-compute/references/ssh.md | 7 +- src/commands/compute.rs | 28 ++- src/commands/supervise.rs | 26 ++- src/commands/tensorboard.rs | 16 +- src/commands/up/compute_settings.rs | 50 ++++- src/compute.rs | 9 +- src/config.rs | 3 + src/jobs/mod.rs | 4 + src/jobs/slurm.rs | 45 +++-- src/jobs/ssh.rs | 186 +++++++++++++++--- src/jobs/ssh/container.rs | 11 +- src/jobs/ssh/container/tests.rs | 3 +- src/local/hf.rs | 1 + src/local/k8s.rs | 1 + src/local/localrun.rs | 1 + src/local/modal.rs | 1 + src/local/openresearch.rs | 1 + src/local/ray.rs | 1 + src/local/slurm.rs | 3 + src/local/ssh.rs | 3 + src/telemetry.rs | 1 + 23 files changed, 332 insertions(+), 76 deletions(-) diff --git a/agent-skills/orx-compute/references/configuration.md b/agent-skills/orx-compute/references/configuration.md index b9a83bb5c..fb79a26e4 100644 --- a/agent-skills/orx-compute/references/configuration.md +++ b/agent-skills/orx-compute/references/configuration.md @@ -10,6 +10,7 @@ orx compute show ssh --json orx compute configure ssh --default-host lab orx compute configure ssh --host lab --container research orx compute configure ssh --host lab --clear container +orx compute configure ssh --host lab --remote-root /data/me/orx # runs + source cache; default ~/.orx orx compute test ssh --host lab # tests saved container too orx compute test ssh --host lab --no-container # explicitly tests host orx compute connect ssh --host lab # interactive login/MFA @@ -17,6 +18,7 @@ orx compute configure slurm --host cluster --partition gpu --account lab --time- orx compute configure slurm --cpus-per-task 8 --mem 64G # defaults; exp run --cpus/--mem override orx compute configure slurm --clear cpus-per-task --clear mem orx compute configure slurm --clear time-limit # cluster chooses +orx compute configure slurm --remote-root /scratch/me/orx # shared FS; --clear remote-root for ~/.orx orx compute connect slurm --host cluster orx compute configure k8s --context research --namespace experiments orx compute configure ray --address http://127.0.0.1:8265 diff --git a/agent-skills/orx-compute/references/slurm.md b/agent-skills/orx-compute/references/slurm.md index 586ddc535..e891143a7 100644 --- a/agent-skills/orx-compute/references/slurm.md +++ b/agent-skills/orx-compute/references/slurm.md @@ -22,5 +22,10 @@ orx exp run --backend slurm profile—is used as-is. - `--timeout` covers the batch job. Unset falls back to the Slurm time-limit setting, then the cluster's limit. +- Runs and the cached source archive live under `~/.orx` on the cluster unless + `orx compute configure slurm --remote-root /scratch//orx` names another + absolute directory, such as scratch space when home quota is small. It must + be on a filesystem the compute nodes share. Each run keeps the root it was + submitted with; `--clear remote-root` restores the default for new runs. - A detached `orx supervise` process records scheduler status and logs; do not kill it. diff --git a/agent-skills/orx-compute/references/ssh.md b/agent-skills/orx-compute/references/ssh.md index 81b6960d8..ebe056391 100644 --- a/agent-skills/orx-compute/references/ssh.md +++ b/agent-skills/orx-compute/references/ssh.md @@ -32,7 +32,10 @@ orx exp run --backend ssh --host lab --container research orx exp run --backend ssh --host lab --no-container ``` -Machine-local SSH settings can save a container name or ID for each host. +Machine-local SSH settings can save a container name or ID for each host, and +an absolute remote storage root (`--remote-root /data//orx`) used instead +of `~/.orx` for that host's runs and cached source archives. Each run keeps the +root it was launched with. A saved default SSH host lets launches omit `--host`. These defaults apply only to SSH experiments; dashboards, terminals and coding agents still connect directly to the host. @@ -54,7 +57,7 @@ creation or lifecycle management. `--image` and `--flavor` are unsupported. The committed source archive is cached on the host, then streamed into the container's `$HOME/.orx/runs//repo`. Host logs and exit status stay in the -host's `~/.orx/runs//`, so container removal does not erase them. A detached +host's `~/.orx/runs//` (or the saved remote root), so container removal does not erase them. A detached supervisor reattaches using the saved container ID and start time; later settings changes cannot redirect an existing run. Do not kill the supervisor. diff --git a/src/commands/compute.rs b/src/commands/compute.rs index 4a3d19cfa..7014286a3 100644 --- a/src/commands/compute.rs +++ b/src/commands/compute.rs @@ -235,6 +235,9 @@ pub struct ConfigureArgs { host: Option, #[arg(long)] container: Option, + /// Absolute remote dir for runs and the source cache (SSH per host, or Slurm); default ~/.orx. + #[arg(long)] + remote_root: Option, #[arg(long)] default_host: Option, #[arg(long)] @@ -321,7 +324,7 @@ impl ConfigureArgs { fn body(self) -> Result { let backend = self.backend.name(); let allowed: &[&str] = match backend { - "ssh" => &["host", "container", "defaultHost"], + "ssh" => &["host", "container", "remoteRoot", "defaultHost"], "slurm" => &[ "host", "partition", @@ -329,6 +332,7 @@ impl ConfigureArgs { "timeLimit", "cpusPerTask", "mem", + "remoteRoot", ], "k8s" => &["context", "namespace"], "ray" => &["address"], @@ -341,6 +345,7 @@ impl ConfigureArgs { for (key, value) in [ ("host", self.host), ("container", self.container), + ("remoteRoot", self.remote_root), ("defaultHost", self.default_host), ("partition", self.partition), ("account", self.account), @@ -389,6 +394,7 @@ impl ConfigureArgs { "time-limit" => "timeLimit", "cpus-per-task" => "cpusPerTask", "default-host" => "defaultHost", + "remote-root" => "remoteRoot", other => other, }; if !allowed.contains(&key) @@ -421,11 +427,14 @@ impl ConfigureArgs { "Supply settings to change; use compute show {backend} to inspect." )); } - if backend == "ssh" && body.contains_key("host") && !body.contains_key("container") { - return Err(anyhow!("--host selects a container setting; supply --container or --clear container. Use --default-host to select a default.")); + let host_setting = body.contains_key("container") || body.contains_key("remoteRoot"); + if backend == "ssh" && body.contains_key("host") && !host_setting { + return Err(anyhow!("--host selects a host setting; supply --container, --remote-root, or --clear for one of them. Use --default-host to select a default.")); } - if backend == "ssh" && body.contains_key("container") && !body.contains_key("host") { - return Err(anyhow!("A container change requires --host.")); + if backend == "ssh" && host_setting && !body.contains_key("host") { + return Err(anyhow!( + "A container or remote root change requires --host." + )); } Ok(Value::Object(body)) } @@ -766,6 +775,15 @@ mod tests { body(&["ssh", "--host", "lab", "--clear", "container"]).unwrap(), json!({"host":"lab", "container":null}) ); + assert_eq!( + body(&["ssh", "--host", "lab", "--remote-root", "/scratch/me/orx"]).unwrap(), + json!({"host":"lab", "remoteRoot":"/scratch/me/orx"}) + ); + assert_eq!( + body(&["slurm", "--clear", "remote-root"]).unwrap(), + json!({"remoteRoot":""}) + ); + assert!(body(&["ssh", "--remote-root", "/scratch/me/orx"]).is_err()); assert!(crate::Cli::try_parse_from(["orx", "compute", "status", "--json"]).is_ok()); } } diff --git a/src/commands/supervise.rs b/src/commands/supervise.rs index 327b29f38..4be9e58ea 100644 --- a/src/commands/supervise.rs +++ b/src/commands/supervise.rs @@ -787,10 +787,14 @@ async fn watch_ssh_job( let mut job = job; let mut exit_code = None; if descriptor.rental_lease.is_some() && is_terminal_stage(&job.stage) { - exit_code = ssh::ssh_run(&target, &format!("cat \"$HOME/{dir}/exit_code\""), None) - .await - .ok() - .and_then(|value| value.trim().parse::().ok()); + exit_code = ssh::ssh_run( + &target, + &format!("cat {}/exit_code", ssh::shell_path(&dir)), + None, + ) + .await + .ok() + .and_then(|value| value.trim().parse::().ok()); if let Some(code) = exit_code { job.stage = if code == 0 { "COMPLETED" } else { "ERROR" }.into(); job.message = (code != 0).then(|| format!("exited with code {code}")); @@ -840,7 +844,7 @@ async fn watch_ssh_job( { let _ = ssh::ssh_run( &target, - &format!("touch \"$HOME/{dir}/final_collected\""), + &format!("touch {}/final_collected", ssh::shell_path(&dir)), None, ) .await; @@ -855,8 +859,9 @@ async fn watch_ssh_job( job.message.as_deref().unwrap_or("Container running.") ); let command = format!( - "printf '%s\\n' {} >> \"$HOME/{dir}/log\"", - ssh::sh_quote(&message) + "printf '%s\\n' {} >> {}/log", + ssh::sh_quote(&message), + ssh::shell_path(&dir) ); match ssh::ssh_run(&target, &command, None).await { Ok(_) => last_message = job.message, @@ -1131,7 +1136,7 @@ async fn run_openresearch( return Ok(()); } let staged = - ssh::stage_source(&target, &run_id, &source.path, &source.digest, None).await; + ssh::stage_source(&target, &run_id, &source.path, &source.digest, None, None).await; if let Err(err) = staged { eprintln!("supervise {run_id}: source staging failed (will retry): {err}"); launch_err = Some(err); @@ -1144,6 +1149,7 @@ async fn run_openresearch( container: None, target: target.clone(), run_id: run_id.clone(), + remote_root: None, script: script.clone(), env: env.clone(), }) @@ -1358,7 +1364,7 @@ async fn run_slurm( let (host, job_id) = descriptor.slurm_ref()?; let host = host.to_string(); let job_id = job_id.to_string(); - let dir = slurm::run_dir(&run_id); + let dir = slurm::run_dir(descriptor.remote_root.as_deref(), &run_id); eprintln!("supervise {run_id}: watching slurm job {job_id} on {host}"); @@ -1382,7 +1388,7 @@ async fn run_slurm( if !cancel_sent && local_cancel_requested(&store, &run_id) { cancel_slurm(&host, &job_id, &run_id, &mut cancel_sent).await; } - let mut observed = slurm::inspect_job(&host, &run_id, &job_id).await; + let mut observed = slurm::inspect_job(&host, &dir, &job_id).await; if let Ok(job) = &mut observed { if cancel_sent && job.stage == "GONE" { job.stage = "CANCELED".into(); diff --git a/src/commands/tensorboard.rs b/src/commands/tensorboard.rs index bc4971009..8f97ee9e9 100644 --- a/src/commands/tensorboard.rs +++ b/src/commands/tensorboard.rs @@ -25,14 +25,20 @@ fn viewer_command(run_id: &str, backend: &BackendDescriptor) -> Result { let (host, dir) = backend.ssh_ref()?; (host, dir.to_string()) } - "slurm_job" => (backend.slurm_ref()?.0, crate::jobs::slurm::run_dir(run_id)), + "slurm_job" => ( + backend.slurm_ref()?.0, + crate::jobs::slurm::run_dir(backend.remote_root.as_deref(), run_id), + ), _ => { return Err(anyhow!( "TensorBoard logs are kept only for local, Slurm, and SSH runs outside a container" )) } }; - let remote = format!("tensorboard --logdir \"$HOME\"/{dir}/tensorboard --port {PORT}"); + let remote = format!( + "tensorboard --logdir {}/tensorboard --port {PORT}", + crate::jobs::ssh::shell_path(&dir) + ); Ok(format!( "ssh -L {PORT}:localhost:{PORT} {} {}\nThen open http://localhost:{PORT}", sh_quote(host), @@ -65,9 +71,13 @@ mod tests { format!(r#"{{"kind":"{kind}","namespace":"gpu-box","jobId":".orx/runs/r1"}}"#); let command = viewer_command("r1", &backend(&json)).unwrap(); assert!(command.starts_with( - "ssh -L 6006:localhost:6006 'gpu-box' 'tensorboard --logdir \"$HOME\"/.orx/runs/r1/tensorboard --port 6006'" + "ssh -L 6006:localhost:6006 'gpu-box' 'tensorboard --logdir \"$HOME/.orx/runs/r1\"/tensorboard --port 6006'" )); } + let slurm = r#"{"kind":"slurm_job","namespace":"gpu-box","jobId":"7","remoteRoot":"/scratch/me/orx"}"#; + assert!(viewer_command("r1", &backend(slurm)).unwrap().contains( + "'tensorboard --logdir '\\''/scratch/me/orx/runs/r1'\\''/tensorboard --port 6006'" + )); } #[test] diff --git a/src/commands/up/compute_settings.rs b/src/commands/up/compute_settings.rs index 927c27ff3..7d8659b53 100644 --- a/src/commands/up/compute_settings.rs +++ b/src/commands/up/compute_settings.rs @@ -587,6 +587,7 @@ pub(super) async fn ssh_settings() -> ApiResult { .cloned() .unwrap_or_default(); host["container"] = json!(options.container); + host["remoteRoot"] = json!(options.remote_root); } // Best-effort, like the preflight write: a store hiccup shouldn't take // out the host listing — hosts just render as never tested. @@ -622,16 +623,27 @@ pub(super) async fn ssh_settings() -> ApiResult { pub(super) struct SaveSshSettingsReq { host: String, container: Option, + /// `None` keeps the saved root; `Some("")` restores the `~/.orx` default. + #[serde(default)] + remote_root: Option, } pub(super) async fn save_ssh_settings(Json(req): Json) -> ApiResult { let host = req.host.trim().to_string(); require_configured_ssh_host(&host).await?; - let options = crate::config::SshHostSettings { - container: req.container, - }; - crate::jobs::ssh::validate_host_options(&options).map_err(bad_request)?; blocking_api(move || { + let remote_root = match req.remote_root { + Some(root) => Some(root.trim().to_string()).filter(|root| !root.is_empty()), + None => crate::config::ssh_settings()? + .hosts + .get(&host) + .and_then(|saved| saved.remote_root.clone()), + }; + let options = crate::config::SshHostSettings { + container: req.container, + remote_root, + }; + crate::jobs::ssh::validate_host_options(&options).map_err(bad_request)?; crate::config::set_ssh_host(host, options)?; Ok(Json(json!({"ok": true}))) }) @@ -848,6 +860,7 @@ pub(super) fn slurm_settings_json() -> Value { "timeLimit": settings.time_limit, "cpusPerTask": settings.cpus_per_task, "mem": settings.mem, + "remoteRoot": settings.remote_root, "hosts": list_ssh_hosts(), }) } @@ -869,6 +882,7 @@ pub(super) struct SetSlurmSettingsReq { /// Cores per job. `Some(0)` clears it (partition default). cpus_per_task: Option, mem: Option, + remote_root: Option, } pub(super) async fn set_slurm_settings(Json(req): Json) -> ApiResult { @@ -905,6 +919,13 @@ pub(super) async fn set_slurm_settings(Json(req): Json) -> None => None, }; } + if let Some(root) = req.remote_root { + let root = norm(root); + if let Some(root) = &root { + crate::jobs::ssh::validate_remote_root(root).map_err(bad_request)?; + } + settings.remote_root = root; + } slurm::save_settings(&settings)?; Ok(Json(slurm_settings_json())) }) @@ -1447,10 +1468,23 @@ pub(crate) async fn configure(backend: &str, body: Value) -> Result { .await .map_err(|ApiError(_, message)| anyhow!(message))?; } - if body.get("container").is_some() { - let req: SaveSshSettingsReq = serde_json::from_value( - json!({"host": body["host"], "container": body["container"]}), - )?; + if body.get("container").is_some() || body.get("remoteRoot").is_some() { + // A root-only change keeps the host's saved container. + let container = match body.get("container") { + Some(container) => container.clone(), + None => { + let host = body["host"].as_str().unwrap_or_default(); + json!(crate::config::ssh_settings()? + .hosts + .get(host) + .and_then(|saved| saved.container.clone())) + } + }; + let mut req = json!({"host": body["host"], "container": container}); + if let Some(root) = body.get("remoteRoot") { + req["remoteRoot"] = root.clone(); + } + let req: SaveSshSettingsReq = serde_json::from_value(req)?; cli_result(save_ssh_settings(Json(req)).await)?; } if let Some(host) = body.get("defaultHost") { diff --git a/src/compute.rs b/src/compute.rs index 7d1044005..3de0a3ae6 100644 --- a/src/compute.rs +++ b/src/compute.rs @@ -850,6 +850,7 @@ pub async fn submit(args: &crate::ExpRunArgs) -> Result { source_digest: None, source_path: None, source_size: None, + remote_root: None, }; source.0.apply_to_descriptor(&mut descriptor); let now = crate::store::now_ms(); @@ -1190,9 +1191,10 @@ pub(crate) mod tests { "lab".into(), SshHostSettings { container: Some("research".into()), + remote_root: Some("/scratch/lab/orx".into()), }, ), - ("direct".into(), SshHostSettings { container: None }), + ("direct".into(), SshHostSettings::default()), ]), }; let mut args = tinker_args(); @@ -1219,6 +1221,11 @@ pub(crate) mod tests { let (host, options) = crate::jobs::ssh::resolve_options(&args, &settings).unwrap(); assert_eq!(host, expected_host); assert_eq!(options.container.as_deref(), expected_container); + assert_eq!( + options.remote_root.is_some(), + host == "lab", + "the saved root follows the host, not the container" + ); } args.host = None; assert!(crate::jobs::ssh::resolve_options(&args, &SshSettings::default()).is_err()); diff --git a/src/config.rs b/src/config.rs index 4a27fc3f8..adbb02d20 100644 --- a/src/config.rs +++ b/src/config.rs @@ -461,6 +461,9 @@ pub fn write_synced_env_vars(values: &[(&str, &str)]) -> Result<()> { #[serde(rename_all = "camelCase")] pub struct SshHostSettings { pub container: Option, + /// Absolute storage root for runs and the source cache; `None` is `~/.orx`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_root: Option, } #[derive(Debug, Clone, Default, Serialize, Deserialize)] diff --git a/src/jobs/mod.rs b/src/jobs/mod.rs index f200022d9..7996e866b 100644 --- a/src/jobs/mod.rs +++ b/src/jobs/mod.rs @@ -106,6 +106,9 @@ pub struct BackendDescriptor { pub source_path: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub source_size: Option, + /// The host's storage root at submit (ssh/slurm); `None` is `~/.orx`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_root: Option, } impl BackendDescriptor { @@ -295,6 +298,7 @@ mod tests { source_digest: None, source_path: None, source_size: None, + remote_root: None, } } diff --git a/src/jobs/slurm.rs b/src/jobs/slurm.rs index 113150ca3..0ab9a0be5 100644 --- a/src/jobs/slurm.rs +++ b/src/jobs/slurm.rs @@ -27,7 +27,7 @@ use std::collections::HashMap; use serde::{Deserialize, Serialize}; -use super::ssh::{sh_quote, ssh_run, SshTarget}; +use super::ssh::{sh_quote, shell_path, ssh_run, SshTarget}; use crate::error::{anyhow, Result}; // --- settings --------------------------------------------------------------- @@ -58,6 +58,9 @@ pub struct SlurmSettings { /// default. #[serde(default)] pub mem: Option, + /// Absolute shared-filesystem root for runs and the source cache; `None` is `~/.orx`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_root: Option, } fn settings_path() -> std::path::PathBuf { @@ -96,8 +99,10 @@ pub fn save_settings(settings: &SlurmSettings) -> Result<()> { pub struct SlurmJobSpec { /// ssh config host alias of the cluster's login node. pub host: String, - /// Names the remote run dir `~/.orx/runs/`. + /// Names the remote run dir `/runs/`. pub run_id: String, + /// Configured storage root; `None` is `~/.orx`. + pub remote_root: Option, /// Runs on the login node in the run dir at submit time, with `env` /// exported — the clone step. Piped to `bash -s`, never written to disk. pub setup_script: String, @@ -263,28 +268,28 @@ fn parse_job_id(out: &str) -> Result { // --- lifecycle ---------------------------------------------------------------- -/// The remote run dir for a run, relative to `$HOME` (shared convention with -/// the ssh backend; derived, not stored in the descriptor). -pub fn run_dir(run_id: &str) -> String { - format!(".orx/runs/{run_id}") +/// The remote run dir for a run (shared convention with the ssh backend), +/// derived from the root recorded in the descriptor at submit. +pub fn run_dir(root: Option<&str>, run_id: &str) -> String { + super::ssh::remote_run_dir(root, run_id) } /// Submit the job: create the run dir, clone on the login node, write /// job.sbatch, `sbatch --parsable`. Returns the Slurm job id — the reattach /// handle (together with the host). pub async fn run_job(spec: &SlurmJobSpec) -> Result { - let dir = run_dir(&spec.run_id); + let dir = shell_path(&run_dir(spec.remote_root.as_deref(), &spec.run_id)); // Login-node setup (clone). Env + script travel via stdin so tokens never // land on an argv or in a file. let setup = format!( - "{exports}\ncd \"$HOME/{dir}\" || exit 97\n{script}\n", + "{exports}\ncd {dir} || exit 97\n{script}\n", exports = render_exports(&spec.env), script = spec.setup_script, ); ssh_run( &SshTarget::alias(&spec.host), - &format!("mkdir -p \"$HOME/{dir}\" && chmod 700 \"$HOME/{dir}\" && bash -s"), + &format!("mkdir -p {dir} && chmod 700 {dir} && bash -s"), Some(&setup), ) .await @@ -293,13 +298,13 @@ pub async fn run_job(spec: &SlurmJobSpec) -> Result { // Write the batch script (owner-only: it embeds tokens) and submit. ssh_run( &SshTarget::alias(&spec.host), - &format!("umask 077 && cat > \"$HOME/{dir}/job.sbatch\""), + &format!("umask 077 && cat > {dir}/job.sbatch"), Some(&render_sbatch(spec)), ) .await?; let out = ssh_run( &SshTarget::alias(&spec.host), - &format!("cd \"$HOME/{dir}\" && sbatch --parsable job.sbatch"), + &format!("cd {dir} && sbatch --parsable job.sbatch"), None, ) .await @@ -335,11 +340,11 @@ fn queue_probe(job_id: &str) -> String { /// /// `GONE` means monitoring is unavailable, not that the job terminated. /// Accounting can be disabled or unavailable while a job still runs. -pub async fn inspect_job(host: &str, run_id: &str, job_id: &str) -> Result { - let dir = run_dir(run_id); +pub async fn inspect_job(host: &str, dir: &str, job_id: &str) -> Result { + let dir = shell_path(dir); let queue = queue_probe(job_id); let cmd = format!( - "d=\"$HOME/{dir}\"; \ + "d={dir}; \ if [ -f \"$d/exit_code\" ]; then echo \"EXIT $(cat \"$d/exit_code\")\"; \ else queue=$({queue}); \ case \"$queue\" in \ @@ -519,6 +524,7 @@ mod tests { SlurmJobSpec { host: "cluster".into(), run_id: "0123456789abcdef".into(), + remote_root: None, setup_script: "git clone x repo".into(), command: "python train.py".into(), env: HashMap::new(), @@ -648,6 +654,7 @@ mod tests { let mk = |run_id: &str, command: &str| SlurmJobSpec { host: host.clone(), run_id: run_id.into(), + remote_root: None, setup_script: "mkdir -p repo".into(), command: command.into(), env: HashMap::from([("ORX_E2E".into(), "1".into())]), @@ -663,7 +670,9 @@ mod tests { let host = host.clone(); async move { for _ in 0..150 { - let s = inspect_job(&host, &run_id, &job_id).await.unwrap(); + let s = inspect_job(&host, &run_dir(None, &run_id), &job_id) + .await + .unwrap(); if until.contains(&s.stage.as_str()) { return s; } @@ -679,8 +688,8 @@ mod tests { let run_b = format!("e2e-b-{}", uuid::Uuid::new_v4()); let cleanup = format!( "rm -rf \"$HOME/{}\" \"$HOME/{}\"", - run_dir(&run_a), - run_dir(&run_b) + run_dir(None, &run_a), + run_dir(None, &run_b) ); // Happy path: runs, completes, logs arrive. @@ -692,7 +701,7 @@ mod tests { let mut lines = Vec::new(); crate::jobs::ssh::stream_logs( &SshTarget::alias(&host), - &run_dir(&run_a), + &run_dir(None, &run_a), 0, std::time::Duration::from_secs(5), true, diff --git a/src/jobs/ssh.rs b/src/jobs/ssh.rs index 8e23ee77a..2beef8f30 100644 --- a/src/jobs/ssh.rs +++ b/src/jobs/ssh.rs @@ -83,15 +83,67 @@ fn prepare_control_dir() -> Result<()> { pub struct ResolvedLaunch { pub target: SshTarget, pub container: Option, + pub remote_root: Option, } pub fn validate_host_options(options: &crate::config::SshHostSettings) -> Result<()> { if let Some(reference) = &options.container { validate_container_reference(reference)?; } + if let Some(root) = &options.remote_root { + validate_remote_root(root)?; + } Ok(()) } +/// A configured remote storage root: absolute and free of shell metacharacters, +/// so it stays one path however a remote shell or batch script reads it. +pub fn validate_remote_root(root: &str) -> Result<()> { + let components = root.strip_prefix('/').map(|rest| rest.split('/')); + let valid = components.is_some_and(|mut components| { + components.all(|part| { + !matches!(part, "" | "." | "..") + && part + .chars() + .all(|c| c.is_ascii_alphanumeric() || "._-+@".contains(c)) + }) + }); + if !valid { + return Err(anyhow!( + "Remote storage root must be an absolute path such as /scratch/$USER/orx, using letters, numbers, '.', '_', '-', '+' or '@' (no trailing slash)." + )); + } + Ok(()) +} + +/// A run's remote dir: under the configured root, or `$HOME`-relative by default. +pub fn remote_run_dir(root: Option<&str>, run_id: &str) -> String { + match root { + Some(root) => format!("{root}/runs/{run_id}"), + None => format!(".orx/runs/{run_id}"), + } +} + +fn remote_source_cache(root: Option<&str>, digest: &str) -> String { + match root { + Some(root) => format!("{root}/source/{digest}.tar"), + None => format!(".orx/source/{digest}.tar"), + } +} + +/// One shell word for a remote path; relative paths resolve against `$HOME`. +pub(crate) fn shell_path(path: &str) -> String { + if path.starts_with('/') { + sh_quote(path) + } else { + format!("\"$HOME/{path}\"") + } +} + +fn parent_dir(path: &str) -> &str { + path.rsplit_once('/').map_or(path, |(parent, _)| parent) +} + pub fn resolve_options( args: &crate::ExpRunArgs, settings: &crate::config::SshSettings, @@ -111,7 +163,10 @@ pub fn resolve_options( } else { args.container.clone().or_else(|| saved.container.clone()) }; - let options = crate::config::SshHostSettings { container }; + let options = crate::config::SshHostSettings { + container, + remote_root: saved.remote_root, + }; validate_host_options(&options)?; Ok((host, options)) } @@ -133,7 +188,11 @@ pub async fn resolve_launch(args: &crate::ExpRunArgs) -> Result Some(reference) => Some(container::resolve(&target, &reference).await?), None => None, }; - Ok(ResolvedLaunch { target, container }) + Ok(ResolvedLaunch { + target, + container, + remote_root: options.remote_root, + }) } /// An ssh endpoint. The classic ssh backend connects by `~/.ssh/config` alias @@ -689,9 +748,10 @@ pub async fn stage_source( archive: &std::path::Path, digest: &str, container: Option<&ContainerRun>, + root: Option<&str>, ) -> Result { - let dir = format!(".orx/runs/{run_id}"); - let cache = format!(".orx/source/{digest}.tar"); + let dir = remote_run_dir(root, run_id); + let cache = remote_source_cache(root, digest); let size = tokio::fs::metadata(archive).await?.len(); let (check, upload) = source_cache_commands(&cache, size); if ssh_run(target, &check, None).await?.trim() != "present" { @@ -705,15 +765,21 @@ pub async fn stage_source( container.require_running(target).await?; let path = sh_quote(&container.run_dir); let extract = container.exec(&format!("set -e; umask 077; mkdir -p {path}/repo; chmod 700 {path} {path}/repo; tar -xf - -C {path}/repo")); + let (d, cache) = (shell_path(&dir), shell_path(&cache)); // Start the bound here too: submission may fail before detached launch. - ssh_run(target, &format!("umask 077; mkdir -p \"$HOME/{dir}\" && chmod 700 \"$HOME/{dir}\" && {extract} < \"$HOME/{cache}\" && date +%s > \"$HOME/{dir}/launch_time.tmp\" && mv \"$HOME/{dir}/launch_time.tmp\" \"$HOME/{dir}/launch_time\""), None).await?; + ssh_run(target, &format!("umask 077; mkdir -p {d} && chmod 700 {d} && {extract} < {cache} && date +%s > {d}/launch_time.tmp && mv {d}/launch_time.tmp {d}/launch_time"), None).await?; } else { + let (runs, d, cache) = ( + shell_path(parent_dir(&dir)), + shell_path(&dir), + shell_path(&cache), + ); ssh_run( target, &format!( - "umask 077; mkdir -p \"$HOME/.orx/runs\" \"$HOME/{dir}/repo\"; \ - chmod 700 \"$HOME/.orx/runs\" \"$HOME/{dir}\" \"$HOME/{dir}/repo\"; \ - tar -xf \"$HOME/{cache}\" -C \"$HOME/{dir}/repo\"" + "umask 077; mkdir -p {runs} {d}/repo; \ + chmod 700 {runs} {d} {d}/repo; \ + tar -xf {cache} -C {d}/repo" ), None, ) @@ -726,12 +792,13 @@ pub async fn stage_source( /// still lets `cat` exit 0 on EOF, so both sides require the exact byte length; /// a short cached file reads as absent and is replaced by the next upload. fn source_cache_commands(cache: &str, size: u64) -> (String, String) { + let (dir, file) = (shell_path(parent_dir(cache)), shell_path(cache)); let check = format!( - "f=\"$HOME/{cache}\"; test -f \"$f\" && test $(wc -c < \"$f\") -eq {size} && echo present || true" + "f={file}; test -f \"$f\" && test $(wc -c < \"$f\") -eq {size} && echo present || true" ); let upload = format!( - "umask 077; mkdir -p \"$HOME/.orx/source\"; tmp=\"$HOME/{cache}.tmp.$$\"; \ - cat > \"$tmp\" && test $(wc -c < \"$tmp\") -eq {size} && mv \"$tmp\" \"$HOME/{cache}\" \ + "umask 077; mkdir -p {dir}; f={file}; tmp=\"$f.tmp.$$\"; \ + cat > \"$tmp\" && test $(wc -c < \"$tmp\") -eq {size} && mv \"$tmp\" \"$f\" \ || {{ rm -f \"$tmp\"; echo \"source archive upload incomplete\" >&2; exit 1; }}" ); (check, upload) @@ -745,8 +812,10 @@ pub(crate) fn sh_quote(s: &str) -> String { pub struct SshJobSpec { /// Where to run: a config alias (the ssh backend) or an explicit endpoint. pub target: SshTarget, - /// Names the remote run dir `~/.orx/runs/`. + /// Names the remote run dir `/runs/`. pub run_id: String, + /// Configured storage root; `None` is `~/.orx`. + pub remote_root: Option, /// The shared snapshot-and-run payload (`bash` script body). pub script: String, /// Exported inside run.sh on the remote (tokens, synced env). @@ -772,7 +841,8 @@ trap 'stop_child TERM 143' TERM trap 'stop_child INT 130' INT"#; fn host_script(dir: &str, exports: &str, script: &str) -> String { - format!("#!/usr/bin/env bash\ncd \"$HOME/{dir}\" || exit 97\nchild=\n{TERMINAL_TRAPS}\n(\n{exports}\n{script}\n) > log 2>&1 &\nchild=$!\nwait \"$child\"\n") + let dir = shell_path(dir); + format!("#!/usr/bin/env bash\ncd {dir} || exit 97\nchild=\n{TERMINAL_TRAPS}\n(\n{exports}\n{script}\n) > log 2>&1 &\nchild=$!\nwait \"$child\"\n") } #[derive(Debug)] @@ -787,9 +857,9 @@ impl std::fmt::Display for LaunchUncertain { impl std::error::Error for LaunchUncertain {} /// Submit the job: write run.sh, launch it detached, record its pid. Returns -/// the remote run dir (relative to `$HOME`) — the reattach handle. +/// the remote run dir (absolute, or relative to `$HOME`) — the reattach handle. pub async fn run_job(spec: &SshJobSpec) -> Result { - let dir = format!(".orx/runs/{}", spec.run_id); + let dir = remote_run_dir(spec.remote_root.as_deref(), &spec.run_id); let env = super::default_python_env(&spec.env); let exports: String = env .iter() @@ -806,16 +876,15 @@ pub async fn run_job(spec: &SshJobSpec) -> Result { }; // Create the dir (owner-only) and write run.sh from stdin. - let setup = format!( - "umask 077; mkdir -p \"$HOME/{dir}\" && chmod 700 \"$HOME/{dir}\" && cat > \"$HOME/{dir}/run.sh\"", - ); + let d = shell_path(&dir); + let setup = format!("umask 077; mkdir -p {d} && chmod 700 {d} && cat > {d}/run.sh"); ssh_run(&spec.target, &setup, Some(&run_sh)).await?; // The container host wrapper records its own PID; direct runs retain the existing launcher. let launch = if spec.container.is_some() { - format!("cd \"$HOME/{dir}\" && date +%s > launch_time.tmp && mv launch_time.tmp launch_time && {{ if command -v setsid >/dev/null 2>&1; then setsid bash run.sh /dev/null 2>&1 & else nohup bash run.sh /dev/null 2>&1 & fi; }}") + format!("cd {d} && date +%s > launch_time.tmp && mv launch_time.tmp launch_time && {{ if command -v setsid >/dev/null 2>&1; then setsid bash run.sh /dev/null 2>&1 & else nohup bash run.sh /dev/null 2>&1 & fi; }}") } else { - format!("cd \"$HOME/{dir}\" && if command -v setsid >/dev/null 2>&1; then setsid bash run.sh /dev/null 2>&1 & else set -m; trap '' HUP; bash run.sh /dev/null 2>&1 & fi; echo $! > pid") + format!("cd {d} && if command -v setsid >/dev/null 2>&1; then setsid bash run.sh /dev/null 2>&1 & else set -m; trap '' HUP; bash run.sh /dev/null 2>&1 & fi; echo $! > pid") }; ssh_run(&spec.target, &launch, None) .await @@ -851,11 +920,12 @@ async fn inspect_host_job(target: &SshTarget, dir: &str) -> Result { fn host_job_script(dir: &str) -> String { format!( "{HOST_PROCESS_HELPERS}\n\ - d=\"$HOME/{dir}\"; \ + d={}; \ if [ ! -d \"$d\" ]; then echo MISSING; \ elif [ -f \"$d/exit_code\" ]; then echo \"EXIT $(cat \"$d/exit_code\")\"; \ elif [ -f \"$d/pid\" ] && host_process_alive \"$(cat \"$d/pid\")\"; then echo RUNNING; \ elif [ -f \"$d/pid\" ]; then echo DEAD; else echo PENDING; fi", + shell_path(dir), ) } @@ -953,9 +1023,9 @@ pub async fn stream_logs( sink: &mut (dyn FnMut(&str) + Send), ) -> Result { let cmd = format!( - "printf '__ORX_LOG_START__\\n'; tail -n +{} \"$HOME/{}/log\" 2>/dev/null || true", + "printf '__ORX_LOG_START__\\n'; tail -n +{} {}/log 2>/dev/null || true", skip + 1, - dir + shell_path(dir) ); let out = ssh_run(target, &cmd, None).await?; let out = out @@ -995,7 +1065,7 @@ pub async fn cancel_job( fn host_cancel_script(dir: &str) -> String { format!( "{HOST_PROCESS_HELPERS}\n\ - p=$(cat \"$HOME/{dir}/pid\" 2>/dev/null); \ + p=$(cat {}/pid 2>/dev/null); \ [ -n \"$p\" ] || exit 0; \ case \"$p\" in ''|*[!0-9]*) exit 0;; esac; \ kill -TERM -- -\"$p\" 2>/dev/null || kill -TERM \"$p\" 2>/dev/null || true; \ @@ -1003,6 +1073,7 @@ fn host_cancel_script(dir: &str) -> String { kill -KILL -- -\"$p\" 2>/dev/null || true; \ for ((i=0; i<10; i++)); do host_group_alive \"$p\" || exit 0; sleep 0.1; done; \ echo 'Experiment process group is still alive' >&2; exit 1", + shell_path(dir), ) } @@ -1132,6 +1203,73 @@ mod tests { ); } + #[test] + fn remote_roots_are_absolute_plain_paths() { + for root in ["/scratch/me/orx", "/data/user.name/orx-1", "/gpfs/a+b@c"] { + validate_remote_root(root).unwrap(); + } + for root in [ + "", + "/", + "scratch/orx", + "~/orx", + "/scratch/orx/", + "/scratch//orx", + "/scratch/../etc", + "/scratch/$USER", + "/scratch/a b", + "/scratch/`id`", + "/scratch/x;rm", + "/scratch/'q'", + ] { + assert!(validate_remote_root(root).is_err(), "{root}"); + } + assert_eq!(remote_run_dir(None, "r1"), ".orx/runs/r1"); + assert_eq!( + remote_run_dir(Some("/scratch/orx"), "r1"), + "/scratch/orx/runs/r1" + ); + assert_eq!(shell_path(".orx/runs/r1"), "\"$HOME/.orx/runs/r1\""); + assert_eq!(shell_path("/scratch/orx/runs/r1"), "'/scratch/orx/runs/r1'"); + assert_eq!( + remote_source_cache(Some("/scratch/orx"), "d"), + "/scratch/orx/source/d.tar" + ); + } + + #[cfg(unix)] + #[test] + fn configured_root_holds_the_source_cache_and_run_state() { + use std::io::Write; + let home = crate::local::git::TemporaryDirectory::new("orx-ssh-root-home").unwrap(); + let scratch = crate::local::git::TemporaryDirectory::new("orx-ssh-root").unwrap(); + let root = scratch.path().join("orx"); + let root = root.to_str().unwrap(); + let (_, upload) = source_cache_commands(&remote_source_cache(Some(root), "d"), 3); + let mut child = std::process::Command::new("sh") + .args(["-c", &upload]) + .env("HOME", home.path()) + .stdin(Stdio::piped()) + .spawn() + .unwrap(); + child.stdin.take().unwrap().write_all(b"abc").unwrap(); + assert!(child.wait().unwrap().success()); + assert_eq!( + std::fs::read(scratch.path().join("orx/source/d.tar")).unwrap(), + b"abc" + ); + let dir = remote_run_dir(Some(root), "r1"); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write(format!("{dir}/exit_code"), "3\n").unwrap(); + let out = std::process::Command::new("bash") + .args(["-c", &host_job_script(&dir)]) + .env("HOME", home.path()) + .output() + .unwrap(); + assert_eq!(String::from_utf8(out.stdout).unwrap().trim(), "EXIT 3"); + assert!(!home.path().join(".orx").exists()); + } + #[tokio::test] async fn upload_source_retries_a_dropped_attempt_then_succeeds() { use std::sync::atomic::{AtomicUsize, Ordering}; diff --git a/src/jobs/ssh/container.rs b/src/jobs/ssh/container.rs index bdb935bbb..2749914c2 100644 --- a/src/jobs/ssh/container.rs +++ b/src/jobs/ssh/container.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; -use super::{sh_quote, ssh_run, JobState, SshTarget}; +use super::{sh_quote, shell_path, ssh_run, JobState, SshTarget}; use crate::error::{anyhow, Result}; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -215,7 +215,8 @@ pub(super) fn host_script(dir: &str, container: &ContainerRun) -> String { let check = sh_quote("{{.State.StartedAt}} {{.State.Status}}"); let expected = sh_quote(&format!("{} running", container.started_at)); let id = sh_quote(&container.id); - format!("#!/usr/bin/env bash\ncd \"$HOME/{dir}\" || exit 97\nprintf '%s\\n' \"$$\" > pid\n(\nactual=$(docker container inspect --format {check} -- {id}) || exit $?\n[ \"$actual\" = {expected} ] || {{ echo 'Container generation or readiness changed before launch' >&2; exit 97; }}\ndocker exec {id} setsid --wait bash {inner} > log 2>&1\nprintf '%s\\n' \"$?\" > exit_code\n") + let dir = shell_path(dir); + format!("#!/usr/bin/env bash\ncd {dir} || exit 97\nprintf '%s\\n' \"$$\" > pid\n(\nactual=$(docker container inspect --format {check} -- {id}) || exit $?\n[ \"$actual\" = {expected} ] || {{ echo 'Container generation or readiness changed before launch' >&2; exit 97; }}\ndocker exec {id} setsid --wait bash {inner} > log 2>&1\nprintf '%s\\n' \"$?\" > exit_code\n") } fn failed(message: String) -> JobState { @@ -226,7 +227,8 @@ fn failed(message: String) -> JobState { } async fn host_state(target: &SshTarget, dir: &str) -> Result { - ssh_run(target, &format!(r#"d="$HOME/{dir}"; if [ -f "$d/exit_code" ]; then echo "EXIT $(cat "$d/exit_code")"; elif [ -f "$d/pid" ] && kill -0 "$(cat "$d/pid")" 2>/dev/null; then echo RUNNING; else echo DEAD; fi; if [ -f "$d/launch_time" ]; then echo "$(( $(date +%s) - $(cat "$d/launch_time") ))"; else echo 0; fi"#), None).await + let dir = shell_path(dir); + ssh_run(target, &format!(r#"d={dir}; if [ -f "$d/exit_code" ]; then echo "EXIT $(cat "$d/exit_code")"; elif [ -f "$d/pid" ] && kill -0 "$(cat "$d/pid")" 2>/dev/null; then echo RUNNING; else echo DEAD; fi; if [ -f "$d/launch_time" ]; then echo "$(( $(date +%s) - $(cat "$d/launch_time") ))"; else echo 0; fi"#), None).await } pub(super) async fn inspect( @@ -297,7 +299,8 @@ else echo DEAD; fi if final_state.stage != "RUNNING" { return Ok(final_state); } - let waited = ssh_run(target, &format!(r#"d="$HOME/{dir}"; umask 077; if [ ! -f "$d/completion_wait" ]; then date +%s > "$d/completion_wait.tmp" && mv "$d/completion_wait.tmp" "$d/completion_wait"; fi; echo "$(( $(date +%s) - $(cat "$d/completion_wait") ))""#), None).await?; + let d = shell_path(dir); + let waited = ssh_run(target, &format!(r#"d={d}; umask 077; if [ ! -f "$d/completion_wait" ]; then date +%s > "$d/completion_wait.tmp" && mv "$d/completion_wait.tmp" "$d/completion_wait"; fi; echo "$(( $(date +%s) - $(cat "$d/completion_wait") ))""#), None).await?; if waited.trim().parse::()? >= 30 { let final_state = super::inspect_host_job(target, dir).await?; if final_state.stage != "RUNNING" { diff --git a/src/jobs/ssh/container/tests.rs b/src/jobs/ssh/container/tests.rs index 6b09221dc..77b039e63 100644 --- a/src/jobs/ssh/container/tests.rs +++ b/src/jobs/ssh/container/tests.rs @@ -30,12 +30,13 @@ async fn prepare( .status() .unwrap() .success()); - let dir = ssh::stage_source(target, &id, &archive, &id, container.as_ref()) + let dir = ssh::stage_source(target, &id, &archive, &id, container.as_ref(), None) .await .unwrap(); let spec = SshJobSpec { target: target.clone(), run_id: id, + remote_root: None, script: crate::compute::staged_script("bash run.sh"), env: HashMap::from([( "ORX_TEST_SECRET".into(), diff --git a/src/local/hf.rs b/src/local/hf.rs index 31b4f0fe2..7c4797541 100644 --- a/src/local/hf.rs +++ b/src/local/hf.rs @@ -128,6 +128,7 @@ pub async fn submit_local_hf_with_source( source_digest: None, source_path: None, source_size: None, + remote_root: None, }; source.apply_to_descriptor(&mut descriptor); if let Err(error) = crate::compute::record_submission_handle(&run_id, &descriptor) { diff --git a/src/local/k8s.rs b/src/local/k8s.rs index 1ba180cc6..1371cbfc4 100644 --- a/src/local/k8s.rs +++ b/src/local/k8s.rs @@ -172,6 +172,7 @@ pub async fn submit_local_k8s_with_source( source_digest: None, source_path: None, source_size: None, + remote_root: None, }; source.apply_to_descriptor(&mut descriptor); if let Err(error) = crate::compute::record_submission_handle(&run_id, &descriptor) { diff --git a/src/local/localrun.rs b/src/local/localrun.rs index 079a6cf35..47c14e7d1 100644 --- a/src/local/localrun.rs +++ b/src/local/localrun.rs @@ -148,6 +148,7 @@ async fn submit_controller_run( source_digest: None, source_path: None, source_size: None, + remote_root: None, }; source.apply_to_descriptor(&mut descriptor); if let Err(error) = crate::compute::record_submission_handle(&run_id, &descriptor) { diff --git a/src/local/modal.rs b/src/local/modal.rs index fdeee054a..8cf590520 100644 --- a/src/local/modal.rs +++ b/src/local/modal.rs @@ -132,6 +132,7 @@ pub async fn submit_local_modal_with_source( source_digest: None, source_path: None, source_size: None, + remote_root: None, }; source.apply_to_descriptor(&mut descriptor); if let Err(error) = crate::compute::record_submission_handle(&run_id, &descriptor) { diff --git a/src/local/openresearch.rs b/src/local/openresearch.rs index 62b7e6ab3..9691f3436 100644 --- a/src/local/openresearch.rs +++ b/src/local/openresearch.rs @@ -203,6 +203,7 @@ pub async fn submit_local_openresearch_with_source( source_digest: None, source_path: None, source_size: None, + remote_root: None, }; source.apply_to_descriptor(&mut descriptor); if let Err(error) = crate::compute::record_submission_handle(&run_id, &descriptor) { diff --git a/src/local/ray.rs b/src/local/ray.rs index 2714f0695..49e7134e1 100644 --- a/src/local/ray.rs +++ b/src/local/ray.rs @@ -145,6 +145,7 @@ pub async fn submit_local_ray_with_source( source_digest: None, source_path: None, source_size: None, + remote_root: None, }; source.apply_to_descriptor(&mut descriptor); if let Err(error) = crate::compute::record_submission_handle(&run_id, &descriptor) { diff --git a/src/local/slurm.rs b/src/local/slurm.rs index 9815ab27c..72de6dedc 100644 --- a/src/local/slurm.rs +++ b/src/local/slurm.rs @@ -141,11 +141,13 @@ pub async fn submit_local_slurm_with_source( &source.path, &source.digest, None, + settings.remote_root.as_deref(), ) .await?; let job_id = slurm::run_job(&slurm::SlurmJobSpec { host: host.clone(), run_id: run_id.clone(), + remote_root: settings.remote_root.clone(), setup_script: "test -d repo".to_string(), command: run_command.clone(), env, @@ -181,6 +183,7 @@ pub async fn submit_local_slurm_with_source( source_digest: None, source_path: None, source_size: None, + remote_root: settings.remote_root.clone(), }; source.apply_to_descriptor(&mut descriptor); if let Err(error) = crate::compute::record_submission_handle(&run_id, &descriptor) { diff --git a/src/local/ssh.rs b/src/local/ssh.rs index 5392c75bb..7ee127c70 100644 --- a/src/local/ssh.rs +++ b/src/local/ssh.rs @@ -67,6 +67,7 @@ pub async fn submit_local_ssh_with_source( &source.path, &source.digest, container.as_ref(), + launch.remote_root.as_deref(), ) .await?; let script = format!( @@ -105,6 +106,7 @@ pub async fn submit_local_ssh_with_source( source_digest: None, source_path: None, source_size: None, + remote_root: launch.remote_root.clone(), }; source.apply_to_descriptor(&mut descriptor); if container.is_some() { @@ -113,6 +115,7 @@ pub async fn submit_local_ssh_with_source( let submission = ssh::run_job(&ssh::SshJobSpec { target: target.clone(), run_id: run_id.clone(), + remote_root: launch.remote_root.clone(), script, env, container: container.clone(), diff --git a/src/telemetry.rs b/src/telemetry.rs index bdbb87b8f..087e8cd06 100644 --- a/src/telemetry.rs +++ b/src/telemetry.rs @@ -1629,6 +1629,7 @@ mod tests { set_persisted_disabled(true).unwrap(); let options = SshHostSettings { container: Some("research".into()), + remote_root: Some("/scratch/lab/orx".into()), }; set_ssh_host("lab".into(), options.clone()).unwrap(); set_ssh_default(Some("lab".into())).unwrap(); From 020e42fcdb45d49127f854b7355f78a9790784a8 Mon Sep 17 00:00:00 2001 From: Myles Anderson Date: Thu, 22 Oct 2026 14:44:01 -0700 Subject: [PATCH 02/10] Address review nits for remote storage root Co-Authored-By: Claude Opus 5.5 --- agent-skills/orx-compute/references/configuration.md | 1 + agent-skills/orx-compute/references/ssh.md | 4 ++-- src/commands/compute.rs | 2 +- src/commands/up/compute_settings.rs | 5 ++--- src/jobs/mod.rs | 4 ++-- src/jobs/slurm.rs | 10 +++++----- src/jobs/ssh.rs | 4 ++-- src/local/slurm.rs | 3 +++ 8 files changed, 18 insertions(+), 15 deletions(-) diff --git a/agent-skills/orx-compute/references/configuration.md b/agent-skills/orx-compute/references/configuration.md index fb79a26e4..54563ad3e 100644 --- a/agent-skills/orx-compute/references/configuration.md +++ b/agent-skills/orx-compute/references/configuration.md @@ -11,6 +11,7 @@ orx compute configure ssh --default-host lab orx compute configure ssh --host lab --container research orx compute configure ssh --host lab --clear container orx compute configure ssh --host lab --remote-root /data/me/orx # runs + source cache; default ~/.orx +orx compute configure ssh --host lab --clear remote-root orx compute test ssh --host lab # tests saved container too orx compute test ssh --host lab --no-container # explicitly tests host orx compute connect ssh --host lab # interactive login/MFA diff --git a/agent-skills/orx-compute/references/ssh.md b/agent-skills/orx-compute/references/ssh.md index ebe056391..063a4241a 100644 --- a/agent-skills/orx-compute/references/ssh.md +++ b/agent-skills/orx-compute/references/ssh.md @@ -57,8 +57,8 @@ creation or lifecycle management. `--image` and `--flavor` are unsupported. The committed source archive is cached on the host, then streamed into the container's `$HOME/.orx/runs//repo`. Host logs and exit status stay in the -host's `~/.orx/runs//` (or the saved remote root), so container removal does not erase them. A detached -supervisor reattaches using the saved container ID and start time; later settings +host's `/runs//` (`~/.orx` unless a remote root is saved), so +container removal does not erase them. A detached supervisor reattaches using the saved container ID and start time; later settings changes cannot redirect an existing run. Do not kill the supervisor. Cancellation targets only the experiment process group, waits five seconds for diff --git a/src/commands/compute.rs b/src/commands/compute.rs index 7014286a3..79d567909 100644 --- a/src/commands/compute.rs +++ b/src/commands/compute.rs @@ -429,7 +429,7 @@ impl ConfigureArgs { } let host_setting = body.contains_key("container") || body.contains_key("remoteRoot"); if backend == "ssh" && body.contains_key("host") && !host_setting { - return Err(anyhow!("--host selects a host setting; supply --container, --remote-root, or --clear for one of them. Use --default-host to select a default.")); + return Err(anyhow!("--host needs --container, --remote-root, or a matching --clear; use --default-host to pick a default.")); } if backend == "ssh" && host_setting && !body.contains_key("host") { return Err(anyhow!( diff --git a/src/commands/up/compute_settings.rs b/src/commands/up/compute_settings.rs index 7d8659b53..d7a3ccb5b 100644 --- a/src/commands/up/compute_settings.rs +++ b/src/commands/up/compute_settings.rs @@ -623,8 +623,7 @@ pub(super) async fn ssh_settings() -> ApiResult { pub(super) struct SaveSshSettingsReq { host: String, container: Option, - /// `None` keeps the saved root; `Some("")` restores the `~/.orx` default. - #[serde(default)] + /// Unlike `container`, `None` keeps the saved root (the UI never sends it); `Some("")` clears it. remote_root: Option, } @@ -1473,7 +1472,7 @@ pub(crate) async fn configure(backend: &str, body: Value) -> Result { let container = match body.get("container") { Some(container) => container.clone(), None => { - let host = body["host"].as_str().unwrap_or_default(); + let host = body["host"].as_str().unwrap_or_default().trim(); json!(crate::config::ssh_settings()? .hosts .get(host) diff --git a/src/jobs/mod.rs b/src/jobs/mod.rs index 7996e866b..d9c8c7940 100644 --- a/src/jobs/mod.rs +++ b/src/jobs/mod.rs @@ -106,7 +106,7 @@ pub struct BackendDescriptor { pub source_path: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub source_size: Option, - /// The host's storage root at submit (ssh/slurm); `None` is `~/.orx`. + /// The host's storage root at submit; Slurm derives its run dir from it, SSH keeps it for source-cache cleanup. #[serde(default, skip_serializing_if = "Option::is_none")] pub remote_root: Option, } @@ -157,7 +157,7 @@ impl BackendDescriptor { } /// The Slurm (host, job id) handle; the login-node host rides on - /// `namespace`, and the run dir derives from the run id. + /// `namespace`, and the run dir derives from the run id and `remote_root`. pub fn slurm_ref(&self) -> Result<(&str, &str)> { if self.kind != "slurm_job" { return Err(anyhow!("Unsupported backend kind: {}", self.kind)); diff --git a/src/jobs/slurm.rs b/src/jobs/slurm.rs index 0ab9a0be5..bc7df5731 100644 --- a/src/jobs/slurm.rs +++ b/src/jobs/slurm.rs @@ -8,7 +8,7 @@ //! the login node — the same trade SkyPilot makes. //! //! The remote layout extends the ssh backend's convention. A run owns -//! `~/.orx/runs//` on the cluster's shared filesystem: +//! `/runs//` (root defaults to `~/.orx`) on the cluster's shared filesystem: //! repo/ the experiment branch, cloned on the login node at submit //! time (compute nodes often have no internet, so cloning //! inside the job — the other backends' pattern — would break) @@ -17,7 +17,7 @@ //! exit_code written by job.sbatch's closing lines when the payload ends //! //! The reattach handle is (host, slurm job id); the run dir derives from the -//! run id. Job state is read exit_code-first (scheduler-independent truth, +//! run id and the root recorded at submit. Job state is read exit_code-first (scheduler-independent truth, //! like the ssh backend), then `squeue` for live jobs, then `sacct` for jobs //! that left the queue without writing an exit code (scancel/timeout/node //! failure). `sacct` may be disabled cluster-wide, so it's a best-effort @@ -687,9 +687,9 @@ mod tests { let run_a = format!("e2e-a-{}", uuid::Uuid::new_v4()); let run_b = format!("e2e-b-{}", uuid::Uuid::new_v4()); let cleanup = format!( - "rm -rf \"$HOME/{}\" \"$HOME/{}\"", - run_dir(None, &run_a), - run_dir(None, &run_b) + "rm -rf {} {}", + shell_path(&run_dir(None, &run_a)), + shell_path(&run_dir(None, &run_b)) ); // Happy path: runs, completes, logs arrive. diff --git a/src/jobs/ssh.rs b/src/jobs/ssh.rs index 2beef8f30..292749210 100644 --- a/src/jobs/ssh.rs +++ b/src/jobs/ssh.rs @@ -7,7 +7,7 @@ //! polls reuse one TCP session instead of a handshake apiece. //! Win32-OpenSSH cannot, so on Windows background calls need an agent key or no passphrase. //! -//! The handle is a remote run directory `~/.orx/runs//` holding: +//! The handle is a remote run directory `/runs//` (root defaults to `~/.orx`) holding: //! run.sh the launcher (exported env + snapshot-and-run payload) //! log merged stdout/stderr //! pid the detached process-group leader @@ -110,7 +110,7 @@ pub fn validate_remote_root(root: &str) -> Result<()> { }); if !valid { return Err(anyhow!( - "Remote storage root must be an absolute path such as /scratch/$USER/orx, using letters, numbers, '.', '_', '-', '+' or '@' (no trailing slash)." + "Remote storage root must be an absolute path like /scratch//orx (letters, digits, . _ - + @; no trailing slash)." )); } Ok(()) diff --git a/src/local/slurm.rs b/src/local/slurm.rs index 72de6dedc..9cc43c517 100644 --- a/src/local/slurm.rs +++ b/src/local/slurm.rs @@ -88,6 +88,9 @@ pub async fn submit_local_slurm_with_source( } let settings = slurm::load_settings()?.unwrap_or_default(); + if let Some(root) = &settings.remote_root { + crate::jobs::ssh::validate_remote_root(root)?; + } let host = args .host .clone() From f98b36090d64bd9ee5f6998aedce006cb968be47 Mon Sep 17 00:00:00 2001 From: Myles Anderson Date: Thu, 22 Oct 2026 14:45:58 -0700 Subject: [PATCH 03/10] Use literal paths in remote-root doc examples Co-Authored-By: Claude Opus 5.5 --- agent-skills/orx-compute/references/slurm.md | 2 +- agent-skills/orx-compute/references/ssh.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/agent-skills/orx-compute/references/slurm.md b/agent-skills/orx-compute/references/slurm.md index e891143a7..20cd93384 100644 --- a/agent-skills/orx-compute/references/slurm.md +++ b/agent-skills/orx-compute/references/slurm.md @@ -23,7 +23,7 @@ orx exp run --backend slurm - `--timeout` covers the batch job. Unset falls back to the Slurm time-limit setting, then the cluster's limit. - Runs and the cached source archive live under `~/.orx` on the cluster unless - `orx compute configure slurm --remote-root /scratch//orx` names another + `orx compute configure slurm --remote-root /scratch/alice/orx` names another absolute directory, such as scratch space when home quota is small. It must be on a filesystem the compute nodes share. Each run keeps the root it was submitted with; `--clear remote-root` restores the default for new runs. diff --git a/agent-skills/orx-compute/references/ssh.md b/agent-skills/orx-compute/references/ssh.md index 063a4241a..03230a2c7 100644 --- a/agent-skills/orx-compute/references/ssh.md +++ b/agent-skills/orx-compute/references/ssh.md @@ -33,7 +33,7 @@ orx exp run --backend ssh --host lab --no-container ``` Machine-local SSH settings can save a container name or ID for each host, and -an absolute remote storage root (`--remote-root /data//orx`) used instead +an absolute remote storage root (`--remote-root /data/alice/orx`) used instead of `~/.orx` for that host's runs and cached source archives. Each run keeps the root it was launched with. A saved default SSH host lets launches omit `--host`. From 367dd9cc4a8549f53bc7847efe30e5d5de4f4217 Mon Sep 17 00:00:00 2001 From: Myles Anderson Date: Thu, 22 Oct 2026 14:57:32 -0700 Subject: [PATCH 04/10] Remove finished remote run dirs and unused source tars (OR-575) Co-Authored-By: Claude Opus 5.5 --- .../orx-compute/references/configuration.md | 2 + agent-skills/orx-compute/references/slurm.md | 6 + agent-skills/orx-compute/references/ssh.md | 8 + src/commands/compute.rs | 25 +- src/commands/supervise.rs | 22 +- src/commands/up.rs | 62 ++++ src/commands/up/compute_settings.rs | 40 ++- src/compute.rs | 1 + src/config.rs | 3 + src/jobs/slurm.rs | 3 + src/jobs/ssh.rs | 19 +- src/jobs/ssh/cleanup.rs | 283 ++++++++++++++++++ src/local/mod.rs | 1 + src/local/remote_cleanup.rs | 105 +++++++ src/store.rs | 153 ++++++++++ src/telemetry.rs | 1 + 16 files changed, 706 insertions(+), 28 deletions(-) create mode 100644 src/jobs/ssh/cleanup.rs create mode 100644 src/local/remote_cleanup.rs diff --git a/agent-skills/orx-compute/references/configuration.md b/agent-skills/orx-compute/references/configuration.md index 54563ad3e..e340f8bfe 100644 --- a/agent-skills/orx-compute/references/configuration.md +++ b/agent-skills/orx-compute/references/configuration.md @@ -12,6 +12,7 @@ orx compute configure ssh --host lab --container research orx compute configure ssh --host lab --clear container orx compute configure ssh --host lab --remote-root /data/me/orx # runs + source cache; default ~/.orx orx compute configure ssh --host lab --clear remote-root +orx compute configure ssh --host lab --keep-remote-runs true # keep finished run dirs on the host orx compute test ssh --host lab # tests saved container too orx compute test ssh --host lab --no-container # explicitly tests host orx compute connect ssh --host lab # interactive login/MFA @@ -20,6 +21,7 @@ orx compute configure slurm --cpus-per-task 8 --mem 64G # defaults; exp run -- orx compute configure slurm --clear cpus-per-task --clear mem orx compute configure slurm --clear time-limit # cluster chooses orx compute configure slurm --remote-root /scratch/me/orx # shared FS; --clear remote-root for ~/.orx +orx compute configure slurm --keep-remote-runs true # keep finished run dirs on the cluster orx compute connect slurm --host cluster orx compute configure k8s --context research --namespace experiments orx compute configure ray --address http://127.0.0.1:8265 diff --git a/agent-skills/orx-compute/references/slurm.md b/agent-skills/orx-compute/references/slurm.md index 20cd93384..bdf4c2c8d 100644 --- a/agent-skills/orx-compute/references/slurm.md +++ b/agent-skills/orx-compute/references/slurm.md @@ -27,5 +27,11 @@ orx exp run --backend slurm absolute directory, such as scratch space when home quota is small. It must be on a filesystem the compute nodes share. Each run keeps the root it was submitted with; `--clear remote-root` restores the default for new runs. +- Once a run finishes and its log is collected, the running OpenResearch app + removes its remote run dir, along with cached source archives that no active + run uses. Anything the job wrote into `repo/` goes with it. To keep results + on the cluster, run `orx compute configure slurm --keep-remote-runs true`, or + write outputs outside the run dir. Deleting a project removes its runs' + remote dirs regardless. - A detached `orx supervise` process records scheduler status and logs; do not kill it. diff --git a/agent-skills/orx-compute/references/ssh.md b/agent-skills/orx-compute/references/ssh.md index 03230a2c7..57ab52514 100644 --- a/agent-skills/orx-compute/references/ssh.md +++ b/agent-skills/orx-compute/references/ssh.md @@ -36,6 +36,14 @@ Machine-local SSH settings can save a container name or ID for each host, and an absolute remote storage root (`--remote-root /data/alice/orx`) used instead of `~/.orx` for that host's runs and cached source archives. Each run keeps the root it was launched with. + +Once a run finishes and its log is collected, the running OpenResearch app +removes the run's remote dir (including anything written into `repo/`) and +cached source archives that no active run uses. Unreachable hosts are retried +later. `orx compute configure ssh --host lab --keep-remote-runs true` keeps that +host's finished runs. Deleting a project removes its runs' remote dirs +regardless. + A saved default SSH host lets launches omit `--host`. These defaults apply only to SSH experiments; dashboards, terminals and coding agents still connect directly to the host. diff --git a/src/commands/compute.rs b/src/commands/compute.rs index 79d567909..3d7e94fcd 100644 --- a/src/commands/compute.rs +++ b/src/commands/compute.rs @@ -238,6 +238,9 @@ pub struct ConfigureArgs { /// Absolute remote dir for runs and the source cache (SSH per host, or Slurm); default ~/.orx. #[arg(long)] remote_root: Option, + /// Keep finished runs' remote dirs (true) instead of removing them once logs are collected. + #[arg(long)] + keep_remote_runs: Option, #[arg(long)] default_host: Option, #[arg(long)] @@ -324,7 +327,13 @@ impl ConfigureArgs { fn body(self) -> Result { let backend = self.backend.name(); let allowed: &[&str] = match backend { - "ssh" => &["host", "container", "remoteRoot", "defaultHost"], + "ssh" => &[ + "host", + "container", + "remoteRoot", + "keepRemoteRuns", + "defaultHost", + ], "slurm" => &[ "host", "partition", @@ -333,6 +342,7 @@ impl ConfigureArgs { "cpusPerTask", "mem", "remoteRoot", + "keepRemoteRuns", ], "k8s" => &["context", "namespace"], "ray" => &["address"], @@ -359,6 +369,9 @@ impl ConfigureArgs { body.insert(key.into(), json!(value)); } } + if let Some(keep) = self.keep_remote_runs { + body.insert("keepRemoteRuns".into(), json!(keep)); + } if let Some(cpus) = self.cpus_per_task { body.insert("cpusPerTask".into(), json!(cpus)); } @@ -427,9 +440,11 @@ impl ConfigureArgs { "Supply settings to change; use compute show {backend} to inspect." )); } - let host_setting = body.contains_key("container") || body.contains_key("remoteRoot"); + let host_setting = ["container", "remoteRoot", "keepRemoteRuns"] + .iter() + .any(|key| body.contains_key(*key)); if backend == "ssh" && body.contains_key("host") && !host_setting { - return Err(anyhow!("--host needs --container, --remote-root, or a matching --clear; use --default-host to pick a default.")); + return Err(anyhow!("--host needs --container, --remote-root, --keep-remote-runs, or a matching --clear; use --default-host to pick a default.")); } if backend == "ssh" && host_setting && !body.contains_key("host") { return Err(anyhow!( @@ -783,6 +798,10 @@ mod tests { body(&["slurm", "--clear", "remote-root"]).unwrap(), json!({"remoteRoot":""}) ); + assert_eq!( + body(&["ssh", "--host", "lab", "--keep-remote-runs", "true"]).unwrap(), + json!({"host":"lab", "keepRemoteRuns":true}) + ); assert!(body(&["ssh", "--remote-root", "/scratch/me/orx"]).is_err()); assert!(crate::Cli::try_parse_from(["orx", "compute", "status", "--json"]).is_ok()); } diff --git a/src/commands/supervise.rs b/src/commands/supervise.rs index 4be9e58ea..376f9dda8 100644 --- a/src/commands/supervise.rs +++ b/src/commands/supervise.rs @@ -849,6 +849,9 @@ async fn watch_ssh_job( ) .await; } + if descriptor.kind == "ssh_job" && collected { + queue_remote_cleanup(store, run_id); + } eprintln!("supervise {run_id}: finished ({status})"); return Ok(status); } @@ -1436,10 +1439,11 @@ async fn run_slurm( if is_terminal_stage(stage) { let _ = done_tx.send(true); - if tokio::time::timeout(Duration::from_secs(20), &mut log_task) - .await - .is_err() - { + let collected = matches!( + tokio::time::timeout(Duration::from_secs(20), &mut log_task).await, + Ok(Ok(true)) + ); + if !collected { log_task.abort(); } let applied = store.update_status(&run_id, status, Some(now_ms()), None)?; @@ -1452,6 +1456,9 @@ async fn run_slurm( } } } + if collected { + queue_remote_cleanup(&store, &run_id); + } eprintln!("supervise {run_id}: finished ({status})"); return Ok(()); } @@ -1465,6 +1472,13 @@ async fn run_slurm( } } +/// Once logs are local, `orx up` removes the remote run dir unless the target keeps results. +fn queue_remote_cleanup(store: &Store, run_id: &str) { + if let Err(err) = store.queue_finished_remote_cleanup(run_id) { + eprintln!("supervise {run_id}: could not queue remote cleanup: {err}"); + } +} + async fn cancel_slurm(host: &str, job_id: &str, run_id: &str, cancel_sent: &mut bool) { eprintln!("supervise {run_id}: cancel requested — scancel {job_id}"); match slurm::cancel_job(host, job_id).await { diff --git a/src/commands/up.rs b/src/commands/up.rs index 64689a70b..2559236b0 100644 --- a/src/commands/up.rs +++ b/src/commands/up.rs @@ -178,6 +178,10 @@ pub async fn run(args: UpArgs) -> Result<()> { state.data_dir_move_in_progress.clone(), state.data_dir_gate.clone(), )); + tokio::spawn(reclaim_remote_run_dirs( + state.data_dir_move_in_progress.clone(), + state.data_dir_gate.clone(), + )); spawn_claude_auth_monitor( state.chat.clone(), claude.clone(), @@ -4695,6 +4699,64 @@ async fn reclaim_orphaned_run_artifacts( } } +/// Runs `f` on a fresh store under the data-dir gate; `None` while the data dir is moving. +async fn with_gated_store( + moving: &AtomicBool, + gate: &tokio::sync::Mutex<()>, + f: impl FnOnce(&Store) -> Result + Send + 'static, +) -> Option> { + if moving.load(Ordering::SeqCst) { + return None; + } + let _gate = gate.lock().await; + if moving.load(Ordering::SeqCst) { + return None; + } + Some( + tokio::task::spawn_blocking(move || f(&Store::open()?)) + .await + .unwrap_or_else(|error| Err(anyhow!("store task failed: {error}"))), + ) +} + +/// Removes queued remote run dirs; the gate is held only around store access, never over SSH. +async fn reclaim_remote_run_dirs(moving: Arc, gate: Arc>) { + tokio::time::sleep(Duration::from_secs(30)).await; + let mut reported = HashSet::new(); + loop { + let mut report = |failures: Vec<(String, String)>| { + for (run_id, error) in failures { + if reported.insert(run_id.clone()) { + eprintln!("orx up: remote cleanup for run {run_id}: {error}"); + } + } + }; + match with_gated_store(&moving, &gate, local::remote_cleanup::plan).await { + Some(Ok((batches, rejected))) => { + report(rejected); + if !batches.is_empty() { + let results = local::remote_cleanup::clean(batches).await; + match with_gated_store(&moving, &gate, move |store| { + local::remote_cleanup::record(store, results) + }) + .await + { + Some(Ok(failures)) => report(failures), + Some(Err(error)) => { + eprintln!("orx up: could not record remote cleanup: {error}") + } + // Unrecorded entries stay due and removal is idempotent. + None => {} + } + } + } + Some(Err(error)) => eprintln!("orx up: could not plan remote cleanup: {error}"), + None => {} + } + tokio::time::sleep(ORPHANED_RUN_SCAN_INTERVAL).await; + } +} + /// Spawns the missing supervisors and returns the runs whose spawn failed. fn adopt_orphaned_runs_once() -> Result> { let store = Store::open()?; diff --git a/src/commands/up/compute_settings.rs b/src/commands/up/compute_settings.rs index d7a3ccb5b..2fe9be033 100644 --- a/src/commands/up/compute_settings.rs +++ b/src/commands/up/compute_settings.rs @@ -588,6 +588,7 @@ pub(super) async fn ssh_settings() -> ApiResult { .unwrap_or_default(); host["container"] = json!(options.container); host["remoteRoot"] = json!(options.remote_root); + host["keepRemoteRuns"] = json!(options.keep_remote_runs); } // Best-effort, like the preflight write: a store hiccup shouldn't take // out the host listing — hosts just render as never tested. @@ -623,24 +624,27 @@ pub(super) async fn ssh_settings() -> ApiResult { pub(super) struct SaveSshSettingsReq { host: String, container: Option, - /// Unlike `container`, `None` keeps the saved root (the UI never sends it); `Some("")` clears it. + /// Unlike `container`, `None` keeps the saved value (the UI never sends these); `Some("")` clears the root. remote_root: Option, + keep_remote_runs: Option, } pub(super) async fn save_ssh_settings(Json(req): Json) -> ApiResult { let host = req.host.trim().to_string(); require_configured_ssh_host(&host).await?; blocking_api(move || { - let remote_root = match req.remote_root { - Some(root) => Some(root.trim().to_string()).filter(|root| !root.is_empty()), - None => crate::config::ssh_settings()? - .hosts - .get(&host) - .and_then(|saved| saved.remote_root.clone()), - }; + let saved = crate::config::ssh_settings()? + .hosts + .get(&host) + .cloned() + .unwrap_or_default(); let options = crate::config::SshHostSettings { container: req.container, - remote_root, + remote_root: match req.remote_root { + Some(root) => Some(root.trim().to_string()).filter(|root| !root.is_empty()), + None => saved.remote_root, + }, + keep_remote_runs: req.keep_remote_runs.unwrap_or(saved.keep_remote_runs), }; crate::jobs::ssh::validate_host_options(&options).map_err(bad_request)?; crate::config::set_ssh_host(host, options)?; @@ -860,6 +864,7 @@ pub(super) fn slurm_settings_json() -> Value { "cpusPerTask": settings.cpus_per_task, "mem": settings.mem, "remoteRoot": settings.remote_root, + "keepRemoteRuns": settings.keep_remote_runs, "hosts": list_ssh_hosts(), }) } @@ -882,6 +887,7 @@ pub(super) struct SetSlurmSettingsReq { cpus_per_task: Option, mem: Option, remote_root: Option, + keep_remote_runs: Option, } pub(super) async fn set_slurm_settings(Json(req): Json) -> ApiResult { @@ -925,6 +931,9 @@ pub(super) async fn set_slurm_settings(Json(req): Json) -> } settings.remote_root = root; } + if let Some(keep) = req.keep_remote_runs { + settings.keep_remote_runs = keep; + } slurm::save_settings(&settings)?; Ok(Json(slurm_settings_json())) }) @@ -1467,8 +1476,11 @@ pub(crate) async fn configure(backend: &str, body: Value) -> Result { .await .map_err(|ApiError(_, message)| anyhow!(message))?; } - if body.get("container").is_some() || body.get("remoteRoot").is_some() { - // A root-only change keeps the host's saved container. + if ["container", "remoteRoot", "keepRemoteRuns"] + .iter() + .any(|key| body.get(*key).is_some()) + { + // A change without --container keeps the host's saved container. let container = match body.get("container") { Some(container) => container.clone(), None => { @@ -1480,8 +1492,10 @@ pub(crate) async fn configure(backend: &str, body: Value) -> Result { } }; let mut req = json!({"host": body["host"], "container": container}); - if let Some(root) = body.get("remoteRoot") { - req["remoteRoot"] = root.clone(); + for key in ["remoteRoot", "keepRemoteRuns"] { + if let Some(value) = body.get(key) { + req[key] = value.clone(); + } } let req: SaveSshSettingsReq = serde_json::from_value(req)?; cli_result(save_ssh_settings(Json(req)).await)?; diff --git a/src/compute.rs b/src/compute.rs index 3de0a3ae6..281e13840 100644 --- a/src/compute.rs +++ b/src/compute.rs @@ -1192,6 +1192,7 @@ pub(crate) mod tests { SshHostSettings { container: Some("research".into()), remote_root: Some("/scratch/lab/orx".into()), + keep_remote_runs: false, }, ), ("direct".into(), SshHostSettings::default()), diff --git a/src/config.rs b/src/config.rs index adbb02d20..63c9b3cb0 100644 --- a/src/config.rs +++ b/src/config.rs @@ -464,6 +464,9 @@ pub struct SshHostSettings { /// Absolute storage root for runs and the source cache; `None` is `~/.orx`. #[serde(default, skip_serializing_if = "Option::is_none")] pub remote_root: Option, + /// Keep finished runs' remote dirs instead of removing them once logs are collected. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub keep_remote_runs: bool, } #[derive(Debug, Clone, Default, Serialize, Deserialize)] diff --git a/src/jobs/slurm.rs b/src/jobs/slurm.rs index bc7df5731..fcfa6ff92 100644 --- a/src/jobs/slurm.rs +++ b/src/jobs/slurm.rs @@ -61,6 +61,9 @@ pub struct SlurmSettings { /// Absolute shared-filesystem root for runs and the source cache; `None` is `~/.orx`. #[serde(default, skip_serializing_if = "Option::is_none")] pub remote_root: Option, + /// Keep finished runs' remote dirs instead of removing them once logs are collected. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub keep_remote_runs: bool, } fn settings_path() -> std::path::PathBuf { diff --git a/src/jobs/ssh.rs b/src/jobs/ssh.rs index 292749210..993e6ddaa 100644 --- a/src/jobs/ssh.rs +++ b/src/jobs/ssh.rs @@ -14,6 +14,7 @@ //! exit_code written when the payload finishes //! A restarted `orx supervise` reattaches purely from that directory. +pub(crate) mod cleanup; mod container; #[cfg(unix)] mod prepared; @@ -124,13 +125,17 @@ pub fn remote_run_dir(root: Option<&str>, run_id: &str) -> String { } } -fn remote_source_cache(root: Option<&str>, digest: &str) -> String { +fn remote_source_dir(root: Option<&str>) -> String { match root { - Some(root) => format!("{root}/source/{digest}.tar"), - None => format!(".orx/source/{digest}.tar"), + Some(root) => format!("{root}/source"), + None => ".orx/source".into(), } } +fn remote_source_cache(root: Option<&str>, digest: &str) -> String { + format!("{}/{digest}.tar", remote_source_dir(root)) +} + /// One shell word for a remote path; relative paths resolve against `$HOME`. pub(crate) fn shell_path(path: &str) -> String { if path.starts_with('/') { @@ -163,10 +168,7 @@ pub fn resolve_options( } else { args.container.clone().or_else(|| saved.container.clone()) }; - let options = crate::config::SshHostSettings { - container, - remote_root: saved.remote_root, - }; + let options = crate::config::SshHostSettings { container, ..saved }; validate_host_options(&options)?; Ok((host, options)) } @@ -791,10 +793,11 @@ pub async fn stage_source( /// Remote (check, upload) scripts for a cached source tar. An interrupted upload /// still lets `cat` exit 0 on EOF, so both sides require the exact byte length; /// a short cached file reads as absent and is replaced by the next upload. +/// A reused tar is touched so cache cleanup sees it as recently used. fn source_cache_commands(cache: &str, size: u64) -> (String, String) { let (dir, file) = (shell_path(parent_dir(cache)), shell_path(cache)); let check = format!( - "f={file}; test -f \"$f\" && test $(wc -c < \"$f\") -eq {size} && echo present || true" + "f={file}; test -f \"$f\" && test $(wc -c < \"$f\") -eq {size} && touch \"$f\" && echo present || true" ); let upload = format!( "umask 077; mkdir -p {dir}; f={file}; tmp=\"$f.tmp.$$\"; \ diff --git a/src/jobs/ssh/cleanup.rs b/src/jobs/ssh/cleanup.rs new file mode 100644 index 000000000..5bcc0a631 --- /dev/null +++ b/src/jobs/ssh/cleanup.rs @@ -0,0 +1,283 @@ +//! Removing finished runs' remote dirs and unused cached source tars on SSH and Slurm hosts. + +use std::collections::{BTreeMap, BTreeSet}; +use std::time::Duration; + +use super::{ + remote_run_dir, remote_source_dir, sh_quote, shell_path, ssh_run, validate_remote_root, + SshTarget, HOST_PROCESS_HELPERS, +}; +use crate::error::{anyhow, Result}; +use crate::jobs::BackendDescriptor; + +/// Cached tars untouched this long are no longer being staged by any orx install. +const SOURCE_IDLE_MINUTES: u32 = 60; +const HOST_TIMEOUT: Duration = Duration::from_secs(120); + +/// A run's remote dir, checked against the layout orx writes before anything removes it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RemoteRun { + pub run_id: String, + pub host: String, + pub root: Option, + dir: String, + /// Docker id and in-container run dir for container runs. + container: Option<(String, String)>, +} + +impl RemoteRun { + /// `Ok(None)` for backends whose remote state orx does not own. + pub fn from_descriptor(run_id: &str, descriptor: &BackendDescriptor) -> Result> { + if !matches!(descriptor.kind.as_str(), "ssh_job" | "slurm_job") { + return Ok(None); + } + if uuid::Uuid::parse_str(run_id).is_err() { + return Err(anyhow!("Unexpected run id {run_id}")); + } + let root = descriptor.remote_root.clone(); + if let Some(root) = &root { + validate_remote_root(root)?; + } + let dir = remote_run_dir(root.as_deref(), run_id); + // SSH records its dir; a different one was not created by this layout. + if descriptor.kind == "ssh_job" && descriptor.job_id.as_deref().is_some_and(|d| d != dir) { + return Err(anyhow!("Run {run_id} has an unexpected remote dir")); + } + let host = descriptor + .namespace + .clone() + .filter(|host| !host.is_empty()) + .ok_or_else(|| anyhow!("Run {run_id} has no remote host"))?; + let container = match &descriptor.ssh_container { + Some(container) + if container.run_dir.starts_with('/') + && container.run_dir.ends_with(&format!("/.orx/runs/{run_id}")) => + { + Some((container.id.clone(), container.run_dir.clone())) + } + Some(_) => return Err(anyhow!("Run {run_id} has an unexpected container dir")), + None => None, + }; + Ok(Some(Self { + run_id: run_id.into(), + host, + root, + dir, + container, + })) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Outcome { + Removed, + /// The run's process group is still alive on the host. + Busy, + Failed, +} + +/// One script per host: remove each run dir whose process group is gone, then cached tars +/// that no active run references and nothing has touched recently. +fn host_script(runs: &[RemoteRun], keep_digests: &BTreeSet) -> String { + let mut script = format!("{HOST_PROCESS_HELPERS}\n"); + for run in runs { + let inner = run + .container + .as_ref() + .map_or_else(String::new, |(id, dir)| { + format!( + "docker exec {} rm -rf -- {} >/dev/null 2>&1; ", + sh_quote(id), + sh_quote(dir) + ) + }); + script.push_str(&format!( + "d={dir}; p=$(cat \"$d/pid\" 2>/dev/null); case \"$p\" in ''|*[!0-9]*) p=;; esac; \ + if [ -n \"$p\" ] && host_group_alive \"$p\"; then echo 'BUSY {id}'; \ + else {inner}rm -rf -- \"$d\" && echo 'REMOVED {id}'; fi\n", + dir = shell_path(&run.dir), + id = run.run_id, + )); + } + let keep = keep_digests + .iter() + .map(String::as_str) + .collect::>() + .join(" "); + let roots: BTreeSet<_> = runs.iter().map(|run| run.root.as_deref()).collect(); + for root in roots { + script.push_str(&format!( + "s={source}; if [ -d \"$s\" ]; then \ + find \"$s\" -maxdepth 1 -type f \\( -name '*.tar' -o -name '*.tar.tmp.*' \\) -mmin +{SOURCE_IDLE_MINUTES} \ + | while IFS= read -r f; do n=${{f##*/}}; case ' {keep} ' in *\" ${{n%%.tar*}} \"*) ;; *) rm -f -- \"$f\";; esac; done; fi\n", + source = shell_path(&remote_source_dir(root)), + )); + } + script.push_str("exit 0\n"); + script +} + +fn parse_outcomes(runs: &[RemoteRun], output: &str) -> BTreeMap { + runs.iter() + .map(|run| { + let outcome = if output + .lines() + .any(|line| line == format!("REMOVED {}", run.run_id)) + { + Outcome::Removed + } else if output + .lines() + .any(|line| line == format!("BUSY {}", run.run_id)) + { + Outcome::Busy + } else { + Outcome::Failed + }; + (run.run_id.clone(), outcome) + }) + .collect() +} + +/// Removes `runs` (all on `host`) and that host's unreferenced cached tars. +pub async fn clean_host( + host: &str, + runs: &[RemoteRun], + keep_digests: &BTreeSet, +) -> Result> { + let script = host_script(runs, keep_digests); + let command = format!("bash -c {}", sh_quote(&script)); + let output = tokio::time::timeout( + HOST_TIMEOUT, + ssh_run(&SshTarget::alias(host), &command, None), + ) + .await + .map_err(|_| anyhow!("Timed out cleaning up runs on {host}"))??; + Ok(parse_outcomes(runs, &output)) +} + +#[cfg(test)] +mod tests { + use super::*; + + const RUN: &str = "6f1c8f0e-6b5e-4c41-9a0e-6d3f4c1b2a10"; + + fn descriptor(json: &str) -> BackendDescriptor { + BackendDescriptor::parse(json).unwrap() + } + + #[test] + fn only_orx_layout_dirs_are_cleanable() { + let ssh = descriptor(&format!( + r#"{{"kind":"ssh_job","namespace":"lab","jobId":".orx/runs/{RUN}"}}"# + )); + let run = RemoteRun::from_descriptor(RUN, &ssh).unwrap().unwrap(); + assert_eq!(run.dir, format!(".orx/runs/{RUN}")); + let rooted = descriptor(&format!( + r#"{{"kind":"ssh_job","namespace":"lab","jobId":"/scratch/orx/runs/{RUN}","remoteRoot":"/scratch/orx"}}"# + )); + assert!(RemoteRun::from_descriptor(RUN, &rooted).unwrap().is_some()); + let slurm = + descriptor(r#"{"kind":"slurm_job","namespace":"login","remoteRoot":"/scratch/orx"}"#); + assert_eq!( + RemoteRun::from_descriptor(RUN, &slurm) + .unwrap() + .unwrap() + .dir, + format!("/scratch/orx/runs/{RUN}") + ); + for json in [ + r#"{"kind":"ssh_job","namespace":"lab","jobId":"/home/me"}"#.to_string(), + r#"{"kind":"slurm_job","namespace":"login","remoteRoot":"/a b"}"#.to_string(), + format!(r#"{{"kind":"ssh_job","jobId":".orx/runs/{RUN}"}}"#), + format!( + r#"{{"kind":"ssh_job","namespace":"lab","jobId":".orx/runs/{RUN}","sshContainer":{{"reference":"c","id":"i","startedAt":"t","runDir":"/root"}}}}"# + ), + ] { + assert!( + RemoteRun::from_descriptor(RUN, &descriptor(&json)).is_err(), + "{json}" + ); + } + assert!(RemoteRun::from_descriptor("../x", &ssh).is_err()); + for kind in ["openresearch_job", "k8s_job", "local_job"] { + let other = descriptor(&format!(r#"{{"kind":"{kind}","namespace":"x"}}"#)); + assert!(RemoteRun::from_descriptor(RUN, &other).unwrap().is_none()); + } + } + + fn run(id: &str, root: Option<&str>) -> RemoteRun { + RemoteRun { + run_id: id.into(), + host: "lab".into(), + root: root.map(Into::into), + dir: remote_run_dir(root, id), + container: None, + } + } + + #[cfg(unix)] + #[test] + fn host_script_removes_finished_dirs_and_stale_unreferenced_tars() { + use std::os::unix::process::CommandExt as _; + let home = crate::local::git::TemporaryDirectory::new("orx-cleanup-home").unwrap(); + let runs = home.path().join(".orx/runs"); + let source = home.path().join(".orx/source"); + for id in ["done", "busy"] { + std::fs::create_dir_all(runs.join(id).join("repo")).unwrap(); + } + let mut sleeper = std::process::Command::new("sleep") + .arg("30") + .process_group(0) + .spawn() + .unwrap(); + std::fs::write(runs.join("busy/pid"), sleeper.id().to_string()).unwrap(); + std::fs::write(runs.join("done/pid"), "999999").unwrap(); + std::fs::create_dir_all(&source).unwrap(); + for name in ["old.tar", "kept.tar", "fresh.tar", "old.tar.tmp.12"] { + std::fs::write(source.join(name), "x").unwrap(); + } + for name in ["old.tar", "kept.tar", "old.tar.tmp.12"] { + assert!(std::process::Command::new("touch") + .args(["-t", "202001010000"]) + .arg(source.join(name)) + .status() + .unwrap() + .success()); + } + let batch = [run("done", None), run("busy", None), run("never", None)]; + let out = std::process::Command::new("bash") + .args(["-c", &host_script(&batch, &BTreeSet::from(["kept".into()]))]) + .env("HOME", home.path()) + .output() + .unwrap(); + let _ = sleeper.kill(); + let _ = sleeper.wait(); + assert!(out.status.success()); + let outcomes = parse_outcomes(&batch, &String::from_utf8(out.stdout).unwrap()); + assert_eq!(outcomes["done"], Outcome::Removed); + assert_eq!(outcomes["busy"], Outcome::Busy); + // An already-missing dir counts as removed. + assert_eq!(outcomes["never"], Outcome::Removed); + assert!(!runs.join("done").exists()); + assert!(runs.join("busy/repo").exists()); + let mut left: Vec<_> = std::fs::read_dir(&source) + .unwrap() + .map(|e| e.unwrap().file_name().into_string().unwrap()) + .collect(); + left.sort(); + assert_eq!(left, ["fresh.tar", "kept.tar"]); + } + + #[test] + fn host_script_quotes_rooted_and_container_paths() { + let mut rooted = run(RUN, Some("/scratch/orx")); + rooted.container = Some(("abc".into(), format!("/root/.orx/runs/{RUN}"))); + let script = host_script(&[rooted], &BTreeSet::new()); + assert!(script.contains(&format!("d='/scratch/orx/runs/{RUN}';"))); + assert!(script.contains(&format!( + "docker exec 'abc' rm -rf -- '/root/.orx/runs/{RUN}'" + ))); + assert!(script.contains("s='/scratch/orx/source';")); + assert!(parse_outcomes(&[run(RUN, None)], "REMOVED other\n")[RUN] == Outcome::Failed); + } +} diff --git a/src/local/mod.rs b/src/local/mod.rs index 3b3a78502..a8c77bb9a 100644 --- a/src/local/mod.rs +++ b/src/local/mod.rs @@ -40,6 +40,7 @@ pub mod overleaf_live; pub mod portable_git; pub mod projects; pub mod ray; +pub mod remote_cleanup; pub mod resolve; pub mod shell_env; pub mod skills; diff --git a/src/local/remote_cleanup.rs b/src/local/remote_cleanup.rs new file mode 100644 index 000000000..2979212ed --- /dev/null +++ b/src/local/remote_cleanup.rs @@ -0,0 +1,105 @@ +//! Background removal of queued remote run dirs; the remote half lives in `jobs::ssh::cleanup`. + +use std::collections::{BTreeMap, BTreeSet}; + +use crate::error::Result; +use crate::jobs::ssh::cleanup::{self, Outcome, RemoteRun}; +use crate::jobs::BackendDescriptor; +use crate::store::Store; + +/// Bounds one pass; the rest wait for the next. +const BATCH: usize = 50; + +/// (run id, reason) pairs for the log. +pub type Failures = Vec<(String, String)>; + +pub struct HostBatch { + host: String, + runs: Vec, + keep_digests: BTreeSet, +} + +/// Groups due entries by host. Entries that need no cleanup leave the queue; ones orx +/// cannot safely clean leave it too and come back as (run, reason). +pub fn plan(store: &Store) -> Result<(Vec, Failures)> { + let ssh = crate::config::ssh_settings()?; + let slurm_keeps = crate::jobs::slurm::load_settings()? + .unwrap_or_default() + .keep_remote_runs; + let mut hosts: BTreeMap> = BTreeMap::new(); + let mut rejected = Vec::new(); + for entry in store.due_remote_cleanups(BATCH)? { + let parsed = BackendDescriptor::parse(&entry.backend_json).and_then(|descriptor| { + let keeps = match descriptor.kind.as_str() { + "slurm_job" => slurm_keeps, + _ => descriptor + .namespace + .as_deref() + .and_then(|host| ssh.hosts.get(host)) + .is_some_and(|saved| saved.keep_remote_runs), + }; + Ok(RemoteRun::from_descriptor(&entry.run_id, &descriptor)?.zip(Some(keeps))) + }); + match parsed { + Ok(Some((run, keeps))) if entry.deleted || !keeps => { + hosts.entry(run.host.clone()).or_default().push(run) + } + Ok(_) => store.finish_remote_cleanup(&entry.run_id)?, + Err(error) => { + store.finish_remote_cleanup(&entry.run_id)?; + rejected.push((entry.run_id, error.to_string())); + } + } + } + let batches = hosts + .into_iter() + .map(|(host, runs)| { + Ok(HostBatch { + keep_digests: store.active_remote_source_digests(&host)?, + host, + runs, + }) + }) + .collect::>()?; + Ok((batches, rejected)) +} + +pub type Results = Vec<(String, std::result::Result)>; + +pub async fn clean(batches: Vec) -> Results { + let mut results = Vec::new(); + for batch in batches { + match cleanup::clean_host(&batch.host, &batch.runs, &batch.keep_digests).await { + Ok(outcomes) => results.extend(outcomes.into_iter().map(|(id, o)| (id, Ok(o)))), + Err(error) => results.extend( + batch + .runs + .iter() + .map(|run| (run.run_id.clone(), Err(format!("{error:#}")))), + ), + } + } + results +} + +/// Records outcomes and returns failures worth reporting; failed and busy runs retry later. +pub fn record(store: &Store, results: Results) -> Result { + let mut failures = Vec::new(); + for (run_id, result) in results { + let error = match result { + Ok(Outcome::Removed) => { + store.finish_remote_cleanup(&run_id)?; + continue; + } + Ok(Outcome::Busy) => { + store.defer_remote_cleanup(&run_id, "its process group is still running")?; + continue; + } + Ok(Outcome::Failed) => "the remote dir could not be removed".to_string(), + Err(error) => error, + }; + store.defer_remote_cleanup(&run_id, &error)?; + failures.push((run_id, error)); + } + Ok(failures) +} diff --git a/src/store.rs b/src/store.rs index 96b1ea447..3e69ebdd3 100644 --- a/src/store.rs +++ b/src/store.rs @@ -508,6 +508,24 @@ enum OrphanFate { KeepWithAllSnapshots, } +/// Remote dirs waiting for removal; `deleted` entries ignore a target's keep setting. +pub(crate) struct RemoteCleanup { + pub run_id: String, + pub backend_json: String, + pub deleted: bool, +} + +const QUEUE_REMOTE_CLEANUP: &str = + "INSERT INTO remote_cleanups (run_id, backend_json, deleted, next_attempt_at)"; +// Appended to a SELECT's WHERE, which SQLite needs before ON CONFLICT; deletion wins over a pending finish. +const UPSERT_REMOTE_CLEANUP: &str = "AND CASE WHEN json_valid(backend_json) + THEN json_extract(backend_json, '$.kind') END IN ('ssh_job', 'slurm_job') + ON CONFLICT(run_id) DO UPDATE SET backend_json = excluded.backend_json, + deleted = MAX(deleted, excluded.deleted), next_attempt_at = excluded.next_attempt_at"; +/// Failed remote cleanups retry after 5 minutes, doubling up to a day. +const REMOTE_CLEANUP_RETRY_MS: i64 = 5 * 60 * 1000; +const REMOTE_CLEANUP_MAX_RETRY_MS: i64 = 24 * 60 * 60 * 1000; + pub struct Store { conn: Connection, data_dir: PathBuf, @@ -760,6 +778,14 @@ impl Store { error TEXT, tested_at INTEGER NOT NULL ); + CREATE TABLE IF NOT EXISTS remote_cleanups ( + run_id TEXT PRIMARY KEY, + backend_json TEXT NOT NULL, + deleted INTEGER NOT NULL, + attempts INTEGER NOT NULL DEFAULT 0, + next_attempt_at INTEGER NOT NULL, + last_error TEXT + ); CREATE TABLE IF NOT EXISTS overleaf_links ( project_id TEXT NOT NULL, tex_path TEXT NOT NULL, @@ -2121,6 +2147,72 @@ impl Store { Ok(()) } + /// Queues a finished run's remote dir for removal; non-SSH/Slurm runs are ignored. + pub(crate) fn queue_finished_remote_cleanup(&self, run_id: &str) -> Result<()> { + self.conn.execute( + &format!("{QUEUE_REMOTE_CLEANUP} SELECT id, backend_json, 0, ?2 FROM runs WHERE id = ?1 {UPSERT_REMOTE_CLEANUP}"), + params![run_id, now_ms()], + )?; + Ok(()) + } + + pub(crate) fn due_remote_cleanups(&self, limit: usize) -> Result> { + let mut statement = self.conn.prepare( + "SELECT run_id, backend_json, deleted FROM remote_cleanups + WHERE next_attempt_at <= ?1 ORDER BY next_attempt_at LIMIT ?2", + )?; + let rows = statement.query_map(params![now_ms(), limit as i64], |row| { + Ok(RemoteCleanup { + run_id: row.get(0)?, + backend_json: row.get(1)?, + deleted: row.get(2)?, + }) + })?; + Ok(rows.collect::, _>>()?) + } + + pub(crate) fn finish_remote_cleanup(&self, run_id: &str) -> Result<()> { + self.conn.execute( + "DELETE FROM remote_cleanups WHERE run_id = ?1", + params![run_id], + )?; + Ok(()) + } + + pub(crate) fn defer_remote_cleanup(&self, run_id: &str, error: &str) -> Result<()> { + self.conn.execute( + "UPDATE remote_cleanups SET attempts = attempts + 1, last_error = ?2, + next_attempt_at = ?3 + MIN(?4 * (1 << MIN(attempts, 10)), ?5) + WHERE run_id = ?1", + params![ + run_id, + error, + now_ms(), + REMOTE_CLEANUP_RETRY_MS, + REMOTE_CLEANUP_MAX_RETRY_MS + ], + )?; + Ok(()) + } + + /// Source digests that unfinished SSH/Slurm runs on `host` may still stage. + pub(crate) fn active_remote_source_digests( + &self, + host: &str, + ) -> Result> { + let mut digests = std::collections::BTreeSet::new(); + for run in self.list_active_runs()? { + if let Ok(descriptor) = crate::jobs::BackendDescriptor::parse(&run.backend_json) { + if matches!(descriptor.kind.as_str(), "ssh_job" | "slurm_job") + && descriptor.namespace.as_deref() == Some(host) + { + digests.extend(descriptor.source_digest); + } + } + } + Ok(digests) + } + fn data_dir_move_claimed(&self) -> Result { self.clear_stale_data_dir_move_lease()?; Ok(self.conn.query_row( @@ -2182,6 +2274,10 @@ impl Store { "DELETE FROM chat_sessions WHERE project_id = ?1", params![id], )?; + tx.execute( + &format!("{QUEUE_REMOTE_CLEANUP} SELECT id, backend_json, 1, ?2 FROM runs WHERE project_id = ?1 {UPSERT_REMOTE_CLEANUP}"), + params![id, now_ms()], + )?; self.conn .execute("DELETE FROM runs WHERE project_id = ?1", params![id])?; self.conn.execute( @@ -4166,6 +4262,63 @@ mod tests { std::fs::remove_dir_all(dir).unwrap(); } + #[test] + fn remote_cleanups_queue_on_finish_and_project_delete_then_back_off() { + let dir = std::env::temp_dir().join(format!("orx-remote-cleanup-{}", uuid::Uuid::new_v4())); + let store = Store::open_at(dir.clone()).unwrap(); + project_with_run_artifacts(&store, "p", "local", &"d".repeat(64)); + let insert = |id: &str, status: &str, kind: &str| { + let descriptor = + serde_json::json!({"kind": kind, "namespace": "lab", "sourceDigest": id}); + store.conn.execute( + "INSERT INTO runs (id, experiment_id, project_id, status, backend_json, created_at, updated_at) + VALUES (?1, 'experiment', 'p', ?2, ?3, 1, 1)", + params![id, status, descriptor.to_string()], + ).unwrap(); + }; + insert("ssh", "done", "ssh_job"); + insert("slurm", "failed", "slurm_job"); + insert("box", "done", "openresearch_job"); + insert("live", "running", "ssh_job"); + for id in ["ssh", "box", "local"] { + store.queue_finished_remote_cleanup(id).unwrap(); + } + let due = |store: &Store| { + store + .due_remote_cleanups(10) + .unwrap() + .into_iter() + .map(|entry| (entry.run_id, entry.deleted)) + .collect::>() + }; + assert_eq!(due(&store), [("ssh".to_string(), false)]); + assert_eq!( + store.active_remote_source_digests("lab").unwrap(), + std::collections::BTreeSet::from(["live".to_string()]) + ); + store.defer_remote_cleanup("ssh", "unreachable").unwrap(); + assert!(due(&store).is_empty()); + store + .conn + .execute("UPDATE runs SET status = 'done'", []) + .unwrap(); + store.delete_local_project("p").unwrap(); + let mut entries = due(&store); + entries.sort(); + // Deletion marks every SSH/Slurm run, including one already waiting to retry. + assert_eq!( + entries, + [ + ("live".to_string(), true), + ("slurm".to_string(), true), + ("ssh".to_string(), true) + ] + ); + store.finish_remote_cleanup("ssh").unwrap(); + assert_eq!(due(&store).len(), 2); + std::fs::remove_dir_all(dir).unwrap(); + } + #[test] fn orphan_cleanup_preserves_live_local_runs_and_remote_recovery_handles() { let dir = std::env::temp_dir().join(format!("orx-live-orphan-{}", uuid::Uuid::new_v4())); diff --git a/src/telemetry.rs b/src/telemetry.rs index 087e8cd06..fcb7a5396 100644 --- a/src/telemetry.rs +++ b/src/telemetry.rs @@ -1630,6 +1630,7 @@ mod tests { let options = SshHostSettings { container: Some("research".into()), remote_root: Some("/scratch/lab/orx".into()), + keep_remote_runs: true, }; set_ssh_host("lab".into(), options.clone()).unwrap(); set_ssh_default(Some("lab".into())).unwrap(); From 277ad8aa2d86d2f3472501fd889a31247860fd69 Mon Sep 17 00:00:00 2001 From: Myles Anderson Date: Thu, 22 Oct 2026 15:08:23 -0700 Subject: [PATCH 05/10] Keep TensorBoard logs and user tars during remote cleanup Co-Authored-By: Claude Opus 5.5 --- agent-skills/orx-compute/references/slurm.md | 11 +- agent-skills/orx-compute/references/ssh.md | 20 +- src/commands/compute.rs | 5 +- src/commands/supervise.rs | 2 +- src/commands/up.rs | 13 +- src/commands/up/compute_settings.rs | 3 +- src/jobs/ssh/cleanup.rs | 204 ++++++++++++------- src/local/remote_cleanup.rs | 60 ++++-- src/store.rs | 60 ++---- 9 files changed, 227 insertions(+), 151 deletions(-) diff --git a/agent-skills/orx-compute/references/slurm.md b/agent-skills/orx-compute/references/slurm.md index bdf4c2c8d..889a2d433 100644 --- a/agent-skills/orx-compute/references/slurm.md +++ b/agent-skills/orx-compute/references/slurm.md @@ -28,10 +28,11 @@ orx exp run --backend slurm be on a filesystem the compute nodes share. Each run keeps the root it was submitted with; `--clear remote-root` restores the default for new runs. - Once a run finishes and its log is collected, the running OpenResearch app - removes its remote run dir, along with cached source archives that no active - run uses. Anything the job wrote into `repo/` goes with it. To keep results - on the cluster, run `orx compute configure slurm --keep-remote-runs true`, or - write outputs outside the run dir. Deleting a project removes its runs' - remote dirs regardless. + removes its remote run dir except `tensorboard/`, along with orx's cached + source archives that no active run uses and that have been idle for an hour. + Anything the job wrote into `repo/` goes with it. To keep results on the + cluster, set `orx compute configure slurm --keep-remote-runs true` before + runs finish, or write outputs outside the run dir. Deleting a project + removes its runs' remote dirs entirely regardless. - A detached `orx supervise` process records scheduler status and logs; do not kill it. diff --git a/agent-skills/orx-compute/references/ssh.md b/agent-skills/orx-compute/references/ssh.md index 57ab52514..71f8180bf 100644 --- a/agent-skills/orx-compute/references/ssh.md +++ b/agent-skills/orx-compute/references/ssh.md @@ -37,13 +37,6 @@ an absolute remote storage root (`--remote-root /data/alice/orx`) used instead of `~/.orx` for that host's runs and cached source archives. Each run keeps the root it was launched with. -Once a run finishes and its log is collected, the running OpenResearch app -removes the run's remote dir (including anything written into `repo/`) and -cached source archives that no active run uses. Unreachable hosts are retried -later. `orx compute configure ssh --host lab --keep-remote-runs true` keeps that -host's finished runs. Deleting a project removes its runs' remote dirs -regardless. - A saved default SSH host lets launches omit `--host`. These defaults apply only to SSH experiments; dashboards, terminals and coding agents still connect directly to the host. @@ -51,6 +44,14 @@ agents still connect directly to the host. - Omitted container selection uses the host's saved target. `--container` overrides it; `--no-container` explicitly runs on the host. The flags conflict. +Once a run finishes and its log is collected, the running OpenResearch app +removes the run's remote dir except `tensorboard/`, so anything written into +`repo/` goes with it. It also removes orx's cached source archives that no +active run uses and that have been idle for an hour. Unreachable hosts are +retried later. `orx compute configure ssh --host lab --keep-remote-runs true` +keeps that host's runs that finish while it is set. Deleting a project removes +its runs' remote dirs entirely, regardless of that setting. + Keep environment preparation in the project's committed setup/run scripts and invoke them through the fixed run command, just as with other backends. For example, a committed `run.sh` can source Conda, activate the environment, and @@ -66,8 +67,9 @@ creation or lifecycle management. `--image` and `--flavor` are unsupported. The committed source archive is cached on the host, then streamed into the container's `$HOME/.orx/runs//repo`. Host logs and exit status stay in the host's `/runs//` (`~/.orx` unless a remote root is saved), so -container removal does not erase them. A detached supervisor reattaches using the saved container ID and start time; later settings -changes cannot redirect an existing run. Do not kill the supervisor. +container removal does not erase them before they are collected. A detached +supervisor reattaches using the saved container ID and start time; later +settings changes cannot redirect an existing run. Do not kill the supervisor. Cancellation targets only the experiment process group, waits five seconds for TERM, then uses KILL if needed. It leaves the container and unrelated processes diff --git a/src/commands/compute.rs b/src/commands/compute.rs index 3d7e94fcd..fa6ebd53d 100644 --- a/src/commands/compute.rs +++ b/src/commands/compute.rs @@ -238,7 +238,7 @@ pub struct ConfigureArgs { /// Absolute remote dir for runs and the source cache (SSH per host, or Slurm); default ~/.orx. #[arg(long)] remote_root: Option, - /// Keep finished runs' remote dirs (true) instead of removing them once logs are collected. + /// true keeps finished runs' remote dirs; false (the default) removes them once logs are collected. #[arg(long)] keep_remote_runs: Option, #[arg(long)] @@ -408,6 +408,7 @@ impl ConfigureArgs { "cpus-per-task" => "cpusPerTask", "default-host" => "defaultHost", "remote-root" => "remoteRoot", + "keep-remote-runs" => "keepRemoteRuns", other => other, }; if !allowed.contains(&key) @@ -425,6 +426,8 @@ impl ConfigureArgs { Value::Null } else if key == "cpusPerTask" { json!(0) + } else if key == "keepRemoteRuns" { + json!(false) } else { json!("") }, diff --git a/src/commands/supervise.rs b/src/commands/supervise.rs index 376f9dda8..19f41a782 100644 --- a/src/commands/supervise.rs +++ b/src/commands/supervise.rs @@ -849,7 +849,7 @@ async fn watch_ssh_job( ) .await; } - if descriptor.kind == "ssh_job" && collected { + if collected { queue_remote_cleanup(store, run_id); } eprintln!("supervise {run_id}: finished ({status})"); diff --git a/src/commands/up.rs b/src/commands/up.rs index 2559236b0..44fb81007 100644 --- a/src/commands/up.rs +++ b/src/commands/up.rs @@ -4621,6 +4621,7 @@ impl axum::body::HttpBody for ActiveBody { } const ORPHANED_RUN_SCAN_INTERVAL: Duration = Duration::from_secs(60); +const REMOTE_CLEANUP_INTERVAL: Duration = Duration::from_secs(60); /// SSH container launches record their handle before the job starts; leave them this long to finish. const ORPHANED_RUN_SETTLE_MS: i64 = 5 * 60 * 1000; @@ -4721,8 +4722,10 @@ async fn with_gated_store( /// Removes queued remote run dirs; the gate is held only around store access, never over SSH. async fn reclaim_remote_run_dirs(moving: Arc, gate: Arc>) { + // Startup's own SSH traffic (supervisor adoption, settings checks) goes first. tokio::time::sleep(Duration::from_secs(30)).await; let mut reported = HashSet::new(); + let mut plan_error = None; loop { let mut report = |failures: Vec<(String, String)>| { for (run_id, error) in failures { @@ -4750,10 +4753,16 @@ async fn reclaim_remote_run_dirs(moving: Arc, gate: Arc eprintln!("orx up: could not plan remote cleanup: {error}"), + Some(Err(error)) => { + let error = error.to_string(); + if plan_error.as_ref() != Some(&error) { + eprintln!("orx up: could not plan remote cleanup: {error}"); + plan_error = Some(error); + } + } None => {} } - tokio::time::sleep(ORPHANED_RUN_SCAN_INTERVAL).await; + tokio::time::sleep(REMOTE_CLEANUP_INTERVAL).await; } } diff --git a/src/commands/up/compute_settings.rs b/src/commands/up/compute_settings.rs index 2fe9be033..4fd097b68 100644 --- a/src/commands/up/compute_settings.rs +++ b/src/commands/up/compute_settings.rs @@ -624,8 +624,9 @@ pub(super) async fn ssh_settings() -> ApiResult { pub(super) struct SaveSshSettingsReq { host: String, container: Option, - /// Unlike `container`, `None` keeps the saved value (the UI never sends these); `Some("")` clears the root. + /// Unlike `container`, `None` keeps the saved root (the UI never sends it); `Some("")` clears it. remote_root: Option, + /// `None` keeps the saved value. keep_remote_runs: Option, } diff --git a/src/jobs/ssh/cleanup.rs b/src/jobs/ssh/cleanup.rs index 5bcc0a631..5040578a1 100644 --- a/src/jobs/ssh/cleanup.rs +++ b/src/jobs/ssh/cleanup.rs @@ -17,17 +17,27 @@ const HOST_TIMEOUT: Duration = Duration::from_secs(120); /// A run's remote dir, checked against the layout orx writes before anything removes it. #[derive(Debug, Clone, PartialEq, Eq)] pub struct RemoteRun { - pub run_id: String, - pub host: String, - pub root: Option, + pub(crate) run_id: String, + pub(crate) host: String, + root: Option, dir: String, /// Docker id and in-container run dir for container runs. container: Option<(String, String)>, + /// `orx tensorboard` reads a finished run's remote `tensorboard/`; only deletion removes it. + keep_tensorboard: bool, } impl RemoteRun { - /// `Ok(None)` for backends whose remote state orx does not own. - pub fn from_descriptor(run_id: &str, descriptor: &BackendDescriptor) -> Result> { + /// `Ok(None)` when orx owns no remote dir for the run; `deleted` also removes TensorBoard logs. + pub fn from_descriptor( + run_id: &str, + descriptor: &BackendDescriptor, + deleted: bool, + ) -> Result> { + // A submission that failed before choosing a host left nothing remote. + let Some(host) = descriptor.namespace.clone().filter(|host| !host.is_empty()) else { + return Ok(None); + }; if !matches!(descriptor.kind.as_str(), "ssh_job" | "slurm_job") { return Ok(None); } @@ -43,11 +53,6 @@ impl RemoteRun { if descriptor.kind == "ssh_job" && descriptor.job_id.as_deref().is_some_and(|d| d != dir) { return Err(anyhow!("Run {run_id} has an unexpected remote dir")); } - let host = descriptor - .namespace - .clone() - .filter(|host| !host.is_empty()) - .ok_or_else(|| anyhow!("Run {run_id} has no remote host"))?; let container = match &descriptor.ssh_container { Some(container) if container.run_dir.starts_with('/') @@ -64,16 +69,17 @@ impl RemoteRun { root, dir, container, + keep_tensorboard: !deleted, })) } } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] pub enum Outcome { Removed, /// The run's process group is still alive on the host. Busy, - Failed, + Failed(String), } /// One script per host: remove each run dir whose process group is gone, then cached tars @@ -91,10 +97,16 @@ fn host_script(runs: &[RemoteRun], keep_digests: &BTreeSet) -> String { sh_quote(dir) ) }); + let remove = if run.keep_tensorboard { + "if [ -d \"$d/tensorboard\" ]; then find \"$d\" -mindepth 1 -maxdepth 1 ! -name tensorboard -exec rm -rf -- {} +; else rm -rf -- \"$d\"; fi" + } else { + "rm -rf -- \"$d\"" + }; script.push_str(&format!( "d={dir}; p=$(cat \"$d/pid\" 2>/dev/null); case \"$p\" in ''|*[!0-9]*) p=;; esac; \ if [ -n \"$p\" ] && host_group_alive \"$p\"; then echo 'BUSY {id}'; \ - else {inner}rm -rf -- \"$d\" && echo 'REMOVED {id}'; fi\n", + else {inner}if err=$({{ {remove}; }} 2>&1); then echo 'REMOVED {id}'; \ + else printf 'FAILED {id} %s\\n' \"$(printf '%s' \"$err\" | head -n 1)\"; fi; fi\n", dir = shell_path(&run.dir), id = run.run_id, )); @@ -106,10 +118,15 @@ fn host_script(runs: &[RemoteRun], keep_digests: &BTreeSet) -> String { .join(" "); let roots: BTreeSet<_> = runs.iter().map(|run| run.root.as_deref()).collect(); for root in roots { + // Only names orx writes (`.tar` and its upload temps): a root may hold the user's own tars. script.push_str(&format!( "s={source}; if [ -d \"$s\" ]; then \ - find \"$s\" -maxdepth 1 -type f \\( -name '*.tar' -o -name '*.tar.tmp.*' \\) -mmin +{SOURCE_IDLE_MINUTES} \ - | while IFS= read -r f; do n=${{f##*/}}; case ' {keep} ' in *\" ${{n%%.tar*}} \"*) ;; *) rm -f -- \"$f\";; esac; done; fi\n", + find \"$s\" -maxdepth 1 -type f -mmin +{SOURCE_IDLE_MINUTES} | while IFS= read -r f; do \ + n=${{f##*/}}; b=${{n%%.tar*}}; \ + case \"$b\" in ''|*[!0-9a-f]*) continue;; esac; [ ${{#b}} -eq 64 ] || continue; \ + case \"$n\" in \"$b.tar\"|\"$b.tar.tmp.\"*) ;; *) continue;; esac; \ + case ' {keep} ' in *\" $b \"*) continue;; esac; \ + [ -n \"$(find \"$f\" -mmin +{SOURCE_IDLE_MINUTES})\" ] && rm -f -- \"$f\"; done; fi\n", source = shell_path(&remote_source_dir(root)), )); } @@ -118,21 +135,28 @@ fn host_script(runs: &[RemoteRun], keep_digests: &BTreeSet) -> String { } fn parse_outcomes(runs: &[RemoteRun], output: &str) -> BTreeMap { + let mut reported = BTreeMap::new(); + for line in output.lines() { + let mut words = line.splitn(3, ' '); + let outcome = match words.next() { + Some("REMOVED") => Outcome::Removed, + Some("BUSY") => Outcome::Busy, + Some("FAILED") => Outcome::Failed(String::new()), + _ => continue, + }; + if let Some(id) = words.next() { + let outcome = match outcome { + Outcome::Failed(_) => Outcome::Failed(words.next().unwrap_or_default().into()), + outcome => outcome, + }; + reported.insert(id.to_string(), outcome); + } + } runs.iter() .map(|run| { - let outcome = if output - .lines() - .any(|line| line == format!("REMOVED {}", run.run_id)) - { - Outcome::Removed - } else if output - .lines() - .any(|line| line == format!("BUSY {}", run.run_id)) - { - Outcome::Busy - } else { - Outcome::Failed - }; + let outcome = reported + .remove(&run.run_id) + .unwrap_or_else(|| Outcome::Failed("no result from the host".into())); (run.run_id.clone(), outcome) }) .collect() @@ -144,11 +168,11 @@ pub async fn clean_host( runs: &[RemoteRun], keep_digests: &BTreeSet, ) -> Result> { + // Over stdin: a large batch would overrun Windows' command-line limit. let script = host_script(runs, keep_digests); - let command = format!("bash -c {}", sh_quote(&script)); let output = tokio::time::timeout( HOST_TIMEOUT, - ssh_run(&SshTarget::alias(host), &command, None), + ssh_run(&SshTarget::alias(host), "bash -s", Some(&script)), ) .await .map_err(|_| anyhow!("Timed out cleaning up runs on {host}"))??; @@ -165,66 +189,71 @@ mod tests { BackendDescriptor::parse(json).unwrap() } + fn cleanable(json: &str) -> Result> { + RemoteRun::from_descriptor(RUN, &descriptor(json), false) + } + #[test] fn only_orx_layout_dirs_are_cleanable() { - let ssh = descriptor(&format!( - r#"{{"kind":"ssh_job","namespace":"lab","jobId":".orx/runs/{RUN}"}}"# - )); - let run = RemoteRun::from_descriptor(RUN, &ssh).unwrap().unwrap(); - assert_eq!(run.dir, format!(".orx/runs/{RUN}")); - let rooted = descriptor(&format!( + let ssh = format!(r#"{{"kind":"ssh_job","namespace":"lab","jobId":".orx/runs/{RUN}"}}"#); + assert_eq!( + cleanable(&ssh).unwrap().unwrap().dir, + format!(".orx/runs/{RUN}") + ); + let rooted = format!( r#"{{"kind":"ssh_job","namespace":"lab","jobId":"/scratch/orx/runs/{RUN}","remoteRoot":"/scratch/orx"}}"# - )); - assert!(RemoteRun::from_descriptor(RUN, &rooted).unwrap().is_some()); - let slurm = - descriptor(r#"{"kind":"slurm_job","namespace":"login","remoteRoot":"/scratch/orx"}"#); + ); + assert!(cleanable(&rooted).unwrap().is_some()); + let slurm = r#"{"kind":"slurm_job","namespace":"login","remoteRoot":"/scratch/orx"}"#; assert_eq!( - RemoteRun::from_descriptor(RUN, &slurm) - .unwrap() - .unwrap() - .dir, + cleanable(slurm).unwrap().unwrap().dir, format!("/scratch/orx/runs/{RUN}") ); for json in [ r#"{"kind":"ssh_job","namespace":"lab","jobId":"/home/me"}"#.to_string(), r#"{"kind":"slurm_job","namespace":"login","remoteRoot":"/a b"}"#.to_string(), - format!(r#"{{"kind":"ssh_job","jobId":".orx/runs/{RUN}"}}"#), format!( r#"{{"kind":"ssh_job","namespace":"lab","jobId":".orx/runs/{RUN}","sshContainer":{{"reference":"c","id":"i","startedAt":"t","runDir":"/root"}}}}"# ), ] { - assert!( - RemoteRun::from_descriptor(RUN, &descriptor(&json)).is_err(), - "{json}" - ); + assert!(cleanable(&json).is_err(), "{json}"); } - assert!(RemoteRun::from_descriptor("../x", &ssh).is_err()); + assert!(RemoteRun::from_descriptor("../x", &descriptor(&ssh), false).is_err()); + let unsubmitted = format!(r#"{{"kind":"ssh_job","jobId":".orx/runs/{RUN}"}}"#); + assert!(cleanable(&unsubmitted).unwrap().is_none()); for kind in ["openresearch_job", "k8s_job", "local_job"] { - let other = descriptor(&format!(r#"{{"kind":"{kind}","namespace":"x"}}"#)); - assert!(RemoteRun::from_descriptor(RUN, &other).unwrap().is_none()); + assert!( + cleanable(&format!(r#"{{"kind":"{kind}","namespace":"x"}}"#)) + .unwrap() + .is_none() + ); } } - fn run(id: &str, root: Option<&str>) -> RemoteRun { + fn run(id: &str, root: Option<&str>, deleted: bool) -> RemoteRun { RemoteRun { run_id: id.into(), host: "lab".into(), root: root.map(Into::into), dir: remote_run_dir(root, id), container: None, + keep_tensorboard: !deleted, } } #[cfg(unix)] #[test] - fn host_script_removes_finished_dirs_and_stale_unreferenced_tars() { + fn host_script_removes_finished_dirs_and_stale_orx_tars_only() { use std::os::unix::process::CommandExt as _; let home = crate::local::git::TemporaryDirectory::new("orx-cleanup-home").unwrap(); let runs = home.path().join(".orx/runs"); let source = home.path().join(".orx/source"); - for id in ["done", "busy"] { + for id in ["done", "busy", "charted", "deleted"] { std::fs::create_dir_all(runs.join(id).join("repo")).unwrap(); } + for id in ["charted", "deleted"] { + std::fs::create_dir_all(runs.join(id).join("tensorboard")).unwrap(); + } let mut sleeper = std::process::Command::new("sleep") .arg("30") .process_group(0) @@ -233,20 +262,35 @@ mod tests { std::fs::write(runs.join("busy/pid"), sleeper.id().to_string()).unwrap(); std::fs::write(runs.join("done/pid"), "999999").unwrap(); std::fs::create_dir_all(&source).unwrap(); - for name in ["old.tar", "kept.tar", "fresh.tar", "old.tar.tmp.12"] { + let (old, kept, fresh) = ("a".repeat(64), "b".repeat(64), "c".repeat(64)); + let names = [ + format!("{old}.tar"), + format!("{old}.tar.tmp.12"), + format!("{kept}.tar"), + format!("{fresh}.tar"), + "dataset.tar".to_string(), + format!("{old}.tar.bak"), + ]; + for name in &names { std::fs::write(source.join(name), "x").unwrap(); + if !name.starts_with(&fresh) { + assert!(std::process::Command::new("touch") + .args(["-t", "202001010000"]) + .arg(source.join(name)) + .status() + .unwrap() + .success()); + } } - for name in ["old.tar", "kept.tar", "old.tar.tmp.12"] { - assert!(std::process::Command::new("touch") - .args(["-t", "202001010000"]) - .arg(source.join(name)) - .status() - .unwrap() - .success()); - } - let batch = [run("done", None), run("busy", None), run("never", None)]; + let batch = [ + run("done", None, false), + run("busy", None, false), + run("never", None, false), + run("charted", None, false), + run("deleted", None, true), + ]; let out = std::process::Command::new("bash") - .args(["-c", &host_script(&batch, &BTreeSet::from(["kept".into()]))]) + .args(["-c", &host_script(&batch, &BTreeSet::from([kept.clone()]))]) .env("HOME", home.path()) .output() .unwrap(); @@ -254,23 +298,33 @@ mod tests { let _ = sleeper.wait(); assert!(out.status.success()); let outcomes = parse_outcomes(&batch, &String::from_utf8(out.stdout).unwrap()); - assert_eq!(outcomes["done"], Outcome::Removed); + for id in ["done", "never", "charted", "deleted"] { + assert_eq!(outcomes[id], Outcome::Removed, "{id}"); + } assert_eq!(outcomes["busy"], Outcome::Busy); - // An already-missing dir counts as removed. - assert_eq!(outcomes["never"], Outcome::Removed); assert!(!runs.join("done").exists()); assert!(runs.join("busy/repo").exists()); + assert!(runs.join("charted/tensorboard").is_dir()); + assert!(!runs.join("charted/repo").exists()); + assert!(!runs.join("deleted").exists()); let mut left: Vec<_> = std::fs::read_dir(&source) .unwrap() .map(|e| e.unwrap().file_name().into_string().unwrap()) .collect(); left.sort(); - assert_eq!(left, ["fresh.tar", "kept.tar"]); + let mut expected = vec![ + format!("{kept}.tar"), + format!("{fresh}.tar"), + "dataset.tar".to_string(), + format!("{old}.tar.bak"), + ]; + expected.sort(); + assert_eq!(left, expected); } #[test] fn host_script_quotes_rooted_and_container_paths() { - let mut rooted = run(RUN, Some("/scratch/orx")); + let mut rooted = run(RUN, Some("/scratch/orx"), true); rooted.container = Some(("abc".into(), format!("/root/.orx/runs/{RUN}"))); let script = host_script(&[rooted], &BTreeSet::new()); assert!(script.contains(&format!("d='/scratch/orx/runs/{RUN}';"))); @@ -278,6 +332,12 @@ mod tests { "docker exec 'abc' rm -rf -- '/root/.orx/runs/{RUN}'" ))); assert!(script.contains("s='/scratch/orx/source';")); - assert!(parse_outcomes(&[run(RUN, None)], "REMOVED other\n")[RUN] == Outcome::Failed); + let batch = [run(RUN, None, false), run("other", None, false)]; + let outcomes = parse_outcomes(&batch, &format!("FAILED {RUN} rm: Permission denied\n")); + assert_eq!( + outcomes[RUN], + Outcome::Failed("rm: Permission denied".into()) + ); + assert!(matches!(outcomes["other"], Outcome::Failed(_))); } } diff --git a/src/local/remote_cleanup.rs b/src/local/remote_cleanup.rs index 2979212ed..5c1be6add 100644 --- a/src/local/remote_cleanup.rs +++ b/src/local/remote_cleanup.rs @@ -19,8 +19,9 @@ pub struct HostBatch { keep_digests: BTreeSet, } -/// Groups due entries by host. Entries that need no cleanup leave the queue; ones orx -/// cannot safely clean leave it too and come back as (run, reason). +/// Groups due entries by host. The keep setting is read now, so it covers runs finishing while +/// it is on. Entries that need no cleanup leave the queue; ones orx cannot safely clean leave it +/// too and come back as failures. pub fn plan(store: &Store) -> Result<(Vec, Failures)> { let ssh = crate::config::ssh_settings()?; let slurm_keeps = crate::jobs::slurm::load_settings()? @@ -38,7 +39,10 @@ pub fn plan(store: &Store) -> Result<(Vec, Failures)> { .and_then(|host| ssh.hosts.get(host)) .is_some_and(|saved| saved.keep_remote_runs), }; - Ok(RemoteRun::from_descriptor(&entry.run_id, &descriptor)?.zip(Some(keeps))) + Ok( + RemoteRun::from_descriptor(&entry.run_id, &descriptor, entry.deleted)? + .map(|run| (run, keeps)), + ) }); match parsed { Ok(Some((run, keeps))) if entry.deleted || !keeps => { @@ -55,7 +59,7 @@ pub fn plan(store: &Store) -> Result<(Vec, Failures)> { .into_iter() .map(|(host, runs)| { Ok(HostBatch { - keep_digests: store.active_remote_source_digests(&host)?, + keep_digests: active_source_digests(store, &host)?, host, runs, }) @@ -64,18 +68,34 @@ pub fn plan(store: &Store) -> Result<(Vec, Failures)> { Ok((batches, rejected)) } -pub type Results = Vec<(String, std::result::Result)>; +/// Source digests that unfinished SSH/Slurm runs on `host` may still stage. +fn active_source_digests(store: &Store, host: &str) -> Result> { + let mut digests = BTreeSet::new(); + for run in store.list_active_runs()? { + let Ok(descriptor) = BackendDescriptor::parse(&run.backend_json) else { + continue; + }; + if matches!(descriptor.kind.as_str(), "ssh_job" | "slurm_job") + && descriptor.namespace.as_deref() == Some(host) + { + digests.extend(descriptor.source_digest.filter(|digest| { + digest.len() == 64 && digest.bytes().all(|c| c.is_ascii_hexdigit()) + })); + } + } + Ok(digests) +} -pub async fn clean(batches: Vec) -> Results { +pub async fn clean(batches: Vec) -> Vec<(String, Outcome)> { let mut results = Vec::new(); for batch in batches { match cleanup::clean_host(&batch.host, &batch.runs, &batch.keep_digests).await { - Ok(outcomes) => results.extend(outcomes.into_iter().map(|(id, o)| (id, Ok(o)))), + Ok(outcomes) => results.extend(outcomes), Err(error) => results.extend( batch .runs .iter() - .map(|run| (run.run_id.clone(), Err(format!("{error:#}")))), + .map(|run| (run.run_id.clone(), Outcome::Failed(format!("{error:#}")))), ), } } @@ -83,23 +103,19 @@ pub async fn clean(batches: Vec) -> Results { } /// Records outcomes and returns failures worth reporting; failed and busy runs retry later. -pub fn record(store: &Store, results: Results) -> Result { +pub fn record(store: &Store, results: Vec<(String, Outcome)>) -> Result { let mut failures = Vec::new(); - for (run_id, result) in results { - let error = match result { - Ok(Outcome::Removed) => { - store.finish_remote_cleanup(&run_id)?; - continue; + for (run_id, outcome) in results { + match outcome { + Outcome::Removed => store.finish_remote_cleanup(&run_id)?, + Outcome::Busy => { + store.defer_remote_cleanup(&run_id, "its process group is still running")? } - Ok(Outcome::Busy) => { - store.defer_remote_cleanup(&run_id, "its process group is still running")?; - continue; + Outcome::Failed(error) => { + store.defer_remote_cleanup(&run_id, &error)?; + failures.push((run_id, error)); } - Ok(Outcome::Failed) => "the remote dir could not be removed".to_string(), - Err(error) => error, - }; - store.defer_remote_cleanup(&run_id, &error)?; - failures.push((run_id, error)); + } } Ok(failures) } diff --git a/src/store.rs b/src/store.rs index 3e69ebdd3..b9f4325ab 100644 --- a/src/store.rs +++ b/src/store.rs @@ -515,13 +515,26 @@ pub(crate) struct RemoteCleanup { pub deleted: bool, } -const QUEUE_REMOTE_CLEANUP: &str = - "INSERT INTO remote_cleanups (run_id, backend_json, deleted, next_attempt_at)"; -// Appended to a SELECT's WHERE, which SQLite needs before ON CONFLICT; deletion wins over a pending finish. -const UPSERT_REMOTE_CLEANUP: &str = "AND CASE WHEN json_valid(backend_json) - THEN json_extract(backend_json, '$.kind') END IN ('ssh_job', 'slurm_job') - ON CONFLICT(run_id) DO UPDATE SET backend_json = excluded.backend_json, - deleted = MAX(deleted, excluded.deleted), next_attempt_at = excluded.next_attempt_at"; +/// Queues the SSH/Slurm runs whose `column` equals `value`; deletion wins over a pending finish. +fn queue_remote_cleanups( + conn: &Connection, + column: &str, + value: &str, + deleted: bool, +) -> Result<()> { + conn.execute( + &format!( + "INSERT INTO remote_cleanups (run_id, backend_json, deleted, next_attempt_at) + SELECT id, backend_json, ?2, ?3 FROM runs + WHERE {column} = ?1 AND CASE WHEN json_valid(backend_json) + THEN json_extract(backend_json, '$.kind') END IN ('ssh_job', 'slurm_job') + ON CONFLICT(run_id) DO UPDATE SET backend_json = excluded.backend_json, + deleted = MAX(deleted, excluded.deleted), next_attempt_at = excluded.next_attempt_at" + ), + params![value, deleted, now_ms()], + )?; + Ok(()) +} /// Failed remote cleanups retry after 5 minutes, doubling up to a day. const REMOTE_CLEANUP_RETRY_MS: i64 = 5 * 60 * 1000; const REMOTE_CLEANUP_MAX_RETRY_MS: i64 = 24 * 60 * 60 * 1000; @@ -2149,11 +2162,7 @@ impl Store { /// Queues a finished run's remote dir for removal; non-SSH/Slurm runs are ignored. pub(crate) fn queue_finished_remote_cleanup(&self, run_id: &str) -> Result<()> { - self.conn.execute( - &format!("{QUEUE_REMOTE_CLEANUP} SELECT id, backend_json, 0, ?2 FROM runs WHERE id = ?1 {UPSERT_REMOTE_CLEANUP}"), - params![run_id, now_ms()], - )?; - Ok(()) + queue_remote_cleanups(&self.conn, "id", run_id, false) } pub(crate) fn due_remote_cleanups(&self, limit: usize) -> Result> { @@ -2195,24 +2204,6 @@ impl Store { Ok(()) } - /// Source digests that unfinished SSH/Slurm runs on `host` may still stage. - pub(crate) fn active_remote_source_digests( - &self, - host: &str, - ) -> Result> { - let mut digests = std::collections::BTreeSet::new(); - for run in self.list_active_runs()? { - if let Ok(descriptor) = crate::jobs::BackendDescriptor::parse(&run.backend_json) { - if matches!(descriptor.kind.as_str(), "ssh_job" | "slurm_job") - && descriptor.namespace.as_deref() == Some(host) - { - digests.extend(descriptor.source_digest); - } - } - } - Ok(digests) - } - fn data_dir_move_claimed(&self) -> Result { self.clear_stale_data_dir_move_lease()?; Ok(self.conn.query_row( @@ -2274,10 +2265,7 @@ impl Store { "DELETE FROM chat_sessions WHERE project_id = ?1", params![id], )?; - tx.execute( - &format!("{QUEUE_REMOTE_CLEANUP} SELECT id, backend_json, 1, ?2 FROM runs WHERE project_id = ?1 {UPSERT_REMOTE_CLEANUP}"), - params![id, now_ms()], - )?; + queue_remote_cleanups(&tx, "project_id", id, true)?; self.conn .execute("DELETE FROM runs WHERE project_id = ?1", params![id])?; self.conn.execute( @@ -4292,10 +4280,6 @@ mod tests { .collect::>() }; assert_eq!(due(&store), [("ssh".to_string(), false)]); - assert_eq!( - store.active_remote_source_digests("lab").unwrap(), - std::collections::BTreeSet::from(["live".to_string()]) - ); store.defer_remote_cleanup("ssh", "unreachable").unwrap(); assert!(due(&store).is_empty()); store From 642a4e16dc6b069481924d6652387e42ef8c7df9 Mon Sep 17 00:00:00 2001 From: Myles Anderson Date: Thu, 22 Oct 2026 15:13:23 -0700 Subject: [PATCH 06/10] Address round-2 review nits for remote cleanup Co-Authored-By: Claude Opus 5.5 --- agent-skills/orx-compute/references/slurm.md | 10 ++++----- agent-skills/orx-compute/references/ssh.md | 14 ++++++------ src/commands/compute.rs | 2 +- src/commands/up.rs | 15 +++++++------ src/jobs/ssh/cleanup.rs | 23 +++++++++----------- src/store.rs | 6 +++-- 6 files changed, 35 insertions(+), 35 deletions(-) diff --git a/agent-skills/orx-compute/references/slurm.md b/agent-skills/orx-compute/references/slurm.md index 889a2d433..3cb4e1bdd 100644 --- a/agent-skills/orx-compute/references/slurm.md +++ b/agent-skills/orx-compute/references/slurm.md @@ -28,11 +28,11 @@ orx exp run --backend slurm be on a filesystem the compute nodes share. Each run keeps the root it was submitted with; `--clear remote-root` restores the default for new runs. - Once a run finishes and its log is collected, the running OpenResearch app - removes its remote run dir except `tensorboard/`, along with orx's cached - source archives that no active run uses and that have been idle for an hour. - Anything the job wrote into `repo/` goes with it. To keep results on the - cluster, set `orx compute configure slurm --keep-remote-runs true` before - runs finish, or write outputs outside the run dir. Deleting a project + removes its remote run dir except `tensorboard/`, and while there, orx's + cached source archives that no active run uses and that have been idle for an + hour. Anything the job wrote into `repo/` goes with it. To keep results on + the cluster, set `orx compute configure slurm --keep-remote-runs true` + (it covers runs not yet cleaned up), or write outputs outside the run dir. Deleting a project removes its runs' remote dirs entirely regardless. - A detached `orx supervise` process records scheduler status and logs; do not kill it. diff --git a/agent-skills/orx-compute/references/ssh.md b/agent-skills/orx-compute/references/ssh.md index 71f8180bf..6a7aa0e8d 100644 --- a/agent-skills/orx-compute/references/ssh.md +++ b/agent-skills/orx-compute/references/ssh.md @@ -36,7 +36,6 @@ Machine-local SSH settings can save a container name or ID for each host, and an absolute remote storage root (`--remote-root /data/alice/orx`) used instead of `~/.orx` for that host's runs and cached source archives. Each run keeps the root it was launched with. - A saved default SSH host lets launches omit `--host`. These defaults apply only to SSH experiments; dashboards, terminals and coding agents still connect directly to the host. @@ -45,12 +44,13 @@ agents still connect directly to the host. overrides it; `--no-container` explicitly runs on the host. The flags conflict. Once a run finishes and its log is collected, the running OpenResearch app -removes the run's remote dir except `tensorboard/`, so anything written into -`repo/` goes with it. It also removes orx's cached source archives that no -active run uses and that have been idle for an hour. Unreachable hosts are -retried later. `orx compute configure ssh --host lab --keep-remote-runs true` -keeps that host's runs that finish while it is set. Deleting a project removes -its runs' remote dirs entirely, regardless of that setting. +removes the run's host dir except `tensorboard/` (a container's copy goes +entirely), so anything written into `repo/` goes with it. While cleaning a +host it also removes orx's cached source archives that no active run uses and +that have been idle for an hour. Unreachable hosts are retried later. +`orx compute configure ssh --host lab --keep-remote-runs true` keeps that +host's runs not yet cleaned up. Deleting a project removes its runs' remote +dirs entirely, regardless of that setting. Keep environment preparation in the project's committed setup/run scripts and invoke them through the fixed run command, just as with other backends. For diff --git a/src/commands/compute.rs b/src/commands/compute.rs index fa6ebd53d..370680071 100644 --- a/src/commands/compute.rs +++ b/src/commands/compute.rs @@ -451,7 +451,7 @@ impl ConfigureArgs { } if backend == "ssh" && host_setting && !body.contains_key("host") { return Err(anyhow!( - "A container or remote root change requires --host." + "A container, remote root, or keep-remote-runs change requires --host." )); } Ok(Value::Object(body)) diff --git a/src/commands/up.rs b/src/commands/up.rs index 44fb81007..37fef7458 100644 --- a/src/commands/up.rs +++ b/src/commands/up.rs @@ -4726,16 +4726,17 @@ async fn reclaim_remote_run_dirs(moving: Arc, gate: Arc| { - for (run_id, error) in failures { - if reported.insert(run_id.clone()) { - eprintln!("orx up: remote cleanup for run {run_id}: {error}"); - } + let mut report = |failures: Vec<(String, String)>| { + for (run_id, error) in failures { + if reported.insert(run_id.clone()) { + eprintln!("orx up: remote cleanup for run {run_id}: {error}"); } - }; + } + }; + loop { match with_gated_store(&moving, &gate, local::remote_cleanup::plan).await { Some(Ok((batches, rejected))) => { + plan_error = None; report(rejected); if !batches.is_empty() { let results = local::remote_cleanup::clean(batches).await; diff --git a/src/jobs/ssh/cleanup.rs b/src/jobs/ssh/cleanup.rs index 5040578a1..de33e9086 100644 --- a/src/jobs/ssh/cleanup.rs +++ b/src/jobs/ssh/cleanup.rs @@ -92,7 +92,7 @@ fn host_script(runs: &[RemoteRun], keep_digests: &BTreeSet) -> String { .as_ref() .map_or_else(String::new, |(id, dir)| { format!( - "docker exec {} rm -rf -- {} >/dev/null 2>&1; ", + "docker exec {} rm -rf -- {} /dev/null 2>&1; ", sh_quote(id), sh_quote(dir) ) @@ -122,7 +122,7 @@ fn host_script(runs: &[RemoteRun], keep_digests: &BTreeSet) -> String { script.push_str(&format!( "s={source}; if [ -d \"$s\" ]; then \ find \"$s\" -maxdepth 1 -type f -mmin +{SOURCE_IDLE_MINUTES} | while IFS= read -r f; do \ - n=${{f##*/}}; b=${{n%%.tar*}}; \ + case \"$f\" in \"$s\"/*) ;; *) continue;; esac; n=${{f##*/}}; b=${{n%%.tar*}}; \ case \"$b\" in ''|*[!0-9a-f]*) continue;; esac; [ ${{#b}} -eq 64 ] || continue; \ case \"$n\" in \"$b.tar\"|\"$b.tar.tmp.\"*) ;; *) continue;; esac; \ case ' {keep} ' in *\" $b \"*) continue;; esac; \ @@ -138,19 +138,16 @@ fn parse_outcomes(runs: &[RemoteRun], output: &str) -> BTreeMap let mut reported = BTreeMap::new(); for line in output.lines() { let mut words = line.splitn(3, ' '); - let outcome = match words.next() { - Some("REMOVED") => Outcome::Removed, - Some("BUSY") => Outcome::Busy, - Some("FAILED") => Outcome::Failed(String::new()), + let (Some(tag), Some(id)) = (words.next(), words.next()) else { + continue; + }; + let outcome = match tag { + "REMOVED" => Outcome::Removed, + "BUSY" => Outcome::Busy, + "FAILED" => Outcome::Failed(words.next().unwrap_or_default().into()), _ => continue, }; - if let Some(id) = words.next() { - let outcome = match outcome { - Outcome::Failed(_) => Outcome::Failed(words.next().unwrap_or_default().into()), - outcome => outcome, - }; - reported.insert(id.to_string(), outcome); - } + reported.insert(id.to_string(), outcome); } runs.iter() .map(|run| { diff --git a/src/store.rs b/src/store.rs index b9f4325ab..1a110cd6a 100644 --- a/src/store.rs +++ b/src/store.rs @@ -518,7 +518,7 @@ pub(crate) struct RemoteCleanup { /// Queues the SSH/Slurm runs whose `column` equals `value`; deletion wins over a pending finish. fn queue_remote_cleanups( conn: &Connection, - column: &str, + column: &'static str, value: &str, deleted: bool, ) -> Result<()> { @@ -529,12 +529,14 @@ fn queue_remote_cleanups( WHERE {column} = ?1 AND CASE WHEN json_valid(backend_json) THEN json_extract(backend_json, '$.kind') END IN ('ssh_job', 'slurm_job') ON CONFLICT(run_id) DO UPDATE SET backend_json = excluded.backend_json, - deleted = MAX(deleted, excluded.deleted), next_attempt_at = excluded.next_attempt_at" + deleted = MAX(deleted, excluded.deleted), attempts = 0, + next_attempt_at = excluded.next_attempt_at" ), params![value, deleted, now_ms()], )?; Ok(()) } + /// Failed remote cleanups retry after 5 minutes, doubling up to a day. const REMOTE_CLEANUP_RETRY_MS: i64 = 5 * 60 * 1000; const REMOTE_CLEANUP_MAX_RETRY_MS: i64 = 24 * 60 * 60 * 1000; From af548a46211962b6db78571d26bfb91f0417f004 Mon Sep 17 00:00:00 2001 From: Myles Anderson Date: Thu, 22 Oct 2026 16:36:06 -0700 Subject: [PATCH 07/10] Keep each project's newest remote runs, matching local retention Co-Authored-By: Claude Opus 5.5 --- .../orx-compute/references/configuration.md | 4 +- agent-skills/orx-compute/references/slurm.md | 13 +-- agent-skills/orx-compute/references/ssh.md | 11 +- src/commands/compute.rs | 2 +- src/commands/supervise.rs | 3 +- src/config.rs | 2 +- src/jobs/slurm.rs | 2 +- src/store.rs | 101 +++++++++++------- 8 files changed, 82 insertions(+), 56 deletions(-) diff --git a/agent-skills/orx-compute/references/configuration.md b/agent-skills/orx-compute/references/configuration.md index e340f8bfe..81c38199b 100644 --- a/agent-skills/orx-compute/references/configuration.md +++ b/agent-skills/orx-compute/references/configuration.md @@ -12,7 +12,7 @@ orx compute configure ssh --host lab --container research orx compute configure ssh --host lab --clear container orx compute configure ssh --host lab --remote-root /data/me/orx # runs + source cache; default ~/.orx orx compute configure ssh --host lab --clear remote-root -orx compute configure ssh --host lab --keep-remote-runs true # keep finished run dirs on the host +orx compute configure ssh --host lab --keep-remote-runs true # keep every finished run dir on the host orx compute test ssh --host lab # tests saved container too orx compute test ssh --host lab --no-container # explicitly tests host orx compute connect ssh --host lab # interactive login/MFA @@ -21,7 +21,7 @@ orx compute configure slurm --cpus-per-task 8 --mem 64G # defaults; exp run -- orx compute configure slurm --clear cpus-per-task --clear mem orx compute configure slurm --clear time-limit # cluster chooses orx compute configure slurm --remote-root /scratch/me/orx # shared FS; --clear remote-root for ~/.orx -orx compute configure slurm --keep-remote-runs true # keep finished run dirs on the cluster +orx compute configure slurm --keep-remote-runs true # keep every finished run dir on the cluster orx compute connect slurm --host cluster orx compute configure k8s --context research --namespace experiments orx compute configure ray --address http://127.0.0.1:8265 diff --git a/agent-skills/orx-compute/references/slurm.md b/agent-skills/orx-compute/references/slurm.md index 3cb4e1bdd..ee6f36693 100644 --- a/agent-skills/orx-compute/references/slurm.md +++ b/agent-skills/orx-compute/references/slurm.md @@ -27,12 +27,13 @@ orx exp run --backend slurm absolute directory, such as scratch space when home quota is small. It must be on a filesystem the compute nodes share. Each run keeps the root it was submitted with; `--clear remote-root` restores the default for new runs. -- Once a run finishes and its log is collected, the running OpenResearch app +- Each project's three newest SSH and Slurm runs stay whole on the cluster. + Once an older finished run's log is collected, the running OpenResearch app removes its remote run dir except `tensorboard/`, and while there, orx's - cached source archives that no active run uses and that have been idle for an - hour. Anything the job wrote into `repo/` goes with it. To keep results on - the cluster, set `orx compute configure slurm --keep-remote-runs true` - (it covers runs not yet cleaned up), or write outputs outside the run dir. Deleting a project - removes its runs' remote dirs entirely regardless. + cached source archives that no active run uses and that have been idle for + an hour. Anything the job wrote into `repo/` goes with it. To keep every run + on the cluster, set `orx compute configure slurm --keep-remote-runs true` + (it covers runs not yet cleaned up), or write outputs outside the run dir. + Deleting a project removes its runs' remote dirs entirely regardless. - A detached `orx supervise` process records scheduler status and logs; do not kill it. diff --git a/agent-skills/orx-compute/references/ssh.md b/agent-skills/orx-compute/references/ssh.md index 6a7aa0e8d..a18e90289 100644 --- a/agent-skills/orx-compute/references/ssh.md +++ b/agent-skills/orx-compute/references/ssh.md @@ -43,13 +43,14 @@ agents still connect directly to the host. - Omitted container selection uses the host's saved target. `--container` overrides it; `--no-container` explicitly runs on the host. The flags conflict. -Once a run finishes and its log is collected, the running OpenResearch app -removes the run's host dir except `tensorboard/` (a container's copy goes -entirely), so anything written into `repo/` goes with it. While cleaning a +Each project's three newest SSH and Slurm runs stay whole on their hosts. Once +an older finished run's log is collected, the running OpenResearch app removes +the run's host dir except `tensorboard/` (a container's copy goes entirely), +so anything written into `repo/` goes with it. While cleaning a host it also removes orx's cached source archives that no active run uses and that have been idle for an hour. Unreachable hosts are retried later. -`orx compute configure ssh --host lab --keep-remote-runs true` keeps that -host's runs not yet cleaned up. Deleting a project removes its runs' remote +`orx compute configure ssh --host lab --keep-remote-runs true` keeps all of +that host's runs not yet cleaned up. Deleting a project removes its runs' remote dirs entirely, regardless of that setting. Keep environment preparation in the project's committed setup/run scripts and diff --git a/src/commands/compute.rs b/src/commands/compute.rs index 370680071..05ec004cd 100644 --- a/src/commands/compute.rs +++ b/src/commands/compute.rs @@ -238,7 +238,7 @@ pub struct ConfigureArgs { /// Absolute remote dir for runs and the source cache (SSH per host, or Slurm); default ~/.orx. #[arg(long)] remote_root: Option, - /// true keeps finished runs' remote dirs; false (the default) removes them once logs are collected. + /// true keeps every finished run's remote dir; false (the default) keeps each project's newest few. #[arg(long)] keep_remote_runs: Option, #[arg(long)] diff --git a/src/commands/supervise.rs b/src/commands/supervise.rs index 19f41a782..d11afc6c4 100644 --- a/src/commands/supervise.rs +++ b/src/commands/supervise.rs @@ -1472,7 +1472,8 @@ async fn run_slurm( } } -/// Once logs are local, `orx up` removes the remote run dir unless the target keeps results. +/// Once logs are local, `orx up` removes the remote run dir when it is past its project's newest +/// few, unless the target keeps results. fn queue_remote_cleanup(store: &Store, run_id: &str) { if let Err(err) = store.queue_finished_remote_cleanup(run_id) { eprintln!("supervise {run_id}: could not queue remote cleanup: {err}"); diff --git a/src/config.rs b/src/config.rs index 63c9b3cb0..145572035 100644 --- a/src/config.rs +++ b/src/config.rs @@ -464,7 +464,7 @@ pub struct SshHostSettings { /// Absolute storage root for runs and the source cache; `None` is `~/.orx`. #[serde(default, skip_serializing_if = "Option::is_none")] pub remote_root: Option, - /// Keep finished runs' remote dirs instead of removing them once logs are collected. + /// Keep every finished run's remote dir, not just each project's newest few. #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub keep_remote_runs: bool, } diff --git a/src/jobs/slurm.rs b/src/jobs/slurm.rs index fcfa6ff92..ecc09399b 100644 --- a/src/jobs/slurm.rs +++ b/src/jobs/slurm.rs @@ -61,7 +61,7 @@ pub struct SlurmSettings { /// Absolute shared-filesystem root for runs and the source cache; `None` is `~/.orx`. #[serde(default, skip_serializing_if = "Option::is_none")] pub remote_root: Option, - /// Keep finished runs' remote dirs instead of removing them once logs are collected. + /// Keep every finished run's remote dir, not just each project's newest few. #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub keep_remote_runs: bool, } diff --git a/src/store.rs b/src/store.rs index 8456c8e18..e9a43cbe3 100644 --- a/src/store.rs +++ b/src/store.rs @@ -477,8 +477,22 @@ pub(crate) const RUN_ARTIFACTS_WAIT: Duration = Duration::from_secs(1); pub(crate) const RUN_ARTIFACTS_POLL: Duration = Duration::from_millis(250); /// Bounds removal time per orphan pass; classification is unbudgeted but cheap for finished runs. const ORPHAN_PASS_BUDGET: Duration = Duration::from_secs(5); -/// Finished runs older than each project's newest few local runs lose their rebuildable environments. +/// Finished runs older than each project's newest few lose their rebuildable environments +/// (local runs) or their remote dirs (SSH and Slurm runs). const FULL_RUNS_PER_PROJECT: i64 = 3; + +/// `(id, status)` of runs of `kinds` past each project's newest [`FULL_RUNS_PER_PROJECT`]. +fn runs_past_newest(kinds: &str) -> String { + format!( + "SELECT id, status FROM ( + SELECT id, status, ROW_NUMBER() OVER + (PARTITION BY project_id ORDER BY created_at DESC, id DESC) AS position + FROM runs + WHERE CASE WHEN json_valid(backend_json) THEN json_extract(backend_json, '$.kind') END + IN ({kinds})) + WHERE position > {FULL_RUNS_PER_PROJECT}" + ) +} const SLIMMED_RUN_PATHS: [&[&str]; 2] = [&["repo", ".venv"], &["repo", ".cache", "uv"]]; const SUBMISSION_HANDLE: (&str, &str) = ("submission-handles", ".json"); @@ -1921,17 +1935,11 @@ impl Store { started: std::time::Instant, pass: &mut OrphanPass, ) -> Result<()> { - let mut statement = self.conn.prepare( - "SELECT id, status FROM ( - SELECT id, status, ROW_NUMBER() OVER - (PARTITION BY project_id ORDER BY created_at DESC, id DESC) AS position - FROM runs - WHERE CASE WHEN json_valid(backend_json) THEN json_extract(backend_json, '$.kind') END - IN ('local_job', 'tinker_job')) - WHERE position > ?1", - )?; + let mut statement = self + .conn + .prepare(&runs_past_newest("'local_job', 'tinker_job'"))?; let runs = statement - .query_map(params![FULL_RUNS_PER_PROJECT], |row| { + .query_map([], |row| { Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?)) })? .collect::, _>>()?; @@ -2166,11 +2174,15 @@ impl Store { queue_remote_cleanups(&self.conn, "id", run_id, false) } + /// Due entries; a finished run waits until newer runs push it past its project's newest few. pub(crate) fn due_remote_cleanups(&self, limit: usize) -> Result> { - let mut statement = self.conn.prepare( + let mut statement = self.conn.prepare(&format!( "SELECT run_id, backend_json, deleted FROM remote_cleanups - WHERE next_attempt_at <= ?1 ORDER BY next_attempt_at LIMIT ?2", - )?; + WHERE next_attempt_at <= ?1 AND (deleted = 1 OR run_id IN + (SELECT id FROM ({}))) + ORDER BY next_attempt_at LIMIT ?2", + runs_past_newest("'ssh_job', 'slurm_job'") + ))?; let rows = statement.query_map(params![now_ms(), limit as i64], |row| { Ok(RemoteCleanup { run_id: row.get(0)?, @@ -4252,55 +4264,66 @@ mod tests { } #[test] - fn remote_cleanups_queue_on_finish_and_project_delete_then_back_off() { + fn remote_cleanups_keep_each_projects_newest_runs_until_project_delete() { let dir = std::env::temp_dir().join(format!("orx-remote-cleanup-{}", uuid::Uuid::new_v4())); let store = Store::open_at(dir.clone()).unwrap(); project_with_run_artifacts(&store, "p", "local", &"d".repeat(64)); - let insert = |id: &str, status: &str, kind: &str| { + let insert = |id: &str, status: &str, kind: &str, created_at: i64| { let descriptor = serde_json::json!({"kind": kind, "namespace": "lab", "sourceDigest": id}); store.conn.execute( "INSERT INTO runs (id, experiment_id, project_id, status, backend_json, created_at, updated_at) - VALUES (?1, 'experiment', 'p', ?2, ?3, 1, 1)", - params![id, status, descriptor.to_string()], + VALUES (?1, 'experiment', 'p', ?2, ?3, ?4, ?4)", + params![id, status, descriptor.to_string(), created_at], ).unwrap(); }; - insert("ssh", "done", "ssh_job"); - insert("slurm", "failed", "slurm_job"); - insert("box", "done", "openresearch_job"); - insert("live", "running", "ssh_job"); - for id in ["ssh", "box", "local"] { - store.queue_finished_remote_cleanup(id).unwrap(); - } let due = |store: &Store| { - store + let mut entries = store .due_remote_cleanups(10) .unwrap() .into_iter() .map(|entry| (entry.run_id, entry.deleted)) + .collect::>(); + entries.sort(); + entries + }; + let finished = |ids: &[&str]| { + ids.iter() + .map(|id| (id.to_string(), false)) .collect::>() }; - assert_eq!(due(&store), [("ssh".to_string(), false)]); - store.defer_remote_cleanup("ssh", "unreachable").unwrap(); + insert("r1", "done", "ssh_job", 1); + insert("r2", "failed", "slurm_job", 2); + insert("r3", "done", "ssh_job", 3); + insert("box", "done", "openresearch_job", 4); + for id in ["r1", "r2", "r3", "box", "local"] { + store.queue_finished_remote_cleanup(id).unwrap(); + } + // The newest three remote runs stay whole; local and other kinds don't count. assert!(due(&store).is_empty()); + insert("r4", "done", "ssh_job", 5); + store.queue_finished_remote_cleanup("r4").unwrap(); + assert_eq!(due(&store), finished(&["r1"])); + // As locally, an unfinished newer run takes one of the slots. + insert("r5", "running", "slurm_job", 6); + assert_eq!(due(&store), finished(&["r1", "r2"])); + store.defer_remote_cleanup("r1", "unreachable").unwrap(); + assert_eq!(due(&store), finished(&["r2"])); store .conn .execute("UPDATE runs SET status = 'done'", []) .unwrap(); store.delete_local_project("p").unwrap(); - let mut entries = due(&store); - entries.sort(); - // Deletion marks every SSH/Slurm run, including one already waiting to retry. + // Deletion marks every SSH/Slurm run, including the newest and one waiting to retry. assert_eq!( - entries, - [ - ("live".to_string(), true), - ("slurm".to_string(), true), - ("ssh".to_string(), true) - ] + due(&store), + ["r1", "r2", "r3", "r4", "r5"] + .iter() + .map(|id| (id.to_string(), true)) + .collect::>() ); - store.finish_remote_cleanup("ssh").unwrap(); - assert_eq!(due(&store).len(), 2); + store.finish_remote_cleanup("r1").unwrap(); + assert_eq!(due(&store).len(), 4); std::fs::remove_dir_all(dir).unwrap(); } From b61a910a13b9ba7c6b54069e3b440700b7ae000c Mon Sep 17 00:00:00 2001 From: Myles Anderson Date: Thu, 22 Oct 2026 16:39:05 -0700 Subject: [PATCH 08/10] Clarify shared remote run slots and tighten the retention test Co-Authored-By: Claude Opus 5.5 --- agent-skills/orx-compute/references/slurm.md | 10 +++---- agent-skills/orx-compute/references/ssh.md | 17 ++++++------ src/commands/compute.rs | 3 +- src/commands/supervise.rs | 4 +-- src/local/remote_cleanup.rs | 4 +-- src/store.rs | 29 ++++++++++++-------- 6 files changed, 37 insertions(+), 30 deletions(-) diff --git a/agent-skills/orx-compute/references/slurm.md b/agent-skills/orx-compute/references/slurm.md index ee6f36693..8a436a9c5 100644 --- a/agent-skills/orx-compute/references/slurm.md +++ b/agent-skills/orx-compute/references/slurm.md @@ -27,11 +27,11 @@ orx exp run --backend slurm absolute directory, such as scratch space when home quota is small. It must be on a filesystem the compute nodes share. Each run keeps the root it was submitted with; `--clear remote-root` restores the default for new runs. -- Each project's three newest SSH and Slurm runs stay whole on the cluster. - Once an older finished run's log is collected, the running OpenResearch app - removes its remote run dir except `tensorboard/`, and while there, orx's - cached source archives that no active run uses and that have been idle for - an hour. Anything the job wrote into `repo/` goes with it. To keep every run +- Each project's three newest SSH and Slurm runs, counted together across + hosts and including unfinished ones, stay whole. Once an older finished + run's log is collected, the running OpenResearch app removes its remote run + dir except `tensorboard/`, and while there, orx's cached source archives + that no active run uses and that have been idle for an hour. Anything the job wrote into `repo/` goes with it. To keep every run on the cluster, set `orx compute configure slurm --keep-remote-runs true` (it covers runs not yet cleaned up), or write outputs outside the run dir. Deleting a project removes its runs' remote dirs entirely regardless. diff --git a/agent-skills/orx-compute/references/ssh.md b/agent-skills/orx-compute/references/ssh.md index a18e90289..f89c3f9ba 100644 --- a/agent-skills/orx-compute/references/ssh.md +++ b/agent-skills/orx-compute/references/ssh.md @@ -43,15 +43,16 @@ agents still connect directly to the host. - Omitted container selection uses the host's saved target. `--container` overrides it; `--no-container` explicitly runs on the host. The flags conflict. -Each project's three newest SSH and Slurm runs stay whole on their hosts. Once -an older finished run's log is collected, the running OpenResearch app removes -the run's host dir except `tensorboard/` (a container's copy goes entirely), -so anything written into `repo/` goes with it. While cleaning a -host it also removes orx's cached source archives that no active run uses and -that have been idle for an hour. Unreachable hosts are retried later. +Each project's three newest SSH and Slurm runs, counted together across hosts +and including unfinished ones, stay whole. Once an older finished run's log is +collected, the running OpenResearch app removes its host dir except +`tensorboard/` (a container's copy goes entirely), so anything written into +`repo/` goes with it. While cleaning a host it also removes orx's cached source +archives that no active run uses and that have been idle for an hour. +Unreachable hosts are retried later. `orx compute configure ssh --host lab --keep-remote-runs true` keeps all of -that host's runs not yet cleaned up. Deleting a project removes its runs' remote -dirs entirely, regardless of that setting. +that host's runs not yet cleaned up. Deleting a project removes its runs' +remote dirs entirely, regardless of that setting. Keep environment preparation in the project's committed setup/run scripts and invoke them through the fixed run command, just as with other backends. For diff --git a/src/commands/compute.rs b/src/commands/compute.rs index 05ec004cd..e59eaa394 100644 --- a/src/commands/compute.rs +++ b/src/commands/compute.rs @@ -238,7 +238,8 @@ pub struct ConfigureArgs { /// Absolute remote dir for runs and the source cache (SSH per host, or Slurm); default ~/.orx. #[arg(long)] remote_root: Option, - /// true keeps every finished run's remote dir; false (the default) keeps each project's newest few. + /// true keeps every finished run's remote dir; false (the default) keeps each project's 3 newest + /// SSH/Slurm runs and removes older ones except tensorboard/. #[arg(long)] keep_remote_runs: Option, #[arg(long)] diff --git a/src/commands/supervise.rs b/src/commands/supervise.rs index d11afc6c4..d4d2df4ca 100644 --- a/src/commands/supervise.rs +++ b/src/commands/supervise.rs @@ -1472,8 +1472,8 @@ async fn run_slurm( } } -/// Once logs are local, `orx up` removes the remote run dir when it is past its project's newest -/// few, unless the target keeps results. +/// Queues the remote run dir once logs are local; `orx up` removes it when it falls past its +/// project's newest few, unless the target keeps results. fn queue_remote_cleanup(store: &Store, run_id: &str) { if let Err(err) = store.queue_finished_remote_cleanup(run_id) { eprintln!("supervise {run_id}: could not queue remote cleanup: {err}"); diff --git a/src/local/remote_cleanup.rs b/src/local/remote_cleanup.rs index 5c1be6add..c1cf29f0d 100644 --- a/src/local/remote_cleanup.rs +++ b/src/local/remote_cleanup.rs @@ -19,8 +19,8 @@ pub struct HostBatch { keep_digests: BTreeSet, } -/// Groups due entries by host. The keep setting is read now, so it covers runs finishing while -/// it is on. Entries that need no cleanup leave the queue; ones orx cannot safely clean leave it +/// Groups due entries by host. The keep setting is read when a run comes due, so it covers runs +/// not yet cleaned up. Entries that need no cleanup leave the queue; ones orx cannot safely clean leave it /// too and come back as failures. pub fn plan(store: &Store) -> Result<(Vec, Failures)> { let ssh = crate::config::ssh_settings()?; diff --git a/src/store.rs b/src/store.rs index e9a43cbe3..40558139c 100644 --- a/src/store.rs +++ b/src/store.rs @@ -481,18 +481,6 @@ const ORPHAN_PASS_BUDGET: Duration = Duration::from_secs(5); /// (local runs) or their remote dirs (SSH and Slurm runs). const FULL_RUNS_PER_PROJECT: i64 = 3; -/// `(id, status)` of runs of `kinds` past each project's newest [`FULL_RUNS_PER_PROJECT`]. -fn runs_past_newest(kinds: &str) -> String { - format!( - "SELECT id, status FROM ( - SELECT id, status, ROW_NUMBER() OVER - (PARTITION BY project_id ORDER BY created_at DESC, id DESC) AS position - FROM runs - WHERE CASE WHEN json_valid(backend_json) THEN json_extract(backend_json, '$.kind') END - IN ({kinds})) - WHERE position > {FULL_RUNS_PER_PROJECT}" - ) -} const SLIMMED_RUN_PATHS: [&[&str]; 2] = [&["repo", ".venv"], &["repo", ".cache", "uv"]]; const SUBMISSION_HANDLE: (&str, &str) = ("submission-handles", ".json"); @@ -529,6 +517,19 @@ pub(crate) struct RemoteCleanup { pub deleted: bool, } +/// `(id, status)` of runs of `kinds` past each project's newest [`FULL_RUNS_PER_PROJECT`]. +fn runs_past_newest(kinds: &'static str) -> String { + format!( + "SELECT id, status FROM ( + SELECT id, status, ROW_NUMBER() OVER + (PARTITION BY project_id ORDER BY created_at DESC, id DESC) AS position + FROM runs + WHERE CASE WHEN json_valid(backend_json) THEN json_extract(backend_json, '$.kind') END + IN ({kinds})) + WHERE position > {FULL_RUNS_PER_PROJECT}" + ) +} + /// Queues the SSH/Slurm runs whose `column` equals `value`; deletion wins over a pending finish. fn queue_remote_cleanups( conn: &Connection, @@ -4268,6 +4269,10 @@ mod tests { let dir = std::env::temp_dir().join(format!("orx-remote-cleanup-{}", uuid::Uuid::new_v4())); let store = Store::open_at(dir.clone()).unwrap(); project_with_run_artifacts(&store, "p", "local", &"d".repeat(64)); + store + .conn + .execute("UPDATE runs SET created_at = 10 WHERE id = 'local'", []) + .unwrap(); let insert = |id: &str, status: &str, kind: &str, created_at: i64| { let descriptor = serde_json::json!({"kind": kind, "namespace": "lab", "sourceDigest": id}); From fa544b78ccbc07e4946c88ad25e5e8dc32888137 Mon Sep 17 00:00:00 2001 From: Myles Anderson Date: Thu, 22 Oct 2026 18:05:19 -0700 Subject: [PATCH 09/10] Keep remote cleanup retrying for stopped containers and in-flight deletions; revisit source caches Co-Authored-By: Claude Opus 5.5 --- agent-skills/orx-compute/references/slurm.md | 5 +- agent-skills/orx-compute/references/ssh.md | 3 +- src/jobs/ssh/cleanup.rs | 100 ++++++++++++++--- src/local/remote_cleanup.rs | 111 +++++++++++++------ src/store.rs | 92 +++++++++++++-- 5 files changed, 249 insertions(+), 62 deletions(-) diff --git a/agent-skills/orx-compute/references/slurm.md b/agent-skills/orx-compute/references/slurm.md index 8a436a9c5..a08638ced 100644 --- a/agent-skills/orx-compute/references/slurm.md +++ b/agent-skills/orx-compute/references/slurm.md @@ -30,8 +30,9 @@ orx exp run --backend slurm - Each project's three newest SSH and Slurm runs, counted together across hosts and including unfinished ones, stay whole. Once an older finished run's log is collected, the running OpenResearch app removes its remote run - dir except `tensorboard/`, and while there, orx's cached source archives - that no active run uses and that have been idle for an hour. Anything the job wrote into `repo/` goes with it. To keep every run + dir except `tensorboard/`, and while there (and again an hour later), orx's + cached source archives that no active run uses and that have been idle for + an hour. Anything the job wrote into `repo/` goes with it. To keep every run on the cluster, set `orx compute configure slurm --keep-remote-runs true` (it covers runs not yet cleaned up), or write outputs outside the run dir. Deleting a project removes its runs' remote dirs entirely regardless. diff --git a/agent-skills/orx-compute/references/ssh.md b/agent-skills/orx-compute/references/ssh.md index f89c3f9ba..8dc536b34 100644 --- a/agent-skills/orx-compute/references/ssh.md +++ b/agent-skills/orx-compute/references/ssh.md @@ -48,7 +48,8 @@ and including unfinished ones, stay whole. Once an older finished run's log is collected, the running OpenResearch app removes its host dir except `tensorboard/` (a container's copy goes entirely), so anything written into `repo/` goes with it. While cleaning a host it also removes orx's cached source -archives that no active run uses and that have been idle for an hour. +archives that no active run uses and that have been idle for an hour, and +checks again an hour later. Unreachable hosts are retried later. `orx compute configure ssh --host lab --keep-remote-runs true` keeps all of that host's runs not yet cleaned up. Deleting a project removes its runs' diff --git a/src/jobs/ssh/cleanup.rs b/src/jobs/ssh/cleanup.rs index de33e9086..992127c26 100644 --- a/src/jobs/ssh/cleanup.rs +++ b/src/jobs/ssh/cleanup.rs @@ -72,6 +72,14 @@ impl RemoteRun { keep_tensorboard: !deleted, })) } + + pub(crate) fn root(&self) -> Option<&str> { + self.root.as_deref() + } + + pub(crate) fn deleted(&self) -> bool { + !self.keep_tensorboard + } } #[derive(Debug, Clone, PartialEq, Eq)] @@ -83,20 +91,26 @@ pub enum Outcome { } /// One script per host: remove each run dir whose process group is gone, then cached tars -/// that no active run references and nothing has touched recently. -fn host_script(runs: &[RemoteRun], keep_digests: &BTreeSet) -> String { +/// under `roots` that no active run references and nothing has touched recently. +fn host_script( + runs: &[RemoteRun], + roots: &BTreeSet>, + keep_digests: &BTreeSet, +) -> String { let mut script = format!("{HOST_PROCESS_HELPERS}\n"); for run in runs { - let inner = run - .container - .as_ref() - .map_or_else(String::new, |(id, dir)| { + // A container that still exists but can't run the removal (stopped) keeps the run queued. + let container_failed = run.container.as_ref().map_or_else( + || "false".to_string(), + |(id, dir)| { format!( - "docker exec {} rm -rf -- {} /dev/null 2>&1; ", - sh_quote(id), - sh_quote(dir) + "docker container inspect -- {id} /dev/null 2>&1 && \ + ! cerr=$(docker exec {id} rm -rf -- {dir} &1)", + id = sh_quote(id), + dir = sh_quote(dir) ) - }); + }, + ); let remove = if run.keep_tensorboard { "if [ -d \"$d/tensorboard\" ]; then find \"$d\" -mindepth 1 -maxdepth 1 ! -name tensorboard -exec rm -rf -- {} +; else rm -rf -- \"$d\"; fi" } else { @@ -105,8 +119,10 @@ fn host_script(runs: &[RemoteRun], keep_digests: &BTreeSet) -> String { script.push_str(&format!( "d={dir}; p=$(cat \"$d/pid\" 2>/dev/null); case \"$p\" in ''|*[!0-9]*) p=;; esac; \ if [ -n \"$p\" ] && host_group_alive \"$p\"; then echo 'BUSY {id}'; \ - else {inner}if err=$({{ {remove}; }} 2>&1); then echo 'REMOVED {id}'; \ - else printf 'FAILED {id} %s\\n' \"$(printf '%s' \"$err\" | head -n 1)\"; fi; fi\n", + elif {container_failed}; then \ + printf 'FAILED {id} container: %s\\n' \"$(printf '%s' \"$cerr\" | head -n 1)\"; \ + elif err=$({{ {remove}; }} 2>&1); then echo 'REMOVED {id}'; \ + else printf 'FAILED {id} %s\\n' \"$(printf '%s' \"$err\" | head -n 1)\"; fi\n", dir = shell_path(&run.dir), id = run.run_id, )); @@ -116,7 +132,6 @@ fn host_script(runs: &[RemoteRun], keep_digests: &BTreeSet) -> String { .map(String::as_str) .collect::>() .join(" "); - let roots: BTreeSet<_> = runs.iter().map(|run| run.root.as_deref()).collect(); for root in roots { // Only names orx writes (`.tar` and its upload temps): a root may hold the user's own tars. script.push_str(&format!( @@ -127,7 +142,7 @@ fn host_script(runs: &[RemoteRun], keep_digests: &BTreeSet) -> String { case \"$n\" in \"$b.tar\"|\"$b.tar.tmp.\"*) ;; *) continue;; esac; \ case ' {keep} ' in *\" $b \"*) continue;; esac; \ [ -n \"$(find \"$f\" -mmin +{SOURCE_IDLE_MINUTES})\" ] && rm -f -- \"$f\"; done; fi\n", - source = shell_path(&remote_source_dir(root)), + source = shell_path(&remote_source_dir(root.as_deref())), )); } script.push_str("exit 0\n"); @@ -159,14 +174,15 @@ fn parse_outcomes(runs: &[RemoteRun], output: &str) -> BTreeMap .collect() } -/// Removes `runs` (all on `host`) and that host's unreferenced cached tars. +/// Removes `runs` (all on `host`) and that host's unreferenced cached tars under `roots`. pub async fn clean_host( host: &str, runs: &[RemoteRun], + roots: &BTreeSet>, keep_digests: &BTreeSet, ) -> Result> { // Over stdin: a large batch would overrun Windows' command-line limit. - let script = host_script(runs, keep_digests); + let script = host_script(runs, roots, keep_digests); let output = tokio::time::timeout( HOST_TIMEOUT, ssh_run(&SshTarget::alias(host), "bash -s", Some(&script)), @@ -287,7 +303,14 @@ mod tests { run("deleted", None, true), ]; let out = std::process::Command::new("bash") - .args(["-c", &host_script(&batch, &BTreeSet::from([kept.clone()]))]) + .args([ + "-c", + &host_script( + &batch, + &BTreeSet::from([None]), + &BTreeSet::from([kept.clone()]), + ), + ]) .env("HOME", home.path()) .output() .unwrap(); @@ -319,11 +342,52 @@ mod tests { assert_eq!(left, expected); } + #[cfg(unix)] + #[test] + fn a_stopped_container_keeps_its_run_queued() { + use std::os::unix::fs::PermissionsExt as _; + let home = crate::local::git::TemporaryDirectory::new("orx-cleanup-container").unwrap(); + let bin = home.path().join("bin"); + std::fs::create_dir_all(&bin).unwrap(); + let docker = bin.join("docker"); + // Inspect finds the container; exec fails as it does for a stopped one. + std::fs::write( + &docker, + "#!/bin/sh\n[ \"$1\" = container ] && exit 0\necho 'container is not running' >&2\nexit 1\n", + ) + .unwrap(); + std::fs::set_permissions(&docker, std::fs::Permissions::from_mode(0o755)).unwrap(); + let mut stopped = run("stopped", None, false); + stopped.container = Some(("abc".into(), "/root/.orx/runs/stopped".into())); + let dir = home.path().join(".orx/runs/stopped"); + std::fs::create_dir_all(&dir).unwrap(); + let path = format!("{}:{}", bin.display(), std::env::var("PATH").unwrap()); + let out = std::process::Command::new("bash") + .args([ + "-c", + &host_script(&[stopped.clone()], &BTreeSet::new(), &BTreeSet::new()), + ]) + .env("HOME", home.path()) + .env("PATH", path) + .output() + .unwrap(); + let outcomes = parse_outcomes(&[stopped], &String::from_utf8(out.stdout).unwrap()); + assert_eq!( + outcomes["stopped"], + Outcome::Failed("container: container is not running".into()) + ); + assert!(dir.exists()); + } + #[test] fn host_script_quotes_rooted_and_container_paths() { let mut rooted = run(RUN, Some("/scratch/orx"), true); rooted.container = Some(("abc".into(), format!("/root/.orx/runs/{RUN}"))); - let script = host_script(&[rooted], &BTreeSet::new()); + let script = host_script( + &[rooted], + &BTreeSet::from([Some("/scratch/orx".to_string())]), + &BTreeSet::new(), + ); assert!(script.contains(&format!("d='/scratch/orx/runs/{RUN}';"))); assert!(script.contains(&format!( "docker exec 'abc' rm -rf -- '/root/.orx/runs/{RUN}'" diff --git a/src/local/remote_cleanup.rs b/src/local/remote_cleanup.rs index c1cf29f0d..0b808b294 100644 --- a/src/local/remote_cleanup.rs +++ b/src/local/remote_cleanup.rs @@ -13,21 +13,33 @@ const BATCH: usize = 50; /// (run id, reason) pairs for the log. pub type Failures = Vec<(String, String)>; +#[derive(Default)] pub struct HostBatch { host: String, runs: Vec, + /// Storage roots whose source cache this pass sweeps. + roots: BTreeSet>, keep_digests: BTreeSet, } -/// Groups due entries by host. The keep setting is read when a run comes due, so it covers runs -/// not yet cleaned up. Entries that need no cleanup leave the queue; ones orx cannot safely clean leave it -/// too and come back as failures. +pub struct HostResult { + host: String, + roots: BTreeSet>, + swept_only: bool, + /// (run id, deleted flag the pass acted on, outcome). + runs: Vec<(String, bool, Outcome)>, + reached: bool, +} + +/// Groups due entries and source sweeps by host. The keep setting is read when a run comes due, +/// so it covers runs not yet cleaned up. Entries that need no cleanup leave the queue; ones orx +/// cannot safely clean leave it too and come back as failures. pub fn plan(store: &Store) -> Result<(Vec, Failures)> { let ssh = crate::config::ssh_settings()?; let slurm_keeps = crate::jobs::slurm::load_settings()? .unwrap_or_default() .keep_remote_runs; - let mut hosts: BTreeMap> = BTreeMap::new(); + let mut hosts: BTreeMap = BTreeMap::new(); let mut rejected = Vec::new(); for entry in store.due_remote_cleanups(BATCH)? { let parsed = BackendDescriptor::parse(&entry.backend_json).and_then(|descriptor| { @@ -46,23 +58,31 @@ pub fn plan(store: &Store) -> Result<(Vec, Failures)> { }); match parsed { Ok(Some((run, keeps))) if entry.deleted || !keeps => { - hosts.entry(run.host.clone()).or_default().push(run) + let batch = hosts.entry(run.host.clone()).or_default(); + batch.roots.insert(run.root().map(Into::into)); + batch.runs.push(run); } - Ok(_) => store.finish_remote_cleanup(&entry.run_id)?, + Ok(_) => store.finish_remote_cleanup(&entry.run_id, entry.deleted)?, Err(error) => { - store.finish_remote_cleanup(&entry.run_id)?; + store.finish_remote_cleanup(&entry.run_id, entry.deleted)?; rejected.push((entry.run_id, error.to_string())); } } } + for (host, root) in store.due_source_sweeps()? { + match root.as_deref().map(crate::jobs::ssh::validate_remote_root) { + Some(Err(_)) => store.finish_source_sweep(&host, root.as_deref())?, + _ => { + hosts.entry(host).or_default().roots.insert(root); + } + } + } let batches = hosts .into_iter() - .map(|(host, runs)| { - Ok(HostBatch { - keep_digests: active_source_digests(store, &host)?, - host, - runs, - }) + .map(|(host, mut batch)| { + batch.keep_digests = active_source_digests(store, &host)?; + batch.host = host; + Ok(batch) }) .collect::>()?; Ok((batches, rejected)) @@ -86,34 +106,61 @@ fn active_source_digests(store: &Store, host: &str) -> Result> Ok(digests) } -pub async fn clean(batches: Vec) -> Vec<(String, Outcome)> { +pub async fn clean(batches: Vec) -> Vec { let mut results = Vec::new(); for batch in batches { - match cleanup::clean_host(&batch.host, &batch.runs, &batch.keep_digests).await { - Ok(outcomes) => results.extend(outcomes), - Err(error) => results.extend( - batch - .runs - .iter() - .map(|run| (run.run_id.clone(), Outcome::Failed(format!("{error:#}")))), - ), - } + let cleaned = + cleanup::clean_host(&batch.host, &batch.runs, &batch.roots, &batch.keep_digests).await; + let reached = cleaned.is_ok(); + let mut outcomes = cleaned.map_err(|error| format!("{error:#}")); + let runs = batch + .runs + .iter() + .map(|run| { + let outcome = match &mut outcomes { + Ok(outcomes) => outcomes + .remove(&run.run_id) + .unwrap_or_else(|| Outcome::Failed("no result from the host".into())), + Err(error) => Outcome::Failed(error.clone()), + }; + (run.run_id.clone(), run.deleted(), outcome) + }) + .collect(); + results.push(HostResult { + swept_only: batch.runs.is_empty(), + host: batch.host, + roots: batch.roots, + runs, + reached, + }); } results } /// Records outcomes and returns failures worth reporting; failed and busy runs retry later. -pub fn record(store: &Store, results: Vec<(String, Outcome)>) -> Result { +pub fn record(store: &Store, results: Vec) -> Result { let mut failures = Vec::new(); - for (run_id, outcome) in results { - match outcome { - Outcome::Removed => store.finish_remote_cleanup(&run_id)?, - Outcome::Busy => { - store.defer_remote_cleanup(&run_id, "its process group is still running")? + for result in results { + for root in &result.roots { + // A sweep-only pass that reached the host is done; any other pass may have left tars too fresh to remove. + if result.reached && result.swept_only { + store.finish_source_sweep(&result.host, root.as_deref())?; + } else { + store.schedule_source_sweep(&result.host, root.as_deref())?; } - Outcome::Failed(error) => { - store.defer_remote_cleanup(&run_id, &error)?; - failures.push((run_id, error)); + } + for (run_id, deleted, outcome) in result.runs { + match outcome { + Outcome::Removed => store.finish_remote_cleanup(&run_id, deleted)?, + Outcome::Busy => store.defer_remote_cleanup( + &run_id, + deleted, + "its process group is still running", + )?, + Outcome::Failed(error) => { + store.defer_remote_cleanup(&run_id, deleted, &error)?; + failures.push((run_id, error)); + } } } } diff --git a/src/store.rs b/src/store.rs index a96241079..decb90b0d 100644 --- a/src/store.rs +++ b/src/store.rs @@ -555,6 +555,8 @@ fn queue_remote_cleanups( /// Failed remote cleanups retry after 5 minutes, doubling up to a day. const REMOTE_CLEANUP_RETRY_MS: i64 = 5 * 60 * 1000; const REMOTE_CLEANUP_MAX_RETRY_MS: i64 = 24 * 60 * 60 * 1000; +/// Just past `jobs::ssh::cleanup`'s idle window, so tars a pass had to skip are removable. +const SOURCE_SWEEP_DELAY_MS: i64 = 61 * 60 * 1000; pub struct Store { conn: Connection, @@ -816,6 +818,12 @@ impl Store { next_attempt_at INTEGER NOT NULL, last_error TEXT ); + CREATE TABLE IF NOT EXISTS remote_source_sweeps ( + host TEXT NOT NULL, + root TEXT NOT NULL, + due_at INTEGER NOT NULL, + PRIMARY KEY (host, root) + ); CREATE TABLE IF NOT EXISTS overleaf_links ( project_id TEXT NOT NULL, tex_path TEXT NOT NULL, @@ -2194,21 +2202,28 @@ impl Store { Ok(rows.collect::, _>>()?) } - pub(crate) fn finish_remote_cleanup(&self, run_id: &str) -> Result<()> { + /// `deleted` is what the pass acted on; a deletion queued meanwhile stays for the next pass. + pub(crate) fn finish_remote_cleanup(&self, run_id: &str, deleted: bool) -> Result<()> { self.conn.execute( - "DELETE FROM remote_cleanups WHERE run_id = ?1", - params![run_id], + "DELETE FROM remote_cleanups WHERE run_id = ?1 AND deleted = ?2", + params![run_id, deleted], )?; Ok(()) } - pub(crate) fn defer_remote_cleanup(&self, run_id: &str, error: &str) -> Result<()> { + pub(crate) fn defer_remote_cleanup( + &self, + run_id: &str, + deleted: bool, + error: &str, + ) -> Result<()> { self.conn.execute( - "UPDATE remote_cleanups SET attempts = attempts + 1, last_error = ?2, - next_attempt_at = ?3 + MIN(?4 * (1 << MIN(attempts, 10)), ?5) - WHERE run_id = ?1", + "UPDATE remote_cleanups SET attempts = attempts + 1, last_error = ?3, + next_attempt_at = ?4 + MIN(?5 * (1 << MIN(attempts, 10)), ?6) + WHERE run_id = ?1 AND deleted = ?2", params![ run_id, + deleted, error, now_ms(), REMOTE_CLEANUP_RETRY_MS, @@ -2218,6 +2233,40 @@ impl Store { Ok(()) } + /// Revisits a host's source cache once its tars from this pass have gone idle. + pub(crate) fn schedule_source_sweep(&self, host: &str, root: Option<&str>) -> Result<()> { + self.conn.execute( + "INSERT INTO remote_source_sweeps (host, root, due_at) VALUES (?1, ?2, ?3) + ON CONFLICT(host, root) DO UPDATE SET due_at = excluded.due_at", + params![ + host, + root.unwrap_or_default(), + now_ms() + SOURCE_SWEEP_DELAY_MS + ], + )?; + Ok(()) + } + + /// Due (host, root) source-cache sweeps; `None` is the default root. + pub(crate) fn due_source_sweeps(&self) -> Result)>> { + let mut statement = self + .conn + .prepare("SELECT host, root FROM remote_source_sweeps WHERE due_at <= ?1")?; + let rows = statement.query_map(params![now_ms()], |row| { + let root: String = row.get(1)?; + Ok((row.get(0)?, Some(root).filter(|root| !root.is_empty()))) + })?; + Ok(rows.collect::, _>>()?) + } + + pub(crate) fn finish_source_sweep(&self, host: &str, root: Option<&str>) -> Result<()> { + self.conn.execute( + "DELETE FROM remote_source_sweeps WHERE host = ?1 AND root = ?2", + params![host, root.unwrap_or_default()], + )?; + Ok(()) + } + fn data_dir_move_claimed(&self) -> Result { self.clear_stale_data_dir_move_lease()?; Ok(self.conn.query_row( @@ -4329,7 +4378,9 @@ mod tests { // As locally, an unfinished newer run takes one of the slots. insert("r5", "running", "slurm_job", 6); assert_eq!(due(&store), finished(&["r1", "r2"])); - store.defer_remote_cleanup("r1", "unreachable").unwrap(); + store + .defer_remote_cleanup("r1", false, "unreachable") + .unwrap(); assert_eq!(due(&store), finished(&["r2"])); store .conn @@ -4344,8 +4395,31 @@ mod tests { .map(|id| (id.to_string(), true)) .collect::>() ); - store.finish_remote_cleanup("r1").unwrap(); + // A pass that started before the deletion queued only finishes what it acted on. + store.finish_remote_cleanup("r2", false).unwrap(); + store.defer_remote_cleanup("r2", false, "stale").unwrap(); + store.finish_remote_cleanup("r1", true).unwrap(); assert_eq!(due(&store).len(), 4); + store.schedule_source_sweep("lab", None).unwrap(); + store + .schedule_source_sweep("lab", Some("/scratch/orx")) + .unwrap(); + assert!(store.due_source_sweeps().unwrap().is_empty()); + store + .conn + .execute("UPDATE remote_source_sweeps SET due_at = 0", []) + .unwrap(); + let mut sweeps = store.due_source_sweeps().unwrap(); + sweeps.sort(); + assert_eq!( + sweeps, + [ + ("lab".to_string(), None), + ("lab".to_string(), Some("/scratch/orx".to_string())) + ] + ); + store.finish_source_sweep("lab", None).unwrap(); + assert_eq!(store.due_source_sweeps().unwrap().len(), 1); std::fs::remove_dir_all(dir).unwrap(); } From 702139a7a793dca04b32e778d3261d4e880f1731 Mon Sep 17 00:00:00 2001 From: Myles Anderson Date: Thu, 22 Oct 2026 18:10:57 -0700 Subject: [PATCH 10/10] Remove the host copy beside a stopped container; finish best-effort sweeps Co-Authored-By: Claude Opus 5.5 --- agent-skills/orx-compute/references/slurm.md | 6 +- agent-skills/orx-compute/references/ssh.md | 4 +- src/jobs/ssh/cleanup.rs | 131 +++++++++++-------- src/local/remote_cleanup.rs | 57 ++++---- src/store.rs | 11 +- 5 files changed, 116 insertions(+), 93 deletions(-) diff --git a/agent-skills/orx-compute/references/slurm.md b/agent-skills/orx-compute/references/slurm.md index a08638ced..62629ff00 100644 --- a/agent-skills/orx-compute/references/slurm.md +++ b/agent-skills/orx-compute/references/slurm.md @@ -30,9 +30,9 @@ orx exp run --backend slurm - Each project's three newest SSH and Slurm runs, counted together across hosts and including unfinished ones, stay whole. Once an older finished run's log is collected, the running OpenResearch app removes its remote run - dir except `tensorboard/`, and while there (and again an hour later), orx's - cached source archives that no active run uses and that have been idle for - an hour. Anything the job wrote into `repo/` goes with it. To keep every run + dir except `tensorboard/`. It also removes orx's cached source archives that + no active run uses and that have been idle for an hour, then checks again an + hour later. Anything the job wrote into `repo/` goes with it. To keep every run on the cluster, set `orx compute configure slurm --keep-remote-runs true` (it covers runs not yet cleaned up), or write outputs outside the run dir. Deleting a project removes its runs' remote dirs entirely regardless. diff --git a/agent-skills/orx-compute/references/ssh.md b/agent-skills/orx-compute/references/ssh.md index 8dc536b34..9d56bdbf5 100644 --- a/agent-skills/orx-compute/references/ssh.md +++ b/agent-skills/orx-compute/references/ssh.md @@ -46,8 +46,8 @@ agents still connect directly to the host. Each project's three newest SSH and Slurm runs, counted together across hosts and including unfinished ones, stay whole. Once an older finished run's log is collected, the running OpenResearch app removes its host dir except -`tensorboard/` (a container's copy goes entirely), so anything written into -`repo/` goes with it. While cleaning a host it also removes orx's cached source +`tensorboard/`, and the container's copy entirely (retried while the container +is stopped), so anything written into `repo/` goes with it. While cleaning a host it also removes orx's cached source archives that no active run uses and that have been idle for an hour, and checks again an hour later. Unreachable hosts are retried later. diff --git a/src/jobs/ssh/cleanup.rs b/src/jobs/ssh/cleanup.rs index 992127c26..03d6fdb48 100644 --- a/src/jobs/ssh/cleanup.rs +++ b/src/jobs/ssh/cleanup.rs @@ -11,7 +11,7 @@ use crate::error::{anyhow, Result}; use crate::jobs::BackendDescriptor; /// Cached tars untouched this long are no longer being staged by any orx install. -const SOURCE_IDLE_MINUTES: u32 = 60; +pub(crate) const SOURCE_IDLE_MINUTES: i64 = 60; const HOST_TIMEOUT: Duration = Duration::from_secs(120); /// A run's remote dir, checked against the layout orx writes before anything removes it. @@ -23,8 +23,9 @@ pub struct RemoteRun { dir: String, /// Docker id and in-container run dir for container runs. container: Option<(String, String)>, - /// `orx tensorboard` reads a finished run's remote `tensorboard/`; only deletion removes it. - keep_tensorboard: bool, + /// Queued by a project or experiment delete; otherwise `orx tensorboard` still reads the + /// finished run's remote `tensorboard/`, so it stays. + deleted: bool, } impl RemoteRun { @@ -69,7 +70,7 @@ impl RemoteRun { root, dir, container, - keep_tensorboard: !deleted, + deleted, })) } @@ -78,7 +79,7 @@ impl RemoteRun { } pub(crate) fn deleted(&self) -> bool { - !self.keep_tensorboard + self.deleted } } @@ -99,19 +100,22 @@ fn host_script( ) -> String { let mut script = format!("{HOST_PROCESS_HELPERS}\n"); for run in runs { - // A container that still exists but can't run the removal (stopped) keeps the run queued. + // A container that exists but can't run the removal (stopped), or a Docker daemon that + // can't say, keeps the run queued; only a container Docker reports gone counts as cleaned. let container_failed = run.container.as_ref().map_or_else( || "false".to_string(), |(id, dir)| { format!( - "docker container inspect -- {id} /dev/null 2>&1 && \ - ! cerr=$(docker exec {id} rm -rf -- {dir} &1)", + "{{ if docker container inspect -- {id} /dev/null 2>&1; then \ + ! cerr=$(docker exec {id} rm -rf -- {dir} &1); \ + elif ! command -v docker >/dev/null 2>&1 || docker info /dev/null 2>&1; then false; \ + else cerr='Docker is unavailable'; fi; }}", id = sh_quote(id), dir = sh_quote(dir) ) }, ); - let remove = if run.keep_tensorboard { + let remove = if !run.deleted { "if [ -d \"$d/tensorboard\" ]; then find \"$d\" -mindepth 1 -maxdepth 1 ! -name tensorboard -exec rm -rf -- {} +; else rm -rf -- \"$d\"; fi" } else { "rm -rf -- \"$d\"" @@ -119,10 +123,10 @@ fn host_script( script.push_str(&format!( "d={dir}; p=$(cat \"$d/pid\" 2>/dev/null); case \"$p\" in ''|*[!0-9]*) p=;; esac; \ if [ -n \"$p\" ] && host_group_alive \"$p\"; then echo 'BUSY {id}'; \ - elif {container_failed}; then \ - printf 'FAILED {id} container: %s\\n' \"$(printf '%s' \"$cerr\" | head -n 1)\"; \ - elif err=$({{ {remove}; }} 2>&1); then echo 'REMOVED {id}'; \ - else printf 'FAILED {id} %s\\n' \"$(printf '%s' \"$err\" | head -n 1)\"; fi\n", + else cf=; if {container_failed}; then cf=1; fi; \ + if ! err=$({{ {remove}; }} 2>&1); then printf 'FAILED {id} %s\\n' \"$(printf '%s' \"$err\" | head -n 1)\"; \ + elif [ -n \"$cf\" ]; then printf 'FAILED {id} container: %s\\n' \"$(printf '%s' \"$cerr\" | head -n 1)\"; \ + else echo 'REMOVED {id}'; fi; fi\n", dir = shell_path(&run.dir), id = run.run_id, )); @@ -149,7 +153,8 @@ fn host_script( script } -fn parse_outcomes(runs: &[RemoteRun], output: &str) -> BTreeMap { +/// One outcome per run, in `runs` order. +fn parse_outcomes(runs: &[RemoteRun], output: &str) -> Vec { let mut reported = BTreeMap::new(); for line in output.lines() { let mut words = line.splitn(3, ' '); @@ -166,10 +171,9 @@ fn parse_outcomes(runs: &[RemoteRun], output: &str) -> BTreeMap } runs.iter() .map(|run| { - let outcome = reported + reported .remove(&run.run_id) - .unwrap_or_else(|| Outcome::Failed("no result from the host".into())); - (run.run_id.clone(), outcome) + .unwrap_or_else(|| Outcome::Failed("no result from the host".into())) }) .collect() } @@ -180,7 +184,7 @@ pub async fn clean_host( runs: &[RemoteRun], roots: &BTreeSet>, keep_digests: &BTreeSet, -) -> Result> { +) -> Result> { // Over stdin: a large batch would overrun Windows' command-line limit. let script = host_script(runs, roots, keep_digests); let output = tokio::time::timeout( @@ -243,6 +247,13 @@ mod tests { } } + fn outcomes_by_id(runs: &[RemoteRun], output: &str) -> BTreeMap { + runs.iter() + .map(|run| run.run_id.clone()) + .zip(parse_outcomes(runs, output)) + .collect() + } + fn run(id: &str, root: Option<&str>, deleted: bool) -> RemoteRun { RemoteRun { run_id: id.into(), @@ -250,7 +261,7 @@ mod tests { root: root.map(Into::into), dir: remote_run_dir(root, id), container: None, - keep_tensorboard: !deleted, + deleted, } } @@ -317,7 +328,7 @@ mod tests { let _ = sleeper.kill(); let _ = sleeper.wait(); assert!(out.status.success()); - let outcomes = parse_outcomes(&batch, &String::from_utf8(out.stdout).unwrap()); + let outcomes = outcomes_by_id(&batch, &String::from_utf8(out.stdout).unwrap()); for id in ["done", "never", "charted", "deleted"] { assert_eq!(outcomes[id], Outcome::Removed, "{id}"); } @@ -344,39 +355,55 @@ mod tests { #[cfg(unix)] #[test] - fn a_stopped_container_keeps_its_run_queued() { + fn container_runs_stay_queued_until_docker_reports_the_container_gone() { use std::os::unix::fs::PermissionsExt as _; - let home = crate::local::git::TemporaryDirectory::new("orx-cleanup-container").unwrap(); - let bin = home.path().join("bin"); - std::fs::create_dir_all(&bin).unwrap(); - let docker = bin.join("docker"); - // Inspect finds the container; exec fails as it does for a stopped one. - std::fs::write( - &docker, - "#!/bin/sh\n[ \"$1\" = container ] && exit 0\necho 'container is not running' >&2\nexit 1\n", - ) - .unwrap(); - std::fs::set_permissions(&docker, std::fs::Permissions::from_mode(0o755)).unwrap(); - let mut stopped = run("stopped", None, false); - stopped.container = Some(("abc".into(), "/root/.orx/runs/stopped".into())); - let dir = home.path().join(".orx/runs/stopped"); - std::fs::create_dir_all(&dir).unwrap(); - let path = format!("{}:{}", bin.display(), std::env::var("PATH").unwrap()); - let out = std::process::Command::new("bash") - .args([ - "-c", - &host_script(&[stopped.clone()], &BTreeSet::new(), &BTreeSet::new()), - ]) - .env("HOME", home.path()) - .env("PATH", path) - .output() + // (inspect exit, info exit, expected outcome): stopped, removed, daemon down. + for (inspect, info, expected) in [ + ( + 0, + 0, + Outcome::Failed("container: container is not running".into()), + ), + (1, 0, Outcome::Removed), + ( + 1, + 1, + Outcome::Failed("container: Docker is unavailable".into()), + ), + ] { + let home = crate::local::git::TemporaryDirectory::new("orx-cleanup-container").unwrap(); + let bin = home.path().join("bin"); + std::fs::create_dir_all(&bin).unwrap(); + let docker = bin.join("docker"); + std::fs::write( + &docker, + format!( + "#!/bin/sh\ncase \"$1\" in container) exit {inspect};; info) exit {info};; esac\n\ + echo 'container is not running' >&2\nexit 1\n" + ), + ) .unwrap(); - let outcomes = parse_outcomes(&[stopped], &String::from_utf8(out.stdout).unwrap()); - assert_eq!( - outcomes["stopped"], - Outcome::Failed("container: container is not running".into()) - ); - assert!(dir.exists()); + std::fs::set_permissions(&docker, std::fs::Permissions::from_mode(0o755)).unwrap(); + let mut container = run("c", None, false); + container.container = Some(("abc".into(), "/root/.orx/runs/c".into())); + let dir = home.path().join(".orx/runs/c"); + std::fs::create_dir_all(dir.join("repo")).unwrap(); + let path = format!("{}:{}", bin.display(), std::env::var("PATH").unwrap()); + let batch = std::slice::from_ref(&container); + let out = std::process::Command::new("bash") + .args([ + "-c", + &host_script(batch, &BTreeSet::new(), &BTreeSet::new()), + ]) + .env("HOME", home.path()) + .env("PATH", path) + .output() + .unwrap(); + let outcomes = outcomes_by_id(batch, &String::from_utf8(out.stdout).unwrap()); + assert_eq!(outcomes["c"], expected, "inspect {inspect} info {info}"); + // The host copy goes either way; only the container's copy keeps the run queued. + assert!(!dir.exists()); + } } #[test] @@ -394,7 +421,7 @@ mod tests { ))); assert!(script.contains("s='/scratch/orx/source';")); let batch = [run(RUN, None, false), run("other", None, false)]; - let outcomes = parse_outcomes(&batch, &format!("FAILED {RUN} rm: Permission denied\n")); + let outcomes = outcomes_by_id(&batch, &format!("FAILED {RUN} rm: Permission denied\n")); assert_eq!( outcomes[RUN], Outcome::Failed("rm: Permission denied".into()) diff --git a/src/local/remote_cleanup.rs b/src/local/remote_cleanup.rs index 0b808b294..f3882eed8 100644 --- a/src/local/remote_cleanup.rs +++ b/src/local/remote_cleanup.rs @@ -13,7 +13,6 @@ const BATCH: usize = 50; /// (run id, reason) pairs for the log. pub type Failures = Vec<(String, String)>; -#[derive(Default)] pub struct HostBatch { host: String, runs: Vec, @@ -25,10 +24,8 @@ pub struct HostBatch { pub struct HostResult { host: String, roots: BTreeSet>, - swept_only: bool, /// (run id, deleted flag the pass acted on, outcome). runs: Vec<(String, bool, Outcome)>, - reached: bool, } /// Groups due entries and source sweeps by host. The keep setting is read when a run comes due, @@ -39,7 +36,7 @@ pub fn plan(store: &Store) -> Result<(Vec, Failures)> { let slurm_keeps = crate::jobs::slurm::load_settings()? .unwrap_or_default() .keep_remote_runs; - let mut hosts: BTreeMap = BTreeMap::new(); + let mut hosts: BTreeMap, BTreeSet>)> = BTreeMap::new(); let mut rejected = Vec::new(); for entry in store.due_remote_cleanups(BATCH)? { let parsed = BackendDescriptor::parse(&entry.backend_json).and_then(|descriptor| { @@ -58,9 +55,9 @@ pub fn plan(store: &Store) -> Result<(Vec, Failures)> { }); match parsed { Ok(Some((run, keeps))) if entry.deleted || !keeps => { - let batch = hosts.entry(run.host.clone()).or_default(); - batch.roots.insert(run.root().map(Into::into)); - batch.runs.push(run); + let (runs, roots) = hosts.entry(run.host.clone()).or_default(); + roots.insert(run.root().map(Into::into)); + runs.push(run); } Ok(_) => store.finish_remote_cleanup(&entry.run_id, entry.deleted)?, Err(error) => { @@ -70,19 +67,17 @@ pub fn plan(store: &Store) -> Result<(Vec, Failures)> { } } for (host, root) in store.due_source_sweeps()? { - match root.as_deref().map(crate::jobs::ssh::validate_remote_root) { - Some(Err(_)) => store.finish_source_sweep(&host, root.as_deref())?, - _ => { - hosts.entry(host).or_default().roots.insert(root); - } - } + hosts.entry(host).or_default().1.insert(root); } let batches = hosts .into_iter() - .map(|(host, mut batch)| { - batch.keep_digests = active_source_digests(store, &host)?; - batch.host = host; - Ok(batch) + .map(|(host, (runs, roots))| { + Ok(HostBatch { + keep_digests: active_source_digests(store, &host)?, + host, + runs, + roots, + }) }) .collect::>()?; Ok((batches, rejected)) @@ -109,29 +104,23 @@ fn active_source_digests(store: &Store, host: &str) -> Result> pub async fn clean(batches: Vec) -> Vec { let mut results = Vec::new(); for batch in batches { - let cleaned = - cleanup::clean_host(&batch.host, &batch.runs, &batch.roots, &batch.keep_digests).await; - let reached = cleaned.is_ok(); - let mut outcomes = cleaned.map_err(|error| format!("{error:#}")); + let outcomes = + match cleanup::clean_host(&batch.host, &batch.runs, &batch.roots, &batch.keep_digests) + .await + { + Ok(outcomes) => outcomes, + Err(error) => vec![Outcome::Failed(format!("{error:#}")); batch.runs.len()], + }; let runs = batch .runs .iter() - .map(|run| { - let outcome = match &mut outcomes { - Ok(outcomes) => outcomes - .remove(&run.run_id) - .unwrap_or_else(|| Outcome::Failed("no result from the host".into())), - Err(error) => Outcome::Failed(error.clone()), - }; - (run.run_id.clone(), run.deleted(), outcome) - }) + .zip(outcomes) + .map(|(run, outcome)| (run.run_id.clone(), run.deleted(), outcome)) .collect(); results.push(HostResult { - swept_only: batch.runs.is_empty(), host: batch.host, roots: batch.roots, runs, - reached, }); } results @@ -142,8 +131,8 @@ pub fn record(store: &Store, results: Vec) -> Result { let mut failures = Vec::new(); for result in results { for root in &result.roots { - // A sweep-only pass that reached the host is done; any other pass may have left tars too fresh to remove. - if result.reached && result.swept_only { + // Run passes can leave tars too fresh to remove; a sweep-only pass is best effort. + if result.runs.is_empty() { store.finish_source_sweep(&result.host, root.as_deref())?; } else { store.schedule_source_sweep(&result.host, root.as_deref())?; diff --git a/src/store.rs b/src/store.rs index decb90b0d..704045bf9 100644 --- a/src/store.rs +++ b/src/store.rs @@ -555,8 +555,8 @@ fn queue_remote_cleanups( /// Failed remote cleanups retry after 5 minutes, doubling up to a day. const REMOTE_CLEANUP_RETRY_MS: i64 = 5 * 60 * 1000; const REMOTE_CLEANUP_MAX_RETRY_MS: i64 = 24 * 60 * 60 * 1000; -/// Just past `jobs::ssh::cleanup`'s idle window, so tars a pass had to skip are removable. -const SOURCE_SWEEP_DELAY_MS: i64 = 61 * 60 * 1000; +/// Just past the source cache's idle window, so tars a pass had to skip are removable. +const SOURCE_SWEEP_DELAY_MS: i64 = (crate::jobs::ssh::cleanup::SOURCE_IDLE_MINUTES + 1) * 60 * 1000; pub struct Store { conn: Connection, @@ -4400,6 +4400,13 @@ mod tests { store.defer_remote_cleanup("r2", false, "stale").unwrap(); store.finish_remote_cleanup("r1", true).unwrap(); assert_eq!(due(&store).len(), 4); + std::fs::remove_dir_all(dir).unwrap(); + } + + #[test] + fn source_sweeps_come_due_after_the_idle_window() { + let dir = std::env::temp_dir().join(format!("orx-source-sweep-{}", uuid::Uuid::new_v4())); + let store = Store::open_at(dir.clone()).unwrap(); store.schedule_source_sweep("lab", None).unwrap(); store .schedule_source_sweep("lab", Some("/scratch/orx"))