From a9288e840f1400f31698812f6422226575f26cf9 Mon Sep 17 00:00:00 2001 From: Reiase Date: Fri, 24 Jul 2026 00:37:15 +0800 Subject: [PATCH 1/2] Add persisting-compute crate and integrate with CLI - Introduced the `persisting-compute` crate with dependencies and initial structure. - Updated `Cargo.toml` and `Cargo.lock` to include `persisting-compute`. - Enhanced the CLI to support a new `Compute` command for executing compute plans. - Implemented local validation and checkpointing features for compute tasks. - Added proxy configuration arguments for better CLI integration. This commit lays the groundwork for a compute control plane, enabling task execution and management within the existing framework. --- Cargo.lock | 25 + Cargo.toml | 1 + crates/persisting-cli/Cargo.toml | 1 + .../src/capture/proxy_config_args.rs | 165 ++++++ crates/persisting-cli/src/main.rs | 34 +- crates/persisting-compute/Cargo.toml | 35 ++ crates/persisting-compute/src/blocks.rs | 110 ++++ crates/persisting-compute/src/check.rs | 344 ++++++++++++ crates/persisting-compute/src/checkpoint.rs | 288 ++++++++++ crates/persisting-compute/src/cli.rs | 389 +++++++++++++ crates/persisting-compute/src/dist.rs | 220 ++++++++ crates/persisting-compute/src/driver.rs | 354 ++++++++++++ crates/persisting-compute/src/executor.rs | 424 +++++++++++++++ crates/persisting-compute/src/future.rs | 106 ++++ crates/persisting-compute/src/job_control.rs | 374 +++++++++++++ crates/persisting-compute/src/lib.rs | 65 +++ crates/persisting-compute/src/observe.rs | 513 ++++++++++++++++++ crates/persisting-compute/src/plan.rs | 168 ++++++ crates/persisting-compute/src/pulsing_ext.rs | 55 ++ crates/persisting-compute/src/python_env.rs | 91 ++++ crates/persisting-compute/src/result_cache.rs | 99 ++++ crates/persisting-compute/src/runtime.rs | 451 +++++++++++++++ crates/persisting-compute/src/scheduler.rs | 458 ++++++++++++++++ crates/persisting-compute/src/sink.rs | 291 ++++++++++ crates/persisting-compute/src/sink_traj.rs | 138 +++++ crates/persisting-compute/src/sink_writer.rs | 199 +++++++ crates/persisting-compute/src/skip.rs | 80 +++ crates/persisting-compute/src/task.rs | 242 +++++++++ crates/persisting-compute/src/worker.rs | 327 +++++++++++ crates/persisting-compute/tests/common/mod.rs | 31 ++ .../tests/integration_local.rs | 330 +++++++++++ .../tests/integration_pulsing_cluster.rs | 160 ++++++ docs/mkdocs.yml | 2 + docs/src/design/cli_architecture.zh.md | 2 + docs/src/design/compute_control_plane.zh.md | 389 +++++++++++++ docs/src/design/index.md | 1 + docs/src/design/index.zh.md | 2 + examples/compute/README.md | 52 ++ examples/compute/plan_simple.py | 37 ++ 39 files changed, 7043 insertions(+), 10 deletions(-) create mode 100644 crates/persisting-cli/src/capture/proxy_config_args.rs create mode 100644 crates/persisting-compute/Cargo.toml create mode 100644 crates/persisting-compute/src/blocks.rs create mode 100644 crates/persisting-compute/src/check.rs create mode 100644 crates/persisting-compute/src/checkpoint.rs create mode 100644 crates/persisting-compute/src/cli.rs create mode 100644 crates/persisting-compute/src/dist.rs create mode 100644 crates/persisting-compute/src/driver.rs create mode 100644 crates/persisting-compute/src/executor.rs create mode 100644 crates/persisting-compute/src/future.rs create mode 100644 crates/persisting-compute/src/job_control.rs create mode 100644 crates/persisting-compute/src/lib.rs create mode 100644 crates/persisting-compute/src/observe.rs create mode 100644 crates/persisting-compute/src/plan.rs create mode 100644 crates/persisting-compute/src/pulsing_ext.rs create mode 100644 crates/persisting-compute/src/python_env.rs create mode 100644 crates/persisting-compute/src/result_cache.rs create mode 100644 crates/persisting-compute/src/runtime.rs create mode 100644 crates/persisting-compute/src/scheduler.rs create mode 100644 crates/persisting-compute/src/sink.rs create mode 100644 crates/persisting-compute/src/sink_traj.rs create mode 100644 crates/persisting-compute/src/sink_writer.rs create mode 100644 crates/persisting-compute/src/skip.rs create mode 100644 crates/persisting-compute/src/task.rs create mode 100644 crates/persisting-compute/src/worker.rs create mode 100644 crates/persisting-compute/tests/common/mod.rs create mode 100644 crates/persisting-compute/tests/integration_local.rs create mode 100644 crates/persisting-compute/tests/integration_pulsing_cluster.rs create mode 100644 docs/src/design/compute_control_plane.zh.md create mode 100644 examples/compute/README.md create mode 100644 examples/compute/plan_simple.py 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..24ab65d 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -97,6 +97,7 @@ plugins: Traj Capture Subcommands: Traj Capture 子命令 Capture Quick Start: Capture 快速上手 Capture Architecture Design: Capture 架构设计 + Compute Control Plane: Compute 控制面 Search Command: Search 命令 Similar Systems Reference: 类似系统参考 API Reference: API 参考 @@ -130,6 +131,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 Control Plane: 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..6ae739b 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)。 + --- ## 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..8cfbea2 --- /dev/null +++ b/docs/src/design/compute_control_plane.zh.md @@ -0,0 +1,389 @@ +# Compute 控制面设计 + +> 状态：**Phase-1 已落地**（可运行竖切） +> 代码：`crates/persisting-compute` · CLI：`persisting compute` +> 示例：`examples/compute/plan_simple.py` +> 对齐主张：Agentic Infra（L2 编排 · L3 执行缝 · 唯一 sink） + +薄编排层：算法工程师写一个 Python 文件，本地能跑，换启动命令即可扩规模。 +**不是** Ray，**不**嵌用户解释器，**不** DIY `torchrun`。 + +--- + +## 1. 定位与边界 + +### 1.1 在 Agentic Infra 中的位置 + +```text +L3 Workload 复杂负载 → 干净、可调度的一次运行（本阶段 = plan.py::execute） +L2 Compute 把众多 Future 编成吞吐（本 crate） +L1 Storage 就绪 Result 落盘为资产（Phase-1 = JSONL；目标 = Vortex 轨迹） +``` + +三层是**模块边界**；部署上仍是**单一控制面入口**（`persisting compute`），不是三套集群产品。 + +### 1.2 是 / 不是 + +| 是 | 不是 | +|----|------| +| L2 编排：plan 流式 → Driver 分发 → Worker 执行 | 分布式训练框架 | +| L3 执行缝：调用用户 `execute(item)` | 新编程模型 / DSL / 装饰器框架 | +| 用真实 `torchrun` 做多进程 | 自研 launcher / rdzv | +| 用 Rust Pulsing actors 做骨架 | Python Pulsing / Ray 替代品 | +| 控制面 **spawn** `--python` | PyO3 内嵌用户环境 | + +### 1.3 与 CLI 引擎架构的关系 + +`search` / `traj` 走「瘦 CLI + 动态加载 engine」。 +**`compute` 不同**：逻辑在 `persisting-compute` 库内、由 CLI **静态链接并直接 async 跑**，不经 RON/engine ABI。落盘目标接上 Vortex 后，再经 engine 的 `TrajectoryAppend` 写 L1（sink adapter），编排本身仍留在本 crate。 + +--- + +## 2. 用户合约（算法面） + +产品面只有一个文件、两个函数： + +```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() # 推荐：在 plan()/execute 内解析，勿依赖模块 import 副作用 + 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)) +``` + +### 2.1 约定 + +1. **`plan()`**（或常量 `PLAN`）产出任务；每项可 JSON 序列化为 object。 +2. **`execute(item)`** 收到的对象与 `plan()` yield **同形**（`{id, …fields}`），不是内部控制面线格式。 +3. **argv 一致**：`python task.py --n 2` 与 `persisting compute task.py -- --n 2` 看到同一套 `sys.argv`（`--` 之后原样转发）。 +4. **扩规模不改文件**：只换启动命令（本地 `-w` / `torchrun`）。 + +### 2.2 argv 注意 + +- Plan 进程（emit）与 Execute host（worker）都会设置 `sys.argv = [script, *user_args]`。 +- Execute host **按脚本路径缓存**已 import 的模块：若在**模块顶层**解析 argparse，同路径不会因换参数而重载。 + **推荐**在 `plan()` / `execute()` 内调用 `parse_args()`。 + +--- + +## 3. CLI + +```bash +# 本地扩规模 +persisting compute SCRIPT [-w N] [--per-worker N] [--python PY] [-E PATH] \ + [--retries N] [--sink DIR] [--results ndjson|summary|quiet] -- + +# 校验（不替代完整跑） +persisting compute SCRIPT --check [--limit N] -- + +# 内置 smoke（须显式） +persisting compute --self-test + +# 多进程（真实 torchrun） +torchrun --nproc_per_node=K -- persisting compute SCRIPT -- +``` + +| 开关 | 含义 | +|------|------| +| （默认） | **run**：拉起 fleet，跑完退出 | +| `-w N` | 本地 worker 数（非 torchrun 时） | +| `--per-worker N` | 每 worker/rank 的并发槽数（每槽独立 Actor + Python host） | +| `--observe` / `--observe-file` / `--observe-json` | 可观测：stderr `[obs]` 进度；file/json 为 NDJSON。须写在 `--` 前；开启后默认 `--results quiet` | +| `--retries N` | L2 基础设施重试次数（默认 2；仅 worker ask 失败） | +| `--sink DIR` | 唯一落盘目录：`ready.ndjson` + `failures.ndjson` + `checkpoint.json` | +| `--resume` | 从 `--sink` 账本跳过已完成 `task_id`（须同时给 `--sink`） | +| `--traj` | Tee 写 Vortex 轨迹（`compute.result`）；须同时 `--sink` | +| `--results` | 终端展示：`ndjson`（默认）/ `summary` / `quiet` | +| Ctrl-C | Job cancel：未启动任务标 `cancelled`（见 §7.1 限制） | +| `--check` | env → plan emit → `execute` 存在 → 本地试跑 | +| 省略 SCRIPT | **报错**（不再默认跑 smoke） | + +唯一入口：`persisting compute`。无独立 `persisting-compute` 二进制，无 `check/run/emit/worker` 子命令族。 + +--- + +## 4. 数据面 + +### 4.1 任务（内部名 `TaskExpr`） + +线格式：一行一个 JSON（NDJSON）。用户侧推荐平面写法（与 yield 一致）： + +```json +{"id": "t-0", "x": 1} +``` + +控制面归一化： + +```json +{"id": "t-0", "op": "execute", "args": {"x": 1}, "meta": {}} +``` + +| 字段 | 说明 | +|------|------| +| `id` | 任务身份；缺省则生成 UUID | +| `op` | 缺省 `"execute"`；**产品面仅支持 `execute`** | +| `args` | 业务字段；平面 JSON 的剩余键落入此 | +| `meta` | 预留；当前调度/执行不消费 | + +新能力写在用户 `execute` 内，**不**靠扩展 op 插件表。 + +> 命名对照（Agentic Infra）：用户侧 ≈ TaskSpec 载荷；内部仍叫 `TaskExpr`，对外文档逐步对齐。 + +### 4.2 结果（`TaskResult`） + +| 字段 | 说明 | +|------|------| +| `task_id` | 对应任务 id | +| `ok` | 是否成功 | +| `cancelled` | L2 取消（非业务失败） | +| `value` / `error` / `traceback` | 成功值或失败信息 | +| `worker` | 执行 worker id | +| `started_at` / `finished_at` | Unix 秒 | +| `infra_retries` | L2 基础设施重试次数（0 = 首次成功） | + +--- + +## 5. 运行时拓扑 + +```text + plan.py + （spawn --python -c bootstrap） + │ NDJSON TaskExpr + ▼ + Driver（rank0 / 本地控制进程） + │ plan emit + least-loaded dispatch + │（直接 ask Worker，并行） + ┌────────────┼────────────┐ + ▼ ▼ ▼ + slot actors (N workers × P per-worker) + │ each: WorkerActor + Python host + └──── stdin/stdout JSON ────┘ + │ + ▼ + plan.py::execute(item) + │ + ▼ + 控制面 ResultSink（stdout 与/或 --sink） +``` + +每个从 `plan()` 出来的任务在 Driver 调度层对应一个 **`RunFuture`**（`wait` / `cancel` / `task_id`）。 + +### 5.1 本地 `-w N` + +单进程：N 个 `WorkerActor`；本进程即 **Driver**（`Driver::run_plan`），无独立 DriverActor。 + +### 5.2 torchrun + +| Rank | 角色 | +|------|------| +| 0 | **Driver**（plan + dispatch）+ 本机 `per_worker` 个槽位 Actor | +| >0 | 本机 `per_worker` 个槽位 Actor，等 `Shutdown` | + +- 读环境：`RANK` / `WORLD_SIZE` / `MASTER_ADDR` / `MASTER_PORT`（及可选 `LOCAL_RANK`）。 +- Pulsing 种子：`MASTER_ADDR:(MASTER_PORT+17)`，可用 `PERSISTING_PULSING_PORT` 覆盖，避免与 c10d 撞端口。 +- 控制面**不**自己 spawn 进程；由 torchrun 拉齐 ranks。 + +### 5.3 调度参数 + +- **`--per-worker N`**（默认 1）：每个逻辑 worker / torchrun rank 上的**并发 Execute 槽**。 + - 实现：每槽一个 `WorkerActor` + 独立 Python host（Actor mailbox 仍串行，槽与槽之间并行）。 + - 池按 **slot-major** 展平：先铺满各 worker 的 slot0，再 slot1…，least-loaded 仍优先打散到不同 worker。 +- **Least-loaded**：在展平后的槽位上选当前 in-flight 最少者；平手取更小下标。 +- **`max_inflight`**：全局上限（默认 `workers × per-worker`）。 +- 任务 **直接 `ask` 槽位 Actor**（不经 Driver 串行 `await Execute`）。 +- **完成即回调**：`on_result` / sink 按完成序；outstanding ≤ `max_inflight`。 +- Infra retry：失败则释放槽位再 acquire，避免粘在坏节点上。 +- Job 级 `CancellationToken`：Ctrl-C → 停止接新任务 + 子 token 取消在途 acquire。 + +--- + +## 6. L3 执行缝 + +每个 **槽位** 持有一个长驻 `--python` 子进程（行协议 JSON）： + +| cmd | 行为 | +|-----|------| +| `run_plan` | 加载 plan 模块（带 `argv`），调用 `execute(item)`；按路径缓存模块 | +| `shutdown` | 退出 host | + +- Worker **不**解释业务语义；Driver **不**碰 Python。 +- 失败时 traceback 编入 `TaskResult` 回传。 +- 产品面 Executor 路由：**仅** `op=execute` → `PlanExecuteExecutor`。 + +--- + +## 7. Phase-1 合约实现 + +相对 Agentic Infra PPT，本阶段落地三条主契约（算法面保持极轻）。 + +### 7.1 RunFuture（调度原子） + +- 实现：`future.rs`；进程内 `JoinHandle` + `CancellationToken`。 +- **已支持**：`wait`（始终 join）、协作式 `cancel`（未启动 / 等待 acquire）。 +- **在途 cancel**：Worker 共享 job `CancellationToken`；取消时 **kill Python host** 并返回 `cancelled`（下一任务会重新拉起 host）。 +- **跨 rank**：每 rank 有独立 `JobControlActor`（与 Worker **分邮箱**）；rank0 Ctrl-C → 广播 `Cancel` → 远端 token cancel → 在途 host kill。不经过 Worker mailbox，避免被串行 `Execute` 挡住。 +- **取消语义**：`wait` **不**把已成功结果改写成 `cancelled`；取消来迟则保留成功。 +- Driver 在 `ask` 上同样 select cancel，避免 Ctrl-C 后干等 Python。 + +### 7.2 L2 基础设施重试 + +- Driver 在 **worker ask 失败**（节点丢、传输错）时重投；次数 `--retries`。 +- **Sticky**：失败后优先同一槽位，便于命中该 Worker 的 `task_id → TaskResult` 缓存（同 worker 上丢失 reply 可幂等重取，不重跑 `execute`）。 +- **不**解读 `execute` 返回的业务失败（那是 L3 Meta 语义重试，未做）。 +- Infra 耗尽 → 该任务 `TaskResult` 失败（`ok=false`），**不** `bail` 整次 job；sibling 继续跑。 +- `TaskResult.infra_retries` 可观测。 +- **Sink 去重**：`JsonlFileSink` 打开时从已有 JSONL **seed** `task_id`；跨进程重复 append 跳过。Vortex Tee 同样可 seed。 +- **Seen ⊆ durable**：先 reserve `task_id`，写盘/Vortex 失败则 **rollback seen**，避免永久丢账。 +- **Live SkipSet**：dispatch 前 claim `task_id`（resume 种子 + 本 job 已派发）；同 id 不再打到另一 worker。`--check --limit` 用 SkipSet 跳过尾部任务。 +- **Async sink writer**：完成路径 `await` enqueue（队列满则自然背压 Driver drain）；`join` 汇总 persist 失败并 fail job。persist 失败会 `SkipSet::remove`，便于后续 `--resume` 发现未落盘 id。 +- **Tee 尽力而为**：某一路失败不阻断兄弟 sink（仍返回首个错误）。 +- **Sticky-only after contact**：一旦对该槽发起过 `Execute` ask，infra 重试**只**打同一槽（`acquire_sticky`）；若该槽被 quarantine，**拒绝**改投他槽（避免跨槽 at-least-once 重跑 `execute`），任务记 infra 失败。 +- **槽位 quarantine**：连续 infra ask 失败达到阈值后摘槽；**全槽 quarantine → acquire fail-fast**（任务 infra 失败，非整 job 死等）。DeathWatch（仅本地）在 Worker 终止时 `force_quarantine`。Watch 注册必须用 **slot-major 扁平下标**（`DistEnv::slot_flat_index`），不能用本 rank 的 `0..per_worker` 序数。 +- **Pulsing 深化**：Worker `spawn_factory` + `SupervisionSpec`（失败重启上限）；**ResultCache 与槽位同生命周期**（`Arc` 跨重启共享）；统一 `resolve_named` / `ask_timeout`；JobControl 侧信道 cancel + 本地 DeathWatch。 +- **仍须注意**：跨槽拒绝后该任务失败（非整 job bail）；算法若需更强保证应自备幂等或缩小副作用面。 + +### 7.3 唯一 sink（控制面落盘） + +- 只有控制面写就绪结果；Executor **禁止**直写存储。 +- `--sink DIR` → `ready.ndjson`（成功）+ `failures.ndjson`（失败/取消）；**逐条**在 `on_result` 时 append（非跑完批写）。 +- Driver **按完成序**调用 `on_result`（非 plan 提交序）；有界 outstanding（≤ `max_inflight`，`FuturesUnordered` + `RunFuture`），完成即释放，避免大 plan 囤积 JoinHandle。 +- 同 process 内按 `task_id` 跳过重复 append。 +- **`--traj`**（feature `traj-sink`）：Tee 到 Vortex，`CaptureRecord{kind=compute.result|compute.failure}` → `TrajectoryAppend`；默认 `{sink}/traj` / agent=`compute` / session=sink 目录名。JSONL 仍是 `--resume` 账本。 +- 未指定 `--sink` 时：默认仅 stdout NDJSON（开发视图，非耐久 L1）。 +- **限制（已知）**：Vortex 侧跨进程未按 `task_id` 去重；resume 跳过已完成 id 可避免重复写。 + +### 7.4 Checkpoint（断点续跑） + +- 账本 = `--sink` 下已有 `ready.ndjson` / `failures.ndjson` 的 `task_id`。 +- `--resume`（须带 `--sink`）：Driver 跳过已终端 id，不重跑；`plan()` 仍会再 emit 一遍。 +- 同目录 `checkpoint.json`：`ok/fail/cancelled/skipped/dispatched/updated_at`（节流写入）。 +- 失败/取消默认也不重跑；要重跑请编辑 failures 或换 sink 目录。 + +### 7.5 与全图的差距（诚实清单） + +| 主张 | Phase-1 | +|------|---------| +| TaskSpec / RunFuture / RunResult 对外语言 | 内部名仍 `TaskExpr` / 进程内 Future / `TaskResult` | +| 就绪 → L1 轨迹 | JSONL 占位 | +| 可中断 | 调度层 + kill Python host；跨 rank 经 JobControlActor 广播 | +| 可复现 / 位置无关 | 未钉死 | +| Meta 语义重试 | 未做 | +| 模板 fan-out / 配额 / 亲和 / 观测账本 | least-loaded + per-worker；无模板/亲和 | + +--- + +## 8. 模块地图与语义原语 + +产品按**语义原语 / 接口**切开：每块有稳定 id、主模块、同文件合约测。 +注册表见 [`blocks`](../../../crates/persisting-compute/src/blocks.rs)。 + +**测试策略** + +- **单元 / 合约测**：放在原语所在源文件末尾的 `#[cfg(test)]`（与接口同文件）。 +- **`tests/`**：只放跨模块**集成**（如 `tests/integration_local.rs`：fleet 完成序、resume skip、cancel、`--per-worker`、argv）。 + +| Block id | 合约 | 主模块 | 测试 | +|----------|------|--------|------| +| `task_wire` | TaskExpr / TaskResult | `task` | 同文件 | +| `python_env` | PYTHONPATH 合并 | `python_env` | 同文件 | +| `placement` | least-loaded · sticky · slot-major 命名 | `scheduler` · `dist` | 同文件 | +| `run_future` | wait / cancel 语义 | `future` | 同文件 | +| `idempotency` | 同 worker 结果缓存 | `result_cache` · `worker` | 同文件 | +| `sink` / `checkpoint` | 唯一 sink · resume skip · seed 去重 | `sink` · `checkpoint` | 同文件 | +| `observe` | 可选进度事件 | `observe` | 同文件 | +| `plan` / `execute` | plan emit · execute host · cancel kill | `plan` · `executor` | 同文件 | +| `worker` | Actor 缝合 · cache · 多 slot ask | `worker` | 同文件 | +| `driver` / `fleet` | 完成序 drain · skip · job cancel · 装配 | `driver` · `runtime` | `tests/integration_*.rs` | + +### 8.1 文件职责 + +| 路径 | 职责 | +|------|------| +| `persisting-cli` → `Command::Compute` | 唯一 CLI 入口 | +| `cli.rs` | 参数 → `run_fleet` / sink / Ctrl-C | +| `blocks.rs` | 语义原语 id 注册表（文档对齐；非独立测套） | +| `driver.rs` | **Driver**：plan emit + least-loaded dispatch；完成序 `on_result` + 有界 inflight | +| `checkpoint.rs` | `--resume` 账本 + `checkpoint.json` 进度 | +| `future.rs` | `RunFuture` · `wait_all` | +| `result_cache.rs` | 同 worker infra 幂等缓存 | +| `scheduler.rs` | Least-loaded · sticky prefer | +| `observe.rs` | `--observe` 事件 | +| `sink.rs` | `ResultSink` · `JsonlFileSink` · `TeeSink` | +| `sink_traj.rs` | （feature `traj-sink`）`VortexResultSink` | +| `plan.rs` | bootstrap：`sys.argv` + `plan()` → NDJSON | +| `runtime.rs` | local / torchrun fleet 装配（`spawn_rank_slots` / `run_driver_loop` 共享） | +| `worker.rs` | `Execute` / `Shutdown` + cache + cancel token | +| `executor.rs` | PlanExecute host（可取消 kill） | +| `dist.rs` | torchrun 环境 → Pulsing seed · slot 命名 | +| `job_control.rs` | 旁路 cancel + 本地 DeathWatch | +| `pulsing_ext.rs` | 统一 resolve / ask_timeout / spawn_supervised | +| `skip.rs` / `sink_writer.rs` | live claim · 异步 persist | +| `check.rs` | `--check` / `--self-test` | +| `task.rs` | `TaskExpr` / `TaskResult` | +| `python_env.rs` | `PYTHONPATH` 合并 | + +crate 公开面收窄：多数模块 `pub(crate)`；`lib` 只 re-export CLI / 集成测所需类型（`run_compute`、`RunOptions`、`SkipSet`、sink 主类型等）。 + +--- + +## 9. 校验与自测 + +| 模式 | 行为 | +|------|------| +| `cargo test -p persisting-compute` | 同文件单元 / 合约测 + `tests/integration_*.rs` | +| `--check` | python 探针 → plan emit → 确认 `execute` 可调用 → `run_local_fleet` 试跑 | +| `--self-test` | 写临时内置 `plan()`+`execute`，走完整 check（不依赖用户文件） | + +--- + +## 10. 非目标与路线 + +### 10.1 刻意不做（保持薄） + +- 多 op 产品面（`echo` / `call` / `py` 等已移除） +- 独立 Python `persisting.compute` 编排包 +- 控制面内嵌用户解释器 +- 自研进程启动器替代 torchrun +- 在 L2 堆 harness / 判分 / Agent 会话 + +### 10.2 下一阶段优先 + +1. **跨 worker 幂等**（已：sticky-only-after-contact + 同槽 result_cache + sink seed/rollback + 跨 rank cancel；quarantine 后该任务 infra 失败而非改投） +2. **`ResultSink` → Vortex**（真 L1 默认路径） +3. 命名对外对齐（TaskSpec / RunFuture）+ 更多红灯 +4. （可选）Meta Executor 薄包装 + +--- + +## 11. 设计原则 + +1. **算法入口极轻**：两个函数 + 可选 argparse；本地 `for` 循环即真。 +2. **扩规模换命令**：同一文件，`persisting compute` / `torchrun`。 +3. **argv 一致**：`--` 后与 `python task.py …` 对齐。 +4. **能力成本沉在骨架**：Pulsing / torchrun 负责分布；用户不学 actors。 +5. **边界守住**：新产品能力优先进 `execute`，不膨胀 CLI 与 op 表。 +6. **两类重试分开**：L2 管计算单元在不在；L3 Meta 管要不要再产一次。 +7. **唯一 sink**：就绪 Result 只从控制面 append。 +8. **竖切优先**：先打通可跑路径与合约形状，再补工厂能力（模板 / 配额 / 亲和）。 + +--- + +## 12. 相关文档 + +- Agentic Infra 主张：`deeplink.fabric/projects/agent.infra/report-agent-infra.typ` +- 轨迹 L1（Vortex）：[trajectory_storage.zh.md](trajectory_storage.zh.md) +- CLI 整体（engine 路径）：[cli_architecture.zh.md](cli_architecture.zh.md) +- 用法示例：[examples/compute/README.md](../../../examples/compute/README.md) diff --git a/docs/src/design/index.md b/docs/src/design/index.md index 42732d9..7d0450c 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 control plane (Phase-1) | [compute_control_plane.zh.md](compute_control_plane.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..00af573 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) | `plan()` + `execute` 薄编排 / torchrun（Phase-1） | | Search | [cli_search_command.zh.md](cli_search_command.zh.md) | 导入、索引、检索 | ## 参考与分析 diff --git a/examples/compute/README.md b/examples/compute/README.md new file mode 100644 index 0000000..50ecad8 --- /dev/null +++ b/examples/compute/README.md @@ -0,0 +1,52 @@ +# Persisting compute + +一个文件： + +```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)) +``` + +```bash +python3 task.py --n 2 +persisting compute task.py -w 4 --per-worker 1 -- --n 2 +# --per-worker 1（默认）：每 worker 一槽；耗时差大时优先加 -w +# --per-worker N：每 worker/rank N 个槽（各带独立 Python host，真并行） +persisting compute task.py -w 1 --per-worker 4 -- --n 2 +persisting compute task.py -w 4 --retries 2 --sink /tmp/run1 -- --n 2 +# 中断后续跑（跳过 ready/failures 里已有 task_id）；看进度：/tmp/run1/checkpoint.json +persisting compute task.py -w 4 --sink /tmp/run1 --resume --observe -- --n 1000 +# 同时写入 Vortex 轨迹（JSONL 仍用于 resume） +persisting compute task.py -w 2 --sink /tmp/run1 --traj -- --n 4 +# 读回：persisting traj stats /tmp/run1/traj --agent-id compute --session-id run1 +# 可观测：stderr 上 `[obs]` 行（默认同时关掉结果 NDJSON 刷屏） +# 注意：`--observe` 必须写在 `--` 前面 +persisting compute task.py -w 2 --observe -- --n 2 +persisting compute task.py -w 2 --observe --observe-file /tmp/obs.ndjson -- --n 2 +persisting compute task.py -w 2 --observe --results ndjson -- --n 2 # 结果 + obs 都要 +torchrun --nproc_per_node=8 -- persisting compute task.py -- --n 2 +persisting compute task.py --check -- --n 2 +``` + +设计说明见 [docs/src/design/compute_control_plane.zh.md](../../docs/src/design/compute_control_plane.zh.md)。 + +- `execute` 收到的就是 `plan()` yield 出来的 dict +- `--` 后参数进入 `sys.argv` +- `--sink` 由控制面唯一落盘 `ready.ndjson` / `failures.ndjson` +- Ctrl-C 取消未启动的 RunFuture（Phase-1） 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)) From 772e0adea47472b6005c48065ac345c358a05358 Mon Sep 17 00:00:00 2001 From: Reiase Date: Fri, 24 Jul 2026 11:00:13 +0800 Subject: [PATCH 2/2] Update documentation and structure for Compute features - Renamed "Compute Control Plane" to "Compute Architecture" for clarity. - Added a new "Compute Quick Start" guide to assist users in running independent tasks. - Updated navigation links in documentation to reflect the new structure and added quick start references. - Enhanced existing design documents to include links to the new quick start guide and clarified architecture descriptions. These changes improve the usability and accessibility of the Compute features within the documentation. --- docs/mkdocs.yml | 6 +- docs/src/design/cli_architecture.zh.md | 2 +- docs/src/design/compute_control_plane.zh.md | 453 +++++++------------- docs/src/design/index.md | 2 +- docs/src/design/index.zh.md | 4 +- docs/src/guide/compute_quickstart.zh.md | 168 ++++++++ docs/src/guide/index.md | 7 + docs/src/guide/index.zh.md | 7 + examples/compute/README.md | 53 +-- 9 files changed, 347 insertions(+), 355 deletions(-) create mode 100644 docs/src/guide/compute_quickstart.zh.md diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 24ab65d..5a209f1 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -97,7 +97,8 @@ plugins: Traj Capture Subcommands: Traj Capture 子命令 Capture Quick Start: Capture 快速上手 Capture Architecture Design: Capture 架构设计 - Compute Control Plane: Compute 控制面 + Compute Quick Start: Compute 快速上手 + Compute Architecture: Compute 架构 Search Command: Search 命令 Similar Systems Reference: 类似系统参考 API Reference: API 参考 @@ -118,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 @@ -131,7 +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 Control Plane: design/compute_control_plane.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 6ae739b..c385307 100644 --- a/docs/src/design/cli_architecture.zh.md +++ b/docs/src/design/cli_architecture.zh.md @@ -4,7 +4,7 @@ `search`、`traj`（`trajectory`）等子命令共用此架构。 -`compute` 例外：编排逻辑在 `persisting-compute` 内由 CLI 直接调用（不经 engine RON ABI）；设计见 [Compute 控制面](compute_control_plane.zh.md)。 +`compute` 例外：编排逻辑在 `persisting-compute` 内由 CLI 直接调用（不经 engine RON ABI）；见 [Compute 架构](compute_control_plane.zh.md) · [快速上手](../guide/compute_quickstart.zh.md)。 --- diff --git a/docs/src/design/compute_control_plane.zh.md b/docs/src/design/compute_control_plane.zh.md index 8cfbea2..239b801 100644 --- a/docs/src/design/compute_control_plane.zh.md +++ b/docs/src/design/compute_control_plane.zh.md @@ -1,389 +1,236 @@ -# Compute 控制面设计 +# Compute 架构 -> 状态：**Phase-1 已落地**（可运行竖切） +> 状态：**Phase-1 已落地** > 代码：`crates/persisting-compute` · CLI：`persisting compute` -> 示例：`examples/compute/plan_simple.py` -> 对齐主张：Agentic Infra（L2 编排 · L3 执行缝 · 唯一 sink） +> 用法：[Compute 快速上手](../guide/compute_quickstart.zh.md) · 示例：`examples/compute/` -薄编排层：算法工程师写一个 Python 文件，本地能跑，换启动命令即可扩规模。 -**不是** Ray，**不**嵌用户解释器，**不** DIY `torchrun`。 +薄编排层：算法写一个 Python 文件（`plan` + `execute`），控制面负责流式派发、并行执行、落盘与续跑。 +**不是** Ray，**不**嵌用户解释器，**不**自研替代 `torchrun` 的启动器。 --- -## 1. 定位与边界 - -### 1.1 在 Agentic Infra 中的位置 +## 1. 定位 ```text -L3 Workload 复杂负载 → 干净、可调度的一次运行（本阶段 = plan.py::execute） -L2 Compute 把众多 Future 编成吞吐（本 crate） -L1 Storage 就绪 Result 落盘为资产（Phase-1 = JSONL；目标 = Vortex 轨迹） +用户脚本 plan() 产出任务 · execute(item) 做计算 + ↓ +L2 Compute 流式领取 · 有界并发 · sticky 派发 · 取消 / 重试 + ↓ +L3 执行缝 每槽一个长驻 Python host，只调用户 execute + ↓ +唯一 sink 控制面 append 结果（JSONL 账本；可选 Tee Vortex） ``` -三层是**模块边界**；部署上仍是**单一控制面入口**（`persisting compute`），不是三套集群产品。 - -### 1.2 是 / 不是 - | 是 | 不是 | |----|------| -| L2 编排：plan 流式 → Driver 分发 → Worker 执行 | 分布式训练框架 | -| L3 执行缝：调用用户 `execute(item)` | 新编程模型 / DSL / 装饰器框架 | -| 用真实 `torchrun` 做多进程 | 自研 launcher / rdzv | -| 用 Rust Pulsing actors 做骨架 | Python Pulsing / Ray 替代品 | -| 控制面 **spawn** `--python` | PyO3 内嵌用户环境 | - -### 1.3 与 CLI 引擎架构的关系 +| 独立任务的 map 式编排 | 分布式训练框架 | +| 调用用户 `execute(item)` | 新 DSL / 装饰器框架 | +| 真实 `torchrun` 多进程 | 自研 rdzv / launcher | +| Rust Pulsing 做发现与投递 | Python 侧 Actor 编程模型 | +| spawn 外部 `--python` | PyO3 内嵌用户环境 | -`search` / `traj` 走「瘦 CLI + 动态加载 engine」。 -**`compute` 不同**：逻辑在 `persisting-compute` 库内、由 CLI **静态链接并直接 async 跑**，不经 RON/engine ABI。落盘目标接上 Vortex 后，再经 engine 的 `TrajectoryAppend` 写 L1（sink adapter），编排本身仍留在本 crate。 +与 `search` / `traj` 不同：`compute` **静态链接**进 CLI，直接 async 跑，不经 engine RON ABI。落盘若走 Vortex，再经 sink adapter 写轨迹；编排本身仍在本 crate。 --- -## 2. 用户合约（算法面） - -产品面只有一个文件、两个函数： - -```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() # 推荐：在 plan()/execute 内解析，勿依赖模块 import 副作用 - for i in range(args.n): - yield {"id": f"t-{i}", "x": i} +## 2. 系统全景 -def execute(item): - return {"x2": item["x"] * 2} - -if __name__ == "__main__": - for xx in plan(): - print(execute(xx)) +```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） ``` -### 2.1 约定 - -1. **`plan()`**（或常量 `PLAN`）产出任务；每项可 JSON 序列化为 object。 -2. **`execute(item)`** 收到的对象与 `plan()` yield **同形**（`{id, …fields}`），不是内部控制面线格式。 -3. **argv 一致**：`python task.py --n 2` 与 `persisting compute task.py -- --n 2` 看到同一套 `sys.argv`（`--` 之后原样转发）。 -4. **扩规模不改文件**：只换启动命令（本地 `-w` / `torchrun`）。 - -### 2.2 argv 注意 +要点： -- Plan 进程（emit）与 Execute host（worker）都会设置 `sys.argv = [script, *user_args]`。 -- Execute host **按脚本路径缓存**已 import 的模块：若在**模块顶层**解析 argparse，同路径不会因换参数而重载。 - **推荐**在 `plan()` / `execute()` 内调用 `parse_args()`。 +- **Driver 不是 Pulsing Actor**：串行领取 plan、并行 `ask` worker，用 `FuturesUnordered` 有界 drain。 +- **Pulsing** 负责命名解析、mailbox、Supervision、本地 DeathWatch；**不做**业务调度。 +- **任务粒度**由用户决定：适合分片 / 文件级；不适合把每一行都打成控制面任务。 --- -## 3. CLI +## 3. 用户合约（算法面） -```bash -# 本地扩规模 -persisting compute SCRIPT [-w N] [--per-worker N] [--python PY] [-E PATH] \ - [--retries N] [--sink DIR] [--results ndjson|summary|quiet] -- +产品面只有一个文件、两个函数（细节与示例见[快速上手](../guide/compute_quickstart.zh.md)）： -# 校验（不替代完整跑） -persisting compute SCRIPT --check [--limit N] -- +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`。 -# 内置 smoke（须显式） -persisting compute --self-test +内部控制面会把平面 JSON 归一成 `TaskExpr`（`id` / `op=execute` / `args` / `meta`）；产品面 **只支持** `op=execute`。新能力写在用户 `execute` 里，不靠扩展 op 表。 -# 多进程（真实 torchrun） -torchrun --nproc_per_node=K -- persisting compute SCRIPT -- -``` - -| 开关 | 含义 | -|------|------| -| （默认） | **run**：拉起 fleet，跑完退出 | -| `-w N` | 本地 worker 数（非 torchrun 时） | -| `--per-worker N` | 每 worker/rank 的并发槽数（每槽独立 Actor + Python host） | -| `--observe` / `--observe-file` / `--observe-json` | 可观测：stderr `[obs]` 进度；file/json 为 NDJSON。须写在 `--` 前；开启后默认 `--results quiet` | -| `--retries N` | L2 基础设施重试次数（默认 2；仅 worker ask 失败） | -| `--sink DIR` | 唯一落盘目录：`ready.ndjson` + `failures.ndjson` + `checkpoint.json` | -| `--resume` | 从 `--sink` 账本跳过已完成 `task_id`（须同时给 `--sink`） | -| `--traj` | Tee 写 Vortex 轨迹（`compute.result`）；须同时 `--sink` | -| `--results` | 终端展示：`ndjson`（默认）/ `summary` / `quiet` | -| Ctrl-C | Job cancel：未启动任务标 `cancelled`（见 §7.1 限制） | -| `--check` | env → plan emit → `execute` 存在 → 本地试跑 | -| 省略 SCRIPT | **报错**（不再默认跑 smoke） | - -唯一入口：`persisting compute`。无独立 `persisting-compute` 二进制，无 `check/run/emit/worker` 子命令族。 +结果线格式为 `TaskResult`：`task_id`、`ok`、`cancelled`、`value` / `error`、`worker`、时间戳、`infra_retries` 等。 --- -## 4. 数据面 +## 4. 运行时拓扑 -### 4.1 任务（内部名 `TaskExpr`） +### 4.1 本地 `-w N` -线格式：一行一个 JSON（NDJSON）。用户侧推荐平面写法（与 yield 一致）： +单进程：`N × per_worker` 个槽位 Actor；本进程即 Driver。 -```json -{"id": "t-0", "x": 1} -``` - -控制面归一化： - -```json -{"id": "t-0", "op": "execute", "args": {"x": 1}, "meta": {}} -``` +### 4.2 torchrun -| 字段 | 说明 | +| Rank | 角色 | |------|------| -| `id` | 任务身份；缺省则生成 UUID | -| `op` | 缺省 `"execute"`；**产品面仅支持 `execute`** | -| `args` | 业务字段；平面 JSON 的剩余键落入此 | -| `meta` | 预留；当前调度/执行不消费 | - -新能力写在用户 `execute` 内，**不**靠扩展 op 插件表。 +| 0 | Driver（plan + dispatch）+ 本机槽位 | +| >0 | 本机槽位，等 `Shutdown` | -> 命名对照（Agentic Infra）：用户侧 ≈ TaskSpec 载荷；内部仍叫 `TaskExpr`，对外文档逐步对齐。 +- 读 `RANK` / `WORLD_SIZE` / `MASTER_ADDR` / `MASTER_PORT`（及可选 `LOCAL_RANK`）。 +- Pulsing 种子端口默认 `MASTER_PORT+17`，可用 `PERSISTING_PULSING_PORT` 覆盖，避免与 c10d 冲突。 +- 控制面**不**自己拉齐进程；由 torchrun 负责。 -### 4.2 结果（`TaskResult`） +### 4.3 槽位命名 -| 字段 | 说明 | -|------|------| -| `task_id` | 对应任务 id | -| `ok` | 是否成功 | -| `cancelled` | L2 取消（非业务失败） | -| `value` / `error` / `traceback` | 成功值或失败信息 | -| `worker` | 执行 worker id | -| `started_at` / `finished_at` | Unix 秒 | -| `infra_retries` | L2 基础设施重试次数（0 = 首次成功） | +池按 **slot-major** 展平：先各 worker 的 slot0，再 slot1… +扁平下标：`slot * n_workers + worker`（DeathWatch / quarantine 必须用这套下标，不能用本 rank 局部序数）。 --- -## 5. 运行时拓扑 +## 5. 调度与执行 -```text - plan.py - （spawn --python -c bootstrap） - │ NDJSON TaskExpr - ▼ - Driver（rank0 / 本地控制进程） - │ plan emit + least-loaded dispatch - │（直接 ask Worker，并行） - ┌────────────┼────────────┐ - ▼ ▼ ▼ - slot actors (N workers × P per-worker) - │ each: WorkerActor + Python host - └──── stdin/stdout JSON ────┘ - │ - ▼ - plan.py::execute(item) - │ - ▼ - 控制面 ResultSink（stdout 与/或 --sink） -``` +### 5.1 Driver 循环 -每个从 `plan()` 出来的任务在 Driver 调度层对应一个 **`RunFuture`**（`wait` / `cancel` / `task_id`）。 +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.1 本地 `-w N` +### 5.2 sticky-only 与 quarantine -单进程：N 个 `WorkerActor`；本进程即 **Driver**（`Driver::run_plan`），无独立 DriverActor。 - -### 5.2 torchrun - -| Rank | 角色 | +| 情况 | 行为 | |------|------| -| 0 | **Driver**（plan + dispatch）+ 本机 `per_worker` 个槽位 Actor | -| >0 | 本机 `per_worker` 个槽位 Actor，等 `Shutdown` | +| 尚未接触任何槽 | least-loaded 选槽 | +| 已 ask 过某槽 | infra 重试 **只**打该槽 | +| 该槽被 quarantine | **拒绝**改投他槽，任务记 infra 失败（避免跨槽 at-least-once 重跑 `execute`） | +| 全槽 quarantine | acquire fail-fast，不再死等 | -- 读环境：`RANK` / `WORLD_SIZE` / `MASTER_ADDR` / `MASTER_PORT`（及可选 `LOCAL_RANK`）。 -- Pulsing 种子：`MASTER_ADDR:(MASTER_PORT+17)`，可用 `PERSISTING_PULSING_PORT` 覆盖，避免与 c10d 撞端口。 -- 控制面**不**自己 spawn 进程；由 torchrun 拉齐 ranks。 +同槽 **ResultCache**（`task_id → TaskResult`）跨 Supervision 重启共享：丢 reply 时可幂等取回，不重跑 `execute`。 -### 5.3 调度参数 +`--retries` 只覆盖 **worker ask / 基础设施**失败；**不**解读业务 `ok=false`（那是 L3 语义重试，未做）。 -- **`--per-worker N`**（默认 1）：每个逻辑 worker / torchrun rank 上的**并发 Execute 槽**。 - - 实现：每槽一个 `WorkerActor` + 独立 Python host（Actor mailbox 仍串行，槽与槽之间并行）。 - - 池按 **slot-major** 展平：先铺满各 worker 的 slot0，再 slot1…，least-loaded 仍优先打散到不同 worker。 -- **Least-loaded**：在展平后的槽位上选当前 in-flight 最少者；平手取更小下标。 -- **`max_inflight`**：全局上限（默认 `workers × per-worker`）。 -- 任务 **直接 `ask` 槽位 Actor**（不经 Driver 串行 `await Execute`）。 -- **完成即回调**：`on_result` / sink 按完成序；outstanding ≤ `max_inflight`。 -- Infra retry：失败则释放槽位再 acquire，避免粘在坏节点上。 -- Job 级 `CancellationToken`：Ctrl-C → 停止接新任务 + 子 token 取消在途 acquire。 +### 5.3 L3 执行缝 ---- - -## 6. L3 执行缝 - -每个 **槽位** 持有一个长驻 `--python` 子进程（行协议 JSON）： +每槽一个长驻 Python 子进程（行协议 JSON）： | cmd | 行为 | |-----|------| -| `run_plan` | 加载 plan 模块（带 `argv`），调用 `execute(item)`；按路径缓存模块 | +| `run_plan` | 加载脚本（带 argv），调 `execute(item)`；按路径缓存模块 | | `shutdown` | 退出 host | -- Worker **不**解释业务语义；Driver **不**碰 Python。 -- 失败时 traceback 编入 `TaskResult` 回传。 -- 产品面 Executor 路由：**仅** `op=execute` → `PlanExecuteExecutor`。 +Worker **不**解释业务；Driver **不**碰 Python。失败 traceback 编进 `TaskResult`。 ---- +### 5.4 取消 -## 7. Phase-1 合约实现 - -相对 Agentic Infra PPT，本阶段落地三条主契约（算法面保持极轻）。 - -### 7.1 RunFuture（调度原子） - -- 实现：`future.rs`；进程内 `JoinHandle` + `CancellationToken`。 -- **已支持**：`wait`（始终 join）、协作式 `cancel`（未启动 / 等待 acquire）。 -- **在途 cancel**：Worker 共享 job `CancellationToken`；取消时 **kill Python host** 并返回 `cancelled`（下一任务会重新拉起 host）。 -- **跨 rank**：每 rank 有独立 `JobControlActor`（与 Worker **分邮箱**）；rank0 Ctrl-C → 广播 `Cancel` → 远端 token cancel → 在途 host kill。不经过 Worker mailbox，避免被串行 `Execute` 挡住。 -- **取消语义**：`wait` **不**把已成功结果改写成 `cancelled`；取消来迟则保留成功。 -- Driver 在 `ask` 上同样 select cancel，避免 Ctrl-C 后干等 Python。 - -### 7.2 L2 基础设施重试 - -- Driver 在 **worker ask 失败**（节点丢、传输错）时重投；次数 `--retries`。 -- **Sticky**：失败后优先同一槽位，便于命中该 Worker 的 `task_id → TaskResult` 缓存（同 worker 上丢失 reply 可幂等重取，不重跑 `execute`）。 -- **不**解读 `execute` 返回的业务失败（那是 L3 Meta 语义重试，未做）。 -- Infra 耗尽 → 该任务 `TaskResult` 失败（`ok=false`），**不** `bail` 整次 job；sibling 继续跑。 -- `TaskResult.infra_retries` 可观测。 -- **Sink 去重**：`JsonlFileSink` 打开时从已有 JSONL **seed** `task_id`；跨进程重复 append 跳过。Vortex Tee 同样可 seed。 -- **Seen ⊆ durable**：先 reserve `task_id`，写盘/Vortex 失败则 **rollback seen**，避免永久丢账。 -- **Live SkipSet**：dispatch 前 claim `task_id`（resume 种子 + 本 job 已派发）；同 id 不再打到另一 worker。`--check --limit` 用 SkipSet 跳过尾部任务。 -- **Async sink writer**：完成路径 `await` enqueue（队列满则自然背压 Driver drain）；`join` 汇总 persist 失败并 fail job。persist 失败会 `SkipSet::remove`，便于后续 `--resume` 发现未落盘 id。 -- **Tee 尽力而为**：某一路失败不阻断兄弟 sink（仍返回首个错误）。 -- **Sticky-only after contact**：一旦对该槽发起过 `Execute` ask，infra 重试**只**打同一槽（`acquire_sticky`）；若该槽被 quarantine，**拒绝**改投他槽（避免跨槽 at-least-once 重跑 `execute`），任务记 infra 失败。 -- **槽位 quarantine**：连续 infra ask 失败达到阈值后摘槽；**全槽 quarantine → acquire fail-fast**（任务 infra 失败，非整 job 死等）。DeathWatch（仅本地）在 Worker 终止时 `force_quarantine`。Watch 注册必须用 **slot-major 扁平下标**（`DistEnv::slot_flat_index`），不能用本 rank 的 `0..per_worker` 序数。 -- **Pulsing 深化**：Worker `spawn_factory` + `SupervisionSpec`（失败重启上限）；**ResultCache 与槽位同生命周期**（`Arc` 跨重启共享）；统一 `resolve_named` / `ask_timeout`；JobControl 侧信道 cancel + 本地 DeathWatch。 -- **仍须注意**：跨槽拒绝后该任务失败（非整 job bail）；算法若需更强保证应自备幂等或缩小副作用面。 - -### 7.3 唯一 sink（控制面落盘） - -- 只有控制面写就绪结果；Executor **禁止**直写存储。 -- `--sink DIR` → `ready.ndjson`（成功）+ `failures.ndjson`（失败/取消）；**逐条**在 `on_result` 时 append（非跑完批写）。 -- Driver **按完成序**调用 `on_result`（非 plan 提交序）；有界 outstanding（≤ `max_inflight`，`FuturesUnordered` + `RunFuture`），完成即释放，避免大 plan 囤积 JoinHandle。 -- 同 process 内按 `task_id` 跳过重复 append。 -- **`--traj`**（feature `traj-sink`）：Tee 到 Vortex，`CaptureRecord{kind=compute.result|compute.failure}` → `TrajectoryAppend`；默认 `{sink}/traj` / agent=`compute` / session=sink 目录名。JSONL 仍是 `--resume` 账本。 -- 未指定 `--sink` 时：默认仅 stdout NDJSON（开发视图，非耐久 L1）。 -- **限制（已知）**：Vortex 侧跨进程未按 `task_id` 去重；resume 跳过已完成 id 可避免重复写。 - -### 7.4 Checkpoint（断点续跑） - -- 账本 = `--sink` 下已有 `ready.ndjson` / `failures.ndjson` 的 `task_id`。 -- `--resume`（须带 `--sink`）：Driver 跳过已终端 id，不重跑；`plan()` 仍会再 emit 一遍。 -- 同目录 `checkpoint.json`：`ok/fail/cancelled/skipped/dispatched/updated_at`（节流写入）。 -- 失败/取消默认也不重跑；要重跑请编辑 failures 或换 sink 目录。 - -### 7.5 与全图的差距（诚实清单） - -| 主张 | Phase-1 | -|------|---------| -| TaskSpec / RunFuture / RunResult 对外语言 | 内部名仍 `TaskExpr` / 进程内 Future / `TaskResult` | -| 就绪 → L1 轨迹 | JSONL 占位 | -| 可中断 | 调度层 + kill Python host；跨 rank 经 JobControlActor 广播 | -| 可复现 / 位置无关 | 未钉死 | -| Meta 语义重试 | 未做 | -| 模板 fan-out / 配额 / 亲和 / 观测账本 | least-loaded + per-worker；无模板/亲和 | +- Ctrl-C → job `CancellationToken`：停接新任务；在途 acquire / ask 可取消。 +- 在途 execute：**kill Python host**，返回 `cancelled`（下一任务会重新拉起 host）。 +- 跨 rank：每 rank 独立 `JobControlActor`（与 Worker **分邮箱**），rank0 广播 `Cancel`，避免被串行 `Execute` 堵住。 +- 已成功的结果**不会**被改写成 `cancelled`。 --- -## 8. 模块地图与语义原语 +## 6. 耐久、幂等与续跑 -产品按**语义原语 / 接口**切开：每块有稳定 id、主模块、同文件合约测。 -注册表见 [`blocks`](../../../crates/persisting-compute/src/blocks.rs)。 +### 6.1 唯一 sink -**测试策略** +只有控制面写就绪结果；Executor **禁止**直写存储账本。 -- **单元 / 合约测**：放在原语所在源文件末尾的 `#[cfg(test)]`（与接口同文件）。 -- **`tests/`**：只放跨模块**集成**（如 `tests/integration_local.rs`：fleet 完成序、resume skip、cancel、`--per-worker`、argv）。 +| 路径 | 作用 | +|------|------| +| `--sink DIR` | `ready.ndjson` + `failures.ndjson` + `checkpoint.json` | +| 无 `--sink` | 默认仅 stdout NDJSON（开发视图，非耐久） | +| `--traj` | Tee 到 Vortex；**JSONL 仍是 resume 真相** | -| Block id | 合约 | 主模块 | 测试 | -|----------|------|--------|------| -| `task_wire` | TaskExpr / TaskResult | `task` | 同文件 | -| `python_env` | PYTHONPATH 合并 | `python_env` | 同文件 | -| `placement` | least-loaded · sticky · slot-major 命名 | `scheduler` · `dist` | 同文件 | -| `run_future` | wait / cancel 语义 | `future` | 同文件 | -| `idempotency` | 同 worker 结果缓存 | `result_cache` · `worker` | 同文件 | -| `sink` / `checkpoint` | 唯一 sink · resume skip · seed 去重 | `sink` · `checkpoint` | 同文件 | -| `observe` | 可选进度事件 | `observe` | 同文件 | -| `plan` / `execute` | plan emit · execute host · cancel kill | `plan` · `executor` | 同文件 | -| `worker` | Actor 缝合 · cache · 多 slot ask | `worker` | 同文件 | -| `driver` / `fleet` | 完成序 drain · skip · job cancel · 装配 | `driver` · `runtime` | `tests/integration_*.rs` | +`JsonlFileSink`：打开时从已有文件 **seed** `task_id`；先 reserve 再写盘，失败则 **rollback seen**（seen ⊆ durable）。 +异步 `sink_writer`：`join` 汇总 persist 失败并 fail job；失败时 `SkipSet::remove`，便于后续 `--resume` 发现未落盘 id。损坏 / 无 `task_id` 的 JSONL 行：**skip + warn**，不拖垮整本账本。 -### 8.1 文件职责 +### 6.2 Checkpoint / `--resume` -| 路径 | 职责 | -|------|------| -| `persisting-cli` → `Command::Compute` | 唯一 CLI 入口 | -| `cli.rs` | 参数 → `run_fleet` / sink / Ctrl-C | -| `blocks.rs` | 语义原语 id 注册表（文档对齐；非独立测套） | -| `driver.rs` | **Driver**：plan emit + least-loaded dispatch；完成序 `on_result` + 有界 inflight | -| `checkpoint.rs` | `--resume` 账本 + `checkpoint.json` 进度 | -| `future.rs` | `RunFuture` · `wait_all` | -| `result_cache.rs` | 同 worker infra 幂等缓存 | -| `scheduler.rs` | Least-loaded · sticky prefer | -| `observe.rs` | `--observe` 事件 | -| `sink.rs` | `ResultSink` · `JsonlFileSink` · `TeeSink` | -| `sink_traj.rs` | （feature `traj-sink`）`VortexResultSink` | -| `plan.rs` | bootstrap：`sys.argv` + `plan()` → NDJSON | -| `runtime.rs` | local / torchrun fleet 装配（`spawn_rank_slots` / `run_driver_loop` 共享） | -| `worker.rs` | `Execute` / `Shutdown` + cache + cancel token | -| `executor.rs` | PlanExecute host（可取消 kill） | -| `dist.rs` | torchrun 环境 → Pulsing seed · slot 命名 | -| `job_control.rs` | 旁路 cancel + 本地 DeathWatch | -| `pulsing_ext.rs` | 统一 resolve / ask_timeout / spawn_supervised | -| `skip.rs` / `sink_writer.rs` | live claim · 异步 persist | -| `check.rs` | `--check` / `--self-test` | -| `task.rs` | `TaskExpr` / `TaskResult` | -| `python_env.rs` | `PYTHONPATH` 合并 | - -crate 公开面收窄：多数模块 `pub(crate)`；`lib` 只 re-export CLI / 集成测所需类型（`run_compute`、`RunOptions`、`SkipSet`、sink 主类型等）。 +- 账本 = sink 下已终端的 `task_id`（ready ∪ failures）。 +- `--resume`：Driver 跳过这些 id；**`plan()` 仍会再 emit 一遍**（大 plan 的成本是已知限制）。 +- 失败 / 取消默认也不重跑；要重跑请编辑 failures 或换目录。 + +### 6.3 两类重试（边界） + +| 层 | 管什么 | Phase-1 | +|----|--------|---------| +| L2 | 计算单元在不在（ask / 节点） | `--retries` + sticky + quarantine | +| L3 | 要不要再产一次（业务） | 未做；留给用户 `execute` | --- -## 9. 校验与自测 +## 7. 模块地图 -| 模式 | 行为 | +| 模块 | 职责 | |------|------| -| `cargo test -p persisting-compute` | 同文件单元 / 合约测 + `tests/integration_*.rs` | -| `--check` | python 探针 → plan emit → 确认 `execute` 可调用 → `run_local_fleet` 试跑 | -| `--self-test` | 写临时内置 `plan()`+`execute`，走完整 check（不依赖用户文件） | +| `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`。 --- -## 10. 非目标与路线 +## 8. 非目标与下一阶段 -### 10.1 刻意不做（保持薄） +### 刻意不做（保持薄） -- 多 op 产品面（`echo` / `call` / `py` 等已移除） -- 独立 Python `persisting.compute` 编排包 -- 控制面内嵌用户解释器 +- 多 op 产品面、独立 Python 编排包、内嵌用户解释器 - 自研进程启动器替代 torchrun - 在 L2 堆 harness / 判分 / Agent 会话 +- Meta 语义重试、模板 fan-out、配额、亲和调度 -### 10.2 下一阶段优先 +### 已知差距 -1. **跨 worker 幂等**（已：sticky-only-after-contact + 同槽 result_cache + sink seed/rollback + 跨 rank cancel；quarantine 后该任务 infra 失败而非改投） -2. **`ResultSink` → Vortex**（真 L1 默认路径） -3. 命名对外对齐（TaskSpec / RunFuture）+ 更多红灯 -4. （可选）Meta Executor 薄包装 - ---- +| 项 | 说明 | +|----|------| +| 昂贵 plan 无 cursor | resume 仍全量再 emit | +| 真 torchrun CI | 有双 ActorSystem 烟测；真实多进程 e2e 仍薄 | +| 错误分类 | infra / execute / cancel 尚未成稳定枚举字段 | +| DeathWatch | 仅本地；远端死槽靠 ask 失败路径 | +| 对外命名 | 内部仍 `TaskExpr`；文档侧逐步对齐 TaskSpec | -## 11. 设计原则 +### 设计原则（摘要） -1. **算法入口极轻**：两个函数 + 可选 argparse；本地 `for` 循环即真。 -2. **扩规模换命令**：同一文件，`persisting compute` / `torchrun`。 -3. **argv 一致**：`--` 后与 `python task.py …` 对齐。 -4. **能力成本沉在骨架**：Pulsing / torchrun 负责分布；用户不学 actors。 -5. **边界守住**：新产品能力优先进 `execute`，不膨胀 CLI 与 op 表。 -6. **两类重试分开**：L2 管计算单元在不在；L3 Meta 管要不要再产一次。 -7. **唯一 sink**：就绪 Result 只从控制面 append。 -8. **竖切优先**：先打通可跑路径与合约形状，再补工厂能力（模板 / 配额 / 亲和）。 +1. 算法入口极轻：两个函数；本地 `for` 循环即真。 +2. 扩规模换命令，不改文件。 +3. Driver ≠ Actor；Pulsing ≠ 调度器。 +4. 跨槽拒绝优先于跨槽重跑。 +5. persist 失败必须可见；JSONL 是 resume 真相。 +6. 新产品能力优先进 `execute`，不膨胀 CLI。 --- -## 12. 相关文档 +## 9. 相关文档 -- Agentic Infra 主张：`deeplink.fabric/projects/agent.infra/report-agent-infra.typ` -- 轨迹 L1（Vortex）：[trajectory_storage.zh.md](trajectory_storage.zh.md) -- CLI 整体（engine 路径）：[cli_architecture.zh.md](cli_architecture.zh.md) -- 用法示例：[examples/compute/README.md](../../../examples/compute/README.md) +- [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 7d0450c..71b750a 100644 --- a/docs/src/design/index.md +++ b/docs/src/design/index.md @@ -30,7 +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 control plane (Phase-1) | [compute_control_plane.zh.md](compute_control_plane.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 00af573..18b16f1 100644 --- a/docs/src/design/index.zh.md +++ b/docs/src/design/index.zh.md @@ -42,7 +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 已落地 | +| 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) | 规划 | @@ -54,7 +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) | `plan()` + `execute` 薄编排 / torchrun（Phase-1） | +| **`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 index 50ecad8..b10a506 100644 --- a/examples/compute/README.md +++ b/examples/compute/README.md @@ -1,52 +1,13 @@ # Persisting compute -一个文件： - -```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)) -``` +一个 Python 文件：`plan()` 产出任务，`execute(item)` 处理一条。 ```bash -python3 task.py --n 2 -persisting compute task.py -w 4 --per-worker 1 -- --n 2 -# --per-worker 1（默认）：每 worker 一槽；耗时差大时优先加 -w -# --per-worker N：每 worker/rank N 个槽（各带独立 Python host，真并行） -persisting compute task.py -w 1 --per-worker 4 -- --n 2 -persisting compute task.py -w 4 --retries 2 --sink /tmp/run1 -- --n 2 -# 中断后续跑（跳过 ready/failures 里已有 task_id）；看进度：/tmp/run1/checkpoint.json -persisting compute task.py -w 4 --sink /tmp/run1 --resume --observe -- --n 1000 -# 同时写入 Vortex 轨迹（JSONL 仍用于 resume） -persisting compute task.py -w 2 --sink /tmp/run1 --traj -- --n 4 -# 读回：persisting traj stats /tmp/run1/traj --agent-id compute --session-id run1 -# 可观测：stderr 上 `[obs]` 行（默认同时关掉结果 NDJSON 刷屏） -# 注意：`--observe` 必须写在 `--` 前面 -persisting compute task.py -w 2 --observe -- --n 2 -persisting compute task.py -w 2 --observe --observe-file /tmp/obs.ndjson -- --n 2 -persisting compute task.py -w 2 --observe --results ndjson -- --n 2 # 结果 + obs 都要 -torchrun --nproc_per_node=8 -- persisting compute task.py -- --n 2 -persisting compute task.py --check -- --n 2 +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 ``` -设计说明见 [docs/src/design/compute_control_plane.zh.md](../../docs/src/design/compute_control_plane.zh.md)。 - -- `execute` 收到的就是 `plan()` yield 出来的 dict -- `--` 后参数进入 `sys.argv` -- `--sink` 由控制面唯一落盘 `ready.ndjson` / `failures.ndjson` -- Ctrl-C 取消未启动的 RunFuture（Phase-1） +完整用法：[Compute 快速上手](../../docs/src/guide/compute_quickstart.zh.md) +架构：[Compute 架构](../../docs/src/design/compute_control_plane.zh.md)