Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Build: LocalWorker + the Scheduler

Maps to: Phase 4 (scheduler + LocalWorker). Kind: Build.

Objective

Turn the parts into an engine. In the panoptes-control crate (the binary crate Part V started for the API), add two modules: worker.rs with LocalWorker — the first concrete WorkerHandle, running the eval in-process — and scheduler.rs with plan_jobs, run_job_with_retry, the WorkerSource seam, and the Scheduler that claims a queued run, splits it into jobs, and drives them with bounded concurrency and retry-the-retryable. By the end, tick() takes a queued run all the way to Done, with its records in the log and the store — and it does so through dyn WorkerHandle, never learning local from remote.

Scaffold

Create (two new modules in the existing crates/panoptes-control):

  • crates/panoptes-control/src/worker.rsLocalWorker.
  • crates/panoptes-control/src/scheduler.rsRetryPolicy, WorkerSource, plan_jobs, run_job_with_retry, Scheduler.

Modify:

  • crates/panoptes-control/src/lib.rs — add pub mod scheduler; and pub mod worker;, and re-export the public surface: pub use scheduler::{RetryPolicy, Scheduler, WorkerSource, plan_jobs}; and pub use worker::LocalWorker;.
  • crates/panoptes-control/Cargo.toml — ensure [dependencies] has futures = { workspace = true } (for buffer_unordered) and async-trait = { workspace = true } (for the WorkerHandle impl); control-core, control-eval, control-store, tokio, and tracing are already there from Part V. [dev-dependencies] needs pretty_assertions.

Dependencies this chapter exercises: futures (the buffer_unordered bounded-concurrency adapter — see the concept chapter), async-trait (the LocalWorker: WorkerHandle impl behind dyn), control-store (claim_next_run, insert_jobs, record_job_outcome, new_job), control-eval (run_eval, load_manifest, append_records), tokio::fs (creating the out-dir).

Expected result: cargo test -p panoptes-control5 new tests pass — four in scheduler.rs (plan_jobs_is_one_per_model_epoch, tick_processes_a_run_to_done, retryable_failure_is_redispatched_to_the_next_worker, terminal_failure_is_not_retried) and one in worker.rs (local_worker_runs_the_eval), alongside the API tests already passing.

The spec (givens)

plan_jobs — split a run into the unit that fans out

/// One job carrying the whole vignette batch, per model, per epoch.
pub fn plan_jobs(run: &Run, vignettes: &[Vignette]) -> Vec<Job>;

For every model in run.models, for every epoch in 0..run.epochs, produce one Job carrying the full vignettes batch, that model, and that epoch. Order is model-outer, epoch-inner. Build each job with control_store::new_job(run.id, EvalJob { vignettes: vignettes.to_vec(), model: model.clone(), epoch }), which stamps a fresh JobId and JobStatus::Pending. A run with 2 models × 3 epochs yields 6 jobs — this is the model × epoch fan-out unit the whole coordinator is built around.

→ Answer key

RetryPolicy and run_job_with_retry — retry only the retryable, on the next worker

#[derive(Clone, Copy)]
pub struct RetryPolicy {
    pub max_attempts: u32, // Default: 3
}

async fn run_job_with_retry(
    workers: &[Arc<dyn WorkerHandle>],
    policy: RetryPolicy,
    seed: usize,
    job: Job,
) -> Result<JobOutcome, ControlError>;

RetryPolicy derives Default with max_attempts: 3. run_job_with_retry loops: on attempt n, dispatch to workers[(seed + n) % workers.len()]. On Ok, return it. On an error that is retryable and is under the attempt cap (e.is_retryable() && attempt + 1 < policy.max_attempts), bump attempt and try the next worker. On any other error — terminal, or the cap reached — return the error. The seed is the job's index in the run, so different jobs start on different workers and the pool spreads evenly. This is the concept chapter's loop, now over Arc<dyn WorkerHandle>.

→ Answer key

WorkerSource and StaticPool — where the scheduler gets its pool

/// Where the scheduler gets its workers — snapshotted once per run.
pub trait WorkerSource: Send + Sync {
    fn snapshot(&self) -> Vec<Arc<dyn WorkerHandle>>;
}

/// A fixed set of workers — the local, single-process case.
struct StaticPool(Vec<Arc<dyn WorkerHandle>>);

The scheduler does not hold a Vec of workers directly; it holds an Arc<dyn WorkerSource> and calls snapshot() once at the start of each run. StaticPool is the trivial source — it clones its fixed vec. Why the indirection now, when the pool is fixed? Because in Part VIII the remote pool gains and loses workers as connections open and close, and snapshot() is the seam that lets the pool change underneath without the scheduler noticing. Introduce it here so the payoff arc has somewhere to plug in.

→ Answer key

Scheduler — fields, tick, and process_run

pub struct Scheduler {
    store: Store,
    workers: Arc<dyn WorkerSource>,
    policy: RetryPolicy,
    concurrency: usize, // 4
    out_dir: PathBuf,
}

impl Scheduler {
    /// Over a fixed local worker pool.
    pub fn new(store: Store, workers: Vec<Arc<dyn WorkerHandle>>, out_dir: impl Into<PathBuf>) -> Self;
    /// Over any worker source (the Part VIII hook).
    pub fn with_source(store: Store, workers: Arc<dyn WorkerSource>, out_dir: impl Into<PathBuf>) -> Self;

    /// Claim one queued run and process it to completion. `None` if nothing is queued.
    pub async fn tick(&self) -> Result<Option<RunId>, ControlError>;
}

new wraps the vec in a StaticPool and delegates to with_source; both set policy to RetryPolicy::default() and concurrency to 4.

tick calls store.claim_next_run().await?; if that is None, return Ok(None); otherwise remember the run id, run the (private) process_run(run), and return Ok(Some(run_id)).

process_run(run) is the heart:

  1. tokio::fs::create_dir_all(&self.out_dir).await?.
  2. let workers = self.workers.snapshot(); — snapshot the pool once. If it is empty, return ControlError::Worker("no workers available").
  3. load_manifest(&run.manifest).await? → the vignettes; plan_jobs(&run, &vignettes) → the jobs; self.store.insert_jobs(&jobs).await?.
  4. Build one future per job, .enumerate()d so each job's index is its retry seed. Inside each: run_job_with_retry(&workers, policy, i, job).await?, then append_records(&log_path, &outcome.records).await?, then store.record_job_outcome(run_id, &outcome).await?. The log path is out_dir.join(format!("{run_id}.jsonl")).
  5. Drive them: stream::iter(futures).buffer_unordered(self.concurrency), pulling with .next().await and ?-propagating each result.

record_job_outcome (from Part IV) is what advances done_count and flips the run to Done when the last job lands — so process_run never sets run status itself; it just records outcomes and lets the store's transactional bookkeeping close the run out.

→ Answer key

LocalWorker — the first concrete WorkerHandle

pub struct LocalWorker {
    id: String,
    client: Arc<dyn ModelClient>,
}

impl LocalWorker {
    pub fn new(id: impl Into<String>, client: Arc<dyn ModelClient>) -> Self;
}

#[async_trait]
impl WorkerHandle for LocalWorker {
    fn id(&self) -> &str;
    async fn dispatch(&self, job: Job) -> Result<JobOutcome, ControlError>;
}

dispatch is three lines of glue: run_eval(self.client.as_ref(), &job.spec).await? for the records, then wrap them in JobOutcome { job_id: job.id, records }. That is the entire difference between "runs here" and "runs on another machine" — the LocalWorker runs run_eval in-process; the RemoteWorker of Part VIII ships job over a socket. Both are just a WorkerHandle, and the scheduler holds them as Arc<dyn WorkerHandle> without knowing which.

→ Answer key

Concepts exercised

  • Bounded concurrency with futures::stream::buffer_unordered over a stream of per-job futures.
  • A retry loop that re-dispatches retryable failures across a pool and lets terminal ones surface.
  • The dyn WorkerHandle seam paying off: one LocalWorker impl, driven by a scheduler that names no concrete worker type.
  • A WorkerSource indirection that lets the pool be fixed now and live later.
  • Composing the store's atomic claim and transactional recording (Part IV) with the eval workload (Part III) into one drive loop.

The build loop (you drive)

Test 1 — plan_jobs_is_one_per_model_epoch (in scheduler.rs)

  1. Write the failing test. Build a Run with models = ["a", "b"], epochs = 3, one vignette. Call plan_jobs(&run, &[vignette]). Assert jobs.len() == 6 and jobs.iter().all(|j| j.status == JobStatus::Pending).
  2. Predict: if you looped epoch-outer, model-inner instead, would len() change? Would any test in this chapter notice? (What is pinned about the order?)
  3. Run — fails to compile (no plan_jobs).
  4. Implement plan_jobs.
  5. Run green, commit.

Test 2 — local_worker_runs_the_eval (in worker.rs, #[tokio::test])

  1. Write the failing test. Define a StubClient: ModelClient with no network — generate(prompt) returns ModelResponse { text: format!("re:{prompt}"), .. }. Build LocalWorker::new("w1", Arc::new(StubClient)) and a Job with two vignettes ("a"/"x", "b"/"y"), model "claude", epoch 0. dispatch(job.clone()).await.unwrap(). Assert worker.id() == "w1", outcome.job_id == job.id, outcome.records.len() == 2, outcome.records[0].response == "re:x".
  2. Predict: the stub reports model_name() == "stub", but the job's model is "claude". Which lands in each record's model field, and why is that run_eval's doing rather than the worker's?
  3. Run — fails to compile (no LocalWorker).
  4. Implement LocalWorker + the WorkerHandle impl.
  5. Run green, commit.

Test 3 — retryable_failure_is_redispatched_to_the_next_worker (in scheduler.rs, #[tokio::test])

  1. Write the failing test. Define a RecordingWorker (one record per vignette, no network) and a FlakyWorker that always returns ControlError::Worker("boom") and counts calls in an Arc<AtomicUsize>. Pool = [FlakyWorker, RecordingWorker]. Build a one-vignette Job. Call run_job_with_retry(&workers, RetryPolicy::default(), 0, job).await.unwrap(). Assert the flaky worker was called exactly once and outcome.records.len() == 1.
  2. Predict: with seed = 0, attempt 0 hits index 0 (flaky) and attempt 1 hits (0+1) % 2 = 1 (recording). If you passed seed = 1 instead, which worker would attempt 0 hit — and would the test still pass?
  3. Run — fails to compile (no run_job_with_retry).
  4. Implement RetryPolicy + run_job_with_retry.
  5. Run green, commit.

Test 4 — terminal_failure_is_not_retried (in scheduler.rs, #[tokio::test])

  1. Write the failing test. Define a BadInput worker that returns ControlError::Invalid("nope"). Pool = [BadInput]. Call run_job_with_retry and unwrap_err(). Assert !err.is_retryable().
  2. Predict: Invalid is terminal, so the loop returns on the first Err(e) => arm. If dispatch had returned ControlError::Worker instead, how many times would BadInput be called before the loop gave up, and why? (Hint: one worker, max_attempts = 3, modulo wraps.)
  3. Run, check, implement if the arm order was wrong (terminal must fall through the retryable guard).
  4. Run green, commit.

Test 5 — tick_processes_a_run_to_done (in scheduler.rs, #[tokio::test])

  1. Write the failing test. Store::in_memory(). Write a 2-vignette manifest to a temp file. Insert a Queued run over it (models = ["claude"], epochs = 1). Build Scheduler::new(store.clone(), vec![Arc::new(RecordingWorker { .. })], &out_dir). Assert sched.tick().await.unwrap() == Some(run.id). Then read the run back: status == Done, done_count == 1, and store.run_results(run.id).await.unwrap().len() == 2 (1 job × 2 vignettes). Finally assert a second tick() returns None — nothing queued left.
  2. Predict: process_run never sets the run's status to Done itself. What does, and when? (Recall record_job_outcome from Part IV.)
  3. Run — fails to compile (no Scheduler).
  4. Implement WorkerSource, StaticPool, Scheduler::new/with_source, tick, process_run.
  5. Run green, commit.
Why the tests use hand-written workers, not a mock server Every worker in these tests — RecordingWorker, FlakyWorker, BadInput, StubClient — is a few lines implementing the seam with no network. That is the seam paying off in the test suite itself: the scheduler and retry loop are exercised entirely through dyn WorkerHandle, so a worker that "always fails retryably" or "always returns two records" is a struct, not a server. The same code will drive a real RemoteWorker in Part VIII unchanged.

Done when

cargo test -p panoptes-control is green with the five new tests: plan_jobs yields one job per model×epoch, LocalWorker runs the eval through the seam, a retryable failure lands on the next worker while a terminal one surfaces at once, and tick drives a queued run all the way to Done with its records in both the log and the store.