diff --git a/docs/callable-identity-registration.md b/docs/callable-identity-registration.md index c6122ca66e..46a895f17e 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) ``` diff --git a/docs/directed-next-level-scheduling/pr1-explicit-placement-contract.md b/docs/directed-next-level-scheduling/pr1-explicit-placement-contract.md new file mode 100644 index 0000000000..4a72267836 --- /dev/null +++ b/docs/directed-next-level-scheduling/pr1-explicit-placement-contract.md @@ -0,0 +1,93 @@ +# PR1 Plan: Require Explicit NEXT_LEVEL Placement + +## Objective + +Make the public NEXT_LEVEL submit contract require a concrete stable worker ID +for every task. Keep the existing affinity-aware dispatch implementation in +place so this PR changes the contract and validation without changing queue +ownership or dependency scheduling. + +This PR is the compatibility boundary for the stacked series. After it lands, +new public submissions cannot request automatic worker selection. + +## Scope + +- Require `worker=` on `Orchestrator.submit_next_level`. +- Require `workers=` on `Orchestrator.submit_next_level_group`. +- Require one worker ID per group member. +- Reject negative IDs, duplicate group IDs, unknown workers, and targets that + are outside the final callable/data eligibility set. +- Preserve stable NEXT_LEVEL worker-ID semantics for local and remote workers. +- Update in-repository callers to pass their intended worker IDs. +- Add focused API and validation tests. +- Update contract documentation touched by the API change. + +## Explicit Non-Goals + +- Do not change `submit_sub` or `submit_sub_group`. +- Do not add SUB worker IDs. +- Do not change ready-queue topology. +- Do not remove the scheduler's idle-worker fallback yet. +- Do not change dependency inference, group completion, failure poisoning, + buffer lifetime, or callable registration. +- Do not add compatibility flags or environment-variable gates. + +## Planned File Changes + +### Public Python and binding API + +- `python/simpler/orchestrator.py` + - Make `worker` and `workers` required keyword-only arguments. + - Reject negative and duplicate targets before calling C++. + - Describe the values as exact targets rather than affinities. +- `python/bindings/worker_bind.h` + - Remove nanobind defaults for `worker` and `workers`. + - Update binding help text. + +### C++ submit validation + +- `src/common/hierarchical/orchestrator.cpp` + - Reject missing or negative NEXT_LEVEL targets. + - Require group target count to equal group size. + - Reject duplicate group targets. + - Continue validating target membership in `eligible_worker_ids`. +- `src/common/hierarchical/orchestrator.h` + - Remove public default values for NEXT_LEVEL target arguments. + - Update comments to state the exact-placement contract. + +### Call sites and tests + +- Update examples, scene tests, Python unit tests, and C++ unit tests that rely + on the old unconstrained default. +- Use worker `0` only where the fixture creates exactly one NEXT_LEVEL worker. +- Use explicit member IDs for group fixtures. +- Preserve tests for callable/data eligibility by selecting an ID from the + final eligible set. + +### Documentation + +- Update `docs/orchestrator.md`, `docs/scheduler.md`, `docs/task-flow.md`, + `docs/hierarchical_level_runtime.md`, and remote-L3 placement text. +- Remove statements that `-1`, `None`, or an empty list means unconstrained. + +## Validation + +- Python tests prove missing `worker` and `workers` fail at the API boundary. +- C++ tests prove negative, unknown, duplicate, count-mismatched, and + ineligible targets fail before slot allocation. +- Existing local and remote explicit-placement tests continue to pass. +- Relevant examples and scene tests collect with the required arguments. + +## Size Budget + +- Production code: at most 250 added lines. +- Tests and call-site updates: at most 350 added lines. +- Documentation: at most 250 added lines. +- Total target: at most 850 added lines. + +## Completion Criteria + +- Every public NEXT_LEVEL submit carries exact target worker IDs. +- No SUB API or behavior changes. +- No queue or dispatch algorithm changes. +- The PR builds and its focused Python/C++ tests pass independently. diff --git a/docs/directed-next-level-scheduling/pr2-per-worker-single-queues.md b/docs/directed-next-level-scheduling/pr2-per-worker-single-queues.md new file mode 100644 index 0000000000..717098127a --- /dev/null +++ b/docs/directed-next-level-scheduling/pr2-per-worker-single-queues.md @@ -0,0 +1,93 @@ +# PR2 Plan: Route Single Tasks to Per-Worker Queues + +## Objective + +Replace the shared NEXT_LEVEL queue for non-group tasks with one FIFO queue per +stable NEXT_LEVEL worker ID. A single task becomes dispatchable only from the +queue owned by its submitted target worker. + +PR1 is a prerequisite: every NEXT_LEVEL slot must already carry a validated +target worker ID. + +## Scope + +- Add a fixed, initialization-time registry of NEXT_LEVEL single-task queues. +- Route immediately-ready single tasks to their target worker queue. +- Route dependency-released single tasks to the same target worker queue. +- Dispatch each queue only to its owning worker. +- Leave NEXT_LEVEL group tasks on the existing shared NEXT_LEVEL queue. +- Leave the SUB shared queue and SUB dispatch behavior unchanged. + +## Explicit Non-Goals + +- Do not add or change public APIs. +- Do not implement the group ready queue yet. +- Do not add priorities, work stealing, rebinding, queue scanning, aging, or + resource reservation. +- Do not change the DAG state machine or failure propagation. +- Do not remove group fallback helpers still needed before PR3. + +## Planned Design + +Introduce a small queue owner with a fixed mapping: + +```text +stable NEXT_LEVEL worker ID -> ReadyQueue +``` + +The mapping is initialized after NEXT_LEVEL workers are registered and before +the scheduler starts. It is immutable while submissions are allowed, so queue +lookups need no structural synchronization. Each contained `ReadyQueue` +remains internally synchronized. + +A shared routing helper receives a READY slot and applies this rule: + +```text +NEXT_LEVEL single -> queue[target_worker_id] +NEXT_LEVEL group -> legacy NEXT_LEVEL group-capable queue +SUB -> existing SUB queue +``` + +Both submit-time readiness and completion-time dependency release must call +the same helper. Duplicating this routing logic is not allowed. + +## Planned File Changes + +- `src/common/hierarchical/types.h` and `types.cpp` + - Add the fixed per-worker queue owner or equivalent focused abstraction. +- `src/common/hierarchical/worker_manager.h` and `worker_manager.cpp` + - Expose the registered NEXT_LEVEL worker IDs needed for initialization. +- `src/common/hierarchical/worker.h` and `worker.cpp` + - Own and initialize the per-worker queues. + - Pass them to Orchestrator and Scheduler. +- `src/common/hierarchical/orchestrator.h` and `orchestrator.cpp` + - Route newly READY single slots through the shared routing abstraction. +- `src/common/hierarchical/scheduler.h` and `scheduler.cpp` + - Route newly released single slots identically. + - Drain each per-worker queue only when its worker is idle. +- `tests/ut/cpp/hierarchical/test_orchestrator.cpp` + - Verify submit-time routing by exact worker ID. +- `tests/ut/cpp/hierarchical/test_scheduler.cpp` + - Verify dependency-release routing and independent worker progress. + +## Validation + +- A single task targeted to worker A never dispatches on worker B. +- A busy worker A does not block a READY single task for idle worker B. +- A PENDING task enters the correct queue when its final producer completes. +- A failed producer does not enqueue its consumer. +- Existing group and SUB tests retain their pre-PR2 behavior. + +## Size Budget + +- Production code: at most 450 added lines. +- Tests: at most 350 added lines. +- Documentation adjustments: at most 100 added lines. +- Total target: at most 900 added lines. + +## Completion Criteria + +- All NEXT_LEVEL single dispatch is exact-worker and per-worker FIFO. +- Group and SUB behavior remains unchanged. +- Submit-time and completion-time routing share one implementation. +- The PR builds and focused hierarchical C++ tests pass independently. diff --git a/docs/directed-next-level-scheduling/pr3-group-ready-queue.md b/docs/directed-next-level-scheduling/pr3-group-ready-queue.md new file mode 100644 index 0000000000..aadec150e0 --- /dev/null +++ b/docs/directed-next-level-scheduling/pr3-group-ready-queue.md @@ -0,0 +1,91 @@ +# PR3 Plan: Add Directed Group Dispatch + +## Objective + +Move NEXT_LEVEL group tasks to a dedicated FIFO queue and dispatch a group only +when all user-selected workers are idle. Claim and dispatch the complete worker +set as one scheduler action. + +PR1 and PR2 are prerequisites: group members carry validated exact targets, +and single tasks already use per-worker queues. + +## Scope + +- Give NEXT_LEVEL groups a dedicated FIFO ready queue. +- Route immediately-ready and dependency-released groups to that queue. +- Inspect only the queue head. +- Resolve every member to its submitted stable worker ID. +- Dispatch only when all target workers are idle. +- Dispatch a launchable group before per-worker single queues in the same + scheduler iteration. +- If the head group cannot launch, leave it READY and continue with single + queues without reserving any worker. +- Preserve current group completion aggregation and failure poisoning. + +## Explicit Non-Goals + +- Do not scan past the group queue head. +- Do not reserve a subset of group workers. +- Do not add fairness, aging, priorities, quotas, preemption, or starvation + prevention. +- Do not change SUB group scheduling. +- Do not change group dependency semantics: one group remains one DAG node. + +## Planned Dispatch Rule + +```text +group = group_ready_queue.front +if group exists and every target worker is idle: + pop group + mark slot RUNNING + dispatch every member to its exact target + +for each NEXT_LEVEL worker: + if worker is idle: + dispatch the head of that worker's single queue +``` + +`WorkerThread::dispatch` marks a worker non-idle synchronously. The scheduler +is the only dispatch producer, so resolving all workers before the first +dispatch prevents partial launch caused by an ordinary busy-worker check. + +## Planned File Changes + +- `src/common/hierarchical/types.h` and `types.cpp` + - Add the dedicated group queue to the directed queue owner. +- `src/common/hierarchical/orchestrator.cpp` + - Route READY NEXT_LEVEL groups to the group queue. +- `src/common/hierarchical/scheduler.cpp` + - Add launchable-head group dispatch before single dispatch. + - Remove idle-pool filling for NEXT_LEVEL groups. + - Keep SUB dispatch on its existing path. +- `tests/ut/cpp/hierarchical/test_scheduler.cpp` + - Verify exact member-to-worker mapping. + - Verify no member dispatches while one target is busy. + - Verify a launchable group wins over conflicting queued singles. + - Verify a blocked group does not reserve idle workers. + - Verify group completion releases downstream consumers only once. + +## Failure and Concurrency Checks + +- Duplicate target IDs are rejected by PR1, before queue insertion. +- Worker lookup failure is treated as an invariant violation, not as a request + for fallback selection. +- Existing group member failure handling remains responsible for terminal + aggregation and downstream poison. +- No new mutex may be held while an endpoint executes. + +## Size Budget + +- Production code: at most 350 added lines. +- Tests: at most 400 added lines. +- Documentation adjustments: at most 100 added lines. +- Total target: at most 850 added lines. + +## Completion Criteria + +- NEXT_LEVEL groups use only their exact submitted workers. +- Group allocation is all-or-nothing at dispatch time. +- The only cross-queue policy is launchable-group-first. +- No reservation, scan, rebinding, or fairness mechanism is introduced. +- Focused group, dependency, and failure tests pass independently. diff --git a/docs/hierarchical_level_runtime.md b/docs/hierarchical_level_runtime.md index 6ebf6b1a5f..3830f7fc39 100644 --- a/docs/hierarchical_level_runtime.md +++ b/docs/hierarchical_level_runtime.md @@ -86,7 +86,7 @@ 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 diff --git a/docs/orchestrator.md b/docs/orchestrator.md index 2245134815..a723c2e93c 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,11 @@ 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 remain unconstrained. --- @@ -92,7 +92,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]; @@ -172,7 +173,11 @@ 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 @@ -232,9 +237,12 @@ SubmitResult Orchestrator::submit_next_level_group(const CallableIdentity &calla } ``` -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`. +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 +451,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()` / diff --git a/docs/remote-l3-worker-design.md b/docs/remote-l3-worker-design.md index c5edc32911..8e4c410bb3 100644 --- a/docs/remote-l3-worker-design.md +++ b/docs/remote-l3-worker-design.md @@ -137,7 +137,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]`. @@ -234,9 +234,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: @@ -325,13 +324,10 @@ Required contracts: - `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. +- `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/scheduler.md b/docs/scheduler.md index cf1f1d040a..56361b6f06 100644 --- a/docs/scheduler.md +++ b/docs/scheduler.md @@ -20,8 +20,9 @@ The Scheduler's job: - 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 **ready queues** (Phase 1): dispatch each NEXT_LEVEL single task + only to its requested worker; group and SUB dispatch remain on their shared + queues during this transition. - 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 @@ -40,15 +41,15 @@ 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_; + // Owns one single-task FIFO per stable worker id and one group FIFO. + NextLevelReadyQueues *ready_next_level_queues_; + + // SUB scheduling stays unconstrained and shared. ReadyQueue *ready_sub_queue_; + // Shared READY router used by submit, wiring, and dependency release. + std::function enqueue_ready_cb_; + // Producer: WorkerThread (on endpoint->run() return). // Consumer: Scheduler's own loop, Phase 2. std::queue completion_queue_; @@ -76,24 +77,24 @@ 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`: +NEXT_LEVEL single tasks have one FIFO for every stable worker id. A task is +inserted into the FIFO named by its required `worker` argument: ```cpp -ReadyQueue ready_next_level_queue_; // WorkerType::NEXT_LEVEL tasks -ReadyQueue ready_sub_queue_; // WorkerType::SUB tasks +NextLevelReadyQueues ready_next_level_queues_; // worker FIFOs + group FIFO +ReadyQueue ready_sub_queue_; // all SUB tasks ``` -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. +The Orchestrator owns the routing rule. Both immediately-ready submissions +and Scheduler transitions to READY call the same `enqueue_ready` entry point, +so dependency release cannot accidentally return a directed single task to a +shared queue. A busy target blocks only its own FIFO; the Scheduler continues +checking other worker FIFOs. FIFO order is preserved independently per worker. + +NEXT_LEVEL groups use the dedicated group FIFO. The Scheduler examines only +its head. A group leaves READY only when every submitted target worker is idle; +otherwise the group stays at the head, no worker is reserved, and independent +single-task queues can still dispatch. SUB dispatch remains unchanged. ### Completion queue @@ -172,10 +173,7 @@ void Scheduler::wire_fanout(const WiringEntry &w) { // 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); + enqueue_ready_cb_(csid); } } ``` @@ -192,46 +190,44 @@ The `lock_guard(p.fanout_mu)` + `p.state.load()` check ensures we either: ## 5. Phase 1 — dispatch -`dispatch_ready` drains each per-type ready queue with its own -head-of-line break so one saturated pool cannot stall the other: +`dispatch_ready` tries launchable NEXT_LEVEL groups first, then checks each +NEXT_LEVEL worker's single-task FIFO, and finally drains the SUB queue. The +only cross-queue policy is launchable-group-first. + +Group dispatch is all-or-nothing: + +```cpp +TaskSlot group; +if (ready_next_level_queues_->try_front_group(group)) { + resolve every submitted worker id; + if (every target worker is idle) { + ready_next_level_queues_->try_pop_group(group); + mark every group member RUNNING; + dispatch every member to its submitted target; + } +} +``` + +If any target is busy, the head is not removed or moved to the tail. The +Scheduler does not scan later groups and does not reserve the currently idle +members. It continues to the single-task queues: ```cpp -void Scheduler::dispatch_ready() { - auto drain_one = [&](ReadyQueue *q) { +void Scheduler::dispatch_next_level_singles() { + for (int32_t worker_id : ready_next_level_queues_->worker_ids()) { + WorkerThread *worker = + manager_->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (!worker || !worker->idle()) continue; + 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}); - } + while (worker->idle() && + ready_next_level_queues_->try_pop_single(worker_id, slot)) { + TaskSlotState &task = slots_[slot]; + if (task.state.load() != TaskState::READY) continue; + task.state.store(TaskState::RUNNING); + worker->dispatch({slot, 0}); } - }; - drain_one(ready_next_level_queue_); - drain_one(ready_sub_queue_); + } } ``` @@ -241,18 +237,17 @@ Dispatch hands off a `WorkerDispatch {slot, group_index}` to a and encodes it into the per-WT mailbox — see [worker-manager.md](worker-manager.md) §3 for the dispatch protocol. -**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. +**Directed-single back-pressure**: if the requested worker is busy, its FIFO +is left untouched. Other NEXT_LEVEL worker FIFOs and the SUB queue still make +progress. The ring's back-pressure at the Orch side caps the total number of +in-flight tasks. 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. +For a NEXT_LEVEL single task, `s.get_affinity(0)` is the stable requested +worker id and can differ from the `next_level_threads_` vector index. SUB has +no public affinity. --- @@ -287,12 +282,7 @@ void Scheduler::on_task_complete(const WorkerCompletion &completion) { } 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); + enqueue_ready_cb_(csid); } } diff --git a/docs/task-flow.md b/docs/task-flow.md index da6fa9fb88..eba8c91068 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: ... ``` @@ -527,7 +526,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) diff --git a/examples/workers/l3/README.md b/examples/workers/l3/README.md index bee62b4f6a..ac86c17219 100644 --- a/examples/workers/l3/README.md +++ b/examples/workers/l3/README.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..de61224bcf 100644 --- a/examples/workers/l3/per_task_runtime_env/README.md +++ b/examples/workers/l3/per_task_runtime_env/README.md @@ -22,7 +22,7 @@ def orch_fn(orch, _args, _cfg): 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 + 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 2da4007efe..b88230bb95 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..b8b3462d76 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) @@ -154,15 +154,12 @@ 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() expected_namespace = ( @@ -194,6 +191,8 @@ def submit_next_level( 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) + if cpp_worker_id < 0: + raise ValueError("worker must be a non-negative NEXT_LEVEL worker id") worker = self._worker # Do the (fallible) kind4 provenance analysis BEFORE capturing remote slot # refs, so an exception here can never leave captured refs neither @@ -250,16 +249,21 @@ def submit_next_level_group( # noqa: PLR0912 -- linear per-member sidecar + eli 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 = [int(x) for x in workers] + if len(worker_ids) != len(args_list): + raise ValueError("workers length must match args_list length") + if any(worker_id < 0 for worker_id in worker_ids): + raise ValueError("workers must contain only non-negative NEXT_LEVEL worker ids") + 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) @@ -488,9 +492,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 58af03ea68..3b5af902b9 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) diff --git a/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index 926fb35c68..e5266f6aed 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,8 +151,7 @@ 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 affinities{worker_id}; std::vector> worker_id_sets; if (!eligible_worker_ids.empty()) worker_id_sets = {eligible_worker_ids}; std::vector sidecars; @@ -309,7 +308,7 @@ SubmitResult Orchestrator::submit_impl( // saturated sub pool cannot stall next-level dispatch (and vice versa). 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 +317,29 @@ 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.get_affinity(0), slot); + } + return; + } + ready_sub_queue_->push(slot); +} + void Orchestrator::validate_worker_eligibility( WorkerType worker_type, size_t args_count, const std::vector &affinities, const std::vector> &eligible_worker_ids ) const { - if (!affinities.empty() && affinities.size() != args_count) { + if (worker_type == WorkerType::NEXT_LEVEL && affinities.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(affinities.size()) + + " does not match args length " + std::to_string(args_count) ); } if (!eligible_worker_ids.empty() && eligible_worker_ids.size() != args_count) { @@ -335,6 +349,20 @@ void Orchestrator::validate_worker_eligibility( ); } + if (worker_type == WorkerType::NEXT_LEVEL) { + std::unordered_set unique_targets; + for (int32_t worker_id : affinities) { + 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]; diff --git a/src/common/hierarchical/orchestrator.h b/src/common/hierarchical/orchestrator.h index 84727e1c1e..e0a13d5c93 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -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}; diff --git a/src/common/hierarchical/scheduler.cpp b/src/common/hierarchical/scheduler.cpp index 8b1e1a1789..097b8323e8 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -11,6 +11,7 @@ #include "scheduler.h" +#include #include #include @@ -35,13 +36,9 @@ bool is_terminal_group_state(GroupMemberState state) { // Scheduler // ============================================================================= -// ============================================================================= -// 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 +50,9 @@ 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(); + // Shut down every ready queue so any wait_pop waiters unblock. cfg_.ready_sub_queue->shutdown(); + cfg_.ready_next_level_queues->shutdown(); if (sched_thread_.joinable()) sched_thread_.join(); @@ -167,7 +164,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 +241,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(); } } @@ -339,9 +332,10 @@ void Scheduler::try_consume(TaskSlot slot) { // ============================================================================= 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. + dispatch_next_level_group(); + dispatch_next_level_singles(); + + // SUB scheduling remains unconstrained and shared. auto drain_one = [this](ReadyQueue *q) { TaskSlot slot; while (q->try_pop(slot)) { @@ -420,6 +414,80 @@ void Scheduler::dispatch_ready() { } }; - drain_one(cfg_.ready_next_level_queue); drain_one(cfg_.ready_sub_queue); } + +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.get_affinity(i); + WorkerThread *worker = cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (worker == nullptr || !s.worker_allowed(i, worker_id)) { + 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"); + } + + { + std::lock_guard lk(s.group_mu); + s.group_member_states.assign(static_cast(group_size), GroupMemberState::RUNNING); + s.group_member_outcomes.assign(static_cast(group_size), EndpointOutcome::SKIPPED); + s.group_terminal_count.store(0, std::memory_order_relaxed); + s.group_dispatched_count.store(group_size, std::memory_order_relaxed); + s.group_failed = false; + s.group_first_failure_index = -1; + s.group_first_failure_message.clear(); + } + 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.get_affinity(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..3baee85e52 100644 --- a/src/common/hierarchical/scheduler.h +++ b/src/common/hierarchical/scheduler.h @@ -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,6 @@ 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(); }; diff --git a/src/common/hierarchical/types.cpp b/src/common/hierarchical/types.cpp index 6df625ccc1..60f8bc4d4a 100644 --- a/src/common/hierarchical/types.cpp +++ b/src/common/hierarchical/types.cpp @@ -11,6 +11,8 @@ #include "types.h" +#include + // ============================================================================= // TaskSlotState // ============================================================================= @@ -80,6 +82,13 @@ bool ReadyQueue::empty() const { return q_.empty(); } +bool ReadyQueue::try_front(TaskSlot &out) { + std::lock_guard lk(mu_); + if (q_.empty()) return false; + out = q_.front(); + return true; +} + bool ReadyQueue::wait_pop(TaskSlot &out) { std::unique_lock lk(mu_); cv_.wait(lk, [this] { @@ -98,3 +107,58 @@ void ReadyQueue::shutdown() { } cv_.notify_all(); } + +// ============================================================================= +// 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()); + } +} + +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; +} + +void NextLevelReadyQueues::shutdown() { + for (auto &queue : queues_) { + queue->shutdown(); + } + group_queue_.shutdown(); +} diff --git a/src/common/hierarchical/types.h b/src/common/hierarchical/types.h index b788b9f0d7..4366ccf49c 100644 --- a/src/common/hierarchical/types.h +++ b/src/common/hierarchical/types.h @@ -33,6 +33,7 @@ #include #include #include +#include #include #include #include @@ -427,6 +428,9 @@ class ReadyQueue { bool empty() const; + // Non-blocking: copies the front without removing it. + bool try_front(TaskSlot &out); + // 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); @@ -439,3 +443,26 @@ class ReadyQueue { 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; + void shutdown(); + + 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 aeb10414f4..632d18cc41 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -94,23 +94,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 bf7764286c..cb922ac710 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -204,10 +204,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..de1a1a26dc 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -463,6 +463,15 @@ WorkerThread *WorkerManager::get_worker_by_id(WorkerType type, int32_t worker_id return nullptr; } +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( WorkerType type, const std::vector &exclude, const std::vector &eligible_worker_ids ) const { diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index c4668a0dde..332489c285 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -474,6 +474,7 @@ class WorkerManager { // 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. 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..816f1bddef 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 @@ -233,7 +233,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,7 +242,7 @@ 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(post_handle, chip_args_post, config, worker=0) worker.run(orch_post) assert torch.allclose(args_post.f, torch.full_like(args_post.f, expected), rtol=1e-5, atol=1e-5) @@ -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..da2da41ea3 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,14 @@ 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); + rq_next_level.try_pop_single(3, ready); 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 +374,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 +382,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 +413,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..c75718ac7c 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,8 @@ 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(); @@ -563,7 +565,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 +579,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 +701,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 +721,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,6 +736,68 @@ TEST_F(GroupSchedulerFixture, EndpointEligibilityRestrictsIdleSelection) { wait_consumed(slot); } +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, AffinityMustBeInEligibleEndpointSet) { TaskArgs args = single_tensor_args(0xE1, TensorArgType::OUTPUT); EXPECT_THROW((void)orch.submit_next_level(C(56), args, cfg, 0, {1}), std::invalid_argument); @@ -644,15 +805,15 @@ TEST_F(GroupSchedulerFixture, AffinityMustBeInEligibleEndpointSet) { 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) { 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 +824,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 +831,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 +881,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,7 +917,7 @@ 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); } // =========================================================================== @@ -767,8 +931,8 @@ 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 +947,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 +955,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,7 +1003,7 @@ 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()); @@ -866,10 +1034,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 +1051,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 bd4c9ebecc..4c467ae229 100644 --- a/tests/ut/py/test_callable_identity.py +++ b/tests/ut/py/test_callable_identity.py @@ -658,6 +658,29 @@ 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) + 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="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 +705,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 +718,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 +843,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 +883,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 +915,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 +949,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 +988,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..3ae975cc5a 100644 --- a/tests/ut/py/test_worker/test_child_addr_guard.py +++ b/tests/ut/py/test_worker/test_child_addr_guard.py @@ -296,29 +296,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 +303,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 +326,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,14 +334,12 @@ 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. + # Two group members cannot target the same exact worker. _fake_handle["eligible"] = (0,) w = _l3() w._chip_shms = [object(), object()] @@ -386,8 +347,8 @@ def test_group_two_child_members_same_owner_rejected(self, _fake_handle): _record_malloc(w, 0, 0x2000) 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="duplicate"): + o.submit_next_level_group(object(), [_child_args(0x1000), _child_args(0x2000)], None, workers=[0, 0]) fake.submit_next_level_group.assert_not_called() def test_domain_pointer_dispatch_to_owner_then_rejected_after_release(self, _fake_handle): 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