fanout_solution.py

"""INTERVIEWER ONLY — reference implementation of fan_out().

A solid ~35-minute solution using the sync client + ThreadPoolExecutor.
Drop this fan_out into starter/fanout.py to sanity-check the exercise.
Not the only valid shape: asyncio + AsyncRunloop is equally good.
"""

from __future__ import annotations

import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List

# When copied into the starter these imports already exist there.
from fanout import BoxResult, api_errors  # type: ignore


def fan_out(client, command: str, count: int, timeout_s: float, mock: bool) -> List[BoxResult]:
    errors = api_errors(mock)
    provisioned: List[str] = []          # every devbox id we ever created
    results: List[BoxResult] = []
    deadline = time.monotonic() + timeout_s

    def worker(i: int) -> BoxResult:
        start = time.monotonic()
        box = None
        try:
            box = client.devboxes.create_and_await_running(name=f"rlfan-{i}")
            provisioned.append(box.id)  # list.append is atomic under the GIL
            result = client.devboxes.execute_sync(box.id, command=command)
            return BoxResult(
                devbox_id=box.id,
                exit_status=result.exit_status,
                duration_s=time.monotonic() - start,
                stdout_head=(result.stdout or "").splitlines()[0] if result.stdout else "",
                stderr_head=(result.stderr or "").splitlines()[0] if result.stderr else "",
            )
        except errors as e:
            return BoxResult(
                devbox_id=box.id if box else None,
                exit_status=None,
                duration_s=time.monotonic() - start,
                error=f"{type(e).__name__}: {e}",
            )

    try:
        with ThreadPoolExecutor(max_workers=count) as pool:
            futures = {pool.submit(worker, i): i for i in range(count)}
            for fut in as_completed(futures, timeout=max(0.0, deadline - time.monotonic())):
                results.append(fut.result())
    except TimeoutError:
        done = {r.devbox_id for r in results}
        for i in range(count - len(results)):
            results.append(BoxResult(devbox_id=None, exit_status=None,
                                     duration_s=timeout_s, error="run timed out"))
        print(f"WARN: timed out with {len(done)}/{count} complete", file=sys.stderr)
    finally:
        # Teardown EVERYTHING we provisioned, in parallel, tolerating failures.
        # Runs even on timeout or KeyboardInterrupt.
        if provisioned:
            with ThreadPoolExecutor(max_workers=len(provisioned)) as pool:
                def _shutdown(box_id: str) -> None:
                    for attempt in (1, 2, 3):
                        try:
                            client.devboxes.shutdown(box_id)
                            return
                        except errors as e:
                            if attempt == 3:
                                print(f"WARN: could not shut down {box_id}: {e}",
                                      file=sys.stderr)
                            time.sleep(0.2 * attempt)
                list(pool.map(_shutdown, provisioned))

    return results