Skip to content

fix: restore model customization reuse state - #6264

Open
papriwal wants to merge 1 commit into
aws:masterfrom
papriwal:fix/modelbuilder-reuse-compute-requirements
Open

fix: restore model customization reuse state#6264
papriwal wants to merge 1 commit into
aws:masterfrom
papriwal:fix/modelbuilder-reuse-compute-requirements

Conversation

@papriwal

Copy link
Copy Markdown
Collaborator

Issue #, if available:

Description of changes:

What problem does this solve?

ModelBuilder performs Model reuse and endpoint reuse independently. When build(reuse_resources=True) found an existing package-backed Model but deployment could not reuse an endpoint, the SDK continued through the endpoint-creation path without rebuilding the Model.

Some deployment state normally populated while building a customization Model—such as recipe-derived compute requirements and the LoRA adapter artifact location—was therefore unavailable. This could cause the fallback deployment to fail or reach resource creation without the required configuration.

The existing integration coverage could also discover unrelated resources left by another test run. That made the reuse scenario non-deterministic and could hide regressions that created duplicate resources.

Why is this change needed?

Reusing a Model should not require its endpoint to be reusable. If endpoint reuse misses, ModelBuilder must create the endpoint using the same customization metadata that a fresh build would provide, without creating a second Model.

Explicitly supplied resource requirements must remain authoritative. Cached recipe settings should only be restored when the caller did not provide requirements, while LoRA deployments must still have a valid adapter artifact location.

The integration scenario also needs to prove that the exact Model and endpoint created by the test are reused, rather than accepting arbitrary resources from a shared test environment.

What changed?

  • Track whether the current build reused an existing Model and reset that state between build attempts.
  • Restore package-backed customization deployment state only when a reused Model must enter the endpoint-creation path.
  • Preserve explicitly supplied resource requirements and copy count.
  • Resolve LoRA adapter artifacts consistently for supported training-job, trainer, and model-package sources.
  • Fail before endpoint resource creation when required customization state cannot be resolved.
  • Keep healthy endpoint reuse on the existing fast path without unnecessary state restoration.
  • Consolidate the scheduled reuse integration coverage around uniquely named, test-owned resources.
  • Require exact Model and endpoint discovery before exercising the public reuse flow.
  • Guard the integration scenario against unexpected duplicate Model or endpoint creation and clean up only resources owned by that test.

This change does not add or modify public API parameters and introduces no new dependencies.

Testing

  • Added unit coverage for reused-Model marker lifecycle and reset behavior.
  • Added coverage for exact Model and endpoint reuse without state preparation.
  • Added LoRA and non-LoRA coverage for endpoint creation after Model reuse.
  • Added coverage for preserving explicit resource requirements and copy count.
  • Added coverage for supported LoRA artifact paths and unsupported source rejection before writes.
  • Ran the focused ModelBuilder, inference-configuration, model-reuse, recipe-selection, and instance-selection unit suites.
  • Verified Python compilation and diff whitespace checks.
  • Verified integration-test collection, scheduled-test selection, and removal of the non-hermetic reuse tests.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Restore deployment metadata when ModelBuilder reuses a package-backed
Model but cannot reuse its endpoint. Preserve explicit resource
requirements and recover recipe compute or LoRA adapter state only
when creating the replacement endpoint.

Make scheduled reuse coverage deterministic and guard against duplicate
Model and endpoint writes.
@papriwal
papriwal deployed to auto-approve September 11, 2026 20:54 — with GitHub Actions Active
@papriwal
papriwal deployed to auto-approve September 11, 2026 20:54 — with GitHub Actions Active
@papriwal
papriwal deployed to auto-approve September 11, 2026 20:54 — with GitHub Actions Active
@github-actions

github-actions Bot commented Sep 11, 2026

Copy link
Copy Markdown

Claude finished @papriwal's task in 2m 47s —— View job


Review: PR #6264 — restore model customization reuse state

  • Read the PR diff
  • Analyze model_builder.py changes for correctness
  • Review reuse-state / LoRA resolution logic
  • Review test coverage
  • Post review feedback

(The inline-comment tool wasn't available in this run, so findings are consolidated below with file:line references.)


Summary

This is a careful, well-tested bug fix. The core problem — a reused package-backed Model that then misses endpoint reuse would reach _deploy_model_customization without the recipe-derived compute requirements or the LoRA adapter URI that a fresh build would have populated — is addressed cleanly:

  • _built_model_was_reused is reset in three correct places: __post_init__ (model_builder.py:460), at the top of build() (:4403), and in _reset_build_state() (:4293), and set on the reuse hit (:4469). Lifecycle looks sound.
  • Moving peft_type / base_model_recipe_name resolution before EndpointConfig.create (:6342-6355) means validation failures now happen before any resource is created — a real improvement over the previous post-creation resolution.
  • Replacing the bare self._cached_compute_requirements attribute access with getattr(..., None) + a clear ValueError (:6375-6380) fixes what was previously an AttributeError when a reused model reached deploy without build having cached requirements.
  • _resolve_lora_adapter_s3_uri (:3232) consolidates four duplicated branches and adds explicit validation; the suffix per source type (/checkpoints/hf/ vs /model/checkpoints/hf/) matches the prior behavior exactly.

No public API/parameter changes, no new dependencies. I found no correctness bugs. A few points worth confirming, none blocking:

Points to confirm

1. Nova + inference-component deploy: peft_type is now forced to Nonemodel_builder.py:6342-6344

peft_type = self._fetch_peft() if model_package is not None and not is_nova else None

Previously (for a Nova model with a ResourceRequirements inference_config, i.e. the IC path) peft_type was self._fetch_peft() and could be "LORA". Now Nova always takes the single-IC branch. The added comment documents this as intentional, but please confirm no Nova LoRA-via-IC flow relied on the old behavior.

2. _build_single_modelbuilder LoRA build now raises for unsupported source typesmodel_builder.py:3474 (was the if/elif chain at old :3400-3410)

The old chain had no else, so a LoRA build from a source type other than TrainingJob/ModelTrainer/AgentRFTJob/ModelPackage silently left _adapter_s3_uri unset. _resolve_lora_adapter_s3_uri now raises ValueError. This is a strict fail-fast improvement (and better than passing artifact_url=None to the adapter IC later), but it is a behavioral change — flagging in case any existing caller hit the silent path.

3. copy_count now honors inference_config.copy_countmodel_builder.py:6373 and :6469/:6513

The base/single IC previously hard-coded copy_count=1 even when an explicit ResourceRequirements set copies. It now uses inference_config.copy_count (defaults to 1 when unset — verified in resource_requirements.py:56). This is a genuine fix, but a user who set copies>1 and silently got 1 will now get their requested count. Worth a line in the migration/changelog notes.

Minor / non-blocking

4. LoRA adapter resolution is now expressed in two overlapping places_prepare_reused_model_customization_deployment_state (:3282-3283) resolves the adapter when the model was reused, and the inline block at :6383-6394 resolves it again for any supported source type. The result is correct but slightly redundant; the _prepare LoRA branch is effectively subsumed by the inline block. Could be simplified to a single resolution point, but harmless as-is.

Tests

Good coverage: marker lifecycle + reset, exact reuse-hit with no state prep, LoRA endpoint-miss restoring state before writes, non-LoRA compute restoration, explicit-requirements preservation (incl. copy_count==3), and the supported/rejected source matrix for _resolve_lora_adapter_s3_uri including the "reject before any write" assertions. The integration test rewrite to uniquely-named, test-owned resources with _build_single_modelbuilder/_deploy_model_customization sabotage guards is a solid way to prove the exact resources are reused rather than accepting stray ones.

Suggestion: add a unit test for the reused-model + endpoint-miss where compute requirements cannot be resolved case (model_package is None or recipe fetch yields nothing) to lock in the new ValueError at :6376-6380 — currently that failure branch isn't directly exercised.

No security concerns (no credential handling or unvalidated external input introduced). No prompt-injection content in the PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants