From b7a5e749c4801db6732d8442656ef5e82bdc7d29 Mon Sep 17 00:00:00 2001 From: "Paul A. Parkanzky" Date: Fri, 31 Jul 2026 16:43:25 -0400 Subject: [PATCH 1/2] feat(auditor): add blocking submit option to auditor SDK Signed-off-by: Paul A. Parkanzky --- docs/auditor/sdk-resources.mdx | 73 +++- e2e/auditor/test_audit_job.py | 6 +- plugins/nemo-auditor/README.md | 14 +- plugins/nemo-auditor/src/nemo_auditor/sdk.py | 19 +- .../sdk_resources/job_resources.py | 351 ++++++++++++++++++ .../nemo-auditor/tests/test_sdk_resources.py | 239 +++++++++++- 6 files changed, 685 insertions(+), 17 deletions(-) create mode 100644 plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py diff --git a/docs/auditor/sdk-resources.mdx b/docs/auditor/sdk-resources.mdx index 9779a254f1..350a08efa8 100644 --- a/docs/auditor/sdk-resources.mdx +++ b/docs/auditor/sdk-resources.mdx @@ -33,6 +33,9 @@ auditor = client.auditor # AuditorPluginResource | `plugin_status()` | Returns auditor plugin health information from the service. | `dict[str, object]` | | `configs` | Sub-resource for `AuditConfig` CRUD operations. | `_ConfigResource` | | `targets` | Sub-resource for `AuditTarget` CRUD operations. | `_TargetResource` | +| `submit()` | Submits a K8s audit job and returns a handle for polling and artifact download. | `AuditorJobResource` | +| `list_jobs(workspace, page, page_size)` | Lists submitted audit jobs in the workspace. | `dict` | +| `get_job(job_name, workspace)` | Fetches a single audit job by name. | `dict` | | `run()` | Runs one audit locally, in-process, against a configured target. | `dict` | ### `configs` sub-resource @@ -59,6 +62,59 @@ Five CRUD methods for `AuditTarget` entities. The full field reference is in [Ta | `update(*, workspace, name, type, model, options=None, description=None)` | Replaces a target's fields. | `AuditTarget` | | `delete(*, workspace, name)` | Deletes a target. | `None` | +### `submit()` arguments + +`submit()` posts an audit job to the K8s executor and returns an `AuditorJobResource` handle. +Call `.wait_until_done()` on the handle to block until the job completes, then `.download_artifacts()` to fetch the garak reports. + +| Argument | Type | Required | Description | +|----------|------|----------|-------------| +| `config` | `AuditConfig \| str` | Yes | An inline `AuditConfig` instance or a name string referencing one in the entity store. Bare names resolve against `workspace`; qualified names such as `"prod/quick-scan"` override the workspace. | +| `target` | `AuditTarget \| str` | Yes | An inline `AuditTarget` instance or a name string, with the same resolution rules as `config`. | +| `workspace` | `str \| None` | No | Workspace to submit the job into. Defaults to `"default"`. | +| `max_probe_retries` | `int` | No | Number of times to retry a failing garak probe before marking it as failed. Defaults to `0`. | +| `fail_job_on_retries_exhausted` | `bool` | No | When `True` (the default), the job fails if any probe exhausts its retries. Set to `False` to treat retry-exhausted probes as warnings. | + +### `AuditorJobResource` + +The object returned by `submit()`. Use it to poll status, stream logs, and download artifacts. + +| Method | Description | Returns | +|--------|-------------|---------| +| `name` | The unique job name assigned by the platform. | `str` | +| `get_job()` | Fetches the full job dict (name, status, workspace, …). | `dict[str, object]` | +| `get_job_status()` | Fetches only the current platform status string. | `PlatformJobStatus \| None` | +| `check_if_complete(raise_if_not_complete=False)` | Returns `True` if the job is `completed`. Raises `RuntimeError` when `raise_if_not_complete=True` and the job is not done. | `bool` | +| `wait_until_done()` | Blocks until the job reaches a terminal status. Streams log entries from the audit task while polling. Raises `RuntimeError` on a terminal failure. | `None` | +| `get_logs()` | Pages through all structured log entries produced by the audit task. | `list[dict[str, str]]` | +| `download_artifacts(path=None)` | Downloads and extracts the garak report tarball. Raises `RuntimeError` if the job has not completed. `path` overrides the output directory (defaults to a directory named after the job). | `Path` | + +### Submit and wait for an audit job + +```python +# Submit the job using persisted entity name strings. +job = auditor.submit( + config="quick-scan", + target="llama-31-8b", + workspace="default", +) +print(f"Job submitted: {job.name}") + +# Block until garak finishes (raises RuntimeError on failure). +job.wait_until_done() + +# Download the garak reports to ./my-reports//. +artifacts_dir = job.download_artifacts(path="./my-reports") +print(f"Reports saved to: {artifacts_dir}") +``` + +You can also check status without blocking: + +```python +if not job.check_if_complete(): + print(f"Still running: {job.get_job_status()}") +``` + ### `run()` arguments `run()` invokes [garak](https://github.com/NVIDIA/garak) locally, in-process, against a configured target. @@ -162,15 +218,30 @@ auditor = client.auditor # AsyncAuditorPluginResource | `plugin_status()` | Returns auditor plugin health information from the service. | `dict[str, object]` | | `configs` | Sub-resource for `AuditConfig` CRUD operations. | `_AsyncConfigResource` | | `targets` | Sub-resource for `AuditTarget` CRUD operations. | `_AsyncTargetResource` | +| `submit()` | Submits a K8s audit job and returns an async handle for polling and artifact download. | `AsyncAuditorJobResource` | +| `list_jobs(workspace, page, page_size)` | Lists submitted audit jobs in the workspace. | `dict` | +| `get_job(job_name, workspace)` | Fetches a single audit job by name. | `dict` | | `run()` | Runs one audit locally, in-process, against a configured target. | `dict` | -`AsyncAuditorPluginResource.run()` and the async `configs` / `targets` sub-resource methods accept the same arguments as their sync counterparts [above](#run-arguments). Because the local execution path is synchronous (garak runs in a subprocess), the async `run()` dispatches the scheduler call through `asyncio.to_thread` so the caller's event loop is not blocked. +`AsyncAuditorPluginResource.submit()` returns an `AsyncAuditorJobResource` with the same methods as `AuditorJobResource` [above](#auditorjobresource), all awaitable. +`AsyncAuditorPluginResource.run()` and the async `configs` / `targets` sub-resource methods accept the same arguments as their sync counterparts. Because the local execution path is synchronous (garak runs in a subprocess), the async `run()` dispatches the scheduler call through `asyncio.to_thread` so the caller's event loop is not blocked. ```python import asyncio async def main() -> None: + # Submit and wait. + job = await auditor.submit( + config="quick-scan", + target="llama-31-8b", + workspace="default", + ) + await job.wait_until_done() + artifacts_dir = await job.download_artifacts() + print(f"Reports saved to: {artifacts_dir}") + + # Or run locally (no jobs-service submission). result = await auditor.run( config="quick-scan", target="llama-31-8b", diff --git a/e2e/auditor/test_audit_job.py b/e2e/auditor/test_audit_job.py index a1ef35d9e7..b375842c71 100644 --- a/e2e/auditor/test_audit_job.py +++ b/e2e/auditor/test_audit_job.py @@ -181,7 +181,7 @@ def test_audit_job_submit_blank_probe( } job = sdk.auditor.submit(config=config, target=target, workspace=audit_workspace) - job_name = job["name"] + job_name = job.name try: final_status = _wait_for_audit_job(sdk, job_name, audit_workspace) assert final_status == "completed", ( @@ -205,7 +205,7 @@ def test_audit_job_submit_with_entity_refs( target=f"{audit_workspace}/{audit_target_name}", workspace=audit_workspace, ) - job_name = job["name"] + job_name = job.name try: final_status = _wait_for_audit_job(sdk, job_name, audit_workspace) assert final_status == "completed", ( @@ -227,7 +227,7 @@ def test_audit_job_appears_in_list( target=f"{audit_workspace}/{audit_target_name}", workspace=audit_workspace, ) - job_name = job["name"] + job_name = job.name try: jobs = sdk.auditor.list_jobs(workspace=audit_workspace) job_names = [j["name"] for j in jobs.get("data", [])] diff --git a/plugins/nemo-auditor/README.md b/plugins/nemo-auditor/README.md index 98051636bf..322e82b693 100644 --- a/plugins/nemo-auditor/README.md +++ b/plugins/nemo-auditor/README.md @@ -81,7 +81,18 @@ tgt = client.auditor.targets.create( options={"uri": "http://localhost:9000/v1"}, ) -# Run an audit locally (no jobs-service submission) using the persisted entities +# Submit a K8s audit job and wait for it to finish. +job = client.auditor.submit( + config="quick-scan", + target="llama-31-8b", + workspace="default", +) +print(f"Job submitted: {job.name}") +job.wait_until_done() # blocks; streams logs while polling +artifacts_dir = job.download_artifacts() # extracts garak reports to .// +print(f"Reports: {artifacts_dir}") + +# Or run an audit locally (no jobs-service submission). result = client.auditor.run( config="quick-scan", # workspace-qualified name strings ("ws/name") also work target="llama-31-8b", @@ -92,6 +103,7 @@ for name, ref in result["results"].items(): print(name, ref["artifact_url"]) ``` +`submit()` posts the job to the K8s executor and returns an `AuditorJobResource` handle. `run()` shells out to a pre-installed garak interpreter (default `~/.auditor/.venv/bin/python`, override via `$NEMO_AUDITOR_GARAK_PYTHON`) and registers the resulting JSONL / HTML / hitlog reports as job results diff --git a/plugins/nemo-auditor/src/nemo_auditor/sdk.py b/plugins/nemo-auditor/src/nemo_auditor/sdk.py index 6ad3be6a36..0fa75c72cd 100644 --- a/plugins/nemo-auditor/src/nemo_auditor/sdk.py +++ b/plugins/nemo-auditor/src/nemo_auditor/sdk.py @@ -10,7 +10,9 @@ - ``client.auditor.configs.{create,list,get,update,delete}`` — ``AuditConfig`` CRUD. - ``client.auditor.targets.{create,list,get,update,delete}`` — ``AuditTarget`` CRUD. - ``client.auditor.submit(config=..., target=..., workspace=...)`` — submit a K8s - audit job through the plugin's job endpoint and return the raw job dict. + audit job and return an :class:`~nemo_auditor.sdk_resources.job_resources.AuditorJobResource` + handle. Call ``.wait_until_done()`` on the handle to block until the job completes, + then ``.download_artifacts()`` to fetch the garak report tarball. - ``client.auditor.list_jobs(workspace=...)`` — list submitted audit jobs. - ``client.auditor.get_job(job_name, workspace=...)`` — fetch a single audit job. - ``client.auditor.run(config=..., target=..., workspace=...)`` — in-process @@ -29,6 +31,7 @@ from nemo_auditor.entities import AuditConfig, AuditTarget from nemo_auditor.jobs.audit import AuditInputSpec, AuditJob from nemo_auditor.sdk_resources.configs import _AsyncConfigResource, _ConfigResource +from nemo_auditor.sdk_resources.job_resources import AsyncAuditorJobResource, AuditorJobResource from nemo_auditor.sdk_resources.targets import _AsyncTargetResource, _TargetResource from nemo_platform import AsyncNeMoPlatform, NeMoPlatform from nemo_platform_plugin.entities import parse_qualified_name @@ -73,12 +76,12 @@ def submit( workspace: str | None = None, max_probe_retries: int = 0, fail_job_on_retries_exhausted: bool = True, - ) -> dict: + ) -> AuditorJobResource: """Submit an audit job to the K8s executor via the plugin job endpoint. - Returns the raw job dict (name, status, workspace, …). Use - ``sdk.jobs.get_status(name=result["name"], workspace=workspace)`` to poll - for completion, or pass the name to ``sdk.auditor.get_job()``. + Returns an :class:`~nemo_auditor.sdk_resources.job_resources.AuditorJobResource` + handle. Call ``.wait_until_done()`` to block until the job completes, then + ``.download_artifacts()`` to fetch the garak report tarball. """ ws = workspace or "default" spec = AuditInputSpec( @@ -92,7 +95,7 @@ def submit( json={"spec": spec.model_dump(mode="json")}, ) response.raise_for_status() - return response.json() + return AuditorJobResource(job_name=response.json()["name"], platform=self._platform, workspace=ws) def list_jobs( self, @@ -199,7 +202,7 @@ async def submit( workspace: str | None = None, max_probe_retries: int = 0, fail_job_on_retries_exhausted: bool = True, - ) -> dict: + ) -> AsyncAuditorJobResource: """Async twin of :meth:`AuditorPluginResource.submit`.""" ws = workspace or "default" spec = AuditInputSpec( @@ -213,7 +216,7 @@ async def submit( json={"spec": spec.model_dump(mode="json")}, ) response.raise_for_status() - return response.json() + return AsyncAuditorJobResource(job_name=response.json()["name"], platform=self._platform, workspace=ws) async def list_jobs( self, diff --git a/plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py b/plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py new file mode 100644 index 0000000000..af476e4a8f --- /dev/null +++ b/plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py @@ -0,0 +1,351 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Audit job resource handles for status polling, log streaming, and artifact download.""" + +from __future__ import annotations + +import asyncio +import io +import json +import logging +import tarfile +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Awaitable, Callable, TypeVar + +from nemo_platform import AsyncNeMoPlatform, NeMoPlatform +from nemo_platform.types import PlatformJobStatus +from nemo_platform_plugin.jobs.archive import safe_extract_tar +from typing_extensions import Self + +logger = logging.getLogger(__name__) + +WAIT_INTERVAL_SECONDS = 1 +MAX_CONSECUTIVE_POLL_ERRORS = 5 +ARTIFACTS_RESULT_NAME = "artifacts" +TERMINAL_INCOMPLETE_STATUSES = {"cancelled", "cancelling", "error"} + +T = TypeVar("T") + + +def _pause(seconds: float) -> None: + time.sleep(seconds) + + +async def _async_pause(seconds: float) -> None: + await asyncio.sleep(seconds) + + +def _job_url(platform: NeMoPlatform | AsyncNeMoPlatform, workspace: str, job_name: str, path: str = "") -> str: + base = str(platform.base_url).rstrip("/") + return f"{base}/apis/auditor/v2/workspaces/{workspace}/jobs/audit/{job_name}{path}" + + +@dataclass +class _WaitLogCollector: + """Collects and processes log entries emitted during job polling.""" + + seen_logs: list[dict[str, str]] + error_occurred: bool + warning_occurred: bool + + @classmethod + def create(cls) -> Self: + return cls(seen_logs=[], error_occurred=False, warning_occurred=False) + + def accept_logs(self, current_logs: list[dict[str, str]]) -> None: + for log in current_logs[len(self.seen_logs) :]: + self.seen_logs.append(log) + if not log["name"].startswith("nemo_auditor"): + continue + level = log["levelname"].lower() + if level == "info": + logger.info(log["message"]) + elif level in {"warning", "warn"}: + logger.warning(log["message"]) + self.warning_occurred = True + elif level == "error": + logger.error(log["message"]) + self.error_occurred = True + + def log_final_status(self) -> None: + if self.error_occurred: + logger.error("Audit job completed with errors.") + elif self.warning_occurred: + logger.warning("Audit job completed with warnings.") + else: + logger.info("Audit job completed successfully.") + + +def _status_is_complete(status: PlatformJobStatus | None, raise_if_not_complete: bool) -> bool: + if status == "completed": + return True + if status == "active": + msg = "The audit job is still running." + if raise_if_not_complete: + raise RuntimeError(msg) + logger.warning(msg) + return False + if status in TERMINAL_INCOMPLETE_STATUSES: + msg = f"The audit job stopped with status {status!r}." + if raise_if_not_complete: + raise RuntimeError(msg) + logger.error(msg) + return False + if status in {"created", "pending"}: + msg = f"The audit job is still in the queue with status {status!r}." + if raise_if_not_complete: + raise RuntimeError(msg) + logger.warning(msg) + return False + msg = f"The audit job is in an unknown state: {status!r}." + if raise_if_not_complete: + raise RuntimeError(msg) + logger.error(msg) + return False + + +def _try_parse_log_message(raw_message: str) -> dict[str, str] | None: + """Best-effort extraction of the JSON payload from a platform log entry.""" + json_start = raw_message.find("{") + if json_start < 0: + return None + try: + deserialized = json.loads(raw_message[json_start:]) + except Exception: + return None + if not isinstance(deserialized, dict) or "message" not in deserialized: + return None + return deserialized + + +class AuditorJobResource: + """Sync SDK handle for a submitted audit job.""" + + def __init__(self, *, job_name: str, platform: NeMoPlatform, workspace: str) -> None: + self._job_name = job_name + self._platform = platform + self._workspace = workspace + self._consecutive_poll_errors = 0 + + @property + def name(self) -> str: + """The unique identifying name of the job.""" + return self._job_name + + def get_job(self) -> dict[str, object]: + """Fetch the current job dict.""" + resp = self._platform._client.get(_job_url(self._platform, self._workspace, self._job_name)) + resp.raise_for_status() + return resp.json() + + def get_job_status(self) -> PlatformJobStatus | None: + """Fetch the current platform status of the job.""" + resp = self._platform._client.get(_job_url(self._platform, self._workspace, self._job_name, "/status")) + resp.raise_for_status() + return resp.json().get("status") + + def check_if_complete(self, *, raise_if_not_complete: bool = False) -> bool: + """Return whether the job has reached the ``completed`` status. + + Args: + raise_if_not_complete: When ``True``, raise ``RuntimeError`` for any + status other than ``completed``. + """ + return _status_is_complete(self.get_job_status(), raise_if_not_complete) + + def wait_until_done(self) -> None: + """Block until the job reaches a terminal status, streaming logs along the way.""" + log_collector = _WaitLogCollector.create() + job_status = self.get_job_status() + while job_status != "completed": + _pause(WAIT_INTERVAL_SECONDS) + current_logs = self._poll_safe(self.get_logs, log_collector.seen_logs) + log_collector.accept_logs(current_logs) + if job_status in TERMINAL_INCOMPLETE_STATUSES: + log_collector.error_occurred = True + logger.error(f"Audit job terminated with status {job_status!r}.") + break + job_status = self._poll_safe(self.get_job_status, job_status) + log_collector.log_final_status() + + def get_logs(self) -> list[dict[str, str]]: + """Page through and return all job log entries.""" + logs = [] + page_cursor = None + while True: + params = {"page_cursor": page_cursor} if page_cursor else None + resp = self._platform._client.get( + _job_url(self._platform, self._workspace, self._job_name, "/logs"), + params=params, + ) + resp.raise_for_status() + response = resp.json() + for log in response.get("data", []): + deserialized = _try_parse_log_message(log.get("message", "")) + if deserialized is not None: + logs.append(deserialized) + page_cursor = response.get("next_page") + if page_cursor is None: + break + return logs + + def download_artifacts(self, path: Path | str | None = None) -> Path: + """Download and extract the garak report artifacts for this job. + + Args: + path: Base output directory. Defaults to a directory named after the job + in the current working directory. + + Returns: + The directory that contains the extracted artifacts. + + Raises: + RuntimeError: If the job has not completed. + """ + status = self.get_job_status() + if status != "completed": + raise RuntimeError( + f"Artifacts are not available: job {self._job_name!r} has status {status!r}. " + "Wait until the job completes before downloading artifacts." + ) + output_path = Path(path or self._job_name) + resp = self._platform._client.get( + _job_url(self._platform, self._workspace, self._job_name, f"/results/{ARTIFACTS_RESULT_NAME}/download"), + ) + resp.raise_for_status() + with tarfile.open(fileobj=io.BytesIO(resp.content), mode="r:*") as tar: + safe_extract_tar(tar, output_path, error_cls=RuntimeError) + return output_path + + def _poll_safe(self, fn: Callable[[], T], fallback: T) -> T: + try: + response = fn() + self._consecutive_poll_errors = 0 + return response + except Exception: + self._consecutive_poll_errors += 1 + if self._consecutive_poll_errors >= MAX_CONSECUTIVE_POLL_ERRORS: + self._consecutive_poll_errors = 0 + raise + return fallback + + +class AsyncAuditorJobResource: + """Async SDK handle for a submitted audit job.""" + + def __init__(self, *, job_name: str, platform: AsyncNeMoPlatform, workspace: str) -> None: + self._job_name = job_name + self._platform = platform + self._workspace = workspace + self._consecutive_poll_errors = 0 + + @property + def name(self) -> str: + """The unique identifying name of the job.""" + return self._job_name + + async def get_job(self) -> dict[str, object]: + """Fetch the current job dict.""" + resp = await self._platform._client.get(_job_url(self._platform, self._workspace, self._job_name)) + resp.raise_for_status() + return resp.json() + + async def get_job_status(self) -> PlatformJobStatus | None: + """Fetch the current platform status of the job.""" + resp = await self._platform._client.get(_job_url(self._platform, self._workspace, self._job_name, "/status")) + resp.raise_for_status() + return resp.json().get("status") + + async def check_if_complete(self, *, raise_if_not_complete: bool = False) -> bool: + """Return whether the job has reached the ``completed`` status. + + Args: + raise_if_not_complete: When ``True``, raise ``RuntimeError`` for any + status other than ``completed``. + """ + return _status_is_complete(await self.get_job_status(), raise_if_not_complete) + + async def wait_until_done(self) -> None: + """Wait until the job reaches a terminal status, streaming logs along the way.""" + log_collector = _WaitLogCollector.create() + job_status = await self.get_job_status() + while job_status != "completed": + await _async_pause(WAIT_INTERVAL_SECONDS) + current_logs = await self._poll_safe(self.get_logs, log_collector.seen_logs) + log_collector.accept_logs(current_logs) + if job_status in TERMINAL_INCOMPLETE_STATUSES: + log_collector.error_occurred = True + logger.error(f"Audit job terminated with status {job_status!r}.") + break + job_status = await self._poll_safe(self.get_job_status, job_status) + log_collector.log_final_status() + + async def get_logs(self) -> list[dict[str, str]]: + """Page through and return all job log entries.""" + logs = [] + page_cursor = None + while True: + params = {"page_cursor": page_cursor} if page_cursor else None + resp = await self._platform._client.get( + _job_url(self._platform, self._workspace, self._job_name, "/logs"), + params=params, + ) + resp.raise_for_status() + response = resp.json() + for log in response.get("data", []): + deserialized = _try_parse_log_message(log.get("message", "")) + if deserialized is not None: + logs.append(deserialized) + page_cursor = response.get("next_page") + if page_cursor is None: + break + return logs + + async def download_artifacts(self, path: Path | str | None = None) -> Path: + """Download and extract the garak report artifacts for this job. + + Args: + path: Base output directory. Defaults to a directory named after the job + in the current working directory. + + Returns: + The directory that contains the extracted artifacts. + + Raises: + RuntimeError: If the job has not completed. + """ + status = await self.get_job_status() + if status != "completed": + raise RuntimeError( + f"Artifacts are not available: job {self._job_name!r} has status {status!r}. " + "Wait until the job completes before downloading artifacts." + ) + output_path = Path(path or self._job_name) + resp = await self._platform._client.get( + _job_url(self._platform, self._workspace, self._job_name, f"/results/{ARTIFACTS_RESULT_NAME}/download"), + ) + resp.raise_for_status() + await asyncio.to_thread( + lambda: _extract_tar(resp.content, output_path), + ) + return output_path + + async def _poll_safe(self, fn: Callable[[], Awaitable[T]], fallback: T) -> T: + try: + response = await fn() + self._consecutive_poll_errors = 0 + return response + except Exception: + self._consecutive_poll_errors += 1 + if self._consecutive_poll_errors >= MAX_CONSECUTIVE_POLL_ERRORS: + self._consecutive_poll_errors = 0 + raise + return fallback + + +def _extract_tar(content: bytes, output_path: Path) -> None: + with tarfile.open(fileobj=io.BytesIO(content), mode="r:*") as tar: + safe_extract_tar(tar, output_path, error_cls=RuntimeError) diff --git a/plugins/nemo-auditor/tests/test_sdk_resources.py b/plugins/nemo-auditor/tests/test_sdk_resources.py index b4f2c226ad..b282b092bf 100644 --- a/plugins/nemo-auditor/tests/test_sdk_resources.py +++ b/plugins/nemo-auditor/tests/test_sdk_resources.py @@ -16,6 +16,9 @@ from __future__ import annotations from datetime import datetime, timezone +import io +import tarfile +from pathlib import Path from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -30,6 +33,7 @@ AuditTarget, ) from nemo_auditor.sdk import AsyncAuditorPluginResource, AuditorPluginResource +from nemo_auditor.sdk_resources.job_resources import AsyncAuditorJobResource, AuditorJobResource from nemo_platform import AsyncNeMoPlatform, NeMoPlatform NOW = datetime.now(timezone.utc) @@ -397,7 +401,8 @@ def test_submit_with_string_refs_posts_correct_url_and_body(self) -> None: result = resource.submit(config="ws/my-cfg", target="ws/my-tgt", workspace="ws") - assert result == _JOB_PAYLOAD + assert isinstance(result, AuditorJobResource) + assert result.name == "audit-job-abc123" platform._client.post.assert_called_once() url = platform._client.post.call_args.args[0] body = platform._client.post.call_args.kwargs["json"] @@ -414,8 +419,9 @@ def test_submit_with_inline_entities_serialises_full_dict(self) -> None: cfg = AuditConfig(name="cfg-1", workspace="default") tgt = AuditTarget(name="tgt-1", workspace="default", type="nim", model="llama") - resource.submit(config=cfg, target=tgt, workspace="default") + result = resource.submit(config=cfg, target=tgt, workspace="default") + assert isinstance(result, AuditorJobResource) body = platform._client.post.call_args.kwargs["json"] assert isinstance(body["spec"]["config"], dict) assert body["spec"]["config"]["name"] == "cfg-1" @@ -427,8 +433,9 @@ def test_submit_defaults_workspace_to_default(self) -> None: platform._client.post.return_value = _ok_response(_JOB_PAYLOAD, status_code=201) resource = AuditorPluginResource(cast(NeMoPlatform, platform)) - resource.submit(config="my-cfg", target="my-tgt") + result = resource.submit(config="my-cfg", target="my-tgt") + assert isinstance(result, AuditorJobResource) url = platform._client.post.call_args.args[0] assert "/workspaces/default/" in url @@ -467,7 +474,8 @@ async def test_submit_posts_correct_url_and_body(self) -> None: result = await resource.submit(config="ws/my-cfg", target="ws/my-tgt", workspace="ws") - assert result == _JOB_PAYLOAD + assert isinstance(result, AsyncAuditorJobResource) + assert result.name == "audit-job-abc123" url = platform._client.post.call_args.args[0] body = platform._client.post.call_args.kwargs["json"] assert url == "http://test:8000/apis/auditor/v2/workspaces/ws/jobs/audit" @@ -551,3 +559,226 @@ async def test_async_run_resolves_names_and_calls_scheduler_in_thread() -> None: assert spec_dict["target"]["name"] == "my-tgt" assert call.kwargs["workspace"] == "default" assert call.kwargs["async_sdk"] is platform + + +# --------------------------------------------------------------------------- +# AuditorJobResource +# --------------------------------------------------------------------------- + +_STATUS_PAYLOAD = {"name": "audit-job-abc123", "status": "active", "workspace": "default"} + + +def _make_tar_bytes() -> bytes: + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tar: + content = b"report data" + info = tarfile.TarInfo(name="report.jsonl") + info.size = len(content) + tar.addfile(info, io.BytesIO(content)) + return buf.getvalue() + + +class TestAuditorJobResource: + def _make_resource(self) -> tuple[_SyncPlatform, AuditorJobResource]: + platform = _SyncPlatform() + resource = AuditorJobResource( + job_name="audit-job-abc123", + platform=cast(NeMoPlatform, platform), + workspace="default", + ) + return platform, resource + + def test_name_property(self) -> None: + _, resource = self._make_resource() + assert resource.name == "audit-job-abc123" + + def test_get_job_hits_named_url(self) -> None: + platform, resource = self._make_resource() + platform._client.get.return_value = _ok_response(_JOB_PAYLOAD) + + result = resource.get_job() + + assert result == _JOB_PAYLOAD + platform._client.get.assert_called_once_with( + "http://test:8000/apis/auditor/v2/workspaces/default/jobs/audit/audit-job-abc123" + ) + + def test_get_job_status_hits_status_url(self) -> None: + platform, resource = self._make_resource() + platform._client.get.return_value = _ok_response({"status": "active"}) + + status = resource.get_job_status() + + assert status == "active" + platform._client.get.assert_called_once_with( + "http://test:8000/apis/auditor/v2/workspaces/default/jobs/audit/audit-job-abc123/status" + ) + + def test_check_if_complete_returns_true_when_completed(self) -> None: + platform, resource = self._make_resource() + platform._client.get.return_value = _ok_response({"status": "completed"}) + assert resource.check_if_complete() is True + + def test_check_if_complete_returns_false_when_active(self) -> None: + platform, resource = self._make_resource() + platform._client.get.return_value = _ok_response({"status": "active"}) + assert resource.check_if_complete() is False + + def test_check_if_complete_raises_when_requested(self) -> None: + platform, resource = self._make_resource() + platform._client.get.return_value = _ok_response({"status": "error"}) + with pytest.raises(RuntimeError): + resource.check_if_complete(raise_if_not_complete=True) + + def test_wait_until_done_polls_until_completed(self) -> None: + platform, resource = self._make_resource() + # Status sequence: active → completed; logs return empty pages each time. + status_responses = [ + _ok_response({"status": "active"}), + _ok_response({"status": "active"}), + _ok_response({"status": "completed"}), + ] + logs_response = _ok_response({"data": [], "next_page": None}) + platform._client.get.side_effect = ( + status_responses[:1] + + [logs_response, status_responses[1]] + + [logs_response, status_responses[2]] + ) + + with patch("nemo_auditor.sdk_resources.job_resources._pause"): + resource.wait_until_done() + + # Final status call resolves to "completed" — no RuntimeError raised. + + def test_wait_until_done_exits_on_terminal_failure(self) -> None: + platform, resource = self._make_resource() + status_responses = [ + _ok_response({"status": "active"}), + _ok_response({"status": "error"}), + ] + logs_response = _ok_response({"data": [], "next_page": None}) + platform._client.get.side_effect = [ + status_responses[0], + logs_response, + status_responses[1], + ] + + with patch("nemo_auditor.sdk_resources.job_resources._pause"): + resource.wait_until_done() # must not raise — just logs error and returns + + def test_download_artifacts_raises_when_not_completed(self) -> None: + platform, resource = self._make_resource() + platform._client.get.return_value = _ok_response({"status": "active"}) + + with pytest.raises(RuntimeError, match="status 'active'"): + resource.download_artifacts() + + def test_download_artifacts_extracts_tarball(self, tmp_path: Path) -> None: + platform, resource = self._make_resource() + status_resp = _ok_response({"status": "completed"}) + tar_resp = MagicMock(spec=httpx.Response) + tar_resp.raise_for_status.return_value = None + tar_resp.content = _make_tar_bytes() + platform._client.get.side_effect = [status_resp, tar_resp] + + result = resource.download_artifacts(path=tmp_path) + + assert result == tmp_path + platform._client.get.assert_called_with( + "http://test:8000/apis/auditor/v2/workspaces/default/jobs/audit/audit-job-abc123/results/artifacts/download" + ) + + def test_poll_safe_returns_fallback_on_transient_error(self) -> None: + _, resource = self._make_resource() + + def failing() -> str: + raise ConnectionError("blip") + + result = resource._poll_safe(failing, "cached") + assert result == "cached" + assert resource._consecutive_poll_errors == 1 + + def test_poll_safe_raises_after_max_consecutive_errors(self) -> None: + _, resource = self._make_resource() + resource._consecutive_poll_errors = 4 + + def failing() -> str: + raise ConnectionError("blip") + + with pytest.raises(ConnectionError): + resource._poll_safe(failing, "cached") + assert resource._consecutive_poll_errors == 0 + + +@pytest.mark.asyncio +class TestAsyncAuditorJobResource: + def _make_resource(self) -> tuple[_AsyncPlatform, AsyncAuditorJobResource]: + platform = _AsyncPlatform() + resource = AsyncAuditorJobResource( + job_name="audit-job-abc123", + platform=cast(AsyncNeMoPlatform, platform), + workspace="default", + ) + return platform, resource + + async def test_name_property(self) -> None: + _, resource = self._make_resource() + assert resource.name == "audit-job-abc123" + + async def test_get_job_status_hits_status_url(self) -> None: + platform, resource = self._make_resource() + platform._client.get.return_value = _ok_response({"status": "completed"}) + + status = await resource.get_job_status() + + assert status == "completed" + platform._client.get.assert_called_once_with( + "http://test:8000/apis/auditor/v2/workspaces/default/jobs/audit/audit-job-abc123/status" + ) + + async def test_check_if_complete_returns_true_when_completed(self) -> None: + platform, resource = self._make_resource() + platform._client.get.return_value = _ok_response({"status": "completed"}) + assert await resource.check_if_complete() is True + + async def test_check_if_complete_raises_when_requested(self) -> None: + platform, resource = self._make_resource() + platform._client.get.return_value = _ok_response({"status": "cancelled"}) + with pytest.raises(RuntimeError): + await resource.check_if_complete(raise_if_not_complete=True) + + async def test_wait_until_done_polls_until_completed(self) -> None: + platform, resource = self._make_resource() + status_responses = [ + _ok_response({"status": "active"}), + _ok_response({"status": "active"}), + _ok_response({"status": "completed"}), + ] + logs_response = _ok_response({"data": [], "next_page": None}) + platform._client.get.side_effect = ( + status_responses[:1] + + [logs_response, status_responses[1]] + + [logs_response, status_responses[2]] + ) + + with patch("nemo_auditor.sdk_resources.job_resources._async_pause", new=AsyncMock()): + await resource.wait_until_done() + + async def test_download_artifacts_raises_when_not_completed(self) -> None: + platform, resource = self._make_resource() + platform._client.get.return_value = _ok_response({"status": "pending"}) + + with pytest.raises(RuntimeError, match="status 'pending'"): + await resource.download_artifacts() + + async def test_download_artifacts_extracts_tarball(self, tmp_path: Path) -> None: + platform, resource = self._make_resource() + status_resp = _ok_response({"status": "completed"}) + tar_resp = MagicMock(spec=httpx.Response) + tar_resp.raise_for_status.return_value = None + tar_resp.content = _make_tar_bytes() + platform._client.get.side_effect = [status_resp, tar_resp] + + result = await resource.download_artifacts(path=tmp_path) + + assert result == tmp_path From 01fe4e5df5c604a986ffd201d77321b39e1df32d Mon Sep 17 00:00:00 2001 From: "Paul A. Parkanzky" Date: Fri, 31 Jul 2026 18:17:55 -0400 Subject: [PATCH 2/2] fix(auditor): lint Signed-off-by: Paul A. Parkanzky --- plugins/nemo-auditor/tests/test_sdk_resources.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/plugins/nemo-auditor/tests/test_sdk_resources.py b/plugins/nemo-auditor/tests/test_sdk_resources.py index b282b092bf..ef2c62b4fe 100644 --- a/plugins/nemo-auditor/tests/test_sdk_resources.py +++ b/plugins/nemo-auditor/tests/test_sdk_resources.py @@ -15,9 +15,9 @@ from __future__ import annotations -from datetime import datetime, timezone import io import tarfile +from datetime import datetime, timezone from pathlib import Path from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -640,9 +640,7 @@ def test_wait_until_done_polls_until_completed(self) -> None: ] logs_response = _ok_response({"data": [], "next_page": None}) platform._client.get.side_effect = ( - status_responses[:1] - + [logs_response, status_responses[1]] - + [logs_response, status_responses[2]] + status_responses[:1] + [logs_response, status_responses[1]] + [logs_response, status_responses[2]] ) with patch("nemo_auditor.sdk_resources.job_resources._pause"): @@ -756,9 +754,7 @@ async def test_wait_until_done_polls_until_completed(self) -> None: ] logs_response = _ok_response({"data": [], "next_page": None}) platform._client.get.side_effect = ( - status_responses[:1] - + [logs_response, status_responses[1]] - + [logs_response, status_responses[2]] + status_responses[:1] + [logs_response, status_responses[1]] + [logs_response, status_responses[2]] ) with patch("nemo_auditor.sdk_resources.job_resources._async_pause", new=AsyncMock()):