diff --git a/Cargo.lock b/Cargo.lock index 65ddec7..30d0357 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5792,6 +5792,7 @@ dependencies = [ "dirs", "libloading", "persisting-capture", + "persisting-compute", "persisting-dlcapt", "persisting-engine", "persisting-proto", @@ -5807,6 +5808,30 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "persisting-compute" +version = "0.1.0" +dependencies = [ + "anyhow", + "async-trait", + "chrono", + "clap", + "futures", + "persisting-capture", + "persisting-engine", + "persisting-proto", + "pulsing-actor", + "serde", + "serde_json", + "tempfile", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", + "tracing-subscriber", + "uuid", +] + [[package]] name = "persisting-core" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index e29d338..aae774f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ members = [ "crates/persisting-capture", "crates/persisting-cli", "crates/persisting-dlcapt", + "crates/persisting-compute", ] [workspace.package] diff --git a/crates/persisting-cli/Cargo.toml b/crates/persisting-cli/Cargo.toml index 2a0ad6b..b589a42 100644 --- a/crates/persisting-cli/Cargo.toml +++ b/crates/persisting-cli/Cargo.toml @@ -25,6 +25,7 @@ persisting-capture = { path = "../persisting-capture" } persisting-dlcapt = { path = "../persisting-dlcapt", optional = true } persisting-engine = { path = "../persisting-engine" } persisting-proto = { path = "../persisting-proto" } +persisting-compute = { path = "../persisting-compute", features = ["traj-sink"] } ron = "0.8" serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/crates/persisting-cli/src/capture/proxy_config_args.rs b/crates/persisting-cli/src/capture/proxy_config_args.rs new file mode 100644 index 0000000..59c24bd --- /dev/null +++ b/crates/persisting-cli/src/capture/proxy_config_args.rs @@ -0,0 +1,165 @@ +//! Shared proxy config CLI flags + env vars for `traj capture` / `traj proxy`. + +use std::path::PathBuf; + +use anyhow::{Context, Result}; +use clap::Args; +use persisting_capture::config::{ + env, materialize_proxy_config, parse_capture_level, resolve_proxy_config, CaptureLevel, + ProxyConfig, ProxyConfigInput, ProxyConfigOverrides, +}; + +/// Proxy settings overridable via CLI flags and environment variables. +#[derive(Debug, Clone, Default, Args)] +pub struct ProxyConfigOverrideArgs { + /// Proxy listen address (overrides TOML `listen`). + #[arg(long, env = env::LISTEN, value_name = "ADDR")] + pub listen: Option, + /// Admin API listen address (overrides TOML `admin_listen`). + #[arg(long, env = env::ADMIN_LISTEN, value_name = "ADDR")] + pub admin_listen: Option, + /// Trajectory agent id segment (overrides TOML `agent_id`). + #[arg(long, env = env::AGENT_ID, value_name = "SEG")] + pub agent_id: Option, + /// Session id HTTP header name (overrides TOML `session_header`). + #[arg(long, env = env::SESSION_HEADER, value_name = "HEADER")] + pub session_header: Option, + /// Capture granularity: `summary`, `dialogue`, or `full`. + #[arg(long, env = env::CAPTURE_LEVEL, value_name = "LEVEL")] + pub capture_level: Option, + /// Model routes as JSON array (overrides all `[[models]]` entries). + #[arg(long, env = env::MODELS_JSON, value_name = "JSON")] + pub models_json: Option, + /// Model routes as TOML `[[models]]` table(s) (overrides all routes). + #[arg(long, env = env::MODELS_TOML, value_name = "TOML")] + pub models_toml: Option, + /// Model route name to create or patch (default: `*`). + #[arg(long, env = env::MODEL, value_name = "NAME")] + pub model: Option, + /// Upstream OpenAI-compatible base URL for `--model` (e.g. `https://api.deepseek.com/v1`). + #[arg(long, env = env::UPSTREAM, value_name = "URL")] + pub upstream: Option, + /// Anthropic-compatible upstream for `/v1/messages`. + #[arg(long, env = env::UPSTREAM_ANTHROPIC, value_name = "URL")] + pub upstream_anthropic: Option, + /// Provider label for the model route: `openai`, `anthropic`, … + #[arg(long, env = env::PROVIDER, value_name = "NAME")] + pub provider: Option, + /// Env var holding the upstream API key for this route. + #[arg(long, env = env::API_KEY_ENV, value_name = "VAR")] + pub api_key_env: Option, + /// Inline upstream API key (prefer `--api-key-env` in scripts). + #[arg(long, env = env::API_KEY, value_name = "KEY")] + pub api_key: Option, + /// Forward client model to another configured route name. + #[arg(long, env = env::FORWARD, value_name = "NAME")] + pub forward: Option, + /// Legacy API path prefix for `--model` route (prefer full prefix in `--upstream`). + #[arg(long, env = env::PATH_PREFIX, value_name = "PREFIX")] + pub path_prefix: Option, +} + +impl ProxyConfigOverrideArgs { + pub fn to_overrides(&self) -> Result { + let capture_level = match &self.capture_level { + Some(s) => Some(parse_capture_level(s)?), + None => None, + }; + Ok(ProxyConfigOverrides { + listen: self.listen.clone(), + admin_listen: self.admin_listen.clone(), + agent_id: self.agent_id.clone(), + session_header: self.session_header.clone(), + capture_level, + debug: None, + models_json: self.models_json.clone(), + models_toml: self.models_toml.clone(), + model_name: self.model.clone(), + upstream: self.upstream.clone(), + upstream_anthropic: self.upstream_anthropic.clone(), + provider: self.provider.clone(), + api_key_env: self.api_key_env.clone(), + api_key: self.api_key.clone(), + forward: self.forward.clone(), + path_prefix: self.path_prefix.clone(), + }) + } +} + +/// Optional proxy TOML path plus CLI/env overrides. +#[derive(Debug, Clone, Default, Args)] +pub struct ProxyConfigArgs { + /// Proxy config TOML (`listen`, `models`, …). Optional when set via env/flags. + #[arg( + long, + short = 'c', + value_name = "FILE", + env = env::CONFIG_FILE, + conflicts_with = "config_toml" + )] + pub config: Option, + /// Full proxy config as inline TOML (alternative to `-c`; supports every TOML field). + #[arg(long, env = env::CONFIG_TOML, value_name = "TOML", conflicts_with = "config")] + pub config_toml: Option, + #[command(flatten)] + pub overrides: ProxyConfigOverrideArgs, +} + +impl ProxyConfigArgs { + pub fn input(&self) -> Result { + Ok(ProxyConfigInput { + config_file: self.config.clone(), + config_toml: self.config_toml.clone(), + overrides: self.overrides.to_overrides()?, + }) + } + + pub fn resolve(&self) -> Result { + resolve_proxy_config(&self.input()?) + } + + pub fn resolve_with_debug(&self, cli_debug: bool) -> Result { + let mut input = self.input()?; + if cli_debug { + input.overrides.debug = Some(true); + } + resolve_proxy_config(&input) + } +} + +/// Resolved config and on-disk path (explicit `-c` or materialized `{storage}/proxy.toml`). +pub struct ResolvedProxyConfig { + pub config: ProxyConfig, + pub config_path: PathBuf, +} + +impl ProxyConfigArgs { + pub fn materialize(&self, storage: &std::path::Path, cli_debug: bool) -> Result { + let mut input = self.input()?; + if cli_debug { + input.overrides.debug = Some(true); + } + let (config, config_path) = + materialize_proxy_config(storage, &input).context("resolve proxy config")?; + Ok(ResolvedProxyConfig { + config, + config_path, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn override_args_map_capture_level() { + let args = ProxyConfigOverrideArgs { + capture_level: Some("full".into()), + upstream: Some("http://127.0.0.1:1/v1".into()), + ..Default::default() + }; + let o = args.to_overrides().unwrap(); + assert_eq!(o.capture_level, Some(CaptureLevel::Full)); + } +} diff --git a/crates/persisting-cli/src/main.rs b/crates/persisting-cli/src/main.rs index ee63094..47035a9 100644 --- a/crates/persisting-cli/src/main.rs +++ b/crates/persisting-cli/src/main.rs @@ -296,12 +296,11 @@ struct Cli { #[derive(Debug, Subcommand)] enum Command { Search(SearchArgs), - /// Agent trajectory: capture, proxy, inspect, repair - #[command(long_about = TRAJ_LONG_ABOUT)] + /// Agent trajectory: capture, proxy, inspect, repair（短名 `traj`） + #[command(visible_alias = "traj", long_about = TRAJ_LONG_ABOUT)] Trajectory(TrajectoryArgs), - /// Short alias for `trajectory` - #[command(name = "traj", long_about = TRAJ_LONG_ABOUT)] - Traj(TrajectoryArgs), + /// Run a compute plan (default). Use `--check` to validate locally first. + Compute(persisting_compute::ComputeArgs), } #[derive(Debug, Args)] @@ -1035,10 +1034,25 @@ fn engine_lib_names() -> [&'static str; 3] { fn main() -> Result<()> { let cli = Cli::parse_from(normalize_cli_args(std::env::args().collect())); - let mut lazy = LazyEngine::new(&cli); match &cli.command { - Command::Search(args) => run_search(&mut lazy, args)?, - Command::Trajectory(args) | Command::Traj(args) => run_trajectory(&mut lazy, args)?, + Command::Compute(args) => { + persisting_compute::cli::init_tracing_with_verbose(args.verbose); + let args = args.clone(); + let code = tokio::runtime::Runtime::new() + .context("tokio runtime")? + .block_on(persisting_compute::run_compute(args))?; + if code != std::process::ExitCode::SUCCESS { + std::process::exit(1); + } + } + Command::Search(args) => { + let mut lazy = LazyEngine::new(&cli); + run_search(&mut lazy, args)?; + } + Command::Trajectory(args) => { + let mut lazy = LazyEngine::new(&cli); + run_trajectory(&mut lazy, args)?; + } } Ok(()) } @@ -2525,11 +2539,11 @@ mod tests { fn add_args_from_cli(argv: &[&str]) -> TrajectoryAddArgs { let cli = Cli::try_parse_from(argv).unwrap(); - let Command::Traj(TrajectoryArgs { + let Command::Trajectory(TrajectoryArgs { command: TrajectoryCommand::Add(args), }) = cli.command else { - panic!("expected traj add"); + panic!("expected traj/trajectory add"); }; args } diff --git a/crates/persisting-compute/Cargo.toml b/crates/persisting-compute/Cargo.toml new file mode 100644 index 0000000..d9000a8 --- /dev/null +++ b/crates/persisting-compute/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "persisting-compute" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +description = "Thin compute control plane: plan script → streamed tasks → Pulsing workers (via `persisting compute`)" + +[dependencies] +anyhow = "1" +async-trait = "0.1" +chrono = { version = "0.4", optional = true } +clap = { version = "4", features = ["derive", "env"] } +futures = "0.3" +persisting-capture = { path = "../persisting-capture", optional = true } +persisting-engine = { path = "../persisting-engine", optional = true } +persisting-proto = { path = "../persisting-proto", optional = true } +pulsing-actor = { workspace = true } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tokio = { version = "1", features = ["macros", "rt-multi-thread", "process", "io-util", "sync", "time", "signal", "fs"] } +tokio-stream = "0.1" +tokio-util = { version = "0.7", features = ["rt"] } +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } +uuid = { version = "1", features = ["v4"] } + +[features] +default = [] +# Append compute results to Vortex via TrajectoryAppend (Tee with JsonlFileSink). +traj-sink = ["dep:chrono", "dep:persisting-capture", "dep:persisting-engine", "dep:persisting-proto"] + +[dev-dependencies] +tempfile = "3" +tokio = { version = "1", features = ["macros", "rt-multi-thread"] } diff --git a/crates/persisting-compute/src/blocks.rs b/crates/persisting-compute/src/blocks.rs new file mode 100644 index 0000000..b779755 --- /dev/null +++ b/crates/persisting-compute/src/blocks.rs @@ -0,0 +1,110 @@ +//! Index of **semantic primitives** — one contract per owning module. +//! +//! # Testing policy +//! +//! - **Unit / contract tests** live in the **same file** as the primitive +//! (`#[cfg(test)]` at the bottom of `task.rs`, `scheduler.rs`, …). +//! - **`tests/`** holds **integration** only (multi-module paths: fleet, resume +//! through Driver, argv end-to-end, …). +//! +//! | Primitive | Interface (types / traits / fns) | Module | +//! |-----------|----------------------------------|--------| +//! | Task wire | [`TaskExpr`], [`TaskResult`] | [`crate::task`] | +//! | Placement | [`Scheduler`], sticky-only / quarantine | [`crate::scheduler`] | +//! | Slot naming / dist | [`DistEnv`] | [`crate::dist`] | +//! | Run future | [`RunFuture`], [`wait_all`] | [`crate::future`] | +//! | Idempotency cache | [`ResultCache`] | [`crate::result_cache`] | +//! | Live skip / claim | [`SkipSet`] | [`crate::skip`] | +//! | Result sink | [`ResultSink`], [`persist_terminal`] | [`crate::sink`] | +//! | Async sink writer | [`SinkSubmitter`], [`spawn_sink_writer`] | [`crate::sink_writer`] | +//! | Checkpoint | [`CheckpointLedger`], [`CheckpointTracker`] | [`crate::checkpoint`] | +//! | Plan emit | [`stream_plan_tasks`] | [`crate::plan`] | +//! | Execute host | [`Executor`], [`ExecutorRouter`] | [`crate::executor`] | +//! | Worker seam | [`WorkerActor`], supervised spawn | [`crate::worker`] | +//! | Job cancel | [`JobControlActor`], DeathWatch | [`crate::job_control`] | +//! | Pulsing helpers | resolve / ask_timeout / spawn_supervised | [`crate::pulsing_ext`] | +//! | Driver | [`Driver`], [`RunOptions`] | [`crate::driver`] | +//! | Fleet boot | [`run_local_fleet`], [`run_fleet`] | [`crate::runtime`] | +//! | Observe | [`Observer`] | [`crate::observe`] | +//! | Python env | [`merge_pythonpath_parts`] | [`crate::python_env`] | +//! +//! [`TaskExpr`]: crate::task::TaskExpr +//! [`TaskResult`]: crate::task::TaskResult +//! [`Scheduler`]: crate::scheduler::Scheduler +//! [`DistEnv`]: crate::dist::DistEnv +//! [`RunFuture`]: crate::future::RunFuture +//! [`wait_all`]: crate::future::wait_all +//! [`ResultCache`]: crate::result_cache::ResultCache +//! [`SkipSet`]: crate::skip::SkipSet +//! [`ResultSink`]: crate::sink::ResultSink +//! [`persist_terminal`]: crate::sink::persist_terminal +//! [`SinkSubmitter`]: crate::sink_writer::SinkSubmitter +//! [`spawn_sink_writer`]: crate::sink_writer::spawn_sink_writer +//! [`CheckpointLedger`]: crate::checkpoint::CheckpointLedger +//! [`CheckpointTracker`]: crate::checkpoint::CheckpointTracker +//! [`stream_plan_tasks`]: crate::plan::stream_plan_tasks +//! [`Executor`]: crate::executor::Executor +//! [`ExecutorRouter`]: crate::executor::ExecutorRouter +//! [`WorkerActor`]: crate::worker::WorkerActor +//! [`WorkerCommand`]: crate::worker::WorkerCommand +//! [`JobControlActor`]: crate::job_control::JobControlActor +//! [`Driver`]: crate::driver::Driver +//! [`RunOptions`]: crate::driver::RunOptions +//! [`run_local_fleet`]: crate::runtime::run_local_fleet +//! [`run_fleet`]: crate::runtime::run_fleet +//! [`Observer`]: crate::observe::Observer +//! [`merge_pythonpath_parts`]: crate::python_env::merge_pythonpath_parts + +/// Stable ids for docs / observability (not a separate test suite). +pub mod ids { + pub const TASK_WIRE: &str = "task_wire"; + pub const PLACEMENT: &str = "placement"; + pub const RUN_FUTURE: &str = "run_future"; + pub const IDEMPOTENCY: &str = "idempotency"; + pub const SKIP: &str = "skip"; + pub const SINK: &str = "sink"; + pub const SINK_WRITER: &str = "sink_writer"; + pub const CHECKPOINT: &str = "checkpoint"; + pub const PLAN: &str = "plan"; + pub const EXECUTE: &str = "execute"; + pub const WORKER: &str = "worker"; + pub const JOB_CONTROL: &str = "job_control"; + pub const PULSING_EXT: &str = "pulsing_ext"; + pub const DRIVER: &str = "driver"; + pub const FLEET: &str = "fleet"; + pub const OBSERVE: &str = "observe"; + pub const PYTHON_ENV: &str = "python_env"; + + pub const ALL: &[&str] = &[ + TASK_WIRE, + PYTHON_ENV, + PLACEMENT, + RUN_FUTURE, + IDEMPOTENCY, + SKIP, + SINK, + SINK_WRITER, + CHECKPOINT, + OBSERVE, + PLAN, + EXECUTE, + WORKER, + JOB_CONTROL, + PULSING_EXT, + DRIVER, + FLEET, + ]; +} + +#[cfg(test)] +mod tests { + use super::ids; + + #[test] + fn primitive_ids_unique() { + let mut seen = std::collections::HashSet::new(); + for id in ids::ALL { + assert!(seen.insert(*id), "duplicate id {id}"); + } + } +} diff --git a/crates/persisting-compute/src/check.rs b/crates/persisting-compute/src/check.rs new file mode 100644 index 0000000..9395147 --- /dev/null +++ b/crates/persisting-compute/src/check.rs @@ -0,0 +1,344 @@ +//! Local validation: prove env + plan + execute before scale-out. +//! +//! ```text +//! persisting compute plan.py --check --python python3 +//! ``` + +use crate::plan::stream_plan_tasks; +use crate::python_env::{self, pythonpath_for_script}; +use crate::runtime::{run_local_fleet, RunOptions}; +use crate::task::TaskExpr; +use anyhow::{bail, Context, Result}; +use futures::StreamExt; +use serde_json::{json, Value}; +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; +use std::process::Stdio; +use tokio::process::Command; + +#[derive(Debug, Clone)] +pub struct CheckOptions { + pub script: PathBuf, + pub python: PathBuf, + /// Max tasks to actually execute (0 = all). + pub limit: usize, + pub workers: usize, + pub verbose: bool, + pub pythonpath_extra: Vec, + /// Forwarded to `task.py` as `sys.argv[1:]` (after `--`). + pub script_args: Vec, +} + +#[derive(Debug, Default)] +pub struct CheckReport { + pub python_ok: bool, + pub python_version: Option, + pub plan_tasks: usize, + pub plan_ops: BTreeMap, + pub execute_ok: bool, + pub run_ok: usize, + pub run_fail: usize, + pub errors: Vec, +} + +impl CheckReport { + pub fn passed(&self) -> bool { + self.python_ok + && self.errors.is_empty() + && self.plan_tasks > 0 + && self.execute_ok + && self.run_fail == 0 + } + + pub fn to_json(&self) -> Value { + json!({ + "ok": self.passed(), + "python": { "ok": self.python_ok, "version": self.python_version }, + "plan": { "tasks": self.plan_tasks, "ops": self.plan_ops }, + "execute": { "ok": self.execute_ok }, + "run": { "ok": self.run_ok, "failed": self.run_fail }, + "errors": self.errors, + }) + } +} + +/// Full local check pipeline. Progress → stderr; JSON summary → stdout. +pub async fn run_check(opts: CheckOptions) -> Result { + let mut report = CheckReport::default(); + let mut extras = pythonpath_for_script(&opts.script); + extras.extend(opts.pythonpath_extra.iter().cloned()); + + eprint_stage(1, 4, "python env"); + match probe_python(&opts.python).await { + Ok(ver) => { + report.python_ok = true; + report.python_version = Some(ver.clone()); + eprintln!(" OK {} ({})", opts.python.display(), ver.trim()); + } + Err(e) => { + report.python_ok = false; + report.errors.push(format!("python: {e:#}")); + eprintln!(" FAIL {e:#}"); + print_summary(&report); + return Ok(report); + } + } + if let Some(pp) = python_env::merge_pythonpath(&extras) { + eprintln!(" PYTHONPATH+= {}", shorten_pp(&pp)); + } + + eprint_stage(2, 4, "plan emit + schema"); + let tasks = + match collect_plan_tasks(&opts.script, &opts.python, &extras, &opts.script_args).await { + Ok(t) => t, + Err(e) => { + report.errors.push(format!("plan: {e:#}")); + eprintln!(" FAIL {e:#}"); + print_summary(&report); + return Ok(report); + } + }; + if tasks.is_empty() { + report.errors.push("plan emitted zero tasks".into()); + eprintln!(" FAIL no tasks"); + print_summary(&report); + return Ok(report); + } + report.plan_tasks = tasks.len(); + for t in &tasks { + *report.plan_ops.entry(t.op.clone()).or_default() += 1; + } + eprintln!( + " OK {} task(s) ops={}", + tasks.len(), + format_ops(&report.plan_ops) + ); + if opts.verbose { + for t in tasks.iter().take(5) { + eprintln!(" {}", t.to_ndjson().unwrap_or_default()); + } + if tasks.len() > 5 { + eprintln!(" … {} more", tasks.len() - 5); + } + } + + eprint_stage(3, 4, "resolve execute"); + match probe_plan_execute(&opts.python, &opts.script, &extras).await { + Ok(()) => { + report.execute_ok = true; + eprintln!(" OK {}::execute(item)", opts.script.display()); + } + Err(e) => { + report.execute_ok = false; + report.errors.push(format!("execute: {e:#}")); + eprintln!(" FAIL execute: {e:#}"); + print_summary(&report); + return Ok(report); + } + } + + eprint_stage(4, 4, "local run"); + let to_run = if opts.limit == 0 { + tasks.len() + } else { + opts.limit.min(tasks.len()) + }; + // Skip the remainder so --limit actually bounds execute without changing plan(). + let skip = if to_run < tasks.len() { + eprintln!( + " note: --limit {to_run} (skipping {} of {} plan tasks)", + tasks.len() - to_run, + tasks.len() + ); + crate::skip::SkipSet::from_iter(tasks[to_run..].iter().map(|t| t.id.clone())) + } else { + crate::skip::SkipSet::new() + }; + + let run_opts = RunOptions { + script: opts.script.clone(), + python: opts.python.clone(), + workers: opts.workers.max(1), + max_inflight: opts.workers.max(1), + per_worker_inflight: 1, + pythonpath_extra: extras.clone(), + script_args: opts.script_args.clone(), + infra_retries: 2, + job_cancel: tokio_util::sync::CancellationToken::new(), + observer: crate::observe::Observer::disabled(), + skip_task_ids: skip, + checkpoint: None, + sink_submitter: None, + }; + + match run_local_fleet(run_opts, |_| {}).await { + Ok(results) => { + for r in &results { + if r.ok { + report.run_ok += 1; + } else { + report.run_fail += 1; + let msg = format!( + "task {}: {}", + r.task_id, + r.error.clone().unwrap_or_else(|| "failed".into()) + ); + report.errors.push(msg.clone()); + if opts.verbose { + if let Some(tb) = &r.traceback { + eprintln!(" FAIL {msg}\n{tb}"); + } else { + eprintln!(" FAIL {msg}"); + } + } + } + } + eprintln!( + " {} {}/{} ok", + if report.run_fail == 0 { "OK" } else { "FAIL" }, + report.run_ok, + results.len() + ); + } + Err(e) => { + report.errors.push(format!("run: {e:#}")); + eprintln!(" FAIL {e:#}"); + } + } + + print_summary(&report); + Ok(report) +} + +fn eprint_stage(n: u32, total: u32, title: &str) { + eprintln!("[{n}/{total}] {title}"); +} + +fn print_summary(report: &CheckReport) { + println!("{}", report.to_json()); +} + +fn format_ops(ops: &BTreeMap) -> String { + ops.iter() + .map(|(k, v)| format!("{k}×{v}")) + .collect::>() + .join(", ") +} + +fn shorten_pp(pp: &str) -> String { + let parts: Vec = std::env::split_paths(pp) + .map(|p| p.to_string_lossy().into_owned()) + .collect(); + if parts.len() <= 3 { + parts.join(":") + } else { + format!("{} … (+{} paths)", parts[..2].join(":"), parts.len() - 2) + } +} + +async fn probe_python(python: &Path) -> Result { + let out = Command::new(python) + .args(["-c", "import sys; print(sys.version.split()[0])"]) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .output() + .await + .with_context(|| format!("exec {}", python.display()))?; + if !out.status.success() { + bail!( + "exit {:?}: {}", + out.status.code(), + String::from_utf8_lossy(&out.stderr).trim() + ); + } + Ok(String::from_utf8_lossy(&out.stdout).trim().to_string()) +} + +async fn collect_plan_tasks( + script: &Path, + python: &Path, + extras: &[PathBuf], + script_args: &[String], +) -> Result> { + if let Some(pp) = python_env::merge_pythonpath(extras) { + std::env::set_var("PYTHONPATH", pp); + } + let mut stream = stream_plan_tasks( + script.to_path_buf(), + python.to_path_buf(), + script_args.to_vec(), + ); + let mut tasks = Vec::new(); + while let Some(item) = stream.next().await { + tasks.push(item?); + } + Ok(tasks) +} + +async fn probe_plan_execute(python: &Path, script: &Path, extras: &[PathBuf]) -> Result<()> { + let script = script + .canonicalize() + .with_context(|| format!("plan script {}", script.display()))?; + let code = r#" +import importlib.util, sys +from pathlib import Path +path = Path(sys.argv[1]) +spec = importlib.util.spec_from_file_location("user_plan_probe", path) +mod = importlib.util.module_from_spec(spec) +spec.loader.exec_module(mod) +if not hasattr(mod, "execute") or not callable(mod.execute): + raise SystemExit("plan must define execute(item)") +print("ok", flush=True) +"#; + let mut cmd = Command::new(python); + cmd.arg("-c") + .arg(code) + .arg(&script) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + python_env::apply_pythonpath(&mut cmd, extras); + let out = cmd.output().await.context("probe execute")?; + if !out.status.success() { + bail!("{}", String::from_utf8_lossy(&out.stderr).trim()); + } + Ok(()) +} + +/// Built-in smoke: no user plan. Proves `plan()` + `execute(item)` end-to-end. +pub async fn run_self_test(python: PathBuf, workers: usize, verbose: bool) -> Result { + let dir = std::env::temp_dir().join(format!( + "persisting-compute-selftest-{}", + std::process::id() + )); + std::fs::create_dir_all(&dir).with_context(|| format!("mkdir {}", dir.display()))?; + + let plan_py = dir.join("plan_smoke.py"); + std::fs::write( + &plan_py, + r#" +def plan(): + yield {"id": "t0", "x": 1} + yield {"id": "t1", "x": 2} + yield {"id": "t2", "x": 3} + +def execute(item): + x = item["x"] + return {"x": x, "x2": x * 2} +"#, + )?; + + eprintln!("self-test: built-in plan() + execute(item) (no user script)"); + let report = run_check(CheckOptions { + script: plan_py, + python, + limit: 0, + workers: workers.max(1), + verbose, + pythonpath_extra: vec![dir.clone()], + script_args: vec![], + }) + .await; + + let _ = std::fs::remove_dir_all(&dir); + report +} diff --git a/crates/persisting-compute/src/checkpoint.rs b/crates/persisting-compute/src/checkpoint.rs new file mode 100644 index 0000000..59c4f5a --- /dev/null +++ b/crates/persisting-compute/src/checkpoint.rs @@ -0,0 +1,288 @@ +//! Checkpoint ledger: resume unfinished work from `--sink` + progress file. +//! +//! Done set = `task_id`s already in `ready.ndjson` / `failures.ndjson`. +//! Progress snapshot = `checkpoint.json` (throttled writes). + +use crate::task::unix_now; +use anyhow::{Context, Result}; +use serde::{Deserialize, Serialize}; +use std::collections::HashSet; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Mutex; +use std::time::Instant; +use tokio::fs; +use tokio::io::{AsyncBufReadExt, BufReader}; + +/// Terminal task ids already recorded under a sink root. +#[derive(Debug, Clone, Default)] +pub struct CheckpointLedger { + pub ready: HashSet, + pub failed: HashSet, +} + +impl CheckpointLedger { + pub fn skip_ids(&self) -> HashSet { + let mut s = self.ready.clone(); + s.extend(self.failed.iter().cloned()); + s + } + + pub async fn load(root: &Path) -> Result { + let ready = load_task_ids(&root.join("ready.ndjson")).await?; + let failed = load_task_ids(&root.join("failures.ndjson")).await?; + Ok(Self { ready, failed }) + } +} + +async fn load_task_ids(path: &Path) -> Result> { + let mut out = HashSet::new(); + if !path.exists() { + return Ok(out); + } + let f = fs::File::open(path) + .await + .with_context(|| format!("open {}", path.display()))?; + let mut lines = BufReader::new(f).lines(); + let mut line_no = 0u64; + while let Some(line) = lines.next_line().await? { + line_no += 1; + let line = line.trim(); + if line.is_empty() { + continue; + } + let v: serde_json::Value = match serde_json::from_str(line) { + Ok(v) => v, + Err(e) => { + tracing::warn!( + path = %path.display(), + line = line_no, + error = %e, + "skipping corrupt checkpoint JSONL line" + ); + continue; + } + }; + if let Some(id) = v + .get("task_id") + .and_then(|x| x.as_str()) + .map(|s| s.to_string()) + { + out.insert(id); + } else { + tracing::warn!( + path = %path.display(), + line = line_no, + "skipping checkpoint line without task_id" + ); + } + } + Ok(out) +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct CheckpointProgress { + pub ok: u64, + pub fail: u64, + pub cancelled: u64, + pub skipped: u64, + pub dispatched: u64, + pub updated_at: f64, +} + +/// Live progress tracker with throttled `checkpoint.json` writes. +pub struct CheckpointTracker { + root: PathBuf, + state: Mutex, + last_flush: Mutex, + flush_every_ms: u64, + dirty: AtomicU64, +} + +impl CheckpointTracker { + pub fn new(root: impl Into) -> Self { + Self { + root: root.into(), + state: Mutex::new(CheckpointProgress::default()), + last_flush: Mutex::new(Instant::now() - std::time::Duration::from_secs(2)), + flush_every_ms: 1000, + dirty: AtomicU64::new(0), + } + } + + pub fn seed_from_ledger(&self, ledger: &CheckpointLedger) { + let mut g = self.state.lock().unwrap_or_else(|e| e.into_inner()); + g.ok = ledger.ready.len() as u64; + g.fail = ledger.failed.len() as u64; + g.updated_at = unix_now(); + } + + pub fn note_skipped(&self, n: u64) { + let mut g = self.state.lock().unwrap_or_else(|e| e.into_inner()); + g.skipped = g.skipped.saturating_add(n); + g.updated_at = unix_now(); + self.dirty.fetch_add(1, Ordering::Relaxed); + } + + pub fn note_dispatched(&self) { + let mut g = self.state.lock().unwrap_or_else(|e| e.into_inner()); + g.dispatched = g.dispatched.saturating_add(1); + g.updated_at = unix_now(); + self.dirty.fetch_add(1, Ordering::Relaxed); + } + + pub fn note_terminal(&self, ok: bool, cancelled: bool) { + let mut g = self.state.lock().unwrap_or_else(|e| e.into_inner()); + if cancelled { + g.cancelled = g.cancelled.saturating_add(1); + } else if ok { + g.ok = g.ok.saturating_add(1); + } else { + g.fail = g.fail.saturating_add(1); + } + g.updated_at = unix_now(); + self.dirty.fetch_add(1, Ordering::Relaxed); + } + + pub fn snapshot(&self) -> CheckpointProgress { + self.state.lock().unwrap_or_else(|e| e.into_inner()).clone() + } + + pub async fn maybe_flush(&self) -> Result<()> { + if self.dirty.load(Ordering::Relaxed) == 0 { + return Ok(()); + } + let should = { + let last = self.last_flush.lock().unwrap_or_else(|e| e.into_inner()); + last.elapsed().as_millis() as u64 >= self.flush_every_ms + }; + if should { + self.flush().await?; + } + Ok(()) + } + + pub async fn flush(&self) -> Result<()> { + let snap = self.snapshot(); + let path = self.root.join("checkpoint.json"); + let tmp = self.root.join("checkpoint.json.tmp"); + let body = serde_json::to_vec_pretty(&snap)?; + fs::write(&tmp, &body) + .await + .with_context(|| format!("write {}", tmp.display()))?; + fs::rename(&tmp, &path) + .await + .with_context(|| format!("rename {}", path.display()))?; + if let Ok(mut last) = self.last_flush.lock() { + *last = Instant::now(); + } + self.dirty.store(0, Ordering::Relaxed); + Ok(()) + } + + pub fn summary_line(&self) -> String { + let s = self.snapshot(); + format!( + "[ckpt] ok={} fail={} cancelled={} skipped={} dispatched={}", + s.ok, s.fail, s.cancelled, s.skipped, s.dispatched + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::AsyncWriteExt; + + #[tokio::test] + async fn load_ready_and_failures() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + let mut ready = fs::File::create(root.join("ready.ndjson")).await.unwrap(); + ready + .write_all(br#"{"task_id":"t-0","ok":true}"#) + .await + .unwrap(); + ready.write_all(b"\n").await.unwrap(); + let mut fail = fs::File::create(root.join("failures.ndjson")) + .await + .unwrap(); + fail.write_all(br#"{"task_id":"t-1","ok":false}"#) + .await + .unwrap(); + fail.write_all(b"\n").await.unwrap(); + + let ledger = CheckpointLedger::load(root).await.unwrap(); + assert!(ledger.ready.contains("t-0")); + assert!(ledger.failed.contains("t-1")); + assert_eq!(ledger.skip_ids().len(), 2); + } + + #[tokio::test] + async fn tracker_flush_writes_checkpoint_json() { + let dir = tempfile::tempdir().unwrap(); + let tracker = CheckpointTracker::new(dir.path()); + tracker.note_dispatched(); + tracker.note_terminal(true, false); + tracker.note_skipped(2); + tracker.flush().await.unwrap(); + let body = fs::read_to_string(dir.path().join("checkpoint.json")) + .await + .unwrap(); + let snap: CheckpointProgress = serde_json::from_str(&body).unwrap(); + assert_eq!(snap.ok, 1); + assert_eq!(snap.dispatched, 1); + assert_eq!(snap.skipped, 2); + } + + #[tokio::test] + async fn empty_dir_loads_empty_ledger() { + let dir = tempfile::tempdir().unwrap(); + let ledger = CheckpointLedger::load(dir.path()).await.unwrap(); + assert!(ledger.skip_ids().is_empty()); + } + + #[tokio::test] + async fn skip_ids_is_ready_union_failed() { + let dir = tempfile::tempdir().unwrap(); + let mut ready = fs::File::create(dir.path().join("ready.ndjson")) + .await + .unwrap(); + ready.write_all(br#"{"task_id":"ok"}"#).await.unwrap(); + ready.write_all(b"\n").await.unwrap(); + let mut fail = fs::File::create(dir.path().join("failures.ndjson")) + .await + .unwrap(); + fail.write_all(br#"{"task_id":"bad"}"#).await.unwrap(); + fail.write_all(b"\n").await.unwrap(); + let ledger = CheckpointLedger::load(dir.path()).await.unwrap(); + let skip = ledger.skip_ids(); + assert!(skip.contains("ok") && skip.contains("bad")); + let tracker = CheckpointTracker::new(dir.path()); + tracker.seed_from_ledger(&ledger); + assert_eq!(tracker.snapshot().ok, 1); + assert_eq!(tracker.snapshot().fail, 1); + } + + #[tokio::test] + async fn load_skips_corrupt_lines_and_keeps_valid() { + let dir = tempfile::tempdir().unwrap(); + let mut ready = fs::File::create(dir.path().join("ready.ndjson")) + .await + .unwrap(); + ready + .write_all( + br#"{"task_id":"t-0","ok":true} +not-json +{"task_id":"t-1","ok":true} +{"no_id":true} +"#, + ) + .await + .unwrap(); + let ledger = CheckpointLedger::load(dir.path()).await.unwrap(); + assert!(ledger.ready.contains("t-0")); + assert!(ledger.ready.contains("t-1")); + assert_eq!(ledger.ready.len(), 2); + } +} diff --git a/crates/persisting-compute/src/cli.rs b/crates/persisting-compute/src/cli.rs new file mode 100644 index 0000000..5bef09c --- /dev/null +++ b/crates/persisting-compute/src/cli.rs @@ -0,0 +1,389 @@ +//! CLI surface for compute — used by `persisting compute [plan]`. +//! +//! - With script → **run** (default); `--check` validates first +//! - `--self-test` → built-in smoke (no user plan) + +use crate::check::{run_check, run_self_test, CheckOptions}; +use crate::checkpoint::{CheckpointLedger, CheckpointTracker}; +use crate::observe::{Observer, ObserverOptions}; +use crate::runtime::{run_fleet, RunOptions}; +use crate::sink::{JsonlFileSink, ResultSink, TeeSink}; +use crate::sink_writer::spawn_sink_writer; +use crate::skip::SkipSet; +use crate::task::TaskResult; +use anyhow::{bail, Context, Result}; +use clap::{Args, ValueEnum}; +use std::path::PathBuf; +use std::process::ExitCode; +use std::sync::Arc; +use tokio_util::sync::CancellationToken; + +/// `persisting compute [plan.py] …` +#[derive(Debug, Clone, Args)] +#[command(about = "Run a compute plan (`plan()` + `execute(item)`).")] +pub struct ComputeArgs { + /// Plan script (`plan()` / `execute`). Required unless `--self-test`. + #[arg(value_name = "SCRIPT")] + pub script: Option, + + /// Built-in smoke test (no user plan). + #[arg(long)] + pub self_test: bool, + + /// Validate env + plan + execute (+ sample run) instead of a full run. + #[arg(long)] + pub check: bool, + + #[arg(short = 'w', long, default_value_t = 4)] + pub workers: usize, + + /// Concurrent Execute slots per logical worker/rank. Each slot is its own + /// WorkerActor + Python host (real parallelism). Default 1 — best when task + /// times vary widely across workers. + #[arg(long, default_value_t = 1)] + pub per_worker: usize, + + /// Global inflight cap (default: workers × per-worker, or WORLD_SIZE × per-worker under torchrun). + #[arg(long)] + pub max_inflight: Option, + + /// L2 infrastructure retries when a worker ask fails (not semantic retry). + #[arg(long, default_value_t = 2)] + pub retries: u32, + + /// Durable unique sink directory (`ready.ndjson` + `failures.ndjson` + `checkpoint.json`). + #[arg(long, value_name = "DIR")] + pub sink: Option, + + /// Resume from `--sink`: skip task ids already in ready/failures. + #[arg(long)] + pub resume: bool, + + /// Also append terminal results to a Vortex trajectory (requires `--sink`). + /// Writes `compute.result` / `compute.failure` events under traj storage. + #[cfg(feature = "traj-sink")] + #[arg(long)] + pub traj: bool, + + /// Vortex storage root (default: `{sink}/traj`). + #[cfg(feature = "traj-sink")] + #[arg(long, value_name = "DIR")] + pub traj_storage: Option, + + /// Trajectory agent_id (default: `compute`). + #[cfg(feature = "traj-sink")] + #[arg(long, default_value = "compute")] + pub traj_agent: String, + + /// Trajectory session_id (default: sink directory name). + #[cfg(feature = "traj-sink")] + #[arg(long)] + pub traj_session: Option, + + #[arg(long, env = "PERSISTING_PYTHON", default_value = "python3")] + pub python: PathBuf, + + /// Extra PYTHONPATH entries (plan dir is always added). + #[arg(short = 'E', long = "pythonpath")] + pub pythonpath: Vec, + + /// Result stream on stdout. Default: `ndjson`; with `--observe` default becomes `quiet` + /// so progress is not drowned out (pass `--results ndjson` to keep both). + #[arg(long, value_enum)] + pub results: Option, + + /// With `--check`: max tasks to execute (0 = all). Ignored on normal run. + #[arg(long, default_value_t = 0)] + pub limit: usize, + + /// Live queue / placement / duration progress on stderr (`[obs] …`). + /// Must be **before** `--`. Env: `PERSISTING_OBSERVE=1`. + #[arg(long, env = "PERSISTING_OBSERVE")] + pub observe: bool, + + /// Append observe events as NDJSON to FILE (implies observe). + #[arg(long, value_name = "FILE", env = "PERSISTING_OBSERVE_FILE")] + pub observe_file: Option, + + /// Also emit observe NDJSON on stderr (in addition to `[obs]` human lines). + #[arg(long)] + pub observe_json: bool, + + /// More stderr tracing (`persisting_compute` + Pulsing at info). + #[arg(short, long)] + pub verbose: bool, + + /// Args forwarded to the plan script (`sys.argv[1:]`). Put after `--`. + /// Example: `persisting compute task.py -- --model x --n 2` + #[arg(last = true, value_name = "SCRIPT_ARGS")] + pub script_args: Vec, +} + +#[derive(Clone, Copy, Debug, ValueEnum)] +pub enum ResultsFormat { + Ndjson, + Summary, + Quiet, +} + +/// Run compute (async). Caller owns the Tokio runtime. +pub async fn run_compute(args: ComputeArgs) -> Result { + if args.self_test { + let report = run_self_test(args.python, args.workers, args.verbose).await?; + return Ok(if report.passed() { + ExitCode::SUCCESS + } else { + ExitCode::FAILURE + }); + } + + let Some(script) = args.script else { + bail!("missing SCRIPT; pass a plan file, or use --self-test"); + }; + + if args.check { + let report = run_check(CheckOptions { + script, + python: args.python, + limit: args.limit, + workers: args.workers.max(1), + verbose: args.verbose, + pythonpath_extra: args.pythonpath, + script_args: args.script_args, + }) + .await?; + return Ok(if report.passed() { + ExitCode::SUCCESS + } else { + ExitCode::FAILURE + }); + } + + let under_torch = std::env::var_os("RANK").is_some(); + let per_worker = args.per_worker.max(1); + let max_inflight = args.max_inflight.unwrap_or_else(|| { + if under_torch { + let ws: usize = std::env::var("WORLD_SIZE") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(4); + ws.saturating_mul(per_worker).max(1) + } else { + args.workers.saturating_mul(per_worker).max(1) + } + }); + + let job_cancel = CancellationToken::new(); + let cancel_bg = job_cancel.clone(); + tokio::spawn(async move { + let _ = tokio::signal::ctrl_c().await; + tracing::warn!("Ctrl-C: cancelling job (RunFutures)"); + cancel_bg.cancel(); + }); + + let file_sink: Option> = if let Some(dir) = &args.sink { + let mut sinks: Vec> = Vec::new(); + sinks.push(Box::new( + JsonlFileSink::open(dir).await.context("open jsonl sink")?, + )); + #[cfg(feature = "traj-sink")] + if args.traj { + let traj_storage = args + .traj_storage + .clone() + .unwrap_or_else(|| dir.join("traj")); + tokio::fs::create_dir_all(&traj_storage) + .await + .with_context(|| format!("mkdir traj {}", traj_storage.display()))?; + let session = args.traj_session.clone().unwrap_or_else(|| { + dir.file_name() + .and_then(|s| s.to_str()) + .unwrap_or("run") + .to_string() + }); + let vortex = crate::sink_traj::VortexResultSink::new( + traj_storage.display().to_string(), + args.traj_agent.clone(), + session.clone(), + ); + if let Ok(ledger) = CheckpointLedger::load(dir).await { + vortex.seed_seen(ledger.skip_ids()); + } + eprintln!( + "[traj] vortex append → {}/{}/{}", + traj_storage.display(), + args.traj_agent, + session + ); + sinks.push(Box::new(vortex)); + } + Some(Arc::new(TeeSink::new(sinks))) + } else { + None + }; + + if args.resume && file_sink.is_none() { + bail!("--resume requires --sink DIR"); + } + + #[cfg(feature = "traj-sink")] + if args.traj && args.sink.is_none() { + bail!("--traj requires --sink DIR (JSONL ledger for --resume)"); + } + + let (skip_task_ids, checkpoint) = if let Some(dir) = &args.sink { + let tracker = Arc::new(CheckpointTracker::new(dir.clone())); + if args.resume { + let ledger = CheckpointLedger::load(dir) + .await + .context("load checkpoint ledger")?; + let skip = ledger.skip_ids(); + eprintln!( + "[ckpt] resume: ready={} fail={} skip_total={}", + ledger.ready.len(), + ledger.failed.len(), + skip.len() + ); + tracker.seed_from_ledger(&ledger); + (SkipSet::from_iter(skip), Some(tracker)) + } else { + (SkipSet::new(), Some(tracker)) + } + } else { + (SkipSet::new(), None) + }; + + let observe_on = args.observe || args.observe_file.is_some() || args.observe_json; + let observer = if observe_on { + Observer::open(ObserverOptions { + human: true, + json_stderr: args.observe_json, + path: args.observe_file.clone(), + }) + .await + .context("open observe sink")? + } else { + Observer::disabled() + }; + + // With observe, default to quiet results so `[obs]` lines are visible. + let results_fmt = args.results.unwrap_or(if observe_on { + ResultsFormat::Quiet + } else { + ResultsFormat::Ndjson + }); + + let sink_writer = file_sink.as_ref().map(|sink| { + spawn_sink_writer( + Arc::clone(sink), + checkpoint.clone(), + Some(skip_task_ids.clone()), + max_inflight.saturating_mul(2).max(16), + ) + }); + let sink_submit = sink_writer.as_ref().map(|w| w.submitter()); + + let opts = RunOptions { + script, + python: args.python, + workers: args.workers, + max_inflight, + per_worker_inflight: per_worker, + pythonpath_extra: args.pythonpath, + script_args: args.script_args, + infra_retries: args.retries, + job_cancel, + observer, + skip_task_ids, + checkpoint: checkpoint.clone(), + sink_submitter: sink_submit, + }; + + let collected = run_fleet(opts, move |r: TaskResult| { + if matches!(results_fmt, ResultsFormat::Ndjson) { + if let Ok(line) = r.to_ndjson() { + println!("{line}"); + } + } + }) + .await + .context("run fleet")?; + + if let Some(w) = sink_writer { + w.join().await.context("sink persist")?; + } + + if let Some(ckpt) = &checkpoint { + let _ = ckpt.flush().await; + eprintln!("{}", ckpt.summary_line()); + } + + if collected.is_empty() && under_torch { + let rank: usize = std::env::var("RANK") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0); + if rank != 0 { + return Ok(ExitCode::SUCCESS); + } + } + + let failed = collected.iter().filter(|r| !r.ok || r.cancelled).count(); + + if matches!(results_fmt, ResultsFormat::Summary) { + let ok = collected.iter().filter(|r| r.ok && !r.cancelled).count(); + println!( + "{}", + serde_json::json!({ + "total": collected.len(), + "ok": ok, + "failed": failed, + "cancelled": collected.iter().filter(|r| r.cancelled).count(), + "sink": args.sink.as_ref().map(|p| p.display().to_string()), + }) + ); + for r in &collected { + if !r.ok { + if let Ok(line) = r.to_ndjson() { + eprintln!("{line}"); + } + } + } + } + + Ok(if failed == 0 { + ExitCode::SUCCESS + } else { + ExitCode::FAILURE + }) +} + +/// Ensure tracing is initialized once (safe to call from nested CLI). +/// +/// Default is quiet: hush Pulsing actor lifecycle noise. Override with `RUST_LOG`, +/// or pass `--verbose` for `persisting_compute=info,pulsing_actor=info`. +pub fn init_tracing() { + init_tracing_with_verbose(false); +} + +/// Same as [`init_tracing`], with optional verbose default when `RUST_LOG` is unset. +pub fn init_tracing_with_verbose(verbose: bool) { + let _ = tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| { + if verbose { + tracing_subscriber::EnvFilter::new( + "persisting_compute=info,pulsing_actor=info,info", + ) + } else { + // Quiet default: only warn+ from Pulsing; compute stays at warn. + tracing_subscriber::EnvFilter::new( + "persisting_compute=warn,pulsing_actor=warn,warn", + ) + } + }), + ) + .with_writer(std::io::stderr) + .with_target(verbose) + .try_init(); +} diff --git a/crates/persisting-compute/src/dist.rs b/crates/persisting-compute/src/dist.rs new file mode 100644 index 0000000..efcadd0 --- /dev/null +++ b/crates/persisting-compute/src/dist.rs @@ -0,0 +1,220 @@ +//! Read placement from **torchrun** (or compatible) environment. +//! +//! Semantic block: [`crate::blocks::ids::PLACEMENT`] (distributed naming + seed). +//! +//! We do **not** spawn processes ourselves. Launch with: +//! +//! ```bash +//! torchrun --nproc_per_node=4 -- persisting compute plan.py +//! ``` + +use anyhow::{bail, Context, Result}; +use std::collections::HashMap; +use std::env; +use std::net::SocketAddr; + +#[derive(Debug, Clone)] +pub struct DistEnv { + pub rank: usize, + pub world_size: usize, + pub master_addr: String, + /// Where rank0 binds Pulsing so peers can join (no custom rdzv protocol). + pub pulsing_seed: SocketAddr, +} + +impl DistEnv { + /// `Some` when launched under torchrun (`RANK` + `WORLD_SIZE` present). + pub fn from_env() -> Result> { + let map: HashMap = env::vars().collect(); + Self::from_map(&map) + } + + /// Testable entry: build from an explicit env map (no process-global mutation). + pub fn from_map(env: &HashMap) -> Result> { + let rank = match env.get("RANK") { + Some(v) => v.parse::().context("parse RANK")?, + None => return Ok(None), + }; + let world_size = env + .get("WORLD_SIZE") + .context("WORLD_SIZE required when RANK is set (use torchrun)")? + .parse::() + .context("parse WORLD_SIZE")?; + if world_size == 0 { + bail!("WORLD_SIZE must be >= 1"); + } + if rank >= world_size { + bail!("RANK {rank} >= WORLD_SIZE {world_size}"); + } + let master_addr = env + .get("MASTER_ADDR") + .cloned() + .unwrap_or_else(|| "127.0.0.1".into()); + let master_port: u16 = env + .get("MASTER_PORT") + .map(|s| s.as_str()) + .unwrap_or("29500") + .parse() + .context("parse MASTER_PORT")?; + let pulsing_port: u16 = env + .get("PERSISTING_PULSING_PORT") + .and_then(|v| v.parse().ok()) + .unwrap_or(master_port.saturating_add(17)); + let pulsing_seed = Self::pulsing_seed_addr(&master_addr, pulsing_port)?; + Ok(Some(Self { + rank, + world_size, + master_addr, + pulsing_seed, + })) + } + + pub fn pulsing_seed_addr(master_addr: &str, pulsing_port: u16) -> Result { + format!("{master_addr}:{pulsing_port}") + .parse() + .with_context(|| format!("parse Pulsing seed {master_addr}:{pulsing_port}")) + } + + pub fn is_driver(&self) -> bool { + self.rank == 0 + } + + /// Logical worker actor name when `per_worker == 1`. + pub fn worker_name(rank: usize) -> String { + format!("compute/worker/{rank}") + } + + /// One concurrent execute slot under a logical worker/rank. + /// + /// When `per_worker == 1`, uses the legacy name [`Self::worker_name`] (no `/slot/`). + pub fn slot_name(worker: usize, slot: usize, per_worker: usize) -> String { + if per_worker <= 1 { + Self::worker_name(worker) + } else { + format!("compute/worker/{worker}/slot/{slot}") + } + } + + /// Flat pool index in **slot-major** order (matches [`Self::slot_names`]). + /// + /// `index = slot * n_workers + worker` — used by Scheduler / DeathWatch. + pub fn slot_flat_index( + worker: usize, + slot: usize, + n_workers: usize, + per_worker: usize, + ) -> usize { + let per_worker = per_worker.max(1); + let n_workers = n_workers.max(1); + debug_assert!( + worker < n_workers, + "worker {worker} >= n_workers {n_workers}" + ); + debug_assert!(slot < per_worker, "slot {slot} >= per_worker {per_worker}"); + slot.saturating_mul(n_workers).saturating_add(worker) + } + + /// Flat pool names in **slot-major** order: all workers' slot0, then slot1, … + pub fn slot_names(n_workers: usize, per_worker: usize) -> Vec { + let per_worker = per_worker.max(1); + let n_workers = n_workers.max(1); + let mut names = Vec::with_capacity(n_workers.saturating_mul(per_worker)); + for slot in 0..per_worker { + for worker in 0..n_workers { + names.push(Self::slot_name(worker, slot, per_worker)); + } + } + names + } + + pub fn worker_names(world_size: usize) -> Vec { + Self::slot_names(world_size, 1) + } + + /// Side-channel cancel actor for rank (separate mailbox from WorkerActor). + pub fn job_control_name(rank: usize) -> String { + format!("compute/job_control/{rank}") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn map(pairs: &[(&str, &str)]) -> HashMap { + pairs + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect() + } + + #[test] + fn from_map_none_without_rank() { + assert!(DistEnv::from_map(&HashMap::new()).unwrap().is_none()); + } + + #[test] + fn from_map_parses_torchrun_defaults() { + let env = map(&[ + ("RANK", "1"), + ("WORLD_SIZE", "4"), + ("MASTER_ADDR", "10.0.0.1"), + ]); + let d = DistEnv::from_map(&env).unwrap().unwrap(); + assert_eq!(d.rank, 1); + assert_eq!(d.world_size, 4); + assert!(!d.is_driver()); + assert_eq!(d.pulsing_seed.port(), 29500 + 17); + assert_eq!(d.master_addr, "10.0.0.1"); + } + + #[test] + fn from_map_rejects_bad_rank() { + let env = map(&[("RANK", "4"), ("WORLD_SIZE", "4")]); + assert!(DistEnv::from_map(&env).is_err()); + } + + #[test] + fn pulsing_port_override() { + let env = map(&[ + ("RANK", "0"), + ("WORLD_SIZE", "2"), + ("MASTER_PORT", "30000"), + ("PERSISTING_PULSING_PORT", "31000"), + ]); + let d = DistEnv::from_map(&env).unwrap().unwrap(); + assert!(d.is_driver()); + assert_eq!(d.pulsing_seed.port(), 31000); + } + + #[test] + fn slot_names_slot_major() { + assert_eq!( + DistEnv::slot_names(2, 2), + vec![ + "compute/worker/0/slot/0", + "compute/worker/1/slot/0", + "compute/worker/0/slot/1", + "compute/worker/1/slot/1", + ] + ); + assert_eq!( + DistEnv::slot_names(2, 1), + vec!["compute/worker/0", "compute/worker/1"] + ); + } + + #[test] + fn slot_flat_index_matches_slot_names_order() { + let names = DistEnv::slot_names(3, 2); + for worker in 0..3 { + for slot in 0..2 { + let i = DistEnv::slot_flat_index(worker, slot, 3, 2); + assert_eq!(names[i], DistEnv::slot_name(worker, slot, 2)); + } + } + // Regression: rank0's second slot is NOT flat index 1 when world>1. + assert_eq!(DistEnv::slot_flat_index(0, 1, 2, 2), 2); + assert_eq!(DistEnv::slot_flat_index(1, 0, 2, 2), 1); + } +} diff --git a/crates/persisting-compute/src/driver.rs b/crates/persisting-compute/src/driver.rs new file mode 100644 index 0000000..12ba500 --- /dev/null +++ b/crates/persisting-compute/src/driver.rs @@ -0,0 +1,354 @@ +//! Driver — control plane that drives the worker fleet. +//! +//! Conceptual ownership (one Driver per job on rank0 / local control process): +//! 1. **Plan** — stream tasks from the user `plan()` script +//! 2. **Dispatch** — least-loaded placement; `ask` workers **directly** (parallel) +//! 3. **Drain** — complete → `on_result` immediately; at most `max_inflight` JoinHandles +//! +//! This is a Rust control object, not a Pulsing Actor that `await`s Execute in +//! `receive` (that would serialize the fleet). Workers remain Pulsing actors. + +use crate::checkpoint::CheckpointTracker; +use crate::future::RunFuture; +use crate::observe::Observer; +use crate::plan::stream_plan_tasks; +use crate::pulsing_ext::{ask_timeout, ASK_TIMEOUT}; +use crate::scheduler::{AcquireError, Scheduler, StickyLost, WorkerPool}; +use crate::sink_writer::SinkSubmitter; +use crate::skip::SkipSet; +use crate::task::{unix_now, TaskExpr, TaskResult}; +use crate::worker::{WorkerCommand, WorkerReply}; +use anyhow::Result; +use futures::stream::FuturesUnordered; +use futures::StreamExt; +use std::path::PathBuf; +use std::sync::Arc; +use tokio_util::sync::CancellationToken; + +/// Job knobs owned/consumed by the Driver when running a plan. +#[derive(Clone)] +pub struct RunOptions { + pub script: PathBuf, + pub python: PathBuf, + /// Local-only worker count when not under torchrun. + pub workers: usize, + /// Global cap on concurrent tasks (defaults to `workers * per_worker_inflight`). + /// Also bounds outstanding JoinHandles (complete → drop). + pub max_inflight: usize, + /// Max concurrent Executes per worker (least-loaded scheduling). Default 1. + pub per_worker_inflight: usize, + /// Extra entries prepended onto PYTHONPATH for plan / execute host. + pub pythonpath_extra: Vec, + /// Forwarded to task.py as ``sys.argv[1:]`` (argparse-friendly). + pub script_args: Vec, + /// L2 infrastructure retries on worker ask failure (default 2). + pub infra_retries: u32, + /// Job-level cancel (Ctrl-C / external). Child tokens per in-flight task. + pub job_cancel: CancellationToken, + /// Optional observability sink (`--observe`). + pub observer: Arc, + /// Task ids to skip / already claimed (`--resume` seed + live completions). + pub skip_task_ids: SkipSet, + /// Optional checkpoint progress (`checkpoint.json` under `--sink`). + pub checkpoint: Option>, + /// Optional async sink enqueue (awaited on completion — back-pressures drain). + pub sink_submitter: Option, +} + +/// Drives one compute job: emit plan tasks and dispatch them onto the fleet. +pub struct Driver { + pool: WorkerPool, + sched: Arc, +} + +impl Driver { + pub fn new(pool: WorkerPool, sched: Arc) -> Self { + Self { pool, sched } + } + + pub fn scheduler(&self) -> &Arc { + &self.sched + } + + pub fn pool(&self) -> &WorkerPool { + &self.pool + } + + /// Run user `plan()` → place each task on a worker → collect results. + /// + /// Completions call `on_result` in **finish order** (not plan emit order). + /// At most `max_inflight` [`RunFuture`]s are outstanding at once. + pub async fn run_plan( + &self, + opts: &RunOptions, + mut on_result: impl FnMut(TaskResult) + Send, + ) -> Result> { + let global_cap = opts.max_inflight.max(1).min(self.sched.capacity().max(1)); + let mut stream = stream_plan_tasks( + opts.script.clone(), + opts.python.clone(), + opts.script_args.clone(), + ); + // Each element is a RunFuture::wait(); len() bounds outstanding JoinHandles. + let mut inflight: FuturesUnordered< + std::pin::Pin> + Send>>, + > = FuturesUnordered::new(); + let mut out = Vec::new(); + let mut plan_done = false; + let retries = opts.infra_retries; + let observer = Arc::clone(&opts.observer); + + loop { + // Guard before select: empty set + plan_done would disable every arm. + if plan_done && inflight.is_empty() { + break; + } + + tokio::select! { + biased; + + // Prefer draining completions so sink / resume stay current. + Some(joined) = inflight.next(), if !inflight.is_empty() => { + let r: TaskResult = joined?; + on_result(r.clone()); + if let Some(sink) = &opts.sink_submitter { + sink.submit(r.clone()) + .await + .map_err(|e| anyhow::anyhow!("sink enqueue: {e}"))?; + } + out.push(r); + } + + item = stream.next(), if !plan_done + && inflight.len() < global_cap + && !opts.job_cancel.is_cancelled() => + { + match item { + None => plan_done = true, + Some(Err(e)) => return Err(e), + Some(Ok(task)) => { + let task_id = task.id.clone(); + // Claim before dispatch: resume seed, live terminals, and + // duplicate plan() yields share one skip set (cross-worker + // re-dispatch of a finished id is skipped). + if !opts.skip_task_ids.insert(task_id.clone()) { + if let Some(ckpt) = &opts.checkpoint { + ckpt.note_skipped(1); + let _ = ckpt.maybe_flush().await; + } + continue; + } + observer.task_queued(&task_id).await; + if let Some(ckpt) = &opts.checkpoint { + ckpt.note_dispatched(); + } + let pool = Arc::clone(&self.pool); + let sched = Arc::clone(&self.sched); + let observer = Arc::clone(&observer); + let cancel = opts.job_cancel.child_token(); + let cancel_watch = cancel.clone(); + let task_id_for_join = task_id.clone(); + let join = tokio::spawn(async move { + if cancel_watch.is_cancelled() { + observer + .task_finished( + &task_id_for_join, + false, + true, + None, + &sched, + ) + .await; + return Ok(TaskResult::cancelled(task_id_for_join)); + } + execute_with_placement( + pool, + sched, + observer, + task, + retries, + cancel_watch, + ) + .await + }); + inflight.push(Box::pin(RunFuture::new(task_id, join, cancel).wait())); + } + } + } + + _ = opts.job_cancel.cancelled(), if !plan_done => { + tracing::debug!("job cancel: stop accepting new plan tasks"); + plan_done = true; + } + } + } + + Ok(out) + } +} + +async fn execute_with_placement( + pool: WorkerPool, + sched: Arc, + observer: Arc, + task: TaskExpr, + infra_retries: u32, + cancel: CancellationToken, +) -> Result { + let task_id = task.id.clone(); + let started = unix_now(); + let task_json = serde_json::to_vec(&task).map_err(|e| anyhow::anyhow!("encode task: {e}"))?; + let mut last_err = None; + // After first Execute contact: stick forever (result_cache is per-slot). + // Never fall through to another slot — that would be at-least-once re-execute. + let mut sticky: Option = None; + + for attempt in 0..=infra_retries { + if cancel.is_cancelled() { + observer + .task_finished(&task_id, false, true, None, &sched) + .await; + return Ok(TaskResult::cancelled(task_id)); + } + let guard = tokio::select! { + biased; + _ = cancel.cancelled() => { + observer + .task_finished(&task_id, false, true, None, &sched) + .await; + return Ok(TaskResult::cancelled(task_id)); + } + g = async { + match sticky { + Some(slot) => sched + .acquire_guard_sticky(slot) + .await + .map_err(PlacementErr::Sticky), + None => sched + .acquire_guard_prefer(None) + .await + .map_err(PlacementErr::AllGone), + } + } => g, + }; + let guard = match guard { + Ok(g) => g, + Err(PlacementErr::Sticky(StickyLost::Quarantined(slot))) => { + let err = format!( + "sticky slot {slot} quarantined after contact (refuse cross-slot re-execute)" + ); + tracing::warn!(%task_id, %err); + observer + .task_finished(&task_id, false, false, Some(err.clone()), &sched) + .await; + let mut r = + TaskResult::failure(task_id, format!("infra: {err}"), None, "infra", started); + r.infra_retries = attempt; + return Ok(r); + } + Err(PlacementErr::AllGone(AcquireError::AllQuarantined)) => { + let err = "all worker slots quarantined".to_string(); + tracing::error!(%task_id, %err); + observer + .task_finished(&task_id, false, false, Some(err.clone()), &sched) + .await; + let mut r = + TaskResult::failure(task_id, format!("infra: {err}"), None, "infra", started); + r.infra_retries = attempt; + return Ok(r); + } + }; + let idx = guard.index(); + let worker_id = format!("w{idx}"); + observer + .task_assigned(&task_id, idx, &worker_id, attempt, &sched) + .await; + let worker = { + let g = pool + .read() + .map_err(|_| anyhow::anyhow!("worker pool lock"))?; + g.get(idx) + .cloned() + .ok_or_else(|| anyhow::anyhow!("worker index {idx} out of range"))? + }; + + observer.task_running(&task_id).await; + // Contacted this slot: subsequent infra retries must stay here. + sticky = Some(idx); + let ask = ask_timeout::<_, WorkerReply>( + &worker, + WorkerCommand::Execute { + task_json: task_json.clone(), + }, + ASK_TIMEOUT, + ); + tokio::pin!(ask); + let reply = tokio::select! { + biased; + _ = cancel.cancelled() => { + // Local workers share job_cancel; remote ranks get Cancel via JobControlActor. + drop(guard); + observer + .task_finished(&task_id, false, true, None, &sched) + .await; + return Ok(TaskResult::cancelled(task_id)); + } + r = &mut ask => r, + }; + match reply { + Ok(WorkerReply::Result { result_json }) => { + let mut r: TaskResult = serde_json::from_slice(&result_json) + .map_err(|e| anyhow::anyhow!("decode result: {e}"))?; + if attempt > 0 { + r.infra_retries = attempt; + tracing::debug!(%task_id, attempt, worker = idx, "infra retry succeeded"); + } + if r.worker.is_none() { + r.worker = Some(format!("w{idx}")); + } + sched.note_success(idx); + observer + .task_finished(&task_id, r.ok, r.cancelled, r.error.clone(), &sched) + .await; + drop(guard); + return Ok(r); + } + Ok(WorkerReply::Bye) => { + last_err = Some("unexpected Bye on Execute".into()); + sched.note_failure(idx); + drop(guard); + } + Err(e) => { + tracing::warn!( + %task_id, + attempt, + worker = idx, + error = %e, + "infra retry: worker ask failed (sticky-only re-ask)" + ); + last_err = Some(e.to_string()); + sched.note_failure(idx); + drop(guard); + } + } + } + + let err = last_err.unwrap_or_else(|| "unknown".into()); + observer + .task_finished(&task_id, false, false, Some(err.clone()), &sched) + .await; + let mut r = TaskResult::failure( + task_id, + format!("infra retries exhausted: {err}"), + None, + "infra", + started, + ); + r.infra_retries = infra_retries; + Ok(r) +} + +enum PlacementErr { + Sticky(StickyLost), + AllGone(AcquireError), +} diff --git a/crates/persisting-compute/src/executor.rs b/crates/persisting-compute/src/executor.rs new file mode 100644 index 0000000..144223f --- /dev/null +++ b/crates/persisting-compute/src/executor.rs @@ -0,0 +1,424 @@ +//! L3 Executor seam — Worker invokes this; L2 Driver only schedules TaskExpr. +//! +//! **Primitive:** [`Executor`] trait · [`ExecutorRouter`] (product: `op=execute` only). +//! +//! ```text +//! Driver --ask--> WorkerActor -- op=execute --> plan.py::execute(item) +//! ``` + +use crate::python_env; +use crate::task::{unix_now, TaskExpr, TaskResult}; +use anyhow::{bail, Context, Result}; +use async_trait::async_trait; +use serde_json::{json, Value}; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::Arc; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use tokio::process::{Child, ChildStdin, ChildStdout, Command}; +use tokio::sync::Mutex; +use tokio_util::sync::CancellationToken; + +/// Long-lived host: load plan module once, call `execute(item)`. +const PLAN_EXECUTE_HOST: &str = r#" +import importlib.util, json, sys, traceback +from pathlib import Path + +plan_mods = {} + +def load_plan_module(script, argv=None): + script = str(Path(script).resolve()) + if script in plan_mods: + return plan_mods[script] + path = Path(script) + # Match `python task.py --foo bar` so argparse works at import time. + sys.argv = [script, *(argv or [])] + spec = importlib.util.spec_from_file_location("user_plan_exec", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + plan_mods[script] = mod + return mod + +def handle(msg): + cmd = msg.get("cmd") + if cmd == "shutdown": + return {"ok": True, "value": "bye"} + if cmd == "run_plan": + script = msg.get("script") + if not script: + raise ValueError("run_plan requires script= path to plan.py") + mod = load_plan_module(script, msg.get("argv") or []) + if not hasattr(mod, "execute"): + raise AttributeError( + f"{script} must define execute(item) — same object plan() yields" + ) + fn = getattr(mod, "execute") + if not callable(fn): + raise TypeError("execute must be callable") + # Pass the same shape plan() yields: {id, ...fields}, not wire TaskExpr. + task = msg.get("task") or {} + item = dict(task.get("args") or {}) + if task.get("id") is not None: + item["id"] = task["id"] + return {"ok": True, "value": fn(item)} + raise ValueError(f"unknown cmd: {cmd!r}") + +def main(): + for line in sys.stdin: + line = line.strip() + if not line: + continue + req_id = None + try: + msg = json.loads(line) + req_id = msg.get("id") + out = handle(msg) + out["id"] = req_id + print(json.dumps(out, default=str), flush=True) + if msg.get("cmd") == "shutdown": + break + except Exception as e: + print(json.dumps({ + "id": req_id, + "ok": False, + "error": str(e), + "traceback": traceback.format_exc(), + }), flush=True) + +if __name__ == "__main__": + main() +"#; + +/// One-shot run of a TaskExpr. +#[async_trait] +pub trait Executor: Send + Sync { + fn name(&self) -> &str; + async fn run(&self, task: TaskExpr, worker_id: &str) -> TaskResult; +} + +struct PlanHost { + child: Child, + stdin: ChildStdin, + stdout: BufReader, + next_id: u64, +} + +/// Owns one long-lived `--python` process that caches the loaded plan module. +struct PlanHostExecutor { + python: PathBuf, + pythonpath_extra: Vec, + host: Mutex>, +} + +impl PlanHostExecutor { + fn new(python: PathBuf, pythonpath_extra: Vec) -> Self { + Self { + python, + pythonpath_extra, + host: Mutex::new(None), + } + } + + async fn ensure_host( + host: &mut Option, + python: &PathBuf, + pythonpath_extra: &[PathBuf], + ) -> Result<()> { + if host.is_some() { + return Ok(()); + } + let mut cmd = Command::new(python); + cmd.arg("-u") + .arg("-c") + .arg(PLAN_EXECUTE_HOST) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()); + python_env::apply_pythonpath(&mut cmd, pythonpath_extra); + let mut child = cmd + .spawn() + .with_context(|| format!("spawn plan execute host: {}", python.display()))?; + let stdin = child + .stdin + .take() + .ok_or_else(|| anyhow::anyhow!("plan host missing stdin"))?; + let stdout = child + .stdout + .take() + .ok_or_else(|| anyhow::anyhow!("plan host missing stdout"))?; + *host = Some(PlanHost { + child, + stdin, + stdout: BufReader::new(stdout), + next_id: 1, + }); + Ok(()) + } + + async fn request(host: &mut PlanHost, mut msg: Value) -> Result { + let id = host.next_id; + host.next_id += 1; + msg.as_object_mut() + .ok_or_else(|| anyhow::anyhow!("request must be object"))? + .insert("id".into(), json!(id)); + let line = serde_json::to_string(&msg)?; + host.stdin.write_all(line.as_bytes()).await?; + host.stdin.write_all(b"\n").await?; + host.stdin.flush().await?; + + let mut reply = String::new(); + let n = host.stdout.read_line(&mut reply).await?; + if n == 0 { + bail!("plan execute host closed stdout"); + } + let parsed: Value = serde_json::from_str(reply.trim()) + .with_context(|| format!("invalid host reply: {}", reply.trim()))?; + if parsed.get("ok").and_then(|v| v.as_bool()) == Some(true) { + Ok(parsed.get("value").cloned().unwrap_or(Value::Null)) + } else { + let err = parsed + .get("error") + .and_then(|v| v.as_str()) + .unwrap_or("plan execute failed"); + let tb = parsed + .get("traceback") + .and_then(|v| v.as_str()) + .unwrap_or(""); + bail!("{err}\n{tb}") + } + } + + async fn shutdown(&self) { + let mut guard = self.host.lock().await; + if let Some(mut h) = guard.take() { + let _ = Self::request(&mut h, json!({"cmd": "shutdown"})).await; + let _ = h.child.kill().await; + } + } + + async fn run_plan_execute( + &self, + plan_script: &PathBuf, + script_args: &[String], + task: TaskExpr, + worker_id: &str, + cancel: CancellationToken, + ) -> TaskResult { + let started = unix_now(); + if cancel.is_cancelled() { + return TaskResult::cancelled(task.id); + } + let task_json = match serde_json::to_value(&task) { + Ok(v) => v, + Err(e) => { + return TaskResult::failure( + task.id, + format!("encode task: {e}"), + None, + worker_id, + started, + ); + } + }; + let script = match plan_script.canonicalize() { + Ok(p) => p, + Err(_) => plan_script.clone(), + }; + let mut guard = self.host.lock().await; + if let Err(e) = Self::ensure_host(&mut guard, &self.python, &self.pythonpath_extra).await { + return TaskResult::failure( + task.id, + e.to_string(), + Some(format!("{e:#}")), + worker_id, + started, + ); + } + let msg = json!({ + "cmd": "run_plan", + "script": script, + "argv": script_args, + "task": task_json, + }); + // In-flight cancel: kill the Python host so execute does not outlive Ctrl-C. + let result = { + let host = guard.as_mut().expect("host just ensured"); + tokio::select! { + biased; + _ = cancel.cancelled() => { + Err(anyhow::anyhow!("cancelled")) + } + r = Self::request(host, msg) => r, + } + }; + match result { + Ok(value) => TaskResult::success(task.id, value, worker_id, started), + Err(e) if cancel.is_cancelled() || e.to_string().contains("cancelled") => { + if let Some(mut h) = guard.take() { + let _ = h.child.kill().await; + } + TaskResult::cancelled(task.id) + } + Err(e) => { + let tb = format!("{e:#}"); + TaskResult::failure(task.id, e.to_string(), Some(tb), worker_id, started) + } + } + } +} + +/// Default algo path: plan script's `execute(item)`. +pub struct PlanExecuteExecutor { + host: Arc, + plan_script: PathBuf, + script_args: Vec, +} + +#[async_trait] +impl Executor for PlanExecuteExecutor { + fn name(&self) -> &str { + "execute" + } + + async fn run(&self, task: TaskExpr, worker_id: &str) -> TaskResult { + self.run_with_cancel(task, worker_id, CancellationToken::new()) + .await + } +} + +impl PlanExecuteExecutor { + async fn run_with_cancel( + &self, + task: TaskExpr, + worker_id: &str, + cancel: CancellationToken, + ) -> TaskResult { + self.host + .run_plan_execute( + &self.plan_script, + &self.script_args, + task, + worker_id, + cancel, + ) + .await + } +} + +/// Routes `op=execute` to [`PlanExecuteExecutor`]. Unknown ops fail clearly. +pub struct ExecutorRouter { + execute: Arc, + host: Arc, +} + +impl ExecutorRouter { + /// Worker stack for the product surface (plan + execute only). + pub fn local_stack( + python: PathBuf, + pythonpath_extra: Vec, + plan_script: PathBuf, + script_args: Vec, + ) -> Self { + let host = Arc::new(PlanHostExecutor::new(python, pythonpath_extra)); + let execute = Arc::new(PlanExecuteExecutor { + host: Arc::clone(&host), + plan_script, + script_args, + }); + Self { execute, host } + } + + pub async fn run(&self, task: TaskExpr, worker_id: &str) -> TaskResult { + self.run_with_cancel(task, worker_id, CancellationToken::new()) + .await + } + + pub async fn run_with_cancel( + &self, + task: TaskExpr, + worker_id: &str, + cancel: CancellationToken, + ) -> TaskResult { + if task.op != "execute" { + let started = unix_now(); + return TaskResult::failure( + task.id, + format!( + "unknown op {:?}: only op=execute (plan.py::execute) is supported", + task.op + ), + None, + worker_id, + started, + ); + } + self.execute.run_with_cancel(task, worker_id, cancel).await + } + + pub async fn shutdown(&self) { + self.host.shutdown().await; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::task::TaskExpr; + use serde_json::json; + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn run_with_cancel_kills_slow_execute() { + let dir = tempfile::tempdir().unwrap(); + let script = dir.path().join("slow.py"); + std::fs::write( + &script, + r#" +import time + +def plan(): + yield {"id": "t-0"} + +def execute(item): + time.sleep(5) + return {"done": True} +"#, + ) + .unwrap(); + let router = ExecutorRouter::local_stack(PathBuf::from("python3"), vec![], script, vec![]); + let cancel = CancellationToken::new(); + let bg = cancel.clone(); + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_millis(150)).await; + bg.cancel(); + }); + let task = TaskExpr::from_value(json!({"id": "t-0", "x": 1})).unwrap(); + let t0 = std::time::Instant::now(); + let r = router.run_with_cancel(task, "w0", cancel).await; + assert!(r.cancelled, "{r:?}"); + assert!(t0.elapsed().as_secs_f64() < 2.0); + router.shutdown().await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn run_execute_returns_value() { + let dir = tempfile::tempdir().unwrap(); + let script = dir.path().join("ok.py"); + std::fs::write( + &script, + r#" +def plan(): + yield {"id": "t-0"} + +def execute(item): + return {"x2": item["x"] * 2} +"#, + ) + .unwrap(); + let router = ExecutorRouter::local_stack(PathBuf::from("python3"), vec![], script, vec![]); + let task = TaskExpr::from_value(json!({"id": "t-0", "x": 3})).unwrap(); + let r = router.run(task, "w0").await; + assert!(r.ok); + assert_eq!(r.value, Some(json!({"x2": 6}))); + router.shutdown().await; + } +} diff --git a/crates/persisting-compute/src/future.rs b/crates/persisting-compute/src/future.rs new file mode 100644 index 0000000..8430713 --- /dev/null +++ b/crates/persisting-compute/src/future.rs @@ -0,0 +1,106 @@ +//! RunFuture — L2 scheduling atom (TaskSpec → wait / cancel). +//! +//! Cancel is cooperative at dispatch / acquire, and in-flight Python is killed +//! via the shared job [`CancellationToken`] watched by the plan execute host. +//! `wait` always joins and **never rewrites a successful result** as cancelled. + +use crate::task::TaskResult; +use anyhow::{Context, Result}; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; + +/// One submitted task under L2 control. +pub struct RunFuture { + task_id: String, + join: JoinHandle>, + cancel: CancellationToken, +} + +impl RunFuture { + pub(crate) fn new( + task_id: String, + join: JoinHandle>, + cancel: CancellationToken, + ) -> Self { + Self { + task_id, + join, + cancel, + } + } + + pub fn task_id(&self) -> &str { + &self.task_id + } + + /// Best-effort cancel: signals placement + shared job token (host kill). + /// Under torchrun, Driver also broadcasts to each rank's job-control actor. + pub fn cancel(&self) { + self.cancel.cancel(); + } + + pub fn is_cancelled(&self) -> bool { + self.cancel.is_cancelled() + } + + /// Wait until the task finishes. Successful results are kept even if cancel raced. + pub async fn wait(self) -> Result { + let task_id = self.task_id.clone(); + let joined = self.join.await.context("run future join")?; + match joined { + Ok(r) => Ok(r), + Err(e) => Err(e).with_context(|| format!("task {task_id} failed")), + } + } +} + +/// Wait for many futures in **submission** order. +/// +/// Prefer [`crate::driver::Driver::run_plan`] for jobs: it drains in completion +/// order and bounds outstanding work. This helper is for callers that need a +/// stable result order matching the input vec. +pub async fn wait_all(futures: Vec) -> Result> { + let mut out = Vec::with_capacity(futures.len()); + for f in futures { + out.push(f.wait().await?); + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[tokio::test] + async fn wait_keeps_success_even_if_cancelled() { + let token = CancellationToken::new(); + token.cancel(); + let join = + tokio::spawn(async { Ok(TaskResult::success("t-1", json!({"x": 1}), "w0", 0.0)) }); + let fut = RunFuture::new("t-1".into(), join, token); + assert!(fut.is_cancelled()); + let r = fut.wait().await.unwrap(); + assert!(r.ok); + assert!(!r.cancelled); + assert_eq!(r.task_id, "t-1"); + } + + #[tokio::test] + async fn wait_all_preserves_submission_order() { + let mk = |id: &str, delay_ms: u64| { + let id = id.to_string(); + let token = CancellationToken::new(); + let id_for_join = id.clone(); + let join = tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + Ok(TaskResult::success(id_for_join, json!(1), "w0", 0.0)) + }); + RunFuture::new(id, join, token) + }; + // Slow first, fast second — wait_all still returns submission order. + let out = wait_all(vec![mk("a", 40), mk("b", 1)]).await.unwrap(); + assert_eq!(out[0].task_id, "a"); + assert_eq!(out[1].task_id, "b"); + } +} diff --git a/crates/persisting-compute/src/job_control.rs b/crates/persisting-compute/src/job_control.rs new file mode 100644 index 0000000..98faadf --- /dev/null +++ b/crates/persisting-compute/src/job_control.rs @@ -0,0 +1,374 @@ +//! Side-channel job cancel + local DeathWatch for Pulsing fleets. +//! +//! WorkerActor mailboxes are serial: an in-flight `Execute` holds `receive`, so a +//! `Cancel` sitting behind it cannot stop Python. This actor shares the process +//! [`CancellationToken`] on a **separate** mailbox. +//! +//! Pulsing `watch` only supports **local** targets; we register local slot +//! ActorIds here and quarantine them in [`Scheduler`] on termination. + +use crate::dist::DistEnv; +use crate::pulsing_ext::{ask_timeout, resolve_actor, ASK_TIMEOUT}; +use crate::scheduler::Scheduler; +use futures::future::join_all; +use pulsing_actor::actor::{ActorId, StopReason}; +use pulsing_actor::prelude::*; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio_util::sync::CancellationToken; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum JobControlCommand { + /// Cancel the shared job token (in-flight execute hosts select! + kill). + Cancel, + /// Register a local slot for DeathWatch → quarantine on stop. + /// `slot` is the **flat pool index** (slot-major), not a per-rank ordinal. + WatchSlot { slot: usize, actor_id: ActorId }, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum JobControlReply { + Ack { already: bool }, + Watched, +} + +pub struct JobControlActor { + job_cancel: CancellationToken, + sched: Option>, + /// actor_id → slot index (local watches only). + watched: Mutex>, +} + +impl JobControlActor { + pub fn new(job_cancel: CancellationToken) -> Self { + Self { + job_cancel, + sched: None, + watched: Mutex::new(HashMap::new()), + } + } + + pub fn with_scheduler(job_cancel: CancellationToken, sched: Arc) -> Self { + Self { + job_cancel, + sched: Some(sched), + watched: Mutex::new(HashMap::new()), + } + } +} + +#[async_trait] +impl Actor for JobControlActor { + fn metadata(&self) -> HashMap { + HashMap::from([ + ("role".into(), "compute-job-control".into()), + ( + "cancelled".into(), + self.job_cancel.is_cancelled().to_string(), + ), + ]) + } + + async fn receive( + &mut self, + msg: Message, + ctx: &mut ActorContext, + ) -> pulsing_actor::error::Result { + // DeathWatch notification: (ActorId, StopReason) — parse borrows so we + // can still unpack JobControlCommand on the ask path. + if let Ok((dead_id, reason)) = msg.parse::<(ActorId, StopReason)>() { + let slot = self + .watched + .lock() + .ok() + .and_then(|g| g.get(&dead_id).copied()); + if let Some(slot) = slot { + tracing::warn!( + %dead_id, + slot, + %reason, + "DeathWatch: local worker terminated → quarantine" + ); + if let Some(sched) = &self.sched { + sched.force_quarantine(slot); + } + } + return Message::pack(&JobControlReply::Ack { already: true }); + } + + let cmd: JobControlCommand = msg.unpack()?; + let reply = match cmd { + JobControlCommand::Cancel => { + let already = self.job_cancel.is_cancelled(); + if !already { + tracing::warn!("job control: cancelling shared job token"); + self.job_cancel.cancel(); + } + JobControlReply::Ack { already } + } + JobControlCommand::WatchSlot { slot, actor_id } => { + if let Ok(mut g) = self.watched.lock() { + g.insert(actor_id, slot); + } + if let Err(e) = ctx.watch(&actor_id).await { + tracing::warn!(%actor_id, slot, error = %e, "DeathWatch register failed"); + } else { + tracing::debug!(%actor_id, slot, "DeathWatch registered"); + } + JobControlReply::Watched + } + }; + Message::pack(&reply) + } +} + +/// Ask every rank's job-control actor to cancel (best-effort, parallel). +pub async fn broadcast_job_cancel(system: &Arc, world_size: usize) { + let world_size = world_size.max(1); + let mut futs = Vec::with_capacity(world_size); + for rank in 0..world_size { + let name = DistEnv::job_control_name(rank); + let system = Arc::clone(system); + futs.push(async move { + match resolve_actor(system.as_ref(), &name).await { + Ok(ctrl) => match ask_timeout::<_, JobControlReply>( + &ctrl, + JobControlCommand::Cancel, + ASK_TIMEOUT, + ) + .await + { + Ok(JobControlReply::Ack { already }) => { + tracing::debug!(%name, already, "job cancel broadcast ok"); + } + Ok(_) => tracing::debug!(%name, "job cancel unexpected reply"), + Err(e) => tracing::warn!(%name, error = %e, "job cancel ask failed"), + }, + Err(e) => tracing::warn!(%name, error = %e, "job cancel resolve failed"), + } + }); + } + join_all(futs).await; +} + +/// Register DeathWatch for local slots. +/// +/// Each entry is `(actor_ref, flat_pool_index)`. Callers must pass the same +/// slot-major indices used by [`crate::scheduler::Scheduler`] / [`DistEnv::slot_names`] +/// — **not** `0..local_count` when the fleet spans multiple ranks. +pub async fn register_local_watches( + control: &ActorRef, + slots: &[(ActorRef, usize)], +) -> anyhow::Result<()> { + for (wref, flat_idx) in slots { + let _ = ask_timeout::<_, JobControlReply>( + control, + JobControlCommand::WatchSlot { + slot: *flat_idx, + actor_id: *wref.id(), + }, + Duration::from_secs(5), + ) + .await?; + } + Ok(()) +} + +/// Watch local `job_cancel`; when it fires, fan-out to all ranks' control actors. +pub fn spawn_cancel_broadcast( + system: Arc, + job_cancel: CancellationToken, + world_size: usize, +) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + job_cancel.cancelled().await; + tracing::warn!( + world_size, + "local job cancel: broadcasting to job-control actors" + ); + broadcast_job_cancel(&system, world_size).await; + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn cancel_command_trips_shared_token() { + let token = CancellationToken::new(); + let system = Arc::new( + ActorSystem::builder() + .mailbox_capacity(16) + .build() + .await + .unwrap(), + ); + let name = DistEnv::job_control_name(0); + let ctrl = system + .spawn_named(&name, JobControlActor::new(token.clone())) + .await + .unwrap(); + assert!(!token.is_cancelled()); + let ack = ctrl + .ask::<_, JobControlReply>(JobControlCommand::Cancel) + .await + .unwrap(); + assert!(matches!(ack, JobControlReply::Ack { already: false })); + assert!(token.is_cancelled()); + system.shutdown().await.unwrap(); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn broadcast_reaches_peer_rank_control() { + let t0 = CancellationToken::new(); + let t1 = CancellationToken::new(); + let system = Arc::new( + ActorSystem::builder() + .mailbox_capacity(16) + .build() + .await + .unwrap(), + ); + system + .spawn_named( + &DistEnv::job_control_name(0), + JobControlActor::new(t0.clone()), + ) + .await + .unwrap(); + system + .spawn_named( + &DistEnv::job_control_name(1), + JobControlActor::new(t1.clone()), + ) + .await + .unwrap(); + broadcast_job_cancel(&system, 2).await; + assert!(t0.is_cancelled() && t1.is_cancelled()); + system.shutdown().await.unwrap(); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn death_watch_quarantines_local_slot() { + let token = CancellationToken::new(); + let sched = Scheduler::new(1, 1); + let system = Arc::new( + ActorSystem::builder() + .mailbox_capacity(16) + .build() + .await + .unwrap(), + ); + let ctrl = system + .spawn_named( + &DistEnv::job_control_name(0), + JobControlActor::with_scheduler(token.clone(), Arc::clone(&sched)), + ) + .await + .unwrap(); + let worker = system + .spawn_named( + "compute/worker/0", + crate::worker::WorkerActor::with_plan( + "w0", + std::path::PathBuf::from("python3"), + vec![], + std::path::PathBuf::from("/dev/null"), + vec![], + token.clone(), + ), + ) + .await + .unwrap(); + register_local_watches(&ctrl, &[(worker.clone(), 0)]) + .await + .unwrap(); + assert!(!sched.is_quarantined(0)); + system.stop("compute/worker/0").await.unwrap(); + // Allow DeathWatch delivery. + for _ in 0..50 { + if sched.is_quarantined(0) { + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + assert!(sched.is_quarantined(0), "expected DeathWatch quarantine"); + system.shutdown().await.unwrap(); + } + + /// Regression: torchrun-shaped pool (world=2, per_worker=2) must map + /// rank0 slot1 → flat index 2, not local ordinal 1 (which is peer w1s0). + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn death_watch_uses_slot_major_flat_index() { + let token = CancellationToken::new(); + let world = 2; + let per_worker = 2; + let n_slots = world * per_worker; + let sched = Scheduler::new(n_slots, 1); + let system = Arc::new( + ActorSystem::builder() + .mailbox_capacity(16) + .build() + .await + .unwrap(), + ); + let ctrl = system + .spawn_named( + &DistEnv::job_control_name(0), + JobControlActor::with_scheduler(token.clone(), Arc::clone(&sched)), + ) + .await + .unwrap(); + + // Only spawn rank0's two local slots (as driver does before peers join). + let mut watches = Vec::new(); + for slot in 0..per_worker { + let name = DistEnv::slot_name(0, slot, per_worker); + let w = system + .spawn_named( + &name, + crate::worker::WorkerActor::with_plan( + format!("w0s{slot}"), + std::path::PathBuf::from("python3"), + vec![], + std::path::PathBuf::from("/dev/null"), + vec![], + token.clone(), + ), + ) + .await + .unwrap(); + let flat = DistEnv::slot_flat_index(0, slot, world, per_worker); + watches.push((w, flat)); + } + assert_eq!(watches[0].1, 0); + assert_eq!(watches[1].1, 2, "w0s1 must be flat 2, not 1"); + + register_local_watches(&ctrl, &watches).await.unwrap(); + + // Kill rank0 slot1 → must quarantine flat 2, leave 1 (peer) alone. + system + .stop(&DistEnv::slot_name(0, 1, per_worker)) + .await + .unwrap(); + for _ in 0..50 { + if sched.is_quarantined(2) { + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + assert!( + sched.is_quarantined(2), + "expected quarantine of flat index 2 (w0s1)" + ); + assert!( + !sched.is_quarantined(1), + "must not quarantine flat 1 (would be peer w1s0)" + ); + assert!(!sched.is_quarantined(0)); + system.shutdown().await.unwrap(); + } +} diff --git a/crates/persisting-compute/src/lib.rs b/crates/persisting-compute/src/lib.rs new file mode 100644 index 0000000..e1be10e --- /dev/null +++ b/crates/persisting-compute/src/lib.rs @@ -0,0 +1,65 @@ +//! Persisting compute control plane (library; CLI entry is `persisting compute`). +//! +//! # Semantic primitives +//! +//! Contracts are listed in [`blocks`]. Unit / interface tests live in the same +//! source file as each primitive; `tests/` holds multi-module integration only. +//! +//! # Design sketch +//! +//! - **One CLI** = `persisting compute` (thin orchestration, not Ray) +//! - **Driver** owns plan emit + dispatch (least-loaded asks workers directly) +//! - **Launch with torchrun** for multi-process; local `-w N` for single-process +//! - **Plan script** (user Python): `plan()` + `execute(item)`; argv after `--` +//! +//! ```text +//! torchrun --nproc_per_node=4 -- persisting compute plan.py -- --n 2 +//! rank0: Driver (plan + dispatch) + worker slots +//! rankN: worker slots +//! ``` +//! +//! Most modules are `pub(crate)`; only surfaces needed by the CLI binary and +//! integration tests are re-exported or left as `pub mod`. + +pub(crate) mod blocks; +pub(crate) mod check; +pub(crate) mod checkpoint; +pub mod cli; +pub mod dist; +pub mod driver; +pub(crate) mod executor; +pub(crate) mod future; +pub mod job_control; +pub mod observe; +pub(crate) mod plan; +pub mod pulsing_ext; +pub(crate) mod python_env; +pub(crate) mod result_cache; +pub mod runtime; +pub(crate) mod scheduler; +pub(crate) mod sink; +#[cfg(feature = "traj-sink")] +pub(crate) mod sink_traj; +pub(crate) mod sink_writer; +pub(crate) mod skip; +pub mod task; +pub(crate) mod worker; + +// ── Public surface (CLI + integration tests) ───────────────────────── + +pub use check::{run_check, run_self_test, CheckOptions, CheckReport}; +pub use checkpoint::CheckpointLedger; +pub use cli::{init_tracing, init_tracing_with_verbose, run_compute, ComputeArgs, ResultsFormat}; +pub use dist::DistEnv; +pub use driver::{Driver, RunOptions}; +pub use observe::{Observer, ObserverOptions}; +pub use runtime::{run_fleet, run_local_fleet}; +pub use skip::SkipSet; +pub use task::{TaskExpr, TaskResult}; +pub use worker::{WorkerActor, WorkerCommand, WorkerReply}; + +// Sink types used by CLI orchestration (re-exported so `cli` stays thin). +pub use sink::{JsonlFileSink, ResultSink, TeeSink}; +#[cfg(feature = "traj-sink")] +pub use sink_traj::VortexResultSink; +pub use sink_writer::{spawn_sink_writer, SinkSubmitter, SinkWriterHandle}; diff --git a/crates/persisting-compute/src/observe.rs b/crates/persisting-compute/src/observe.rs new file mode 100644 index 0000000..a159a74 --- /dev/null +++ b/crates/persisting-compute/src/observe.rs @@ -0,0 +1,513 @@ +//! Lightweight compute observability: queue / placement / per-task timing. +//! +//! Enabled with `--observe`. Progress lines go to **stderr** (prefix `[obs]`). +//! Machine NDJSON goes to `--observe-file` (and optionally stderr with `--observe-json`). + +use crate::scheduler::Scheduler; +use crate::task::unix_now; +use serde::Serialize; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; +use std::time::Instant; +use tokio::fs::OpenOptions; +use tokio::io::AsyncWriteExt; +use tokio::sync::Mutex as AsyncMutex; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum TaskPhase { + Queued, + Assigned, + Running, + Finished, + Failed, + Cancelled, +} + +#[derive(Debug, Clone, Serialize)] +pub struct ObsEvent { + pub kind: &'static str, + pub ts: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub task_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub worker: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub worker_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub attempt: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub phase: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub ok: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub duration_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub elapsed_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub loads: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub inflight: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub queued: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +#[derive(Debug, Clone, Serialize)] +pub struct InflightView { + pub task_id: String, + pub worker: usize, + pub worker_id: String, + pub phase: TaskPhase, + pub elapsed_ms: u64, + pub attempt: u32, +} + +struct InflightRec { + worker: usize, + worker_id: String, + phase: TaskPhase, + since: Instant, + attempt: u32, +} + +/// Shared live view of in-flight tasks + event sink. +pub struct Observer { + enabled: bool, + /// Human-readable `[obs] …` lines on stderr. + human: bool, + /// NDJSON on stderr (in addition to optional file). + json_stderr: bool, + file: AsyncMutex>, + inflight: Mutex>, + queued: Mutex, +} + +pub struct ObserverOptions { + pub human: bool, + pub json_stderr: bool, + pub path: Option, +} + +impl Observer { + pub fn disabled() -> Arc { + Arc::new(Self { + enabled: false, + human: false, + json_stderr: false, + file: AsyncMutex::new(None), + inflight: Mutex::new(HashMap::new()), + queued: Mutex::new(0), + }) + } + + pub async fn open(opts: ObserverOptions) -> anyhow::Result> { + let file = if let Some(p) = opts.path.as_deref() { + if let Some(parent) = p.parent() { + if !parent.as_os_str().is_empty() { + tokio::fs::create_dir_all(parent).await?; + } + } + Some(OpenOptions::new().create(true).append(true).open(p).await?) + } else { + None + }; + let obs = Arc::new(Self { + enabled: true, + human: opts.human, + json_stderr: opts.json_stderr, + file: AsyncMutex::new(file), + inflight: Mutex::new(HashMap::new()), + queued: Mutex::new(0), + }); + if obs.human { + let mut msg = String::from( + "[obs] enabled — progress on stderr (results default to --results quiet)", + ); + if let Some(p) = opts.path.as_deref() { + msg.push_str(&format!("; NDJSON → {}", p.display())); + } + if opts.json_stderr { + msg.push_str("; NDJSON also on stderr (--observe-json)"); + } + eprintln!("{msg}"); + } + Ok(obs) + } + + pub fn enabled(&self) -> bool { + self.enabled + } + + fn human_line(&self, line: &str) { + if self.human { + eprintln!("[obs] {line}"); + } + } + + async fn emit_json(&self, event: &ObsEvent) { + if !self.enabled { + return; + } + let mut guard = self.file.lock().await; + if !self.json_stderr && guard.is_none() { + return; + } + let Ok(line) = serde_json::to_string(event) else { + return; + }; + if self.json_stderr { + eprintln!("{line}"); + } + if let Some(f) = guard.as_mut() { + let _ = f.write_all(line.as_bytes()).await; + let _ = f.write_all(b"\n").await; + let _ = f.flush().await; + } + } + + async fn emit(&self, event: ObsEvent, human: Option) { + if !self.enabled { + return; + } + if let Some(h) = human { + self.human_line(&h); + } + self.emit_json(&event).await; + } + + pub async fn task_queued(&self, task_id: &str) { + if !self.enabled { + return; + } + let q = { + let mut g = self.queued.lock().unwrap_or_else(|e| e.into_inner()); + *g += 1; + *g + }; + self.emit( + ObsEvent { + kind: "task.queued", + ts: unix_now(), + task_id: Some(task_id.into()), + worker: None, + worker_id: None, + attempt: None, + phase: Some(TaskPhase::Queued), + ok: None, + duration_ms: None, + elapsed_ms: None, + loads: None, + inflight: None, + queued: Some(q), + error: None, + }, + Some(format!("queued {task_id} queue={q}")), + ) + .await; + } + + pub async fn task_assigned( + &self, + task_id: &str, + worker: usize, + worker_id: &str, + attempt: u32, + sched: &Scheduler, + ) { + if !self.enabled { + return; + } + { + let mut q = self.queued.lock().unwrap_or_else(|e| e.into_inner()); + *q = q.saturating_sub(1); + } + { + let mut map = self.inflight.lock().unwrap_or_else(|e| e.into_inner()); + map.insert( + task_id.to_string(), + InflightRec { + worker, + worker_id: worker_id.to_string(), + phase: TaskPhase::Assigned, + since: Instant::now(), + attempt, + }, + ); + } + let loads = sched.load_snapshot(); + let q = self.queued_count().unwrap_or(0); + self.emit( + ObsEvent { + kind: "task.assigned", + ts: unix_now(), + task_id: Some(task_id.into()), + worker: Some(worker), + worker_id: Some(worker_id.into()), + attempt: Some(attempt), + phase: Some(TaskPhase::Assigned), + ok: None, + duration_ms: None, + elapsed_ms: None, + loads: Some(loads.clone()), + inflight: None, + queued: Some(q), + error: None, + }, + Some(format!( + "assigned {task_id} → {worker_id} loads={loads:?} queue={q}" + )), + ) + .await; + } + + pub async fn task_running(&self, task_id: &str) { + if !self.enabled { + return; + } + let (worker, worker_id, attempt, elapsed_ms) = { + let mut map = self.inflight.lock().unwrap_or_else(|e| e.into_inner()); + let Some(rec) = map.get_mut(task_id) else { + return; + }; + rec.phase = TaskPhase::Running; + ( + rec.worker, + rec.worker_id.clone(), + rec.attempt, + rec.since.elapsed().as_millis() as u64, + ) + }; + self.emit( + ObsEvent { + kind: "task.running", + ts: unix_now(), + task_id: Some(task_id.into()), + worker: Some(worker), + worker_id: Some(worker_id.clone()), + attempt: Some(attempt), + phase: Some(TaskPhase::Running), + ok: None, + duration_ms: None, + elapsed_ms: Some(elapsed_ms), + loads: None, + inflight: None, + queued: None, + error: None, + }, + Some(format!("running {task_id} on {worker_id}")), + ) + .await; + } + + pub async fn task_finished( + &self, + task_id: &str, + ok: bool, + cancelled: bool, + error: Option, + sched: &Scheduler, + ) { + if !self.enabled { + return; + } + let (worker, worker_id, attempt, duration_ms, phase) = { + let mut map = self.inflight.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(rec) = map.remove(task_id) { + let phase = if cancelled { + TaskPhase::Cancelled + } else if ok { + TaskPhase::Finished + } else { + TaskPhase::Failed + }; + ( + Some(rec.worker), + Some(rec.worker_id), + Some(rec.attempt), + Some(rec.since.elapsed().as_millis() as u64), + Some(phase), + ) + } else { + ( + None, + None, + None, + None, + Some(if cancelled { + TaskPhase::Cancelled + } else if ok { + TaskPhase::Finished + } else { + TaskPhase::Failed + }), + ) + } + }; + let loads = sched.load_snapshot(); + let q = self.queued_count().unwrap_or(0); + let status = if cancelled { + "cancelled" + } else if ok { + "ok" + } else { + "fail" + }; + let wid = worker_id.as_deref().unwrap_or("?").to_string(); + let dur = duration_ms + .map(|d| format!("{d}ms")) + .unwrap_or_else(|| "?".into()); + let err_s = error + .as_ref() + .map(|e| format!(" err={e}")) + .unwrap_or_default(); + let human = format!("finished {task_id} on {wid} {dur} {status}{err_s}"); + self.emit( + ObsEvent { + kind: "task.finished", + ts: unix_now(), + task_id: Some(task_id.into()), + worker, + worker_id, + attempt, + phase, + ok: Some(ok && !cancelled), + duration_ms, + elapsed_ms: duration_ms, + loads: Some(loads), + inflight: None, + queued: Some(q), + error, + }, + Some(human), + ) + .await; + } + + pub async fn fleet_snapshot(&self, sched: &Scheduler) { + if !self.enabled { + return; + } + let inflight = { + let map = self.inflight.lock().unwrap_or_else(|e| e.into_inner()); + map.iter() + .map(|(id, rec)| InflightView { + task_id: id.clone(), + worker: rec.worker, + worker_id: rec.worker_id.clone(), + phase: rec.phase, + elapsed_ms: rec.since.elapsed().as_millis() as u64, + attempt: rec.attempt, + }) + .collect::>() + }; + let loads = sched.load_snapshot(); + let q = self.queued_count().unwrap_or(0); + let n = inflight.len(); + // Keep human fleet lines short: only when something is in flight or queued. + let human = if n > 0 || q > 0 { + let sample: Vec = inflight + .iter() + .take(4) + .map(|v| format!("{}@{}:{}ms", v.task_id, v.worker_id, v.elapsed_ms)) + .collect(); + let more = if n > 4 { + format!(" +{}", n - 4) + } else { + String::new() + }; + Some(format!( + "fleet loads={loads:?} queue={q} inflight={n} [{}{more}]", + sample.join(", ") + )) + } else { + None + }; + self.emit( + ObsEvent { + kind: "fleet.snapshot", + ts: unix_now(), + task_id: None, + worker: None, + worker_id: None, + attempt: None, + phase: None, + ok: None, + duration_ms: None, + elapsed_ms: None, + loads: Some(loads), + inflight: Some(inflight), + queued: Some(q), + error: None, + }, + human, + ) + .await; + } + + fn queued_count(&self) -> Option { + self.queued.lock().ok().map(|g| *g).or_else(|| Some(0)) + } +} + +/// Background ticker for [`Observer::fleet_snapshot`]. +pub fn spawn_snapshot_loop( + obs: Arc, + sched: Arc, + cancel: tokio_util::sync::CancellationToken, + every: std::time::Duration, +) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + if !obs.enabled() { + return; + } + let mut interval = tokio::time::interval(every); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + biased; + _ = cancel.cancelled() => break, + _ = interval.tick() => { + obs.fleet_snapshot(&sched).await; + } + } + } + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::scheduler::Scheduler; + + #[tokio::test] + async fn disabled_observer_is_noop() { + let obs = Observer::disabled(); + assert!(!obs.enabled()); + obs.task_queued("t-0").await; + let sched = Scheduler::new(1, 1); + obs.task_assigned("t-0", 0, "w0", 0, &sched).await; + obs.task_running("t-0").await; + obs.task_finished("t-0", true, false, None, &sched).await; + obs.fleet_snapshot(&sched).await; + } + + #[tokio::test] + async fn enabled_tracks_queued_count() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("obs.ndjson"); + let obs = Observer::open(ObserverOptions { + human: false, + json_stderr: false, + path: Some(path.clone()), + }) + .await + .unwrap(); + obs.task_queued("t-0").await; + obs.task_queued("t-1").await; + let body = tokio::fs::read_to_string(&path).await.unwrap(); + assert!(body.contains("task.queued")); + assert_eq!(body.lines().count(), 2); + } +} diff --git a/crates/persisting-compute/src/plan.rs b/crates/persisting-compute/src/plan.rs new file mode 100644 index 0000000..4869285 --- /dev/null +++ b/crates/persisting-compute/src/plan.rs @@ -0,0 +1,168 @@ +//! Run a user plan script and stream TaskExpr lines (NDJSON on stdout). +//! +//! The control plane never embeds the user's interpreter (PyO3). It **invokes** +//! `--python` so quirky envs stay isolated; stacks stay in that process. +//! +//! User CLI args after `--` become ``sys.argv`` for the plan module (argparse-friendly). + +use crate::task::TaskExpr; +use anyhow::{bail, Context, Result}; +use futures::Stream; +use std::path::PathBuf; +use std::pin::Pin; +use std::process::Stdio; +use tokio::io::{AsyncBufReadExt, BufReader}; +use tokio::process::Command; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; + +/// Bootstrap: set ``sys.argv = [script, *user_args]`` then import plan(). +const PLAN_BOOTSTRAP: &str = r#" +import asyncio, json, sys +from pathlib import Path +path = Path(sys.argv[1]).resolve() +user_args = sys.argv[2:] +# So argparse in task.py sees the same argv as `python task.py --foo bar` +sys.argv = [str(path), *user_args] +import importlib.util +spec = importlib.util.spec_from_file_location("user_plan", path) +mod = importlib.util.module_from_spec(spec) +spec.loader.exec_module(mod) + +def _dump(item): + if hasattr(item, "to_dict"): + item = item.to_dict() + print(json.dumps(item, ensure_ascii=False), flush=True) + +async def _emit(): + if hasattr(mod, "plan"): + out = mod.plan() + if asyncio.iscoroutine(out): + out = await out + if hasattr(out, "__aiter__"): + async for item in out: + _dump(item) + return + for item in out: + _dump(item) + return + if hasattr(mod, "PLAN"): + for item in mod.PLAN: + _dump(item) + return + raise SystemExit("plan script must define plan() or PLAN") + +asyncio.run(_emit()) +"#; + +/// Stream tasks from a plan script under `python`. +pub fn stream_plan_tasks( + script: PathBuf, + python: PathBuf, + script_args: Vec, +) -> Pin> + Send>> { + let (tx, rx) = mpsc::channel::>(64); + tokio::spawn(async move { + if let Err(e) = run_plan_process(script, python, script_args, tx.clone()).await { + let _ = tx.send(Err(e)).await; + } + }); + Box::pin(ReceiverStream::new(rx)) +} + +async fn run_plan_process( + script: PathBuf, + python: PathBuf, + script_args: Vec, + tx: mpsc::Sender>, +) -> Result<()> { + let script = script + .canonicalize() + .with_context(|| format!("plan script not found: {}", script.display()))?; + + let mut cmd = Command::new(&python); + cmd.arg("-c") + .arg(PLAN_BOOTSTRAP) + .arg(&script) + .args(&script_args) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .env("PERSISTING_COMPUTE_EMIT_PLAN", "1"); + let mut child = cmd + .spawn() + .with_context(|| format!("spawn plan python: {}", python.display()))?; + + let stdout = child + .stdout + .take() + .ok_or_else(|| anyhow::anyhow!("missing plan stdout"))?; + let stderr = child + .stderr + .take() + .ok_or_else(|| anyhow::anyhow!("missing plan stderr"))?; + + let stderr_task = tokio::spawn(async move { + let mut lines = BufReader::new(stderr).lines(); + let mut buf = String::new(); + while let Ok(Some(line)) = lines.next_line().await { + buf.push_str(&line); + buf.push('\n'); + } + buf + }); + + let mut lines = BufReader::new(stdout).lines(); + while let Some(line) = lines.next_line().await? { + let line = line.trim(); + if line.is_empty() || line.starts_with('#') { + continue; + } + let parsed = serde_json::from_str::(line) + .with_context(|| format!("invalid NDJSON from plan: {line}")) + .and_then(TaskExpr::from_value); + if tx.send(parsed).await.is_err() { + break; + } + } + + let status = child.wait().await.context("wait plan process")?; + let err = stderr_task.await.unwrap_or_default(); + if !status.success() { + bail!( + "plan script exited {}: {}", + status.code().unwrap_or(-1), + err.trim() + ); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::StreamExt; + + #[tokio::test] + async fn stream_plan_tasks_emits_flat_items() { + let dir = tempfile::tempdir().unwrap(); + let script = dir.path().join("plan.py"); + std::fs::write( + &script, + r#" +def plan(): + for i in range(3): + yield {"id": f"t-{i}", "x": i} + +def execute(item): + return item +"#, + ) + .unwrap(); + let mut stream = stream_plan_tasks(script, PathBuf::from("python3"), vec![]); + let mut ids = Vec::new(); + while let Some(item) = stream.next().await { + ids.push(item.unwrap().id); + } + assert_eq!(ids, vec!["t-0", "t-1", "t-2"]); + } +} diff --git a/crates/persisting-compute/src/pulsing_ext.rs b/crates/persisting-compute/src/pulsing_ext.rs new file mode 100644 index 0000000..004f4d0 --- /dev/null +++ b/crates/persisting-compute/src/pulsing_ext.rs @@ -0,0 +1,55 @@ +//! Thin helpers over Pulsing so compute uses one resolve/spawn style. + +use anyhow::{Context, Result}; +use pulsing_actor::prelude::*; +use std::sync::Arc; +use std::time::Duration; + +/// Default infra `ask` deadline (select! / timeout wrapper). +pub const ASK_TIMEOUT: Duration = Duration::from_secs(120); + +/// Worker slot supervision: restart on failure, capped. +pub fn worker_supervision() -> SupervisionSpec { + SupervisionSpec::on_failure().with_max_restarts(3) +} + +/// Unified named resolve (prefer this over mixing `resolve` / `resolve_named`). +pub async fn resolve_actor(system: &ActorSystem, name: &str) -> Result { + system + .resolve_named(name, None) + .await + .with_context(|| format!("resolve Pulsing actor {name}")) +} + +/// `ask` with an explicit deadline (Pulsing has no ask_timeout yet). +pub async fn ask_timeout(actor: &ActorRef, msg: M, timeout: Duration) -> Result +where + M: serde::Serialize + 'static, + R: serde::de::DeserializeOwned, +{ + match tokio::time::timeout(timeout, actor.ask::(msg)).await { + Ok(Ok(r)) => Ok(r), + Ok(Err(e)) => Err(anyhow::anyhow!("pulsing ask: {e}")), + Err(_) => Err(anyhow::anyhow!("pulsing ask timed out after {timeout:?}")), + } +} + +/// Spawn a named actor with supervision via factory (enables restart). +pub async fn spawn_supervised( + system: &Arc, + name: &str, + factory: F, +) -> Result +where + F: FnMut() -> pulsing_actor::error::Result + Send + 'static, + A: Actor, +{ + system + .spawning() + .name(name) + .supervision(worker_supervision()) + .mailbox_capacity(256) + .spawn_factory(factory) + .await + .with_context(|| format!("spawn supervised {name}")) +} diff --git a/crates/persisting-compute/src/python_env.rs b/crates/persisting-compute/src/python_env.rs new file mode 100644 index 0000000..b28da69 --- /dev/null +++ b/crates/persisting-compute/src/python_env.rs @@ -0,0 +1,91 @@ +//! Shared helpers for invoking user `--python` with a usable module path. +//! +//! Semantic block: [`crate::blocks::ids::PYTHON_ENV`]. + +use std::env; +use std::ffi::OsStr; +use std::path::{Path, PathBuf}; + +/// Build PYTHONPATH from an explicit existing value + extras (testable without mutating env). +pub fn merge_pythonpath_parts(existing: Option<&OsStr>, extras: &[PathBuf]) -> Option { + let mut parts: Vec = Vec::new(); + if let Some(cur) = existing { + for p in env::split_paths(cur) { + if !p.as_os_str().is_empty() && !parts.iter().any(|x| x == &p) { + parts.push(p); + } + } + } + for e in extras { + let p = e.canonicalize().unwrap_or_else(|_| e.to_path_buf()); + if !parts.iter().any(|x| x == &p) { + parts.push(p); + } + } + if parts.is_empty() { + None + } else { + env::join_paths(&parts) + .ok() + .map(|os| os.to_string_lossy().into_owned()) + } +} + +/// Build PYTHONPATH: existing env + extras (plan dir, `-E` paths, …). +pub fn merge_pythonpath(extras: &[PathBuf]) -> Option { + merge_pythonpath_parts(env::var_os("PYTHONPATH").as_deref(), extras) +} + +/// Default extras for a plan script: its parent directory (so sibling modules import). +pub fn pythonpath_for_script(script: &Path) -> Vec { + let mut out = Vec::new(); + if let Some(parent) = script.parent() { + out.push(parent.to_path_buf()); + } + out +} + +pub fn apply_pythonpath(cmd: &mut tokio::process::Command, extras: &[PathBuf]) { + if let Some(pp) = merge_pythonpath(extras) { + cmd.env("PYTHONPATH", pp); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::ffi::OsString; + + #[test] + fn merge_dedups_and_preserves_order() { + let existing = OsString::from("/a:/b"); + let extras = vec![PathBuf::from("/b"), PathBuf::from("/c")]; + let merged = merge_pythonpath_parts(Some(&existing), &extras).unwrap(); + let parts: Vec<_> = env::split_paths(&merged).collect(); + assert!(parts + .iter() + .any(|p| p.ends_with("a") || p == Path::new("/a"))); + assert_eq!( + parts + .iter() + .filter(|p| p.ends_with("b") || **p == PathBuf::from("/b")) + .count(), + 1 + ); + assert!(parts + .iter() + .any(|p| p.ends_with("c") || p == Path::new("/c"))); + } + + #[test] + fn merge_empty_returns_none() { + assert!(merge_pythonpath_parts(None, &[]).is_none()); + } + + #[test] + fn pythonpath_for_script_uses_parent() { + let p = Path::new("/tmp/plans/task.py"); + let extras = pythonpath_for_script(p); + assert_eq!(extras, vec![PathBuf::from("/tmp/plans")]); + } +} diff --git a/crates/persisting-compute/src/result_cache.rs b/crates/persisting-compute/src/result_cache.rs new file mode 100644 index 0000000..e61667c --- /dev/null +++ b/crates/persisting-compute/src/result_cache.rs @@ -0,0 +1,99 @@ +//! Infra-retry idempotency cache: `task_id` → last terminal [`TaskResult`]. +//! +//! Semantic block: **L2 idempotency (same worker)**. +//! Kept separate from [`crate::worker::WorkerActor`] so the policy is unit-testable +//! without Pulsing / Python. + +use crate::task::TaskResult; +use std::collections::{HashMap, VecDeque}; + +/// Default cap for cached results per worker slot. +pub const DEFAULT_RESULT_CACHE_CAP: usize = 4096; + +/// Bounded LRU-ish cache (evict oldest insert order). +#[derive(Debug, Default)] +pub struct ResultCache { + map: HashMap, + order: VecDeque, + cap: usize, +} + +impl ResultCache { + pub fn new(cap: usize) -> Self { + Self { + map: HashMap::new(), + order: VecDeque::new(), + cap: cap.max(1), + } + } + + pub fn get(&self, task_id: &str) -> Option<&TaskResult> { + self.map.get(task_id) + } + + /// Insert or replace. Skips caching cancelled results (caller should not insert them). + pub fn put(&mut self, task_id: impl Into, result: TaskResult) { + let task_id = task_id.into(); + if result.cancelled { + return; + } + if self.map.contains_key(&task_id) { + self.map.insert(task_id, result); + return; + } + while self.order.len() >= self.cap { + if let Some(old) = self.order.pop_front() { + self.map.remove(&old); + } else { + break; + } + } + self.order.push_back(task_id.clone()); + self.map.insert(task_id, result); + } + + pub fn len(&self) -> usize { + self.map.len() + } + + pub fn is_empty(&self) -> bool { + self.map.is_empty() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn put_get_and_skip_cancelled() { + let mut c = ResultCache::new(8); + c.put("t-0", TaskResult::success("t-0", json!(1), "w0", 0.0)); + assert!(c.get("t-0").unwrap().ok); + c.put("t-1", TaskResult::cancelled("t-1")); + assert!(c.get("t-1").is_none()); + assert_eq!(c.len(), 1); + } + + #[test] + fn evicts_oldest_when_over_cap() { + let mut c = ResultCache::new(2); + c.put("a", TaskResult::success("a", json!(1), "w0", 0.0)); + c.put("b", TaskResult::success("b", json!(2), "w0", 0.0)); + c.put("c", TaskResult::success("c", json!(3), "w0", 0.0)); + assert!(c.get("a").is_none()); + assert!(c.get("b").is_some()); + assert!(c.get("c").is_some()); + assert_eq!(c.len(), 2); + } + + #[test] + fn replace_keeps_cap() { + let mut c = ResultCache::new(2); + c.put("a", TaskResult::success("a", json!(1), "w0", 0.0)); + c.put("a", TaskResult::success("a", json!(9), "w0", 0.0)); + assert_eq!(c.len(), 1); + assert_eq!(c.get("a").unwrap().value, Some(json!(9))); + } +} diff --git a/crates/persisting-compute/src/runtime.rs b/crates/persisting-compute/src/runtime.rs new file mode 100644 index 0000000..2851512 --- /dev/null +++ b/crates/persisting-compute/src/runtime.rs @@ -0,0 +1,451 @@ +//! Boot Pulsing fleet, then hand the job to [`crate::driver::Driver`]. +//! +//! - Local `-w N`: spawn N×`--per-worker` slot actors (each owns a Python host). +//! - torchrun: each rank spawns `--per-worker` slot actors; rank0 is Driver. +//! +//! `--per-worker N` = N concurrent Execute slots per logical worker/rank. Each +//! slot is its own [`WorkerActor`] (serial mailbox) + dedicated plan host, so +//! concurrency is real — not mailbox queuing on one actor. +//! +//! Workers are spawned via [`crate::pulsing_ext::spawn_supervised`] (factory + +//! `SupervisionSpec`) so Pulsing can restart a failed slot. + +use crate::dist::DistEnv; +use crate::driver::Driver; +use crate::job_control::{register_local_watches, spawn_cancel_broadcast, JobControlActor}; +use crate::observe::spawn_snapshot_loop; +use crate::pulsing_ext::{ask_timeout, resolve_actor, spawn_supervised, ASK_TIMEOUT}; +use crate::python_env::pythonpath_for_script; +use crate::scheduler::{Scheduler, WorkerPool}; +use crate::task::TaskResult; +use crate::worker::{ShutdownGate, WorkerCommand, WorkerConfig, WorkerReply}; +use anyhow::{bail, Context, Result}; +use pulsing_actor::prelude::*; +use std::path::PathBuf; +use std::sync::{Arc, RwLock}; +use std::time::Duration; +use tokio::time::sleep; +use tokio_util::sync::CancellationToken; + +pub use crate::driver::RunOptions; + +/// Entry: torchrun env → distributed; else in-process local fleet. +pub async fn run_fleet( + opts: RunOptions, + on_result: impl FnMut(TaskResult) + Send, +) -> Result> { + if let Some(dist) = DistEnv::from_env()? { + tracing::debug!( + rank = dist.rank, + world_size = dist.world_size, + seed = %dist.pulsing_seed, + "torchrun placement detected" + ); + if dist.is_driver() { + run_driver_rank(dist, opts, on_result).await + } else { + let pp = apply_pythonpath(&opts); + run_worker_rank(dist, &opts, pp).await?; + Ok(Vec::new()) + } + } else { + run_local_fleet(opts, on_result).await + } +} + +/// Single-process: spawn N×P slot workers; this process runs the Driver. +pub async fn run_local_fleet( + opts: RunOptions, + on_result: impl FnMut(TaskResult) + Send, +) -> Result> { + let pythonpath = apply_pythonpath(&opts); + let system: Arc = ActorSystem::builder() + .mailbox_capacity(256) + .build() + .await + .context("build ActorSystem")?; + let per_worker = opts.per_worker_inflight.max(1); + let n_workers = opts.workers.max(1); + let names = DistEnv::slot_names(n_workers, per_worker); + let n_slots = names.len(); + // Second arg is always 1: pool is already flattened one-actor-per-slot. + let sched = Scheduler::new(n_slots, 1); + + let control = spawn_job_control( + &system, + 0, + opts.job_cancel.clone(), + Some(Arc::clone(&sched)), + ) + .await?; + let cancel_fanout = spawn_cancel_broadcast(Arc::clone(&system), opts.job_cancel.clone(), 1); + + let watches = + spawn_local_fleet_slots(&system, n_workers, per_worker, &opts, &pythonpath).await?; + let pool: WorkerPool = Arc::new(RwLock::new( + watches.iter().map(|(r, _)| r.clone()).collect(), + )); + register_local_watches(&control, &watches).await?; + + tracing::debug!( + workers = n_workers, + per_worker, + slots = n_slots, + capacity = sched.capacity(), + "driver ready (local fleet)" + ); + + run_driver_loop(pool, sched, &opts, on_result, cancel_fanout, system, None).await +} + +/// Rank0 under torchrun: bind Pulsing seed, wait for peer slots, run Driver. +async fn run_driver_rank( + dist: DistEnv, + opts: RunOptions, + on_result: impl FnMut(TaskResult) + Send, +) -> Result> { + let bind = format!("0.0.0.0:{}", dist.pulsing_seed.port()); + let system: Arc = ActorSystem::builder() + .mailbox_capacity(256) + .addr(bind.as_str()) + .build() + .await + .context("build driver ActorSystem")?; + tracing::debug!( + advertised = %dist.pulsing_seed, + bound = %system.addr(), + "driver Pulsing listening (peers join via MASTER_ADDR seed)" + ); + + let per_worker = opts.per_worker_inflight.max(1); + let names = DistEnv::slot_names(dist.world_size, per_worker); + let n_slots = names.len(); + // Second arg is always 1: pool is already flattened one-actor-per-slot. + let sched = Scheduler::new(n_slots, 1); + + let control = spawn_job_control( + &system, + 0, + opts.job_cancel.clone(), + Some(Arc::clone(&sched)), + ) + .await?; + + let pythonpath = apply_pythonpath(&opts); + let local_watches = spawn_rank_slots( + &system, + 0, + dist.world_size, + per_worker, + &opts, + &pythonpath, + None, + ) + .await?; + let pool: WorkerPool = Arc::new(RwLock::new( + local_watches.iter().map(|(r, _)| r.clone()).collect(), + )); + register_local_watches(&control, &local_watches).await?; + + wait_and_fill_workers(&system, &pool, &names, Duration::from_secs(120)).await?; + + let cancel_fanout = spawn_cancel_broadcast( + Arc::clone(&system), + opts.job_cancel.clone(), + dist.world_size, + ); + + tracing::debug!( + world_size = dist.world_size, + per_worker, + slots = n_slots, + capacity = sched.capacity(), + "driver ready (torchrun)" + ); + + let world_size = dist.world_size; + let per_worker_shutdown = per_worker; + run_driver_loop( + pool, + sched, + &opts, + on_result, + cancel_fanout, + system, + Some(DriverPostShutdown { + names, + world_size, + per_worker: per_worker_shutdown, + }), + ) + .await +} + +/// Rank > 0: join Pulsing, serve `--per-worker` slot actors until Shutdown. +async fn run_worker_rank(dist: DistEnv, opts: &RunOptions, pythonpath: Vec) -> Result<()> { + let seed = dist.pulsing_seed.to_string(); + let mut last = None; + let system = { + let mut built = None; + for attempt in 1..=60 { + match ActorSystem::builder() + .mailbox_capacity(256) + .addr("0.0.0.0:0") + .seeds([seed.as_str()]) + .build() + .await + { + Ok(s) => { + tracing::debug!(attempt, %seed, "joined Pulsing cluster"); + built = Some(s); + break; + } + Err(e) => { + last = Some(e.to_string()); + tracing::debug!(attempt, error = %e, "waiting for driver Pulsing seed"); + sleep(Duration::from_millis(250)).await; + } + } + } + match built { + Some(s) => s, + None => bail!("failed to join Pulsing at {seed}: {last:?}"), + } + }; + + if let Some(pp) = crate::python_env::merge_pythonpath(&pythonpath) { + std::env::set_var("PYTHONPATH", pp); + } + + spawn_job_control(&system, dist.rank, opts.job_cancel.clone(), None).await?; + + let per_worker = opts.per_worker_inflight.max(1); + let gate = ShutdownGate::new(per_worker); + let _slots = spawn_rank_slots( + &system, + dist.rank, + dist.world_size, + per_worker, + opts, + &pythonpath, + Some(Arc::clone(&gate)), + ) + .await?; + + gate.wait().await; + tracing::debug!(rank = dist.rank, "all worker slots shutdown"); + system + .shutdown() + .await + .map_err(|e| anyhow::anyhow!("shutdown: {e}"))?; + Ok(()) +} + +// ── shared boot helpers ─────────────────────────────────────────────── + +/// Resolve PYTHONPATH extras and optionally set the process env. +fn apply_pythonpath(opts: &RunOptions) -> Vec { + let mut extras = pythonpath_for_script(&opts.script); + extras.extend(opts.pythonpath_extra.iter().cloned()); + if let Some(pp) = crate::python_env::merge_pythonpath(&extras) { + std::env::set_var("PYTHONPATH", pp); + } + extras +} + +fn worker_slot_id(worker: usize, slot: usize, per_worker: usize) -> String { + if per_worker <= 1 { + format!("w{worker}") + } else { + format!("w{worker}s{slot}") + } +} + +async fn spawn_job_control( + system: &Arc, + rank: usize, + job_cancel: CancellationToken, + sched: Option>, +) -> Result { + let name = DistEnv::job_control_name(rank); + let actor = match sched { + Some(s) => JobControlActor::with_scheduler(job_cancel, s), + None => JobControlActor::new(job_cancel), + }; + system + .spawn_named(&name, actor) + .await + .map_err(|e| anyhow::anyhow!("spawn job control {name}: {e}")) +} + +async fn spawn_one_slot( + system: &Arc, + worker: usize, + slot: usize, + n_workers: usize, + per_worker: usize, + opts: &RunOptions, + pythonpath: &[PathBuf], + gate: Option>, +) -> Result<(ActorRef, usize)> { + let name = DistEnv::slot_name(worker, slot, per_worker); + let cfg = WorkerConfig::with_fresh_cache( + worker_slot_id(worker, slot, per_worker), + opts.python.clone(), + pythonpath.to_vec(), + opts.script.clone(), + opts.script_args.clone(), + opts.job_cancel.clone(), + gate, + ); + let wref = spawn_supervised(system, &name, move || Ok(cfg.build())).await?; + let flat = DistEnv::slot_flat_index(worker, slot, n_workers, per_worker); + tracing::debug!(%name, worker, slot, flat, "worker slot ready"); + Ok((wref, flat)) +} + +/// Spawn all slots for one rank (torchrun driver local / worker rank). +/// Returns `(ActorRef, slot-major flat index)` pairs. +async fn spawn_rank_slots( + system: &Arc, + rank: usize, + n_workers: usize, + per_worker: usize, + opts: &RunOptions, + pythonpath: &[PathBuf], + gate: Option>, +) -> Result> { + let mut out = Vec::with_capacity(per_worker); + for slot in 0..per_worker { + out.push( + spawn_one_slot( + system, + rank, + slot, + n_workers, + per_worker, + opts, + pythonpath, + gate.clone(), + ) + .await?, + ); + } + Ok(out) +} + +/// Local fleet: slot-major order over all workers (matches [`DistEnv::slot_names`]). +async fn spawn_local_fleet_slots( + system: &Arc, + n_workers: usize, + per_worker: usize, + opts: &RunOptions, + pythonpath: &[PathBuf], +) -> Result> { + let mut out = Vec::with_capacity(n_workers.saturating_mul(per_worker)); + for slot in 0..per_worker { + for worker in 0..n_workers { + out.push( + spawn_one_slot( + system, worker, slot, n_workers, per_worker, opts, pythonpath, None, + ) + .await?, + ); + } + } + Ok(out) +} + +struct DriverPostShutdown { + names: Vec, + world_size: usize, + per_worker: usize, +} + +/// Shared Driver + observe snapshot + system shutdown. +async fn run_driver_loop( + pool: WorkerPool, + sched: Arc, + opts: &RunOptions, + on_result: impl FnMut(TaskResult) + Send, + cancel_fanout: tokio::task::JoinHandle<()>, + system: Arc, + post: Option, +) -> Result> { + let driver = Driver::new(Arc::clone(&pool), Arc::clone(&sched)); + let snap = spawn_snapshot_loop( + Arc::clone(&opts.observer), + Arc::clone(&sched), + opts.job_cancel.clone(), + Duration::from_secs(1), + ); + let out = driver.run_plan(opts, on_result).await?; + snap.abort(); + cancel_fanout.abort(); + if let Some(p) = post { + if opts.job_cancel.is_cancelled() { + crate::job_control::broadcast_job_cancel(&system, p.world_size).await; + } + shutdown_workers_resolved(&system, &p.names, p.per_worker).await?; + } + system + .shutdown() + .await + .map_err(|e| anyhow::anyhow!("shutdown: {e}"))?; + Ok(out) +} + +async fn wait_and_fill_workers( + system: &Arc, + pool: &WorkerPool, + names: &[String], + timeout: Duration, +) -> Result<()> { + let deadline = tokio::time::Instant::now() + timeout; + loop { + let mut refs = Vec::with_capacity(names.len()); + let mut missing = Vec::new(); + for name in names { + match resolve_actor(system.as_ref(), name).await { + Ok(r) => refs.push(r), + Err(_) => missing.push(name.clone()), + } + } + if missing.is_empty() { + let mut g = pool + .write() + .map_err(|_| anyhow::anyhow!("worker pool lock"))?; + *g = refs; + tracing::debug!(n = g.len(), "all worker slots resolved"); + return Ok(()); + } + if tokio::time::Instant::now() >= deadline { + bail!("timed out waiting for workers: missing {missing:?}"); + } + tracing::debug!(?missing, "waiting for worker gossip/resolve"); + sleep(Duration::from_millis(300)).await; + } +} + +async fn shutdown_workers_resolved( + system: &Arc, + names: &[String], + per_worker: usize, +) -> Result<()> { + for name in names { + let is_local = (0..per_worker.max(1)) + .any(|slot| name == &DistEnv::slot_name(0, slot, per_worker.max(1))); + if is_local { + continue; + } + match resolve_actor(system.as_ref(), name).await { + Ok(w) => { + let _ = + ask_timeout::<_, WorkerReply>(&w, WorkerCommand::Shutdown, ASK_TIMEOUT).await; + } + Err(e) => tracing::warn!(%name, error = %e, "shutdown: resolve failed"), + } + } + sleep(Duration::from_millis(200)).await; + Ok(()) +} diff --git a/crates/persisting-compute/src/scheduler.rs b/crates/persisting-compute/src/scheduler.rs new file mode 100644 index 0000000..81b6494 --- /dev/null +++ b/crates/persisting-compute/src/scheduler.rs @@ -0,0 +1,458 @@ +//! L2 task placement: least-loaded slots with optional sticky prefer. +//! +//! **Primitive:** [`Scheduler`] · [`SlotGuard`] · [`WorkerPool`]. +//! +//! When the runtime flattens `--per-worker` into one actor+host per slot, +//! construct with `Scheduler::new(n_slots, 1)`. Slot-major pool ordering +//! (see [`crate::dist::DistEnv::slot_names`]) keeps least-loaded spreading +//! across logical workers before filling a second slot on the same worker. +//! +//! Consecutive infra ask failures quarantine a slot for the rest of the job +//! so placement skips known-bad hosts. + +use pulsing_actor::prelude::ActorRef; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, RwLock}; +use tokio::sync::Notify; + +pub type WorkerPool = Arc>>; + +/// Consecutive infra failures before a slot is quarantined (job-local). +pub const DEFAULT_QUARANTINE_AFTER: usize = 3; + +/// Shared placement state for one fleet run. +pub struct Scheduler { + loads: Vec, + /// Max concurrent Executes per worker (typically 1 for long/uneven tasks). + per_worker: usize, + notify: Notify, + consecutive_failures: Vec, + quarantined: Vec, + quarantine_after: usize, +} + +impl Scheduler { + pub fn new(n_workers: usize, per_worker: usize) -> Arc { + Self::with_quarantine_after(n_workers, per_worker, DEFAULT_QUARANTINE_AFTER) + } + + pub fn with_quarantine_after( + n_workers: usize, + per_worker: usize, + quarantine_after: usize, + ) -> Arc { + let per_worker = per_worker.max(1); + let n_workers = n_workers.max(1); + let quarantine_after = quarantine_after.max(1); + Arc::new(Self { + loads: (0..n_workers).map(|_| AtomicUsize::new(0)).collect(), + per_worker, + notify: Notify::new(), + consecutive_failures: (0..n_workers).map(|_| AtomicUsize::new(0)).collect(), + quarantined: (0..n_workers).map(|_| AtomicBool::new(false)).collect(), + quarantine_after, + }) + } + + pub fn worker_count(&self) -> usize { + self.loads.len() + } + + pub fn per_worker(&self) -> usize { + self.per_worker + } + + pub fn capacity(&self) -> usize { + self.active_slots().saturating_mul(self.per_worker).max(1) + } + + /// Slots not quarantined. + pub fn active_slots(&self) -> usize { + self.quarantined + .iter() + .filter(|q| !q.load(Ordering::Acquire)) + .count() + } + + pub fn is_quarantined(&self, index: usize) -> bool { + self.quarantined + .get(index) + .map(|q| q.load(Ordering::Acquire)) + .unwrap_or(true) + } + + /// Record a successful ask — clears consecutive failure streak. + pub fn note_success(&self, index: usize) { + if index < self.consecutive_failures.len() { + self.consecutive_failures[index].store(0, Ordering::Release); + } + } + + /// Record an infra ask failure; may quarantine the slot. + pub fn note_failure(&self, index: usize) { + if index >= self.consecutive_failures.len() { + return; + } + let n = self.consecutive_failures[index].fetch_add(1, Ordering::AcqRel) + 1; + if n >= self.quarantine_after && !self.quarantined[index].swap(true, Ordering::AcqRel) { + tracing::warn!( + slot = index, + failures = n, + "quarantining slot after consecutive infra failures" + ); + self.notify.notify_waiters(); + } + } + + /// Immediately quarantine a slot (DeathWatch / explicit drop). + pub fn force_quarantine(&self, index: usize) { + if index >= self.quarantined.len() { + return; + } + if !self.quarantined[index].swap(true, Ordering::AcqRel) { + tracing::warn!(slot = index, "quarantining slot (forced)"); + self.notify.notify_waiters(); + } + } + + /// Reserve a worker slot (least in-flight among those under capacity). + pub async fn acquire(&self) -> Result { + loop { + if let Some(i) = self.try_acquire() { + return Ok(i); + } + if self.active_slots() == 0 { + tracing::error!("all slots quarantined; placement fail-fast"); + return Err(AcquireError::AllQuarantined); + } + self.notify.notified().await; + } + } + + fn slot_usable(&self, i: usize, cur: usize) -> bool { + !self.quarantined[i].load(Ordering::Acquire) && cur < self.per_worker + } + + fn try_acquire(&self) -> Option { + loop { + let mut best: Option<(usize, usize)> = None; + for (i, cell) in self.loads.iter().enumerate() { + let cur = cell.load(Ordering::Acquire); + if !self.slot_usable(i, cur) { + continue; + } + match best { + None => best = Some((i, cur)), + Some((_, b)) if cur < b => best = Some((i, cur)), + Some((bi, b)) if cur == b && i < bi => best = Some((i, cur)), + _ => {} + } + } + let (i, expect) = best?; + match self.loads[i].compare_exchange( + expect, + expect + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return Some(i), + Err(_) => continue, + } + } + } + + /// Prefer `prefer` when that slot still has capacity; otherwise least-loaded. + pub async fn acquire_prefer(&self, prefer: Option) -> Result { + loop { + if let Some(i) = prefer { + if self.try_acquire_index(i) { + return Ok(i); + } + } + if let Some(i) = self.try_acquire() { + return Ok(i); + } + if self.active_slots() == 0 { + tracing::error!("all slots quarantined; placement fail-fast"); + return Err(AcquireError::AllQuarantined); + } + self.notify.notified().await; + } + } + + /// Acquire **only** `slot` (sticky-after-contact). Err if quarantined. + pub async fn acquire_sticky(&self, slot: usize) -> Result { + loop { + if self.is_quarantined(slot) { + return Err(StickyLost::Quarantined(slot)); + } + if self.try_acquire_index(slot) { + return Ok(slot); + } + self.notify.notified().await; + } + } + + fn try_acquire_index(&self, index: usize) -> bool { + if index >= self.loads.len() || self.quarantined[index].load(Ordering::Acquire) { + return false; + } + loop { + let cur = self.loads[index].load(Ordering::Acquire); + if cur >= self.per_worker { + return false; + } + match self.loads[index].compare_exchange( + cur, + cur + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return true, + Err(_) => continue, + } + } + } + + pub fn release(&self, index: usize) { + if index >= self.loads.len() { + return; + } + let prev = self.loads[index].fetch_sub(1, Ordering::AcqRel); + debug_assert!(prev > 0, "scheduler release underflow"); + self.notify.notify_waiters(); + } + + pub fn load_snapshot(&self) -> Vec { + self.loads + .iter() + .map(|c| c.load(Ordering::Relaxed)) + .collect() + } + + pub fn quarantine_snapshot(&self) -> Vec { + self.quarantined + .iter() + .map(|q| q.load(Ordering::Relaxed)) + .collect() + } +} + +/// RAII guard that releases a scheduler slot when dropped. +pub struct SlotGuard { + sched: Arc, + index: usize, +} + +impl SlotGuard { + pub fn index(&self) -> usize { + self.index + } +} + +impl Drop for SlotGuard { + fn drop(&mut self) { + self.sched.release(self.index); + } +} + +impl Scheduler { + pub async fn acquire_guard(self: &Arc) -> Result { + let index = self.acquire().await?; + Ok(SlotGuard { + sched: Arc::clone(self), + index, + }) + } + + pub async fn acquire_guard_prefer( + self: &Arc, + prefer: Option, + ) -> Result { + let index = self.acquire_prefer(prefer).await?; + Ok(SlotGuard { + sched: Arc::clone(self), + index, + }) + } + + pub async fn acquire_guard_sticky( + self: &Arc, + slot: usize, + ) -> Result { + let index = self.acquire_sticky(slot).await?; + Ok(SlotGuard { + sched: Arc::clone(self), + index, + }) + } +} + +/// No usable slots remain (all quarantined). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AcquireError { + AllQuarantined, +} + +impl std::fmt::Display for AcquireError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + AcquireError::AllQuarantined => write!(f, "all worker slots quarantined"), + } + } +} + +/// Sticky-after-contact placement cannot continue on this slot. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StickyLost { + Quarantined(usize), +} + +impl std::fmt::Display for StickyLost { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + StickyLost::Quarantined(s) => { + write!( + f, + "sticky slot {s} quarantined (refuse cross-slot re-execute)" + ) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn prefers_idle_worker() { + let s = Scheduler::new(3, 2); + let a = s.acquire().await.unwrap(); + assert_eq!(a, 0); + let b = s.acquire().await.unwrap(); + assert_eq!(b, 1); + let c = s.acquire().await.unwrap(); + assert_eq!(c, 2); + let d = s.acquire().await.unwrap(); + assert_eq!(d, 0); + s.release(b); + let e = s.acquire().await.unwrap(); + assert_eq!(e, 1); + s.release(a); + s.release(c); + s.release(d); + s.release(e); + assert_eq!(s.load_snapshot(), vec![0, 0, 0]); + } + + #[tokio::test] + async fn acquire_prefer_sticky_when_free() { + let s = Scheduler::new(3, 1); + let a = s.acquire_prefer(None).await.unwrap(); + assert_eq!(a, 0); + s.release(a); + let b = s.acquire_prefer(Some(2)).await.unwrap(); + assert_eq!(b, 2); + s.release(b); + } + + #[tokio::test] + async fn blocks_when_full_then_unblocks() { + let s = Scheduler::new(1, 1); + let g = s.acquire_guard().await.unwrap(); + let s2 = Arc::clone(&s); + let handle = tokio::spawn(async move { s2.acquire().await }); + tokio::task::yield_now().await; + assert!(!handle.is_finished()); + drop(g); + let idx = handle.await.unwrap().unwrap(); + assert_eq!(idx, 0); + s.release(idx); + } + + #[tokio::test] + async fn flat_pool_first_wave_spreads_across_slot0() { + let s = Scheduler::new(6, 1); + let mut got = Vec::new(); + for _ in 0..3 { + got.push(s.acquire().await.unwrap()); + } + assert_eq!(got, vec![0, 1, 2]); + for i in got { + s.release(i); + } + } + + #[tokio::test] + async fn concurrent_acquire_fills_distinct_slots() { + let s = Arc::clone(&Scheduler::new(4, 1)); + let mut handles = Vec::new(); + for _ in 0..4 { + let s2 = Arc::clone(&s); + handles.push(tokio::spawn(async move { s2.acquire().await })); + } + let mut got: Vec = futures::future::join_all(handles) + .await + .into_iter() + .map(|r| r.unwrap().unwrap()) + .collect(); + got.sort(); + assert_eq!(got, vec![0, 1, 2, 3]); + for i in got { + s.release(i); + } + } + + #[tokio::test] + async fn quarantine_skips_slot_on_acquire() { + let s = Scheduler::with_quarantine_after(2, 1, 2); + s.note_failure(0); + s.note_failure(0); + assert!(s.is_quarantined(0)); + let g = s.acquire_guard().await.unwrap(); + assert_eq!(g.index(), 1); + drop(g); + let g2 = s.acquire_guard_prefer(Some(0)).await.unwrap(); + assert_eq!(g2.index(), 1); + } + + #[tokio::test] + async fn acquire_sticky_errors_when_quarantined() { + let s = Scheduler::new(2, 1); + s.force_quarantine(0); + assert!(matches!( + s.acquire_sticky(0).await, + Err(StickyLost::Quarantined(0)) + )); + assert_eq!(s.acquire_sticky(1).await.unwrap(), 1); + s.release(1); + } + + #[tokio::test] + async fn all_quarantined_acquire_fail_fast() { + let s = Scheduler::with_quarantine_after(2, 1, 1); + s.note_failure(0); + s.note_failure(1); + assert_eq!(s.active_slots(), 0); + assert!(matches!( + s.acquire().await, + Err(AcquireError::AllQuarantined) + )); + assert!(matches!( + s.acquire_prefer(Some(0)).await, + Err(AcquireError::AllQuarantined) + )); + } + + #[tokio::test] + async fn success_resets_failure_streak() { + let s = Scheduler::with_quarantine_after(1, 1, 3); + s.note_failure(0); + s.note_failure(0); + s.note_success(0); + s.note_failure(0); + s.note_failure(0); + assert!(!s.is_quarantined(0)); + } +} diff --git a/crates/persisting-compute/src/sink.rs b/crates/persisting-compute/src/sink.rs new file mode 100644 index 0000000..731514f --- /dev/null +++ b/crates/persisting-compute/src/sink.rs @@ -0,0 +1,291 @@ +//! Unique control-plane sink: ready results append here (not from Executors). +//! +//! **Primitive:** [`ResultSink`] · [`persist_terminal`] · [`JsonlFileSink`] (task_id dedup). +//! +//! - Phase-1 ledger: append-only JSONL under `--sink` (`task_id` for `--resume`). +//! - Optional L1: feature `traj-sink` → [`crate::sink_traj::VortexResultSink`] via Tee. + +use crate::checkpoint::CheckpointLedger; +use crate::task::TaskResult; +use anyhow::{Context, Result}; +use async_trait::async_trait; +use std::collections::HashSet; +use std::path::{Path, PathBuf}; +use std::sync::Mutex; +use tokio::fs::OpenOptions; +use tokio::io::AsyncWriteExt; + +#[async_trait] +pub trait ResultSink: Send + Sync { + /// Ready / terminal result that belongs in the asset ledger. + async fn append_ready(&self, result: &TaskResult) -> Result<()>; + /// Side ledger (infra / execute failure). Must not mix into training pool. + async fn append_failure(&self, result: &TaskResult) -> Result<()>; +} + +/// Append-only JSONL under `{root}/ready.ndjson` + `{root}/failures.ndjson`. +/// +/// Idempotent by `task_id`: seeds from existing files on open, then skips duplicates. +/// `seen` is reserved before write and **rolled back** if the durable append fails, +/// so a failed write never permanently blocks a later retry of the same id. +pub struct JsonlFileSink { + root: PathBuf, + seen_ready: Mutex>, + seen_fail: Mutex>, +} + +impl JsonlFileSink { + pub async fn open(root: impl Into) -> Result { + let root = root.into(); + tokio::fs::create_dir_all(&root) + .await + .with_context(|| format!("mkdir sink {}", root.display()))?; + // Seed from disk so infra-retry / crash-resume duplicates do not re-append. + let ledger = CheckpointLedger::load(&root).await?; + Ok(Self { + root, + seen_ready: Mutex::new(ledger.ready), + seen_fail: Mutex::new(ledger.failed), + }) + } + + pub fn root(&self) -> &Path { + &self.root + } + + async fn append_line(&self, file: &str, result: &TaskResult) -> Result<()> { + let path = self.root.join(file); + let line = result.to_ndjson()?; + let mut f = OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .await + .with_context(|| format!("open {}", path.display()))?; + f.write_all(line.as_bytes()).await?; + f.write_all(b"\n").await?; + f.flush().await?; + Ok(()) + } + + /// Reserve `task_id` in `seen`. Returns `false` if already present (caller skips). + fn reserve(seen: &Mutex>, task_id: &str) -> Result { + let mut g = seen + .lock() + .map_err(|_| anyhow::anyhow!("sink lock poisoned"))?; + Ok(g.insert(task_id.to_string())) + } + + fn unreserve(seen: &Mutex>, task_id: &str) { + if let Ok(mut g) = seen.lock() { + g.remove(task_id); + } + } +} + +#[async_trait] +impl ResultSink for JsonlFileSink { + async fn append_ready(&self, result: &TaskResult) -> Result<()> { + if !Self::reserve(&self.seen_ready, &result.task_id)? { + return Ok(()); + } + match self.append_line("ready.ndjson", result).await { + Ok(()) => Ok(()), + Err(e) => { + Self::unreserve(&self.seen_ready, &result.task_id); + Err(e) + } + } + } + + async fn append_failure(&self, result: &TaskResult) -> Result<()> { + if !Self::reserve(&self.seen_fail, &result.task_id)? { + return Ok(()); + } + match self.append_line("failures.ndjson", result).await { + Ok(()) => Ok(()), + Err(e) => { + Self::unreserve(&self.seen_fail, &result.task_id); + Err(e) + } + } + } +} + +/// Fan-out to several sinks (e.g. stdout view + durable JSONL). +pub struct TeeSink { + sinks: Vec>, +} + +impl TeeSink { + pub fn new(sinks: Vec>) -> Self { + Self { sinks } + } +} + +#[async_trait] +impl ResultSink for TeeSink { + async fn append_ready(&self, result: &TaskResult) -> Result<()> { + let mut first_err: Option = None; + for s in &self.sinks { + if let Err(e) = s.append_ready(result).await { + tracing::error!( + task_id = %result.task_id, + error = %e, + "tee sink append_ready failed (continuing siblings)" + ); + if first_err.is_none() { + first_err = Some(e); + } + } + } + match first_err { + Some(e) => Err(e), + None => Ok(()), + } + } + + async fn append_failure(&self, result: &TaskResult) -> Result<()> { + let mut first_err: Option = None; + for s in &self.sinks { + if let Err(e) = s.append_failure(result).await { + tracing::error!( + task_id = %result.task_id, + error = %e, + "tee sink append_failure failed (continuing siblings)" + ); + if first_err.is_none() { + first_err = Some(e); + } + } + } + match first_err { + Some(e) => Err(e), + None => Ok(()), + } + } +} + +/// Route by result status: cancelled/failed → failure ledger; ok → ready. +pub async fn persist_terminal(sink: &dyn ResultSink, result: &TaskResult) -> Result<()> { + if result.ok && !result.cancelled { + sink.append_ready(result).await + } else { + sink.append_failure(result).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + use tokio::io::AsyncWriteExt; + + #[tokio::test] + async fn open_seeds_seen_and_skips_duplicate_append() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + let mut ready = tokio::fs::File::create(root.join("ready.ndjson")) + .await + .unwrap(); + ready + .write_all(br#"{"task_id":"t-0","ok":true}"#) + .await + .unwrap(); + ready.write_all(b"\n").await.unwrap(); + + let sink = JsonlFileSink::open(root).await.unwrap(); + let dup = TaskResult::success("t-0", json!({"x": 1}), "w0", 0.0); + sink.append_ready(&dup).await.unwrap(); + + let body = tokio::fs::read_to_string(root.join("ready.ndjson")) + .await + .unwrap(); + assert_eq!(body.lines().count(), 1, "duplicate must not append: {body}"); + } + + #[tokio::test] + async fn persist_terminal_routes_ok_and_fail() { + let dir = tempfile::tempdir().unwrap(); + let sink = JsonlFileSink::open(dir.path()).await.unwrap(); + persist_terminal(&sink, &TaskResult::success("ok-1", json!(1), "w0", 0.0)) + .await + .unwrap(); + persist_terminal(&sink, &TaskResult::cancelled("c-1")) + .await + .unwrap(); + persist_terminal(&sink, &TaskResult::failure("f-1", "e", None, "w0", 0.0)) + .await + .unwrap(); + let ready = tokio::fs::read_to_string(dir.path().join("ready.ndjson")) + .await + .unwrap(); + let fail = tokio::fs::read_to_string(dir.path().join("failures.ndjson")) + .await + .unwrap(); + assert!(ready.contains("ok-1")); + assert!(fail.contains("c-1") && fail.contains("f-1")); + assert!(!ready.contains("c-1")); + } + + #[tokio::test] + async fn tee_fanout_to_jsonl() { + let dir = tempfile::tempdir().unwrap(); + let a = JsonlFileSink::open(dir.path().join("a")).await.unwrap(); + let b = JsonlFileSink::open(dir.path().join("b")).await.unwrap(); + let tee = TeeSink::new(vec![Box::new(a), Box::new(b)]); + persist_terminal(&tee, &TaskResult::success("t", json!(1), "w0", 0.0)) + .await + .unwrap(); + for sub in ["a", "b"] { + let body = tokio::fs::read_to_string(dir.path().join(sub).join("ready.ndjson")) + .await + .unwrap(); + assert!(body.contains("\"task_id\":\"t\"")); + } + } + + #[tokio::test] + async fn reopen_skips_duplicate_ready_from_disk() { + let dir = tempfile::tempdir().unwrap(); + { + let sink = JsonlFileSink::open(dir.path()).await.unwrap(); + persist_terminal(&sink, &TaskResult::success("t-0", json!(1), "w0", 0.0)) + .await + .unwrap(); + } + let sink = JsonlFileSink::open(dir.path()).await.unwrap(); + persist_terminal(&sink, &TaskResult::success("t-0", json!(99), "w0", 0.0)) + .await + .unwrap(); + let body = tokio::fs::read_to_string(dir.path().join("ready.ndjson")) + .await + .unwrap(); + assert_eq!(body.lines().count(), 1); + } + + #[tokio::test] + async fn write_failure_unreserves_seen_so_retry_can_persist() { + // P0 regression: seen-before-write must not permanently drop a task_id. + let dir = tempfile::tempdir().unwrap(); + let sink = JsonlFileSink::open(dir.path()).await.unwrap(); + // Make ready.ndjson a directory so OpenOptions::open fails. + tokio::fs::create_dir(dir.path().join("ready.ndjson")) + .await + .unwrap(); + let r = TaskResult::success("t-stuck", json!(1), "w0", 0.0); + assert!( + sink.append_ready(&r).await.is_err(), + "first append must fail" + ); + tokio::fs::remove_dir(dir.path().join("ready.ndjson")) + .await + .unwrap(); + sink.append_ready(&r).await.expect("retry after fix path"); + let body = tokio::fs::read_to_string(dir.path().join("ready.ndjson")) + .await + .unwrap(); + assert_eq!(body.lines().count(), 1); + assert!(body.contains("t-stuck")); + } +} diff --git a/crates/persisting-compute/src/sink_traj.rs b/crates/persisting-compute/src/sink_traj.rs new file mode 100644 index 0000000..228400f --- /dev/null +++ b/crates/persisting-compute/src/sink_traj.rs @@ -0,0 +1,138 @@ +//! Vortex trajectory sink: TaskResult → CaptureRecord → TrajectoryAppend. +//! +//! Enabled with feature `traj-sink`. Always Tee with [`JsonlFileSink`] so +//! `--resume` keeps using the JSONL task_id ledger. + +use crate::sink::ResultSink; +use crate::task::TaskResult; +use anyhow::{Context, Result}; +use async_trait::async_trait; +use persisting_capture::record::{now_rfc3339, record_to_engine_line, CaptureRecord}; +use persisting_proto::{TrajectoryAppendRequest, TrajectoryStorageFormat}; +use std::collections::HashSet; +use std::sync::Mutex; + +/// Append terminal compute results as `compute.result` / `compute.failure` events. +pub struct VortexResultSink { + storage: String, + agent_id: String, + session_id: String, + seen: Mutex>, +} + +impl VortexResultSink { + pub fn new( + storage: impl Into, + agent_id: impl Into, + session_id: impl Into, + ) -> Self { + Self { + storage: storage.into(), + agent_id: agent_id.into(), + session_id: session_id.into(), + seen: Mutex::new(HashSet::new()), + } + } + + pub fn storage(&self) -> &str { + &self.storage + } + + pub fn agent_id(&self) -> &str { + &self.agent_id + } + + pub fn session_id(&self) -> &str { + &self.session_id + } + + /// Seed dedup set (e.g. from JSONL ledger) so re-appends skip known `task_id`s. + pub fn seed_seen(&self, ids: impl IntoIterator) { + if let Ok(mut seen) = self.seen.lock() { + seen.extend(ids); + } + } + + fn to_record(&self, result: &TaskResult) -> Result { + let kind = if result.ok && !result.cancelled { + "compute.result" + } else { + "compute.failure" + }; + let payload = serde_json::to_value(result).context("TaskResult to JSON")?; + Ok(CaptureRecord { + seq: 0, + source: "persisting-compute".into(), + kind: kind.into(), + timestamp: Some(now_rfc3339()), + session_id: Some(self.session_id.clone()), + agent_id: Some(self.agent_id.clone()), + parent_uuid: None, + trace_id: None, + call_id: Some(result.task_id.clone()), + subagent_id: None, + parent_agent_id: None, + branch: None, + parent_call_id: None, + payload, + }) + } + + async fn append_record(&self, result: &TaskResult) -> Result<()> { + { + let mut seen = self + .seen + .lock() + .map_err(|_| anyhow::anyhow!("vortex sink lock poisoned"))?; + if !seen.insert(result.task_id.clone()) { + return Ok(()); + } + } + let write = async { + let rec = self.to_record(result)?; + let line = record_to_engine_line(&rec).context("encode CaptureRecord")?; + let req = TrajectoryAppendRequest { + storage: self.storage.clone(), + agent_id: self.agent_id.clone(), + session_id: self.session_id.clone(), + root_session_id: None, + records_ronl: line, + storage_format: TrajectoryStorageFormat::Vortex, + }; + // bridge is sync (blocks on runtime inside engine); ok from async via spawn_blocking. + let resp = + tokio::task::spawn_blocking(move || persisting_engine::trajectory_append(req)) + .await + .context("join trajectory_append")? + .context("trajectory_append")?; + tracing::debug!( + task_id = %result.task_id, + accepted = resp.accepted_records, + dataset = %resp.dataset, + "compute result appended to vortex" + ); + Ok(()) + }; + match write.await { + Ok(()) => Ok(()), + Err(e) => { + // Roll back reservation so a later retry is not permanently skipped. + if let Ok(mut seen) = self.seen.lock() { + seen.remove(&result.task_id); + } + Err(e) + } + } + } +} + +#[async_trait] +impl ResultSink for VortexResultSink { + async fn append_ready(&self, result: &TaskResult) -> Result<()> { + self.append_record(result).await + } + + async fn append_failure(&self, result: &TaskResult) -> Result<()> { + self.append_record(result).await + } +} diff --git a/crates/persisting-compute/src/sink_writer.rs b/crates/persisting-compute/src/sink_writer.rs new file mode 100644 index 0000000..5eec634 --- /dev/null +++ b/crates/persisting-compute/src/sink_writer.rs @@ -0,0 +1,199 @@ +//! Bounded async sink writer — keeps Driver `on_result` off the disk path. +//! +//! Completions are enqueued via async `send` (natural backpressure — no +//! `block_in_place`). A single background task runs [`persist_terminal`] + +//! checkpoint notes. Persist failures are counted and surface from +//! [`SinkWriterHandle::join`]; the task id is unclaimed from [`SkipSet`] so a +//! later `--resume` can rediscover work that never hit durable storage. + +use crate::checkpoint::CheckpointTracker; +use crate::sink::{persist_terminal, ResultSink}; +use crate::skip::SkipSet; +use crate::task::TaskResult; +use anyhow::{bail, Context, Result}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; + +/// Cloneable enqueue handle for Driver (awaited on the completion path). +#[derive(Clone)] +pub struct SinkSubmitter { + tx: mpsc::Sender, +} + +impl SinkSubmitter { + /// Async enqueue — when the bound is full, waits (back-pressures Driver). + pub async fn submit(&self, result: TaskResult) -> Result<()> { + self.tx + .send(result) + .await + .map_err(|_| anyhow::anyhow!("sink writer closed; cannot enqueue")) + } +} + +#[derive(Default)] +struct PersistErrors { + count: AtomicUsize, + first: Mutex>, +} + +impl PersistErrors { + fn note(&self, task_id: &str, err: impl std::fmt::Display) { + let prev = self.count.fetch_add(1, Ordering::AcqRel); + if prev == 0 { + if let Ok(mut g) = self.first.lock() { + *g = Some(format!("{task_id}: {err}")); + } + } + } + + fn count(&self) -> usize { + self.count.load(Ordering::Acquire) + } + + fn first_msg(&self) -> Option { + self.first.lock().ok().and_then(|g| g.clone()) + } +} + +/// Handle returned by [`spawn_sink_writer`]. Await [`SinkWriterHandle::join`] to drain. +pub struct SinkWriterHandle { + submitter: SinkSubmitter, + join: JoinHandle<()>, + errors: Arc, +} + +impl SinkWriterHandle { + pub fn submitter(&self) -> SinkSubmitter { + self.submitter.clone() + } + + pub async fn submit(&self, result: TaskResult) -> Result<()> { + self.submitter.submit(result).await + } + + /// Drop the queue sender, wait for the writer, then fail if any persist errored. + pub async fn join(self) -> Result<()> { + drop(self.submitter); + self.join.await.context("sink writer task join")?; + let n = self.errors.count(); + if n > 0 { + let first = self.errors.first_msg().unwrap_or_else(|| "unknown".into()); + bail!("sink persist failed for {n} result(s); first: {first}"); + } + Ok(()) + } +} + +/// Spawn a dedicated persist task. `capacity` bounds queued completions. +pub fn spawn_sink_writer( + sink: Arc, + checkpoint: Option>, + skip: Option, + capacity: usize, +) -> SinkWriterHandle { + let capacity = capacity.max(1); + let (tx, mut rx) = mpsc::channel::(capacity); + let errors = Arc::new(PersistErrors::default()); + let errors_bg = Arc::clone(&errors); + let join = tokio::spawn(async move { + while let Some(r) = rx.recv().await { + if let Err(e) = persist_terminal(sink.as_ref(), &r).await { + tracing::error!(task_id = %r.task_id, error = %e, "sink persist failed"); + errors_bg.note(&r.task_id, &e); + // Not durable — allow a future resume / duplicate plan yield to reclaim. + if let Some(skip) = &skip { + skip.remove(&r.task_id); + } + continue; + } + if let Some(ckpt) = &checkpoint { + ckpt.note_terminal(r.ok, r.cancelled); + let _ = ckpt.maybe_flush().await; + } + } + }); + SinkWriterHandle { + submitter: SinkSubmitter { tx }, + join, + errors, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sink::JsonlFileSink; + use async_trait::async_trait; + use serde_json::json; + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn writer_persists_and_drains() { + let dir = tempfile::tempdir().unwrap(); + let sink: Arc = Arc::new(JsonlFileSink::open(dir.path()).await.unwrap()); + let w = spawn_sink_writer(Arc::clone(&sink), None, None, 8); + w.submit(TaskResult::success("t-0", json!(1), "w0", 0.0)) + .await + .unwrap(); + w.submit(TaskResult::failure("t-1", "e", None, "w0", 0.0)) + .await + .unwrap(); + w.join().await.unwrap(); + let ready = tokio::fs::read_to_string(dir.path().join("ready.ndjson")) + .await + .unwrap(); + let fail = tokio::fs::read_to_string(dir.path().join("failures.ndjson")) + .await + .unwrap(); + assert!(ready.contains("t-0")); + assert!(fail.contains("t-1")); + } + + struct AlwaysFailSink; + + #[async_trait] + impl ResultSink for AlwaysFailSink { + async fn append_ready(&self, _: &TaskResult) -> Result<()> { + bail!("disk full") + } + async fn append_failure(&self, _: &TaskResult) -> Result<()> { + bail!("disk full") + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn persist_failure_surfaces_on_join_and_unclaims_skip() { + let skip = SkipSet::from_iter(["t-0".into()]); + assert!(skip.contains("t-0")); + let w = spawn_sink_writer(Arc::new(AlwaysFailSink), None, Some(skip.clone()), 4); + w.submit(TaskResult::success("t-0", json!(1), "w0", 0.0)) + .await + .unwrap(); + let err = w.join().await.unwrap_err(); + assert!(err.to_string().contains("sink persist failed")); + assert!( + !skip.contains("t-0"), + "failed persist must unclaim so resume can rediscover" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn submit_is_async_backpressure_not_block_in_place() { + // Smoke: async send path works under modest load (no block_in_place). + let dir = tempfile::tempdir().unwrap(); + let sink: Arc = Arc::new(JsonlFileSink::open(dir.path()).await.unwrap()); + let w = spawn_sink_writer(Arc::clone(&sink), None, None, 2); + for i in 0..4 { + w.submit(TaskResult::success(format!("t-{i}"), json!(i), "w0", 0.0)) + .await + .unwrap(); + } + // Must not retain submitter clones across join (keeps writer channel open). + w.join().await.unwrap(); + let ready = tokio::fs::read_to_string(dir.path().join("ready.ndjson")) + .await + .unwrap(); + assert_eq!(ready.lines().count(), 4); + } +} diff --git a/crates/persisting-compute/src/skip.rs b/crates/persisting-compute/src/skip.rs new file mode 100644 index 0000000..c81ad32 --- /dev/null +++ b/crates/persisting-compute/src/skip.rs @@ -0,0 +1,80 @@ +//! Live skip set: task_ids that must not be dispatched. +//! +//! Seeded from `--resume` (ready ∪ failures) and grown as tasks are claimed / +//! completed in the current job — so mid-run sink persistence and duplicate +//! `plan()` yields do not re-dispatch the same id onto another worker. + +use std::collections::HashSet; +use std::sync::{Arc, Mutex}; + +/// Shared, mutable set of task ids to skip (or already claimed). +#[derive(Clone, Default)] +pub struct SkipSet { + inner: Arc>>, +} + +impl SkipSet { + pub fn new() -> Self { + Self::default() + } + + pub fn from_iter(ids: impl IntoIterator) -> Self { + Self { + inner: Arc::new(Mutex::new(ids.into_iter().collect())), + } + } + + pub fn contains(&self, id: &str) -> bool { + self.inner.lock().map(|g| g.contains(id)).unwrap_or(false) + } + + /// Insert `id`. Returns `true` if it was newly claimed (caller may dispatch). + pub fn insert(&self, id: impl Into) -> bool { + self.inner + .lock() + .map(|mut g| g.insert(id.into())) + .unwrap_or(false) + } + + /// Drop a claim (e.g. sink persist failed — id was never durable). + pub fn remove(&self, id: &str) -> bool { + self.inner.lock().map(|mut g| g.remove(id)).unwrap_or(false) + } + + pub fn len(&self) -> usize { + self.inner.lock().map(|g| g.len()).unwrap_or(0) + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn claim_once() { + let s = SkipSet::new(); + assert!(s.insert("t-0")); + assert!(!s.insert("t-0")); + assert!(s.contains("t-0")); + } + + #[test] + fn from_iter_seeds() { + let s = SkipSet::from_iter(["a".into(), "b".into()]); + assert!(s.contains("a") && s.contains("b")); + assert!(!s.insert("a")); + } + + #[test] + fn remove_unclaims() { + let s = SkipSet::new(); + assert!(s.insert("t-0")); + assert!(s.remove("t-0")); + assert!(!s.contains("t-0")); + assert!(s.insert("t-0")); + } +} diff --git a/crates/persisting-compute/src/task.rs b/crates/persisting-compute/src/task.rs new file mode 100644 index 0000000..0cf44af --- /dev/null +++ b/crates/persisting-compute/src/task.rs @@ -0,0 +1,242 @@ +//! Task expression / result wire format (one JSON object per line). +//! +//! **Primitive:** [`TaskExpr`] (plan item) · [`TaskResult`] (terminal outcome). +//! Product contract: `plan()` yield shape `{id, …fields}` ↔ `execute(item)`. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; +use std::time::{SystemTime, UNIX_EPOCH}; + +fn now_secs() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs_f64()) + .unwrap_or(0.0) +} + +/// One unit of work in an execution plan. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskExpr { + pub id: String, + #[serde(default = "default_op")] + pub op: String, + #[serde(default)] + pub args: HashMap, + #[serde(default)] + pub meta: HashMap, +} + +fn default_op() -> String { + "execute".into() +} + +impl TaskExpr { + pub fn from_value(mut v: Value) -> anyhow::Result { + let obj = v + .as_object_mut() + .ok_or_else(|| anyhow::anyhow!("task must be a JSON object"))?; + + let id = obj + .remove("id") + .or_else(|| obj.remove("task_id")) + .and_then(|x| match x { + Value::String(s) => Some(s), + Value::Number(n) => Some(n.to_string()), + _ => None, + }) + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + + let op = obj + .remove("op") + .or_else(|| obj.remove("type")) + .and_then(|x| x.as_str().map(|s| s.to_string())) + .unwrap_or_else(|| "execute".into()); + + let meta = match obj.remove("meta") { + Some(Value::Object(m)) => m.into_iter().collect(), + Some(_) => { + return Err(anyhow::anyhow!("task.meta must be an object")); + } + None => HashMap::new(), + }; + + let args = match obj.remove("args") { + Some(Value::Object(m)) => m.into_iter().collect(), + Some(_) => { + return Err(anyhow::anyhow!("task.args must be an object")); + } + None => { + // Flat payload: remaining fields become args. + std::mem::take(obj).into_iter().collect() + } + }; + + Ok(Self { id, op, args, meta }) + } + + pub fn to_ndjson(&self) -> anyhow::Result { + Ok(serde_json::to_string(self)?) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskResult { + pub task_id: String, + pub ok: bool, + /// L2 cancel (not an execute failure). + #[serde(default)] + pub cancelled: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub value: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub traceback: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub worker: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub started_at: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub finished_at: Option, + /// How many L2 infra retries before this terminal result (0 = first try). + #[serde(default, skip_serializing_if = "is_zero")] + pub infra_retries: u32, +} + +fn is_zero(v: &u32) -> bool { + *v == 0 +} + +impl TaskResult { + pub fn success( + task_id: impl Into, + value: Value, + worker: &str, + started_at: f64, + ) -> Self { + Self { + task_id: task_id.into(), + ok: true, + cancelled: false, + value: Some(value), + error: None, + traceback: None, + worker: Some(worker.to_string()), + started_at: Some(started_at), + finished_at: Some(now_secs()), + infra_retries: 0, + } + } + + pub fn failure( + task_id: impl Into, + error: impl Into, + traceback: Option, + worker: &str, + started_at: f64, + ) -> Self { + Self { + task_id: task_id.into(), + ok: false, + cancelled: false, + value: None, + error: Some(error.into()), + traceback, + worker: Some(worker.to_string()), + started_at: Some(started_at), + finished_at: Some(now_secs()), + infra_retries: 0, + } + } + + pub fn cancelled(task_id: impl Into) -> Self { + let started = now_secs(); + Self { + task_id: task_id.into(), + ok: false, + cancelled: true, + value: None, + error: Some("cancelled".into()), + traceback: None, + worker: None, + started_at: Some(started), + finished_at: Some(started), + infra_retries: 0, + } + } + + pub fn to_ndjson(&self) -> anyhow::Result { + Ok(serde_json::to_string(self)?) + } +} + +pub fn unix_now() -> f64 { + now_secs() +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn flat_payload_becomes_args() { + let t = TaskExpr::from_value(json!({"id": "t-0", "x": 1, "y": "a"})).unwrap(); + assert_eq!(t.id, "t-0"); + assert_eq!(t.op, "execute"); + assert_eq!(t.args.get("x"), Some(&json!(1))); + assert_eq!(t.args.get("y"), Some(&json!("a"))); + } + + #[test] + fn nested_args_and_task_id_alias() { + let t = TaskExpr::from_value(json!({ + "task_id": 7, + "type": "execute", + "args": {"n": 3}, + "meta": {"prio": 1} + })) + .unwrap(); + assert_eq!(t.id, "7"); + assert_eq!(t.op, "execute"); + assert_eq!(t.args.get("n"), Some(&json!(3))); + assert_eq!(t.meta.get("prio"), Some(&json!(1))); + } + + #[test] + fn rejects_non_object_meta() { + assert!(TaskExpr::from_value(json!({"id": "t", "meta": 1})).is_err()); + } + + #[test] + fn result_roundtrip_flags() { + let ok = TaskResult::success("t", json!({"v": 1}), "w0", 1.0); + let line = ok.to_ndjson().unwrap(); + let back: TaskResult = serde_json::from_str(&line).unwrap(); + assert!(back.ok && !back.cancelled); + + let c = TaskResult::cancelled("t"); + assert!(c.cancelled && !c.ok); + let f = TaskResult::failure("t", "boom", Some("tb".into()), "w0", 0.0); + assert!(!f.ok && f.traceback.as_deref() == Some("tb")); + } + + #[test] + fn missing_id_gets_uuid() { + let t = TaskExpr::from_value(json!({"x": 1})).unwrap(); + assert!(!t.id.is_empty()); + assert_eq!(t.args.get("x"), Some(&json!(1))); + } + + #[test] + fn flat_yield_shape_roundtrips_to_result_ndjson() { + // plan() yield {"id", ...fields} → TaskExpr → TaskResult ndjson for sink. + let wire = TaskExpr::from_value(json!({"id": "t-0", "x": 2, "tag": "a"})).unwrap(); + assert_eq!(wire.op, "execute"); + assert_eq!(wire.args.get("x"), Some(&json!(2))); + let ok = TaskResult::success(&wire.id, json!({"x2": 4}), "w0", 0.0); + let line = ok.to_ndjson().unwrap(); + assert!(line.contains("\"task_id\":\"t-0\"") && line.contains("\"ok\":true")); + } +} diff --git a/crates/persisting-compute/src/worker.rs b/crates/persisting-compute/src/worker.rs new file mode 100644 index 0000000..8a6be07 --- /dev/null +++ b/crates/persisting-compute/src/worker.rs @@ -0,0 +1,327 @@ +//! Pulsing WorkerActor — dispatches TaskExpr to plan.py `execute(item)`. +//! +//! Semantic block: **fleet worker seam** (see [`crate::blocks`]). +//! - Shares a job [`CancellationToken`]: in-flight execute kills the Python host. +//! - Uses [`crate::result_cache::ResultCache`] for same-worker infra-retry idempotency. +//! - Spawned via factory + [`SupervisionSpec`] so Pulsing can restart a failed slot. + +use crate::executor::ExecutorRouter; +use crate::result_cache::{ResultCache, DEFAULT_RESULT_CACHE_CAP}; +use crate::task::{TaskExpr, TaskResult}; +use pulsing_actor::prelude::*; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum WorkerCommand { + Execute { task_json: Vec }, + Shutdown, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum WorkerReply { + Result { result_json: Vec }, + Bye, +} + +/// Shared gate so supervised restarts still signal rank shutdown. +#[derive(Debug)] +pub struct ShutdownGate { + remaining: AtomicUsize, + done: Notify, +} + +impl ShutdownGate { + pub fn new(slots: usize) -> Arc { + Arc::new(Self { + remaining: AtomicUsize::new(slots.max(1)), + done: Notify::new(), + }) + } + + pub fn note_shutdown(&self) { + let prev = self.remaining.fetch_sub(1, Ordering::AcqRel); + if prev <= 1 { + self.remaining.store(0, Ordering::Release); + self.done.notify_waiters(); + } + } + + pub async fn wait(&self) { + loop { + if self.remaining.load(Ordering::Acquire) == 0 { + return; + } + self.done.notified().await; + } + } +} + +/// Cloneable config for supervised `spawn_factory`. +/// +/// [`result_cache`] is shared across supervised restarts so sticky infra +/// re-ask still hits cached TaskResults after the actor is rebuilt. +#[derive(Clone)] +pub struct WorkerConfig { + pub worker_id: String, + pub python: PathBuf, + pub pythonpath_extra: Vec, + pub plan_script: PathBuf, + pub script_args: Vec, + pub job_cancel: CancellationToken, + pub shutdown_gate: Option>, + /// Slot-scoped cache; one Arc per logical slot, shared by factory rebuilds. + pub result_cache: Arc>, +} + +impl WorkerConfig { + pub fn with_fresh_cache( + worker_id: impl Into, + python: PathBuf, + pythonpath_extra: Vec, + plan_script: PathBuf, + script_args: Vec, + job_cancel: CancellationToken, + shutdown_gate: Option>, + ) -> Self { + Self { + worker_id: worker_id.into(), + python, + pythonpath_extra, + plan_script, + script_args, + job_cancel, + shutdown_gate, + result_cache: Arc::new(Mutex::new(ResultCache::new(DEFAULT_RESULT_CACHE_CAP))), + } + } + + pub fn build(&self) -> WorkerActor { + WorkerActor { + worker_id: self.worker_id.clone(), + executors: Arc::new(ExecutorRouter::local_stack( + self.python.clone(), + self.pythonpath_extra.clone(), + self.plan_script.clone(), + self.script_args.clone(), + )), + done: 0, + shutdown_gate: self.shutdown_gate.clone(), + job_cancel: self.job_cancel.clone(), + result_cache: Arc::clone(&self.result_cache), + } + } +} + +pub struct WorkerActor { + pub worker_id: String, + executors: Arc, + done: u64, + shutdown_gate: Option>, + job_cancel: CancellationToken, + result_cache: Arc>, +} + +impl WorkerActor { + pub fn with_plan( + worker_id: impl Into, + python: PathBuf, + pythonpath_extra: Vec, + plan_script: PathBuf, + script_args: Vec, + job_cancel: CancellationToken, + ) -> Self { + WorkerConfig::with_fresh_cache( + worker_id, + python, + pythonpath_extra, + plan_script, + script_args, + job_cancel, + None, + ) + .build() + } + + pub fn from_config(cfg: &WorkerConfig) -> Self { + cfg.build() + } + + async fn execute(&mut self, task: TaskExpr) -> TaskResult { + if let Ok(g) = self.result_cache.lock() { + if let Some(cached) = g.get(&task.id) { + tracing::debug!( + task_id = %task.id, + worker = %self.worker_id, + "infra idempotency: returning cached TaskResult" + ); + return cached.clone(); + } + } + let task_id = task.id.clone(); + let r = self + .executors + .run_with_cancel(task, &self.worker_id, self.job_cancel.clone()) + .await; + if r.ok { + self.done += 1; + } + if let Ok(mut g) = self.result_cache.lock() { + g.put(task_id, r.clone()); + } + r + } +} + +#[async_trait] +impl Actor for WorkerActor { + fn metadata(&self) -> HashMap { + HashMap::from([ + ("role".into(), "compute-worker".into()), + ("worker_id".into(), self.worker_id.clone()), + ("done".into(), self.done.to_string()), + ]) + } + + async fn receive( + &mut self, + msg: Message, + _ctx: &mut ActorContext, + ) -> pulsing_actor::error::Result { + let cmd: WorkerCommand = msg.unpack()?; + let reply = match cmd { + WorkerCommand::Shutdown => { + self.executors.shutdown().await; + if let Some(gate) = &self.shutdown_gate { + gate.note_shutdown(); + } + WorkerReply::Bye + } + WorkerCommand::Execute { task_json } => { + let task: TaskExpr = serde_json::from_slice(&task_json).map_err(|e| { + pulsing_actor::error::PulsingError::from( + pulsing_actor::error::RuntimeError::Serialization(e.to_string()), + ) + })?; + let result = self.execute(task).await; + let result_json = serde_json::to_vec(&result).map_err(|e| { + pulsing_actor::error::PulsingError::from( + pulsing_actor::error::RuntimeError::Serialization(e.to_string()), + ) + })?; + WorkerReply::Result { result_json } + } + }; + Message::pack(&reply) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn result_cache_avoids_second_execute() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("count.py"); + let counter = dir.path().join("counter.txt"); + let counter_lit = counter.display().to_string(); + std::fs::write( + &path, + format!( + r#" +COUNTER = {counter_lit:?} + +def plan(): + yield {{"id": "unused"}} + +def execute(item): + with open(COUNTER, "a") as f: + f.write("run\n") + return {{"ok": True}} +"# + ), + ) + .unwrap(); + let token = CancellationToken::new(); + let system = Arc::new( + ActorSystem::builder() + .mailbox_capacity(16) + .build() + .await + .unwrap(), + ); + let cfg = WorkerConfig { + worker_id: "w0".into(), + python: PathBuf::from("python3"), + pythonpath_extra: vec![], + plan_script: path, + script_args: vec![], + job_cancel: token, + shutdown_gate: None, + result_cache: Arc::new(Mutex::new(ResultCache::new(DEFAULT_RESULT_CACHE_CAP))), + }; + let w = crate::pulsing_ext::spawn_supervised(&system, "compute/worker/0", move || { + Ok(cfg.build()) + }) + .await + .unwrap(); + let task = TaskExpr::from_value(json!({"id": "t-0", "x": 1})).unwrap(); + let task_json = serde_json::to_vec(&task).unwrap(); + for _ in 0..2 { + let reply = w + .ask::<_, WorkerReply>(WorkerCommand::Execute { + task_json: task_json.clone(), + }) + .await + .unwrap(); + assert!(matches!(reply, WorkerReply::Result { .. })); + } + let runs = std::fs::read_to_string(&counter).unwrap(); + assert_eq!(runs.lines().count(), 1, "second ask must hit result cache"); + system.shutdown().await.unwrap(); + } + + #[tokio::test] + async fn shutdown_gate_fires_when_all_slots_done() { + let gate = ShutdownGate::new(2); + let g = Arc::clone(&gate); + let h = tokio::spawn(async move { + g.wait().await; + }); + gate.note_shutdown(); + assert!(!h.is_finished()); + gate.note_shutdown(); + h.await.unwrap(); + } + + #[test] + fn shared_result_cache_survives_rebuild() { + let cache = Arc::new(Mutex::new(ResultCache::new(8))); + let cfg = WorkerConfig { + worker_id: "w0".into(), + python: PathBuf::from("python3"), + pythonpath_extra: vec![], + plan_script: PathBuf::from("/dev/null"), + script_args: vec![], + job_cancel: CancellationToken::new(), + shutdown_gate: None, + result_cache: Arc::clone(&cache), + }; + let _a = cfg.build(); + cache + .lock() + .unwrap() + .put("t-0", TaskResult::success("t-0", json!(1), "w0", 0.0)); + // Supervised restart = factory rebuild; same Arc must still hold entries. + let _b = cfg.build(); + assert!(cache.lock().unwrap().get("t-0").is_some()); + assert!(Arc::ptr_eq(&cfg.result_cache, &cache)); + } +} diff --git a/crates/persisting-compute/tests/common/mod.rs b/crates/persisting-compute/tests/common/mod.rs new file mode 100644 index 0000000..128b508 --- /dev/null +++ b/crates/persisting-compute/tests/common/mod.rs @@ -0,0 +1,31 @@ +//! Helpers for crate integration tests (`tests/integration_*.rs`). + +#![allow(dead_code)] + +use persisting_compute::{Observer, RunOptions, SkipSet}; +use std::path::{Path, PathBuf}; +use tokio_util::sync::CancellationToken; + +pub fn write_plan(dir: &Path, name: &str, body: &str) -> PathBuf { + let path = dir.join(name); + std::fs::write(&path, body).expect("write plan"); + path +} + +pub fn run_opts(script: PathBuf) -> RunOptions { + RunOptions { + script, + python: PathBuf::from("python3"), + workers: 2, + max_inflight: 4, + per_worker_inflight: 1, + pythonpath_extra: vec![], + script_args: vec![], + infra_retries: 0, + job_cancel: CancellationToken::new(), + observer: Observer::disabled(), + skip_task_ids: SkipSet::new(), + checkpoint: None, + sink_submitter: None, + } +} diff --git a/crates/persisting-compute/tests/integration_local.rs b/crates/persisting-compute/tests/integration_local.rs new file mode 100644 index 0000000..a4bb5f6 --- /dev/null +++ b/crates/persisting-compute/tests/integration_local.rs @@ -0,0 +1,330 @@ +//! Integration: local fleet paths that cross Driver + Runtime + Worker + L3. +//! +//! Contract unit tests live next to their modules (`src/*.rs`). This file only +//! covers multi-module behavior. + +mod common; + +use persisting_compute::{ + run_local_fleet, spawn_sink_writer, CheckpointLedger, JsonlFileSink, Observer, RunOptions, + SkipSet, +}; +use std::collections::HashSet; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; +use std::time::Instant; +use tokio_util::sync::CancellationToken; + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn completion_order_across_fleet() { + let dir = tempfile::tempdir().unwrap(); + let script = common::write_plan( + dir.path(), + "delay.py", + r#" +import time + +def plan(): + for i in range(4): + yield {"id": f"t-{i}", "delay": 0.3 if i == 0 else 0.01} + +def execute(item): + time.sleep(float(item["delay"])) + return {"x": item["id"]} +"#, + ); + let order = Arc::new(Mutex::new(Vec::new())); + let order_cb = Arc::clone(&order); + let mut opts = common::run_opts(script); + opts.workers = 2; + opts.max_inflight = 4; + run_local_fleet(opts, move |r| { + order_cb.lock().unwrap().push(r.task_id.clone()); + }) + .await + .unwrap(); + let seen = order.lock().unwrap().clone(); + assert_ne!( + seen[0], "t-0", + "expected fast task before slow t-0: {seen:?}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn resume_skip_ids_not_dispatched() { + let dir = tempfile::tempdir().unwrap(); + let script = common::write_plan( + dir.path(), + "skip.py", + r#" +def plan(): + for i in range(3): + yield {"id": f"t-{i}", "x": i} + +def execute(item): + return {"x": item["x"]} +"#, + ); + let mut opts = common::run_opts(script); + opts.workers = 1; + opts.skip_task_ids = SkipSet::from_iter(["t-0".into(), "t-1".into()]); + let results = run_local_fleet(opts, |_| {}).await.unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].task_id, "t-2"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn ctrl_c_cancels_in_flight_execute() { + let dir = tempfile::tempdir().unwrap(); + let script = common::write_plan( + dir.path(), + "slow.py", + r#" +import time + +def plan(): + yield {"id": "slow"} + +def execute(item): + time.sleep(5) + return {} +"#, + ); + let cancel = CancellationToken::new(); + let c = cancel.clone(); + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + c.cancel(); + }); + let opts = RunOptions { + script, + python: PathBuf::from("python3"), + workers: 1, + max_inflight: 1, + per_worker_inflight: 1, + pythonpath_extra: vec![], + script_args: vec![], + infra_retries: 0, + job_cancel: cancel, + observer: Observer::disabled(), + skip_task_ids: SkipSet::new(), + checkpoint: None, + sink_submitter: None, + }; + let t0 = Instant::now(); + let results = run_local_fleet(opts, |_| {}).await.unwrap(); + assert!(results.iter().any(|r| r.cancelled)); + assert!(t0.elapsed().as_secs_f64() < 2.5); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn per_worker_slots_run_in_parallel() { + let dir = tempfile::tempdir().unwrap(); + let script = common::write_plan( + dir.path(), + "par.py", + r#" +import time + +def plan(): + for i in range(2): + yield {"id": f"t-{i}"} + +def execute(item): + time.sleep(0.45) + return {} +"#, + ); + let mut opts = common::run_opts(script); + opts.workers = 1; + opts.per_worker_inflight = 2; + opts.max_inflight = 2; + let started = Instant::now(); + let results = run_local_fleet(opts, |_| {}).await.unwrap(); + let elapsed = started.elapsed().as_secs_f64(); + assert_eq!(results.len(), 2); + assert!(results.iter().all(|r| r.ok)); + let workers: HashSet<_> = results.iter().filter_map(|r| r.worker.clone()).collect(); + assert_eq!(workers.len(), 2, "expected two slot ids, got {workers:?}"); + assert!(elapsed < 0.9, "took {elapsed:.2}s (looks serial)"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn script_args_reach_plan_and_execute() { + let dir = tempfile::tempdir().unwrap(); + let script = common::write_plan( + dir.path(), + "args.py", + r#" +import argparse + +def _parse(argv=None): + p = argparse.ArgumentParser() + p.add_argument("--n", type=int, default=1) + return p.parse_args(argv) + +def plan(): + args = _parse() + for i in range(args.n): + yield {"id": f"t-{i}", "n": args.n} + +def execute(item): + return {"n": item["n"]} +"#, + ); + let mut opts = common::run_opts(script); + opts.workers = 1; + opts.script_args = vec!["--n".into(), "3".into()]; + let results = run_local_fleet(opts, |_| {}).await.unwrap(); + assert_eq!(results.len(), 3); + assert_eq!(results[0].value.as_ref().unwrap()["n"], 3); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn duplicate_plan_ids_dispatch_once() { + let dir = tempfile::tempdir().unwrap(); + let script = common::write_plan( + dir.path(), + "dup.py", + r#" +def plan(): + yield {"id": "same", "x": 1} + yield {"id": "same", "x": 2} + +def execute(item): + return {"x": item["x"]} +"#, + ); + let mut opts = common::run_opts(script); + opts.workers = 2; + let results = run_local_fleet(opts, |_| {}).await.unwrap(); + assert_eq!( + results.len(), + 1, + "second identical id must be claimed/skipped" + ); + assert_eq!(results[0].task_id, "same"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn bounded_inflight_handles_many_tasks() { + let dir = tempfile::tempdir().unwrap(); + let script = common::write_plan( + dir.path(), + "many.py", + r#" +def plan(): + for i in range(40): + yield {"id": f"t-{i}", "x": i} + +def execute(item): + return {"x": item["x"] * 2} +"#, + ); + let mut opts = common::run_opts(script); + opts.workers = 2; + opts.max_inflight = 2; + opts.per_worker_inflight = 1; + let results = run_local_fleet(opts, |_| {}).await.unwrap(); + assert_eq!(results.len(), 40); + assert!(results.iter().all(|r| r.ok)); +} + +/// `--sink` persist then `--resume` from ledger, including corrupt JSONL lines. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn sink_then_resume_skips_done_and_tolerates_bad_lines() { + let dir = tempfile::tempdir().unwrap(); + let sink_root = dir.path().join("sink"); + let script = common::write_plan( + dir.path(), + "resume.py", + r#" +def plan(): + for i in range(3): + yield {"id": f"t-{i}", "x": i} + +def execute(item): + return {"x": item["x"]} +"#, + ); + + // Run 1: all three tasks → ready.ndjson + { + let sink: Arc = + Arc::new(JsonlFileSink::open(&sink_root).await.unwrap()); + let writer = spawn_sink_writer(sink, None, None, 32); + let mut opts = common::run_opts(script.clone()); + opts.workers = 1; + opts.sink_submitter = Some(writer.submitter()); + let results = run_local_fleet(opts, |_| {}).await.unwrap(); + writer.join().await.unwrap(); + assert_eq!(results.len(), 3); + assert!(results.iter().all(|r| r.ok)); + } + + // Corrupt a line; valid terminals must still load for resume. + { + let ready_path = sink_root.join("ready.ndjson"); + let mut body = std::fs::read_to_string(&ready_path).unwrap(); + body.push_str("this-is-not-json\n"); + body.push_str("{\"no_id\":true}\n"); + std::fs::write(&ready_path, body).unwrap(); + } + + let ledger = CheckpointLedger::load(&sink_root).await.unwrap(); + assert_eq!( + ledger.ready.len(), + 3, + "corrupt lines must not drop valid ids" + ); + assert!(ledger.skip_ids().contains("t-0")); + assert!(ledger.skip_ids().contains("t-1")); + assert!(ledger.skip_ids().contains("t-2")); + + // Run 2: resume — nothing left to dispatch. + { + let sink: Arc = + Arc::new(JsonlFileSink::open(&sink_root).await.unwrap()); + let writer = spawn_sink_writer(sink, None, None, 32); + let mut opts = common::run_opts(script.clone()); + opts.workers = 1; + opts.skip_task_ids = SkipSet::from_iter(ledger.skip_ids()); + opts.sink_submitter = Some(writer.submitter()); + let results = run_local_fleet(opts, |_| {}).await.unwrap(); + writer.join().await.unwrap(); + assert!( + results.is_empty(), + "resume must skip all ledger terminals, got {results:?}" + ); + } + + // Partial resume: only t-0 done on disk → run t-1/t-2. + let partial = dir.path().join("sink_partial"); + std::fs::create_dir_all(&partial).unwrap(); + std::fs::write( + partial.join("ready.ndjson"), + "{\"task_id\":\"t-0\",\"ok\":true}\nnot-json\n", + ) + .unwrap(); + let partial_ledger = CheckpointLedger::load(&partial).await.unwrap(); + assert_eq!(partial_ledger.ready, HashSet::from(["t-0".into()])); + + let sink: Arc = + Arc::new(JsonlFileSink::open(&partial).await.unwrap()); + let writer = spawn_sink_writer(sink, None, None, 32); + let mut opts = common::run_opts(script); + opts.workers = 1; + opts.skip_task_ids = SkipSet::from_iter(partial_ledger.skip_ids()); + opts.sink_submitter = Some(writer.submitter()); + let results = run_local_fleet(opts, |_| {}).await.unwrap(); + writer.join().await.unwrap(); + let ids: HashSet<_> = results.iter().map(|r| r.task_id.clone()).collect(); + assert_eq!(ids, HashSet::from(["t-1".into(), "t-2".into()])); + assert!(results.iter().all(|r| r.ok)); + + let after = CheckpointLedger::load(&partial).await.unwrap(); + assert_eq!(after.ready.len(), 3); + assert!(after.skip_ids().contains("t-0")); + assert!(after.skip_ids().contains("t-1")); + assert!(after.skip_ids().contains("t-2")); +} diff --git a/crates/persisting-compute/tests/integration_pulsing_cluster.rs b/crates/persisting-compute/tests/integration_pulsing_cluster.rs new file mode 100644 index 0000000..305e9d7 --- /dev/null +++ b/crates/persisting-compute/tests/integration_pulsing_cluster.rs @@ -0,0 +1,160 @@ +//! Integration: two in-process Pulsing ActorSystems (seed + peer). +//! +//! Covers named resolve, job-control cancel, and **cross-node Execute** +//! without requiring torchrun. + +mod common; + +use persisting_compute::dist::DistEnv; +use persisting_compute::job_control::{ + broadcast_job_cancel, JobControlActor, JobControlCommand, JobControlReply, +}; +use persisting_compute::pulsing_ext::{ask_timeout, resolve_actor, ASK_TIMEOUT}; +use persisting_compute::{TaskExpr, WorkerActor, WorkerCommand, WorkerReply}; +use pulsing_actor::prelude::*; +use serde_json::json; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; +use tokio_util::sync::CancellationToken; + +async fn join_peer(seed_addr: &str) -> Arc { + let mut last = None; + for _ in 1..=40 { + match ActorSystem::builder() + .mailbox_capacity(64) + .addr("127.0.0.1:0") + .seeds([seed_addr]) + .build() + .await + { + Ok(s) => return s, + Err(e) => { + last = Some(e.to_string()); + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + } + panic!("peer join failed: {last:?}"); +} + +async fn resolve_eventually(system: &ActorSystem, name: &str) -> ActorRef { + for _ in 0..80 { + if let Ok(r) = resolve_actor(system, name).await { + return r; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + panic!("timed out resolving {name}"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn two_systems_resolve_and_cancel() { + let seed: Arc = ActorSystem::builder() + .mailbox_capacity(64) + .addr("127.0.0.1:0") + .build() + .await + .expect("seed system"); + let seed_addr = seed.addr().to_string(); + + let t0 = CancellationToken::new(); + seed.spawn_named( + &DistEnv::job_control_name(0), + JobControlActor::new(t0.clone()), + ) + .await + .unwrap(); + + let peer = join_peer(&seed_addr).await; + + let t1 = CancellationToken::new(); + peer.spawn_named( + &DistEnv::job_control_name(1), + JobControlActor::new(t1.clone()), + ) + .await + .unwrap(); + + let remote_ctrl = resolve_eventually(seed.as_ref(), &DistEnv::job_control_name(1)).await; + + let ack = + ask_timeout::<_, JobControlReply>(&remote_ctrl, JobControlCommand::Cancel, ASK_TIMEOUT) + .await + .expect("ask cancel on remote"); + assert!(matches!(ack, JobControlReply::Ack { already: false })); + assert!(t1.is_cancelled()); + + broadcast_job_cancel(&seed, 2).await; + assert!(t0.is_cancelled()); + + peer.shutdown().await.unwrap(); + seed.shutdown().await.unwrap(); +} + +/// Cross-node Execute: peer hosts a WorkerActor; seed resolves + ask. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn two_systems_remote_execute() { + let dir = tempfile::tempdir().unwrap(); + let script = common::write_plan( + dir.path(), + "remote_exec.py", + r#" +def plan(): + yield {"id": "unused"} + +def execute(item): + return {"echo": item.get("x", 0), "from": "peer"} +"#, + ); + + let seed: Arc = ActorSystem::builder() + .mailbox_capacity(64) + .addr("127.0.0.1:0") + .build() + .await + .expect("seed system"); + let seed_addr = seed.addr().to_string(); + + let peer = join_peer(&seed_addr).await; + let job_cancel = CancellationToken::new(); + let worker_name = DistEnv::slot_name(1, 0, 1); + peer.spawn_named( + &worker_name, + WorkerActor::with_plan( + "w1", + PathBuf::from("python3"), + vec![], + script, + vec![], + job_cancel, + ), + ) + .await + .unwrap(); + + let remote = resolve_eventually(seed.as_ref(), &worker_name).await; + let task = TaskExpr::from_value(json!({"id": "t-remote", "x": 42})).unwrap(); + let task_json = serde_json::to_vec(&task).unwrap(); + let reply = ask_timeout::<_, WorkerReply>( + &remote, + WorkerCommand::Execute { task_json }, + Duration::from_secs(30), + ) + .await + .expect("remote Execute"); + + let WorkerReply::Result { result_json } = reply else { + panic!("unexpected reply: {reply:?}"); + }; + let result: persisting_compute::TaskResult = serde_json::from_slice(&result_json).unwrap(); + assert!(result.ok, "remote execute failed: {:?}", result.error); + assert_eq!(result.task_id, "t-remote"); + assert_eq!( + result.value.as_ref().and_then(|v| v.get("echo")), + Some(&json!(42)) + ); + + peer.shutdown().await.unwrap(); + seed.shutdown().await.unwrap(); +} diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 365fa74..5a209f1 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -97,6 +97,8 @@ plugins: Traj Capture Subcommands: Traj Capture 子命令 Capture Quick Start: Capture 快速上手 Capture Architecture Design: Capture 架构设计 + Compute Quick Start: Compute 快速上手 + Compute Architecture: Compute 架构 Search Command: Search 命令 Similar Systems Reference: 类似系统参考 API Reference: API 参考 @@ -117,6 +119,7 @@ nav: - Guide: - guide/index.md - Capture Quick Start: guide/capture_quickstart.md + - Compute Quick Start: guide/compute_quickstart.zh.md - Queue Backends: guide/backends.md - Lance Backend: guide/lance.md - Persisting Backend: guide/persisting.md @@ -130,6 +133,7 @@ nav: - CLI Architecture: design/cli_architecture.zh.md - Traj Capture Subcommands: design/cli_capture_command.zh.md - Capture Architecture Design: design/capture_design.zh.md + - Compute Architecture: design/compute_control_plane.zh.md - Traj Command: design/cli_trajectory_command.zh.md - Search Command: design/cli_search_command.zh.md - Similar Systems Reference: design/similar_systems_reference.md diff --git a/docs/src/design/cli_architecture.zh.md b/docs/src/design/cli_architecture.zh.md index b224ccf..c385307 100644 --- a/docs/src/design/cli_architecture.zh.md +++ b/docs/src/design/cli_architecture.zh.md @@ -4,6 +4,8 @@ `search`、`traj`（`trajectory`）等子命令共用此架构。 +`compute` 例外：编排逻辑在 `persisting-compute` 内由 CLI 直接调用（不经 engine RON ABI）；见 [Compute 架构](compute_control_plane.zh.md) · [快速上手](../guide/compute_quickstart.zh.md)。 + --- ## 1. 核心思想 diff --git a/docs/src/design/compute_control_plane.zh.md b/docs/src/design/compute_control_plane.zh.md new file mode 100644 index 0000000..239b801 --- /dev/null +++ b/docs/src/design/compute_control_plane.zh.md @@ -0,0 +1,236 @@ +# Compute 架构 + +> 状态：**Phase-1 已落地** +> 代码：`crates/persisting-compute` · CLI：`persisting compute` +> 用法：[Compute 快速上手](../guide/compute_quickstart.zh.md) · 示例：`examples/compute/` + +薄编排层：算法写一个 Python 文件（`plan` + `execute`），控制面负责流式派发、并行执行、落盘与续跑。 +**不是** Ray，**不**嵌用户解释器，**不**自研替代 `torchrun` 的启动器。 + +--- + +## 1. 定位 + +```text +用户脚本 plan() 产出任务 · execute(item) 做计算 + ↓ +L2 Compute 流式领取 · 有界并发 · sticky 派发 · 取消 / 重试 + ↓ +L3 执行缝 每槽一个长驻 Python host，只调用户 execute + ↓ +唯一 sink 控制面 append 结果（JSONL 账本；可选 Tee Vortex） +``` + +| 是 | 不是 | +|----|------| +| 独立任务的 map 式编排 | 分布式训练框架 | +| 调用用户 `execute(item)` | 新 DSL / 装饰器框架 | +| 真实 `torchrun` 多进程 | 自研 rdzv / launcher | +| Rust Pulsing 做发现与投递 | Python 侧 Actor 编程模型 | +| spawn 外部 `--python` | PyO3 内嵌用户环境 | + +与 `search` / `traj` 不同：`compute` **静态链接**进 CLI，直接 async 跑，不经 engine RON ABI。落盘若走 Vortex，再经 sink adapter 写轨迹；编排本身仍在本 crate。 + +--- + +## 2. 系统全景 + +```text + plan.py + （spawn python · bootstrap） + │ NDJSON 任务流 + ▼ + Driver（rank0 / 本地） + plan emit + 有界 inflight + least-loaded / sticky ask + ┌────────────┼────────────┐ + ▼ ▼ ▼ + WorkerActor 槽位（workers × per-worker） + │ 每槽：独立 Python host + └──── stdin/stdout JSON ────┘ + │ + ▼ + plan.py::execute(item) + │ + ▼ + ResultSink（--sink JSONL · 可选 --traj） +``` + +要点： + +- **Driver 不是 Pulsing Actor**：串行领取 plan、并行 `ask` worker，用 `FuturesUnordered` 有界 drain。 +- **Pulsing** 负责命名解析、mailbox、Supervision、本地 DeathWatch；**不做**业务调度。 +- **任务粒度**由用户决定：适合分片 / 文件级；不适合把每一行都打成控制面任务。 + +--- + +## 3. 用户合约（算法面） + +产品面只有一个文件、两个函数（细节与示例见[快速上手](../guide/compute_quickstart.zh.md)）： + +1. **`plan()`**（或常量 `PLAN`）流式产出可 JSON 序列化的 object。 +2. **`execute(item)`** 收到与 yield **同形**的 dict（`{id, …fields}`）。 +3. **argv 一致**：`python task.py --n 2` 与 `persisting compute task.py -- --n 2` 看到同一套 `sys.argv`。 +4. **扩规模不改文件**：只换 `-w` / `--per-worker` / `torchrun`。 + +内部控制面会把平面 JSON 归一成 `TaskExpr`（`id` / `op=execute` / `args` / `meta`）；产品面 **只支持** `op=execute`。新能力写在用户 `execute` 里，不靠扩展 op 表。 + +结果线格式为 `TaskResult`：`task_id`、`ok`、`cancelled`、`value` / `error`、`worker`、时间戳、`infra_retries` 等。 + +--- + +## 4. 运行时拓扑 + +### 4.1 本地 `-w N` + +单进程：`N × per_worker` 个槽位 Actor；本进程即 Driver。 + +### 4.2 torchrun + +| Rank | 角色 | +|------|------| +| 0 | Driver（plan + dispatch）+ 本机槽位 | +| >0 | 本机槽位，等 `Shutdown` | + +- 读 `RANK` / `WORLD_SIZE` / `MASTER_ADDR` / `MASTER_PORT`（及可选 `LOCAL_RANK`）。 +- Pulsing 种子端口默认 `MASTER_PORT+17`，可用 `PERSISTING_PULSING_PORT` 覆盖，避免与 c10d 冲突。 +- 控制面**不**自己拉齐进程；由 torchrun 负责。 + +### 4.3 槽位命名 + +池按 **slot-major** 展平：先各 worker 的 slot0，再 slot1… +扁平下标：`slot * n_workers + worker`（DeathWatch / quarantine 必须用这套下标，不能用本 rank 局部序数）。 + +--- + +## 5. 调度与执行 + +### 5.1 Driver 循环 + +1. 从 `plan()` 流式读任务（NDJSON）。 +2. 全局 `max_inflight`（默认 `workers × per_worker`）限制在途量。 +3. 派发前 **SkipSet claim** `task_id`（resume 种子 + 本 job 已派发）；同 id 不二次派发。 +4. 首触：**least-loaded** 选槽；对该槽发起过 `Execute` 之后 → **sticky-only**（只打同一槽）。 +5. 完成即 `on_result`，并 `await` 异步 sink enqueue（队列满则背压 Driver）。 + +### 5.2 sticky-only 与 quarantine + +| 情况 | 行为 | +|------|------| +| 尚未接触任何槽 | least-loaded 选槽 | +| 已 ask 过某槽 | infra 重试 **只**打该槽 | +| 该槽被 quarantine | **拒绝**改投他槽，任务记 infra 失败（避免跨槽 at-least-once 重跑 `execute`） | +| 全槽 quarantine | acquire fail-fast，不再死等 | + +同槽 **ResultCache**（`task_id → TaskResult`）跨 Supervision 重启共享：丢 reply 时可幂等取回，不重跑 `execute`。 + +`--retries` 只覆盖 **worker ask / 基础设施**失败；**不**解读业务 `ok=false`（那是 L3 语义重试，未做）。 + +### 5.3 L3 执行缝 + +每槽一个长驻 Python 子进程（行协议 JSON）： + +| cmd | 行为 | +|-----|------| +| `run_plan` | 加载脚本（带 argv），调 `execute(item)`；按路径缓存模块 | +| `shutdown` | 退出 host | + +Worker **不**解释业务；Driver **不**碰 Python。失败 traceback 编进 `TaskResult`。 + +### 5.4 取消 + +- Ctrl-C → job `CancellationToken`：停接新任务；在途 acquire / ask 可取消。 +- 在途 execute：**kill Python host**，返回 `cancelled`（下一任务会重新拉起 host）。 +- 跨 rank：每 rank 独立 `JobControlActor`（与 Worker **分邮箱**），rank0 广播 `Cancel`，避免被串行 `Execute` 堵住。 +- 已成功的结果**不会**被改写成 `cancelled`。 + +--- + +## 6. 耐久、幂等与续跑 + +### 6.1 唯一 sink + +只有控制面写就绪结果；Executor **禁止**直写存储账本。 + +| 路径 | 作用 | +|------|------| +| `--sink DIR` | `ready.ndjson` + `failures.ndjson` + `checkpoint.json` | +| 无 `--sink` | 默认仅 stdout NDJSON（开发视图，非耐久） | +| `--traj` | Tee 到 Vortex；**JSONL 仍是 resume 真相** | + +`JsonlFileSink`：打开时从已有文件 **seed** `task_id`；先 reserve 再写盘，失败则 **rollback seen**（seen ⊆ durable）。 +异步 `sink_writer`：`join` 汇总 persist 失败并 fail job；失败时 `SkipSet::remove`，便于后续 `--resume` 发现未落盘 id。损坏 / 无 `task_id` 的 JSONL 行：**skip + warn**，不拖垮整本账本。 + +### 6.2 Checkpoint / `--resume` + +- 账本 = sink 下已终端的 `task_id`（ready ∪ failures）。 +- `--resume`：Driver 跳过这些 id；**`plan()` 仍会再 emit 一遍**（大 plan 的成本是已知限制）。 +- 失败 / 取消默认也不重跑；要重跑请编辑 failures 或换目录。 + +### 6.3 两类重试（边界） + +| 层 | 管什么 | Phase-1 | +|----|--------|---------| +| L2 | 计算单元在不在（ask / 节点） | `--retries` + sticky + quarantine | +| L3 | 要不要再产一次（业务） | 未做；留给用户 `execute` | + +--- + +## 7. 模块地图 + +| 模块 | 职责 | +|------|------| +| `cli` | 参数 → fleet / sink / Ctrl-C | +| `runtime` | local / torchrun 装配（共享 spawn / driver loop） | +| `driver` | plan 流 + 有界 drain + skip + sticky 派发 | +| `scheduler` | least-loaded · sticky · quarantine | +| `worker` | `Execute` / `Shutdown` · ResultCache · Supervision | +| `executor` / `plan` | Python host · plan NDJSON bootstrap | +| `dist` | torchrun 环境 · slot 命名 / 扁平下标 | +| `job_control` | 旁路 cancel · 本地 DeathWatch | +| `sink` / `sink_writer` / `checkpoint` / `skip` | 唯一落盘 · 异步 persist · resume 账本 · live claim | +| `observe` | 可选进度事件 | +| `task` | `TaskExpr` / `TaskResult` 线格式 | + +多数模块 `pub(crate)`；公开面主要服务 CLI 与集成测。 + +语义原语 id 见 `blocks.rs`。合约测与源码同文件；跨模块行为在 `tests/integration_*.rs`。 + +--- + +## 8. 非目标与下一阶段 + +### 刻意不做（保持薄） + +- 多 op 产品面、独立 Python 编排包、内嵌用户解释器 +- 自研进程启动器替代 torchrun +- 在 L2 堆 harness / 判分 / Agent 会话 +- Meta 语义重试、模板 fan-out、配额、亲和调度 + +### 已知差距 + +| 项 | 说明 | +|----|------| +| 昂贵 plan 无 cursor | resume 仍全量再 emit | +| 真 torchrun CI | 有双 ActorSystem 烟测；真实多进程 e2e 仍薄 | +| 错误分类 | infra / execute / cancel 尚未成稳定枚举字段 | +| DeathWatch | 仅本地；远端死槽靠 ask 失败路径 | +| 对外命名 | 内部仍 `TaskExpr`；文档侧逐步对齐 TaskSpec | + +### 设计原则（摘要） + +1. 算法入口极轻：两个函数；本地 `for` 循环即真。 +2. 扩规模换命令，不改文件。 +3. Driver ≠ Actor；Pulsing ≠ 调度器。 +4. 跨槽拒绝优先于跨槽重跑。 +5. persist 失败必须可见；JSONL 是 resume 真相。 +6. 新产品能力优先进 `execute`，不膨胀 CLI。 + +--- + +## 9. 相关文档 + +- [Compute 快速上手](../guide/compute_quickstart.zh.md) +- [CLI 整体架构](cli_architecture.zh.md)（`compute` 为例外静态路径） +- [轨迹存储](trajectory_storage.zh.md)（`--traj` / L1） +- 示例：[`examples/compute/`](../../../examples/compute/) diff --git a/docs/src/design/index.md b/docs/src/design/index.md index 42732d9..71b750a 100644 --- a/docs/src/design/index.md +++ b/docs/src/design/index.md @@ -30,6 +30,7 @@ See [index.zh.md](index.zh.md) for the full Chinese index with document links an | Trajectory storage model | [trajectory_storage.zh.md](trajectory_storage.zh.md) (中文) | | Trajectory Markdown format | [trajectory_tlv_format.zh.md](trajectory_tlv_format.zh.md) | | Capture architecture | [capture_design.zh.md](capture_design.zh.md) (中文) | +| Compute architecture (Phase-1) | [compute_control_plane.zh.md](compute_control_plane.zh.md) (中文) · [Quick start](../guide/compute_quickstart.zh.md) | | CLI architecture | [cli_architecture.zh.md](cli_architecture.zh.md) | | Queue persistence | [architecture.md](architecture.md) | | TTAS addressing | [tensor_address_algebra.md](tensor_address_algebra.md) | diff --git a/docs/src/design/index.zh.md b/docs/src/design/index.zh.md index 7dac0fc..18b16f1 100644 --- a/docs/src/design/index.zh.md +++ b/docs/src/design/index.zh.md @@ -42,6 +42,7 @@ Persisting 的设计文档：寻址与分层存储（演进中）、队列持久 | 轨迹存储模型 | [trajectory_storage.zh.md](trajectory_storage.zh.md) | 已实现 | | 轨迹 Markdown 格式 | [trajectory_tlv_format.zh.md](trajectory_tlv_format.zh.md) | 已实现 | | Capture 架构设计（对外） | [capture_design.zh.md](capture_design.zh.md) | 已实现 | +| Compute 架构 | [compute_control_plane.zh.md](compute_control_plane.zh.md) | Phase-1 已落地 | | 队列持久化架构 | [architecture.zh.md](architecture.zh.md) | 已实现 | | TTAS 寻址模型 | [tensor_address_algebra.md](tensor_address_algebra.md) | 设计 | | 分布式分层存储 | [distributed_tiered_storage.md](distributed_tiered_storage.md) | 规划 | @@ -53,6 +54,7 @@ Persisting 的设计文档：寻址与分层存储（演进中）、队列持久 | 整体架构 | [cli_architecture.zh.md](cli_architecture.zh.md) | 瘦 CLI + 动态引擎、异步任务、版本化协议 | | **`traj`（轨迹）** | [cli_trajectory_command.zh.md](cli_trajectory_command.zh.md) | 统一入口：capture / proxy / import / stats / replay / materialize / … | | Capture 子命令 | [cli_capture_command.zh.md](cli_capture_command.zh.md) | `traj capture`、`traj proxy`、`traj import` 设计说明 | +| **`compute`** | [compute_control_plane.zh.md](compute_control_plane.zh.md) · [快速上手](../guide/compute_quickstart.zh.md) | `plan()` + `execute` 薄编排 / torchrun（Phase-1） | | Search | [cli_search_command.zh.md](cli_search_command.zh.md) | 导入、索引、检索 | ## 参考与分析 diff --git a/docs/src/guide/compute_quickstart.zh.md b/docs/src/guide/compute_quickstart.zh.md new file mode 100644 index 0000000..d82a624 --- /dev/null +++ b/docs/src/guide/compute_quickstart.zh.md @@ -0,0 +1,168 @@ +# Compute 快速上手 + +用 **`persisting compute`** 跑一批独立任务：写一个 Python 文件，本地能跑，换命令即可加并行。 + +> 架构说明见 [Compute 架构](../design/compute_control_plane.zh.md)。示例：[`examples/compute/`](../../examples/compute/)。 + +--- + +## 你需要准备什么 + +- 已编译的 `persisting` CLI（含 `persisting-compute`） +- 本机可用的 `python3`（或 `--python` 指定解释器） + +```bash +cargo build -p persisting-cli +export PERSISTING="$(pwd)/target/debug/persisting" +"$PERSISTING" compute --help +``` + +--- + +## 1. 写一个文件 + +只要两个函数：`plan()` 产出任务，`execute(item)` 处理一条。 + +```python +import argparse + +def _parse_args(argv=None): + p = argparse.ArgumentParser() + p.add_argument("-n", "--n", type=int, default=10) + return p.parse_args(argv) + +def plan(): + args = _parse_args() + for i in range(args.n): + yield {"id": f"t-{i}", "x": i} + +def execute(item): + return {"x2": item["x"] * 2} + +if __name__ == "__main__": + for xx in plan(): + print(execute(xx)) +``` + +约定很简单： + +- 每条任务是一个带 **`id`** 的 dict（缺 `id` 时控制面会生成 UUID） +- `execute` 收到的就是 `plan()` yield 出来的那条 +- 业务参数写在 `--` **后面**，和直接 `python task.py …` 一样进 `sys.argv` + +把参数解析放在 `plan()` / `execute()` **里面**（不要只在模块 import 时解析一次）。 + +--- + +## 2. 先本地验证 + +```bash +python3 task.py --n 2 +persisting compute task.py --check -- --n 2 +``` + +`--check`：检查环境、能否 emit、能否调用 `execute`，并小规模试跑。 + +--- + +## 3. 并行跑起来 + +```bash +# 本机 4 个 worker +persisting compute task.py -w 4 -- --n 100 + +# 单 worker、每 worker 4 个槽（更适合单机多核、任务偏 CPU） +persisting compute task.py -w 1 --per-worker 4 -- --n 100 + +# 多进程（真实 torchrun；同一文件不用改） +torchrun --nproc_per_node=8 -- persisting compute task.py -- --n 100 +``` + +`--` 前面是 compute 自己的开关；`--` 后面原样转给脚本。 + +| 常用开关 | 作用 | +|----------|------| +| `-w N` | 本地 worker 数（非 torchrun 时） | +| `--per-worker N` | 每个 worker/rank 的并发槽（默认 1） | +| `--python PATH` | Python 解释器（也可用环境变量 `PERSISTING_PYTHON`） | +| `-E PATH` | 追加到 `PYTHONPATH` | +| `--retries N` | 基础设施重试次数（默认 2；不是业务失败重试） | +| `--results ndjson\|summary\|quiet` | 终端怎么打印结果（默认 `ndjson`） | + +--- + +## 4. 落盘与中断后续跑 + +```bash +persisting compute task.py -w 4 --sink /tmp/run1 -- --n 1000 + +# 中断后接着跑：跳过 ready / failures 里已经有的 id +persisting compute task.py -w 4 --sink /tmp/run1 --resume -- --n 1000 +``` + +`--sink DIR` 下会看到： + +| 文件 | 内容 | +|------|------| +| `ready.ndjson` | 成功结果（一行一条） | +| `failures.ndjson` | 失败 / 取消 | +| `checkpoint.json` | 进度摘要（节流更新） | + +说明： + +- `--resume` **必须**同时带 `--sink` +- 已出现在账本里的 `id` 不会再跑（成功和失败都不重跑；要重跑请改 failures 或换目录） +- resume 时 `plan()` 仍会再扫一遍，只是已完成的 id 不会再派发 + +看进度还可以加： + +```bash +persisting compute task.py -w 4 --sink /tmp/run1 --resume --observe -- --n 1000 +``` + +`--observe` 必须写在 `--` **前面**；开启后默认少刷结果 NDJSON（需要两者都要时再加 `--results ndjson`）。 + +--- + +## 5. 可选：写入轨迹 + +若 CLI 编译了 `traj-sink` 相关能力： + +```bash +persisting compute task.py -w 2 --sink /tmp/run1 --traj -- --n 4 +``` + +JSONL 仍是 resume 账本；Vortex 轨迹是额外一份。读回可用 `persisting traj stats`（目录默认在 `{sink}/traj`）。 + +--- + +## 6. 内置自测 + +```bash +persisting compute --self-test +``` + +不依赖你的脚本，用来确认本机 CLI / Python 通路正常。 + +--- + +## 常见问题 + +**任务没跑 / id 重复？** +同一 `id` 在一次 job 里只会派发一次。resume 时账本里已有的 id 也会被跳过。 + +**Ctrl-C 之后？** +未开始的任务会标为取消；正在跑的 execute 会被打断。已成功写入 sink 的不会被改成取消。 + +**业务失败会不会自动重试？** +不会。`--retries` 只覆盖「连不上 worker」这类基础设施问题。业务上要重试，请在 `execute` 里自己做，或清掉 failures 后再跑。 + +**和直接写 Ray / 自己 multiprocessing 比？** +你只维护 `plan` + `execute`；并行、取消、落盘、续跑由 CLI 管。不适合把每一行数据都当成一个任务——任务粒度用文件 / 分片更合适。 + +--- + +## 下一步 + +- 示例脚本：[`examples/compute/plan_simple.py`](../../examples/compute/plan_simple.py) +- 运行时与调度：[Compute 架构](../design/compute_control_plane.zh.md) diff --git a/docs/src/guide/index.md b/docs/src/guide/index.md index 483c92c..1ba7e1c 100644 --- a/docs/src/guide/index.md +++ b/docs/src/guide/index.md @@ -44,6 +44,13 @@ Proxy and record LLM traffic under **`persisting traj`**: - [Capture architecture](../design/capture_design.zh.md) (中文) - **Standalone dlcapt (`persisting-dlcapt`)** — independent OpenAI-compatible proxy migrated from Capture `external/dlcapt`; see `crates/persisting-dlcapt/README.md`. Phase-1 does not replace `persisting-capture`. +### Compute (available) + +Run independent tasks with **`persisting compute`**: one Python file (`plan` + `execute`), local parallelism or `torchrun`. + +- **[Compute Quick Start](compute_quickstart.zh.md)** (中文) — script, parallelism, sink / resume +- [Compute architecture](../design/compute_control_plane.zh.md) (中文) + ### Streaming Append (available now) Lance storage engine's append-only access pattern — for event streaming and durable queues (trajectory capture uses Vortex separately; see Capture quick start). diff --git a/docs/src/guide/index.zh.md b/docs/src/guide/index.zh.md index 3af15cd..d882b64 100644 --- a/docs/src/guide/index.zh.md +++ b/docs/src/guide/index.zh.md @@ -44,6 +44,13 @@ arr = kv["s1", 0, 2, 0:512].tensor() - [Capture 架构设计](../design/capture_design.zh.md) — 概念与数据流 - **独立 dlcapt（`persisting-dlcapt`）** — 自 Capture `external/dlcapt` 迁入的独立 OpenAI 兼容代理；见 `crates/persisting-dlcapt/README.md`。第一阶段不替代 `persisting-capture`。 +### Compute 任务编排（已可用） + +用 **`persisting compute`** 跑一批独立任务：一个 Python 文件里的 `plan()` + `execute`，本地并行或 `torchrun` 扩规模。 + +- **[Compute 快速上手](compute_quickstart.zh.md)** — 写脚本、并行、落盘与续跑 +- [Compute 架构](../design/compute_control_plane.zh.md) — Driver / Worker / sink / sticky + ### 流式追加（已可用） 流式追加与队列持久化基于 Lance 存储引擎——与轨迹 Vortex 层独立。这是当前已可用的生产能力。 diff --git a/examples/compute/README.md b/examples/compute/README.md new file mode 100644 index 0000000..b10a506 --- /dev/null +++ b/examples/compute/README.md @@ -0,0 +1,13 @@ +# Persisting compute + +一个 Python 文件：`plan()` 产出任务，`execute(item)` 处理一条。 + +```bash +python3 examples/compute/plan_simple.py --n 2 +persisting compute examples/compute/plan_simple.py -w 4 -- --n 2 +persisting compute examples/compute/plan_simple.py -w 4 --sink /tmp/run1 --resume -- --n 100 +torchrun --nproc_per_node=8 -- persisting compute examples/compute/plan_simple.py -- --n 2 +``` + +完整用法：[Compute 快速上手](../../docs/src/guide/compute_quickstart.zh.md) +架构：[Compute 架构](../../docs/src/design/compute_control_plane.zh.md) diff --git a/examples/compute/plan_simple.py b/examples/compute/plan_simple.py new file mode 100644 index 0000000..8c9737f --- /dev/null +++ b/examples/compute/plan_simple.py @@ -0,0 +1,37 @@ +"""One file: plan() + execute(item), with optional argparse. + +Local:: + + python3 examples/compute/plan_simple.py + python3 examples/compute/plan_simple.py --n 2 + +Scale (same argv after ``--``):: + + persisting compute examples/compute/plan_simple.py -w 4 -- --n 2 +""" + +from __future__ import annotations + +import argparse + + +def _parse_args(argv=None): + p = argparse.ArgumentParser(description="demo plan") + p.add_argument("-n", "--n", type=int, default=4, help="number of tasks") + return p.parse_args(argv) + + +def plan(): + args = _parse_args() + for i in range(args.n): + yield {"id": f"t-{i}", "x": i} + + +def execute(item): + x = item["x"] + return {"x": x, "x2": x * 2} + + +if __name__ == "__main__": + for xx in plan(): + print(execute(xx))