nemo-platform-plugin 0.1.0__py3-none-any.whl

nemo_platform_plugin/jobs/result_manager.py

@@ -0,0 +1,291 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+import tarfile
+from abc import ABC
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Generic, Literal, Type, TypeVar, overload
+
+from nemo_platform import APIError, AsyncNeMoPlatform, ConflictError, NeMoPlatform
+from nemo_platform.filesets import parse_fileset_ref
+from nemo_platform.types import PlatformJobResultResponse as SDKPlatformJobResult
+from nemo_platform_plugin.jobs.constants import NEMO_JOB_WORKSPACE_ENVVAR
+from nemo_platform_plugin.jobs.file_manager import (
+    AsyncFilesetFileManager,
+    FilesetFileManager,
+    TmpDirPath,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class CreateJobResultError(Exception): ...
+
+
+class FileDoesNotExist(Exception): ...
+
+
+FileManagerClsT = TypeVar("FileManagerClsT", Type["FilesetFileManager"], Type["AsyncFilesetFileManager"])
+PlatformSDKT = TypeVar("PlatformSDKT", "NeMoPlatform", "AsyncNeMoPlatform")
+
+
+@dataclass
+class BaseResultManager(Generic[FileManagerClsT, PlatformSDKT], ABC):
+    """
+    Base class for sync and async versions of the ResultManager. If there is any common code that can be used across
+    both versions of the ResultManager, try to lift it up into this class to reduce duplication.
+    """
+
+    job_name: str
+    workspace: str
+    file_manager_cls: FileManagerClsT
+    files_sdk: PlatformSDKT
+    jobs_sdk: PlatformSDKT
+    attempt_id: str | None = field(default=None)
+
+    def _validate_local_path(self, artifact_local_path: str | Path) -> Path:
+        if isinstance(artifact_local_path, str):
+            artifact_local_path = Path(artifact_local_path)
+
+        if not artifact_local_path.exists():
+            raise FileDoesNotExist(f"No file exists at: {artifact_local_path}")
+
+        return artifact_local_path
+
+    def _result_remote_path(self, attempt_id: str, result_name: str) -> str:
+        """Build the remote path for a result artifact."""
+        return f"results/{attempt_id}/{result_name}"
+
+
+@dataclass
+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)
+        attempt_id = self.attempt_id if self.attempt_id is not None else job.attempt_id
+        return attempt_id, job.fileset
+
+    def _create_file_manager(self, fileset_name: str) -> FilesetFileManager:
+        """Create a file manager for the given fileset."""
+        return self.file_manager_cls(
+            workspace=self.workspace,
+            fileset_name=fileset_name,
+            sdk=self.files_sdk,
+        )
+
+    def create_result(
+        self,
+        result_name: str,
+        artifact_local_path: str | Path,
+        ignore_patterns: list[str] | str | None = None,
+    ) -> SDKPlatformJobResult:
+        attempt_id, fileset_name = self._fetch_job_metadata()
+        file_manager = self._create_file_manager(fileset_name)
+        file_manager.validate_storage()
+        artifact_local_path = self._validate_local_path(artifact_local_path)
+        remote_path = self._result_remote_path(attempt_id, result_name)
+        artifact_url = file_manager.upload(
+            local_path=artifact_local_path, remote_path=remote_path, ignore_patterns=ignore_patterns
+        )
+        try:
+            job_result = 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:
+            # 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(
+                name=result_name,
+                job=self.job_name,
+                workspace=self.workspace,
+            )
+        except APIError as e:
+            msg = f"Error creating job result: {str(e)}"
+            logger.exception(msg)
+            raise CreateJobResultError(msg) from e
+
+        return job_result
+
+    def download_artifact(self, artifact_url: str, local_dir: str | Path | None = None) -> TmpDirPath:
+        _, fileset_name, _ = parse_fileset_ref(artifact_url, workspace_fallback=self.workspace)
+        if fileset_name is None:
+            raise ValueError(f"Could not extract fileset name from URL: {artifact_url}")
+        file_manager = self._create_file_manager(fileset_name)
+        return file_manager.download_from_url(url=artifact_url, local_dir=local_dir)
+
+
+@dataclass
+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)
+        attempt_id = self.attempt_id if self.attempt_id is not None else job.attempt_id
+        return attempt_id, job.fileset
+
+    def _create_file_manager(self, fileset_name: str) -> AsyncFilesetFileManager:
+        """Create a file manager for the given fileset."""
+        return self.file_manager_cls(
+            workspace=self.workspace,
+            fileset_name=fileset_name,
+            sdk=self.files_sdk,
+        )
+
+    async def create_result(
+        self,
+        result_name: str,
+        artifact_local_path: str | Path,
+        ignore_patterns: list[str] | str | None = None,
+    ) -> SDKPlatformJobResult:
+        attempt_id, fileset_name = await self._fetch_job_metadata()
+        file_manager = self._create_file_manager(fileset_name)
+        await file_manager.validate_storage()
+        artifact_local_path = self._validate_local_path(artifact_local_path)
+        remote_path = self._result_remote_path(attempt_id, result_name)
+        artifact_url = await file_manager.upload(
+            local_path=artifact_local_path, remote_path=remote_path, ignore_patterns=ignore_patterns
+        )
+        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:
+            # 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:
+            msg = f"Error creating job result: {str(e)}"
+            logger.exception(msg)
+            raise CreateJobResultError(msg) from e
+
+        return job_result
+
+    async def download_artifact(self, artifact_url: str, local_dir: str | Path | None = None) -> TmpDirPath:
+        _, fileset_name, _ = parse_fileset_ref(artifact_url, workspace_fallback=self.workspace)
+        if fileset_name is None:
+            raise ValueError(f"Could not extract fileset name from URL: {artifact_url}")
+        file_manager = self._create_file_manager(fileset_name)
+        return await file_manager.download_from_url(url=artifact_url, local_dir=local_dir)
+
+
+@overload
+def result_manager_factory(
+    job_name: str,
+    *,
+    attempt_id: str | None = None,
+    workspace: str | None = None,
+    files_sdk: AsyncNeMoPlatform,
+    jobs_sdk: AsyncNeMoPlatform | None = None,
+    is_async: Literal[True] = True,
+) -> AsyncResultManager: ...
+
+
+@overload
+def result_manager_factory(
+    job_name: str,
+    *,
+    attempt_id: str | None = None,
+    workspace: str | None = None,
+    files_sdk: NeMoPlatform,
+    jobs_sdk: NeMoPlatform | None = None,
+    is_async: Literal[False],
+) -> ResultManager: ...
+
+
+def result_manager_factory(
+    job_name: str,
+    *,
+    attempt_id: str | None = None,
+    workspace: str | None = None,
+    files_sdk: NeMoPlatform | AsyncNeMoPlatform,
+    jobs_sdk: NeMoPlatform | AsyncNeMoPlatform | None = None,
+    is_async: bool = True,
+) -> ResultManager | AsyncResultManager:
+    """Create a ResultManager for uploading job results.
+
+    Args:
+        job_name: Name of the job to create results for.
+        attempt_id: Optional attempt ID for the job.
+        workspace: Job workspace. If not provided, reads from NEMO_JOB_WORKSPACE env var.
+        files_sdk: SDK instance for file operations (required).
+        jobs_sdk: SDK instance for job operations. If not provided, uses files_sdk.
+        is_async: Whether to create an async result manager.
+
+    Returns:
+        ResultManager or AsyncResultManager instance.
+    """
+
+    if workspace is None:
+        workspace = _get_job_workspace()
+
+    if jobs_sdk is None:
+        jobs_sdk = files_sdk
+
+    file_manager_cls = AsyncFilesetFileManager if is_async else FilesetFileManager
+    result_manager_cls = AsyncResultManager if is_async else ResultManager
+    return result_manager_cls(
+        job_name=job_name,
+        workspace=workspace,
+        attempt_id=attempt_id,
+        file_manager_cls=file_manager_cls,  # type: ignore
+        files_sdk=files_sdk,  # type: ignore
+        jobs_sdk=jobs_sdk,  # type: ignore
+    )
+
+
+def _get_job_workspace() -> str:
+    workspace = os.getenv(NEMO_JOB_WORKSPACE_ENVVAR)
+    if not workspace:
+        raise ValueError(f"{NEMO_JOB_WORKSPACE_ENVVAR} environment variable is not set")
+    return workspace
+
+
+async def download_from_result_info(
+    result_name: str,
+    job_name: str,
+    *,
+    artifact_url: str,
+    workspace: str | None = None,
+    files_sdk: AsyncNeMoPlatform,
+) -> tuple[str, TmpDirPath]:
+    """
+    This is a helper composition function that creates a result_manager
+    and prepares the artifact for download.
+
+    Returns:
+        A tuple containing:
+          - filename (str): The final filename of the artifact (may include .tar.gz extension if compressed)
+          - tmp_dir_path (TmpDirPath): Path object pointing to the downloaded artifact
+    """
+    result_manager = result_manager_factory(
+        job_name=job_name,
+        workspace=workspace,
+        files_sdk=files_sdk,
+    )
+
+    tmp_dir_path = await result_manager.download_artifact(artifact_url=artifact_url)
+    filename = result_name
+
+    if tmp_dir_path.path.is_dir():
+        filename = f"{filename}.tar.gz"
+        tar_path = tmp_dir_path.tmp_dir / filename
+        with tarfile.open(tar_path, "w:gz") as tar:
+            tar.add(tmp_dir_path.path, arcname=os.path.basename(tmp_dir_path.path))
+
+        tmp_dir_path.path = tar_path
+
+    return filename, tmp_dir_path

nemo_platform_plugin/jobs/routes.py

@@ -0,0 +1,297 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""High-level router wiring for ``NemoJob`` subclasses.
+
+``add_job_routes(job_cls)`` is the one-line replacement for the verbose
+``job_route_factory(...)`` call pattern in plugin services. It derives
+``service_name`` / ``job_type`` / ``job_input`` / ``job_output`` /
+``input_to_output`` / ``platform_job_config_compiler`` from the job
+class's declared attributes and delegates to :func:`job_route_factory`
+for the actual route generation.
+
+Usage::
+
+    from nemo_platform_plugin.jobs.routes import add_job_routes
+
+    router = add_job_routes(MyJob)
+    app.include_router(
+        router,
+        prefix="/v2/workspaces/{workspace}",
+        tags=["My Plugin"],
+    )
+
+Compare to the legacy pattern this replaces::
+
+    router = job_route_factory(
+        service_name="my-plugin",
+        job_type="MyJob",
+        job_input=MyInputSchema,
+        job_output=MyCanonicalSchema,
+        input_to_output=my_to_spec,
+        platform_job_config_compiler=my_compiler,
+        generate_job_name=my_name_generator,
+    )
+
+Scope note (phase-1, MR 1.1b): ``add_job_routes`` applies
+:func:`stamp_profile` to the compiled ``PlatformJobSpec`` using
+``default_profile``. It does *not* yet thread ``profile`` / ``options``
+body fields through :class:`BaseJobRequest` — those land together with
+the ``BaseJobRequest`` extension in a follow-up MR. The
+submitter-facing CLI flags (``--profile``, ``-o``) from MR 1.3b reach
+``submit_remote``, which POSTs them in the body; the server currently
+silently drops them (Pydantic ``extra="ignore"`` default). The wrapper
+passes ``profile=None`` / ``options=None`` to ``NemoJob.compile``
+until the body shape is extended.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from fastapi import APIRouter
+from fastapi.routing import APIRoute
+from nemo_platform_plugin.job import job_collection_path_for
+from nemo_platform_plugin.jobs.api_factory import (
+    JobRouteOption,
+    PlatformJobResultRoute,
+    job_route_factory,
+)
+from nemo_platform_plugin.jobs.exceptions import PlatformJobCompilationError
+from nemo_platform_plugin.jobs.profile import stamp_profile
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from nemo_platform_plugin.job import NemoJob
+
+
+def add_job_routes(
+    job_cls: type["NemoJob"],
+    *,
+    service_name: str | None = None,
+    route_options: list[JobRouteOption] | None = None,
+    job_result_routes: list[PlatformJobResultRoute] | None = None,
+    generate_job_name: "Callable[..., str] | None" = None,
+    default_profile: str = "default",
+) -> APIRouter:
+    """Mount submit/list/get/delete routes for *job_cls* on a fresh router.
+
+    Returns an :class:`APIRouter` produced by :func:`job_route_factory`,
+    rebased onto the job collection path. The caller mounts it with the
+    service/workspace URL prefix (typically
+    ``/apis/{api}/v2/workspaces/{workspace}``).
+
+    Args:
+        job_cls: :class:`~nemo_platform_plugin.job.NemoJob` subclass. Must declare
+            ``spec_schema``; must override ``compile`` to be remote-capable.
+            When ``input_spec_schema`` is also declared, :meth:`to_spec` is
+            used to transform the submitted spec into the canonical shape
+            that ``compile`` sees.
+        service_name: Override the factory's ``service_name``. Defaults to
+            the top-level module name of *job_cls*, converted from
+            underscores to hyphens (e.g. ``my_plugin.jobs.widget`` →
+            ``"my-plugin"``). The factory uses this in OpenAPI schema names
+            and error messages.
+        route_options: Which :class:`JobRouteOption` routes to enable.
+            Defaults to ``[JobRouteOption.CORE]`` inside the factory.
+        job_result_routes: Custom :class:`PlatformJobResultRoute` entries
+            for result download endpoints.
+        generate_job_name: Callable invoked when the submitter didn't
+            provide a ``name``. Passthrough to the factory.
+        default_profile: Profile label stamped onto each step of the
+            compiled ``PlatformJobSpec`` when the plugin's ``compile``
+            didn't set one explicitly. Matches the plan's default-profile
+            behavior; submitter-chosen ``--profile`` plumbing lands in a
+            follow-up MR that extends ``BaseJobRequest``.
+
+    Returns:
+        An :class:`APIRouter` with the standard job endpoints mounted.
+
+    Raises:
+        At request time, ``compile()`` on a class that hasn't overridden
+        the base marker raises :class:`NotImplementedError`; the wrapper
+        catches it and re-raises as
+        :class:`PlatformJobCompilationError` so the factory surfaces a
+        clean 422 response instead of a 500 traceback.
+    """
+    if job_cls.spec_schema is None:
+        raise TypeError(
+            f"{job_cls.__name__}.spec_schema is None; add_job_routes requires a declared spec_schema. "
+            f"Set it to a Pydantic BaseModel subclass on the NemoJob class."
+        )
+
+    resolved_service_name = service_name or _derive_service_name(job_cls)
+    job_type = _derive_job_type(job_cls)
+    job_input = job_cls.input_spec_schema or job_cls.spec_schema
+    # When input_spec_schema is declared, the stored shape is the canonical
+    # spec_schema; otherwise input and stored shapes coincide.
+    job_output = job_cls.spec_schema if job_cls.input_spec_schema is not None else None
+    input_to_output = _adapt_to_spec(job_cls) if job_cls.input_spec_schema is not None else None
+
+    router = job_route_factory(
+        service_name=resolved_service_name,
+        job_type=job_type,
+        job_input=job_input,
+        job_output=job_output,
+        input_to_output=input_to_output,
+        platform_job_config_compiler=_adapt_compile(job_cls, default_profile),
+        route_options=route_options,
+        job_result_routes=job_result_routes,
+        generate_job_name=generate_job_name,
+    )
+    return _rebase_job_collection_routes(router, job_collection_path_for(job_cls))
+
+
+def _rebase_job_collection_routes(router: APIRouter, collection_path: str) -> APIRouter:
+    """Move factory routes from ``/jobs`` to *collection_path*.
+
+    ``job_route_factory`` remains the low-level legacy primitive that emits a
+    generic ``/jobs`` collection. ``add_job_routes`` is the NemoJob-aware
+    helper, so it owns rebasing those routes onto the job type's collection
+    path (``/jobs/{job_cls.name}`` by default).
+    """
+    rebased = APIRouter()
+    for route in router.routes:
+        if not isinstance(route, APIRoute):
+            rebased.routes.append(route)
+            continue
+
+        path = route.path
+        if path == "/jobs" or path.startswith("/jobs/"):
+            path = f"{collection_path}{path[len('/jobs') :]}"
+
+        rebased.add_api_route(
+            path=path,
+            endpoint=route.endpoint,
+            methods=route.methods,
+            name=route.name,
+            response_model=route.response_model,
+            status_code=route.status_code,
+            tags=route.tags,
+            dependencies=route.dependencies,
+            summary=route.summary,
+            description=route.description,
+            response_description=route.response_description,
+            responses=route.responses,
+            deprecated=route.deprecated,
+            operation_id=route.operation_id,
+            response_model_include=route.response_model_include,
+            response_model_exclude=route.response_model_exclude,
+            response_model_by_alias=route.response_model_by_alias,
+            response_model_exclude_unset=route.response_model_exclude_unset,
+            response_model_exclude_defaults=route.response_model_exclude_defaults,
+            response_model_exclude_none=route.response_model_exclude_none,
+            include_in_schema=route.include_in_schema,
+            response_class=route.response_class,
+            openapi_extra=route.openapi_extra,
+        )
+    return rebased
+
+
+# ---------------------------------------------------------------------------
+# Derivation helpers
+# ---------------------------------------------------------------------------
+
+
+def _derive_service_name(job_cls: type["NemoJob"]) -> str:
+    """Top-level package of *job_cls*, hyphen-normalised.
+
+    Matches the convention used by
+    :func:`nemo_platform_plugin.scheduler._api_segment_for`, so the factory's
+    ``service_name`` and the submit URL's ``{api}`` segment stay in
+    lockstep for a given plugin.
+    """
+    return job_cls.__module__.split(".")[0].replace("_", "-")
+
+
+def _derive_job_type(job_cls: type["NemoJob"]) -> str:
+    """PascalCase form of ``job_cls.name`` for OpenAPI schema names.
+
+    ``"generate"`` → ``"Generate"``; ``"metric-eval"`` → ``"MetricEval"``;
+    ``"raw_job"`` → ``"RawJob"``.
+    """
+    parts = job_cls.name.replace("_", "-").split("-")
+    return "".join(part[:1].upper() + part[1:] for part in parts if part)
+
+
+# ---------------------------------------------------------------------------
+# Signature adapters
+# ---------------------------------------------------------------------------
+
+
+def _adapt_to_spec(job_cls: type["NemoJob"]) -> "Callable[..., Any]":
+    """Bridge ``NemoJob.to_spec`` to the factory's ``input_to_output`` shape.
+
+    The factory calls ``input_to_output(original, workspace, entity_client,
+    job_name, sdk)`` with an :class:`AsyncNeMoPlatform` in the ``sdk``
+    slot. :meth:`NemoJob.to_spec` takes the same data but organises it as
+    ``(input_spec, *, workspace, entity_client, async_sdk)`` (no
+    ``job_name``); ``async_sdk`` matches the codebase-wide convention
+    that this name carries an async client. The shim adapts and awaits
+    the ``async classmethod``.
+    """
+
+    async def to_spec_adapter(
+        original_spec: Any,
+        workspace: str,
+        entity_client: Any,
+        job_name: str | None,
+        sdk: Any,
+    ) -> Any:
+        del job_name  # NemoJob.to_spec doesn't use it (names belong to the Jobs service)
+        return await job_cls.to_spec(
+            original_spec,
+            workspace=workspace,
+            entity_client=entity_client,
+            async_sdk=sdk,
+            is_local=False,
+        )
+
+    return to_spec_adapter
+
+
+def _adapt_compile(
+    job_cls: type["NemoJob"],
+    default_profile: str,
+) -> "Callable[..., Any]":
+    """Bridge ``NemoJob.compile`` to the factory's ``platform_job_config_compiler`` shape.
+
+    The factory calls ``compiler(workspace, original_spec, transformed_spec,
+    entity_client, job_name, sdk)``. :meth:`NemoJob.compile` is an
+    ``async classmethod`` that uses kwargs and also accepts
+    ``profile`` / ``options`` — phase 1 MR 1.1b passes ``None`` for
+    both (body-field wiring is a follow-up). After ``compile`` returns,
+    the adapter applies :func:`stamp_profile` with ``default_profile``.
+
+    Missing-override errors from the ``NemoJob.compile`` base marker
+    become :class:`PlatformJobCompilationError` so the factory's
+    existing handler surfaces them as 422 instead of 500.
+    """
+
+    async def compile_adapter(
+        workspace: str,
+        original_spec: Any,
+        transformed_spec: Any,
+        entity_client: Any,
+        job_name: str | None,
+        sdk: Any,
+    ) -> Any:
+        del original_spec  # NemoJob.compile only needs the canonical (transformed) spec
+        try:
+            result = await job_cls.compile(
+                workspace=workspace,
+                spec=transformed_spec,
+                entity_client=entity_client,
+                job_name=job_name,
+                async_sdk=sdk,
+                profile=None,
+                options=None,
+            )
+        except NotImplementedError as exc:
+            raise PlatformJobCompilationError(str(exc)) from exc
+
+        stamp_profile(result, default_profile)
+        return result
+
+    return compile_adapter