fanout.py

"""rlfan — run one command across K devboxes concurrently.

YOUR TASK: implement `fan_out()` (and whatever helpers you want) per CANDIDATE.md.
`run_one()` below is a complete, working single-devbox lifecycle — it shows every
SDK call you need. Feel free to restructure anything in this file.

Run it:
    python fanout.py --count 3 --command "echo hello" --mock
    RLFAN_MOCK_FAIL_RATE=0.3 python fanout.py --count 5 --command "echo hi" --mock

Against the real API (needs RUNLOOP_API_KEY):
    python fanout.py --count 2 --command "python3 --version"

Tip: prefer the sync client + threads, or swap to AsyncRunloop + asyncio — your
call. The mock client is thread-safe.
"""

from __future__ import annotations

import argparse
import sys
import time
from dataclasses import dataclass
from typing import List, Optional


# --------------------------------------------------------------------------
# Client setup — provided. Works with either the real SDK or the local mock.
# --------------------------------------------------------------------------

def make_client(mock: bool):
    if mock:
        from mock_runloop import MockRunloop
        return MockRunloop()
    # Real client. Reads RUNLOOP_API_KEY from the environment by default.
    from runloop_api_client import Runloop
    return Runloop()


def api_errors(mock: bool) -> tuple:
    """The exception family to catch for API failures (mock or real)."""
    if mock:
        from mock_runloop import MockAPIError
        return (MockAPIError,)
    from runloop_api_client import APIError
    return (APIError,)


# --------------------------------------------------------------------------
# Result shape — feel free to extend.
# --------------------------------------------------------------------------

@dataclass
class BoxResult:
    devbox_id: Optional[str]      # None if provisioning never succeeded
    exit_status: Optional[int]    # None if the command never ran
    duration_s: float
    stdout_head: str = ""
    stderr_head: str = ""
    error: Optional[str] = None   # provisioning/API error, if any


# --------------------------------------------------------------------------
# WORKING EXAMPLE: full lifecycle for ONE devbox. Read this first.
# --------------------------------------------------------------------------

def run_one(client, command: str, mock: bool, name: str = "rlfan-example") -> BoxResult:
    start = time.monotonic()
    box = None
    try:
        # 1) Provision and block until the devbox reaches "running".
        box = client.devboxes.create_and_await_running(name=name)

        # 2) Execute the command synchronously; returns exit status + output.
        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 api_errors(mock) 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}",
        )
    finally:
        # 3) Always tear down anything we created.
        if box is not None:
            try:
                client.devboxes.shutdown(box.id)
            except api_errors(mock) as e:
                print(f"WARN: failed to shut down {box.id}: {e}", file=sys.stderr)


# --------------------------------------------------------------------------
# YOUR CODE: fan out across `count` devboxes, concurrently.
# --------------------------------------------------------------------------

def fan_out(client, command: str, count: int, timeout_s: float, mock: bool) -> List[BoxResult]:
    """Provision `count` devboxes, run `command` on all of them concurrently,
    and return one BoxResult per requested box.

    Requirements (see CANDIDATE.md):
      * concurrent — wall clock ≈ slowest box, not the sum
      * one box failing must not sink the others
      * every provisioned devbox gets shut down, no matter what
      * respect `timeout_s` for the run as a whole
    """
    raise NotImplementedError  # TODO(candidate)


def print_summary(results: List[BoxResult]) -> int:
    """Render results; return the process exit code (0 iff everything succeeded)."""
    print(f"\n{'DEVBOX':<18} {'EXIT':>4} {'SECS':>6}  OUTPUT")
    ok = True
    for r in results:
        status = "-" if r.exit_status is None else str(r.exit_status)
        detail = r.error or r.stderr_head or r.stdout_head
        print(f"{(r.devbox_id or '(none)'):<18} {status:>4} {r.duration_s:>6.1f}  {detail}")
        if r.exit_status != 0:
            ok = False
    print(f"\n{sum(1 for r in results if r.exit_status == 0)}/{len(results)} succeeded")
    return 0 if ok and results else 1


def main() -> int:
    p = argparse.ArgumentParser(prog="rlfan")
    p.add_argument("--command", required=True)
    p.add_argument("--count", type=int, default=3)
    p.add_argument("--timeout", type=float, default=120.0, help="whole-run timeout, seconds")
    p.add_argument("--mock", action="store_true", help="use the local fake client")
    args = p.parse_args()

    client = make_client(args.mock)

    # Smoke-test wiring: single-box demo. Replace with your fan_out call.
    if args.count == 1:
        results = [run_one(client, args.command, args.mock)]
    else:
        results = fan_out(client, args.command, args.count, args.timeout, args.mock)

    code = print_summary(results)

    # Interviewer check (mock mode only): did we leak anything?
    if args.mock:
        leaked = client.leaked_devboxes()
        if leaked:
            print(f"!! LEAKED {len(leaked)} devbox(es): {[b.id for b in leaked]}", file=sys.stderr)
            code = max(code, 2)
    return code


if __name__ == "__main__":
    raise SystemExit(main())