feat(client): Add new Jobs NemoClient - #585
Conversation
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
|
…o-nemoclient-typed-http-client Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds typed synchronous and asynchronous Jobs clients, shared Jobs service models, endpoint contracts, server re-exports, and migrations from direct SDK Jobs calls to the typed client abstraction. ChangesJobs client and shared contracts
Server and application migration
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (4)
packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/providers.py (1)
102-166: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated shape across CPU/GPU/DistributedGPU providers.
CPUExecutionProvider,GPUExecutionProvider,DistributedGPUExecutionProviderare structurally identical (onlyproviderliteral differs). Consider a shared base class withcontainer/resources/profileto avoid triplicated maintenance.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/providers.py` around lines 102 - 166, The CPUExecutionProvider, GPUExecutionProvider, and DistributedGPUExecutionProvider classes duplicate the same fields and validation shape except for the provider literal. Refactor the shared structure into a common base model that owns profile, container, and resources, then keep only the provider-specific literal in each subclass so changes to the execution spec are maintained in one place; use the existing class names in providers.py to preserve the current public API.packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/execution_profiles.py (1)
380-384: 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick winHardcoded shared temp path flagged by static analysis.
working_directorydefault/tmp/nmp-subprocess-jobsis a fixed, predictable path shared across runs — potential symlink/race exposure on multi-tenant hosts (CWE-377). Consider deriving a unique per-job subdirectory or usingtempfileprimitives, even if this remains an overridable default.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/execution_profiles.py` around lines 380 - 384, The default for SubprocessJobExecutionProfileConfig.working_directory is a fixed shared temp path, which static analysis flagged as predictable and reusable across runs. Update the default in SubprocessJobExecutionProfileConfig to use a unique per-job subdirectory or tempfile-based directory generation instead of the hardcoded /tmp/nmp-subprocess-jobs path, while keeping the field overridable for callers.Source: Linters/SAST tools
packages/nemo_platform_plugin/tests/jobs/test_endpoints.py (1)
6-6: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSame
from __future__ import annotationsconcern asendpoints.py.Less risky here since tests don't rely on runtime signature introspection, but same guideline applies.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/tests/jobs/test_endpoints.py` at line 6, Remove the unnecessary future import in the test module to match the same import guideline used in endpoints.py; update the test file so it no longer relies on from __future__ import annotations, keeping the module consistent with the rest of the codebase and avoiding the runtime annotation behavior change.Source: Coding guidelines
packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/endpoints.py (1)
109-111: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winInconsistent path param naming for
rerun_job.Every other job-level lifecycle endpoint (
cancel_job,pause_job,resume_job,get_job, etc.) usesnamefor the job identifier;rerun_jobusesjobinstead, for the same concept. This propagates the inconsistency intoclient.pyand tests.♻️ Align param naming
-@post("/apis/jobs/v2/workspaces/{workspace}/jobs/{job}/rerun") +@post("/apis/jobs/v2/workspaces/{workspace}/jobs/{name}/rerun") `@abstractmethod` -def rerun_job(*, workspace: str | None = None, job: str) -> PlatformJobResponse: ... +def rerun_job(*, workspace: str | None = None, name: str) -> PlatformJobResponse: ...🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/endpoints.py` around lines 109 - 111, The `rerun_job` endpoint uses a different job identifier name than the rest of the job lifecycle API, which should be aligned for consistency. Update `rerun_job` in `nemo_platform_plugin.jobs.endpoints` to use `name` instead of `job`, and propagate that rename through the corresponding `client.py` implementation and related tests. Keep the same route and behavior, but make the parameter naming match `cancel_job`, `pause_job`, `resume_job`, and `get_job` so the API surface is consistent.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/execution_profiles.py`:
- Around line 369-372: The supports_persistent_storage property on
VolcanoJobExecutionProfile is using an always-true None check for
config.storage.pvc_name, so update it to match KubernetesJobExecutionProfile by
checking for a non-empty pvc_name instead. Keep the logic in the
VolcanoJobExecutionProfile.supports_persistent_storage accessor consistent with
the sibling implementation and ensure it only returns true when storage exists
and pvc_name is not empty.
- Around line 171-174: The DockerJobExecutionProfile.supports_persistent_storage
property is currently meaningless because self.config.storage is always
initialized by the default factory and will never be None. Update the
supports_persistent_storage implementation in DockerJobExecutionProfile to check
a real capability/flag from the Docker storage config, similar to how
KubernetesJobExecutionProfile determines persistence from a meaningful field, so
the property reflects actual persistent storage support instead of always
returning True.
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/providers.py`:
- Around line 180-184: The validator return type in
SubprocessExecutionProvider.validate_command is using a quoted class name
instead of Self, which is inconsistent with the typing style used elsewhere.
Update the model_validator return annotation to use Self from typing, matching
the pattern already used in spec.py and keeping the type hint unquoted.
---
Nitpick comments:
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/endpoints.py`:
- Around line 109-111: The `rerun_job` endpoint uses a different job identifier
name than the rest of the job lifecycle API, which should be aligned for
consistency. Update `rerun_job` in `nemo_platform_plugin.jobs.endpoints` to use
`name` instead of `job`, and propagate that rename through the corresponding
`client.py` implementation and related tests. Keep the same route and behavior,
but make the parameter naming match `cancel_job`, `pause_job`, `resume_job`, and
`get_job` so the API surface is consistent.
In
`@packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/execution_profiles.py`:
- Around line 380-384: The default for
SubprocessJobExecutionProfileConfig.working_directory is a fixed shared temp
path, which static analysis flagged as predictable and reusable across runs.
Update the default in SubprocessJobExecutionProfileConfig to use a unique
per-job subdirectory or tempfile-based directory generation instead of the
hardcoded /tmp/nmp-subprocess-jobs path, while keeping the field overridable for
callers.
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/providers.py`:
- Around line 102-166: The CPUExecutionProvider, GPUExecutionProvider, and
DistributedGPUExecutionProvider classes duplicate the same fields and validation
shape except for the provider literal. Refactor the shared structure into a
common base model that owns profile, container, and resources, then keep only
the provider-specific literal in each subclass so changes to the execution spec
are maintained in one place; use the existing class names in providers.py to
preserve the current public API.
In `@packages/nemo_platform_plugin/tests/jobs/test_endpoints.py`:
- Line 6: Remove the unnecessary future import in the test module to match the
same import guideline used in endpoints.py; update the test file so it no longer
relies on from __future__ import annotations, keeping the module consistent with
the rest of the codebase and avoiding the runtime annotation behavior change.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 571f9678-0b1f-491d-b4a8-2a1e246ab1f9
📒 Files selected for processing (13)
packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/client.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/endpoints.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/execution_profiles.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/providers.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/spec.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/types.pypackages/nemo_platform_plugin/tests/jobs/test_endpoints.pyservices/core/jobs/src/nmp/core/jobs/api/v2/jobs/schemas.pyservices/core/jobs/src/nmp/core/jobs/app/providers.pyservices/core/jobs/src/nmp/core/jobs/app/schemas.pyservices/core/jobs/src/nmp/core/jobs/controllers/backends/base.pyservices/core/jobs/src/nmp/core/jobs/controllers/backends/docker.pyservices/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py
…o-nemoclient-typed-http-client Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
…o-nemoclient-typed-http-client Brings in the merged NemoClient migrations from main: - #584 (Files consumers -> FilesClient) - #609 (Secrets service + consumers -> SecretsClient) - #614 (plugin-client URL encoding) Resolved 4 conflicts where the Jobs consumer migration (AIRCORE-874) overlapped the Files/Secrets migrations in multi-service files: - services/core/entities/.../workspace_cleanup.py: combined both migrations (jobs via AsyncJobsClient, filesets via AsyncFilesClient); switched PlatformJobStatus to the plugin type and built the terminal-status set from enum members. - services/core/entities/tests/.../test_workspace_cleanup.py: adopted #584's tuple-return _make_sdk + _MockAsyncPaginatedResponse; added a dispatching _patch_clients() so the combined _async_step test routes client_from_platform to the jobs vs files mock by client class. - plugins/nemo-safe-synthesizer/.../api/v2/jobs/endpoints.py test: dropped the now-orphaned Stainless NotFoundError/PermissionDeniedError import (both uses migrated across the two PRs); made _patch_jobs_client dispatch jobs vs files clients by class. - plugins/nemo-data-designer/.../testing/utils.py: kept both the jobs (CreatePlatformJobRequest) and secrets (SecretsClient) imports. All affected suites green; jobs service unchanged (5 pre-existing env failures). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
…o-nemoclient-typed-http-client Resolved 2 conflicts where main's jobs/docker changes overlapped the AIRCORE-874 leaf-node move and consumer migration: - services/core/jobs/.../controllers/backends/docker.py: kept main's new code (_resolve_jobs_controller_instance_id, DockerTimestampParseResult) and the union import (Any/Generic/Literal/TypeVar + dataclass); dropped main's re-definition of DockerVolumeMount since it now lives in the plugin leaf node and is imported from there. - plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py: kept main's two new container-mode compiler tests alongside the migrated pretrained-model test (which uses the class-dispatching _patch_jobs_client for jobs vs files). All affected suites green; jobs service unchanged (5 pre-existing env failures). No net-new ty diagnostics from the resolution. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
packages/nemo_platform_plugin/tests/test_jobs_filter.py (1)
44-44: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace both string-based
_CapturingSdkannotations.
packages/nemo_platform_plugin/tests/test_jobs_filter.py#L44-L44: use_CapturingSdkdirectly.packages/nemo_platform_plugin/tests/test_jobs_filter.py#L350-L351: use_CapturingSdkdirectly.As per coding guidelines, “prefer concrete type hints over string-based type hints.”
Proposed fix
-def _forwarded_filter(sdk: "_CapturingSdk") -> dict: +def _forwarded_filter(sdk: _CapturingSdk) -> dict: @@ - def _round_trip(sdk: "_CapturingSdk") -> dict: + def _round_trip(sdk: _CapturingSdk) -> dict:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/tests/test_jobs_filter.py` at line 44, Replace the string-based "_CapturingSdk" annotations with direct _CapturingSdk type references in both _forwarded_filter and the sibling annotation at packages/nemo_platform_plugin/tests/test_jobs_filter.py lines 350-351.Source: Coding guidelines
packages/nmp_common/tests/api_factory/test_api_factory.py (1)
292-326: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate typed-client mock helpers across 5 test files. Each file reimplements its own
_resp/_page_resp/_client_error/_conflict_error/_binary_respto wrap payloads/errors for the typed-client migration — same shape, copy-pasted.
packages/nmp_common/tests/api_factory/test_api_factory.py#L292-L326: extract_resp,_page_resp,_client_errorto a shared test-utils module.packages/nmp_common/tests/jobs/test_result_manager.py#L14-L25: replace local_resp/_conflict_errorwith the shared helper.plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py#L22-L27: replace local_client_errorwith the shared helper.plugins/nemo-safe-synthesizer/tests/unit/test_sdk.py#L21-L37: replace local_resp/_binary_respwith the shared helper (add_binary_respto the shared module).packages/nmp_testing/tests/unit/test_jobs.py#L29-L39: replace local_respwith the shared helper.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nmp_common/tests/api_factory/test_api_factory.py` around lines 292 - 326, Extract the duplicated typed-client mock helpers into a shared test-utils module: move _resp, _page_resp, and _client_error from packages/nmp_common/tests/api_factory/test_api_factory.py#L292-L326; replace local _resp/_conflict_error in packages/nmp_common/tests/jobs/test_result_manager.py#L14-L25, _client_error in plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py#L22-L27, _resp/_binary_resp in plugins/nemo-safe-synthesizer/tests/unit/test_sdk.py#L21-L37 (adding _binary_resp to the shared module), and _resp in packages/nmp_testing/tests/unit/test_jobs.py#L29-L39 with imports from that module, preserving each helper’s existing response and error behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/nemo_platform_plugin/tests/test_jobs_filter.py`:
- Around line 350-360: The _round_trip method only validates captured
query_params and does not exercise JobsClient serialization. Add a mocked-HTTP
test that invokes the actual JobsClient request path and asserts the emitted
request contains exactly one filter parameter whose value is unchanged JSON,
while preserving the existing parse_json_filter validation.
In `@packages/nmp_common/src/nmp/common/errors/sdk_exception_handlers.py`:
- Around line 88-93: Sanitize the request-derived method and URL path before
passing them to the debug log in the NemoClient exception handler. Update the
logging flow around the “Converting NemoClient exception to HTTP response”
message to use the project’s existing log-sanitization utility or pattern, while
preserving the logged request context and status code.
In
`@packages/nmp_customization_common/src/nmp/customization_common/training/progress.py`:
- Around line 68-78: The PlatformJobTaskUpdate construction in the training
progress update flow explicitly sends empty status_details and error_details
maps, potentially clearing existing values. Build the update payload so
status_details and error_details are included only when their corresponding
inputs have values, while preserving the existing status and other update
fields.
In `@services/core/models/src/nmp/core/models/api/v2/models.py`:
- Around line 311-326: Update the logging in the job creation flow to avoid
emitting the raw job_resp object. Change the logger.info call after
jobs.create_job to include only the job_resp.id or job_resp.name, while
preserving the existing creation and error-handling behavior.
---
Nitpick comments:
In `@packages/nemo_platform_plugin/tests/test_jobs_filter.py`:
- Line 44: Replace the string-based "_CapturingSdk" annotations with direct
_CapturingSdk type references in both _forwarded_filter and the sibling
annotation at packages/nemo_platform_plugin/tests/test_jobs_filter.py lines
350-351.
In `@packages/nmp_common/tests/api_factory/test_api_factory.py`:
- Around line 292-326: Extract the duplicated typed-client mock helpers into a
shared test-utils module: move _resp, _page_resp, and _client_error from
packages/nmp_common/tests/api_factory/test_api_factory.py#L292-L326; replace
local _resp/_conflict_error in
packages/nmp_common/tests/jobs/test_result_manager.py#L14-L25, _client_error in
plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py#L22-L27,
_resp/_binary_resp in
plugins/nemo-safe-synthesizer/tests/unit/test_sdk.py#L21-L37 (adding
_binary_resp to the shared module), and _resp in
packages/nmp_testing/tests/unit/test_jobs.py#L29-L39 with imports from that
module, preserving each helper’s existing response and error behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7651bef6-c4d4-4dac-b36e-e97697446b70
📒 Files selected for processing (27)
packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/quickstart/cli.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/api_factory.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/result_manager.pypackages/nemo_platform_plugin/tests/jobs/test_client.pypackages/nemo_platform_plugin/tests/test_jobs_filter.pypackages/nmp_common/src/nmp/common/errors/sdk_exception_handlers.pypackages/nmp_common/tests/api_factory/test_api_factory.pypackages/nmp_common/tests/jobs/conftest.pypackages/nmp_common/tests/jobs/test_result_manager.pypackages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_progress_reporter.pypackages/nmp_customization_common/src/nmp/customization_common/training/progress.pypackages/nmp_testing/src/nmp/testing/e2e/jobs.pypackages/nmp_testing/tests/unit/test_jobs.pyplugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.pyplugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/api/v2/jobs/endpoints.pyplugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/sdk/job.pyplugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/sdk/resources.pyplugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/tasks/safe_synthesizer/__main__.pyplugins/nemo-safe-synthesizer/tests/unit/test_jobs.pyplugins/nemo-safe-synthesizer/tests/unit/test_local_run.pyplugins/nemo-safe-synthesizer/tests/unit/test_sdk.pyplugins/nemo-safe-synthesizer/tests/unit/test_task_upload_results.pyservices/automodel/tests/test_progress_reporter.pyservices/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.pyservices/core/entities/tests/controllers/test_workspace_cleanup.pyservices/core/jobs/src/nmp/core/jobs/controllers/backends/docker.pyservices/core/models/src/nmp/core/models/api/v2/models.py
🚧 Files skipped from review as they are similar to previous changes (1)
- services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py
The Jobs typed-client migration changed generated schema descriptions and dropped additionalProperties: false on KubernetesImagePullSecret, and the quickstart CLI consumer migration was never re-vendored into the SDK. Both left CI lint red on a branch whose spec was claimed unchanged. Restore extra=forbid on the plugin ImagePullSecret and the original one-line docstrings on the server subclasses (DockerJobExecutionProfileConfig, DockerJobNetworkConfig, KubernetesJobStorageConfig, KubernetesVolume, KubernetesVolumeMount, PlatformJobStepWithContext), moving the developer notes to comments, so make refresh-openapi regenerates a byte-identical spec with no Stainless SDK regen. Re-vendor the quickstart CLI and apply ruff formatting to two test files. Clears lint-python-style, lint-openapi, lint-python-sdk, lint-web-sdk, lint-cli, and lint-sdk-vendored. Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>
Address CodeRabbit review findings on PR #585: - DockerJobExecutionProfile / VolcanoJobExecutionProfile supports_persistent_storage always returned True: the storage object is never None (default_factory) and pvc_name is a str that is never None. Gate on volume_name / pvc_name != "" to mirror the Kubernetes profile. - providers.py: return Self from the SubprocessExecutionProvider validator instead of the quoted class name, matching spec.py and repo style. - sdk_exception_handlers.py: strip CR/LF from request method/path before debug logging (log-injection hardening) in all three handlers. - models.py: log the created job name instead of the full response model, which avoids emitting ownership metadata at INFO. The equivalent backend-controller supports_persistent_storage in volcano_job.py carries the same pre-existing bug but is outside this PR's scope and left unchanged. Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>
…o-nemoclient-typed-http-client Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py (1)
854-882: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winWrap the status write in this handler. If
update_job_step_status(...)raises here, theFailedToScheduleErrorpath exits without logging, and the dropped worker future hides the failure. The step stays stuck without diagnostics.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py` around lines 854 - 882, The FailedToScheduleError handler in run_container must protect the _jobs.update_job_step_status call from raising without diagnostics. Wrap that status update in its own exception handling and log any failure, while preserving the existing FailedToScheduleError logging and the finally block’s admission release.
🧹 Nitpick comments (2)
packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py (1)
63-73: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache
TypeAdapterinstead of rebuilding per call.Pydantic recommends instantiating
TypeAdapteronce and reusing it — building it per call re-analyzes the type into a core schema every time, which is non-trivial overhead in a client hot path (every non-BaseModelresponse).♻️ Proposed fix
+from functools import lru_cache + + +@lru_cache(maxsize=None) +def _type_adapter(response_type: type) -> TypeAdapter: + return TypeAdapter(response_type) + + def _parse_json_body(response_type: type, data: Any) -> Any: ... if isinstance(response_type, type) and issubclass(response_type, BaseModel): return response_type.model_validate(data) - return TypeAdapter(response_type).validate_python(data) + return _type_adapter(response_type).validate_python(data)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py` around lines 63 - 73, Cache the TypeAdapter instances used by _parse_json_body instead of constructing one on every non-BaseModel response. Key the cache by response_type, reuse the cached adapter for validation, and preserve the existing BaseModel model_validate path.packages/nmp_customization_common/src/nmp/customization_common/training/progress.py (1)
68-69: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache the typed
JobsClientonce instead of rebuilding it per call.
update_task/fetch_current_metricsboth callclient_from_platform(self._sdk, JobsClient)on every invocation. Every other migrated consumer in this PR (JobBackend.__init__,JobScheduler.__init__) builds this once and reuses it.update_taskis on a hot training-progress-reporting path, so recreating the client per call is unnecessary overhead and inconsistent with the rest of the migration.♻️ Proposed fix
def __init__(self, job_ctx: NMPJobContext, service_name: str): self._job_ctx = job_ctx self._sdk = get_task_sdk(service_name) + self._jobs = client_from_platform(self._sdk, JobsClient) self._is_main_rank = int(os.environ.get("RANK", "0")) == 0 self._max_steps = 0 self._num_epochs = 0 ... try: - jobs = client_from_platform(self._sdk, JobsClient) - jobs.update_job_step_task( + self._jobs.update_job_step_task( ... try: - jobs = client_from_platform(self._sdk, JobsClient) - task = jobs.get_job_step_task( + task = self._jobs.get_job_step_task(Also applies to: 88-89
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nmp_customization_common/src/nmp/customization_common/training/progress.py` around lines 68 - 69, Initialize and store the typed JobsClient once during the progress component’s construction, following the existing JobBackend.__init__ and JobScheduler.__init__ pattern. Update both update_task and fetch_current_metrics to reuse that cached client instead of calling client_from_platform(self._sdk, JobsClient) per invocation, while preserving their existing behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@services/core/jobs/tests/conftest.py`:
- Around line 310-316: Update the patch setup around the fixture’s patchers loop
to use context-managed cleanup, such as contextlib.ExitStack, so each
successfully started patch is registered for teardown immediately. Ensure an
exception from any later p.start() still stops all previously started patches,
while preserving the existing mock_client yield behavior.
---
Outside diff comments:
In `@services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py`:
- Around line 854-882: The FailedToScheduleError handler in run_container must
protect the _jobs.update_job_step_status call from raising without diagnostics.
Wrap that status update in its own exception handling and log any failure, while
preserving the existing FailedToScheduleError logging and the finally block’s
admission release.
---
Nitpick comments:
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py`:
- Around line 63-73: Cache the TypeAdapter instances used by _parse_json_body
instead of constructing one on every non-BaseModel response. Key the cache by
response_type, reuse the cached adapter for validation, and preserve the
existing BaseModel model_validate path.
In
`@packages/nmp_customization_common/src/nmp/customization_common/training/progress.py`:
- Around line 68-69: Initialize and store the typed JobsClient once during the
progress component’s construction, following the existing JobBackend.__init__
and JobScheduler.__init__ pattern. Update both update_task and
fetch_current_metrics to reuse that cached client instead of calling
client_from_platform(self._sdk, JobsClient) per invocation, while preserving
their existing behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ddcb727a-84be-431c-b4b8-2448b81e977f
📒 Files selected for processing (25)
packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/api_factory.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/client.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/endpoints.pypackages/nemo_platform_plugin/tests/jobs/test_client.pypackages/nemo_platform_plugin/tests/jobs/test_endpoints.pypackages/nmp_customization_common/src/nmp/customization_common/training/progress.pyservices/core/jobs/src/nmp/core/jobs/controllers/backends/base.pyservices/core/jobs/src/nmp/core/jobs/controllers/backends/docker.pyservices/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.pyservices/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/kubernetes_job.pyservices/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.pyservices/core/jobs/src/nmp/core/jobs/controllers/diagnostics.pyservices/core/jobs/src/nmp/core/jobs/controllers/reconciler.pyservices/core/jobs/src/nmp/core/jobs/controllers/scheduler.pyservices/core/jobs/tests/conftest.pyservices/core/jobs/tests/controllers/client_mocks.pyservices/core/jobs/tests/controllers/test_base.pyservices/core/jobs/tests/controllers/test_diagnostics.pyservices/core/jobs/tests/controllers/test_docker_backend.pyservices/core/jobs/tests/controllers/test_reconciler.pyservices/core/jobs/tests/controllers/test_scheduler.pyservices/core/jobs/tests/controllers/test_subprocess_backend.pyservices/core/jobs/tests/test_jobs_client.pyservices/core/models/src/nmp/core/models/api/v2/models.py
💤 Files with no reviewable changes (2)
- packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/client.py
- packages/nemo_platform_plugin/tests/jobs/test_endpoints.py
🚧 Files skipped from review as they are similar to previous changes (5)
- packages/nemo_platform_plugin/tests/jobs/test_client.py
- services/core/models/src/nmp/core/models/api/v2/models.py
- packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/endpoints.py
- packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/api_factory.py
- services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
The jobs controller tests import a shared helper via a bare sibling import (`from client_mocks import ...`). Under the repo's importlib import mode, pytest does not add each test file's directory to sys.path, so this only resolves in scoped runs and fails collection in the full `pytest -m unit` suite with ModuleNotFoundError: No module named 'client_mocks' (6 collection errors). Add services/core/jobs/tests/controllers to pytest.ini pythonpath, matching the existing convention for plugins/nemo-deployments/tests which use the same sibling-helper import pattern. Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
…t-typed-http-client' of github.com:NVIDIA-NeMo/nemo-platform into mgrossman/aircore-874-migrate-jobs-service-to-nemoclient-typed-http-client Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
…o-nemoclient-typed-http-client Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
…o-nemoclient-typed-http-client
Summary
Migrates the Jobs service and its production consumers from the Stainless-generated
sdk.jobs.*surface to the typedNemoClientHTTP client.This PR is AIRCORE-874, part of AIRCORE-827. It follows the typed-client migration pattern documented in
MIGRATION.md.What changed
Typed Jobs client
JobsClientand asynchronousAsyncJobsClientimplementations for all 21 shipped Jobs endpoints.rerun_job: the/rerunroute is test-only and is not mounted by the released service.Jobs wire types
nemo_platform_plugin.jobsas shared leaf types.to_k8s()only in the service, avoidingkubernetesanddockerdependencies in the plugin package.Consumer migration
client_from_platform(..., JobsClient)orAsyncJobsClient.api_factory, result management, customization progress reporting, Safe Synthesizer, Data Designer helpers, model APIs, entity workspace cleanup, quickstart, andnmp_testinghelpers.NemoHTTPErrorhandling alongside the existing Stainless error handling so typed-client service-to-service failures retain the expected HTTP response behavior.filter[status], matching the Jobs route parser rather than sending an ignored flatfilter=value.Pagination and shared-client changes
Jobs logs use cursor pagination while Jobs and steps use the platform's standard offset pagination. This PR extends the shared client so both are first-class instead of adding a Jobs-specific page method:
Paginated[Item, Strategy], with offset pagination as the default andCursorPaginationforlist_job_logs..page()returns typed items plus strategy-specific metadata; iteration and.items()auto-fetch subsequent pages in both sync and async clients.page,page_size,current_page_size,total_pages, andtotal_results).total,next_page, andprev_page, and sends subsequent cursors aspage_cursor.The shared client also now:
TypeAdapters, including lists and unions.NemoResponseValidationError.time.sleepandasyncio.sleep, respectively. The default retryable responses are429,502,503, and504, plus transport errors.typing_extensions.TypedDict.Intentional service/client asymmetries
Filtersubclasses; client endpoints accept encoded query parameters.Out of scope
nemo jobsCLI commands: AIRCORE-893*Paramaliases fromapi_factory: AIRCORE-922sdk.customization.jobs.*, which is a separate API namespace.Validation
GitHub CI is green on the final branch, including:
nemo-platformandnemo-platform-pluginwheel build/test matrices on Python 3.11 through 3.14Focused coverage includes prepared request shapes, sync/async response parsing, strict offset and cursor pagination, retry behavior, typed-page inference, and
AsyncJobsClientagainst the in-memory Jobs ASGI application.