fix(training,models,web): slug mapping, offline peft, stacked-load admission guard, node-targeted launch, stacked visibility

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

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()

ainode/models/api_routes.py

@@ -209,18 +209,58 @@ 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()
         except Exception:
             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}"

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/<slug>;
+    # 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:

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,

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/<hash>/``; 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/<slug>, 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