Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from typing import ClassVar, Literal
from urllib.parse import urlsplit

import httpx
import typer
import yaml
from nemo_experimentalist_plugin.client import make_client
Expand Down Expand Up @@ -47,6 +48,7 @@
load_env_file,
resolve_base_url,
)
from nemo_platform import NeMoPlatformError
from nemo_platform_plugin.cli import NemoCLI

DEFAULT_WORKSPACE = "default"
Expand All @@ -59,6 +61,8 @@
# Tests monkeypatch this global with a recorder, which bypasses the lazy import.
run_experimentalist = None

_PLATFORM_CLIENT_ERRORS = (NeMoPlatformError, httpx.HTTPError, OSError, RuntimeError, ValueError)

# TODO: Add remote train/validation dataset support when remote experiment mode is implemented.


Expand Down Expand Up @@ -420,6 +424,7 @@ def doctor(
base_url=base_url_resolved,
probes=_PREFLIGHT_PROBES,
)
results.append(asyncio.run(_check_platform_client_bootstrap(base_url_resolved)))
if profile_obj is not None:
if plan is not None:
results += check_artifacts(
Expand All @@ -438,6 +443,39 @@ def doctor(
return app


async def _check_platform_client_bootstrap(base_url: str) -> CheckResult:
"""Verify doctor can construct the same Platform client used by a run."""
try:
client = make_client(base_url)
except _PLATFORM_CLIENT_ERRORS as exc:
return CheckResult(
name="platform-client-bootstrap",
group="platform",
status="fail",
severity="required",
message=f"Platform client initialization failed ({type(exc).__name__})",
hint="check --base-url/NMP_BASE_URL and the active authentication context",
)
try:
await client.close()
except _PLATFORM_CLIENT_ERRORS as exc:
return CheckResult(
name="platform-client-bootstrap",
group="platform",
status="fail",
severity="required",
message=f"Platform client cleanup failed ({type(exc).__name__})",
hint="check the active authentication context and retry",
)
return CheckResult(
name="platform-client-bootstrap",
group="platform",
status="pass",
severity="required",
message="Platform client initialized with the effective authentication path",
)


def _load_profile_or_error(profile_path: Path | None) -> tuple[AgentProfile | None, str | None]:
"""Load the explicit or discovered profile; announce a discovered one.

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from urllib.parse import urlparse

from nemo_platform import AsyncNeMoPlatform
from nemo_platform.auth.helpers import discover_nmp_config
from nemo_platform.config.config import Config

# Loopback hosts are served by an unauthenticated local platform; attaching
Expand All @@ -33,8 +34,10 @@ def make_client(base_url: str | None) -> AsyncNeMoPlatform:

- No ``base_url``: use the active nmp context for both URL and auth.
- Loopback ``base_url``: direct mode (local platform is unauthenticated).
- Remote ``base_url`` with an nmp config present: combine the URL with the
context's auth so the SDK injects and refreshes a Bearer token.
- Authenticated remote ``base_url`` with an nmp config present: combine the
URL with the context's auth so the SDK injects and refreshes a Bearer token.
- Unauthenticated remote ``base_url``: direct mode, even when an unrelated
OAuth context exists locally.
- Remote ``base_url`` without an nmp config: direct mode (no credentials to
use; the request will surface a clear auth error).
"""
Expand All @@ -46,4 +49,7 @@ def make_client(base_url: str | None) -> AsyncNeMoPlatform:
if host in LOOPBACK_HOSTS or not config_path.exists():
return AsyncNeMoPlatform(base_url=base_url)

if not discover_nmp_config(base_url).auth_enabled:
return AsyncNeMoPlatform(base_url=base_url)

return AsyncNeMoPlatform(base_url=base_url, config_path=config_path)
36 changes: 36 additions & 0 deletions plugins/nemo-experimentalist/tests/test_cli_profile.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from dataclasses import dataclass
from pathlib import Path

import httpx
import pytest
from nemo_experimentalist_plugin import cli
from nemo_experimentalist_plugin.preflight import Probes
Expand Down Expand Up @@ -304,6 +305,41 @@ def test_doctor_healthy_exits_zero(app, profile_tree: Path, monkeypatch) -> None
assert "✓" in result.output


def test_doctor_fails_when_run_client_bootstrap_fails(app, profile_tree: Path, monkeypatch) -> None:
write_task_toml(profile_tree)

def fail_client(_base_url: str) -> FakePlatformClient:
raise httpx.UnsupportedProtocol("invalid cached auth context")

monkeypatch.setattr(cli, "make_client", fail_client)
monkeypatch.chdir(profile_tree)

result = runner.invoke(app, ["doctor", "--base-url", "https://platform.example"])

assert result.exit_code == 1
assert "Platform client initialization failed (UnsupportedProtocol)" in result.output
assert "invalid cached auth context" not in result.output


def test_doctor_bootstraps_and_closes_run_client(app, profile_tree: Path, monkeypatch) -> None:
write_task_toml(profile_tree)
clients: list[FakePlatformClient] = []

def record_client(base_url: str) -> FakePlatformClient:
client = FakePlatformClient(base_url)
clients.append(client)
return client

monkeypatch.setattr(cli, "make_client", record_client)
monkeypatch.chdir(profile_tree)

result = runner.invoke(app, ["doctor", "--base-url", "https://platform.example"])

assert result.exit_code == 0, result.output
assert clients == [FakePlatformClient("https://platform.example", closed=True)]
assert "Platform client initialized with the effective authentication path" in result.output


def test_doctor_no_profile_exits_one_with_skeleton(app, tmp_path: Path, monkeypatch) -> None:
monkeypatch.setattr(cli, "_PREFLIGHT_PROBES", quiet_probes())
monkeypatch.chdir(tmp_path)
Expand Down
98 changes: 98 additions & 0 deletions plugins/nemo-experimentalist/tests/test_client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,98 @@
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from unittest.mock import MagicMock, patch

import pytest
from nemo_experimentalist_plugin.client import make_client
from nemo_platform.auth.helpers import NMPOIDCConfig

REMOTE_URL = "https://nemo-platform.example.com"


def test_no_base_url_uses_active_context() -> None:
with (
patch("nemo_experimentalist_plugin.client.discover_nmp_config") as discover,
patch("nemo_experimentalist_plugin.client.AsyncNeMoPlatform") as client_cls,
):
client = make_client(None)

discover.assert_not_called()
client_cls.assert_called_once_with()
assert client is client_cls.return_value


@pytest.mark.parametrize("host", ["localhost", "127.0.0.1", "::1", "0.0.0.0"])
def test_loopback_uses_direct_mode_without_auth_discovery(host: str) -> None:
config_path = MagicMock()
config_path.exists.return_value = True
base_url = f"http://[{host}]:8080" if host == "::1" else f"http://{host}:8080"

with (
patch("nemo_experimentalist_plugin.client.Config.get_default_config_path", return_value=config_path),
patch("nemo_experimentalist_plugin.client.discover_nmp_config") as discover,
patch("nemo_experimentalist_plugin.client.AsyncNeMoPlatform") as client_cls,
):
client = make_client(base_url)

discover.assert_not_called()
client_cls.assert_called_once_with(base_url=base_url)
assert client is client_cls.return_value


def test_remote_without_local_config_uses_direct_mode_without_auth_discovery() -> None:
config_path = MagicMock()
config_path.exists.return_value = False

with (
patch("nemo_experimentalist_plugin.client.Config.get_default_config_path", return_value=config_path),
patch("nemo_experimentalist_plugin.client.discover_nmp_config") as discover,
patch("nemo_experimentalist_plugin.client.AsyncNeMoPlatform") as client_cls,
):
client = make_client(REMOTE_URL)

discover.assert_not_called()
client_cls.assert_called_once_with(base_url=REMOTE_URL)
assert client is client_cls.return_value


def test_remote_no_auth_ignores_unrelated_local_oauth_context() -> None:
config_path = MagicMock()
config_path.exists.return_value = True

with (
patch("nemo_experimentalist_plugin.client.Config.get_default_config_path", return_value=config_path),
patch(
"nemo_experimentalist_plugin.client.discover_nmp_config",
return_value=NMPOIDCConfig(auth_enabled=False),
) as discover,
patch("nemo_experimentalist_plugin.client.AsyncNeMoPlatform") as client_cls,
):
client = make_client(REMOTE_URL)

discover.assert_called_once_with(REMOTE_URL)
client_cls.assert_called_once_with(base_url=REMOTE_URL)
assert client is client_cls.return_value


def test_remote_auth_uses_local_oauth_context() -> None:
config_path = MagicMock()
config_path.exists.return_value = True

with (
patch("nemo_experimentalist_plugin.client.Config.get_default_config_path", return_value=config_path),
patch(
"nemo_experimentalist_plugin.client.discover_nmp_config",
return_value=NMPOIDCConfig(
auth_enabled=True,
client_id="nemo-cli",
token_endpoint="https://auth.example.com/token",
),
) as discover,
patch("nemo_experimentalist_plugin.client.AsyncNeMoPlatform") as client_cls,
):
client = make_client(REMOTE_URL)

discover.assert_called_once_with(REMOTE_URL)
client_cls.assert_called_once_with(base_url=REMOTE_URL, config_path=config_path)
assert client is client_cls.return_value
Loading