Skip to content

feat: Enhanced LoRA transforms with flexible path mapping and Pydantic model support - #46

Draft
zhaozuy wants to merge 13 commits into
aws:mainfrom
zhaozuy:better-transforms
Draft

feat: Enhanced LoRA transforms with flexible path mapping and Pydantic model support#46
zhaozuy wants to merge 13 commits into
aws:mainfrom
zhaozuy:better-transforms

Conversation

@zhaozuy

@zhaozuy zhaozuy commented Dec 19, 2025

Copy link
Copy Markdown
Contributor

Summary

Introduces enhanced LoRA adapter transforms that address architectural limitations in the current implementation while maintaining full backward compatibility.

Motivation

The current LoRA transform implementation has several limitations:

  • Missing support for engine request Pydantic model validation
  • Complex inheritance hierarchy that's hard to extend
  • Limited flexibility in request path mapping

Changes Made

🆕 Enhanced Transform Architecture:

  • BaseLoRAApiTransform: Clean ABC-based architecture with separation of concerns
  • LoadLoraApiTransform & UnloadLoraApiTransform: Modular, specialized transform classes
  • lora_factory.py: Enhanced decorator factory with flexible configuration
  • Comprehensive debug logging throughout transformation pipeline

Enhanced Public API:

  • Added new parameters to register_load_adapter_handler() and register_unload_adapter_handler()
  • engine_request_lora_name_path & engine_request_lora_src_path: Flexible path mapping
  • engine_request_model_cls: Pydantic model support for request validation
  • engine_request_defaults: Automatic default value injection

Infrastructure Improvements:

  • Enhanced set_value() utility with better typing and optional depth limits

API Examples

Old API (still supported):

@register_load_adapter_handler(
    request_shape={"adapter_name": "body.name", "adapter_path": "body.src"}
)

New API:

@register_load_adapter_handler(
    engine_request_lora_name_path="body.adapter_name",
    engine_request_lora_src_path="body.source_path",
    engine_request_model_cls=MyEngineRequestModel,
    engine_request_defaults={"config.rank": 4, "config.alpha": 16}
)

Backward Compatibility

  • Existing request_shape API continues to work unchanged
  • Deprecation warnings guide users toward new API
  • Gradual migration path with no breaking changes

Testing

  • Core implementation complete and functional
  • Integration tests are work in progress

Files Changed

  • python/model_hosting_container_standards/sagemaker/lora2/ - New enhanced transforms
  • python/model_hosting_container_standards/sagemaker/__init__.py - Enhanced decorators
  • python/model_hosting_container_standards/common/transforms/utils.py - Utility improvements

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

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.

1 participant