def setUp(self) -> None:
self._backend = FakeRendezvousBackend()
mock_get_state = MagicMock(wraps=self._backend.get_state)
mock_set_state = MagicMock(wraps=self._backend.set_state)
self._mock_backend = Mock()
self._mock_backend.get_state = mock_get_state
self._mock_backend.set_state = mock_set_state
setattr(self._backend, "get_state", mock_get_state) # noqa: B010
setattr(self._backend, "set_state", mock_set_state) # noqa: B010
self._settings = RendezvousSettings(
run_id="dummy_run_id",
min_nodes=1,
max_nodes=1,
timeout=RendezvousTimeout(),
keep_alive_interval=timedelta(seconds=30),
keep_alive_max_attempt=3,
)
self._cache_duration = 0
self._now = datetime(2000, 1, 1, hour=0, minute=0)
self._datetime_patch = patch(
"torch.distributed.elastic.rendezvous.dynamic_rendezvous.datetime")
mock_datetime = self._datetime_patch.start()
mock_datetime.utcnow.return_value = self._now
def _create_settings(self) -> RendezvousSettings:
return RendezvousSettings(
run_id="dummy_run_id",
min_nodes=self._min_nodes,
max_nodes=self._max_nodes,
timeout=self._timeout,
keep_alive_interval=timedelta(seconds=30),
keep_alive_max_attempt=3,
)
def _get_next_action(self) -> _Action:
op = self._create_op()
settings = RendezvousSettings(
run_id="dummy_run_id",
min_nodes=self._min_nodes,
max_nodes=self._max_nodes,
timeout=RendezvousTimeout(),
keep_alive_interval=self._keep_alive_interval,
keep_alive_max_attempt=3,
)
ctx = _RendezvousContext(self._node, self._state, settings)
return op(ctx, self._deadline)
def _create_handler(self) -> DynamicRendezvousHandler:
settings = RendezvousSettings(
run_id="dummy_run_id",
min_nodes=self._min_nodes,
max_nodes=self._max_nodes,
timeout=RendezvousTimeout(
join=self._join_timeout,
close=self._close_timeout,
heartbeat=self._heartbeat_timeout,
),
keep_alive_interval=self._keep_alive_interval,
keep_alive_max_attempt=3,
)
self._state_holder.state = self._state
return DynamicRendezvousHandler(self._node, settings, "dummy_backend",
self._store, self._state_holder)