1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
|
# 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, cast, Optional, Tuple
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."""
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)
|
