diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/quickstart/cli.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/quickstart/cli.py index 03e1d59e0c..1ab5d00c06 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/quickstart/cli.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/quickstart/cli.py @@ -967,6 +967,9 @@ def _run_job_diagnostic(port: int, registry: str, tag: str, *, admin_email: str import uuid from nemo_platform import NeMoPlatform + from nemo_platform_plugin.client.adapter import client_from_platform + from nemo_platform_plugin.jobs.client import JobsClient + from nemo_platform_plugin.jobs.types import CreatePlatformJobRequest # When auth is enabled, use an unsigned JWT for the admin principal. default_headers = None @@ -997,29 +1000,32 @@ def _run_job_diagnostic(port: int, registry: str, tag: str, *, admin_email: str job_name = f"diagnostic-{uuid.uuid4().hex[:8]}" console.print(f" • Creating diagnostic job: {job_name}") - job = client.jobs.create( - platform_spec={ - "steps": [ - { - "name": "diagnostic", - "executor": { - "provider": "cpu", - "container": { - "image": cpu_image, - "entrypoint": [ - "python", - "-c", - "import sys; print(f'Python {sys.version}'); print('Job system is working correctly!')", - ], + jobs_client = client_from_platform(client, JobsClient) + job = jobs_client.create_job( + body=CreatePlatformJobRequest( + platform_spec={ + "steps": [ + { + "name": "diagnostic", + "executor": { + "provider": "cpu", + "container": { + "image": cpu_image, + "entrypoint": [ + "python", + "-c", + "import sys; print(f'Python {sys.version}'); print('Job system is working correctly!')", + ], + }, }, - }, - } - ] - }, - source="quickstart-doctor", - spec={}, - name=job_name, - ) + } + ] + }, + source="quickstart-doctor", + spec={}, + name=job_name, + ) + ).data() console.print(" • Waiting for job to complete...") @@ -1028,10 +1034,10 @@ def _run_job_diagnostic(port: int, registry: str, tag: str, *, admin_email: str poll_interval = 2 elapsed = 0 status = "pending" - job_status = client.jobs.retrieve(job.name) + job_status = jobs_client.get_job(name=job.name).data() while elapsed < max_wait: - job_status = client.jobs.retrieve(job.name) + job_status = jobs_client.get_job(name=job.name).data() status = job_status.status if status in ("completed", "error", "cancelled"): @@ -1058,9 +1064,9 @@ def _run_job_diagnostic(port: int, registry: str, tag: str, *, admin_email: str # Fetch and display logs console.print("\n [bold]Job output:[/bold]") try: - logs = client.jobs.get_logs(job.name) + logs = jobs_client.list_job_logs(name=job.name) log_lines = [] - for log_entry in logs: + for log_entry in logs.items(): if hasattr(log_entry, "message"): log_lines.append(log_entry.message) @@ -1075,7 +1081,7 @@ def _run_job_diagnostic(port: int, registry: str, tag: str, *, admin_email: str # Clean up the job (only if successful) if status == "completed": try: - client.jobs.delete(job.name) + jobs_client.delete_job(name=job.name) except Exception: pass # Ignore cleanup errors else: diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py index e62fb3443c..23747bb25c 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py @@ -21,6 +21,7 @@ def make_sync_resource(platform: NeMoPlatform) -> NemoClient: from nemo_platform import AsyncNeMoPlatform, NeMoPlatform from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient +from nemo_platform_plugin.client.types import RetryPolicy SyncT = TypeVar("SyncT", bound=NemoClient) AsyncT = TypeVar("AsyncT", bound=AsyncNemoClient) @@ -50,9 +51,23 @@ def client_from_platform( _skip = {"accept", "accept-encoding", "connection", "user-agent", "host"} headers = {k: v for k, v in platform._client.headers.items() if k.lower() not in _skip} # type: ignore[union-attr] + retry = RetryPolicy(max_retries=platform.max_retries) + if isinstance(platform, AsyncNeMoPlatform): + if not issubclass(client_cls, AsyncNemoClient): + raise TypeError("AsyncNeMoPlatform requires an AsyncNemoClient class") + return client_cls( + base_url=str(platform.base_url).rstrip("/"), + workspace=platform.workspace, + default_headers=headers or None, + retry=retry, + http_client=platform._client, + ) + if not issubclass(client_cls, NemoClient): + raise TypeError("NeMoPlatform requires a NemoClient class") return client_cls( base_url=str(platform.base_url).rstrip("/"), workspace=platform.workspace, default_headers=headers or None, - http_client=platform._client, # type: ignore[arg-type] + retry=retry, + http_client=platform._client, ) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py index f570aa2f07..daa3b90e68 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py @@ -21,9 +21,11 @@ import inspect import json import time -from collections.abc import Mapping +from collections.abc import AsyncIterator, Iterator, Mapping +from contextlib import asynccontextmanager, contextmanager +from functools import cache from pathlib import Path -from typing import Any, Self, TypeVar, get_args, get_origin, overload +from typing import Any, Self, TypeVar, cast, get_args, get_origin, overload from urllib.parse import quote import httpx @@ -31,7 +33,11 @@ StaticToken, TokenProvider, ) -from nemo_platform_plugin.client.errors import raise_for_status +from nemo_platform_plugin.client.errors import ( + NemoResponseValidationError, + NemoTransportError, + raise_for_status, +) from nemo_platform_plugin.client.response import ( AsyncNemoBinaryResponse, AsyncNemoPaginatedResponse, @@ -51,31 +57,57 @@ PreparedRequest, ResponseT, RetryPolicy, + StrategyT, Stream, ) -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter, ValidationError ModelT = TypeVar("ModelT", bound=BaseModel) DEFAULT_TIMEOUT = 60.0 -def _get_stream_model_type(response_type: type) -> type[BaseModel]: +@cache +def _type_adapter(response_type: type[ResponseT]) -> TypeAdapter[ResponseT]: + """Build each response annotation's validation schema once.""" + return TypeAdapter(response_type) + + +def _parse_json_body(response_type: type[ResponseT], data: object) -> ResponseT: + """Parse a decoded JSON body against an endpoint's return annotation. + + ``TypeAdapter`` handles both model classes and arbitrary annotations such as + ``list[Profile]`` while preserving the annotation's type for callers. + """ + return _type_adapter(response_type).validate_python(data) + + +def _parse_response_body(response_type: type[ResponseT], response: httpx.Response) -> ResponseT: + """Decode and validate a response, normalizing contract failures.""" + try: + return _parse_json_body(response_type, response.json()) + except (ValueError, ValidationError) as exc: + raise NemoResponseValidationError(response, exc) from exc + + +def _get_stream_model_type(response_type: type[Stream[ModelT]]) -> type[ModelT]: """Extract the ModelT from a Stream[ModelT] generic alias.""" args = get_args(response_type) if not args: raise TypeError(f"Stream response type must be parameterized, got {response_type}") - return args[0] + return cast(type[ModelT], args[0]) -def _get_paginated_types(response_type: type) -> tuple[type[BaseModel], type]: +def _get_paginated_types( + response_type: type[Paginated[ModelT, StrategyT]], +) -> tuple[type[ModelT], type[StrategyT]]: """Extract (ModelT, StrategyT) from a Paginated[ModelT, StrategyT] generic alias.""" args = get_args(response_type) if not args: raise TypeError(f"Paginated response type must be parameterized, got {response_type}") model_type = args[0] strategy_type = args[1] if len(args) > 1 else OffsetPagination - return model_type, strategy_type + return cast(type[ModelT], model_type), cast(type[StrategyT], strategy_type) # --------------------------------------------------------------------------- @@ -147,7 +179,7 @@ def __init__( self._auth: TokenProvider | None = StaticToken(auth) if isinstance(auth, str) else auth self._retry = retry self._default_headers = dict(default_headers) if default_headers else {} - self._timeout: float | None = None + self._timeout: float | httpx.Timeout | None = None @property def base_url(self) -> str: @@ -209,7 +241,7 @@ def with_options( *, headers: Mapping[str, str] | None = None, retry: RetryPolicy | None = None, - timeout: float | None = None, + timeout: float | httpx.Timeout | None = None, ) -> Self: """Return a copy of this client with the given options merged in. @@ -305,11 +337,11 @@ def send( @overload def send( self, - request: PreparedRequest[Paginated[ModelT, Any]], + request: PreparedRequest[Paginated[ModelT, StrategyT]], *, headers: dict[str, str] | None = None, retry: RetryPolicy | None = None, - ) -> NemoPaginatedResponse[ModelT]: ... + ) -> NemoPaginatedResponse[ModelT, StrategyT]: ... @overload def send( self, @@ -381,16 +413,12 @@ def send( resolved_retry = self._resolve_retry(retry) if self._is_binary(request): - stream_ctx = self._http.stream( - request.method, url, content=request.content, headers=req_headers, params=params - ) + stream_ctx = self._stream_with_retry(request, url, req_headers, params, resolved_retry) return NemoBinaryResponse(stream_ctx, request) if self._is_stream(request): assert request.response_type is not None - stream_ctx = self._http.stream( - request.method, url, content=request.content, headers=req_headers, params=params - ) + stream_ctx = self._stream_with_retry(request, url, req_headers, params, resolved_retry) model_type = _get_stream_model_type(request.response_type) return NemoStreamResponse(stream_ctx, model_type, request) @@ -409,7 +437,7 @@ def send( raise_for_status(raw) body = None if request.response_type is not None: - body = request.response_type.model_validate(raw.json()) + body = _parse_response_body(request.response_type, raw) return NemoResponse(http_response=raw, body=body, request=request) def _request_with_retry( @@ -433,7 +461,7 @@ def _request_with_retry( if backoff is not None: time.sleep(backoff) continue - raise + raise NemoTransportError(exc) from exc if retry: backoff = _should_retry(raw, None, attempt, retry) if backoff is not None: @@ -445,12 +473,45 @@ def _request_with_retry( assert last_response is not None return last_response + @contextmanager + def _stream_with_retry( + self, + request: PreparedRequest, + url: str, + headers: dict[str, str] | None, + params: dict | None, + retry: RetryPolicy | None, + ) -> Iterator[httpx.Response]: + """Open a stream, retrying failures before handing it to the caller.""" + for attempt in range(retry.max_retries + 1 if retry else 1): + yielded = False + try: + kwargs: dict = {"content": request.content, "headers": headers, "params": params} + if self._timeout is not None: + kwargs["timeout"] = self._timeout + with self._http.stream(request.method, url, **kwargs) as raw: + backoff = _should_retry(raw, None, attempt, retry) if retry else None + if backoff is not None: + time.sleep(backoff) + continue + yielded = True + yield raw + return + except httpx.TransportError as exc: + if yielded: + raise NemoTransportError(exc) from exc + backoff = _should_retry(None, exc, attempt, retry) if retry else None + if backoff is not None: + time.sleep(backoff) + continue + raise NemoTransportError(exc) from exc + def _make_page_fetcher( - self, strategy: type[PaginationStrategy], retry: RetryPolicy | None = None + self, strategy: type[PaginationStrategy[Any, Any]], retry: RetryPolicy | None = None ) -> SyncPageFetcher: """Create a page-fetching callback bound to this client and strategy.""" - def fetch(request: PreparedRequest, page: int | str) -> httpx.Response: + def fetch(request: PreparedRequest, page: Any) -> httpx.Response: url = self._resolve_path(request) req_headers = self._request_headers(request) existing_params = self._resolve_query_params(request) or {} @@ -517,11 +578,11 @@ async def send( @overload async def send( self, - request: PreparedRequest[Paginated[ModelT, Any]], + request: PreparedRequest[Paginated[ModelT, StrategyT]], *, headers: dict[str, str] | None = None, retry: RetryPolicy | None = None, - ) -> AsyncNemoPaginatedResponse[ModelT]: ... + ) -> AsyncNemoPaginatedResponse[ModelT, StrategyT]: ... @overload async def send( self, @@ -585,16 +646,12 @@ async def send( resolved_retry = self._resolve_retry(retry) if self._is_binary(request): - stream_ctx = self._http.stream( - request.method, url, content=request.content, headers=req_headers, params=params - ) + stream_ctx = self._stream_with_retry(request, url, req_headers, params, resolved_retry) return AsyncNemoBinaryResponse(stream_ctx, request) if self._is_stream(request): assert request.response_type is not None - stream_ctx = self._http.stream( - request.method, url, content=request.content, headers=req_headers, params=params - ) + stream_ctx = self._stream_with_retry(request, url, req_headers, params, resolved_retry) model_type = _get_stream_model_type(request.response_type) return AsyncNemoStreamResponse(stream_ctx, model_type, request) @@ -613,7 +670,7 @@ async def send( raise_for_status(raw) body = None if request.response_type is not None: - body = request.response_type.model_validate(raw.json()) + body = _parse_response_body(request.response_type, raw) return NemoResponse(http_response=raw, body=body, request=request) async def _request_with_retry( @@ -637,7 +694,7 @@ async def _request_with_retry( if backoff is not None: await asyncio.sleep(backoff) continue - raise + raise NemoTransportError(exc) from exc if retry: backoff = _should_retry(raw, None, attempt, retry) if backoff is not None: @@ -649,12 +706,45 @@ async def _request_with_retry( assert last_response is not None return last_response + @asynccontextmanager + async def _stream_with_retry( + self, + request: PreparedRequest, + url: str, + headers: dict[str, str] | None, + params: dict | None, + retry: RetryPolicy | None, + ) -> AsyncIterator[httpx.Response]: + """Open an async stream, retrying failures before handing it to the caller.""" + for attempt in range(retry.max_retries + 1 if retry else 1): + yielded = False + try: + kwargs: dict = {"content": request.content, "headers": headers, "params": params} + if self._timeout is not None: + kwargs["timeout"] = self._timeout + async with self._http.stream(request.method, url, **kwargs) as raw: + backoff = _should_retry(raw, None, attempt, retry) if retry else None + if backoff is not None: + await asyncio.sleep(backoff) + continue + yielded = True + yield raw + return + except httpx.TransportError as exc: + if yielded: + raise NemoTransportError(exc) from exc + backoff = _should_retry(None, exc, attempt, retry) if retry else None + if backoff is not None: + await asyncio.sleep(backoff) + continue + raise NemoTransportError(exc) from exc + def _make_page_fetcher( - self, strategy: type[PaginationStrategy], retry: RetryPolicy | None = None + self, strategy: type[PaginationStrategy[Any, Any]], retry: RetryPolicy | None = None ) -> AsyncPageFetcher: """Create an async page-fetching callback bound to this client and strategy.""" - async def fetch(request: PreparedRequest, page: int | str) -> httpx.Response: + async def fetch(request: PreparedRequest, page: Any) -> httpx.Response: url = self._resolve_path(request) req_headers = self._request_headers(request) existing_params = self._resolve_query_params(request) or {} diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/endpoint.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/endpoint.py index 2d3c77019d..3d9117b7a9 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/endpoint.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/endpoint.py @@ -34,7 +34,7 @@ def hello(self, *, name: str) -> HelloResponse: import inspect import string from collections.abc import AsyncIterable, Callable, Iterable -from typing import Any, get_type_hints +from typing import get_type_hints from nemo_platform_plugin.client.types import ( BLESSED_CLIENT_PARAMS, @@ -121,7 +121,7 @@ def _build_prepared_request( query_params: dict[str, str | int | bool | None] | None = None content: bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None content_type: str | None = None - client_options: dict[str, Any] | None = None + client_options: dict[str, object] | None = None body_model: BaseModel | None = None for name, value in bound.arguments.items(): diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/errors.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/errors.py index e8d67b1c9f..9062dc28fc 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/errors.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/errors.py @@ -1,11 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""HTTP error hierarchy for the NemoClient. +"""Error hierarchy for the NemoClient. -Provides :class:`NemoHTTPError` and status-code-specific subclasses -(e.g. :class:`NotFoundError`, :class:`ConflictError`) raised by -:func:`raise_for_status` on non-2xx responses. +All request/response failures derive from :class:`NemoClientError`. +Non-2xx responses additionally derive from :class:`NemoHTTPError` so callers +can distinguish an HTTP status from transport and response-validation errors. """ from __future__ import annotations @@ -13,7 +13,41 @@ import httpx -class NemoHTTPError(Exception): +class NemoClientError(Exception): + """Base class for failures while executing or parsing a client request.""" + + +class NemoTransportError(NemoClientError): + """Raised when the HTTP transport fails after retries are exhausted.""" + + def __init__(self, error: httpx.TransportError) -> None: + self.error = error + try: + self.request = error.request + except RuntimeError: + self.request = None + super().__init__(str(error)) + + +class NemoResponseValidationError(NemoClientError): + """Raised when a successful response does not match the endpoint contract.""" + + def __init__(self, http_response: httpx.Response, error: Exception) -> None: + self.http_response = http_response + self.status_code = http_response.status_code + self.body = self._extract_body(http_response) + self.error = error + super().__init__("Data returned by API is invalid for the expected schema") + + @staticmethod + def _extract_body(resp: httpx.Response) -> object | None: + try: + return resp.json() + except Exception: + return None + + +class NemoHTTPError(NemoClientError): """Raised on non-2xx HTTP responses. Attributes: diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/method.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/method.py index a167f7a8ff..58c9184feb 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/method.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/method.py @@ -29,8 +29,8 @@ class attributes (astral-sh/ty#3254). The types themselves are correct from __future__ import annotations import functools -from collections.abc import Callable -from typing import Any, Coroutine, Generic, TypeVar, overload +from collections.abc import Awaitable, Callable +from typing import Generic, TypeVar, overload from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient from nemo_platform_plugin.client.response import ( @@ -49,6 +49,7 @@ class attributes (astral-sh/ty#3254). The types themselves are correct Paginated, PreparedRequest, ResponseT, + StrategyT, Stream, ) @@ -76,9 +77,7 @@ def __init__(self, endpoint_fn: Callable[P, PreparedRequest]) -> None: @overload def __get__(self, obj: NemoClient, objtype: type | None = None) -> Callable[P, SyncReturnT]: ... @overload - def __get__( - self, obj: AsyncNemoClient, objtype: type | None = None - ) -> Callable[P, Coroutine[Any, Any, AsyncReturnT]]: ... + def __get__(self, obj: AsyncNemoClient, objtype: type | None = None) -> Callable[P, Awaitable[AsyncReturnT]]: ... def __get__(self, obj: NemoClient | AsyncNemoClient | None, objtype: type | None = None) -> object: assert obj is not None @@ -116,8 +115,12 @@ def method( @overload def method( - endpoint_fn: Callable[P, PreparedRequest[Paginated[ModelT, Any]]], -) -> EndpointMethod[P, NemoPaginatedResponse[ModelT], AsyncNemoPaginatedResponse[ModelT]]: ... + endpoint_fn: Callable[P, PreparedRequest[Paginated[ModelT, StrategyT]]], +) -> EndpointMethod[ + P, + NemoPaginatedResponse[ModelT, StrategyT], + AsyncNemoPaginatedResponse[ModelT, StrategyT], +]: ... @overload diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.py index 0bdc675709..fae9515f3c 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/response.py @@ -5,15 +5,42 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Callable, Coroutine, Iterator +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator from contextlib import AbstractAsyncContextManager, AbstractContextManager, asynccontextmanager, contextmanager from dataclasses import dataclass -from typing import Any, Generic, TypeVar +from typing import Any, Generic, TypeVar, overload import httpx -from nemo_platform_plugin.client.errors import raise_for_status +from nemo_platform_plugin.client.errors import NemoResponseValidationError, NemoTransportError, raise_for_status from nemo_platform_plugin.client.types import OffsetPagination, PaginationStrategy, PreparedRequest -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError +from typing_extensions import TypeVar as TypeVarExt + +ModelT = TypeVar("ModelT", bound=BaseModel) + + +def _validated_page( + response: httpx.Response, + model_type: type[ModelT], + strategy: type[PaginationStrategy[Any, Any]], +) -> tuple[list[ModelT], dict, Any]: + """Decode one page and normalize response-contract failures.""" + try: + body = response.json() + except ValueError as exc: + raise NemoResponseValidationError(response, exc) from exc + + if not isinstance(body, dict): + exc = ValueError("Paginated responses must be JSON objects") + raise NemoResponseValidationError(response, exc) from exc + + try: + raw_items = strategy.extract_items(body) + items = [model_type.model_validate(item) for item in raw_items] + metadata = strategy.extract_metadata(body) + except (KeyError, TypeError, ValueError, ValidationError) as exc: + raise NemoResponseValidationError(response, exc) from exc + return items, body, metadata def _parse_stream_line(line: str, headers: httpx.Headers) -> str | None: @@ -33,7 +60,6 @@ def _parse_stream_line(line: str, headers: httpx.Headers) -> str | None: ResponseT = TypeVar("ResponseT") -ModelT = TypeVar("ModelT", bound=BaseModel) @dataclass(frozen=True, slots=True) @@ -51,7 +77,7 @@ class NemoResponse(Generic[ResponseT]): http_response: httpx.Response body: ResponseT - request: PreparedRequest + request: PreparedRequest[ResponseT] def data(self) -> ResponseT: """Return the parsed response body. @@ -96,10 +122,13 @@ def http_response(self) -> httpx.Response: def read(self) -> bytes: """Read and return the entire response body as bytes.""" - with self._stream_ctx as raw: - data = raw.read() - raise_for_status(raw) - return data + try: + with self._stream_ctx as raw: + data = raw.read() + raise_for_status(raw) + return data + except httpx.TransportError as exc: + raise NemoTransportError(exc) from exc @contextmanager def stream(self, chunk_size: int | None = None) -> Iterator[Iterator[bytes]]: @@ -117,10 +146,13 @@ def stream(self, chunk_size: int | None = None) -> Iterator[Iterator[bytes]]: for chunk in chunks: ... """ - with self._stream_ctx as raw: - self._http_response = raw - raise_for_status(raw) - yield raw.iter_raw(chunk_size) if chunk_size else raw.iter_raw() + try: + with self._stream_ctx as raw: + self._http_response = raw + raise_for_status(raw) + yield raw.iter_raw(chunk_size) if chunk_size else raw.iter_raw() + except httpx.TransportError as exc: + raise NemoTransportError(exc) from exc class NemoStreamResponse(Generic[ModelT]): @@ -151,16 +183,22 @@ def __init__( @contextmanager def stream(self) -> Iterator[Iterator[ModelT]]: """Yield an iterator of parsed model objects.""" - with self._stream_ctx as raw: - raise_for_status(raw) + try: + with self._stream_ctx as raw: + raise_for_status(raw) - def _iter() -> Iterator[ModelT]: - for line in raw.iter_lines(): - payload = _parse_stream_line(line, raw.headers) - if payload is not None: - yield self._model_type.model_validate_json(payload) + def _iter() -> Iterator[ModelT]: + for line in raw.iter_lines(): + payload = _parse_stream_line(line, raw.headers) + if payload is not None: + try: + yield self._model_type.model_validate_json(payload) + except (ValueError, ValidationError) as exc: + raise NemoResponseValidationError(raw, exc) from exc - yield _iter() + yield _iter() + except httpx.TransportError as exc: + raise NemoTransportError(exc) from exc # --------------------------------------------------------------------------- @@ -197,10 +235,13 @@ def http_response(self) -> httpx.Response: async def read(self) -> bytes: """Read and return the entire response body as bytes.""" - async with self._stream_ctx as raw: - data = await raw.aread() - raise_for_status(raw) - return data + try: + async with self._stream_ctx as raw: + data = await raw.aread() + raise_for_status(raw) + return data + except httpx.TransportError as exc: + raise NemoTransportError(exc) from exc @asynccontextmanager async def stream(self, chunk_size: int | None = None) -> AsyncIterator[AsyncIterator[bytes]]: @@ -218,10 +259,13 @@ async def stream(self, chunk_size: int | None = None) -> AsyncIterator[AsyncIter async for chunk in chunks: ... """ - async with self._stream_ctx as raw: - self._http_response = raw - raise_for_status(raw) - yield raw.aiter_raw(chunk_size) if chunk_size else raw.aiter_raw() + try: + async with self._stream_ctx as raw: + self._http_response = raw + raise_for_status(raw) + yield raw.aiter_raw(chunk_size) if chunk_size else raw.aiter_raw() + except httpx.TransportError as exc: + raise NemoTransportError(exc) from exc class AsyncNemoStreamResponse(Generic[ModelT]): @@ -250,16 +294,22 @@ def __init__( @asynccontextmanager async def stream(self) -> AsyncIterator[AsyncIterator[ModelT]]: """Yield an async iterator of parsed model objects.""" - async with self._stream_ctx as raw: - raise_for_status(raw) + try: + async with self._stream_ctx as raw: + raise_for_status(raw) - async def _iter() -> AsyncIterator[ModelT]: - async for line in raw.aiter_lines(): - payload = _parse_stream_line(line, raw.headers) - if payload is not None: - yield self._model_type.model_validate_json(payload) + async def _iter() -> AsyncIterator[ModelT]: + async for line in raw.aiter_lines(): + payload = _parse_stream_line(line, raw.headers) + if payload is not None: + try: + yield self._model_type.model_validate_json(payload) + except (ValueError, ValidationError) as exc: + raise NemoResponseValidationError(raw, exc) from exc - yield _iter() + yield _iter() + except httpx.TransportError as exc: + raise NemoTransportError(exc) from exc # --------------------------------------------------------------------------- @@ -268,13 +318,17 @@ async def _iter() -> AsyncIterator[ModelT]: # Type aliases for the page-fetching callbacks used by paginated responses. -# The page value is int for offset-based or str for cursor-based pagination. +# Fetchers receive the strategy-specific page token returned by ``next_page()``. SyncPageFetcher = Callable[[PreparedRequest, Any], httpx.Response] -AsyncPageFetcher = Callable[[PreparedRequest, Any], Coroutine[Any, Any, httpx.Response]] +AsyncPageFetcher = Callable[[PreparedRequest, Any], Awaitable[httpx.Response]] + + +PageModelT = TypeVar("PageModelT", bound=BaseModel) +PageMetadataT = TypeVar("PageMetadataT") @dataclass(frozen=True, slots=True) -class PageResult(Generic[ModelT]): +class PageResult(Generic[PageModelT, PageMetadataT]): """A single page of results with pagination metadata. Returned by :meth:`NemoPaginatedResponse.page` for callers who want @@ -282,19 +336,30 @@ class PageResult(Generic[ModelT]): resp = client.send(list_items()) page = resp.page() - print(f"Page {page.page} of {page.total_pages} ({page.total_results} total)") + print( + f"Page {page.metadata['page']} of {page.metadata['total_pages']} " + f"({page.metadata['total_results']} total)" + ) for item in page.items: print(item.name) """ - items: list[ModelT] - page: int | None = None - page_size: int | None = None - total_pages: int | None = None - total_results: int | None = None + items: list[PageModelT] + metadata: PageMetadataT + + +PaginatedModelT = TypeVar("PaginatedModelT", bound=BaseModel) +PaginatedStrategyT_co = TypeVarExt( + "PaginatedStrategyT_co", + bound=PaginationStrategy[Any, Any], + default=OffsetPagination, + covariant=True, +) +PageTokenT = TypeVar("PageTokenT") +MetadataT = TypeVar("MetadataT") -class NemoPaginatedResponse(Generic[ModelT]): +class NemoPaginatedResponse(Generic[PaginatedModelT, PaginatedStrategyT_co]): """Sync paginated API response. Provides two iteration modes:: @@ -305,75 +370,87 @@ class NemoPaginatedResponse(Generic[ModelT]): # Iterate page by page with metadata for page in response.pages(): - print(f"Page {page.page}/{page.total_pages}") + print(f"Page {page.metadata['page']}/{page.metadata['total_pages']}") for item in page.items: process(item) For single-page access, use :meth:`page`:: page = response.page() - print(f"{page.total_results} total across {page.total_pages} pages") + print( + f"{page.metadata['total_results']} total across " + f"{page.metadata['total_pages']} pages" + ) """ def __init__( self, first_http_response: httpx.Response, - model_type: type[ModelT], + model_type: type[PaginatedModelT], request: PreparedRequest, fetch_page: SyncPageFetcher, - strategy: type[PaginationStrategy] | None = None, + strategy: type[PaginationStrategy[Any, Any]] | None = None, ) -> None: self._first_response = first_http_response self._model_type = model_type self.request = request self._fetch_page = fetch_page - self._strategy: type[PaginationStrategy] = strategy or OffsetPagination + self._strategy: type[PaginationStrategy[Any, Any]] = strategy or OffsetPagination @property def http_response(self) -> httpx.Response: return self._first_response - def _parse_page(self, raw: httpx.Response) -> tuple[list[ModelT], dict]: - """Parse a page response into (items, raw_body).""" + def _parse_page(self, raw: httpx.Response) -> tuple[list[PaginatedModelT], dict, Any]: + """Parse a page response into items, its raw body, and typed metadata.""" raise_for_status(raw) - body = raw.json() - items = [self._model_type.model_validate(item) for item in self._strategy.extract_items(body)] - return items, body + return _validated_page(raw, self._model_type, self._strategy) + + @overload + def page( + self: NemoPaginatedResponse[PaginatedModelT, PaginationStrategy[PageTokenT, MetadataT]], + ) -> PageResult[PaginatedModelT, MetadataT]: ... + + @overload + def page(self) -> PageResult[PaginatedModelT, Any]: ... - def page(self) -> PageResult[ModelT]: + def page(self) -> PageResult[PaginatedModelT, Any]: """Return the first page as a :class:`PageResult` with metadata.""" - items, body = self._parse_page(self._first_response) - metadata = self._strategy.extract_metadata(body) - return PageResult(items=items, **metadata) + items, _, metadata = self._parse_page(self._first_response) + return PageResult(items=items, metadata=metadata) - def items(self) -> Iterator[ModelT]: + def items(self) -> Iterator[PaginatedModelT]: """Iterate all items across all pages, fetching subsequent pages lazily.""" - items, body = self._parse_page(self._first_response) + items, body, _ = self._parse_page(self._first_response) yield from items - next_page = self._strategy.next_page(body, 1) + next_page = self._strategy.next_page(body) while next_page is not None: - items, body = self._parse_page(self._fetch_page(self.request, next_page)) + items, body, _ = self._parse_page(self._fetch_page(self.request, next_page)) yield from items - current = next_page - next_page = self._strategy.next_page(body, current) + next_page = self._strategy.next_page(body) + + @overload + def pages( + self: NemoPaginatedResponse[PaginatedModelT, PaginationStrategy[PageTokenT, MetadataT]], + ) -> Iterator[PageResult[PaginatedModelT, MetadataT]]: ... + + @overload + def pages(self) -> Iterator[PageResult[PaginatedModelT, Any]]: ... - def pages(self) -> Iterator[PageResult[ModelT]]: + def pages(self) -> Iterator[PageResult[PaginatedModelT, Any]]: """Iterate page by page, yielding :class:`PageResult` objects with metadata.""" - items, body = self._parse_page(self._first_response) - metadata = self._strategy.extract_metadata(body) - yield PageResult(items=items, **metadata) + items, body, metadata = self._parse_page(self._first_response) + yield PageResult(items=items, metadata=metadata) - next_page = self._strategy.next_page(body, 1) + next_page = self._strategy.next_page(body) while next_page is not None: - items, body = self._parse_page(self._fetch_page(self.request, next_page)) - metadata = self._strategy.extract_metadata(body) - yield PageResult(items=items, **metadata) - current = next_page - next_page = self._strategy.next_page(body, current) + items, body, metadata = self._parse_page(self._fetch_page(self.request, next_page)) + yield PageResult(items=items, metadata=metadata) + next_page = self._strategy.next_page(body) -class AsyncNemoPaginatedResponse(Generic[ModelT]): +class AsyncNemoPaginatedResponse(Generic[PaginatedModelT, PaginatedStrategyT_co]): """Async paginated API response. Async twin of :class:`NemoPaginatedResponse`:: @@ -382,66 +459,75 @@ class AsyncNemoPaginatedResponse(Generic[ModelT]): print(item.name) async for page in response.pages(): - print(f"Page {page.page}/{page.total_pages}") + print(f"Page {page.metadata['page']}/{page.metadata['total_pages']}") """ def __init__( self, first_http_response: httpx.Response, - model_type: type[ModelT], + model_type: type[PaginatedModelT], request: PreparedRequest, fetch_page: AsyncPageFetcher, - strategy: type[PaginationStrategy] | None = None, + strategy: type[PaginationStrategy[Any, Any]] | None = None, ) -> None: self._first_response = first_http_response self._model_type = model_type self.request = request self._fetch_page = fetch_page - self._strategy: type[PaginationStrategy] = strategy or OffsetPagination + self._strategy: type[PaginationStrategy[Any, Any]] = strategy or OffsetPagination @property def http_response(self) -> httpx.Response: return self._first_response - def _parse_page(self, raw: httpx.Response) -> tuple[list[ModelT], dict]: - """Parse a page response into (items, raw_body).""" + def _parse_page(self, raw: httpx.Response) -> tuple[list[PaginatedModelT], dict, Any]: + """Parse a page response into items, its raw body, and typed metadata.""" raise_for_status(raw) - body = raw.json() - items = [self._model_type.model_validate(item) for item in self._strategy.extract_items(body)] - return items, body + return _validated_page(raw, self._model_type, self._strategy) + + @overload + def page( + self: AsyncNemoPaginatedResponse[PaginatedModelT, PaginationStrategy[PageTokenT, MetadataT]], + ) -> PageResult[PaginatedModelT, MetadataT]: ... + + @overload + def page(self) -> PageResult[PaginatedModelT, Any]: ... - def page(self) -> PageResult[ModelT]: + def page(self) -> PageResult[PaginatedModelT, Any]: """Return the first page as a :class:`PageResult` with metadata.""" - items, body = self._parse_page(self._first_response) - metadata = self._strategy.extract_metadata(body) - return PageResult(items=items, **metadata) + items, _, metadata = self._parse_page(self._first_response) + return PageResult(items=items, metadata=metadata) - async def items(self) -> AsyncIterator[ModelT]: + async def items(self) -> AsyncIterator[PaginatedModelT]: """Iterate all items across all pages, fetching subsequent pages lazily.""" - items, body = self._parse_page(self._first_response) + items, body, _ = self._parse_page(self._first_response) for item in items: yield item - next_page = self._strategy.next_page(body, 1) + next_page = self._strategy.next_page(body) while next_page is not None: raw = await self._fetch_page(self.request, next_page) - items, body = self._parse_page(raw) + items, body, _ = self._parse_page(raw) for item in items: yield item - current = next_page - next_page = self._strategy.next_page(body, current) + next_page = self._strategy.next_page(body) + + @overload + def pages( + self: AsyncNemoPaginatedResponse[PaginatedModelT, PaginationStrategy[PageTokenT, MetadataT]], + ) -> AsyncIterator[PageResult[PaginatedModelT, MetadataT]]: ... + + @overload + def pages(self) -> AsyncIterator[PageResult[PaginatedModelT, Any]]: ... - async def pages(self) -> AsyncIterator[PageResult[ModelT]]: + async def pages(self) -> AsyncIterator[PageResult[PaginatedModelT, Any]]: """Iterate page by page, yielding :class:`PageResult` objects with metadata.""" - items, body = self._parse_page(self._first_response) - metadata = self._strategy.extract_metadata(body) - yield PageResult(items=items, **metadata) + items, body, metadata = self._parse_page(self._first_response) + yield PageResult(items=items, metadata=metadata) - next_page = self._strategy.next_page(body, 1) + next_page = self._strategy.next_page(body) while next_page is not None: raw = await self._fetch_page(self.request, next_page) - items, body = self._parse_page(raw) - metadata = self._strategy.extract_metadata(body) - yield PageResult(items=items, **metadata) - current = next_page - next_page = self._strategy.next_page(body, current) + items, body, metadata = self._parse_page(raw) + yield PageResult(items=items, metadata=metadata) + next_page = self._strategy.next_page(body) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py index 15c5ef6f61..524a7533a3 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/types.py @@ -13,7 +13,8 @@ from dataclasses import dataclass, replace from typing import Any, ClassVar, Generic, ParamSpec, Protocol, TypeVar -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter +from typing_extensions import TypedDict from typing_extensions import TypeVar as TypeVarExt P = ParamSpec("P") @@ -46,8 +47,12 @@ def ChatEndpoint(body: ChatRequest, *, workspace: str) -> Stream[ChatChunk]: ... # --------------------------------------------------------------------------- -class PaginationStrategy(Protocol): - """Protocol for pagination strategies. +PageTokenT = TypeVar("PageTokenT") +MetadataT = TypeVar("MetadataT") + + +class PaginationStrategy(Generic[PageTokenT, MetadataT]): + """Base class associating a pagination strategy with its cursor and metadata types. Pagination strategies control how the client extracts items from a page response, determines the next page identifier, builds query params @@ -55,19 +60,34 @@ class PaginationStrategy(Protocol): """ @classmethod - def extract_items(cls, response_body: dict) -> list[dict]: ... + def extract_items(cls, response_body: dict) -> list[dict]: + raise NotImplementedError @classmethod - def next_page(cls, response_body: dict, current_page: Any) -> Any | None: ... + def next_page(cls, response_body: dict) -> PageTokenT | None: + raise NotImplementedError @classmethod - def page_query_params(cls, page: Any) -> dict[str, Any]: ... + def page_query_params(cls, page: PageTokenT) -> dict[str, Any]: + raise NotImplementedError @classmethod - def extract_metadata(cls, response_body: dict) -> dict[str, Any]: ... + def extract_metadata(cls, response_body: dict) -> MetadataT: + raise NotImplementedError + + +class OffsetPaginationMetadata(TypedDict): + page: int + page_size: int + current_page_size: int + total_pages: int + total_results: int + + +_OFFSET_PAGINATION_METADATA_ADAPTER = TypeAdapter(OffsetPaginationMetadata) -class OffsetPagination: +class OffsetPagination(PaginationStrategy[int, OffsetPaginationMetadata]): """Offset-based pagination using ``page`` query parameter. This is the default strategy, matching the standard ``NemoListResponse`` @@ -87,6 +107,7 @@ class MyPagination(OffsetPagination): pagination_field: ClassVar[str] = "pagination" page_field: ClassVar[str] = "page" page_size_field: ClassVar[str] = "page_size" + current_page_size_field: ClassVar[str] = "current_page_size" total_pages_field: ClassVar[str] = "total_pages" total_results_field: ClassVar[str] = "total_results" @@ -95,10 +116,11 @@ def extract_items(cls, response_body: dict) -> list[dict]: return response_body.get(cls.items_field, []) @classmethod - def next_page(cls, response_body: dict, current_page: int) -> int | None: + def next_page(cls, response_body: dict) -> int | None: pagination = response_body.get(cls.pagination_field) if pagination is None: return None + current_page = pagination.get(cls.page_field, 1) total = pagination.get(cls.total_pages_field, 1) if current_page < total: return current_page + 1 @@ -109,20 +131,71 @@ def page_query_params(cls, page: int) -> dict[str, int]: return {cls.page_param: page} @classmethod - def extract_metadata(cls, response_body: dict) -> dict[str, Any]: - pagination = response_body.get(cls.pagination_field) or {} + def extract_metadata(cls, response_body: dict) -> OffsetPaginationMetadata: + pagination = response_body.get(cls.pagination_field) + return _OFFSET_PAGINATION_METADATA_ADAPTER.validate_python( + { + "page": pagination[cls.page_field], + "page_size": pagination[cls.page_size_field], + "current_page_size": pagination[cls.current_page_size_field], + "total_pages": pagination[cls.total_pages_field], + "total_results": pagination[cls.total_results_field], + } + if isinstance(pagination, dict) + else pagination + ) + + +class CursorPaginationMetadata(TypedDict): + total: int + next_page: str | None + prev_page: str | None + + +class CursorPagination(PaginationStrategy[str, CursorPaginationMetadata]): + """Cursor pagination matching the Jobs log response envelope. + + The response contains items and cursor metadata at the top level:: + + {"data": [...], "total": 42, "next_page": "...", "prev_page": null} + """ + + items_field: ClassVar[str] = "data" + cursor_param: ClassVar[str] = "page_cursor" + total_field: ClassVar[str] = "total" + next_page_field: ClassVar[str] = "next_page" + prev_page_field: ClassVar[str] = "prev_page" + + @classmethod + def extract_items(cls, response_body: dict) -> list[dict]: + return response_body.get(cls.items_field, []) + + @classmethod + def next_page(cls, response_body: dict) -> str | None: + return response_body.get(cls.next_page_field) + + @classmethod + def page_query_params(cls, page: str) -> dict[str, str]: + return {cls.cursor_param: page} + + @classmethod + def extract_metadata(cls, response_body: dict) -> CursorPaginationMetadata: return { - "page": pagination.get(cls.page_field), - "page_size": pagination.get(cls.page_size_field), - "total_pages": pagination.get(cls.total_pages_field), - "total_results": pagination.get(cls.total_results_field), + "total": response_body.get(cls.total_field, 0), + "next_page": response_body.get(cls.next_page_field), + "prev_page": response_body.get(cls.prev_page_field), } -StrategyT = TypeVarExt("StrategyT", default=OffsetPagination) +PaginatedModelT = TypeVar("PaginatedModelT", bound=BaseModel) +StrategyT = TypeVarExt( + "StrategyT", + bound=PaginationStrategy[Any, Any], + default=OffsetPagination, +) -class Paginated(Generic[ModelT, StrategyT]): +class Paginated(Generic[PaginatedModelT, StrategyT]): """Marker type: endpoint returns paginated results of ``ModelT``. The second type parameter selects the pagination strategy. It defaults @@ -246,12 +319,12 @@ class PreparedRequest(Generic[ResponseT]): response_type: type[ResponseT] | None query_params: dict[str, str | int | bool | None] | None = None extra_headers: dict[str, str] | None = None - client_options: dict[str, Any] | None = None + client_options: dict[str, object] | None = None # Prebuilt GET to replay on a 409 when ``exist_ok`` is set. Produced by a # ``get_on_conflict`` resolver at request-build time (the resolver needs the # live ``body`` model, which is serialised away by the time this request is # sent). ``send()`` replays it instead of raising ``ConflictError``. - on_conflict_get: PreparedRequest | None = None + on_conflict_get: PreparedRequest[ResponseT] | None = None def with_headers(self, headers: dict[str, str]) -> PreparedRequest[ResponseT]: """Return a new PreparedRequest with additional headers merged in.""" diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/api_factory.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/api_factory.py index 0957db8e2a..1c0f49de13 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/api_factory.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/api_factory.py @@ -45,15 +45,14 @@ StepLifecycleParam, SubprocessExecutionProviderParam, ) -from nemo_platform.types.jobs import ( - PlatformJobResponse as PlatformJob, -) from nemo_platform.types.jobs.platform_job_step_spec_param import Executor from nemo_platform_plugin.api.filter import ComparisonOperation, FilterOperation, FilterOperator, LogicalOperation from nemo_platform_plugin.api.parsed_filter import ParsedFilter, make_filter_dep from nemo_platform_plugin.authz import AuthzScope, CallerKind, path_rule +from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.dependencies import get_entity_client, get_sdk_client from nemo_platform_plugin.entities import EntityClient +from nemo_platform_plugin.jobs.client import AsyncJobsClient from nemo_platform_plugin.jobs.docker import validate_gpu_available_for_docker from nemo_platform_plugin.jobs.exceptions import PlatformJobCompilationError from nemo_platform_plugin.jobs.openapi_utils import generate_openapi_extra_params @@ -65,6 +64,14 @@ PlatformJobStatus, PlatformJobStatusResponse, ) +from nemo_platform_plugin.jobs.types import ( + CreatePlatformJobRequest, + JobLogsQueryParams, + ListJobsQueryParams, +) +from nemo_platform_plugin.jobs.types import ( + PlatformJobResponse as PlatformJob, +) from nemo_platform_plugin.schema import DatetimeFilter, Filter, Page, PaginationData, StringFilter from pydantic import BaseModel, Field, TypeAdapter @@ -72,6 +79,11 @@ # This type is aliased to ensure we don't expose internal stainless # type paths to services integrating the job service. +# +# TODO(AIRCORE-922): remove these Stainless-generated ``*Param`` TypedDict aliases. +# The plugin now owns pydantic equivalents in ``jobs/spec.py`` and ``jobs/providers.py``, +# but ~10 consuming plugins construct these as dict literals (TypedDict), so repointing +# them to the pydantic models is a cross-cutting change tracked by AIRCORE-922. PlatformJobSpec = PlatformJobSpecParam PlatformJobStep = PlatformJobStepSpecParam StepLifecycle = StepLifecycleParam @@ -857,25 +869,35 @@ async def create_job( # Build SDK call kwargs, only including optional fields when they have values # (passing None explicitly causes different serialization than omitting) # Note: We store transformed_spec (not input), which includes auto-generated fields. - sdk_kwargs: dict = { + # ``job_spec`` may be a Pydantic model (the transformed job output); + # the request body's ``spec`` is a plain dict on the wire. The + # Stainless SDK serialized models implicitly — the typed client + # validates the body first, so coerce to a dict here. + spec_dict = job_spec.model_dump() if isinstance(job_spec, BaseModel) else job_spec + + # Only include optional fields when they have values — passing None + # explicitly serializes differently than omitting (exclude_unset). + create_fields: dict = { "source": service_name, - "spec": job_spec, + "spec": spec_dict, "platform_spec": platform_spec, - "workspace": workspace, } # Use the resolved job_name (user-provided or generated) if job_name is not None: - sdk_kwargs["name"] = job_name + create_fields["name"] = job_name if request.description is not None: - sdk_kwargs["description"] = request.description + create_fields["description"] = request.description if request.ownership is not None: - sdk_kwargs["ownership"] = request.ownership + create_fields["ownership"] = request.ownership if request.custom_fields is not None: - sdk_kwargs["custom_fields"] = request.custom_fields + create_fields["custom_fields"] = request.custom_fields if request.project: - sdk_kwargs["extra_body"] = {"project": request.project} + create_fields["project"] = request.project - job_resp = await sdk.jobs.create(**sdk_kwargs) + jobs = client_from_platform(sdk, AsyncJobsClient) + job_resp = ( + await jobs.create_job(workspace=workspace, body=CreatePlatformJobRequest(**create_fields)) + ).data() return from_response(job_resp) @router.get( @@ -930,17 +952,17 @@ async def list_jobs( # accepts raw JSON in ``filter=`` and routes it through # parse_json_filter, so a single JSON-string param survives the # round trip cleanly. - sdk_list_kwargs: dict = { - "workspace": workspace, + list_query: ListJobsQueryParams = { "page": page, "page_size": page_size, "sort": str(sort), - "extra_query": {"filter": json.dumps(parsed.to_response())}, + "filter": json.dumps(parsed.to_response()), } - list_jobs_resp = await sdk.jobs.list(**sdk_list_kwargs) + jobs = client_from_platform(sdk, AsyncJobsClient) + list_page = (await jobs.list_jobs(workspace=workspace, query_params=list_query)).page() return Page( - data=[from_response(job) for job in list_jobs_resp.data], - pagination=PaginationData(**list_jobs_resp.pagination.model_dump()), + data=[from_response(job) for job in list_page.items], + pagination=PaginationData.model_validate(list_page.metadata), sort=sort, filter=user_filter or None, ) @@ -955,7 +977,7 @@ async def get_job( ) -> TypedJobResponse: f"""Get a job by name for the {service_name} microservice.""" - job_resp = await sdk.jobs.retrieve(name=name, workspace=workspace) + job_resp = (await client_from_platform(sdk, AsyncJobsClient).get_job(name=name, workspace=workspace)).data() return from_response(job_resp) # Status @@ -968,7 +990,9 @@ async def get_job_status( sdk: AsyncNeMoPlatform = Depends(get_sdk_client), ) -> PlatformJobStatusResponse: f"""Get the status of a job by name for the {service_name} microservice.""" - job_resp = await sdk.jobs.get_status(name=name, workspace=workspace) + job_resp = ( + await client_from_platform(sdk, AsyncJobsClient).get_job_status(name=name, workspace=workspace) + ).data() return PlatformJobStatusResponse(**job_resp.model_dump()) @router.delete( @@ -981,7 +1005,7 @@ async def delete_job( sdk: AsyncNeMoPlatform = Depends(get_sdk_client), ) -> None: f"""Delete a job by name for the {service_name} microservice.""" - await sdk.jobs.delete(name=name, workspace=workspace) + await client_from_platform(sdk, AsyncJobsClient).delete_job(name=name, workspace=workspace) return None @router.post( @@ -994,7 +1018,9 @@ async def cancel_job( ) -> TypedJobResponse: f"""Cancel a job by name for the {service_name} microservice.""" - job_resp = await sdk.jobs.cancel(name=name, workspace=workspace) + job_resp = ( + await client_from_platform(sdk, AsyncJobsClient).cancel_job(name=name, workspace=workspace) + ).data() return from_response(job_resp) # Logs @@ -1010,8 +1036,17 @@ async def get_job_logs( ) -> PlatformJobLogPage: f"""Get the logs of a job by name for the {service_name} microservice.""" - logs = await sdk.jobs.get_logs(workspace=workspace, name=name, limit=limit, page_cursor=page_cursor) - return PlatformJobLogPage(**logs.model_dump()) + logs_query: JobLogsQueryParams = {} + if limit is not None: + logs_query["limit"] = limit + if page_cursor is not None: + logs_query["page_cursor"] = page_cursor + logs_page = ( + await client_from_platform(sdk, AsyncJobsClient).list_job_logs( + workspace=workspace, name=name, query_params=logs_query + ) + ).page() + return PlatformJobLogPage(data=logs_page.items, **logs_page.metadata) # Results @router.get( @@ -1025,7 +1060,9 @@ async def list_job_results( ) -> PlatformJobListResultResponse: f"""Get the results of a job by name for the {service_name} microservice.""" - results = await sdk.jobs.results.list(name=name, workspace=workspace) + results = ( + await client_from_platform(sdk, AsyncJobsClient).list_job_results(name=name, workspace=workspace) + ).data() result_dicts = [result.model_dump() for result in results.data] list_results = [] for result_dict in result_dicts: @@ -1047,7 +1084,9 @@ async def get_job_result( ) -> PlatformJobResultResponse: f"""Get the result of a job by name for the {service_name} microservice.""" - result_obj = await sdk.jobs.results.retrieve(name=name, job=job, workspace=workspace) + result_obj = ( + await client_from_platform(sdk, AsyncJobsClient).get_job_result(name=name, job=job, workspace=workspace) + ).data() # Construct the URL for downloading this result result_dict = result_obj.model_dump() @@ -1091,7 +1130,9 @@ async def _download_route_helper( - Use the `result_serializer` to know how to properly serialize the output """ - result_info = await sdk.jobs.results.retrieve(name=name, job=job, workspace=workspace) + result_info = ( + await client_from_platform(sdk, AsyncJobsClient).get_job_result(name=name, job=job, workspace=workspace) + ).data() _, tmp_dir_path = await download_from_result_info( result_name=name, job_name=job, @@ -1203,7 +1244,9 @@ async def pause_job( ) -> TypedJobResponse: f"""Pause a job by name for the {service_name} microservice.""" - job_resp = await sdk.jobs.pause(name=name, workspace=workspace) + job_resp = ( + await client_from_platform(sdk, AsyncJobsClient).pause_job(name=name, workspace=workspace) + ).data() return from_response(job_resp) @router.post( @@ -1216,7 +1259,9 @@ async def resume_job( ) -> TypedJobResponse: f"""Resume a job by name for the {service_name} microservice.""" - job_resp = await sdk.jobs.resume(name=name, workspace=workspace) + job_resp = ( + await client_from_platform(sdk, AsyncJobsClient).resume_job(name=name, workspace=workspace) + ).data() return from_response(job_resp) _stamp(pause_job, perm="pause", write=True) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/client.py new file mode 100644 index 0000000000..f1862c2f2b --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/client.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed HTTP clients for the Jobs service. + +Wraps the endpoint functions from ``jobs.endpoints`` as direct methods using +the ``method()`` descriptor, following the example-plugin / Files pattern. + +Usage:: + + from nemo_platform_plugin.jobs.client import JobsClient + from nemo_platform_plugin.jobs.types import CreatePlatformJobRequest + + client = JobsClient(base_url="...", workspace="default") + resp = client.create_job(body=CreatePlatformJobRequest(...)) + job = resp.data() + + for job in client.list_jobs().items(): + print(job.name) + + with client.download_job_result(job="j-1", name="out").stream() as chunks: + for chunk in chunks: + ... +""" + +from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient +from nemo_platform_plugin.client.method import method +from nemo_platform_plugin.jobs import endpoints + + +class _JobsMethods: + # Execution profiles + get_execution_profiles = method(endpoints.get_execution_profiles) + + # Job CRUD + lifecycle + create_job = method(endpoints.create_job) + list_jobs = method(endpoints.list_jobs) + get_job = method(endpoints.get_job) + delete_job = method(endpoints.delete_job) + cancel_job = method(endpoints.cancel_job) + pause_job = method(endpoints.pause_job) + resume_job = method(endpoints.resume_job) + + # Job status + get_job_status = method(endpoints.get_job_status) + update_job_status_details = method(endpoints.update_job_status_details) + + # Job logs + list_job_logs = method(endpoints.list_job_logs) + + # Job results + create_job_result = method(endpoints.create_job_result) + list_job_results = method(endpoints.list_job_results) + get_job_result = method(endpoints.get_job_result) + download_job_result = method(endpoints.download_job_result) + + # Job steps + list_steps = method(endpoints.list_steps) + get_job_step = method(endpoints.get_job_step) + update_job_step_status = method(endpoints.update_job_step_status) + + # Job tasks + list_job_step_tasks = method(endpoints.list_job_step_tasks) + update_job_step_task = method(endpoints.update_job_step_task) + get_job_step_task = method(endpoints.get_job_step_task) + + +class JobsClient(_JobsMethods, NemoClient): + """Sync client for the Jobs service API.""" + + +class AsyncJobsClient(_JobsMethods, AsyncNemoClient): + """Async client for the Jobs service API.""" diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/endpoints.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/endpoints.py new file mode 100644 index 0000000000..cf553152fb --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/endpoints.py @@ -0,0 +1,213 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed endpoint definitions for the Jobs service. + +These are the single source of truth for the HTTP contract. Each function +is decorated with an HTTP-method decorator and ``@abstractmethod``; the +decorator turns the signature into a :class:`PreparedRequest` builder. +""" + +from __future__ import annotations + +from abc import abstractmethod + +from nemo_platform_plugin.client.endpoint import delete, get, patch, post, put +from nemo_platform_plugin.client.types import BinaryContent, CursorPagination, Paginated +from nemo_platform_plugin.jobs.execution_profiles import ( + DockerJobExecutionProfile, + E2EJobExecutionProfile, + KubernetesJobExecutionProfile, + SubprocessJobExecutionProfile, + VolcanoJobExecutionProfile, +) +from nemo_platform_plugin.jobs.schemas import ( + PlatformJobLog, + PlatformJobResultCreateRequest, + PlatformJobResultResponse, + PlatformJobStatusResponse, +) +from nemo_platform_plugin.jobs.types import ( + CreatePlatformJobRequest, + JobLogsQueryParams, + JobStatusDetailsUpdate, + ListJobResultsQueryParams, + ListJobsQueryParams, + ListStepsQueryParams, + PlatformJobListResultResponse, + PlatformJobListTaskResponse, + PlatformJobResponse, + PlatformJobStatusUpdateRequest, + PlatformJobStepResponse, + PlatformJobStepWithContext, + PlatformJobTaskResponse, + PlatformJobTaskUpdate, +) + +# The execution-profiles endpoint returns a union over all configured backend +# profile types (matches the Stainless ``JobListExecutionProfilesResponseItem``). +ExecutionProfile = ( + DockerJobExecutionProfile + | KubernetesJobExecutionProfile + | VolcanoJobExecutionProfile + | SubprocessJobExecutionProfile + | E2EJobExecutionProfile +) + +# --------------------------------------------------------------------------- +# Execution profiles +# --------------------------------------------------------------------------- + + +@get("/apis/jobs/v2/execution-profiles") +@abstractmethod +def get_execution_profiles() -> list[ExecutionProfile]: ... + + +# --------------------------------------------------------------------------- +# Job CRUD + lifecycle +# --------------------------------------------------------------------------- + + +@post("/apis/jobs/v2/workspaces/{workspace}/jobs") +@abstractmethod +def create_job(*, workspace: str | None = None, body: CreatePlatformJobRequest) -> PlatformJobResponse: ... + + +@get("/apis/jobs/v2/workspaces/{workspace}/jobs") +@abstractmethod +def list_jobs( + *, workspace: str | None = None, query_params: ListJobsQueryParams | None = None +) -> Paginated[PlatformJobResponse]: ... + + +@get("/apis/jobs/v2/workspaces/{workspace}/jobs/{name}") +@abstractmethod +def get_job(*, workspace: str | None = None, name: str) -> PlatformJobResponse: ... + + +@delete("/apis/jobs/v2/workspaces/{workspace}/jobs/{name}") +@abstractmethod +def delete_job(*, workspace: str | None = None, name: str) -> None: ... + + +@post("/apis/jobs/v2/workspaces/{workspace}/jobs/{name}/cancel") +@abstractmethod +def cancel_job(*, workspace: str | None = None, name: str) -> PlatformJobResponse: ... + + +@post("/apis/jobs/v2/workspaces/{workspace}/jobs/{name}/pause") +@abstractmethod +def pause_job(*, workspace: str | None = None, name: str) -> PlatformJobResponse: ... + + +@post("/apis/jobs/v2/workspaces/{workspace}/jobs/{name}/resume") +@abstractmethod +def resume_job(*, workspace: str | None = None, name: str) -> PlatformJobResponse: ... + + +# NOTE: no ``rerun_job`` — the server's ``/rerun`` route is test-only and not +# mounted in the release service (see services/core/jobs/.../api/v2/jobs/rerun.py), +# so exposing it on the client would 404 in production. + + +# --------------------------------------------------------------------------- +# Job status +# --------------------------------------------------------------------------- + + +@get("/apis/jobs/v2/workspaces/{workspace}/jobs/{name}/status") +@abstractmethod +def get_job_status(*, workspace: str | None = None, name: str) -> PlatformJobStatusResponse: ... + + +@patch("/apis/jobs/v2/workspaces/{workspace}/jobs/{name}/status-details") +@abstractmethod +def update_job_status_details(*, workspace: str | None = None, name: str, body: JobStatusDetailsUpdate) -> None: ... + + +# --------------------------------------------------------------------------- +# Job logs +# --------------------------------------------------------------------------- + + +@get("/apis/jobs/v2/workspaces/{workspace}/jobs/{name}/logs") +@abstractmethod +def list_job_logs( + *, workspace: str | None = None, name: str, query_params: JobLogsQueryParams | None = None +) -> Paginated[PlatformJobLog, CursorPagination]: ... + + +# --------------------------------------------------------------------------- +# Job results +# --------------------------------------------------------------------------- + + +@post("/apis/jobs/v2/workspaces/{workspace}/jobs/{job}/results/{name}") +@abstractmethod +def create_job_result( + *, workspace: str | None = None, job: str, name: str, body: PlatformJobResultCreateRequest +) -> PlatformJobResultResponse: ... + + +@get("/apis/jobs/v2/workspaces/{workspace}/jobs/{name}/results") +@abstractmethod +def list_job_results( + *, workspace: str | None = None, name: str, query_params: ListJobResultsQueryParams | None = None +) -> PlatformJobListResultResponse: ... + + +@get("/apis/jobs/v2/workspaces/{workspace}/jobs/{job}/results/{name}") +@abstractmethod +def get_job_result(*, workspace: str | None = None, job: str, name: str) -> PlatformJobResultResponse: ... + + +@get("/apis/jobs/v2/workspaces/{workspace}/jobs/{job}/results/{name}/download") +@abstractmethod +def download_job_result(*, workspace: str | None = None, job: str, name: str) -> BinaryContent: ... + + +# --------------------------------------------------------------------------- +# Job steps +# --------------------------------------------------------------------------- + + +@get("/apis/jobs/v2/workspaces/{workspace}/jobs/{name}/steps") +@abstractmethod +def list_steps( + *, workspace: str | None = None, name: str, query_params: ListStepsQueryParams | None = None +) -> Paginated[PlatformJobStepWithContext]: ... + + +@get("/apis/jobs/v2/workspaces/{workspace}/jobs/{job}/steps/{name}") +@abstractmethod +def get_job_step(*, workspace: str | None = None, job: str, name: str) -> PlatformJobStepResponse: ... + + +@patch("/apis/jobs/v2/workspaces/{workspace}/jobs/{job}/steps/{name}/status") +@abstractmethod +def update_job_step_status( + *, workspace: str | None = None, job: str, name: str, body: PlatformJobStatusUpdateRequest +) -> PlatformJobStepResponse: ... + + +# --------------------------------------------------------------------------- +# Job tasks +# --------------------------------------------------------------------------- + + +@get("/apis/jobs/v2/workspaces/{workspace}/jobs/{job}/steps/{name}/tasks") +@abstractmethod +def list_job_step_tasks(*, workspace: str | None = None, job: str, name: str) -> PlatformJobListTaskResponse: ... + + +@put("/apis/jobs/v2/workspaces/{workspace}/jobs/{job}/steps/{step}/tasks/{name}") +@abstractmethod +def update_job_step_task( + *, workspace: str | None = None, job: str, step: str, name: str, body: PlatformJobTaskUpdate +) -> PlatformJobTaskResponse: ... + + +@get("/apis/jobs/v2/workspaces/{workspace}/jobs/{job}/steps/{step}/tasks/{name}") +@abstractmethod +def get_job_step_task(*, workspace: str | None = None, job: str, step: str, name: str) -> PlatformJobTaskResponse: ... diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/execution_profiles.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/execution_profiles.py new file mode 100644 index 0000000000..54319cc166 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/execution_profiles.py @@ -0,0 +1,432 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Execution-profile types for the Jobs service. + +An *execution profile* describes a configured backend (docker, kubernetes, +volcano, subprocess, e2e) that the jobs controller can schedule steps onto. +These are returned by the ``get_execution_profiles`` endpoint. + +This module holds the **data shapes** as pure pydantic — no docker or +kubernetes runtime dependencies. Server-side behaviour that needs those +libraries (``KubernetesVolume.to_k8s()`` etc.) lives in the Jobs service, +which subclasses these models. Both the server and the typed HTTP client +share these definitions. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from nemo_platform_plugin.config import NMP_CONFIG_WARNINGS_DISABLED_ENV_VAR +from nemo_platform_plugin.jobs.constants import ( + CONFIG_TASK_STORAGE_PATH_ENVVAR, + EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, + NEMO_JOB_ATTEMPT_ID_ENVVAR, + NEMO_JOB_FILESET_ENVVAR, + NEMO_JOB_ID_ENVVAR, + NEMO_JOB_SECRETS_ENVVAR, + NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, + NEMO_JOB_STEP_ENVVAR, + NEMO_JOB_TASK_ENVVAR, + NEMO_JOB_WORKSPACE_ENVVAR, + PERSISTENT_JOB_STORAGE_PATH_ENVVAR, + TASK_CONFIG_ENVVAR, +) +from nemo_platform_plugin.jobs.providers import ComputeResources +from nemo_platform_plugin.jobs.spec import BaseExecutionProfile, ProviderRef +from pydantic import BaseModel, ConfigDict, Field, model_validator + +# Default image used to set filesystem permissions on job storage volumes. +DEFAULT_VOLUME_PERMISSIONS_IMAGE = "busybox" + +# Env var names set by the platform during job creation; user-provided profile +# environment must not conflict. The job-scoped names come from the shared +# ``jobs.constants`` leaf; the auth / config / telemetry names are stable env +# var strings kept here to avoid importing server-side auth/config modules. +RESERVED_JOB_ENVIRONMENT_VARIABLE_NAMES: frozenset[str] = frozenset( + { + # Job runtime (from nemo_platform_plugin.jobs.constants) + CONFIG_TASK_STORAGE_PATH_ENVVAR, + EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, + NEMO_JOB_ATTEMPT_ID_ENVVAR, + NEMO_JOB_FILESET_ENVVAR, + NEMO_JOB_ID_ENVVAR, + NEMO_JOB_SECRETS_ENVVAR, + NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, + NEMO_JOB_STEP_ENVVAR, + NEMO_JOB_TASK_ENVVAR, + NEMO_JOB_WORKSPACE_ENVVAR, + PERSISTENT_JOB_STORAGE_PATH_ENVVAR, + TASK_CONFIG_ENVVAR, + # Auth + "NMP_PRINCIPAL", + # OTEL (telemetry) + "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT", + "OTEL_LOGS_EXPORTER", + "OTEL_SERVICE_NAME", + "OTEL_EXPORTER_OTLP_LOGS_HEADERS", + # Platform shared envvars (to_shared_envvars with NMP_ prefix) + NMP_CONFIG_WARNINGS_DISABLED_ENV_VAR, + "NMP_BASE_URL", + "NMP_JOBS_URL", + "NMP_FILES_URL", + "NMP_MODELS_URL", + "NMP_SECRETS_URL", + } +) + + +class ImagePullSecret(BaseModel): + """Kubernetes image pull secret reference.""" + + # extra=forbid keeps additionalProperties: false on the generated schema. + model_config = ConfigDict(extra="forbid") + + name: str = Field(description="Kubernetes Secret name for pulling images") + + +class JobExecutionProfileConfig(BaseModel): + ttl_seconds_before_active: int = 30 * 60 # 30 minutes + ttl_seconds_active: int = 24 * 60 * 60 # 24 hours + ttl_seconds_after_finished: int = 60 * 60 # 1 hour + cleanup_completed_jobs_immediately: bool = True + launcher_tool_path: str = Field(default="/tools/jobs-launcher", description="Path to the jobs launcher tool") + default_task_image: str | None = Field( + default=None, + min_length=1, + description="Default container image for job task pods. Used when a job step omits container.image. " + "When unset, falls back to the platform CPU tasks image (platform.image_registry/nmp-cpu-tasks:platform.image_tag).", + ) + env: dict[str, str] = Field( + default_factory=dict, + description="Optional env vars applied to all jobs (e.g. HOME=/tmp). Keys must not conflict with platform-reserved names. Job steps may override these variables.", + ) + + @model_validator(mode="after") + def validate_env_no_reserved_names(self) -> JobExecutionProfileConfig: + conflicting = [k for k in self.env if k in RESERVED_JOB_ENVIRONMENT_VARIABLE_NAMES] + if conflicting: + raise ValueError( + f"Profile environment keys must not conflict with platform-reserved names: {sorted(conflicting)}" + ) + return self + + +# --------------------------------------------------------------------------- +# Docker +# --------------------------------------------------------------------------- + + +class DockerVolumeMount(BaseModel): + volume_name: str = Field(description="Name of the Docker volume to mount") + mount_path: str = Field(description="Path inside the container where the volume will be mounted") + kind: Literal["volume", "tmpfs"] = Field( + default="volume", + description="Type of the Docker volume to mount. Options are 'volume' or 'tmpfs' (default: 'volume'). tmpfs volumes are only supported on Linux hosts.", + ) + options: dict | None = Field(default=None, description="Additional options for the volume") + allow_create_volume: bool = Field( + default=False, description="Whether to allow the creation of the volume if it does not exist (default: false)." + ) + + +class DockerJobStorageConfig(BaseModel): + """Configuration for persistent storage in Docker jobs.""" + + volume_name: str = Field( + default="nemo-jobs-storage", description="Name of the Docker volume for persistent storage" + ) + volume_permissions_image: str = Field( + default=DEFAULT_VOLUME_PERMISSIONS_IMAGE, description="Docker image used to set permissions on the volume" + ) + additional_volume_mounts: list[DockerVolumeMount] = Field( + default_factory=list, + description="List of additional Docker volume mounts for the job", + ) + + +class DockerJobNetworkConfig(BaseModel): + job_container_network: str = Field(default="host", description="Docker network for the job container") + + +class DockerJobExecutionProfileConfig(JobExecutionProfileConfig): + """Configuration for Docker Job execution profile.""" + + storage: DockerJobStorageConfig = Field( + default_factory=DockerJobStorageConfig, description="Docker storage configuration" + ) + networking: DockerJobNetworkConfig = Field( + default_factory=DockerJobNetworkConfig, description="Docker networking configuration" + ) + + +class DockerJobExecutionProfile(BaseExecutionProfile): + """ + Execution configuration for a Docker Job. + This is used to define the executor type, provider, profile, and any additional configuration + required for the executor to run the job on Docker + """ + + backend: Literal["docker"] = "docker" + config: DockerJobExecutionProfileConfig = Field(description="Additional configuration for the docker executor") + + @property + def supports_persistent_storage(self) -> bool: + """Indicates if the execution profile supports persistent storage.""" + return self.config.storage is not None and self.config.storage.volume_name != "" + + +# --------------------------------------------------------------------------- +# Kubernetes (shared) +# --------------------------------------------------------------------------- + + +class KubernetesObjectMetadata(BaseModel): + labels: dict[str, str] = Field(default_factory=dict) + annotations: dict[str, str] = Field(default_factory=dict) + + +class KubernetesPersistentVolumeClaim(BaseModel): + """Kubernetes Persistent Volume Claim definition.""" + + claim_name: str = Field(description="Persistent Volume Claim Name") + read_only: bool = Field(default=False, description="Whether the volume is mounted read-only") + + +class KubernetesEmptyDirVolume(BaseModel): + """Kubernetes EmptyDir Volume definition.""" + + medium: str | None = Field(default=None, description="The medium of the emptyDir volume (e.g., 'Memory')") + size_limit: str | None = Field(default=None, description="The size limit of the emptyDir volume (e.g., '1Gi')") + + +class KubernetesVolume(BaseModel): + """Kubernetes Volume definition. + + Data shape only. The server subclass adds ``to_k8s()`` which requires the + ``kubernetes`` client library. + """ + + name: str = Field(description="Volume Name") + persistent_volume_claim: KubernetesPersistentVolumeClaim | None = Field( + default=None, description="Persistent Volume Claim configuration" + ) + empty_dir: KubernetesEmptyDirVolume | None = Field(default=None, description="EmptyDir Volume configuration") + + @model_validator(mode="after") + def validate_self(self) -> KubernetesVolume: + """Ensure that exactly one volume source is specified.""" + if sum(source is not None for source in [self.persistent_volume_claim, self.empty_dir]) != 1: + raise ValueError("Exactly one of 'persistent_volume_claim' or 'empty_dir' must be specified.") + return self + + +class KubernetesVolumeMount(BaseModel): + """Kubernetes Volume Mount definition. + + Data shape only. The server subclass adds ``to_k8s()``. + """ + + name: str = Field(description="Volume Name") + mount_path: str = Field(description="Mount Path in the container") + sub_path: str | None = Field(default=None, description="Sub-path within the volume to mount") + read_only: bool = Field(default=False, description="Whether the volume mount is read-only") + + +class KubernetesJobStorageConfig(BaseModel): + """Configuration for persistent storage in Kubernetes jobs.""" + + pvc_name: str = Field(default="", description="Persistent Volume Claim Name to use for job storage.") + volume_permissions_image: str = Field( + default=DEFAULT_VOLUME_PERMISSIONS_IMAGE, description="Image used to set volume permissions" + ) + additional_volumes: list[KubernetesVolume] = Field(default_factory=list, description="Additional volumes to mount") + additional_volume_mounts: list[KubernetesVolumeMount] = Field( + default_factory=list, description="Additional volume mounts" + ) + + +class BaseKubernetesExecutionProfileConfig(JobExecutionProfileConfig): + """Common configuration for Kubernetes execution environment.""" + + namespace: str | None = Field( + default=None, + description="Kubernetes namespace to submit the job to. If not set, it will be determined from the environment.", + ) + + service_account_name: str = Field( + default="default", + description="Kubernetes service account name for job pods. Uses the Kubernetes default service account when set to 'default'.", + ) + + # Scheduling and resource configuration + tolerations: list[dict[str, Any]] = Field( + default_factory=list, description="Tolerations for the Kubernetes job pods." + ) + node_selector: dict[str, str] = Field( + default_factory=dict, description="Node selector for the Kubernetes job pods." + ) + affinity: dict[str, Any] = Field(default_factory=dict, description="Affinity for the Kubernetes job pods.") + resources: ComputeResources = Field( + default_factory=ComputeResources, description="Resource requests and limits for the Kubernetes job pods." + ) + pod_security_context: dict[str, Any] = Field( + default_factory=dict, description="Pod security context for the Kubernetes job pods." + ) + + # Image pull secrets + image_pull_secrets: list[ImagePullSecret] = Field( + default_factory=list, description="Image pull secrets for the Kubernetes job pods." + ) + + # Optional metadata to add to each job object + job_metadata: KubernetesObjectMetadata = Field( + default_factory=KubernetesObjectMetadata, + description="Metadata to add to each job object in the Kubernetes job.", + ) + + # Optional metadata to add to each pod in the job + pod_metadata: KubernetesObjectMetadata = Field( + default_factory=KubernetesObjectMetadata, description="Metadata to add to each pod in the Kubernetes job." + ) + + # Storage configurations for the job + storage: KubernetesJobStorageConfig = Field( + default_factory=KubernetesJobStorageConfig, description="Storage configuration for the Kubernetes job pods." + ) + + num_gpus: int = Field(default=1, description="Number of GPUs to request for the job") + + scheduler_name: str = Field( + default="", + description="The scheduler name to use for the pod spec. When non-empty, this value is applied to the pod's schedulerName field, enabling custom schedulers such as KAI Scheduler. Empty string omits schedulerName so the cluster default scheduler is used.", + ) + + launcher_image: str = Field( + default="nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest", + description="Container image that contains the jobs-launcher binary.", + ) + + +class KubernetesJobExecutionProfileConfig(BaseKubernetesExecutionProfileConfig): + """Configuration for Kubernetes execution environment.""" + + +class KubernetesJobExecutionProfile(BaseExecutionProfile): + """ + Execution configuration for a Kubernetes Job. + This is used to define the executor type, provider, profile, and any additional configuration + required for the executor to run the job on Kubernetes + """ + + backend: Literal["kubernetes_job"] = "kubernetes_job" + config: KubernetesJobExecutionProfileConfig = Field( + description="Additional configuration for the kubernetes executor", + ) + + @property + def supports_persistent_storage(self) -> bool: + """Indicates if the execution profile supports persistent storage.""" + return self.config.storage is not None and self.config.storage.pvc_name != "" + + +# --------------------------------------------------------------------------- +# Volcano +# --------------------------------------------------------------------------- + + +class VolcanoJobExecutionProfileConfig(BaseKubernetesExecutionProfileConfig): + """Configuration for Volcano Job Execution Profile""" + + queue: str = Field( + default="default", + description="The Volcano queue to submit the job to.", + ) + scheduler_name: str = Field( + default="volcano", + description="The scheduler name to use for the Volcano job.", + ) + + max_retry: int = Field(default=0, description="maxRetry indicates the maximum number of retries allowed by the job") + + plugins: dict[str, Any] = Field( + default_factory=dict, + description="plugins indicates the plugins used by Volcano when the job is scheduled. We always add the pytorch plugin if more than one node.", + ) + + enable_multi_node_networking: bool = Field( + default=True, + description="Enable multi-node networking injection. Sets annotations to trigger Kyverno policy mutations.", + ) + + +class VolcanoJobExecutionProfile(BaseExecutionProfile): + """Volcano Job Execution Profile""" + + backend: Literal["volcano_job"] = "volcano_job" + config: VolcanoJobExecutionProfileConfig = Field( + description="Additional configuration for the kubernetes executor", + ) + + @property + def supports_persistent_storage(self) -> bool: + """Indicates if the execution profile supports persistent storage.""" + return self.config.storage is not None and self.config.storage.pvc_name != "" + + +# --------------------------------------------------------------------------- +# Subprocess +# --------------------------------------------------------------------------- + + +class SubprocessJobExecutionProfileConfig(JobExecutionProfileConfig): + working_directory: str = Field( + default="/tmp/nmp-subprocess-jobs", + description="Root directory for subprocess job state, config, storage, and logs.", + ) + graceful_shutdown_timeout_seconds: int = Field( + default=10, + description="How long to wait after SIGTERM before force killing the process group.", + ) + cleanup_completed_jobs_immediately: bool = Field( + default=False, + description="Keep subprocess working directories by default so runs remain inspectable.", + ) + + +class SubprocessJobExecutionProfile(BaseExecutionProfile): + provider: ProviderRef = Field(default="subprocess") + backend: Literal["subprocess"] = "subprocess" + config: SubprocessJobExecutionProfileConfig = Field( + default_factory=SubprocessJobExecutionProfileConfig, + description="Additional configuration for the subprocess executor", + ) + + @property + def supports_persistent_storage(self) -> bool: + return True + + +# --------------------------------------------------------------------------- +# E2E test backend +# --------------------------------------------------------------------------- + + +class E2EJobExecutionProfile(BaseExecutionProfile): + """ + Execution configuration for E2E testing. + This backend auto-completes jobs without actually running containers, + making tests fast and deterministic. + """ + + backend: Literal["e2e"] = "e2e" + config: JobExecutionProfileConfig = Field( + default_factory=JobExecutionProfileConfig, + description="Configuration for the e2e test executor", + ) + + @property + def supports_persistent_storage(self) -> bool: + """E2E backend claims to support persistent storage since jobs auto-complete without execution.""" + return True diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/providers.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/providers.py new file mode 100644 index 0000000000..c913d87477 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/providers.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Execution provider and container spec types for the Jobs service. + +These types describe *how* a job step runs (which executor, what container, +what resources). They are part of the request body for job creation +(``CreatePlatformJobRequest.platform_spec``) and are pure pydantic — no +server, docker, or kubernetes runtime dependencies. + +This module is the single source of truth for the executor tree. The Jobs +service (``nmp.core.jobs.app.providers``) re-exports from here so both the +server and the typed HTTP client share one definition. +""" + +from __future__ import annotations + +import re +from typing import Annotated, Literal, Self, Union + +from pydantic import BaseModel, Field, field_validator, model_validator + +# SHM: megabyte/gigabyte scale only — Mi, Gi (binary) or M, G (decimal SI). +# Ki / Ti / Pi / Ei and other suffixes are not accepted for /dev/shm. +_SHM_QUANTITY_RE = re.compile(r"^([+-]?(?:\d+|\d*\.\d+)(?:[eE][+-]?\d+)?)(Mi|Gi|M|G)$") + + +class ContainerSpec(BaseModel): + """ + Specification for a container configuration. + + Defines the container image and related configuration for job execution. + """ + + image: str | None = Field(default=None, min_length=1) + """The container image to use for execution. When omitted, resolved from the execution profile's default_task_image or the platform CPU tasks image.""" + + entrypoint: list[str] = Field(default_factory=list) + """The entrypoint for the container as a list of strings (e.g., ['python', 'script.py']). This overrides a container's default entrypoint (e.g. ENTRYPOINT in Docker) if provided.""" + + command: list[str] = Field(default_factory=list) + """The command to execute as a list of strings (e.g., ['python', 'script.py']). This overrides a container's default commands (e.g. CMD in Docker) if provided.""" + + +class ComputeResourceSpec(BaseModel): + """Resource specification.""" + + cpu: str | None = Field(default=None, description="CPU specification (e.g., '250m', '1', '2.5').") + memory: str | None = Field(default=None, description="Memory specification (e.g., '128Mi', '1Gi', '512M').") + + +class ComputeResources(BaseModel): + """Resource requirements matching k8s ResourceRequirements format.""" + + requests: ComputeResourceSpec = Field( + default_factory=ComputeResourceSpec, description="Minimum resources requested for the container." + ) + + limits: ComputeResourceSpec = Field( + default_factory=ComputeResourceSpec, description="Maximum resources the container can use." + ) + + num_nodes: int = Field(default=1, ge=1, description="Number of nodes to use.") + + num_gpus: int | None = Field(default=None, description="Step requesting number of GPUs.") + + shm_size: str | None = Field( + default=None, + description="Shared memory (/dev/shm) size as a Kubernetes quantity (e.g. '1Gi', '4Gi'). " + "Used for GPU and distributed-GPU job executors. When unset, defaults to 1Gi per allocated GPU.", + ) + + @field_validator("shm_size") + @classmethod + def validate_shm_size_quantity(cls, v: str | None) -> str | None: + if v is None: + return None + s = v.strip() + if not s: + raise ValueError("shm_size cannot be empty or whitespace-only") + if not _SHM_QUANTITY_RE.fullmatch(s): + raise ValueError( + "shm_size must use a megabyte/gigabyte-scale suffix: Mi, Gi, M, or G (e.g. '1Gi', '512Mi', '2G')." + ) + return s + + +class TaskSpec(BaseModel): + """ + Specification for a task to be executed. + + Defines the command and arguments for a job task. + """ + + command: list[str] + """The command to execute as a list of strings (e.g., ['python', 'script.py']).""" + + args: list[str] | str + """Arguments to pass to the command. Can be a list of strings or a single string.""" + + +class CPUExecutionProvider(BaseModel): + """ + CPU-based execution provider. + + Provides configuration for running jobs on CPU resources with + resource requests and limits. + """ + + provider: Literal["cpu"] = "cpu" + """The provider type, always 'cpu' for CPU execution.""" + + profile: str = "default" + """The execution profile to use. Defaults to 'default'.""" + + container: ContainerSpec + """Container specification defining the execution environment.""" + + resources: ComputeResources = Field( + default_factory=ComputeResources, description="Resource requests and limits for CPU execution." + ) + + +class GPUExecutionProvider(BaseModel): + """ + GPU-based execution provider. + + Provides configuration for running jobs on GPU resources with + resource requests and limits. + """ + + provider: Literal["gpu"] = "gpu" + """The provider type, always 'gpu' for GPU execution.""" + + profile: str = "default" + """The execution profile to use. Defaults to 'default'.""" + + container: ContainerSpec + """Container specification defining the execution environment.""" + + resources: ComputeResources = Field( + default_factory=ComputeResources, description="Resource requests and limits for GPU execution." + ) + + +class DistributedGPUExecutionProvider(BaseModel): + """ + GPU-based execution provider. + + Provides configuration for running jobs on GPU resources with + resource requests and limits. + """ + + provider: Literal["gpu_distributed"] = "gpu_distributed" + """The provider type, always 'gpu_distributed' for distributed GPU execution.""" + + profile: str = "default" + """The execution profile to use. Defaults to 'default'.""" + + container: ContainerSpec + """Container specification defining the execution environment.""" + + resources: ComputeResources = Field( + default_factory=ComputeResources, description="Resource requests and limits for distributed GPU execution." + ) + + +class SubprocessExecutionProvider(BaseModel): + """Host subprocess execution provider.""" + + provider: Literal["subprocess"] = "subprocess" + """The provider type, always 'subprocess' for host subprocess execution.""" + + profile: str = "default" + """The execution profile to use. Defaults to 'default'.""" + + command: list[str] + """The host command to execute as a list of strings (e.g., ['python', '-m', 'my_task']).""" + + @model_validator(mode="after") + def validate_command(self) -> Self: + if not self.command: + raise ValueError("subprocess execution requires command to be set") + return self + + +# Type alias for the current execution provider implementation +ExecutionProviderT = Union[ + CPUExecutionProvider, GPUExecutionProvider, DistributedGPUExecutionProvider, SubprocessExecutionProvider +] +"""Type alias representing the current execution provider type.""" + +# Discriminated union type for execution providers +Provider = Annotated[ + ExecutionProviderT, + Field(discriminator="provider"), +] +""" +Discriminated union type for execution providers. + +Uses the 'provider' field to determine the specific provider type. +Currently supports CPU execution providers, with extensibility for future provider types. +""" diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/result_manager.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/result_manager.py index b65a31de3f..b4c910e467 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/result_manager.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/result_manager.py @@ -9,15 +9,19 @@ from pathlib import Path from typing import Generic, Literal, Type, TypeVar, overload -from nemo_platform import APIError, AsyncNeMoPlatform, ConflictError, NeMoPlatform +from nemo_platform import AsyncNeMoPlatform, NeMoPlatform from nemo_platform.filesets import parse_fileset_ref -from nemo_platform.types import PlatformJobResultResponse as SDKPlatformJobResult +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ConflictError as ClientConflictError +from nemo_platform_plugin.client.errors import NemoClientError +from nemo_platform_plugin.jobs.client import AsyncJobsClient, JobsClient from nemo_platform_plugin.jobs.constants import NEMO_JOB_WORKSPACE_ENVVAR from nemo_platform_plugin.jobs.file_manager import ( AsyncFilesetFileManager, FilesetFileManager, TmpDirPath, ) +from nemo_platform_plugin.jobs.schemas import PlatformJobResultCreateRequest, PlatformJobResultResponse logger = logging.getLogger(__name__) @@ -64,7 +68,8 @@ def _result_remote_path(self, attempt_id: str, result_name: str) -> str: class ResultManager(BaseResultManager[Type[FilesetFileManager], NeMoPlatform]): def _fetch_job_metadata(self) -> tuple[str, str]: """Fetch job and return (attempt_id, fileset_name).""" - job = self.jobs_sdk.jobs.retrieve(name=self.job_name, workspace=self.workspace) + jobs = client_from_platform(self.jobs_sdk, JobsClient) + job = jobs.get_job(name=self.job_name, workspace=self.workspace).data() attempt_id = self.attempt_id if self.attempt_id is not None else job.attempt_id return attempt_id, job.fileset @@ -81,7 +86,7 @@ def create_result( result_name: str, artifact_local_path: str | Path, ignore_patterns: list[str] | str | None = None, - ) -> SDKPlatformJobResult: + ) -> PlatformJobResultResponse: attempt_id, fileset_name = self._fetch_job_metadata() file_manager = self._create_file_manager(fileset_name) file_manager.validate_storage() @@ -90,23 +95,26 @@ def create_result( artifact_url = file_manager.upload( local_path=artifact_local_path, remote_path=remote_path, ignore_patterns=ignore_patterns ) + jobs = client_from_platform(self.jobs_sdk, JobsClient) try: - job_result = self.jobs_sdk.jobs.results.create( + job_result = jobs.create_job_result( name=result_name, job=self.job_name, workspace=self.workspace, - artifact_url=artifact_url, - artifact_storage_type=file_manager.storage_type().value, - ) - except ConflictError: + body=PlatformJobResultCreateRequest( + artifact_url=artifact_url, + artifact_storage_type=file_manager.storage_type(), + ), + ).data() + except ClientConflictError: # Result already exists - fetch and return the existing one # This supports the use case of saving partial results across multiple batches - job_result = self.jobs_sdk.jobs.results.retrieve( + job_result = jobs.get_job_result( name=result_name, job=self.job_name, workspace=self.workspace, - ) - except APIError as e: + ).data() + except NemoClientError as e: msg = f"Error creating job result: {str(e)}" logger.exception(msg) raise CreateJobResultError(msg) from e @@ -125,7 +133,8 @@ def download_artifact(self, artifact_url: str, local_dir: str | Path | None = No class AsyncResultManager(BaseResultManager[Type[AsyncFilesetFileManager], AsyncNeMoPlatform]): async def _fetch_job_metadata(self) -> tuple[str, str]: """Fetch job and return (attempt_id, fileset_name).""" - job = await self.jobs_sdk.jobs.retrieve(name=self.job_name, workspace=self.workspace) + jobs = client_from_platform(self.jobs_sdk, AsyncJobsClient) + job = (await jobs.get_job(name=self.job_name, workspace=self.workspace)).data() attempt_id = self.attempt_id if self.attempt_id is not None else job.attempt_id return attempt_id, job.fileset @@ -142,7 +151,7 @@ async def create_result( result_name: str, artifact_local_path: str | Path, ignore_patterns: list[str] | str | None = None, - ) -> SDKPlatformJobResult: + ) -> PlatformJobResultResponse: attempt_id, fileset_name = await self._fetch_job_metadata() file_manager = self._create_file_manager(fileset_name) await file_manager.validate_storage() @@ -151,23 +160,30 @@ async def create_result( artifact_url = await file_manager.upload( local_path=artifact_local_path, remote_path=remote_path, ignore_patterns=ignore_patterns ) + jobs = client_from_platform(self.jobs_sdk, AsyncJobsClient) try: - job_result = await self.jobs_sdk.jobs.results.create( - name=result_name, - job=self.job_name, - workspace=self.workspace, - artifact_url=artifact_url, - artifact_storage_type=file_manager.storage_type().value, - ) - except ConflictError: + job_result = ( + await jobs.create_job_result( + name=result_name, + job=self.job_name, + workspace=self.workspace, + body=PlatformJobResultCreateRequest( + artifact_url=artifact_url, + artifact_storage_type=file_manager.storage_type(), + ), + ) + ).data() + except ClientConflictError: # Result already exists - fetch and return the existing one # This supports the use case of saving partial results across multiple batches - job_result = await self.jobs_sdk.jobs.results.retrieve( - name=result_name, - job=self.job_name, - workspace=self.workspace, - ) - except APIError as e: + job_result = ( + await jobs.get_job_result( + name=result_name, + job=self.job_name, + workspace=self.workspace, + ) + ).data() + except NemoClientError as e: msg = f"Error creating job result: {str(e)}" logger.exception(msg) raise CreateJobResultError(msg) from e diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/spec.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/spec.py new file mode 100644 index 0000000000..0f66b75d5f --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/spec.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Job specification types for the Jobs service. + +``PlatformJobSpec`` and its children describe *what* a job runs: an ordered +list of steps, each with an executor (:mod:`nemo_platform_plugin.jobs.providers`), +environment, and configuration. This is the core of the job-creation request +body (``CreatePlatformJobRequest.platform_spec``). + +Pure pydantic — no server, docker, or kubernetes runtime dependencies. The +Jobs service (``nmp.core.jobs.app.schemas``) re-exports from here so both the +server and the typed HTTP client share one definition. +""" + +from __future__ import annotations + +from typing import Optional, Self + +from nemo_platform_plugin.jobs.constants import PERSISTENT_JOB_STORAGE_PATH_ENVVAR +from nemo_platform_plugin.jobs.providers import Provider +from pydantic import BaseModel, ConfigDict, Field, model_validator + +# RFC 1035 compliant pattern with temporary support for special characters. +# Mirrors ``nmp.common.entities.constants.NAME_PATTERN`` — inlined so this +# module stays a dependency-free leaf node (see files/types.py for the same +# pattern of inlining name constraints into the plugin). +NAME_PATTERN = r"^[a-z](?!.*--)[a-z0-9\-@.+_]{1,62}(? Self: + # Ensure one of value or from_secret is provided + if self.value is None and self.from_secret is None: + raise ValueError("Either value or from_secret must be provided for environment variables.") + + # Ensure only one of value or from_secret is provided + if self.value is not None and self.from_secret is not None: + raise ValueError("Only one of value or from_secret can be provided for environment variables.") + + return self + + +class StepLifecycle(BaseModel): + """Controller-level lifecycle configuration for a job step. + + These settings control how the jobs controller manages the step, + as opposed to ``config`` which is the task payload forwarded to + the container. + """ + + staleness_timeout_seconds: int = Field( + default=0, + description="If every active task in the step goes this many seconds without an update, the step is terminated. " + "A value of 0 disables staleness detection.", + ) + + +class PlatformJobStepSpec(BaseModel): + """Specification for a single step in a platform job.""" + + name: str = Field( + description=f"The name of the step. Must be unique for all steps in a job. {NAME_PATTERN_DESCRIPTION}", + pattern=NAME_PATTERN, + examples=["preprocess", "train-model", "eval-step-v1"], + ) + environment: Optional[list[PlatformJobEnvironmentVariable]] = Field( + default=None, description="Environment variables for the step" + ) + executor: Provider = Field(description="The executor for the step") + config: dict = Field(default_factory=dict, description="Configuration for the step") + lifecycle: StepLifecycle = Field( + default_factory=StepLifecycle, description="Lifecycle configuration settings for the step" + ) + + @property + def requires_persistent_storage(self) -> bool: + """ + Determine if the step requires persistent storage. + + This is determined by checking if the step has an environment variable + matching the value of PERSISTENT_JOB_STORAGE_PATH_ENVVAR. + """ + for envvar in self.environment or []: + if envvar.name == PERSISTENT_JOB_STORAGE_PATH_ENVVAR: + return True + return False + + model_config = ConfigDict(regex_engine="python-re") + + +class PlatformJobSpec(BaseModel): + """Specification for a platform job, containing steps and secrets.""" + + steps: list[PlatformJobStepSpec] = Field(description="List of steps to be executed in the job") + + @model_validator(mode="after") + def validate_steps(self) -> Self: + # Ensure there is at least one step. + if not self.steps: + raise ValueError("At least one step is required in the job specification.") + + # Ensure that each step has a unique name. + step_names = [step.name for step in self.steps] + if len(step_names) != len(set(step_names)): + raise ValueError("Each step must have a unique name.") + return self + + +# String aliases for provider / profile / backend references. +ProviderRef = str +ProfileRef = str +BackendRef = str + + +class BaseExecutionProfile(BaseModel): + """Execution configuration for a job. + + Base class for the concrete execution profiles in + :mod:`nemo_platform_plugin.jobs.execution_profiles`. + """ + + provider: ProviderRef = Field( + default="cpu", + description="The compute provider for the executor, e.g., cpu, gpu", + ) + profile: str = Field( + default="default", + description="The profile name for the executor, e.g., high_priority_a100, low_priority, etc.", + ) + + @property + def supports_persistent_storage(self) -> bool: + """Indicates if the execution profile supports persistent storage.""" + return False + + def __str__(self) -> str: + return f"{self.profile}:{self.provider}" diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/types.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/types.py new file mode 100644 index 0000000000..0ab3757f88 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/types.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Request/response DTOs for the Jobs service HTTP contract. + +These types define what job endpoints accept and return. Both the server +(FastAPI routes in ``nmp.core.jobs.api``) and the typed HTTP client import +from here — one source of truth, no Stainless-generated duplicates. + +The deep spec types live in sibling modules: +- :mod:`nemo_platform_plugin.jobs.spec` — ``PlatformJobSpec`` and children +- :mod:`nemo_platform_plugin.jobs.providers` — the executor tree +- :mod:`nemo_platform_plugin.jobs.execution_profiles` — backend profiles +- :mod:`nemo_platform_plugin.jobs.schemas` — status/result/log DTOs +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from enum import Enum +from typing import Any, NotRequired, Optional, TypedDict + +from nemo_platform_plugin.jobs.schemas import ( + PlatformJobResultResponse, + PlatformJobStatus, +) +from nemo_platform_plugin.jobs.spec import PlatformJobSpec, PlatformJobStepSpec +from nemo_platform_plugin.schema import Value +from pydantic import BaseModel, Field, RootModel + +# --------------------------------------------------------------------------- +# Auth context (data-only mirror of nmp.common.auth.AuthContext) +# --------------------------------------------------------------------------- + + +class AuthContext(BaseModel): + """Auth context captured at resource creation for delegated access. + + Stores a snapshot of the creating principal's identity so that controllers + can later act on their behalf (e.g., accessing secrets). + + This is the wire/data shape. The server's ``nmp.common.auth.AuthContext`` + adds ``from_principal`` / ``to_principal`` behaviour on top of the same + fields. + """ + + principal_id: str = Field(..., description="The principal's unique identifier") + principal_email: Optional[str] = Field(default=None, description="The principal's email address") + principal_groups: list[str] = Field(default_factory=list, description="Groups the principal belongs to") + principal_on_behalf_of: Optional[str] = Field( + default=None, description="If acting on behalf of another principal, their principal ID" + ) + principal_on_behalf_of_groups: Optional[list[str]] = Field( + default=None, description="Groups the on-behalf-of principal belongs to" + ) + principal_on_behalf_of_email: Optional[str] = Field( + default=None, description="The on-behalf-of principal's email address" + ) + + +# --------------------------------------------------------------------------- +# Sort fields +# --------------------------------------------------------------------------- + + +class PlatformJobLogSortField(str, Enum): + TIMESTAMP_ASC = "timestamp" + TIMESTAMP_DESC = "-timestamp" + + def get_field_name(self) -> str: + return self.value.lstrip("-") + + def get_sort_direction(self) -> str: + return "desc" if self.value.startswith("-") else "asc" + + +class PlatformJobSortField(str, Enum): + CREATED_AT_ASC = "created_at" + CREATED_AT_DESC = "-created_at" + UPDATED_AT_ASC = "updated_at" + UPDATED_AT_DESC = "-updated_at" + + def get_field_name(self) -> str: + return self.value.lstrip("-") + + def get_sort_direction(self) -> str: + return "desc" if self.value.startswith("-") else "asc" + + +class PlatformJobAttemptSortField(str, Enum): + SEQ_ASC = "seq" + SEQ_DESC = "-seq" + + def get_field_name(self) -> str: + return self.value.lstrip("-") + + def get_sort_direction(self) -> str: + return "desc" if self.value.startswith("-") else "asc" + + +# --------------------------------------------------------------------------- +# Response DTOs +# --------------------------------------------------------------------------- + + +class PlatformJobResponse(BaseModel): + """Response model for a platform job.""" + + id: str + attempt_id: str + name: str + workspace: str = Field(..., description="Workspace identifier") + project: Optional[str] = Field(default=None, description="Project URN") + description: str | None = None + source: str + spec: dict[str, Any] = Field(default_factory=dict, description="Job Spec") + platform_spec: PlatformJobSpec + fileset: str = Field(..., description="Fileset ID for storing job artifacts") + status: PlatformJobStatus + status_details: dict[str, Any] = Field(default_factory=dict, description="Details about the job status") + error_details: dict[str, Any] | None = None + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + ownership: Optional[dict[str, Any]] = None + custom_fields: Optional[dict[str, Any]] = Field(default=None, description="Custom Fields") + + +class PlatformJobStepResponse(BaseModel): + """Response model for a job step (wire shape of the ``PlatformJobStep`` entity).""" + + id: str + entity_id: str + parent: str = Field(..., description="Parent entity ID (the attempt ID)") + attempt_id: str = Field(..., description="Parent attempt ID") + name: str | None = None + workspace: str + project: str | None = None + config: dict[str, Any] = Field(default_factory=dict, description="Configuration for the step") + status: PlatformJobStatus = PlatformJobStatus.CREATED + status_details: dict[str, Any] = Field(default_factory=dict, description="Status details") + error_details: dict[str, Any] | None = None + created_at: datetime | None = None + created_by: str | None = None + updated_at: datetime | None = None + updated_by: str | None = None + + +class PlatformJobStepWithContext(BaseModel): + """Step with additional context from parent job/attempt.""" + + id: str + job: str + attempt_id: str + fileset: str + workspace: str + name: str + step_spec: PlatformJobStepSpec | None = None + status: PlatformJobStatus = PlatformJobStatus.CREATED + status_details: dict[str, Any] | None = None + error_details: dict[str, Any] | None = None + auth_context: Optional[AuthContext] = Field(default=None, description="Auth context for task execution") + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + + +class PlatformJobTaskResponse(BaseModel): + """Response model for a job task (wire shape of the ``PlatformJobTask`` entity).""" + + id: str + entity_id: str + parent: str = Field(..., description="Parent entity ID (the step ID)") + step_id: str = Field(..., description="Parent step ID") + name: str | None = None + workspace: str + project: str | None = None + status: PlatformJobStatus = PlatformJobStatus.PENDING + status_details: dict[str, Any] = Field(default_factory=dict, description="Details about the task status") + error_details: dict[str, Any] | None = None + error_stack: str | None = None + created_at: datetime | None = None + created_by: str | None = None + updated_at: datetime | None = None + updated_by: str | None = None + + +class PlatformJobListResultResponse(Value): + """Response model for listing job results.""" + + data: list[PlatformJobResultResponse] + + +class PlatformJobListTaskResponse(Value): + """Response model for listing job tasks.""" + + data: list[PlatformJobTaskResponse] + + +# --------------------------------------------------------------------------- +# Request DTOs +# --------------------------------------------------------------------------- + + +class CreatePlatformJobRequest(BaseModel): + """Request model for creating a new platform job.""" + + name: Optional[str] = None + description: Optional[str] = None + project: Optional[str] = None + spec: dict + platform_spec: PlatformJobSpec + source: str + ownership: Optional[dict] = None + custom_fields: Optional[dict] = None + + +class PlatformJobTaskUpdate(BaseModel): + """Request model for updating a platform job task.""" + + status: PlatformJobStatus = PlatformJobStatus.PENDING + status_details: dict[str, Any] | None = None + error_details: dict[str, Any] | None = None + error_stack: str | None = None + + +class PlatformJobStatusUpdateRequest(BaseModel): + """Request model for updating job status.""" + + status: PlatformJobStatus = Field(..., description="The new status to set for the job.") + status_details: dict[str, Any] | None = Field( + default_factory=dict, description="Optional status details related to the status update." + ) + error_details: dict[str, Any] | None = Field( + default_factory=dict, description="Optional error details related to the status update." + ) + + +# Status-details PATCH body: a free-form dict of status details. The server +# accepts a bare JSON object (typed as ``dict[str, Any]``); the client uses the +# ``JobStatusDetailsUpdate`` RootModel wrapper so it can be passed as a typed +# request ``body`` (it serialises to the same bare object). +PlatformJobStatusDetailsUpdateRequest = dict[str, Any] + + +class JobStatusDetailsUpdate(RootModel[dict[str, Any]]): + """Client request body for ``update_job_status_details`` (a bare JSON object).""" + + +# NB: list *filter* models (``PlatformJobsListFilter`` etc.) are intentionally +# NOT defined here. They subclass the entity-store ``Filter`` (with field +# mapping / translation) and are server-side only. Clients pass a ``filter`` +# query-param string via the query-param TypedDicts below. + + +# --------------------------------------------------------------------------- +# Query parameter types (client-side) +# --------------------------------------------------------------------------- + + +class ListJobsQueryParams(TypedDict, total=False): + page: NotRequired[int] + page_size: NotRequired[int] + sort: NotRequired[str] + filter: NotRequired[str] + + +class ListStepsQueryParams(TypedDict, total=False): + page: NotRequired[int] + page_size: NotRequired[int] + sort: NotRequired[str] + filter: NotRequired[str] + + +class ListJobResultsQueryParams(TypedDict, total=False): + sort: NotRequired[str] + + +class JobLogsQueryParams(TypedDict, total=False): + limit: NotRequired[int] + page_cursor: NotRequired[str] + attempt_id: NotRequired[int] + step_id: NotRequired[str] + task_id: NotRequired[str] diff --git a/packages/nemo_platform_plugin/tests/client/test_adapter.py b/packages/nemo_platform_plugin/tests/client/test_adapter.py new file mode 100644 index 0000000000..48ca65656e --- /dev/null +++ b/packages/nemo_platform_plugin/tests/client/test_adapter.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import httpx +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.jobs.client import JobsClient + + +def test_client_from_platform_preserves_retry_count_with_nemoclient_defaults() -> None: + http_client = httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request))) + platform = NeMoPlatform( + base_url="http://test", + workspace="default", + max_retries=4, + http_client=http_client, + ) + + client = client_from_platform(platform, JobsClient) + + assert client.retry is not None + assert client.retry.max_retries == 4 + assert client.retry.retryable_status_codes == (502, 503, 504, 429) diff --git a/packages/nemo_platform_plugin/tests/client/test_client.py b/packages/nemo_platform_plugin/tests/client/test_client.py index cdcdd1c5ca..cbd1ee079e 100644 --- a/packages/nemo_platform_plugin/tests/client/test_client.py +++ b/packages/nemo_platform_plugin/tests/client/test_client.py @@ -7,9 +7,9 @@ import httpx import pytest -from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient +from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient, _type_adapter from nemo_platform_plugin.client.endpoint import delete, get, post -from nemo_platform_plugin.client.errors import NemoHTTPError, NotFoundError +from nemo_platform_plugin.client.errors import NemoHTTPError, NemoResponseValidationError, NotFoundError from nemo_platform_plugin.client.response import NemoResponse from pydantic import BaseModel @@ -25,6 +25,16 @@ class ItemResponse(BaseModel): name: str +def test_response_type_adapters_are_cached() -> None: + _type_adapter.cache_clear() + + first = _type_adapter(ItemResponse) + second = _type_adapter(ItemResponse) + + assert first is second + assert _type_adapter.cache_info().misses == 1 + + @post("/apis/test/v2/items") def CREATE_ITEM(body: ItemRequest) -> ItemResponse: raise NotImplementedError @@ -386,6 +396,23 @@ def test_error_response_raises_specific_subclass() -> None: assert exc_info.value.detail == "Not found" +def test_success_response_validation_error_uses_client_error_contract() -> None: + mock_http = MagicMock(spec=httpx.Client) + response = httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice"), + json={"id": "not-an-integer", "name": "alice"}, + ) + mock_http.request.return_value = response + + with pytest.raises(NemoResponseValidationError) as exc_info: + NemoClient(base_url=BASE, http_client=mock_http).send(GET_ITEM(name="alice")) + + assert exc_info.value.http_response is response + assert exc_info.value.status_code == 200 + assert exc_info.value.body == {"id": "not-an-integer", "name": "alice"} + + # --------------------------------------------------------------------------- # Per-request headers # --------------------------------------------------------------------------- diff --git a/packages/nemo_platform_plugin/tests/client/test_client_options.py b/packages/nemo_platform_plugin/tests/client/test_client_options.py index 391d87cb24..b0dee0fdf9 100644 --- a/packages/nemo_platform_plugin/tests/client/test_client_options.py +++ b/packages/nemo_platform_plugin/tests/client/test_client_options.py @@ -11,9 +11,15 @@ import pytest from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient from nemo_platform_plugin.client.endpoint import delete, get, post -from nemo_platform_plugin.client.errors import ConflictError, NemoHTTPError, NotFoundError +from nemo_platform_plugin.client.errors import ( + ConflictError, + NemoHTTPError, + NemoResponseValidationError, + NemoTransportError, + NotFoundError, +) from nemo_platform_plugin.client.method import method -from nemo_platform_plugin.client.types import PreparedRequest, RetryPolicy +from nemo_platform_plugin.client.types import BinaryContent, PreparedRequest, RetryPolicy, Stream from pydantic import BaseModel BASE = "http://test:8000" @@ -43,6 +49,16 @@ def DELETE_ITEM(*, name: str) -> None: raise NotImplementedError +@get("/apis/test/v2/download") +def DOWNLOAD() -> BinaryContent: + raise NotImplementedError + + +@get("/apis/test/v2/events") +def EVENTS() -> Stream[ItemResponse]: + raise NotImplementedError + + def _get_item_on_conflict(body: ItemRequest, workspace: str | None) -> PreparedRequest[ItemResponse]: """Resolver: on a create 409, retrieve the existing item by name.""" return GET_ITEM(name=body.name) @@ -392,6 +408,22 @@ def test_retry_on_transport_error(self) -> None: assert resp.body.name == "alice" assert mock_http.request.call_count == 2 + def test_exhausted_transport_error_is_wrapped(self) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = httpx.ConnectError("Connection refused") + client = NemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=2, backoff_base=0.0), + ) + + with pytest.raises(NemoTransportError) as exc_info: + client.send(GET_ITEM(name="alice")) + + assert isinstance(exc_info.value.error, httpx.ConnectError) + assert exc_info.value.request is None + assert mock_http.request.call_count == 3 + def test_per_request_retry_overrides_client_default(self) -> None: mock_http = MagicMock(spec=httpx.Client) mock_http.request.return_value = httpx.Response( @@ -427,6 +459,53 @@ def test_no_retry_without_policy(self) -> None: assert exc_info.value.status_code == 503 assert mock_http.request.call_count == 1 + def test_binary_stream_retries_before_returning_content(self) -> None: + attempts = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + if attempts == 1: + return httpx.Response(503, request=request, json={"detail": "unavailable"}) + return httpx.Response(200, request=request, content=b"artifact") + + client = NemoClient( + base_url=BASE, + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + retry=RetryPolicy(max_retries=1, backoff_base=0.0), + ) + + assert client.send(DOWNLOAD()).read() == b"artifact" + assert attempts == 2 + + def test_binary_stream_transport_failure_is_wrapped(self) -> None: + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("Connection refused", request=request) + + client = NemoClient( + base_url=BASE, + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + retry=RetryPolicy(max_retries=1, backoff_base=0.0), + ) + + with pytest.raises(NemoTransportError): + client.send(DOWNLOAD()).read() + + def test_stream_item_validation_failure_is_wrapped(self) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + request=request, + headers={"content-type": "application/x-ndjson"}, + content=b'{"id":"bad","name":"alice"}\n', + ) + + client = NemoClient(base_url=BASE, http_client=httpx.Client(transport=httpx.MockTransport(handler))) + + with pytest.raises(NemoResponseValidationError): + with client.send(EVENTS()).stream() as events: + list(events) + # --------------------------------------------------------------------------- # Async: RetryPolicy @@ -460,3 +539,42 @@ async def test_retry_on_503_async(self) -> None: assert resp.http_response.status_code == 200 assert resp.body.name == "alice" assert mock_http.request.call_count == 2 + + @pytest.mark.asyncio + async def test_exhausted_transport_error_is_wrapped_async(self) -> None: + mock_http = AsyncMock(spec=httpx.AsyncClient) + request = httpx.Request("GET", f"{BASE}/apis/test/v2/items/alice") + mock_http.request.side_effect = httpx.ConnectError("Connection refused", request=request) + client = AsyncNemoClient( + base_url=BASE, + http_client=mock_http, + retry=RetryPolicy(max_retries=1, backoff_base=0.0), + ) + + with pytest.raises(NemoTransportError) as exc_info: + await client.send(GET_ITEM(name="alice")) + + assert exc_info.value.request is request + assert mock_http.request.call_count == 2 + + @pytest.mark.asyncio + async def test_binary_stream_retries_async(self) -> None: + attempts = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise httpx.ConnectError("Connection refused", request=request) + return httpx.Response(200, request=request, content=b"artifact") + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + client = AsyncNemoClient( + base_url=BASE, + http_client=http_client, + retry=RetryPolicy(max_retries=1, backoff_base=0.0), + ) + response = await client.send(DOWNLOAD()) + assert await response.read() == b"artifact" + + assert attempts == 2 diff --git a/packages/nemo_platform_plugin/tests/client/test_pagination.py b/packages/nemo_platform_plugin/tests/client/test_pagination.py index 5e5caf0003..d08aa7d038 100644 --- a/packages/nemo_platform_plugin/tests/client/test_pagination.py +++ b/packages/nemo_platform_plugin/tests/client/test_pagination.py @@ -11,9 +11,10 @@ import pytest from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient from nemo_platform_plugin.client.endpoint import get +from nemo_platform_plugin.client.errors import NemoResponseValidationError from nemo_platform_plugin.client.method import method from nemo_platform_plugin.client.response import AsyncNemoPaginatedResponse, NemoPaginatedResponse -from nemo_platform_plugin.client.types import OffsetPagination, Paginated, RetryPolicy +from nemo_platform_plugin.client.types import CursorPagination, OffsetPagination, Paginated, RetryPolicy from pydantic import BaseModel BASE = "http://test:8000" @@ -101,10 +102,11 @@ def test_data_returns_page_result_with_metadata(self) -> None: page = resp.page() assert len(page.items) == 1 assert page.items[0].name == "a" - assert page.page == 1 - assert page.total_pages == 5 - assert page.total_results == 10 - assert page.page_size == 2 + assert page.metadata["page"] == 1 + assert page.metadata["total_pages"] == 5 + assert page.metadata["total_results"] == 10 + assert page.metadata["page_size"] == 2 + assert page.metadata["current_page_size"] == 1 # No additional requests for data() assert mock_http.request.call_count == 1 @@ -119,8 +121,8 @@ def test_empty_page(self) -> None: items = list(resp.items()) assert items == [] - def test_no_pagination_metadata(self) -> None: - """When pagination is None, treat as single page.""" + @pytest.mark.parametrize("iteration", ["page", "items", "pages"]) + def test_missing_pagination_metadata_is_invalid(self, iteration: str) -> None: mock_http = MagicMock(spec=httpx.Client) mock_http.request.return_value = httpx.Response( 200, @@ -131,9 +133,26 @@ def test_no_pagination_metadata(self) -> None: client = NemoClient(base_url=BASE, workspace="default", http_client=mock_http) resp = client.send(LIST_ITEMS()) - items = list(resp.items()) - assert len(items) == 1 - assert mock_http.request.call_count == 1 + with pytest.raises(NemoResponseValidationError): + if iteration == "page": + resp.page() + elif iteration == "items": + list(resp.items()) + else: + list(resp.pages()) + + def test_partial_pagination_metadata_is_invalid(self) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/workspaces/default/items"), + json={"data": [{"id": 1, "name": "a"}], "pagination": {"page": 1}}, + ) + + resp = NemoClient(base_url=BASE, workspace="default", http_client=mock_http).send(LIST_ITEMS()) + + with pytest.raises(NemoResponseValidationError): + resp.page() def test_page_query_param_passed_on_subsequent_pages(self) -> None: """Subsequent page fetches should include page=N in query params.""" @@ -151,6 +170,25 @@ def test_page_query_param_passed_on_subsequent_pages(self) -> None: second_call_params = mock_http.request.call_args_list[1][1]["params"] assert second_call_params["page"] == 2 + @pytest.mark.parametrize("iteration", ["items", "pages"]) + def test_iteration_continues_after_response_page(self, iteration: str) -> None: + """A request beginning after page one must not fetch that page again.""" + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + _page_response([{"id": 3, "name": "c"}], page=2, total_pages=3), + _page_response([{"id": 4, "name": "d"}], page=3, total_pages=3), + ] + response = NemoClient(base_url=BASE, workspace="default", http_client=mock_http).send(LIST_ITEMS()) + + if iteration == "items": + names = [item.name for item in response.items()] + else: + names = [item.name for page in response.pages() for item in page.items] + + assert names == ["c", "d"] + assert mock_http.request.call_count == 2 + assert mock_http.request.call_args_list[1].kwargs["params"]["page"] == 3 + # --------------------------------------------------------------------------- # Via method() descriptor @@ -213,10 +251,29 @@ async def test_async_data_returns_page_result(self) -> None: page = resp.page() assert len(page.items) == 1 - assert page.page == 1 - assert page.total_pages == 3 + assert page.metadata["page"] == 1 + assert page.metadata["total_pages"] == 3 assert mock_http.request.call_count == 1 + @pytest.mark.asyncio + @pytest.mark.parametrize("iteration", ["items", "pages"]) + async def test_async_iteration_continues_after_response_page(self, iteration: str) -> None: + mock_http = AsyncMock(spec=httpx.AsyncClient) + mock_http.request.side_effect = [ + _page_response([{"id": 3, "name": "c"}], page=2, total_pages=3), + _page_response([{"id": 4, "name": "d"}], page=3, total_pages=3), + ] + response = await AsyncNemoClient(base_url=BASE, workspace="default", http_client=mock_http).send(LIST_ITEMS()) + + if iteration == "items": + names = [item.name async for item in response.items()] + else: + names = [item.name async for page in response.pages() for item in page.items] + + assert names == ["c", "d"] + assert mock_http.request.call_count == 2 + assert mock_http.request.call_args_list[1].kwargs["params"]["page"] == 3 + # --------------------------------------------------------------------------- # Retry on subsequent pages @@ -337,3 +394,88 @@ def test_custom_page_param(self) -> None: second_call_params = mock_http.request.call_args_list[1][1]["params"] assert "offset" in second_call_params assert second_call_params["offset"] == 2 + + +# --------------------------------------------------------------------------- +# Cursor pagination +# --------------------------------------------------------------------------- + + +@get("/apis/test/v2/workspaces/{workspace}/logs") +def LIST_LOGS( + *, workspace: str | None = None, query_params: dict[str, str | int] | None = None +) -> Paginated[Item, CursorPagination]: + raise NotImplementedError + + +def _cursor_response( + items: list[dict], *, total: int, next_page: str | None, prev_page: str | None = None +) -> httpx.Response: + return httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/test/v2/workspaces/default/logs"), + json={ + "data": items, + "total": total, + "next_page": next_page, + "prev_page": prev_page, + }, + ) + + +class TestCursorPagination: + def test_page_exposes_typed_cursor_metadata(self) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.return_value = _cursor_response([{"id": 1, "name": "a"}], total=2, next_page="cursor-2") + + page = NemoClient(base_url=BASE, workspace="default", http_client=mock_http).send(LIST_LOGS()).page() + + assert [item.name for item in page.items] == ["a"] + assert page.metadata == {"total": 2, "next_page": "cursor-2", "prev_page": None} + + def test_items_follow_cursors_and_preserve_filters(self) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + _cursor_response([{"id": 1, "name": "a"}], total=2, next_page="cursor-2"), + _cursor_response([{"id": 2, "name": "b"}], total=2, next_page=None, prev_page="cursor-1"), + ] + response = NemoClient(base_url=BASE, workspace="default", http_client=mock_http).send( + LIST_LOGS(query_params={"limit": 1, "attempt_id": 3}) + ) + + assert [item.name for item in response.items()] == ["a", "b"] + assert mock_http.request.call_count == 2 + assert mock_http.request.call_args_list[1].kwargs["params"] == { + "limit": 1, + "attempt_id": 3, + "page_cursor": "cursor-2", + } + + def test_pages_follow_cursors_from_explicit_start(self) -> None: + mock_http = MagicMock(spec=httpx.Client) + mock_http.request.side_effect = [ + _cursor_response([{"id": 2, "name": "b"}], total=3, next_page="cursor-3", prev_page="cursor-1"), + _cursor_response([{"id": 3, "name": "c"}], total=3, next_page=None, prev_page="cursor-2"), + ] + response = NemoClient(base_url=BASE, workspace="default", http_client=mock_http).send( + LIST_LOGS(query_params={"page_cursor": "cursor-2"}) + ) + + pages = list(response.pages()) + + assert [[item.name for item in page.items] for page in pages] == [["b"], ["c"]] + assert pages[0].metadata["prev_page"] == "cursor-1" + assert pages[1].metadata["next_page"] is None + assert mock_http.request.call_args_list[1].kwargs["params"]["page_cursor"] == "cursor-3" + + @pytest.mark.asyncio + async def test_async_items_follow_cursors(self) -> None: + mock_http = AsyncMock(spec=httpx.AsyncClient) + mock_http.request.side_effect = [ + _cursor_response([{"id": 1, "name": "a"}], total=2, next_page="cursor-2"), + _cursor_response([{"id": 2, "name": "b"}], total=2, next_page=None, prev_page="cursor-1"), + ] + response = await AsyncNemoClient(base_url=BASE, workspace="default", http_client=mock_http).send(LIST_LOGS()) + + assert [item.name async for item in response.items()] == ["a", "b"] + assert mock_http.request.call_args_list[1].kwargs["params"]["page_cursor"] == "cursor-2" diff --git a/packages/nemo_platform_plugin/tests/client/test_typing.py b/packages/nemo_platform_plugin/tests/client/test_typing.py new file mode 100644 index 0000000000..19a45da4a4 --- /dev/null +++ b/packages/nemo_platform_plugin/tests/client/test_typing.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Static return-type contracts for the typed client. + +The functions in this module are checked by ``ty`` but intentionally not +collected by pytest. They ensure endpoint annotations flow through prepared +requests and both client implementations without being erased to ``Any``. +""" + +from collections.abc import AsyncIterator, Iterator +from typing import assert_type + +from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient, _parse_json_body +from nemo_platform_plugin.client.endpoint import get +from nemo_platform_plugin.client.response import ( + AsyncNemoPaginatedResponse, + AsyncNemoStreamResponse, + NemoPaginatedResponse, + NemoResponse, + NemoStreamResponse, + PageResult, +) +from nemo_platform_plugin.client.types import ( + CursorPagination, + CursorPaginationMetadata, + OffsetPagination, + OffsetPaginationMetadata, + Paginated, + PreparedRequest, + Stream, +) +from nemo_platform_plugin.files.client import AsyncFilesClient, FilesClient +from nemo_platform_plugin.files.types import FilesetOutput +from nemo_platform_plugin.jobs.client import AsyncJobsClient, JobsClient +from nemo_platform_plugin.jobs.schemas import PlatformJobLog +from nemo_platform_plugin.jobs.types import PlatformJobResponse, PlatformJobStepWithContext +from nemo_platform_plugin.secrets.client import AsyncSecretsClient, SecretsClient +from nemo_platform_plugin.secrets.types import PlatformSecretResponse +from pydantic import BaseModel + + +class Item(BaseModel): + name: str + + +@get("/items/{name}") +def get_item(*, name: str) -> Item: + raise NotImplementedError + + +@get("/items") +def get_item_list() -> list[Item]: + raise NotImplementedError + + +@get("/items/pages") +def get_item_pages() -> Paginated[Item]: + raise NotImplementedError + + +@get("/items/cursor-pages") +def get_cursor_item_pages() -> Paginated[Item, CursorPagination]: + raise NotImplementedError + + +@get("/items/stream") +def get_item_stream() -> Stream[Item]: + raise NotImplementedError + + +def _check_sync_return_types(client: NemoClient) -> None: + item_request = get_item(name="one") + assert_type(item_request, PreparedRequest[Item]) + assert_type(client.send(item_request), NemoResponse[Item]) + + list_request = get_item_list() + assert_type(list_request, PreparedRequest[list[Item]]) + assert_type(client.send(list_request), NemoResponse[list[Item]]) + pages = client.send(get_item_pages()) + assert_type(pages, NemoPaginatedResponse[Item, OffsetPagination]) + assert_type(pages.page(), PageResult[Item, OffsetPaginationMetadata]) + assert_type(pages.items(), Iterator[Item]) + assert_type(pages.pages(), Iterator[PageResult[Item, OffsetPaginationMetadata]]) + cursor_pages = client.send(get_cursor_item_pages()) + assert_type(cursor_pages, NemoPaginatedResponse[Item, CursorPagination]) + assert_type(cursor_pages.page(), PageResult[Item, CursorPaginationMetadata]) + assert_type(cursor_pages.items(), Iterator[Item]) + assert_type(cursor_pages.pages(), Iterator[PageResult[Item, CursorPaginationMetadata]]) + assert_type(client.send(get_item_stream()), NemoStreamResponse[Item]) + + assert_type(_parse_json_body(Item, {"name": "one"}), Item) + assert_type(_parse_json_body(list[Item], [{"name": "one"}]), list[Item]) + + +async def _check_async_return_types(client: AsyncNemoClient) -> None: + assert_type(await client.send(get_item(name="one")), NemoResponse[Item]) + assert_type(await client.send(get_item_list()), NemoResponse[list[Item]]) + pages = await client.send(get_item_pages()) + assert_type(pages, AsyncNemoPaginatedResponse[Item, OffsetPagination]) + assert_type(pages.page(), PageResult[Item, OffsetPaginationMetadata]) + assert_type(pages.items(), AsyncIterator[Item]) + assert_type(pages.pages(), AsyncIterator[PageResult[Item, OffsetPaginationMetadata]]) + cursor_pages = await client.send(get_cursor_item_pages()) + assert_type(cursor_pages, AsyncNemoPaginatedResponse[Item, CursorPagination]) + assert_type(cursor_pages.page(), PageResult[Item, CursorPaginationMetadata]) + assert_type(cursor_pages.items(), AsyncIterator[Item]) + assert_type(cursor_pages.pages(), AsyncIterator[PageResult[Item, CursorPaginationMetadata]]) + assert_type(await client.send(get_item_stream()), AsyncNemoStreamResponse[Item]) + + +def _check_jobs_return_types(client: JobsClient) -> None: + logs = client.list_job_logs(name="job") + assert_type(logs, NemoPaginatedResponse[PlatformJobLog, CursorPagination]) + assert_type(logs.page(), PageResult[PlatformJobLog, CursorPaginationMetadata]) + + +async def _check_async_jobs_return_types(client: AsyncJobsClient) -> None: + logs = await client.list_job_logs(name="job") + assert_type(logs, AsyncNemoPaginatedResponse[PlatformJobLog, CursorPagination]) + assert_type(logs.page(), PageResult[PlatformJobLog, CursorPaginationMetadata]) + + +def _check_offset_client_return_types( + files: FilesClient, + secrets: SecretsClient, + jobs: JobsClient, +) -> None: + filesets = files.list_filesets() + assert_type(filesets, NemoPaginatedResponse[FilesetOutput, OffsetPagination]) + assert_type(filesets.page(), PageResult[FilesetOutput, OffsetPaginationMetadata]) + + secret_pages = secrets.list_secrets() + assert_type(secret_pages, NemoPaginatedResponse[PlatformSecretResponse, OffsetPagination]) + assert_type(secret_pages.page(), PageResult[PlatformSecretResponse, OffsetPaginationMetadata]) + + job_pages = jobs.list_jobs() + assert_type(job_pages, NemoPaginatedResponse[PlatformJobResponse, OffsetPagination]) + assert_type(job_pages.page(), PageResult[PlatformJobResponse, OffsetPaginationMetadata]) + + step_pages = jobs.list_steps(name="job") + assert_type(step_pages, NemoPaginatedResponse[PlatformJobStepWithContext, OffsetPagination]) + assert_type(step_pages.page(), PageResult[PlatformJobStepWithContext, OffsetPaginationMetadata]) + + +async def _check_async_offset_client_return_types( + files: AsyncFilesClient, + secrets: AsyncSecretsClient, + jobs: AsyncJobsClient, +) -> None: + filesets = await files.list_filesets() + assert_type(filesets, AsyncNemoPaginatedResponse[FilesetOutput, OffsetPagination]) + assert_type(filesets.page(), PageResult[FilesetOutput, OffsetPaginationMetadata]) + + secret_pages = await secrets.list_secrets() + assert_type(secret_pages, AsyncNemoPaginatedResponse[PlatformSecretResponse, OffsetPagination]) + assert_type(secret_pages.page(), PageResult[PlatformSecretResponse, OffsetPaginationMetadata]) + + job_pages = await jobs.list_jobs() + assert_type(job_pages, AsyncNemoPaginatedResponse[PlatformJobResponse, OffsetPagination]) + assert_type(job_pages.page(), PageResult[PlatformJobResponse, OffsetPaginationMetadata]) + + step_pages = await jobs.list_steps(name="job") + assert_type(step_pages, AsyncNemoPaginatedResponse[PlatformJobStepWithContext, OffsetPagination]) + assert_type(step_pages.page(), PageResult[PlatformJobStepWithContext, OffsetPaginationMetadata]) diff --git a/packages/nemo_platform_plugin/tests/jobs/test_client.py b/packages/nemo_platform_plugin/tests/jobs/test_client.py new file mode 100644 index 0000000000..4162d02d6d --- /dev/null +++ b/packages/nemo_platform_plugin/tests/jobs/test_client.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the JobsClient / AsyncJobsClient via a mocked httpx transport. + +Drives ``send()`` end-to-end (path resolution, body serialization, response +unwrapping, pagination, binary, error mapping) without a network.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +from nemo_platform_plugin.client.errors import NotFoundError +from nemo_platform_plugin.jobs.client import AsyncJobsClient, JobsClient +from nemo_platform_plugin.jobs.types import CreatePlatformJobRequest + +BASE = "http://test:8000" + +_JOB_JSON = { + "id": "job-1", + "attempt_id": "att-1", + "name": "my-job", + "workspace": "default", + "source": "test", + "spec": {}, + "platform_spec": {"steps": [{"name": "step-one", "executor": {"provider": "cpu", "container": {"image": "x"}}}]}, + "fileset": "fs-1", + "status": "created", +} + + +def _mock_http(response: httpx.Response) -> MagicMock: + mock = MagicMock(spec=httpx.Client) + mock.request.return_value = response + return mock + + +def test_create_job_serializes_body_and_unwraps() -> None: + mock_http = _mock_http( + httpx.Response( + 201, + request=httpx.Request("POST", f"{BASE}/apis/jobs/v2/workspaces/default/jobs"), + json=_JOB_JSON, + ) + ) + client = JobsClient(base_url=BASE, workspace="default", http_client=mock_http) + body = CreatePlatformJobRequest( + spec={}, + source="test", + platform_spec={"steps": [{"name": "step-one", "executor": {"provider": "cpu", "container": {"image": "x"}}}]}, + ) + resp = client.create_job(body=body) + + assert resp.data().name == "my-job" + _, kwargs = mock_http.request.call_args + assert b'"source":"test"' in kwargs["content"] + + +def test_get_job_resolves_path() -> None: + mock_http = _mock_http( + httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/jobs/v2/workspaces/default/jobs/my-job"), + json=_JOB_JSON, + ) + ) + client = JobsClient(base_url=BASE, workspace="default", http_client=mock_http) + job = client.get_job(name="my-job").data() + + assert job.id == "job-1" + args, _ = mock_http.request.call_args + assert "/apis/jobs/v2/workspaces/default/jobs/my-job" in str(mock_http.request.call_args) + + +def test_list_jobs_paginated_items() -> None: + mock_http = _mock_http( + httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/jobs/v2/workspaces/default/jobs"), + json={ + "data": [_JOB_JSON], + "pagination": { + "page": 1, + "page_size": 10, + "current_page_size": 1, + "total_pages": 1, + "total_results": 1, + }, + }, + ) + ) + client = JobsClient(base_url=BASE, workspace="default", http_client=mock_http) + jobs = list(client.list_jobs().items()) + + assert len(jobs) == 1 + assert jobs[0].name == "my-job" + + +def test_delete_job_returns_none() -> None: + mock_http = _mock_http( + httpx.Response( + 204, + request=httpx.Request("DELETE", f"{BASE}/apis/jobs/v2/workspaces/default/jobs/my-job"), + ) + ) + client = JobsClient(base_url=BASE, workspace="default", http_client=mock_http) + assert client.delete_job(name="my-job").data() is None + + +def test_download_job_result_reads_bytes() -> None: + mock_http = MagicMock(spec=httpx.Client) + stream_ctx = MagicMock() + raw = httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/jobs/v2/workspaces/default/jobs/j/results/out/download"), + content=b"artifact-bytes", + ) + stream_ctx.__enter__.return_value = raw + stream_ctx.__exit__.return_value = False + mock_http.stream.return_value = stream_ctx + + client = JobsClient(base_url=BASE, workspace="default", http_client=mock_http) + data = client.download_job_result(job="j", name="out").read() + + assert data == b"artifact-bytes" + + +def test_get_job_not_found_maps_error() -> None: + mock_http = _mock_http( + httpx.Response( + 404, + request=httpx.Request("GET", f"{BASE}/apis/jobs/v2/workspaces/default/jobs/missing"), + json={"detail": "Job not found"}, + ) + ) + client = JobsClient(base_url=BASE, workspace="default", http_client=mock_http) + with pytest.raises(NotFoundError) as exc: + client.get_job(name="missing") + assert exc.value.status_code == 404 + + +# A JSON array response (not an object) — the shape that broke ``send()`` when +# the endpoint's return annotation is a bare ``list[...]`` generic. +_PROFILES_JSON = [ + {"backend": "subprocess", "provider": "subprocess", "profile": "default", "config": {}}, + {"backend": "e2e", "provider": "cpu", "profile": "default", "config": {}}, +] + + +def test_get_execution_profiles_parses_list_response() -> None: + mock_http = _mock_http( + httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/jobs/v2/execution-profiles"), + json=_PROFILES_JSON, + ) + ) + client = JobsClient(base_url=BASE, workspace="default", http_client=mock_http) + + profiles = client.get_execution_profiles().data() + + assert isinstance(profiles, list) + assert len(profiles) == 2 + assert {p.backend for p in profiles} == {"subprocess", "e2e"} + + +@pytest.mark.asyncio +async def test_async_get_execution_profiles_parses_list_response() -> None: + mock_http = MagicMock(spec=httpx.AsyncClient) + mock_http.request = AsyncMock( + return_value=httpx.Response( + 200, + request=httpx.Request("GET", f"{BASE}/apis/jobs/v2/execution-profiles"), + json=_PROFILES_JSON, + ) + ) + client = AsyncJobsClient(base_url=BASE, workspace="default", http_client=mock_http) + + profiles = (await client.get_execution_profiles()).data() + + assert isinstance(profiles, list) + assert {p.backend for p in profiles} == {"subprocess", "e2e"} diff --git a/packages/nemo_platform_plugin/tests/jobs/test_endpoints.py b/packages/nemo_platform_plugin/tests/jobs/test_endpoints.py new file mode 100644 index 0000000000..bcb400be96 --- /dev/null +++ b/packages/nemo_platform_plugin/tests/jobs/test_endpoints.py @@ -0,0 +1,280 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for Jobs service endpoint definitions.""" + +from __future__ import annotations + +import json +from typing import get_args, get_origin + +from nemo_platform_plugin.client.types import BinaryContent, CursorPagination, Paginated, PreparedRequest +from nemo_platform_plugin.jobs import endpoints +from nemo_platform_plugin.jobs.schemas import ( + PlatformJobLog, + PlatformJobResultCreateRequest, + PlatformJobStatusResponse, +) +from nemo_platform_plugin.jobs.types import ( + CreatePlatformJobRequest, + JobStatusDetailsUpdate, + PlatformJobResponse, + PlatformJobStatusUpdateRequest, + PlatformJobStepResponse, + PlatformJobStepWithContext, + PlatformJobTaskResponse, + PlatformJobTaskUpdate, +) + + +def _create_request() -> CreatePlatformJobRequest: + return CreatePlatformJobRequest( + spec={}, + source="test", + platform_spec={"steps": [{"name": "step-one", "executor": {"provider": "cpu", "container": {"image": "img"}}}]}, + ) + + +# --------------------------------------------------------------------------- +# Execution profiles +# --------------------------------------------------------------------------- + + +def test_get_execution_profiles() -> None: + prepared = endpoints.get_execution_profiles() + + assert isinstance(prepared, PreparedRequest) + assert prepared.method == "GET" + assert prepared.path_template == "/apis/jobs/v2/execution-profiles" + assert prepared.path_params == {} + assert prepared.content is None + + +# --------------------------------------------------------------------------- +# Job CRUD + lifecycle +# --------------------------------------------------------------------------- + + +def test_create_job() -> None: + body = _create_request() + prepared = endpoints.create_job(workspace="default", body=body) + + assert prepared.method == "POST" + assert prepared.path_template == "/apis/jobs/v2/workspaces/{workspace}/jobs" + assert prepared.path_params == {"workspace": "default"} + assert prepared.content == body.model_dump_json(exclude_unset=True).encode() + assert prepared.content_type == "application/json" + assert prepared.response_type is PlatformJobResponse + + +def test_create_job_workspace_optional() -> None: + prepared = endpoints.create_job(body=_create_request()) + assert prepared.path_params == {} + + +def test_list_jobs() -> None: + prepared = endpoints.list_jobs(workspace="default") + + assert prepared.method == "GET" + assert prepared.path_params == {"workspace": "default"} + assert prepared.content is None + assert get_origin(prepared.response_type) is Paginated + + +def test_list_jobs_with_query_params() -> None: + prepared = endpoints.list_jobs(workspace="default", query_params={"page": 2, "page_size": 10}) + assert prepared.query_params == {"page": 2, "page_size": 10} + + +def test_get_job() -> None: + prepared = endpoints.get_job(workspace="default", name="j-1") + + assert prepared.method == "GET" + assert prepared.path_params == {"workspace": "default", "name": "j-1"} + assert prepared.response_type is PlatformJobResponse + + +def test_delete_job() -> None: + prepared = endpoints.delete_job(workspace="default", name="j-1") + + assert prepared.method == "DELETE" + assert prepared.path_params == {"workspace": "default", "name": "j-1"} + assert prepared.content is None + assert prepared.response_type is None + + +def test_cancel_job() -> None: + prepared = endpoints.cancel_job(workspace="default", name="j-1") + assert prepared.method == "POST" + assert prepared.path_template.endswith("/jobs/{name}/cancel") + assert prepared.response_type is PlatformJobResponse + + +def test_pause_job() -> None: + prepared = endpoints.pause_job(workspace="default", name="j-1") + assert prepared.method == "POST" + assert prepared.path_template.endswith("/jobs/{name}/pause") + + +def test_resume_job() -> None: + prepared = endpoints.resume_job(workspace="default", name="j-1") + assert prepared.method == "POST" + assert prepared.path_template.endswith("/jobs/{name}/resume") + + +# --------------------------------------------------------------------------- +# Job status +# --------------------------------------------------------------------------- + + +def test_get_job_status() -> None: + prepared = endpoints.get_job_status(workspace="default", name="j-1") + assert prepared.method == "GET" + assert prepared.path_template.endswith("/jobs/{name}/status") + assert prepared.response_type is PlatformJobStatusResponse + + +def test_update_job_status_details() -> None: + prepared = endpoints.update_job_status_details( + workspace="default", name="j-1", body=JobStatusDetailsUpdate({"note": "x"}) + ) + assert prepared.method == "PATCH" + assert prepared.path_template.endswith("/jobs/{name}/status-details") + assert prepared.response_type is None + # RootModel serialises to the bare JSON object + assert json.loads(prepared.content) == {"note": "x"} + + +# --------------------------------------------------------------------------- +# Job logs +# --------------------------------------------------------------------------- + + +def test_list_job_logs() -> None: + prepared = endpoints.list_job_logs( + workspace="default", name="j-1", query_params={"limit": 50, "page_cursor": "abc"} + ) + assert prepared.method == "GET" + assert prepared.path_template.endswith("/jobs/{name}/logs") + assert prepared.query_params == {"limit": 50, "page_cursor": "abc"} + assert get_origin(prepared.response_type) is Paginated + assert get_args(prepared.response_type) == (PlatformJobLog, CursorPagination) + + +# --------------------------------------------------------------------------- +# Job results +# --------------------------------------------------------------------------- + + +def test_create_job_result() -> None: + body = PlatformJobResultCreateRequest(artifact_url="s3://x", artifact_storage_type="fileset") + prepared = endpoints.create_job_result(workspace="default", job="j-1", name="out", body=body) + assert prepared.method == "POST" + assert prepared.path_params == {"workspace": "default", "job": "j-1", "name": "out"} + assert prepared.content == body.model_dump_json(exclude_unset=True).encode() + + +def test_list_job_results() -> None: + prepared = endpoints.list_job_results(workspace="default", name="j-1") + assert prepared.method == "GET" + assert prepared.path_template.endswith("/jobs/{name}/results") + + +def test_get_job_result() -> None: + prepared = endpoints.get_job_result(workspace="default", job="j-1", name="out") + assert prepared.method == "GET" + assert prepared.path_params == {"workspace": "default", "job": "j-1", "name": "out"} + + +def test_download_job_result() -> None: + prepared = endpoints.download_job_result(workspace="default", job="j-1", name="out") + assert prepared.method == "GET" + assert prepared.path_template.endswith("/results/{name}/download") + assert prepared.path_params == {"workspace": "default", "job": "j-1", "name": "out"} + assert prepared.content is None + assert prepared.response_type is BinaryContent + + +# --------------------------------------------------------------------------- +# Job steps +# --------------------------------------------------------------------------- + + +def test_list_steps() -> None: + prepared = endpoints.list_steps(workspace="default", name="j-1") + assert prepared.method == "GET" + assert prepared.path_template.endswith("/jobs/{name}/steps") + assert get_origin(prepared.response_type) is Paginated + + +def test_list_steps_with_query_params() -> None: + prepared = endpoints.list_steps(workspace="default", name="-", query_params={"page": 1, "sort": "created_at"}) + assert prepared.query_params == {"page": 1, "sort": "created_at"} + + +def test_get_job_step() -> None: + prepared = endpoints.get_job_step(workspace="default", job="j-1", name="step-one") + assert prepared.method == "GET" + assert prepared.path_params == {"workspace": "default", "job": "j-1", "name": "step-one"} + assert prepared.response_type is PlatformJobStepResponse + + +def test_update_job_step_status() -> None: + body = PlatformJobStatusUpdateRequest(status="active") + prepared = endpoints.update_job_step_status(workspace="default", job="j-1", name="step-one", body=body) + assert prepared.method == "PATCH" + assert prepared.path_template.endswith("/steps/{name}/status") + assert prepared.response_type is PlatformJobStepResponse + + +# --------------------------------------------------------------------------- +# Job tasks +# --------------------------------------------------------------------------- + + +def test_list_job_step_tasks() -> None: + prepared = endpoints.list_job_step_tasks(workspace="default", job="j-1", name="step-one") + assert prepared.method == "GET" + assert prepared.path_template.endswith("/steps/{name}/tasks") + assert prepared.path_params == {"workspace": "default", "job": "j-1", "name": "step-one"} + + +def test_update_job_step_task() -> None: + body = PlatformJobTaskUpdate(status="completed") + prepared = endpoints.update_job_step_task(workspace="default", job="j-1", step="step-one", name="task-1", body=body) + assert prepared.method == "PUT" + assert prepared.path_params == {"workspace": "default", "job": "j-1", "step": "step-one", "name": "task-1"} + assert prepared.response_type is PlatformJobTaskResponse + + +def test_get_job_step_task() -> None: + prepared = endpoints.get_job_step_task(workspace="default", job="j-1", step="step-one", name="task-1") + assert prepared.method == "GET" + assert prepared.path_params == {"workspace": "default", "job": "j-1", "step": "step-one", "name": "task-1"} + assert prepared.response_type is PlatformJobTaskResponse + + +# --------------------------------------------------------------------------- +# Request-body serialisation +# --------------------------------------------------------------------------- + + +def test_create_job_body_roundtrip() -> None: + body = _create_request() + prepared = endpoints.create_job(workspace="default", body=body) + content = json.loads(prepared.content) + assert content["source"] == "test" + assert content["platform_spec"]["steps"][0]["name"] == "step-one" + + +def test_step_status_update_excludes_unset() -> None: + body = PlatformJobStatusUpdateRequest(status="active") + prepared = endpoints.update_job_step_status(workspace="default", job="j-1", name="s", body=body) + content = json.loads(prepared.content) + assert content["status"] == "active" + + +def test_step_with_context_response_type() -> None: + prepared = endpoints.list_steps(workspace="default", name="j-1") + # Paginated marker parametrised with the step-with-context model. + assert prepared.response_type.__args__[0] is PlatformJobStepWithContext # type: ignore[attr-defined] diff --git a/packages/nemo_platform_plugin/tests/test_jobs_filter.py b/packages/nemo_platform_plugin/tests/test_jobs_filter.py index 1fafe7d9e5..4272b02751 100644 --- a/packages/nemo_platform_plugin/tests/test_jobs_filter.py +++ b/packages/nemo_platform_plugin/tests/test_jobs_filter.py @@ -19,7 +19,9 @@ from datetime import datetime from types import SimpleNamespace from typing import Any +from unittest.mock import patch +import pytest from fastapi import FastAPI from nemo_platform_plugin.dependencies import get_entity_client, get_sdk_client from nemo_platform_plugin.jobs.api_factory import job_route_factory @@ -27,15 +29,27 @@ from starlette.testclient import TestClient -def _forwarded_filter(sdk: "_CapturingSdk") -> dict: - """Decode the JSON ``filter`` param the factory pushed through ``extra_query``. +@pytest.fixture(autouse=True) +def _patch_client_from_platform(): + """The factory calls ``client_from_platform(sdk, AsyncJobsClient)``; the test's + ``_CapturingSdk`` exposes the captured client as ``sdk.jobs_client``, so route + it through here.""" + with patch( + "nemo_platform_plugin.jobs.api_factory.client_from_platform", + side_effect=lambda sdk, _cls: sdk.jobs_client, + ): + yield - The factory bypasses the SDK's typed ``filter`` kwarg because the bundled + +def _forwarded_filter(sdk: _CapturingSdk) -> dict: + """Decode the JSON ``filter`` param the factory pushed through ``query_params``. + + The factory bypasses the client's typed ``filter`` handling because the querystring serializer mangles ``$and``-style list-of-dict values. It sends - a JSON-encoded filter via ``extra_query`` instead — tests parse it back - here so assertions stay shape-driven, not string-driven. + a JSON-encoded filter via the ``filter`` query param instead — tests parse + it back here so assertions stay shape-driven, not string-driven. """ - return json.loads(sdk.list_kwargs["extra_query"]["filter"]) + return json.loads(sdk.list_kwargs["query_params"]["filter"]) class _Spec(BaseModel): @@ -46,33 +60,40 @@ def _fake_compiler(workspace, original_spec, transformed_spec, entity_client, jo return {"steps": []} -def _fake_pagination(): +def _fake_page(): + """A ``PageResult``-like object for the typed client's ``list_jobs().page()``.""" return SimpleNamespace( - model_dump=lambda: { + items=[], + metadata={ "page": 1, "page_size": 10, "current_page_size": 0, "total_pages": 1, "total_results": 0, - } + }, ) def _fake_list_response(): - return SimpleNamespace(data=[], pagination=_fake_pagination()) + return SimpleNamespace(page=_fake_page) class _CapturingSdk: - """Captures kwargs passed to ``sdk.jobs.list(...)`` for assertion.""" + """Captures kwargs passed to ``JobsClient.list_jobs(...)`` for assertion. + + The factory now calls ``client_from_platform(sdk, AsyncJobsClient).list_jobs(...)``. + ``_build_app`` patches ``client_from_platform`` to return this object's + ``jobs_client`` so the ``list_jobs`` kwargs are captured here. + """ def __init__(self) -> None: self.list_kwargs: dict[str, Any] = {} - async def _list(**kwargs: Any) -> Any: + async def _list_jobs(**kwargs: Any) -> Any: self.list_kwargs = kwargs return _fake_list_response() - self.jobs = SimpleNamespace(list=_list) + self.jobs_client = SimpleNamespace(list_jobs=_list_jobs) def _build_app() -> tuple[FastAPI, _CapturingSdk]: @@ -324,30 +345,28 @@ class TestForwardedFilterSurvivesSdkSerialization: *what* the factory hands to ``sdk.jobs.list``, but not whether the SDK's querystring serializer can encode it onto the wire without mangling. - These tests run the forwarded ``extra_query`` value through the bundled - SDK's actual ``Querystring`` (with the platform client's ``array_format`` - setting), then through ``make_filter_dep``'s parsing path, and assert the + The typed client forwards ``filter`` as a single JSON-string query param + (not through the Stainless deep-object serializer, which mangled + ``$and``-style list-of-dict values). These tests take that forwarded value + and run it through ``make_filter_dep``'s parsing path, asserting the resulting ``FilterOperation`` tree matches what the plugin composed. - A regression that broke the SDK encoding (e.g., reverting to a typed - ``filter=`` dict with logical-array values) would produce repr-joined - garbage on the wire and fail to round-trip here. + A regression that reverted to a typed ``filter=`` dict with logical-array + values would produce repr-joined garbage on the wire and fail to round-trip + here. """ @staticmethod - def _round_trip(sdk: "_CapturingSdk") -> dict: - from nemo_platform._qs import Querystring - - qs = Querystring(array_format="comma") # matches AsyncNeMoPlatform's qs - items = qs.stringify_items(sdk.list_kwargs.get("extra_query") or {}) - # Find the encoded ``filter`` value as it would appear in the URL. - filter_values = [v for k, v in items if k == "filter"] - assert len(filter_values) == 1, f"expected one filter param, got {filter_values!r}" + def _round_trip(sdk: _CapturingSdk) -> dict: + # The migrated factory forwards a single JSON-encoded ``filter`` string + # via ``query_params`` — no deep-object array serialization involved. + filter_value = (sdk.list_kwargs.get("query_params") or {}).get("filter") + assert filter_value is not None, "expected a forwarded filter param" # Decode just like core jobs make_filter_dep would: see leading ``{``, # route through parse_json_filter. from nemo_platform_plugin.api.filter import parse_json_filter - operation = parse_json_filter(filter_values[0]) + operation = parse_json_filter(filter_value) return operation.to_dict() def test_user_filter_round_trip(self): diff --git a/packages/nmp_common/src/nmp/common/errors/sdk_exception_handlers.py b/packages/nmp_common/src/nmp/common/errors/sdk_exception_handlers.py index 2c31130829..55e17ff26c 100644 --- a/packages/nmp_common/src/nmp/common/errors/sdk_exception_handlers.py +++ b/packages/nmp_common/src/nmp/common/errors/sdk_exception_handlers.py @@ -25,6 +25,7 @@ from fastapi import FastAPI, Request from fastapi.responses import JSONResponse from nemo_platform import APIStatusError +from nemo_platform_plugin.client.errors import NemoHTTPError from nmp.common.entities.client import ( EntityConflictError, EntityNotFoundError, @@ -34,6 +35,12 @@ logger = logging.getLogger(__name__) + +def _scrub_crlf(value: str) -> str: + """Strip CR/LF so user-controlled request fields cannot forge log lines.""" + return value.replace("\r", " ").replace("\n", " ") + + # Map entity store exceptions to HTTP status codes ENTITY_ERROR_STATUS_CODES: dict[type[EntityStoreError], int] = { EntityNotFoundError: 404, @@ -65,8 +72,8 @@ async def sdk_status_error_handler(request: Request, exc: APIStatusError) -> JSO logger.debug( "Converting SDK exception to HTTP response: %s %s -> %d", - request.method, - request.url.path, + _scrub_crlf(request.method), + _scrub_crlf(request.url.path), exc.status_code, ) @@ -76,6 +83,27 @@ async def sdk_status_error_handler(request: Request, exc: APIStatusError) -> JSO ) +async def nemo_client_error_handler(request: Request, exc: NemoHTTPError) -> JSONResponse: + """Convert NemoClient HTTP exceptions back to HTTP responses. + + The typed ``NemoClient`` raises ``NemoHTTPError`` (and subclasses) on + non-2xx service-to-service responses; convert them to proper HTTP + responses the same way :func:`sdk_status_error_handler` does for the + Stainless SDK. + """ + logger.debug( + "Converting NemoClient exception to HTTP response: %s %s -> %d", + _scrub_crlf(request.method), + _scrub_crlf(request.url.path), + exc.status_code, + ) + + return JSONResponse( + status_code=exc.status_code, + content={"detail": exc.detail}, + ) + + async def entity_store_error_handler(request: Request, exc: EntityStoreError) -> JSONResponse: """Convert EntityStoreError exceptions to HTTP responses. @@ -99,8 +127,8 @@ async def entity_store_error_handler(request: Request, exc: EntityStoreError) -> "Converting %s to %d: %s %s", type(exc).__name__, status_code, - request.method, - request.url.path, + _scrub_crlf(request.method), + _scrub_crlf(request.url.path), ) return JSONResponse( @@ -119,12 +147,17 @@ def register_sdk_exception_handlers(app: FastAPI) -> None: Args: app: The FastAPI application to register handlers on """ - app.add_exception_handler(APIStatusError, sdk_status_error_handler) - app.add_exception_handler(EntityStoreError, entity_store_error_handler) + # Handlers are annotated with the specific exception subtype they handle; + # Starlette's stub types the callback against the base ``Exception``, so ty + # flags the narrower signature. The handlers are correct at runtime. + app.add_exception_handler(APIStatusError, sdk_status_error_handler) # ty: ignore[invalid-argument-type] + app.add_exception_handler(NemoHTTPError, nemo_client_error_handler) # ty: ignore[invalid-argument-type] + app.add_exception_handler(EntityStoreError, entity_store_error_handler) # ty: ignore[invalid-argument-type] __all__ = [ "sdk_status_error_handler", + "nemo_client_error_handler", "entity_store_error_handler", "register_sdk_exception_handlers", ] diff --git a/packages/nmp_common/tests/api_factory/test_api_factory.py b/packages/nmp_common/tests/api_factory/test_api_factory.py index cff87af28f..ab70063bc9 100644 --- a/packages/nmp_common/tests/api_factory/test_api_factory.py +++ b/packages/nmp_common/tests/api_factory/test_api_factory.py @@ -2,20 +2,26 @@ # SPDX-License-Identifier: Apache-2.0 import json +from contextlib import contextmanager from datetime import datetime, timedelta from pathlib import Path from typing import Annotated -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import FastAPI from fastapi.openapi.utils import get_openapi from fastapi.routing import APIRoute from fastapi.testclient import TestClient -from httpx import Request, Response -from nemo_platform import APIStatusError from nemo_platform.types.jobs import PlatformJobResponse as PlatformJob from nemo_platform.types.shared.platform_job_status import PlatformJobStatus +from nemo_platform_plugin.client.errors import ( + ConflictError as ClientConflictError, +) +from nemo_platform_plugin.client.errors import ( + NotFoundError as ClientNotFoundError, +) from nemo_platform_plugin.entities import EntityClient from nemo_platform_plugin.jobs.api_factory import ( ContainerSpec, @@ -38,7 +44,6 @@ _validate_job_spec, job_route_factory, ) -from nmp.common.api.common import Page, PaginationData from nmp.common.errors.sdk_exception_handlers import register_sdk_exception_handlers from nmp.common.jobs.exceptions import PlatformJobCompilationError from nmp.common.jobs.schemas import ( @@ -284,12 +289,71 @@ class TestModel(BaseModel): assert json.loads(lines[0])["error"]["line"] == 1 -@pytest.fixture -def mock_service_with_job_routes(): - """Create a test FastAPI app with job routes for testing job route functionality.""" +def _resp(data): + """Wrap a payload in a NemoResponse-like object whose ``.data()`` returns it. + + Production now consumes typed-client responses via ``(await client.(...)).data()``, + so mocked jobs-client methods must return an object with a ``.data()`` accessor + rather than the payload directly. + """ + m = MagicMock() + m.data.return_value = data + return m + + +def _page_resp(items, *, page=1, page_size=10, total_pages=1, total_results=None): + """Wrap a list of jobs in a response whose ``.page()`` returns a PageResult-like object. + + The list handler calls ``.page()`` and reads its items and offset metadata. + """ + page_result = MagicMock() + page_result.items = items + page_result.metadata = { + "page": page, + "page_size": page_size, + "current_page_size": len(items), + "total_pages": total_pages, + "total_results": total_results if total_results is not None else len(items), + } + m = MagicMock() + m.page.return_value = page_result + return m + + +def _cursor_page_resp(page: PlatformJobLogPage): + """Wrap a Jobs log envelope as a cursor-paginated typed-client response.""" + page_result = MagicMock() + page_result.items = page.data + page_result.metadata = { + "total": page.total, + "next_page": page.next_page, + "prev_page": page.prev_page, + } + response = MagicMock() + response.page.return_value = page_result + return response + + +def _client_error(error_cls, status_code: int, detail: str): + """Build a NemoHTTPError subclass from an httpx.Response, as the typed client raises.""" + request = httpx.Request("POST", "http://test") + response = httpx.Response(status_code=status_code, json={"detail": detail}, request=request) + return error_cls(response) + + +@contextmanager +def _job_routes_app(): + """Create a test FastAPI app with job routes, patching the typed jobs client. + + The generated handlers call ``client_from_platform(sdk, AsyncJobsClient)`` and then + invoke methods on the returned client. We patch ``client_from_platform`` in the + api_factory module to return a single ``MagicMock`` jobs client that tests configure + (each method as an ``AsyncMock`` returning a ``_resp(...)`` / ``_page_resp(...)``). + """ from nmp.common.service.dependencies import get_sdk_client mock_sdk = MagicMock() + mock_jobs = MagicMock() router = job_route_factory( service_name="test_service", job_type="TestJob", @@ -306,7 +370,16 @@ def mock_service_with_job_routes(): # Include a prefix to indicate the location of the jobs router app.include_router(router, prefix="/v2/workspaces/{workspace}/test") - return app, mock_sdk + + with patch("nemo_platform_plugin.jobs.api_factory.client_from_platform", return_value=mock_jobs): + yield app, mock_jobs + + +@pytest.fixture +def mock_service_with_job_routes(): + """Yield a test FastAPI app plus the mocked typed jobs client for job-route tests.""" + with _job_routes_app() as (app, mock_jobs): + yield app, mock_jobs def create_mock_platform_job( @@ -365,12 +438,12 @@ def create_mock_log_page( def test_get_job_logs_default_parameters(mock_service_with_job_routes): """Test get_job_logs with default parameters (no limit or page_cursor).""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - # Mock the SDK response + # Mock the jobs-client response mock_log_page = create_mock_log_page(num_logs=5, total=10) - mock_sdk.jobs.get_logs = AsyncMock(return_value=mock_log_page) + mock_jobs.list_job_logs = AsyncMock(return_value=_cursor_page_resp(mock_log_page)) # Make the request response = client.get("/v2/workspaces/default/test/jobs/test-job-123/logs") @@ -381,23 +454,22 @@ def test_get_job_logs_default_parameters(mock_service_with_job_routes): assert len(response_data["data"]) == 5 assert response_data["total"] == 10 - # Verify SDK was called with correct parameters - mock_sdk.jobs.get_logs.assert_called_once_with( + # None-valued limit/page_cursor are omitted from query_params. + mock_jobs.list_job_logs.assert_called_once_with( workspace="default", name="test-job-123", - limit=None, - page_cursor=None, + query_params={}, ) def test_get_job_logs_with_limit(mock_service_with_job_routes): """Test get_job_logs with limit parameter.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - # Mock the SDK response + # Mock the jobs-client response mock_log_page = create_mock_log_page(num_logs=2, total=10, next_page="next_cursor_123") - mock_sdk.jobs.get_logs = AsyncMock(return_value=mock_log_page) + mock_jobs.list_job_logs = AsyncMock(return_value=_cursor_page_resp(mock_log_page)) # Make the request with limit response = client.get("/v2/workspaces/default/test/jobs/test-job-123/logs?limit=2") @@ -409,23 +481,21 @@ def test_get_job_logs_with_limit(mock_service_with_job_routes): assert response_data["total"] == 10 assert response_data["next_page"] == "next_cursor_123" - # Verify SDK was called with correct parameters - mock_sdk.jobs.get_logs.assert_called_once_with( + mock_jobs.list_job_logs.assert_called_once_with( workspace="default", name="test-job-123", - limit=2, - page_cursor=None, + query_params={"limit": 2}, ) def test_get_job_logs_with_page_cursor(mock_service_with_job_routes): """Test get_job_logs with page_cursor parameter.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - # Mock the SDK response + # Mock the jobs-client response mock_log_page = create_mock_log_page(num_logs=3, total=10, next_page="next_cursor_456", prev_page="prev_cursor_789") - mock_sdk.jobs.get_logs = AsyncMock(return_value=mock_log_page) + mock_jobs.list_job_logs = AsyncMock(return_value=_cursor_page_resp(mock_log_page)) # Make the request with page_cursor response = client.get("/v2/workspaces/default/test/jobs/test-job-123/logs?page_cursor=cursor_abc_123") @@ -438,25 +508,23 @@ def test_get_job_logs_with_page_cursor(mock_service_with_job_routes): assert response_data["next_page"] == "next_cursor_456" assert response_data["prev_page"] == "prev_cursor_789" - # Verify SDK was called with correct parameters - mock_sdk.jobs.get_logs.assert_called_once_with( + mock_jobs.list_job_logs.assert_called_once_with( workspace="default", name="test-job-123", - limit=None, - page_cursor="cursor_abc_123", + query_params={"page_cursor": "cursor_abc_123"}, ) def test_get_job_logs_with_both_parameters(mock_service_with_job_routes): """Test get_job_logs with both limit and page_cursor parameters.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - # Mock the SDK response + # Mock the jobs-client response mock_log_page = create_mock_log_page( num_logs=5, total=100, next_page="next_cursor_combined", prev_page="prev_cursor_combined" ) - mock_sdk.jobs.get_logs = AsyncMock(return_value=mock_log_page) + mock_jobs.list_job_logs = AsyncMock(return_value=_cursor_page_resp(mock_log_page)) # Make the request with both parameters response = client.get("/v2/workspaces/default/test/jobs/test-job-123/logs?limit=5&page_cursor=combined_cursor_xyz") @@ -469,23 +537,21 @@ def test_get_job_logs_with_both_parameters(mock_service_with_job_routes): assert response_data["next_page"] == "next_cursor_combined" assert response_data["prev_page"] == "prev_cursor_combined" - # Verify SDK was called with correct parameters - mock_sdk.jobs.get_logs.assert_called_once_with( + mock_jobs.list_job_logs.assert_called_once_with( workspace="default", name="test-job-123", - limit=5, - page_cursor="combined_cursor_xyz", + query_params={"limit": 5, "page_cursor": "combined_cursor_xyz"}, ) def test_get_job_logs_with_zero_limit(mock_service_with_job_routes): """Test get_job_logs with limit=0.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - # Mock the SDK response + # Mock the jobs-client response mock_log_page = create_mock_log_page(num_logs=0, total=10, next_page="next_cursor_zero") - mock_sdk.jobs.get_logs = AsyncMock(return_value=mock_log_page) + mock_jobs.list_job_logs = AsyncMock(return_value=_cursor_page_resp(mock_log_page)) # Make the request with limit=0 response = client.get("/v2/workspaces/default/test/jobs/test-job-123/logs?limit=0") @@ -497,23 +563,22 @@ def test_get_job_logs_with_zero_limit(mock_service_with_job_routes): assert response_data["total"] == 10 assert response_data["next_page"] == "next_cursor_zero" - # Verify SDK was called with correct parameters - mock_sdk.jobs.get_logs.assert_called_once_with( + # limit=0 is not None, so it is included in query_params. + mock_jobs.list_job_logs.assert_called_once_with( workspace="default", name="test-job-123", - limit=0, - page_cursor=None, + query_params={"limit": 0}, ) def test_get_job_logs_empty_page_cursor(mock_service_with_job_routes): """Test get_job_logs with empty string page_cursor.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - # Mock the SDK response + # Mock the jobs-client response mock_log_page = create_mock_log_page(num_logs=3, total=10) - mock_sdk.jobs.get_logs = AsyncMock(return_value=mock_log_page) + mock_jobs.list_job_logs = AsyncMock(return_value=_cursor_page_resp(mock_log_page)) # Make the request with empty page_cursor response = client.get("/v2/workspaces/default/test/jobs/test-job-123/logs?page_cursor=") @@ -523,29 +588,21 @@ def test_get_job_logs_empty_page_cursor(mock_service_with_job_routes): response_data = response.json() assert len(response_data["data"]) == 3 - # Verify SDK was called with empty string (which should be treated as None by FastAPI) - mock_sdk.jobs.get_logs.assert_called_once_with( + # Empty string is not None, so it is included in query_params. + mock_jobs.list_job_logs.assert_called_once_with( workspace="default", name="test-job-123", - limit=None, - page_cursor="", + query_params={"page_cursor": ""}, ) def test_get_job_logs_job_not_found(mock_service_with_job_routes): """Test get_job_logs when job is not found.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - # Mock the SDK to raise an exception - request = Request("GET", "https://api.example.com/jobs/nonexistent-job/logs") - mock_sdk.jobs.get_logs = AsyncMock( - side_effect=APIStatusError( - message="Job not found", - response=Response(status_code=404, request=request), - body={"detail": "Job not found"}, - ) - ) + # The typed client raises a NemoHTTPError subclass on non-2xx responses. + mock_jobs.list_job_logs = AsyncMock(side_effect=_client_error(ClientNotFoundError, 404, "Job not found")) # Make the request response = client.get("/v2/workspaces/default/test/jobs/nonexistent-job/logs") @@ -555,23 +612,22 @@ def test_get_job_logs_job_not_found(mock_service_with_job_routes): response_data = response.json() assert "Job not found" in response_data["detail"] - # Verify SDK was called - mock_sdk.jobs.get_logs.assert_called_once_with( + # Verify the jobs client was called + mock_jobs.list_job_logs.assert_called_once_with( workspace="default", name="nonexistent-job", - limit=None, - page_cursor=None, + query_params={}, ) def test_get_job_logs_large_limit(mock_service_with_job_routes): """Test get_job_logs with a large limit value.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - # Mock the SDK response + # Mock the jobs-client response mock_log_page = create_mock_log_page(num_logs=1000, total=1000) - mock_sdk.jobs.get_logs = AsyncMock(return_value=mock_log_page) + mock_jobs.list_job_logs = AsyncMock(return_value=_cursor_page_resp(mock_log_page)) # Make the request with large limit response = client.get("/v2/workspaces/default/test/jobs/test-job-123/logs?limit=1000") @@ -582,12 +638,10 @@ def test_get_job_logs_large_limit(mock_service_with_job_routes): assert len(response_data["data"]) == 1000 assert response_data["total"] == 1000 - # Verify SDK was called with correct parameters - mock_sdk.jobs.get_logs.assert_called_once_with( + mock_jobs.list_job_logs.assert_called_once_with( workspace="default", name="test-job-123", - limit=1000, - page_cursor=None, + query_params={"limit": 1000}, ) @@ -595,12 +649,12 @@ def test_get_job_logs_special_characters_in_cursor(mock_service_with_job_routes) """Test get_job_logs with special characters in page_cursor.""" import urllib.parse - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - # Mock the SDK response + # Mock the jobs-client response mock_log_page = create_mock_log_page(num_logs=2, total=10) - mock_sdk.jobs.get_logs = AsyncMock(return_value=mock_log_page) + mock_jobs.list_job_logs = AsyncMock(return_value=_cursor_page_resp(mock_log_page)) # Make the request with special characters in cursor (URL encoded) special_cursor = "cursor_with_special_chars_!@#$%^&*()_+-=" @@ -612,18 +666,17 @@ def test_get_job_logs_special_characters_in_cursor(mock_service_with_job_routes) response_data = response.json() assert len(response_data["data"]) == 2 - # Verify SDK was called with correct parameters (decoded) - mock_sdk.jobs.get_logs.assert_called_once_with( + # Verify the jobs client was called with correct parameters (decoded) + mock_jobs.list_job_logs.assert_called_once_with( workspace="default", name="test-job-123", - limit=None, - page_cursor=special_cursor, + query_params={"page_cursor": special_cursor}, ) def test_get_job_logs_response_structure(mock_service_with_job_routes): """Test that get_job_logs returns proper PlatformJobLogPage structure.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) # Create a detailed mock log page @@ -651,7 +704,7 @@ def test_get_job_logs_response_structure(mock_service_with_job_routes): prev_page="cursor_prev_page_123", ) - mock_sdk.jobs.get_logs = AsyncMock(return_value=mock_log_page) + mock_jobs.list_job_logs = AsyncMock(return_value=_cursor_page_resp(mock_log_page)) # Make the request response = client.get("/v2/workspaces/default/test/jobs/test-job-456/logs?limit=2&page_cursor=middle_cursor") @@ -674,30 +727,19 @@ def test_get_job_logs_response_structure(mock_service_with_job_routes): assert log_entry["job_task"] == "data_validation" assert log_entry["message"] == "Starting data validation process" - # Verify SDK was called with correct parameters - mock_sdk.jobs.get_logs.assert_called_once_with( + # Verify the jobs client was called with correct parameters + mock_jobs.list_job_logs.assert_called_once_with( workspace="default", name="test-job-456", - limit=2, - page_cursor="middle_cursor", + query_params={"limit": 2, "page_cursor": "middle_cursor"}, ) def test_list_jobs_no_results(mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - mock_sdk.jobs.list = AsyncMock( - return_value=Page( - data=[], - pagination=PaginationData(page=1, page_size=5, current_page_size=0, total_pages=0, total_results=0), - sort="-created_at", - filter={ - "source": "test_service", - "status": "completed", - }, - ) - ) + mock_jobs.list_jobs = AsyncMock(return_value=_page_resp([], page=1, page_size=5, total_pages=0, total_results=0)) response = client.get("/v2/workspaces/default/test/jobs?page=1&page_size=5&filter[status]=completed") @@ -709,14 +751,14 @@ def test_list_jobs_no_results(mock_service_with_job_routes): # The user filter is AND-composed with the service-source predicate so # logical roots ($or/$and/$not) stay scoped — a flat dict merge would # silently drop the source clause under a logical root. The composed tree - # is forwarded as a JSON string via ``extra_query`` so the SDK querystring + # is forwarded as a JSON string in ``query_params`` so the querystring # serializer doesn't mangle list-of-dict values. - mock_sdk.jobs.list.assert_called_once_with( + mock_jobs.list_jobs.assert_called_once_with( workspace="default", - page=1, - page_size=5, - sort="-created_at", - extra_query={ + query_params={ + "page": 1, + "page_size": 5, + "sort": "-created_at", "filter": json.dumps( { "$and": [ @@ -724,25 +766,26 @@ def test_list_jobs_no_results(mock_service_with_job_routes): {"source": {"$eq": "test_service"}}, ] } - ) + ), }, ) def test_list_jobs_with_results(mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - mock_sdk.jobs.list = AsyncMock( - return_value=Page( - data=[ + mock_jobs.list_jobs = AsyncMock( + return_value=_page_resp( + [ create_mock_platform_job("job-1", "completed", "2023-01-01T10:00:00Z"), create_mock_platform_job("job-2", "completed", "2023-01-01T11:00:00Z"), create_mock_platform_job("job-3", "completed", "2023-01-01T12:00:00Z"), ], - pagination=PaginationData(page=1, page_size=5, current_page_size=3, total_pages=1, total_results=3), - sort="-created_at", - filter={"source": "test_service"}, + page=1, + page_size=5, + total_pages=1, + total_results=3, ) ) @@ -759,38 +802,42 @@ def test_list_jobs_with_results(mock_service_with_job_routes): assert response_data["pagination"]["total_results"] == 3 # No user filter — source predicate stands alone (no $and wrap needed). - mock_sdk.jobs.list.assert_called_once_with( + mock_jobs.list_jobs.assert_called_once_with( workspace="default", - page=1, - page_size=5, - sort="-created_at", - extra_query={"filter": json.dumps({"source": {"$eq": "test_service"}})}, + query_params={ + "page": 1, + "page_size": 5, + "sort": "-created_at", + "filter": json.dumps({"source": {"$eq": "test_service"}}), + }, ) def test_list_jobs_with_multiple_pages(mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - mock_sdk.jobs.list = AsyncMock( + mock_jobs.list_jobs = AsyncMock( side_effect=[ - Page( - data=[ + _page_resp( + [ create_mock_platform_job("job-1", "active", "2023-01-01T10:00:00Z"), create_mock_platform_job("job-2", "active", "2023-01-01T11:00:00Z"), ], - pagination=PaginationData(page=1, page_size=2, current_page_size=2, total_pages=3, total_results=6), - sort="-created_at", - filter={"source": "test_service", "status": "active"}, + page=1, + page_size=2, + total_pages=3, + total_results=6, ), - Page( - data=[ + _page_resp( + [ create_mock_platform_job("job-3", "active", "2023-01-01T10:00:00Z"), create_mock_platform_job("job-4", "active", "2023-01-01T11:00:00Z"), ], - pagination=PaginationData(page=2, page_size=2, current_page_size=2, total_pages=3, total_results=6), - sort="-created_at", - filter={"source": "test_service", "status": "active"}, + page=2, + page_size=2, + total_pages=3, + total_results=6, ), ] ) @@ -812,16 +859,18 @@ def test_list_jobs_with_multiple_pages(mock_service_with_job_routes): ] } ) - mock_sdk.jobs.list.assert_called_once_with( + mock_jobs.list_jobs.assert_called_once_with( workspace="default", - page=1, - page_size=2, - sort="-created_at", - extra_query={"filter": expected_filter}, + query_params={ + "page": 1, + "page_size": 2, + "sort": "-created_at", + "filter": expected_filter, + }, ) # Fetch the second page - mock_sdk.reset_mock() + mock_jobs.list_jobs.reset_mock() response = client.get("/v2/workspaces/default/test/jobs?page=2&page_size=2&filter[status]=active") response_data = response.json() @@ -830,29 +879,32 @@ def test_list_jobs_with_multiple_pages(mock_service_with_job_routes): assert response_data["pagination"]["total_pages"] == 3 assert response_data["pagination"]["total_results"] == 6 - mock_sdk.jobs.list.assert_called_once_with( + mock_jobs.list_jobs.assert_called_once_with( workspace="default", - page=2, - page_size=2, - sort="-created_at", - extra_query={"filter": expected_filter}, + query_params={ + "page": 2, + "page_size": 2, + "sort": "-created_at", + "filter": expected_filter, + }, ) def test_list_jobs_with_custom_sort(mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - mock_sdk.jobs.list = AsyncMock( - return_value=Page( - data=[ + mock_jobs.list_jobs = AsyncMock( + return_value=_page_resp( + [ create_mock_platform_job("job-3", "completed", "2023-01-03T10:00:00Z"), create_mock_platform_job("job-2", "completed", "2023-01-02T10:00:00Z"), create_mock_platform_job("job-1", "completed", "2023-01-01T10:00:00Z"), ], - pagination=PaginationData(page=1, page_size=10, current_page_size=3, total_pages=1, total_results=3), - sort="created_at", - filter={"source": "test_service"}, + page=1, + page_size=10, + total_pages=1, + total_results=3, ) ) @@ -864,32 +916,34 @@ def test_list_jobs_with_custom_sort(mock_service_with_job_routes): assert response_data["data"][1]["id"] == "job-2" assert response_data["data"][2]["id"] == "job-1" - mock_sdk.jobs.list.assert_called_once_with( + mock_jobs.list_jobs.assert_called_once_with( workspace="default", - page=1, - page_size=10, - sort="created_at", - extra_query={"filter": json.dumps({"source": {"$eq": "test_service"}})}, + query_params={ + "page": 1, + "page_size": 10, + "sort": "created_at", + "filter": json.dumps({"source": {"$eq": "test_service"}}), + }, ) def test_list_jobs_invalid_filter(mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) # ``source`` is not in BaseJobsListFilter — make_filter_dep's allowlist - # rejects unknown fields with a 400 before any SDK call is made. + # rejects unknown fields with a 400 before any jobs-client call is made. response = client.get("/v2/workspaces/default/test/jobs?page=1&page_size=5&filter[source]=foo") assert response.status_code == 400 response_data = response.json() assert "source" in response_data["detail"] - mock_sdk.jobs.list.assert_not_called() + mock_jobs.list_jobs.assert_not_called() def test_list_jobs_invalid_sort(mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) response = client.get("/v2/workspaces/default/test/jobs?page=1&page_size=5&sort=foo") @@ -901,12 +955,12 @@ def test_list_jobs_invalid_sort(mock_service_with_job_routes): == "Input should be 'created_at', '-created_at', 'updated_at' or '-updated_at'" ) - mock_sdk.jobs.list.assert_not_called() + mock_jobs.list_jobs.assert_not_called() def test_list_job_results_includes_download_url(mock_service_with_job_routes): """Test that list_job_results includes download_url for each result.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) mock_results = [ @@ -926,11 +980,7 @@ def test_list_job_results_includes_download_url(mock_service_with_job_routes): ), ] - mock_sdk.jobs.results.list = AsyncMock( - return_value=PlatformJobListResultResponse( - data=mock_results, - ) - ) + mock_jobs.list_job_results = AsyncMock(return_value=_resp(PlatformJobListResultResponse(data=mock_results))) # Make the request response = client.get("/v2/workspaces/default/test/jobs/test-job-123/results") @@ -960,20 +1010,16 @@ def test_list_job_results_includes_download_url(mock_service_with_job_routes): == "http://testserver/v2/workspaces/default/test/jobs/test-job-123/results/result2/download" ) - # Verify SDK was called correctly - mock_sdk.jobs.results.list.assert_called_once_with(name="test-job-123", workspace="default") + # Verify the jobs client was called correctly + mock_jobs.list_job_results.assert_called_once_with(name="test-job-123", workspace="default") def test_list_job_results_empty_list(mock_service_with_job_routes): """Test that list_job_results works correctly with empty results.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - mock_sdk.jobs.results.list = AsyncMock( - return_value=PlatformJobListResultResponse( - data=[], - ) - ) + mock_jobs.list_job_results = AsyncMock(return_value=_resp(PlatformJobListResultResponse(data=[]))) # Make the request response = client.get("/v2/workspaces/default/test/jobs/test-job-456/results") @@ -984,13 +1030,13 @@ def test_list_job_results_empty_list(mock_service_with_job_routes): assert "data" in response_data assert len(response_data["data"]) == 0 - # Verify SDK was called correctly - mock_sdk.jobs.results.list.assert_called_once_with(name="test-job-456", workspace="default") + # Verify the jobs client was called correctly + mock_jobs.list_job_results.assert_called_once_with(name="test-job-456", workspace="default") def test_get_job_result_includes_download_url(mock_service_with_job_routes): """Test that get_job_result includes download_url in the response.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) # Create mock result without download_url @@ -1002,7 +1048,7 @@ def test_get_job_result_includes_download_url(mock_service_with_job_routes): artifact_url="default/test-fileset#test_result", ) - mock_sdk.jobs.results.retrieve = AsyncMock(return_value=mock_result) + mock_jobs.get_job_result = AsyncMock(return_value=_resp(mock_result)) # Make the request response = client.get("/v2/workspaces/default/test/jobs/test-job-789/results/test_result") @@ -1025,13 +1071,13 @@ def test_get_job_result_includes_download_url(mock_service_with_job_routes): assert response_data["artifact_storage_type"] == "fileset" assert response_data["artifact_url"] == "default/test-fileset#test_result" - # Verify SDK was called correctly - mock_sdk.jobs.results.retrieve.assert_called_once_with(name="test_result", job="test-job-789", workspace="default") + # Verify the jobs client was called correctly + mock_jobs.get_job_result.assert_called_once_with(name="test_result", job="test-job-789", workspace="default") def test_get_job_result_download_url_format(mock_service_with_job_routes): """Test that download_url has the correct format with special characters in result name.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) # Create mock result with a result name that has underscores @@ -1043,7 +1089,7 @@ def test_get_job_result_download_url_format(mock_service_with_job_routes): artifact_url="default/test-fileset#evaluation_report_v1", ) - mock_sdk.jobs.results.retrieve = AsyncMock(return_value=mock_result) + mock_jobs.get_job_result = AsyncMock(return_value=_resp(mock_result)) # Make the request response = client.get("/v2/workspaces/default/test/jobs/test-job-abc/results/evaluation_report_v1") @@ -1060,15 +1106,15 @@ def test_get_job_result_download_url_format(mock_service_with_job_routes): == "http://testserver/v2/workspaces/default/test/jobs/test-job-abc/results/evaluation_report_v1/download" ) - # Verify SDK was called correctly - mock_sdk.jobs.results.retrieve.assert_called_once_with( + # Verify the jobs client was called correctly + mock_jobs.get_job_result.assert_called_once_with( name="evaluation_report_v1", job="test-job-abc", workspace="default" ) def test_list_job_results_download_url_different_job_ids(mock_service_with_job_routes): """Test that download_url correctly reflects different job IDs.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) # Test with different job IDs @@ -1083,11 +1129,7 @@ def test_list_job_results_download_url_different_job_ids(mock_service_with_job_r artifact_url=f"default/test-fileset#{job_id}/output", ) - mock_sdk.jobs.results.list = AsyncMock( - return_value=PlatformJobListResultResponse( - data=[mock_result], - ) - ) + mock_jobs.list_job_results = AsyncMock(return_value=_resp(PlatformJobListResultResponse(data=[mock_result]))) # Make the request response = client.get(f"/v2/workspaces/default/test/jobs/{job_id}/results") @@ -1102,26 +1144,16 @@ def test_list_job_results_download_url_different_job_ids(mock_service_with_job_r ) # Reset mock for next iteration - mock_sdk.reset_mock() + mock_jobs.reset_mock() def test_get_job_result_not_found(mock_service_with_job_routes): """Test that get_job_result handles not found errors correctly.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app) - from httpx import Request, Response - from nemo_platform import APIStatusError - - # Mock the SDK to raise a 404 error - request = Request("GET", "https://api.example.com/jobs/nonexistent-job/results/nonexistent-result") - mock_sdk.jobs.results.retrieve = AsyncMock( - side_effect=APIStatusError( - message="Result not found", - response=Response(status_code=404, request=request), - body={"detail": "Result not found"}, - ) - ) + # The typed client raises a NemoHTTPError subclass on non-2xx responses. + mock_jobs.get_job_result = AsyncMock(side_effect=_client_error(ClientNotFoundError, 404, "Result not found")) # Make the request response = client.get("/v2/workspaces/default/test/jobs/nonexistent-job/results/nonexistent-result") @@ -1131,8 +1163,8 @@ def test_get_job_result_not_found(mock_service_with_job_routes): response_data = response.json() assert "Result not found" in response_data["detail"] - # Verify SDK was called - mock_sdk.jobs.results.retrieve.assert_called_once_with( + # Verify the jobs client was called + mock_jobs.get_job_result.assert_called_once_with( name="nonexistent-result", job="nonexistent-job", workspace="default" ) @@ -1418,10 +1450,12 @@ def compiler( app.dependency_overrides[get_sdk_client] = lambda: mock_sdk app.dependency_overrides[get_entity_client] = lambda: mock_entity_client - mock_sdk.jobs.create = AsyncMock(return_value=create_mock_platform_job("test-job-123", "pending")) + mock_jobs = MagicMock() + mock_jobs.create_job = AsyncMock(return_value=_resp(create_mock_platform_job("test-job-123", "pending"))) client = TestClient(app) - response = client.post("/v2/workspaces/default/test/jobs", json={"spec": {"foo": "test", "bar": 42}}) + with patch("nemo_platform_plugin.jobs.api_factory.client_from_platform", return_value=mock_jobs): + response = client.post("/v2/workspaces/default/test/jobs", json={"spec": {"foo": "test", "bar": 42}}) assert response.status_code == 201 assert received_workspace == "default" @@ -1464,13 +1498,15 @@ def sync_compiler( app.include_router(router, prefix="/v2/workspaces/{workspace}/test") mock_sdk = MagicMock() - mock_sdk.jobs.create = AsyncMock(return_value=create_mock_platform_job("test-job-123", "pending")) + mock_jobs = MagicMock() + mock_jobs.create_job = AsyncMock(return_value=_resp(create_mock_platform_job("test-job-123", "pending"))) app.dependency_overrides[get_sdk_client] = lambda: mock_sdk app.dependency_overrides[get_entity_client] = lambda: mock_entity_client client = TestClient(app) - response = client.post("/v2/workspaces/default/test/jobs", json={"spec": {"foo": "test", "bar": 42}}) + with patch("nemo_platform_plugin.jobs.api_factory.client_from_platform", return_value=mock_jobs): + response = client.post("/v2/workspaces/default/test/jobs", json={"spec": {"foo": "test", "bar": 42}}) assert response.status_code == 201 assert compiler_called, "Sync compiler should have been called" @@ -1736,31 +1772,21 @@ def compiler_bad_config(workspace, input_spec, output_spec, entity_client, job_n assert "not json serializable" in exc_info.value.detail.lower() -def _make_api_status_error(status_code: int, detail: str) -> APIStatusError: - """Create an APIStatusError with a realistic body structure.""" - body = {"detail": detail} - return APIStatusError( - message=f"Error code: {status_code} - {body}", - response=Response(status_code=status_code, json=body, request=Request("POST", "http://test")), - body=body, - ) - - class TestSDKExceptionHandling: - """Tests that APIStatusError from SDK calls is handled by the global exception handler. + """Tests that NemoHTTPError from typed-client calls is handled by the global exception handler. - The api_factory endpoints should NOT catch APIStatusError themselves — the global - sdk_status_error_handler registered on the app extracts the detail cleanly from - the exception body, avoiding the ugly stringified format like: + The api_factory endpoints should NOT catch the client errors themselves — the global + nemo_client_error_handler registered on the app extracts the detail cleanly from + the exception, avoiding an ugly stringified format like: {"detail": "Error code: 409 - {'detail': 'Job already exists'}"} """ def test_create_job_duplicate_name_returns_clean_409(self, mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app, raise_server_exceptions=False) error_detail = "Unable to create job: Job with name 'my-job' already exists in workspace 'default'." - mock_sdk.jobs.create = AsyncMock(side_effect=_make_api_status_error(409, error_detail)) + mock_jobs.create_job = AsyncMock(side_effect=_client_error(ClientConflictError, 409, error_detail)) response = client.post( "/v2/workspaces/default/test/jobs", @@ -1771,11 +1797,11 @@ def test_create_job_duplicate_name_returns_clean_409(self, mock_service_with_job assert response.json()["detail"] == error_detail def test_get_job_not_found_returns_clean_404(self, mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app, raise_server_exceptions=False) error_detail = "Job 'nonexistent' not found in workspace 'default'." - mock_sdk.jobs.retrieve = AsyncMock(side_effect=_make_api_status_error(404, error_detail)) + mock_jobs.get_job = AsyncMock(side_effect=_client_error(ClientNotFoundError, 404, error_detail)) response = client.get("/v2/workspaces/default/test/jobs/nonexistent") @@ -1783,11 +1809,11 @@ def test_get_job_not_found_returns_clean_404(self, mock_service_with_job_routes) assert response.json()["detail"] == error_detail def test_delete_job_not_found_returns_clean_404(self, mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app, raise_server_exceptions=False) error_detail = "Job 'nonexistent' not found." - mock_sdk.jobs.delete = AsyncMock(side_effect=_make_api_status_error(404, error_detail)) + mock_jobs.delete_job = AsyncMock(side_effect=_client_error(ClientNotFoundError, 404, error_detail)) response = client.delete("/v2/workspaces/default/test/jobs/nonexistent") @@ -1795,11 +1821,11 @@ def test_delete_job_not_found_returns_clean_404(self, mock_service_with_job_rout assert response.json()["detail"] == error_detail def test_cancel_job_bad_state_returns_clean_409(self, mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app, raise_server_exceptions=False) error_detail = "Job 'my-job' cannot be cancelled in its current state." - mock_sdk.jobs.cancel = AsyncMock(side_effect=_make_api_status_error(409, error_detail)) + mock_jobs.cancel_job = AsyncMock(side_effect=_client_error(ClientConflictError, 409, error_detail)) response = client.post("/v2/workspaces/default/test/jobs/my-job/cancel") @@ -1807,11 +1833,11 @@ def test_cancel_job_bad_state_returns_clean_409(self, mock_service_with_job_rout assert response.json()["detail"] == error_detail def test_pause_job_returns_clean_error(self, mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app, raise_server_exceptions=False) error_detail = "Job 'my-job' cannot be paused." - mock_sdk.jobs.pause = AsyncMock(side_effect=_make_api_status_error(409, error_detail)) + mock_jobs.pause_job = AsyncMock(side_effect=_client_error(ClientConflictError, 409, error_detail)) response = client.post("/v2/workspaces/default/test/jobs/my-job/pause") @@ -1819,11 +1845,11 @@ def test_pause_job_returns_clean_error(self, mock_service_with_job_routes): assert response.json()["detail"] == error_detail def test_resume_job_returns_clean_error(self, mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app, raise_server_exceptions=False) error_detail = "Job 'my-job' cannot be resumed." - mock_sdk.jobs.resume = AsyncMock(side_effect=_make_api_status_error(409, error_detail)) + mock_jobs.resume_job = AsyncMock(side_effect=_client_error(ClientConflictError, 409, error_detail)) response = client.post("/v2/workspaces/default/test/jobs/my-job/resume") @@ -1831,11 +1857,11 @@ def test_resume_job_returns_clean_error(self, mock_service_with_job_routes): assert response.json()["detail"] == error_detail def test_get_logs_returns_clean_error(self, mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app, raise_server_exceptions=False) error_detail = "Job 'nonexistent' not found." - mock_sdk.jobs.get_logs = AsyncMock(side_effect=_make_api_status_error(404, error_detail)) + mock_jobs.list_job_logs = AsyncMock(side_effect=_client_error(ClientNotFoundError, 404, error_detail)) response = client.get("/v2/workspaces/default/test/jobs/nonexistent/logs") @@ -1843,11 +1869,11 @@ def test_get_logs_returns_clean_error(self, mock_service_with_job_routes): assert response.json()["detail"] == error_detail def test_list_results_returns_clean_error(self, mock_service_with_job_routes): - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app, raise_server_exceptions=False) error_detail = "Job 'nonexistent' not found." - mock_sdk.jobs.results.list = AsyncMock(side_effect=_make_api_status_error(404, error_detail)) + mock_jobs.list_job_results = AsyncMock(side_effect=_client_error(ClientNotFoundError, 404, error_detail)) response = client.get("/v2/workspaces/default/test/jobs/nonexistent/results") @@ -1857,11 +1883,11 @@ def test_list_results_returns_clean_error(self, mock_service_with_job_routes): def test_error_detail_is_not_stringified(self, mock_service_with_job_routes): """Verify the response detail is the clean string, not a stringified repr like "Error code: 409 - {'detail': '...'}" which was the old behavior.""" - app, mock_sdk = mock_service_with_job_routes + app, mock_jobs = mock_service_with_job_routes client = TestClient(app, raise_server_exceptions=False) error_detail = "Job already exists." - mock_sdk.jobs.create = AsyncMock(side_effect=_make_api_status_error(409, error_detail)) + mock_jobs.create_job = AsyncMock(side_effect=_client_error(ClientConflictError, 409, error_detail)) response = client.post( "/v2/workspaces/default/test/jobs", diff --git a/packages/nmp_common/tests/jobs/conftest.py b/packages/nmp_common/tests/jobs/conftest.py index f50f9461b7..19ac3c1bab 100644 --- a/packages/nmp_common/tests/jobs/conftest.py +++ b/packages/nmp_common/tests/jobs/conftest.py @@ -8,6 +8,7 @@ from nmp.common.entities import DEFAULT_WORKSPACE from nmp.common.entities.utils import get_random_bytes from nmp.common.jobs.file_manager import FilesetFileManager +from nmp.common.jobs.schemas import FileStorageType NMP_URL = "http://localhost:8080" @@ -115,7 +116,7 @@ def mock_sync_file_manager(): mock_fm = MagicMock() mock_fm.validate_storage.return_value = None mock_fm.upload.return_value = "test-ws/test-fileset#results/att-123/my-result" - mock_fm.storage_type.return_value = MagicMock(value="fileset") + mock_fm.storage_type.return_value = FileStorageType.FILESET return mock_fm @@ -128,5 +129,5 @@ def mock_async_file_manager(): mock_fm = MagicMock() mock_fm.validate_storage = AsyncMock(return_value=None) mock_fm.upload = AsyncMock(return_value="test-ws/test-fileset#results/att-123/my-result") - mock_fm.storage_type.return_value = MagicMock(value="fileset") + mock_fm.storage_type.return_value = FileStorageType.FILESET return mock_fm diff --git a/packages/nmp_common/tests/jobs/test_result_manager.py b/packages/nmp_common/tests/jobs/test_result_manager.py index e813795d5d..37d8bdf7cf 100644 --- a/packages/nmp_common/tests/jobs/test_result_manager.py +++ b/packages/nmp_common/tests/jobs/test_result_manager.py @@ -3,12 +3,29 @@ from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest -from nemo_platform import ConflictError +from nemo_platform_plugin.client.errors import ConflictError as ClientConflictError +from nemo_platform_plugin.client.errors import NemoTransportError +from nemo_platform_plugin.jobs.result_manager import CreateJobResultError from nmp.common.config import Configuration from nmp.common.jobs import result_manager as rm from nmp.common.jobs.file_manager import AsyncFilesetFileManager, FilesetFileManager, TmpDirPath + +def _resp(data): + """Wrap a payload in a NemoResponse-like object whose ``.data()`` returns it.""" + m = MagicMock() + m.data.return_value = data + return m + + +def _conflict_error() -> ClientConflictError: + """Build the ConflictError the typed client raises on a 409 response.""" + request = httpx.Request("POST", "http://test") + return ClientConflictError(httpx.Response(status_code=409, json={}, request=request)) + + # ============================================================================= # FilesetFileManager Factory Tests # ============================================================================= @@ -81,13 +98,13 @@ def test_create_result_returns_existing_on_conflict_sync(tmp_path, mock_sdk, moc test_file = tmp_path / "artifact.bin" test_file.write_bytes(b"test content") - # Configure SDK to raise ConflictError on create, return existing on retrieve + # Configure the typed jobs client to raise ConflictError on create, return existing on retrieve. existing_result = MagicMock(name="my-result") - mock_nmp_sdk.jobs.results.create.side_effect = ConflictError(message="conflict", response=MagicMock(), body=None) - mock_nmp_sdk.jobs.results.retrieve.return_value = existing_result - - # Mock job retrieval to return fileset name - mock_nmp_sdk.jobs.retrieve.return_value = MagicMock(attempt_id="att-123", fileset="test-fileset") + mock_jobs = MagicMock() + mock_jobs.create_job_result.side_effect = _conflict_error() + mock_jobs.get_job_result.return_value = _resp(existing_result) + # Job retrieval provides the fileset name. + mock_jobs.get_job.return_value = _resp(MagicMock(attempt_id="att-123", fileset="test-fileset")) mgr = rm.ResultManager( job_name="test-job", @@ -97,14 +114,43 @@ def test_create_result_returns_existing_on_conflict_sync(tmp_path, mock_sdk, moc jobs_sdk=mock_nmp_sdk, ) - # Patch the _create_file_manager method to return our mock - with patch.object(mgr, "_create_file_manager", return_value=mock_sync_file_manager): + # Patch the _create_file_manager method to return our mock, and the typed-client + # factory to return our mock jobs client. + with ( + patch.object(mgr, "_create_file_manager", return_value=mock_sync_file_manager), + patch("nemo_platform_plugin.jobs.result_manager.client_from_platform", return_value=mock_jobs), + ): result = mgr.create_result("my-result", test_file) - mock_nmp_sdk.jobs.results.retrieve.assert_called_once_with(name="my-result", job="test-job", workspace="test-ws") + mock_jobs.get_job_result.assert_called_once_with(name="my-result", job="test-job", workspace="test-ws") assert result is existing_result +def test_create_result_wraps_transport_errors_sync(tmp_path, mock_sdk, mock_nmp_sdk, mock_sync_file_manager): + test_file = tmp_path / "artifact.bin" + test_file.write_bytes(b"test content") + request = httpx.Request("POST", "http://test/apis/jobs/v2/workspaces/test-ws/jobs/test-job/results/my-result") + mock_jobs = MagicMock() + mock_jobs.get_job.return_value = _resp(MagicMock(attempt_id="att-123", fileset="test-fileset")) + mock_jobs.create_job_result.side_effect = NemoTransportError( + httpx.ConnectError("Connection refused", request=request) + ) + mgr = rm.ResultManager( + job_name="test-job", + workspace="test-ws", + file_manager_cls=FilesetFileManager, + files_sdk=mock_sdk, + jobs_sdk=mock_nmp_sdk, + ) + + with ( + patch.object(mgr, "_create_file_manager", return_value=mock_sync_file_manager), + patch("nemo_platform_plugin.jobs.result_manager.client_from_platform", return_value=mock_jobs), + pytest.raises(CreateJobResultError, match="Error creating job result"), + ): + mgr.create_result("my-result", test_file) + + @pytest.mark.asyncio async def test_create_result_returns_existing_on_conflict_async( tmp_path, mock_sdk, mock_async_nmp_sdk, mock_async_file_manager @@ -113,15 +159,13 @@ async def test_create_result_returns_existing_on_conflict_async( test_file = tmp_path / "artifact.bin" test_file.write_bytes(b"test content") - # Configure async SDK to raise ConflictError on create, return existing on retrieve + # Configure the typed async jobs client to raise ConflictError on create, return existing on retrieve. existing_result = MagicMock(name="my-result") - mock_async_nmp_sdk.jobs.results.create.side_effect = ConflictError( - message="conflict", response=MagicMock(), body=None - ) - mock_async_nmp_sdk.jobs.results.retrieve.return_value = existing_result - - # Mock job retrieval to return fileset name - mock_async_nmp_sdk.jobs.retrieve.return_value = MagicMock(attempt_id="att-123", fileset="test-fileset") + mock_jobs = MagicMock() + mock_jobs.create_job_result = AsyncMock(side_effect=_conflict_error()) + mock_jobs.get_job_result = AsyncMock(return_value=_resp(existing_result)) + # Job retrieval provides the fileset name. + mock_jobs.get_job = AsyncMock(return_value=_resp(MagicMock(attempt_id="att-123", fileset="test-fileset"))) mgr = rm.AsyncResultManager( job_name="test-job", @@ -131,16 +175,46 @@ async def test_create_result_returns_existing_on_conflict_async( jobs_sdk=mock_async_nmp_sdk, ) - # Patch the _create_file_manager method to return our mock - with patch.object(mgr, "_create_file_manager", return_value=mock_async_file_manager): + # Patch the _create_file_manager method to return our mock, and the typed-client + # factory to return our mock jobs client. + with ( + patch.object(mgr, "_create_file_manager", return_value=mock_async_file_manager), + patch("nemo_platform_plugin.jobs.result_manager.client_from_platform", return_value=mock_jobs), + ): result = await mgr.create_result("my-result", test_file) - mock_async_nmp_sdk.jobs.results.retrieve.assert_called_once_with( - name="my-result", job="test-job", workspace="test-ws" - ) + mock_jobs.get_job_result.assert_called_once_with(name="my-result", job="test-job", workspace="test-ws") assert result is existing_result +@pytest.mark.asyncio +async def test_create_result_wraps_transport_errors_async( + tmp_path, mock_sdk, mock_async_nmp_sdk, mock_async_file_manager +): + test_file = tmp_path / "artifact.bin" + test_file.write_bytes(b"test content") + request = httpx.Request("POST", "http://test/apis/jobs/v2/workspaces/test-ws/jobs/test-job/results/my-result") + mock_jobs = MagicMock() + mock_jobs.get_job = AsyncMock(return_value=_resp(MagicMock(attempt_id="att-123", fileset="test-fileset"))) + mock_jobs.create_job_result = AsyncMock( + side_effect=NemoTransportError(httpx.ConnectError("Connection refused", request=request)) + ) + mgr = rm.AsyncResultManager( + job_name="test-job", + workspace="test-ws", + file_manager_cls=AsyncFilesetFileManager, + files_sdk=mock_sdk, + jobs_sdk=mock_async_nmp_sdk, + ) + + with ( + patch.object(mgr, "_create_file_manager", return_value=mock_async_file_manager), + patch("nemo_platform_plugin.jobs.result_manager.client_from_platform", return_value=mock_jobs), + pytest.raises(CreateJobResultError, match="Error creating job result"), + ): + await mgr.create_result("my-result", test_file) + + @pytest.mark.asyncio @patch("nmp.common.jobs.result_manager.result_manager_factory") @patch("nmp.common.jobs.result_manager.get_async_platform_sdk") diff --git a/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_progress_reporter.py b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_progress_reporter.py index afb4c3963c..9a2acaa31b 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_progress_reporter.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_progress_reporter.py @@ -6,8 +6,11 @@ import logging from typing import Any, Protocol -from nemo_platform import NeMoPlatform, omit -from nemo_platform._exceptions import APIError +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoHTTPError +from nemo_platform_plugin.jobs.client import JobsClient +from nemo_platform_plugin.jobs.types import PlatformJobTaskUpdate from nmp.common.jobs.schemas import PlatformJobStatus from nmp.customization_common.schemas.file_io import ProgressReportError from nmp.customization_common.service.context import NMPJobContext @@ -64,17 +67,25 @@ def update_progress( with sdk_error_handler( ProgressReportError, f"update progress for task: {self.task_id}, job: {self.job_id}, step: {self.step_name}", - passthrough=(APIError,), + passthrough=(NemoHTTPError,), ): - self.sdk.jobs.tasks.create_or_update( - self.task_id, + # Only set fields that have values (mirrors the SDK's ``omit`` + # sentinels — ``exclude_unset`` keeps them off the wire). + task_update: dict[str, Any] = {"status": status} + if status_details: + task_update["status_details"] = status_details + if error_details: + task_update["error_details"] = error_details + if error_stack: + task_update["error_stack"] = error_stack + + jobs = client_from_platform(self.sdk, JobsClient) + jobs.update_job_step_task( + name=self.task_id, workspace=self.workspace, job=self.job_id, step=self.step_name, - status=status.value, - status_details=status_details if status_details else omit, - error_details=error_details if error_details else omit, - error_stack=error_stack if error_stack else omit, + body=PlatformJobTaskUpdate(**task_update), ) logger.debug(f"Progress updated: {status} - {status_details}") except Exception as e: diff --git a/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py b/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py index 8ab1381be1..f56221da11 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py @@ -16,6 +16,10 @@ import os from typing import Any, cast +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.jobs.client import JobsClient +from nemo_platform_plugin.jobs.schemas import PlatformJobStatus +from nemo_platform_plugin.jobs.types import PlatformJobTaskUpdate from nmp.common.sdk_factory import get_task_sdk from nmp.customization_common.service.context import NMPJobContext @@ -62,14 +66,17 @@ def update_task( return try: - self._sdk.jobs.tasks.create_or_update( + jobs = client_from_platform(self._sdk, JobsClient) + jobs.update_job_step_task( name=self._job_ctx.normalized_task, workspace=self._job_ctx.workspace, job=self._job_ctx.job_id, step=self._job_ctx.step, - status=status, # ty: ignore[invalid-argument-type] - status_details=status_details or {}, - error_details=error_details or {}, + body=PlatformJobTaskUpdate( + status=PlatformJobStatus(status), + status_details=status_details or {}, + error_details=error_details or {}, + ), ) except Exception as e: logger.warning(f"Failed to update task progress: {e}") @@ -79,12 +86,13 @@ def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: return {"train_loss": [], "val_loss": []} try: - task = self._sdk.jobs.tasks.retrieve( + jobs = client_from_platform(self._sdk, JobsClient) + task = jobs.get_job_step_task( name=self._job_ctx.normalized_task, workspace=self._job_ctx.workspace, job=self._job_ctx.job_id, step=self._job_ctx.step, - ) + ).data() metrics = cast(dict[str, Any], (task.status_details or {}).get("metrics", {}) or {}) return { "train_loss": metrics.get("train_loss", []), diff --git a/packages/nmp_testing/src/nmp/testing/e2e/jobs.py b/packages/nmp_testing/src/nmp/testing/e2e/jobs.py index 764ee815b0..5dd12838e4 100644 --- a/packages/nmp_testing/src/nmp/testing/e2e/jobs.py +++ b/packages/nmp_testing/src/nmp/testing/e2e/jobs.py @@ -12,7 +12,10 @@ from collections.abc import Callable from nemo_platform import NeMoPlatform -from nemo_platform.types.jobs.platform_job_response import PlatformJobResponse +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.jobs.client import JobsClient +from nemo_platform_plugin.jobs.schemas import PlatformJobLogPage +from nemo_platform_plugin.jobs.types import PlatformJobResponse logger = logging.getLogger(__name__) @@ -125,7 +128,7 @@ def wait_for_platform_job( def get_status() -> str: nonlocal last_job - last_job = sdk.jobs.retrieve(job_name, workspace=workspace) + last_job = client_from_platform(sdk, JobsClient).get_job(name=job_name, workspace=workspace).data() current = last_job.status.lower() if last_job.status else "" if not status_history or status_history[-1] != current: status_history.append(current) @@ -143,7 +146,7 @@ def get_status() -> str: except TimeoutError as e: error_parts = [str(e), f"Status history: {' -> '.join(status_history)}"] try: - job_status = sdk.jobs.get_status(job_name, workspace=workspace) + job_status = client_from_platform(sdk, JobsClient).get_job_status(name=job_name, workspace=workspace).data() error_parts.append(f"Job status details: {job_status.model_dump()}") except Exception as detail_err: error_parts.append(f"Failed to get job status: {detail_err}") @@ -257,7 +260,8 @@ def wait_for_job_logs( logs = None while time.time() - start_time < timeout: - logs = sdk.jobs.get_logs(workspace=workspace, name=job_name) + page = client_from_platform(sdk, JobsClient).list_job_logs(workspace=workspace, name=job_name).page() + logs = PlatformJobLogPage(data=page.items, **page.metadata) if len(logs.data) >= min_log_count: return logs time.sleep(poll_interval) diff --git a/packages/nmp_testing/tests/unit/test_jobs.py b/packages/nmp_testing/tests/unit/test_jobs.py index 61092e5960..8b43435cdc 100644 --- a/packages/nmp_testing/tests/unit/test_jobs.py +++ b/packages/nmp_testing/tests/unit/test_jobs.py @@ -15,6 +15,7 @@ counted against the main job-execution timeout. """ +from contextlib import contextmanager from unittest.mock import MagicMock, patch import pytest @@ -25,17 +26,46 @@ # --------------------------------------------------------------------------- -def _make_sdk(*statuses: str) -> MagicMock: - """Return an SDK mock whose retrieve() cycles through *statuses* on each call.""" - sdk = MagicMock() - jobs = [] +def _resp(data): + """Wrap a payload in a NemoResponse-like object whose ``.data()`` returns it. + + Production now consumes typed-client responses via ``client.(...).data()``, + so mocked jobs-client methods must return an object with a ``.data()`` accessor + rather than the payload directly. + """ + m = MagicMock() + m.data.return_value = data + return m + + +def _make_jobs_client(*statuses: str) -> MagicMock: + """Return a typed jobs-client mock whose get_job() cycles through *statuses*. + + Each ``get_job`` call returns a ``_resp(job)`` where ``job.status`` is the next + status in *statuses*. ``get_job_status`` returns a ``_resp`` around a model with + an empty ``model_dump``. + """ + jobs_client = MagicMock() + responses = [] for s in statuses: j = MagicMock() j.status = s - jobs.append(j) - sdk.jobs.retrieve.side_effect = jobs - sdk.jobs.get_status.return_value = MagicMock(model_dump=MagicMock(return_value={})) - return sdk + responses.append(_resp(j)) + jobs_client.get_job.side_effect = responses + jobs_client.get_job_status.return_value = _resp(MagicMock(model_dump=MagicMock(return_value={}))) + return jobs_client + + +@contextmanager +def _patch_client(jobs_client: MagicMock): + """Patch ``client_from_platform`` in the production module to return *jobs_client*.""" + with patch("nmp.testing.e2e.jobs.client_from_platform", return_value=jobs_client): + yield + + +def _make_sdk(*statuses: str) -> MagicMock: + """Return an SDK mock (unused by production routing, kept for call signatures).""" + return MagicMock() # --------------------------------------------------------------------------- @@ -48,30 +78,34 @@ class TestWaitForPlatformJobTerminalStatus: def test_returns_immediately_on_completed(self): """Returns as soon as job is 'completed'.""" - sdk = _make_sdk("completed") - job = wait_for_platform_job(sdk, "my-job", "ws", timeout=5.0) + jobs_client = _make_jobs_client("completed") + with _patch_client(jobs_client): + job = wait_for_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0) assert job.status == "completed" def test_returns_immediately_on_error(self): """Returns (without raising) when job is 'error'.""" - sdk = _make_sdk("error") - job = wait_for_platform_job(sdk, "my-job", "ws", timeout=5.0) + jobs_client = _make_jobs_client("error") + with _patch_client(jobs_client): + job = wait_for_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0) assert job.status == "error" def test_returns_immediately_on_cancelled(self): """Returns when job is 'cancelled'.""" - sdk = _make_sdk("cancelled") - job = wait_for_platform_job(sdk, "my-job", "ws", timeout=5.0) + jobs_client = _make_jobs_client("cancelled") + with _patch_client(jobs_client): + job = wait_for_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0) assert job.status == "cancelled" def test_returns_job_object(self): - """Returns the actual job object from sdk.jobs.retrieve (not a copy).""" - sdk = MagicMock() + """Returns the actual job object from get_job().data() (not a copy).""" expected_job = MagicMock() expected_job.status = "completed" - sdk.jobs.retrieve.return_value = expected_job - sdk.jobs.get_status.return_value = MagicMock(model_dump=MagicMock(return_value={})) - job = wait_for_platform_job(sdk, "my-job", "ws", timeout=5.0) + jobs_client = MagicMock() + jobs_client.get_job.return_value = _resp(expected_job) + jobs_client.get_job_status.return_value = _resp(MagicMock(model_dump=MagicMock(return_value={}))) + with _patch_client(jobs_client): + job = wait_for_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0) assert job is expected_job @@ -85,26 +119,28 @@ class TestWaitForPlatformJobStatusToCheck: def test_stops_on_status_to_check(self): """Returns when the job reaches status_to_check before terminal.""" - sdk = _make_sdk("created", "pending", "active") - job = wait_for_platform_job(sdk, "my-job", "ws", timeout=5.0, status_to_check="active") + jobs_client = _make_jobs_client("created", "pending", "active") + with _patch_client(jobs_client): + job = wait_for_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0, status_to_check="active") assert job.status == "active" def test_also_stops_on_terminal_when_status_to_check_set(self): """If job reaches a terminal status before status_to_check, still returns.""" - sdk = _make_sdk("created", "error") - job = wait_for_platform_job(sdk, "my-job", "ws", timeout=5.0, status_to_check="active") + jobs_client = _make_jobs_client("created", "error") + with _patch_client(jobs_client): + job = wait_for_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0, status_to_check="active") assert job.status == "error" def test_terminal_set_includes_status_to_check(self): """poll_until_terminal is called with status_to_check merged into terminal set.""" - sdk = _make_sdk("paused") - with patch("nmp.testing.e2e.jobs.poll_until_terminal") as mock_poll: + jobs_client = _make_jobs_client("paused") + with _patch_client(jobs_client), patch("nmp.testing.e2e.jobs.poll_until_terminal") as mock_poll: # Simulate poll_until_terminal calling get_status once def fake_poll(get_status, label, terminal, timeout, image_pull_timeout, poll_interval): get_status() mock_poll.side_effect = fake_poll - wait_for_platform_job(sdk, "my-job", "ws", status_to_check="paused") + wait_for_platform_job(_make_sdk(), "my-job", "ws", status_to_check="paused") _, kwargs = mock_poll.call_args terminal_used = mock_poll.call_args[1]["terminal"] if mock_poll.call_args[1] else mock_poll.call_args[0][2] @@ -124,20 +160,21 @@ class TestWaitForPlatformJobImagePullTimeout: def test_pending_status_does_not_consume_main_timeout(self): """A job stuck in pending does not exhaust the execution timeout.""" # pending -> completed: pending time should NOT count against timeout=5.0 - sdk = _make_sdk("pending", "completed") - job = wait_for_platform_job(sdk, "my-job", "ws", timeout=5.0, image_pull_timeout=60.0) + jobs_client = _make_jobs_client("pending", "completed") + with _patch_client(jobs_client): + job = wait_for_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0, image_pull_timeout=60.0) assert job.status == "completed" def test_image_pull_timeout_parameter_passed_to_poll_until_terminal(self): """image_pull_timeout is forwarded to poll_until_terminal.""" - sdk = _make_sdk("completed") - with patch("nmp.testing.e2e.jobs.poll_until_terminal") as mock_poll: + jobs_client = _make_jobs_client("completed") + with _patch_client(jobs_client), patch("nmp.testing.e2e.jobs.poll_until_terminal") as mock_poll: def fake_poll(get_status, label, terminal, timeout, image_pull_timeout, poll_interval): get_status() mock_poll.side_effect = fake_poll - wait_for_platform_job(sdk, "my-job", "ws", timeout=30.0, image_pull_timeout=999.0) + wait_for_platform_job(_make_sdk(), "my-job", "ws", timeout=30.0, image_pull_timeout=999.0) args = mock_poll.call_args # image_pull_timeout may be positional or keyword @@ -148,14 +185,14 @@ def fake_poll(get_status, label, terminal, timeout, image_pull_timeout, poll_int def test_default_image_pull_timeout_is_600(self): """Default image_pull_timeout is 600 seconds.""" - sdk = _make_sdk("completed") - with patch("nmp.testing.e2e.jobs.poll_until_terminal") as mock_poll: + jobs_client = _make_jobs_client("completed") + with _patch_client(jobs_client), patch("nmp.testing.e2e.jobs.poll_until_terminal") as mock_poll: def fake_poll(get_status, label, terminal, timeout, image_pull_timeout, poll_interval): get_status() mock_poll.side_effect = fake_poll - wait_for_platform_job(sdk, "my-job", "ws") + wait_for_platform_job(_make_sdk(), "my-job", "ws") args = mock_poll.call_args if args[1]: @@ -174,15 +211,15 @@ class TestWaitForPlatformJobTimeoutError: def test_raises_timeout_error_when_poll_times_out(self): """TimeoutError propagates when poll_until_terminal raises it.""" - sdk = _make_sdk("created") - with patch("nmp.testing.e2e.jobs.poll_until_terminal") as mock_poll: + jobs_client = _make_jobs_client("created") + with _patch_client(jobs_client), patch("nmp.testing.e2e.jobs.poll_until_terminal") as mock_poll: mock_poll.side_effect = TimeoutError("'my-job' timed out after 5.0s. Status: created") with pytest.raises(TimeoutError): - wait_for_platform_job(sdk, "my-job", "ws", timeout=5.0) + wait_for_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0) def test_timeout_error_includes_status_history(self): """TimeoutError message includes the accumulated status history.""" - sdk = _make_sdk("created", "pending") + jobs_client = _make_jobs_client("created", "pending") def fake_poll(get_status, label, terminal, timeout, image_pull_timeout, poll_interval): # Call get_status twice to populate history, then timeout @@ -190,27 +227,27 @@ def fake_poll(get_status, label, terminal, timeout, image_pull_timeout, poll_int get_status() raise TimeoutError(f"'{label}' timed out after {timeout}s. Status: pending") - with patch("nmp.testing.e2e.jobs.poll_until_terminal", side_effect=fake_poll): + with _patch_client(jobs_client), patch("nmp.testing.e2e.jobs.poll_until_terminal", side_effect=fake_poll): with pytest.raises(TimeoutError) as exc_info: - wait_for_platform_job(sdk, "my-job", "ws", timeout=5.0) + wait_for_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0) assert "Status history:" in str(exc_info.value) assert "created" in str(exc_info.value) assert "pending" in str(exc_info.value) def test_timeout_error_includes_job_status_details(self): - """TimeoutError message includes detailed job status from get_status API.""" - sdk = _make_sdk("pending") - sdk.jobs.get_status.return_value = MagicMock( - model_dump=MagicMock(return_value={"status": "pending", "message": "pulling image"}) + """TimeoutError message includes detailed job status from get_job_status API.""" + jobs_client = _make_jobs_client("pending") + jobs_client.get_job_status.return_value = _resp( + MagicMock(model_dump=MagicMock(return_value={"status": "pending", "message": "pulling image"})) ) def fake_poll(get_status, label, terminal, timeout, image_pull_timeout, poll_interval): get_status() raise TimeoutError(f"'{label}' timed out") - with patch("nmp.testing.e2e.jobs.poll_until_terminal", side_effect=fake_poll): + with _patch_client(jobs_client), patch("nmp.testing.e2e.jobs.poll_until_terminal", side_effect=fake_poll): with pytest.raises(TimeoutError) as exc_info: - wait_for_platform_job(sdk, "my-job", "ws", timeout=5.0) + wait_for_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0) assert "Job status details:" in str(exc_info.value) diff --git a/plugins/example-plugin/tests/test_sdk.py b/plugins/example-plugin/tests/test_sdk.py index 7953180822..861839084f 100644 --- a/plugins/example-plugin/tests/test_sdk.py +++ b/plugins/example-plugin/tests/test_sdk.py @@ -116,7 +116,19 @@ def test_sync_get_item() -> None: def test_sync_list_items() -> None: client, mock_http = _sync_client() mock_http.request.return_value = _resp( - 200, {"data": [ITEM_PAYLOAD], "pagination": None, "sort": None, "filter": None} + 200, + { + "data": [ITEM_PAYLOAD], + "pagination": { + "page": 1, + "page_size": 20, + "current_page_size": 1, + "total_pages": 1, + "total_results": 1, + }, + "sort": None, + "filter": None, + }, ) resp = client.list_items() @@ -189,7 +201,19 @@ async def test_async_get_item() -> None: async def test_async_list_items() -> None: client, mock_http = _async_client() mock_http.request.return_value = _resp( - 200, {"data": [ITEM_PAYLOAD], "pagination": None, "sort": None, "filter": None} + 200, + { + "data": [ITEM_PAYLOAD], + "pagination": { + "page": 1, + "page_size": 20, + "current_page_size": 1, + "total_pages": 1, + "total_results": 1, + }, + "sort": None, + "filter": None, + }, ) resp = await client.list_items() diff --git a/plugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.py b/plugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.py index 9c613efbb3..02b06fca43 100644 --- a/plugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.py +++ b/plugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.py @@ -36,7 +36,9 @@ from nemo_platform_plugin.job_context import JobContext, StoragePaths from nemo_platform_plugin.job_results import PlatformJobResults from nemo_platform_plugin.jobs.api_factory import PlatformJobSpec +from nemo_platform_plugin.jobs.client import JobsClient from nemo_platform_plugin.jobs.result_manager import ResultManager +from nemo_platform_plugin.jobs.types import CreatePlatformJobRequest from nemo_platform_plugin.secrets.client import SecretsClient from nemo_platform_plugin.secrets.types import PlatformSecretCreateRequest from nmp.core.files.service import FilesService @@ -431,16 +433,19 @@ def __init__(self) -> None: workspace="default", ) as client_context, ): - job = client_context.sdk.jobs.create( + jobs_client = client_from_platform(client_context.sdk, JobsClient) + job = jobs_client.create_job( workspace="default", - name=job_name, - source="data-designer", - # Store the canonical DataDesignerStepConfig as the job's spec so that - # downstream Data Designer routes (e.g. ``GET /jobs/create/{name}``, - # which deserializes the stored spec back through the schema) succeed. - spec=step_config, - platform_spec=job_config_dict, - ) + body=CreatePlatformJobRequest( + name=job_name, + source="data-designer", + # Store the canonical DataDesignerStepConfig as the job's spec so that + # downstream Data Designer routes (e.g. ``GET /jobs/create/{name}``, + # which deserializes the stored spec back through the schema) succeed. + spec=step_config if isinstance(step_config, dict) else step_config.model_dump(), + platform_spec=job_config_dict, + ), + ).data() job_ctx = JobContext( workspace="default", storage=StoragePaths(ephemeral=ephemeral, persistent=persistent), diff --git a/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/api/v2/jobs/endpoints.py b/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/api/v2/jobs/endpoints.py index 61e057bac8..78b4fc3a45 100644 --- a/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/api/v2/jobs/endpoints.py +++ b/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/api/v2/jobs/endpoints.py @@ -35,6 +35,7 @@ SubprocessExecutionProviderSpec, job_route_factory, ) +from nemo_platform_plugin.jobs.client import AsyncJobsClient from nemo_platform_plugin.jobs.exceptions import PlatformJobCompilationError from nemo_platform_plugin.jobs.image import get_qualified_image from nemo_safe_synthesizer.config.external_results import SafeSynthesizerSummary @@ -176,12 +177,13 @@ async def job_config_compiler( transformed_spec.pretrained_model_job, workspace_fallback=workspace ) try: - await sdk.jobs.results.retrieve(name="adapter", job=model_job, workspace=model_workspace) - except NotFoundError as e: + jobs = client_from_platform(sdk, AsyncJobsClient) + await jobs.get_job_result(name="adapter", job=model_job, workspace=model_workspace) + except ClientNotFoundError as e: raise PlatformJobCompilationError( f"Could not find adapter result for NSS job {model_workspace}/{model_job!r}" ) from e - except PermissionDeniedError as e: + except ClientPermissionDeniedError as e: raise PlatformJobCompilationError( f"Failed to retrieve adapter result for NSS job {model_workspace}/{model_job!r}: " f"access denied to workspace {model_workspace!r}" diff --git a/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/sdk/job.py b/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/sdk/job.py index ffd36c4cb4..c9c3a1a2a9 100644 --- a/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/sdk/job.py +++ b/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/sdk/job.py @@ -16,8 +16,11 @@ import httpx import pandas as pd from nemo_platform import NeMoPlatform -from nemo_platform._types import Omit, omit -from nemo_platform.types import PlatformJobLog, PlatformJobStatusResponse +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoClientError +from nemo_platform_plugin.jobs.client import JobsClient +from nemo_platform_plugin.jobs.schemas import PlatformJobLog, PlatformJobStatusResponse +from nemo_platform_plugin.jobs.types import JobLogsQueryParams from nemo_safe_synthesizer.config.external_results import SafeSynthesizerSummary from typing_extensions import Self @@ -58,6 +61,7 @@ def __init__(self, job_name: str, client: NeMoPlatform, workspace: str = "defaul self.job_name = job_name self._client = client self._workspace = workspace + self._jobs = client_from_platform(client, JobsClient) def fetch_status(self) -> str: """Fetch the current job status.""" @@ -65,7 +69,7 @@ def fetch_status(self) -> str: def fetch_status_info(self) -> PlatformJobStatusResponse: """Fetch the current job status response.""" - return self._client.jobs.get_status(self.job_name, workspace=self._workspace) + return self._jobs.get_job_status(name=self.job_name, workspace=self._workspace).data() def wait_for_completion( self, poll_interval: int = 10, verbose: bool = True, log_timeout: float | None = None @@ -90,7 +94,7 @@ def wait_for_completion( if log_key not in seen_log_keys: print(new_log.message.strip()) seen_log_keys.add(log_key) - except httpx.HTTPError as e: + except (NemoClientError, httpx.HTTPError) as e: logger.warning("Error fetching logs while waiting for job completion: %s", e) finally: if logging_level is not None: @@ -115,13 +119,15 @@ def wait_for_completion( def fetch_summary(self) -> SafeSynthesizerSummary: """Fetch the machine-readable job summary.""" - response = self._client.jobs.results.download("summary", job=self.job_name, workspace=self._workspace) - return SafeSynthesizerSummary.model_validate(json.loads(response.read().decode("utf-8"))) + data = self._jobs.download_job_result(name="summary", job=self.job_name, workspace=self._workspace).read() + return SafeSynthesizerSummary.model_validate(json.loads(data.decode("utf-8"))) def fetch_report(self) -> ReportHtml: """Fetch the evaluation report as HTML.""" - response = self._client.jobs.results.download("evaluation-report", job=self.job_name, workspace=self._workspace) - return ReportHtml(html=response.read().decode("utf-8")) + data = self._jobs.download_job_result( + name="evaluation-report", job=self.job_name, workspace=self._workspace + ).read() + return ReportHtml(html=data.decode("utf-8")) def display_report_in_notebook(self, width: str = "100%", height: int = 1000) -> None: """Display the evaluation report in a Jupyter notebook.""" @@ -133,8 +139,10 @@ def save_report(self, path: str | Path) -> None: def fetch_data(self) -> pd.DataFrame: """Fetch generated synthetic data as a pandas DataFrame.""" - response = self._client.jobs.results.download("synthetic-data", job=self.job_name, workspace=self._workspace) - return pd.read_csv(BytesIO(response.read())) + data = self._jobs.download_job_result( + name="synthetic-data", job=self.job_name, workspace=self._workspace + ).read() + return pd.read_csv(BytesIO(data)) def _fetch_logs_incremental( self, page_cursor: str | None = None, timeout: float | None = None @@ -142,39 +150,36 @@ def _fetch_logs_incremental( """Fetch logs incrementally starting from a page cursor.""" timeout = 300.0 if timeout is None else timeout all_logs: list[PlatformJobLog] = [] - current_cursor: str | Omit = omit if page_cursor is None else page_cursor + current_cursor: str | None = page_cursor last_cursor_with_data: str | None = page_cursor while True: - response = self._client.with_options(timeout=timeout).jobs.get_logs( - self.job_name, - page_cursor=current_cursor, - workspace=self._workspace, + logs_query: JobLogsQueryParams = {} + if current_cursor is not None: + logs_query["page_cursor"] = current_cursor + page = ( + self._jobs.with_options(timeout=timeout) + .list_job_logs(name=self.job_name, workspace=self._workspace, query_params=logs_query) + .page() ) - if response.data: - all_logs.extend(response.data) - if isinstance(current_cursor, str): + if page.items: + all_logs.extend(page.items) + if current_cursor is not None: last_cursor_with_data = current_cursor - if response.next_page is None: + if page.metadata["next_page"] is None: return all_logs, last_cursor_with_data - current_cursor = response.next_page + current_cursor = page.metadata["next_page"] def fetch_logs(self, timeout: float | None = None) -> Iterator[PlatformJobLog]: """Fetch job logs as an iterator over log objects.""" timeout = 300.0 if timeout is None else timeout - page_cursor: str | Omit = omit - while True: - response = self._client.with_options(timeout=timeout).jobs.get_logs( - self.job_name, - page_cursor=page_cursor, - workspace=self._workspace, - ) - yield from response.data - if response.next_page is None: - break - page_cursor = response.next_page + yield from ( + self._jobs.with_options(timeout=timeout) + .list_job_logs(name=self.job_name, workspace=self._workspace) + .items() + ) def print_logs(self, timeout: float | None = None) -> None: """Print job logs to stdout.""" diff --git a/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/sdk/resources.py b/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/sdk/resources.py index 910e86b8b7..f4505096f1 100644 --- a/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/sdk/resources.py +++ b/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/sdk/resources.py @@ -5,11 +5,15 @@ from __future__ import annotations -from typing import Any +from typing import Any, cast from urllib.parse import quote import httpx from nemo_platform import AsyncNeMoPlatform, NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.jobs.client import AsyncJobsClient, JobsClient +from nemo_platform_plugin.jobs.schemas import PlatformJobLogPage +from nemo_platform_plugin.jobs.types import JobLogsQueryParams from nemo_platform_plugin.sdk import NemoPluginSDKResources from nemo_safe_synthesizer_plugin.sdk import http_utils @@ -80,11 +84,27 @@ def retrieve(self, name: str, *, workspace: str | None = None) -> Any: def get_status(self, name: str, *, workspace: str | None = None) -> Any: """Retrieve Safe Synthesizer job status.""" - return self._platform.jobs.get_status(name, workspace=workspace) + return client_from_platform(self._platform, JobsClient).get_job_status(name=name, workspace=workspace).data() - def get_logs(self, name: str, *, workspace: str | None = None, **kwargs: Any) -> Any: + def get_logs( + self, + name: str, + *, + workspace: str | None = None, + **params: Any, + ) -> Any: """Retrieve paginated Safe Synthesizer job logs from the Jobs service.""" - return self._platform.jobs.get_logs(name, workspace=workspace, **kwargs) + query_params = {key: value for key, value in params.items() if value is not None} + page = ( + client_from_platform(self._platform, JobsClient) + .list_job_logs( + name=name, + workspace=workspace, + query_params=cast(JobLogsQueryParams, query_params) or None, + ) + .page() + ) + return PlatformJobLogPage(data=page.items, **page.metadata) class SafeSynthesizerResource: @@ -161,11 +181,25 @@ async def retrieve(self, name: str, *, workspace: str | None = None) -> Any: async def get_status(self, name: str, *, workspace: str | None = None) -> Any: """Retrieve Safe Synthesizer job status.""" - return await self._platform.jobs.get_status(name, workspace=workspace) + jobs = client_from_platform(self._platform, AsyncJobsClient) + return (await jobs.get_job_status(name=name, workspace=workspace)).data() - async def get_logs(self, name: str, *, workspace: str | None = None, **kwargs: Any) -> Any: + async def get_logs( + self, + name: str, + *, + workspace: str | None = None, + **params: Any, + ) -> Any: """Retrieve paginated Safe Synthesizer job logs from the Jobs service.""" - return await self._platform.jobs.get_logs(name, workspace=workspace, **kwargs) + query_params = {key: value for key, value in params.items() if value is not None} + response = await client_from_platform(self._platform, AsyncJobsClient).list_job_logs( + name=name, + workspace=workspace, + query_params=cast(JobLogsQueryParams, query_params) or None, + ) + page = response.page() + return PlatformJobLogPage(data=page.items, **page.metadata) class AsyncSafeSynthesizerResource: diff --git a/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/tasks/safe_synthesizer/__main__.py b/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/tasks/safe_synthesizer/__main__.py index 1f585ef928..83b26c6e94 100644 --- a/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/tasks/safe_synthesizer/__main__.py +++ b/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/tasks/safe_synthesizer/__main__.py @@ -27,7 +27,9 @@ from datasets import Dataset, DatasetDict, load_dataset from nemo_platform import NeMoPlatform from nemo_platform.filesets import parse_fileset_ref +from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.config import get_platform_config +from nemo_platform_plugin.jobs.client import JobsClient from nemo_platform_plugin.jobs.constants import ( DEFAULT_TASK_STORAGE_PATH, EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, @@ -36,6 +38,7 @@ NEMO_JOB_WORKSPACE_ENVVAR, ) from nemo_platform_plugin.jobs.file_manager import FilesetFileManager +from nemo_platform_plugin.jobs.schemas import PlatformJobResultCreateRequest from nemo_platform_plugin.sdk_provider import get_platform_sdk from nemo_safe_synthesizer.config.internal_results import SafeSynthesizerResults from nemo_safe_synthesizer.observability import initialize_observability @@ -154,7 +157,7 @@ def upload_results(result: SafeSynthesizerResults, adapter_path: Path | None = N ) file_manager.validate_storage() - job = sdk.jobs.retrieve(name=job_id, workspace=workspace) + job = client_from_platform(sdk, JobsClient).get_job(name=job_id, workspace=workspace).data() attempt_id = job.attempt_id with tempfile.TemporaryDirectory() as temp_dir: @@ -209,12 +212,11 @@ def write_results_local(result: SafeSynthesizerResults, output_dir: Path, adapte def _create_job_result(sdk: NeMoPlatform, workspace: str, job_name: str, result_name: str, artifact_url: str): """Create a job result record.""" - sdk.jobs.results.create( + client_from_platform(sdk, JobsClient).create_job_result( name=result_name, job=job_name, workspace=workspace, - artifact_url=artifact_url, - artifact_storage_type="fileset", + body=PlatformJobResultCreateRequest(artifact_url=artifact_url, artifact_storage_type="fileset"), ) logger.info("Created job result: %s", result_name) @@ -237,7 +239,11 @@ def _resolve_pretrained_model( workspace_fallback=workspace, ) try: - adapter_result = sdk.jobs.results.retrieve(name="adapter", job=model_job, workspace=model_workspace) + adapter_result = ( + client_from_platform(sdk, JobsClient) + .get_job_result(name="adapter", job=model_job, workspace=model_workspace) + .data() + ) except Exception as e: raise RuntimeError( f"Failed to resolve adapter result for pretrained_model_job={job_config.pretrained_model_job!r}" diff --git a/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py b/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py index e1165525a3..fbeb3de847 100644 --- a/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py +++ b/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py @@ -3,8 +3,8 @@ from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest -from nemo_platform import NotFoundError from nemo_platform_plugin.client.errors import NotFoundError as ClientNotFoundError from nemo_platform_plugin.client.errors import PermissionDeniedError as ClientPermissionDeniedError from nemo_platform_plugin.jobs.exceptions import PlatformJobCompilationError @@ -19,6 +19,38 @@ DEFAULT_DATA_SOURCE = "default/test-data#file.csv" +def _client_error(error_cls, status_code: int, detail: str): + """Build a NemoHTTPError subclass from an httpx.Response, as the typed client raises.""" + request = httpx.Request("GET", "http://test") + response = httpx.Response(status_code=status_code, json={"detail": detail}, request=request) + return error_cls(response) + + +def _patch_jobs_client(jobs_client: MagicMock, files_client: MagicMock): + """Patch ``client_from_platform`` in the endpoints module, dispatching by class. + + The compiler validates the ``data_source`` fileset via + ``client_from_platform(sdk, AsyncFilesClient).get_fileset(...)`` and then resolves + the pretrained-model adapter via + ``client_from_platform(sdk, AsyncJobsClient).get_job_result(...)`` — return the + matching mock for each. + """ + from nemo_platform_plugin.files.client import AsyncFilesClient + from nemo_platform_plugin.jobs.client import AsyncJobsClient + + def _dispatch(_sdk, client_cls): + if client_cls is AsyncJobsClient: + return jobs_client + if client_cls is AsyncFilesClient: + return files_client + raise AssertionError(f"unexpected client class: {client_cls!r}") + + return patch( + "nemo_safe_synthesizer_plugin.api.v2.jobs.endpoints.client_from_platform", + side_effect=_dispatch, + ) + + @pytest.fixture def mock_files_client(): mock_client = MagicMock() @@ -154,8 +186,9 @@ async def test_job_config_compiler_container_mode_uses_image_ref_override(mock_s @pytest.mark.asyncio -async def test_job_config_compiler_validates_pretrained_model_job(mock_sdk): - mock_sdk.jobs.results.retrieve = AsyncMock( +async def test_job_config_compiler_validates_pretrained_model_job(mock_sdk, mock_files_client): + jobs_client = MagicMock() + jobs_client.get_job_result = AsyncMock( return_value=MagicMock(artifact_url="default/job-results-prior#results/attempt-1/adapter") ) spec = PluginJobConfig.model_validate( @@ -166,9 +199,10 @@ async def test_job_config_compiler_validates_pretrained_model_job(mock_sdk): } ) - await _compile(spec, mock_sdk) + with _patch_jobs_client(jobs_client, mock_files_client): + await _compile(spec, mock_sdk) - mock_sdk.jobs.results.retrieve.assert_awaited_once_with( + jobs_client.get_job_result.assert_awaited_once_with( name="adapter", job="prior-safe-synth-job", workspace=DEFAULT_WORKSPACE, @@ -176,8 +210,9 @@ async def test_job_config_compiler_validates_pretrained_model_job(mock_sdk): @pytest.mark.asyncio -async def test_plugin_job_config_allows_pretrained_model_job_runtime_config(mock_sdk): - mock_sdk.jobs.results.retrieve = AsyncMock( +async def test_plugin_job_config_allows_pretrained_model_job_runtime_config(mock_sdk, mock_files_client): + jobs_client = MagicMock() + jobs_client.get_job_result = AsyncMock( return_value=MagicMock(artifact_url="default/job-results-prior#results/attempt-1/adapter") ) spec = PluginJobConfig.model_validate( @@ -188,7 +223,8 @@ async def test_plugin_job_config_allows_pretrained_model_job_runtime_config(mock } ) - compiled = await _compile(spec, mock_sdk) + with _patch_jobs_client(jobs_client, mock_files_client): + compiled = await _compile(spec, mock_sdk) step = next(iter(compiled["steps"])) reparsed = PluginJobConfig.model_validate(step["config"]) @@ -238,10 +274,9 @@ def test_runtime_job_config_preserves_pretrained_model_without_pretrained_model_ @pytest.mark.asyncio -async def test_job_config_compiler_pretrained_model_job_not_found(mock_sdk): - mock_sdk.jobs.results.retrieve = AsyncMock( - side_effect=NotFoundError(message="not found", response=MagicMock(status_code=404), body=None) - ) +async def test_job_config_compiler_pretrained_model_job_not_found(mock_sdk, mock_files_client): + jobs_client = MagicMock() + jobs_client.get_job_result = AsyncMock(side_effect=_client_error(ClientNotFoundError, 404, "not found")) spec = PluginJobConfig.model_validate( { "data_source": DEFAULT_DATA_SOURCE, @@ -250,8 +285,9 @@ async def test_job_config_compiler_pretrained_model_job_not_found(mock_sdk): } ) - with pytest.raises(PlatformJobCompilationError, match="Could not find adapter result"): - await _compile(spec, mock_sdk) + with _patch_jobs_client(jobs_client, mock_files_client): + with pytest.raises(PlatformJobCompilationError, match="Could not find adapter result"): + await _compile(spec, mock_sdk) def test_plugin_job_config_rejects_conflicting_pretrained_model_sources(): diff --git a/plugins/nemo-safe-synthesizer/tests/unit/test_local_run.py b/plugins/nemo-safe-synthesizer/tests/unit/test_local_run.py index d5ff3b9db1..89e4de5dde 100644 --- a/plugins/nemo-safe-synthesizer/tests/unit/test_local_run.py +++ b/plugins/nemo-safe-synthesizer/tests/unit/test_local_run.py @@ -11,6 +11,16 @@ import pytest +def _resp(data): + """Wrap a payload in a NemoResponse-like object whose ``.data()`` returns it. + + Production consumes typed-client responses via ``client.(...).data()``. + """ + m = MagicMock() + m.data.return_value = data + return m + + def import_task_main_without_heavy_runtime(monkeypatch): pytest.importorskip("nemo_safe_synthesizer.config.job") library_builder = ModuleType("nemo_safe_synthesizer.sdk.library_builder") @@ -128,11 +138,16 @@ def test_run_local_resolves_pretrained_model_job_before_run(tmp_path, monkeypatc ) sdk = MagicMock() - sdk.jobs.results.retrieve.return_value = SimpleNamespace( - artifact_url="default/job-results-prior#results/attempt-1/adapter" - ) monkeypatch.setattr(task_main, "get_platform_sdk", lambda: sdk) + # Production resolves the prior adapter via + # client_from_platform(sdk, JobsClient).get_job_result(...).data(). + jobs_client = MagicMock() + jobs_client.get_job_result.return_value = _resp( + SimpleNamespace(artifact_url="default/job-results-prior#results/attempt-1/adapter") + ) + monkeypatch.setattr(task_main, "client_from_platform", lambda _sdk, _cls: jobs_client) + file_manager = MagicMock() file_manager.download_from_url.return_value = SimpleNamespace( path=adapter_dir, @@ -157,7 +172,7 @@ def fake_run_config(job_config, data_source, save_path, *, adapter_location=None task_main.run_local(spec_file=spec_file, workspace="default", output_dir=output_dir, data_source=data_file) assert captured["adapter_location"] == adapter_dir - sdk.jobs.results.retrieve.assert_called_once_with( + jobs_client.get_job_result.assert_called_once_with( name="adapter", job="prior-safe-synth-job", workspace="default", @@ -191,11 +206,16 @@ def test_run_local_cleans_pretrained_model_tmp_when_run_config_raises(tmp_path, ) sdk = MagicMock() - sdk.jobs.results.retrieve.return_value = SimpleNamespace( - artifact_url="default/job-results-prior#results/attempt-1/adapter" - ) monkeypatch.setattr(task_main, "get_platform_sdk", lambda: sdk) + # Production resolves the prior adapter via + # client_from_platform(sdk, JobsClient).get_job_result(...).data(). + jobs_client = MagicMock() + jobs_client.get_job_result.return_value = _resp( + SimpleNamespace(artifact_url="default/job-results-prior#results/attempt-1/adapter") + ) + monkeypatch.setattr(task_main, "client_from_platform", lambda _sdk, _cls: jobs_client) + pretrained_model_tmp = SimpleNamespace( path=adapter_dir, cleanup_tmp_dir=MagicMock(), diff --git a/plugins/nemo-safe-synthesizer/tests/unit/test_sdk.py b/plugins/nemo-safe-synthesizer/tests/unit/test_sdk.py index 81e073072b..08267a0b1e 100644 --- a/plugins/nemo-safe-synthesizer/tests/unit/test_sdk.py +++ b/plugins/nemo-safe-synthesizer/tests/unit/test_sdk.py @@ -5,18 +5,50 @@ import json from datetime import datetime, timezone -from io import BytesIO from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pandas as pd import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.errors import NemoTransportError from nemo_platform_plugin.discovery import discover, discover_entry_points from nemo_safe_synthesizer_plugin.sdk.job import SafeSynthesizerJob from nemo_safe_synthesizer_plugin.sdk.job_builder import SafeSynthesizerJobBuilder -from nemo_safe_synthesizer_plugin.sdk.resources import AsyncSafeSynthesizerJobsResource, SafeSynthesizerResource +from nemo_safe_synthesizer_plugin.sdk.resources import ( + AsyncSafeSynthesizerJobsResource, + SafeSynthesizerJobsResource, + SafeSynthesizerResource, +) + + +def _resp(data): + """Wrap a payload in a NemoResponse-like object whose ``.data()`` returns it. + + Production now consumes typed-client responses via ``client.(...).data()``, + so mocked jobs-client methods return an object with a ``.data()`` accessor. + """ + m = MagicMock() + m.data.return_value = data + return m + + +def _binary_resp(data: bytes): + """Wrap bytes in a binary NemoResponse-like object whose ``.read()`` returns them.""" + m = MagicMock() + m.read.return_value = data + return m + + +def _paginated_resp(items, *, total: int, next_page: str | None, prev_page: str | None = None): + response = MagicMock() + response.page.return_value = SimpleNamespace( + items=items, + metadata={"total": total, "next_page": next_page, "prev_page": prev_page}, + ) + response.items.return_value = iter(items) + return response def _mock_platform(requests: list[httpx.Request]) -> NeMoPlatform: @@ -84,15 +116,50 @@ def handler(_request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio -async def test_async_safe_synthesizer_resource_get_logs_awaits_platform_jobs() -> None: +async def test_async_safe_synthesizer_resource_get_logs_forwards_query_params() -> None: platform = MagicMock() - platform.jobs.get_logs = AsyncMock(return_value=SimpleNamespace(data=[])) resource = AsyncSafeSynthesizerJobsResource(platform) - response = await resource.get_logs("safe-synth-job", workspace="default", limit=10) + mock_jobs = MagicMock() + mock_jobs.list_job_logs = AsyncMock(return_value=_paginated_resp([], total=0, next_page=None)) + with patch("nemo_safe_synthesizer_plugin.sdk.resources.client_from_platform", return_value=mock_jobs): + response = await resource.get_logs( + "safe-synth-job", + workspace="default", + limit=10, + page_cursor="next-page", + step_id=None, + ) + + assert response.data == [] + mock_jobs.list_job_logs.assert_awaited_once_with( + name="safe-synth-job", + workspace="default", + query_params={"limit": 10, "page_cursor": "next-page"}, + ) + + +def test_safe_synthesizer_resource_get_logs_forwards_query_params() -> None: + platform = MagicMock() + resource = SafeSynthesizerJobsResource(platform) + mock_jobs = MagicMock() + mock_jobs.list_job_logs.return_value = _paginated_resp([], total=0, next_page=None) + + with patch("nemo_safe_synthesizer_plugin.sdk.resources.client_from_platform", return_value=mock_jobs): + response = resource.get_logs( + "safe-synth-job", + workspace="default", + attempt_id=2, + step_id="step-1", + task_id=None, + ) assert response.data == [] - platform.jobs.get_logs.assert_awaited_once_with("safe-synth-job", workspace="default", limit=10) + mock_jobs.list_job_logs.assert_called_once_with( + name="safe-synth-job", + workspace="default", + query_params={"attempt_id": 2, "step_id": "step-1"}, + ) def test_job_builder_uploads_dataframe_and_submits_spec() -> None: @@ -147,37 +214,70 @@ def test_job_builder_submits_pretrained_model_job_for_adapter_reuse() -> None: assert "pretrained_model" not in create_kwargs["spec"]["config"]["training"] +def _make_job(mock_jobs: MagicMock, name: str = "safe-synth-job", workspace: str = "default") -> SafeSynthesizerJob: + """Build a SafeSynthesizerJob whose typed jobs client is *mock_jobs*. + + ``SafeSynthesizerJob.__init__`` resolves ``self._jobs = client_from_platform(client, JobsClient)``, + so we patch that lookup in the job module during construction. + """ + with patch("nemo_safe_synthesizer_plugin.sdk.job.client_from_platform", return_value=mock_jobs): + return SafeSynthesizerJob(name, MagicMock(), workspace=workspace) + + @pytest.mark.parametrize("status", ["error", "cancelled"]) def test_safe_synthesizer_job_wait_for_completion_raises_on_terminal_failure(status: str) -> None: - client = MagicMock() - client.jobs.get_status.return_value = SimpleNamespace( - status=status, - status_details={"reason": "failed"}, - error_details={"message": "boom"}, + mock_jobs = MagicMock() + mock_jobs.get_job_status.return_value = _resp( + SimpleNamespace( + status=status, + status_details={"reason": "failed"}, + error_details={"message": "boom"}, + ) ) - job = SafeSynthesizerJob("safe-synth-job", client, workspace="default") + job = _make_job(mock_jobs) with pytest.raises(RuntimeError, match=f"ended with status '{status}'"): job.wait_for_completion(poll_interval=0, verbose=False) - client.jobs.get_status.assert_called_once_with("safe-synth-job", workspace="default") + mock_jobs.get_job_status.assert_called_once_with(name="safe-synth-job", workspace="default") + + +def test_safe_synthesizer_job_wait_ignores_typed_client_log_failures() -> None: + mock_jobs = MagicMock() + mock_jobs.get_job_status.side_effect = [ + _resp(SimpleNamespace(status="active", status_details={}, error_details={})), + _resp(SimpleNamespace(status="completed", status_details={}, error_details={})), + ] + job = _make_job(mock_jobs) + request = httpx.Request("GET", "http://test/apis/jobs/v2/workspaces/default/jobs/safe-synth-job/logs") + + with patch.object( + job, + "_fetch_logs_incremental", + side_effect=NemoTransportError(httpx.ConnectError("Connection refused", request=request)), + ): + job.wait_for_completion(poll_interval=0, verbose=True) + + assert mock_jobs.get_job_status.call_count == 2 def test_safe_synthesizer_job_fetch_data_reads_synthetic_csv() -> None: - client = MagicMock() - client.jobs.results.download.return_value = BytesIO(b"name,value\nalice,1\nbob,2\n") - job = SafeSynthesizerJob("safe-synth-job", client, workspace="default") + mock_jobs = MagicMock() + mock_jobs.download_job_result.return_value = _binary_resp(b"name,value\nalice,1\nbob,2\n") + job = _make_job(mock_jobs) result = job.fetch_data() - client.jobs.results.download.assert_called_once_with("synthetic-data", job="safe-synth-job", workspace="default") + mock_jobs.download_job_result.assert_called_once_with( + name="synthetic-data", job="safe-synth-job", workspace="default" + ) assert list(result.columns) == ["name", "value"] assert result["value"].tolist() == [1, 2] def test_safe_synthesizer_job_fetch_summary_parses_json() -> None: - client = MagicMock() - client.jobs.results.download.return_value = BytesIO( + mock_jobs = MagicMock() + mock_jobs.download_job_result.return_value = _binary_resp( json.dumps( { "synthetic_data_quality_score": 8.5, @@ -188,20 +288,20 @@ def test_safe_synthesizer_job_fetch_summary_parses_json() -> None: } ).encode() ) - job = SafeSynthesizerJob("safe-synth-job", client, workspace="default") + job = _make_job(mock_jobs) summary = job.fetch_summary() - client.jobs.results.download.assert_called_once_with("summary", job="safe-synth-job", workspace="default") + mock_jobs.download_job_result.assert_called_once_with(name="summary", job="safe-synth-job", workspace="default") assert summary.synthetic_data_quality_score == 8.5 assert summary.data_privacy_score == 9.0 assert summary.timing.total_time_sec == 12.5 def test_safe_synthesizer_job_fetch_logs_follows_pagination() -> None: - client = MagicMock() - log_client = MagicMock() - client.with_options.return_value = log_client + mock_jobs = MagicMock() + options_client = MagicMock() + mock_jobs.with_options.return_value = options_client first_log = SimpleNamespace( job="safe-synth-job", job_step="safe-synthesizer", @@ -216,14 +316,13 @@ def test_safe_synthesizer_job_fetch_logs_follows_pagination() -> None: message="second", timestamp=datetime(2026, 1, 1, 0, 0, 1, tzinfo=timezone.utc), ) - log_client.jobs.get_logs.side_effect = [ - SimpleNamespace(data=[first_log], next_page="cursor-2"), - SimpleNamespace(data=[second_log], next_page=None), - ] - job = SafeSynthesizerJob("safe-synth-job", client, workspace="default") + paginated_response = _paginated_resp([first_log, second_log], total=2, next_page=None) + options_client.list_job_logs.return_value = paginated_response + job = _make_job(mock_jobs) logs = list(job.fetch_logs(timeout=5.0)) assert [log.message for log in logs] == ["first", "second"] - assert client.with_options.call_count == 2 - assert log_client.jobs.get_logs.call_args_list[1].kwargs["page_cursor"] == "cursor-2" + mock_jobs.with_options.assert_called_once_with(timeout=5.0) + options_client.list_job_logs.assert_called_once_with(name="safe-synth-job", workspace="default") + paginated_response.items.assert_called_once_with() diff --git a/plugins/nemo-safe-synthesizer/tests/unit/test_task_upload_results.py b/plugins/nemo-safe-synthesizer/tests/unit/test_task_upload_results.py index ef19975b2c..7e9410f09b 100644 --- a/plugins/nemo-safe-synthesizer/tests/unit/test_task_upload_results.py +++ b/plugins/nemo-safe-synthesizer/tests/unit/test_task_upload_results.py @@ -11,6 +11,16 @@ from nemo_platform_plugin.jobs.constants import NEMO_JOB_ID_ENVVAR, NEMO_JOB_WORKSPACE_ENVVAR +def _resp(data): + """Wrap a payload in a NemoResponse-like object whose ``.data()`` returns it. + + Production consumes typed-client responses via ``client.(...).data()``. + """ + m = MagicMock() + m.data.return_value = data + return m + + def import_task_main_without_heavy_runtime(monkeypatch): pytest.importorskip("nemo_safe_synthesizer.config.job") library_builder = ModuleType("nemo_safe_synthesizer.sdk.library_builder") @@ -30,9 +40,15 @@ def test_upload_results_uploads_and_registers_adapter(tmp_path, monkeypatch): (adapter_path / "adapter_config.json").write_text("{}", encoding="utf-8") sdk = MagicMock() - sdk.jobs.retrieve.return_value = SimpleNamespace(attempt_id="attempt-123") monkeypatch.setattr(task_main, "get_platform_sdk", lambda: sdk) + # Production routes both the job lookup and result creation through + # client_from_platform(sdk, JobsClient): get_job(...).data() for attempt_id, + # then create_job_result(..., body=PlatformJobResultCreateRequest(...)). + jobs_client = MagicMock() + jobs_client.get_job.return_value = _resp(SimpleNamespace(attempt_id="attempt-123")) + monkeypatch.setattr(task_main, "client_from_platform", lambda _sdk, _cls: jobs_client) + file_manager = MagicMock() file_manager.upload.side_effect = lambda _local_path, remote_path: ( f"test-workspace/job-results-safe-synth-job#{remote_path}" @@ -56,10 +72,17 @@ def test_upload_results_uploads_and_registers_adapter(tmp_path, monkeypatch): ) file_manager.validate_storage.assert_called_once_with() file_manager.upload.assert_has_calls([call(adapter_path, "results/attempt-123/adapter")], any_order=True) - sdk.jobs.results.create.assert_any_call( - name="adapter", - job="safe-synth-job", - workspace="test-workspace", - artifact_url="test-workspace/job-results-safe-synth-job#results/attempt-123/adapter", - artifact_storage_type="fileset", - ) + + jobs_client.get_job.assert_called_once_with(name="safe-synth-job", workspace="test-workspace") + + # create_job_result is called once per result; assert the adapter call, checking the + # artifact fields on the PlatformJobResultCreateRequest body object (which replaced the + # flat artifact_url/artifact_storage_type kwargs of the old Stainless call). + adapter_calls = [c for c in jobs_client.create_job_result.call_args_list if c.kwargs.get("name") == "adapter"] + assert len(adapter_calls) == 1 + adapter_call = adapter_calls[0] + assert adapter_call.kwargs["job"] == "safe-synth-job" + assert adapter_call.kwargs["workspace"] == "test-workspace" + body = adapter_call.kwargs["body"] + assert body.artifact_url == "test-workspace/job-results-safe-synth-job#results/attempt-123/adapter" + assert body.artifact_storage_type == "fileset" diff --git a/pytest.ini b/pytest.ini index 1576c26d12..29033c341a 100644 --- a/pytest.ini +++ b/pytest.ini @@ -13,6 +13,7 @@ pythonpath = plugins/nemo-deployments/tests/unit plugins/nemo-deployments/tests/integration plugins/nemo-safe-synthesizer/src + services/core/jobs/tests/controllers # Test discovery paths - packages and stable services # Option A: Centralized Testing - All tests runnable from root diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/quickstart/cli.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/quickstart/cli.py index 456b074727..59174901e6 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/quickstart/cli.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/quickstart/cli.py @@ -967,6 +967,9 @@ def _run_job_diagnostic(port: int, registry: str, tag: str, *, admin_email: str import uuid from nemo_platform import NeMoPlatform + from nemo_platform_plugin.client.adapter import client_from_platform + from nemo_platform_plugin.jobs.client import JobsClient + from nemo_platform_plugin.jobs.types import CreatePlatformJobRequest # When auth is enabled, use an unsigned JWT for the admin principal. default_headers = None @@ -997,29 +1000,32 @@ def _run_job_diagnostic(port: int, registry: str, tag: str, *, admin_email: str job_name = f"diagnostic-{uuid.uuid4().hex[:8]}" console.print(f" • Creating diagnostic job: {job_name}") - job = client.jobs.create( - platform_spec={ - "steps": [ - { - "name": "diagnostic", - "executor": { - "provider": "cpu", - "container": { - "image": cpu_image, - "entrypoint": [ - "python", - "-c", - "import sys; print(f'Python {sys.version}'); print('Job system is working correctly!')", - ], + jobs_client = client_from_platform(client, JobsClient) + job = jobs_client.create_job( + body=CreatePlatformJobRequest( + platform_spec={ + "steps": [ + { + "name": "diagnostic", + "executor": { + "provider": "cpu", + "container": { + "image": cpu_image, + "entrypoint": [ + "python", + "-c", + "import sys; print(f'Python {sys.version}'); print('Job system is working correctly!')", + ], + }, }, - }, - } - ] - }, - source="quickstart-doctor", - spec={}, - name=job_name, - ) + } + ] + }, + source="quickstart-doctor", + spec={}, + name=job_name, + ) + ).data() console.print(" • Waiting for job to complete...") @@ -1028,10 +1034,10 @@ def _run_job_diagnostic(port: int, registry: str, tag: str, *, admin_email: str poll_interval = 2 elapsed = 0 status = "pending" - job_status = client.jobs.retrieve(job.name) + job_status = jobs_client.get_job(name=job.name).data() while elapsed < max_wait: - job_status = client.jobs.retrieve(job.name) + job_status = jobs_client.get_job(name=job.name).data() status = job_status.status if status in ("completed", "error", "cancelled"): @@ -1058,9 +1064,9 @@ def _run_job_diagnostic(port: int, registry: str, tag: str, *, admin_email: str # Fetch and display logs console.print("\n [bold]Job output:[/bold]") try: - logs = client.jobs.get_logs(job.name) + logs = jobs_client.list_job_logs(name=job.name) log_lines = [] - for log_entry in logs: + for log_entry in logs.items(): if hasattr(log_entry, "message"): log_lines.append(log_entry.message) @@ -1075,7 +1081,7 @@ def _run_job_diagnostic(port: int, registry: str, tag: str, *, admin_email: str # Clean up the job (only if successful) if status == "completed": try: - client.jobs.delete(job.name) + jobs_client.delete_job(name=job.name) except Exception: pass # Ignore cleanup errors else: diff --git a/services/automodel/tests/test_progress_reporter.py b/services/automodel/tests/test_progress_reporter.py index 90b20ec372..66fcc5dd73 100644 --- a/services/automodel/tests/test_progress_reporter.py +++ b/services/automodel/tests/test_progress_reporter.py @@ -4,9 +4,8 @@ from __future__ import annotations from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch -from nemo_platform import omit from nmp.common.jobs.schemas import PlatformJobStatus from nmp.customization_common.service.context import NMPJobContext from nmp.customization_common.tasks.file_io_progress_reporter import JobsServiceProgressReporter @@ -14,6 +13,7 @@ def test_progress_reporter_calls_sdk_create_or_update() -> None: sdk = MagicMock() + mock_jobs = MagicMock() ctx = NMPJobContext( workspace="ws-a", job_id="job-1", @@ -26,15 +26,23 @@ def test_progress_reporter_calls_sdk_create_or_update() -> None: config_path=Path("/tmp/job/config.json"), ) reporter = JobsServiceProgressReporter(sdk, ctx.workspace, ctx.job_id, ctx.step, ctx.normalized_task) - reporter.update_progress(PlatformJobStatus.ACTIVE, status_details={"phase": "training"}) - sdk.jobs.tasks.create_or_update.assert_called_once_with( - ctx.normalized_task, - workspace=ctx.workspace, - job=ctx.job_id, - step=ctx.step, - status=PlatformJobStatus.ACTIVE.value, - status_details={"phase": "training"}, - error_details=omit, - error_stack=omit, - ) + with patch( + "nmp.customization_common.tasks.file_io_progress_reporter.client_from_platform", + return_value=mock_jobs, + ): + reporter.update_progress(PlatformJobStatus.ACTIVE, status_details={"phase": "training"}) + + mock_jobs.update_job_step_task.assert_called_once() + call_kwargs = mock_jobs.update_job_step_task.call_args.kwargs + assert call_kwargs["name"] == ctx.normalized_task + assert call_kwargs["workspace"] == ctx.workspace + assert call_kwargs["job"] == ctx.job_id + assert call_kwargs["step"] == ctx.step + # status + status_details now travel on the PlatformJobTaskUpdate body; unset + # fields (error_details/error_stack) are omitted, leaving their model defaults. + body = call_kwargs["body"] + assert body.status == PlatformJobStatus.ACTIVE + assert body.status_details == {"phase": "training"} + assert body.error_details is None + assert body.error_stack is None diff --git a/services/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.py b/services/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.py index 565283d600..4cf5b42366 100644 --- a/services/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.py +++ b/services/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.py @@ -6,9 +6,10 @@ import threading from nemo_platform import AsyncNeMoPlatform -from nemo_platform.types import PlatformJobStatus from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.files.client import AsyncFilesClient +from nemo_platform_plugin.jobs.client import AsyncJobsClient +from nemo_platform_plugin.jobs.schemas import PlatformJobStatus from nmp.common.api.filter import ComparisonOperation, FilterOperator from nmp.common.controller.controller import Controller from nmp.common.observability import start_span_with_ctx @@ -21,7 +22,9 @@ meter = metrics.get_meter(__name__) logger = logging.getLogger(__name__) -_TERMINAL_JOB_STATUSES: frozenset[PlatformJobStatus] = frozenset({"completed", "error", "cancelled"}) +_TERMINAL_JOB_STATUSES: frozenset[PlatformJobStatus] = frozenset( + {PlatformJobStatus.COMPLETED, PlatformJobStatus.ERROR, PlatformJobStatus.CANCELLED} +) class WorkspaceCleanup(Controller): @@ -122,14 +125,14 @@ async def _async_step(self): async def _cleanup_jobs(self, workspace: Workspace) -> None: logger.info(f"Cleaning up jobs for workspace: {workspace.name}") try: - jobs_response = await self._nmp_sdk.jobs.list(workspace=workspace.name) - jobs = [job async for job in jobs_response] + jobs_client = client_from_platform(self._nmp_sdk, AsyncJobsClient) + jobs = [job async for job in (await jobs_client.list_jobs(workspace=workspace.name)).items()] for job in jobs: if job.status not in _TERMINAL_JOB_STATUSES: try: logger.info(f"Cancelling job: {job.name}") - await self._nmp_sdk.jobs.cancel( + await jobs_client.cancel_job( name=job.name, workspace=workspace.name, ) @@ -138,7 +141,7 @@ async def _cleanup_jobs(self, workspace: Workspace) -> None: try: logger.info(f"Deleting job: {job.name}") - await self._nmp_sdk.jobs.delete( + await jobs_client.delete_job( name=job.name, workspace=workspace.name, ) diff --git a/services/core/entities/tests/controllers/test_workspace_cleanup.py b/services/core/entities/tests/controllers/test_workspace_cleanup.py index e156ce6ec5..2e0f294a39 100644 --- a/services/core/entities/tests/controllers/test_workspace_cleanup.py +++ b/services/core/entities/tests/controllers/test_workspace_cleanup.py @@ -56,6 +56,21 @@ def _make_mock_files_client(filesets: list | None = None) -> AsyncMock: return mock_files +def _make_jobs_client(jobs: list | None = None) -> MagicMock: + """Build a mock typed AsyncJobsClient. + + Production routes jobs calls through ``client_from_platform(sdk, AsyncJobsClient)`` + and iterates ``(await jobs_client.list_jobs(...)).items()``. So ``list_jobs`` is an + ``AsyncMock`` returning a paginated response whose ``.items()`` yields an async + iterator over the jobs. + """ + jobs_client = MagicMock() + jobs_client.list_jobs = AsyncMock(return_value=_MockAsyncPaginatedResponse(jobs or [])) + jobs_client.cancel_job = AsyncMock() + jobs_client.delete_job = AsyncMock() + return jobs_client + + def _make_sdk( jobs: list | None = None, deployments: list | None = None, @@ -64,17 +79,45 @@ def _make_sdk( """Build a MagicMock SDK with async mocks wired to the correct paths. Returns (sdk, mock_files_client) so tests can assert on files client calls. + Only deployments/filesets stay on the ``sdk.*`` / files-client accessors; jobs + are handled via the typed jobs client patched onto ``client_from_platform`` + (see ``_patch_jobs_client``). """ sdk = MagicMock() - sdk.jobs.list = AsyncMock(return_value=_AsyncIterator(jobs or [])) - sdk.jobs.cancel = AsyncMock() - sdk.jobs.delete = AsyncMock() sdk.inference.deployments.list = AsyncMock(return_value=_AsyncIterator(deployments or [])) sdk.inference.deployments.delete = AsyncMock() mock_files = _make_mock_files_client(filesets) return sdk, mock_files +_CLIENT_FROM_PLATFORM_PATCH = "nmp.core.entities.controllers.workspace_cleanup.client_from_platform" + + +def _patch_jobs_client(jobs_client: MagicMock): + """Patch ``client_from_platform`` in the workspace_cleanup module to return *jobs_client*.""" + return patch(_CLIENT_FROM_PLATFORM_PATCH, return_value=jobs_client) + + +def _patch_clients(jobs_client: MagicMock, files_client: MagicMock): + """Patch ``client_from_platform`` to dispatch by requested client class. + + ``_async_step`` cleans up both jobs and filesets, so it calls + ``client_from_platform(sdk, AsyncJobsClient)`` and + ``client_from_platform(sdk, AsyncFilesClient)`` — return the matching mock. + """ + from nemo_platform_plugin.files.client import AsyncFilesClient + from nemo_platform_plugin.jobs.client import AsyncJobsClient + + def _dispatch(_sdk, client_cls): + if client_cls is AsyncFilesClient: + return files_client + if client_cls is AsyncJobsClient: + return jobs_client + raise AssertionError(f"unexpected client class: {client_cls!r}") + + return patch(_CLIENT_FROM_PLATFORM_PATCH, side_effect=_dispatch) + + def _make_job(name: str, status: str = "completed") -> MagicMock: job = MagicMock() job.name = name @@ -170,13 +213,12 @@ async def test_successful_workspace_deletion(self): repo.mark_workspace_for_deletion.return_value = True sdk = MagicMock() - sdk.jobs.list = AsyncMock(return_value=_AsyncIterator([])) sdk.inference.deployments.list = AsyncMock(return_value=_AsyncIterator([])) mock_files = _make_mock_files_client([]) controller = _make_controller(workspace_repo=repo, nmp_sdk=sdk) - with patch(_FILES_CLIENT_PATCH, return_value=mock_files): + with _patch_clients(_make_jobs_client([]), mock_files): await controller._async_step() repo.mark_workspace_for_deletion.assert_any_call( @@ -205,12 +247,13 @@ async def test_cleanup_failure_marks_workspace_failed(self): repo.list_workspaces.return_value = ([workspace], None) repo.mark_workspace_for_deletion.return_value = True - sdk = MagicMock() - sdk.jobs.list = AsyncMock(side_effect=Exception("jobs service down")) + jobs_client = MagicMock() + jobs_client.list_jobs = AsyncMock(side_effect=Exception("jobs service down")) - controller = _make_controller(workspace_repo=repo, nmp_sdk=sdk) + controller = _make_controller(workspace_repo=repo) - await controller._async_step() + with _patch_jobs_client(jobs_client): + await controller._async_step() repo.mark_workspace_for_deletion.assert_any_call( name="test-workspace", @@ -225,12 +268,12 @@ async def test_cleanup_failure_increments_error_counter(self): repo.list_workspaces.return_value = ([workspace], None) repo.mark_workspace_for_deletion.return_value = True - sdk = MagicMock() - sdk.jobs.list = AsyncMock(side_effect=Exception("boom")) + jobs_client = MagicMock() + jobs_client.list_jobs = AsyncMock(side_effect=Exception("boom")) - controller = _make_controller(workspace_repo=repo, nmp_sdk=sdk) + controller = _make_controller(workspace_repo=repo) - with patch.object(controller._cleanup_errors, "add") as mock_add: + with _patch_jobs_client(jobs_client), patch.object(controller._cleanup_errors, "add") as mock_add: await controller._async_step() mock_add.assert_called_once_with(1, attributes={"error_type": "cleanup_failed"}) @@ -243,19 +286,17 @@ async def test_cancels_running_jobs_before_deleting(self): running_job.name = "running-job" running_job.status = "active" - sdk = MagicMock() - sdk.jobs.list = AsyncMock(return_value=_AsyncIterator([running_job])) - sdk.jobs.cancel = AsyncMock() - sdk.jobs.delete = AsyncMock() + jobs_client = _make_jobs_client([running_job]) - controller = _make_controller(nmp_sdk=sdk) - await controller._cleanup_jobs(workspace) + controller = _make_controller() + with _patch_jobs_client(jobs_client): + await controller._cleanup_jobs(workspace) - sdk.jobs.cancel.assert_awaited_once_with( + jobs_client.cancel_job.assert_awaited_once_with( name="running-job", workspace="test-workspace", ) - sdk.jobs.delete.assert_awaited_once_with( + jobs_client.delete_job.assert_awaited_once_with( name="running-job", workspace="test-workspace", ) @@ -267,16 +308,14 @@ async def test_deletes_completed_jobs_without_cancelling(self): completed_job.name = "completed-job" completed_job.status = "completed" - sdk = MagicMock() - sdk.jobs.list = AsyncMock(return_value=_AsyncIterator([completed_job])) - sdk.jobs.cancel = AsyncMock() - sdk.jobs.delete = AsyncMock() + jobs_client = _make_jobs_client([completed_job]) - controller = _make_controller(nmp_sdk=sdk) - await controller._cleanup_jobs(workspace) + controller = _make_controller() + with _patch_jobs_client(jobs_client): + await controller._cleanup_jobs(workspace) - sdk.jobs.cancel.assert_not_awaited() - sdk.jobs.delete.assert_awaited_once() + jobs_client.cancel_job.assert_not_awaited() + jobs_client.delete_job.assert_awaited_once() @pytest.mark.asyncio async def test_continues_on_individual_job_failure(self): @@ -288,25 +327,26 @@ async def test_continues_on_individual_job_failure(self): job2.name = "ok-job" job2.status = "completed" - sdk = MagicMock() - sdk.jobs.list = AsyncMock(return_value=_AsyncIterator([job1, job2])) - sdk.jobs.delete = AsyncMock(side_effect=[Exception("fail"), None]) + jobs_client = _make_jobs_client([job1, job2]) + jobs_client.delete_job = AsyncMock(side_effect=[Exception("fail"), None]) - controller = _make_controller(nmp_sdk=sdk) - await controller._cleanup_jobs(workspace) + controller = _make_controller() + with _patch_jobs_client(jobs_client): + await controller._cleanup_jobs(workspace) - assert sdk.jobs.delete.await_count == 2 + assert jobs_client.delete_job.await_count == 2 @pytest.mark.asyncio async def test_raises_on_list_failure(self): workspace = _make_workspace() - sdk = MagicMock() - sdk.jobs.list = AsyncMock(side_effect=Exception("unavailable")) + jobs_client = MagicMock() + jobs_client.list_jobs = AsyncMock(side_effect=Exception("unavailable")) - controller = _make_controller(nmp_sdk=sdk) + controller = _make_controller() with pytest.raises(Exception, match="unavailable"): - await controller._cleanup_jobs(workspace) + with _patch_jobs_client(jobs_client): + await controller._cleanup_jobs(workspace) class TestWorkspaceCleanupDeployments: @@ -388,24 +428,26 @@ class TestJobCancellationBranches: @pytest.mark.asyncio async def test_cancels_pending_jobs(self): workspace = _make_workspace() - sdk, _ = _make_sdk(jobs=[_make_job("pending-job", status="pending")]) + jobs_client = _make_jobs_client([_make_job("pending-job", status="pending")]) - controller = _make_controller(nmp_sdk=sdk) - await controller._cleanup_jobs(workspace) + controller = _make_controller() + with _patch_jobs_client(jobs_client): + await controller._cleanup_jobs(workspace) - sdk.jobs.cancel.assert_awaited_once_with(name="pending-job", workspace="test-workspace") - sdk.jobs.delete.assert_awaited_once_with(name="pending-job", workspace="test-workspace") + jobs_client.cancel_job.assert_awaited_once_with(name="pending-job", workspace="test-workspace") + jobs_client.delete_job.assert_awaited_once_with(name="pending-job", workspace="test-workspace") @pytest.mark.asyncio async def test_cancels_created_jobs(self): workspace = _make_workspace() - sdk, _ = _make_sdk(jobs=[_make_job("created-job", status="created")]) + jobs_client = _make_jobs_client([_make_job("created-job", status="created")]) - controller = _make_controller(nmp_sdk=sdk) - await controller._cleanup_jobs(workspace) + controller = _make_controller() + with _patch_jobs_client(jobs_client): + await controller._cleanup_jobs(workspace) - sdk.jobs.cancel.assert_awaited_once() - sdk.jobs.delete.assert_awaited_once() + jobs_client.cancel_job.assert_awaited_once() + jobs_client.delete_job.assert_awaited_once() @pytest.mark.asyncio async def test_does_not_cancel_terminal_jobs(self): @@ -415,26 +457,28 @@ async def test_does_not_cancel_terminal_jobs(self): _make_job("failed", status="error"), _make_job("stopped", status="cancelled"), ] - sdk, _ = _make_sdk(jobs=jobs) + jobs_client = _make_jobs_client(jobs) - controller = _make_controller(nmp_sdk=sdk) - await controller._cleanup_jobs(workspace) + controller = _make_controller() + with _patch_jobs_client(jobs_client): + await controller._cleanup_jobs(workspace) - sdk.jobs.cancel.assert_not_awaited() - assert sdk.jobs.delete.await_count == 3 + jobs_client.cancel_job.assert_not_awaited() + assert jobs_client.delete_job.await_count == 3 @pytest.mark.asyncio async def test_cancel_failure_still_deletes(self): """Regression: cancel() throwing must not prevent delete().""" workspace = _make_workspace() - sdk, _ = _make_sdk(jobs=[_make_job("flaky-job", status="active")]) - sdk.jobs.cancel = AsyncMock(side_effect=Exception("cancel failed")) + jobs_client = _make_jobs_client([_make_job("flaky-job", status="active")]) + jobs_client.cancel_job = AsyncMock(side_effect=Exception("cancel failed")) - controller = _make_controller(nmp_sdk=sdk) - await controller._cleanup_jobs(workspace) + controller = _make_controller() + with _patch_jobs_client(jobs_client): + await controller._cleanup_jobs(workspace) - sdk.jobs.cancel.assert_awaited_once() - sdk.jobs.delete.assert_awaited_once_with(name="flaky-job", workspace="test-workspace") + jobs_client.cancel_job.assert_awaited_once() + jobs_client.delete_job.assert_awaited_once_with(name="flaky-job", workspace="test-workspace") @pytest.mark.asyncio async def test_mixed_statuses(self): @@ -444,11 +488,12 @@ async def test_mixed_statuses(self): _make_job("done-job", status="completed"), _make_job("pending-job", status="pending"), ] - sdk, _ = _make_sdk(jobs=jobs) + jobs_client = _make_jobs_client(jobs) - controller = _make_controller(nmp_sdk=sdk) - await controller._cleanup_jobs(workspace) + controller = _make_controller() + with _patch_jobs_client(jobs_client): + await controller._cleanup_jobs(workspace) - cancel_calls = [c.kwargs["name"] for c in sdk.jobs.cancel.call_args_list] + cancel_calls = [c.kwargs["name"] for c in jobs_client.cancel_job.call_args_list] assert set(cancel_calls) == {"active-job", "pending-job"} - assert sdk.jobs.delete.await_count == 3 + assert jobs_client.delete_job.await_count == 3 diff --git a/services/core/files/tests/integration/test_files_basic.py b/services/core/files/tests/integration/test_files_basic.py index 267361b223..be7b3189ba 100644 --- a/services/core/files/tests/integration/test_files_basic.py +++ b/services/core/files/tests/integration/test_files_basic.py @@ -147,8 +147,8 @@ def test_fileset_list_pagination(self, sdk: NeMoPlatform): ) page1 = resp1.page() assert len(page1.items) == 2 - assert page1.page == 1 - assert page1.page_size == 2 + assert page1.metadata["page"] == 1 + assert page1.metadata["page_size"] == 2 # Test second page resp2 = files.list_filesets( @@ -157,7 +157,7 @@ def test_fileset_list_pagination(self, sdk: NeMoPlatform): ) page2 = resp2.page() assert len(page2.items) == 2 - assert page2.page == 2 + assert page2.metadata["page"] == 2 # Verify pages have different data page1_ids = {fs.id for fs in page1.items} diff --git a/services/core/jobs/src/nmp/core/jobs/api/v2/jobs/schemas.py b/services/core/jobs/src/nmp/core/jobs/api/v2/jobs/schemas.py index 9fd0070285..31ca0b1978 100644 --- a/services/core/jobs/src/nmp/core/jobs/api/v2/jobs/schemas.py +++ b/services/core/jobs/src/nmp/core/jobs/api/v2/jobs/schemas.py @@ -3,16 +3,23 @@ """Schemas for the v2 Jobs Service. -This module contains: -- View models (PlatformJob, PlatformJobStepWithContext) -- Request/response schemas -- Filter schemas +Most request/response/filter/sort types now live in +:mod:`nemo_platform_plugin.jobs.types` so that both the server and the typed +HTTP client (``JobsClient``) share one source of truth; this module +re-exports them. + +Two list-response wrappers (``PlatformJobListResultResponse``, +``PlatformJobListTaskResponse``) remain defined here because they wrap raw +entity instances server-side; the plugin exposes DTO equivalents for clients. """ -from datetime import datetime, timezone -from enum import Enum -from typing import Any, Dict, List, Optional +from typing import List, Optional +# Re-exported shared types (single source of truth in the plugin). +# NB: ``AuthContext`` is intentionally NOT re-exported from the plugin here — +# this module exposes the behaviour-carrying ``nmp.common.auth.AuthContext`` +# (imported above), which the ``PlatformJobStepWithContext`` subclass uses. +from nemo_platform_plugin.jobs import types as _types from nmp.common.auth import AuthContext from nmp.common.entities import ( DatetimeFilter, @@ -21,9 +28,19 @@ Value, get_random_id, ) -from nmp.common.jobs.schemas import PlatformJobResultResponse, PlatformJobStatus -from nmp.core.jobs.entities import PlatformJobSpec, PlatformJobStepSpec, PlatformJobTask -from pydantic import BaseModel, Field +from nmp.common.jobs.schemas import PlatformJobResultResponse +from nmp.common.jobs.schemas import PlatformJobStatus as PlatformJobStatus +from nmp.core.jobs.entities import PlatformJobTask +from pydantic import Field + +CreatePlatformJobRequest = _types.CreatePlatformJobRequest +PlatformJobAttemptSortField = _types.PlatformJobAttemptSortField +PlatformJobLogSortField = _types.PlatformJobLogSortField +PlatformJobResponse = _types.PlatformJobResponse +PlatformJobSortField = _types.PlatformJobSortField +PlatformJobStatusDetailsUpdateRequest = _types.PlatformJobStatusDetailsUpdateRequest +PlatformJobStatusUpdateRequest = _types.PlatformJobStatusUpdateRequest +PlatformJobTaskUpdate = _types.PlatformJobTaskUpdate # ============================================================================= # Utilities @@ -41,137 +58,7 @@ def get_model_id(prefix: str) -> str: # ============================================================================= -# Sort Fields -# ============================================================================= - - -class PlatformJobLogSortField(str, Enum): - TIMESTAMP_ASC = "timestamp" - TIMESTAMP_DESC = "-timestamp" - - def get_field_name(self) -> str: - return self.value.lstrip("-") - - def get_sort_direction(self) -> str: - return "desc" if self.value.startswith("-") else "asc" - - -class PlatformJobSortField(str, Enum): - CREATED_AT_ASC = "created_at" - CREATED_AT_DESC = "-created_at" - UPDATED_AT_ASC = "updated_at" - UPDATED_AT_DESC = "-updated_at" - - def get_field_name(self) -> str: - return self.value.lstrip("-") - - def get_sort_direction(self) -> str: - return "desc" if self.value.startswith("-") else "asc" - - -class PlatformJobAttemptSortField(str, Enum): - SEQ_ASC = "seq" - SEQ_DESC = "-seq" - - def get_field_name(self) -> str: - return self.value.lstrip("-") - - def get_sort_direction(self) -> str: - return "desc" if self.value.startswith("-") else "asc" - - -# ============================================================================= -# View Models (composite models for API responses) -# ============================================================================= - -# Import entities here to avoid circular imports at module level -# These are used in view models and will be imported when needed - - -class PlatformJobResponse(BaseModel): - """Response model for a platform job.""" - - id: str - attempt_id: str - name: str - workspace: str = Field(..., description="Workspace identifier") - project: Optional[str] = Field(default=None, description="Project URN") - description: str | None = None - source: str - spec: Dict[str, Any] = Field(default_factory=dict, description="Job Spec") - platform_spec: PlatformJobSpec - fileset: str = Field(..., description="Fileset ID for storing job artifacts") - status: PlatformJobStatus - status_details: Dict[str, Any] = Field(default_factory=dict, description="Details about the job status") - error_details: Dict[str, Any] | None = None - created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) - updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) - ownership: Optional[Dict[str, Any]] = None - custom_fields: Optional[Dict[str, Any]] = Field(default=None, description="Custom Fields") - - -class PlatformJobStepWithContext(BaseModel): - """Step with additional context from parent job/attempt.""" - - id: str - job: str - attempt_id: str - fileset: str - workspace: str - name: str - step_spec: PlatformJobStepSpec | None = None - status: PlatformJobStatus = PlatformJobStatus.CREATED - status_details: Dict[str, Any] | None = None - error_details: Dict[str, Any] | None = None - auth_context: Optional[AuthContext] = Field(default=None, description="Auth context for task execution") - created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) - updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) - - -# ============================================================================= -# Request Schemas -# ============================================================================= - - -class CreatePlatformJobRequest(BaseModel): - """Request model for creating a new platform job.""" - - name: Optional[str] = None - description: Optional[str] = None - project: Optional[str] = None - spec: dict - platform_spec: PlatformJobSpec - source: str - ownership: Optional[dict] = None - custom_fields: Optional[dict] = None - - -class PlatformJobTaskUpdate(BaseModel): - """Request model for updating a platform job task.""" - - status: PlatformJobStatus = PlatformJobStatus.PENDING - status_details: Dict[str, Any] | None = None - error_details: Dict[str, Any] | None = None - error_stack: str | None = None - - -class PlatformJobStatusUpdateRequest(BaseModel): - """Request model for updating job status.""" - - status: PlatformJobStatus = Field(..., description="The new status to set for the job.") - status_details: Dict[str, Any] | None = Field( - default_factory=dict, description="Optional status details related to the status update." - ) - error_details: Dict[str, Any] | None = Field( - default_factory=dict, description="Optional error details related to the status update." - ) - - -PlatformJobStatusDetailsUpdateRequest = Dict[str, Any] - - -# ============================================================================= -# Response Schemas +# Response Schemas (server-side — wrap raw entity instances) # ============================================================================= @@ -187,8 +74,18 @@ class PlatformJobListTaskResponse(Value): data: List[PlatformJobTask] +class PlatformJobStepWithContext(_types.PlatformJobStepWithContext): + """Step with additional context from parent job/attempt.""" + + # Overrides ``auth_context`` with the behaviour-carrying + # ``nmp.common.auth.AuthContext`` (``to_principal`` / ``from_principal``); + # the plugin base uses the data-only mirror for the wire shape. + auth_context: Optional[AuthContext] = Field(default=None, description="Auth context for task execution") + + # ============================================================================= -# Filter Schemas +# Filter Schemas (server-side — subclass the entity-store ``Filter`` for +# field-mapping / translation support; not part of the client wire contract) # ============================================================================= diff --git a/services/core/jobs/src/nmp/core/jobs/app/providers.py b/services/core/jobs/src/nmp/core/jobs/app/providers.py index c57f6c08bd..f70adceb25 100644 --- a/services/core/jobs/src/nmp/core/jobs/app/providers.py +++ b/services/core/jobs/src/nmp/core/jobs/app/providers.py @@ -1,189 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import re -from typing import Annotated, Literal, Union +"""Execution provider types for the Jobs service. -from pydantic import BaseModel, Field, field_validator, model_validator - -# SHM: megabyte/gigabyte scale only — Mi, Gi (binary) or M, G (decimal SI). -# Ki / Ti / Pi / Ei and other suffixes are not accepted for /dev/shm. -_SHM_QUANTITY_RE = re.compile(r"^([+-]?(?:\d+|\d*\.\d+)(?:[eE][+-]?\d+)?)(Mi|Gi|M|G)$") - - -class ContainerSpec(BaseModel): - """ - Specification for a container configuration. - - Defines the container image and related configuration for job execution. - """ - - image: str | None = Field(default=None, min_length=1) - """The container image to use for execution. When omitted, resolved from the execution profile's default_task_image or the platform CPU tasks image.""" - - entrypoint: list[str] = Field(default_factory=list) - """The entrypoint for the container as a list of strings (e.g., ['python', 'script.py']). This overrides a container's default entrypoint (e.g. ENTRYPOINT in Docker) if provided.""" - - command: list[str] = Field(default_factory=list) - """The command to execute as a list of strings (e.g., ['python', 'script.py']). This overrides a container's default commands (e.g. CMD in Docker) if provided.""" - - -class ComputeResourceSpec(BaseModel): - """Resource specification.""" - - cpu: str | None = Field(default=None, description="CPU specification (e.g., '250m', '1', '2.5').") - memory: str | None = Field(default=None, description="Memory specification (e.g., '128Mi', '1Gi', '512M').") - - -class ComputeResources(BaseModel): - """Resource requirements matching k8s ResourceRequirements format.""" - - requests: ComputeResourceSpec = Field( - default_factory=ComputeResourceSpec, description="Minimum resources requested for the container." - ) - - limits: ComputeResourceSpec = Field( - default_factory=ComputeResourceSpec, description="Maximum resources the container can use." - ) - - num_nodes: int = Field(default=1, ge=1, description="Number of nodes to use.") - - num_gpus: int | None = Field(default=None, description="Step requesting number of GPUs.") - - shm_size: str | None = Field( - default=None, - description="Shared memory (/dev/shm) size as a Kubernetes quantity (e.g. '1Gi', '4Gi'). " - "Used for GPU and distributed-GPU job executors. When unset, defaults to 1Gi per allocated GPU.", - ) - - @field_validator("shm_size") - @classmethod - def validate_shm_size_quantity(cls, v: str | None) -> str | None: - if v is None: - return None - s = v.strip() - if not s: - raise ValueError("shm_size cannot be empty or whitespace-only") - if not _SHM_QUANTITY_RE.fullmatch(s): - raise ValueError( - "shm_size must use a megabyte/gigabyte-scale suffix: Mi, Gi, M, or G (e.g. '1Gi', '512Mi', '2G')." - ) - return s - - -class TaskSpec(BaseModel): - """ - Specification for a task to be executed. - - Defines the command and arguments for a job task. - """ - - command: list[str] - """The command to execute as a list of strings (e.g., ['python', 'script.py']).""" - - args: list[str] | str - """Arguments to pass to the command. Can be a list of strings or a single string.""" - - -class CPUExecutionProvider(BaseModel): - """ - CPU-based execution provider. - - Provides configuration for running jobs on CPU resources with - resource requests and limits. - """ - - provider: Literal["cpu"] = "cpu" - """The provider type, always 'cpu' for CPU execution.""" - - profile: str = "default" - """The execution profile to use. Defaults to 'default'.""" - - container: ContainerSpec - """Container specification defining the execution environment.""" - - resources: ComputeResources = Field( - default_factory=ComputeResources, description="Resource requests and limits for CPU execution." - ) - - -class GPUExecutionProvider(BaseModel): - """ - GPU-based execution provider. - - Provides configuration for running jobs on GPU resources with - resource requests and limits. - """ - - provider: Literal["gpu"] = "gpu" - """The provider type, always 'gpu' for GPU execution.""" - - profile: str = "default" - """The execution profile to use. Defaults to 'default'.""" - - container: ContainerSpec - """Container specification defining the execution environment.""" - - resources: ComputeResources = Field( - default_factory=ComputeResources, description="Resource requests and limits for GPU execution." - ) - - -class DistributedGPUExecutionProvider(BaseModel): - """ - GPU-based execution provider. - - Provides configuration for running jobs on GPU resources with - resource requests and limits. - """ - - provider: Literal["gpu_distributed"] = "gpu_distributed" - """The provider type, always 'gpu_distributed' for distributed GPU execution.""" - - profile: str = "default" - """The execution profile to use. Defaults to 'default'.""" - - container: ContainerSpec - """Container specification defining the execution environment.""" - - resources: ComputeResources = Field( - default_factory=ComputeResources, description="Resource requests and limits for distributed GPU execution." - ) - - -class SubprocessExecutionProvider(BaseModel): - """Host subprocess execution provider.""" - - provider: Literal["subprocess"] = "subprocess" - """The provider type, always 'subprocess' for host subprocess execution.""" - - profile: str = "default" - """The execution profile to use. Defaults to 'default'.""" - - command: list[str] - """The host command to execute as a list of strings (e.g., ['python', '-m', 'my_task']).""" - - @model_validator(mode="after") - def validate_command(self) -> "SubprocessExecutionProvider": - if not self.command: - raise ValueError("subprocess execution requires command to be set") - return self - - -# Type alias for the current execution provider implementation -ExecutionProviderT = Union[ - CPUExecutionProvider, GPUExecutionProvider, DistributedGPUExecutionProvider, SubprocessExecutionProvider -] -"""Type alias representing the current execution provider type.""" - -# Discriminated union type for execution providers -Provider = Annotated[ - ExecutionProviderT, - Field(discriminator="provider"), -] +The definitions now live in :mod:`nemo_platform_plugin.jobs.providers` so that +both the server and the typed HTTP client (``JobsClient``) share one source of +truth. This module re-exports them for backward compatibility. """ -Discriminated union type for execution providers. -Uses the 'provider' field to determine the specific provider type. -Currently supports CPU execution providers, with extensibility for future provider types. -""" +from nemo_platform_plugin.jobs import providers as _providers + +ComputeResources = _providers.ComputeResources +ComputeResourceSpec = _providers.ComputeResourceSpec +ContainerSpec = _providers.ContainerSpec +CPUExecutionProvider = _providers.CPUExecutionProvider +DistributedGPUExecutionProvider = _providers.DistributedGPUExecutionProvider +ExecutionProviderT = _providers.ExecutionProviderT +GPUExecutionProvider = _providers.GPUExecutionProvider +Provider = _providers.Provider +SubprocessExecutionProvider = _providers.SubprocessExecutionProvider +TaskSpec = _providers.TaskSpec diff --git a/services/core/jobs/src/nmp/core/jobs/app/schemas.py b/services/core/jobs/src/nmp/core/jobs/app/schemas.py index d5e8b75992..9a8c4cf2ed 100644 --- a/services/core/jobs/src/nmp/core/jobs/app/schemas.py +++ b/services/core/jobs/src/nmp/core/jobs/app/schemas.py @@ -1,140 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Shared schemas for the Jobs service.""" - -from typing import Optional, Self - -from nmp.common.entities.constants import NAME_PATTERN, NAME_PATTERN_DESCRIPTION -from nmp.common.jobs.constants import PERSISTENT_JOB_STORAGE_PATH_ENVVAR -from nmp.core.jobs.app.providers import Provider -from pydantic import BaseModel, ConfigDict, Field, model_validator - -# ============================================================================= -# Domain Models (shared by entities, dispatcher, and API) -# ============================================================================= - - -class PlatformJobSecretEnvironmentVariableRef(BaseModel): - """Reference to a secret to populate an environment variable for a job step.""" - - name: str = Field(description="The name of the secret to reference") - - -class PlatformJobEnvironmentVariable(BaseModel): - """Environment variable for a job step""" - - name: str = Field(description="The environment variable name") - value: Optional[str] = Field(default=None, description="The environment variable value") - from_secret: Optional[PlatformJobSecretEnvironmentVariableRef] = Field( - default=None, description="Reference to a secret environment variable to populate the environment variable" - ) - - @model_validator(mode="after") - def validate_self(self) -> Self: - # Ensure one of value or from_secret is provided - if self.value is None and self.from_secret is None: - raise ValueError("Either value or from_secret must be provided for environment variables.") - - # Ensure only one of value or from_secret is provided - if self.value is not None and self.from_secret is not None: - raise ValueError("Only one of value or from_secret can be provided for environment variables.") - - return self - - -class StepLifecycle(BaseModel): - """Controller-level lifecycle configuration for a job step. - - These settings control how the jobs controller manages the step, - as opposed to ``config`` which is the task payload forwarded to - the container. - """ - - staleness_timeout_seconds: int = Field( - default=0, - description="If every active task in the step goes this many seconds without an update, the step is terminated. " - "A value of 0 disables staleness detection.", - ) - - -class PlatformJobStepSpec(BaseModel): - """Specification for a single step in a platform job.""" - - name: str = Field( - description=f"The name of the step. Must be unique for all steps in a job. {NAME_PATTERN_DESCRIPTION}", - pattern=NAME_PATTERN, - examples=["preprocess", "train-model", "eval-step-v1"], - ) - environment: Optional[list[PlatformJobEnvironmentVariable]] = Field( - default=None, description="Environment variables for the step" - ) - executor: Provider = Field(description="The executor for the step") - config: dict = Field(default_factory=dict, description="Configuration for the step") - lifecycle: StepLifecycle = Field( - default_factory=StepLifecycle, description="Lifecycle configuration settings for the step" - ) - - @property - def requires_persistent_storage(self) -> bool: - """ - Determine if the step requires persistent storage. - - This is determined by checking if the step has an environment variable - matching the value of PERSISTENT_JOB_STORAGE_PATH_ENVVAR. - """ - for envvar in self.environment or []: - if envvar.name == PERSISTENT_JOB_STORAGE_PATH_ENVVAR: - return True - return False - - model_config = ConfigDict(regex_engine="python-re") - - -class PlatformJobSpec(BaseModel): - """Specification for a platform job, containing steps and secrets.""" - - steps: list[PlatformJobStepSpec] = Field(description="List of steps to be executed in the job") - - @model_validator(mode="after") - def validate_steps(self) -> Self: - # Ensure there is at least one step. - if not self.steps: - raise ValueError("At least one step is required in the job specification.") - - # Ensure that each step has a unique name. - step_names = [step.name for step in self.steps] - if len(step_names) != len(set(step_names)): - raise ValueError("Each step must have a unique name.") - return self - - -# ============================================================================= -# Misc Schemas -# ============================================================================= - - -ProviderRef = str -ProfileRef = str -BackendRef = str - - -class BaseExecutionProfile(BaseModel): - """Execution configuration for a job.""" - - provider: ProviderRef = Field( - default="cpu", - description="The compute provider for the executor, e.g., cpu, gpu", - ) - profile: str = Field( - default="default", - description="The profile name for the executor, e.g., high_priority_a100, low_priority, etc.", - ) - - @property - def supports_persistent_storage(self) -> bool: - """Indicates if the execution profile supports persistent storage.""" - return False - - def __str__(self) -> str: - return f"{self.profile}:{self.provider}" +"""Shared schemas for the Jobs service. + +The job specification types now live in :mod:`nemo_platform_plugin.jobs.spec` +and the ``BaseExecutionProfile`` base in the same package, so that both the +server and the typed HTTP client (``JobsClient``) share one source of truth. +This module re-exports them for backward compatibility. +""" + +from nemo_platform_plugin.jobs import spec as _spec + +BackendRef = _spec.BackendRef +BaseExecutionProfile = _spec.BaseExecutionProfile +PlatformJobEnvironmentVariable = _spec.PlatformJobEnvironmentVariable +PlatformJobSecretEnvironmentVariableRef = _spec.PlatformJobSecretEnvironmentVariableRef +PlatformJobSpec = _spec.PlatformJobSpec +PlatformJobStepSpec = _spec.PlatformJobStepSpec +ProfileRef = _spec.ProfileRef +ProviderRef = _spec.ProviderRef +StepLifecycle = _spec.StepLifecycle diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/base.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/base.py index 0bf9088313..52adba9e0a 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/base.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/base.py @@ -9,33 +9,21 @@ from enum import Enum from typing import Generic, Optional, TypeVar -from nemo_platform import NeMoPlatform, NotFoundError -from nemo_platform.types import PlatformJobStatus -from nemo_platform.types.jobs import PlatformJobStep, PlatformJobStepWithContext -from nmp.common.auth.models import NMP_PRINCIPAL_ENVVAR +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError as ClientNotFoundError +from nemo_platform_plugin.jobs import execution_profiles as _execution_profiles +from nemo_platform_plugin.jobs.client import JobsClient +from nemo_platform_plugin.jobs.schemas import PlatformJobStatus +from nemo_platform_plugin.jobs.types import PlatformJobStepResponse, PlatformJobStepWithContext from nmp.common.config.base import ( LOOPBACK_ADDRESSES, - NMP_CONFIG_WARNINGS_DISABLED_ENV_VAR, PlatformConfig, determine_loopback_override, ) -from nmp.common.jobs.constants import ( - CONFIG_TASK_STORAGE_PATH_ENVVAR, - EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, - NEMO_JOB_ATTEMPT_ID_ENVVAR, - NEMO_JOB_FILESET_ENVVAR, - NEMO_JOB_ID_ENVVAR, - NEMO_JOB_SECRETS_ENVVAR, - NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, - NEMO_JOB_STEP_ENVVAR, - NEMO_JOB_TASK_ENVVAR, - NEMO_JOB_WORKSPACE_ENVVAR, - PERSISTENT_JOB_STORAGE_PATH_ENVVAR, - TASK_CONFIG_ENVVAR, -) from nmp.common.sdk_factory import get_entity_parts from nmp.core.jobs.app.providers import ComputeResources -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel logger = logging.getLogger(__name__) @@ -44,39 +32,12 @@ DEFAULT_PROFILE = "default" DEFAULT_PROVIDER = "cpu" +JobExecutionProfileConfig = _execution_profiles.JobExecutionProfileConfig +RESERVED_JOB_ENVIRONMENT_VARIABLE_NAMES = _execution_profiles.RESERVED_JOB_ENVIRONMENT_VARIABLE_NAMES -# Env var names set by the platform during job creation; user-provided profile environment must not conflict. -RESERVED_JOB_ENVIRONMENT_VARIABLE_NAMES: frozenset[str] = frozenset( - { - # From nmp.common.jobs.constants - CONFIG_TASK_STORAGE_PATH_ENVVAR, - EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, - NEMO_JOB_ATTEMPT_ID_ENVVAR, - NEMO_JOB_FILESET_ENVVAR, - NEMO_JOB_ID_ENVVAR, - NEMO_JOB_SECRETS_ENVVAR, - NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, - NEMO_JOB_STEP_ENVVAR, - NEMO_JOB_TASK_ENVVAR, - NEMO_JOB_WORKSPACE_ENVVAR, - PERSISTENT_JOB_STORAGE_PATH_ENVVAR, - TASK_CONFIG_ENVVAR, - # Auth - NMP_PRINCIPAL_ENVVAR, - # OTEL (telemetry) - "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT", - "OTEL_LOGS_EXPORTER", - "OTEL_SERVICE_NAME", - "OTEL_EXPORTER_OTLP_LOGS_HEADERS", - # Platform shared envvars (to_shared_envvars with NMP_ prefix) - NMP_CONFIG_WARNINGS_DISABLED_ENV_VAR, - "NMP_BASE_URL", - "NMP_JOBS_URL", - "NMP_FILES_URL", - "NMP_MODELS_URL", - "NMP_SECRETS_URL", - } -) +# The env-var-name reserved set and the base ``JobExecutionProfileConfig`` now +# live in the shared plugin leaf node (imported above) so that both the server +# and the typed HTTP client agree on the wire shape and validation. class JobUpdate(BaseModel): @@ -85,33 +46,6 @@ class JobUpdate(BaseModel): error_details: dict | None = None -class JobExecutionProfileConfig(BaseModel): - ttl_seconds_before_active: int = 30 * 60 # 30 minutes - ttl_seconds_active: int = 24 * 60 * 60 # 24 hours - ttl_seconds_after_finished: int = 60 * 60 # 1 hour - cleanup_completed_jobs_immediately: bool = True - launcher_tool_path: str = Field(default="/tools/jobs-launcher", description="Path to the jobs launcher tool") - default_task_image: str | None = Field( - default=None, - min_length=1, - description="Default container image for job task pods. Used when a job step omits container.image. " - "When unset, falls back to the platform CPU tasks image (platform.image_registry/nmp-cpu-tasks:platform.image_tag).", - ) - env: dict[str, str] = Field( - default_factory=dict, - description="Optional env vars applied to all jobs (e.g. HOME=/tmp). Keys must not conflict with platform-reserved names. Job steps may override these variables.", - ) - - @model_validator(mode="after") - def validate_env_no_reserved_names(self) -> JobExecutionProfileConfig: - conflicting = [k for k in self.env if k in RESERVED_JOB_ENVIRONMENT_VARIABLE_NAMES] - if conflicting: - raise ValueError( - f"Profile environment keys must not conflict with platform-reserved names: {sorted(conflicting)}" - ) - return self - - _DEFAULT_TASK_IMAGE_NAME = "nmp-cpu-tasks" @@ -150,6 +84,10 @@ def __init__( profile_name: str, ): self._nmp_sdk = nmp_sdk + # Typed Jobs client sharing the SDK's transport/headers. Built once; every + # call passes ``workspace=`` explicitly (including the cross-workspace "-"), + # so the client's default workspace is never relied upon. + self._jobs = client_from_platform(nmp_sdk, JobsClient) self._execution_profile_config = execution_profile_config self._profile_name = profile_name self.init() @@ -193,15 +131,15 @@ def get_secrets_environment_variable_for_injection(self, step: PlatformJobStepWi env_var_str += f"{envvar.name}={workspace}/{secret_name}" return env_var_str - def get_step(self, job: str, step_name: str, workspace: str) -> PlatformJobStep: - """Fetch the latest state of a job step from the NeMo Platform SDK.""" - return self._nmp_sdk.jobs.steps.retrieve(name=step_name, workspace=workspace, job=job) + def get_step(self, job: str, step_name: str, workspace: str) -> PlatformJobStepResponse: + """Fetch the latest state of a job step via the typed Jobs client.""" + return self._jobs.get_job_step(name=step_name, workspace=workspace, job=job).data() - def get_step_safe(self, job: str, step_name: str, workspace: str) -> Optional[PlatformJobStep]: - """Fetch the job step from the NeMo Platform SDK, or None if not found (e.g. 404, workspace deleted).""" + def get_step_safe(self, job: str, step_name: str, workspace: str) -> Optional[PlatformJobStepResponse]: + """Fetch the job step, or None if not found (e.g. 404, workspace deleted).""" try: return self.get_step(job=job, step_name=step_name, workspace=workspace) - except NotFoundError: + except ClientNotFoundError: return None except Exception as e: raise e @@ -215,7 +153,7 @@ def check_step_is_terminal(self, job: str, step_name: str, workspace: str) -> bo try: step = self.get_step(job=job, step_name=step_name, workspace=workspace) return step.status in ("cancelled", "error", "completed") - except NotFoundError: + except ClientNotFoundError: # If the job step entity is not found, we treat it as terminal so cleanup can proceed. return True except Exception as e: @@ -224,9 +162,9 @@ def check_step_is_terminal(self, job: str, step_name: str, workspace: str) -> bo def check_job_is_terminal(self, job: str, workspace: str) -> bool: """Check if a job is in a terminal state.""" try: - job_response = self._nmp_sdk.jobs.retrieve(name=job, workspace=workspace) + job_response = self._jobs.get_job(name=job, workspace=workspace).data() return job_response.status in ("cancelled", "error", "completed") - except NotFoundError: + except ClientNotFoundError: # If the job entity is not found, we treat it as terminal so cleanup can proceed. return True except Exception as e: @@ -294,11 +232,11 @@ def check_step_is_stale(self, step: PlatformJobStepWithContext) -> bool: return False try: - tasks = self._nmp_sdk.jobs.tasks.list( + tasks = self._jobs.list_job_step_tasks( name=step.name, job=step.job, workspace=step.workspace, - ) + ).data() except Exception: logger.warning("Failed to fetch tasks for staleness check", extra={"step": step.name, "job": step.job}) return False diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py index c4a0e1e343..6c536dd8bd 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py @@ -14,13 +14,32 @@ from abc import abstractmethod from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass -from typing import Any, Generic, Literal, TypeVar +from typing import Any, Generic, TypeVar import docker.types from docker.errors import APIError, ImageNotFound, NotFound from docker.models.containers import Container from docker.types import LogConfig, Mount -from nemo_platform.types.jobs import PlatformJobStepWithContext +from nemo_platform_plugin.jobs.execution_profiles import ( + DockerJobExecutionProfile as PluginDockerJobExecutionProfile, +) +from nemo_platform_plugin.jobs.execution_profiles import ( + DockerJobExecutionProfileConfig as PluginDockerJobExecutionProfileConfig, +) +from nemo_platform_plugin.jobs.execution_profiles import ( + DockerJobNetworkConfig as PluginDockerJobNetworkConfig, +) +from nemo_platform_plugin.jobs.execution_profiles import ( + DockerJobStorageConfig as DockerJobStorageConfig, +) +from nemo_platform_plugin.jobs.execution_profiles import ( + DockerVolumeMount as DockerVolumeMount, +) +from nemo_platform_plugin.jobs.types import ( + PlatformJobStatusUpdateRequest, + PlatformJobStepWithContext, + PlatformJobTaskUpdate, +) from nmp.common.auth import AuthContext from nmp.common.config import get_platform_config, nmp_user_data_dir from nmp.common.docker.gpu_pool import GPUAllocationError @@ -69,10 +88,8 @@ ExecutionProviderT, GPUExecutionProvider, ) -from nmp.core.jobs.app.schemas import BaseExecutionProfile from nmp.core.jobs.controllers.backends.base import ( JobBackend, - JobExecutionProfileConfig, JobUpdate, get_logs_endpoint_from_fileset, resolve_gpu_job_shm_size, @@ -86,7 +103,7 @@ SchedulingDeferred, ) from opentelemetry import trace -from pydantic import BaseModel, Field +from pydantic import Field import docker @@ -123,6 +140,10 @@ def k8s_shm_quantity_to_docker(quantity: str) -> str: ProviderT = TypeVar("ProviderT", bound=ExecutionProviderT) +# DockerVolumeMount and DockerJobStorageConfig are pure data shapes shared with +# the typed HTTP client — imported from the plugin leaf node (see imports). + + def _resolve_jobs_controller_instance_id() -> str: configured = os.getenv(NMP_JOBS_DOCKER_OWNER_ID_ENVVAR) if configured: @@ -139,66 +160,34 @@ class DockerTimestampParseResult: is_zero: bool -class DockerVolumeMount(BaseModel): - volume_name: str = Field(description="Name of the Docker volume to mount") - mount_path: str = Field(description="Path inside the container where the volume will be mounted") - kind: Literal["volume", "tmpfs"] = Field( - default="volume", - description="Type of the Docker volume to mount. Options are 'volume' or 'tmpfs' (default: 'volume'). tmpfs volumes are only supported on Linux hosts.", - ) - options: dict | None = Field(default=None, description="Additional options for the volume") - allow_create_volume: bool = Field( - default=False, description="Whether to allow the creation of the volume if it does not exist (default: false)." - ) - - -class DockerJobStorageConfig(BaseModel): - """Configuration for persistent storage in Docker jobs.""" - - volume_name: str = Field( - default="nemo-jobs-storage", description="Name of the Docker volume for persistent storage" - ) - volume_permissions_image: str = Field( - default=DEFAULT_VOLUME_PERMISSIONS_IMAGE, description="Docker image used to set permissions on the volume" - ) - additional_volume_mounts: list[DockerVolumeMount] = Field( - default_factory=list, - description="List of additional Docker volume mounts for the job", - ) - - -class DockerJobNetworkConfig(BaseModel): +# Server-side override: the default network name comes from the +# ``NEMO_JOBS_DEFAULT_DOCKER_NETWORK`` env var (used by quickstart and e2e). +# No docstring on purpose — a docstring would surface as the schema +# ``description``, and this type carries none on the wire. +class DockerJobNetworkConfig(PluginDockerJobNetworkConfig): job_container_network: str = Field( default=NEMO_JOBS_DEFAULT_DOCKER_NETWORK, description="Docker network for the job container" ) -class DockerJobExecutionProfileConfig(JobExecutionProfileConfig): +class DockerJobExecutionProfileConfig(PluginDockerJobExecutionProfileConfig): """Configuration for Docker Job execution profile.""" - storage: DockerJobStorageConfig = Field( - default_factory=DockerJobStorageConfig, description="Docker storage configuration" - ) + # ``networking`` re-typed to the server ``DockerJobNetworkConfig`` (env-var default). networking: DockerJobNetworkConfig = Field( default_factory=DockerJobNetworkConfig, description="Docker networking configuration" ) -class DockerJobExecutionProfile(BaseExecutionProfile): +class DockerJobExecutionProfile(PluginDockerJobExecutionProfile): """ Execution configuration for a Docker Job. This is used to define the executor type, provider, profile, and any additional configuration required for the executor to run the job on Docker """ - backend: Literal["docker"] = "docker" config: DockerJobExecutionProfileConfig = Field(description="Additional configuration for the docker executor") - @property - def supports_persistent_storage(self) -> bool: - """Indicates if the execution profile supports persistent storage.""" - return self.config.storage is not None - class DockerJobBackend(JobBackend[ProviderT, DockerJobExecutionProfileConfig], Generic[ProviderT]): BACKEND_NAME: str = "docker" @@ -836,12 +825,11 @@ def cancel_scheduling(self, step: PlatformJobStepWithContext) -> bool: status = PlatformJobStatus.CANCELLED status_details["message"] = "Job is cancelled, not creating container" logger.info("Job step is not scheduling container", extra={"status": updated_step.status}) - self._nmp_sdk.jobs.steps.update_status( - step.name, + self._jobs.update_job_step_status( + name=step.name, workspace=step.workspace, job=step.job, - status=status.value, - status_details=status_details, + body=PlatformJobStatusUpdateRequest(status=status, status_details=status_details), ) return is_cancelling_or_pausing @@ -879,15 +867,17 @@ def run_container( try: self._run_container_in_thread(step, container_args) except FailedToScheduleError as e: - status = PlatformJobStatus.ERROR - self._nmp_sdk.jobs.steps.update_status( - step.name, - workspace=step.workspace, - job=step.job, - status=status.value, - error_details=e.error_details, # type: ignore - ) logger.exception("Failed to schedule container for job step") + status = PlatformJobStatus.ERROR + try: + self._jobs.update_job_step_status( + name=step.name, + workspace=step.workspace, + job=step.job, + body=PlatformJobStatusUpdateRequest(status=status, error_details=e.error_details), + ) + except Exception: + logger.exception("Failed to persist scheduling error for job step") except Exception: logger.exception("Unexpected error while scheduling container for job step") finally: @@ -895,7 +885,7 @@ def run_container( def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_args: dict): status_details = {} - status = PlatformJobStatus.PENDING.value + status = PlatformJobStatus.PENDING # If a request to pause or cancel came in while we were waiting for scheduling loop, # cancel scheduling the container @@ -987,12 +977,11 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a status_details["message"] = f"Pulling image {container_args['image']} from registry" # Send the status update to indicate we are pulling the image - self._nmp_sdk.jobs.steps.update_status( - step.name, + self._jobs.update_job_step_status( + name=step.name, workspace=step.workspace, job=step.job, - status=status, - status_details=status_details, + body=PlatformJobStatusUpdateRequest(status=status, status_details=status_details), ) try: pull_start = time.time() @@ -1016,12 +1005,11 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a # Send the status update to indicate we are starting the container status_details["message"] = f"Creating container with image {container_args['image']}" - self._nmp_sdk.jobs.steps.update_status( - step.name, + self._jobs.update_job_step_status( + name=step.name, workspace=step.workspace, job=step.job, - status=status, - status_details=status_details, + body=PlatformJobStatusUpdateRequest(status=status, status_details=status_details), ) # Now create it with the pulled container image @@ -1105,12 +1093,11 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a # If no errors to this point, start the container status_details["message"] = "Starting container" pre_start_status_write_started_at = time.monotonic() - self._nmp_sdk.jobs.steps.update_status( - step.name, + self._jobs.update_job_step_status( + name=step.name, workspace=step.workspace, job=step.job, - status=status, - status_details=status_details, + body=PlatformJobStatusUpdateRequest(status=status, status_details=status_details), ) logger.debug( "Docker pre-start status update succeeded", @@ -1224,7 +1211,7 @@ def enforce_sync_ttl( if container is None: return JobUpdate( - status=PlatformJobStatus.ERROR.value, + status=PlatformJobStatus.ERROR, status_details={"message": message}, error_details={"message": message}, ) @@ -1247,7 +1234,7 @@ def _kill_container_with_error( }, ) return JobUpdate( - status=PlatformJobStatus.ERROR.value, + status=PlatformJobStatus.ERROR, status_details=status_details, error_details=error_details, ) @@ -1261,14 +1248,16 @@ def _kill_container_with_error( raise task_id = self.get_label_from_container(container, JOB_TASK_ID_LABEL) - self._nmp_sdk.jobs.tasks.create_or_update( - task_id, + self._jobs.update_job_step_task( + name=task_id, workspace=step.workspace, job=step.job, step=step.name, - status=PlatformJobStatus.ERROR.value, - status_details=status_details, # type: ignore - error_details=error_details, # type: ignore + body=PlatformJobTaskUpdate( + status=PlatformJobStatus.ERROR, + status_details=status_details, + error_details=error_details, + ), ) logger.info( "Updated task", @@ -1280,9 +1269,7 @@ def _kill_container_with_error( "error_details": error_details, }, ) - return JobUpdate( - status=PlatformJobStatus.ERROR.value, status_details=status_details, error_details=error_details - ) + return JobUpdate(status=PlatformJobStatus.ERROR, status_details=status_details, error_details=error_details) def sync_pending(self, step: PlatformJobStepWithContext, container: Container | None) -> JobUpdate: if container is None: @@ -1427,18 +1414,20 @@ def create_step_update(self, step: PlatformJobStepWithContext, container: Contai ) # Upsert the task against the Jobs API. - self._nmp_sdk.jobs.tasks.create_or_update( - task_id, + self._jobs.update_job_step_task( + name=task_id, workspace=step.workspace, job=step.job, step=step.name, - status=status.value, - status_details=status_details, - error_details=error_details, - error_stack=error_stack, + body=PlatformJobTaskUpdate( + status=status, + status_details=status_details, + error_details=error_details, + error_stack=error_stack, + ), ) logger.info("Updated task", extra={"task_id": task_id, "status": status}) - return JobUpdate(status=status.value, status_details=status_details, error_details=error_details) + return JobUpdate(status=status, status_details=status_details, error_details=error_details) def map_docker_container_status_to_platform_status( self, step: PlatformJobStepWithContext, container: Container diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py index a2b9a706d2..7f6a8a574e 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py @@ -12,8 +12,35 @@ from kubernetes.client.models import V1Pod from kubernetes.client.rest import ApiException from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.jobs.client import JobsClient +from nemo_platform_plugin.jobs.execution_profiles import ( + BaseKubernetesExecutionProfileConfig as PluginBaseKubernetesExecutionProfileConfig, +) +from nemo_platform_plugin.jobs.execution_profiles import ( + ImagePullSecret as ImagePullSecret, +) +from nemo_platform_plugin.jobs.execution_profiles import ( + KubernetesEmptyDirVolume as KubernetesEmptyDirVolume, +) +from nemo_platform_plugin.jobs.execution_profiles import ( + KubernetesJobStorageConfig as PluginKubernetesJobStorageConfig, +) +from nemo_platform_plugin.jobs.execution_profiles import ( + KubernetesObjectMetadata as KubernetesObjectMetadata, +) +from nemo_platform_plugin.jobs.execution_profiles import ( + KubernetesPersistentVolumeClaim as KubernetesPersistentVolumeClaim, +) +from nemo_platform_plugin.jobs.execution_profiles import ( + KubernetesVolume as PluginKubernetesVolume, +) +from nemo_platform_plugin.jobs.execution_profiles import ( + KubernetesVolumeMount as PluginKubernetesVolumeMount, +) +from nemo_platform_plugin.jobs.types import PlatformJobStepWithContext, PlatformJobTaskUpdate from nmp.common.auth import AuthContext -from nmp.common.config import ImagePullSecret, get_platform_config +from nmp.common.config import get_platform_config from nmp.common.jobs.constants import ( DEFAULT_CONFIG_STORAGE_PATH, DEFAULT_NEMO_JOB_STEP_CONFIG_FILE_PATH, @@ -32,9 +59,7 @@ TERMINAL_EXIT_CODES, ) from nmp.common.jobs.schemas import PlatformJobStatus -from nmp.core.jobs.api.v2.jobs.schemas import PlatformJobStepWithContext from nmp.core.jobs.app.constants import ( - DEFAULT_VOLUME_PERMISSIONS_IMAGE, JOB_ATTEMPT_ID_LABEL, JOB_EXECUTION_BACKEND_LABEL, JOB_EXECUTION_PROFILE_LABEL, @@ -53,13 +78,12 @@ ) from nmp.core.jobs.app.providers import ComputeResources, ContainerSpec from nmp.core.jobs.controllers.backends.base import ( - JobExecutionProfileConfig, get_logs_endpoint_from_fileset, resolve_gpu_job_shm_size, resolve_task_image, ) from nmp.core.jobs.controllers.backends.exceptions import FailedToScheduleError, JobStorageError -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, Field logger = logging.getLogger(__name__) @@ -125,15 +149,16 @@ def build_image_pull_secrets( """Build Kubernetes image pull secrets from configuration.""" global_image_pull_secrets = get_platform_config().image_pull_secrets - # Merge the two lists, avoiding duplicates - merged_secrets = {secret.name: secret for secret in global_image_pull_secrets} + # Merge the two lists by secret name, avoiding duplicates. Both the global + # (platform config) and profile (plugin) ImagePullSecret models expose + # ``name``; only the name is needed to build the reference. + merged_names: dict[str, None] = {} + for secret in global_image_pull_secrets: + merged_names[secret.name] = None for secret in image_pull_secrets: - merged_secrets[secret.name] = secret + merged_names[secret.name] = None - result = [] - for secret in list(merged_secrets.values()): - result.append(client.V1LocalObjectReference(name=secret.name)) - return result + return [client.V1LocalObjectReference(name=name) for name in merged_names] def build_resource_requirements( @@ -206,11 +231,6 @@ def build_pod_security_context(security_context: dict[str, Any] | None) -> clien raise ValueError(f"Invalid security context configuration: {e}") from e -class KubernetesObjectMetadata(BaseModel): - labels: dict[str, str] = Field(default_factory=dict) - annotations: dict[str, str] = Field(default_factory=dict) - - def build_metadata(labels: dict[str, str] | None, metadata: KubernetesObjectMetadata | None) -> client.V1ObjectMeta: """Build Kubernetes metadata from configuration.""" if not metadata: @@ -579,36 +599,16 @@ def __init__(self, *args): return obj -class KubernetesPersistentVolumeClaim(BaseModel): - """Kubernetes Persistent Volume Claim definition.""" - - claim_name: str = Field(description="Persistent Volume Claim Name") - read_only: bool = Field(default=False, description="Whether the volume is mounted read-only") - - -class KubernetesEmptyDirVolume(BaseModel): - """Kubernetes EmptyDir Volume definition.""" - - medium: str | None = Field(default=None, description="The medium of the emptyDir volume (e.g., 'Memory')") - size_limit: str | None = Field(default=None, description="The size limit of the emptyDir volume (e.g., '1Gi')") +# The Kubernetes volume/config data shapes live in the shared plugin leaf node +# (imported as ``Plugin*`` below) so the typed HTTP client and the server agree +# on the wire shape. The server subclasses add the ``to_k8s()`` behaviour that +# requires the ``kubernetes`` client library, and re-type the volume fields so +# ``to_k8s()`` is available on nested volumes. -class KubernetesVolume(BaseModel): +class KubernetesVolume(PluginKubernetesVolume): """Kubernetes Volume definition.""" - name: str = Field(description="Volume Name") - persistent_volume_claim: KubernetesPersistentVolumeClaim | None = Field( - default=None, description="Persistent Volume Claim configuration" - ) - empty_dir: KubernetesEmptyDirVolume | None = Field(default=None, description="EmptyDir Volume configuration") - - @model_validator(mode="after") - def validate_self(self): - """Ensure that exactly one volume source is specified.""" - if sum(source is not None for source in [self.persistent_volume_claim, self.empty_dir]) != 1: - raise ValueError("Exactly one of 'persistent_volume_claim' or 'empty_dir' must be specified.") - return self - def to_k8s(self) -> client.V1Volume: """Convert to Kubernetes V1Volume object.""" volume = client.V1Volume(name=self.name) @@ -625,14 +625,9 @@ def to_k8s(self) -> client.V1Volume: return volume -class KubernetesVolumeMount(BaseModel): +class KubernetesVolumeMount(PluginKubernetesVolumeMount): """Kubernetes Volume Mount definition.""" - name: str = Field(description="Volume Name") - mount_path: str = Field(description="Mount Path in the container") - sub_path: str | None = Field(default=None, description="Sub-path within the volume to mount") - read_only: bool = Field(default=False, description="Whether the volume mount is read-only") - def to_k8s(self) -> client.V1VolumeMount: """Convert to Kubernetes V1VolumeMount object.""" return client.V1VolumeMount( @@ -643,80 +638,23 @@ def to_k8s(self) -> client.V1VolumeMount: ) -class KubernetesJobStorageConfig(BaseModel): +class KubernetesJobStorageConfig(PluginKubernetesJobStorageConfig): """Configuration for persistent storage in Kubernetes jobs.""" - pvc_name: str = Field(default="", description="Persistent Volume Claim Name to use for job storage.") - volume_permissions_image: str = Field( - default=DEFAULT_VOLUME_PERMISSIONS_IMAGE, description="Image used to set volume permissions" - ) + # Volume fields re-typed to the server subclasses so nested volumes carry ``to_k8s()``. additional_volumes: list[KubernetesVolume] = Field(default_factory=list, description="Additional volumes to mount") additional_volume_mounts: list[KubernetesVolumeMount] = Field( default_factory=list, description="Additional volume mounts" ) -class BaseKubernetesExecutionProfileConfig(JobExecutionProfileConfig): - """Common configuration for Kubernetes execution environment.""" +class BaseKubernetesExecutionProfileConfig(PluginBaseKubernetesExecutionProfileConfig): + """Kubernetes execution config whose storage carries ``to_k8s()`` (server-side).""" - namespace: str | None = Field( - default=None, - description="Kubernetes namespace to submit the job to. If not set, it will be determined from the environment.", - ) - - service_account_name: str = Field( - default="default", - description="Kubernetes service account name for job pods. Uses the Kubernetes default service account when set to 'default'.", - ) - - # Scheduling and resource configuration - tolerations: list[dict[str, Any]] = Field( - default_factory=list, description="Tolerations for the Kubernetes job pods." - ) - node_selector: dict[str, str] = Field( - default_factory=dict, description="Node selector for the Kubernetes job pods." - ) - affinity: dict[str, Any] = Field(default_factory=dict, description="Affinity for the Kubernetes job pods.") - resources: ComputeResources = Field( - default_factory=ComputeResources, description="Resource requests and limits for the Kubernetes job pods." - ) - pod_security_context: dict[str, Any] = Field( - default_factory=dict, description="Pod security context for the Kubernetes job pods." - ) - - # Image pull secrets - image_pull_secrets: list[ImagePullSecret] = Field( - default_factory=list, description="Image pull secrets for the Kubernetes job pods." - ) - - # Optional metadata to add to each job object - job_metadata: KubernetesObjectMetadata = Field( - default_factory=KubernetesObjectMetadata, - description="Metadata to add to each job object in the Kubernetes job.", - ) - - # Optional metadata to add to each pod in the job - pod_metadata: KubernetesObjectMetadata = Field( - default_factory=KubernetesObjectMetadata, description="Metadata to add to each pod in the Kubernetes job." - ) - - # Storage configurations for the job storage: KubernetesJobStorageConfig = Field( default_factory=KubernetesJobStorageConfig, description="Storage configuration for the Kubernetes job pods." ) - num_gpus: int = Field(default=1, description="Number of GPUs to request for the job") - - scheduler_name: str = Field( - default="", - description="The scheduler name to use for the pod spec. When non-empty, this value is applied to the pod's schedulerName field, enabling custom schedulers such as KAI Scheduler. Empty string omits schedulerName so the cluster default scheduler is used.", - ) - - launcher_image: str = Field( - default="nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest", - description="Container image that contains the jobs-launcher binary.", - ) - # This is the name of the shared volume used to inject the launcher binary into the main job container LAUNCHER_VOLUME_NAME = "launcher" @@ -1192,15 +1130,17 @@ def update_all_tasks( error_details["message"] = f"Pod {pod_status.name} is in error state" # Upsert the task against the Jobs API. - nmp_sdk.jobs.tasks.create_or_update( - pod_status.task_id, + client_from_platform(nmp_sdk, JobsClient).update_job_step_task( + name=pod_status.task_id, workspace=step.workspace, job=step.job, step=step.name, - status=status.value, - status_details=status_details, - error_details=error_details, # type: ignore - error_stack=error_stack, + body=PlatformJobTaskUpdate( + status=status, + status_details=status_details, + error_details=error_details, + error_stack=error_stack, + ), ) logger.info(f"updated task '{pod_status.task_id}'") diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/kubernetes_job.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/kubernetes_job.py index dcac06443b..684755d4ec 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/kubernetes_job.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/kubernetes_job.py @@ -7,8 +7,8 @@ from kubernetes import client from kubernetes.client.models import V1Job, V1JobStatus from kubernetes.client.rest import ApiException +from nemo_platform_plugin.jobs.types import PlatformJobStepWithContext, PlatformJobTaskUpdate from nmp.common.jobs.schemas import PlatformJobStatus -from nmp.core.jobs.api.v2.jobs.schemas import PlatformJobStepWithContext from nmp.core.jobs.app.constants import ( JOB_EXECUTION_BACKEND_LABEL, JOB_EXECUTION_PROFILE_LABEL, @@ -359,7 +359,7 @@ def enforce_sync_ttl( error_details = {"message": f"Job timed out after reaching max TTL of {ttl_seconds} seconds"} status_details["events"] = self.get_kube_job_events(k8s_job) update_all_tasks(self._nmp_sdk, self._core_v1, self.namespace, step) - return JobUpdate(status=status.value, status_details=status_details, error_details=error_details) + return JobUpdate(status=status, status_details=status_details, error_details=error_details) def sync_active(self, step: PlatformJobStepWithContext, job: V1Job | None) -> JobUpdate: job_name = name_for_step(step) @@ -389,20 +389,22 @@ def sync_terminate_job(self, step: PlatformJobStepWithContext, job: V1Job | None # Job already deleted # List all the tasks on the step that are ACTIVE and mark them as CANCELLED too, # since at this point all those pods should be deleted. - tasks = self._nmp_sdk.jobs.tasks.list( + tasks = self._jobs.list_job_step_tasks( name=step.name, job=step.job, workspace=step.workspace, - ) + ).data() for task in tasks.data: if task.status == PlatformJobStatus.ACTIVE: - self._nmp_sdk.jobs.tasks.create_or_update( + self._jobs.update_job_step_task( name=task.name, workspace=step.workspace, job=step.job, step=step.name, - status=PlatformJobStatus.CANCELLED.value, - status_details={"message": "Task cancelled as part of job cancellation"}, + body=PlatformJobTaskUpdate( + status=PlatformJobStatus.CANCELLED, + status_details={"message": "Task cancelled as part of job cancellation"}, + ), ) return JobUpdate( diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/volcano_job.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/volcano_job.py index 0d07a89aa0..1ce9ead0e2 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/volcano_job.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/volcano_job.py @@ -7,8 +7,8 @@ from kubernetes import client from kubernetes.client.rest import ApiException +from nemo_platform_plugin.jobs.types import PlatformJobStepWithContext from nmp.common.jobs.schemas import PlatformJobStatus -from nmp.core.jobs.api.v2.jobs.schemas import PlatformJobStepWithContext from nmp.core.jobs.app.constants import ( JOB_EXECUTION_BACKEND_LABEL, JOB_EXECUTION_PROFILE_LABEL, @@ -627,14 +627,14 @@ def sync_remove_job_with_status( status_details["events"] = events self.terminate_job(job) if error_details is not None: - status = PlatformJobStatus.ERROR.value + status = PlatformJobStatus.ERROR else: - status = step.status.value + status = step.status return JobUpdate(status=status, status_details=status_details, error_details=error_details) else: # If the job was not found, then it has been successfully stopped and removed # Transition into terminal state - return JobUpdate(status=stop_status.value) + return JobUpdate(status=stop_status) def terminate_job(self, job: dict): labels = job.get("metadata", {}).get("labels", {}) or {} diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py index 6944e0cf21..2fb747ab1e 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py @@ -17,7 +17,7 @@ from pathlib import Path from typing import Literal -from nemo_platform.types.jobs import PlatformJobStepWithContext +from nemo_platform_plugin.jobs.types import PlatformJobStepWithContext, PlatformJobTaskUpdate from nmp.common.auth import AuthContext from nmp.common.config import get_platform_config from nmp.common.jobs.constants import ( @@ -163,7 +163,7 @@ def shutdown(self) -> None: def schedule(self, executor_config: SubprocessExecutionProvider, step: PlatformJobStepWithContext) -> JobUpdate: if not executor_config.command: return JobUpdate( - status=PlatformJobStatus.ERROR.value, + status=PlatformJobStatus.ERROR, status_details={"message": _ERR_COMMAND_REQUIRED}, error_details={"message": _ERR_COMMAND_REQUIRED}, ) @@ -172,7 +172,7 @@ def schedule(self, executor_config: SubprocessExecutionProvider, step: PlatformJ existing = self._process_registry.get(key) if existing is not None and existing.process.poll() is None: return JobUpdate( - status=PlatformJobStatus.PENDING.value, + status=PlatformJobStatus.PENDING, status_details={ "message": "Subprocess already running", **self._task_status_details(existing), @@ -203,7 +203,7 @@ def schedule(self, executor_config: SubprocessExecutionProvider, step: PlatformJ message = f"Failed to start subprocess: executable not found: {command[0]}" logger.exception(message, extra=log_extra) return JobUpdate( - status=PlatformJobStatus.ERROR.value, + status=PlatformJobStatus.ERROR, status_details={"message": message}, error_details={"message": message, "error": str(exc)}, ) @@ -212,7 +212,7 @@ def schedule(self, executor_config: SubprocessExecutionProvider, step: PlatformJ message = f"Failed to start subprocess: {exc}" logger.exception(message) return JobUpdate( - status=PlatformJobStatus.ERROR.value, + status=PlatformJobStatus.ERROR, status_details={"message": message}, error_details={"message": message}, ) @@ -232,7 +232,7 @@ def schedule(self, executor_config: SubprocessExecutionProvider, step: PlatformJ message = f"Failed to initialize subprocess runtime: {exc}" logger.exception(message) return JobUpdate( - status=PlatformJobStatus.ERROR.value, + status=PlatformJobStatus.ERROR, status_details={"message": message}, error_details={"message": message}, ) @@ -251,17 +251,19 @@ def schedule(self, executor_config: SubprocessExecutionProvider, step: PlatformJ self._start_log_capture(step, metadata, "stderr") status_details = {"message": "Subprocess scheduled", **self._task_status_details(metadata)} - self._nmp_sdk.jobs.tasks.create_or_update( - metadata.task_id, + self._jobs.update_job_step_task( + name=metadata.task_id, workspace=step.workspace, job=step.job, step=step.name, - status=PlatformJobStatus.PENDING.value, - status_details=status_details, - error_details={}, + body=PlatformJobTaskUpdate( + status=PlatformJobStatus.PENDING, + status_details=status_details, + error_details={}, + ), ) - return JobUpdate(status=PlatformJobStatus.PENDING.value, status_details=status_details) + return JobUpdate(status=PlatformJobStatus.PENDING, status_details=status_details) def sync(self, step: PlatformJobStepWithContext) -> JobUpdate: key = SubprocessProcessKey(step.workspace, step.job, str(step.attempt_id), step.name) @@ -270,12 +272,12 @@ def sync(self, step: PlatformJobStepWithContext) -> JobUpdate: if metadata is None: if step.status == PlatformJobStatus.CANCELLING: return JobUpdate( - status=PlatformJobStatus.CANCELLED.value, + status=PlatformJobStatus.CANCELLED, status_details={"message": "Subprocess not found, job cancelled"}, ) if step.status == PlatformJobStatus.PAUSING: return JobUpdate( - status=PlatformJobStatus.PAUSED.value, + status=PlatformJobStatus.PAUSED, status_details={"message": "Subprocess not found, job paused"}, ) task_fallback = self._get_task_fallback_update(step) @@ -292,12 +294,12 @@ def sync(self, step: PlatformJobStepWithContext) -> JobUpdate: # serialized, jobs-backed state so reconciliation does not depend on process-local # memory. return JobUpdate( - status=PlatformJobStatus.PENDING.value, + status=PlatformJobStatus.PENDING, status_details=step.status_details or {"message": "Awaiting subprocess metadata"}, error_details=step.error_details or {}, ) return JobUpdate( - status=PlatformJobStatus.ERROR.value, + status=PlatformJobStatus.ERROR, error_details={"message": "Local subprocess metadata not found"}, ) @@ -324,7 +326,7 @@ def sync(self, step: PlatformJobStepWithContext) -> JobUpdate: self._tail_log_file(metadata.log_path), ) return JobUpdate( - status=PlatformJobStatus.ERROR.value, + status=PlatformJobStatus.ERROR, status_details={"message": message}, error_details={"message": message}, ) @@ -355,11 +357,11 @@ def cleanup_steps(self) -> None: def _get_task_fallback_update(self, step: PlatformJobStepWithContext) -> JobUpdate | None: try: - tasks = self._nmp_sdk.jobs.tasks.list( + tasks = self._jobs.list_job_step_tasks( name=step.name, job=step.job, workspace=step.workspace, - ) + ).data() except Exception: logger.warning( "Failed to fetch tasks for subprocess metadata fallback", @@ -514,7 +516,7 @@ def _start_log_capture( def _create_step_update(self, step: PlatformJobStepWithContext, metadata: SubprocessProcessMetadata) -> JobUpdate: status, status_details, error_details, error_stack = self._map_process_status(step, metadata) self._update_task(step, metadata, status, status_details, error_details, error_stack) - return JobUpdate(status=status.value, status_details=status_details, error_details=error_details) + return JobUpdate(status=status, status_details=status_details, error_details=error_details) def _map_process_status( self, step: PlatformJobStepWithContext, metadata: SubprocessProcessMetadata @@ -575,15 +577,17 @@ def _update_task( error_details: dict, error_stack: str = "", ) -> None: - self._nmp_sdk.jobs.tasks.create_or_update( - metadata.task_id, + self._jobs.update_job_step_task( + name=metadata.task_id, workspace=step.workspace, job=step.job, step=step.name, - status=status.value, - status_details=status_details, - error_details=error_details, - error_stack=error_stack, + body=PlatformJobTaskUpdate( + status=status, + status_details=status_details, + error_details=error_details, + error_stack=error_stack, + ), ) @staticmethod diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/test.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/test.py index 0604f82657..0a38b5a965 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/test.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/test.py @@ -2,39 +2,19 @@ # SPDX-License-Identifier: Apache-2.0 import logging -from typing import Generic, Literal, TypeVar +from typing import Generic, TypeVar -from nemo_platform.types.jobs import PlatformJobStepWithContext +from nemo_platform_plugin.jobs.execution_profiles import E2EJobExecutionProfile as E2EJobExecutionProfile +from nemo_platform_plugin.jobs.types import PlatformJobStepWithContext from nmp.common.jobs.schemas import PlatformJobStatus from nmp.core.jobs.app.providers import CPUExecutionProvider, ExecutionProviderT, GPUExecutionProvider -from nmp.core.jobs.app.schemas import BaseExecutionProfile from nmp.core.jobs.controllers.backends.base import JobBackend, JobExecutionProfileConfig, JobUpdate from nmp.core.jobs.controllers.backends.docker import DockerJobExecutionProfileConfig from nmp.core.jobs.controllers.backends.kubernetes import KubernetesJobExecutionProfileConfig -from pydantic import Field ProviderT = TypeVar("ProviderT", bound=ExecutionProviderT) -class E2EJobExecutionProfile(BaseExecutionProfile): - """ - Execution configuration for E2E testing. - This backend auto-completes jobs without actually running containers, - making tests fast and deterministic. - """ - - backend: Literal["e2e"] = "e2e" - config: JobExecutionProfileConfig = Field( - default_factory=JobExecutionProfileConfig, - description="Configuration for the e2e test executor", - ) - - @property - def supports_persistent_storage(self) -> bool: - """E2E backend claims to support persistent storage since jobs auto-complete without execution.""" - return True - - class MockJobBackend(Generic[ProviderT]): """ Provides a backend that can be used for testing diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/diagnostics.py b/services/core/jobs/src/nmp/core/jobs/controllers/diagnostics.py index 50249680d6..b4fc9df90e 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/diagnostics.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/diagnostics.py @@ -8,6 +8,8 @@ from typing import Any, Protocol from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.jobs.client import JobsClient from nmp.core.jobs.config import config _MAX_LOG_ENTRIES = 20 @@ -95,14 +97,16 @@ def collect_job_diagnostics( "step_name": step_ref.name, } + jobs = client_from_platform(sdk, JobsClient) + try: - job = sdk.jobs.retrieve(step_ref.job, workspace=step_ref.workspace) + job = jobs.get_job(name=step_ref.job, workspace=step_ref.workspace).data() diagnostics["job"] = _job_dict(job) except Exception as exc: diagnostics["job_error"] = str(exc) try: - status = sdk.jobs.get_status(step_ref.job, workspace=step_ref.workspace) + status = jobs.get_job_status(name=step_ref.job, workspace=step_ref.workspace).data() diagnostics["status_api"] = { "status": status.status, "status_details": status.status_details, @@ -119,21 +123,21 @@ def collect_job_diagnostics( diagnostics["status_api_error"] = str(exc) try: - refreshed_step = sdk.jobs.steps.retrieve(step_ref.name, job=step_ref.job, workspace=step_ref.workspace) + refreshed_step = jobs.get_job_step(name=step_ref.name, job=step_ref.job, workspace=step_ref.workspace).data() diagnostics["step"] = _step_dict(refreshed_step) except Exception as exc: diagnostics["step_error"] = str(exc) try: - tasks = sdk.jobs.tasks.list(step_ref.name, job=step_ref.job, workspace=step_ref.workspace) + tasks = jobs.list_job_step_tasks(name=step_ref.name, job=step_ref.job, workspace=step_ref.workspace).data() diagnostics["tasks_api"] = [_task_dict(task) for task in tasks.data] except Exception as exc: diagnostics["tasks_api_error"] = str(exc) try: if config.include_job_logs_in_diagnostics: - logs = sdk.jobs.get_logs(workspace=step_ref.workspace, name=step_ref.job) - diagnostics["job_logs"] = [entry.message for entry in logs.data[-_MAX_LOG_ENTRIES:]] + logs = jobs.list_job_logs(workspace=step_ref.workspace, name=step_ref.job).page() + diagnostics["job_logs"] = [entry.message for entry in logs.items[-_MAX_LOG_ENTRIES:]] except Exception as exc: diagnostics["job_logs_error"] = str(exc) diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/reconciler.py b/services/core/jobs/src/nmp/core/jobs/controllers/reconciler.py index 423370132e..00e04d3631 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/reconciler.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/reconciler.py @@ -4,11 +4,17 @@ import logging import threading import time +from typing import cast -from nemo_platform import APIError, APIStatusError, NeMoPlatform -from nemo_platform.types.jobs import PlatformJobStepWithContext -from nemo_platform.types.jobs.platform_job_steps_list_filter_param import PlatformJobStepsListFilterParam -from nemo_platform.types.shared.platform_job_status import PlatformJobStatus as SDKPlatformJobStatus +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoClientError, NemoHTTPError +from nemo_platform_plugin.jobs.client import JobsClient +from nemo_platform_plugin.jobs.types import ( + ListStepsQueryParams, + PlatformJobStatusUpdateRequest, + PlatformJobStepWithContext, +) from nmp.common.controller import Controller from nmp.common.jobs.schemas import PlatformJobStatus from nmp.common.observability import scoped_app_ctx, start_span_with_ctx @@ -32,6 +38,10 @@ def __init__( ) -> None: self._backend_registry = backend_registry self._nmp_sdk = nmp_sdk + # Typed Jobs client sharing the SDK's transport; every call passes + # ``workspace=`` explicitly (incl. cross-workspace "-"), so the client's + # default workspace is never relied upon. + self._jobs = client_from_platform(nmp_sdk, JobsClient) self._stop_signal = stop_signal self._is_healthy = False self._logger = logger @@ -58,7 +68,7 @@ def step(self): fetch_started_at = time.monotonic() with tracer.start_as_current_span("jobs_reconciler/fetch_steps_for_reconciliation"): try: - statuses: list[SDKPlatformJobStatus] = [ + statuses: list[str] = [ PlatformJobStatus.PENDING.value, PlatformJobStatus.ACTIVE.value, PlatformJobStatus.CANCELLING.value, @@ -66,7 +76,7 @@ def step(self): ] steps_to_reconcile = self.get_steps_for_reconciliation(statuses) self._is_healthy = True - except APIError: + except NemoClientError: self._is_healthy = False logger.exception("Could not fetch job steps for reconciliation", exc_info=True) return @@ -108,10 +118,7 @@ def step(self): }, ) logger.info(f"Updating job step status from '{step.status}' to '{job_update.status}'") - if ( - job_update.status == PlatformJobStatus.ERROR.value - and step.status != PlatformJobStatus.ERROR - ): + if job_update.status == PlatformJobStatus.ERROR and step.status != PlatformJobStatus.ERROR: log_job_diagnostics_if_debug( self._nmp_sdk, step, @@ -126,7 +133,7 @@ def step(self): status_details=job_update.status_details, error_details=job_update.error_details, ) - except APIStatusError as e: + except NemoHTTPError as e: # In cases when attempting to update job step status results in a conflict (409), # log a warning and continue processing other steps. if e.status_code == 409: @@ -175,20 +182,23 @@ def _update_step_status_with_timing( step: PlatformJobStepWithContext, provider: str, profile: str, - status: str, + status: PlatformJobStatus, status_details: dict | None = None, error_details: dict | None = None, ): started_at = time.monotonic() + update_fields: dict = {"status": status} + if status_details is not None: + update_fields["status_details"] = status_details + if error_details is not None: + update_fields["error_details"] = error_details try: - response = self._nmp_sdk.jobs.steps.update_status( - step.name, + response = self._jobs.update_job_step_status( + name=step.name, workspace=step.workspace, job=step.job, - status=status, - status_details=status_details, # type: ignore - error_details=error_details, # type: ignore - ) + body=PlatformJobStatusUpdateRequest(**update_fields), + ).data() except Exception: logger.warning( "Reconciler step status update failed", @@ -219,18 +229,26 @@ def _update_step_status_with_timing( ) return response - def get_steps_for_reconciliation(self, statuses: list[SDKPlatformJobStatus]) -> list[PlatformJobStepWithContext]: + def get_steps_for_reconciliation(self, statuses: list[str]) -> list[PlatformJobStepWithContext]: """ Return the list of steps to reconcile. """ - # Iterate through all pages to get all steps + # Iterate through all pages to get all steps. + # deepObject query param: sent as ``filter[status]=pending,active,...`` + # (comma form), which the steps-list route splits back into a list. The + # bracketed key isn't expressible as a TypedDict field, so cast the dict. steps = [] - filter_params: PlatformJobStepsListFilterParam = {"status": statuses} - for step in self._nmp_sdk.jobs.steps.list( + query = cast( + ListStepsQueryParams, + { + "filter[status]": ",".join(statuses), + "sort": "updated_at", + }, + ) + for step in self._jobs.list_steps( name="-", # Use "-" to indicate all jobs workspace="-", # Cross-workspace query - filter=filter_params, - sort="updated_at", - ): + query_params=query, + ).items(): steps.append(step) return steps diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/scheduler.py b/services/core/jobs/src/nmp/core/jobs/controllers/scheduler.py index e2f2f2a495..347f09df10 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/scheduler.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/scheduler.py @@ -5,13 +5,17 @@ import threading import time import traceback -from typing import TypedDict, cast, get_args +from typing import cast -import nemo_platform -from nemo_platform import APIStatusError, NeMoPlatform -from nemo_platform.types import PlatformJobStatus as SDKPlatformJobStatus -from nemo_platform.types.jobs import PlatformJobStepWithContext -from nemo_platform.types.jobs.platform_job_steps_list_filter_param import PlatformJobStepsListFilterParam +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoClientError, NemoHTTPError +from nemo_platform_plugin.jobs.client import JobsClient +from nemo_platform_plugin.jobs.types import ( + ListStepsQueryParams, + PlatformJobStatusUpdateRequest, + PlatformJobStepWithContext, +) from nmp.common.controller import Controller from nmp.common.jobs.schemas import PlatformJobStatus from nmp.common.observability import start_span_with_ctx @@ -28,18 +32,6 @@ DEFAULT_PROFILE = "default" DEFAULT_PROVIDER = "cpu" -SDK_PLATFORM_JOB_STATUSES = frozenset(get_args(SDKPlatformJobStatus)) - - -class StepStatusDetailParams(TypedDict, total=False): - status_details: dict[str, object] - error_details: dict[str, object] - - -def as_sdk_platform_job_status(status: str) -> SDKPlatformJobStatus: - if status not in SDK_PLATFORM_JOB_STATUSES: - raise ValueError(f"Unsupported platform job status: {status}") - return cast(SDKPlatformJobStatus, status) class JobScheduler(Controller): @@ -51,6 +43,10 @@ def __init__( ) -> None: self._backend_registry = backend_registry self._nmp_sdk = nmp_sdk + # Typed Jobs client sharing the SDK's transport; every call passes + # ``workspace=`` explicitly (incl. cross-workspace "-"), so the client's + # default workspace is never relied upon. + self._jobs = client_from_platform(nmp_sdk, JobsClient) self._stop_signal = stop_signal self._is_healthy = False self._logger = logger @@ -80,7 +76,7 @@ def step(self): try: steps = self.get_steps_for_scheduling() self._is_healthy = True - except nemo_platform.APIError: + except NemoClientError: self._is_healthy = False logger.exception("Could not fetch job steps for scheduling", exc_info=True) return @@ -118,7 +114,7 @@ def step(self): status_details=update.status_details, error_details=update.error_details, ) - except APIStatusError as e: + except NemoHTTPError as e: # Stopgap for a scheduler/reconciler race: by the time the scheduler persists # CREATED -> PENDING, another controller pass may already have advanced the # step to ACTIVE (or later). In that case, treating the stale PENDING write @@ -151,7 +147,7 @@ def step(self): self._update_step_status_with_timing( step=step, phase="resource_allocation_error", - status=PlatformJobStatus.ERROR.value, + status=PlatformJobStatus.ERROR, status_details={"message": e.message}, error_details={"message": e.message}, ) @@ -177,7 +173,7 @@ def step(self): self._update_step_status_with_timing( step=step, phase="unexpected_error", - status=PlatformJobStatus.ERROR.value, + status=PlatformJobStatus.ERROR, status_details={"message": str(e)}, error_details={"message": str(e), "error": traceback.format_exc()}, ) @@ -187,24 +183,23 @@ def _update_step_status_with_timing( *, step: PlatformJobStepWithContext, phase: str, - status: str, + status: PlatformJobStatus, status_details: dict[str, object] | None = None, error_details: dict[str, object] | None = None, ): started_at = time.monotonic() - detail_params: StepStatusDetailParams = {} + update_fields: dict = {"status": status} if status_details is not None: - detail_params["status_details"] = status_details + update_fields["status_details"] = status_details if error_details is not None: - detail_params["error_details"] = error_details + update_fields["error_details"] = error_details try: - response = self._nmp_sdk.jobs.steps.update_status( - step.name, + response = self._jobs.update_job_step_status( + name=step.name, workspace=step.workspace, job=step.job, - status=as_sdk_platform_job_status(status), - **detail_params, - ) + body=PlatformJobStatusUpdateRequest(**update_fields), + ).data() except Exception: logger.warning( "Scheduler step status update failed", @@ -238,17 +233,22 @@ def get_steps_for_scheduling(self) -> list[PlatformJobStepWithContext]: Return the oldest set of steps to schedule. We using the set of pending steps as our queue for what to schedule next. """ - # Iterate through all pages to get all steps + # The steps-list route parses ``filter`` as a deepObject query param, so + # the status list is sent as ``filter[status]=created,resuming`` (comma + # form), which the server splits back into a list. steps = [] - filter_params: PlatformJobStepsListFilterParam = { - "status": [PlatformJobStatus.CREATED.value, PlatformJobStatus.RESUMING.value] - } - for step in self._nmp_sdk.jobs.steps.list( + query = cast( + ListStepsQueryParams, + { + "filter[status]": f"{PlatformJobStatus.CREATED.value},{PlatformJobStatus.RESUMING.value}", + "sort": "created_at", + }, + ) + for step in self._jobs.list_steps( name="-", # Use "-" to indicate all jobs workspace="-", # Cross-workspace query - filter=filter_params, - sort="created_at", - ): + query_params=query, + ).items(): steps.append(step) return steps @@ -268,16 +268,16 @@ def _should_ignore_conflicting_pending_update( self, step: PlatformJobStepWithContext, update: JobUpdate, - error: APIStatusError, + error: NemoHTTPError, ) -> bool: - if error.status_code != 409 or update.status != PlatformJobStatus.PENDING.value: + if error.status_code != 409 or update.status != PlatformJobStatus.PENDING: return False - current_step = self._nmp_sdk.jobs.steps.retrieve( - step.name, + current_step = self._jobs.get_job_step( + name=step.name, workspace=step.workspace, job=step.job, - ) + ).data() original_status = PlatformJobStatus(step.status) current_status = PlatformJobStatus(current_step.status) return current_status != original_status and original_status.can_transition_to(current_status) diff --git a/services/core/jobs/tests/conftest.py b/services/core/jobs/tests/conftest.py index dc36f6da96..9791e085f1 100644 --- a/services/core/jobs/tests/conftest.py +++ b/services/core/jobs/tests/conftest.py @@ -3,6 +3,7 @@ import datetime import tempfile +from contextlib import ExitStack from pathlib import Path from typing import AsyncGenerator from unittest.mock import AsyncMock, MagicMock, patch @@ -255,22 +256,61 @@ def _mock_files_client(): return mock_files +# Controller modules that import ``client_from_platform`` to build a typed Jobs +# client. The fixture patches it in each so the shared ``mock_jobs`` client is +# returned for ``JobsClient`` requests. +# +# The backends (docker/subprocess/kubernetes_job) build their client once in +# ``JobBackend.__init__`` (base module) and reuse it via ``self._jobs``, so they +# no longer import ``client_from_platform`` directly — patching ``base`` covers +# them. ``common`` has a standalone helper that still builds its own client. +_JOBS_CLIENT_CONTROLLER_MODULES = ( + "nmp.core.jobs.controllers.scheduler", + "nmp.core.jobs.controllers.reconciler", + "nmp.core.jobs.controllers.diagnostics", + "nmp.core.jobs.controllers.backends.base", + "nmp.core.jobs.controllers.backends.kubernetes.common", +) + + @fixture -def mock_nmp_client(_mock_files_client): - """Create a flexible mock of NeMoPlatform for testing.""" - mock_client = MagicMock() +def mock_jobs_client(): + """Mock of the typed ``JobsClient`` used by the controllers. + + Methods return ``.data()``/``.items()``-aware responses so call sites like + ``client_from_platform(sdk, JobsClient).get_job_step(...).data()`` work. Tests + set ``.return_value`` on the individual methods and assert against them. + """ + return MagicMock() + - # Set up the nested structure that tests expect +@fixture +def mock_nmp_client(_mock_files_client, mock_jobs_client): + """Create a flexible mock of NeMoPlatform for testing. + + ``client_from_platform`` is patched in the dispatcher (returns the files client) + and in every controller module that builds a typed Jobs client. The controller + patches dispatch on the requested client type: ``JobsClient`` requests resolve to + ``mock_jobs_client``; anything else falls back to the files client. + """ + mock_client = MagicMock() mock_client.beta = MagicMock() mock_client.jobs = MagicMock() - mock_client.jobs.list = MagicMock() - mock_client.jobs.update_status = MagicMock() - mock_client.jobs.steps = MagicMock() - mock_client.jobs.steps.list = MagicMock() - mock_client.jobs.steps.retrieve = MagicMock() - mock_client.jobs.steps.update_status = MagicMock() - - with patch("nmp.core.jobs.app.dispatcher.client_from_platform", return_value=_mock_files_client): + + from nemo_platform_plugin.jobs.client import JobsClient + + def _dispatch(_sdk, client_type): + if client_type is JobsClient: + return mock_jobs_client + return _mock_files_client + + patchers = [patch("nmp.core.jobs.app.dispatcher.client_from_platform", return_value=_mock_files_client)] + patchers += [ + patch(f"{module}.client_from_platform", side_effect=_dispatch) for module in _JOBS_CLIENT_CONTROLLER_MODULES + ] + with ExitStack() as stack: + for patcher in patchers: + stack.enter_context(patcher) yield mock_client diff --git a/services/core/jobs/tests/controllers/client_mocks.py b/services/core/jobs/tests/controllers/client_mocks.py new file mode 100644 index 0000000000..508d77eb81 --- /dev/null +++ b/services/core/jobs/tests/controllers/client_mocks.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Builders for mock typed-client responses used by controller tests. + +The Jobs controllers call the typed ``JobsClient`` and chain ``.data()`` (single +resource) or ``.items()`` (paginated). These helpers wrap plain values in mock +responses that mimic ``NemoResponse`` / ``NemoPaginatedResponse`` so tests can set +``mock_jobs_client..return_value`` without threading the client's real +response objects through. +""" + +from typing import Any +from unittest.mock import MagicMock + + +def data_response(value: Any) -> MagicMock: + """Build a mock typed-client response whose ``.data()`` returns ``value``.""" + resp = MagicMock() + resp.data.return_value = value + return resp + + +def paginated_response(items: Any) -> MagicMock: + """Build a mock paginated response: ``.items()`` iterates, ``.data()`` returns the list.""" + resp = MagicMock() + materialized = list(items) + resp.items.return_value = iter(materialized) + resp.data.return_value = materialized + return resp diff --git a/services/core/jobs/tests/controllers/test_base.py b/services/core/jobs/tests/controllers/test_base.py index 6d85637939..4a728249b9 100644 --- a/services/core/jobs/tests/controllers/test_base.py +++ b/services/core/jobs/tests/controllers/test_base.py @@ -4,6 +4,8 @@ """Unit tests for jobs controller backends base module.""" import datetime +from contextlib import contextmanager +from types import SimpleNamespace from unittest.mock import MagicMock, patch from nmp.common.config import PlatformConfig @@ -14,6 +16,8 @@ from nmp.core.jobs.controllers.backends.base import get_logs_endpoint_from_fileset, resolve_task_image from nmp.core.jobs.controllers.backends.test import MockKubernetesCPUJobBackend +from services.core.jobs.tests.controllers.client_mocks import data_response + class TestGetLogsEndpointFromFileset: """Tests for get_logs_endpoint_from_fileset.""" @@ -193,6 +197,29 @@ def _make_backend(mock_sdk: MagicMock | None = None) -> MockKubernetesCPUJobBack return MockKubernetesCPUJobBackend(nmp_sdk=sdk, execution_profile_config=MagicMock(), profile_name="default") +@contextmanager +def _patched_jobs_client(backend): + """Stub the backend's held ``self._jobs`` handle and yield the mock ``JobsClient``. + + The backend builds its typed Jobs client once in ``JobBackend.__init__`` and + reuses it as ``self._jobs``, so tests stub that handle directly (rather than + patching ``client_from_platform``). ``check_step_is_stale`` fetches tasks via + ``self._jobs.list_job_step_tasks(...).data()``; the returned page exposes the + task list on its ``.data`` attribute. + """ + mock_jobs = MagicMock() + original = backend._jobs + backend._jobs = mock_jobs + try: + yield mock_jobs + finally: + backend._jobs = original + + +def _set_tasks(mock_jobs: MagicMock, tasks: list) -> None: + mock_jobs.list_job_step_tasks.return_value = data_response(SimpleNamespace(data=tasks)) + + class TestCheckTaskStaleness: """Tests for JobBackend.check_step_is_stale.""" @@ -216,106 +243,107 @@ def test_not_stale_when_step_too_young(self): created_at=datetime.datetime.now(datetime.timezone.utc), ) - assert backend.check_step_is_stale(step) is False - backend._nmp_sdk.jobs.tasks.list.assert_not_called() + with _patched_jobs_client(backend) as mock_jobs: + assert backend.check_step_is_stale(step) is False + mock_jobs.list_job_step_tasks.assert_not_called() def test_not_stale_when_no_active_tasks(self): - mock_sdk = MagicMock() - mock_sdk.jobs.tasks.list.return_value.data = [ - _make_task(status="completed"), - _make_task(status="error"), - ] - backend = _make_backend(mock_sdk) + backend = _make_backend() step = _make_step( staleness_timeout=60, created_at=datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=120), ) - assert backend.check_step_is_stale(step) is False + with _patched_jobs_client(backend) as mock_jobs: + _set_tasks(mock_jobs, [_make_task(status="completed"), _make_task(status="error")]) + assert backend.check_step_is_stale(step) is False def test_not_stale_when_task_recently_updated(self): - mock_sdk = MagicMock() - mock_sdk.jobs.tasks.list.return_value.data = [ - _make_task(status="active", updated_at=datetime.datetime.now(datetime.timezone.utc)), - ] - backend = _make_backend(mock_sdk) + backend = _make_backend() step = _make_step( staleness_timeout=60, created_at=datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=120), ) - assert backend.check_step_is_stale(step) is False + with _patched_jobs_client(backend) as mock_jobs: + _set_tasks( + mock_jobs, + [_make_task(status="active", updated_at=datetime.datetime.now(datetime.timezone.utc))], + ) + assert backend.check_step_is_stale(step) is False def test_stale_when_all_active_tasks_exceed_threshold(self): now = datetime.datetime.now(datetime.timezone.utc) - mock_sdk = MagicMock() - mock_sdk.jobs.tasks.list.return_value.data = [ - _make_task(status="active", updated_at=now - datetime.timedelta(seconds=200)), - _make_task(status="active", updated_at=now - datetime.timedelta(seconds=300)), - _make_task(status="completed"), - ] - backend = _make_backend(mock_sdk) + backend = _make_backend() step = _make_step( staleness_timeout=60, created_at=now - datetime.timedelta(seconds=120), ) - assert backend.check_step_is_stale(step) is True + with _patched_jobs_client(backend) as mock_jobs: + _set_tasks( + mock_jobs, + [ + _make_task(status="active", updated_at=now - datetime.timedelta(seconds=200)), + _make_task(status="active", updated_at=now - datetime.timedelta(seconds=300)), + _make_task(status="completed"), + ], + ) + assert backend.check_step_is_stale(step) is True def test_not_stale_when_one_active_task_is_fresh(self): now = datetime.datetime.now(datetime.timezone.utc) - mock_sdk = MagicMock() - mock_sdk.jobs.tasks.list.return_value.data = [ - _make_task(status="active", updated_at=now - datetime.timedelta(seconds=200)), - _make_task(status="active", updated_at=now - datetime.timedelta(seconds=10)), - ] - backend = _make_backend(mock_sdk) + backend = _make_backend() step = _make_step( staleness_timeout=60, created_at=now - datetime.timedelta(seconds=120), ) - assert backend.check_step_is_stale(step) is False + with _patched_jobs_client(backend) as mock_jobs: + _set_tasks( + mock_jobs, + [ + _make_task(status="active", updated_at=now - datetime.timedelta(seconds=200)), + _make_task(status="active", updated_at=now - datetime.timedelta(seconds=10)), + ], + ) + assert backend.check_step_is_stale(step) is False def test_returns_false_on_api_failure(self): - mock_sdk = MagicMock() - mock_sdk.jobs.tasks.list.side_effect = RuntimeError("connection error") - backend = _make_backend(mock_sdk) + backend = _make_backend() step = _make_step( staleness_timeout=60, created_at=datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=120), ) - assert backend.check_step_is_stale(step) is False + with _patched_jobs_client(backend) as mock_jobs: + mock_jobs.list_job_step_tasks.side_effect = RuntimeError("connection error") + assert backend.check_step_is_stale(step) is False def test_handles_naive_updated_at_as_utc(self): now = datetime.datetime.now(datetime.timezone.utc) naive_old = (now - datetime.timedelta(seconds=200)).replace(tzinfo=None) - mock_sdk = MagicMock() - mock_sdk.jobs.tasks.list.return_value.data = [ - _make_task(status="active", updated_at=naive_old), - ] - backend = _make_backend(mock_sdk) + backend = _make_backend() step = _make_step( staleness_timeout=60, created_at=now - datetime.timedelta(seconds=120), ) - assert backend.check_step_is_stale(step) is True + with _patched_jobs_client(backend) as mock_jobs: + _set_tasks(mock_jobs, [_make_task(status="active", updated_at=naive_old)]) + assert backend.check_step_is_stale(step) is True def test_returns_false_when_task_missing_updated_at(self): now = datetime.datetime.now(datetime.timezone.utc) - mock_sdk = MagicMock() - mock_sdk.jobs.tasks.list.return_value.data = [ - _make_task(status="active", updated_at=None), - ] - backend = _make_backend(mock_sdk) + backend = _make_backend() step = _make_step( staleness_timeout=60, created_at=now - datetime.timedelta(seconds=120), ) - assert backend.check_step_is_stale(step) is False + with _patched_jobs_client(backend) as mock_jobs: + _set_tasks(mock_jobs, [_make_task(status="active", updated_at=None)]) + assert backend.check_step_is_stale(step) is False class TestResolveTaskImage: diff --git a/services/core/jobs/tests/controllers/test_diagnostics.py b/services/core/jobs/tests/controllers/test_diagnostics.py index b3bac55821..3b3de63dcf 100644 --- a/services/core/jobs/tests/controllers/test_diagnostics.py +++ b/services/core/jobs/tests/controllers/test_diagnostics.py @@ -6,38 +6,51 @@ from nmp.core.jobs.controllers.diagnostics import _MAX_ERROR_STACK_CHARS, collect_job_diagnostics - -def _make_sdk_with_logs() -> Mock: - sdk = Mock() - sdk.jobs.retrieve.return_value = SimpleNamespace( - name="job-1", - status="error", - status_details={}, - error_details={}, +from services.core.jobs.tests.controllers.client_mocks import data_response + + +def _make_jobs_client_with_logs() -> Mock: + """Build a mock typed ``JobsClient`` with response-shaped method results.""" + jobs = Mock() + jobs.get_job.return_value = data_response( + SimpleNamespace( + name="job-1", + status="error", + status_details={}, + error_details={}, + ) ) - sdk.jobs.get_status.return_value = SimpleNamespace( - status="error", - status_details={}, - error_details={}, - steps=[], + jobs.get_job_status.return_value = data_response( + SimpleNamespace( + status="error", + status_details={}, + error_details={}, + steps=[], + ) ) - sdk.jobs.steps.retrieve.return_value = SimpleNamespace( - name="step-1", - status="error", - status_details={}, - error_details={}, + jobs.get_job_step.return_value = data_response( + SimpleNamespace( + name="step-1", + status="error", + status_details={}, + error_details={}, + ) ) - sdk.jobs.tasks.list.return_value = SimpleNamespace(data=[]) - sdk.jobs.get_logs.return_value = SimpleNamespace( - data=[SimpleNamespace(message="secret-token=abc123"), SimpleNamespace(message="another line")] + jobs.list_job_step_tasks.return_value = data_response(SimpleNamespace(data=[])) + jobs.list_job_logs.return_value.page.return_value = SimpleNamespace( + items=[SimpleNamespace(message="secret-token=abc123"), SimpleNamespace(message="another line")] ) - return sdk + return jobs def test_collect_job_diagnostics_omits_raw_job_logs_by_default() -> None: - sdk = _make_sdk_with_logs() + sdk = Mock() + jobs = _make_jobs_client_with_logs() - with patch("nmp.core.jobs.controllers.diagnostics.config.include_job_logs_in_diagnostics", False): + with ( + patch("nmp.core.jobs.controllers.diagnostics.client_from_platform", return_value=jobs), + patch("nmp.core.jobs.controllers.diagnostics.config.include_job_logs_in_diagnostics", False), + ): diagnostics = collect_job_diagnostics( sdk, workspace="default", @@ -47,13 +60,17 @@ def test_collect_job_diagnostics_omits_raw_job_logs_by_default() -> None: ) assert "job_logs" not in diagnostics - sdk.jobs.get_logs.assert_not_called() + jobs.list_job_logs.assert_not_called() def test_collect_job_diagnostics_includes_raw_job_logs_when_enabled() -> None: - sdk = _make_sdk_with_logs() + sdk = Mock() + jobs = _make_jobs_client_with_logs() - with patch("nmp.core.jobs.controllers.diagnostics.config.include_job_logs_in_diagnostics", True): + with ( + patch("nmp.core.jobs.controllers.diagnostics.client_from_platform", return_value=jobs), + patch("nmp.core.jobs.controllers.diagnostics.config.include_job_logs_in_diagnostics", True), + ): diagnostics = collect_job_diagnostics( sdk, workspace="default", @@ -63,67 +80,77 @@ def test_collect_job_diagnostics_includes_raw_job_logs_when_enabled() -> None: ) assert diagnostics["job_logs"] == ["secret-token=abc123", "another line"] - sdk.jobs.get_logs.assert_called_once_with(workspace="default", name="job-1") + jobs.list_job_logs.assert_called_once_with(workspace="default", name="job-1") def test_collect_job_diagnostics_trims_long_error_details_tracebacks() -> None: sdk = Mock() + jobs = Mock() long_error = "traceback-" + ("x" * (_MAX_ERROR_STACK_CHARS + 50)) expected_trimmed = long_error[-_MAX_ERROR_STACK_CHARS:] error_details = {"message": "boom", "error": long_error, "other": "keep"} - sdk.jobs.retrieve.return_value = SimpleNamespace( - name="job-1", - status="error", - status_details={}, - error_details=error_details, + jobs.get_job.return_value = data_response( + SimpleNamespace( + name="job-1", + status="error", + status_details={}, + error_details=error_details, + ) ) - sdk.jobs.get_status.return_value = SimpleNamespace( - status="error", - status_details={}, - error_details=error_details, - steps=[ - SimpleNamespace( - name="step-1", - status="error", - status_details={}, - error_details=error_details, - tasks=[ - SimpleNamespace( - name="task-1", - status="error", - status_details={}, - error_details=error_details, - error_stack=long_error, - ) - ], - ) - ], + jobs.get_job_status.return_value = data_response( + SimpleNamespace( + status="error", + status_details={}, + error_details=error_details, + steps=[ + SimpleNamespace( + name="step-1", + status="error", + status_details={}, + error_details=error_details, + tasks=[ + SimpleNamespace( + name="task-1", + status="error", + status_details={}, + error_details=error_details, + error_stack=long_error, + ) + ], + ) + ], + ) ) - sdk.jobs.steps.retrieve.return_value = SimpleNamespace( - name="step-1", - status="error", - status_details={}, - error_details=error_details, + jobs.get_job_step.return_value = data_response( + SimpleNamespace( + name="step-1", + status="error", + status_details={}, + error_details=error_details, + ) ) - sdk.jobs.tasks.list.return_value = SimpleNamespace( - data=[ - SimpleNamespace( - name="task-1", - status="error", - status_details={}, - error_details=error_details, - error_stack=long_error, - ) - ] + jobs.list_job_step_tasks.return_value = data_response( + SimpleNamespace( + data=[ + SimpleNamespace( + name="task-1", + status="error", + status_details={}, + error_details=error_details, + error_stack=long_error, + ) + ] + ) ) - diagnostics = collect_job_diagnostics( - sdk, - workspace="default", - job_name="job-1", - step_name="step-1", - context="test", - ) + with patch("nmp.core.jobs.controllers.diagnostics.client_from_platform", return_value=jobs): + diagnostics = collect_job_diagnostics( + sdk, + workspace="default", + job_name="job-1", + step_name="step-1", + context="test", + ) assert diagnostics["job"]["error_details"]["error"] == expected_trimmed assert diagnostics["status_api"]["error_details"]["error"] == expected_trimmed diff --git a/services/core/jobs/tests/controllers/test_docker_backend.py b/services/core/jobs/tests/controllers/test_docker_backend.py index b6859358ba..1c828edbd9 100644 --- a/services/core/jobs/tests/controllers/test_docker_backend.py +++ b/services/core/jobs/tests/controllers/test_docker_backend.py @@ -58,9 +58,15 @@ DockerVolumeMount, GPUDockerJobBackend, ) -from nmp.core.jobs.controllers.backends.exceptions import ResourceAllocationError, SchedulingDeferred +from nmp.core.jobs.controllers.backends.exceptions import ( + FailedToScheduleError, + ResourceAllocationError, + SchedulingDeferred, +) from pydantic import ValidationError +from services.core.jobs.tests.controllers.client_mocks import data_response + TEST_JOBS_CONTROLLER_INSTANCE_ID = "test-owner" @@ -357,6 +363,8 @@ def test_docker_job_sync(docker_job, docker_client_mock, test_job_step): JOB_STEP_NAME_LABEL: test_job_step.name, JOB_TASK_ID_LABEL: task_id, } + # Error-path status derivation reads container.logs(...).decode() for error_stack. + container_mock.logs.return_value = b"boom" # Clear side_effect so return_value takes precedence docker_client_mock.containers.get.side_effect = None docker_client_mock.containers.get.return_value = container_mock @@ -1027,6 +1035,8 @@ def test_gpu_cleanup_on_job_error(mock_nmp_client, docker_client_mock): JOB_CONTROLLER_INSTANCE_ID_LABEL: TEST_JOBS_CONTROLLER_INSTANCE_ID, JOB_TYPE_LABEL: JOB_TYPE_JOB, } + # Error-path status derivation reads container.logs(...).decode() for error_stack. + container_mock.logs.return_value = b"boom" # Clear side_effect so return_value takes precedence docker_client_mock.containers.get.side_effect = None @@ -1352,6 +1362,7 @@ def test_schedule_additional_volume_mounts(docker_job: CPUDockerJobBackend, dock ) def test_cancel_scheduling( docker_job, + mock_jobs_client, test_job_step, step_status, expected_result, @@ -1360,33 +1371,34 @@ def test_cancel_scheduling( expected_message, ): """Test cancel_scheduling behavior for different step statuses.""" - # Mock the retrieved step with the specified status + # Mock the retrieved step with the specified status. get_step fetches it via + # the typed client's get_job_step(...).data(). mock_refreshed_step = MagicMock() mock_refreshed_step.status = step_status.value - docker_job._nmp_sdk.jobs.steps.retrieve.return_value = mock_refreshed_step + mock_jobs_client.get_job_step.return_value = data_response(mock_refreshed_step) result = docker_job.cancel_scheduling(test_job_step) # Verify result matches expectation assert result is expected_result - # Verify API retrieve was called (get_step calls retrieve with keyword args) - docker_job._nmp_sdk.jobs.steps.retrieve.assert_called_once_with( + # Verify the step was fetched via the typed client (get_step -> get_job_step) + mock_jobs_client.get_job_step.assert_called_once_with( name=test_job_step.name, workspace=test_job_step.workspace, job=test_job_step.job ) if should_update_status: - # Verify update_status was called with expected parameters (name as positional to match backend call) - docker_job._nmp_sdk.jobs.steps.update_status.assert_called_once_with( - test_job_step.name, - workspace=test_job_step.workspace, - job=test_job_step.job, - status=expected_final_status.value, - status_details={"message": expected_message}, - ) + # Verify update_job_step_status was called with the expected status/details in the body. + mock_jobs_client.update_job_step_status.assert_called_once() + call = mock_jobs_client.update_job_step_status.call_args + assert call.kwargs["name"] == test_job_step.name + assert call.kwargs["workspace"] == test_job_step.workspace + assert call.kwargs["job"] == test_job_step.job + assert call.kwargs["body"].status == expected_final_status + assert call.kwargs["body"].status_details == {"message": expected_message} else: - # Verify update_status was NOT called - docker_job._nmp_sdk.jobs.steps.update_status.assert_not_called() + # Verify update_job_step_status was NOT called + mock_jobs_client.update_job_step_status.assert_not_called() @pytest.mark.parametrize("cleanup_completed_jobs_immediately", [True, False]) @@ -1678,6 +1690,23 @@ def test_docker_schedule_defers_when_start_admission_full(docker_job, docker_cli docker_job._container_start_admission.release() +def test_failed_schedule_logs_status_update_failure_and_releases_admission(docker_job, test_job_step): + assert docker_job._container_start_admission.acquire(blocking=False) + docker_job._run_container_in_thread = MagicMock( + side_effect=FailedToScheduleError("container failed", error_details={"message": "container failed"}) + ) + docker_job._jobs.update_job_step_status.side_effect = RuntimeError("jobs service unavailable") + + with patch("nmp.core.jobs.controllers.backends.docker.logger.exception") as log_exception: + docker_job.run_container(test_job_step, {}) + + log_exception.assert_any_call("Failed to schedule container for job step") + log_exception.assert_any_call("Failed to persist scheduling error for job step") + docker_job._jobs.update_job_step_status.assert_called_once() + assert docker_job._container_start_admission.acquire(blocking=False) + docker_job._container_start_admission.release() + + def test_resuming_step_skips_before_active_ttl_enforcement(docker_job, test_job_step): """RESUMING must not apply ttl_seconds_before_active (pause/resume rebasing).""" ttl_seconds = docker_job._execution_profile_config.ttl_seconds_before_active @@ -1698,7 +1727,7 @@ def test_before_active_ttl_uses_latest_of_created_and_updated(docker_job, test_j assert docker_job.check_step_ttl_before_active(test_job_step, ttl_seconds) is False -def test_cleanup_pending_created_container_by_ttl(docker_job, docker_client_mock, test_job_step): +def test_cleanup_pending_created_container_by_ttl(docker_job, docker_client_mock, mock_jobs_client, test_job_step): """A stale PENDING step with a Docker-created container transitions to ERROR.""" # Get the TTL configuration (default is 30 minutes) ttl_seconds = docker_job._execution_profile_config.ttl_seconds_before_active @@ -1740,15 +1769,17 @@ def test_cleanup_pending_created_container_by_ttl(docker_job, docker_client_mock container_mock.kill.assert_called_once() # Verify that the task was updated via the API - docker_job._nmp_sdk.jobs.tasks.create_or_update.assert_called_once_with( - task_id, - workspace=test_job_step.workspace, - job=test_job_step.job, - step=test_job_step.name, - status=PlatformJobStatus.ERROR.value, - status_details={"message": "Job timed out after reaching max TTL of 1800 seconds"}, - error_details={"message": "Job timed out after reaching max TTL of 1800 seconds"}, - ) + mock_jobs_client.update_job_step_task.assert_called_once() + task_call = mock_jobs_client.update_job_step_task.call_args + assert task_call.kwargs["name"] == task_id + assert task_call.kwargs["workspace"] == test_job_step.workspace + assert task_call.kwargs["job"] == test_job_step.job + assert task_call.kwargs["step"] == test_job_step.name + assert task_call.kwargs["body"].status == PlatformJobStatus.ERROR + assert task_call.kwargs["body"].status_details == { + "message": "Job timed out after reaching max TTL of 1800 seconds" + } + assert task_call.kwargs["body"].error_details == {"message": "Job timed out after reaching max TTL of 1800 seconds"} def test_pending_running_container_preempts_before_active_ttl(docker_job, docker_client_mock, test_job_step): @@ -1813,7 +1844,7 @@ def test_pending_exited_container_preempts_before_active_ttl(docker_job, docker_ container_mock.kill.assert_not_called() -def test_cleanup_active_by_ttl(docker_job, docker_client_mock, test_job_step): +def test_cleanup_active_by_ttl(docker_job, docker_client_mock, mock_jobs_client, test_job_step): """Test that sync of an ACTIVE step transitions to an ERROR state when step's created_at exceeds TTL.""" # Get the TTL configuration for active jobs (default is 24 hours) ttl_seconds = docker_job._execution_profile_config.ttl_seconds_active @@ -1854,19 +1885,23 @@ def test_cleanup_active_by_ttl(docker_job, docker_client_mock, test_job_step): container_mock.kill.assert_called_once() # Verify that the task was updated via the API - docker_job._nmp_sdk.jobs.tasks.create_or_update.assert_called_once_with( - task_id, - workspace=test_job_step.workspace, - job=test_job_step.job, - step=test_job_step.name, - status=PlatformJobStatus.ERROR.value, - status_details={"message": "Job timed out after reaching max TTL of 86400 seconds"}, - error_details={"message": "Job timed out after reaching max TTL of 86400 seconds"}, - ) + mock_jobs_client.update_job_step_task.assert_called_once() + task_call = mock_jobs_client.update_job_step_task.call_args + assert task_call.kwargs["name"] == task_id + assert task_call.kwargs["workspace"] == test_job_step.workspace + assert task_call.kwargs["job"] == test_job_step.job + assert task_call.kwargs["step"] == test_job_step.name + assert task_call.kwargs["body"].status == PlatformJobStatus.ERROR + assert task_call.kwargs["body"].status_details == { + "message": "Job timed out after reaching max TTL of 86400 seconds" + } + assert task_call.kwargs["body"].error_details == { + "message": "Job timed out after reaching max TTL of 86400 seconds" + } def test_ttl_enforcement_handles_409_when_kill_races_with_stopped_container( - docker_job, docker_client_mock, test_job_step + docker_job, docker_client_mock, mock_jobs_client, test_job_step ): """Test TTL enforcement handles gracefully when container.kill() races with a stopped container.""" # Get the TTL configuration for pending/created jobs @@ -1915,18 +1950,20 @@ def test_ttl_enforcement_handles_409_when_kill_races_with_stopped_container( container_mock.kill.assert_called_once() # Verify that the task was updated via the API despite the 409 error - docker_job._nmp_sdk.jobs.tasks.create_or_update.assert_called_once_with( - task_id, - workspace=test_job_step.workspace, - job=test_job_step.job, - step=test_job_step.name, - status=PlatformJobStatus.ERROR.value, - status_details={"message": "Job timed out after reaching max TTL of 1800 seconds"}, - error_details={"message": "Job timed out after reaching max TTL of 1800 seconds"}, - ) + mock_jobs_client.update_job_step_task.assert_called_once() + task_call = mock_jobs_client.update_job_step_task.call_args + assert task_call.kwargs["name"] == task_id + assert task_call.kwargs["workspace"] == test_job_step.workspace + assert task_call.kwargs["job"] == test_job_step.job + assert task_call.kwargs["step"] == test_job_step.name + assert task_call.kwargs["body"].status == PlatformJobStatus.ERROR + assert task_call.kwargs["body"].status_details == { + "message": "Job timed out after reaching max TTL of 1800 seconds" + } + assert task_call.kwargs["body"].error_details == {"message": "Job timed out after reaching max TTL of 1800 seconds"} -def test_sync_stop_container_already_stopped(docker_job, docker_client_mock, test_job_step): +def test_sync_stop_container_already_stopped(docker_job, docker_client_mock, mock_jobs_client, test_job_step): """Test that sync handles gracefully when container.stop() is called on already stopped container.""" # Set the step to CANCELLING status (which triggers sync_stop_container) test_job_step.status = PlatformJobStatus.CANCELLING @@ -1967,16 +2004,16 @@ def test_sync_stop_container_already_stopped(docker_job, docker_client_mock, tes assert result.error_details == {} # Verify that the task was updated via the API - docker_job._nmp_sdk.jobs.tasks.create_or_update.assert_called_once_with( - task_id, - workspace=test_job_step.workspace, - job=test_job_step.job, - step=test_job_step.name, - status=PlatformJobStatus.CANCELLED.value, - status_details={"message": "Job was cancelled successfully with exit code 0"}, - error_details={}, - error_stack="", - ) + mock_jobs_client.update_job_step_task.assert_called_once() + task_call = mock_jobs_client.update_job_step_task.call_args + assert task_call.kwargs["name"] == task_id + assert task_call.kwargs["workspace"] == test_job_step.workspace + assert task_call.kwargs["job"] == test_job_step.job + assert task_call.kwargs["step"] == test_job_step.name + assert task_call.kwargs["body"].status == PlatformJobStatus.CANCELLED + assert task_call.kwargs["body"].status_details == {"message": "Job was cancelled successfully with exit code 0"} + assert task_call.kwargs["body"].error_details == {} + assert task_call.kwargs["body"].error_stack == "" def test_sync_stop_container_skips_when_not_owned_by_jobs_controller(docker_job, docker_client_mock, test_job_step): diff --git a/services/core/jobs/tests/controllers/test_reconciler.py b/services/core/jobs/tests/controllers/test_reconciler.py index 94c04e19bd..50641264cb 100644 --- a/services/core/jobs/tests/controllers/test_reconciler.py +++ b/services/core/jobs/tests/controllers/test_reconciler.py @@ -1,8 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from unittest.mock import MagicMock, call, patch +from unittest.mock import patch +import httpx +from nemo_platform_plugin.client.errors import NemoTransportError from nmp.common.jobs.schemas import PlatformJobStatus from nmp.core.jobs.api.v2.jobs.schemas import PlatformJobStepWithContext from nmp.core.jobs.controllers.backends import JobUpdate @@ -10,24 +12,19 @@ from nmp.core.jobs.controllers.backends.test import MockDockerCPUJobBackend from nmp.core.jobs.controllers.reconciler import JobReconciler +from services.core.jobs.tests.controllers.client_mocks import paginated_response + def test_job_reconciler_syncs_active_job( backend_registry: BackendRegistry, + mock_nmp_client, + mock_jobs_client, test_step_active: PlatformJobStepWithContext, ): - mock_client = MagicMock() - - # Set up the nested structure that tests expect - mock_client.jobs = MagicMock() - mock_client.jobs.list = MagicMock() - mock_client.jobs.update_status = MagicMock() - mock_client.jobs.steps = MagicMock() - mock_client.jobs.steps.list = MagicMock() - mock_client.jobs.steps.update_status = MagicMock() - job_reconciler = JobReconciler(backend_registry, mock_client) + job_reconciler = JobReconciler(backend_registry, mock_nmp_client) # Mock the jobs list response - mock_client.jobs.steps.list.return_value = [test_step_active] + mock_jobs_client.list_steps.return_value = paginated_response([test_step_active]) # Get the test backend from the registry test_backend = job_reconciler._backend_registry.get_backend(provider="cpu", profile="default") @@ -36,54 +33,70 @@ def test_job_reconciler_syncs_active_job( # Run reconciler step job_reconciler.step() - # Verify the NeMo Platform client was called with the correct filter for active steps and pending steps - # Note: MARK_INTERNAL_REQUEST_HEADERS are now set at SDK initialization, not per-request - assert mock_client.jobs.steps.list.mock_calls == [ - call( - workspace="-", - name="-", - filter={"status": ["pending", "active", "cancelling", "pausing"]}, - sort="updated_at", - ), - ] + # Verify the typed Jobs client was called with the correct deepObject filter for + # active/pending steps. The status list is encoded as a comma-joined filter value. + mock_jobs_client.list_steps.assert_called_once_with( + workspace="-", + name="-", + query_params={ + "filter[status]": "pending,active,cancelling,pausing", + "sort": "updated_at", + }, + ) # Test backend should have received one sync call for our test job assert len(test_backend.mock.sync_calls) == 1 assert test_backend.mock.sync_calls[0]["step"].id == test_step_active.id # Verify the status update was called with the correct job ID and status - mock_client.jobs.steps.update_status.assert_called_with( - "test-step", - workspace="default", - job="test-job-id", - status="completed", - error_details=None, - status_details=None, - ) + mock_jobs_client.update_job_step_status.assert_called() + update_call = mock_jobs_client.update_job_step_status.call_args + assert update_call.kwargs["name"] == "test-step" + assert update_call.kwargs["workspace"] == "default" + assert update_call.kwargs["job"] == "test-job-id" + assert update_call.kwargs["body"].status == PlatformJobStatus.COMPLETED def test_job_reconciler_logs_diagnostics_for_error_transition_in_debug_mode( backend_registry: BackendRegistry, + mock_nmp_client, + mock_jobs_client, test_step_active: PlatformJobStepWithContext, ): - mock_client = MagicMock() - mock_client.jobs = MagicMock() - mock_client.jobs.steps = MagicMock() - mock_client.jobs.steps.list.return_value = [test_step_active] + mock_jobs_client.list_steps.return_value = paginated_response([test_step_active]) - job_reconciler = JobReconciler(backend_registry, mock_client) + job_reconciler = JobReconciler(backend_registry, mock_nmp_client) test_backend = job_reconciler._backend_registry.get_backend(provider="cpu", profile="default") assert isinstance(test_backend, MockDockerCPUJobBackend) with ( - patch.object(test_backend, "sync", return_value=JobUpdate(status=PlatformJobStatus.ERROR.value)), + patch.object(test_backend, "sync", return_value=JobUpdate(status=PlatformJobStatus.ERROR)), patch("nmp.core.jobs.controllers.reconciler.log_job_diagnostics_if_debug") as log_diagnostics, ): job_reconciler.step() log_diagnostics.assert_called_once_with( - mock_client, + mock_nmp_client, test_step_active, logger=job_reconciler._logger, context="step transitioned to error during reconciliation", ) + + +def test_job_reconciler_marks_itself_unhealthy_after_transport_failure( + backend_registry: BackendRegistry, + mock_nmp_client, + mock_jobs_client, +): + job_reconciler = JobReconciler(backend_registry, mock_nmp_client) + mock_jobs_client.list_steps.return_value = paginated_response([]) + job_reconciler.step() + assert job_reconciler.is_healthy + + request = httpx.Request("GET", "http://localhost/apis/jobs/v2/workspaces/-/jobs/-/steps") + mock_jobs_client.list_steps.side_effect = NemoTransportError( + httpx.ConnectError("Connection refused", request=request) + ) + job_reconciler.step() + + assert not job_reconciler.is_healthy diff --git a/services/core/jobs/tests/controllers/test_scheduler.py b/services/core/jobs/tests/controllers/test_scheduler.py index 889f1bb420..932da1b03b 100644 --- a/services/core/jobs/tests/controllers/test_scheduler.py +++ b/services/core/jobs/tests/controllers/test_scheduler.py @@ -4,7 +4,7 @@ from unittest.mock import patch import httpx -from nemo_platform import ConflictError +from nemo_platform_plugin.client.errors import ConflictError, NemoTransportError from nmp.common.jobs.schemas import PlatformJobStatus from nmp.core.jobs.api.v2.jobs.schemas import PlatformJobStepWithContext from nmp.core.jobs.controllers.backends.exceptions import ResourceAllocationError, SchedulingDeferred @@ -13,6 +13,17 @@ from nmp.core.jobs.controllers.scheduler import JobScheduler from pytest import fixture +from services.core.jobs.tests.controllers.client_mocks import data_response, paginated_response + + +def _conflict_error(detail: str) -> ConflictError: + """Build a client ``ConflictError`` (HTTP 409) with the given detail message.""" + request = httpx.Request( + "PATCH", "http://localhost/apis/jobs/v2/workspaces/default/jobs/test-job-id/steps/test-step/status" + ) + response = httpx.Response(409, request=request, json={"detail": detail}) + return ConflictError(response) + @fixture def job_scheduler(backend_registry: BackendRegistry, mock_nmp_client) -> JobScheduler: @@ -23,10 +34,11 @@ def test_does_schedule_job( job_scheduler: JobScheduler, backend_registry: BackendRegistry, mock_nmp_client, + mock_jobs_client, test_step_pending: PlatformJobStepWithContext, ): # Mock the jobs list response - mock_nmp_client.jobs.steps.list.return_value = [test_step_pending] + mock_jobs_client.list_steps.return_value = paginated_response([test_step_pending]) # Get the test backend from the registry backend = backend_registry.get_backend(provider="cpu", profile="default") @@ -36,13 +48,15 @@ def test_does_schedule_job( # Run scheduler step job_scheduler.step() - # Verify the NeMo Platform client was called with the correct filter - # Note: MARK_INTERNAL_REQUEST_HEADERS are now set at SDK initialization, not per-request - mock_nmp_client.jobs.steps.list.assert_called_once_with( + # Verify the typed Jobs client was called with the correct deepObject filter. + # The status list is encoded as ``filter[status]=created,resuming`` (comma form). + mock_jobs_client.list_steps.assert_called_once_with( workspace="-", name="-", - filter={"status": ["created", "resuming"]}, - sort="created_at", + query_params={ + "filter[status]": "created,resuming", + "sort": "created_at", + }, ) # Test backend should have received one schedule call for our test job @@ -54,44 +68,48 @@ def test_does_schedule_job( def test_scheduling_deferred_leaves_step_created( job_scheduler: JobScheduler, mock_nmp_client, + mock_jobs_client, test_step_pending: PlatformJobStepWithContext, ): - mock_nmp_client.jobs.steps.list.return_value = [test_step_pending] + mock_jobs_client.list_steps.return_value = paginated_response([test_step_pending]) with patch.object(job_scheduler, "schedule_step", side_effect=SchedulingDeferred("capacity full")): job_scheduler.step() - mock_nmp_client.jobs.steps.update_status.assert_not_called() + mock_jobs_client.update_job_step_status.assert_not_called() def test_resource_allocation_error_marks_step_as_error( job_scheduler: JobScheduler, mock_nmp_client, + mock_jobs_client, test_step_pending: PlatformJobStepWithContext, ): """When ResourceAllocationError is raised (e.g. no GPUs), scheduler marks step as error with error_details.""" - mock_nmp_client.jobs.steps.list.return_value = [test_step_pending] + mock_jobs_client.list_steps.return_value = paginated_response([test_step_pending]) error_message = "No GPUs available on this system. GPU jobs require a system with NVIDIA GPUs." with patch.object(job_scheduler, "schedule_step", side_effect=ResourceAllocationError(error_message)): job_scheduler.step() - mock_nmp_client.jobs.steps.update_status.assert_called_once_with( - test_step_pending.name, - workspace=test_step_pending.workspace, - job=test_step_pending.job, - status=PlatformJobStatus.ERROR, - status_details={"message": error_message}, - error_details={"message": error_message}, - ) + mock_jobs_client.update_job_step_status.assert_called_once() + call = mock_jobs_client.update_job_step_status.call_args + assert call.kwargs["name"] == test_step_pending.name + assert call.kwargs["workspace"] == test_step_pending.workspace + assert call.kwargs["job"] == test_step_pending.job + body = call.kwargs["body"] + assert body.status == PlatformJobStatus.ERROR + assert body.status_details == {"message": error_message} + assert body.error_details == {"message": error_message} def test_scheduler_logs_diagnostics_for_unexpected_schedule_error_in_debug_mode( job_scheduler: JobScheduler, mock_nmp_client, + mock_jobs_client, test_step_pending: PlatformJobStepWithContext, ): - mock_nmp_client.jobs.steps.list.return_value = [test_step_pending] + mock_jobs_client.list_steps.return_value = paginated_response([test_step_pending]) with ( patch.object(job_scheduler, "schedule_step", side_effect=RuntimeError("boom")), @@ -111,43 +129,30 @@ def test_scheduler_logs_diagnostics_for_unexpected_schedule_error_in_debug_mode( def test_scheduler_does_not_mark_step_error_when_pending_update_conflicts_with_concurrent_advance( job_scheduler: JobScheduler, mock_nmp_client, + mock_jobs_client, test_step_pending: PlatformJobStepWithContext, ): - mock_nmp_client.jobs.steps.list.return_value = [test_step_pending] + mock_jobs_client.list_steps.return_value = paginated_response([test_step_pending]) - request = httpx.Request( - "PATCH", "http://localhost/apis/jobs/v2/workspaces/default/jobs/test-job-id/steps/test-step/status" - ) - response = httpx.Response( - 409, - request=request, - json={ - "detail": ( - "Invalid status transition from PlatformJobStatus.ACTIVE to " - "PlatformJobStatus.PENDING for step test-step-id" - ) - }, - ) - conflict = ConflictError( - "Error code: 409 - {'detail': 'Invalid status transition from PlatformJobStatus.ACTIVE " - "to PlatformJobStatus.PENDING for step test-step-id'}", - response=response, - body=response.json(), + conflict = _conflict_error( + "Invalid status transition from PlatformJobStatus.ACTIVE to PlatformJobStatus.PENDING for step test-step-id" ) active_step = test_step_pending.model_copy(update={"status": PlatformJobStatus.ACTIVE}) - mock_nmp_client.jobs.steps.update_status.side_effect = [conflict] - mock_nmp_client.jobs.steps.retrieve.return_value = active_step + mock_jobs_client.update_job_step_status.side_effect = [conflict] + get_step_resp = active_step + mock_jobs_client.get_job_step.return_value.data.return_value = get_step_resp job_scheduler.step() - mock_nmp_client.jobs.steps.update_status.assert_called_once_with( - test_step_pending.name, - workspace=test_step_pending.workspace, - job=test_step_pending.job, - status=PlatformJobStatus.PENDING, - ) - mock_nmp_client.jobs.steps.retrieve.assert_called_once_with( - test_step_pending.name, + mock_jobs_client.update_job_step_status.assert_called_once() + update_call = mock_jobs_client.update_job_step_status.call_args + assert update_call.kwargs["name"] == test_step_pending.name + assert update_call.kwargs["workspace"] == test_step_pending.workspace + assert update_call.kwargs["job"] == test_step_pending.job + assert update_call.kwargs["body"].status == PlatformJobStatus.PENDING + + mock_jobs_client.get_job_step.assert_called_once_with( + name=test_step_pending.name, workspace=test_step_pending.workspace, job=test_step_pending.job, ) @@ -156,41 +161,46 @@ def test_scheduler_does_not_mark_step_error_when_pending_update_conflicts_with_c def test_scheduler_does_not_ignore_pending_update_conflict_when_step_remains_resuming( job_scheduler: JobScheduler, mock_nmp_client, + mock_jobs_client, test_step_pending: PlatformJobStepWithContext, ): resuming_step = test_step_pending.model_copy(update={"status": PlatformJobStatus.RESUMING}) - mock_nmp_client.jobs.steps.list.return_value = [resuming_step] + mock_jobs_client.list_steps.return_value = paginated_response([resuming_step]) - request = httpx.Request( - "PATCH", "http://localhost/apis/jobs/v2/workspaces/default/jobs/test-job-id/steps/test-step/status" + conflict = _conflict_error( + "Invalid status transition from PlatformJobStatus.RESUMING to PlatformJobStatus.PENDING for step test-step-id" ) - response = httpx.Response( - 409, - request=request, - json={ - "detail": ( - "Invalid status transition from PlatformJobStatus.RESUMING to " - "PlatformJobStatus.PENDING for step test-step-id" - ) - }, - ) - conflict = ConflictError( - "Error code: 409 - {'detail': 'Invalid status transition from PlatformJobStatus.RESUMING " - "to PlatformJobStatus.PENDING for step test-step-id'}", - response=response, - body=response.json(), - ) - mock_nmp_client.jobs.steps.update_status.side_effect = [conflict, None] - mock_nmp_client.jobs.steps.retrieve.return_value = resuming_step + # First call (CREATED->PENDING) conflicts; the second call (marking ERROR) succeeds + # and its response is chained with ``.data()`` by the scheduler. + mock_jobs_client.update_job_step_status.side_effect = [conflict, data_response(None)] + mock_jobs_client.get_job_step.return_value.data.return_value = resuming_step job_scheduler.step() - assert mock_nmp_client.jobs.steps.update_status.call_count == 2 - error_call = mock_nmp_client.jobs.steps.update_status.call_args_list[1] - assert error_call.kwargs["status"] == PlatformJobStatus.ERROR.value - assert "409" in error_call.kwargs["status_details"]["message"] - mock_nmp_client.jobs.steps.retrieve.assert_called_once_with( - resuming_step.name, + assert mock_jobs_client.update_job_step_status.call_count == 2 + error_call = mock_jobs_client.update_job_step_status.call_args_list[1] + error_body = error_call.kwargs["body"] + assert error_body.status == PlatformJobStatus.ERROR + assert "409" in error_body.status_details["message"] + mock_jobs_client.get_job_step.assert_called_once_with( + name=resuming_step.name, workspace=resuming_step.workspace, job=resuming_step.job, ) + + +def test_scheduler_marks_itself_unhealthy_after_transport_failure( + job_scheduler: JobScheduler, + mock_jobs_client, +): + mock_jobs_client.list_steps.return_value = paginated_response([]) + job_scheduler.step() + assert job_scheduler.is_healthy + + request = httpx.Request("GET", "http://localhost/apis/jobs/v2/workspaces/-/jobs/-/steps") + mock_jobs_client.list_steps.side_effect = NemoTransportError( + httpx.ConnectError("Connection refused", request=request) + ) + job_scheduler.step() + + assert not job_scheduler.is_healthy diff --git a/services/core/jobs/tests/controllers/test_subprocess_backend.py b/services/core/jobs/tests/controllers/test_subprocess_backend.py index db7f945b84..a32d8d220b 100644 --- a/services/core/jobs/tests/controllers/test_subprocess_backend.py +++ b/services/core/jobs/tests/controllers/test_subprocess_backend.py @@ -16,6 +16,8 @@ SubprocessProcessKey, ) +from services.core.jobs.tests.controllers.client_mocks import data_response + def _subprocess_backend(mock_nmp_client, tmp_path, mock_platform_config) -> SubprocessJobBackend: with patch("nmp.core.jobs.controllers.backends.subprocess.get_platform_config", return_value=mock_platform_config): @@ -56,7 +58,7 @@ def _schedule_without_otel_export(backend: SubprocessJobBackend, step): def test_schedule_starts_process_and_stages_environment( - mock_nmp_client, tmp_path, mock_platform_config, test_step_pending + mock_nmp_client, mock_jobs_client, tmp_path, mock_platform_config, test_step_pending ): backend = _subprocess_backend(mock_nmp_client, tmp_path, mock_platform_config) step = _step_with_command(test_step_pending, ["/bin/sh", "-c", "printf 'hello local\\n'"]) @@ -70,7 +72,7 @@ def test_schedule_starts_process_and_stages_environment( assert metadata.process.wait(timeout=5) == 0 assert metadata.work_dir.is_relative_to(tmp_path) assert metadata.persistent_dir.is_relative_to(tmp_path) - mock_nmp_client.jobs.tasks.create_or_update.assert_called() + mock_jobs_client.update_job_step_task.assert_called() def test_created_step_does_not_ttl_before_backend_acceptance( @@ -176,7 +178,9 @@ def test_schedule_terminates_process_when_post_popen_setup_fails( assert not any(tmp_path.iterdir()) -def test_sync_completed_closes_logs(mock_nmp_client, tmp_path, mock_platform_config, test_step_pending): +def test_sync_completed_closes_logs( + mock_nmp_client, mock_jobs_client, tmp_path, mock_platform_config, test_step_pending +): backend = _subprocess_backend(mock_nmp_client, tmp_path, mock_platform_config) step = _step_with_command(test_step_pending, ["/bin/sh", "-c", "printf 'hello logs\\n'"]) @@ -192,8 +196,8 @@ def test_sync_completed_closes_logs(mock_nmp_client, tmp_path, mock_platform_con assert update.status == PlatformJobStatus.COMPLETED assert metadata.closed_logs is True assert "hello logs" in metadata.log_path.read_text(encoding="utf-8") - last_call = mock_nmp_client.jobs.tasks.create_or_update.call_args - assert last_call.kwargs["status"] == PlatformJobStatus.COMPLETED.value + last_call = mock_jobs_client.update_job_step_task.call_args + assert last_call.kwargs["body"].status == PlatformJobStatus.COMPLETED def test_shutdown_finishes_logs(mock_nmp_client, tmp_path, mock_platform_config, test_step_pending): @@ -332,7 +336,7 @@ def test_cancelling_terminates_running_process(mock_nmp_client, tmp_path, mock_p def test_sync_uses_persisted_task_when_local_metadata_is_missing( - mock_nmp_client, tmp_path, mock_platform_config, test_step_active + mock_nmp_client, mock_jobs_client, tmp_path, mock_platform_config, test_step_active ): backend = _subprocess_backend(mock_nmp_client, tmp_path, mock_platform_config) step = _step_with_command(test_step_active, ["/bin/sh", "-c", "sleep 10"]) @@ -344,16 +348,18 @@ def test_sync_uses_persisted_task_when_local_metadata_is_missing( metadata.process.terminate() assert metadata.process.wait(timeout=5) is not None backend._process_registry.pop(key) - mock_nmp_client.jobs.tasks.list.return_value = SimpleNamespace( - data=[ - SimpleNamespace( - status=PlatformJobStatus.ACTIVE.value, - status_details={"message": "Job is running"}, - error_details={}, - created_at=step.created_at, - updated_at=step.updated_at, - ) - ] + mock_jobs_client.list_job_step_tasks.return_value = data_response( + SimpleNamespace( + data=[ + SimpleNamespace( + status=PlatformJobStatus.ACTIVE.value, + status_details={"message": "Job is running"}, + error_details={}, + created_at=step.created_at, + updated_at=step.updated_at, + ) + ] + ) ) update = backend.sync(step) @@ -364,27 +370,29 @@ def test_sync_uses_persisted_task_when_local_metadata_is_missing( def test_get_task_fallback_update_handles_missing_task_timestamps( - mock_nmp_client, tmp_path, mock_platform_config, test_step_active + mock_nmp_client, mock_jobs_client, tmp_path, mock_platform_config, test_step_active ): backend = _subprocess_backend(mock_nmp_client, tmp_path, mock_platform_config) step = _step_with_command(test_step_active, ["/bin/sh", "-c", "true"]) - mock_nmp_client.jobs.tasks.list.return_value = SimpleNamespace( - data=[ - SimpleNamespace( - status=PlatformJobStatus.PENDING.value, - status_details={"message": "missing timestamps"}, - error_details={}, - created_at=None, - updated_at=None, - ), - SimpleNamespace( - status=PlatformJobStatus.ACTIVE.value, - status_details={"message": "latest"}, - error_details={}, - created_at=step.created_at, - updated_at=step.updated_at, - ), - ] + mock_jobs_client.list_job_step_tasks.return_value = data_response( + SimpleNamespace( + data=[ + SimpleNamespace( + status=PlatformJobStatus.PENDING.value, + status_details={"message": "missing timestamps"}, + error_details={}, + created_at=None, + updated_at=None, + ), + SimpleNamespace( + status=PlatformJobStatus.ACTIVE.value, + status_details={"message": "latest"}, + error_details={}, + created_at=step.created_at, + updated_at=step.updated_at, + ), + ] + ) ) update = backend._get_task_fallback_update(step) @@ -396,11 +404,11 @@ def test_get_task_fallback_update_handles_missing_task_timestamps( def test_sync_keeps_recent_pending_step_pending_when_local_metadata_is_missing( - mock_nmp_client, tmp_path, mock_platform_config, test_step_pending + mock_nmp_client, mock_jobs_client, tmp_path, mock_platform_config, test_step_pending ): backend = _subprocess_backend(mock_nmp_client, tmp_path, mock_platform_config) step = _step_with_command(test_step_pending, ["/bin/sh", "-c", "true"]) - mock_nmp_client.jobs.tasks.list.return_value = SimpleNamespace(data=[]) + mock_jobs_client.list_job_step_tasks.return_value = data_response(SimpleNamespace(data=[])) update = backend.sync(step) diff --git a/services/core/jobs/tests/test_jobs_client.py b/services/core/jobs/tests/test_jobs_client.py new file mode 100644 index 0000000000..e327ada804 --- /dev/null +++ b/services/core/jobs/tests/test_jobs_client.py @@ -0,0 +1,297 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""End-to-end tests that drive the typed ``JobsClient`` against the real Jobs +service routes (in-memory ASGI app). + +Unlike ``tests/jobs/test_endpoints.py`` (which only asserts ``PreparedRequest`` +shape) these exercise ``send()`` all the way through path resolution, HTTP, +and response parsing — the layer where response-type bugs actually surface. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, call, patch + +import pytest +from httpx import AsyncClient +from nemo_platform_plugin.jobs.client import AsyncJobsClient +from nemo_platform_plugin.jobs.schemas import ( + FileStorageType, + PlatformJobLog, + PlatformJobLogPage, + PlatformJobResultCreateRequest, + PlatformJobStatus, +) +from nemo_platform_plugin.jobs.types import ( + CreatePlatformJobRequest, + JobStatusDetailsUpdate, + PlatformJobStatusUpdateRequest, + PlatformJobTaskUpdate, +) +from nmp.common.jobs.file_manager import TmpDirPath +from nmp.common.jobs.log_client import dep_job_logs_client + + +@pytest.fixture +def jobs_client(test_client: AsyncClient) -> AsyncJobsClient: + """A typed AsyncJobsClient bound to the in-memory Jobs app. + + Mirrors how ``test_sdk`` builds the Stainless SDK, but returns the new + typed client so responses flow through ``NemoClient.send()``. + """ + return AsyncJobsClient(base_url=str(test_client.base_url), http_client=test_client) + + +async def _create_job( + jobs_client: AsyncJobsClient, + request: CreatePlatformJobRequest, + name: str, +): + body = request.model_copy(update={"name": name}) + return (await jobs_client.create_job(workspace="default", body=body)).data() + + +@pytest.mark.asyncio +async def test_get_execution_profiles_parses_response(jobs_client: AsyncJobsClient, test_client: AsyncClient): + """``get_execution_profiles`` parses the route's JSON array response.""" + # Sanity: the raw route really does return a JSON list (server side is fine). + raw = await test_client.get("/apis/jobs/v2/execution-profiles") + assert raw.status_code == 200 + assert isinstance(raw.json(), list) + + # The actual regression: the typed client must not crash parsing it. + resp = await jobs_client.get_execution_profiles() + profiles = resp.data() + assert isinstance(profiles, list) + + +async def _create_hello_world_job(test_client: AsyncClient, name: str = "e2e-client-job") -> None: + """Create a job via the hello-world factory route (service-specific body).""" + resp = await test_client.post( + "/apis/jobs/v2/workspaces/default/hello-world/jobs", + json={ + "name": name, + "description": "typed-client e2e", + "spec": {"config": {"key": "Value"}, "target": "str"}, + "ownership": {"user": "u", "service": "s"}, + }, + ) + assert resp.status_code == 201, f"create failed: {resp.status_code} {resp.text}" + + +@pytest.mark.asyncio +async def test_list_jobs_round_trips_through_client(jobs_client: AsyncJobsClient, test_client: AsyncClient): + """``list_jobs`` must page + parse real ``PlatformJobResponse`` items.""" + await _create_hello_world_job(test_client, name="list-me") + + page = (await jobs_client.list_jobs(workspace="default")).page() + assert page.metadata["total_results"] is not None and page.metadata["total_results"] >= 1 + names = [j.name for j in page.items] + assert "list-me" in names + # items are the plugin DTO, fully parsed + job = next(j for j in page.items if j.name == "list-me") + assert job.workspace == "default" + assert job.status is not None + + +@pytest.mark.asyncio +async def test_get_job_and_status_round_trip(jobs_client: AsyncJobsClient, test_client: AsyncClient): + """``get_job`` and ``get_job_status`` must parse their real responses.""" + await _create_hello_world_job(test_client, name="get-me") + + job = (await jobs_client.get_job(name="get-me", workspace="default")).data() + assert job.name == "get-me" + assert job.fileset # non-empty + + status = (await jobs_client.get_job_status(name="get-me", workspace="default")).data() + assert status.status is not None + + +@pytest.mark.asyncio +async def test_job_lifecycle_methods_round_trip( + jobs_client: AsyncJobsClient, + sample_platform_job_request: CreatePlatformJobRequest, +): + paused_job = await _create_job(jobs_client, sample_platform_job_request, "typed-lifecycle") + active_step = ( + await jobs_client.update_job_step_status( + workspace="default", + job=paused_job.name, + name="basic", + body=PlatformJobStatusUpdateRequest(status=PlatformJobStatus.ACTIVE), + ) + ).data() + assert active_step.status == PlatformJobStatus.ACTIVE + + pausing = (await jobs_client.pause_job(workspace="default", name=paused_job.name)).data() + assert pausing.status == PlatformJobStatus.PAUSING + await jobs_client.update_job_step_status( + workspace="default", + job=paused_job.name, + name="basic", + body=PlatformJobStatusUpdateRequest(status=PlatformJobStatus.PAUSED), + ) + resuming = (await jobs_client.resume_job(workspace="default", name=paused_job.name)).data() + assert resuming.status == PlatformJobStatus.RESUMING + + cancelled_job = await _create_job(jobs_client, sample_platform_job_request, "typed-cancel") + cancelled = (await jobs_client.cancel_job(workspace="default", name=cancelled_job.name)).data() + assert cancelled.status == PlatformJobStatus.CANCELLED + + deleted_job = await _create_job(jobs_client, sample_platform_job_request, "typed-delete") + deleted = await jobs_client.delete_job(workspace="default", name=deleted_job.name) + assert deleted.http_response.status_code == 204 + + +@pytest.mark.asyncio +async def test_status_steps_and_tasks_round_trip( + jobs_client: AsyncJobsClient, + sample_platform_job_request: CreatePlatformJobRequest, +): + job = await _create_job(jobs_client, sample_platform_job_request, "typed-state") + status_update = await jobs_client.update_job_status_details( + workspace="default", + name=job.name, + body=JobStatusDetailsUpdate(root={"progress": 25}), + ) + assert status_update.http_response.status_code == 200 + updated_job = (await jobs_client.get_job(workspace="default", name=job.name)).data() + assert updated_job.status_details == {"progress": 25} + + step_page = (await jobs_client.list_steps(workspace="default", name=job.name)).page() + assert [step.name for step in step_page.items] == ["basic"] + step = (await jobs_client.get_job_step(workspace="default", job=job.name, name="basic")).data() + assert step.name == "basic" + + task = ( + await jobs_client.update_job_step_task( + workspace="default", + job=job.name, + step="basic", + name="task-1", + body=PlatformJobTaskUpdate( + status=PlatformJobStatus.ACTIVE, + status_details={"message": "running"}, + ), + ) + ).data() + assert task.status == PlatformJobStatus.ACTIVE + tasks = (await jobs_client.list_job_step_tasks(workspace="default", job=job.name, name="basic")).data() + assert [item.name for item in tasks.data] == ["task-1"] + fetched_task = ( + await jobs_client.get_job_step_task( + workspace="default", + job=job.name, + step="basic", + name="task-1", + ) + ).data() + assert fetched_task.status_details == {"message": "running"} + + +@pytest.mark.asyncio +async def test_logs_round_trip( + jobs_client: AsyncJobsClient, + test_client: AsyncClient, + sample_platform_job_request: CreatePlatformJobRequest, +): + job = await _create_job(jobs_client, sample_platform_job_request, "typed-logs") + timestamp = datetime(2026, 1, 1, tzinfo=timezone.utc) + logs_client = AsyncMock() + logs_client.query_logs.side_effect = [ + PlatformJobLogPage( + data=[ + PlatformJobLog( + timestamp=timestamp, + job=job.name, + job_step="basic", + job_task="task-1", + message="hello", + ) + ], + total=2, + next_page="cursor-2", + prev_page=None, + ), + PlatformJobLogPage( + data=[ + PlatformJobLog( + timestamp=timestamp, + job=job.name, + job_step="basic", + job_task="task-1", + message="world", + ) + ], + total=2, + next_page=None, + prev_page="cursor-1", + ), + ] + app = test_client._transport.app # type: ignore[attr-defined] + app.dependency_overrides[dep_job_logs_client] = lambda: logs_client + try: + response = await jobs_client.list_job_logs( + workspace="default", + name=job.name, + query_params={"limit": 5, "step_id": "basic", "task_id": "task-1"}, + ) + page = response.page() + logs = [log async for log in response.items()] + finally: + app.dependency_overrides.pop(dep_job_logs_client) + + assert [log.message for log in page.items] == ["hello"] + assert page.metadata == {"total": 2, "next_page": "cursor-2", "prev_page": None} + assert [log.message for log in logs] == ["hello", "world"] + filters = { + "job": job.name, + "job_attempt": job.attempt_id, + "job_step": "basic", + "job_task": "task-1", + } + assert logs_client.query_logs.await_args_list == [ + call(job.fileset, workspace="default", filters=filters, page_size=5, page_cursor=None), + call(job.fileset, workspace="default", filters=filters, page_size=5, page_cursor="cursor-2"), + ] + + +@pytest.mark.asyncio +async def test_result_methods_round_trip( + jobs_client: AsyncJobsClient, + sample_platform_job_request: CreatePlatformJobRequest, + tmp_path, +): + job = await _create_job(jobs_client, sample_platform_job_request, "typed-results") + result = ( + await jobs_client.create_job_result( + workspace="default", + job=job.name, + name="output", + body=PlatformJobResultCreateRequest( + artifact_url="default/test-fileset#output.txt", + artifact_storage_type=FileStorageType.FILESET, + ), + ) + ).data() + assert result.name == "output" + + listed = (await jobs_client.list_job_results(workspace="default", name=job.name)).data() + assert [item.name for item in listed.data] == ["output"] + fetched = (await jobs_client.get_job_result(workspace="default", job=job.name, name="output")).data() + assert fetched.artifact_url == "default/test-fileset#output.txt" + + result_dir = tmp_path / "typed-result" + result_dir.mkdir() + result_path = result_dir / "output.txt" + result_path.write_bytes(b"typed result") + downloaded = TmpDirPath(path=result_path, tmp_dir=result_dir) + with patch( + "nmp.core.jobs.api.v2.jobs.endpoints.download_from_result_info", + new=AsyncMock(return_value=("output.txt", downloaded)), + ): + content = await (await jobs_client.download_job_result(workspace="default", job=job.name, name="output")).read() + + assert content == b"typed result" diff --git a/services/core/models/src/nmp/core/models/api/v2/models.py b/services/core/models/src/nmp/core/models/api/v2/models.py index 4a8d80fdef..4b7a449b7d 100644 --- a/services/core/models/src/nmp/core/models/api/v2/models.py +++ b/services/core/models/src/nmp/core/models/api/v2/models.py @@ -4,7 +4,9 @@ import logging from fastapi import APIRouter, Depends, HTTPException, Query, status -from nemo_platform import APIError, AsyncNeMoPlatform +from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoClientError from nemo_platform_plugin.jobs.api_factory import ( ContainerSpec, CPUExecutionProviderSpec, @@ -15,7 +17,10 @@ ResourcesRequestsSpec, ResourcesSpec, ) +from nemo_platform_plugin.jobs.client import AsyncJobsClient from nemo_platform_plugin.jobs.image import get_qualified_image +from nemo_platform_plugin.jobs.spec import PlatformJobSpec as PlatformJobSpecModel +from nemo_platform_plugin.jobs.types import CreatePlatformJobRequest from nmp.common.api.common import Page from nmp.common.api.parsed_filter import ParsedFilter, make_filter_dep from nmp.common.api.utils import generate_openapi_extra_params @@ -304,17 +309,25 @@ async def start_update_model_spec_job(model_entity: ModelEntity): ] ) try: - job_resp = await sdk.jobs.create( - source="models-system", - workspace=model_entity.workspace, - platform_spec=task_spec, - spec={}, - description=f"Model Spec Analyzer for model {model_entity.workspace}/{model_entity.name}", - ownership=model_entity.ownership, - project=model_entity.project, - ) - logger.info(f"Job Created - {job_resp}") - except APIError as err: + jobs = client_from_platform(sdk, AsyncJobsClient) + job_resp = ( + await jobs.create_job( + workspace=model_entity.workspace, + body=CreatePlatformJobRequest( + source="models-system", + # ``task_spec`` is built from the api_factory ``*Param`` TypedDict + # aliases (see AIRCORE-922); validate it into the plugin pydantic + # ``PlatformJobSpec`` the request model expects. + platform_spec=PlatformJobSpecModel.model_validate(task_spec), + spec={}, + description=f"Model Spec Analyzer for model {model_entity.workspace}/{model_entity.name}", + ownership=model_entity.ownership, + project=model_entity.project, + ), + ) + ).data() + logger.info("Job Created - %s", job_resp.name) + except NemoClientError as err: logger.warning(f"Failed to create model spec job. {err}") diff --git a/services/core/models/tests/unit/api/test_models_api.py b/services/core/models/tests/unit/api/test_models_api.py index 4cc12b11f8..3b0a9c2a03 100644 --- a/services/core/models/tests/unit/api/test_models_api.py +++ b/services/core/models/tests/unit/api/test_models_api.py @@ -6,15 +6,17 @@ from datetime import datetime from unittest.mock import AsyncMock, MagicMock, Mock, patch +import httpx import pytest from fastapi import FastAPI from fastapi.testclient import TestClient +from nemo_platform_plugin.client.errors import NemoTransportError from nmp.common.api.common import Page, PaginationData from nmp.common.auth import AuthClient, Principal, get_auth_client from nmp.common.entities.client import EntityValidationError from nmp.core.models.api.service.adapter_entity_service import AdapterEntityService from nmp.core.models.api.service.model_entity_service import ModelEntityService -from nmp.core.models.api.v2.models import router +from nmp.core.models.api.v2.models import router, start_update_model_spec_job from nmp.core.models.schemas import ModelEntity @@ -408,6 +410,23 @@ def test_create_model_entity_validation_error_returns_422(client, mock_model_ent assert "name must match pattern" in response.json()["detail"] +@pytest.mark.asyncio +async def test_model_spec_job_transport_failure_does_not_fail_persisted_model(sample_model_entity): + request = httpx.Request("POST", "http://test/apis/jobs/v2/workspaces/nvidia/jobs") + jobs = MagicMock() + jobs.create_job = AsyncMock( + side_effect=NemoTransportError(httpx.ConnectError("Connection refused", request=request)) + ) + + with ( + patch("nmp.core.models.api.v2.models.get_async_platform_sdk"), + patch("nmp.core.models.api.v2.models.client_from_platform", return_value=jobs), + ): + await start_update_model_spec_job(sample_model_entity) + + jobs.create_job.assert_awaited_once() + + def test_update_model_entity_validation_error_returns_422(client, mock_model_entity_service, sample_model_entity): """Test that entity store validation errors during model update return 422.""" mock_model_entity_service.entity_client.get.return_value = Mock() diff --git a/services/unsloth/tests/test_progress.py b/services/unsloth/tests/test_progress.py index 3ccd80c6b3..430dff10e6 100644 --- a/services/unsloth/tests/test_progress.py +++ b/services/unsloth/tests/test_progress.py @@ -21,16 +21,25 @@ def test_progress_reporter_calls_sdk_create_or_update() -> None: config_path=Path("/tmp/job/config.json"), ) mock_sdk = MagicMock() + mock_jobs = MagicMock() - with patch("nmp.customization_common.training.progress.get_task_sdk", return_value=mock_sdk): + with ( + patch("nmp.customization_common.training.progress.get_task_sdk", return_value=mock_sdk), + patch( + "nmp.customization_common.training.progress.client_from_platform", + return_value=mock_jobs, + ), + ): reporter = JobsServiceProgressReporter(ctx) reporter.report_running(phase="training", step=1, train_loss=2.5, backend="unsloth") - mock_sdk.jobs.tasks.create_or_update.assert_called_once() - call_kwargs = mock_sdk.jobs.tasks.create_or_update.call_args.kwargs + mock_jobs.update_job_step_task.assert_called_once() + call_kwargs = mock_jobs.update_job_step_task.call_args.kwargs assert call_kwargs["name"] == ctx.normalized_task assert call_kwargs["workspace"] == ctx.workspace assert call_kwargs["job"] == ctx.job_id assert call_kwargs["step"] == ctx.step - assert call_kwargs["status_details"]["train_loss"] == 2.5 - assert call_kwargs["status_details"]["backend"] == "unsloth" + # status_details is now carried on the PlatformJobTaskUpdate body object. + body = call_kwargs["body"] + assert body.status_details["train_loss"] == 2.5 + assert body.status_details["backend"] == "unsloth"