[Core]Extract is_last_rank in Ray for tpu to override (#33012)
Signed-off-by: Chenyaaang <chenyangli@google.com>
This commit is contained in:
@@ -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:
|
||||||
|
|||||||
Reference in New Issue
Block a user