diff --git a/agent-skills/orx-compute/references/configuration.md b/agent-skills/orx-compute/references/configuration.md index 54563ad3e..81c38199b 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 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 @@ -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 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 20cd93384..62629ff00 100644 --- a/agent-skills/orx-compute/references/slurm.md +++ b/agent-skills/orx-compute/references/slurm.md @@ -27,5 +27,14 @@ 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, 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/`. 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. - 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..9d56bdbf5 100644 --- a/agent-skills/orx-compute/references/ssh.md +++ b/agent-skills/orx-compute/references/ssh.md @@ -43,6 +43,18 @@ 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, 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/`, 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. +`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 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 @@ -58,8 +70,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 79d567909..e59eaa394 100644 --- a/src/commands/compute.rs +++ b/src/commands/compute.rs @@ -238,6 +238,10 @@ 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 3 newest + /// SSH/Slurm runs and removes older ones except tensorboard/. + #[arg(long)] + keep_remote_runs: Option, #[arg(long)] default_host: Option, #[arg(long)] @@ -324,7 +328,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 +343,7 @@ impl ConfigureArgs { "cpusPerTask", "mem", "remoteRoot", + "keepRemoteRuns", ], "k8s" => &["context", "namespace"], "ray" => &["address"], @@ -359,6 +370,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)); } @@ -395,6 +409,7 @@ impl ConfigureArgs { "cpus-per-task" => "cpusPerTask", "default-host" => "defaultHost", "remote-root" => "remoteRoot", + "keep-remote-runs" => "keepRemoteRuns", other => other, }; if !allowed.contains(&key) @@ -412,6 +427,8 @@ impl ConfigureArgs { Value::Null } else if key == "cpusPerTask" { json!(0) + } else if key == "keepRemoteRuns" { + json!(false) } else { json!("") }, @@ -427,13 +444,15 @@ 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!( - "A container or remote root change requires --host." + "A container, remote root, or keep-remote-runs change requires --host." )); } Ok(Value::Object(body)) @@ -783,6 +802,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..d4d2df4ca 100644 --- a/src/commands/supervise.rs +++ b/src/commands/supervise.rs @@ -849,6 +849,9 @@ async fn watch_ssh_job( ) .await; } + if 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,14 @@ async fn run_slurm( } } +/// 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}"); + } +} + 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 c753d5d4a..beface201 100644 --- a/src/commands/up.rs +++ b/src/commands/up.rs @@ -179,6 +179,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(), @@ -4616,6 +4620,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; @@ -4705,6 +4710,73 @@ 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>) { + // 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; + 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; + 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)) => { + 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(REMOTE_CLEANUP_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 f7476be1c..9479c50f3 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. @@ -625,22 +626,26 @@ pub(super) struct SaveSshSettingsReq { container: Option, /// 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, } 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 +865,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 +888,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 +932,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())) }) @@ -1475,8 +1485,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 => { @@ -1488,8 +1501,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 8dd21e8af..80ffa37d0 100644 --- a/src/config.rs +++ b/src/config.rs @@ -514,6 +514,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 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, } #[derive(Debug, Clone, Default, Serialize, Deserialize)] diff --git a/src/jobs/slurm.rs b/src/jobs/slurm.rs index bc7df5731..ecc09399b 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 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, } 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..03d6fdb48 --- /dev/null +++ b/src/jobs/ssh/cleanup.rs @@ -0,0 +1,431 @@ +//! 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. +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. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RemoteRun { + 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)>, + /// 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 { + /// `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); + } + 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 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, + deleted, + })) + } + + pub(crate) fn root(&self) -> Option<&str> { + self.root.as_deref() + } + + pub(crate) fn deleted(&self) -> bool { + self.deleted + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Outcome { + Removed, + /// The run's process group is still alive on the host. + Busy, + Failed(String), +} + +/// One script per host: remove each run dir whose process group is gone, then cached tars +/// 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 { + // 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!( + "{{ 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.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\"" + }; + 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 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, + )); + } + let keep = keep_digests + .iter() + .map(String::as_str) + .collect::>() + .join(" "); + 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 -mmin +{SOURCE_IDLE_MINUTES} | while IFS= read -r f; do \ + 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; \ + [ -n \"$(find \"$f\" -mmin +{SOURCE_IDLE_MINUTES})\" ] && rm -f -- \"$f\"; done; fi\n", + source = shell_path(&remote_source_dir(root.as_deref())), + )); + } + script.push_str("exit 0\n"); + script +} + +/// 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, ' '); + 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, + }; + reported.insert(id.to_string(), outcome); + } + runs.iter() + .map(|run| { + reported + .remove(&run.run_id) + .unwrap_or_else(|| Outcome::Failed("no result from the host".into())) + }) + .collect() +} + +/// 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, roots, keep_digests); + let output = tokio::time::timeout( + HOST_TIMEOUT, + ssh_run(&SshTarget::alias(host), "bash -s", Some(&script)), + ) + .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() + } + + fn cleanable(json: &str) -> Result> { + RemoteRun::from_descriptor(RUN, &descriptor(json), false) + } + + #[test] + fn only_orx_layout_dirs_are_cleanable() { + 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!(cleanable(&rooted).unwrap().is_some()); + let slurm = r#"{"kind":"slurm_job","namespace":"login","remoteRoot":"/scratch/orx"}"#; + assert_eq!( + 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","namespace":"lab","jobId":".orx/runs/{RUN}","sshContainer":{{"reference":"c","id":"i","startedAt":"t","runDir":"/root"}}}}"# + ), + ] { + assert!(cleanable(&json).is_err(), "{json}"); + } + 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"] { + assert!( + cleanable(&format!(r#"{{"kind":"{kind}","namespace":"x"}}"#)) + .unwrap() + .is_none() + ); + } + } + + 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(), + host: "lab".into(), + root: root.map(Into::into), + dir: remote_run_dir(root, id), + container: None, + deleted, + } + } + + #[cfg(unix)] + #[test] + 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", "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) + .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(); + 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()); + } + } + 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([None]), + &BTreeSet::from([kept.clone()]), + ), + ]) + .env("HOME", home.path()) + .output() + .unwrap(); + let _ = sleeper.kill(); + let _ = sleeper.wait(); + assert!(out.status.success()); + 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}"); + } + assert_eq!(outcomes["busy"], Outcome::Busy); + 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(); + let mut expected = vec![ + format!("{kept}.tar"), + format!("{fresh}.tar"), + "dataset.tar".to_string(), + format!("{old}.tar.bak"), + ]; + expected.sort(); + assert_eq!(left, expected); + } + + #[cfg(unix)] + #[test] + fn container_runs_stay_queued_until_docker_reports_the_container_gone() { + use std::os::unix::fs::PermissionsExt as _; + // (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(); + 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] + 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::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}'" + ))); + assert!(script.contains("s='/scratch/orx/source';")); + let batch = [run(RUN, None, false), run("other", None, false)]; + let outcomes = outcomes_by_id(&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/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..f3882eed8 --- /dev/null +++ b/src/local/remote_cleanup.rs @@ -0,0 +1,157 @@ +//! 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, + /// Storage roots whose source cache this pass sweeps. + roots: BTreeSet>, + keep_digests: BTreeSet, +} + +pub struct HostResult { + host: String, + roots: BTreeSet>, + /// (run id, deleted flag the pass acted on, outcome). + runs: Vec<(String, bool, Outcome)>, +} + +/// 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, 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| { + 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, entry.deleted)? + .map(|run| (run, keeps)), + ) + }); + match parsed { + Ok(Some((run, keeps))) if entry.deleted || !keeps => { + 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) => { + 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()? { + hosts.entry(host).or_default().1.insert(root); + } + let batches = hosts + .into_iter() + .map(|(host, (runs, roots))| { + Ok(HostBatch { + keep_digests: active_source_digests(store, &host)?, + host, + runs, + roots, + }) + }) + .collect::>()?; + Ok((batches, rejected)) +} + +/// 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) -> Vec { + let mut results = Vec::new(); + for batch in batches { + 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() + .zip(outcomes) + .map(|(run, outcome)| (run.run_id.clone(), run.deleted(), outcome)) + .collect(); + results.push(HostResult { + host: batch.host, + roots: batch.roots, + runs, + }); + } + results +} + +/// Records outcomes and returns failures worth reporting; failed and busy runs retry later. +pub fn record(store: &Store, results: Vec) -> Result { + let mut failures = Vec::new(); + for result in results { + for root in &result.roots { + // 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())?; + } + } + 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)); + } + } + } + } + Ok(failures) +} diff --git a/src/store.rs b/src/store.rs index 64d49a3e8..333ee52c7 100644 --- a/src/store.rs +++ b/src/store.rs @@ -477,7 +477,8 @@ 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; /// Uploads are written before the message naming them, so the sweep spares recent files. const CHAT_ATTACHMENT_SETTLE: Duration = Duration::from_secs(24 * 60 * 60); @@ -510,6 +511,54 @@ 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, +} + +/// `(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, + column: &'static 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), 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; +/// 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, data_dir: PathBuf, @@ -762,6 +811,20 @@ 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 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, @@ -1882,17 +1945,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::, _>>()?; @@ -2122,6 +2179,95 @@ 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<()> { + 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(&format!( + "SELECT run_id, backend_json, deleted FROM remote_cleanups + 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)?, + backend_json: row.get(1)?, + deleted: row.get(2)?, + }) + })?; + Ok(rows.collect::, _>>()?) + } + + /// `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 AND deleted = ?2", + params![run_id, deleted], + )?; + Ok(()) + } + + 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 = ?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, + REMOTE_CLEANUP_MAX_RETRY_MS + ], + )?; + 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()?; self.data_dir_move_lease_exists() @@ -2210,6 +2356,7 @@ impl Store { "DELETE FROM chat_sessions WHERE project_id = ?1", params![id], )?; + queue_remote_cleanups(&tx, "project_id", id, true)?; self.conn .execute("DELETE FROM runs WHERE project_id = ?1", params![id])?; self.conn.execute( @@ -4299,6 +4446,106 @@ mod tests { std::fs::remove_dir_all(dir).unwrap(); } + #[test] + 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)); + 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}); + store.conn.execute( + "INSERT INTO runs (id, experiment_id, project_id, status, backend_json, created_at, updated_at) + VALUES (?1, 'experiment', 'p', ?2, ?3, ?4, ?4)", + params![id, status, descriptor.to_string(), created_at], + ).unwrap(); + }; + let due = |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::>() + }; + 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", false, "unreachable") + .unwrap(); + assert_eq!(due(&store), finished(&["r2"])); + store + .conn + .execute("UPDATE runs SET status = 'done'", []) + .unwrap(); + store.delete_local_project("p").unwrap(); + // Deletion marks every SSH/Slurm run, including the newest and one waiting to retry. + assert_eq!( + due(&store), + ["r1", "r2", "r3", "r4", "r5"] + .iter() + .map(|id| (id.to_string(), true)) + .collect::>() + ); + // 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); + 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")) + .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(); + } + #[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();