Skip to content

[Bugfix] Misc bug fixes#208

Merged
majunze2001 merged 3 commits into
masterfrom
jm-fix-ref-count
May 29, 2026
Merged

[Bugfix] Misc bug fixes#208
majunze2001 merged 3 commits into
masterfrom
jm-fix-ref-count

Conversation

@majunze2001

Copy link
Copy Markdown
Collaborator
  • Fix task refcount for sharing
  • Fix typos and unify example apps
  • Fix GPU allocation for TP task executors

@majunze2001
majunze2001 merged commit d1e3732 into master May 29, 2026
2 checks passed
@majunze2001
majunze2001 deleted the jm-fix-ref-count branch May 29, 2026 05:50

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request updates example models, re-enables actual GPU allocation in the resource manager, and implements a mechanism to allocate GPUs with consecutive local ranks to support tensor parallelism. It also refactors task registration to ensure that the cluster-wide refcount is bumped for all shared tasks. The review feedback highlights two critical issues: first, skipping the refcount bump when an instance name is not yet committed can lead to premature teardown of shared tasks, which should be resolved by polling and waiting; second, raising an allocation error immediately on the first node that lacks consecutive free GPUs causes premature failure, and the allocator should instead continue searching other nodes in the priority queue.

Comment on lines +794 to +804
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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

If the instance name is not yet committed because a concurrent deployment is still in Phase 2, skipping the refcount bump entirely will cause the cluster-wide refcount to be undercounted. This can lead to premature teardown of the shared task when one of the active users unregisters it. Instead of skipping, we should poll and wait for the concurrent deployment to commit the instance name.

Suggested change
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
instance_name = None
for _ in range(100): # Wait up to 10 seconds (100 * 0.1s)
instance_name = self.unit_task_instance_names.get(task_id)
if instance_name is not None:
break
await asyncio.sleep(0.1)
if instance_name is None:
logger.warning(
"Skipping CR refcount bump for task %s: instance name not yet committed after timeout",
task_id,
)
continue

Comment on lines +234 to +238
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.",
)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Raising CannotColocateError immediately when the first node in node_priority does not have consecutive free GPUs will cause the allocation to fail, even if there are other nodes in the priority queue that have consecutive free GPUs. Instead of failing immediately, we should continue to the next node in node_priority and only raise the error if we have exhausted all available nodes.

Suggested change
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.",
)
if len(run) != num_gpus:
if not node_priority:
raise CannotColocateError(
f"Cannot allocate {num_gpus} GPUs with consecutive local ranks on a "
f"single node. Free GPUs on {node} are too fragmented.",
)
continue

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant