Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions crates/aprender-distribute/src/executor/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@
//! This module defines the `Executor` trait and provides implementations
//! for different backends (CPU, GPU, Remote).

// `cpu` is the only executor module that was NOT feature-gated, and it is the
// one that needs tokio (`tokio::process::Command`, `AsyncReadExt`, `timeout`).
// Its five siblings below are all gated; this makes it consistent. #3176.
#[cfg(feature = "tokio")]
pub mod cpu;

#[cfg(feature = "gpu")]
Expand Down
32 changes: 32 additions & 0 deletions crates/aprender-distribute/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,19 @@ mod generated_contracts;
pub mod checkpoint;
pub mod error;
pub mod executor;
// TOKIO, NOT `cpu`. #3176 proposed `#[cfg(feature = "cpu")]` and the intent is
// right — gate the modules, do NOT make tokio a hard dependency — but `cpu` is
// only one of SIX features that pull tokio in: `cpu`, `remote`, `checkpoint`,
// `tui`, `serverless`, `microvm`. Gating on `cpu` would leave
// `--no-default-features --features remote` with tokio present and the modules
// that need it compiled out, which trades one broken configuration for another.
//
// Cargo creates an implicit `tokio` feature for the optional dependency, and it
// is true exactly when some feature brought tokio in. That is the predicate the
// code actually depends on.
#[cfg(feature = "tokio")]
pub mod messaging;
#[cfg(feature = "tokio")]
pub mod scheduler;
#[cfg(feature = "serverless")]
pub mod serverless;
Expand All @@ -90,16 +102,33 @@ pub mod task;
pub mod tensor;
// tui module removed — was ratatui-only

// Only the Pool API uses these, and the Pool API is tokio-gated below.
#[cfg(feature = "tokio")]
use error::{RepartirError, Result};
// `Pool` is the crate's high-level API and it is async to its core --
// tokio::sync::RwLock, tokio::task::JoinSet, tokio::time::sleep, and a
// CpuExecutor that spawns processes through tokio. There is no meaningful
// tokio-free Pool to preserve, so the whole API surface is gated with the
// runtime it requires rather than half-gated into something that compiles and
// cannot run. #3176.
#[cfg(feature = "tokio")]
use executor::cpu::CpuExecutor;
#[cfg(feature = "tokio")]
use executor::Executor;
#[cfg(feature = "tokio")]
use scheduler::Scheduler;
#[cfg(feature = "tokio")]
use std::sync::Arc;
#[cfg(feature = "tokio")]
use task::{ExecutionResult, Task};
#[cfg(feature = "tokio")]
use tokio::sync::RwLock;
#[cfg(feature = "tokio")]
use tokio::task::JoinSet;
#[cfg(feature = "tokio")]
use tracing::{debug, info};

#[cfg(feature = "tokio")]
/// High-level API for distributed task execution.
///
/// The `Pool` manages executors, schedules tasks, and coordinates results.
Expand Down Expand Up @@ -131,6 +160,7 @@ pub struct Pool {
workers: Arc<RwLock<JoinSet<()>>>,
}

#[cfg(feature = "tokio")]
impl Pool {
/// Creates a new pool builder.
#[must_use]
Expand Down Expand Up @@ -182,13 +212,15 @@ impl Pool {
}
}

#[cfg(feature = "tokio")]
/// Builder for `Pool`.
#[derive(Default)]
pub struct PoolBuilder {
cpu_workers: Option<usize>,
max_queue_size: Option<usize>,
}

#[cfg(feature = "tokio")]
impl PoolBuilder {
/// Sets the number of CPU workers.
#[must_use]
Expand Down
Loading