def _mock_train(*args, _training_state, **kwargs):
assert ray.get(_training_state.queue.actor.get_node_id.remote(
)) == ray.state.current_node_id()
assert ray.get(
_training_state.stop_event.actor.get_node_id.remote()) == \
ray.state.current_node_id()
return _train(*args, _training_state=_training_state, **kwargs)
def _mock_train(*args, _training_state, **kwargs):
try:
results = _train(
*args, _training_state=_training_state, **kwargs)
return results
except Exception:
raise
finally:
assert len(_training_state.actors) == num_actors
if not any(a is None for a in _training_state.actors):
actor_infos = ray.state.actors()
actor_nodes = []
for a in _training_state.actors:
actor_info = actor_infos.get(a._actor_id.hex())
actor_node = actor_info["Address"]["NodeID"]
actor_nodes.append(actor_node)
assert actor_nodes[0] == actor_nodes[1]