diff --git a/serve/replicas.py b/serve/replicas.py index 5a139efa8a..6f8406b3d7 100644 --- a/serve/replicas.py +++ b/serve/replicas.py @@ -56,11 +56,12 @@ def gpu_groups(cfg: dict) -> list[list[int]]: raise ReplicaConfigError(f'"replicas": {spec} does not divide the {len(cards)} cards {cards} evenly: name the ' f'groups ("replicas": [{{"gpus": [0, 1]}}, ...]) or change the count') n = len(cards) // spec - return [cards[i * n:(i + 1) * n] for i in range(spec)] - groups = [] - for e in spec: - g = e.get("gpus") if isinstance(e, dict) else e - groups.append(_cards(g)) + groups = [cards[i * n:(i + 1) * n] for i in range(spec)] + else: + groups = [] + for e in spec: + g = e.get("gpus") if isinstance(e, dict) else e + groups.append(_cards(g)) if any(not g for g in groups): raise ReplicaConfigError('every replica needs its cards: {"gpus": [0, 1]}') flat = [c for g in groups for c in g] diff --git a/serve/test_replicas.py b/serve/test_replicas.py index 95f1a12e65..00c0c0dfe3 100644 --- a/serve/test_replicas.py +++ b/serve/test_replicas.py @@ -27,6 +27,25 @@ def test_named_groups(self): self.assertEqual(gpu_groups({"replicas": [{"gpus": [0, 1]}, {"gpus": [2]}]}), [[0, 1], [2]]) self.assertEqual(gpu_groups({"replicas": [{"gpus": [0, 1]}]}), []) + def test_automatic_groups_reject_duplicate_cards(self): + for cards in ([0, 0], [0, 1, 1, 2], [0, 0, 1, 2], "0,0", "0,1,1,2", "0,00"): + with self.subTest(cards=cards): + with self.assertRaisesRegex(ReplicaConfigError, "a card is in two replicas"): + gpu_groups({"gpu": cards, "replicas": 2}) + + def test_named_groups_still_reject_duplicate_cards(self): + for groups in ([[0], [0]], [[0, 1], [1, 2]], [[0, 0], [1, 2]], [[0, 0]]): + with self.subTest(groups=groups): + with self.assertRaisesRegex(ReplicaConfigError, "a card is in two replicas"): + gpu_groups({"replicas": [{"gpus": g} for g in groups]}) + + def test_valid_card_order_and_disabled_replicas_are_unchanged(self): + for cards in ([3, 1, 2, 0], "3,1,2,0"): + self.assertEqual(gpu_groups({"gpu": cards, "replicas": 2}), [[3, 1], [2, 0]]) + for count in (None, False, 0, 1): + with self.subTest(count=count): + self.assertEqual(gpu_groups({"gpu": [0, 0], "replicas": count}), []) + def test_mistakes_are_said(self): for cfg in ({"gpu": [0, 1, 2], "replicas": 2}, {"replicas": 2}, {"replicas": "2"}, {"replicas": [{"gpus": [0, 1]}, {"gpus": [1, 2]}]}, {"replicas": [{"gpus": []}, {"gpus": [1]}]}):