From d235fba59fab6860f535147e2aa17e16691e4560 Mon Sep 17 00:00:00 2001 From: Jeff Ma Date: Fri, 29 May 2026 05:44:49 +0000 Subject: [PATCH 1/3] fix task ref count --- .../services/gateway/task_manager.py | 32 +++++++++++++++---- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/python/cornserve/services/gateway/task_manager.py b/python/cornserve/services/gateway/task_manager.py index 97d17a66..a5d5ef15 100644 --- a/python/cornserve/services/gateway/task_manager.py +++ b/python/cornserve/services/gateway/task_manager.py @@ -697,10 +697,15 @@ async def declare_used(self, tasks: list[UnitTask]) -> None: if existing_task.is_equivalent_to(task): logger.info("Skipping already deployed task: %r", task) task_ids.append(task_id) - # If this replica hasn't claimed this task yet (watch-synced), - # we need to bump the CR refcount in Phase 2. - if self.task_usage_counter[task_id] == 0: - to_bump_refcount.append(task_id) + # Every declare_used reference to an already-deployed task + # must bump the CR refcount, mirroring declare_not_used which + # always releases one. This covers both watch-synced tasks + # (local counter 0) and tasks already used by another app on + # this replica (counter > 0). Without bumping in the latter + # case, a shared task's CR refcount undercounts its real users + # and the first unregister tears it down while other apps are + # still using it. + to_bump_refcount.append(task_id) break else: # If the task is not already deployed, deploy it @@ -781,16 +786,29 @@ async def declare_used(self, tasks: list[UnitTask]) -> None: if errors: raise RuntimeError("Error while deploying tasks") - # Bump CR refcount for watch-synced tasks being used locally for the first time + # Bump the CR refcount for every already-deployed task reference so the + # cluster-wide refcount tracks the real number of users (each + # declare_not_used always releases one). Covers watch-synced tasks and + # additional local apps sharing an already-deployed task. for task_id in to_bump_refcount: + instance_name = self.unit_task_instance_names.get(task_id) + if instance_name is None: + # The matched task is still being deployed by a concurrent + # declare_used and has not committed its instance name yet + # (Phase 3). Skip the bump rather than crash — this is the same + # rare window the rollback path below also tolerates with .get(). + logger.warning( + "Skipping CR refcount bump for task %s: instance name not yet committed", + task_id, + ) + continue task = self.tasks[task_id] - instance_name = self.unit_task_instance_names[task_id] lock_key = self._compute_canonical_task_key(task) async with self._refcount_lock(lock_key): current = await self._get_usage_refcount(instance_name) await self._patch_usage_refcount(instance_name, current + 1) logger.info( - "Bumped CR refcount for watch-synced task %s (%s): %d -> %d", + "Bumped CR refcount for shared task %s (%s): %d -> %d", task_id, instance_name, current, From 0098f5932c73014c6d12d092713db8e39431acbd Mon Sep 17 00:00:00 2001 From: Jeff Ma Date: Fri, 29 May 2026 05:47:40 +0000 Subject: [PATCH 2/3] consecutive gpu allocation for tp --- python/cornserve/services/resource.py | 34 +++++++++++++++++-- .../services/resource_manager/manager.py | 4 +-- 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/python/cornserve/services/resource.py b/python/cornserve/services/resource.py index 793001f2..eeb1d3d1 100644 --- a/python/cornserve/services/resource.py +++ b/python/cornserve/services/resource.py @@ -127,6 +127,7 @@ def allocate( owner: str, must_colocate: bool = True, node_selection_policy: Literal["pack", "spread"] = "spread", + require_consecutive: bool = True, ) -> list[GPU]: """Allocate a number of GPUs to a owner. @@ -149,6 +150,14 @@ def allocate( owner: Owner name must_colocate: Whether GPU colocation is required node_selection_policy: Node selection policy ("pack" or "spread") + require_consecutive: If True (the default; only takes effect when + `must_colocate` and `num_gpus > 1`), the allocated GPUs must have + consecutive local ranks on the node. A task executor with tensor + parallelism shares a single contiguous slice of the per-node + sidecar shared-memory slab keyed by local rank (see + `sidecar.utils.init_shmem`), so its GPUs must be consecutive -- + not merely colocated. Set to False only for allocations that do + not back a single TP group. Raises: CannotColocateError: If `must_colocate` is True and GPUs cannot @@ -207,8 +216,29 @@ def allocate( allocated_gpus: list[GPU] = [] while num_allocated < num_gpus and node_priority: _, _, node = heapq.heappop(node_priority) - gpus = self.node_to_gpus[node] - gpus = [gpu for gpu in gpus if gpu.is_free][: num_gpus - num_allocated] + # `node_to_gpus[node]` is sorted by local rank. + free_gpus = [gpu for gpu in self.node_to_gpus[node] if gpu.is_free] + if require_consecutive and num_gpus > 1 and must_colocate: + # Scan for the first run of `num_gpus` consecutive local ranks. + # Unlike taking the lowest-N free GPUs, this avoids straddling + # GPUs owned by other task managers (e.g. picking local ranks + # 2 and 6 when 3, 4, 5 are busy), which would violate the + # consecutive-local-rank invariant of the sidecar shared memory. + run: list[GPU] = [] + for gpu in free_gpus: + if run and gpu.local_rank != run[-1].local_rank + 1: + run = [] + run.append(gpu) + if len(run) == num_gpus: + break + if len(run) != num_gpus: + raise CannotColocateError( + f"Cannot allocate {num_gpus} GPUs with consecutive local ranks on a " + f"single node. Free GPUs on {node} are too fragmented.", + ) + gpus = run + else: + gpus = free_gpus[: num_gpus - num_allocated] for gpu in gpus: allocated_gpus.append(gpu.allocate_to(owner)) num_allocated += 1 diff --git a/python/cornserve/services/resource_manager/manager.py b/python/cornserve/services/resource_manager/manager.py index ecba2aa3..c0b1a9b6 100644 --- a/python/cornserve/services/resource_manager/manager.py +++ b/python/cornserve/services/resource_manager/manager.py @@ -1008,9 +1008,9 @@ async def _spawn_task_manager( logger.info("Allocating %d GPUs for task %s based on profile", num_gpus, task) # Allocate resource starter pack for the task manager - # state.resources = self.resource.allocate(num_gpus=num_gpus, owner=state.id) + state.resources = self.resource.allocate(num_gpus=num_gpus, owner=state.id) # uncomment below during benchmarking to speedup - state.resources = [] + # state.resources = [] span.set_attribute( "resource_manager._spawn_task_manager.gpu_global_ranks", [gpu.global_rank for gpu in state.resources], From 7952fdd0e23dd676814f48c7e0311a0b39b0cfdc Mon Sep 17 00:00:00 2001 From: Jeff Ma Date: Fri, 29 May 2026 05:48:16 +0000 Subject: [PATCH 3/3] fix example apps --- examples/macro_unittask.py | 1 + examples/mllm.py | 2 +- examples/router.py | 32 +++++++++++++------------------- 3 files changed, 15 insertions(+), 20 deletions(-) diff --git a/examples/macro_unittask.py b/examples/macro_unittask.py index 14571e14..efb026a2 100644 --- a/examples/macro_unittask.py +++ b/examples/macro_unittask.py @@ -21,6 +21,7 @@ llm_tp_size=1, llm_max_num_seqs=32, llm_gpu_memory_utilization=0.9, + macro_ut_deployment_id="qwen3_omni_macro_unittask", ) diff --git a/examples/mllm.py b/examples/mllm.py index 711e89d9..50c8448c 100644 --- a/examples/mllm.py +++ b/examples/mllm.py @@ -4,7 +4,7 @@ $ cornserve register examples/mllm.py $ cornserve invoke mllm --aggregate-keys choices.0.delta.content --data - <