[Core]Extract is_last_rank in Ray for tpu to override (#33012)

Signed-off-by: Chenyaaang <chenyangli@google.com>
This commit is contained in:
Chenyaaang
2026-02-26 19:18:52 -08:00
committed by GitHub
parent c29ee9c326
commit 06be53563b

View File

@@ -108,7 +108,7 @@ try:
if isinstance(output, AsyncModelRunnerOutput): if isinstance(output, AsyncModelRunnerOutput):
output = output.get_output() output = output.get_output()
if not get_pp_group().is_last_rank: if not self._is_last_rank():
# Case where there are no scheduled requests # Case where there are no scheduled requests
# but may still be finished requests. # but may still be finished requests.
assert not output or not output.req_ids assert not output or not output.req_ids
@@ -128,6 +128,9 @@ try:
def _is_intermediate_tensors(self, output) -> bool: def _is_intermediate_tensors(self, output) -> bool:
return isinstance(output, IntermediateTensors) return isinstance(output, IntermediateTensors)
def _is_last_rank(self) -> bool:
return get_pp_group().is_last_rank
ray_import_err = None ray_import_err = None
except ImportError as e: except ImportError as e: