bug: ModelBuilder overwrites user-provided HF_MODEL_ID for DJL Serving, preventing S3 model loading

[michaelraskansky]

Description

When using ModelBuilder with ModelServer.DJL_SERVING, the SDK unconditionally overwrites HF_MODEL_ID with the model parameter value. This prevents users from loading models from S3, even when explicitly setting HF_MODEL_ID in env_vars.

Expected Behavior

Users should be able to specify an S3 URI for model loading via env_vars:

env_vars={"HF_MODEL_ID": "s3://bucket/model/"}

Per DJL documentation, option.model_id accepts both HuggingFace Hub IDs and S3 URIs.

Actual Behavior

The SDK ignores user-provided values and always sets HF_MODEL_ID to the HuggingFace model ID from the model parameter.

Root Cause

In sagemaker-serve/src/sagemaker/serve/model_builder_servers.py, multiple build methods use:

self.env_vars.update({"HF_MODEL_ID": self.model})

This unconditionally overwrites any user-provided HF_MODEL_ID. Affected lines: 139, 215, 323, 429, 535.

Suggested Fix

Use setdefault() to preserve user-provided values:

self.env_vars.setdefault("HF_MODEL_ID", self.model)

Reproduction

from sagemaker.serve.model_builder import ModelBuilder
from sagemaker.serve.utils.types import ModelServer
from sagemaker.serve.builder.schema_builder import SchemaBuilder
from sagemaker.core.image_uris import retrieve as get_image_uri

S3_PATH = 's3://my-bucket/models/Qwen/'

mb = ModelBuilder(
 model='Qwen/Qwen3-VL-4B-Instruct',
 image_uri=get_image_uri(framework='djl-lmi', version='0.30.0', region='us-west-2'),
 model_server=ModelServer.DJL_SERVING,
 schema_builder=SchemaBuilder({'inputs': 'Hello', 'parameters': {}}, [{'generated_text': 'Hi'}]),
 role_arn='arn:aws:iam::123456789012:role/SageMakerRole',
 env_vars={'HF_MODEL_ID': S3_PATH},
 instance_type='ml.g5.2xlarge',
)

print('Before build:', mb.env_vars.get('HF_MODEL_ID')) # s3://my-bucket/models/Qwen/
mb.build()
print('After build:', mb.env_vars.get('HF_MODEL_ID')) # Qwen/Qwen3-VL-4B-Instruct (WRONG)

Use Case

Loading models from S3 is essential for:

• Pre-synced models to avoid runtime download latency
• Models with architectures not supported by TGI (e.g., qwen3_vl) that require vLLM backend

Environment

• SageMaker Python SDK: 3.4.0 (also affects 3.3.x)
• Python: 3.12
• DJL LMI container: 0.30.0

I'm happy to submit a PR with the fix.

[sagarneeldubey]

Filed #5943 as the TGI counterpart to this. Same HF_MODEL_ID / S3-weight-source theme that shows up here (#5529, DJL env-var overwrite) and in #5588 (DJL passthrough injecting a bad ModelDataUrl), but on the TGI path the failure differs: model_path="s3://..." silently mkdirs a junk local s3:/... dir, and s3_model_data_url is consumed as an upload destination — both ignored as weight sources.

Key distinction: unlike DJL/LMI, TGI's HF_MODEL_ID doesn't accept an S3 URI, so the setdefault fix proposed here is necessary but not sufficient for TGI — it also needs the weights mounted at /opt/ml/model via ModelDataSource. Given #5529, #5588, and #5943 are three facets of the same gap (inconsistent S3-as-weight-source handling across model_server branches), might be worth a unified fix rather than three point patches.

[mohamedzeidan2021]

Fixed in #5734, released in v3.11.0. All the DJL build paths now use self.env_vars.setdefault("HF_MODEL_ID", self.model) instead of unconditionally overwriting it, so an S3 URI (or any HF_MODEL_ID) you set in env_vars is preserved through build(). Covered by unit tests.

The TGI-side S3-weight-source handling you flagged (@sagarneeldubey) is tracked separately in #5943 (now fixed) and #5588. Closing this one as resolved for the DJL/LMI case — please upgrade to v3.11.0+.