From a25c6b9559707d37a0714f68671d8490866353db Mon Sep 17 00:00:00 2001 From: puddingfjz <2811443837@qq.com> Date: Tue, 21 Jul 2026 22:23:41 -0700 Subject: [PATCH 1/2] Update: require explicit NEXT_LEVEL placement - Add four executable rollout plans with bounded PR scopes - Require exact worker ids for single and group NEXT_LEVEL submits - Validate malformed, unknown, and ineligible worker targets - Migrate maintained call sites while keeping SUB placement unchanged --- docs/callable-identity-registration.md | 12 +-- .../pr1-explicit-placement-contract.md | 93 +++++++++++++++++++ docs/hierarchical_level_runtime.md | 2 +- docs/orchestrator.md | 19 ++-- docs/remote-l3-worker-design.md | 18 ++-- docs/scheduler.md | 8 +- docs/task-flow.md | 11 +-- examples/workers/l3/README.md | 2 +- .../workers/l3/per_task_runtime_env/README.md | 2 +- .../workers/l3/per_task_runtime_env/main.py | 2 +- python/bindings/worker_bind.h | 9 +- python/simpler/orchestrator.py | 34 ++++--- python/simpler/worker.py | 2 +- src/common/hierarchical/orchestrator.cpp | 23 ++++- src/common/hierarchical/orchestrator.h | 8 +- .../dynamic_register/test_dynamic_register.py | 26 +++--- .../test_l3_dependency.py | 2 +- .../tensormap_and_ringbuffer/test_l3_group.py | 2 +- .../ut/cpp/hierarchical/test_orchestrator.cpp | 55 ++++++----- tests/ut/cpp/hierarchical/test_scheduler.cpp | 36 +++---- tests/ut/py/test_callable_identity.py | 37 ++++++-- .../py/test_worker/test_child_addr_guard.py | 49 +--------- .../py/test_worker/test_error_propagation.py | 4 +- tests/ut/py/test_worker/test_l4_recursive.py | 30 +++--- .../py/test_worker/test_startup_readiness.py | 12 +-- 25 files changed, 297 insertions(+), 201 deletions(-) create mode 100644 docs/directed-next-level-scheduling/pr1-explicit-placement-contract.md 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/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..6642324f9f 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]; @@ -443,11 +444,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..dc324651b7 100644 --- a/docs/scheduler.md +++ b/docs/scheduler.md @@ -242,7 +242,7 @@ 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 +workers that also satisfy placement 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. @@ -250,9 +250,9 @@ caps the total number of in-flight tasks across both types. 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 NEXT_LEVEL placement, `s.get_affinity(i)` is the required 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. --- 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..196838e06d 100644 --- a/src/common/hierarchical/orchestrator.cpp +++ b/src/common/hierarchical/orchestrator.cpp @@ -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; @@ -322,10 +321,10 @@ 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 +334,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..4dcfeec46d 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -98,17 +98,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 = {} ); 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..cf3e676a16 100644 --- a/tests/ut/cpp/hierarchical/test_orchestrator.cpp +++ b/tests/ut/cpp/hierarchical/test_orchestrator.cpp @@ -76,7 +76,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 +85,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 +111,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 +130,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 +141,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 +159,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 +168,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 +184,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 +208,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 +222,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 +240,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 +261,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 +293,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 +307,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 +328,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 +349,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); 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 +371,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 +379,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 +410,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..538a799ff1 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -373,7 +373,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 +387,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 +431,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 +444,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(); @@ -538,7 +538,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(); @@ -563,7 +563,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(); @@ -580,7 +580,7 @@ TEST_F(GroupSchedulerFixture, GroupCompletesOnlyWhenAllDone) { 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 +602,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 +622,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(); @@ -644,7 +644,7 @@ 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) { @@ -717,7 +717,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 +753,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); } // =========================================================================== @@ -835,7 +835,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 +866,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(); 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 From bdb6fe579fb757d7e5a07291aa5cd26f3648f0a0 Mon Sep 17 00:00:00 2001 From: puddingfjz <2811443837@qq.com> Date: Tue, 21 Jul 2026 21:33:05 -0700 Subject: [PATCH 2/2] Update: route NEXT_LEVEL singles per worker Add fixed per-worker ready queues for directed NEXT_LEVEL single tasks and share the READY routing path between submission and dependency release. Keep group and SUB scheduling behavior unchanged, document the transitional queue topology, and cover FIFO and independent-worker progress. --- .../pr2-per-worker-single-queues.md | 93 +++++++++++++ docs/orchestrator.md | 6 +- docs/scheduler.md | 128 +++++++----------- src/common/hierarchical/orchestrator.cpp | 18 ++- src/common/hierarchical/orchestrator.h | 14 +- src/common/hierarchical/scheduler.cpp | 43 ++++-- src/common/hierarchical/scheduler.h | 9 +- src/common/hierarchical/types.cpp | 49 +++++++ src/common/hierarchical/types.h | 21 +++ src/common/hierarchical/worker.cpp | 17 ++- src/common/hierarchical/worker.h | 6 +- src/common/hierarchical/worker_manager.cpp | 9 ++ src/common/hierarchical/worker_manager.h | 1 + .../ut/cpp/hierarchical/test_orchestrator.cpp | 11 +- tests/ut/cpp/hierarchical/test_scheduler.cpp | 115 ++++++++++++++-- 15 files changed, 416 insertions(+), 124 deletions(-) create mode 100644 docs/directed-next-level-scheduling/pr2-per-worker-single-queues.md 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/orchestrator.md b/docs/orchestrator.md index 6642324f9f..ed593ff978 100644 --- a/docs/orchestrator.md +++ b/docs/orchestrator.md @@ -173,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. NEXT_LEVEL groups and SUB tasks remain on their shared queues in +this transition. **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 diff --git a/docs/scheduler.md b/docs/scheduler.md index dc324651b7..0cfa9ef206 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,19 @@ 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). + // NEXT_LEVEL single tasks use one FIFO per stable worker id. + PerWorkerReadyQueues *ready_next_level_single_queues_; + + // Transitional shared queue: NEXT_LEVEL groups only. PR3 gives this + // queue an explicit group contract and name. ReadyQueue *ready_next_level_queue_; + + // 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 +81,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 +PerWorkerReadyQueues ready_next_level_single_queues_; // worker id -> FIFO +ReadyQueue ready_next_level_queue_; // NEXT_LEVEL groups +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. + +This PR intentionally leaves NEXT_LEVEL groups in the previous shared queue +and leaves SUB dispatch unchanged. The next PR replaces the transitional group +path with a dedicated all-targets-ready dispatcher. ### Completion queue @@ -172,10 +177,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 +194,26 @@ 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` first drains the transitional NEXT_LEVEL group queue, then +checks each NEXT_LEVEL worker's single-task FIFO, and finally drains the SUB +queue. The single-task path has no worker selection: ```cpp -void Scheduler::dispatch_ready() { - auto drain_one = [&](ReadyQueue *q) { +void Scheduler::dispatch_next_level_singles() { + for (int32_t worker_id : ready_next_level_single_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_single_queues_->try_pop(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 +223,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 placement 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 placement, `s.get_affinity(i)` is the required 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 +268,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/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index 196838e06d..a5e9c65295 100644 --- a/src/common/hierarchical/orchestrator.cpp +++ b/src/common/hierarchical/orchestrator.cpp @@ -19,13 +19,15 @@ 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 + ReadyQueue *ready_sub_queue, PerWorkerReadyQueues *ready_next_level_single_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_single_queues_ = ready_next_level_single_queues; manager_ = manager; ready_notify_cb_ = std::move(ready_notify_cb); active_tasks_.store(0, std::memory_order_relaxed); @@ -308,7 +310,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); @@ -317,6 +319,18 @@ 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 && !s.is_group()) { + if (ready_next_level_single_queues_ == nullptr) { + throw std::runtime_error("Orchestrator::enqueue_ready: NEXT_LEVEL single queues are not initialized"); + } + ready_next_level_single_queues_->push(s.get_affinity(0), slot); + return; + } + ready_queue_for(s.worker_type)->push(slot); +} + void Orchestrator::validate_worker_eligibility( WorkerType worker_type, size_t args_count, const std::vector &affinities, const std::vector> &eligible_worker_ids diff --git a/src/common/hierarchical/orchestrator.h b/src/common/hierarchical/orchestrator.h index 4dcfeec46d..516d83b925 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -68,13 +68,12 @@ 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. + // NEXT_LEVEL singles use one FIFO per stable worker id. NEXT_LEVEL groups + // and SUB tasks retain their shared queues in this stage of the stack. 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 = {} + ReadyQueue *ready_sub_queue, PerWorkerReadyQueues *ready_next_level_single_queues, + WorkerManager *manager = nullptr, std::function ready_notify_cb = {} ); // Allocate an intermediate buffer from the Worker's HeapRing (MAP_SHARED, @@ -158,6 +157,10 @@ 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; @@ -169,6 +172,7 @@ class Orchestrator { // without being blocked by another pool's saturation. ReadyQueue *ready_next_level_queue_ = nullptr; ReadyQueue *ready_sub_queue_ = nullptr; + PerWorkerReadyQueues *ready_next_level_single_queues_ = nullptr; // Returns the ready queue that owns tasks of the given worker type. // The method itself does not mutate the Orchestrator (hence `const`); diff --git a/src/common/hierarchical/scheduler.cpp b/src/common/hierarchical/scheduler.cpp index 8b1e1a1789..364dc7d28a 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -35,13 +35,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) + cfg.ready_next_level_single_queues == nullptr || cfg.manager == nullptr || !cfg.enqueue_ready_cb) throw std::invalid_argument("Scheduler::start: null config fields"); cfg_ = cfg; @@ -53,9 +49,10 @@ 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. + // Shut down every ready queue so any wait_pop waiters unblock. cfg_.ready_next_level_queue->shutdown(); cfg_.ready_sub_queue->shutdown(); + cfg_.ready_next_level_single_queues->shutdown(); if (sched_thread_.joinable()) sched_thread_.join(); @@ -168,7 +165,8 @@ void Scheduler::run() { std::unique_lock lk(completion_mu_); completion_cv_.wait(lk, [this] { return !completion_queue_.empty() || !cfg_.ready_next_level_queue->empty() || - !cfg_.ready_sub_queue->empty() || stop_requested_.load(std::memory_order_acquire); + !cfg_.ready_next_level_single_queues->empty() || !cfg_.ready_sub_queue->empty() || + stop_requested_.load(std::memory_order_acquire); }); } @@ -244,11 +242,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(); } } @@ -421,5 +415,30 @@ void Scheduler::dispatch_ready() { }; drain_one(cfg_.ready_next_level_queue); + dispatch_next_level_singles(); drain_one(cfg_.ready_sub_queue); } + +void Scheduler::dispatch_next_level_singles() { + for (int32_t worker_id : cfg_.ready_next_level_single_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_single_queues->try_pop(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..b0262306bb 100644 --- a/src/common/hierarchical/scheduler.h +++ b/src/common/hierarchical/scheduler.h @@ -54,12 +54,14 @@ 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. + // NEXT_LEVEL singles use one FIFO per stable worker id. NEXT_LEVEL + // groups and SUB tasks use their shared queues. ReadyQueue *ready_next_level_queue; ReadyQueue *ready_sub_queue; + PerWorkerReadyQueues *ready_next_level_single_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 +103,5 @@ class Scheduler { void poison_task(TaskSlot slot, const std::string &root_message); void try_consume(TaskSlot slot); void dispatch_ready(); + void dispatch_next_level_singles(); }; diff --git a/src/common/hierarchical/types.cpp b/src/common/hierarchical/types.cpp index 6df625ccc1..a9569229aa 100644 --- a/src/common/hierarchical/types.cpp +++ b/src/common/hierarchical/types.cpp @@ -11,6 +11,8 @@ #include "types.h" +#include + // ============================================================================= // TaskSlotState // ============================================================================= @@ -98,3 +100,50 @@ void ReadyQueue::shutdown() { } cv_.notify_all(); } + +// ============================================================================= +// PerWorkerReadyQueues +// ============================================================================= + +void PerWorkerReadyQueues::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("PerWorkerReadyQueues::reset: negative worker id"); + for (int32_t existing : worker_ids_) { + if (existing == worker_id) { + throw std::invalid_argument("PerWorkerReadyQueues::reset: duplicate worker id"); + } + } + worker_ids_.push_back(worker_id); + queues_.push_back(std::make_unique()); + } +} + +size_t PerWorkerReadyQueues::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("PerWorkerReadyQueues: unknown worker id " + std::to_string(worker_id)); +} + +void PerWorkerReadyQueues::push(int32_t worker_id, TaskSlot slot) { queues_[index_for(worker_id)]->push(slot); } + +bool PerWorkerReadyQueues::try_pop(int32_t worker_id, TaskSlot &out) { + return queues_[index_for(worker_id)]->try_pop(out); +} + +bool PerWorkerReadyQueues::empty() const { + for (const auto &queue : queues_) { + if (!queue->empty()) return false; + } + return true; +} + +void PerWorkerReadyQueues::shutdown() { + for (auto &queue : queues_) { + queue->shutdown(); + } +} diff --git a/src/common/hierarchical/types.h b/src/common/hierarchical/types.h index b788b9f0d7..c1318548db 100644 --- a/src/common/hierarchical/types.h +++ b/src/common/hierarchical/types.h @@ -33,6 +33,7 @@ #include #include #include +#include #include #include #include @@ -439,3 +440,23 @@ class ReadyQueue { std::condition_variable cv_; bool shutdown_{false}; }; + +// Fixed NEXT_LEVEL worker-id to single-task ready-queue mapping. Worker +// registration is complete before reset() and the mapping is immutable while +// Orchestrator submissions and Scheduler dispatch are active. +class PerWorkerReadyQueues { +public: + void reset(const std::vector &worker_ids); + void push(int32_t worker_id, TaskSlot slot); + bool try_pop(int32_t worker_id, 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_; +}; diff --git a/src/common/hierarchical/worker.cpp b/src/common/hierarchical/worker.cpp index aeb10414f4..b56529e458 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -94,23 +94,28 @@ 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_single_queues_.reset(manager_.next_level_worker_ids()); + orchestrator_.init( + &tensormap_, &allocator_, &scope_, &ready_next_level_queue_, &ready_sub_queue_, + &ready_next_level_single_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_single_queues = &ready_next_level_single_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..b88b29c5be 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -204,9 +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. + // NEXT_LEVEL singles use one FIFO per stable worker id. Groups remain on + // the shared NEXT_LEVEL queue until directed group dispatch is installed. + PerWorkerReadyQueues ready_next_level_single_queues_; ReadyQueue ready_next_level_queue_; ReadyQueue ready_sub_queue_; Orchestrator orchestrator_; 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/ut/cpp/hierarchical/test_orchestrator.cpp b/tests/ut/cpp/hierarchical/test_orchestrator.cpp index cf3e676a16..672611a454 100644 --- a/tests/ut/cpp/hierarchical/test_orchestrator.cpp +++ b/tests/ut/cpp/hierarchical/test_orchestrator.cpp @@ -32,6 +32,7 @@ struct OrchestratorFixture : public ::testing::Test { // Strict-4: per-type ready queues. ReadyQueue rq_next_level; ReadyQueue rq_sub; + PerWorkerReadyQueues rq_next_level_single; Orchestrator orch; CallConfig cfg; @@ -39,11 +40,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 { + PerWorkerReadyQueues *queues; + bool try_pop(TaskSlot &out) { return queues->try_pop(0, out); } + } rq{&rq_next_level_single}; void SetUp() override { allocator.init(/*heap_bytes=*/1ULL << 20); - orch.init(&tm, &allocator, &scope, &rq_next_level, &rq_sub); + rq_next_level_single.reset({0, 1, 3}); + orch.init(&tm, &allocator, &scope, &rq_next_level, &rq_sub, &rq_next_level_single); } void TearDown() override { allocator.shutdown(); } @@ -351,7 +356,7 @@ TEST_F(OrchestratorFixture, RemoteInputSidecarUsesRemoteTensorMapKey) { output_sidecar.tensors[0].desc.nbytes = 1; auto producer = orch.submit_next_level(C(42), output_args, cfg, 3, {3}, output_sidecar); TaskSlot ready; - rq.try_pop(ready); + rq_next_level_single.try_pop(3, ready); TaskArgs input_args; Tensor in = out; diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index 538a799ff1..4f9356c5de 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -269,6 +269,7 @@ struct SchedulerFixture : public ::testing::Test { // Strict-4: per-type ready queues. ReadyQueue rq_next_level; ReadyQueue rq_sub; + PerWorkerReadyQueues rq_next_level_single; Orchestrator orch; MockMailboxWorker mock_worker; WorkerManager manager; @@ -282,21 +283,27 @@ 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_single.reset(manager.next_level_worker_ids()); + orch.init( + &tm, &allocator, &scope, &rq_next_level, &rq_sub, &rq_next_level_single, &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_single_queues = &rq_next_level_single; 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); @@ -475,6 +482,7 @@ struct GroupSchedulerFixture : public ::testing::Test { // Strict-4: per-type ready queues. ReadyQueue rq_next_level; ReadyQueue rq_sub; + PerWorkerReadyQueues rq_next_level_single; Orchestrator orch; MockMailboxWorker worker_a; MockMailboxWorker worker_b; @@ -489,9 +497,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 +505,21 @@ struct GroupSchedulerFixture : public ::testing::Test { manager.start(&allocator, [this](WorkerCompletion completion) { sched.worker_done(std::move(completion)); }); + rq_next_level_single.reset(manager.next_level_worker_ids()); + orch.init( + &tm, &allocator, &scope, &rq_next_level, &rq_sub, &rq_next_level_single, &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_single_queues = &rq_next_level_single; 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); @@ -637,6 +651,68 @@ TEST_F(GroupSchedulerFixture, ExplicitTargetWithinEligibilityIsUsed) { 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); @@ -653,6 +729,7 @@ TEST(SchedulerWorkerAffinityTest, NextLevelAffinityUsesWorkerIdNotVectorIndex) { Scope scope; ReadyQueue rq_next_level; ReadyQueue rq_sub; + PerWorkerReadyQueues rq_next_level_single; Orchestrator orch; MockMailboxWorker worker_a; MockMailboxWorker worker_b; @@ -663,10 +740,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 +747,21 @@ TEST(SchedulerWorkerAffinityTest, NextLevelAffinityUsesWorkerIdNotVectorIndex) { manager.start(&allocator, [&sched](WorkerCompletion completion) { sched.worker_done(std::move(completion)); }); + rq_next_level_single.reset(manager.next_level_worker_ids()); + orch.init( + &tm, &allocator, &scope, &rq_next_level, &rq_sub, &rq_next_level_single, &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_single_queues = &rq_next_level_single; 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); @@ -769,6 +851,7 @@ struct MixedTypeSchedulerFixture : public ::testing::Test { Scope scope; ReadyQueue rq_next_level; ReadyQueue rq_sub; + PerWorkerReadyQueues rq_next_level_single; Orchestrator orch; MockMailboxWorker next_level_worker; MockMailboxWorker sub_worker; @@ -783,9 +866,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 +874,21 @@ struct MixedTypeSchedulerFixture : public ::testing::Test { manager.start(&allocator, [this](WorkerCompletion completion) { sched.worker_done(std::move(completion)); }); + rq_next_level_single.reset(manager.next_level_worker_ids()); + orch.init( + &tm, &allocator, &scope, &rq_next_level, &rq_sub, &rq_next_level_single, &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_single_queues = &rq_next_level_single; 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);