diff --git a/charts/flame/values.yaml b/charts/flame/values.yaml index 2b5be09d..7292d9e2 100644 --- a/charts/flame/values.yaml +++ b/charts/flame/values.yaml @@ -12,7 +12,6 @@ cluster: resreq: "cpu=1,mem=2g" policies: - drf - - gang scheduleInterval: 100 storage: "fs:///var/lib/flame/session" executors: diff --git a/ci/flame-cluster-benchmark.yaml b/ci/flame-cluster-benchmark.yaml index cee5c162..522a03d2 100644 --- a/ci/flame-cluster-benchmark.yaml +++ b/ci/flame-cluster-benchmark.yaml @@ -5,7 +5,6 @@ cluster: resreq: "cpu=1,mem=1g" policies: - priority - - gang storage: none schedule_interval: 100 executors: diff --git a/ci/flame-cluster-docker.yaml b/ci/flame-cluster-docker.yaml index 18b8a18f..4238ea31 100644 --- a/ci/flame-cluster-docker.yaml +++ b/ci/flame-cluster-docker.yaml @@ -5,7 +5,6 @@ cluster: resreq: "cpu=1,mem=1g" policies: - priority - - gang storage: fs://data/ schedule_interval: 500 # Scheduler loop interval in milliseconds (default: 500) executors: diff --git a/ci/flame-cluster-local.yaml b/ci/flame-cluster-local.yaml index c3dfee60..a8025e2d 100644 --- a/ci/flame-cluster-local.yaml +++ b/ci/flame-cluster-local.yaml @@ -5,7 +5,6 @@ cluster: resreq: "cpu=1,mem=1g" policies: - priority - - gang storage: fs://data/ schedule_interval: 100 # Scheduler loop interval in milliseconds (default: 500) executors: diff --git a/ci/flame-cluster.yaml b/ci/flame-cluster.yaml index 96eda841..9f409998 100644 --- a/ci/flame-cluster.yaml +++ b/ci/flame-cluster.yaml @@ -5,7 +5,6 @@ cluster: resreq: "cpu=1,mem=2g" policies: - priority - - gang storage: none schedule_interval: 100 executors: diff --git a/ci/k8s/e2e.sh b/ci/k8s/e2e.sh index a8ccbe0d..82996ee7 100755 --- a/ci/k8s/e2e.sh +++ b/ci/k8s/e2e.sh @@ -24,7 +24,6 @@ HELM_E2E_ARGS=( --set "global.imagePullPolicy=${IMAGE_PULL_POLICY}" --set "cluster.storage=${SESSION_MANAGER_STORAGE}" --set "cluster.policies[0]=drf" - --set "cluster.policies[1]=gang" --set sessionManager.persistence.enabled=false --set objectCache.persistence.enabled=false --set "objectCache.replicas=${OBJECT_CACHE_REPLICAS}" diff --git a/common/src/ctx.rs b/common/src/ctx.rs index 090872d6..37123b7d 100644 --- a/common/src/ctx.rs +++ b/common/src/ctx.rs @@ -27,9 +27,10 @@ const DEFAULT_FLAME_CONF: &str = "flame-cluster.yaml"; const DEFAULT_CONTEXT_NAME: &str = "flame"; const DEFAULT_FLAME_ENDPOINT: &str = "http://127.0.0.1:8080"; /// Default policies to enable when none specified in config. -/// Available configurable policies: "priority", "drf", "gang" -/// Note: "shim" plugin is always enabled (required for executor matching) -pub const DEFAULT_POLICIES: &[&str] = &["priority", "drf", "gang"]; +/// +/// Effective default scheduler stack: `priority + batch + shim`. +/// `batch` and `shim` are always enabled, while DRF remains available as an opt-in policy. +pub const DEFAULT_POLICIES: &[&str] = &["priority"]; const DEFAULT_STORAGE: &str = "sqlite://flame.db"; const DEFAULT_MAX_EXECUTORS_PER_NODE: u32 = 128; pub const DEFAULT_SESSION_RETRY_LIMITS: u32 = 5; @@ -597,6 +598,15 @@ mod tests { use super::*; use tempfile::TempDir; + #[test] + fn test_default_scheduler_policies() { + assert_eq!(DEFAULT_POLICIES, &["priority"]); + assert_eq!( + FlameCluster::default().policies, + vec!["priority".to_string()] + ); + } + #[test] fn test_flame_context_from_file() -> Result<(), FlameError> { // New config structure: max_executors is under cluster.limits @@ -607,7 +617,6 @@ cluster: resreq: "cpu=1,mem=1g" policies: - priority - - gang storage: sqlite://flame.db executors: shim: host @@ -628,7 +637,7 @@ cluster: ctx.cluster.resreq, Some(ResourceRequirement::from("cpu=1,mem=1g")) ); - assert_eq!(ctx.cluster.policies, vec!["priority", "gang"]); + assert_eq!(ctx.cluster.policies, vec!["priority"]); assert_eq!(ctx.cluster.storage, "sqlite://flame.db"); assert_eq!(ctx.cluster.executors.shim, Shim::Host); assert_eq!(ctx.cluster.limits.max_executors, 10); diff --git a/docs/api/types.md b/docs/api/types.md index d36937d2..0682b873 100644 --- a/docs/api/types.md +++ b/docs/api/types.md @@ -107,7 +107,7 @@ message SessionSpec { | `common_data` | bytes | Data shared across all tasks (optional) | | `min_instances` | uint32 | Minimum executor instances (default: 0) | | `max_instances` | uint32 | Maximum executor instances (optional, unlimited if not set) | -| `batch_size` | uint32 | Executors per batch for gang scheduling (default: 1) | +| `batch_size` | uint32 | Executors per batch for batch scheduling (default: 1) | | `priority` | uint32 | Session priority; higher = more important (default: 0) | | `resreq` | ResourceRequirement | Per-task resource request (optional); when absent, the server applies `cluster.resreq`, with a hardcoded `cpu=1,mem=1g,gpu=0` fallback. | diff --git a/docs/designs/RFE323-runner-v2/FS.md b/docs/designs/RFE323-runner-v2/FS.md index 88534078..5d3d221c 100644 --- a/docs/designs/RFE323-runner-v2/FS.md +++ b/docs/designs/RFE323-runner-v2/FS.md @@ -365,14 +365,14 @@ No changes required. The executor manager continues to handle executor lifecycle - Executor allocation happens asynchronously in the scheduler loop (not immediately on session creation): - **Scheduler Loop**: Runs periodically (every `schedule_interval` ms, typically 1000ms) - **AllocateAction** (`scheduler/actions/allocate.rs`): Executed as part of scheduler actions - - Gets all open sessions, orders them by the scheduler plugins' `ssn_order_fn` (PriorityPlugin → DRFPlugin → GangPlugin) + - Gets all open sessions, orders them by the scheduler plugins' `ssn_order_fn` (PriorityPlugin → DRFPlugin → BatchPlugin) - For each session, checks `plugins.is_underused(ssn)` - **NEW**: Add explicit `max_instances` check by counting actual executors from snapshot - This prevents over-allocation when multiple executors are created in one cycle - The plugin's cached `allocated` count doesn't update within a cycle - If underused and below `max_instances`, allocates executors - Calls `ctx.create_executor(node, ssn)` to actually create executor - - Delegate allocation decisions to the scheduler plugins (PriorityPlugin / DRFPlugin / GangPlugin) + - Delegate allocation decisions to the scheduler plugins (PriorityPlugin / DRFPlugin / BatchPlugin) - Plugins respect both `min_instances` (guaranteed) and `max_instances` (limit) - Note: Sessions with `min_instances > 0` will be allocated executors in the first scheduler cycle (within ~1 second) @@ -548,7 +548,7 @@ The enhancement spans multiple layers from SDK to session manager to executors: - Executes scheduling actions (AllocateAction, DispatchAction, etc.) - **AllocateAction** (`actions/allocate.rs`): - Gets all open sessions from snapshot - - Orders sessions by the scheduler plugins' `ssn_order_fn` (Priority → DRF → Gang) + - Orders sessions by the scheduler plugins' `ssn_order_fn` (Priority → DRF → Batch) - For each session: - Checks `plugins.is_underused(ssn)` - returns true if `allocated < deserved` in any resource dimension - **NEW**: Explicit `max_instances` check - counts actual executors from snapshot to prevent over-allocation @@ -1201,7 +1201,7 @@ if let Some(max_instances) = ssn.max_instances { - This is needed because the plugin's cached `allocated` count doesn't update within a single scheduler cycle - Prevents allocating beyond `max_instances` when creating multiple executors in one cycle - Creates one executor per iteration, then re-evaluates -- Sessions are processed in priority order (PriorityPlugin → DRFPlugin → GangPlugin) +- Sessions are processed in priority order (PriorityPlugin → DRFPlugin → BatchPlugin) ### System Considerations diff --git a/docs/designs/RFE400-batch-session/FS.md b/docs/designs/RFE400-batch-session/FS.md index 6d2f35b8..5ea4337d 100644 --- a/docs/designs/RFE400-batch-session/FS.md +++ b/docs/designs/RFE400-batch-session/FS.md @@ -151,15 +151,15 @@ flmctl list -s **Updates Required:** - `Plugin trait`: Add `is_ready()`, `on_pipeline_executor()`, `on_discard_executor()` methods -- `GangPlugin`: New scheduler plugin implementing these methods for batch tracking +- `BatchPlugin`: New scheduler plugin implementing these methods for batch tracking - `Statement`: New struct for accumulating pending allocations (pipeline/commit/discard) -- `AllocateAction`: Use Statement pattern - action is NOT aware of GangPlugin +- `AllocateAction`: Use Statement pattern - action is NOT aware of BatchPlugin - `Controller::launch_task()`: Batch-indexed task assignment - Storage schemas: Add `batch_size` to sessions, `batch_index` to executors **Design Principle: Actions Unaware of Gang Logic** -Actions use Statement to accumulate allocations. Statement calls Plugin trait callbacks (`on_pipeline_executor`, `on_discard_executor`). GangPlugin implements these callbacks to track pipelined count. Actions only call `ctx.is_ready()` - they don't know about GangPlugin. +Actions use Statement to accumulate allocations. Statement calls Plugin trait callbacks (`on_pipeline_executor`, `on_discard_executor`). BatchPlugin implements these callbacks to track pipelined count. Actions only call `ctx.is_ready()` - they don't know about BatchPlugin. **Compatibility:** - Fully backward compatible: `batch_size=1` (default) maintains existing behavior @@ -196,7 +196,7 @@ Actions use Statement to accumulate allocations. Statement calls Plugin trait ca │ │ │ │ │ │ plugins: HashMap │ │ │ │ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │ │ -│ │ │ Priority + DRF │ │ ShimPlugin │ │ GangPlugin │ │ │ +│ │ │ Priority + DRF │ │ ShimPlugin │ │ BatchPlugin │ │ │ │ │ │ (Plugin trait) │ │ (Plugin trait) │ │ (Plugin trait) │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ - is_underused │ │ - is_available │ │ - is_underused │ │ │ @@ -255,7 +255,7 @@ Actions use Statement to accumulate allocations. Statement calls Plugin trait ca **Design Principle: Extend Existing Plugin Trait** -The existing Plugin trait is extended with a single new method `is_ready()` for gang readiness checks. No other changes to the plugin machinery. GangPlugin implements this method alongside existing trait methods. +The existing Plugin trait is extended with a single new method `is_ready()` for gang readiness checks. No other changes to the plugin machinery. BatchPlugin implements this method alongside existing trait methods. **1. Plugin Trait (session_manager/src/scheduler/plugins/mod.rs)** @@ -323,23 +323,23 @@ impl PluginManager { } ``` -**3. GangPlugin (session_manager/src/scheduler/plugins/gang.rs)** +**3. BatchPlugin (session_manager/src/scheduler/plugins/batch.rs)** Implements Plugin trait including `is_ready()`, `on_pipeline_executor()`, `on_discard_executor()`: ```rust -pub struct GangPlugin { - ssn_state: HashMap, +pub struct BatchPlugin { + ssn_state: HashMap, } -struct GangState { +struct BatchState { batch_size: u32, allocated: u32, pipelined: u32, // Count of reserved (but not committed) executors max_instances: Option, } -impl Plugin for GangPlugin { +impl Plugin for BatchPlugin { fn setup(&mut self, ss: &SnapShot) -> Result<(), FlameError> { self.ssn_state.clear(); for ssn in ss.sessions.values() { @@ -347,7 +347,7 @@ impl Plugin for GangPlugin { .filter(|e| e.ssn_id.as_ref() == Some(&ssn.id)) .count() as u32; - self.ssn_state.insert(ssn.id.clone(), GangState { + self.ssn_state.insert(ssn.id.clone(), BatchState { batch_size: ssn.batch_size.max(1), allocated, pipelined: 0, @@ -434,7 +434,7 @@ impl Context { **5. Statement (session_manager/src/scheduler/statement.rs) - NEW** -Statement accumulates pending allocations without committing them. It uses Plugin trait callbacks to notify plugins - actions are NOT aware of GangPlugin. +Statement accumulates pending allocations without committing them. It uses Plugin trait callbacks to notify plugins - actions are NOT aware of BatchPlugin. ```rust pub struct Statement { @@ -456,12 +456,12 @@ impl Statement { } /// Reserve resources for an executor (in-memory only, no actual creation) - /// Calls on_pipeline_executor callback - GangPlugin tracks pipelined count + /// Calls on_pipeline_executor callback - BatchPlugin tracks pipelined count pub fn pipeline(&mut self, node: &NodeInfoPtr, ssn: &SessionInfoPtr) -> Result<(), FlameError> { // Reserve `ssn.resreq` resources on node in snapshot self.ctx.snapshot.reserve(node, ssn)?; - // Notify plugins via callback (GangPlugin increments pipelined count) + // Notify plugins via callback (BatchPlugin increments pipelined count) self.ctx.plugins.on_pipeline_executor(node.clone(), ssn.clone())?; // Record operation for later commit/discard @@ -482,14 +482,14 @@ impl Statement { } /// Discard all pending operations - release reserved resources - /// Calls on_discard_executor callback - GangPlugin decrements pipelined count + /// Calls on_discard_executor callback - BatchPlugin decrements pipelined count pub fn discard(self) -> Result<(), FlameError> { // Reverse order to properly unwind for op in self.operations.into_iter().rev() { // Release `ssn.resreq` resources on node self.ctx.snapshot.release(&op.node, &op.ssn)?; - // Notify plugins via callback (GangPlugin decrements pipelined count) + // Notify plugins via callback (BatchPlugin decrements pipelined count) self.ctx.plugins.on_discard_executor(op.node, op.ssn)?; } Ok(()) @@ -628,7 +628,7 @@ pub struct Executor { **Statement Pattern for Gang Scheduling:** -The Statement pattern ensures executors are only created when a full batch can be scheduled. Actions use Statement without knowing about GangPlugin - all gang logic is in Plugin trait callbacks. +The Statement pattern ensures executors are only created when a full batch can be scheduled. Actions use Statement without knowing about BatchPlugin - all gang logic is in Plugin trait callbacks. ``` ┌────────────────────────────────────────────────────────────────────┐ @@ -643,31 +643,31 @@ The Statement pattern ensures executors are only created when a full batch can b │ │ stmt.pipeline(node, ssn) │ │ │ │ → Reserves resources on node (in-memory) │ │ │ │ → Calls plugins.on_pipeline_executor() │ │ -│ │ → GangPlugin increments pipelined count │ │ +│ │ → BatchPlugin increments pipelined count │ │ │ │ → NO executor created yet │ │ │ └─────────────────────────────────────────────────────────┘ │ │ │ │ 3. Check Readiness (action doesn't know about gang) │ │ if ctx.is_ready(&ssn) { │ │ → plugins.is_ready() iterates all plugins │ -│ → GangPlugin checks: (allocated + pipelined) % batch == 0 │ +│ → BatchPlugin checks: (allocated + pipelined) % batch == 0 │ │ } │ │ │ │ 4a. READY → Commit All │ │ stmt.commit() │ │ → for each op: create_executor() │ -│ → on_create_executor() → GangPlugin clears pipelined │ +│ → on_create_executor() → BatchPlugin clears pipelined │ │ │ │ 4b. NOT READY → Discard All │ │ stmt.discard() │ │ → for each op (reverse): release resources │ │ → Calls plugins.on_discard_executor() │ -│ → GangPlugin decrements pipelined count │ +│ → BatchPlugin decrements pipelined count │ │ │ └────────────────────────────────────────────────────────────────────┘ ``` -**AllocateAction with Statement Pattern (no GangPlugin awareness):** +**AllocateAction with Statement Pattern (no BatchPlugin awareness):** ```rust // session_manager/src/scheduler/actions/allocate.rs @@ -692,7 +692,7 @@ async fn execute(&self, ctx: &mut Context) -> Result<(), FlameError> { // Reserve resources - plugins notified via on_pipeline_executor stmt.pipeline(node, &ssn)?; - // Check if ready - GangPlugin checks pipelined >= batch_size + // Check if ready - BatchPlugin checks pipelined >= batch_size if ctx.is_ready(&ssn)? { // Commit: create all executors stmt.commit().await?; @@ -876,7 +876,7 @@ for worker_id in range(8): | `on_discard_executor` | Gang | **NEW**: Track discarded executor | **Implementation References:** -- `session_manager/src/scheduler/plugins/gang.rs` - GangPlugin implementation +- `session_manager/src/scheduler/plugins/batch.rs` - BatchPlugin implementation - `session_manager/src/scheduler/plugins/mod.rs` - Plugin trait and PluginManager - `common/src/apis/types.rs` - Session and Task definitions - `session_manager/src/scheduler/actions/` - Scheduler actions @@ -886,7 +886,7 @@ for worker_id in range(8): ## Design Decisions -1. **Extend Plugin trait with `is_ready`**: The existing Plugin trait is extended with a single `is_ready()` method. GangPlugin implements this to check if enough resources are pipelined for a complete batch. +1. **Extend Plugin trait with `is_ready`**: The existing Plugin trait is extended with a single `is_ready()` method. BatchPlugin implements this to check if enough resources are pipelined for a complete batch. 2. **Statement pattern for batch allocation**: Executors are NOT created until a full batch can be scheduled. The Statement accumulates pending allocations (pipeline), checks gang readiness via `is_ready()`, then either commits all (creates executors) or discards all (releases reservations). diff --git a/docs/designs/RFE408-fairshare-batch-size/FS.md b/docs/designs/RFE408-fairshare-batch-size/FS.md index e764203f..46fcbf49 100644 --- a/docs/designs/RFE408-fairshare-batch-size/FS.md +++ b/docs/designs/RFE408-fairshare-batch-size/FS.md @@ -17,7 +17,7 @@ priority tier is now provided by `DRFPlugin` (see RFE433). ### Background -The `batch_size` feature (RFE400) introduced gang scheduling, enabling coordinated multi-executor workloads such as multi-node LLM inference with tensor parallelism. The Gang plugin ensures executors are committed in complete batches; however, the Fairshare plugin's resource distribution algorithm does not fully account for batch boundaries. +The `batch_size` feature (RFE400) introduced gang scheduling, enabling coordinated multi-executor workloads such as multi-node LLM inference with tensor parallelism. The Batch plugin ensures executors are committed in complete batches; however, the Fairshare plugin's resource distribution algorithm does not fully account for batch boundaries. ### Current Limitations @@ -25,7 +25,7 @@ The `batch_size` feature (RFE400) introduced gang scheduling, enabling coordinat 2. **Inefficient `is_underused` check**: The current implementation (`allocated < deserved`) does not consider batch boundaries. A session with `allocated=4` and `deserved=6` appears "underused" by 2 slots, but those 2 slots cannot form a complete batch, leading to wasted scheduling cycles. -3. **Fairshare-Gang plugin mismatch**: Fairshare calculates what a session "deserves" while Gang enforces batch atomicity at commit time. When Fairshare grants a non-batch-aligned allocation, Gang rejects partial batches, causing scheduling thrashing. +3. **Fairshare-Batch plugin mismatch**: Fairshare calculates what a session "deserves" while Gang enforces batch atomicity at commit time. When Fairshare grants a non-batch-aligned allocation, Gang rejects partial batches, causing scheduling thrashing. ### Problem Example @@ -43,7 +43,7 @@ Problem: - Session A with deserved=6 can only run 1 complete batch (4 executors) - 2 slots are "deserved" but can never be used (incomplete batch) - Allocate action keeps trying to allocate for Session A -- Gang plugin's is_ready() blocks until a full batch is pipelined +- Batch plugin's is_ready() blocks until a full batch is pipelined - Those 2 slots could have gone to Session B instead ``` @@ -98,7 +98,7 @@ No CLI changes are required. Validation errors are returned when invalid paramet **Out of Scope:** -- Changes to Gang plugin (already handles batch atomicity correctly) +- Changes to Batch plugin (already handles batch atomicity correctly) - New metrics or observability (may be added in follow-up work) **Limitations:** @@ -110,8 +110,8 @@ No CLI changes are required. Validation errors are returned when invalid paramet **Related Features:** -- **RFE400 (Batch Session)**: Introduced `batch_size` concept and Gang plugin -- **Gang Plugin**: Enforces batch atomicity at allocation/binding time +- **RFE400 (Batch Session)**: Introduced `batch_size` concept and Batch plugin +- **Batch Plugin**: Enforces batch atomicity at allocation/binding time - **Allocate/Dispatch Actions**: Use Fairshare's `is_underused` to determine scheduling decisions **Required Updates:** @@ -152,7 +152,7 @@ This enhancement modifies the Fairshare plugin's internal algorithm while mainta │ ┌─────────────────────────────────────────────────────────────────┐ │ │ │ PluginManager │ │ │ │ ┌─────────────────────┐ ┌─────────────────────┐ │ │ -│ │ │ FairShare │ │ GangPlugin │ │ │ +│ │ │ FairShare │ │ BatchPlugin │ │ │ │ │ │ │ │ │ │ │ │ │ │ ENHANCED: │ │ UNCHANGED: │ │ │ │ │ │ - setup() aligns │ │ - is_ready() │ │ │ @@ -529,7 +529,7 @@ Cluster: 10 slots | File | Description | | ---------------------------------------------------- | ---------------- | | `session_manager/src/scheduler/plugins/fairshare.rs` | Fairshare plugin | -| `session_manager/src/scheduler/plugins/gang.rs` | Gang plugin | +| `session_manager/src/scheduler/plugins/batch.rs` | Batch plugin | | `session_manager/src/scheduler/plugins/mod.rs` | Plugin trait | | `session_manager/src/scheduler/actions/allocate.rs` | Allocate action | | `session_manager/src/scheduler/actions/dispatch.rs` | Dispatch action | @@ -543,7 +543,7 @@ Cluster: 10 slots | ----------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | **Batch-unit distribution** | Distribute in batch units during the fair-share algorithm rather than aligning after distribution. This ensures fairness is calculated on usable resources. | | **Incomplete batch handling** | When cluster capacity cannot accommodate a full batch, those resources go to other sessions or remain unallocated. Allocating unusable resources wastes scheduling cycles. | -| **No Gang plugin changes** | The Gang plugin already correctly enforces batch atomicity. This enhancement ensures Fairshare uses consistent batch boundaries. | +| **No Batch plugin changes** | The Batch plugin already correctly enforces batch atomicity. This enhancement ensures Fairshare uses consistent batch boundaries. | | **Greedy batch allocation** | Sessions receive one batch unit at a time via priority queue. This ensures fairness across sessions with different batch sizes. | | **Frontend validation** | Validate `min_instances` and `max_instances` at the API boundary rather than silently aligning them in the scheduler. This provides clear feedback to users about invalid configurations. | | **Validation rules** | When `batch_size > 1`, `min_instances` must be 0 or a multiple of `batch_size`, and `max_instances` must be None or a multiple of `batch_size`. | diff --git a/docs/designs/RFE413-priority-scheduling/FS.md b/docs/designs/RFE413-priority-scheduling/FS.md index c365cee9..42beb00a 100644 --- a/docs/designs/RFE413-priority-scheduling/FS.md +++ b/docs/designs/RFE413-priority-scheduling/FS.md @@ -56,7 +56,7 @@ Legacy proportional-share behavior (pre-PriorityPlugin): | ---------- | -------- | ------- | -------------------------------------------------------------------------------------------------- | | `priority` | `uint32` | `0` | Session priority. Higher value = higher priority. Sessions with `priority = 0` use default priority. | -**Scheduling policy**: Priority-based scheduling is always active when the `PriorityPlugin` is registered. No configuration flag is required to enable it; the plugin is registered by default alongside `DRFPlugin` and `GangPlugin`. When all sessions share the same priority, the priority dimension has no effect and ordering falls through to the downstream deferral plugins. +**Scheduling policy**: Priority-based scheduling is always active when the `PriorityPlugin` is registered. No configuration flag is required to enable it; the plugin is registered by default alongside `DRFPlugin` and `BatchPlugin`. When all sessions share the same priority, the priority dimension has no effect and ordering falls through to the downstream deferral plugins. ### API @@ -179,7 +179,7 @@ session = flame.open_session( | Feature | Interaction | | ------- | ----------- | | **DRF (RFE433)** | Priority ordering overlays DRF's dominant-resource fairness ordering. Within the same priority level, DRF determines per-session resource shares and session ordering by dominant resource fraction. | -| **GangPlugin (RFE400)** | Orthogonal. `batch_size` constraints apply within a priority level independently of priority ordering. | +| **BatchPlugin (RFE400)** | Orthogonal. `batch_size` constraints apply within a priority level independently of priority ordering. | | **AllocateAction** | Consults `ssn_order_fn` from all plugins; PriorityPlugin's ordering ensures high-priority sessions are processed before low-priority ones. | | **DispatchAction** | Consults `is_underused` from all plugins; PriorityPlugin blocks dispatch to lower-priority sessions when any higher-priority session is needy. | | **ShuffleAction** | Unchanged in V1. Priority-based preemption (reclaiming executors via ShuffleAction) is deferred. | @@ -220,7 +220,7 @@ session = flame.open_session( │ │ PluginManager │ │ │ │ │ │ │ │ ┌──────────────────┐ ┌───────────────────┐ ┌──────────────────┐ │ │ -│ │ │ PriorityPlugin │ │ DRFPlugin │ │ GangPlugin │ │ │ +│ │ │ PriorityPlugin │ │ DRFPlugin │ │ BatchPlugin │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ setup(): │ │ setup(): │ │ setup(): │ │ │ │ │ │ read cluster │ │ compute │ │ track batch │ │ │ @@ -647,11 +647,11 @@ pub fn ssn_order_fn(&self, s1: &SessionInfo, s2: &SessionInfo) -> Ordering { | ----- | ------ | --------------------- | | 1 | `PriorityPlugin` | Returns `Some(ord)` when priorities differ; `None` when equal | | 2 | `DRFPlugin` | Returns `Some(ord)` based on dominant resource share | -| 3 | `GangPlugin` | Returns `None` (no ordering opinion) | +| 3 | `BatchPlugin` | Returns `None` (no ordering opinion) | This chain ensures priority is the primary sort key, with DRF providing tiebreaking within a priority level. -The `is_underused` aggregation rule is **ANY** (any plugin returning `Some(true)` makes a session underused). `PriorityPlugin` returns `Some(false)` for blocked lower-priority sessions, and `None` for sessions at or above the needy priority tier. The `None` return defers to DRF and GangPlugin, preserving all existing underuse logic for unblocked sessions. +The `is_underused` aggregation rule is **ANY** (any plugin returning `Some(true)` makes a session underused). `PriorityPlugin` returns `Some(false)` for blocked lower-priority sessions, and `None` for sessions at or above the needy priority tier. The `None` return defers to DRF and BatchPlugin, preserving all existing underuse logic for unblocked sessions. #### 6. flmctl Changes @@ -741,7 +741,7 @@ PriorityPlugin.ssn_order_fn(s1, s2): if p1 == p2 → return None (defer to the next plugin in the chain) PluginManager chain (first non-None wins): - PriorityPlugin → DRFPlugin → GangPlugin → Ordering::Equal + PriorityPlugin → DRFPlugin → BatchPlugin → Ordering::Equal ``` #### Priority-Aware Resource Distribution (PriorityPlugin.setup) @@ -821,12 +821,12 @@ For each session in is_underused(ssn): if desired != ResourceRequirement::default() and !allocated.great_equal(&desired): return Some(true) → overrides any downstream-plugin veto else: - return None → let DRF and GangPlugin decide + return None → let DRF and BatchPlugin decide Aggregation (ANY semantics): - is_underused(ssn) = PriorityPlugin OR DRFPlugin OR GangPlugin + is_underused(ssn) = PriorityPlugin OR DRFPlugin OR BatchPlugin If PriorityPlugin returns Some(false), session is not underused - regardless of what DRF/GangPlugin return. + regardless of what DRF/BatchPlugin return. ``` **Why `pending > 0` as the "needy" criterion:** @@ -969,7 +969,7 @@ The `ssn_priority` map grows linearly with the number of open sessions. Session | `common/src/apis/types.rs` | Internal | Session and SessionAttributes structs | | `session_manager/src/model/mod.rs` | Internal | SessionInfo (scheduler snapshot type) | | `session_manager/src/scheduler/plugins/mod.rs` | Internal | Plugin trait, PluginManager, registration | -| `session_manager/src/scheduler/plugins/gang.rs` | Internal | GangPlugin (unchanged; orthogonal) | +| `session_manager/src/scheduler/plugins/batch.rs` | Internal | BatchPlugin (unchanged; orthogonal) | | SQLite | External | Session storage; `ALTER TABLE` for `priority` column | --- @@ -1079,7 +1079,7 @@ PriorityPlugin: llm-inference (priority=100) is needy → max_needy_priority=100 batch-work (priority=0) is blocked AllocateAction: - llm-inference: GangPlugin requires batches of 4 executors × cpu=4 = 16 CPU + llm-inference: BatchPlugin requires batches of 4 executors × cpu=4 = 16 CPU Statement pipelines 4 executors → is_ready() true → commit llm-inference gets 16 CPU in one complete batch @@ -1113,7 +1113,7 @@ llm-inference pending=0, batch-work becomes eligible. | `session_manager/src/model/mod.rs` | SessionInfo — add `priority` field | | `session_manager/src/scheduler/plugins/ty.rs` | **NEW** PriorityPlugin implementation | | `session_manager/src/scheduler/plugins/mod.rs` | Plugin trait, PluginManager, registration order | -| `session_manager/src/scheduler/plugins/gang.rs` | GangPlugin (unchanged) | +| `session_manager/src/scheduler/plugins/batch.rs` | BatchPlugin (unchanged) | | `session_manager/src/scheduler/actions/allocate.rs` | AllocateAction (unchanged) | | `session_manager/src/scheduler/actions/dispatch.rs` | DispatchAction (unchanged) | | `session_manager/src/apiserver/frontend.rs` | Frontend — pass `priority` from proto to attributes | @@ -1128,7 +1128,7 @@ llm-inference pending=0, batch-work becomes eligible. | -------- | --------- | | **Higher `priority` value = higher priority** | Consistent with Kubernetes `PriorityClass.value` and Volcano queue priority. Allows future expansion (e.g., system sessions at priority > 1000, user session–999). | | **PriorityPlugin as a separate plugin, not a modification of the proportional-share plugin** | Keeps the proportional-share / DRF plugin focused on resource fairness. Priority is an independent scheduling dimension that composes with, rather than replaces, the downstream distribution. | -| **Plugin consultation order: Priority → DRF → Gang** | `ssn_order_fn` returns the first non-`None` result. PriorityPlugin gives a definitive answer when priorities differ; DRF breaks ties within a priority level; GangPlugin has no ordering opinion. This ensures priority is the primary sort key globally. | +| **Plugin consultation order: Priority → DRF → Batch** | `ssn_order_fn` returns the first non-`None` result. PriorityPlugin gives a definitive answer when priorities differ; DRF breaks ties within a priority level; BatchPlugin has no ordering opinion. This ensures priority is the primary sort key globally. | | **Blocking via `is_underused = Some(false)`, not preemption** | V1 focuses on controlling new resource allocation. Preemption (reclaiming executors from running low-priority sessions) requires executor lifecycle management, session state coordination, and policy decisions about partial reclaim — deferred to a future RFE. | | **"Needy" criterion: `pending > 0`** | A session with no pending tasks cannot benefit from additional executors regardless of `allocated vs. deserved`. Using task backlog directly avoids dependency on the downstream plugin's internal `deserved` computation while correctly capturing whether the session needs more resources. | | **Default `priority = 0`** | Proto3 default for `uint32` is `0`. All existing sessions automatically start at the lowest priority without any migration or explicit opt-in. Clusters that don't use the priority feature are unaffected. | diff --git a/docs/designs/RFE433-gpu-drf/FS.md b/docs/designs/RFE433-gpu-drf/FS.md index 4f79e88e..daac27a5 100644 --- a/docs/designs/RFE433-gpu-drf/FS.md +++ b/docs/designs/RFE433-gpu-drf/FS.md @@ -80,7 +80,6 @@ scheduler: policies: - priority - drf # NEW: Dominant Resource Fairness - - gang ``` **Examples:** @@ -321,7 +320,7 @@ pub struct SessionAttributes { | Feature | Interaction | |---------|-------------| | **PriorityPlugin (RFE413)** | Orthogonal; priority ordering applies to GPU sessions equally | -| **GangPlugin (RFE400)** | Orthogonal; batch constraints apply to GPU sessions | +| **BatchPlugin (RFE400)** | Orthogonal; batch constraints apply to GPU sessions | **Required Updates:** @@ -704,8 +703,8 @@ const PLUGIN_REGISTRY: &[PluginInfo] = &[ configurable: true, }, PluginInfo { - name: "gang", - constructor: GangPlugin::new_ptr, + name: "batch", + constructor: BatchPlugin::new_ptr, configurable: true, }, PluginInfo { diff --git a/docs/designs/RFE478-helm-install/FS.md b/docs/designs/RFE478-helm-install/FS.md index 883c73d7..85928461 100644 --- a/docs/designs/RFE478-helm-install/FS.md +++ b/docs/designs/RFE478-helm-install/FS.md @@ -102,7 +102,6 @@ cluster: policies: - priority - drf - - gang scheduleInterval: 100 storage: "sqlite:///var/lib/flame/session/flame.db" executors: @@ -212,7 +211,6 @@ cluster: policies: - priority - drf - - gang storage: "sqlite:///var/lib/flame/session/flame.db" schedule_interval: 100 executors: diff --git a/docs/designs/RFE498-scheduler-readiness/FS.md b/docs/designs/RFE498-scheduler-readiness/FS.md new file mode 100644 index 00000000..04275e90 --- /dev/null +++ b/docs/designs/RFE498-scheduler-readiness/FS.md @@ -0,0 +1,327 @@ +# Design: Scheduler `is_ready` / `is_fulfilled` Semantics + +**Status:** Draft + +**Author:** Klaus Ma + +**Created:** 2026-07-08 + +**Issues:** [#498](https://github.com/xflops/flame/issues/498), [#500](https://github.com/xflops/flame/pull/500), [#508](https://github.com/xflops/flame/pull/508) + +--- + +## 1. Motivation + +Flame's scheduler runs **Allocate → Dispatch → Shuffle** once per cycle. The existing APIs answer two action-specific questions: + +| API | Existing meaning retained by this design | +|-----|------------------------------------------| +| `is_fulfilled(session)` | Allocate-side fulfillment: enough executor supply exists for the session's incomplete tasks, including associated executors and reusable Void/Idle executors selected by the current Statement. | +| `is_ready(session)` | Dispatch-side readiness: enough executors are ready, bound, or committed in flight for the session's tasks. | + +[Issue #498](https://github.com/xflops/flame/issues/498) exposed duplicate binding while an executor remained in `Binding`. [PR #500](https://github.com/xflops/flame/pull/500) addressed that and two related regressions, but added separate non-batch counters and mixed rate limiting into `PluginManager`. [PR #508](https://github.com/xflops/flame/pull/508) reverted #500 so the behavior could be redesigned. + +This design targets the tree after merged #508. It keeps `is_ready` / `is_fulfilled`, the existing plugin callbacks, `Statement`, and the existing Batch state fields. It introduces **no new scheduler function and no new scheduler state field**. + +### Problems to solve + +| Problem | Symptom | Root cause | +|---------|---------|------------| +| Batch under-provision | Only one executor is created for multiple incomplete tasks | Modulo readiness becomes true at every batch multiple instead of at the session's demand target. | +| Duplicate dispatch/binding | Extra executors are selected while one is already in `Binding` | Dispatch has no ready pre-check, and old readiness requires a new bind in the current cycle. | +| Duplicate allocation | An additional batch can be created after demand is already met | Old fulfillment requires a new pipeline/allocation operation before it can be true. | +| Cap overshoot | One batch allocation pass can cross `max_instances` | The action checks the cap only before entering the allocation loop. | + +### Objectives + +1. Fix #498 and the batch demand regressions described by #500 while treating an all-abstaining policy set as already satisfied. +2. Preserve `Context::is_ready` for Dispatch and `Context::is_fulfilled` for Allocate. +3. Keep `Option` inside plugins, while PluginManager and Context return a concrete `bool` after ignoring abstaining plugins. +4. Reuse only `BatchState.batch_size`, `allocated`, `pipelined`, and `bound`; derive task demand and caps from the supplied `SessionInfo`. +5. Make Batch always enabled so every valid policy stack has a readiness/fulfillment owner, while preserving explicit `batch` and `shim` entries as harmless compatibility no-ops. +6. Restore #500's valid non-scheduler changes after #508: Priority as the default configurable policy, DRF opt-in, and the process-wide `SHIMS_TEST_LOCK`. + +--- + +## 2. Function Specification + +### Configuration + +No new configuration is introduced. + +| Field | Role | +|-------|------| +| `cluster.policies` | Selects ordering/fairness policies such as Priority or DRF. Batch is always enabled and owns absolute batch readiness. | +| `session.batch_size` | Batch batch width, normalized with `max(1)`. | +| `session.max_instances` | Hard upper bound for batch allocation demand. | + +On the #508 baseline, restore `priority` as the default configurable policy. DRF remains opt-in. Batch and Shim are always enabled and therefore omitted from defaults and examples; explicitly listing either is accepted and ignored, while unknown policies remain errors. The effective default stack is Priority + Batch + Shim. + +### Existing API contract + +No new method is added. The existing methods retain their names and delegate to `PluginManager`: + +```rust +impl Context { + pub fn is_ready(&self, ssn: &SessionInfoPtr) -> Result; + + pub fn is_fulfilled(&self, ssn: &SessionInfoPtr) -> Result; +} +``` + +The existing `Plugin::is_ready` / `Plugin::is_fulfilled` methods remain optional. PluginManager ignores `None` and combines only concrete plugin opinions: + +| Plugin opinions | Context result | +|-----------------|----------------| +| One or more concrete opinions | `true` only when every concrete opinion is `true`; `None` opinions are ignored. | +| Every plugin returns `None` | `true`; no configured plugin requires more work from the action. | + +This keeps abstention inside the Plugin API while exposing a simple boolean contract to Context and actions. + +### Action contract + +**AllocateAction** + +1. Before creating a `Statement`, call `is_fulfilled(session)` and skip when it returns true. +2. Consider existing Void and Idle executors before creating a new executor. Record selected reusable executors through the existing pipeline operation so Batch's existing `pipelined` field reflects them. +3. After each operation: + - stop when `is_fulfilled` returns true; + - continue when it returns false. +4. Commit a non-empty statement when fulfillment is true; discard a non-empty false statement as an incomplete batch batch. + +**DispatchAction** + +1. Runs after Allocate in the same `Context`. +2. Before creating a `Statement`, call `is_ready(session)` and skip when it returns true. +3. Bind available idle executors. +4. After each bind: + - stop when `is_ready` returns true; + - continue when it returns false. +5. Commit a non-empty statement when readiness is true; discard a non-empty false statement as an incomplete batch batch. + +**ShuffleAction** runs last and keeps its existing behavior. + +All three actions share one `Context` and PluginManager for the cycle. Allocate's speculative `pipelined` updates remain visible to later actions. Reusable Idle executors selected by Allocate remain eligible for Dispatch in the same cycle. Newly created executors are still Void and become dispatchable only after registration in a later cycle. + +An Idle pipeline reservation updates Batch fulfillment only. Priority and DRF do not charge an Idle executor during Allocate because it already exists and consumes node resources; they assign its session share when Dispatch actually binds it. This prevents the reservation from making Dispatch incorrectly consider the session no longer underused. + +### Scope + +**In scope** + +- Keep optional opinions inside plugins and return resolved booleans from PluginManager and Context +- Set the cycle order to Allocate → Dispatch → Shuffle +- Update Allocate and Dispatch to use the resolved boolean checks without a progress argument +- Correct Batch formulas using existing state fields plus `SessionInfo` +- Include reusable Void/Idle executors in Allocate's fulfillment calculation through existing Statement pipeline operations +- Add Dispatch's pre-session `is_ready` guard +- Restore the valid configuration and shared shim-test-lock changes removed by #508 +- Add regression coverage for #498 and every scheduler configuration described by #500 + +**Out of scope** + +- New demand trackers, counters, helper methods, or state fields +- Priority and DRF policy algorithms +- `SessionInfo::is_ready(retry_limits)`, which is a separate bind-failure recovery gate + +--- + +## 3. Batch Semantics Using Existing Fields + +`BatchState` remains exactly: + +```rust +struct BatchState { + batch_size: u32, + allocated: u32, + pipelined: u32, + bound: u32, +} +``` + +`setup()` continues to initialize `batch_size`, count snapshot executors with the session ID into `allocated`, and reset `pipelined` / `bound` to zero. Existing Statement callbacks continue to increment and roll back `pipelined` and `bound`. + +### Derived target + +Both predicates derive the target directly from the supplied `SessionInfo`; no target is cached: + +```text +incomplete_tasks = Pending + Running +batch_size = max(session.batch_size, 1) +uncapped = div_ceil(incomplete_tasks, batch_size) * batch_size +aligned_max = floor(max_instances / batch_size) * batch_size, or the largest u32 batch multiple +needed = min(uncapped, aligned_max) +``` + +Perform summation and rounding in `u64`, apply the cap, and convert the final value to `u32`. The API already rejects non-batch-aligned limits; defensive alignment keeps legacy persisted sessions within both the cap and the full-batch invariant. + +### `is_fulfilled` + +```text +total = allocated + pipelined +is_fulfilled = needed == 0 || total >= needed +``` + +Allocate records reusable Void/Idle executors and new allocations through the existing pipeline/allocation callbacks, so they contribute through `pipelined`. This fixes batch under-provisioning, duplicate allocation, and `max_instances` overshoot without storing `incomplete_tasks`, `needed`, or `max_instances` in `BatchState`. + +### `is_ready` + +```text +total = allocated + bound +is_ready = needed == 0 || total >= needed +``` + +Dispatch uses this predicate before and during task/executor dispatch. An executor in `Binding` has a session ID and is included in `allocated`; it counts as committed in flight even though it is not yet runnable, preventing another dispatch/bind while `on_session_enter` is running. Idle executors selected by a speculative bind contribute through `bound`. + +### Rollback + +The existing callbacks remain the only speculative-state mechanism: + +```text +allocate/pipeline → pipelined += 1 +discard → pipelined -= 1 +bind → bound += 1 +unbind/discard → bound -= 1 +``` + +No state ownership or lifecycle is added. + +--- + +## 4. Use Cases + +### UC1: #498 executor remains in Binding + +Batch enabled, `batch_size=1`, one incomplete task, and one snapshot executor in `Binding`: + +```text +needed = 1 +allocated = 1 +pipelined = 0 +bound = 0 + +is_fulfilled = 1 >= 1 → true +is_ready = 1 >= 1 → true +``` + +Allocate runs first and skips because `is_fulfilled` is satisfied by the snapshot executor. Dispatch then skips because `is_ready` is satisfied. Repeating the cycle does not change executor or bind counts. + +### UC2: Priority-only session + +Priority does not own readiness or fulfillment, so it returns `None` for both. Batch is loaded automatically even when the configured policy list contains only Priority, supplies the concrete demand decision, and allows Allocate/Dispatch to make progress. + +### UC3: Batch batch_size=1 with three tasks + +```text +needed = div_ceil(3,1)*1 = 3 +``` + +Allocate continues until `allocated + pipelined == 3` and `is_fulfilled` returns true, then commits. + +### UC4: Batch batch_size=2 with insufficient resources + +With `needed=2`, one speculative operation leaves `is_fulfilled` false. The non-empty incomplete Statement is discarded and the existing callback rolls back `pipelined`. + +### UC5: Batch demand exceeds max_instances + +For `batch_size=2`, eight incomplete tasks, and `max_instances=4`: + +```text +uncapped = 8 +aligned_max = 4 +needed = 4 +``` + +Allocate stops at four total associated/speculative executors. The fulfillment target itself prevents a fifth operation inside the loop. + +### UC6: Dispatch with reusable Idle supply + +With Batch enabled, an idle executor has no session ID. Allocate first includes it as reusable supply rather than creating another executor. Batch's fulfillment becomes true, and Dispatch then records one bind. Batch's readiness becomes true, so the Statement commits and the executor transitions from `Idle` to `Binding`. + +--- + +## 5. Test Plan + +### Batch unit tests + +- Allocate-side `is_fulfilled` uses `allocated + pipelined` and reaches true at the executor-supply target +- Dispatch-side `is_ready` uses `allocated + bound` and reaches true at the dispatch target +- Existing snapshot executors can satisfy both predicates without a new operation +- Reusable Void and Idle executors selected by Allocate contribute through the existing `pipelined` field +- Reserving an Idle executor does not advance Priority/DRF session allocation until Dispatch binds it +- `batch_size=1` with three incomplete tasks requires three total executors +- Partial `batch_size=2` Statements remain unsatisfied and roll back +- Demand above an aligned `max_instances` stops exactly at the cap +- Large task counts use overflow-safe target arithmetic +- Existing allocate/pipeline/bind rollback tests remain valid with no new state fields + +### PluginManager tests + +- No plugin opinion returns true +- One `Some(false)` opinion returns false +- All concrete opinions true returns true +- Mixed true and false opinions return false + +### Scheduler integration tests + +- DRF-only and priority-only configured policy lists still allocate because Batch is always enabled +- Allocate does not create an executor when sufficient compatible Void/Idle supply already exists +- Batch `batch_size=1` allocates one executor per incomplete task +- Batch `batch_size=2` commits only complete batches +- Batch allocation never exceeds `max_instances` +- Two independent scheduler cycles with an executor held in `Binding` keep both executor count and bind-call count at one +- `Context.actions` executes Allocate, Dispatch, then Shuffle in that exact order +- Allocate, Dispatch, and Shuffle share the same PluginManager so Statement callbacks remain visible across actions + +### Configuration and shim tests + +- Default context and generated `flmadm` configuration list only `priority`; Batch and Shim are implicit +- Batch is loaded when omitted from the configured policy list +- Explicit `batch` and `shim` entries are accepted as no-ops because both plugins are always enabled +- Explicit DRF remains accepted +- Unknown policies still fail +- Every shim test that mutates process-global environment or working directory uses the same `SHIMS_TEST_LOCK` + +--- + +## 6. Migration / Rollout + +1. Start from merged #508, which establishes the pre-#500 implementation baseline. +2. Restore the shared `SHIMS_TEST_LOCK`, Priority defaults, and DRF opt-in configuration. +3. Make Batch and Shim always enabled while accepting explicit entries for compatibility. +4. Keep optional opinions in the Plugin trait; resolve them to boolean in existing PluginManager and Context methods, returning true when every plugin abstains. +5. Set and test the action sequence as Allocate → Dispatch → Shuffle, with one shared Context and PluginManager. +6. Update Allocate to use boolean `is_fulfilled` and include Void/Idle supply. +7. Update Dispatch to use boolean `is_ready`. +8. Replace Batch's modulo checks with the derived, max-capped absolute target using only existing fields. +9. Close #498 only after the two-cycle Binding-stall integration test passes. + +**Compatibility:** No scheduler method or state field is added. Context returns boolean readiness/fulfillment without a progress argument; an all-abstaining plugin set resolves to true. Batch is always loaded, so valid sessions receive a concrete readiness opinion without listing `batch` in `cluster.policies`; an explicit `batch` entry remains accepted. The old `gang` policy name is removed rather than retained as an alias. The action order changes from the post-#508 baseline's Dispatch → Allocate → Shuffle to Allocate → Dispatch → Shuffle. An executor created by Allocate is not eligible for Dispatch until it registers and appears as Idle in a later cycle. Clusters that require DRF must add `drf` explicitly because the default configurable policy is Priority. + +**Rollback:** Revert this replacement on top of #508. This restores the pre-#500 optional plugin API and modulo behavior without a data migration. + +--- + +## 7. Files and References + +| File | Change | +|------|--------| +| `session_manager/src/scheduler/plugins/batch.rs` | Change formulas only; keep the existing struct fields and callbacks. | +| `session_manager/src/scheduler/plugins/mod.rs` | Always load Batch, accept explicit always-on plugin entries, ignore abstaining opinions, and resolve the all-`None` case to true; add no counter or field. | +| `session_manager/src/scheduler/plugins/priority.rs`, `drf.rs` | Defer Idle-executor session accounting until Dispatch binds it. | +| `session_manager/src/scheduler/ctx.rs` | Return resolved booleans from the existing methods. | +| `session_manager/src/scheduler/actions/allocate.rs` | Use `is_fulfilled` and consume reusable Void/Idle supply first. | +| `session_manager/src/scheduler/actions/dispatch.rs` | Use `is_ready` and add the pre-check. | +| `common/src/ctx.rs`, `flmadm/src/managers/config.rs` | List only Priority by default; Batch and Shim remain implicit and DRF stays opt-in. | +| `executor_manager/src/shims/` | Restore the process-wide test lock. | + +**Related designs** + +- [RFE400 Batch Session](../RFE400-batch-session/FS.md) +- [RFE413 Priority Scheduling](../RFE413-priority-scheduling/FS.md) +- [RFE408 Fairshare batch_size](../RFE408-fairshare-batch-size/FS.md) + +**External** + +- [PR #500](https://github.com/xflops/flame/pull/500) — prior fix being replaced +- [PR #508](https://github.com/xflops/flame/pull/508) — merged revert and implementation baseline diff --git a/docs/tutorials/local-development.md b/docs/tutorials/local-development.md index 08275187..471cbf93 100644 --- a/docs/tutorials/local-development.md +++ b/docs/tutorials/local-development.md @@ -181,7 +181,6 @@ cluster: policies: - priority - drf - - gang storage: "fs:///tmp/flame-dev/data" executors: shim: host diff --git a/e2e/tests/test_session.py b/e2e/tests/test_session.py index 2175455f..a9a72c08 100644 --- a/e2e/tests/test_session.py +++ b/e2e/tests/test_session.py @@ -208,7 +208,7 @@ def test_session_with_min_max_instances(self): session.close() def test_session_batch_size(self): - """Test creating session with batch_size for gang scheduling.""" + """Test creating session with batch_size for batch scheduling.""" session_id = f"test-batch-{random_string(8)}" session = flamepy.create_session( application=FLM_TEST_SVC_APP, diff --git a/executor_manager/src/shims/grpc_shim.rs b/executor_manager/src/shims/grpc_shim.rs index f7334692..b27aed59 100644 --- a/executor_manager/src/shims/grpc_shim.rs +++ b/executor_manager/src/shims/grpc_shim.rs @@ -202,15 +202,13 @@ impl Future for WaitForSvcSocketFuture { #[cfg(test)] mod tests { + use super::super::SHIMS_TEST_LOCK; use super::*; use common::apis::{ApplicationContext, Shim as ShimType}; use std::collections::HashMap; use std::path::PathBuf; - use std::sync::Mutex; use tempfile::tempdir; - static TEST_LOCK: Mutex<()> = Mutex::new(()); - fn setup_test_env(temp: &tempfile::TempDir) -> PathBuf { let socket_dir = temp.path().join("sockets"); std::fs::create_dir_all(&socket_dir).unwrap(); @@ -237,7 +235,7 @@ mod tests { #[test] fn test_grpc_shim_new() { - let _guard = TEST_LOCK.lock().unwrap(); + let _guard = SHIMS_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let temp = tempdir().unwrap(); setup_test_env(&temp); let work_dir = create_test_work_dir("exec-grpc-test", &temp); @@ -251,7 +249,7 @@ mod tests { #[test] fn test_grpc_shim_endpoint() { - let _guard = TEST_LOCK.lock().unwrap(); + let _guard = SHIMS_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let temp = tempdir().unwrap(); setup_test_env(&temp); let work_dir = create_test_work_dir("exec-endpoint-test", &temp); @@ -263,7 +261,7 @@ mod tests { #[test] fn test_grpc_shim_close_without_connection() { - let _guard = TEST_LOCK.lock().unwrap(); + let _guard = SHIMS_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let temp = tempdir().unwrap(); setup_test_env(&temp); let work_dir = create_test_work_dir("exec-close-test", &temp); @@ -285,7 +283,7 @@ mod tests { #[tokio::test] #[allow(clippy::await_holding_lock)] async fn test_on_session_enter_without_connection() { - let _guard = TEST_LOCK.lock().unwrap(); + let _guard = SHIMS_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let temp = tempdir().unwrap(); setup_test_env(&temp); let work_dir = create_test_work_dir("exec-session-test", &temp); @@ -318,7 +316,7 @@ mod tests { #[tokio::test] #[allow(clippy::await_holding_lock)] async fn test_on_task_invoke_without_connection() { - let _guard = TEST_LOCK.lock().unwrap(); + let _guard = SHIMS_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let temp = tempdir().unwrap(); setup_test_env(&temp); let work_dir = create_test_work_dir("exec-task-test", &temp); @@ -341,7 +339,7 @@ mod tests { #[tokio::test] #[allow(clippy::await_holding_lock)] async fn test_on_session_leave_without_connection() { - let _guard = TEST_LOCK.lock().unwrap(); + let _guard = SHIMS_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let temp = tempdir().unwrap(); setup_test_env(&temp); let work_dir = create_test_work_dir("exec-leave-test", &temp); diff --git a/executor_manager/src/shims/mod.rs b/executor_manager/src/shims/mod.rs index 9938bf2a..eeb99d87 100644 --- a/executor_manager/src/shims/mod.rs +++ b/executor_manager/src/shims/mod.rs @@ -225,16 +225,17 @@ pub trait Shim: Send + 'static { async fn on_session_leave(&mut self) -> Result<(), FlameError>; } +/// Process-wide lock for shim tests that mutate environment variables and the working directory. +#[cfg(test)] +pub(crate) static SHIMS_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + #[cfg(test)] mod tests { use super::*; use std::collections::HashMap; use std::fs::File; - use std::sync::Mutex; use tempfile::tempdir; - static TEST_LOCK: Mutex<()> = Mutex::new(()); - fn create_test_app(name: &str, working_directory: Option) -> ApplicationContext { ApplicationContext { name: name.to_string(), @@ -259,7 +260,7 @@ mod tests { #[test] fn test_executor_work_dir_with_auto_dir() { - let _guard = TEST_LOCK.lock().unwrap(); + let _guard = SHIMS_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let temp = tempdir().unwrap(); let socket_dir = setup_test_env(&temp); @@ -280,7 +281,7 @@ mod tests { #[test] fn test_executor_work_dir_with_custom_working_directory() { - let _guard = TEST_LOCK.lock().unwrap(); + let _guard = SHIMS_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let temp = tempdir().unwrap(); let socket_dir = setup_test_env(&temp); let custom_dir = temp.path().join("custom-workdir"); @@ -301,7 +302,7 @@ mod tests { #[test] fn test_executor_work_dir_socket_path_length() { - let _guard = TEST_LOCK.lock().unwrap(); + let _guard = SHIMS_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let temp = tempdir().unwrap(); setup_test_env(&temp); @@ -321,7 +322,7 @@ mod tests { #[test] fn test_executor_work_dir_cleanup_on_drop() { - let _guard = TEST_LOCK.lock().unwrap(); + let _guard = SHIMS_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let temp = tempdir().unwrap(); setup_test_env(&temp); @@ -355,7 +356,7 @@ mod tests { #[test] fn test_executor_work_dir_no_cleanup_custom_dir_on_drop() { - let _guard = TEST_LOCK.lock().unwrap(); + let _guard = SHIMS_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let temp = tempdir().unwrap(); setup_test_env(&temp); let custom_dir = temp.path().join("persistent-workdir"); @@ -382,7 +383,7 @@ mod tests { #[test] fn test_socket_path_is_fixed_location() { - let _guard = TEST_LOCK.lock().unwrap(); + let _guard = SHIMS_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let temp = tempdir().unwrap(); let socket_dir = setup_test_env(&temp); diff --git a/flmadm/src/managers/config.rs b/flmadm/src/managers/config.rs index f3699149..752fdf61 100644 --- a/flmadm/src/managers/config.rs +++ b/flmadm/src/managers/config.rs @@ -44,8 +44,6 @@ cluster: resreq: "cpu=1,mem=2g" policies: - priority - - drf - - gang storage: "fs://{prefix}/data" executors: shim: host diff --git a/flmctl/src/main.rs b/flmctl/src/main.rs index c3962c9b..e4f52fac 100644 --- a/flmctl/src/main.rs +++ b/flmctl/src/main.rs @@ -101,7 +101,7 @@ enum Commands { /// The name of Application #[arg(short, long)] app: String, - /// Number of executors per batch for gang scheduling + /// Number of executors per batch for batch scheduling #[arg(short, long, default_value = "1")] batch_size: u32, /// Session priority (higher = more important, default: 0) diff --git a/rpc/protos/types.proto b/rpc/protos/types.proto index 28b126ee..6c92a3de 100644 --- a/rpc/protos/types.proto +++ b/rpc/protos/types.proto @@ -37,7 +37,7 @@ message SessionSpec { optional bytes common_data = 4; uint32 min_instances = 5; // Minimum number of instances (default: 0) optional uint32 max_instances = 6; // Maximum number of instances (null means unlimited) - uint32 batch_size = 7; // Number of executors per batch for gang scheduling (default: 1) + uint32 batch_size = 7; // Number of executors per batch for batch scheduling (default: 1) uint32 priority = 8; // Session priority (default: 0 = lowest) optional ResourceRequirement resreq = 9; // Explicit resource request; if absent, server applies cluster default } diff --git a/sdk/python/protos/types.proto b/sdk/python/protos/types.proto index 28b126ee..6c92a3de 100644 --- a/sdk/python/protos/types.proto +++ b/sdk/python/protos/types.proto @@ -37,7 +37,7 @@ message SessionSpec { optional bytes common_data = 4; uint32 min_instances = 5; // Minimum number of instances (default: 0) optional uint32 max_instances = 6; // Maximum number of instances (null means unlimited) - uint32 batch_size = 7; // Number of executors per batch for gang scheduling (default: 1) + uint32 batch_size = 7; // Number of executors per batch for batch scheduling (default: 1) uint32 priority = 8; // Session priority (default: 0 = lowest) optional ResourceRequirement resreq = 9; // Explicit resource request; if absent, server applies cluster default } diff --git a/sdk/python/src/flamepy/core/client.py b/sdk/python/src/flamepy/core/client.py index 2e2ff527..4b1a8d98 100644 --- a/sdk/python/src/flamepy/core/client.py +++ b/sdk/python/src/flamepy/core/client.py @@ -149,7 +149,7 @@ def create_session( session_id: Optional session ID min_instances: Minimum number of instances (default: 0) max_instances: Maximum number of instances (None = unlimited) - batch_size: Number of executors per batch for gang scheduling (default: 1) + batch_size: Number of executors per batch for batch scheduling (default: 1) resreq: Optional explicit resource requirements. When omitted, the server applies cluster.resource_requirement (or a hardcoded fallback when that is unset). diff --git a/sdk/rust/protos/types.proto b/sdk/rust/protos/types.proto index 28b126ee..6c92a3de 100644 --- a/sdk/rust/protos/types.proto +++ b/sdk/rust/protos/types.proto @@ -37,7 +37,7 @@ message SessionSpec { optional bytes common_data = 4; uint32 min_instances = 5; // Minimum number of instances (default: 0) optional uint32 max_instances = 6; // Maximum number of instances (null means unlimited) - uint32 batch_size = 7; // Number of executors per batch for gang scheduling (default: 1) + uint32 batch_size = 7; // Number of executors per batch for batch scheduling (default: 1) uint32 priority = 8; // Session priority (default: 0 = lowest) optional ResourceRequirement resreq = 9; // Explicit resource request; if absent, server applies cluster default } diff --git a/sdk/rust/tests/benchmark_test.rs b/sdk/rust/tests/benchmark_test.rs index a6dd47b5..248160c7 100644 --- a/sdk/rust/tests/benchmark_test.rs +++ b/sdk/rust/tests/benchmark_test.rs @@ -16,7 +16,7 @@ limitations under the License. //! This test creates 10 concurrent sessions, each running 1000 tasks, //! for a total of 10,000 tasks. The benchmark must complete within 10 minutes. //! -//! Also includes gang scheduling tests to verify tasks don't start partially +//! Also includes batch scheduling tests to verify tasks don't start partially //! when batch_size > 1 (all executors in a batch must be ready before tasks start). use std::sync::atomic::{AtomicU64, Ordering}; @@ -199,11 +199,11 @@ const BATCH_SIZE: u32 = 2; const BATCH_TIMEOUT_SECS: u64 = 120; #[tokio::test] -async fn benchmark_gang_scheduling_no_partial_start() -> Result<(), FlameError> { +async fn benchmark_batch_scheduling_no_partial_start() -> Result<(), FlameError> { tracing_subscriber::fmt::try_init().ok(); println!("\n============================================================"); - println!("GANG SCHEDULING TEST: batch_size={}", BATCH_SIZE); + println!("BATCH SCHEDULING TEST: batch_size={}", BATCH_SIZE); println!("Verifying single task stays Pending until batch is complete..."); println!("============================================================\n"); @@ -213,7 +213,7 @@ async fn benchmark_gang_scheduling_no_partial_start() -> Result<(), FlameError> let conn = flame::client::connect_with_tls(FLAME_ADDR, Some(&tls_config)).await?; let ssn_attr = SessionAttributes { - id: format!("gang-test-{}", stdng::rand::short_name()), + id: format!("batch-test-{}", stdng::rand::short_name()), application: FLAME_APP.to_string(), common_data: None, min_instances: 0, @@ -293,7 +293,7 @@ async fn benchmark_gang_scheduling_no_partial_start() -> Result<(), FlameError> let duration = start.elapsed(); println!("\n============================================================"); - println!("GANG SCHEDULING RESULTS"); + println!("BATCH SCHEDULING RESULTS"); println!("============================================================"); println!("Duration: {:.2}s", duration.as_secs_f64()); println!("Single task correctly stayed Pending until batch filled"); @@ -302,7 +302,7 @@ async fn benchmark_gang_scheduling_no_partial_start() -> Result<(), FlameError> assert!( duration < Duration::from_secs(BATCH_TIMEOUT_SECS), - "Gang scheduling test exceeded {} second timeout: {:.2}s", + "Batch scheduling test exceeded {} second timeout: {:.2}s", BATCH_TIMEOUT_SECS, duration.as_secs_f64() ); diff --git a/session_manager/migrations/sqlite/20260413000000_add_batch_fields.sql b/session_manager/migrations/sqlite/20260413000000_add_batch_fields.sql index d36c381f..ab0d917f 100644 --- a/session_manager/migrations/sqlite/20260413000000_add_batch_fields.sql +++ b/session_manager/migrations/sqlite/20260413000000_add_batch_fields.sql @@ -1,5 +1,5 @@ -- Add batch scheduling fields (RFE400-batch-session) --- batch_size: number of executors per batch for gang scheduling +-- batch_size: number of executors per batch for batch scheduling -- batch_index: executor's index within its batch (0 to batch_size-1) ALTER TABLE sessions ADD COLUMN batch_size INTEGER NOT NULL DEFAULT 1; diff --git a/session_manager/src/controller/executors/mod.rs b/session_manager/src/controller/executors/mod.rs index a6caef00..cb5e0aa4 100644 --- a/session_manager/src/controller/executors/mod.rs +++ b/session_manager/src/controller/executors/mod.rs @@ -129,7 +129,7 @@ mod tests { endpoint: "http://localhost:8080".to_string(), storage: url, resreq: None, - policies: vec!["priority".to_string(), "gang".to_string()], + policies: vec!["priority".to_string()], schedule_interval: 1000, executors: FlameExecutors { shim: Shim::default(), diff --git a/session_manager/src/controller/mod.rs b/session_manager/src/controller/mod.rs index 1b9e7594..a2242700 100644 --- a/session_manager/src/controller/mod.rs +++ b/session_manager/src/controller/mod.rs @@ -957,7 +957,7 @@ mod tests { endpoint: "http://localhost:8080".to_string(), storage: url, resreq: None, - policies: vec!["priority".to_string(), "gang".to_string()], + policies: vec!["priority".to_string()], schedule_interval: 1000, executors: FlameExecutors { shim: Shim::default(), diff --git a/session_manager/src/controller/nodes/mod.rs b/session_manager/src/controller/nodes/mod.rs index 7ee23b88..c6494d92 100644 --- a/session_manager/src/controller/nodes/mod.rs +++ b/session_manager/src/controller/nodes/mod.rs @@ -123,7 +123,7 @@ mod tests { endpoint: "http://localhost:8080".to_string(), storage: url, resreq: None, - policies: vec!["priority".to_string(), "gang".to_string()], + policies: vec!["priority".to_string()], schedule_interval: 1000, executors: FlameExecutors { shim: Shim::default(), diff --git a/session_manager/src/scheduler/actions/allocate.rs b/session_manager/src/scheduler/actions/allocate.rs index a0375ad8..88dbd724 100644 --- a/session_manager/src/scheduler/actions/allocate.rs +++ b/session_manager/src/scheduler/actions/allocate.rs @@ -16,7 +16,7 @@ use std::sync::Arc; use stdng::collections::{BinaryHeap, Cmp}; use stdng::{logs::TraceFn, trace_fn}; -use crate::model::{ALL_NODE, READY_SESSION, UNBINDING_EXECUTOR, VOID_EXECUTOR}; +use crate::model::{ALL_NODE, IDLE_EXECUTOR, READY_SESSION, UNBINDING_EXECUTOR, VOID_EXECUTOR}; use crate::scheduler::actions::{Action, ActionPtr}; use crate::scheduler::plugins::node_order_fn; use crate::scheduler::plugins::ssn_order_fn; @@ -61,11 +61,13 @@ impl Action for AllocateAction { let mut void_executors = ss.find_executors(VOID_EXECUTOR)?; let mut unbinding_executors = ss.find_executors(UNBINDING_EXECUTOR)?; + let mut idle_executors = ss.find_executors(IDLE_EXECUTOR)?; tracing::debug!( - "AllocateAction: {} void executors, {} unbinding executors", + "AllocateAction: {} void executors, {} unbinding executors, {} idle executors", void_executors.len(), - unbinding_executors.len() + unbinding_executors.len(), + idle_executors.len() ); let node_order_fn = node_order_fn(ctx); @@ -114,55 +116,54 @@ impl Action for AllocateAction { } } - // Dispatch runs first in this cycle. `PluginManager` is not re-setup between actions - // (see `Context`): Gang's counters include binds Dispatch committed via - // `Statement` in this same `Context`. If those already satisfy gang scheduling, do not - // create more executors here. - let fulfilled = ctx.is_fulfilled(&ssn)?; - let ready = ctx.is_ready(&ssn)?; - if fulfilled || ready { + if ctx.is_fulfilled(&ssn)? { tracing::debug!( - "Skip allocate resources for session <{}>: is_fulfilled={}, is_ready={}", + "Skip allocate resources for session <{}>: is_fulfilled=true", ssn.id, - fulfilled, - ready ); continue; } let mut stmt = Statement::new(ss.clone(), ctx.plugins.clone(), ctx.controller.clone()); - let pipelineable = void_executors + let pipelineable: Vec<_> = void_executors .values() .chain(unbinding_executors.values()) - .filter(|e| ctx.is_available(e, &ssn).unwrap_or(false)); + .chain(idle_executors.values()) + .filter(|e| ctx.is_available(e, &ssn).unwrap_or(false)) + .cloned() + .collect(); - for exec in pipelineable { + for exec in &pipelineable { stmt.pipeline(exec, &ssn)?; - if ctx.is_ready(&ssn)? { + if ctx.is_fulfilled(&ssn)? { break; } } - for node in nodes.iter() { - if ctx.is_ready(&ssn)? { - break; - } - while ctx.is_allocatable(node, &ssn)? && !ctx.is_ready(&ssn)? { + 'nodes: for node in nodes.iter() { + loop { + if ctx.is_fulfilled(&ssn)? { + break 'nodes; + } + if !ctx.is_allocatable(node, &ssn)? { + break; + } stmt.allocate(node, &ssn)?; } } - if ctx.is_ready(&ssn)? { + let fulfilled = ctx.is_fulfilled(&ssn)?; + if !stmt.is_empty() && fulfilled { let op_count = stmt.len(); let pipelined_ids = stmt.commit().await?; for id in &pipelined_ids { void_executors.remove(id); unbinding_executors.remove(id); + idle_executors.remove(id); } tracing::debug!("Committed {} op(s) for session <{}>", op_count, ssn.id); nodes.sort_by(|a, b| node_order_fn.cmp(a, b)); - open_ssns.push(ssn.clone()); } else if !stmt.is_empty() { tracing::debug!( "Discarding incomplete batch for session <{}>: not enough resources", diff --git a/session_manager/src/scheduler/actions/dispatch.rs b/session_manager/src/scheduler/actions/dispatch.rs index e7b11f60..e8aa6bb7 100644 --- a/session_manager/src/scheduler/actions/dispatch.rs +++ b/session_manager/src/scheduler/actions/dispatch.rs @@ -65,6 +65,11 @@ impl Action for DispatchAction { continue; } + if ctx.is_ready(&ssn)? { + tracing::debug!("Session <{}> is already ready, skip dispatch.", ssn.id); + continue; + } + tracing::debug!( "Session <{}> is underused, start to allocate resources.", &ssn.id @@ -75,26 +80,19 @@ impl Action for DispatchAction { for (_, exec) in idle_executors.iter() { if ctx.is_available(exec, &ssn)? { stmt.bind(exec, &ssn)?; - if ctx.is_fulfilled(&ssn)? { + if ctx.is_ready(&ssn)? { break; } } } - if ctx.is_fulfilled(&ssn)? { + let ready = ctx.is_ready(&ssn)?; + if !stmt.is_empty() && ready { tracing::debug!("Bind executor for session <{}>.", ssn.id); let bound_ids = stmt.commit().await?; for id in &bound_ids { idle_executors.remove(id); } - - // Re-queue only if we actually bound executors; otherwise the session - // would loop infinitely when is_underused() keeps returning true but - // no idle executors remain. - if !bound_ids.is_empty() { - open_ssns.push(ssn); - } - continue; } else if !stmt.is_empty() { tracing::debug!( "Discarding unfulfilled binding for session <{}>: no available idle executors", diff --git a/session_manager/src/scheduler/ctx.rs b/session_manager/src/scheduler/ctx.rs index a91dd563..38c0b662 100644 --- a/session_manager/src/scheduler/ctx.rs +++ b/session_manager/src/scheduler/ctx.rs @@ -24,7 +24,7 @@ use common::apis::ExecutorState; use common::FlameError; /// One scheduling cycle: a single `Context` (one [`PluginManager::setup`] on the current -/// snapshot) is shared by Dispatch → Allocate → Shuffle. In-memory plugin counters (e.g. Gang) +/// snapshot) is shared by Allocate → Dispatch → Shuffle. In-memory plugin counters (e.g. Batch) /// accumulate across those actions; do not re-run `setup` between them. pub struct Context { pub snapshot: SnapShotPtr, @@ -43,8 +43,8 @@ impl Context { plugins, controller, actions: vec![ - DispatchAction::new_ptr(), AllocateAction::new_ptr(), + DispatchAction::new_ptr(), ShuffleAction::new_ptr(), ], }) @@ -74,15 +74,14 @@ impl Context { self.plugins.is_available(exec, ssn) } - /// Allocation-side batch readiness (e.g. Gang: pipelined + allocated executors form full - /// batches). Reflects in-memory plugin state, including `Statement` ops earlier in this - /// same cycle (typically pipeline/allocate before commit). + /// Dispatch-side readiness: enough executors are associated with the session or selected by + /// speculative bind operations. If no plugin has an opinion, the session is considered ready. pub fn is_ready(&self, ssn: &SessionInfoPtr) -> Result { self.plugins.is_ready(ssn) } - /// Binding-side batch readiness (e.g. Gang: bound + on-session executors form full batches). - /// After Dispatch commits binds, this can be true so Allocate skips provisioning. + /// Allocate-side fulfillment: enough associated, reusable, or speculative executors exist + /// for the session. If no plugin has an opinion, the session is considered fulfilled. pub fn is_fulfilled(&self, ssn: &SessionInfoPtr) -> Result { self.plugins.is_fulfilled(ssn) } diff --git a/session_manager/src/scheduler/mod.rs b/session_manager/src/scheduler/mod.rs index a10e7e18..c8dbd7bc 100644 --- a/session_manager/src/scheduler/mod.rs +++ b/session_manager/src/scheduler/mod.rs @@ -50,7 +50,7 @@ impl FlameThread for ScheduleRunner { let mut ctx = Context::new(self.controller.clone(), policies)?; // Same `ctx` (and thus same in-memory `plugins`) for every action: Dispatch mutations - // are visible to Allocate (e.g. Gang `is_fulfilled` / `is_ready` after binds). + // are visible to Allocate (e.g. Batch `is_fulfilled` / `is_ready` after binds). for action in ctx.actions.clone() { if let Err(e) = action.execute(&mut ctx).await { tracing::error!("Failed to run scheduling: {e}"); @@ -125,6 +125,41 @@ mod tests { } } + fn setup_test_session( + controller: &ControllerPtr, + ssn_id: &str, + batch_size: u32, + max_instances: Option, + task_count: usize, + ) -> Result<(), FlameError> { + tokio_test::block_on( + controller.register_application("flmtest".to_string(), new_test_application()), + )?; + tokio_test::block_on( + controller + .storage() + .register_node(&new_test_node("node_1".to_string())), + )?; + tokio_test::block_on(controller.create_session(common::apis::SessionAttributes { + id: ssn_id.to_string(), + application: "flmtest".to_string(), + common_data: None, + min_instances: 0, + max_instances, + batch_size, + priority: 0, + resreq: Some(common::apis::ResourceRequirement { + cpu: 1, + memory: 1024 * 1024 * 1024, + gpu: 0, + }), + }))?; + for _ in 0..task_count { + tokio_test::block_on(controller.create_task(ssn_id.to_string(), None))?; + } + Ok(()) + } + struct TestEnv { url: String, pub controller: ControllerPtr, @@ -175,7 +210,7 @@ mod tests { } } - /// Test the allocation of void executors to underused sessions. + /// Allocate enough executors for all incomplete tasks and avoid duplicates in later cycles. #[test] fn test_allocate_executors() -> Result<(), FlameError> { let env = TestEnv::new()?; @@ -229,12 +264,12 @@ mod tests { actions: vec![], }; - let dispatch = DispatchAction::new_ptr(); - tokio_test::block_on(dispatch.execute(&mut ctx))?; - let alloc = AllocateAction::new_ptr(); tokio_test::block_on(alloc.execute(&mut ctx))?; + let dispatch = DispatchAction::new_ptr(); + tokio_test::block_on(dispatch.execute(&mut ctx))?; + let ssn_list = snapshot.find_sessions(OPEN_SESSION)?; assert_eq!(ssn_list.len(), 1); assert_eq!(ssn_list.values().next().unwrap().id, ssn_1.id.clone()); @@ -244,14 +279,14 @@ mod tests { assert_eq!(node_list.values().next().unwrap().name, "node_1"); let exec_list = controller.list_executor()?; - assert_eq!(exec_list.len(), 1); + assert_eq!(exec_list.len(), task_num); } Ok(()) } /// One scheduling cycle must keep the same in-memory [`crate::scheduler::plugins::PluginManager`] - /// so Gang (and similar) state from Dispatch is visible to Allocate. + /// across Allocate, Dispatch, and Shuffle. #[test] fn test_scheduler_cycle_reuses_plugin_manager_across_actions() -> Result<(), FlameError> { let env = TestEnv::new()?; @@ -339,4 +374,128 @@ mod tests { Ok(()) } + + #[test] + fn test_priority_policy_uses_always_enabled_batch() -> Result<(), FlameError> { + let env = TestEnv::new()?; + let controller = env.controller.clone(); + let ssn_id = format!("priority-only-{}", Uuid::new_v4()); + setup_test_session(&controller, &ssn_id, 1, None, 3)?; + + let policies = vec!["priority".to_string()]; + let mut ctx = Context::new(controller.clone(), &policies)?; + assert!(!ctx.is_fulfilled( + ctx.snapshot + .find_sessions(OPEN_SESSION)? + .values() + .next() + .unwrap(), + )?); + + tokio_test::block_on(AllocateAction::new_ptr().execute(&mut ctx))?; + assert_eq!(controller.list_executor()?.len(), 3); + Ok(()) + } + + #[test] + fn test_drf_policy_uses_always_enabled_batch() -> Result<(), FlameError> { + let env = TestEnv::new()?; + let controller = env.controller.clone(); + let ssn_id = format!("drf-only-{}", Uuid::new_v4()); + setup_test_session(&controller, &ssn_id, 1, None, 3)?; + + let policies = vec!["drf".to_string()]; + let mut ctx = Context::new(controller.clone(), &policies)?; + tokio_test::block_on(AllocateAction::new_ptr().execute(&mut ctx))?; + assert_eq!(controller.list_executor()?.len(), 3); + Ok(()) + } + + #[test] + fn test_batch_allocates_for_all_incomplete_tasks() -> Result<(), FlameError> { + let env = TestEnv::new()?; + let controller = env.controller.clone(); + let ssn_id = format!("batch-all-tasks-{}", Uuid::new_v4()); + setup_test_session(&controller, &ssn_id, 1, None, 3)?; + + let policies = vec!["priority".to_string()]; + let mut ctx = Context::new(controller.clone(), &policies)?; + tokio_test::block_on(AllocateAction::new_ptr().execute(&mut ctx))?; + + assert_eq!(controller.list_executor()?.len(), 3); + Ok(()) + } + + #[test] + fn test_batch_allocation_respects_max_instances() -> Result<(), FlameError> { + let env = TestEnv::new()?; + let controller = env.controller.clone(); + let ssn_id = format!("batch-max-{}", Uuid::new_v4()); + setup_test_session(&controller, &ssn_id, 2, Some(4), 8)?; + + let policies = vec!["priority".to_string()]; + let mut ctx = Context::new(controller.clone(), &policies)?; + tokio_test::block_on(AllocateAction::new_ptr().execute(&mut ctx))?; + + assert_eq!(controller.list_executor()?.len(), 4); + Ok(()) + } + + #[test] + fn test_allocate_reuses_idle_executor() -> Result<(), FlameError> { + let env = TestEnv::new()?; + let controller = env.controller.clone(); + let ssn_id = format!("reuse-idle-{}", Uuid::new_v4()); + setup_test_session(&controller, &ssn_id, 1, None, 1)?; + let executor = + tokio_test::block_on(controller.create_executor("node_1".to_string(), ssn_id.clone()))?; + tokio_test::block_on(controller.register_executor(&executor))?; + + let policies = vec!["priority".to_string()]; + let mut ctx = Context::new(controller.clone(), &policies)?; + for action in ctx.actions.clone() { + tokio_test::block_on(action.execute(&mut ctx))?; + } + + let executors = controller.list_executor()?; + assert_eq!(executors.len(), 1); + assert_eq!(executors[0].state, common::apis::ExecutorState::Binding); + assert_eq!(executors[0].ssn_id.as_deref(), Some(ssn_id.as_str())); + Ok(()) + } + + #[test] + fn test_binding_executor_prevents_duplicate_across_cycles() -> Result<(), FlameError> { + let env = TestEnv::new()?; + let controller = env.controller.clone(); + let ssn_id = format!("binding-stall-{}", Uuid::new_v4()); + setup_test_session(&controller, &ssn_id, 1, None, 1)?; + let executor = + tokio_test::block_on(controller.create_executor("node_1".to_string(), ssn_id.clone()))?; + tokio_test::block_on(controller.register_executor(&executor))?; + tokio_test::block_on(controller.bind_session(executor.id.clone(), ssn_id.clone()))?; + + let policies = vec!["priority".to_string()]; + for _ in 0..2 { + let mut ctx = Context::new(controller.clone(), &policies)?; + let ssn = ctx + .snapshot + .find_sessions(OPEN_SESSION)? + .values() + .next() + .cloned() + .expect("test session must exist"); + assert!(ctx.is_fulfilled(&ssn)?); + assert!(ctx.is_ready(&ssn)?); + for action in ctx.actions.clone() { + tokio_test::block_on(action.execute(&mut ctx))?; + } + } + + let executors = controller.list_executor()?; + assert_eq!(executors.len(), 1); + assert_eq!(executors[0].state, common::apis::ExecutorState::Binding); + assert_eq!(executors[0].ssn_id.as_deref(), Some(ssn_id.as_str())); + Ok(()) + } } diff --git a/session_manager/src/scheduler/plugins/gang.rs b/session_manager/src/scheduler/plugins/batch.rs similarity index 81% rename from session_manager/src/scheduler/plugins/gang.rs rename to session_manager/src/scheduler/plugins/batch.rs index 51eec191..d2ff6a5f 100644 --- a/session_manager/src/scheduler/plugins/gang.rs +++ b/session_manager/src/scheduler/plugins/batch.rs @@ -13,34 +13,34 @@ limitations under the License. use std::collections::HashMap; -use common::apis::SessionID; +use common::apis::{SessionID, TaskState}; use common::FlameError; use crate::model::{ExecutorInfoPtr, NodeInfoPtr, SessionInfoPtr, SnapShot}; use crate::scheduler::plugins::{Plugin, PluginPtr}; -struct GangState { +struct BatchState { batch_size: u32, allocated: u32, pipelined: u32, bound: u32, } -pub struct GangPlugin { - ssn_state: HashMap, +pub struct BatchPlugin { + ssn_state: HashMap, } -impl GangPlugin { +impl BatchPlugin { pub fn new_ptr() -> PluginPtr { - Box::new(GangPlugin { + Box::new(BatchPlugin { ssn_state: HashMap::new(), }) } } -impl Plugin for GangPlugin { +impl Plugin for BatchPlugin { fn name(&self) -> &'static str { - "gang" + "batch" } fn setup(&mut self, ss: &SnapShot) -> Result<(), FlameError> { @@ -55,7 +55,7 @@ impl Plugin for GangPlugin { for ssn in sessions.values() { self.ssn_state.insert( ssn.id.clone(), - GangState { + BatchState { batch_size: ssn.batch_size.max(1), allocated: 0, pipelined: 0, @@ -79,14 +79,48 @@ impl Plugin for GangPlugin { fn is_ready(&self, ssn: &SessionInfoPtr) -> Option { let state = self.ssn_state.get(&ssn.id)?; - let total = state.allocated + state.pipelined; - Some(state.pipelined > 0 && total % state.batch_size == 0) + let incomplete_tasks = [TaskState::Pending, TaskState::Running] + .iter() + .map(|task_state| { + ssn.tasks_status + .get(task_state) + .copied() + .unwrap_or(0) + .max(0) as u64 + }) + .sum::(); + let batch_size = state.batch_size as u64; + let uncapped = incomplete_tasks.div_ceil(batch_size) * batch_size; + let aligned_max = ssn + .max_instances + .map(|max| (max as u64 / batch_size) * batch_size) + .unwrap_or((u32::MAX as u64 / batch_size) * batch_size); + let needed = uncapped.min(aligned_max); + let total = state.allocated as u64 + state.bound as u64; + Some(needed == 0 || total >= needed) } fn is_fulfilled(&self, ssn: &SessionInfoPtr) -> Option { let state = self.ssn_state.get(&ssn.id)?; - let total = state.allocated + state.bound; - Some(state.bound > 0 && total % state.batch_size == 0) + let incomplete_tasks = [TaskState::Pending, TaskState::Running] + .iter() + .map(|task_state| { + ssn.tasks_status + .get(task_state) + .copied() + .unwrap_or(0) + .max(0) as u64 + }) + .sum::(); + let batch_size = state.batch_size as u64; + let uncapped = incomplete_tasks.div_ceil(batch_size) * batch_size; + let aligned_max = ssn + .max_instances + .map(|max| (max as u64 / batch_size) * batch_size) + .unwrap_or((u32::MAX as u64 / batch_size) * batch_size); + let needed = uncapped.min(aligned_max); + let total = state.allocated as u64 + state.pipelined as u64; + Some(needed == 0 || total >= needed) } fn on_executor_allocate(&mut self, _node: NodeInfoPtr, ssn: SessionInfoPtr) { @@ -191,7 +225,7 @@ mod tests { let ssn = create_test_session("ssn-1", 1); ss.add_session(ssn.clone()).unwrap(); - let mut plugin = GangPlugin { + let mut plugin = BatchPlugin { ssn_state: HashMap::new(), }; plugin.setup(&ss).unwrap(); @@ -199,7 +233,7 @@ mod tests { assert!(!plugin.is_fulfilled(&ssn).unwrap()); let node = create_test_node("node-1"); - plugin.on_session_bind(ssn.clone()); + plugin.on_executor_allocate(node, ssn.clone()); assert!(plugin.is_fulfilled(&ssn).unwrap()); } @@ -211,7 +245,7 @@ mod tests { let ssn = create_test_session("ssn-1", 2); ss.add_session(ssn.clone()).unwrap(); - let mut plugin = GangPlugin { + let mut plugin = BatchPlugin { ssn_state: HashMap::new(), }; plugin.setup(&ss).unwrap(); @@ -219,11 +253,11 @@ mod tests { assert!(!plugin.is_fulfilled(&ssn).unwrap()); let node = create_test_node("node-1"); - plugin.on_session_bind(ssn.clone()); + plugin.on_executor_allocate(node.clone(), ssn.clone()); assert!(!plugin.is_fulfilled(&ssn).unwrap()); - plugin.on_session_bind(ssn.clone()); + plugin.on_executor_allocate(node, ssn.clone()); assert!(plugin.is_fulfilled(&ssn).unwrap()); } @@ -238,7 +272,7 @@ mod tests { let exec = create_test_executor("exec-1", Some("ssn-1")); ss.add_executor(exec).unwrap(); - let mut plugin = GangPlugin { + let mut plugin = BatchPlugin { ssn_state: HashMap::new(), }; plugin.setup(&ss).unwrap(); @@ -246,7 +280,7 @@ mod tests { assert!(!plugin.is_fulfilled(&ssn).unwrap()); let node = create_test_node("node-1"); - plugin.on_session_bind(ssn.clone()); + plugin.on_executor_allocate(node, ssn.clone()); assert!(plugin.is_fulfilled(&ssn).unwrap()); } @@ -258,15 +292,14 @@ mod tests { let ssn = create_test_session("ssn-1", 1); ss.add_session(ssn.clone()).unwrap(); - let mut plugin = GangPlugin { + let mut plugin = BatchPlugin { ssn_state: HashMap::new(), }; plugin.setup(&ss).unwrap(); assert!(!plugin.is_ready(&ssn).unwrap()); - let node = create_test_node("node-1"); - plugin.on_executor_allocate(node, ssn.clone()); + plugin.on_session_bind(ssn.clone()); assert!(plugin.is_ready(&ssn).unwrap()); } @@ -278,19 +311,18 @@ mod tests { let ssn = create_test_session("ssn-1", 2); ss.add_session(ssn.clone()).unwrap(); - let mut plugin = GangPlugin { + let mut plugin = BatchPlugin { ssn_state: HashMap::new(), }; plugin.setup(&ss).unwrap(); assert!(!plugin.is_ready(&ssn).unwrap()); - let node = create_test_node("node-1"); - plugin.on_executor_allocate(node.clone(), ssn.clone()); + plugin.on_session_bind(ssn.clone()); assert!(!plugin.is_ready(&ssn).unwrap()); - plugin.on_executor_allocate(node, ssn.clone()); + plugin.on_session_bind(ssn.clone()); assert!(plugin.is_ready(&ssn).unwrap()); } @@ -305,15 +337,14 @@ mod tests { let exec = create_test_executor("exec-1", Some("ssn-1")); ss.add_executor(exec).unwrap(); - let mut plugin = GangPlugin { + let mut plugin = BatchPlugin { ssn_state: HashMap::new(), }; plugin.setup(&ss).unwrap(); assert!(!plugin.is_ready(&ssn).unwrap()); - let node = create_test_node("node-1"); - plugin.on_executor_allocate(node, ssn.clone()); + plugin.on_session_bind(ssn.clone()); assert!(plugin.is_ready(&ssn).unwrap()); } @@ -325,7 +356,7 @@ mod tests { let ssn = create_test_session("ssn-1", 2); ss.add_session(ssn.clone()).unwrap(); - let mut plugin = GangPlugin { + let mut plugin = BatchPlugin { ssn_state: HashMap::new(), }; plugin.setup(&ss).unwrap(); @@ -333,19 +364,19 @@ mod tests { let exec = create_test_executor("exec-1", None); plugin.on_executor_pipeline(exec.clone(), ssn.clone()); - assert!(!plugin.is_ready(&ssn).unwrap()); + assert!(!plugin.is_fulfilled(&ssn).unwrap()); plugin.on_executor_pipeline(exec.clone(), ssn.clone()); - assert!(plugin.is_ready(&ssn).unwrap()); + assert!(plugin.is_fulfilled(&ssn).unwrap()); plugin.on_executor_discard(exec.clone(), ssn.clone()); - assert!(!plugin.is_ready(&ssn).unwrap()); + assert!(!plugin.is_fulfilled(&ssn).unwrap()); plugin.on_executor_discard(exec, ssn.clone()); - assert!(!plugin.is_ready(&ssn).unwrap()); + assert!(!plugin.is_fulfilled(&ssn).unwrap()); } #[test] @@ -355,25 +386,25 @@ mod tests { let ssn = create_test_session("ssn-1", 2); ss.add_session(ssn.clone()).unwrap(); - let mut plugin = GangPlugin { + let mut plugin = BatchPlugin { ssn_state: HashMap::new(), }; plugin.setup(&ss).unwrap(); plugin.on_session_bind(ssn.clone()); - assert!(!plugin.is_fulfilled(&ssn).unwrap()); + assert!(!plugin.is_ready(&ssn).unwrap()); plugin.on_session_bind(ssn.clone()); - assert!(plugin.is_fulfilled(&ssn).unwrap()); + assert!(plugin.is_ready(&ssn).unwrap()); plugin.on_session_unbind(ssn.clone()); - assert!(!plugin.is_fulfilled(&ssn).unwrap()); + assert!(!plugin.is_ready(&ssn).unwrap()); plugin.on_session_unbind(ssn.clone()); - assert!(!plugin.is_fulfilled(&ssn).unwrap()); + assert!(!plugin.is_ready(&ssn).unwrap()); } } diff --git a/session_manager/src/scheduler/plugins/drf.rs b/session_manager/src/scheduler/plugins/drf.rs index 0ca39f10..45609c01 100644 --- a/session_manager/src/scheduler/plugins/drf.rs +++ b/session_manager/src/scheduler/plugins/drf.rs @@ -19,7 +19,7 @@ use crate::model::{ OPEN_SESSION, }; use crate::scheduler::plugins::{Plugin, PluginPtr}; -use common::apis::{ResourceRequirement, SessionID, TaskState}; +use common::apis::{ExecutorState, ResourceRequirement, SessionID, TaskState}; use common::FlameError; #[derive(Default, Clone)] @@ -252,6 +252,11 @@ impl Plugin for DRFPlugin { } fn on_executor_pipeline(&mut self, exec: ExecutorInfoPtr, ssn: SessionInfoPtr) { + // Idle executors already consume node resources. Allocate reserves them for Batch + // fulfillment, while DRF assigns their session share when Dispatch binds them. + if exec.state == ExecutorState::Idle { + return; + } let ssn_resreq = self.get_session_resreq(&ssn); if let Some(entry) = self.ssn_map.get_mut(&ssn.id) { @@ -274,6 +279,9 @@ impl Plugin for DRFPlugin { } fn on_executor_discard(&mut self, exec: ExecutorInfoPtr, ssn: SessionInfoPtr) { + if exec.state == ExecutorState::Idle { + return; + } let ssn_resreq = self.get_session_resreq(&ssn); if let Some(entry) = self.ssn_map.get_mut(&ssn.id) { diff --git a/session_manager/src/scheduler/plugins/mod.rs b/session_manager/src/scheduler/plugins/mod.rs index 3edc78ca..43c949b6 100644 --- a/session_manager/src/scheduler/plugins/mod.rs +++ b/session_manager/src/scheduler/plugins/mod.rs @@ -19,16 +19,16 @@ use stdng::collections; use stdng::{lock_ptr, new_ptr, MutexPtr}; use crate::model::{ExecutorInfoPtr, NodeInfo, NodeInfoPtr, SessionInfo, SessionInfoPtr, SnapShot}; +use crate::scheduler::plugins::batch::BatchPlugin; use crate::scheduler::plugins::drf::DRFPlugin; -use crate::scheduler::plugins::gang::GangPlugin; use crate::scheduler::plugins::priority::PriorityPlugin; use crate::scheduler::plugins::shim::ShimPlugin; use crate::scheduler::Context; use common::FlameError; +mod batch; mod drf; -mod gang; mod priority; mod shim; @@ -132,9 +132,9 @@ const PLUGIN_REGISTRY: &[PluginInfo] = &[ configurable: true, }, PluginInfo { - name: "gang", - constructor: GangPlugin::new_ptr, - configurable: true, + name: "batch", + constructor: BatchPlugin::new_ptr, + configurable: false, }, PluginInfo { name: "shim", @@ -161,13 +161,21 @@ impl PluginManager { enabled_policies: &[String], ) -> Result { let valid_names = configurable_policy_names(); + let all_plugin_names: Vec<&str> = PLUGIN_REGISTRY.iter().map(|p| p.name).collect(); for p in enabled_policies { if !valid_names.contains(&p.as_str()) { - return Err(FlameError::InvalidConfig(format!( - "unknown policy: {}. available: {:?}", - p, valid_names - ))); + if all_plugin_names.contains(&p.as_str()) { + tracing::info!( + "Policy '{}' is always enabled and does not need to be listed explicitly; ignoring", + p + ); + } else { + return Err(FlameError::InvalidConfig(format!( + "unknown policy: {}. configurable policies: {:?}", + p, valid_names + ))); + } } } @@ -194,7 +202,7 @@ impl PluginManager { /// Returns whether the session is underused (needs more executors). /// /// Uses "first non-`None` wins" ordering, identical to `ssn_order_fn`. - /// Plugins are consulted in registration order (Priority → DRF → Gang → Shim). + /// Plugins are consulted in registration order (Priority → DRF → Batch → Shim). /// The first plugin that returns `Some(result)` wins; `None` means "no opinion, ask the /// next plugin". If no plugin has an opinion, the session is considered NOT underused. /// @@ -332,27 +340,36 @@ impl PluginManager { Ok(()) } - /// True if every plugin that implements [`Plugin::is_ready`] reports readiness (no opinion - /// defaults to true). Counters are in-memory and advance when [`crate::scheduler::Statement`] - /// records `allocate` / `pipeline` without `discard`. Dispatch and Allocate share one - /// `PluginManager` per scheduling cycle. + /// Returns the conjunction of plugin readiness opinions. Plugins returning `None` are skipped; + /// if every plugin abstains, the session is considered ready. pub fn is_ready(&self, ssn: &SessionInfoPtr) -> Result { let plugins = lock_ptr!(self.plugins)?; - - Ok(plugins - .iter() - .all(|(_, plugin)| plugin.is_ready(ssn).unwrap_or(true))) + let mut result = None; + for (_, plugin) in plugins.iter() { + if let Some(ready) = plugin.is_ready(ssn) { + if !ready { + return Ok(false); + } + result = Some(true); + } + } + Ok(result.unwrap_or(true)) } - /// True if every plugin that implements [`Plugin::is_fulfilled`] reports fulfillment (no opinion - /// defaults to true). Updates when [`crate::scheduler::Statement`] records `bind`; after - /// Dispatch commits, Allocate uses this to skip redundant provisioning. + /// Returns the conjunction of plugin fulfillment opinions. Plugins returning `None` are + /// skipped; if every plugin abstains, the session is considered fulfilled. pub fn is_fulfilled(&self, ssn: &SessionInfoPtr) -> Result { let plugins = lock_ptr!(self.plugins)?; - - Ok(plugins - .iter() - .all(|(_, plugin)| plugin.is_fulfilled(ssn).unwrap_or(true))) + let mut result = None; + for (_, plugin) in plugins.iter() { + if let Some(fulfilled) = plugin.is_fulfilled(ssn) { + if !fulfilled { + return Ok(false); + } + result = Some(true); + } + } + Ok(result.unwrap_or(true)) } pub fn on_executor_pipeline( @@ -496,7 +513,7 @@ impl collections::Cmp for SsnOrderFn { #[cfg(test)] mod tests { use super::*; - use crate::model::ExecutorInfo; + use crate::model::{ExecutorInfo, SessionInfo}; use chrono::Utc; use common::apis::{ExecutorState, ResourceRequirement, Shim}; @@ -542,6 +559,53 @@ mod tests { assert_eq!(exec_void.state, ExecutorState::Void); } + #[test] + fn test_batch_is_always_enabled() { + let ss = SnapShot::new(); + let ssn = Arc::new(SessionInfo { + id: "ssn-always-batch".to_string(), + resreq: Some(ResourceRequirement::default()), + ..Default::default() + }); + ss.add_session(ssn.clone()).unwrap(); + + let plugins = PluginManager::setup(&ss, &["priority".to_string()]).unwrap(); + let plugin_names = plugins + .plugins + .lock() + .unwrap() + .iter() + .map(|(name, _)| name.clone()) + .collect::>(); + assert_eq!(plugin_names, vec!["priority", "batch", "shim"]); + assert!(plugins.is_ready(&ssn).unwrap()); + assert!(plugins.is_fulfilled(&ssn).unwrap()); + } + + #[test] + fn test_explicit_batch_policy_is_accepted() { + let ss = SnapShot::new(); + PluginManager::setup(&ss, &["batch".to_string()]).unwrap(); + } + + #[test] + fn test_removed_gang_policy_is_rejected() { + let ss = SnapShot::new(); + assert!(PluginManager::setup(&ss, &["gang".to_string()]).is_err()); + } + + #[test] + fn test_explicit_shim_policy_is_accepted() { + let ss = SnapShot::new(); + PluginManager::setup(&ss, &["shim".to_string()]).unwrap(); + } + + #[test] + fn test_unknown_policy_is_rejected() { + let ss = SnapShot::new(); + assert!(PluginManager::setup(&ss, &["unknown".to_string()]).is_err()); + } + /// Test documentation for plugin fallback behavior. #[test] fn test_plugin_fallback_behavior_documentation() { diff --git a/session_manager/src/scheduler/plugins/priority.rs b/session_manager/src/scheduler/plugins/priority.rs index d719a579..f64b95d8 100644 --- a/session_manager/src/scheduler/plugins/priority.rs +++ b/session_manager/src/scheduler/plugins/priority.rs @@ -14,7 +14,7 @@ limitations under the License. use std::cmp::Ordering; use std::collections::HashMap; -use common::apis::{ResourceRequirement, SessionID, TaskState}; +use common::apis::{ExecutorState, ResourceRequirement, SessionID, TaskState}; use common::FlameError; use crate::model::{ @@ -334,7 +334,12 @@ impl Plugin for PriorityPlugin { } } - fn on_executor_pipeline(&mut self, _exec: ExecutorInfoPtr, ssn: SessionInfoPtr) { + fn on_executor_pipeline(&mut self, exec: ExecutorInfoPtr, ssn: SessionInfoPtr) { + // Idle executors are existing reusable supply. Allocate reserves them for Batch + // fulfillment, but Priority accounts them only when Dispatch actually binds them. + if exec.state == ExecutorState::Idle { + return; + } let unit = match self.ssn_unit.get(&ssn.id) { Some(u) => u.clone(), None => return, @@ -344,7 +349,10 @@ impl Plugin for PriorityPlugin { } } - fn on_executor_discard(&mut self, _exec: ExecutorInfoPtr, ssn: SessionInfoPtr) { + fn on_executor_discard(&mut self, exec: ExecutorInfoPtr, ssn: SessionInfoPtr) { + if exec.state == ExecutorState::Idle { + return; + } let unit = match self.ssn_unit.get(&ssn.id) { Some(u) => u.clone(), None => return, diff --git a/session_manager/src/storage/engine/filesystem.rs b/session_manager/src/storage/engine/filesystem.rs index 1114ac8b..559b5dbe 100644 --- a/session_manager/src/storage/engine/filesystem.rs +++ b/session_manager/src/storage/engine/filesystem.rs @@ -1071,7 +1071,7 @@ impl Engine for FilesystemEngine { // If spec provided, validate full session attributes (same as sqlite engine // and in-memory cache). Only checking application was insufficient: // a persisted session could have batch_size/min_instances/max_instances that - // differ from the client spec, leading to gang scheduling deadlocks (tasks + // differ from the client spec, leading to batch scheduling deadlocks (tasks // never allocated) without a clear error. if let Some(ref attr) = spec { ssn.validate_spec(attr)?; diff --git a/session_manager/src/storage/load_data_tests.rs b/session_manager/src/storage/load_data_tests.rs index e5047524..7c3ffbf8 100644 --- a/session_manager/src/storage/load_data_tests.rs +++ b/session_manager/src/storage/load_data_tests.rs @@ -29,7 +29,7 @@ mod tests { endpoint: "http://localhost:8080".to_string(), storage: db_url.to_string(), resreq: None, - policies: vec!["priority".to_string(), "gang".to_string()], + policies: vec!["priority".to_string()], schedule_interval: 1000, executors: FlameExecutors { shim: Shim::default(),