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 - < 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, 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],