diff --git a/lab/README.md b/lab/README.md index a8a5cc6..f066bff 100644 --- a/lab/README.md +++ b/lab/README.md @@ -89,7 +89,7 @@ unchanged in stats). ## GPU arm One persistent VM, started and stopped by hand, on the cloud `LAB_VM_PROVIDER` names: -`verda` (default; 1x H200, REST API) or `crusoe` (1x A100 PCIe, `crusoe` CLI). Same +`verda` (default; 1x H200, REST API), `nebius` (1x H200, `nebius` CLI) or `crusoe` (1x A100 PCIe, `crusoe` CLI). Same commands either way: python -m lab.vm status @@ -106,7 +106,13 @@ here is what gets measured, secrets and weights stay home), runs the command in VM is left running; `stop` is yours. Do not run `uv sync` on the box: it would put CPU torch back. Credentials never live in the repo. Verda: `VERDA_CLIENT_ID` and `VERDA_CLIENT_SECRET` (console > -Keys > Cloud API credentials) in your shell. Crusoe: `crusoe config init`. +Keys > Cloud API credentials) in your shell. Nebius: `nebius profile create`. Crusoe: `crusoe config init`. + +Nebius: `stop` is a real stop (boot disk kept, GPU billing ends) and `start` resumes it. The VM is +created with a `lab` user from cloud-init (Nebius refuses root). Platform `gpu-h200-sxm`, preset +`1gpu-16vcpu-200gb` (`LAB_VM_PLATFORM`, `LAB_VM_TYPE`), image family `ubuntu24.04-cuda13.0` +(`LAB_VM_IMAGE`), the project's first subnet unless `LAB_VM_SUBNET` is set, `LAB_VM_PROJECT` for a +project other than the profile's. Verda's `stop` deletes the instance and its OS volume: Verda's hibernate hides the instance from its own API so it cannot be restored, and its shutdown keeps billing the GPU. Each session therefore diff --git a/lab/providers/__init__.py b/lab/providers/__init__.py index 571e21c..1ca3a87 100644 --- a/lab/providers/__init__.py +++ b/lab/providers/__init__.py @@ -39,4 +39,7 @@ def load(name: str): if name == "crusoe": from lab.providers import crusoe return crusoe.Crusoe() - raise ProviderError(f"unknown provider {name!r}; LAB_VM_PROVIDER is verda or crusoe") + if name == "nebius": + from lab.providers import nebius + return nebius.Nebius() + raise ProviderError(f"unknown provider {name!r}; LAB_VM_PROVIDER is verda, nebius or crusoe") diff --git a/lab/providers/nebius.py b/lab/providers/nebius.py new file mode 100644 index 0000000..70b0324 --- /dev/null +++ b/lab/providers/nebius.py @@ -0,0 +1,115 @@ +"""Nebius AI Cloud through the `nebius` CLI (auth: `nebius profile create`). A stopped VM keeps its boot disk; GPU billing ends.""" + +from __future__ import annotations + +import json +import os +import subprocess +from dataclasses import dataclass +from pathlib import Path + +from lab.providers import ProviderError, Record, env_field + +USER = "lab" # created by cloud-init; Nebius refuses root and admin as login names + + +@dataclass(frozen=True) +class Nebius: + type: str = env_field("LAB_VM_TYPE", "1gpu-16vcpu-200gb") # a preset of `platform` + platform: str = env_field("LAB_VM_PLATFORM", "gpu-h200-sxm") + location: str = env_field("LAB_VM_LOCATION", "") # the CLI profile's project decides; informational + image: str = env_field("LAB_VM_IMAGE", "ubuntu24.04-cuda13.0") # an image family; CUDA 13 driver, as on Verda (setup installs cu130 torch) + disk_gb: str = env_field("LAB_VM_DISK_GB", "400") + project: str = env_field("LAB_VM_PROJECT", "") # --parent-id; empty uses the profile's + subnet: str = env_field("LAB_VM_SUBNET", "") # empty: the project's first subnet + default_user: str = USER + stop_word: str = "stop" + + def _cli(self, *args: str) -> str: + argv = ["nebius", *args] + if self.project and args[:1] in (("compute",), ("vpc",)) and args[2:3] in (("list",), ("create",)): + argv += ["--parent-id", self.project] + try: + out = subprocess.run(argv, capture_output=True, text=True) + except FileNotFoundError: + raise ProviderError("`nebius` CLI not installed: https://docs.nebius.com/cli/install") + if out.returncode != 0: + raise ProviderError(f"nebius {' '.join(args[:4])}: {out.stderr.strip() or out.stdout.strip()}") + return out.stdout + + def _json(self, *args: str): + text = self._cli(*args, "--format", "json") + try: + return json.loads(text) if text.strip() else {} + except json.JSONDecodeError: + raise ProviderError(f"nebius {' '.join(args[:4])}: not JSON: {text[:200]}") + + def _raw(self, name: str) -> dict | None: + items = self._json("compute", "instance", "list").get("items") or [] + return next((i for i in items if (i.get("metadata") or {}).get("name") == name), None) + + def get(self, name: str) -> Record | None: + raw = self._raw(name) + if raw is None: + return None + status = raw.get("status") or {} + state = str(status.get("state", "")).lower() + ip = None + for nic in status.get("network_interfaces") or []: + addr = (nic.get("public_ip_address") or {}).get("address") + if addr: + ip = addr.split("/")[0] + break + res = (raw.get("spec") or {}).get("resources") or {} + kind = f"{res.get('platform', '')}/{res.get('preset', '')}".strip("/") + return Record(name=name, state=state, ip=ip, type=kind, detail=state) + + def _subnet(self) -> str: + if self.subnet: + return self.subnet + items = self._json("vpc", "subnet", "list").get("items") or [] + if not items: + raise ProviderError("no VPC subnet in this project; set LAB_VM_SUBNET") + return items[0]["metadata"]["id"] + + def create(self, name: str, keyfile: str) -> None: + try: + public = Path(os.path.expanduser(keyfile)).read_text().strip() + except OSError as e: + raise ProviderError(f"ssh public key {keyfile}: {e.strerror}; set LAB_VM_KEYFILE") + user_data = ("#cloud-config\nusers:\n" + f" - name: {USER}\n sudo: ALL=(ALL) NOPASSWD:ALL\n shell: /bin/bash\n" + f" ssh_authorized_keys:\n - {public}\n") + nics = json.dumps([{"name": "eth0", "ip_address": {}, "public_ip_address": {}, "subnet_id": self._subnet()}]) + self._cli("compute", "instance", "create", "--name", name, + "--resources-platform", self.platform, "--resources-preset", self.type, + "--boot-disk-managed-disk-name", f"{name}-boot", + "--boot-disk-managed-disk-type", "network_ssd", + "--boot-disk-managed-disk-size-gibibytes", str(self.disk_gb), + "--boot-disk-managed-disk-source-image-family-image-family", self.image, + "--boot-disk-attach-mode", "READ_WRITE", + "--network-interfaces", nics, "--cloud-init-user-data", user_data) + + def _id(self, name: str) -> str: + raw = self._raw(name) + if raw is None: + raise ProviderError(f"{name}: no instance") + return raw["metadata"]["id"] + + def start(self, name: str) -> None: + self._cli("compute", "instance", "start", self._id(name)) + + def stop(self, name: str) -> None: + self._cli("compute", "instance", "stop", self._id(name)) + + def delete(self, name: str) -> None: + """Instance and its managed boot disk; `lab.vm` never calls this, it is for ending for good.""" + self._cli("compute", "instance", "delete", self._id(name)) + + def types(self) -> str: + items = self._json("compute", "platform", "list").get("items") or [] + lines = [] + for p in items: + presets = [x.get("name") for x in (p.get("spec") or {}).get("presets") or []] + lines.append(f"{(p.get('metadata') or {}).get('name')}: {', '.join(filter(None, presets))}") + return "\n".join(lines) + "\n" diff --git a/tests/test_lab_nebius.py b/tests/test_lab_nebius.py new file mode 100644 index 0000000..18d4956 --- /dev/null +++ b/tests/test_lab_nebius.py @@ -0,0 +1,116 @@ +"""lab.providers.nebius against a fake `nebius` CLI on PATH: lifecycle, ip parsing, cloud-init key.""" + +from __future__ import annotations + +import json +import os +import stat + +import pytest + +from lab import vm as labvm +from lab.providers import ProviderError, nebius + +FAKE = r'''#!/usr/bin/env python3 +import json, os, sys +args = sys.argv[1:] +open(os.environ["FAKE_LOG"], "a").write(json.dumps(args) + "\n") +p = os.environ["FAKE_STATE"] +db = json.load(open(p)) if os.path.exists(p) else {} +def save(): json.dump(db, open(p, "w")) +def inst(n, v): + return {"metadata": {"id": v["id"], "name": n}, + "spec": {"resources": {"platform": "gpu-h200-sxm", "preset": "1gpu-16vcpu-200gb"}}, + "status": {"state": v["state"], "network_interfaces": [{"public_ip_address": {"address": "203.0.113.9/32"}}]}} +if args[:3] == ["compute", "instance", "list"]: + print(json.dumps({"items": [inst(n, v) for n, v in db.items()]})) +elif args[:3] == ["compute", "instance", "create"]: + n = args[args.index("--name") + 1] + db[n] = {"id": "computeinstance-" + n, "state": "RUNNING", "user_data": args[args.index("--cloud-init-user-data") + 1]}; save() +elif args[:3] in (["compute", "instance", "stop"], ["compute", "instance", "start"], ["compute", "instance", "delete"]): + n = next(k for k, v in db.items() if v["id"] == args[3]) + if args[2] == "delete": db.pop(n) + else: db[n]["state"] = "STOPPED" if args[2] == "stop" else "RUNNING" + save() +elif args[:3] == ["vpc", "subnet", "list"]: + print(json.dumps({"items": [{"metadata": {"id": "vpcsubnet-1"}}]})) +elif args[:3] == ["compute", "platform", "list"]: + print(json.dumps({"items": [{"metadata": {"name": "gpu-h200-sxm"}, "spec": {"presets": [{"name": "1gpu-16vcpu-200gb"}]}}]})) +else: + sys.stderr.write("unknown " + " ".join(args)); sys.exit(2) +''' + + +@pytest.fixture +def fake(tmp_path, monkeypatch): + b = tmp_path / "bin" + b.mkdir() + f = b / "nebius" + f.write_text(FAKE) + f.chmod(f.stat().st_mode | stat.S_IEXEC) + monkeypatch.setenv("PATH", f"{b}{os.pathsep}{os.environ['PATH']}") + monkeypatch.setenv("FAKE_STATE", str(tmp_path / "state.json")) + monkeypatch.setenv("FAKE_LOG", str(tmp_path / "calls.log")) + monkeypatch.setenv("LAB_VM_PROVIDER", "nebius") + key = tmp_path / "key.pub" + key.write_text("ssh-ed25519 AAAAneb lab@mac\n") + monkeypatch.setenv("LAB_VM_KEYFILE", str(key)) + monkeypatch.setattr(labvm, "KNOWN_HOSTS_DIR", tmp_path / "kh") + monkeypatch.setattr(labvm, "POLL_S", 0.0) + return tmp_path + + +def calls(tmp): + return [json.loads(line) for line in (tmp / "calls.log").read_text().splitlines()] + + +def test_create_runs_with_h200_preset_cloud_init_key_and_subnet(fake): + rec = labvm.start(labvm.VM(name="g")) + assert rec.state == "running" and rec.ip == "203.0.113.9" and rec.type == "gpu-h200-sxm/1gpu-16vcpu-200gb" + create = next(c for c in calls(fake) if c[:3] == ["compute", "instance", "create"]) + assert create[create.index("--resources-platform") + 1] == "gpu-h200-sxm" + assert create[create.index("--resources-preset") + 1] == "1gpu-16vcpu-200gb" + assert create[create.index("--boot-disk-managed-disk-size-gibibytes") + 1] == "400" + assert json.loads(create[create.index("--network-interfaces") + 1])[0]["subnet_id"] == "vpcsubnet-1" + ud = json.load(open(fake / "state.json"))["g"]["user_data"] + assert ud.startswith("#cloud-config") and "name: lab" in ud and "ssh-ed25519 AAAAneb" in ud + + +def test_stop_keeps_the_vm_and_start_resumes_it(fake): + vm = labvm.VM(name="g") + labvm.start(vm) + labvm.stop(vm) + assert vm.provider.get("g").state == "stopped" + labvm.start(vm) + kinds = [c[2] for c in calls(fake) if c[:2] == ["compute", "instance"]] + assert kinds.count("create") == 1 and kinds.count("stop") == 1 and kinds.count("start") == 1 + + +def test_login_is_the_cloud_init_user(fake): + vm = labvm.VM(name="g") + assert labvm.ssh_target(vm, labvm.start(vm)) == "lab@203.0.113.9" + + +def test_types_lists_platform_presets(fake): + assert "gpu-h200-sxm: 1gpu-16vcpu-200gb" in nebius.Nebius().types() + + +def test_delete_removes_instance(fake): + labvm.start(labvm.VM(name="g")) + nebius.Nebius().delete("g") + assert nebius.Nebius().get("g") is None + + +def test_missing_cli_and_missing_key(tmp_path, monkeypatch): + monkeypatch.setenv("PATH", str(tmp_path)) + with pytest.raises(ProviderError, match="not installed"): + nebius.Nebius().get("g") + with pytest.raises(ProviderError, match="LAB_VM_KEYFILE"): + nebius.Nebius(subnet="s").create("g", "/nope/key.pub") + + +def test_project_is_passed_as_parent_id_on_list_and_create(fake, monkeypatch): + monkeypatch.setenv("LAB_VM_PROJECT", "project-x") + labvm.start(labvm.VM(name="g")) + scoped = [c for c in calls(fake) if c[2:3] in (["list"], ["create"])] + assert scoped and all(c[c.index("--parent-id") + 1] == "project-x" for c in scoped)