feat(cli): preflight host resource limits before workers (#234)

areno/cli/diagnostics.py

@@ -32,6 +32,20 @@
     "HF_HUB_CACHE",
 )
 
+# Resource demand estimates for multi-process train/serve runs. These are
+# deliberately conservative upper bounds -- a false warning is cheap, a silent
+# FD/shm exhaustion mid-run is not. Constants are module-level so they are easy
+# to tune and surface in diagnostics output; they are not public API.
+_BASE_FDS_PER_WORKER = 64
+_SHM_BASELINE_BYTES = 1 << 30  # 1 GiB per tensor-parallel group; NCCL + CUDA IPC.
+_SHMMAX_PATH = "/proc/sys/kernel/shmmax"
+
+# Severity reused across the host-resource preflight; the diagnostics `check`
+# command uses OK/WARN/FAIL, which maps cleanly to normal/warning/blocking.
+RESOURCE_OK = "OK"
+RESOURCE_WARN = "WARN"
+RESOURCE_FAIL = "FAIL"
+
 
 @click.command(name="env", context_settings={"help_option_names": ["-h", "--help"]})
 @click.option("--json", "as_json", is_flag=True, help="Emit a machine-readable JSON support report.")
@@ -482,3 +496,232 @@ def _print_env_report(report: dict[str, Any]) -> None:
     click.echo("  Environment variables:")
     for name, value in report["env"].items():
         click.echo(f"    {name}={value if value is not None else '<unset>'}")
+
+
+def collect_host_limits() -> dict[str, Any]:
+    """Probe OS-level per-process resource limits without touching the engine.
+
+    Returns a dict with `file_descriptors`, `processes`, and `shared_memory`
+    entries. Each entry carries `available` (bool), `soft`/`hard` (for rlimits),
+    `value` (the effective limit used for comparison), and `error` when a probe
+    could not run. A probe that is unavailable on the current platform degrades
+    to a warning rather than a failure -- the preflight must not block runs on
+    platforms (e.g. macOS/Windows) where a limit simply does not exist.
+    """
+
+    return {
+        "file_descriptors": _fd_limit(),
+        "processes": _nproc_limit(),
+        "shared_memory": _shmmax_limit(),
+    }
+
+
+def _fd_limit() -> dict[str, Any]:
+    try:
+        import resource
+
+        soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+    except (AttributeError, ImportError, OSError, ValueError) as exc:
+        return _unavailable(f"{type(exc).__name__}: {exc}")
+    return _rlimit_entry(soft, hard)
+
+
+def _nproc_limit() -> dict[str, Any]:
+    try:
+        import resource
+
+        soft, hard = resource.getrlimit(resource.RLIMIT_NPROC)
+    except (ImportError, OSError, ValueError) as exc:
+        return _unavailable(f"{type(exc).__name__}: {exc}")
+    except AttributeError:
+        # RLIMIT_NPROC is not defined on every platform; degrade to a warning.
+        return _unavailable("RLIMIT_NPROC unavailable")
+    return _rlimit_entry(soft, hard)
+
+
+def _rlimit_entry(soft: Any, hard: Any) -> dict[str, Any]:
+    """Build a probe entry from a (soft, hard) rlimit pair.
+
+    `RLIM_INFINITY` means the limit is unbounded -- that satisfies any demand,
+    so the probe is `available` with `unbounded=True` rather than unavailable.
+    A finite soft limit is the effective value used for comparison.
+    """
+
+    import resource
+
+    soft_finite = soft != resource.RLIM_INFINITY
+    hard_finite = hard != resource.RLIM_INFINITY
+    if not soft_finite and not hard_finite:
+        return {
+            "available": True,
+            "unbounded": True,
+            "soft": soft,
+            "hard": hard,
+            "value": None,
+            "error": None,
+        }
+    value = soft if soft_finite else hard
+    return {
+        "available": True,
+        "unbounded": False,
+        "soft": soft,
+        "hard": hard,
+        "value": int(value),
+        "error": None,
+    }
+
+
+def _unavailable(error: str) -> dict[str, Any]:
+    return {
+        "available": False,
+        "unbounded": False,
+        "soft": None,
+        "hard": None,
+        "value": None,
+        "error": error,
+    }
+
+
+def _shmmax_limit() -> dict[str, Any]:
+    try:
+        text = Path(_SHMMAX_PATH).read_text(encoding="utf-8").strip()
+        value = int(text)
+    except (OSError, ValueError) as exc:
+        # /proc/sys/kernel/shmmax is Linux-only; absence is expected elsewhere.
+        return _unavailable(f"{type(exc).__name__}: {exc}")
+    return {"available": True, "unbounded": False, "soft": value, "hard": value, "value": value, "error": None}
+
+
+def estimate_resource_demand(world_size: int, tp_size: int) -> dict[str, int]:
+    """Return a documented upper-bound estimate of resource demand for one run.
+
+    - file descriptors: per-worker base plus one socket per cross-rank peer for
+      the NCCL/tensor-parallel mesh. `RLIMIT_NOFILE` is a per-process limit, so
+      this is the demand on a single worker process, not the fleet total.
+    - processes: `world_size` worker ranks plus the driver process. `RLIMIT_NPROC`
+      is a per-user limit, so the fleet total is the right quantity here.
+    - shared memory: baseline per tensor-parallel group, scaled by `tp_size`.
+
+    The estimate is intentionally conservative -- it triggers a warning rather
+    than undershooting a real NCCL/CUDA IPC requirement mid-run.
+    """
+
+    if world_size < 1:
+        raise ValueError("world_size must be >= 1")
+    if tp_size < 1:
+        raise ValueError("tp_size must be >= 1")
+    cross_rank_peers_per_worker = world_size - 1
+    return {
+        "file_descriptors": _BASE_FDS_PER_WORKER + cross_rank_peers_per_worker,
+        "processes": world_size + 1,
+        "shared_memory": _SHM_BASELINE_BYTES * tp_size,
+    }
+
+
+def preflight_host_resources(
+    world_size: int,
+    tp_size: int,
+    *,
+    policy: str = "warn",
+    limits: dict[str, Any] | None = None,
+) -> list[CheckResult]:
+    """Compare observed host limits against documented per-run demand.
+
+    Returns one `CheckResult` per resource dimension (file descriptors,
+    processes, shared memory). `policy` is one of skip/warn/block; it does not
+    change the returned severities -- the caller decides whether a FAIL under
+    `block` should abort. Probes that are unavailable degrade to WARN so the
+    preflight never blocks a run on a platform that simply lacks the limit.
+
+    `limits` is injectable for deterministic tests; production callers omit it
+    so `collect_host_limits()` probes the real host.
+    """
+
+    if policy not in {"skip", "warn", "block"}:
+        raise ValueError(f"resource-check policy must be skip/warn/block, got {policy!r}")
+    if policy == "skip":
+        # `skip` disables the preflight entirely; do not probe the host.
+        return []
+    if limits is None:
+        limits = collect_host_limits()
+    demand = estimate_resource_demand(world_size, tp_size)
+    return [
+        _fd_result(limits["file_descriptors"], demand["file_descriptors"]),
+        _nproc_result(limits["processes"], demand["processes"]),
+        _shmmax_result(limits["shared_memory"], demand["shared_memory"]),
+    ]
+
+
+def _fd_result(observed: dict[str, Any], required: int) -> CheckResult:
+    return _resource_result(
+        observed,
+        required,
+        name="file descriptors (RLIMIT_NOFILE)",
+        unit="",
+        adjust="raise the soft limit before launching workers, e.g. `ulimit -n 65536`",
+    )
+
+
+def _nproc_result(observed: dict[str, Any], required: int) -> CheckResult:
+    return _resource_result(
+        observed,
+        required,
+        name="processes (RLIMIT_NPROC)",
+        unit="",
+        adjust="raise the soft limit or run from a shell without a low nproc ulimit, e.g. `ulimit -u 32768`",
+    )
+
+
+def _shmmax_result(observed: dict[str, Any], required: int) -> CheckResult:
+    return _resource_result(
+        observed,
+        required,
+        name="shared memory (kernel.shmmax)",
+        unit=" bytes",
+        adjust=f"raise the system limit, e.g. `sudo sysctl -w kernel.shmmax={required}`",
+    )
+
+
+def _resource_result(
+    observed: dict[str, Any],
+    required: int,
+    *,
+    name: str,
+    unit: str,
+    adjust: str,
+) -> CheckResult:
+    if not observed.get("available"):
+        error = observed.get("error") or "probe unavailable"
+        return CheckResult(
+            RESOURCE_WARN,
+            name,
+            f"probe unavailable ({error}); required{unit}={required}",
+            f"Limit not observable on this platform; {adjust} if the run fails with fd/shmem exhaustion.",
+        )
+    if observed.get("unbounded"):
+        return CheckResult(RESOURCE_OK, name, f"observed=unbounded required{unit}={required}")
+    value = int(observed["value"])
+    delta = value - required
+    detail = f"observed{unit}={value} required{unit}={required} delta{unit}={delta}"
+    if value >= required:
+        return CheckResult(RESOURCE_OK, name, detail)
+    return CheckResult(RESOURCE_FAIL, name, detail, adjust)
+
+
+def format_resource_preflight(results: list[CheckResult]) -> str:
+    """Render resource preflight results as a concise human-readable block."""
+
+    lines = ["Host resource preflight:"]
+    for result in results:
+        lines.append(f"  {result.status:<4} {result.name}")
+        if result.detail:
+            lines.append(f"       {result.detail}")
+        if result.next_step and result.status in {RESOURCE_WARN, RESOURCE_FAIL}:
+            lines.append(f"       -> {result.next_step}")
+    return "\n".join(lines)
+
+
+def should_block_on_resources(results: list[CheckResult]) -> bool:
+    """True if any resource preflight result is a blocking failure."""
+
+    return any(result.status == RESOURCE_FAIL for result in results)

areno/cli/serve.py

@@ -30,6 +30,7 @@
 from areno.api.openai_chat import build_chat_completion_response, messages_to_prompt_tokens
 from areno.api.tokenizer import apply_chat_template_with_options, configure_chat_template_enable_thinking
 from areno.api.tool_call_parser import ToolCallParser, get_tool_call_parser, infer_tool_call_parser_name
+from areno.cli.diagnostics import RESOURCE_FAIL, format_resource_preflight, preflight_host_resources
 from areno.cli.model_refs import resolve_model_ref
 from areno.engine.data.tokenizer import load_processor, load_tokenizer
 
@@ -523,6 +524,29 @@ def _resolve_serve_attn_backend(
     return "native", warning
 
 
+def _preflight_host_resources(world_size: int, tp_size: int, *, policy: str) -> None:
+    """Run host resource preflight before engine/worker initialization.
+
+    Mirrors the train CLI: `skip` is a no-op, `warn` (default) emits a stderr
+    diagnostics block only on a FAIL and never aborts (preserving existing
+    behavior); `block` raises before workers start.
+    """
+
+    if policy == "skip":
+        return
+    results = preflight_host_resources(world_size, tp_size, policy=policy)
+    failures = [r for r in results if r.status == RESOURCE_FAIL]
+    if failures:
+        click.echo(format_resource_preflight(results), err=True)
+    if policy == "block" and failures:
+        failed = [r.name for r in failures]
+        raise click.UsageError(
+            "host resource limits are below the estimated demand for this serve run "
+            f"(world_size={world_size}, tp_size={tp_size}); failing probes: {failed}. "
+            "Raise the limits, re-run with --resource-check warn to ignore, or see the next-step hints above."
+        )
+
+
 async def _run_request_task(app: FastAPI, item: PendingRequest) -> None:
     """Run one HTTP request as an independent concurrent rollout call."""
 
@@ -1000,6 +1024,14 @@ def _normalize_stop(stop: str | list[str] | None) -> list[str]:
     is_flag=True,
     help="Pass enable_thinking=False to tokenizer chat templates when supported.",
 )
+@click.option(
+    "--resource-check",
+    type=click.Choice(["skip", "warn", "block"], case_sensitive=False),
+    default="warn",
+    show_default=True,
+    help="Preflight host fd/process/shm limits vs run demand before workers start. "
+    "warn (default) prints and never aborts; block aborts on a failed probe; skip disables.",
+)
 @click.option("--lora-rank", type=int, default=None, help="Enable native LoRA with this rank.")
 @click.option("--lora-alpha", type=float, default=16.0, show_default=True, help="Native LoRA alpha.")
 @click.option("--lora-dropout", type=float, default=0.0, show_default=True, help="Native LoRA dropout (must be 0).")
@@ -1027,6 +1059,7 @@ def serve_command(
     eager_decode: bool,
     attn_backend: Literal["flash", "native"],
     disable_thinking: bool,
+    resource_check: str,
     lora_rank: int | None,
     lora_alpha: float,
     lora_dropout: float,
@@ -1036,6 +1069,8 @@ def serve_command(
     """Click entry point: build the app and hand it to uvicorn."""
     import uvicorn
 
+    if default_backend_type() != MLX:
+        _preflight_host_resources(world_size, tp_size, policy=resource_check)
     if base_model_name_or_path is None:
         base_model_name_or_path = model_path
     model_path = resolve_model_ref(model_path, model_hub=model_hub)