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: ...@@ -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: ...@@ -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:
......
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