From 9f0f22dc82e134907492d71d58b91dd8d1ba6d43 Mon Sep 17 00:00:00 2001 From: puddingfjz <2811443837@qq.com> Date: Wed, 22 Jul 2026 17:58:20 -0700 Subject: [PATCH] Update: simplify NEXT_LEVEL scheduling with explicit targets Replace unconstrained NEXT_LEVEL scheduling with explicit placement: every NEXT_LEVEL task names its exact worker at submit time and the Scheduler dispatches only to that worker. The runtime no longer selects a NEXT_LEVEL worker on the caller's behalf. - submit_next_level requires a stable worker id; submit_next_level_group requires one distinct id per member. Booleans, floats, numeric strings, and other coercible non-integers are rejected. - READY singles route through a per-worker FIFO; groups use one shared FIFO whose head launches only when every target worker is idle (all-or-nothing, no partial reservation). SUB placement is unchanged and takes no worker id. - A local (LOCAL_PYTHON / LOCAL_CHIP) callable is rejected when pinned to a remote worker id: a remote endpoint installs only its own dispatcher callables, so a local digest would fail asynchronously as an unknown hashid. - Remove the shared NEXT_LEVEL queue, idle-pool fallback selection, and the queue wait/shutdown/condition-variable state; retain idle-worker selection only for SUB. Docs: rewrite scheduler.md for the directed model, add directed-next-level-scheduling.md as the placement-contract reference, and refresh the maintained call sites and API summaries for the required worker argument. Tests: cover required/malformed target rejection, per-worker FIFO routing, exact member-to-worker mapping (including reversed target order), group-first dispatch, blocked-group non-reservation, and the local-callable remote-target guard. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/callable-identity-registration.md | 21 +- docs/directed-next-level-scheduling.md | 79 +++ docs/hierarchical_level_runtime.md | 28 +- docs/orchestrator.md | 114 ++--- docs/remote-l3-worker-design.md | 36 +- .../implementation-plan.md | 11 +- .../pr-split-and-audit-artifacts.md | 2 +- .../pr-split-and-audit-plan.md | 4 +- docs/scheduler.md | 482 +++++------------- docs/task-flow.md | 26 +- docs/worker-manager.md | 34 +- examples/workers/l3/README.md | 6 +- .../workers/l3/per_task_runtime_env/README.md | 12 +- .../workers/l3/per_task_runtime_env/main.py | 2 +- python/bindings/worker_bind.h | 9 +- python/simpler/orchestrator.py | 126 ++--- python/simpler/worker.py | 54 +- simpler_setup/scene_test.py | 2 +- .../docs/MULTI_RING.md | 4 +- .../docs/MULTI_RING.md | 4 +- src/common/hierarchical/orchestrator.cpp | 109 ++-- src/common/hierarchical/orchestrator.h | 46 +- src/common/hierarchical/scheduler.cpp | 208 ++++---- src/common/hierarchical/scheduler.h | 20 +- src/common/hierarchical/types.cpp | 72 ++- src/common/hierarchical/types.h | 70 +-- src/common/hierarchical/worker.cpp | 17 +- src/common/hierarchical/worker.h | 7 +- src/common/hierarchical/worker_manager.cpp | 33 +- src/common/hierarchical/worker_manager.h | 11 +- .../dynamic_register/test_dynamic_register.py | 40 +- .../test_l3_dependency.py | 2 +- .../tensormap_and_ringbuffer/test_l3_group.py | 2 +- .../ut/cpp/hierarchical/test_orchestrator.cpp | 69 +-- tests/ut/cpp/hierarchical/test_scheduler.cpp | 307 ++++++++--- tests/ut/py/test_callable_identity.py | 49 +- .../py/test_worker/test_child_addr_guard.py | 99 +--- .../py/test_worker/test_error_propagation.py | 4 +- tests/ut/py/test_worker/test_l4_recursive.py | 30 +- .../py/test_worker/test_startup_readiness.py | 12 +- 40 files changed, 1146 insertions(+), 1117 deletions(-) create mode 100644 docs/directed-next-level-scheduling.md diff --git a/docs/callable-identity-registration.md b/docs/callable-identity-registration.md index c6122ca66e..816034dcdc 100644 --- a/docs/callable-identity-registration.md +++ b/docs/callable-identity-registration.md @@ -342,11 +342,11 @@ every active child endpoint in the handle's `target_namespace` for this `Worker` has installed the callable identity. Registering to a user-selected worker subset is not part of this contract. `orch.submit_next_level(..., worker=...)` and -`orch.submit_next_level_group(..., workers=...)` are submit-time affinity -controls; they do not define registration scope. NEXT_LEVEL affinity consumes -stable worker ids. Local Python Worker children and remote L3 workers use the -worker ids returned by `add_worker(...)` / `add_remote_worker(...)`; L3 chip -worker ids are the existing chip worker ids. +`orch.submit_next_level_group(..., workers=...)` require exact submit-time +placement; they do not define registration scope. NEXT_LEVEL placement uses +stable worker ids. Local Python Worker children and remote L3 workers use ids +returned by `add_worker(...)` / `add_remote_worker(...)`; L3 chip worker ids +are the existing chip worker ids. `CallableHandle` is the public callable token returned by registration: @@ -381,7 +381,7 @@ matmul = worker.register(chip_callable) postprocess = worker.register(py_callable) def parent_orch(orch, args, config): - orch.submit_next_level(matmul, args, config) + orch.submit_next_level(matmul, args, config, worker=0) orch.submit_sub(postprocess, args) ``` @@ -488,11 +488,10 @@ Worker teardown is deferred to a later design. ### Dispatch Contract -Parent-side scheduling assumes the handle's `hashid` is installed on every -active target in its registration scope. Dispatch choices are constrained by -the handle namespace, submit-time affinity, and tensor/buffer accessibility. -Submit-time live validation is a preflight check only. It does not pin the -target identity through later drain or child dispatch. Callers must not +Parent-side scheduling requires the handle's `hashid` on the submitted exact +target. The Orchestrator validates the handle namespace and tensor/buffer +accessibility for that target before committing the slot. The target identity +is then fixed in `TaskSlotState` through dispatch. Callers must not concurrently unregister a handle while `Worker.run()` or any in-flight task may submit or use that handle; wait for the relevant run/drain to return before unregistering it. diff --git a/docs/directed-next-level-scheduling.md b/docs/directed-next-level-scheduling.md new file mode 100644 index 0000000000..acb368d5b0 --- /dev/null +++ b/docs/directed-next-level-scheduling.md @@ -0,0 +1,79 @@ +# Directed NEXT_LEVEL Scheduling + +Every NEXT_LEVEL task names its exact worker at submission time, and the +Scheduler dispatches it only to that worker. The runtime never selects a +NEXT_LEVEL worker on the caller's behalf. SUB tasks keep free scheduling and +take no worker id. + +This page is the placement contract. For the queue topology, dispatch loop, +and dependency/failure handling see [scheduler.md](scheduler.md); for the +submit API see [orchestrator.md](orchestrator.md). + +## Placement contract + +- `Orchestrator.submit_next_level(callable, args, config, worker=id)` targets + exactly one stable NEXT_LEVEL worker id. +- `Orchestrator.submit_next_level_group(callable, args_list, config, + workers=[...])` takes one stable worker id per member. The ids must be + distinct and their count must equal the member count. +- A worker id must be non-negative and name a registered NEXT_LEVEL worker. + Booleans, floats, numeric strings, and other coercible non-integers are + rejected. +- A local (`LOCAL_PYTHON` / `LOCAL_CHIP`) callable must target a local child; + a remote worker id is rejected because a remote endpoint installs only its + own dispatcher callables. +- When a callable or its tensor data constrains eligibility (remote callables, + remote-buffer data), the target must lie in that eligible set. Eligibility + bounds which targets are *valid*; it is not a scheduler selection policy. +- `submit_sub` / `submit_sub_group` are unchanged and take no worker id. + +## Queue and dispatch model + +The worker owns a directed-queue registry fixed at `init()` from the stable +NEXT_LEVEL worker ids: + +| Task | Queue | +| ---- | ----- | +| NEXT_LEVEL single | `FIFO[target_worker_id]` | +| NEXT_LEVEL group | one shared group FIFO | +| SUB | one shared SUB FIFO | + +A task's target and queue never change after submission. Tasks that are READY +at submission and tasks released by a completing dependency use the same +routing operation. + +Each Scheduler iteration: + +- launches the group FIFO head only when every one of its target workers is + idle — all members or none, with no partial reservation; +- dispatches the head of each idle worker's single FIFO; +- dispatches READY SUB work through the existing free-selection path. + +Only the group FIFO head is examined, and a launchable group is tried before +conflicting singles. The runtime adds no fairness, aging, priority, or +reservation policy; callers choose worker sets that make acceptable progress. + +## Invariants + +- PENDING tasks live in task-slot state, not in a ready queue; a task is routed + only once all producers complete successfully. +- A failed producer poisons its dependents instead of enqueuing them. +- A group is one DAG node and completes only after every member reaches a + terminal state. +- The NEXT_LEVEL worker set is fixed after `init()`; a worker lookup that fails + after submit-time validation is an invariant violation, not a fallback. +- No queue or scheduler mutex is held while an endpoint executes. + +## Non-goals + +- No SUB worker selection. +- No priorities, work stealing, rebinding, queue scanning, aging, quotas, + preemption, starvation prevention, or partial group reservation. +- No compatibility flags, environment variables, or macros. + +## Related documents + +- [scheduler.md](scheduler.md) — queue topology, dispatch loop, completion +- [orchestrator.md](orchestrator.md) — submit API and DAG construction +- [remote-l3-worker-design.md](remote-l3-worker-design.md) — remote NEXT_LEVEL + workers diff --git a/docs/hierarchical_level_runtime.md b/docs/hierarchical_level_runtime.md index 6ebf6b1a5f..adbd5823bc 100644 --- a/docs/hierarchical_level_runtime.md +++ b/docs/hierarchical_level_runtime.md @@ -86,23 +86,24 @@ Owns: identity and logical offset. - `Scope` — lifetime management for intermediate tensors -One `submit_next_level(callable, task_args, config)` call: +One `submit_next_level(callable, task_args, config, worker=worker_id)` call: 1. allocates a slot 2. moves task data into the slot 3. walks `TaskArgs` tags (INPUT/OUTPUT/INOUT/OUTPUT_EXISTING/NO_DEP) to lookup/insert TensorMap entries 4. records fanin metadata on producer slots -5. pushes the new slot onto the scheduler's wiring queue +5. attaches fanout edges and routes READY work to its target queue See [orchestrator.md](orchestrator.md) for the 7-step submit flow and state machine. ### Scheduler (Scheduler thread) -The **DAG executor**. A dedicated C++ thread that drains three queues: +The **DAG executor**. A dedicated C++ thread that handles: -- **wiring queue** — slots just submitted; wire fanout edges, compute readiness -- **ready queue** — slots with all fanin satisfied; pick an idle WorkerThread and dispatch +- **directed NEXT_LEVEL queues** — one single-task FIFO per stable worker ID + plus one all-targets-ready group FIFO +- **shared SUB queue** — freely select idle SUB WorkerThreads - **completion queue** — slots whose worker finished; release fanout, wake downstream consumers, retire slot The Scheduler never inspects task data — it just moves slot ids between queues @@ -141,13 +142,13 @@ what flows through `ChipWorker::run`. │ submit(callable, args, config) │ │ 1. ring.alloc() │ │ 2. TensorMap lookup/insert │ - │ 3. record fanin │ - │ 4. push wiring_queue ───────►│ - │ │ Phase 0: drain wiring_queue - │ │ wire fanout edges - │ │ if ready → ready_queue - │ │ pop ready_queue - │ │ pick idle WorkerThread + │ 3. attach fanout + fanin │ + │ 4. if READY, enqueue directed/shared ready queue + │ 5. notify Scheduler ─────────►│ wake Scheduler + │ │ drain completion_queue + │ │ release/poison fanout + │ │ pop directed/shared queue + │ │ resolve target (NL) or idle SUB │ │ wt.dispatch(slot_id) ──────► WorkerThread │ │ encode mailbox → spin-poll TASK_DONE │ │ (blocking; child runs the kernel) @@ -163,7 +164,8 @@ Communication channels: | Path | Mechanism | Payload | | ---- | --------- | ------- | -| Orch → Scheduler | wiring_queue (mutex + CV) | slot id | +| Orchestrator/Scheduler → ready queues | direct `enqueue_ready` queue push | slot id | +| Orchestrator/Scheduler → Scheduler loop | condition-variable notification | ready wake-up | | Scheduler → WorkerThread | WorkerThread internal queue | slot id | | WorkerThread → Scheduler | completion_queue (mutex + CV) | slot id + group index + outcome | | WorkerThread ↔ child | shm mailbox (state + error + task data) | encoded blob | diff --git a/docs/orchestrator.md b/docs/orchestrator.md index 2245134815..4c8f46e384 100644 --- a/docs/orchestrator.md +++ b/docs/orchestrator.md @@ -34,13 +34,13 @@ public: SubmitResult submit_next_level(const CallableIdentity &callable, const TaskArgs &args, const CallConfig &config, - int32_t worker = -1, + int32_t worker, const std::vector &eligible_worker_ids = {}, const RemoteTaskArgsSidecar &remote_sidecar = {}); SubmitResult submit_next_level_group(const CallableIdentity &callable, const std::vector &args_list, const CallConfig &config, - const std::vector &workers = {}, + const std::vector &workers, const std::vector> &eligible_worker_ids = {}, const std::vector &remote_sidecars = {}); SubmitResult submit_sub(const CallableIdentity &callable, @@ -75,11 +75,12 @@ Remote L3 submit adds two hidden pieces of metadata: final eligible worker-id sets and optional `RemoteTaskArgsSidecar` entries aligned by tensor index. Python `RemoteCallable` handles supply callable eligibility, and `TaskArgs.add_tensor(RemoteTensorRef(...), tag)` supplies tensor sidecars. The -Orchestrator validates affinity, worker existence, local-vs-remote +Orchestrator validates exact placement, worker existence, local-vs-remote compatibility, remote handle access rights for the tensor tag, bare host pointers, and remote null OUTPUT tensors before committing the slot. -For NEXT_LEVEL tasks, `worker`/`workers` are stable worker ids rather than -C++ worker-thread vector indices. +For NEXT_LEVEL tasks, `worker`/`workers` are required stable worker ids rather +than C++ worker-thread vector indices. SUB submit APIs expose no worker +selection and use the shared SUB ready queue. --- @@ -92,7 +93,8 @@ how the slot is set up. ```cpp SubmitResult Orchestrator::submit_next_level(const CallableIdentity &callable, TaskArgs args, - const CallConfig &config) { + const CallConfig &config, + int32_t worker) { // 1. Alloc slot (blocks on back-pressure if ring full) TaskSlot sid = ring_.alloc(); TaskSlotState &s = slots_[sid]; @@ -103,6 +105,7 @@ SubmitResult Orchestrator::submit_next_level(const CallableIdentity &callable, s.callable = callable; s.task_args = std::move(args); s.config = config; + s.target_worker_ids = {worker}; // 3. Walk task_args tags, derive dependencies // (dedup producers: same producer may appear on multiple input tensors) @@ -130,10 +133,16 @@ SubmitResult Orchestrator::submit_next_level(const CallableIdentity &callable, // 5. Register with scope (holds slot open until scope_end releases ref) scope_.register_task(sid); // increments s.fanout_total by 1 - // 6. Push fanout edges onto scheduler's wiring queue - // (Scheduler wires producer→consumer asynchronously; avoids blocking - // the Orch thread on fanout_mu) - scheduler_.enqueue_wiring(sid, std::move(producers)); + // 6. Attach fanout edges under each producer's mutex. Producers already + // completed do not count as live fanins; failed producers poison this + // slot. Route an immediately READY slot through enqueue_ready(). + attach_fanout_and_count_live_producers(sid, producers); + if (s.fanin_count == 0) { + s.state = TaskState::READY; + enqueue_ready(sid); + } else { + s.state = TaskState::PENDING; + } // 7. Return handle return {sid}; @@ -172,16 +181,21 @@ remote buffer identity and logical offset. **Step 4 — fanin count**: The number of live producers. Decremented by `fanin_released++` each time a producer completes; when `fanin_released == -fanin_count`, the slot is ready. +fanin_count`, the slot is ready. A ready NEXT_LEVEL single task is routed to +the FIFO for its required stable worker id. The same routing function is used +for immediately-ready submissions and tasks released by Scheduler dependency +processing. A ready NEXT_LEVEL group is routed to the dedicated group FIFO; +SUB tasks remain on their shared queue. **Step 5 — scope ref**: Each slot starts with one "scope reference" in its fanout_total. Without this, a task with no downstream consumer would never be reclaimable. See [§6 Scope](#6-scope). -**Step 6 — wiring queue**: Fanout edges (producer knows its consumers) are -wired **asynchronously** by the Scheduler thread. This decouples submit from -`fanout_mu` contention. See [scheduler.md](scheduler.md) §2 for the wiring -phase. +**Step 6 — fanout attachment and READY routing**: Submission synchronously +locks each producer's `fanout_mu`, attaches the consumer, and counts only live +producers. An immediately READY task is routed to its exact NEXT_LEVEL worker +FIFO, the NEXT_LEVEL group FIFO, or the shared SUB FIFO. See +[scheduler.md](scheduler.md) §1. --- @@ -192,49 +206,29 @@ Each worker gets its own `TaskArgs`; the node only reaches COMPLETED when all N finish. ```cpp -SubmitResult Orchestrator::submit_next_level_group(const CallableIdentity &callable, - std::vector args_list, - const CallConfig &config) { - TaskSlot sid = ring_.alloc(); - TaskSlotState &s = slots_[sid]; - s.reset(); - s.worker_type = WorkerType::NEXT_LEVEL; - s.callable = callable; - s.config = config; - s.group_size = args_list.size(); - s.sub_complete_count = 0; - s.task_args_list = std::move(args_list); - - // Tag walk unions all entries in args_list (any input in any member → fanin) - // Dedup both producers and outputs across all args_list entries. - std::vector producers; - std::unordered_set producers_seen; - std::unordered_set outputs_seen; - for (auto &a : s.task_args_list) { - for (int i = 0; i < a.tensor_count(); i++) { - TensorArgType tag = a.tag(i); - uint64_t ptr = a.tensor(i).data; - if (tag == INPUT || tag == INOUT) - if (auto prod = tensormap_.lookup(ptr); prod != INVALID) - if (producers_seen.insert(prod).second) - producers.push_back(prod); - if (tag == OUTPUT || tag == INOUT || tag == OUTPUT_EXISTING) - if (outputs_seen.insert(ptr).second) - tensormap_.insert(ptr, sid); - } - } - - s.fanin_count = static_cast(producers.size()); - s.fanin_released = 0; - scope_.register_task(sid); - scheduler_.enqueue_wiring(sid, std::move(producers)); - return {sid}; +SubmitResult Orchestrator::submit_next_level_group( + const CallableIdentity &callable, const std::vector &args_list, + const CallConfig &config, const std::vector &worker_ids, + const std::vector> &eligible_worker_ids, + const std::vector &remote_sidecars +) { + return submit_impl( + WorkerType::NEXT_LEVEL, callable, config, args_list, worker_ids, + eligible_worker_ids, remote_sidecars + ); } ``` -At dispatch time the Scheduler reserves `group_size` idle WorkerThreads, and -each WorkerThread runs `worker->run` with its own `task_args_list[i]`. -Completion is gated on `sub_complete_count.fetch_add(1) + 1 == group_size`. +`submit_impl` validates that `worker_ids` contains one unique, eligible target +per group member before it performs shared dependency inference and READY +routing. + +At dispatch time the Scheduler checks the group FIFO head and resolves every +entry in `workers` to that exact stable worker ID. It dispatches only if the +entire target set is idle; a blocked group reserves no partial worker set and +does not cause a scan past the FIFO head. Each WorkerThread runs `worker->run` +with its own `task_args_list[i]`. Completion remains aggregated at the group +slot, so downstream consumers are released once after every member is terminal. --- @@ -443,11 +437,11 @@ Flow: ```python def my_orch(orch, args, cfg): with orch.scope(): # ring 1 - orch.submit_next_level(chip_a, a_args, cfg) - orch.submit_next_level(chip_b, b_args, cfg) + orch.submit_next_level(chip_a, a_args, cfg, worker=0) + orch.submit_next_level(chip_b, b_args, cfg, worker=1) # Inner tasks are now eligible for reclamation on ring 1, # without waiting for any outer-scope task. - orch.submit_next_level(chip_c, c_args, cfg) # ring 0 (outer) + orch.submit_next_level(chip_c, c_args, cfg, worker=0) # ring 0 ``` `with orch.scope():` is the recommended form. Raw `orch.scope_begin()` / @@ -722,8 +716,8 @@ instead of stalling forever. Default timeout: 10 s. - [hierarchical_level_runtime.md](hierarchical_level_runtime.md) — how Orchestrator fits alongside Scheduler and Worker -- [scheduler.md](scheduler.md) — what happens to slots after they're pushed - onto the wiring queue +- [scheduler.md](scheduler.md) — READY dispatch and completion-time dependency + release - [task-flow.md](task-flow.md) — the data (Callable / TaskArgs / CallConfig) being moved by `submit_*` - [comm-domain.md](comm-domain.md) — `orch.allocate_domain` dynamic diff --git a/docs/remote-l3-worker-design.md b/docs/remote-l3-worker-design.md index f981e99e7a..66e774a110 100644 --- a/docs/remote-l3-worker-design.md +++ b/docs/remote-l3-worker-design.md @@ -43,10 +43,9 @@ Implemented: - Explicit endpoint outcomes: success, task failure, endpoint failure, and skipped group members. Failed producers poison downstream consumers instead of completing successfully. -- Stable NEXT_LEVEL `worker_id` metadata shared by local and remote - children, submit-time worker eligibility, worker affinity validation - against eligibility, and Scheduler selection from only eligible idle - workers. +- Stable NEXT_LEVEL `worker_id` metadata shared by local and remote children, + submit-time validation of the exact target against worker eligibility, and + directed Scheduler dispatch to that target only. - C++ remote tensor sidecars, remote-aware `TensorKey` values, and submit-time rejection of remote sidecars against local endpoints, bare host pointers, and remote null OUTPUT tensors without a sidecar. @@ -137,7 +136,7 @@ Relevant code paths: `MAILBOX_OFF_ERROR_MSG`. - `src/common/hierarchical/orchestrator.{h,cpp}` - `submit_next_level()` stores `TaskArgs`, `CallConfig`, - `CallableIdentity`, and optional worker affinity in a parent-side slot. + `CallableIdentity`, and the required target worker in a parent-side slot. - Dependency inference happens before dispatch from tags in `TaskArgs`. - `src/common/task_interface/task_args.h` - Process dispatch writes `[T][S][Tensor x T][uint64 x S]`. @@ -178,8 +177,8 @@ On dispatch, `WorkerThread` builds a task packet from `TaskSlotState`, calls the endpoint, reports endpoint errors, and notifies the Scheduler with an explicit success/failure outcome. -Ready queues, group dispatch, affinities, fanin/fanout, and ring release remain -in the existing runtime. The first-error-wins policy remains only as the error +Directed ready queues, exact group dispatch, fanin/fanout, and ring release +remain in the common runtime. The first-error-wins policy remains only as the error reporting policy for choosing which root error `drain()` raises. The important change is that completion is no longer implicitly success; every endpoint, including `LocalMailboxEndpoint`, must report an explicit success/failure @@ -234,9 +233,8 @@ initialization has completed. ## Worker Identity and Callable Routing Remote scheduling needs explicit callable resolver scopes and an explicit -mapping from callable identities to eligible NEXT_LEVEL workers. The current -scheduler can otherwise choose any idle worker, which is only correct when -every NEXT_LEVEL child has the same callable registry. +mapping from callable identities to eligible NEXT_LEVEL workers. Every submit +also names the exact worker that will execute the task. Required contracts: @@ -322,16 +320,14 @@ Required contracts: is confirmed or the worker is removed from eligibility and marked failed. Failed-register rollback whose cleanup cannot be confirmed marks that worker/hashid pair cleanup-uncertain and unusable for the current session. -- `TaskSlotState` stores the final eligible worker-id set for the slot. This is - the intersection of workers that can resolve the callable hashid and workers - that can access every tensor/buffer referenced by the slot. -- If the user passes `worker=worker_id`, submit-time validation checks that - the worker id is eligible for that hashid and for the slot's tensor - sidecars. -- If `worker=-1`, the Scheduler chooses only from idle workers in the - slot's eligible set. -- Group submit validates each affinity independently. Unconstrained group - members are assigned distinct idle eligible endpoints. +- `TaskSlotState` stores the exact target worker id for each slot member. + Submit-time validation checks that target against the callable's eligible + worker set — the intersection of workers that can resolve the callable hashid + and workers that can access every tensor/buffer referenced by the slot. +- `worker=worker_id` is required. Submit-time validation checks that the worker + id is eligible for that hashid and for the slot's tensor sidecars. +- Group submit requires one distinct worker id per member and validates each + target independently. - Mixed local + remote NEXT_LEVEL pools are allowed only when the callable hashid is registered on every endpoint that can receive the slot and the slot's tensors are materialized in a representation those endpoints can diff --git a/docs/remote-l3-worker-design/implementation-plan.md b/docs/remote-l3-worker-design/implementation-plan.md index 3e89815d59..a8452004d1 100644 --- a/docs/remote-l3-worker-design/implementation-plan.md +++ b/docs/remote-l3-worker-design/implementation-plan.md @@ -40,10 +40,11 @@ Status for the local PR #866 cut: - Assign each NEXT_LEVEL child a stable `worker_id`. - Store `callable hashid -> eligible worker ids` in parent runtime metadata. - - Extend submit slots with final eligible worker-id sets computed as - callable eligibility intersected with tensor/buffer data eligibility. - - Teach Scheduler/WorkerManager to pick only eligible idle workers. - - Validate `worker=` affinity against the slot's final eligible set. + - Compute final eligible worker-id sets from callable eligibility intersected + with tensor/buffer data eligibility before C++ submission. + - Validate the exact `worker=` target against that final eligible set. + - Store only the exact target on the slot; Scheduler dispatch does not + repeat eligibility selection. - Keep current docs in sync: describe local fork/shm as `LocalMailboxEndpoint` and remote L3 as a framed endpoint, not as another mailbox child loop. @@ -237,7 +238,7 @@ Status for the local PR #866 cut: | Test | Expected result | | ---- | --------------- | | Local adapter regression | Existing L3/L4 fork/shm behavior unchanged. | -| Endpoint eligibility | Scheduler never picks an ineligible endpoint. | +| Endpoint eligibility | Exact target is rejected before enqueue when ineligible. | | Frame fuzz/bounds | Corrupt lengths and counts are rejected. | | Remote sim hello | Parent bootstraps remote L3 and shuts down cleanly. | | Manifest handoff | Runner reads manifest before transport starts. | diff --git a/docs/remote-l3-worker-design/pr-split-and-audit-artifacts.md b/docs/remote-l3-worker-design/pr-split-and-audit-artifacts.md index efecc9e169..d30b2776fb 100644 --- a/docs/remote-l3-worker-design/pr-split-and-audit-artifacts.md +++ b/docs/remote-l3-worker-design/pr-split-and-audit-artifacts.md @@ -67,7 +67,7 @@ been created by this audit. | --- | ------ | ----- | -------------- | ----- | --------- | -- | | R1 Endpoint abstraction keeps local mailbox local-only and routes remote L3 through `WorkerEndpoint` / `RemoteL3Endpoint` | `remote-l3-worker-design.md` §Target Architecture; `implementation-plan.md` step 1; `worker-manager.md` §4 | required in this PR cut | `WorkerEndpoint`, `LocalMailboxEndpoint`, `RemoteL3Endpoint`, `WorkerManager::add_next_level_endpoint`, `Worker::add_remote_l3_socket` | `test_remote_endpoint`, `test_scheduler`, `test_remote_sim_noop_task_roundtrip` | None found | PR 4 / PR 5 | | R2 Endpoint outcomes distinguish success, task failure, endpoint failure, and skipped group members | `remote-l3-worker-design.md` §Failure Semantics; `implementation-plan.md` steps 1, 4 and Failure Poisoning Contract; `scheduler.md` §§2,6,9 | required in this PR cut | `EndpointOutcome`; endpoint completion mapping; scheduler failure poisoning and group skip state | `RemoteTaskErrorMapsToTaskFailure`, `FailedProducerPoisonsDependentTask`, `GroupFailureWaitsForRunningMembersThenConsumes`, remote sim error/exit tests | None found | PR 4 / PR 5 | -| R3 Scheduler dispatch uses stable worker ids and final eligible worker-id sets; `worker=` affinity is validated | `remote-l3-worker-design.md` §Worker Identity and Callable Routing; `implementation-plan.md` step 2; `scheduler.md` §5 | required in this PR cut | `Orchestrator::validate_worker_eligibility`, `WorkerManager::pick_idle`, Python worker-id set intersection | `EndpointEligibilityRestrictsIdleSelection`, `AffinityMustBeInEligibleEndpointSet`, Python remote callable worker-id intersection tests | None found | PR 4 | +| R3 Scheduler dispatch preserves stable worker ids and `worker=` is validated against final eligibility | `remote-l3-worker-design.md` §Worker Identity and Callable Routing; `implementation-plan.md` step 2; `scheduler.md` §3 | required in this PR cut | `Orchestrator::validate_worker_eligibility`, directed NEXT_LEVEL queues, Python worker-id set intersection | `TargetMustBeInEligibleEndpointSet`, target-ID mapping tests, Python remote callable worker-id intersection tests | None found | PR 4 | | R4 Mixed local/remote pools are allowed only when callable and tensor representations are consumable by the selected worker | `remote-l3-worker-design.md` §Worker Identity and Callable Routing; `buffers-and-transports.md` §TaskArgs Sidecar Contract | required in this PR cut | Python `Orchestrator` only allows `RemoteTensorRef` for `RemoteCallable`; C++ rejects remote sidecars on local workers and non-owner remote-device dispatch without import | `RemoteSidecarRejectsLocalEndpointEligibility`, `RemoteSidecarRejectsNonOwnerEligibleEndpointWithoutImport`, Python RemoteCallable sidecar tests | No blocking drift; add a future end-to-end mixed local+remote smoke during split validation | PR 4 | | R5 Remote sidecars are hidden metadata aligned by tensor index; local endpoints reject sidecars; remote endpoints reject bare host pointers | `remote-l3-worker-design.md` §Remote TaskArgs Representation; `buffers-and-transports.md` §§Public Memory API, TaskArgs Sidecar Contract; `protocol.md` §TASK Payload | required in this PR cut | `RemoteTaskArgsSidecar`, Python `_remote_sidecar_for`, C++ `validate_remote_sidecars`, `LocalMailboxEndpoint::run`, `RemoteL3Endpoint::build_task_payload` | `TestRemoteTaskArgsSidecar`, `RemoteBarePayloadFailsBeforeSlotCommit`, `BareHostPointerWithoutSidecarIsEndpointFailure` | None found | PR 4 / PR 6 | | R6 Remote null `OUTPUT` tensors fail fast unless the caller supplies an explicit `RemoteTensorRef` | `remote-l3-worker-design.md` §Remote TaskArgs Representation; `buffers-and-transports.md` §Remote OUTPUT Allocation Policy; `implementation-plan.md` step 3 | required in this PR cut | `Orchestrator::validate_remote_sidecars` requires sidecar for remote OUTPUT; `reserve_outputs_and_slot` skips local HeapRing only when sidecar present | `RemoteOutputSidecarSkipsLocalAutoAllocAndRegistersRemoteKey`, remote sim OUTPUT buffer tests | No code drift; add a narrow negative test for null remote OUTPUT without sidecar when carving PR 4 | PR 4 / PR 6 | diff --git a/docs/remote-l3-worker-design/pr-split-and-audit-plan.md b/docs/remote-l3-worker-design/pr-split-and-audit-plan.md index da94ae15f2..b082934c37 100644 --- a/docs/remote-l3-worker-design/pr-split-and-audit-plan.md +++ b/docs/remote-l3-worker-design/pr-split-and-audit-plan.md @@ -283,8 +283,8 @@ Scope: Acceptance criteria: -- Scheduler chooses only eligible idle workers. -- Worker affinity is validated against eligibility. +- The Orchestrator validates each exact target against eligibility. +- Scheduler dispatch preserves the validated stable worker ID. - Group partial failure and downstream poison are tested. - Slot release and `drain()` behavior are correct after success and failure. diff --git a/docs/scheduler.md b/docs/scheduler.md index cf1f1d040a..ad119c33ce 100644 --- a/docs/scheduler.md +++ b/docs/scheduler.md @@ -1,406 +1,192 @@ # Scheduler — DAG Dispatch Internals -The Scheduler is the **DAG executor**. A dedicated C++ thread that consumes -submitted slots, wires fanout edges, dispatches ready tasks to worker threads, -and handles completion callbacks. It is the bridge between the Orchestrator -(producer of DAG nodes) and the WorkerManager (consumer of ready nodes). +The Scheduler is the single-threaded DAG executor for one hierarchical +`Worker`. The Orchestrator constructs slots and dependencies; the Scheduler +dispatches READY slots, processes worker completions, releases downstream +dependencies, and retires terminal slots. -For the high-level role of the Scheduler among the three engine components, -see [hierarchical_level_runtime.md](hierarchical_level_runtime.md). For the DAG -construction side (what feeds the Scheduler), see -[orchestrator.md](orchestrator.md). For dispatch mechanics (how -`WorkerThread::dispatch` actually runs a task), see -[worker-manager.md](worker-manager.md). +See [orchestrator.md](orchestrator.md) for submission and dependency inference, +and [worker-manager.md](worker-manager.md) for endpoint execution. ---- +## 1. Queue topology -## 1. Role +The final queue layout separates directed NEXT_LEVEL work from freely +scheduled SUB work: -The Scheduler's job: +```text +NextLevelReadyQueues +├── group FIFO +└── stable worker id -> single-task FIFO -- Drain the **wiring queue** (Phase 0): wire fanout edges for newly - submitted slots; if all producers are already done, promote to the ready queue. -- Drain the **ready queue** (Phase 1): for each ready slot, pick an idle - `WorkerThread` from the appropriate pool and hand off. -- Drain the **completion queue** (Phase 2): for each worker completion, - transition the slot to `COMPLETED` or `FAILED`, release fanout references, - wake or poison downstream consumers, and (if all refs released) retire the - ring slot. +ReadyQueue +└── shared SUB FIFO -One Scheduler per `Worker` instance, one thread per Scheduler. The Scheduler -**does not inspect task data** — it moves slot ids between queues and -consults scheduling metadata (`fanin_count`, `fanout_consumers`, `state`). - ---- - -## 2. The queues - -```cpp -class Scheduler { - // Producer: Orchestrator.submit_*. Consumer: Scheduler's own loop, Phase 0. - LockFreeQueue wiring_queue_; // {slot, producers} - - // Strict-4 — per-worker-type ready queues. - // Producers: Orchestrator.submit_* (routes by slot.worker_type) + - // Scheduler Phase 0 / Phase 2 (fanout-released; routes by - // consumer worker_type). - // Consumer: Scheduler's own loop, Phase 1 (one drain loop per queue, - // each with its own head-of-line break). - ReadyQueue *ready_next_level_queue_; - ReadyQueue *ready_sub_queue_; - - // Producer: WorkerThread (on endpoint->run() return). - // Consumer: Scheduler's own loop, Phase 2. - std::queue completion_queue_; -}; +Scheduler +└── completion FIFO ``` -### Wiring queue - -Introduced so that `Orchestrator::submit_*` does not need to acquire -`fanout_mu` on every producer slot at submit time (see -[orchestrator.md](orchestrator.md) §2 step 6). - -Each entry: +`Orchestrator::enqueue_ready` is the only READY router: -```cpp -struct WiringEntry { - TaskSlot consumer; - std::vector producers; // producers this consumer depends on -}; +```text +NEXT_LEVEL single -> next_level.single[target_worker_id] +NEXT_LEVEL group -> next_level.group +SUB -> ready_sub ``` -### Ready queue - -Slots whose `fanin_count == fanin_released` are ready to dispatch. The queue -holds just the slot id; dispatch reads task data from the -`ring.slot_state(sid)` pool. - -**Strict-4 — per-worker-type split.** In practice the ready queue is two -`ReadyQueue` instances, one per `WorkerType`: +Both immediately-ready submissions and dependency-released consumers use +this function. A directed task therefore cannot re-enter a shared +NEXT_LEVEL queue after waiting in PENDING. -```cpp -ReadyQueue ready_next_level_queue_; // WorkerType::NEXT_LEVEL tasks -ReadyQueue ready_sub_queue_; // WorkerType::SUB tasks -``` +Each `ReadyQueue` is a mutex-protected non-blocking FIFO. Root submission, +worker completion, and stop requests notify the Scheduler condition variable; +its wait predicate checks the completion FIFO, every ready queue, and the stop +flag. Ready queues have no blocking pop or shutdown state. -Matching L2's per-shape ready queues (the shared MPMC `ready_queues[]` split -by AIC / AIV / MIX), with the L3+ exception that we use `std::queue` -(Allowed Exception 3: dynamic data structures on host) and only two -worker types (Allowed Exception 2: `NEXT_LEVEL` + `SUB` at L3+, not -AIC / AIV / MIX). `Orchestrator::submit_*` routes each slot to the queue -matching `slot.worker_type`; `Scheduler::on_task_complete` routes a -newly-ready consumer the same way, based on the *consumer's* worker -type. `Scheduler::dispatch_ready` drains each queue with its own -head-of-line break so a saturated pool of one type cannot stall dispatch -for the other. +## 2. Scheduler loop -### Completion queue +The Scheduler drains completions before dispatching new work: -Endpoint completions whose worker returned or failed. Each -`WorkerCompletion` carries `{slot, group_index, outcome, error_message}`; -`outcome` is success, task failure, endpoint failure, or skipped. The -Scheduler runs completion handling (fanout release, downstream wake/poison, -try_consume) in its own thread so that WorkerThreads can immediately return to -their next task. +```cpp +while (true) { + wait_until_completion_ready_or_stop(); ---- + while (completion_queue has an item) { + on_task_complete(item); + } -## 3. Scheduler loop (pseudocode) + dispatch_next_level_group(); + dispatch_next_level_singles(); + dispatch_sub_ready(); -```cpp -void Scheduler::run() { - while (running_) { - // Phase 0: wiring - WiringEntry w; - while (wiring_queue_.try_pop(w)) { - wire_fanout(w); // see §4 - } - - // Phase 1: dispatch (drains BOTH per-type queues; see §5) - dispatch_ready(); - - // Phase 2: completion - WorkerCompletion c; - while (completion_queue_.try_pop(c)) { - on_task_complete(c); // see §6 - } - - // If all three queues empty, block on a condition variable until - // any producer signals work. - wait_for_work(); + if (stop_requested && all workers are idle) { + drain final completions; + dispatch one final pass; + break; } } ``` -Phase order matters: +`loop_mutex` covers completion and dispatch slot access. `Orchestrator::drain` +uses the same mutex while releasing/resetting slots, preventing slot reuse +while the Scheduler still holds a reference. -- Wiring before dispatch: a task may become ready during wiring (all its - producers already completed); wiring promotes it to ready_queue in the - same Scheduler iteration. -- Dispatch before completion: dispatch the backlog first to keep workers - busy; completion handling is not time-critical (fanout release just - queues more work for the next iteration). +## 3. Directed NEXT_LEVEL dispatch ---- +### Single tasks -## 4. Phase 0 — wiring - -```cpp -void Scheduler::wire_fanout(const WiringEntry &w) { - TaskSlot csid = w.consumer; - TaskSlotState &c = slots_[csid]; - int32_t actual_live = 0; - - for (TaskSlot psid : w.producers) { - TaskSlotState &p = slots_[psid]; - std::lock_guard lk(p.fanout_mu); - // COMPLETED producers are already done; FAILED producers poison this - // consumer instead of making it ready. - if (p.state.load() == TaskState::COMPLETED || - p.state.load() == TaskState::CONSUMED) continue; - if (p.state.load() == TaskState::FAILED) { - poison_task(csid, p.failure_message); - continue; - } - p.fanout_consumers.push_back(csid); - p.fanout_total++; - actual_live++; - } +Every single task contains one exact stable `target_worker_id`. For each +registered NEXT_LEVEL worker, the Scheduler checks only that worker's FIFO: - // Update consumer's fanin to the actual live count (producers already - // finished don't count). - c.fanin_count = actual_live; - if (actual_live == 0) { - // Strict-4: wiring promotes directly to the per-type queue. - auto *q = (c.worker_type == WorkerType::NEXT_LEVEL) ? ready_next_level_queue_ - : ready_sub_queue_; - q->push(csid); - } -} +```text +if target worker is idle and its FIFO is non-empty: + pop FIFO head + READY -> RUNNING + dispatch to that exact worker ``` -**Race with completion**: a producer may finish between submit and wiring. -The `lock_guard(p.fanout_mu)` + `p.state.load()` check ensures we either: - -- wire an edge and the producer's future completion will fire `fanin_released++` - for this consumer, or -- see "already completed" and skip, correctly counting this producer as not - contributing to fanin. - ---- +There is no idle-worker search, rebinding, work stealing, or scan into another +worker's queue. FIFO is independent per worker, so a busy worker A does not +block READY work for worker B. -## 5. Phase 1 — dispatch +### Group tasks -`dispatch_ready` drains each per-type ready queue with its own -head-of-line break so one saturated pool cannot stall the other: +A group is one DAG node with one exact stable worker ID per member. The +Scheduler examines only the group FIFO head: -```cpp -void Scheduler::dispatch_ready() { - auto drain_one = [&](ReadyQueue *q) { - TaskSlot slot; - while (q->try_pop(slot)) { - TaskSlotState &s = slots_[slot]; - int N = s.group_size(); // 1 for single-task slots - - std::vector workers; - for (int i = 0; i < N; i++) { - WorkerThread *wt = nullptr; - int32_t affinity = s.get_affinity(i); - if (affinity >= 0) { - wt = (s.worker_type == WorkerType::NEXT_LEVEL) - ? manager_->get_worker_by_id(s.worker_type, affinity) - : manager_->get_worker_by_index(s.worker_type, affinity); - if (wt && (!wt->idle() || !s.worker_allowed(i, wt->worker_id()))) - wt = nullptr; - } else { - wt = manager_->pick_idle( - s.worker_type, workers, s.eligible_workers_for(i)); - } - if (!wt) break; - workers.push_back(wt); - } - if (static_cast(workers.size()) < N) { - q->push(slot); // put back; try again after a completion - break; - } - s.state.store(TaskState::RUNNING); - for (int i = 0; i < N; i++) { - workers[i]->dispatch({slot, i}); - } - } - }; - drain_one(ready_next_level_queue_); - drain_one(ready_sub_queue_); -} +```text +resolve every target worker +if every target is idle: + pop the group + initialize all member states + READY -> RUNNING + dispatch member i to target_worker_ids[i] +else: + leave the group at the FIFO head ``` -Dispatch hands off a `WorkerDispatch {slot, group_index}` to a -`WorkerThread`. The WorkerThread reads -`ring.slot_state(slot).{callable, task_args, config}` on its own thread -and encodes it into the per-WT mailbox — see -[worker-manager.md](worker-manager.md) §3 for the dispatch protocol. +The check is all-or-nothing: a blocked group reserves no partial worker set. +The Scheduler does not scan later groups. It continues to single-task queues, +so the runtime adds no fairness, aging, priority, or reservation policy beyond +trying a launchable group before singles in each iteration. Users are +responsible for choosing worker sets that make acceptable progress. -**Pick-idle back-pressure**: when the manager cannot provide enough idle -workers that also satisfy affinity and worker eligibility, the slot is -pushed back onto *its* queue and that queue's drain halts; the other-type -queue's drain continues. The ring's back-pressure at the Orch side already -caps the total number of in-flight tasks across both types. +## 4. SUB dispatch -Worker eligibility is opaque scheduling metadata. The Scheduler compares -worker ids and capability bits exposed through `WorkerEndpoint::caps()`, but -does not inspect HCOMM, RDMA, socket, or remote buffer internals. -For NEXT_LEVEL affinity, `s.get_affinity(i)` is a stable worker id and can -be different from the `next_level_threads_` vector index. SUB affinity is not -public and keeps the internal index semantics. +SUB has no public worker-ID selection. All READY SUB tasks share one FIFO. +For a single task the Scheduler chooses any idle SUB worker. For a group it +chooses the required number of distinct idle SUB workers; if there are not +enough, it returns the task to the queue and stops that drain pass. ---- +This is intentionally separate from NEXT_LEVEL placement. SUB free scheduling +does not provide a compatibility path for omitted NEXT_LEVEL targets. -## 6. Phase 2 — completion +## 5. Dependency release and PENDING updates -Called by `WorkerThread::on_complete_(completion)` which pushes to -`completion_queue_`. The Scheduler then: +Submission records each live producer in the consumer's `fanin_count`. A +consumer with live producers enters PENDING and is not placed on a ready +queue. When a producer completes successfully, the Scheduler increments each +consumer's `fanin_released`: ```cpp -void Scheduler::on_task_complete(const WorkerCompletion &completion) { - TaskSlot sid = completion.task_slot; - TaskSlotState &s = slots_[sid]; - - // Group tasks aggregate per-member outcomes before the slot is terminal. - if (s.group_size > 0) { - if (!record_group_member_completion(completion)) return; - } - - bool failed = completion.outcome != EndpointOutcome::SUCCESS; - s.state.store(failed ? TaskState::FAILED : TaskState::COMPLETED); - - // Release fanout refs on downstream consumers - std::vector consumers; - { - std::lock_guard lk(s.fanout_mu); - consumers = s.fanout_consumers; // snapshot (mutex protects vector) +if (++consumer.fanin_released >= consumer.fanin_count) { + if (CAS(consumer.state, PENDING, READY)) { + enqueue_ready_cb(consumer_slot); } - for (TaskSlot csid : consumers) { - if (failed) { - poison_task(csid, completion.error_message); - continue; - } - TaskSlotState &c = slots_[csid]; - if (++c.fanin_released == c.fanin_count) { - // Strict-4: push to the queue matching the *consumer's* - // worker type. A consumer of a NEXT_LEVEL producer can itself - // be SUB, so we pick based on `c.worker_type`, not `s`. - auto *q = (c.worker_type == WorkerType::NEXT_LEVEL) ? ready_next_level_queue_ - : ready_sub_queue_; - q->push(csid); - } - } - - // Also: this task itself may now be CONSUMED - try_consume(sid); } ``` -### `try_consume` +The compare-and-swap ensures that exactly one producer completion performs +the PENDING-to-READY transition. The callback routes by the consumer's own +type and exact target, not by the producer. -```cpp -void Scheduler::try_consume(TaskSlot sid) { - TaskSlotState &s = slots_[sid]; - if (s.state.load() != TaskState::COMPLETED && - s.state.load() != TaskState::FAILED) return; - if (s.fanout_released.load() != s.fanout_total) return; - - s.state.store(TaskState::CONSUMED); - - // Erase tensormap entries this task produced - for (int i = 0; i < s.task_args.tensor_count(); i++) { - // only erase entries still pointing at this slot - uint64_t ptr = s.task_args.tensor(i).data; - if (orchestrator_->tensormap_lookup(ptr) == sid) - orchestrator_->tensormap_erase(ptr); - } +If a producer fails, `poison_task` moves not-yet-running downstream consumers +to FAILED instead of READY. Poisoned tasks are never dispatched but still +release references and retire normally. - // Return slot to ring pool - ring_->release(sid); - s.state.store(TaskState::FREE); -} -``` +## 6. Completion and group aggregation -Scope release (when `scope_end` runs) calls back into the Scheduler to bump -`fanout_released` by 1 on each scope-registered slot, triggering -`try_consume`. This is how leaf tasks get reclaimed. +Each `WorkerThread` reports a `WorkerCompletion` containing the slot, +`group_index`, outcome, and error message. A single-task completion goes +directly to the Scheduler completion FIFO. ---- +Group members update per-member terminal state under `group_mu`. Only when all +members are terminal does `worker_done` enqueue one aggregate completion for +the group slot. The first member failure determines the stored failure; +members already running finish, and not-yet-dispatched SUB members are marked +skipped. Directed NEXT_LEVEL groups are launched as a complete set. -## 7. Start / Stop +On aggregate completion the Scheduler: -```cpp -void Scheduler::start(Config cfg) { - manager_ = cfg.manager; - orchestrator_ = cfg.orchestrator; - running_.store(true); - thread_ = std::thread([this] { run(); }); -} +1. moves the slot to COMPLETED or FAILED; +2. releases or poisons downstream consumers; +3. releases references held on this task's producers; +4. calls `try_consume` when all fanout references are released. -void Scheduler::stop() { - running_.store(false); - wake(); - thread_.join(); -} -``` +The Orchestrator's consume callback erases TensorMap entries that still point +to the slot, releases its Ring storage, and decrements the active-task count. -`Worker::init` calls `start` after all children are registered and -`WorkerManager::start` has spawned the WorkerThread pool. +## 7. Lifecycle and invariants ---- +`Worker::init` registers and starts WorkerManager endpoints first, freezes the +stable NEXT_LEVEL worker-ID queue map, initializes the Orchestrator, and then +starts the Scheduler. `Scheduler::stop` requests termination, wakes the loop, +waits for workers to become idle, drains final completions, and joins its +thread. -## 8. Completion channel from WorkerThread +The scheduling invariants are: -```cpp -// In WorkerThread, after endpoint->run() returns: -void WorkerThread::loop() { - for (;;) { - TaskSlot sid = queue_.pop(); - WorkerCompletion c = endpoint_->run(ring_, {sid, group_index}); - scheduler_->completion_queue_.push(c); // notify Scheduler - } -} -``` +1. A NEXT_LEVEL slot always has exactly one target per member. +2. A NEXT_LEVEL single is present only in its target worker's FIFO. +3. A NEXT_LEVEL group is present only in the group FIFO and launches only on + its complete target set. +4. SUB slots never carry target-worker metadata. +5. Only the Scheduler calls `WorkerThread::dispatch`. +6. Only one successful PENDING-to-READY transition enqueues a consumer. +7. A group produces one aggregate DAG completion regardless of member count. + +## 8. Related documents -The completion path is one-way and asynchronous: the WorkerThread returns to -its own queue immediately, and the Scheduler handles completion in its own -loop. This keeps worker dispatch latency bounded by dispatch cost alone, not -by completion-handling cost. - ---- - -## 9. Invariants - -1. **Scheduler is single-threaded**: all three phase handlers run in the - Scheduler's own thread. Atomics/mutexes on slot state are only needed for - Orch/WorkerThread ↔ Scheduler coordination. -2. **Slot transitions are monotonic**: success follows - `FREE → PENDING → READY → RUNNING → COMPLETED → CONSUMED`; failure follows - `FREE → PENDING/READY/RUNNING → FAILED → CONSUMED`. -3. **Dispatch consumes one ready entry**: every `ready_queue.push` is - matched by exactly one `pick_idle + dispatch`. Group tasks push once, - dispatch N times via `pick_n_idle`. -4. **Completion is per-worker for groups**: `worker_done` is called - `group_size` times; only the terminal aggregate pushes one slot completion. - If any member fails, not-yet-dispatched members become skipped and already - running members are allowed to finish. -5. **Failed producers poison consumers**: consumers of a failed producer move - to `FAILED`, are never dispatched, and still run normal cleanup. -6. **`try_consume` is idempotent on CONSUMED**: a repeated call after - CONSUMED is a no-op. - ---- - -## 10. Related - -- [hierarchical_level_runtime.md](hierarchical_level_runtime.md) — high-level - three-component picture -- [orchestrator.md](orchestrator.md) — the producer feeding the wiring queue -- [worker-manager.md](worker-manager.md) — where dispatched slots go -- [task-flow.md](task-flow.md) — the data (Callable / TaskArgs / CallConfig) - that the Scheduler moves around, opaquely, by slot id +- [hierarchical_level_runtime.md](hierarchical_level_runtime.md) +- [orchestrator.md](orchestrator.md) +- [worker-manager.md](worker-manager.md) +- [task-flow.md](task-flow.md) diff --git a/docs/task-flow.md b/docs/task-flow.md index da6fa9fb88..64509c2433 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -314,13 +314,12 @@ and calls `submit_next_level` / `submit_sub`. These Python methods return ```python class Orchestrator: - # `worker=-1` means unconstrained. Non-negative affinities are stable - # NEXT_LEVEL worker ids. For local Python Worker children and remote - # L3 dispatch, these are returned by add_worker(...) or + # NEXT_LEVEL placement is required. For local Python Worker children and + # remote L3 dispatch, stable ids are returned by add_worker(...) or # add_remote_worker(...). For L3 ChipCallable dispatch, worker ids are # the existing chip worker ids. - def submit_next_level(self, handle, args, config=None, *, worker=-1) -> None: ... - def submit_next_level_group(self, handle, args_list, config=None, *, workers=None) -> None: ... + def submit_next_level(self, handle, args, config=None, *, worker) -> None: ... + def submit_next_level_group(self, handle, args_list, config=None, *, workers) -> None: ... def submit_sub(self, handle, args=None) -> None: ... def submit_sub_group(self, handle, args_list) -> None: ... ``` @@ -343,9 +342,14 @@ fanout wiring), see [orchestrator.md](orchestrator.md). ## 7. Data flow through dispatch -After the scheduler picks an idle `WorkerThread` and calls `wt->dispatch(sid)`, -the parent-side WorkerThread encodes `(callable digest, CallConfig, TaskArgs)` -into the per-WT shm mailbox and the forked child decodes it: +For local endpoints, after the Scheduler resolves the submitted NEXT_LEVEL +target (or chooses an idle SUB worker), `LocalMailboxEndpoint` encodes +`(callable digest, CallConfig, TaskArgs)` into the per-worker shm mailbox and +the forked child decodes it. Remote NEXT_LEVEL dispatch through +`RemoteL3Endpoint` serializes the same logical payload into a framed TASK +request instead. + +Local mailbox path: ```text slot.callable.digest ─┐ @@ -492,7 +496,7 @@ L4 parent process | ---- | ----- | ------------ | | 1 | L4 parent Python | `w4.run(my_l4_orch)` → `scope_begin` → `my_l4_orch(orch4, ...)` | | 2 | L4 `Orchestrator.submit_next_level` | the L3 callable handle digest is stored in the slot's callable identity; slot pushed to L4's ready queue | -| 3 | L4 Scheduler | pop slot; pick idle WorkerThread → the L3 child's mailbox | +| 3 | L4 Scheduler | pop the target worker's FIFO → that L3 child's mailbox | | 4 | L4 WorkerThread (PROCESS) | encode `(callable digest, config, args_blob)` into mailbox; write `TASK_READY`; spin-poll | | 5 | L3 child `_child_worker_loop` | wake on `TASK_READY`; read digest → child-local slot → `my_l3_orch` | | 6 | L3 child | `inner_worker.run(my_l3_orch, args, cfg)` → `scope_begin` → `my_l3_orch(orch3, ...)` | @@ -527,7 +531,7 @@ def my_orch(orch, view, cfg): chip_args = TaskArgs() for i in range(view.tensor_count): chip_args.add_tensor(view.tensors[i], IN if i < 2 else OUT) - orch.submit_next_level(chip_kernel_handle, chip_args, cfg) + orch.submit_next_level(chip_kernel_handle, chip_args, cfg, worker=0) w3 = Worker(level=3, child_mode=PROCESS) w3.add_worker(NEXT_LEVEL, chip_worker_0) @@ -543,7 +547,7 @@ Step-by-step (one chip worker): | 1 | parent Python | user builds `args: TaskArgs`, calls `w3.run(my_orch, args, config)` | | 2 | `Worker::run` | `scope_begin` → call `my_orch(&orch_, args.view(), cfg)` | | 3 | `Orchestrator::submit_next_level` | `slot = ring.alloc()`; move `chip_args` into `slot.task_args`; walk tags → `tensormap.lookup(a.data)`, `tensormap.lookup(b.data)`, `tensormap.insert(c.data, slot)`; push ready | -| 4 | Scheduler thread | pop `slot`; `wt = manager.pick_idle(NEXT_LEVEL)` (WT_chip_0); `wt->dispatch(slot)` | +| 4 | Scheduler thread | pop `slot` from worker 0's FIFO; resolve stable worker ID 0 to WT_chip_0; dispatch | | 5 | WT_chip_0 parent side | encode mailbox: write reserved callable field, `config`, digest prefix, `write_blob` of task_args; set `TASK_READY`; spin-poll | | 6 | chip_0 child process | wake on `TASK_READY`; resolve digest to local slot; `read_blob` → `view`; call `ChipWorker::run(local_slot, view, cfg)` | | 7 | `ChipWorker::run` | assemble `ChipStorageTaskArgs` POD (memcpy view); call `pto2_run_runtime(local_slot, &chip_storage, &cfg)` | diff --git a/docs/worker-manager.md b/docs/worker-manager.md index 7be9c19d1c..976c137116 100644 --- a/docs/worker-manager.md +++ b/docs/worker-manager.md @@ -50,11 +50,10 @@ public: void stop(); // Scheduler API - WorkerThread *get_worker_by_index(WorkerType type, int worker_index) const; WorkerThread *get_worker_by_id(WorkerType type, int32_t worker_id) const; - WorkerThread *pick_idle(WorkerType type, - const std::vector &exclude, - const std::vector &eligible_worker_ids) const; + std::vector next_level_worker_ids() const; + WorkerThread *pick_idle_sub_excluding( + const std::vector &exclude) const; private: struct LocalNextLevelEntry { @@ -77,12 +76,15 @@ the public worker id from the local worker vector index. - **Pool ownership**: two `std::vector` pools, sized at init from `add_*` calls -- **Idle selection & eligibility**: `pick_idle(type, exclude, eligible_worker_ids)` - finds an idle WorkerThread whose queue is empty, skipping any in `exclude` - and (when `eligible_worker_ids` is non-empty) restricted to that set; returns - nullptr if none available. Remote-aware NEXT_LEVEL slots carry final eligible - worker ids, so a task cannot land on a worker that lacks the callable or - tensor sidecars. +- **Directed NEXT_LEVEL lookup**: `get_worker_by_id` resolves the exact stable + target selected by the user; the Scheduler never asks the manager to choose + another NEXT_LEVEL worker +- **SUB-only idle selection**: `pick_idle_sub_excluding` chooses an idle SUB + worker not already used by the same SUB group + +Callable and remote-buffer eligibility is validated against the exact target +during Orchestrator submission. It is not scheduling metadata and is not +stored on the task slot. --- @@ -131,7 +133,8 @@ children. `slot.callable` / `slot.task_args` / `slot.config` on each dispatch via `ring->slot_state(slot_id)`. For a group slot with `group_size() == N`, the Scheduler pushes N `WorkerDispatch` entries (one per member) onto N -idle threads; each thread's `group_index` selects which +exact target threads for NEXT_LEVEL, or N freely selected SUB threads. Each +thread's `group_index` selects which `task_args_list[i]` view to hand to the worker. There is no `WorkerPayload` — the per-dispatch carrier is just the slot id plus the group sub-index. @@ -240,8 +243,8 @@ forked worker kind still follows the existing pattern: 2. Write a child-process loop that polls the mailbox, decodes the args blob, and invokes that entry point. 3. Register the mailbox via `manager.add_next_level_at(worker_id, mailbox)` - for explicit NEXT_LEVEL worker ids, `manager.add_next_level(mailbox)` for - the legacy local default, or `manager.add_sub(mailbox)`. + for an explicit NEXT_LEVEL worker id, `manager.add_next_level(mailbox)` to + allocate the next stable local id, or `manager.add_sub(mailbox)`. Remote L3 is different. It cannot reuse the mailbox wire format because the remote side does not share virtual addresses, fork-time COW registries, POSIX @@ -313,8 +316,9 @@ and shm-safe containers for no benefit. See Alternative: N children share one dispatch queue. Rejected because: -- `WorkerThread` queue is the natural unit of backpressure — if child `i` is - slow, its queue fills up and scheduler falls back to another +- `WorkerThread` is the natural execution unit. Directed NEXT_LEVEL work waits + in child `i`'s ready FIFO if that child is busy; SUB work may use another + idle SUB child - Simpler mental model: one child = one thread that drives it - Zero contention on queue access (only one producer, one consumer per queue) diff --git a/examples/workers/l3/README.md b/examples/workers/l3/README.md index bee62b4f6a..54ce58eec1 100644 --- a/examples/workers/l3/README.md +++ b/examples/workers/l3/README.md @@ -3,8 +3,8 @@ **L3 = HOST**: one host machine that drives multiple L2 chips plus M SubWorkers (plain Python callables), coordinated by an Orchestrator running in the host process. This is where you first see the *DAG* model — you submit a -task per chip, each task carries a dependency graph via `orchestrator` APIs, -and the runtime schedules them onto available devices. +task per chip — each pinned to its target chip worker id — and each task +carries a dependency graph via `orchestrator` APIs. See [`docs/hierarchical_level_runtime.md`](../../../docs/hierarchical_level_runtime.md) for the full L0–L6 diagram and [`docs/task-flow.md`](../../../docs/task-flow.md) @@ -36,7 +36,7 @@ worker.init() # forks chip child processes + sub children, def my_orch(orch, args, cfg): # orch is the Orchestrator. Submit one task per chip + any sub work. - # orch.submit_next_level(...) schedules a ChipCallable onto a free chip. + # orch.submit_next_level(..., worker=chip_id) targets one chip. # orch.submit_sub(postprocess_handle, sub_args) schedules a Python callable. ... diff --git a/examples/workers/l3/per_task_runtime_env/README.md b/examples/workers/l3/per_task_runtime_env/README.md index 086aefccd8..a7bfcdb06b 100644 --- a/examples/workers/l3/per_task_runtime_env/README.md +++ b/examples/workers/l3/per_task_runtime_env/README.md @@ -12,17 +12,15 @@ Before this knob, every L2 dispatched from one L3 shared the process-wide `submit_next_level` gets its own `CallConfig`: Each spec sets `ring_task_window` / `ring_heap` / `ring_dep_pool` to a scalar -(broadcast to every ring) or a 4-entry list (per-ring), so the loop just sets -whichever keys the spec contains: +(broadcast to every ring) or a 4-entry list (per-ring). `_l2_config` preserves +the orchestration's base diagnostics and overrides whichever ring keys the spec +contains: ```python def orch_fn(orch, _args, _cfg): for spec in L2_TASKS: # one entry per L2 task - cfg = CallConfig() - for key in RING_FIELDS: # ring_task_window / ring_heap / ring_dep_pool - if key in spec: # value is a scalar or a 4-entry list - setattr(cfg.runtime_env, key, spec[key]) - orch.submit_next_level(chip_handle, chip_args, cfg) # per-task config + cfg = _l2_config(_cfg, spec) + orch.submit_next_level(chip_handle, chip_args, cfg, worker=0) # per-task config ``` The per-task config travels through the mailbox to the chip child, so each L2 diff --git a/examples/workers/l3/per_task_runtime_env/main.py b/examples/workers/l3/per_task_runtime_env/main.py index d8a0ce2287..905d9174af 100644 --- a/examples/workers/l3/per_task_runtime_env/main.py +++ b/examples/workers/l3/per_task_runtime_env/main.py @@ -205,7 +205,7 @@ def orch_fn(orch, _args, _cfg): chip_args.add_tensor(make_tensor_arg(host_out[i]), TensorArgType.OUTPUT_EXISTING) cfg = _l2_config(_cfg, spec) print(f"[per_task_runtime_env] submit '{spec['label']}': runtime_env={cfg.runtime_env!r}") - orch.submit_next_level(chip_handle, chip_args, cfg) + orch.submit_next_level(chip_handle, chip_args, cfg, worker=0) print(f"[per_task_runtime_env] running DAG ({len(L2_TASKS)} L2 tasks, distinct rings)...") worker.run(orch_fn, args=None, config=CallConfig()) diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 06b834a619..a74ee3dd84 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -268,10 +268,10 @@ inline void bind_worker(nb::module_ &m) { ); }, nb::arg("digest"), nb::arg("kind"), nb::arg("target_namespace"), nb::arg("args"), nb::arg("config"), - nb::arg("worker") = int32_t(-1), nb::arg("eligible_worker_ids") = std::vector{}, + nb::arg("worker"), nb::arg("eligible_worker_ids") = std::vector{}, nb::arg("remote_sidecar") = nb::none(), "Submit a NEXT_LEVEL task by registered callable digest. " - "worker= pins to a stable NEXT_LEVEL worker id (-1 = any)." + "worker= selects the exact stable NEXT_LEVEL worker id." ) .def( "submit_next_level_group", @@ -284,11 +284,10 @@ inline void bind_worker(nb::module_ &m) { ); }, nb::arg("digest"), nb::arg("kind"), nb::arg("target_namespace"), nb::arg("args_list"), nb::arg("config"), - nb::arg("workers") = std::vector{}, - nb::arg("eligible_worker_ids") = std::vector>{}, + nb::arg("workers"), nb::arg("eligible_worker_ids") = std::vector>{}, nb::arg("remote_sidecars") = nb::none(), "Submit a group of NEXT_LEVEL tasks by registered callable digest. " - "workers= per-args stable NEXT_LEVEL worker id affinity (empty = any)." + "workers= selects one exact stable NEXT_LEVEL worker id per member." ) .def( "submit_sub", diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index b013624bdd..1d8ced9e32 100644 --- a/python/simpler/orchestrator.py +++ b/python/simpler/orchestrator.py @@ -18,7 +18,7 @@ def my_orch(orch, args, cfg): a = TaskArgs() a.add_tensor(make_tensor_arg(input_tensor), TensorArgType.INPUT) a.add_tensor(make_tensor_arg(output_tensor), TensorArgType.OUTPUT) - orch.submit_next_level(chip_handle, a, cfg) # handle from Worker.register(chip_callable) + orch.submit_next_level(chip_handle, a, cfg, worker=0) sub_args = TaskArgs() sub_args.add_tensor(make_tensor_arg(output_tensor), TensorArgType.INPUT) @@ -33,6 +33,7 @@ def my_orch(orch, args, cfg): from __future__ import annotations import contextlib +import operator from collections.abc import Iterator, Sequence from typing import Any @@ -88,6 +89,19 @@ def _require_handle( return callable_or_handle.digest, callable_or_handle.kind, callable_or_handle.target_namespace, () +def _require_next_level_worker_id(value: Any, *, argument: str) -> int: + """Return an exact integer worker ID without accepting coercible values.""" + if isinstance(value, bool): + raise TypeError(f"{argument} must be an integer NEXT_LEVEL worker id") + try: + worker_id = operator.index(value) + except TypeError as exc: + raise TypeError(f"{argument} must be an integer NEXT_LEVEL worker id") from exc + if worker_id < 0: + raise ValueError(f"{argument} must be a non-negative NEXT_LEVEL worker id") + return worker_id + + def _split_next_level_args(args: TaskArgs) -> tuple[TaskArgs, _RemoteTaskArgsSidecar | None]: if isinstance(args, TaskArgs): return args, _remote_sidecar_for(args) @@ -154,17 +168,15 @@ def _expected_next_level_namespace(self) -> str | None: # User-facing submit API # ------------------------------------------------------------------ - def submit_next_level( - self, callable_handle: Any, args: TaskArgs, config: CallConfig | None = None, *, worker: int = -1 - ): + def submit_next_level(self, callable_handle: Any, args: TaskArgs, config: CallConfig | None = None, *, worker: int): """Submit a NEXT_LEVEL task by registered callable handle. ``callable_handle`` must be returned by ``Worker.register``. Tags inside ``args`` drive deps. - ``worker``: stable NEXT_LEVEL worker id for affinity - (-1 = unconstrained). For L3 chip dispatch, worker ids are the - existing chip worker ids. + ``worker`` is the exact stable NEXT_LEVEL worker id that runs the + task. For L3 chip dispatch, these are the existing chip worker ids. """ cfg = config if config is not None else CallConfig() + cpp_worker_id = _require_next_level_worker_id(worker, argument="worker") expected_namespace = ( None if isinstance(callable_handle, CallableHandle) @@ -177,6 +189,8 @@ def submit_next_level( worker=self._worker, expected_namespace=expected_namespace, ) + if target_namespace != "REMOTE_TASK_DISPATCHER" and self._worker is not None: + self._worker._require_local_next_level_target(cpp_worker_id, api="submit_next_level") c_args, explicit_remote_sidecar = _split_next_level_args(args) if target_namespace == "REMOTE_TASK_DISPATCHER": remote_sidecar = ( @@ -193,7 +207,6 @@ def submit_next_level( if target_namespace == "LOCAL_CHIP" and self._worker is not None: self._worker._stage_host_buffers_for_chip_submit(c_args) final_worker_ids = _remote_data_eligible_worker_ids(remote_sidecar, eligible_worker_ids) - cpp_worker_id = int(worker) worker = self._worker # Do the (fallible) kind4 provenance analysis BEFORE capturing remote slot # refs, so an exception here can never leave captured refs neither @@ -201,23 +214,13 @@ def submit_next_level( # is the last step before the rollback try. child_ptrs = worker._child_ptrs_in_args(c_args) if worker is not None else [] prov_guard: Any = contextlib.nullcontext() - candidates: set[int] = set() if child_ptrs and worker is not None: - candidates = self._child_dispatch_candidates(cpp_worker_id, final_worker_ids) prov_guard = worker._child_prov_lock captured_refs = worker._capture_remote_sidecar_refs(remote_sidecar) if worker is not None else [] try: with prov_guard: if child_ptrs and worker is not None: - worker._child_prov_check_dispatch(child_ptrs, candidates, api="submit_next_level") - # The child_memory arg resolved to a unique owner; pass that - # worker as the effective affinity so C++ keys the child - # TensorKey by its owner (worker_id), not the raw -1. - # Otherwise the same buffer submitted once as -1 and once as W - # yields two keys (ptr,-1) / (ptr,W) and its dependency is - # missed. The unique target is exactly what the scheduler - # would have picked, so pinning it changes no scheduling. - cpp_worker_id = next(iter(candidates)) + worker._child_prov_check_dispatch(child_ptrs, cpp_worker_id, api="submit_next_level") self._o.submit_next_level( digest, kind, target_namespace, c_args, cfg, cpp_worker_id, final_worker_ids, remote_sidecar ) @@ -228,38 +231,25 @@ def submit_next_level( if self._worker is not None: self._worker._adopt_remote_slot_refs(captured_refs) - def _child_dispatch_candidates(self, cpp_worker_id: int, eligible_ids: Any) -> set[int]: - """Resolve the set of eligible target workers for a kind4 dispatch. - - A pinned ``worker`` is the sole candidate; an unpinned ``-1`` falls back - to the callable's eligible set, or the full next-level pool when the - callable is unconstrained. ``_child_prov_check_dispatch`` rejects any - result that is not a single unique target. - """ - if cpp_worker_id >= 0: - return {cpp_worker_id} - if eligible_ids: - return {int(w) for w in eligible_ids} - if self._worker is None: - return set() - return set(self._worker._next_level_target_ids()) - def submit_next_level_group( # noqa: PLR0912 -- linear per-member sidecar + eligibility + kind4-provenance passes, one branch each self, callable_handle: Any, args_list: list, config: CallConfig | None = None, *, - workers: list | None = None, + workers: list, ): """Submit a group of NEXT_LEVEL tasks (N TaskArgs → N worker selections, 1 DAG node). - ``workers``: per-args stable NEXT_LEVEL worker ids. For L3 chip - dispatch, worker ids are the existing chip worker ids. - None/empty = all unconstrained. + ``workers`` contains the exact stable NEXT_LEVEL worker id for each + member. For L3 chip dispatch, these are the existing chip worker ids. """ cfg = config if config is not None else CallConfig() - worker_ids = [int(x) for x in workers] if workers else [] + worker_ids = [_require_next_level_worker_id(value, argument="workers entries") for value in workers] + if len(worker_ids) != len(args_list): + raise ValueError("workers length must match args_list length") + if len(set(worker_ids)) != len(worker_ids): + raise ValueError("workers must not contain duplicate NEXT_LEVEL worker ids") expected_namespace = ( None if isinstance(callable_handle, CallableHandle) @@ -272,6 +262,9 @@ def submit_next_level_group( # noqa: PLR0912 -- linear per-member sidecar + eli worker=self._worker, expected_namespace=expected_namespace, ) + if target_namespace != "REMOTE_TASK_DISPATCHER" and self._worker is not None: + for worker_id in worker_ids: + self._worker._require_local_next_level_target(worker_id, api="submit_next_level_group") c_args_list = [] explicit_remote_sidecars = [] has_explicit_remote_sidecar = False @@ -307,21 +300,18 @@ def submit_next_level_group( # noqa: PLR0912 -- linear per-member sidecar + eli if eligible_worker_ids else [] ) - cpp_worker_ids = worker_ids # Per-member kind4 dispatch guard: each member's child_memory pointers - # must resolve to that member's unique eligible target and be live there. + # must be live on that member's exact submitted target. # Run this (fallible) analysis BEFORE capturing remote slot refs, so an # exception here can never strand captured refs outside the rollback try. worker = self._worker - member_checks: list[tuple[int, list[tuple[int, int]], set[int]]] = [] + member_checks: list[tuple[list[tuple[int, int]], int]] = [] if worker is not None: for g, c_args in enumerate(c_args_list): child_ptrs = worker._child_ptrs_in_args(c_args) if not child_ptrs: continue - worker_pin = cpp_worker_ids[g] if g < len(cpp_worker_ids) else -1 - eligible_g = worker_id_sets[g] if g < len(worker_id_sets) else [] - member_checks.append((g, child_ptrs, self._child_dispatch_candidates(int(worker_pin), eligible_g))) + member_checks.append((child_ptrs, worker_ids[g])) prov_guard: Any = ( worker._child_prov_lock if (worker is not None and member_checks) else contextlib.nullcontext() ) @@ -331,43 +321,11 @@ def submit_next_level_group( # noqa: PLR0912 -- linear per-member sidecar + eli captured_refs.extend(self._worker._capture_remote_sidecar_refs(sidecar)) try: with prov_guard: - if member_checks: - # Materialise a full per-member affinity so each child member's - # resolved owner is the effective affinity C++ keys its child - # TensorKey by (see the single-submit note); non-child members - # keep their original affinity / -1. Only an empty/None - # ``workers`` is padded — a non-empty list must already be one - # per member (C++ enforces this), so padding a short one would - # silently bypass that length check. - if cpp_worker_ids: - if len(cpp_worker_ids) != len(c_args_list): - raise ValueError( - f"submit_next_level_group: workers length {len(cpp_worker_ids)} " - f"!= {len(c_args_list)} args" - ) - cpp_worker_ids = list(cpp_worker_ids) - else: - cpp_worker_ids = [-1] * len(c_args_list) - for g, child_ptrs, candidates in member_checks: + for child_ptrs, target_worker_id in member_checks: assert worker is not None # member_checks is only populated when worker is present - worker._child_prov_check_dispatch(child_ptrs, candidates, api="submit_next_level_group") - cpp_worker_ids[g] = next(iter(candidates)) - # A group dispatches its members in parallel to distinct workers. - # Two members pinned to the same owner (e.g. two child args on the - # same chip) would give the same affinity twice; the scheduler sees - # that WorkerThread idle for both and serializes them on one thread. - # This rejects that duplicate-pinned case only — full injective - # feasibility (e.g. a wildcard member left with no free worker once - # a child member is pinned) is the scheduler's pre-existing capacity - # concern, not narrowed here. - pinned = [wid for wid in cpp_worker_ids if wid >= 0] - if len(pinned) != len(set(pinned)): - raise ValueError( - f"submit_next_level_group: members resolve to duplicate target workers " - f"{cpp_worker_ids} — a group must dispatch to distinct workers" - ) + worker._child_prov_check_dispatch(child_ptrs, target_worker_id, api="submit_next_level_group") self._o.submit_next_level_group( - digest, kind, target_namespace, c_args_list, cfg, cpp_worker_ids, worker_id_sets, remote_sidecars + digest, kind, target_namespace, c_args_list, cfg, worker_ids, worker_id_sets, remote_sidecars ) except BaseException: if self._worker is not None: @@ -488,9 +446,9 @@ def create_l3_l2_queue(self, *, worker_id: int, depth: int, input_arena_bytes: i # # def my_orch(orch, args): # with orch.scope(): - # orch.submit_next_level(a, ...) - # orch.submit_next_level(b, ...) - # orch.submit_next_level(c, ...) # back on outer-scope ring + # orch.submit_next_level(a, ..., worker=0) + # orch.submit_next_level(b, ..., worker=0) + # orch.submit_next_level(c, ..., worker=0) # outer-scope ring def scope_begin(self) -> None: self._o.scope_begin() diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 475f8e7a26..3236185cd7 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -34,7 +34,7 @@ w.init() def my_orch(orch, args, cfg): - r = orch.submit_next_level(chip_handle, chip_args_ptr, cfg) + r = orch.submit_next_level(chip_handle, chip_args_ptr, cfg, worker=0) orch.submit_sub(sub_handle, sub_args) w.run(my_orch, my_args, my_config) @@ -75,7 +75,6 @@ def my_l4_orch(orch, args, config): import threading import time import uuid -from collections.abc import Sequence from dataclasses import dataclass, field from multiprocessing.shared_memory import SharedMemory from typing import Any, cast @@ -5208,45 +5207,34 @@ def _child_ptrs_in_args(args: Any) -> list[tuple[int, int]]: out.append((int(tensor.data), i)) return out - def _next_level_target_ids(self) -> Sequence[int]: - """The full pool of dispatchable next-level worker ids. - - Chip ids ``0..N`` at L3; the stable ``_next_level_worker_ids`` at L4+ (an - index range would not match the local/remote stable worker ids). - """ - if self._chip_shms: - return range(len(self._chip_shms)) - return self._next_level_worker_ids - - def _child_prov_check_dispatch( - self, child_ptrs: list[tuple[int, int]], candidate_worker_ids: Any, *, api: str - ) -> None: - """Validate every child_memory pointer against its unique target worker. - - A child_memory argument must resolve to exactly one eligible target - worker; ``0`` or ``>= 2`` candidates is ambiguous and rejected (judged on - the resolved eligibility, not the raw ``worker=-1``). The pointer must be - a live allocation on that target — else it is being routed to the wrong - worker, or is stale. - """ + def _child_prov_check_dispatch(self, child_ptrs: list[tuple[int, int]], target_worker_id: int, *, api: str) -> None: + """Validate every child_memory pointer against its exact target worker.""" if not child_ptrs: return - candidates = set(candidate_worker_ids) - if len(candidates) != 1: - arg_index = child_ptrs[0][1] - raise ValueError( - f"orch.{api}: child_memory argument (arg {arg_index}) cannot resolve a unique " - f"target worker (eligible={sorted(candidates)}); pin worker= explicitly" - ) - target = next(iter(candidates)) for ptr, arg_index in child_ptrs: - entry = self._child_alloc_prov.get((target, ptr)) + entry = self._child_alloc_prov.get((target_worker_id, ptr)) if entry is None or not entry.is_live(): raise ValueError( f"orch.{api}: child_memory argument (arg {arg_index}, ptr 0x{ptr:x}) is not a " - f"live allocation on target worker {target} (wrong worker, stale, or interior pointer)" + f"live allocation on target worker {target_worker_id} (wrong worker, stale, or interior pointer)" ) + def _require_local_next_level_target(self, worker_id: int, *, api: str) -> None: + """Reject a local callable pinned to a remote NEXT_LEVEL worker. + + A LOCAL_PYTHON / LOCAL_CHIP callable is installed only in the local + children's registries; a remote-L3 worker's manifest carries only its + dispatcher callables, so routing a local digest to a remote worker id + fails asynchronously with an unknown-hashid on the remote endpoint. The + C++ target check only rejects unregistered ids, so a registered remote + worker slips through — this guards that hole. + """ + if worker_id in set(self._remote_worker_ids): + raise ValueError( + f"orch.{api}: worker {worker_id} is a remote NEXT_LEVEL worker; a local callable " + f"must target a local child (remote workers only run RemoteCallable dispatches)" + ) + def _clear_child_prov(self) -> None: """Drop the whole child-pointer provenance table (close-path hygiene).""" with self._child_prov_lock: diff --git a/simpler_setup/scene_test.py b/simpler_setup/scene_test.py index 308d4f6b83..216a1671ef 100644 --- a/simpler_setup/scene_test.py +++ b/simpler_setup/scene_test.py @@ -415,7 +415,7 @@ def _build_l3_task_args(test_args: TaskArgsBuilder, orch_signature: list): """Build a tagged `TaskArgs` (vector-backed, with `TensorArgType` tags) from `TaskArgsBuilder`. - Used by the L3 path (`orch.submit_next_level(callable, args, config)`): + Used by the L3 path (`orch.submit_next_level(callable, args, config, worker=chip_id)`): the orchestrator reads the tags to drive dependency inference. Returns: diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md index 1f97121a89..4b88487493 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md @@ -266,7 +266,7 @@ cfg = CallConfig() cfg.runtime_env.ring_task_window = 128 # power of 2, >= 4 cfg.runtime_env.ring_heap = 262144 # bytes/ring, >= 1024 cfg.runtime_env.ring_dep_pool = 256 # 4 .. INT32_MAX -orchestrator.submit_next_level(handle, args, cfg) +orchestrator.submit_next_level(handle, args, cfg, worker=0) ``` Assign a four-entry list to tune the scope-depth rings independently. The list @@ -284,7 +284,7 @@ cfg.runtime_env.ring_heap = [ 512 * 1024 * 1024, ] cfg.runtime_env.ring_dep_pool = [4096, 8192, 16384, 32768] -orchestrator.submit_next_level(handle, args, cfg) +orchestrator.submit_next_level(handle, args, cfg, worker=0) ``` Scene tests set the same keys under a nested `runtime_env` block in the diff --git a/src/a5/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md b/src/a5/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md index 96ed560740..35b0391ffb 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md +++ b/src/a5/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md @@ -266,7 +266,7 @@ cfg = CallConfig() cfg.runtime_env.ring_task_window = 128 # power of 2, >= 4 cfg.runtime_env.ring_heap = 262144 # bytes/ring, >= 1024 cfg.runtime_env.ring_dep_pool = 256 # 4 .. INT32_MAX -orchestrator.submit_next_level(handle, args, cfg) +orchestrator.submit_next_level(handle, args, cfg, worker=0) ``` Assign a four-entry list to tune the scope-depth rings independently. The list @@ -284,7 +284,7 @@ cfg.runtime_env.ring_heap = [ 512 * 1024 * 1024, ] cfg.runtime_env.ring_dep_pool = [4096, 8192, 16384, 32768] -orchestrator.submit_next_level(handle, args, cfg) +orchestrator.submit_next_level(handle, args, cfg, worker=0) ``` Scene tests set the same keys under a nested `runtime_env` block in the diff --git a/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index 926fb35c68..dda1db8107 100644 --- a/src/common/hierarchical/orchestrator.cpp +++ b/src/common/hierarchical/orchestrator.cpp @@ -18,14 +18,14 @@ #include "worker_manager.h" void Orchestrator::init( - TensorMap *tensormap, Ring *allocator, Scope *scope, ReadyQueue *ready_next_level_queue, - ReadyQueue *ready_sub_queue, WorkerManager *manager, std::function ready_notify_cb + TensorMap *tensormap, Ring *allocator, Scope *scope, ReadyQueue *ready_sub_queue, + NextLevelReadyQueues *ready_next_level_queues, WorkerManager *manager, std::function ready_notify_cb ) { tensormap_ = tensormap; allocator_ = allocator; scope_ = scope; - ready_next_level_queue_ = ready_next_level_queue; ready_sub_queue_ = ready_sub_queue; + ready_next_level_queues_ = ready_next_level_queues; manager_ = manager; ready_notify_cb_ = std::move(ready_notify_cb); active_tasks_.store(0, std::memory_order_relaxed); @@ -151,14 +151,13 @@ SubmitResult Orchestrator::submit_next_level( const CallableIdentity &callable, const TaskArgs &args, const CallConfig &config, int32_t worker_id, const std::vector &eligible_worker_ids, const RemoteTaskArgsSidecar &remote_sidecar ) { - std::vector affinities; - if (worker_id >= 0) affinities = {worker_id}; + std::vector target_worker_ids{worker_id}; std::vector> worker_id_sets; if (!eligible_worker_ids.empty()) worker_id_sets = {eligible_worker_ids}; std::vector sidecars; if (!remote_sidecar.tensors.empty() || !remote_sidecar.inline_payload.empty()) sidecars = {remote_sidecar}; return submit_impl( - WorkerType::NEXT_LEVEL, callable, config, {args}, std::move(affinities), std::move(worker_id_sets), + WorkerType::NEXT_LEVEL, callable, config, {args}, std::move(target_worker_ids), std::move(worker_id_sets), std::move(sidecars) ); } @@ -187,12 +186,12 @@ SubmitResult Orchestrator::submit_sub_group(const CallableIdentity &callable, co SubmitResult Orchestrator::submit_impl( WorkerType worker_type, const CallableIdentity &callable, const CallConfig &config, std::vector args_list, - std::vector affinities, std::vector> eligible_worker_ids, + std::vector target_worker_ids, std::vector> eligible_worker_ids, std::vector remote_sidecars ) { if (args_list.empty()) throw std::invalid_argument("Orchestrator: args_list must not be empty"); config.validate(); - validate_worker_eligibility(worker_type, args_list.size(), affinities, eligible_worker_ids); + validate_worker_eligibility(worker_type, args_list.size(), target_worker_ids, eligible_worker_ids); validate_remote_sidecars(args_list, remote_sidecars, eligible_worker_ids); // Fail-fast: if a previously-dispatched task has already failed, abort @@ -226,13 +225,11 @@ SubmitResult Orchestrator::submit_impl( s.worker_type = worker_type; s.callable = callable; s.config = config; - s.eligible_worker_ids = std::move(eligible_worker_ids); - // --- Step 2: Walk tags → tensormap.lookup (deps) + tensormap.insert // (outputs). Must happen before we move args_list into the slot because // infer_deps reads tensor data pointers and tags from it. std::vector producers; - infer_deps(slot, args_list, affinities, remote_sidecars, producers, s.output_keys); + infer_deps(slot, args_list, target_worker_ids, remote_sidecars, producers, s.output_keys); // --- Step 3: Store TaskArgs directly (no chip-storage pre-build) --- // Dispatch builds a TaskArgsView on demand via `slot.args_view(i)` @@ -248,7 +245,7 @@ SubmitResult Orchestrator::submit_impl( s.task_args_list = std::move(args_list); s.remote_sidecars = std::move(remote_sidecars); } - s.affinities = std::move(affinities); + s.target_worker_ids = std::move(target_worker_ids); // --- Step 5: Finalize fanin — lock each producer's fanout_mu, attach --- // @@ -304,12 +301,10 @@ SubmitResult Orchestrator::submit_impl( return SubmitResult{slot}; } - // --- Step 6: If no live fanins → READY --- - // Strict-4: push to the queue dedicated to this task's worker type so a - // saturated sub pool cannot stall next-level dispatch (and vice versa). + // --- Step 6: If no live fanins → READY and route by final placement --- if (live_fanins == 0) { s.state.store(TaskState::READY, std::memory_order_release); - ready_queue_for(worker_type)->push(slot); + enqueue_ready(slot); if (ready_notify_cb_) ready_notify_cb_(); } else { s.state.store(TaskState::PENDING, std::memory_order_release); @@ -318,14 +313,36 @@ SubmitResult Orchestrator::submit_impl( return SubmitResult{slot}; } +void Orchestrator::enqueue_ready(TaskSlot slot) { + TaskSlotState &s = slot_state(slot); + if (s.worker_type == WorkerType::NEXT_LEVEL) { + if (ready_next_level_queues_ == nullptr) + throw std::runtime_error("Orchestrator::enqueue_ready: NEXT_LEVEL queues are not initialized"); + if (s.is_group()) { + ready_next_level_queues_->push_group(slot); + } else { + ready_next_level_queues_->push_single(s.target_worker_id(0), slot); + } + return; + } + ready_sub_queue_->push(slot); +} + void Orchestrator::validate_worker_eligibility( - WorkerType worker_type, size_t args_count, const std::vector &affinities, + WorkerType worker_type, size_t args_count, const std::vector &target_worker_ids, const std::vector> &eligible_worker_ids ) const { - if (!affinities.empty() && affinities.size() != args_count) { + if (worker_type == WorkerType::SUB) { + if (!target_worker_ids.empty() || !eligible_worker_ids.empty()) { + throw std::invalid_argument("Orchestrator: SUB tasks do not accept worker-selection metadata"); + } + return; + } + + if (target_worker_ids.size() != args_count) { throw std::invalid_argument( - "Orchestrator: affinity length " + std::to_string(affinities.size()) + " does not match args length " + - std::to_string(args_count) + "Orchestrator: NEXT_LEVEL target count " + std::to_string(target_worker_ids.size()) + + " does not match args length " + std::to_string(args_count) ); } if (!eligible_worker_ids.empty() && eligible_worker_ids.size() != args_count) { @@ -335,6 +352,16 @@ void Orchestrator::validate_worker_eligibility( ); } + std::unordered_set unique_targets; + for (int32_t worker_id : target_worker_ids) { + if (worker_id < 0) { + throw std::invalid_argument("Orchestrator: NEXT_LEVEL worker id must be non-negative"); + } + if (!unique_targets.insert(worker_id).second) { + throw std::invalid_argument("Orchestrator: duplicate NEXT_LEVEL worker id " + std::to_string(worker_id)); + } + } + const std::vector empty_eligible; for (size_t i = 0; i < args_count; ++i) { const auto &eligible = eligible_worker_ids.empty() ? empty_eligible : eligible_worker_ids[i]; @@ -343,51 +370,33 @@ void Orchestrator::validate_worker_eligibility( "Orchestrator: final eligible worker-id set is empty for member " + std::to_string(i) ); } - if (manager_ != nullptr && !eligible_worker_ids.empty()) { + if (manager_ != nullptr) { for (int32_t worker_id : eligible) { - if (manager_->get_worker_by_id(worker_type, worker_id) == nullptr) { + if (manager_->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id) == nullptr) { throw std::invalid_argument( "Orchestrator: eligible worker-id " + std::to_string(worker_id) + " is not a registered worker" ); } } - } - int32_t affinity = affinities.empty() ? -1 : affinities[i]; - if (affinity < 0) continue; - - if (manager_ != nullptr) { - auto *wt = worker_type == WorkerType::NEXT_LEVEL ? manager_->get_worker_by_id(worker_type, affinity) : - manager_->get_worker_by_index(worker_type, affinity); - if (wt == nullptr) { + if (manager_->get_worker_by_id(WorkerType::NEXT_LEVEL, target_worker_ids[i]) == nullptr) { throw std::invalid_argument( - "Orchestrator: worker affinity " + std::to_string(affinity) + " is not a registered worker" + "Orchestrator: target worker " + std::to_string(target_worker_ids[i]) + + " is not a registered worker" ); } - int32_t worker_id = wt->worker_id(); - bool allowed = eligible_worker_ids.empty(); - for (int32_t id : eligible) { - if (id == worker_id) { - allowed = true; - break; - } - } - if (!allowed) { - throw std::invalid_argument( - "Orchestrator: worker affinity " + std::to_string(affinity) + - " is not in the slot's final eligible worker-id set" - ); - } - } else if (affinity >= 0 && !eligible_worker_ids.empty()) { + } + + if (!eligible_worker_ids.empty()) { bool allowed = false; for (int32_t id : eligible) { - if (id == affinity) { + if (id == target_worker_ids[i]) { allowed = true; break; } } if (!allowed) { throw std::invalid_argument( - "Orchestrator: worker affinity " + std::to_string(affinity) + + "Orchestrator: target worker " + std::to_string(target_worker_ids[i]) + " is not in the slot's final eligible worker-id set" ); } @@ -528,7 +537,7 @@ AllocResult Orchestrator::reserve_outputs_and_slot( // ============================================================================= void Orchestrator::infer_deps( - TaskSlot slot, const std::vector &args_list, const std::vector &affinities, + TaskSlot slot, const std::vector &args_list, const std::vector &target_worker_ids, const std::vector &remote_sidecars, std::vector &producers, std::vector &output_keys ) { @@ -562,7 +571,7 @@ void Orchestrator::infer_deps( // reserve_outputs_and_slot before this step) // NO_DEP → skip for (size_t g = 0; g < args_list.size(); ++g) { - int32_t worker_id = (g < affinities.size()) ? affinities[g] : -1; + int32_t worker_id = (g < target_worker_ids.size()) ? target_worker_ids[g] : -1; const TaskArgs &a = args_list[g]; for (int32_t i = 0; i < a.tensor_count(); ++i) { const Tensor &t = a.tensor(i); diff --git a/src/common/hierarchical/orchestrator.h b/src/common/hierarchical/orchestrator.h index 84727e1c1e..cd6d266065 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -13,8 +13,8 @@ * Orchestrator — DAG builder. * * Public API (called by the user's orch fn during Worker::run): - * - submit_next_level(CallableIdentity, TaskArgs, CallConfig) - * - submit_next_level_group(CallableIdentity, vector, CallConfig) + * - submit_next_level(CallableIdentity, TaskArgs, CallConfig, worker_id) + * - submit_next_level_group(CallableIdentity, vector, CallConfig, worker_ids) * - submit_sub(CallableIdentity, TaskArgs) * - submit_sub_group(CallableIdentity, vector) * - alloc(shape, dtype) — runtime-owned intermediate buffer @@ -68,13 +68,10 @@ struct SubmitResult { class Orchestrator { public: - // Strict-4: the engine keeps one ReadyQueue per WorkerType so a - // saturated sub pool cannot head-of-line-block chip dispatch (and vice - // versa). Submit routes to the queue matching the task's worker_type; - // the Scheduler's dispatch_ready walks each queue independently. void init( - TensorMap *tensormap, Ring *allocator, Scope *scope, ReadyQueue *ready_next_level_queue, - ReadyQueue *ready_sub_queue, WorkerManager *manager = nullptr, std::function ready_notify_cb = {} + TensorMap *tensormap, Ring *allocator, Scope *scope, ReadyQueue *ready_sub_queue, + NextLevelReadyQueues *ready_next_level_queues, WorkerManager *manager = nullptr, + std::function ready_notify_cb = {} ); // Allocate an intermediate buffer from the Worker's HeapRing (MAP_SHARED, @@ -98,17 +95,17 @@ class Orchestrator { // by Worker.register(); the child resolves its digest to a private slot. // Tags inside `args` drive dependency inference; OUTPUT tensors with // null data are auto-allocated from the HeapRing. - // `worker_id`: stable NEXT_LEVEL worker id for affinity (-1 = unconstrained). + // `worker_id`: exact stable NEXT_LEVEL worker id that runs this task. SubmitResult submit_next_level( - const CallableIdentity &callable, const TaskArgs &args, const CallConfig &config, int32_t worker_id = -1, + const CallableIdentity &callable, const TaskArgs &args, const CallConfig &config, int32_t worker_id, const std::vector &eligible_worker_ids = {}, const RemoteTaskArgsSidecar &remote_sidecar = {} ); // Submit a group of NEXT_LEVEL tasks: N args -> N worker selections, 1 DAG node. - // `worker_ids`: per-args stable NEXT_LEVEL worker id affinity. + // `worker_ids`: one exact stable NEXT_LEVEL worker id per member. SubmitResult submit_next_level_group( const CallableIdentity &callable, const std::vector &args_list, const CallConfig &config, - const std::vector &worker_ids = {}, const std::vector> &eligible_worker_ids = {}, + const std::vector &worker_ids, const std::vector> &eligible_worker_ids = {}, const std::vector &remote_sidecars = {} ); @@ -158,24 +155,18 @@ class Orchestrator { // CAS — only the winner returns true and runs cleanup; losers return false. bool on_consumed(TaskSlot slot); + // Route a slot whose state is already READY to the queue that owns it. + // Scheduler uses the same path after releasing the final dependency. + void enqueue_ready(TaskSlot slot); + private: TensorMap *tensormap_ = nullptr; Ring *allocator_ = nullptr; Scope *scope_ = nullptr; WorkerManager *manager_ = nullptr; std::function ready_notify_cb_; - // Strict-4 per-worker-type ready queues. Each queue handles tasks of - // exactly one WorkerType so the Scheduler can dispatch from an idle pool - // without being blocked by another pool's saturation. - ReadyQueue *ready_next_level_queue_ = nullptr; ReadyQueue *ready_sub_queue_ = nullptr; - - // Returns the ready queue that owns tasks of the given worker type. - // The method itself does not mutate the Orchestrator (hence `const`); - // the returned pointer is non-const because callers push into the queue. - ReadyQueue *ready_queue_for(WorkerType t) const { - return t == WorkerType::NEXT_LEVEL ? ready_next_level_queue_ : ready_sub_queue_; - } + NextLevelReadyQueues *ready_next_level_queues_ = nullptr; // --- Drain support (owned here, not on Worker) --- std::atomic active_tasks_{0}; @@ -194,7 +185,7 @@ class Orchestrator { // can patch `tensor.data` on OUTPUT tensors flagged for auto-allocation. SubmitResult submit_impl( WorkerType worker_type, const CallableIdentity &callable, const CallConfig &config, - std::vector args_list, std::vector affinities = {}, + std::vector args_list, std::vector target_worker_ids = {}, std::vector> eligible_worker_ids = {}, std::vector remote_sidecars = {} ); @@ -214,14 +205,15 @@ class Orchestrator { // Walk the tags of each TaskArgs in `args_list`, accumulating producer // slots (for INPUT/INOUT tags) and registering outputs in the tensormap // (for OUTPUT/INOUT/OUTPUT_EXISTING tags). NO_DEP tags are skipped. - // `affinities` maps args_list[i] to worker id for TensorKey construction. + // `target_worker_ids` maps NEXT_LEVEL args_list[i] to its exact worker for + // TensorKey construction. It is empty for SUB tasks. void infer_deps( - TaskSlot slot, const std::vector &args_list, const std::vector &affinities, + TaskSlot slot, const std::vector &args_list, const std::vector &target_worker_ids, const std::vector &remote_sidecars, std::vector &producers, std::vector &output_keys ); void validate_worker_eligibility( - WorkerType worker_type, size_t args_count, const std::vector &affinities, + WorkerType worker_type, size_t args_count, const std::vector &target_worker_ids, const std::vector> &eligible_worker_ids ) const; void validate_remote_sidecars( diff --git a/src/common/hierarchical/scheduler.cpp b/src/common/hierarchical/scheduler.cpp index 8b1e1a1789..fcba6e7a65 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -11,6 +11,7 @@ #include "scheduler.h" +#include #include #include @@ -29,19 +30,25 @@ bool is_terminal_group_state(GroupMemberState state) { state == GroupMemberState::SKIPPED; } -} // namespace +void reset_group_state(TaskSlotState &state, int32_t group_size, GroupMemberState initial_member_state) { + std::lock_guard lk(state.group_mu); + state.group_member_states.assign(static_cast(group_size), initial_member_state); + state.group_member_outcomes.assign(static_cast(group_size), EndpointOutcome::SKIPPED); + state.group_terminal_count.store(0, std::memory_order_relaxed); + state.group_failed = false; + state.group_first_failure_index = -1; + state.group_first_failure_message.clear(); +} -// ============================================================================= -// Scheduler -// ============================================================================= +} // namespace // ============================================================================= // Scheduler // ============================================================================= void Scheduler::start(const Config &cfg) { - if (cfg.ring == nullptr || cfg.ready_next_level_queue == nullptr || cfg.ready_sub_queue == nullptr || - cfg.manager == nullptr) + if (cfg.ring == nullptr || cfg.ready_sub_queue == nullptr || cfg.ready_next_level_queues == nullptr || + cfg.manager == nullptr || !cfg.enqueue_ready_cb) throw std::invalid_argument("Scheduler::start: null config fields"); cfg_ = cfg; @@ -53,9 +60,6 @@ void Scheduler::start(const Config &cfg) { void Scheduler::stop() { stop_requested_.store(true, std::memory_order_release); completion_cv_.notify_all(); - // Shut down both per-type ready queues so any wait_pop waiters unblock. - cfg_.ready_next_level_queue->shutdown(); - cfg_.ready_sub_queue->shutdown(); if (sched_thread_.joinable()) sched_thread_.join(); @@ -72,7 +76,6 @@ void Scheduler::worker_done(WorkerCompletion completion) { // Group aggregation: only push to completion queue when ALL workers done if (s.is_group()) { WorkerCompletion terminal = completion; - bool group_terminal = false; { std::lock_guard lk(s.group_mu); const int32_t group_size = s.group_size(); @@ -126,7 +129,6 @@ void Scheduler::worker_done(WorkerCompletion completion) { if (s.group_terminal_count.load(std::memory_order_acquire) < group_size) return; - group_terminal = true; if (s.group_failed) { int32_t failure_index = s.group_first_failure_index; terminal.group_index = failure_index; @@ -141,7 +143,6 @@ void Scheduler::worker_done(WorkerCompletion completion) { terminal.error_message.clear(); } } - if (!group_terminal) return; completion = std::move(terminal); } @@ -167,7 +168,7 @@ void Scheduler::run() { { std::unique_lock lk(completion_mu_); completion_cv_.wait(lk, [this] { - return !completion_queue_.empty() || !cfg_.ready_next_level_queue->empty() || + return !completion_queue_.empty() || !cfg_.ready_next_level_queues->empty() || !cfg_.ready_sub_queue->empty() || stop_requested_.load(std::memory_order_acquire); }); } @@ -244,11 +245,7 @@ void Scheduler::on_task_complete(const WorkerCompletion &completion) { if (released >= cs.fanin_count) { TaskState expected = TaskState::PENDING; if (cs.state.compare_exchange_strong(expected, TaskState::READY, std::memory_order_acq_rel)) { - // Strict-4: route the freshly-ready consumer to the queue - // matching its own worker type. - auto *q = - (cs.worker_type == WorkerType::NEXT_LEVEL) ? cfg_.ready_next_level_queue : cfg_.ready_sub_queue; - q->push(consumer); + cfg_.enqueue_ready_cb(consumer); completion_cv_.notify_one(); } } @@ -338,88 +335,117 @@ void Scheduler::try_consume(TaskSlot slot) { // Dispatch — delegates to WorkerManager // ============================================================================= +// The NEXT_LEVEL worker-id set is fixed by Worker::init() (NextLevelReadyQueues +// is reset once, before scheduling starts) and every queued/targeted worker id +// is validated in Orchestrator::submit_impl. The `throw`s in the dispatch +// helpers below are therefore unreachable invariant checks; because they run on +// sched_thread_ with no surrounding handler, any throw is fatal to the whole +// worker tree (std::terminate), not a per-task failure. void Scheduler::dispatch_ready() { - // Strict-4: drain each per-type queue with its OWN head-of-line break. - // A saturated pool of one type only stalls its own queue; the other - // type continues to dispatch from its pool of idle workers. - auto drain_one = [this](ReadyQueue *q) { - TaskSlot slot; - while (q->try_pop(slot)) { - TaskSlotState &s = *cfg_.ring->slot_state(slot); - if (s.state.load(std::memory_order_acquire) != TaskState::READY) { - continue; - } - int N = s.group_size(); // 1 for normal tasks - - // Affinity-aware dispatch: pin args[i] to affinities[i] when set, - // fill remaining slots from the idle pool. NEXT_LEVEL affinity - // values are stable worker ids; SUB keeps index semantics for - // internal callers. - std::vector workers(static_cast(N), nullptr); - bool ok = true; - - // Pass 1: satisfy affinity constraints - for (int i = 0; i < N; i++) { - int32_t aff = s.get_affinity(i); - if (aff >= 0) { - auto *wt = s.worker_type == WorkerType::NEXT_LEVEL ? - cfg_.manager->get_worker_by_id(s.worker_type, aff) : - cfg_.manager->get_worker_by_index(s.worker_type, aff); - if (!wt || !wt->idle() || !s.worker_allowed(i, wt->worker_id())) { - ok = false; - break; - } - workers[static_cast(i)] = wt; - } - } + dispatch_next_level_group(); + dispatch_next_level_singles(); + dispatch_sub_ready(); +} - // Pass 2: fill unconstrained slots from idle pool - if (ok) { - for (int i = 0; i < N; i++) { - if (workers[static_cast(i)] != nullptr) continue; - auto *wt = cfg_.manager->pick_idle(s.worker_type, workers, s.eligible_workers_for(i)); - if (!wt) { - ok = false; - break; - } - workers[static_cast(i)] = wt; - } - } +void Scheduler::dispatch_sub_ready() { + TaskSlot slot; + while (cfg_.ready_sub_queue->try_pop(slot)) { + TaskSlotState &s = *cfg_.ring->slot_state(slot); + if (s.state.load(std::memory_order_acquire) != TaskState::READY) continue; + if (s.worker_type != WorkerType::SUB) { + throw std::runtime_error("Scheduler::dispatch_sub_ready: misrouted task slot"); + } - if (!ok) { - q->push(slot); - break; + const int32_t group_size = s.group_size(); + std::vector workers; + workers.reserve(static_cast(group_size)); + for (int32_t i = 0; i < group_size; ++i) { + WorkerThread *worker = cfg_.manager->pick_idle_sub_excluding(workers); + if (worker == nullptr) { + cfg_.ready_sub_queue->push(slot); + return; } + workers.push_back(worker); + } - s.state.store(TaskState::RUNNING, std::memory_order_release); + s.state.store(TaskState::RUNNING, std::memory_order_release); + if (s.is_group()) { + reset_group_state(s, group_size, GroupMemberState::NOT_DISPATCHED); + } + for (int32_t i = 0; i < group_size; ++i) { if (s.is_group()) { std::lock_guard lk(s.group_mu); - s.group_member_states.assign(static_cast(N), GroupMemberState::NOT_DISPATCHED); - s.group_member_outcomes.assign(static_cast(N), EndpointOutcome::SKIPPED); - s.group_terminal_count.store(0, std::memory_order_relaxed); - s.group_dispatched_count.store(0, std::memory_order_relaxed); - s.group_failed = false; - s.group_first_failure_index = -1; - s.group_first_failure_message.clear(); + GroupMemberState &member_state = s.group_member_states[static_cast(i)]; + if (member_state != GroupMemberState::NOT_DISPATCHED || s.group_failed) continue; + member_state = GroupMemberState::RUNNING; } - for (int i = 0; i < N; i++) { - if (s.is_group()) { - std::lock_guard lk(s.group_mu); - GroupMemberState &member_state = s.group_member_states[static_cast(i)]; - if (member_state != GroupMemberState::NOT_DISPATCHED || s.group_failed) { - continue; - } - member_state = GroupMemberState::RUNNING; - s.group_dispatched_count.fetch_add(1, std::memory_order_relaxed); - } - WorkerDispatch d; - d.task_slot = slot; - d.group_index = i; - workers[static_cast(i)]->dispatch(d); + workers[static_cast(i)]->dispatch(WorkerDispatch{slot, i}); + } + } +} + +void Scheduler::dispatch_next_level_group() { + TaskSlot slot; + while (cfg_.ready_next_level_queues->try_front_group(slot)) { + TaskSlotState &s = *cfg_.ring->slot_state(slot); + if (s.state.load(std::memory_order_acquire) != TaskState::READY) { + TaskSlot stale; + cfg_.ready_next_level_queues->try_pop_group(stale); + continue; + } + if (s.worker_type != WorkerType::NEXT_LEVEL || !s.is_group()) { + throw std::runtime_error("Scheduler::dispatch_next_level_group: misrouted task slot"); + } + + const int32_t group_size = s.group_size(); + std::vector workers; + workers.reserve(static_cast(group_size)); + for (int32_t i = 0; i < group_size; ++i) { + const int32_t worker_id = s.target_worker_id(i); + WorkerThread *worker = cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (worker == nullptr) { + throw std::runtime_error("Scheduler::dispatch_next_level_group: invalid target worker"); } + if (std::find(workers.begin(), workers.end(), worker) != workers.end()) { + throw std::runtime_error("Scheduler::dispatch_next_level_group: duplicate target worker"); + } + if (!worker->idle()) return; + workers.push_back(worker); + } + + TaskSlot popped; + if (!cfg_.ready_next_level_queues->try_pop_group(popped) || popped != slot) { + throw std::runtime_error("Scheduler::dispatch_next_level_group: group queue changed unexpectedly"); } - }; - drain_one(cfg_.ready_next_level_queue); - drain_one(cfg_.ready_sub_queue); + reset_group_state(s, group_size, GroupMemberState::RUNNING); + s.state.store(TaskState::RUNNING, std::memory_order_release); + for (int32_t i = 0; i < group_size; ++i) { + workers[static_cast(i)]->dispatch(WorkerDispatch{slot, i}); + } + } +} + +void Scheduler::dispatch_next_level_singles() { + for (int32_t worker_id : cfg_.ready_next_level_queues->worker_ids()) { + WorkerThread *worker = cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (worker == nullptr) { + throw std::runtime_error( + "Scheduler::dispatch_next_level_singles: unknown worker id " + std::to_string(worker_id) + ); + } + if (!worker->idle()) continue; + + TaskSlot slot; + while (cfg_.ready_next_level_queues->try_pop_single(worker_id, slot)) { + TaskSlotState &s = *cfg_.ring->slot_state(slot); + if (s.state.load(std::memory_order_acquire) != TaskState::READY) continue; + if (s.worker_type != WorkerType::NEXT_LEVEL || s.is_group() || s.target_worker_id(0) != worker_id) { + throw std::runtime_error("Scheduler::dispatch_next_level_singles: misrouted task slot"); + } + s.state.store(TaskState::RUNNING, std::memory_order_release); + worker->dispatch(WorkerDispatch{slot, 0}); + break; + } + } } diff --git a/src/common/hierarchical/scheduler.h b/src/common/hierarchical/scheduler.h index 563a642e02..328dd9a57c 100644 --- a/src/common/hierarchical/scheduler.h +++ b/src/common/hierarchical/scheduler.h @@ -15,16 +15,16 @@ * The Scheduler thread routes tasks through the DAG lifecycle: * ready_queue → dispatch (via WorkerManager) → completion → fanout release → new ready * - * Worker pool management (WorkerThread creation, idle selection, dispatch) is - * delegated to WorkerManager. The Scheduler only drives the DAG state machine. + * Worker pool management (WorkerThread creation and dispatch) is delegated to + * WorkerManager. NEXT_LEVEL placement is fixed at submit; SUB remains free. * * Flow: - * Orch: submit() → ready_queue.push(slot) + Scheduler::notify_ready() + * Orch: submit() → directed NEXT_LEVEL queue or shared SUB queue + notify * * Scheduler thread: - * wait on cv (ready_queue OR completion_queue non-empty) + * wait on cv (ready queue OR completion queue OR stop requested) * drain completion_queue → on_task_complete → fanout release → ready_queue - * drain ready_queue → manager.pick_idle → dispatch + * launch directed NEXT_LEVEL tasks, then freely scheduled SUB tasks * * WorkerThread (managed by WorkerManager): * loop: task_queue.pop() → endpoint.run(dispatch) → @@ -54,12 +54,11 @@ class Scheduler { public: struct Config { Ring *ring; // owns slot state storage; Scheduler reads via ring->slot_state(id) - // Strict-4 per-worker-type ready queues. `dispatch_ready` walks each - // queue independently so a saturated pool of one worker type cannot - // head-of-line-block dispatch for the other. - ReadyQueue *ready_next_level_queue; ReadyQueue *ready_sub_queue; + NextLevelReadyQueues *ready_next_level_queues; WorkerManager *manager; // not owned — Scheduler calls manager for dispatch + // Shared READY routing path owned by Orchestrator. + std::function enqueue_ready_cb; // Called when a task reaches CONSUMED (TensorMap cleanup + ring release). std::function on_consumed_cb; }; @@ -101,4 +100,7 @@ class Scheduler { void poison_task(TaskSlot slot, const std::string &root_message); void try_consume(TaskSlot slot); void dispatch_ready(); + void dispatch_next_level_group(); + void dispatch_next_level_singles(); + void dispatch_sub_ready(); }; diff --git a/src/common/hierarchical/types.cpp b/src/common/hierarchical/types.cpp index 6df625ccc1..7d060c5e02 100644 --- a/src/common/hierarchical/types.cpp +++ b/src/common/hierarchical/types.cpp @@ -11,6 +11,8 @@ #include "types.h" +#include + // ============================================================================= // TaskSlotState // ============================================================================= @@ -26,7 +28,6 @@ void TaskSlotState::reset() { } fanout_released.store(0, std::memory_order_relaxed); output_keys.clear(); - eligible_worker_ids.clear(); fanin_producers.clear(); failure_message.clear(); worker_type = WorkerType::NEXT_LEVEL; @@ -37,12 +38,11 @@ void TaskSlotState::reset() { is_group_ = false; remote_sidecar.clear(); remote_sidecars.clear(); - affinities.clear(); + target_worker_ids.clear(); // ring_idx / ring_slot_idx are deliberately NOT cleared here: Ring // stamps them at alloc() before the Orchestrator ever calls reset(), // and Ring::release() needs to read them for the FIFO advance. The // fields are rewritten on every alloc, so stale values never escape. - sub_complete_count.store(0, std::memory_order_relaxed); { std::lock_guard lk(group_mu); group_member_states.clear(); @@ -52,7 +52,6 @@ void TaskSlotState::reset() { group_first_failure_message.clear(); } group_terminal_count.store(0, std::memory_order_relaxed); - group_dispatched_count.store(0, std::memory_order_relaxed); } // ============================================================================= @@ -60,11 +59,8 @@ void TaskSlotState::reset() { // ============================================================================= void ReadyQueue::push(TaskSlot slot) { - { - std::lock_guard lk(mu_); - q_.push(slot); - } - cv_.notify_one(); + std::lock_guard lk(mu_); + q_.push(slot); } bool ReadyQueue::try_pop(TaskSlot &out) { @@ -80,21 +76,57 @@ bool ReadyQueue::empty() const { return q_.empty(); } -bool ReadyQueue::wait_pop(TaskSlot &out) { - std::unique_lock lk(mu_); - cv_.wait(lk, [this] { - return !q_.empty() || shutdown_; - }); +bool ReadyQueue::try_front(TaskSlot &out) { + std::lock_guard lk(mu_); if (q_.empty()) return false; out = q_.front(); - q_.pop(); return true; } -void ReadyQueue::shutdown() { - { - std::lock_guard lk(mu_); - shutdown_ = true; +// ============================================================================= +// NextLevelReadyQueues +// ============================================================================= + +void NextLevelReadyQueues::reset(const std::vector &worker_ids) { + worker_ids_.clear(); + queues_.clear(); + worker_ids_.reserve(worker_ids.size()); + queues_.reserve(worker_ids.size()); + for (int32_t worker_id : worker_ids) { + if (worker_id < 0) throw std::invalid_argument("NextLevelReadyQueues::reset: negative worker id"); + for (int32_t existing : worker_ids_) { + if (existing == worker_id) { + throw std::invalid_argument("NextLevelReadyQueues::reset: duplicate worker id"); + } + } + worker_ids_.push_back(worker_id); + queues_.push_back(std::make_unique()); } - cv_.notify_all(); +} + +size_t NextLevelReadyQueues::index_for(int32_t worker_id) const { + for (size_t i = 0; i < worker_ids_.size(); ++i) { + if (worker_ids_[i] == worker_id) return i; + } + throw std::out_of_range("NextLevelReadyQueues: unknown worker id " + std::to_string(worker_id)); +} + +void NextLevelReadyQueues::push_single(int32_t worker_id, TaskSlot slot) { queues_[index_for(worker_id)]->push(slot); } + +bool NextLevelReadyQueues::try_pop_single(int32_t worker_id, TaskSlot &out) { + return queues_[index_for(worker_id)]->try_pop(out); +} + +void NextLevelReadyQueues::push_group(TaskSlot slot) { group_queue_.push(slot); } + +bool NextLevelReadyQueues::try_front_group(TaskSlot &out) { return group_queue_.try_front(out); } + +bool NextLevelReadyQueues::try_pop_group(TaskSlot &out) { return group_queue_.try_pop(out); } + +bool NextLevelReadyQueues::empty() const { + if (!group_queue_.empty()) return false; + for (const auto &queue : queues_) { + if (!queue->empty()) return false; + } + return true; } diff --git a/src/common/hierarchical/types.h b/src/common/hierarchical/types.h index b788b9f0d7..346757c561 100644 --- a/src/common/hierarchical/types.h +++ b/src/common/hierarchical/types.h @@ -30,9 +30,9 @@ #include #include -#include #include #include +#include #include #include #include @@ -305,21 +305,11 @@ struct TaskSlotState { // --- TensorMap keys registered by this task (for cleanup on CONSUMED) --- std::vector output_keys; - // Empty outer vector means legacy/unconstrained dispatch. When present, - // each group member's vector is the final callable/data worker-id - // intersection and must be non-empty. - std::vector> eligible_worker_ids; + // Exact stable NEXT_LEVEL worker id for each task/group member. SUB slots + // leave this empty because SUB has no worker-selection API. + std::vector target_worker_ids; - // --- Submit affinity (set by submit_next_level with worker= parameter) --- - // Empty = unconstrained. For NEXT_LEVEL, non-negative values are stable - // worker ids. SUB has no public affinity surface and keeps index - // semantics for any internal callers that populate this vector. - std::vector affinities; - - int32_t get_affinity(int i) const { - if (affinities.empty()) return -1; - return affinities[static_cast(i)]; - } + int32_t target_worker_id(int32_t i) const { return target_worker_ids.at(static_cast(i)); } // --- Producer tasks this task depends on (for deferred release) --- // When this task reaches COMPLETED, the Scheduler releases one fanout ref @@ -362,36 +352,16 @@ struct TaskSlotState { int32_t ring_slot_idx{0}; // --- Group bookkeeping --- - std::atomic sub_complete_count{0}; std::mutex group_mu; std::vector group_member_states; std::vector group_member_outcomes; std::atomic group_terminal_count{0}; - std::atomic group_dispatched_count{0}; bool group_failed{false}; int32_t group_first_failure_index{-1}; std::string group_first_failure_message; bool is_group() const { return is_group_; } int32_t group_size() const { return is_group_ ? static_cast(task_args_list.size()) : 1; } - bool has_worker_constraints() const { return !eligible_worker_ids.empty(); } - - const std::vector &eligible_workers_for(int32_t i) const { - static const std::vector empty; - if (eligible_worker_ids.empty()) return empty; - if (i < 0 || static_cast(i) >= eligible_worker_ids.size()) return empty; - return eligible_worker_ids[static_cast(i)]; - } - - bool worker_allowed(int32_t i, int32_t worker_id) const { - if (eligible_worker_ids.empty()) return true; - const auto &eligible = eligible_workers_for(i); - for (int32_t id : eligible) { - if (id == worker_id) return true; - } - return false; - } - const RemoteTaskArgsSidecar &remote_sidecar_for(int32_t i) const { static const RemoteTaskArgsSidecar empty; if (is_group_) { @@ -427,15 +397,31 @@ class ReadyQueue { bool empty() const; - // Blocking: waits until a slot is available or shutdown() is called. - // Returns false only when shutdown and queue is empty. - bool wait_pop(TaskSlot &out); - - void shutdown(); + // Non-blocking: copies the front without removing it. + bool try_front(TaskSlot &out); private: std::queue q_; mutable std::mutex mu_; - std::condition_variable cv_; - bool shutdown_{false}; +}; + +// Directed NEXT_LEVEL queues. Worker registration is complete before reset() +// and the worker-id mapping is immutable while scheduling is active. +class NextLevelReadyQueues { +public: + void reset(const std::vector &worker_ids); + void push_single(int32_t worker_id, TaskSlot slot); + bool try_pop_single(int32_t worker_id, TaskSlot &out); + void push_group(TaskSlot slot); + bool try_front_group(TaskSlot &out); + bool try_pop_group(TaskSlot &out); + bool empty() const; + const std::vector &worker_ids() const { return worker_ids_; } + +private: + size_t index_for(int32_t worker_id) const; + + std::vector worker_ids_; + std::vector> queues_; + ReadyQueue group_queue_; }; diff --git a/src/common/hierarchical/worker.cpp b/src/common/hierarchical/worker.cpp index 23629a2862..e061248eea 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -96,23 +96,26 @@ void Worker::add_remote_l3_socket( void Worker::init() { if (initialized_) throw std::runtime_error("Worker: already initialized"); - orchestrator_.init( - &tensormap_, &allocator_, &scope_, &ready_next_level_queue_, &ready_sub_queue_, &manager_, [this] { - scheduler_.notify_ready(); - } - ); - // Start WorkerManager first — creates WorkerThreads. // The on_complete callback routes through the Scheduler's worker_done(). manager_.start(&allocator_, [this](WorkerCompletion completion) { scheduler_.worker_done(std::move(completion)); }); + ready_next_level_queues_.reset(manager_.next_level_worker_ids()); + orchestrator_.init( + &tensormap_, &allocator_, &scope_, &ready_sub_queue_, &ready_next_level_queues_, &manager_, [this] { + scheduler_.notify_ready(); + } + ); Scheduler::Config cfg; cfg.ring = &allocator_; - cfg.ready_next_level_queue = &ready_next_level_queue_; cfg.ready_sub_queue = &ready_sub_queue_; + cfg.ready_next_level_queues = &ready_next_level_queues_; cfg.manager = &manager_; + cfg.enqueue_ready_cb = [this](TaskSlot slot) { + orchestrator_.enqueue_ready(slot); + }; cfg.on_consumed_cb = [this](TaskSlot slot) { orchestrator_.on_consumed(slot); }; diff --git a/src/common/hierarchical/worker.h b/src/common/hierarchical/worker.h index fd65eaa108..1ced949bd5 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -205,10 +205,9 @@ class Worker { TensorMap tensormap_; Ring allocator_; Scope scope_; - // One ready queue per WorkerType. Submit routes by s.worker_type; - // the Scheduler drains each queue independently so saturation of - // one pool cannot head-of-line-block the other. - ReadyQueue ready_next_level_queue_; + // NEXT_LEVEL singles use one FIFO per stable worker id; groups use one + // FIFO whose head launches only when every requested worker is idle. + NextLevelReadyQueues ready_next_level_queues_; ReadyQueue ready_sub_queue_; Orchestrator orchestrator_; Scheduler scheduler_; diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index 6bfd83b318..5044b37f6a 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -449,12 +449,6 @@ void WorkerManager::stop() { sub_threads_.clear(); } -WorkerThread *WorkerManager::get_worker_by_index(WorkerType type, int worker_index) const { - auto &threads = (type == WorkerType::NEXT_LEVEL) ? next_level_threads_ : sub_threads_; - if (worker_index < 0 || static_cast(worker_index) >= threads.size()) return nullptr; - return threads[static_cast(worker_index)].get(); -} - WorkerThread *WorkerManager::get_worker_by_id(WorkerType type, int32_t worker_id) const { auto &threads = (type == WorkerType::NEXT_LEVEL) ? next_level_threads_ : sub_threads_; for (auto &wt : threads) { @@ -463,23 +457,18 @@ WorkerThread *WorkerManager::get_worker_by_id(WorkerType type, int32_t worker_id return nullptr; } -WorkerThread *WorkerManager::pick_idle( - WorkerType type, const std::vector &exclude, const std::vector &eligible_worker_ids -) const { - auto &threads = (type == WorkerType::NEXT_LEVEL) ? next_level_threads_ : sub_threads_; - for (auto &wt : threads) { +std::vector WorkerManager::next_level_worker_ids() const { + std::vector worker_ids; + worker_ids.reserve(next_level_threads_.size()); + for (const auto &worker : next_level_threads_) { + worker_ids.push_back(worker->worker_id()); + } + return worker_ids; +} + +WorkerThread *WorkerManager::pick_idle_sub_excluding(const std::vector &exclude) const { + for (const auto &wt : sub_threads_) { if (!wt->idle()) continue; - if (!eligible_worker_ids.empty()) { - bool eligible = false; - int32_t worker_id = wt->worker_id(); - for (int32_t id : eligible_worker_ids) { - if (id == worker_id) { - eligible = true; - break; - } - } - if (!eligible) continue; - } bool excluded = false; for (auto *ex : exclude) { if (ex == wt.get()) { diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index c4668a0dde..7d4cfe1fd8 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -471,15 +471,12 @@ class WorkerManager { void start(Ring *ring, const OnCompleteFn &on_complete); void stop(); - // Direct index into the worker pool (0-based). - WorkerThread *get_worker_by_index(WorkerType type, int worker_index) const; WorkerThread *get_worker_by_id(WorkerType type, int32_t worker_id) const; + std::vector next_level_worker_ids() const; - // Pick one idle worker NOT in `exclude`, restricted to `eligible_worker_ids` - // when that list is non-empty. Returns nullptr if none available. - WorkerThread *pick_idle( - WorkerType type, const std::vector &exclude, const std::vector &eligible_worker_ids - ) const; + // SUB has no worker-selection API. Pick any idle SUB worker not already + // selected for the same group dispatch. + WorkerThread *pick_idle_sub_excluding(const std::vector &exclude) const; bool any_busy() const; diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py b/tests/st/a2a3/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py index 6b9add99a6..936cdde19a 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py +++ b/tests/st/a2a3/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py @@ -163,7 +163,7 @@ def test_prepare_new_identity_after_start_then_run(st_platform, st_device_ids): # snapshot, so they are in _run_chip_main_loop — the only state in # which the CTRL_REGISTER broadcast in step 2 can be ACKed. def orch_pre(o, _args, _cfg): - o.submit_next_level(pre_handle, chip_args_pre, config) + o.submit_next_level(pre_handle, chip_args_pre, config, worker=0) worker.run(orch_pre) got_pre = args_pre.f @@ -182,7 +182,7 @@ def orch_pre(o, _args, _cfg): # 3. Run with post_handle. If CTRL_REGISTER delivered correctly, the # child has the identity prepared; otherwise dispatch will fail. def orch_post(o, _args, _cfg): - o.submit_next_level(post_handle, chip_args_post, config) + o.submit_next_level(post_handle, chip_args_post, config, worker=0) worker.run(orch_post) got_post = args_post.f @@ -217,14 +217,13 @@ def test_prepare_new_identity_after_start_parallel_broadcast(st_platform, st_dev pre_handle = worker.register(chip_callable) a, b = 2.0, 3.0 expected = _golden(a, b) - # Pre-allocate args for each chip (chip_id = block group). The - # vector_example orchestration partitions the input across cores, so a - # single args bundle works for both chips' first-run trigger; the - # second-run uses the post-start prepared handle. + # Use one first-run bundle to start the control-capable child loop, then one + # post-start bundle per chip so the dynamically prepared handle is executed + # by both broadcast targets. args_pre = _make_args(a, b) - args_post = _make_args(a, b) + args_post = [_make_args(a, b) for _ in device_ids] chip_args_pre, _ = _build_l3_task_args(args_pre, _ORCH_SIG) - chip_args_post, _ = _build_l3_task_args(args_post, _ORCH_SIG) + chip_args_post = [_build_l3_task_args(args, _ORCH_SIG)[0] for args in args_post] worker.init() try: @@ -233,7 +232,7 @@ def test_prepare_new_identity_after_start_parallel_broadcast(st_platform, st_dev config.aicpu_thread_num = 4 def orch_pre(o, _a, _c): - o.submit_next_level(pre_handle, chip_args_pre, config) + o.submit_next_level(pre_handle, chip_args_pre, config, worker=0) worker.run(orch_pre) assert torch.allclose(args_pre.f, torch.full_like(args_pre.f, expected), rtol=1e-5, atol=1e-5) @@ -242,10 +241,11 @@ def orch_pre(o, _a, _c): post_handle = worker.register(post_callable) def orch_post(o, _a, _c): - o.submit_next_level(post_handle, chip_args_post, config) + o.submit_next_level_group(post_handle, chip_args_post, config, workers=[0, 1]) worker.run(orch_post) - assert torch.allclose(args_post.f, torch.full_like(args_post.f, expected), rtol=1e-5, atol=1e-5) + for args in args_post: + assert torch.allclose(args.f, torch.full_like(args.f, expected), rtol=1e-5, atol=1e-5) finally: worker.close() @@ -289,7 +289,7 @@ def test_prepare_capacity_overflow_post_start(st_platform, st_device_ids): config.aicpu_thread_num = 4 def orch_pre(o, _a, _c): - o.submit_next_level(chip_handle, chip_args_pre, config) + o.submit_next_level(chip_handle, chip_args_pre, config, worker=0) worker.run(orch_pre) @@ -340,7 +340,7 @@ def test_duplicate_prepare_same_hashid_survives_one_unregister(st_platform, st_d # 1. Trigger fork via pre_handle to put the chip child into the main loop. def orch_one(o, _args, _cfg): - o.submit_next_level(pre_handle, chip_args_one, config) + o.submit_next_level(pre_handle, chip_args_one, config, worker=0) worker.run(orch_one) assert torch.allclose(args_one.f, torch.full_like(args_one.f, expected), rtol=1e-5, atol=1e-5) @@ -357,10 +357,10 @@ def orch_one(o, _args, _cfg): worker.unregister(pre_handle) with pytest.raises(KeyError, match="not live"): - worker.run(lambda o, _args, _cfg: o.submit_next_level(pre_handle, chip_args_one, config)) + worker.run(lambda o, _args, _cfg: o.submit_next_level(pre_handle, chip_args_one, config, worker=0)) def orch_two(o, _args, _cfg): - o.submit_next_level(duplicate_handle, chip_args_two, config) + o.submit_next_level(duplicate_handle, chip_args_two, config, worker=0) worker.run(orch_two) assert torch.allclose(args_two.f, torch.full_like(args_two.f, expected), rtol=1e-5, atol=1e-5) @@ -368,7 +368,7 @@ def orch_two(o, _args, _cfg): # 4. Dropping the final handle invalidates it through the public API. worker.unregister(duplicate_handle) with pytest.raises(KeyError, match="not live"): - worker.run(lambda o, _args, _cfg: o.submit_next_level(duplicate_handle, chip_args_two, config)) + worker.run(lambda o, _args, _cfg: o.submit_next_level(duplicate_handle, chip_args_two, config, worker=0)) finally: worker.close() @@ -410,7 +410,7 @@ def test_unregister_last_handle_allows_reprepare_same_hashid(st_platform, st_dev config.aicpu_thread_num = 4 def orch_one(o, _args, _cfg): - o.submit_next_level(pre_handle, chip_args_one, config) + o.submit_next_level(pre_handle, chip_args_one, config, worker=0) worker.run(orch_one) assert torch.allclose(args_one.f, torch.full_like(args_one.f, expected), rtol=1e-5, atol=1e-5) @@ -418,14 +418,14 @@ def orch_one(o, _args, _cfg): dyn_handle = worker.register(post_callable) def orch_two(o, _args, _cfg): - o.submit_next_level(dyn_handle, chip_args_two, config) + o.submit_next_level(dyn_handle, chip_args_two, config, worker=0) worker.run(orch_two) assert torch.allclose(args_two.f, torch.full_like(args_two.f, expected), rtol=1e-5, atol=1e-5) worker.unregister(dyn_handle) with pytest.raises(KeyError, match="not live"): - worker.run(lambda o, _args, _cfg: o.submit_next_level(dyn_handle, chip_args_two, config)) + worker.run(lambda o, _args, _cfg: o.submit_next_level(dyn_handle, chip_args_two, config, worker=0)) # Re-prepare the same hashid after its final handle was dropped. again_handle = worker.register(post_callable) @@ -434,7 +434,7 @@ def orch_two(o, _args, _cfg): assert again_handle._handle_id != dyn_handle._handle_id def orch_three(o, _args, _cfg): - o.submit_next_level(again_handle, chip_args_three, config) + o.submit_next_level(again_handle, chip_args_three, config, worker=0) worker.run(orch_three) assert torch.allclose(args_three.f, torch.full_like(args_three.f, expected), rtol=1e-5, atol=1e-5) diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_dependency.py b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_dependency.py index f2fd8ebf94..ae14b9dd71 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_dependency.py +++ b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_dependency.py @@ -34,7 +34,7 @@ def run_dag(orch, callables, task_args, config): chip_args, _ = _build_l3_task_args(task_args, callables.vector_kernel_sig) callables.keep(chip_args) # prevent GC before drain - orch.submit_next_level(callables.vector_kernel, chip_args, config) + orch.submit_next_level(callables.vector_kernel, chip_args, config, worker=0) # SubTask: tag the chip output as INPUT — Orchestrator wires the dep via TensorMap. sub_args = TaskArgs() diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_group.py b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_group.py index 5dfe1e6593..025d0cb674 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_group.py +++ b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_group.py @@ -43,7 +43,7 @@ def run_dag(orch, callables, task_args, config): args1 = _chip_args(task_args.a1, task_args.b1, task_args.f1) callables.keep(args0, args1) # prevent GC before drain - orch.submit_next_level_group(callables.vector_kernel, [args0, args1], config) + orch.submit_next_level_group(callables.vector_kernel, [args0, args1], config, workers=[0, 1]) # SubTask depends on both group outputs (f0, f1) — tag both as INPUT. sub_args = TaskArgs() diff --git a/tests/ut/cpp/hierarchical/test_orchestrator.cpp b/tests/ut/cpp/hierarchical/test_orchestrator.cpp index 4aa019d040..f456d18e58 100644 --- a/tests/ut/cpp/hierarchical/test_orchestrator.cpp +++ b/tests/ut/cpp/hierarchical/test_orchestrator.cpp @@ -29,8 +29,7 @@ struct OrchestratorFixture : public ::testing::Test { TensorMap tm; Ring allocator; Scope scope; - // Strict-4: per-type ready queues. - ReadyQueue rq_next_level; + NextLevelReadyQueues rq_next_level; ReadyQueue rq_sub; Orchestrator orch; CallConfig cfg; @@ -39,11 +38,15 @@ struct OrchestratorFixture : public ::testing::Test { // convenience alias for the next-level queue. Kept public so existing // `rq.try_pop(...)` / `EXPECT_TRUE(rq.try_pop(...))` lines continue to // work without rewriting every assertion. - ReadyQueue &rq = rq_next_level; + struct WorkerQueueView { + NextLevelReadyQueues *queues; + bool try_pop(TaskSlot &out) { return queues->try_pop_single(0, out); } + } rq{&rq_next_level}; void SetUp() override { allocator.init(/*heap_bytes=*/1ULL << 20); - orch.init(&tm, &allocator, &scope, &rq_next_level, &rq_sub); + rq_next_level.reset({0, 1, 3}); + orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level); } void TearDown() override { allocator.shutdown(); } @@ -76,7 +79,7 @@ struct OrchestratorFixture : public ::testing::Test { TEST_F(OrchestratorFixture, IndependentTaskIsImmediatelyReady) { auto a = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); - auto res = orch.submit_next_level(C(42), a, cfg); + auto res = orch.submit_next_level(C(42), a, cfg, 0); EXPECT_NE(res.task_slot, INVALID_SLOT); TaskSlot slot; @@ -85,16 +88,23 @@ TEST_F(OrchestratorFixture, IndependentTaskIsImmediatelyReady) { EXPECT_EQ(S(slot).state.load(), TaskState::READY); } +TEST_F(OrchestratorFixture, NextLevelTargetsMustBeValidAndDistinct) { + auto args = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); + EXPECT_THROW((void)orch.submit_next_level(C(42), args, cfg, -1), std::invalid_argument); + EXPECT_THROW((void)orch.submit_next_level_group(C(42), {args, args}, cfg, {0}), std::invalid_argument); + EXPECT_THROW((void)orch.submit_next_level_group(C(42), {args, args}, cfg, {0, 0}), std::invalid_argument); +} + TEST_F(OrchestratorFixture, DependentTaskIsPending) { // Task A produces an OUTPUT at key 0xBEEF auto args_a = single_tensor_args(0xBEEF, TensorArgType::OUTPUT); - auto a = orch.submit_next_level(C(42), args_a, cfg); + auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot a_slot; rq.try_pop(a_slot); // Task B reads INPUT at the same key -- depends on A auto args_b = single_tensor_args(0xBEEF, TensorArgType::INPUT); - auto b = orch.submit_next_level(C(42), args_b, cfg); + auto b = orch.submit_next_level(C(42), args_b, cfg, 0); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::PENDING); EXPECT_EQ(S(b.task_slot).fanin_count, 1); @@ -104,14 +114,14 @@ TEST_F(OrchestratorFixture, DependentTaskIsPending) { TEST_F(OrchestratorFixture, GroupDuplicateInputsKeepOneProducer) { auto producer_args = single_tensor_args(0xBEEF, TensorArgType::OUTPUT); - auto producer = orch.submit_next_level(C(42), producer_args, cfg); + auto producer = orch.submit_next_level(C(42), producer_args, cfg, 0); TaskSlot ready; ASSERT_TRUE(rq.try_pop(ready)); ASSERT_EQ(ready, producer.task_slot); TaskArgs input0 = single_tensor_args(0xBEEF, TensorArgType::INPUT); TaskArgs input1 = single_tensor_args(0xBEEF, TensorArgType::INPUT); - auto consumer = orch.submit_next_level_group(C(43), {input0, input1}, cfg); + auto consumer = orch.submit_next_level_group(C(43), {input0, input1}, cfg, {0, 1}); EXPECT_EQ(S(consumer.task_slot).state.load(), TaskState::PENDING); EXPECT_EQ(S(consumer.task_slot).fanin_count, 1); @@ -123,7 +133,7 @@ TEST_F(OrchestratorFixture, SubmitAfterFailedProducerPoisonsConsumer) { orch.scope_begin(); auto args_a = single_tensor_args(0xD00D, TensorArgType::OUTPUT); - auto a = orch.submit_next_level(C(42), args_a, cfg); + auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot ready; ASSERT_TRUE(rq.try_pop(ready)); ASSERT_EQ(ready, a.task_slot); @@ -134,7 +144,7 @@ TEST_F(OrchestratorFixture, SubmitAfterFailedProducerPoisonsConsumer) { producer.fanout_released.store(1, std::memory_order_release); auto args_b = single_tensor_args(0xD00D, TensorArgType::INPUT); - auto b = orch.submit_next_level(C(43), args_b, cfg); + auto b = orch.submit_next_level(C(43), args_b, cfg, 0); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::FAILED); EXPECT_EQ(S(b.task_slot).failure_message, "producer failed"); @@ -152,7 +162,7 @@ TEST_F(OrchestratorFixture, SubmitAfterFailedProducerPoisonsConsumer) { TEST_F(OrchestratorFixture, TensorMapTracksProducer) { auto args_a = single_tensor_args(0x1234, TensorArgType::OUTPUT); - auto a = orch.submit_next_level(C(42), args_a, cfg); + auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot drain_slot; rq.try_pop(drain_slot); @@ -161,7 +171,7 @@ TEST_F(OrchestratorFixture, TensorMapTracksProducer) { TEST_F(OrchestratorFixture, OnConsumedCleansUpTensorMap) { auto args_a = single_tensor_args(0x42, TensorArgType::OUTPUT); - auto a = orch.submit_next_level(C(42), args_a, cfg); + auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot slot; rq.try_pop(slot); @@ -177,7 +187,7 @@ TEST_F(OrchestratorFixture, OnConsumedCleansUpTensorMap) { TEST_F(OrchestratorFixture, ScopeRegistersAndReleasesRef) { orch.scope_begin(); auto args_a = single_tensor_args(0x77, TensorArgType::OUTPUT); - auto a = orch.submit_next_level(C(42), args_a, cfg); + auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot slot; rq.try_pop(slot); @@ -201,13 +211,13 @@ TEST_F(OrchestratorFixture, ScopeRegistersAndReleasesRef) { TEST_F(OrchestratorFixture, NoDepTagSkipsDependencyTracking) { // OUTPUT-tagged input registers a producer auto args_a = single_tensor_args(0xAAAA, TensorArgType::OUTPUT); - auto a = orch.submit_next_level(C(42), args_a, cfg); + auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot drain_slot; rq.try_pop(drain_slot); // Second task references same key but tagged NO_DEP -- should be independent auto args_b = single_tensor_args(0xAAAA, TensorArgType::NO_DEP); - auto b = orch.submit_next_level(C(42), args_b, cfg); + auto b = orch.submit_next_level(C(42), args_b, cfg, 0); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::READY); EXPECT_EQ(S(b.task_slot).fanin_count, 0); } @@ -215,7 +225,7 @@ TEST_F(OrchestratorFixture, NoDepTagSkipsDependencyTracking) { TEST_F(OrchestratorFixture, GroupTaskStoresArgsListPerMember) { TaskArgs a0 = single_tensor_args(0xA0, TensorArgType::OUTPUT); TaskArgs a1 = single_tensor_args(0xA1, TensorArgType::OUTPUT); - auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg); + auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {0, 1}); EXPECT_NE(res.task_slot, INVALID_SLOT); EXPECT_TRUE(S(res.task_slot).is_group()); @@ -233,7 +243,7 @@ TEST_F(OrchestratorFixture, GroupTaskStoresArgsListPerMember) { TEST_F(OrchestratorFixture, SingleTaskStoresTaskArgsDirectly) { TaskArgs a0 = single_tensor_args(0xC0, TensorArgType::OUTPUT); - auto res = orch.submit_next_level(C(42), a0, cfg); + auto res = orch.submit_next_level(C(42), a0, cfg, 0); ASSERT_NE(res.task_slot, INVALID_SLOT); EXPECT_FALSE(S(res.task_slot).is_group()); EXPECT_EQ(S(res.task_slot).group_size(), 1); @@ -254,7 +264,7 @@ TEST_F(OrchestratorFixture, OutputAutoAllocsFromHeapRing) { t.dtype = DataType::UINT8; args.add_tensor(t, TensorArgType::OUTPUT); - auto res = orch.submit_next_level(C(42), args, cfg); + auto res = orch.submit_next_level(C(42), args, cfg, 0); ASSERT_NE(res.task_slot, INVALID_SLOT); uint64_t data = S(res.task_slot).task_args.tensor(0).buffer.addr; @@ -286,7 +296,7 @@ TEST_F(OrchestratorFixture, RemoteOutputSidecarSkipsLocalAutoAllocAndRegistersRe sidecar.tensors[0].desc.offset = 64; sidecar.tensors[0].desc.nbytes = 1024; - auto res = orch.submit_next_level(C(42), args, cfg, -1, {3}, sidecar); + auto res = orch.submit_next_level(C(42), args, cfg, 3, {3}, sidecar); ASSERT_NE(res.task_slot, INVALID_SLOT); EXPECT_EQ(S(res.task_slot).task_args.tensor(0).buffer.addr, 0u); @@ -300,7 +310,7 @@ TEST_F(OrchestratorFixture, RemoteBarePayloadFailsBeforeSlotCommit) { RemoteTaskArgsSidecar sidecar; sidecar.tensors.resize(1); - EXPECT_THROW({ (void)orch.submit_next_level(C(42), args, cfg, -1, {3}, sidecar); }, std::invalid_argument); + EXPECT_THROW({ (void)orch.submit_next_level(C(42), args, cfg, 3, {3}, sidecar); }, std::invalid_argument); } TEST_F(OrchestratorFixture, RemoteSidecarRejectsNonOwnerEligibleEndpointWithoutImport) { @@ -321,7 +331,7 @@ TEST_F(OrchestratorFixture, RemoteSidecarRejectsNonOwnerEligibleEndpointWithoutI sidecar.tensors[0].desc.generation = 2; sidecar.tensors[0].desc.nbytes = 1; - EXPECT_THROW({ (void)orch.submit_next_level(C(42), args, cfg, -1, {3, 4}, sidecar); }, std::invalid_argument); + EXPECT_THROW({ (void)orch.submit_next_level(C(42), args, cfg, 3, {3, 4}, sidecar); }, std::invalid_argument); } TEST_F(OrchestratorFixture, RemoteInputSidecarUsesRemoteTensorMapKey) { @@ -342,14 +352,15 @@ TEST_F(OrchestratorFixture, RemoteInputSidecarUsesRemoteTensorMapKey) { output_sidecar.tensors[0].desc.generation = 2; output_sidecar.tensors[0].desc.offset = 0; output_sidecar.tensors[0].desc.nbytes = 1; - auto producer = orch.submit_next_level(C(42), output_args, cfg, -1, {3}, output_sidecar); + auto producer = orch.submit_next_level(C(42), output_args, cfg, 3, {3}, output_sidecar); TaskSlot ready; - rq.try_pop(ready); + ASSERT_TRUE(rq_next_level.try_pop_single(3, ready)); + ASSERT_EQ(ready, producer.task_slot); TaskArgs input_args; Tensor in = out; input_args.add_tensor(in, TensorArgType::INPUT); - auto consumer = orch.submit_next_level(C(43), input_args, cfg, -1, {3}, output_sidecar); + auto consumer = orch.submit_next_level(C(43), input_args, cfg, 3, {3}, output_sidecar); EXPECT_EQ(S(consumer.task_slot).state.load(), TaskState::PENDING); EXPECT_EQ(S(consumer.task_slot).fanin_count, 1); @@ -364,7 +375,7 @@ TEST_F(OrchestratorFixture, InoutWiresCreatorAsFanin) { // the alloc-slot (so its HeapRing slab stays live while they write) // must tag the buffer INOUT. auto creator_args = single_tensor_args(0xFEED, TensorArgType::OUTPUT); - auto creator = orch.submit_next_level(C(42), creator_args, cfg); + auto creator = orch.submit_next_level(C(42), creator_args, cfg, 0); TaskSlot drain; rq.try_pop(drain); // Mark the creator COMPLETED so the new submit mimics the alloc-slot @@ -372,7 +383,7 @@ TEST_F(OrchestratorFixture, InoutWiresCreatorAsFanin) { S(creator.task_slot).state.store(TaskState::COMPLETED, std::memory_order_relaxed); auto writer_args = single_tensor_args(0xFEED, TensorArgType::INOUT); - auto writer = orch.submit_next_level(C(42), writer_args, cfg); + auto writer = orch.submit_next_level(C(42), writer_args, cfg, 0); TaskSlot writer_slot; rq.try_pop(writer_slot); @@ -403,13 +414,13 @@ TEST_F(OrchestratorFixture, OutputAndOutputExistingAreInsertOnly) { }; for (Case c : {Case{0xABCD, TensorArgType::OUTPUT}, Case{0xBEEF, TensorArgType::OUTPUT_EXISTING}}) { auto prior_args = single_tensor_args(c.key, TensorArgType::OUTPUT); - auto prior = orch.submit_next_level(C(42), prior_args, cfg); + auto prior = orch.submit_next_level(C(42), prior_args, cfg, 0); TaskSlot drain; rq.try_pop(drain); S(prior.task_slot).state.store(TaskState::COMPLETED, std::memory_order_relaxed); auto writer_args = single_tensor_args(c.key, c.tag); - auto writer = orch.submit_next_level(C(42), writer_args, cfg); + auto writer = orch.submit_next_level(C(42), writer_args, cfg, 0); EXPECT_EQ(tm.lookup(TensorKey{c.key, -1}), writer.task_slot); EXPECT_EQ(S(writer.task_slot).fanin_count, 0); diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index eb577c648a..7383744692 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -266,9 +266,8 @@ struct SchedulerFixture : public ::testing::Test { TensorMap tm; Ring allocator; Scope scope; - // Strict-4: per-type ready queues. - ReadyQueue rq_next_level; ReadyQueue rq_sub; + NextLevelReadyQueues rq_next_level; Orchestrator orch; MockMailboxWorker mock_worker; WorkerManager manager; @@ -282,21 +281,25 @@ struct SchedulerFixture : public ::testing::Test { void SetUp() override { allocator.init(/*heap_bytes=*/1ULL << 20); - orch.init(&tm, &allocator, &scope, &rq_next_level, &rq_sub, &manager, [this] { - sched.notify_ready(); - }); mock_worker.start(); manager.add_next_level(mock_worker.mailbox_ptr()); manager.start(&allocator, [this](WorkerCompletion completion) { sched.worker_done(std::move(completion)); }); + rq_next_level.reset(manager.next_level_worker_ids()); + orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level, &manager, [this] { + sched.notify_ready(); + }); Scheduler::Config c; c.ring = &allocator; - c.ready_next_level_queue = &rq_next_level; c.ready_sub_queue = &rq_sub; + c.ready_next_level_queues = &rq_next_level; c.manager = &manager; + c.enqueue_ready_cb = [this](TaskSlot slot) { + orch.enqueue_ready(slot); + }; c.on_consumed_cb = [this](TaskSlot s) { orch.on_consumed(s); std::lock_guard lk(consumed_mu); @@ -373,7 +376,7 @@ TEST(WorkerManagerTest, ControlPrepareUsesStableNextLevelWorkerId) { TEST_F(SchedulerFixture, IndependentTaskDispatchedAndConsumed) { auto args_a = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); - auto res = orch.submit_next_level(C(42), args_a, cfg); + auto res = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot slot = res.task_slot; mock_worker.wait_running(); @@ -387,10 +390,10 @@ TEST_F(SchedulerFixture, IndependentTaskDispatchedAndConsumed) { TEST_F(SchedulerFixture, DependentTaskDispatchedAfterProducerCompletes) { auto args_a = single_tensor_args(0xBEEF, TensorArgType::OUTPUT); - auto a = orch.submit_next_level(C(10), args_a, cfg); + auto a = orch.submit_next_level(C(10), args_a, cfg, 0); auto args_b = single_tensor_args(0xBEEF, TensorArgType::INPUT); - auto b = orch.submit_next_level(C(11), args_b, cfg); + auto b = orch.submit_next_level(C(11), args_b, cfg, 0); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::PENDING); mock_worker.wait_running(); @@ -431,7 +434,7 @@ TEST_F(SchedulerFixture, ComposedKernelArgsBlobFitsMailbox) { args.add_scalar(1); args.add_scalar(2); - auto res = orch.submit_next_level(C(76), args, cfg); + auto res = orch.submit_next_level(C(76), args, cfg, 0); mock_worker.wait_running(); ASSERT_GE(mock_worker.dispatched_count(), 1) @@ -444,10 +447,10 @@ TEST_F(SchedulerFixture, ComposedKernelArgsBlobFitsMailbox) { TEST_F(SchedulerFixture, FailedProducerPoisonsDependentTask) { auto args_a = single_tensor_args(0xD00D, TensorArgType::OUTPUT); - auto a = orch.submit_next_level(C(21), args_a, cfg); + auto a = orch.submit_next_level(C(21), args_a, cfg, 0); auto args_b = single_tensor_args(0xD00D, TensorArgType::INPUT); - auto b = orch.submit_next_level(C(22), args_b, cfg); + auto b = orch.submit_next_level(C(22), args_b, cfg, 0); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::PENDING); mock_worker.wait_running(); @@ -472,9 +475,8 @@ struct GroupSchedulerFixture : public ::testing::Test { TensorMap tm; Ring allocator; Scope scope; - // Strict-4: per-type ready queues. - ReadyQueue rq_next_level; ReadyQueue rq_sub; + NextLevelReadyQueues rq_next_level; Orchestrator orch; MockMailboxWorker worker_a; MockMailboxWorker worker_b; @@ -489,9 +491,6 @@ struct GroupSchedulerFixture : public ::testing::Test { void SetUp() override { allocator.init(/*heap_bytes=*/1ULL << 20); - orch.init(&tm, &allocator, &scope, &rq_next_level, &rq_sub, &manager, [this] { - sched.notify_ready(); - }); worker_a.start(); worker_b.start(); @@ -500,12 +499,19 @@ struct GroupSchedulerFixture : public ::testing::Test { manager.start(&allocator, [this](WorkerCompletion completion) { sched.worker_done(std::move(completion)); }); + rq_next_level.reset(manager.next_level_worker_ids()); + orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level, &manager, [this] { + sched.notify_ready(); + }); Scheduler::Config c; c.ring = &allocator; - c.ready_next_level_queue = &rq_next_level; c.ready_sub_queue = &rq_sub; + c.ready_next_level_queues = &rq_next_level; c.manager = &manager; + c.enqueue_ready_cb = [this](TaskSlot slot) { + orch.enqueue_ready(slot); + }; c.on_consumed_cb = [this](TaskSlot s) { orch.on_consumed(s); std::lock_guard lk(consumed_mu); @@ -538,7 +544,7 @@ TEST_F(GroupSchedulerFixture, GroupDispatchesToNWorkers) { TaskArgs a0 = single_tensor_args(0xA0, TensorArgType::OUTPUT); TaskArgs a1 = single_tensor_args(0xA1, TensorArgType::OUTPUT); - auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg); + auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {0, 1}); TaskSlot slot = res.task_slot; worker_a.wait_running(); @@ -547,12 +553,33 @@ TEST_F(GroupSchedulerFixture, GroupDispatchesToNWorkers) { EXPECT_EQ(worker_a.dispatched_count(), 1); EXPECT_EQ(worker_b.dispatched_count(), 1); - // Each worker got a different TaskArgs from the slot's task_args_list — - // proven by the keys 0xA0 and 0xA1 each landing on exactly one worker. - uint64_t keys[2] = {worker_a.dispatched[0].tensor_key, worker_b.dispatched[0].tensor_key}; - std::sort(std::begin(keys), std::end(keys)); - EXPECT_EQ(keys[0], 0xA0u); - EXPECT_EQ(keys[1], 0xA1u); + EXPECT_EQ(worker_a.dispatched[0].tensor_key, 0xA0u); + EXPECT_EQ(worker_b.dispatched[0].tensor_key, 0xA1u); + (void)slot; + + worker_a.complete(); + worker_b.complete(); + wait_consumed(slot); +} + +TEST_F(GroupSchedulerFixture, GroupMapsEachMemberToItsTargetWorkerIdNotIndex) { + // Reversed target order: member 0 -> worker id 1 (worker_b), member 1 -> + // worker id 0 (worker_a). A map-by-registration-index bug would instead + // send member 0 (a0) to worker_a; the reversed keys catch it. + TaskArgs a0 = single_tensor_args(0xA0, TensorArgType::OUTPUT); + TaskArgs a1 = single_tensor_args(0xA1, TensorArgType::OUTPUT); + + auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {1, 0}); + TaskSlot slot = res.task_slot; + + worker_a.wait_running(); + worker_b.wait_running(); + + EXPECT_EQ(worker_a.dispatched_count(), 1); + EXPECT_EQ(worker_b.dispatched_count(), 1); + + EXPECT_EQ(worker_b.dispatched[0].tensor_key, 0xA0u); + EXPECT_EQ(worker_a.dispatched[0].tensor_key, 0xA1u); (void)slot; worker_a.complete(); @@ -563,7 +590,7 @@ TEST_F(GroupSchedulerFixture, GroupDispatchesToNWorkers) { TEST_F(GroupSchedulerFixture, GroupCompletesOnlyWhenAllDone) { TaskArgs a0 = single_tensor_args(0xB0, TensorArgType::OUTPUT); TaskArgs a1 = single_tensor_args(0xB1, TensorArgType::OUTPUT); - auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg); + auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {0, 1}); TaskSlot slot = res.task_slot; worker_a.wait_running(); @@ -577,10 +604,107 @@ TEST_F(GroupSchedulerFixture, GroupCompletesOnlyWhenAllDone) { wait_consumed(slot); } +TEST_F(GroupSchedulerFixture, BlockedGroupDoesNotDispatchPartiallyOrReserveIdleWorker) { + auto running = orch.submit_next_level(C(70), single_tensor_args(0xF0, TensorArgType::OUTPUT), cfg, 0); + worker_a.wait_running(); + ASSERT_TRUE(worker_a.is_running.load()); + + TaskArgs group_a = single_tensor_args(0xF1, TensorArgType::OUTPUT); + TaskArgs group_b = single_tensor_args(0xF2, TensorArgType::OUTPUT); + auto group = orch.submit_next_level_group(C(71), {group_a, group_b}, cfg, {0, 1}); + + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + EXPECT_EQ(S(group.task_slot).state.load(), TaskState::READY); + EXPECT_FALSE(worker_b.is_running.load()); + EXPECT_EQ(worker_b.dispatched_count(), 0); + + auto independent = orch.submit_next_level(C(72), single_tensor_args(0xF3, TensorArgType::OUTPUT), cfg, 1); + worker_b.wait_running(); + ASSERT_TRUE(worker_b.is_running.load()); + EXPECT_EQ(worker_b.dispatched[0].callable_hash0, 72u); + EXPECT_EQ(S(group.task_slot).state.load(), TaskState::READY); + + worker_b.complete(); + wait_consumed(independent.task_slot); + worker_a.complete(); + + auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(500); + while ((!worker_a.is_running.load() || !worker_b.is_running.load()) && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_TRUE(worker_a.is_running.load()); + ASSERT_TRUE(worker_b.is_running.load()); + EXPECT_EQ(worker_a.dispatched[1].callable_hash0, 71u); + EXPECT_EQ(worker_b.dispatched[1].callable_hash0, 71u); + + worker_a.complete(); + worker_b.complete(); + wait_consumed(running.task_slot); + wait_consumed(group.task_slot); +} + +TEST_F(GroupSchedulerFixture, LaunchableGroupPrecedesConflictingSingles) { + auto running_a = orch.submit_next_level(C(73), single_tensor_args(0xF4, TensorArgType::OUTPUT), cfg, 0); + auto running_b = orch.submit_next_level(C(74), single_tensor_args(0xF5, TensorArgType::OUTPUT), cfg, 1); + worker_a.wait_running(); + worker_b.wait_running(); + ASSERT_TRUE(worker_a.is_running.load()); + ASSERT_TRUE(worker_b.is_running.load()); + + TaskArgs group_a = single_tensor_args(0xF6, TensorArgType::OUTPUT); + TaskArgs group_b = single_tensor_args(0xF7, TensorArgType::OUTPUT); + SubmitResult group; + SubmitResult single_a; + SubmitResult single_b; + { + std::lock_guard scheduler_pause(sched.loop_mutex()); + group = orch.submit_next_level_group(C(75), {group_a, group_b}, cfg, {0, 1}); + single_a = orch.submit_next_level(C(76), single_tensor_args(0xF8, TensorArgType::OUTPUT), cfg, 0); + single_b = orch.submit_next_level(C(77), single_tensor_args(0xF9, TensorArgType::OUTPUT), cfg, 1); + worker_a.complete(); + worker_b.complete(); + + auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(500); + while (manager.any_busy() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_FALSE(manager.any_busy()); + } + + worker_a.wait_running(); + worker_b.wait_running(); + ASSERT_EQ(worker_a.dispatched_count(), 2); + ASSERT_EQ(worker_b.dispatched_count(), 2); + EXPECT_EQ(worker_a.dispatched[1].callable_hash0, 75u); + EXPECT_EQ(worker_b.dispatched[1].callable_hash0, 75u); + + worker_a.complete(); + worker_b.complete(); + wait_consumed(group.task_slot); + + auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(500); + while ((worker_a.dispatched_count() < 3 || worker_b.dispatched_count() < 3) && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_EQ(worker_a.dispatched_count(), 3); + ASSERT_EQ(worker_b.dispatched_count(), 3); + EXPECT_EQ(worker_a.dispatched[2].callable_hash0, 76u); + EXPECT_EQ(worker_b.dispatched[2].callable_hash0, 77u); + worker_a.complete(); + worker_b.complete(); + + wait_consumed(running_a.task_slot); + wait_consumed(running_b.task_slot); + wait_consumed(single_a.task_slot); + wait_consumed(single_b.task_slot); +} + TEST_F(GroupSchedulerFixture, GroupFailureWaitsForRunningMembersThenConsumes) { TaskArgs a0 = single_tensor_args(0xC0, TensorArgType::OUTPUT); TaskArgs a1 = single_tensor_args(0xC1, TensorArgType::OUTPUT); - auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg); + auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {0, 1}); TaskSlot slot = res.task_slot; worker_a.wait_running(); @@ -602,7 +726,7 @@ TEST_F(GroupSchedulerFixture, GroupFailureWaitsForRunningMembersThenConsumes) { TEST_F(GroupSchedulerFixture, InvalidGroupIndexFailsAndConsumesGroup) { TaskArgs a0 = single_tensor_args(0xD0, TensorArgType::OUTPUT); TaskArgs a1 = single_tensor_args(0xD1, TensorArgType::OUTPUT); - auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg); + auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {0, 1}); TaskSlot slot = res.task_slot; worker_a.wait_running(); @@ -622,9 +746,9 @@ TEST_F(GroupSchedulerFixture, InvalidGroupIndexFailsAndConsumesGroup) { worker_b.complete(); } -TEST_F(GroupSchedulerFixture, EndpointEligibilityRestrictsIdleSelection) { +TEST_F(GroupSchedulerFixture, ExplicitTargetWithinEligibilityIsUsed) { TaskArgs args = single_tensor_args(0xE0, TensorArgType::OUTPUT); - auto res = orch.submit_next_level(C(55), args, cfg, -1, {1}); + auto res = orch.submit_next_level(C(55), args, cfg, 1, {1}); TaskSlot slot = res.task_slot; worker_b.wait_running(); @@ -637,22 +761,84 @@ TEST_F(GroupSchedulerFixture, EndpointEligibilityRestrictsIdleSelection) { wait_consumed(slot); } -TEST_F(GroupSchedulerFixture, AffinityMustBeInEligibleEndpointSet) { +TEST_F(GroupSchedulerFixture, BusyTargetDoesNotBlockAnotherWorkerQueue) { + auto running_args = single_tensor_args(0xE4, TensorArgType::OUTPUT); + auto running = orch.submit_next_level(C(62), running_args, cfg, 0); + worker_a.wait_running(); + ASSERT_TRUE(worker_a.is_running.load()); + + auto blocked_args = single_tensor_args(0xE5, TensorArgType::OUTPUT); + auto blocked = orch.submit_next_level(C(63), blocked_args, cfg, 0); + auto blocked_second_args = single_tensor_args(0xE8, TensorArgType::OUTPUT); + auto blocked_second = orch.submit_next_level(C(67), blocked_second_args, cfg, 0); + auto independent_args = single_tensor_args(0xE6, TensorArgType::OUTPUT); + auto independent = orch.submit_next_level(C(64), independent_args, cfg, 1); + + worker_b.wait_running(); + ASSERT_TRUE(worker_b.is_running.load()); + EXPECT_EQ(worker_b.dispatched_count(), 1); + EXPECT_EQ(worker_b.dispatched[0].callable_hash0, 64u); + + worker_b.complete(); + wait_consumed(independent.task_slot); + worker_a.complete(); + + auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(500); + while (worker_a.dispatched_count() < 2 && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_EQ(worker_a.dispatched_count(), 2); + EXPECT_EQ(worker_a.dispatched[1].callable_hash0, 63u); + worker_a.complete(); + + deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(500); + while (worker_a.dispatched_count() < 3 && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_EQ(worker_a.dispatched_count(), 3); + EXPECT_EQ(worker_a.dispatched[2].callable_hash0, 67u); + worker_a.complete(); + wait_consumed(running.task_slot); + wait_consumed(blocked.task_slot); + wait_consumed(blocked_second.task_slot); +} + +TEST_F(GroupSchedulerFixture, DependencyReleaseUsesConsumerWorkerQueue) { + auto producer_args = single_tensor_args(0xE7, TensorArgType::OUTPUT); + auto producer = orch.submit_next_level(C(65), producer_args, cfg, 0); + auto consumer_args = single_tensor_args(0xE7, TensorArgType::INPUT); + auto consumer = orch.submit_next_level(C(66), consumer_args, cfg, 1); + EXPECT_EQ(S(consumer.task_slot).state.load(), TaskState::PENDING); + + worker_a.wait_running(); + ASSERT_TRUE(worker_a.is_running.load()); + EXPECT_FALSE(worker_b.is_running.load()); + worker_a.complete(); + + worker_b.wait_running(); + ASSERT_TRUE(worker_b.is_running.load()); + EXPECT_EQ(worker_b.dispatched[0].callable_hash0, 66u); + worker_b.complete(); + wait_consumed(producer.task_slot); + wait_consumed(consumer.task_slot); +} + +TEST_F(GroupSchedulerFixture, TargetMustBeInEligibleEndpointSet) { TaskArgs args = single_tensor_args(0xE1, TensorArgType::OUTPUT); EXPECT_THROW((void)orch.submit_next_level(C(56), args, cfg, 0, {1}), std::invalid_argument); } TEST_F(GroupSchedulerFixture, UnknownEligibleWorkerIdIsRejectedBeforeScheduling) { TaskArgs args = single_tensor_args(0xE3, TensorArgType::OUTPUT); - EXPECT_THROW((void)orch.submit_next_level(C(59), args, cfg, -1, {99}), std::invalid_argument); + EXPECT_THROW((void)orch.submit_next_level(C(59), args, cfg, 99, {99}), std::invalid_argument); } -TEST(SchedulerWorkerAffinityTest, NextLevelAffinityUsesWorkerIdNotVectorIndex) { +TEST(SchedulerWorkerTargetTest, NextLevelTargetUsesWorkerIdNotVectorIndex) { TensorMap tm; Ring allocator; Scope scope; - ReadyQueue rq_next_level; ReadyQueue rq_sub; + NextLevelReadyQueues rq_next_level; Orchestrator orch; MockMailboxWorker worker_a; MockMailboxWorker worker_b; @@ -663,10 +849,6 @@ TEST(SchedulerWorkerAffinityTest, NextLevelAffinityUsesWorkerIdNotVectorIndex) { std::mutex consumed_mu; allocator.init(/*heap_bytes=*/1ULL << 20); - orch.init(&tm, &allocator, &scope, &rq_next_level, &rq_sub, &manager, [&sched] { - sched.notify_ready(); - }); - worker_a.start(); worker_b.start(); manager.add_next_level_at(7, worker_a.mailbox_ptr()); @@ -674,12 +856,19 @@ TEST(SchedulerWorkerAffinityTest, NextLevelAffinityUsesWorkerIdNotVectorIndex) { manager.start(&allocator, [&sched](WorkerCompletion completion) { sched.worker_done(std::move(completion)); }); + rq_next_level.reset(manager.next_level_worker_ids()); + orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level, &manager, [&sched] { + sched.notify_ready(); + }); Scheduler::Config c; c.ring = &allocator; - c.ready_next_level_queue = &rq_next_level; c.ready_sub_queue = &rq_sub; + c.ready_next_level_queues = &rq_next_level; c.manager = &manager; + c.enqueue_ready_cb = [&orch](TaskSlot slot) { + orch.enqueue_ready(slot); + }; c.on_consumed_cb = [&orch, &consumed_slots, &consumed_mu](TaskSlot s) { orch.on_consumed(s); std::lock_guard lk(consumed_mu); @@ -717,7 +906,7 @@ TEST(SchedulerWorkerAffinityTest, NextLevelAffinityUsesWorkerIdNotVectorIndex) { TaskArgs a0 = single_tensor_args(0xE6, TensorArgType::OUTPUT); TaskArgs a1 = single_tensor_args(0xE7, TensorArgType::OUTPUT); - auto group_res = orch.submit_next_level_group(C(61), {a0, a1}, cfg, {}, {{7}, {9}}); + auto group_res = orch.submit_next_level_group(C(61), {a0, a1}, cfg, {7, 9}, {{7}, {9}}); worker_a.wait_running(); worker_b.wait_running(); @@ -753,22 +942,21 @@ TEST_F(GroupSchedulerFixture, RemoteSidecarRejectsLocalEndpointEligibility) { sidecar.tensors[0].desc.generation = 1; sidecar.tensors[0].desc.nbytes = 1; - EXPECT_THROW((void)orch.submit_next_level(C(57), args, cfg, -1, {0}, sidecar), std::invalid_argument); + EXPECT_THROW((void)orch.submit_next_level(C(57), args, cfg, 0, {0}, sidecar), std::invalid_argument); } // =========================================================================== -// Strict-4: per-worker-type ready queues (no head-of-line blocking across -// types). Covered here with one NEXT_LEVEL worker + one SUB worker: with a -// saturated NEXT_LEVEL pool, a SUB task submitted afterwards must still -// dispatch immediately instead of waiting behind the stuck next-level task. +// Directed NEXT_LEVEL and shared SUB queues do not block each other. Covered +// here with one worker of each type: a SUB task submitted while the exact +// NEXT_LEVEL target is busy must still dispatch immediately. // =========================================================================== struct MixedTypeSchedulerFixture : public ::testing::Test { TensorMap tm; Ring allocator; Scope scope; - ReadyQueue rq_next_level; ReadyQueue rq_sub; + NextLevelReadyQueues rq_next_level; Orchestrator orch; MockMailboxWorker next_level_worker; MockMailboxWorker sub_worker; @@ -783,9 +971,6 @@ struct MixedTypeSchedulerFixture : public ::testing::Test { void SetUp() override { allocator.init(/*heap_bytes=*/1ULL << 20); - orch.init(&tm, &allocator, &scope, &rq_next_level, &rq_sub, &manager, [this] { - sched.notify_ready(); - }); next_level_worker.start(); sub_worker.start(); @@ -794,12 +979,19 @@ struct MixedTypeSchedulerFixture : public ::testing::Test { manager.start(&allocator, [this](WorkerCompletion completion) { sched.worker_done(std::move(completion)); }); + rq_next_level.reset(manager.next_level_worker_ids()); + orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level, &manager, [this] { + sched.notify_ready(); + }); Scheduler::Config c; c.ring = &allocator; - c.ready_next_level_queue = &rq_next_level; c.ready_sub_queue = &rq_sub; + c.ready_next_level_queues = &rq_next_level; c.manager = &manager; + c.enqueue_ready_cb = [this](TaskSlot slot) { + orch.enqueue_ready(slot); + }; c.on_consumed_cb = [this](TaskSlot s) { orch.on_consumed(s); std::lock_guard lk(consumed_mu); @@ -835,15 +1027,14 @@ TEST_F(MixedTypeSchedulerFixture, SubTaskDispatchesWhileNextLevelPoolSaturated) // Submit a next-level task; the only chip worker begins running it and // stays blocked until we call complete() on it. auto chip_args = single_tensor_args(0xAAA, TensorArgType::OUTPUT); - auto chip = orch.submit_next_level(C(20), chip_args, cfg); + auto chip = orch.submit_next_level(C(20), chip_args, cfg, 0); next_level_worker.wait_running(); ASSERT_TRUE(next_level_worker.is_running.load()); // Now submit a sub task while the chip pool is saturated. With a single // shared ready queue this would block behind any next-level task sitting - // at the queue head waiting for a free chip worker. With per-type - // queues (Strict-4) it must dispatch immediately to the idle sub - // worker. + // in worker 0's directed FIFO. The independent shared SUB queue must + // dispatch immediately to the idle SUB worker. auto sub_args = single_tensor_args(0xBBB, TensorArgType::OUTPUT); auto sub = orch.submit_sub(C(7), sub_args); @@ -852,7 +1043,7 @@ TEST_F(MixedTypeSchedulerFixture, SubTaskDispatchesWhileNextLevelPoolSaturated) EXPECT_TRUE(next_level_worker.is_running.load()) << "chip worker must still be busy"; // Complete the sub task first; it reaches CONSUMED while the chip task - // is still running -- demonstrating independent per-type dispatch. + // is still running -- demonstrating independent queue dispatch. sub_worker.complete(); wait_consumed(sub.task_slot); EXPECT_FALSE(is_consumed(chip.task_slot)); @@ -866,10 +1057,10 @@ TEST_F(GroupSchedulerFixture, GroupDependencyChain) { // Task B reads INPUT at the same key -- depends on group A. TaskArgs a0 = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); TaskArgs a1 = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); - auto a = orch.submit_next_level_group(C(42), {a0, a1}, cfg); + auto a = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {0, 1}); auto args_b = single_tensor_args(0xCAFE, TensorArgType::INPUT); - auto b = orch.submit_next_level(C(42), args_b, cfg); + auto b = orch.submit_next_level(C(42), args_b, cfg, 0); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::PENDING); worker_a.wait_running(); @@ -883,7 +1074,7 @@ TEST_F(GroupSchedulerFixture, GroupDependencyChain) { std::this_thread::sleep_for(std::chrono::milliseconds(1)); } int total = worker_a.dispatched_count() + worker_b.dispatched_count(); - EXPECT_GE(total, 3); // 2 from group A + 1 from B + EXPECT_EQ(total, 3); // 2 from group A + exactly 1 downstream dispatch if (worker_a.is_running.load()) worker_a.complete(); if (worker_b.is_running.load()) worker_b.complete(); diff --git a/tests/ut/py/test_callable_identity.py b/tests/ut/py/test_callable_identity.py index 61e2a5cd9b..fd867767b9 100644 --- a/tests/ut/py/test_callable_identity.py +++ b/tests/ut/py/test_callable_identity.py @@ -590,7 +590,7 @@ def test_remote_session_manifest_requires_wildcard_bind_opt_in(): worker.close() -def test_remote_submit_worker_affinity_uses_stable_worker_id_after_mixed_add_order(): +def test_remote_submit_target_uses_stable_worker_id_after_mixed_add_order(): class FakeCOrchestrator: def submit_next_level(self, *args): self.submit_next_level_args = args @@ -614,7 +614,7 @@ def submit_next_level(self, *args): worker.close() -def test_local_submit_worker_affinity_maps_stable_worker_id_after_mixed_add_order(): +def test_local_submit_target_maps_stable_worker_id_after_mixed_add_order(): class FakeCOrchestrator: def submit_next_level(self, *args): self.submit_next_level_args = args @@ -658,6 +658,37 @@ def submit_next_level(self, *args): assert call[6] == [] +def test_next_level_submit_requires_valid_explicit_targets(): + class FakeCOrchestrator: + def submit_next_level(self, *args): + self.submit_next_level_args = args + + def submit_next_level_group(self, *args): + self.submit_next_level_group_args = args + + handle = CallableHandle("sha256:" + "00" * 32, "CHIP_CALLABLE", "LOCAL_CHIP") + orch = Orchestrator(FakeCOrchestrator()) + + with pytest.raises(TypeError, match="worker"): + orch.submit_next_level(handle, TaskArgs()) + with pytest.raises(ValueError, match="non-negative"): + orch.submit_next_level(handle, TaskArgs(), worker=-1) + for worker in (True, 1.0, "1"): + with pytest.raises(TypeError, match="integer"): + orch.submit_next_level(handle, TaskArgs(), worker=cast(int, worker)) + with pytest.raises(TypeError, match="workers"): + orch.submit_next_level_group(handle, [TaskArgs()]) + with pytest.raises(ValueError, match="length"): + orch.submit_next_level_group(handle, [TaskArgs(), TaskArgs()], workers=[0]) + with pytest.raises(ValueError, match="non-negative"): + orch.submit_next_level_group(handle, [TaskArgs()], workers=[-1]) + for worker in (True, 1.0, "1"): + with pytest.raises(TypeError, match="integer"): + orch.submit_next_level_group(handle, [TaskArgs()], workers=[worker]) + with pytest.raises(ValueError, match="duplicate"): + orch.submit_next_level_group(handle, [TaskArgs(), TaskArgs()], workers=[0, 0]) + + def test_remote_callable_submit_passes_remote_sidecar_to_cpp_facade(): class FakeCOrchestrator: def submit_next_level(self, *args): @@ -682,7 +713,7 @@ def submit_next_level_group(self, *args): ) args = TaskArgs() args.add_tensor(RemoteTensorRef(buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) - orch.submit_next_level(handle, args) + orch.submit_next_level(handle, args, worker=worker_id) call = fake.submit_next_level_args assert call[1] == "PYTHON_IMPORT" @@ -695,7 +726,7 @@ def submit_next_level_group(self, *args): bare = TaskArgs() bare.add_tensor(Tensor.make(0x1234, (1,), DataType.UINT8), TensorArgType.INPUT) - orch.submit_next_level(handle, bare) + orch.submit_next_level(handle, bare, worker=worker_id) bare_sidecar = fake.submit_next_level_args[7] assert len(bare_sidecar.tensors) == 1 @@ -820,7 +851,7 @@ def submit_next_level(self, *args): args.add_tensor(RemoteTensorRef(remote_buf, shape=(4,), dtype=DataType.UINT8), tag) with pytest.raises(ValueError, match="remote tensor .* access"): - orch.submit_next_level(handle, args) + orch.submit_next_level(handle, args, worker=worker_id) finally: worker.close() @@ -860,7 +891,7 @@ def submit_next_level(self, *args): args = TaskArgs() args.add_tensor(RemoteTensorRef(remote_buf, shape=(4,), dtype=DataType.UINT8), tag) - orch.submit_next_level(handle, args) + orch.submit_next_level(handle, args, worker=worker_id) assert fake.submit_next_level_args[6] == [worker_id] finally: @@ -892,7 +923,7 @@ def submit_next_level_group(self, *args): ) args = TaskArgs() args.add_tensor(RemoteTensorRef(buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) - orch.submit_next_level(handle, args) + orch.submit_next_level(handle, args, worker=worker_id1) assert fake.submit_next_level_args[6] == [worker_id1] finally: @@ -926,7 +957,7 @@ def submit_next_level_group(self, *args): args.add_tensor(RemoteTensorRef(buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) with pytest.raises(ValueError, match="no eligible remote worker"): - orch.submit_next_level(handle, args) + orch.submit_next_level(handle, args, worker=worker_id0) finally: worker.close() @@ -965,7 +996,7 @@ def submit_next_level_group(self, *args): args0.add_tensor(RemoteTensorRef(buf0, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) args1 = TaskArgs() args1.add_tensor(RemoteTensorRef(buf1, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) - orch.submit_next_level_group(handle, [args0, args1]) + orch.submit_next_level_group(handle, [args0, args1], workers=[worker_id0, worker_id1]) assert fake.submit_next_level_group_args[6] == [[worker_id0], [worker_id1]] finally: diff --git a/tests/ut/py/test_worker/test_child_addr_guard.py b/tests/ut/py/test_worker/test_child_addr_guard.py index f4c392ced9..5196443dc8 100644 --- a/tests/ut/py/test_worker/test_child_addr_guard.py +++ b/tests/ut/py/test_worker/test_child_addr_guard.py @@ -133,7 +133,7 @@ def test_empty_entry_is_not_live_fail_closed(self): with pytest.raises(ValueError, match="not a live malloc base"): w._child_prov_require_malloc_base(0, 0x1000, api="free") with pytest.raises(ValueError, match="not a live allocation on target worker 0"): - w._child_prov_check_dispatch([(0x1000, 0)], {0}, api="submit_next_level") + w._child_prov_check_dispatch([(0x1000, 0)], 0, api="submit_next_level") def test_drop_last_role_deletes_entry_no_empty_state(self): # Dropping the last role deletes the key outright — it never leaves a @@ -167,37 +167,17 @@ def test_strict_aba_is_explicitly_not_covered(self): class TestDispatchResolution: - def test_candidates_pinned_worker(self): - o = Orchestrator(MagicMock(), _l3()) - assert o._child_dispatch_candidates(2, [0, 1, 2]) == {2} - - def test_candidates_wildcard_uses_eligible_set(self): - o = Orchestrator(MagicMock(), _l3()) - assert o._child_dispatch_candidates(-1, [0, 1]) == {0, 1} - - def test_candidates_wildcard_unconstrained_uses_full_pool(self): - w = _l3() - w._chip_shms = [object(), object(), object()] # 3 chips - o = Orchestrator(MagicMock(), w) - assert o._child_dispatch_candidates(-1, []) == {0, 1, 2} - def test_unique_target_and_live_passes(self): w = _l3() _record_malloc(w, 0, 0x1000) with w._child_prov_lock: - w._child_prov_check_dispatch([(0x1000, 0)], {0}, api="submit_next_level") - - def test_ambiguous_target_rejected(self): - w = _l3() - _record_malloc(w, 0, 0x1000) - with w._child_prov_lock, pytest.raises(ValueError, match="cannot resolve a unique"): - w._child_prov_check_dispatch([(0x1000, 0)], {0, 1}, api="submit_next_level") + w._child_prov_check_dispatch([(0x1000, 0)], 0, api="submit_next_level") def test_unique_target_but_wrong_worker_rejected(self): w = _l3() _record_malloc(w, 0, 0x1000) # lives on worker 0 with w._child_prov_lock, pytest.raises(ValueError, match="not a live allocation on target worker 1"): - w._child_prov_check_dispatch([(0x1000, 0)], {1}, api="submit_next_level") + w._child_prov_check_dispatch([(0x1000, 0)], 1, api="submit_next_level") # ---------------------------------------------------------------------------- @@ -296,29 +276,6 @@ def test_child_arg_to_wrong_worker_rejected(self, _fake_handle): o.submit_next_level(object(), _child_args(0x1000), None, worker=1) fake.submit_next_level.assert_not_called() - def test_child_arg_wildcard_ambiguous_rejected(self, _fake_handle): - _fake_handle["eligible"] = (0, 1) # two eligible targets - w = _l3() - _record_malloc(w, 0, 0x1000) - fake = MagicMock() - o = Orchestrator(fake, w) - with pytest.raises(ValueError, match="cannot resolve a unique"): - o.submit_next_level(object(), _child_args(0x1000), None, worker=-1) - fake.submit_next_level.assert_not_called() - - def test_child_arg_wildcard_narrowed_passes_resolved_affinity_to_cpp(self, _fake_handle): - # worker=-1 narrowed to a unique owner must be passed to C++ as that - # worker (not the raw -1), so the child TensorKey is keyed by its owner - # and deps against an explicit-worker submit of the same buffer are seen. - _fake_handle["eligible"] = (0,) - w = _l3() - _record_malloc(w, 0, 0x1000) - fake = MagicMock() - o = Orchestrator(fake, w) - o.submit_next_level(object(), _child_args(0x1000), None, worker=-1) - fake.submit_next_level.assert_called_once() - assert fake.submit_next_level.call_args.args[5] == 0 # cpp_worker_id, not -1 - def test_host_only_args_are_not_guarded(self, _fake_handle): # A submit with no child_memory tensor never touches provenance. w = _l3() @@ -326,7 +283,7 @@ def test_host_only_args_are_not_guarded(self, _fake_handle): o = Orchestrator(fake, w) args = TaskArgs() args.add_tensor(Tensor.make(0, (16,), DataType.FLOAT32, child_memory=False), TensorArgType.INPUT) - o.submit_next_level(object(), args, None, worker=-1) + o.submit_next_level(object(), args, None, worker=0) fake.submit_next_level.assert_called_once() def test_group_member_child_arg_wrong_worker_rejected(self, _fake_handle): @@ -349,20 +306,6 @@ def test_group_member_child_arg_correct_worker_passes(self, _fake_handle): o.submit_next_level_group(object(), [TaskArgs(), _child_args(0x1000)], None, workers=[0, 1]) fake.submit_next_level_group.assert_called_once() - def test_group_child_member_wildcard_passes_resolved_affinity_to_cpp(self, _fake_handle): - # A single-member group (schedulable) whose child member's eligibility - # narrows to a unique owner materialises a full per-member affinity - # pinning it to that owner. - _fake_handle["eligible"] = (1,) - w = _l3() - w._chip_shms = [object(), object()] - _record_malloc(w, 1, 0x1000) - fake = MagicMock() - o = Orchestrator(fake, w) - o.submit_next_level_group(object(), [_child_args(0x1000)], None, workers=None) - fake.submit_next_level_group.assert_called_once() - assert fake.submit_next_level_group.call_args.args[5] == [1] # cpp_worker_ids - def test_group_child_member_rejects_mismatched_workers_length(self, _fake_handle): # A non-empty workers list must be one-per-member; a short list must NOT # be silently padded (that would bypass the C++ length check). @@ -371,23 +314,33 @@ def test_group_child_member_rejects_mismatched_workers_length(self, _fake_handle _record_malloc(w, 0, 0x1000) fake = MagicMock() o = Orchestrator(fake, w) - with pytest.raises(ValueError, match="workers length 1 != 2 args"): + with pytest.raises(ValueError, match="workers length must match"): o.submit_next_level_group(object(), [_child_args(0x1000), TaskArgs()], None, workers=[0]) fake.submit_next_level_group.assert_not_called() - def test_group_two_child_members_same_owner_rejected(self, _fake_handle): - # Two child members owned by the same worker resolve to the same - # affinity; a group must dispatch to distinct workers, so reject rather - # than silently serialize them on one WorkerThread. - _fake_handle["eligible"] = (0,) + def test_local_callable_rejects_remote_worker_target(self, _fake_handle): + # A LOCAL_CHIP callable pinned to a remote worker id would enqueue on the + # remote endpoint, whose manifest lacks the local digest -> async unknown + # hashid. The Python guard rejects it up front; a local chip id passes. w = _l3() - w._chip_shms = [object(), object()] - _record_malloc(w, 0, 0x1000) - _record_malloc(w, 0, 0x2000) + w._chip_shms = [object()] + w._remote_worker_ids = [7] + fake = MagicMock() + o = Orchestrator(fake, w) + with pytest.raises(ValueError, match="remote NEXT_LEVEL worker"): + o.submit_next_level(object(), TaskArgs(), None, worker=7) + fake.submit_next_level.assert_not_called() + o.submit_next_level(object(), TaskArgs(), None, worker=0) + fake.submit_next_level.assert_called_once() + + def test_local_group_rejects_remote_worker_target(self, _fake_handle): + w = _l3() + w._chip_shms = [object()] + w._remote_worker_ids = [7] fake = MagicMock() o = Orchestrator(fake, w) - with pytest.raises(ValueError, match="distinct workers"): - o.submit_next_level_group(object(), [_child_args(0x1000), _child_args(0x2000)], None, workers=None) + with pytest.raises(ValueError, match="remote NEXT_LEVEL worker"): + o.submit_next_level_group(object(), [TaskArgs(), TaskArgs()], None, workers=[0, 7]) fake.submit_next_level_group.assert_not_called() def test_domain_pointer_dispatch_to_owner_then_rejected_after_release(self, _fake_handle): @@ -594,7 +547,7 @@ def test_release_vs_dispatch_rejects_during_native_release(self, monkeypatch): def _fake_dispatch(**kwargs): try: with w._child_prov_lock: - w._child_prov_check_dispatch([(0x5000, 0)], {0}, api="submit_next_level") + w._child_prov_check_dispatch([(0x5000, 0)], 0, api="submit_next_level") outcome["dispatch"] = "allowed" except ValueError: outcome["dispatch"] = "rejected" diff --git a/tests/ut/py/test_worker/test_error_propagation.py b/tests/ut/py/test_worker/test_error_propagation.py index cfe5725ff6..85bbed400d 100644 --- a/tests/ut/py/test_worker/test_error_propagation.py +++ b/tests/ut/py/test_worker/test_error_propagation.py @@ -215,12 +215,12 @@ def l3_orch(orch, args, config): w4 = Worker(level=4, num_sub_workers=0) l3_handle = w4.register(l3_orch) - w4.add_worker(l3) + l3_worker_id = w4.add_worker(l3) w4.init() try: def l4_orch(orch, args, config): - orch.submit_next_level(l3_handle, TaskArgs(), CallConfig()) + orch.submit_next_level(l3_handle, TaskArgs(), CallConfig(), worker=l3_worker_id) with pytest.raises(RuntimeError) as info: w4.run(l4_orch) diff --git a/tests/ut/py/test_worker/test_l4_recursive.py b/tests/ut/py/test_worker/test_l4_recursive.py index bf130050c7..ba9fb3b7e9 100644 --- a/tests/ut/py/test_worker/test_l4_recursive.py +++ b/tests/ut/py/test_worker/test_l4_recursive.py @@ -200,11 +200,11 @@ def test_l4_register_python_orch_after_start_succeeds(self): w4 = Worker(level=4, num_sub_workers=0) bootstrap_handle = w4.register(lambda orch, args, config: None) - w4.add_worker(l3) + l3_worker_id = w4.add_worker(l3) w4.init() def bootstrap(orch, args, config): - orch.submit_next_level(bootstrap_handle, TaskArgs(), CallConfig()) + orch.submit_next_level(bootstrap_handle, TaskArgs(), CallConfig(), worker=l3_worker_id) w4.run(bootstrap) @@ -214,7 +214,7 @@ def dynamic_l3_orch(orch, args, config): dynamic_handle = w4.register(dynamic_l3_orch) def l4_orch(orch, args, config): - orch.submit_next_level(dynamic_handle, TaskArgs(), CallConfig()) + orch.submit_next_level(dynamic_handle, TaskArgs(), CallConfig(), worker=l3_worker_id) w4.run(l4_orch) w4.close() @@ -250,11 +250,11 @@ def l3_orch(orch, args, config): # L4 parent: one next-level child, register L3 orch fn w4 = Worker(level=4, num_sub_workers=0) l3_handle = w4.register(l3_orch) - w4.add_worker(l3) + l3_worker_id = w4.add_worker(l3) w4.init() def l4_orch(orch, args, config): - orch.submit_next_level(l3_handle, TaskArgs(), CallConfig()) + orch.submit_next_level(l3_handle, TaskArgs(), CallConfig(), worker=l3_worker_id) w4.run(l4_orch) w4.close() @@ -284,12 +284,12 @@ def l3_orch(orch, args, config): w4 = Worker(level=4, num_sub_workers=0) l3_handle = w4.register(l3_orch) - w4.add_worker(l3) + l3_worker_id = w4.add_worker(l3) w4.init() def l4_orch(orch, args, config): for _ in range(3): - orch.submit_next_level(l3_handle, TaskArgs(), CallConfig()) + orch.submit_next_level(l3_handle, TaskArgs(), CallConfig(), worker=l3_worker_id) w4.run(l4_orch) w4.close() @@ -326,11 +326,11 @@ def l3_orch(orch, args, config): w4 = Worker(level=4, num_sub_workers=1) l3_handle = w4.register(l3_orch) l4_verify_handle = w4.register(lambda args: _increment_counter(counter_buf)) - w4.add_worker(l3) + l3_worker_id = w4.add_worker(l3) w4.init() def l4_orch(orch, args, config): - orch.submit_next_level(l3_handle, TaskArgs(), CallConfig()) + orch.submit_next_level(l3_handle, TaskArgs(), CallConfig(), worker=l3_worker_id) orch.submit_sub(l4_verify_handle) w4.run(l4_orch) @@ -362,11 +362,11 @@ def l3_orch(orch, args, config): w4 = Worker(level=4, num_sub_workers=0) l3_handle = w4.register(l3_orch) - w4.add_worker(l3) + l3_worker_id = w4.add_worker(l3) w4.init() def l4_orch(orch, args, config): - orch.submit_next_level(l3_handle, TaskArgs(), CallConfig()) + orch.submit_next_level(l3_handle, TaskArgs(), CallConfig(), worker=l3_worker_id) for _ in range(5): w4.run(l4_orch) @@ -403,11 +403,11 @@ def l3_orch(orch, args, config): w4 = Worker(level=4, num_sub_workers=0) l3_handle = w4.register(l3_orch) - w4.add_worker(l3) + l3_worker_id = w4.add_worker(l3) w4.init() def l4_orch(orch, args, config): - orch.submit_next_level(l3_handle, TaskArgs(), CallConfig()) + orch.submit_next_level(l3_handle, TaskArgs(), CallConfig(), worker=l3_worker_id) w4.run(l4_orch) w4.close() @@ -439,11 +439,11 @@ def l3_orch(orch, args, config): w4 = Worker(level=4, num_sub_workers=0) l3_handle = w4.register(l3_orch) - w4.add_worker(l3) + l3_worker_id = w4.add_worker(l3) w4.init() def l4_orch(orch, args, config): - orch.submit_next_level(l3_handle, TaskArgs(), CallConfig()) + orch.submit_next_level(l3_handle, TaskArgs(), CallConfig(), worker=l3_worker_id) w4.run(l4_orch) w4.close() diff --git a/tests/ut/py/test_worker/test_startup_readiness.py b/tests/ut/py/test_worker/test_startup_readiness.py index e0a87b3e47..abaa718c69 100644 --- a/tests/ut/py/test_worker/test_startup_readiness.py +++ b/tests/ut/py/test_worker/test_startup_readiness.py @@ -322,11 +322,11 @@ def l3_orch(orch, args, config): w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=30.0) l3_handle = w4.register(l3_orch) - w4.add_worker(l3) + l3_worker_id = w4.add_worker(l3) w4.init() def l4_orch(orch, args, config): - orch.submit_next_level(l3_handle, TaskArgs(), CallConfig()) + orch.submit_next_level(l3_handle, TaskArgs(), CallConfig(), worker=l3_worker_id) w4.run(l4_orch) assert _read_counter(counter_buf) == 1 @@ -361,13 +361,13 @@ def l3b_orch(orch, args, config): w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=30.0) ha = w4.register(l3a_orch) hb = w4.register(l3b_orch) - w4.add_worker(l3a) - w4.add_worker(l3b) + l3a_worker_id = w4.add_worker(l3a) + l3b_worker_id = w4.add_worker(l3b) w4.init() def l4_orch(orch, args, config): - orch.submit_next_level(ha, TaskArgs(), CallConfig()) - orch.submit_next_level(hb, TaskArgs(), CallConfig()) + orch.submit_next_level(ha, TaskArgs(), CallConfig(), worker=l3a_worker_id) + orch.submit_next_level(hb, TaskArgs(), CallConfig(), worker=l3b_worker_id) w4.run(l4_orch) assert _read_counter(a_buf) == 1