Question: Find bad GPU nodes with group tests
Companies: xAI
Difficulty: Medium
Problem statement
A cluster has n GPU nodes with ids 0 to n - 1. Each node is either good or bad, and you do not know which in advance. The only way to learn anything is a test service:
class TestServer:
def test(self, nodes: list[int]) -> bool:
# True if every node in `nodes` is good; False if at least one is bad.
...
A failing test does not say which node is bad or how many are bad. Each call runs a real health job on the hardware, so calls are expensive.
Implement find_bad_nodes(n: int, server: TestServer) -> list[int]. It returns the ids of all bad nodes in ascending order and uses as few test calls as you can.
Constraints
0 <= n <= 10^5- Any number of nodes may be bad, including none or all of them. In practice bad nodes are rare.
- Testing every node on its own (
ncalls) is correct but too slow whennis large. Aim for a call count that depends mainly on the number of bad nodesk.
Example 1
n = 8, bad nodes = {2, 5}
output: [2, 5]
Testing [0..7] fails, so split. [0..3] fails and [4..7] fails, so keep splitting each half until single nodes remain. Halves that pass, such as [0, 1] and [6, 7], are cleared with one call each.
Example 2
n = 6, bad nodes = {}
output: []
One call on [0..5] passes, so every node is good.
Be ready for follow-ups: there is exactly one bad node, tests run in parallel rounds, you know an upper bound on the number of bad nodes, or tests are sometimes wrong.
Approach
This is group testing. One call to test(nodes) tells you only whether a whole group is clean. It does not say which node is bad or how many. The goal is to find every bad node with as few calls as possible.
Brute force. Test each node alone. That always costs n calls. It is correct and a useful oracle for tests, but it wastes calls when bad nodes are rare. A healthy fleet with 10,000 nodes and 3 bad ones should not need 10,000 calls.
Adaptive binary splitting. Test the whole range. If it passes, every node in it is good and you are done with it. If it fails, split it in half and handle each half the same way. A failing block of size 1 is a bad node. Each bad node costs about one test per level on its path down the recursion tree. So the total is about k * log2(n / k) plus O(k) calls for k bad nodes, instead of n.
Skip the test you can infer. Suppose a block is known to contain a bad node and its left half passes. Then the right half must contain the bad node, so you do not test it. You go straight to splitting it. Pass a known_bad flag down the recursion to track this. This saves up to one call per level. With one bad node, the cost drops close to a plain binary search.
Invariant. search(lo, hi, known_bad=True) is only called when [lo, hi) contains at least one bad node. That is why a single node reached with known_bad=True can be reported without a test. Every node not reported sits inside some block that passed a test.
Why the answer is correct. A node is reported only when a test proves it is bad: either its size-1 block failed, or it is the only node left in a block known to be bad. A node is skipped only when a block containing it passed. So the output is exactly the set of bad nodes. Ranges are visited left to right, so the output is already sorted.
If the interviewer says there is exactly one bad node, use plain binary search: ceil(log2 n) calls after the first check, which is optimal.
Code
from typing import Callable, List, Optional, Sequence, Set
class TestServer:
"""Mock oracle. test(nodes) is True when every listed node is good."""
def __init__(self, n: int, bad: Set[int]) -> None:
self.n = n
self._bad = set(bad)
self.calls = 0
def test(self, nodes: Sequence[int]) -> bool:
self.calls += 1
return not any(x in self._bad for x in nodes)
def find_bad_nodes_naive(n: int, server: TestServer) -> List[int]:
"""One test per node: always exactly n calls."""
return [i for i in range(n) if not server.test([i])]
def find_bad_nodes(n: int, server: TestServer) -> List[int]:
"""Adaptive binary splitting. Returns bad node ids in ascending order."""
bad: List[int] = []
def search(lo: int, hi: int, known_bad: bool) -> None:
# Invariant: if known_bad, [lo, hi) contains at least one bad node.
if lo >= hi:
return
if not known_bad and server.test(list(range(lo, hi))):
return # the whole block is good
if hi - lo == 1:
bad.append(lo) # a failing single node is bad
return
mid = (lo + hi) // 2
left_good = server.test(list(range(lo, mid)))
if not left_good:
search(lo, mid, known_bad=True)
# If the left half passed, the bad node must be on the right,
# so the right half needs no test of its own.
search(mid, hi, known_bad=left_good)
search(0, n, known_bad=False)
return bad
def find_one_bad_node(n: int, server: TestServer) -> Optional[int]:
"""Follow-up: exactly one bad node (or none). ceil(log2 n) + 1 calls."""
if n == 0 or server.test(list(range(n))):
return None
lo, hi = 0, n # invariant: bad node is in [lo, hi)
while hi - lo > 1:
mid = (lo + hi) // 2
if server.test(list(range(lo, mid))):
lo = mid
else:
hi = mid
return lo
def find_bad_nodes_by_rounds(
n: int, batch_test: Callable[[List[List[int]]], List[bool]]
) -> List[int]:
"""Follow-up: tests run in parallel rounds. Each round tests every
open block at once, so the number of rounds is ceil(log2 n) + 1."""
bad: List[int] = []
frontier = [list(range(n))] if n else []
while frontier:
results = batch_test(frontier)
nxt: List[List[int]] = []
for block, ok in zip(frontier, results):
if ok:
continue
if len(block) == 1:
bad.append(block[0])
else:
mid = len(block) // 2
nxt += [block[:mid], block[mid:]]
frontier = nxt
return sorted(bad)
Complexity
Let n be the number of nodes and k the number of bad nodes.
-
find_bad_nodes_naive: exactlyncalls. -
find_bad_nodes: 1 call whenk = 0. At most about2k * log2(n / k) + O(k)calls in general, and never more than about2n. Fork = 1it uses betweenlog2 n + 1and2 log2 n + 1calls. This is within a factor of 2 of the information-theoretic lower bound,log2 C(n, k), which is aboutk * log2(n / k). Local time is O(n log n) to build the index lists, since each level of the recursion lists every node in a failing block once. Space is O(log n) for the recursion plus O(n) for one test list. -
find_one_bad_node:ceil(log2 n) + 1calls. O(n) time to build lists, O(n) space. -
find_bad_nodes_by_rounds:ceil(log2 n) + 1rounds. Each round tests at most2kblocks, so total calls are O(k log n).
Test cases
-
n = 8, bad ={2, 5}->[2, 5], with 9 calls (naive uses 8; splitting pays off at largern). -
n = 8, bad ={}->[], with exactly 1 call. -
n = 0->[], with 0 calls. -
n = 1, bad ={0}->[0], with 1 call. -
n = 5, all nodes bad ->[0, 1, 2, 3, 4]. This is the worst case for splitting; it uses 9 calls, versus 5 for naive. -
n = 1024, bad ={700}->[700]with 15 calls.find_one_bad_nodefinds it with 11. -
n = 16, bad ={3}->[3]with 7 calls. - Randomized: for
nup to 60 and bad density 0, 5%, 30% and 100%, all three methods agree with the naive answer.
Follow-ups
-
Exactly one bad node. Use binary search over a range known to be bad. That takes
ceil(log2 n)calls, which is optimal, because each call returns one bit and there arenpossible answers. -
Tests run in parallel, and wall-clock time matters. Each test is a slow burn-in job, so count rounds instead of calls. Test all open blocks in one round, as in
find_bad_nodes_by_rounds. That takesceil(log2 n) + 1rounds. If only one round is allowed, you need a non-adaptive design. For at most one bad node, give each node a binary code and test bitj= "all nodes whose bitjis 1". The failing bits spell the bad index. -
You know
k, or an upper bound on it. Start with blocks of size aboutn / kinstead of the full range (Hwang's generalized binary splitting). This avoids testing large blocks that are almost sure to fail, and it reaches aboutk * log2(n / k) + kcalls. Whenkis close ton, more than aboutn / 3, testing nodes one by one is as good as anything. - A test's cost grows with group size, or a group has a maximum size. Cap the top-level block size at the limit and run the same splitting inside each block. With a cost per node, weigh fewer large tests against more small ones, and pick the block size that minimizes expected cost for the observed failure rate.
- Tests are noisy or nodes fail intermittently. A single failing test is no longer proof. Repeat a test and take the majority, or re-confirm each suspect alone before reporting it. Keep good-looking blocks in a re-test queue rather than trusting one pass. Report suspects with a confidence instead of a hard verdict.
Interviewer expectations
Mid-level (L4)
- Gives the one-test-per-node baseline, then the recursive halving idea without prompting.
- Writes correct recursion with clean base cases (empty range, single node).
- Can count calls for zero bad nodes and for one bad node.
Senior (L5)
- Adds the "left half passed, so the right half is bad" inference and explains the invariant behind it.
- States the bound in terms of
k, notes when splitting beats or loses to the naive method, and handles the exactly-one-bad case with binary search. - Writes a mock
TestServerthat counts calls and tests against the naive answer.
Staff (L6+)
- Frames it as group testing and compares to the
log2 C(n, k)lower bound. - Moves to rounds of parallel tests, non-adaptive designs, and generalized binary splitting when
kis known. - Discusses noisy or flaky tests, maximum group sizes, and how this fits a real cluster burn-in pipeline (draining nodes, retries, topology-aware grouping).