| 44 | |
| 45 | |
| 46 | def assign_workers(workers: List[WorkerConfig], nodes: List[str]) -> Dict[str, List[WorkerConfig]]: |
| 47 | round_robin_node = 0 |
| 48 | used_round_robin = set() |
| 49 | |
| 50 | node_assignments = defaultdict(list) |
| 51 | for index, worker in enumerate(workers): |
| 52 | node = worker.node |
| 53 | if node is not None: |
| 54 | if not (0 <= node < len(nodes)): |
| 55 | raise Exception( |
| 56 | f"Invalid node assignment. Worker {index} wants to be on node " |
| 57 | f"{node}, but there are only {len(nodes)} worker nodes" |
| 58 | ) |
| 59 | else: |
| 60 | node = round_robin_node |
| 61 | round_robin_node = (round_robin_node + 1) % len(nodes) |
| 62 | if node in used_round_robin: |
| 63 | raise Exception(f"There are more workers ({len(workers)}) than worker nodes ({len(nodes)})") |
| 64 | used_round_robin.add(node) |
| 65 | if node >= len(nodes): |
| 66 | raise Exception(f"Selected worker node is {node}, but there are only {len(nodes)} worker node(s)") |
| 67 | node_assignments[nodes[node]].append(worker) |
| 68 | return dict(node_assignments) |