feat(client): NemoClient platform provider + request-scoped FastAPI DI - #800
Merged
Merged
Conversation
maxdubrinsky
changed the title
maxdubrinsky
changed the base branch from
aircore-883-nemoclient-provider-foundation/md
to
main
July 21, 2026 15:49
Contributor
|
maxdubrinsky
marked this pull request as ready for review
July 21, 2026 16:25
Contributor
📝 WalkthroughWalkthroughNemo client creation now resolves pluggable providers through entry points or explicit overrides. Platform factories support routing, authentication, workspace propagation, and task delegation. Service dependencies expose request-scoped clients with synchronized transport lifecycle management. ChangesNemo client provider flow
Sequence Diagram(s)sequenceDiagram
participant ServiceRequest
participant FastAPI
participant DependencyProvider
participant PlatformNemoClientProvider
participant AsyncNemoClient
ServiceRequest->>FastAPI: resolve get_nemo_client
FastAPI->>DependencyProvider: request-scoped client dependency
DependencyProvider->>PlatformNemoClientProvider: create async client
PlatformNemoClientProvider->>AsyncNemoClient: construct with workspace, headers, and routing
AsyncNemoClient-->>FastAPI: return request-scoped client
FastAPI-->>ServiceRequest: inject client
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
packages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.py (1)
187-253: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract shared entry-point resolution helper.
_resolve_provideris duplicated verbatim between the two files (dedup, deterministic-conflict RuntimeError, load/construct, protocol validation) — only the group name, protocol type, and default class differ. One shared helper avoids future divergence when either copy gets a bug fix, and is a natural place to add a lock around the scan+load+construct+cache sequence.
packages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.py#L187-L253: replace with a call to a shared_resolve_provider(group, protocol, default_factory, kind_label)helper.packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.py#L298-L337: replace with a call to the same shared helper, passing"nemo.sdk_provider",SDKProvider,DefaultSDKProvider.🤖 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_provider.py` around lines 187 - 253, Extract the duplicated entry-point resolution logic from _resolve_provider in packages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.py#L187-L253 and packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.py#L298-L337 into one shared _resolve_provider(group, protocol, default_factory, kind_label) helper. Preserve deterministic conflict detection, entry-point loading and construction, protocol validation, default fallback, caching, and add locking around the scan/load/construct/cache sequence; update both providers to call it with their respective group, protocol, and default factory.
🤖 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.
Nitpick comments:
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.py`:
- Around line 187-253: Extract the duplicated entry-point resolution logic from
_resolve_provider in
packages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.py#L187-L253
and
packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.py#L298-L337
into one shared _resolve_provider(group, protocol, default_factory, kind_label)
helper. Preserve deterministic conflict detection, entry-point loading and
construction, protocol validation, default fallback, caching, and add locking
around the scan/load/construct/cache sequence; update both providers to call it
with their respective group, protocol, and default factory.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 72cf3fb6-a3a0-4a76-b118-1f0adb68251b
📒 Files selected for processing (14)
packages/nemo_platform/pyproject.tomlpackages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.pypackages/nemo_platform_plugin/tests/test_client_provider.pypackages/nemo_platform_plugin/tests/test_dependencies.pypackages/nemo_platform_plugin/tests/test_sdk_provider.pypackages/nmp_common/pyproject.tomlpackages/nmp_common/src/nmp/common/client_factory.pypackages/nmp_common/src/nmp/common/service/__init__.pypackages/nmp_common/src/nmp/common/service/base.pypackages/nmp_common/src/nmp/common/service/dependencies.pypackages/nmp_common/tests/client_factory/test_client_factory.pypackages/nmp_common/tests/nmp_common/test_common_service.py
Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>
Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>
Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>
maxdubrinsky
force-pushed
the
aircore-904-add-nemoclient-fastapi-di/md
branch
from
August 4, 2026 20:25
fa95fa3 to
fe224c3
Compare
…ty, UDS)
The NemoClient providers advertised task-container support but had no
get_task_sdk() equivalent, so task callers lost job-creator delegation and
workload-identity auth, and the platform provider mis-routed UDS endpoints.
- Add get_task_nemo_client / get_async_task_nemo_client to the
NemoClientProvider protocol and both providers, mirroring get_task_sdk:
service:{name} + X-NMP-Principal-On-Behalf-Of* from NMP_PRINCIPAL, or
workload-identity bearer-token exchange when NMP_WORKLOAD_IDENTITY_TOKEN_FILE
is set (dropping trusted X-NMP-* headers). Generic get_nemo_client(as_service)
stays undelegated on purpose for background controllers.
- Make the platform get_nemo_client / get_async_nemo_client workload-identity
aware via _should_bootstrap_workload_identity, matching get_platform_sdk.
- Route UDS endpoints through resolve_platform_endpoint().connect_base_url and
an endpoint-appropriate socket transport instead of the raw unix:// base_url
and a generic shared TCP client.
Add regression tests for delegation, workload identity, and UDS transport.
Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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_client_provider.py`:
- Around line 326-340: Update the _CapturingProvider test double in
test_public_functions_pass_workspace_through to implement the missing
get_task_nemo_client and get_async_task_nemo_client methods required by
NemoClientProvider. Match the existing client-provider method behavior and
return the appropriate NemoClient or AsyncNemoClient while preserving captured
kwargs.
🪄 Autofix
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: 6b3f8370-8af0-4185-87f3-394e904abccc
📒 Files selected for processing (16)
packages/nemo_platform/pyproject.tomlpackages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.pypackages/nemo_platform_plugin/tests/test_client_provider.pypackages/nemo_platform_plugin/tests/test_dependencies.pypackages/nemo_platform_plugin/tests/test_nemo_client_task_auth.pypackages/nemo_platform_plugin/tests/test_sdk_provider.pypackages/nmp_common/pyproject.tomlpackages/nmp_common/src/nmp/common/client_factory.pypackages/nmp_common/src/nmp/common/service/__init__.pypackages/nmp_common/src/nmp/common/service/base.pypackages/nmp_common/src/nmp/common/service/dependencies.pypackages/nmp_common/tests/client_factory/test_client_factory.pypackages/nmp_common/tests/nmp_common/test_common_service.pypackages/nmp_common/tests/nmp_common/test_dependency_provider.py
🚧 Files skipped from review as they are similar to previous changes (8)
- packages/nemo_platform/pyproject.toml
- packages/nmp_common/src/nmp/common/service/dependencies.py
- packages/nmp_common/src/nmp/common/service/init.py
- packages/nemo_platform_plugin/tests/test_dependencies.py
- packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py
- packages/nmp_common/pyproject.toml
- packages/nmp_common/src/nmp/common/service/base.py
- packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.py
get_http_client() hardcoded DefaultAsyncHttpxClient(), a TCP client, and injected it into the SDK and NemoClient factories. Those factories skip their own transport selection when a client is supplied, so under a unix:// endpoint the request-scoped clients got the UDS base_url (http://nemo-platform.local) paired with a TCP transport (_uds=None) and service-to-service requests failed. Build the cached client through resolve_platform_endpoint().async_sdk_http_client() so it binds the socket on UDS and stays the default TCP client otherwise. Add provider-layer UDS/TCP regression coverage. Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/nmp_common/tests/nmp_common/test_dependency_provider.py`:
- Around line 39-53: Convert the provider-owned client tests to async tests and
await DependencyProvider.close() in each test’s finally block before calling
Configuration.clear_cache(). Apply this to
packages/nmp_common/tests/nmp_common/test_dependency_provider.py lines 39-53,
56-64, and 67-74, preserving each test’s existing assertions and setup.
🪄 Autofix
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: 6233d8c7-43b0-4149-a9dd-91a443eac717
📒 Files selected for processing (3)
packages/nmp_common/src/nmp/common/service/base.pypackages/nmp_common/tests/nmp_common/test_common_service.pypackages/nmp_common/tests/nmp_common/test_dependency_provider.py
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/nmp_common/src/nmp/common/service/base.py
This was referenced Aug 6, 2026
marcusds
approved these changes
Aug 12, 2026
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Aug 12, 2026
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Aug 12, 2026
maxdubrinsky
enabled auto-merge
August 12, 2026 17:55
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
packages/nmp_common/src/nmp/common/client_factory.py (1)
21-21: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse concrete runtime annotations.
Remove postponed annotations. All referenced types are available at runtime.
packages/nmp_common/src/nmp/common/client_factory.py#L21-L21: removefrom __future__ import annotations.packages/nemo_platform_plugin/tests/test_nemo_client_task_auth.py#L14-L14: removefrom __future__ import annotations.As per coding guidelines: “Prefer concrete type hints over string-based annotations.”
🤖 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/src/nmp/common/client_factory.py` at line 21, Remove the postponed-annotations import from client_factory.py and test_nemo_client_task_auth.py so their type annotations are evaluated concretely at runtime; no other changes are required.Source: Coding guidelines
packages/nmp_common/tests/nmp_common/test_dependency_provider.py (1)
34-36: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNo shared helper for httpx UDS inspection. Both files reach into the private chain
_transport._pool._udsto assert socket binding. An httpx internal rename breaks every site independently.
packages/nmp_common/tests/nmp_common/test_dependency_provider.py#L34-L36: move_uds_ofinto a shared test helper module so other packages can import it.packages/nmp_common/tests/client_factory/test_client_factory.py#L374-L388: replace the two inlinetransport._pool._udstraversals with the shared_uds_ofhelper.🤖 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/nmp_common/test_dependency_provider.py` around lines 34 - 36, Move the _uds_of helper from packages/nmp_common/tests/nmp_common/test_dependency_provider.py:34-36 into a shared test helper module and update imports accordingly. In packages/nmp_common/tests/client_factory/test_client_factory.py:374-388, replace both direct transport._pool._uds traversals with _uds_of, preserving the existing assertions.packages/nemo_platform_plugin/tests/test_sdk_provider.py (1)
253-271: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_FakeEntryPointduplicates the identical class inpackages/nemo_platform_plugin/tests/test_client_provider.py(lines 183-200). Both live in the same test package. Move it toconftest.pyand import it in both modules.🤖 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_sdk_provider.py` around lines 253 - 271, Move the shared _FakeEntryPoint test helper from test_sdk_provider.py into the package conftest.py, preserving its constructor fields and load() behavior. Import and reuse that conftest definition in both test_sdk_provider.py and test_client_provider.py, removing their duplicate local class definitions.
🤖 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/nmp_common/tests/client_factory/test_client_factory.py`:
- Around line 173-193: Update the internal-header assertions in
test_merges_otel_propagation_headers_without_adding_internal_auth and
test_explicit_auth_headers_win_over_conflicting_otel_context to compare header
names case-insensitively. Assert that no key whose lowercase form is
x-nmp-internal remains in the first test, and that only the expected internal
value is present without any conflicting casing variant in the second.
In `@packages/nmp_common/tests/nmp_common/test_dependency_provider.py`:
- Around line 87-94: The test exercising shared HTTP client injection must close
the materialized transport. Convert the test containing
provider.get_http_client() to async, add the project’s pytest asyncio marker,
and await provider.close() after the existing assertions.
---
Nitpick comments:
In `@packages/nemo_platform_plugin/tests/test_sdk_provider.py`:
- Around line 253-271: Move the shared _FakeEntryPoint test helper from
test_sdk_provider.py into the package conftest.py, preserving its constructor
fields and load() behavior. Import and reuse that conftest definition in both
test_sdk_provider.py and test_client_provider.py, removing their duplicate local
class definitions.
In `@packages/nmp_common/src/nmp/common/client_factory.py`:
- Line 21: Remove the postponed-annotations import from client_factory.py and
test_nemo_client_task_auth.py so their type annotations are evaluated concretely
at runtime; no other changes are required.
In `@packages/nmp_common/tests/nmp_common/test_dependency_provider.py`:
- Around line 34-36: Move the _uds_of helper from
packages/nmp_common/tests/nmp_common/test_dependency_provider.py:34-36 into a
shared test helper module and update imports accordingly. In
packages/nmp_common/tests/client_factory/test_client_factory.py:374-388, replace
both direct transport._pool._uds traversals with _uds_of, preserving the
existing assertions.
🪄 Autofix
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: df78cd4b-728b-4365-a270-1c453e57271a
📒 Files selected for processing (16)
packages/nemo_platform/pyproject.tomlpackages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.pypackages/nemo_platform_plugin/tests/test_client_provider.pypackages/nemo_platform_plugin/tests/test_dependencies.pypackages/nemo_platform_plugin/tests/test_nemo_client_task_auth.pypackages/nemo_platform_plugin/tests/test_sdk_provider.pypackages/nmp_common/pyproject.tomlpackages/nmp_common/src/nmp/common/client_factory.pypackages/nmp_common/src/nmp/common/service/__init__.pypackages/nmp_common/src/nmp/common/service/base.pypackages/nmp_common/src/nmp/common/service/dependencies.pypackages/nmp_common/tests/client_factory/test_client_factory.pypackages/nmp_common/tests/nmp_common/test_common_service.pypackages/nmp_common/tests/nmp_common/test_dependency_provider.py
🚧 Files skipped from review as they are similar to previous changes (9)
- packages/nmp_common/src/nmp/common/service/dependencies.py
- packages/nmp_common/pyproject.toml
- packages/nmp_common/src/nmp/common/service/init.py
- packages/nemo_platform/pyproject.toml
- packages/nemo_platform_plugin/tests/test_dependencies.py
- packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.py
- packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py
- packages/nmp_common/src/nmp/common/service/base.py
- packages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Consolidates the NemoClient foundation into one PR: the platform provider (AIRCORE-883) plus the request-scoped FastAPI DI seam (AIRCORE-904). This was previously split across #799 (provider) and #800 (DI, stacked on #799's branch). #800 is now retargeted to
mainand carries both halves; #799 is closed in favor of this PR.Provider foundation (AIRCORE-883)
NemoClientProviderprotocol with explicit override, entry-point discovery, and a zero-nmp-commonenv-var default (DefaultNemoClientProvider).PlatformNemoClientProvider) backed by per-service URL routing, shared HTTP transports, request principal headers, and OTEL propagation.nmp-commonand the bundlednemo-platformpackage.Request-scoped DI (AIRCORE-904)
get_nemo_clientdependency + exact common-package re-export.AsyncNemoClientwith current principal + trace headers, sharing one provider-owned transport.DependencyProvidertransport/SDK init made thread-safe and order-independent; provider-owned transport closed exactly once on shutdown.How it's used (downstream)
In-process (service / middleware / controller), async:
In a job's task container the same call falls back to the env-var
DefaultNemoClientProvider(NMP_BASE_URL/NMP_PRINCIPALinjected by the jobs backend). The plugin imports onlynemo_platform_plugin.client_provider— nevernmp-common, never Stainless.Per-service migration playbook:
packages/nemo_platform_plugin/src/nemo_platform_plugin/client/MIGRATION.md.Unblocks
client_from_platform(...)forDepends(get_nemo_client).endpoints.py), AIRCORE-893 (CLI generator → typed clients).Follow-up (not in this PR)
client_clsparam soget_async_nemo_client(AsyncFilesClient)returns a typed client directly (today it returns a base client; use.send(endpoints.*)).close()/ context-manager support — AIRCORE-933.Validation
git diff --check, focusedty, and wheel entry-point inspection passed. Independent adversarial review found no blocker on either half.main(while stacked on a feature branch, most CI jobs were suppressed).Linear: https://linear.app/nvidia/issue/AIRCORE-883 , https://linear.app/nvidia/issue/AIRCORE-904
Summary by CodeRabbit
New Features
Bug Fixes
Tests