0

Find bad GPU nodes with group tests

by · 1 day ago· ⑂ 0 remixes

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 (n calls) is correct but too slow when n is large. Aim for a call count that depends mainly on the number of bad nodes k.

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: exactly n calls.
  • find_bad_nodes: 1 call when k = 0. At most about 2k * log2(n / k) + O(k) calls in general, and never more than about 2n. For k = 1 it uses between log2 n + 1 and 2 log2 n + 1 calls. This is within a factor of 2 of the information-theoretic lower bound, log2 C(n, k), which is about k * 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) + 1 calls. O(n) time to build lists, O(n) space.
  • find_bad_nodes_by_rounds: ceil(log2 n) + 1 rounds. Each round tests at most 2k blocks, 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 larger n).
  • 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_node finds it with 11.
  • n = 16, bad = {3} -> [3] with 7 calls.
  • Randomized: for n up to 60 and bad density 0, 5%, 30% and 100%, all three methods agree with the naive answer.

Follow-ups

  1. 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 are n possible answers.
  2. 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 takes ceil(log2 n) + 1 rounds. 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 bit j = "all nodes whose bit j is 1". The failing bits spell the bad index.
  3. You know k, or an upper bound on it. Start with blocks of size about n / k instead of the full range (Hwang's generalized binary splitting). This avoids testing large blocks that are almost sure to fail, and it reaches about k * log2(n / k) + k calls. When k is close to n, more than about n / 3, testing nodes one by one is as good as anything.
  4. 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.
  5. 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 TestServer that 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 k is 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).

Sign in to comment, vote or remix.