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.TestCase to plain pytest classes and fixtures
  • Replaced self.assertEqual with assert statements
  • Used @pytest.fixture for 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:

  1. User-provided HF_MODEL_ID (e.g., S3 URI) is preserved during build
  2. HF_MODEL_ID defaults to self.model when 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