diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/cli.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/cli.py index 3e324d6a77..54c1e5d443 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/cli.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/cli.py @@ -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 @@ -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" @@ -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. @@ -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( @@ -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. diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/client.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/client.py index 686d4cc332..cbaa1b8e07 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/client.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/client.py @@ -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 @@ -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). """ @@ -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) diff --git a/plugins/nemo-experimentalist/tests/test_cli_profile.py b/plugins/nemo-experimentalist/tests/test_cli_profile.py index c48ce1172e..bdb311d63e 100644 --- a/plugins/nemo-experimentalist/tests/test_cli_profile.py +++ b/plugins/nemo-experimentalist/tests/test_cli_profile.py @@ -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 @@ -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) diff --git a/plugins/nemo-experimentalist/tests/test_client.py b/plugins/nemo-experimentalist/tests/test_client.py new file mode 100644 index 0000000000..7c43e2bc58 --- /dev/null +++ b/plugins/nemo-experimentalist/tests/test_client.py @@ -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