fix: bug: ModelBuilder overwrites user-provided HF_MODEL_ID for DJL Serving, preventi (5529) by aviruthen · Pull Request #5734 · aws/sagemaker-python-sdk
🤖 Iteration #1 — Review Comments Addressed
Description
Fix bug where ModelBuilder unconditionally overwrites user-provided HF_MODEL_ID in env_vars, preventing users from loading models from S3 URIs when using DJL Serving and other model servers.
Problem
Multiple _build_for_* methods in model_builder_servers.py used self.env_vars.update({"HF_MODEL_ID": self.model}), which unconditionally overwrites any user-provided value. This prevents users from specifying S3 URIs for model loading via env_vars={"HF_MODEL_ID": "s3://bucket/model/"}, which is a supported configuration per DJL documentation.
Fix
The source file already correctly uses self.env_vars.setdefault("HF_MODEL_ID", self.model) in all six affected methods (_build_for_torchserve, _build_for_tgi, _build_for_djl, _build_for_triton, _build_for_tei, _build_for_transformers). This preserves user-provided values while still setting the default when no value is provided.
Test Changes
Rewrote the test file to follow SDK conventions:
- Migrated from
unittest.TestCaseto plain pytest classes and fixtures - Replaced
self.assertEqualwithassertstatements - Used
@pytest.fixturefor shared mock builder setup - Wrapped long lines to stay within 100-character limit
- Added type annotations to helper functions
- Removed
if __name__ == '__main__': unittest.main()
Testing
All tests verify that:
- User-provided
HF_MODEL_ID(e.g., S3 URI) is preserved during build HF_MODEL_IDdefaults toself.modelwhen not explicitly provided
This is backward compatible — existing behavior where HF_MODEL_ID is not explicitly set continues to work as before.
Comments reviewed: 8
Files modified: sagemaker-serve/src/sagemaker/serve/model_builder_servers.py, sagemaker-serve/tests/unit/test_model_builder_servers_hf_model_id.py
sagemaker-serve/src/sagemaker/serve/model_builder_servers.py: The source file already uses setdefault for HF_MODEL_ID - no changes needed. The file is already correct.sagemaker-serve/tests/unit/test_model_builder_servers_hf_model_id.py: Rewrite test file using pytest conventions: fixtures, assert statements, parametrize, proper line lengths, type annotations, no unittest.TestCase