Skip to content
Merged
Show file tree
Hide file tree
Changes from 18 commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions api/install.sh
Original file line number Diff line number Diff line change
Expand Up @@ -593,6 +593,18 @@ multiuser_setup() {
echo "Installing SkyPilot with Kubernetes support..."
uv pip install --python "${GENERAL_UV_ENV_DIR}/bin/python" "skypilot[kubernetes]==0.12.1"

# Nebius AI Cloud CLI (used by the Nebius compute provider; default install path ~/.nebius/bin/nebius).
if command -v curl &> /dev/null; then
if [[ -x "${HOME}/.nebius/bin/nebius" ]]; then
ohai "Nebius CLI already installed at ${HOME}/.nebius/bin/nebius"
else
ohai "Installing Nebius CLI..."
curl -sSL https://storage.eu-north1.nebius.cloud/cli/install.sh | bash
fi
else
warn "curl is not available; skipping Nebius CLI install. Install manually if you use Nebius providers."
fi

echo "Multiuser setup complete."

}
Expand Down
1 change: 1 addition & 0 deletions api/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ dependencies = [
"fastapi==0.125.0",
"fastapi-users[sqlalchemy,oauth]==15.0.5",
"httpx-oauth==0.16.1",
"nebius==0.3.65",
"packaging==26.2",
"psutil==7.2.2",
"pydantic>=2.13.4,<3.0",
Expand Down
71 changes: 71 additions & 0 deletions api/test/compute_providers/test_nebius_subnet.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
"""Tests for Nebius automatic subnet resolution."""

from __future__ import annotations

from unittest.mock import MagicMock

import pytest

from transformerlab.compute_providers.models import ClusterConfig
from transformerlab.compute_providers.nebius import NEBIUS_AUTO_SUBNET_NAME, NebiusProvider


def test_resolve_subnet_uses_configured_id() -> None:
p = NebiusProvider(team_id="t1", parent_id="proj", subnet_id="subnet-fixed")
assert p._resolve_subnet_id_for_launch() == "subnet-fixed"


def test_resolve_subnet_reuses_existing_list() -> None:
p = NebiusProvider(team_id="t1", parent_id="proj-a", subnet_id=None)
p._run_nebius = MagicMock(return_value={"subnets": [{"metadata": {"id": "sn-1", "name": "s1"}}]})
assert p._resolve_subnet_id_for_launch() == "sn-1"
p._run_nebius.assert_called_once()
assert p._resolved_subnet_id == "sn-1"


def test_resolve_subnet_creates_network_then_subnet() -> None:
p = NebiusProvider(team_id="t1", parent_id="proj-b", subnet_id=None)
subnet_list_returns: list = [[], []]

def fake_run(args: list, stdin_json=None, timeout=120):
joined = " ".join(args)
if "subnet list" in joined:
return {"subnets": subnet_list_returns.pop(0)}
if "network list" in joined:
return {"networks": []}
if "create-default" in joined:
return {"metadata": {"id": "net-new"}}
if "subnet create" in joined:
return {"metadata": {"id": "sn-created"}}
raise AssertionError(f"unexpected nebius call: {joined}")

p._run_nebius = MagicMock(side_effect=fake_run)
assert p._resolve_subnet_id_for_launch() == "sn-created"
create_calls = [c for c in p._run_nebius.call_args_list if "subnet create" in " ".join(c[0][0])]
assert len(create_calls) == 1
create_args = create_calls[0][0][0]
assert NEBIUS_AUTO_SUBNET_NAME in create_args


def test_resolve_subnet_requires_parent_when_no_subnet() -> None:
p = NebiusProvider(team_id="t1", parent_id=None, subnet_id=None)
with pytest.raises(ValueError, match="parent_id"):
p._resolve_subnet_id_for_launch()


def test_cloud_init_includes_self_termination_trap() -> None:
p = NebiusProvider(team_id="t1")
user_data = p._build_cloud_init("cluster-a", ClusterConfig(run="echo hello"), public_key="ssh-ed25519 AAAATEST")
assert "_tfl_self_terminate" in user_data
assert "trap _tfl_self_terminate EXIT" in user_data


def test_cloud_init_self_termination_fallbacks() -> None:
p = NebiusProvider(team_id="t1")
user_data = p._build_cloud_init("cluster-a", ClusterConfig(run="echo hello"), public_key="ssh-ed25519 AAAATEST")
assert "_tfl_self_delete_instance" in user_data
assert "_tfl_ensure_nebius_cli" in user_data
assert "https://storage.eu-north1.nebius.cloud/cli/install.sh | bash" in user_data
assert 'export PATH="$HOME/.nebius/bin:$PATH"' in user_data
assert 'nebius compute instance delete --id "$_iid" --async' in user_data
assert "shutdown -h now || poweroff || true" in user_data
34 changes: 34 additions & 0 deletions api/transformerlab/compute_providers/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,16 @@ class ComputeProviderConfig(BaseModel):
service_account_json: Optional[Dict[str, Any]] = None
service_account_email: Optional[str] = None

# Nebius-specific config
nebius_profile: Optional[str] = None # Nebius CLI profile name
nebius_config_path: Optional[str] = None # Provider-scoped Nebius CLI config path
parent_id: Optional[str] = None # Nebius project/parent ID
subnet_id: Optional[str] = None # Nebius VPC subnet ID
default_platform: Optional[str] = None # e.g. "gpu-h100-sxm" or "cpu-d3"
default_preset: Optional[str] = None # e.g. "1gpu-16vcpu-200gb"
boot_image_family: Optional[str] = None # e.g. "ubuntu24.04-cuda13.0"
disk_size_gib: Optional[int] = None

# Additional provider-specific config
extra_config: Dict[str, Any] = Field(default_factory=dict)

Expand Down Expand Up @@ -265,6 +275,30 @@ def create_compute_provider(config: ComputeProviderConfig) -> "ComputeProvider":
team_id=config.team_id,
extra_config=config.extra_config,
)
elif config.type == "nebius":
from .nebius import NebiusProvider

if not config.team_id:
raise ValueError("Nebius provider requires team_id in config")
if not config.subnet_id and not config.parent_id:
raise ValueError(
"Nebius provider requires parent_id (Nebius project) for automatic default network/subnet, "
"or subnet_id if you use your own VPC subnet"
)
return NebiusProvider(
team_id=config.team_id,
parent_id=config.parent_id,
subnet_id=config.subnet_id,
profile=config.nebius_profile,
config_path=config.nebius_config_path,
region=config.region,
default_platform=config.default_platform,
default_preset=config.default_preset,
boot_image_family=config.boot_image_family,
disk_size_gib=config.disk_size_gib,
ssh_user=config.ssh_user,
extra_config=config.extra_config,
)
elif config.type == "vastai":
from .vastai import VastAIProvider

Expand Down
Loading
Loading