refactor(models): migrate downstream consumers to typed NemoClient - #1277
Draft
maxdubrinsky wants to merge 2 commits into
Draft
refactor(models): migrate downstream consumers to typed NemoClient#1277maxdubrinsky wants to merge 2 commits into
maxdubrinsky wants to merge 2 commits into
Conversation
Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>
Contributor
|
The NemoClient migration rewrote fetch_model_entity to resolve models via client_from_platform(sdk, AsyncModelsClient).get_model(...).data() instead of sdk.models.retrieve(...). The unsloth plugin tests were not updated: they still stubbed sdk.models.retrieve and patched client_from_platform with a files-only mock returned for every call, so the new AsyncModelsClient call got the files mock back, get_model() returned an unawaited AsyncMock coroutine, and response.data() yielded that coroutine. transform_input_to_output then hit AttributeError: 'coroutine' object has no attribute 'spec'. Dispatch client_from_platform by client class: return a models mock (whose get_model(...).data() yields the stub model entity) for AsyncModelsClient and the files mock for AsyncFilesClient. Mirrors the pattern in the PR's new test_models_client_migration.py. Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>
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
1 participant
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
Migrates the downstream model consumers off the Stainless-generated
sdk.inference.*/sdk.models.*CRUD surface onto the typedAsyncModelsClientfrom the foundation landed in #993. This is the "downstream consumers" half of AIRCORE-876; the coreservices/core/modelscontrollers are migrated in a separate sibling PR.Related Issue
Linear: AIRCORE-876 (Migrate Models service to NemoClient typed HTTP client), parent epic AIRCORE-827.
Changes
services/platform-seed:seed_model_providernow routes throughclient_from_platform(sdk, AsyncModelsClient)+models.create_provider(..., body=CreateModelProviderRequest(...))instead ofsdk.inference.providers.create(...).services/core/entitiesWorkspaceCleanup._cleanup_deployments: routes deployment listing/deletion throughAsyncModelsClient.list_deployments/delete_deployment(preserving theHeartbeatMixinheartbeat calls recently added on main).packages/nmp_customization_common(model entity run task + platform client) andpackages/data_designer_nemomodel provider: switch to the typed models client.services/automodel,services/rl,services/unslothjob compilers: adopt the typed client.services/guardrailsmodel routing / config utils: adopt the typed client.test_models_client_migration.pyfocused tests fordata_designer_nemo,nmp_customization_common, andautomodel; updates existing suites (test_workspace_cleanup,test_platform_seed_runner,test_model_routing,test_compiler,test_model_entity, etc.) to the typed client mocking style.The branch was rebased onto
origin/mainso the typed-client foundation comes from #993; the divergent in-branch foundation copy was dropped.Type of Change
Quality Gates
Verification
Signed-off-by:traileruv run pre-commit run -apasses, or any blocked checks are identified belowTargeted validation:
uv run ruff check <changed .py files>— All checks passed.uv run --frozen pytest -q <affected test files>— 101 passed in 0.45s (migration tests + updated existing suites for data_designer_nemo, nmp_customization_common, automodel, core/entities workspace_cleanup, guardrails, platform-seed, rl, unsloth).uv run --frozen ty check <changed packages>— 300 diagnostics, all pre-existing onorigin/main(308 on main for the same set); the migration introduces no new type errors and removes 8 by replacing untyped SDK dict access with typedAsyncModelsClientcalls.uv run pre-commit run --files <changed files>— ruff, ruff format, ty, copyright headers, merge-conflict checks all Passed.af4f86533,Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>matches author email.