Unverified Commit 06be5356 authored by Chenyaaang's avatar Chenyaaang Committed by GitHub
Browse files

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


Signed-off-by: default avatarChenyaaang <chenyangli@google.com>
parent c29ee9c3
......@@ -108,7 +108,7 @@ try:
if isinstance(output, AsyncModelRunnerOutput):
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
# but may still be finished requests.
assert not output or not output.req_ids
......@@ -128,6 +128,9 @@ try:
def _is_intermediate_tensors(self, output) -> bool:
return isinstance(output, IntermediateTensors)
def _is_last_rank(self) -> bool:
return get_pp_group().is_last_rank
ray_import_err = None
except ImportError as e:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment