diff --git a/ainode/api/server.py b/ainode/api/server.py index 34c9e11..5b97c7a 100644 --- a/ainode/api/server.py +++ b/ainode/api/server.py @@ -50,6 +50,23 @@ logger = logging.getLogger(__name__) +def _client_max_bytes(config) -> int: + """Inbound request-body ceiling for the API server, in bytes. + + NOT optional: aiohttp defaults to 1 MB, which rejects long-context prompts + at the proxy. A model advertising 262k context can only be fed ~190k tokens + through our own endpoint before the caller gets an opaque 413 that says + nothing about which hop refused it (found 2026-08-25 benchmarking decode + against context depth). A bad value falls back to the default rather than + producing a server that rejects every body. + """ + try: + mb = int(getattr(config, "max_request_mb", 64)) + except (TypeError, ValueError): + mb = 64 + return max(1, mb) * 1024 * 1024 + + def create_app( config: Optional[NodeConfig] = None, engine=None, @@ -71,7 +88,10 @@ def create_app( auth_config = AuthConfig.load() - app = web.Application(middlewares=[cors_middleware, request_log_middleware, auth_middleware]) + app = web.Application( + middlewares=[cors_middleware, request_log_middleware, auth_middleware], + client_max_size=_client_max_bytes(config), + ) init_server_state(app) # Instantiate shared services collector = MetricsCollector() diff --git a/ainode/core/config.py b/ainode/core/config.py index 2965f3f..29796d0 100644 --- a/ainode/core/config.py +++ b/ainode/core/config.py @@ -4,7 +4,7 @@ import json from pathlib import Path from dataclasses import dataclass, asdict, field -from typing import List, Optional +from typing import Dict, List, Optional AINODE_HOME = Path(os.environ.get("AINODE_HOME", Path.home() / ".ainode")) CONFIG_FILE = AINODE_HOME / "config.json" @@ -73,6 +73,20 @@ class NodeConfig: # Setting this also disables the 0.17-era GB10 workarounds that would # otherwise be forced on (see NvidiaBackend._is_pinned_default_image). engine_image: str = "" + # Per-instance environment for the engine container. Some engine features + # are selected by env var, not by a `vllm serve` flag — the b12x FP4 kernel + # path is VLLM_NVFP4_GEMM_BACKEND + friends, with no CLI equivalent. Merged + # OVER the computed NCCL env at launch, so a recipe can also correct an + # autodetected NCCL value when a model needs it. Deliberately unvalidated, + # same as extra_vllm_args: the engine is the authority on what it accepts. + extra_env: Dict[str, str] = field(default_factory=dict) + # Max inbound request body for the API server, in MB. aiohttp defaults to + # 1 MB, which silently caps a 262k-context model at roughly 190k tokens of + # prompt: the proxy 413s the request before the engine ever sees it, and the + # caller gets "Request Entity Too Large" with nothing pointing at us. Sized + # for a 1M-token context (~5 MB of text) plus base64 image/video parts on + # the multimodal models, with headroom. + max_request_mb: int = 64 # Cluster cluster_enabled: bool = True diff --git a/ainode/engine/backends/nvidia.py b/ainode/engine/backends/nvidia.py index 2ceb515..58589ad 100644 --- a/ainode/engine/backends/nvidia.py +++ b/ainode/engine/backends/nvidia.py @@ -93,6 +93,9 @@ # NCCL tuning from Phase 1 floor verification — see # ops/slices/nvidia-vllm-engine/runbooks/01-nccl-floor-verification.md. +# A cold engine image is ~20 GB; a first pull on a slow link needs real room. +IMAGE_PULL_TIMEOUT = 3600 + NCCL_IB_GID_INDEX = "3" MASTER_PORT = "29501" @@ -160,6 +163,12 @@ def start_solo(self) -> bool: # Conflict. The head path already does this (see _launch_head_container); # solo needs it too. self._docker_stop_and_rm_best_effort(container_name) + # Before _build_solo_docker_cmd, because the argv prefix is derived from + # the image's ENTRYPOINT and inspecting a missing image yields nothing. + if not self.ensure_image(self._engine_image()): + logger.error("Engine image %s unavailable; not launching %s", + self._engine_image(), self.config.model) + return False cmd = self._build_solo_docker_cmd(container_name) env = self._build_env_for_subprocess() @@ -722,6 +731,20 @@ def _effective_kv_cache_dtype(self) -> str: return "auto" return dtype + def _engine_env(self, nccl_env: Dict[str, str]) -> Dict[str, str]: + """Env for the engine container: computed NCCL env + per-instance extras. + + ``extra_env`` is applied LAST so a recipe wins over an autodetected + value. Some engine features have no CLI flag at all (the b12x FP4 path + is selected purely by VLLM_NVFP4_GEMM_BACKEND and friends), so without + this they can only be reached by hand-rolling a container — which is + exactly what the launch path exists to avoid. + """ + env = dict(nccl_env or {}) + for key, value in (getattr(self.config, "extra_env", None) or {}).items(): + env[str(key)] = str(value) + return env + def _engine_image(self) -> str: """Container image for THIS instance — per-load override, else the fleet default. Lets one node run a 0.17 model and a 0.27 model side by @@ -754,6 +777,46 @@ def _serve_argv_prefix(self, image: str) -> List[str]: return ["serve"] return ["vllm", "serve"] # shim/shell entrypoint + def _image_present(self, image: str) -> bool: + """True if the image is already on this host.""" + try: + out = subprocess.run(["docker", "image", "inspect", image], + capture_output=True, text=True, timeout=20) + return out.returncode == 0 + except Exception: + return False + + def ensure_image(self, image: str, timeout: float = IMAGE_PULL_TIMEOUT) -> bool: + """Make sure ``image`` is on this host, pulling it if it isn't. + + Per-model ``engine_image`` means a node can be asked for an image it has + never run. Without this the launch just fails: docker starts an implicit + pull, the launch confirmation times out underneath it, and the caller + gets a bare "Failed to launch engine" with no mention of an image. The + entrypoint probe degrades too, since ``docker inspect`` on a missing + image returns nothing and we silently fall back to a default prefix. + Observed 2026-08-25 loading a recipe model onto a fresh node. + """ + if self._image_present(image): + return True + logger.info("Engine image %s not present; pulling (this can take several " + "minutes for a ~20 GB image)", image) + try: + out = subprocess.run(["docker", "pull", image], + capture_output=True, text=True, timeout=timeout) + except subprocess.TimeoutExpired: + logger.error("Timed out pulling engine image %s after %ss", image, timeout) + return False + except Exception as exc: + logger.error("Could not pull engine image %s: %s", image, exc) + return False + if out.returncode != 0: + logger.error("Failed to pull engine image %s: %s", image, + (out.stderr or out.stdout or "").strip()[-400:]) + return False + logger.info("Pulled engine image %s", image) + return True + def _image_entrypoint(self, image: str) -> List[str]: """The image's configured ENTRYPOINT, or [] when unknown.""" try: @@ -873,7 +936,7 @@ def _build_solo_docker_cmd(self, container_name: str) -> List[str]: "-v", f"{models_src}:{self.MODELS_MOUNT}:ro", ] - for key, value in nccl_env.items(): + for key, value in self._engine_env(nccl_env).items(): cmd.extend(["-e", f"{key}={value}"]) image = self._engine_image() @@ -948,7 +1011,7 @@ def _build_ray_docker_cmd( # peer's home-dir cache, which isn't under AINODE_HOME. "-v", f"{self._host_path(hf_cache_dir)}:/root/.cache/huggingface", ] - for key, value in nccl_env.items(): + for key, value in self._engine_env(nccl_env).items(): cmd.extend(["-e", f"{key}={value}"]) cmd.extend([self._engine_image(), "-c", ray_cmd]) return cmd @@ -1134,7 +1197,7 @@ def _build_run_cluster_cmd( f"--{role}", hf_cache_dir, ] - for key, value in nccl_env.items(): + for key, value in self._engine_env(nccl_env).items(): cmd.extend(["-e", f"{key}={value}"]) return cmd diff --git a/ainode/engine/instance_manager.py b/ainode/engine/instance_manager.py index 887cf89..9e459fd 100644 --- a/ainode/engine/instance_manager.py +++ b/ainode/engine/instance_manager.py @@ -11,6 +11,7 @@ from __future__ import annotations +import socket from dataclasses import dataclass from typing import Dict, List, Optional @@ -55,10 +56,35 @@ def is_empty(self) -> bool: def used_ports(self) -> set: return {i.record.api_port for i in self._instances.values()} - def allocate_port(self) -> int: - """Lowest free port from base_port up, not held by an existing instance.""" + @staticmethod + def _port_bindable(port: int, host: str = "0.0.0.0") -> bool: + """True if nothing on the HOST is already listening on this port. + + The engine container runs on the host network, so a port owned by any + other process (not just another AINode instance) collides. Without this + check the engine launches, fails deep in startup with + ``OSError: [Errno 98] Address already in use``, and the caller only sees + a generic launch failure. Observed 2026-08-25 on a node where an + unrelated service had held 8000 for weeks. + """ + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind((host, port)) + return True + except OSError: + return False + + def allocate_port(self, probe: bool = True) -> int: + """Lowest free port from base_port up. + + Skips ports held by an existing instance AND, unless ``probe`` is off, + ports already bound by anything else on the host. + """ used = self.used_ports() port = self._base_port - while port in used: + for _ in range(256): + if port not in used and (not probe or self._port_bindable(port)): + return port port += 1 return port diff --git a/ainode/models/api_routes.py b/ainode/models/api_routes.py index 2956379..da9853e 100644 --- a/ainode/models/api_routes.py +++ b/ainode/models/api_routes.py @@ -145,7 +145,7 @@ def load_instance_manifest() -> list: _OVERRIDE_KEYS = ("served_model_name", "max_model_len", "kv_cache_dtype", "kv_cache_dtype_explicit", "quantization", "trust_remote_code", - "extra_vllm_args", "engine_image") + "extra_vllm_args", "engine_image", "extra_env") def catalog_recipe(model: str) -> dict: @@ -155,7 +155,8 @@ def catalog_recipe(model: str) -> dict: flag set (spec-decode, MoE/mamba backends, reasoning + tool-call parsers). Carrying that in the catalog is what makes them a one-click load instead of a hand-rolled container. Returns {} for anything not curated. Keys: - ``engine_image``, ``extra_vllm_args``, and ``gpu_memory_utilization``. + ``engine_image``, ``extra_vllm_args``, ``extra_env``, and + ``gpu_memory_utilization``. """ from ainode.models.registry import CURATED_CLUSTER_MODELS m = (model or "").strip() @@ -168,6 +169,8 @@ def catalog_recipe(model: str) -> dict: recipe["engine_image"] = info.engine_image if getattr(info, "extra_vllm_args", None): recipe["extra_vllm_args"] = list(info.extra_vllm_args) + if getattr(info, "extra_env", None): + recipe["extra_env"] = dict(info.extra_env) if getattr(info, "recommended_gmu", 0): recipe["gpu_memory_utilization"] = info.recommended_gmu return recipe @@ -568,6 +571,16 @@ async def handle_model_load(request: web.Request) -> web.Response: "(e.g. [\"--moe-backend\", \"marlin\"]) or a shell-style string"}, status=400) overrides["extra_vllm_args"] = [str(a) for a in raw] + if body.get("extra_env") is not None: + raw = body["extra_env"] + if not isinstance(raw, dict) or not all( + isinstance(k, str) and k and isinstance(v, (str, int, float, bool)) + for k, v in raw.items()): + return web.json_response( + {"error": "extra_env must be an object of NAME -> value " + "(e.g. {\"VLLM_NVFP4_GEMM_BACKEND\": \"flashinfer-b12x\"})"}, + status=400) + overrides["extra_env"] = {k: str(v) for k, v in raw.items()} if body.get("engine_image") is not None: img = str(body["engine_image"]).strip() if " " in img: @@ -580,7 +593,7 @@ async def handle_model_load(request: web.Request) -> web.Response: # engine image and flags it actually needs. Anything the caller stated # explicitly above wins; the recipe only fills the gaps. recipe = catalog_recipe(model) - for key in ("engine_image", "extra_vllm_args"): + for key in ("engine_image", "extra_vllm_args", "extra_env"): if key in recipe and key not in overrides: overrides[key] = recipe[key] if gmu is None and "gpu_memory_utilization" in recipe: diff --git a/ainode/models/registry.py b/ainode/models/registry.py index c8ec416..2318875 100644 --- a/ainode/models/registry.py +++ b/ainode/models/registry.py @@ -88,6 +88,7 @@ class ModelInfo: # wins over the recipe; the recipe only fills what wasn't specified. engine_image: str = "" # "" = fleet default engine image extra_vllm_args: list = None # verbatim `vllm serve` flags + extra_env: dict = None # engine-container env (e.g. b12x kernel selection) recommended_gmu: float = 0.0 # 0 = use node default gpu_memory_utilization def __post_init__(self): @@ -95,6 +96,8 @@ def __post_init__(self): self.capabilities = [] if self.extra_vllm_args is None: self.extra_vllm_args = [] + if self.extra_env is None: + self.extra_env = {} def to_dict(self) -> dict: return asdict(self) diff --git a/scripts/bench-serve.py b/scripts/bench-serve.py new file mode 100755 index 0000000..8a596e2 --- /dev/null +++ b/scripts/bench-serve.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +"""Baseline serve benchmark for the AINode OpenAI-compatible endpoint. + +Measures TTFT + decode tok/s single-stream at several prompt sizes, then a +concurrency sweep (the metric that actually matters for batched/agent workloads). +Non-destructive — pure inference load. stdlib only. + + python3 scripts/bench-serve.py --url http://100.122.26.9:3000 \ + --model nvidia/Qwen3-235B-A22B-NVFP4 + +Depth mode answers the question the peak number can't: what does decode actually +do as the KV cache fills? Decode reads the occupied KV every token, so tok/s at +4k and tok/s at 128k are different numbers and get quoted interchangeably. + + python3 scripts/bench-serve.py --url http://100.122.26.9:3000 \ + --model unsloth/Qwen3.8-27B-NVFP4 --mode depth \ + --depths 4000,16000,32000,64000,128000 --reps 3 + +Two things keep depth numbers honest. Each request carries a unique nonce at the +FRONT of the prompt so `--enable-prefix-caching` can't serve a cached prefill and +make depth look free. And the x-axis is the server's own `usage.prompt_tokens`, +never our estimate, so a bad chars-per-token guess shifts nothing. + +# ponytail: stdlib urllib + threads; no httpx/asyncio dep for a lab bench. +""" +import argparse +import json +import statistics +import time +import urllib.request +import uuid +from concurrent.futures import ThreadPoolExecutor + +PROMPTS = { + "short": "Say hello in one sentence.", + "med_2k": "Summarize the following, then list 5 implications.\n" + ("lorem ipsum dolor sit amet " * 300), + "long_8k": "Read this and answer: what is the main theme?\n" + ("the quick brown fox jumps over the lazy dog " * 1500), +} + + +def one_request(url, model, prompt, max_tokens, want_usage=False, no_think=False): + """Stream a completion. + + Returns (ttft_s, decode_toks_per_s, n_tokens, ok, prompt_tokens). The 5th + element is appended rather than inserted so existing positional callers + (the concurrency sweep) keep working unchanged. + """ + payload = { + "model": model, "stream": True, "max_tokens": max_tokens, + "messages": [{"role": "user", "content": prompt}], + } + if no_think: + payload["chat_template_kwargs"] = {"enable_thinking": False} + if want_usage: + # vLLM emits a final chunk with `choices: []` and a usage block. That + # empty choices list is why the parse loop below can't assume [0]. + payload["stream_options"] = {"include_usage": True} + body = json.dumps(payload).encode() + req = urllib.request.Request(url.rstrip("/") + "/v1/chat/completions", + data=body, headers={"Content-Type": "application/json"}) + t0 = time.monotonic() + ttft = None + n = 0 + prompt_tokens = None + completion_tokens = None + try: + with urllib.request.urlopen(req, timeout=600) as r: + for raw in r: + line = raw.decode("utf-8", "ignore").strip() + if not line.startswith("data:"): + continue + data = line[5:].strip() + if data == "[DONE]": + break + chunk = json.loads(data) + usage = chunk.get("usage") + if usage: + prompt_tokens = usage.get("prompt_tokens") or prompt_tokens + completion_tokens = usage.get("completion_tokens") or completion_tokens + choices = chunk.get("choices") or [] + if not choices: + continue # usage-only chunk + d_obj = choices[0].get("delta", {}) or {} + # A reasoning parser splits output into reasoning_content and + # content. Both are generated tokens and both cost decode time, + # so counting only `content` undercounts throughput and can + # leave ttft unset entirely on a reply that is all thinking. + # Key name varies by build: vLLM 0.27.1 emits `reasoning`, + # other versions/parsers use `reasoning_content`. Check all three. + delta = (d_obj.get("content") or d_obj.get("reasoning") + or d_obj.get("reasoning_content")) + if delta: + if ttft is None: + ttft = time.monotonic() - t0 + n += 1 + except Exception as e: + return (None, None, 0, f"ERR {e}", None) + total = time.monotonic() - t0 + # Token count comes from the server's usage block, NOT from counting SSE + # chunks. Under speculative decoding a chunk can carry several accepted + # tokens at once, so chunk-counting reports roughly rate/acceptance-length + # and silently halves the number on an MTP or DFlash config. + gen = completion_tokens if completion_tokens else n + # Decode rate EXCLUDES prefill: the clock starts at the first token, not at + # request send. Including TTFT is what turns a decode number into an + # agent-loop number. + decode = (gen - 1) / (total - ttft) if ttft and gen > 1 and total > ttft else 0.0 + return (ttft, decode, gen, "ok", prompt_tokens) + + +# ---------------------------------------------------------------- depth sweep + +FILLER = ( + "The memory bandwidth of a device sets a hard ceiling on single-stream decode, " + "because every generated token requires reading the active weights out of memory. " + "Quantization shrinks that read. Tensor parallelism splits it across nodes and " + "charges for it in interconnect traffic. Speculative decoding is the only lever " + "that changes the equation itself, by producing more than one token per pass. " +) + + +def calibrate_chars_per_token(url, model): + """Measure this model's chars-per-token on FILLER instead of assuming ~4. + + One cheap probe. If it fails we fall back to 4.0, and it costs us nothing + either way because the reported x-axis is the server's own prompt_tokens. + """ + probe = FILLER * 40 + r = one_request(url, model, probe, 1, want_usage=True) # 1 tok: we only want usage + if r[3] == "ok" and r[4]: + return len(probe) / r[4] + return 4.0 + + +def make_prompt(target_tokens, cpt): + """Prompt of roughly target_tokens, with a unique nonce FIRST. + + The nonce leads so it invalidates the shared prefix. A trailing nonce would + still let prefix caching serve almost the entire prefill and we would be + timing a cache hit while calling it a depth measurement. + """ + nonce = f"[run {uuid.uuid4().hex}] " + reps = max(1, int((target_tokens * cpt - len(nonce)) / len(FILLER))) + # Ask for a long answer on purpose: a 7-token reply makes the decode rate + # noise. We want enough generated tokens for the rate to mean something. + return nonce + (FILLER * reps) + ( + "\n\nWrite a detailed 400-word explanation of the tradeoffs described above. " + "Be thorough and do not stop early.") + + +def depth_sweep(url, model, depths, reps, max_tokens, no_think=False): + cpt = calibrate_chars_per_token(url, model) + print(f"== decode vs context depth ({max_tokens} tok gen, {reps} rep(s), {cpt:.2f} chars/token) ==") + print(f" {'target':>8} {'actual':>8} {'TTFT':>9} {'decode':>12} {'gen':>5}") + rows = [] + for d in depths: + prompt = None + decodes, ttfts, actuals, ntoks = [], [], [], [] + for _ in range(reps): + prompt = make_prompt(d, cpt) # new nonce each rep + ttft, dec, n, ok, ptok = one_request(url, model, prompt, max_tokens, + want_usage=True, no_think=no_think) + if ok != "ok": + print(f" {d:>8} {ok}") + break + if ttft is None or n < 2: + print(f" {d:>8} no usable tokens (n={n}); raise --max-tokens or use --no-think") + break + decodes.append(dec) + ttfts.append(ttft) + ntoks.append(n) + if ptok: + actuals.append(ptok) + if not decodes: + continue + med_dec = statistics.median(decodes) + med_ttft = statistics.median(ttfts) + actual = int(statistics.median(actuals)) if actuals else 0 + rows.append((d, actual, med_ttft, med_dec)) + spread = f" (n={len(decodes)}, {min(decodes):.1f}-{max(decodes):.1f})" if len(decodes) > 1 else "" + print(f" {d:>8} {actual:>8} {med_ttft*1000:>7.0f}ms {med_dec:>8.1f} t/s {int(statistics.median(ntoks)):>5}{spread}") + if len(rows) > 1: + first, last = rows[0], rows[-1] + drop = (1 - last[3] / first[3]) * 100 if first[3] else 0 + print(f"\n decode falls {drop:.0f}% from {first[1]} to {last[1]} prompt tokens " + f"({first[3]:.1f} -> {last[3]:.1f} tok/s)") + return rows + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--url", required=True) + ap.add_argument("--model", required=True) + ap.add_argument("--max-tokens", type=int, default=200) + ap.add_argument("--concurrency", default="1,2,4,8,16") + ap.add_argument("--mode", choices=["legacy", "depth", "all"], default="legacy", + help="legacy = original single-stream + concurrency sections (default)") + ap.add_argument("--depths", default="4000,16000,32000,64000,128000", + help="prompt sizes in tokens for --mode depth/all") + ap.add_argument("--reps", type=int, default=3, help="repetitions per depth; median is reported") + ap.add_argument("--no-think", action="store_true", + help="send chat_template_kwargs.enable_thinking=false for stable generation length") + args = ap.parse_args() + + if args.mode in ("depth", "all"): + depth_sweep(args.url, args.model, [int(x) for x in args.depths.split(",")], + args.reps, args.max_tokens, no_think=args.no_think) + if args.mode == "depth": + return + print() + + print(f"== single-stream ({args.max_tokens} tok) ==") + for name, p in PROMPTS.items(): + ttft, dec, n, ok, _ = one_request(args.url, args.model, p, args.max_tokens) + if ok != "ok": + print(f" {name:8} {ok}") + else: + print(f" {name:8} TTFT {ttft*1000:6.0f}ms decode {dec:5.1f} tok/s ({n} tok)") + + print("\n== concurrency sweep (short prompt, aggregate throughput) ==") + for c in [int(x) for x in args.concurrency.split(",")]: + t0 = time.monotonic() + with ThreadPoolExecutor(max_workers=c) as ex: + res = list(ex.map(lambda _: one_request(args.url, args.model, PROMPTS["short"], args.max_tokens), range(c))) + wall = time.monotonic() - t0 + ok = [r for r in res if r[3] == "ok"] + toks = sum(r[2] for r in ok) + ttfts = [r[0] for r in ok if r[0]] + agg = toks / wall if wall else 0 + med_ttft = statistics.median(ttfts) * 1000 if ttfts else 0 + print(f" c={c:2} {len(ok)}/{c} ok agg {agg:6.1f} tok/s median TTFT {med_ttft:6.0f}ms wall {wall:4.1f}s") + + +if __name__ == "__main__": + main() diff --git a/tests/test_api.py b/tests/test_api.py index 2a1335c..7862ee1 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -6,7 +6,7 @@ import pytest_asyncio from aiohttp.test_utils import TestClient, TestServer -from ainode.api.server import create_app, _build_announcement +from ainode.api.server import create_app, _build_announcement, _client_max_bytes from ainode.core.config import NodeConfig @@ -385,3 +385,40 @@ async def test_nodes_include_stacked_instances(client, app): n2 = next(n for n in data["nodes"] if n["node_id"] == "n2") assert "instances" in n2 assert any(i["model"] == "stacked-model" and i["api_port"] == 8001 for i in n2["instances"]) + + +# --- request body ceiling ------------------------------------------------- +# aiohttp defaults client_max_size to 1 MB. That silently caps long-context +# models: a 262k-context model can only be fed ~190k tokens of prompt through +# our endpoint before the proxy 413s it, and nothing in the error says the +# proxy (not the engine) refused. Found 2026-08-25 while benchmarking decode +# vs context depth — a 200k-token prompt is ~1.01 MB and died on the default. +# +# These assert on _client_max_bytes directly rather than building extra apps: +# create_app() outside a running loop binds module-level asyncio primitives +# (the download semaphore) to the wrong loop and made an unrelated download +# test flaky. + +def test_request_ceiling_clears_the_1mb_default(): + from ainode.api.server import _client_max_bytes + one_mb = 1024 * 1024 + assert _client_max_bytes(NodeConfig()) > one_mb + # A 1M-token context is roughly 5 MB of text; leave room for that plus + # base64 image/video parts on the multimodal models. + assert _client_max_bytes(NodeConfig()) >= 32 * one_mb + + +def test_request_ceiling_is_configurable(): + from ainode.api.server import _client_max_bytes + assert _client_max_bytes(NodeConfig(max_request_mb=8)) == 8 * 1024 * 1024 + + +def test_request_ceiling_never_collapses_to_zero(): + from ainode.api.server import _client_max_bytes + assert _client_max_bytes(NodeConfig(max_request_mb=0)) >= 1024 * 1024 + assert _client_max_bytes(NodeConfig(max_request_mb="nonsense")) >= 1024 * 1024 + + +def test_app_is_actually_wired_to_the_ceiling(app): + # Uses the existing fixture app rather than constructing another one. + assert app._client_max_size == _client_max_bytes(app["config"]) diff --git a/tests/test_engine_recipe_passthrough.py b/tests/test_engine_recipe_passthrough.py index db02ba0..c330491 100644 --- a/tests/test_engine_recipe_passthrough.py +++ b/tests/test_engine_recipe_passthrough.py @@ -202,3 +202,59 @@ def test_qwen38_recipe_pins_kv_cache_auto_for_vision(): assert a[a.index("--kv-cache-dtype") + 1] == "auto" built = args_for(extra_vllm_args=a) assert built.count("--kv-cache-dtype") == 1 and "fp8" not in built + + +# --- extra_env passthrough (b12x and friends are env-selected, not flag-selected) --- + +def _env_pairs(cmd): + """Extract {NAME: value} from the `-e NAME=value` pairs of a docker cmd.""" + out = {} + for i, tok in enumerate(cmd): + if tok == "-e" and i + 1 < len(cmd) and "=" in cmd[i + 1]: + k, _, v = cmd[i + 1].partition("=") + out[k] = v + return out + + +def test_extra_env_reaches_the_engine_container(): + b = NvidiaBackend(NodeConfig(model="m", extra_env={ + "VLLM_NVFP4_GEMM_BACKEND": "flashinfer-b12x", + "VLLM_USE_FLASHINFER_MOE_FP4": "1", + })) + env = _env_pairs(b._build_solo_docker_cmd("c")) + assert env["VLLM_NVFP4_GEMM_BACKEND"] == "flashinfer-b12x" + assert env["VLLM_USE_FLASHINFER_MOE_FP4"] == "1" + + +def test_extra_env_overrides_a_computed_value(): + # The pinned default image forces VLLM_NVFP4_GEMM_BACKEND=marlin. A recipe + # selecting the b12x kernel path must win, or b12x is unreachable on it. + b = NvidiaBackend(NodeConfig(model="some/model-NVFP4", + extra_env={"VLLM_NVFP4_GEMM_BACKEND": "flashinfer-b12x"})) + assert b._nvfp4_serve_env().get("VLLM_NVFP4_GEMM_BACKEND") == "marlin" + assert _env_pairs(b._build_solo_docker_cmd("c"))["VLLM_NVFP4_GEMM_BACKEND"] == "flashinfer-b12x" + + +def test_no_extra_env_leaves_the_command_untouched(): + base = NvidiaBackend(NodeConfig(model="m"))._build_solo_docker_cmd("c") + same = NvidiaBackend(NodeConfig(model="m", extra_env={}))._build_solo_docker_cmd("c") + assert base == same + + +def test_extra_env_values_are_stringified(): + b = NvidiaBackend(NodeConfig(model="m", extra_env={"FLASHINFER_DISABLE_VERSION_CHECK": 1})) + assert _env_pairs(b._build_solo_docker_cmd("c"))["FLASHINFER_DISABLE_VERSION_CHECK"] == "1" + + +def test_catalog_recipe_surfaces_extra_env_when_a_model_carries_it(): + from ainode.models.registry import ModelInfo + import ainode.models.registry as reg + probe = ModelInfo(id="probe-b12x", name="probe", hf_repo="org/probe-b12x", + size_gb=1.0, description="d", + extra_env={"VLLM_NVFP4_GEMM_BACKEND": "flashinfer-b12x"}) + reg.CURATED_CLUSTER_MODELS["probe-b12x"] = probe + try: + assert catalog_recipe("org/probe-b12x")["extra_env"] == { + "VLLM_NVFP4_GEMM_BACKEND": "flashinfer-b12x"} + finally: + del reg.CURATED_CLUSTER_MODELS["probe-b12x"] diff --git a/tests/test_instance_manager.py b/tests/test_instance_manager.py index 68a2943..71d112e 100644 --- a/tests/test_instance_manager.py +++ b/tests/test_instance_manager.py @@ -127,3 +127,31 @@ def test_eject_stops_one_instance(): assert b.stopped is True assert mgr.by_model("A") is None assert app["engine"] is None # primary cleared + + +# --- port allocation must respect the whole host, not just our own instances --- + +def test_allocate_port_skips_a_port_held_by_another_process(): + # The engine container shares the host network, so a port owned by an + # unrelated service collides. Before this, allocate_port() handed it out + # anyway and the engine died with EADDRINUSE deep in startup, surfacing to + # the caller as a generic "Failed to launch engine". + import socket + from ainode.engine.instance_manager import InstanceManager + + squatter = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + squatter.bind(("0.0.0.0", 0)) + squatter.listen(1) + taken = squatter.getsockname()[1] + try: + mgr = InstanceManager(base_port=taken) + assert mgr.allocate_port() != taken + finally: + squatter.close() + + +def test_allocate_port_probe_can_be_disabled(): + from ainode.engine.instance_manager import InstanceManager + mgr = InstanceManager(base_port=8000) + assert mgr.allocate_port(probe=False) == 8000 diff --git a/tests/test_launch_robustness.py b/tests/test_launch_robustness.py index c86fa00..530b995 100644 --- a/tests/test_launch_robustness.py +++ b/tests/test_launch_robustness.py @@ -129,3 +129,61 @@ def test_cli_prefers_the_container_backend_when_vllm_is_missing(self): assert body.index('importlib.util.find_spec("vllm")') < body.index( "from ainode.engine.vllm_engine import VLLMEngine" ), "the viability check must come BEFORE falling back to the legacy engine" + + +# --- a per-model engine_image the node has never seen must be pulled, not fail --- + +def test_ensure_image_is_a_noop_when_present(monkeypatch): + from ainode.core.config import NodeConfig + from ainode.engine.backends.nvidia import NvidiaBackend + b = NvidiaBackend(NodeConfig(model="m")) + monkeypatch.setattr(b, "_image_present", lambda img: True) + calls = [] + monkeypatch.setattr("subprocess.run", lambda *a, **k: calls.append(a) or None) + assert b.ensure_image("some/image:1") is True + assert not calls, "must not pull an image that is already here" + + +def test_ensure_image_pulls_when_missing(monkeypatch): + import subprocess + from ainode.core.config import NodeConfig + from ainode.engine.backends.nvidia import NvidiaBackend + b = NvidiaBackend(NodeConfig(model="m")) + monkeypatch.setattr(b, "_image_present", lambda img: False) + seen = {} + + def fake_run(cmd, **kw): + seen["cmd"] = cmd + return subprocess.CompletedProcess(cmd, 0, "ok", "") + + monkeypatch.setattr("subprocess.run", fake_run) + assert b.ensure_image("some/image:1") is True + assert seen["cmd"][:2] == ["docker", "pull"] + assert seen["cmd"][2] == "some/image:1" + + +def test_ensure_image_reports_failure_instead_of_launching_blind(monkeypatch): + import subprocess + from ainode.core.config import NodeConfig + from ainode.engine.backends.nvidia import NvidiaBackend + b = NvidiaBackend(NodeConfig(model="m")) + monkeypatch.setattr(b, "_image_present", lambda img: False) + monkeypatch.setattr("subprocess.run", + lambda cmd, **kw: subprocess.CompletedProcess(cmd, 1, "", "no such image")) + assert b.ensure_image("bogus/nope:1") is False + + +def test_start_solo_refuses_when_the_image_cannot_be_had(monkeypatch): + # Previously this launched anyway: docker began an implicit pull, the launch + # confirmation timed out under it, and the caller saw a bare failure with no + # mention of an image. + from ainode.core.config import NodeConfig + from ainode.engine.backends.nvidia import NvidiaBackend + b = NvidiaBackend(NodeConfig(model="m")) + monkeypatch.setattr(b, "is_running", lambda: False) + monkeypatch.setattr(b, "_docker_stop_and_rm_best_effort", lambda name: None) + monkeypatch.setattr(b, "ensure_image", lambda img, **kw: False) + launched = [] + monkeypatch.setattr("subprocess.Popen", lambda *a, **k: launched.append(a)) + assert b.start_solo() is False + assert not launched, "must not run the container when the image is unavailable" diff --git a/tests/test_nvidia_backend.py b/tests/test_nvidia_backend.py index aeb90b6..3d61352 100644 --- a/tests/test_nvidia_backend.py +++ b/tests/test_nvidia_backend.py @@ -113,6 +113,10 @@ def test_nvidia_backend_is_no_longer_notimplemented(self): # reporting success (0.5.4) — there is no real docker here, so stub # the state probe. See TestLaunchConfirmation for that contract. NvidiaBackend, "_docker_container_state", return_value="running", + ), mock.patch.object( + # Same reason as above: no real docker, so the engine-image probe + # would refuse the launch. Covered by test_launch_robustness.py. + NvidiaBackend, "ensure_image", return_value=True, ): result = backend.start() assert result is True @@ -286,6 +290,12 @@ def test_start_solo_invokes_docker_run_with_expected_args(self): # test keeps asserting on the docker-run argv alone. backend, "_image_entrypoint", return_value=["/opt/nvidia/nvidia_entrypoint.sh"], + ), mock.patch.object( + # start_solo() now pulls a missing engine image before launching. + # There is no docker here, so the probe would report it absent and + # refuse. Image handling has its own tests in + # test_launch_robustness.py; this one is about the argv. + backend, "ensure_image", return_value=True, ): result = backend.start_solo()