mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
eb8b80b76f
Summary: Action following discussion with distributed and r2p team--the tests under elastic in distributed should be owned by oncall: r2p and not distributed. cc pietern mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse SciPioneer H-Huang Pull Request resolved: https://github.com/pytorch/pytorch/pull/67293 Reviewed By: jbschlosser Differential Revision: D31973779 Pulled By: janeyx99 fbshipit-source-id: 05875a7600c6eb1da1310a48e1e32a1a69461c55
106 lines
3.3 KiB
Python
106 lines
3.3 KiB
Python
# Owner(s): ["oncall: r2p"]
|
|
|
|
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the BSD-style license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import Any, Callable, Optional, Tuple, cast
|
|
|
|
from torch.distributed.elastic.rendezvous import RendezvousStateError
|
|
from torch.distributed.elastic.rendezvous.dynamic_rendezvous import RendezvousBackend, Token
|
|
|
|
|
|
class RendezvousBackendTestMixin(ABC):
|
|
_backend: RendezvousBackend
|
|
|
|
# Type hints
|
|
assertEqual: Callable
|
|
assertNotEqual: Callable
|
|
assertIsNone: Callable
|
|
assertIsNotNone: Callable
|
|
assertRaises: Callable
|
|
|
|
@abstractmethod
|
|
def _corrupt_state(self) -> None:
|
|
"""Corrupts the state stored in the backend."""
|
|
pass
|
|
|
|
def _set_state(self, state: bytes, token: Optional[Any] = None) -> Tuple[bytes, Token, bool]:
|
|
result = self._backend.set_state(state, token)
|
|
|
|
self.assertIsNotNone(result)
|
|
|
|
return cast(Tuple[bytes, Token, bool], result)
|
|
|
|
def test_get_state_returns_backend_state(self) -> None:
|
|
self._backend.set_state(b"x")
|
|
|
|
result = self._backend.get_state()
|
|
|
|
self.assertIsNotNone(result)
|
|
|
|
state, token = cast(Tuple[bytes, Token], result)
|
|
|
|
self.assertEqual(b"x", state)
|
|
self.assertIsNotNone(token)
|
|
|
|
def test_get_state_returns_none_if_backend_state_does_not_exist(self) -> None:
|
|
result = self._backend.get_state()
|
|
|
|
self.assertIsNone(result)
|
|
|
|
def test_get_state_raises_error_if_backend_state_is_corrupt(self) -> None:
|
|
self._corrupt_state()
|
|
|
|
with self.assertRaises(RendezvousStateError):
|
|
self._backend.get_state()
|
|
|
|
def test_set_state_sets_backend_state_if_it_does_not_exist(self) -> None:
|
|
state, token, has_set = self._set_state(b"x")
|
|
|
|
self.assertEqual(b"x", state)
|
|
self.assertIsNotNone(token)
|
|
self.assertTrue(has_set)
|
|
|
|
def test_set_state_sets_backend_state_if_token_is_current(self) -> None:
|
|
state1, token1, has_set1 = self._set_state(b"x")
|
|
|
|
state2, token2, has_set2 = self._set_state(b"y", token1)
|
|
|
|
self.assertEqual(b"y", state2)
|
|
self.assertNotEqual(token1, token2)
|
|
self.assertTrue(has_set1)
|
|
self.assertTrue(has_set2)
|
|
|
|
def test_set_state_returns_current_backend_state_if_token_is_old(self) -> None:
|
|
state1, token1, _ = self._set_state(b"x")
|
|
|
|
state2, token2, _ = self._set_state(b"y", token1)
|
|
|
|
state3, token3, has_set = self._set_state(b"z", token1)
|
|
|
|
self.assertEqual(state2, state3)
|
|
self.assertEqual(token2, token3)
|
|
self.assertFalse(has_set)
|
|
|
|
def test_set_state_returns_current_backend_state_if_token_is_none(self) -> None:
|
|
state1, token1, _ = self._set_state(b"x")
|
|
|
|
state2, token2, has_set = self._set_state(b"y")
|
|
|
|
self.assertEqual(state1, state2)
|
|
self.assertEqual(token1, token2)
|
|
self.assertFalse(has_set)
|
|
|
|
def test_set_state_returns_current_backend_state_if_token_is_invalid(self) -> None:
|
|
state1, token1, _ = self._set_state(b"x")
|
|
|
|
state2, token2, has_set = self._set_state(b"y", token="invalid")
|
|
|
|
self.assertEqual(state1, state2)
|
|
self.assertEqual(token1, token2)
|
|
self.assertFalse(has_set)
|