From 8d6f0461f486e0af25edf4870cb97b736f930571 Mon Sep 17 00:00:00 2001 From: webdevtodayjason Date: Mon, 6 Jul 2026 14:13:35 -0500 Subject: [PATCH] fix(training,models,web): slug mapping, offline peft, stacked-load admission guard, node-targeted launch, stacked visibility Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_017NziXzqT1L9kj2T1byA3Ak --- ainode/api/server.py | 11 ++ ainode/engine/sharding_routes.py | 16 ++- ainode/models/api_routes.py | 46 +++++++- ainode/training/_run_training.py | 17 +++ ainode/training/api_routes.py | 8 +- ainode/training/engine.py | 189 +++++++++++++++++++++++++++++-- ainode/web/static/js/app.js | 71 +++++++++--- ainode/web/templates/index.html | 2 + tests/test_api.py | 60 ++++++++++ tests/test_cluster_load.py | 14 ++- tests/test_models.py | 97 ++++++++++++++++ tests/test_sharding_launch.py | 22 ++++ tests/test_training_command.py | 152 ++++++++++++++++++++++++- 13 files changed, 667 insertions(+), 38 deletions(-) diff --git a/ainode/api/server.py b/ainode/api/server.py index 7d207ea..a63da2f 100644 --- a/ainode/api/server.py +++ b/ainode/api/server.py @@ -673,6 +673,17 @@ async def handle_nodes(request: web.Request) -> web.Response: "distributed_mode": dmode, "distributed_instance_id": getattr(n, "distributed_instance_id", None), "distributed_peers": list(getattr(n, "distributed_peers", []) or []), + # Per-node instance list (primary + any stacked models on ports + # 8001+). The node card renders a sub-row per stacked instance so + # they're no longer invisible in the dashboard. Same source the + # proxy's _routing_candidates uses, so views and routing agree. + "instances": [ + {"model": inst.get("model"), + "api_port": inst.get("api_port"), + "status": inst.get("status")} + for inst in (getattr(n, "instances", []) or []) + if isinstance(inst, dict) and inst.get("model") + ], }) else: # Fallback: return this node diff --git a/ainode/engine/sharding_routes.py b/ainode/engine/sharding_routes.py index 154bb9c..564f060 100644 --- a/ainode/engine/sharding_routes.py +++ b/ainode/engine/sharding_routes.py @@ -209,8 +209,22 @@ def fabric_of(n): name_token = "" if port == config.api_port else str(port) # primary keeps legacy names instance_id = f"{config.node_id or 'head'}:{model}" + # Honour a per-launch GPU memory fraction if the UI supplied one (the same + # #launch-gmu box the solo path uses). Without this the distributed backend + # always renders the shared NodeConfig default (0.5), silently dropping the + # value the user typed for a TP>1 launch. Ignore junk / out-of-range input. + overrides: dict = {} + gmu_raw = body.get("gpu_memory_utilization") + if gmu_raw is not None: + try: + gmu = float(gmu_raw) + except (TypeError, ValueError): + gmu = None + if gmu is not None and 0.0 < gmu <= 1.0: + overrides["gpu_memory_utilization"] = gmu + inst_config = replace(config, model=model, distributed_mode="head", - peer_ips=chosen_peers, api_port=port) + peer_ips=chosen_peers, api_port=port, **overrides) backend = get_backend(inst_config, instance_id=name_token) try: started = backend.start_distributed() diff --git a/ainode/models/api_routes.py b/ainode/models/api_routes.py index 37404db..e1be56e 100644 --- a/ainode/models/api_routes.py +++ b/ainode/models/api_routes.py @@ -209,10 +209,51 @@ def append_solo_instance(app, model: str, gmu=None, *, overrides=None, persist: manager = InstanceManager(base_port=config.api_port) app["instances"] = manager - # Re-loading a model already up replaces THAT instance (stop it first), not - # the whole node — other stacked instances are untouched. + # Re-loading a model already up replaces THAT instance — other stacked + # instances are untouched. We must NOT stop it yet: the admission gate below + # can still reject the request (400/409), and killing the live instance + # BEFORE that check would leave the model unloaded on a "failed" request with + # no automatic restore. So decide admission first, then destroy. existing = manager.by_model(model) replaced_primary = existing is not None and app.get("engine") is existing.backend + + def _existing_id(): + return existing.record.instance_id if existing is not None else None + + # This load becomes the primary iff no OTHER instance remains once `existing` + # (if any) is replaced — i.e. reloading the sole instance, or the very first + # load. Reloading a stacked model while the primary is up is NOT primary. + others = [i for i in manager.instances() if i.record.instance_id != _existing_id()] + is_primary = len(others) == 0 + + # Stacked-load admission control (unified-memory safety). A 2nd+ model on a + # busy node with the node default gpu_memory_utilization (0.5) can push total + # reservation past capacity and crash the whole GB10 (host death, power cycle). + # Reloading the primary keeps today's behavior. Reloading an existing stacked + # model excludes its own (about-to-be-freed) reservation from the running + # total. Runs BEFORE stop/remove so a rejection leaves the live model intact. + if not is_primary and not replaced_primary: + if gmu is None: + return {"ok": False, "status": 400, + "error": ("A stacked load (2nd+ model on this node) must specify " + "gpu_memory_utilization explicitly (e.g. 0.4) — refusing " + "to inherit the node default and risk overcommitting " + "unified memory.")} + existing_total = 0.0 + for inst in others: + g = getattr(getattr(inst.backend, "config", None), "gpu_memory_utilization", None) + if g is not None: + existing_total += float(g) + projected = existing_total + gmu + if projected > 0.9: + return {"ok": False, "status": 409, + "error": (f"Refusing stacked load: this node already reserves " + f"{existing_total:.2f} of GPU memory across " + f"{len(others)} instance(s); the requested " + f"{gmu:.2f} would total {projected:.2f} (> 0.90 cap). " + f"Unload a model or lower gpu_memory_utilization.")} + + # Admission passed (or N/A) — NOW it's safe to tear down the old instance. if existing is not None: try: existing.backend.stop() @@ -220,7 +261,6 @@ def append_solo_instance(app, model: str, gmu=None, *, overrides=None, persist: pass manager.remove(existing.record.instance_id) - is_primary = manager.is_empty() port = manager.allocate_port() name_token = "" if port == config.api_port else str(port) # primary keeps legacy names instance_id = f"{config.node_id or 'head'}:{model}" diff --git a/ainode/training/_run_training.py b/ainode/training/_run_training.py index 18db557..af52e7e 100644 --- a/ainode/training/_run_training.py +++ b/ainode/training/_run_training.py @@ -63,6 +63,23 @@ def main() -> None: if resolved.exists(): config["dataset_path"] = str(resolved) + # Resolve an on-disk base_model slug to a local directory. The container path + # (engine._build_container_command) already rewrites this to /ainode-models/; + # this keeps the host-venv path consistent so a slug like + # "qwen--qwen2.5-0.5b-instruct" loads from the models store instead of choking + # AutoTokenizer.from_pretrained on the '--' (HFValidationError). Only rewrites + # when a matching directory exists, else the hub repo id passes through. + bm = config.get("base_model", "") + if bm and "/" not in bm and not bm.startswith("/") and not bm.startswith("~"): + try: + from ainode.core.config import AINODE_HOME as _bm_home + _home = Path(_bm_home) + except Exception: + _home = Path(os.environ.get("AINODE_HOME", str(Path.home() / ".ainode"))) + _cand = _home / "models" / bm + if _cand.is_dir(): + config["base_model"] = str(_cand) + # Inject HF token if provided in config — needed for gated repos (Llama etc.) hf_token = config.get("hf_token") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get("HF_TOKEN") if hf_token: diff --git a/ainode/training/api_routes.py b/ainode/training/api_routes.py index 88b42e3..3e4692b 100644 --- a/ainode/training/api_routes.py +++ b/ainode/training/api_routes.py @@ -363,7 +363,13 @@ async def _merge_in_container() -> None: """Slim orchestrator (no peft/torch): spawn a GPU container to merge, streaming its stdout so the AINODE_PROGRESS protocol keeps working.""" from ainode.training.engine import build_merge_command - cmd = build_merge_command(merge_job, job.config.base_model, adapter_dir, merged_dir, job.config.hf_token) + # build_merge_command can shell out to a blocking `pip download` (peft + # wheel vendoring, up to 120s on a cold cache) — build it OFF the loop so + # a slow/unreachable network can't freeze the whole API server. + cmd = await loop.run_in_executor( + None, build_merge_command, merge_job, job.config.base_model, + adapter_dir, merged_dir, job.config.hf_token, + ) merge_job._log("Merge command: " + " ".join(cmd)) proc = await asyncio.create_subprocess_exec( *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT, diff --git a/ainode/training/engine.py b/ainode/training/engine.py index 9459946..002f43a 100644 --- a/ainode/training/engine.py +++ b/ainode/training/engine.py @@ -59,6 +59,156 @@ def _host_path(container_path: str) -> str: return container_path +def _loadable_dir(d: Path) -> Optional[Path]: + """Return the directory ``from_pretrained`` should actually load from, or None. + + A direct-download / flat dir holds ``config.json`` at its top level — load it + as-is. An HF-cache-format dir (``models--org--name``) instead nests the real + weights under ``snapshots//``; return that snapshot subdir (its relative + symlinks into ``../../blobs`` still resolve because the whole models tree is + mounted). Prefer a snapshot that actually carries a ``config.json``.""" + if not d.is_dir(): + return None + if (d / "config.json").exists(): + return d + snap = d / "snapshots" + if snap.is_dir(): + subs = sorted(s for s in snap.iterdir() if s.is_dir()) + for s in subs: + if (s / "config.json").exists(): + return s + if subs: + return subs[0] + return None + + +def _resolve_base_model_mount(base_model: str) -> Optional[str]: + """If ``base_model`` names a model already on disk under the models store, + return its CONTAINER mount path under ``/ainode-models/...``; else None. + + The training wizard's downloaded-model cards submit the ON-DISK slug (e.g. + ``qwen--qwen2.5-0.5b-instruct``). Handed straight to + ``AutoTokenizer.from_pretrained`` that raises HFValidationError ("Cannot have + -- or .. in repo_id") and the job dies instantly. Rewriting it to the mounted + directory path makes HF load from local weights (also offline-safe — no hub + round-trip). Accepts both the raw slug and a canonical HF repo id + (``Qwen/Qwen2.5-0.5B-Instruct``). + + Recognizes ALL FOUR on-disk layouts the registry tracks (mirrors + ``ModelManager._find_model_dir``): direct ``org--name`` (our downloader), flat + HF ``models--org--name``, HF cache ``hub/models--org--name``, and out-of-band + ``hf-cache/hub/models--org--name`` (HF_HOME downloads, e.g. from the eugr + distributed-serving backend). Missing the cache layouts silently fell back to + a live hub round-trip that fails on air-gapped nodes for models that ARE on + disk. Returns None for a plain hub repo id with no local copy so it passes + through to load from the hub as before.""" + if not base_model: + return None + models_root = AINODE_HOME / "models" + + def _to_mount(p: Path) -> Optional[str]: + try: + rel = p.relative_to(models_root) + except ValueError: + return None + return "/ainode-models/" + str(rel).replace(os.sep, "/") + + # Flat/direct forms: the raw slug, and (for a repo id) its org--name dir. + # Lenient — the weights of a direct download sit at the dir's top level, so a + # bare existing dir maps straight to its mount (matches the downloader layout + # even before any config.json probe). + flat_slug = base_model.replace("/", "--") + for slug in (base_model, flat_slug): + if not slug or "/" in slug or slug.startswith("."): + continue + d = models_root / slug + if d.is_dir(): + return _to_mount(_loadable_dir(d) or d) + + # HF-cache forms: models--org--name under the store root, hub/, and + # hf-cache/hub/. `flat_slug` is already org--name here (repo id or slug). + hf_slug = "models--" + flat_slug + for cache_dir in (models_root / hf_slug, + models_root / "hub" / hf_slug, + models_root / "hf-cache" / "hub" / hf_slug): + loadable = _loadable_dir(cache_dir) + if loadable is not None: + return _to_mount(loadable) + return None + + +def _vendor_wheel(pkg: str, job_dir: Path) -> Optional[str]: + """Ensure a wheel for ``pkg`` is available in ``job_dir`` (mounted at /job) so + the spawned container can ``pip install --no-index`` it with NO network. + + Wheels are cached once under ``AINODE_HOME/wheels`` and copied into each job + dir. When the cache is empty we fetch it with the orchestrator's own pip + (``pip download --no-deps``) — best-effort, short timeout — unless + ``AINODE_NO_WHEEL_FETCH`` is set (air-gapped nodes pre-seed the cache). + Returns the wheel FILENAME (basename) if vendored, else None so the caller + falls back to online pip. Only sensible for pure-python packages (peft).""" + import glob + import shutil + + cache = AINODE_HOME / "wheels" + norm = pkg.replace("-", "_") + + def _find(directory: Path) -> Optional[str]: + for pat in (f"{norm}-*.whl", f"{pkg}-*.whl"): + hits = sorted(glob.glob(str(directory / pat))) + if hits: + return hits[0] + return None + + try: + cache.mkdir(parents=True, exist_ok=True) + except Exception: + return None + + wheel = _find(cache) + if wheel is None and not os.environ.get("AINODE_NO_WHEEL_FETCH"): + try: + subprocess.run( + [sys.executable, "-m", "pip", "download", "--no-deps", + "--dest", str(cache), pkg], + capture_output=True, text=True, timeout=120, + ) + except Exception: + pass + wheel = _find(cache) + if wheel is None: + return None + try: + dest = job_dir / Path(wheel).name + if not dest.exists(): + shutil.copy2(wheel, dest) + return dest.name + except Exception: + return None + + +def _pip_install_step(pkg: str, job_dir: Path, *, vendor: bool) -> str: + """One tolerant shell step that makes ``pkg`` importable in the spawned + container. Import-guarded (a future baked image satisfies it with no install), + then — for vendored pure-python deps — an offline ``--no-index`` install from + the mounted wheel, falling back to online pip only if the wheel is absent. + Joined with ``;`` (never ``&&``) so a failed install never blocks the runner; + the runner itself reports a clean error if the dep is truly missing.""" + mod = pkg.replace("-", "_") + guard = f"python3 -c 'import {mod}' 2>/dev/null" + if vendor: + wheel = _vendor_wheel(pkg, job_dir) + if wheel: + install = (f"pip install -q --no-index --find-links /job /job/{wheel} " + f"|| pip install -q --no-deps {pkg}") + else: + install = f"pip install -q --no-deps {pkg}" + else: + # Network-only best-effort (e.g. bitsandbytes has no pure-python wheel). + install = f"pip install -q --no-deps {pkg}" + return f"{guard} || {install}" + + @dataclass class TrainingConfig: """Configuration for a training/fine-tuning job.""" @@ -213,8 +363,13 @@ async def start(self) -> None: config_path = self._job_dir / "config.json" config_path.write_text(json.dumps(self.config.to_dict(), indent=2)) - # Build the training command - cmd = self._build_command(config_path) + # Build the training command OFF the event loop. _build_command can + # shell out to a blocking `pip download` (peft wheel vendoring, up to a + # 120s timeout) on a cold cache — running that inline on aiohttp's single + # loop would freeze every concurrent request (live inference proxying + # included) until it returns. run_in_executor keeps the loop responsive. + loop = asyncio.get_event_loop() + cmd = await loop.run_in_executor(None, self._build_command, config_path) self._log(f"Command: {' '.join(cmd)}") try: @@ -424,6 +579,12 @@ def _build_container_command(self) -> list[str]: # via the /job mount) so handle_get_output/download resolve unchanged. container_cfg = dict(c.to_dict()) container_cfg["output_dir"] = "/job/output" + # base_model may be an on-disk slug (what the GUI submits) — rewrite it to + # the mounted weights path so AutoTokenizer.from_pretrained loads locally + # instead of raising HFValidationError on the '--' in the slug. + base_mount = _resolve_base_model_mount(c.base_model) + if base_mount: + container_cfg["base_model"] = base_mount datasets_dir = str(AINODE_HOME / "datasets") ds = c.dataset_path or "" if ds.startswith(datasets_dir): @@ -461,13 +622,17 @@ def _build_container_command(self) -> list[str]: # ponytail: peft (and bitsandbytes for qlora) aren't baked into the train # image yet — pip-shim them at launch. TODO(ponytail): bake peft + # bitsandbytes into the next quant/train-image build and drop this shim. - # bitsandbytes has no aarch64 wheel on some indexes; if the install fails the - # container exits non-zero and the pip error lands in the job log (readable). - pip_pkgs = "peft" + (" bitsandbytes" if c.method == "qlora" else "") + # peft is pure-python → vendor a wheel (offline-safe); bitsandbytes is + # network-only best-effort (no aarch64 pure-python wheel). Steps are ';' + # separated so a failed install never blocks the runner (a fatal '&&' here + # killed jobs on nodes with broken DNS — the merge/train couldn't pip peft). + steps = [_pip_install_step("peft", self._job_dir, vendor=True)] + if c.method == "qlora": + steps.append(_pip_install_step("bitsandbytes", self._job_dir, vendor=False)) + prep = " ; ".join(steps) cmd += [ TRAIN_IMAGE, "sh", "-c", - f"pip install -q --no-deps {pip_pkgs} && " - "python3 /job/_run_training.py --config /job/config.container.json", + f"{prep} ; python3 /job/_run_training.py --config /job/config.container.json", ] return cmd @@ -581,8 +746,10 @@ def build_merge_command( shutil.copy2(Path(__file__).parent / "_run_merge.py", job_dir / "_run_merge.py") token = hf_token or os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or "" + # Same slug→mount rewrite as training: a downloaded base passed as its on-disk + # slug must load from /ainode-models/, not choke AutoTokenizer on the '--'. merge_cfg = { - "base_model": base_model, + "base_model": _resolve_base_model_mount(base_model) or base_model, "adapter_dir": "/adapter", "output_dir": f"/out/{merged_dir.name}", "hf_token": token, @@ -602,10 +769,12 @@ def build_merge_command( if token: cmd += ["-e", f"HF_TOKEN={token}", "-e", f"HUGGING_FACE_HUB_TOKEN={token}"] # ponytail: peft pip-shim — bake it into the next train-image build and drop this. + # Vendor the (pure-python) peft wheel so the merge runs offline; ';' not '&&' + # so a pip hiccup never blocks the runner (broken DNS killed a live merge here). + peft_step = _pip_install_step("peft", job_dir, vendor=True) cmd += [ TRAIN_IMAGE, "sh", "-c", - "pip install -q --no-deps peft && " - "python3 /job/_run_merge.py --config /job/merge_config.json", + f"{peft_step} ; python3 /job/_run_merge.py --config /job/merge_config.json", ] return cmd diff --git a/ainode/web/static/js/app.js b/ainode/web/static/js/app.js index d63944c..cb11bd1 100644 --- a/ainode/web/static/js/app.js +++ b/ainode/web/static/js/app.js @@ -713,17 +713,38 @@ const AINode = { // the federated /v1/models the chat dropdown uses. Skip the distributed one. var seen = {}; (this.state.nodes || []).forEach(function (n) { + var host = n.hostname || n.node_id; var modelName = n.model; - if (!modelName || (distModel && modelName === distModel)) return; - var key = modelName + '@' + (n.node_id || n.hostname); - if (seen[key]) return; - seen[key] = true; - instances.push({ - model: modelName, - strategy: 'single', - nodes: [n.hostname || n.node_id], - status: n.engine_ready ? 'READY' : 'STARTING', - badge: 'SINGLE', + if (modelName && !(distModel && modelName === distModel)) { + var key = modelName + '@' + (n.node_id || n.hostname); + if (!seen[key]) { + seen[key] = true; + instances.push({ + model: modelName, + strategy: 'single', + nodes: [host], + status: n.engine_ready ? 'READY' : 'STARTING', + badge: 'SINGLE', + }); + } + } + // Stacked instances (2nd+ model on this node, ports 8001+) — invisible + // before D5. Render a card per stacked instance (model, port, status). + (n.instances || []).forEach(function (inst) { + var im = inst.model; + if (!im || (distModel && im === distModel)) return; + // Skip the primary (already shown above): same model on the node's main port. + if (im === n.model && (inst.api_port == null || inst.api_port === n.api_port)) return; + var skey = im + '@' + (n.node_id || n.hostname) + ':' + (inst.api_port || ''); + if (seen[skey]) return; + seen[skey] = true; + instances.push({ + model: im, + strategy: 'stacked', + nodes: [host + (inst.api_port ? ':' + inst.api_port : '')], + status: (inst.status === 'serving' || n.engine_ready) ? 'READY' : 'STARTING', + badge: 'STACKED' + (inst.api_port ? ' · :' + inst.api_port : ''), + }); }); }); @@ -860,12 +881,15 @@ const AINode = { return ''; }).join(''); nodeSelector.querySelectorAll('.node-dot').forEach(function (dot) { dot.addEventListener('click', function () { - if (dot.dataset.head) return; // head can't be deselected + // A single active node = solo load on THAT node (head or peer); 2+ = + // distributed. Keep at least one selected so a load always has a target. + if (dot.classList.contains('active') && + nodeSelector.querySelectorAll('.node-dot.active').length <= 1) return; dot.classList.toggle('active'); updateLaunchHint(); }); @@ -1085,6 +1109,10 @@ const AINode = { }); } + var gmuInput = document.getElementById('launch-gmu'); + var gmu = gmuInput && gmuInput.value !== '' ? parseFloat(gmuInput.value) : null; + if (gmu != null && isNaN(gmu)) gmu = null; + var launchBtn = document.getElementById('launch-btn'); if (launchBtn) { launchBtn.disabled = true; launchBtn.textContent = 'LAUNCHING...'; } @@ -1094,10 +1122,15 @@ const AINode = { // Distributed launch on the chosen nodes (head = this node + the rest). endpoint = '/api/sharding/launch'; body = { model: model, strategy: strategy, node_ids: nodeIds }; + if (gmu != null) body.gpu_memory_utilization = gmu; } else { - // Single node launch. - endpoint = '/api/models/load'; - body = { model: model }; + // Single node launch — route to the CHOSEN node via the cluster load + // route (node_id == this node dispatches locally), instead of always + // hitting the head's local engine regardless of the picked node. + endpoint = '/api/cluster/load'; + var target = nodeIds[0] || (this.state.status && this.state.status.node_id); + body = { model: model, node_id: target }; + if (gmu != null) body.gpu_memory_utilization = gmu; } var resp = await fetch(endpoint, { method: 'POST', @@ -3528,7 +3561,13 @@ const AINode = { if (modelsResp && Array.isArray(modelsResp.models)) { modelList = modelsResp.models .filter(function (m) { return m.downloaded || m.is_downloaded || m.status === 'downloaded'; }) - .map(function (m) { return { id: m.id || m.name || m.repo_id, name: m.name || m.id, size: m.size || '' }; }); + // base_model = the canonical HF repo id (cards know it), NOT the on-disk + // slug — the containerized runner hands base_model to + // AutoTokenizer.from_pretrained, which rejects a '--' slug (HFValidationError). + .map(function (m) { + var repo = m.hf_repo || m.id || m.name || m.repo_id; + return { id: repo, name: m.name || repo, size: m.size || m.size_gb || '' }; + }); } if (!modelList.length) modelList = this.trainingModels.slice(); diff --git a/ainode/web/templates/index.html b/ainode/web/templates/index.html index e4a2fe3..bf5d39a 100644 --- a/ainode/web/templates/index.html +++ b/ainode/web/templates/index.html @@ -163,6 +163,8 @@

◇ LAUNCH INSTANCE

Solo — runs on this node only (TP=1).
+ + diff --git a/tests/test_api.py b/tests/test_api.py index 26a546d..2a1335c 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -325,3 +325,63 @@ async def test_engine_update_skips_stop_when_unit_not_swappable(client, tmp_path if c.args and "stop" in c.args[0] ] assert stop_cmds == [] + + +# ---- D4: /api/cluster/load routes by node_id ------------------------------- + +@pytest.mark.asyncio +async def test_cluster_load_unknown_node_404(client): + """A node_id we don't know is honored (looked up) — not silently loaded local.""" + resp = await client.post("/api/cluster/load", json={"model": "m", "node_id": "ghost-node"}) + assert resp.status == 404 + data = await resp.json() + assert "ghost-node" in data["error"] + + +@pytest.mark.asyncio +async def test_cluster_load_local_dispatch_passes_gmu(client, app, monkeypatch): + """node_id == this node dispatches to the LOCAL model handler, forwarding gmu.""" + import ainode.models.api_routes as mr + + calls = {} + + def fake_append(app, model, gmu=None, *, overrides=None, persist=True): + calls["model"] = model + calls["gmu"] = gmu + return {"ok": True, "model": model, "instance_id": "n:x", + "api_port": 8000, "stacked": False} + + monkeypatch.setattr(mr, "append_solo_instance", fake_append) + resp = await client.post("/api/cluster/load", json={ + "model": "m", "node_id": "test-node-1", "gpu_memory_utilization": 0.4, + }) + assert resp.status == 200 + data = await resp.json() + assert data["model"] == "m" + assert calls["model"] == "m" + assert calls["gmu"] == 0.4 + + +# ---- D5: /api/nodes includes per-node instances ---------------------------- + +@pytest.mark.asyncio +async def test_nodes_include_stacked_instances(client, app): + from ainode.discovery.cluster import ClusterNode + from ainode.discovery.broadcast import NodeStatus + + cluster = app["cluster_state"] + cluster.add_node(ClusterNode( + node_id="n2", node_name="N2", gpu_name="GB10", gpu_memory_gb=128.0, + unified_memory=True, model="primary-model", status=NodeStatus.ONLINE, + api_port=8000, web_port=8080, last_seen=0.0, + instances=[ + {"model": "primary-model", "api_port": 8000, "status": "serving"}, + {"model": "stacked-model", "api_port": 8001, "status": "serving"}, + ], + )) + resp = await client.get("/api/nodes") + assert resp.status == 200 + data = await resp.json() + 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"]) diff --git a/tests/test_cluster_load.py b/tests/test_cluster_load.py index 00fb0ea..b08d399 100644 --- a/tests/test_cluster_load.py +++ b/tests/test_cluster_load.py @@ -181,7 +181,8 @@ def test_solo_load_appends_not_replaces(monkeypatch): "ray_autostart_state": None} asyncio.run(mr.handle_model_load(_Req(app, {"model": "model-A"}))) - asyncio.run(mr.handle_model_load(_Req(app, {"model": "model-B"}))) + # A stacked load must now name an explicit gpu_memory_utilization (D3). + asyncio.run(mr.handle_model_load(_Req(app, {"model": "model-B", "gpu_memory_utilization": 0.3}))) mgr = app["instances"] recs = {r.model: r for r in mgr.records()} @@ -226,7 +227,7 @@ def test_load_appends_when_boot_engine_already_seeded(monkeypatch): app = {"engine": boot, "instances": mgr, "config": cfg, "cluster_state": ClusterState(), "ray_autostart_state": None} - asyncio.run(mr.handle_model_load(_Req(app, {"model": "new-model"}))) + asyncio.run(mr.handle_model_load(_Req(app, {"model": "new-model", "gpu_memory_utilization": 0.3}))) recs = {r.model: r.api_port for r in mgr.records()} assert recs == {"boot-model": 8000, "new-model": 8001} # appended on 8001 @@ -338,9 +339,10 @@ def test_next_model_launch_config_does_not_inherit_primary_overrides(monkeypatch "served_model_name": "vl-alias", "trust_remote_code": True, }))) - # Bare stacked load of a DIFFERENT model B — no overrides supplied. + # Stacked load of a DIFFERENT model B — only its own explicit gmu (D3), no + # other overrides supplied; none of A's overrides may leak onto B. asyncio.run(mr.handle_model_load( - _Req(app, {"model": "meta-llama/Llama-3.2-3B-Instruct"}))) + _Req(app, {"model": "meta-llama/Llama-3.2-3B-Instruct", "gpu_memory_utilization": 0.3}))) b_cfg = made["backends"][-1].config defaults = NodeConfig() @@ -350,7 +352,7 @@ def test_next_model_launch_config_does_not_inherit_primary_overrides(monkeypatch assert b_cfg.quantization is None # NOT "awq" assert b_cfg.trust_remote_code is False # NOT True assert b_cfg.served_model_name is None # NOT ["vl-alias"] - assert b_cfg.gpu_memory_utilization == defaults.gpu_memory_utilization # NOT 0.6 + assert b_cfg.gpu_memory_utilization == 0.3 # B's own, NOT A's 0.6 def test_reload_same_primary_keeps_live_and_persisted_in_sync(monkeypatch): @@ -411,7 +413,7 @@ def test_unload_one_stacked_instance_leaves_the_other(monkeypatch): app = {"engine": None, "config": cfg, "cluster_state": ClusterState(), "ray_autostart_state": None} asyncio.run(mr.handle_model_load(_Req(app, {"model": "model-A"}))) - asyncio.run(mr.handle_model_load(_Req(app, {"model": "model-B"}))) + asyncio.run(mr.handle_model_load(_Req(app, {"model": "model-B", "gpu_memory_utilization": 0.3}))) mgr = app["instances"] b_backend = mgr.by_model("model-B").backend diff --git a/tests/test_models.py b/tests/test_models.py index 1f8eae8..b4576e7 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -433,3 +433,100 @@ async def test_refresh_catalog(self, app, aiohttp_client): data = await resp.json() assert data["status"] == "refreshed" assert "count" in data + + +# --------------------------------------------------------------------------- +# D3: stacked-load admission control (append_solo_instance) +# --------------------------------------------------------------------------- + +from types import SimpleNamespace # noqa: E402 + +from ainode.models.api_routes import append_solo_instance # noqa: E402 +from ainode.engine.instance_manager import InstanceManager # noqa: E402 +from ainode.discovery.instance import InstanceRecord # noqa: E402 +from ainode.core.config import NodeConfig # noqa: E402 + + +def _stacked_app(*insts): + """Build a minimal app dict with an InstanceManager pre-loaded with fake + instances. Each `insts` entry is (model, gmu, port).""" + config = NodeConfig(node_id="n1", api_port=8000, model=insts[0][0] if insts else None) + mgr = InstanceManager(base_port=8000) + for model, gmu, port in insts: + backend = MagicMock() + backend.config = SimpleNamespace(gpu_memory_utilization=gmu, distributed_mode="solo") + mgr.add(InstanceRecord(instance_id="n1:" + model, model=model, api_port=port), backend) + return {"config": config, "instances": mgr} + + +class TestStackedLoadAdmission: + def test_stacked_load_missing_gmu_rejected_400(self): + app = _stacked_app(("A", 0.5, 8000)) + res = append_solo_instance(app, "B", None) + assert res["ok"] is False + assert res["status"] == 400 + assert "gpu_memory_utilization" in res["error"] + + def test_stacked_load_overcommit_rejected_409(self): + app = _stacked_app(("A", 0.5, 8000)) + res = append_solo_instance(app, "B", 0.5) + assert res["ok"] is False + assert res["status"] == 409 + assert "0.50" in res["error"] # current total + assert "1.00" in res["error"] # projected total + + def test_stacked_load_under_cap_admitted(self): + app = _stacked_app(("A", 0.5, 8000)) + with patch("ainode.engine.backends.get_backend") as gb: + backend = MagicMock() + backend.start.return_value = True + gb.return_value = backend + res = append_solo_instance(app, "B", 0.4, persist=False) # 0.5 + 0.4 = 0.9 ≤ cap + assert res["ok"] is True + assert res["stacked"] is True + + def test_primary_load_unaffected_by_gmu_rule(self): + app = _stacked_app() # empty manager → primary/solo load + with patch("ainode.engine.backends.get_backend") as gb: + backend = MagicMock() + backend.start.return_value = True + gb.return_value = backend + res = append_solo_instance(app, "A", None, persist=False) # gmu None allowed for primary + assert res["ok"] is True + assert res["stacked"] is False + + # -- reload of an ALREADY-stacked model: admission runs BEFORE teardown ---- + # A rejection must leave the live instance intact (not kill-then-reject). + + def _reload_app(self): + """Primary A (0.5) + stacked B (0.3). Returns (app, B_backend).""" + app = _stacked_app(("A", 0.5, 8000), ("B", 0.3, 8001)) + b_backend = app["instances"].by_model("B").backend + return app, b_backend + + def test_reload_stacked_missing_gmu_400_does_not_stop_instance(self): + app, b_backend = self._reload_app() + res = append_solo_instance(app, "B", None) # omit gmu + assert res["ok"] is False and res["status"] == 400 + b_backend.stop.assert_not_called() # live model untouched + assert app["instances"].by_model("B") is not None # still registered + + def test_reload_stacked_overcommit_409_does_not_stop_instance(self): + app, b_backend = self._reload_app() + res = append_solo_instance(app, "B", 0.6) # 0.5 (A) + 0.6 = 1.10 > cap + assert res["ok"] is False and res["status"] == 409 + assert "0.50" in res["error"] # A's reservation only (B excluded) + b_backend.stop.assert_not_called() + assert app["instances"].by_model("B") is not None + + def test_reload_stacked_under_cap_excludes_own_reservation(self): + # A(0.5)+B(0.3) live; reloading B at 0.4 must count only A (0.5)+0.4=0.9 ≤ cap, + # NOT 0.5+0.3+0.4. Old B is torn down; a fresh B backend replaces it. + app, b_backend = self._reload_app() + with patch("ainode.engine.backends.get_backend") as gb: + new_backend = MagicMock() + new_backend.start.return_value = True + gb.return_value = new_backend + res = append_solo_instance(app, "B", 0.4, persist=False) + assert res["ok"] is True and res["stacked"] is True + b_backend.stop.assert_called_once() # old instance freed AFTER admission diff --git a/tests/test_sharding_launch.py b/tests/test_sharding_launch.py index 06744ae..3fcdb74 100644 --- a/tests/test_sharding_launch.py +++ b/tests/test_sharding_launch.py @@ -115,6 +115,28 @@ def test_unknown_node_id_is_rejected(): assert resp.status == 422 +def test_distributed_launch_honours_gpu_memory_utilization(): + # The #launch-gmu box is sent for BOTH branches; a TP>1 launch must apply the + # typed value to the distributed backend, not silently drop it to the 0.5 default. + config, resp = _run( + {"model": "m", "node_ids": ["head", "m1"], "gpu_memory_utilization": 0.4}, + [_member("m1", "10.100.0.13")], + ) + assert resp.status == 200 + assert _FakeBackend.last["config"].gpu_memory_utilization == 0.4 + + +def test_distributed_launch_ignores_out_of_range_gmu(): + # Junk / out-of-range input falls back to the config default rather than + # forwarding a nonsense fraction to vLLM. + config, resp = _run( + {"model": "m", "node_ids": ["head", "m1"], "gpu_memory_utilization": 5}, + [_member("m1", "10.100.0.13")], + ) + assert resp.status == 200 + assert _FakeBackend.last["config"].gpu_memory_utilization == NodeConfig().gpu_memory_utilization + + @pytest.mark.xfail(reason="F1 (0.4.26) removed proxy_to_vllm's 503 visible-loading guard " "because engine.ready is stale-False even while serving; a correct " "loading-state is an owned follow-up (restore once readiness is fixed).", diff --git a/tests/test_training_command.py b/tests/test_training_command.py index aee42a7..8873fd3 100644 --- a/tests/test_training_command.py +++ b/tests/test_training_command.py @@ -2,11 +2,18 @@ from __future__ import annotations +import json import sys from pathlib import Path from unittest.mock import patch -from ainode.training.engine import TrainingConfig, TrainingJob +from ainode.training import engine as tr_engine +from ainode.training.engine import ( + TrainingConfig, + TrainingJob, + _resolve_base_model_mount, + build_merge_command, +) def _job(**overrides) -> TrainingJob: @@ -18,6 +25,149 @@ def _job(**overrides) -> TrainingJob: return TrainingJob(TrainingConfig(**defaults)) +def _container_job(monkeypatch, tmp_path, **overrides) -> TrainingJob: + """A TrainingJob wired to a tmp AINODE_HOME for container-command tests.""" + monkeypatch.setenv("AINODE_HOST_HOME", str(tmp_path / "host")) + monkeypatch.setenv("AINODE_NO_WHEEL_FETCH", "1") # hermetic: no pip download + monkeypatch.setattr(tr_engine, "AINODE_HOME", tmp_path) + monkeypatch.setattr(tr_engine, "JOBS_DIR", tmp_path / "training" / "jobs") + defaults = dict(base_model="org/model", dataset_path="user/data.jsonl", method="lora") + defaults.update(overrides) + return TrainingJob(TrainingConfig(**defaults)) + + +def _seed_peft_wheel(tmp_path) -> str: + wheels = tmp_path / "wheels" + wheels.mkdir(parents=True, exist_ok=True) + name = "peft-0.11.0-py3-none-any.whl" + (wheels / name).write_bytes(b"fake-wheel") + return name + + +# ---- D1: on-disk base_model slug → mounted path ----------------------------- + +def test_container_command_rewrites_ondisk_slug_base_model(monkeypatch, tmp_path): + """The GUI submits the on-disk slug; the runner must load it from the mount, + not choke AutoTokenizer on the '--' (HFValidationError).""" + slug = "qwen--qwen2.5-0.5b-instruct" + (tmp_path / "models" / slug).mkdir(parents=True) + job = _container_job(monkeypatch, tmp_path, base_model=slug) + job._build_container_command() + cfg = json.loads((job._job_dir / "config.container.json").read_text()) + assert cfg["base_model"] == "/ainode-models/" + slug + + +def test_container_command_rewrites_repo_id_when_slug_on_disk(monkeypatch, tmp_path): + """A canonical HF repo id whose org--name dir exists on disk also maps to the mount.""" + slug = "Qwen--Qwen2.5-0.5B-Instruct" + (tmp_path / "models" / slug).mkdir(parents=True) + job = _container_job(monkeypatch, tmp_path, base_model="Qwen/Qwen2.5-0.5B-Instruct") + job._build_container_command() + cfg = json.loads((job._job_dir / "config.container.json").read_text()) + assert cfg["base_model"] == "/ainode-models/" + slug + + +def test_container_command_passes_through_hub_repo_id(monkeypatch, tmp_path): + """A plain hub repo id with no local copy is left untouched (loads from hub).""" + job = _container_job(monkeypatch, tmp_path, base_model="meta-llama/Llama-3.2-3B-Instruct") + job._build_container_command() + cfg = json.loads((job._job_dir / "config.container.json").read_text()) + assert cfg["base_model"] == "meta-llama/Llama-3.2-3B-Instruct" + + +# ---- D1b: on-disk layout coverage — all 4 layouts the registry tracks ------- + +def test_resolve_mount_hf_cache_hub_snapshot(monkeypatch, tmp_path): + """A model cached via the HF cache layout (hub/models--org--name/snapshots/ + ) — populated by the eugr distributed backend's HF_HOME=models_dir — + must resolve to the SNAPSHOT dir, not fall through to a live hub round-trip.""" + monkeypatch.setattr(tr_engine, "AINODE_HOME", tmp_path) + hf_slug = "models--Qwen--Qwen2.5-0.5B-Instruct" + snap = tmp_path / "models" / "hub" / hf_slug / "snapshots" / "abc123" + snap.mkdir(parents=True) + (snap / "config.json").write_text("{}") + assert (_resolve_base_model_mount("Qwen/Qwen2.5-0.5B-Instruct") + == "/ainode-models/hub/" + hf_slug + "/snapshots/abc123") + + +def test_resolve_mount_hf_cache_out_of_band_snapshot(monkeypatch, tmp_path): + """Out-of-band HF_HOME layout: hf-cache/hub/models--org--name/snapshots/.""" + monkeypatch.setattr(tr_engine, "AINODE_HOME", tmp_path) + hf_slug = "models--meta-llama--Llama-3.2-3B-Instruct" + snap = tmp_path / "models" / "hf-cache" / "hub" / hf_slug / "snapshots" / "deadbeef" + snap.mkdir(parents=True) + (snap / "config.json").write_text("{}") + assert (_resolve_base_model_mount("meta-llama/Llama-3.2-3B-Instruct") + == "/ainode-models/hf-cache/hub/" + hf_slug + "/snapshots/deadbeef") + + +def test_resolve_mount_flat_hf_models_snapshot(monkeypatch, tmp_path): + """Flat HF cache dir at the store root: models--org--name/snapshots/.""" + monkeypatch.setattr(tr_engine, "AINODE_HOME", tmp_path) + hf_slug = "models--Qwen--Qwen2.5-0.5B-Instruct" + snap = tmp_path / "models" / hf_slug / "snapshots" / "cafef00d" + snap.mkdir(parents=True) + (snap / "config.json").write_text("{}") + assert (_resolve_base_model_mount("Qwen/Qwen2.5-0.5B-Instruct") + == "/ainode-models/" + hf_slug + "/snapshots/cafef00d") + + +def test_resolve_mount_absent_returns_none(monkeypatch, tmp_path): + """A hub repo id with no local copy still passes through (None → load from hub).""" + monkeypatch.setattr(tr_engine, "AINODE_HOME", tmp_path) + (tmp_path / "models").mkdir(parents=True) + assert _resolve_base_model_mount("meta-llama/Llama-3.2-3B-Instruct") is None + + +# ---- D2: offline peft — vendored wheel + tolerant fallback ------------------ + +def test_container_command_vendors_peft_wheel_offline(monkeypatch, tmp_path): + wheel = _seed_peft_wheel(tmp_path) + job = _container_job(monkeypatch, tmp_path) + cmd = job._build_container_command() + sh = cmd[-1] + assert f"pip install -q --no-index --find-links /job /job/{wheel}" in sh + assert "|| pip install -q --no-deps peft" in sh # fallback if offline install fails + assert " ; python3 /job/_run_training.py" in sh # ';' not '&&' — never blocks runner + assert " && python3 /job/_run_training.py" not in sh + assert sh.startswith("python3 -c 'import peft'") # tolerant: skip if already baked in + assert (job._job_dir / wheel).exists() # wheel copied into the job dir + + +def test_container_command_falls_back_to_online_pip_without_wheel(monkeypatch, tmp_path): + """No vendored wheel + fetch disabled → plain online pip, still ';'-joined.""" + job = _container_job(monkeypatch, tmp_path) # no wheel seeded + cmd = job._build_container_command() + sh = cmd[-1] + assert "--no-index" not in sh + assert "pip install -q --no-deps peft" in sh + assert " ; python3 /job/_run_training.py" in sh + assert sh.startswith("python3 -c 'import peft'") + + +def test_container_command_qlora_adds_bitsandbytes_step(monkeypatch, tmp_path): + job = _container_job(monkeypatch, tmp_path, method="qlora") + sh = job._build_container_command()[-1] + assert "import bitsandbytes" in sh + assert "pip install -q --no-deps bitsandbytes" in sh + + +def test_merge_command_rewrites_slug_and_vendors_peft(monkeypatch, tmp_path): + slug = "qwen--qwen2.5-0.5b-instruct" + (tmp_path / "models" / slug).mkdir(parents=True) + wheel = _seed_peft_wheel(tmp_path) + merge_job = _container_job(monkeypatch, tmp_path) + cmd = build_merge_command( + merge_job, slug, tmp_path / "adapter", tmp_path / "out" / "merged" + ) + cfg = json.loads((merge_job._job_dir / "merge_config.json").read_text()) + assert cfg["base_model"] == "/ainode-models/" + slug + sh = cmd[-1] + assert f"pip install -q --no-index --find-links /job /job/{wheel}" in sh + assert " ; python3 /job/_run_merge.py" in sh + assert " && python3 /job/_run_merge.py" not in sh + + def test_lora_single_gpu_uses_plain_python(): """LoRA on one GPU does NOT need torchrun — it bloats logs.""" job = _job(method="lora")