bug: ModelBuilder TGI path ignores S3 model inputs (model_path and s3_model_data_url), forcing HuggingFace Hub download

[sagarneeldubey]

Description

PySDK Version

• PySDK V2 (2.x)
• PySDK V3 (3.x)

Describe the bug

When using ModelBuilder with ModelServer.TGI, there is no working way to
load model weights from S3 — the deployed TGI container always pulls weights
from huggingface.co at startup. Both parameters that look like they should
enable S3 loading are silently ignored (neither raises at build time):

1. model_path="s3://..." — the TGI build path calls
   _create_dir_structure(self.model_path), which runs
   Path(model_path).mkdir(parents=True, exist_ok=True). On a POSIX filesystem
   this does not fail on an S3 URI; it silently creates a literal local
   directory tree named s3:/<bucket>/<prefix> in the working directory and
   continues. The S3 URI never reaches the container, and HF_MODEL_ID still
   resolves to the HF repo id.
   Source: sagemaker/serve/model_server/tgi/prepare.py::_create_dir_structure

2. s3_model_data_url="s3://..." — on the TGI path this is treated as an
   upload destination prefix for SDK scaffolding, not as a weight source. The
   build method sets self.s3_upload_path = None (comment: "TGI downloads
   models directly from HuggingFace Hub") and overwrites the attribute via
   self.s3_model_data_url, _ = self._prepare_for_mode(). The container ends up
   with HF_MODEL_ID=<hf_repo_id> regardless.
   Source: sagemaker/serve/model_builder_servers.py::_build_for_tgi

Impact: a 7B model cold-starts in ~10–12 min on every scale-out (re-downloading
from HF), which makes scale-to-zero async endpoints impractical and creates a
hard runtime dependency on huggingface.co. It also blocks deploying custom
fine-tuned weights that are not on the HF Hub at all.

This is closely related to #5529 (DJL Serving overwrites a user-provided
HF_MODEL_ID) and #5588 (DJL passthrough injects a bad ModelDataUrl), but
this report is the TGI path and the required fix differs — see Additional
context.

To reproduce

Build-only repro (no GPU, no deploy, ungated model so no token needed).
sagemaker==3.12.0.

from pathlib import Path

from sagemaker.serve import ModelBuilder
from sagemaker.serve.builder.schema_builder import SchemaBuilder
from sagemaker.serve.utils.types import ModelServer

ROLE = "arn:aws:iam::123456789012:role/SageMakerRole"   # any valid role ARN
S3_URI = "s3://my-bucket/models/qwen2.5-7b-instruct/"    # need not exist for the repro
schema = SchemaBuilder(
    sample_input={"inputs": "Hello", "parameters": {"max_new_tokens": 16}},
    sample_output=[{"generated_text": "Hi"}],
)

# --- Case 1: model_path as an S3 URI ---
mb = ModelBuilder(
    model="Qwen/Qwen2.5-7B-Instruct",
    model_path=S3_URI,
    model_server=ModelServer.TGI,
    schema_builder=schema,
    role_arn=ROLE,
    instance_type="ml.g5.2xlarge",
)
mb.build()  # does NOT raise
print("Case 1 — local junk dir created:", Path(S3_URI).exists())   # True (e.g. ./s3:/my-bucket/...)
print("Case 1 — HF_MODEL_ID:", mb.env_vars.get("HF_MODEL_ID"))     # 'Qwen/Qwen2.5-7B-Instruct' (WRONG)

# --- Case 2: s3_model_data_url as an S3 URI ---
mb2 = ModelBuilder(
    model="Qwen/Qwen2.5-7B-Instruct",
    s3_model_data_url=S3_URI,
    model_server=ModelServer.TGI,
    schema_builder=schema,
    role_arn=ROLE,
    instance_type="ml.g5.2xlarge",
)
mb2.build()  # does NOT raise
print("Case 2 — HF_MODEL_ID:", mb2.env_vars.get("HF_MODEL_ID"))    # 'Qwen/Qwen2.5-7B-Instruct' (WRONG)

Expected behavior

With weights pre-staged in S3, a TGI endpoint should load them from S3 (mounted
to /opt/ml/model) instead of downloading from huggingface.co — and at minimum
the SDK should not silently ignore an S3 value passed via model_path /
s3_model_data_url (either honor it or raise a clear error).

Screenshots or logs

Output of the repro above (SDK 3.12.0):

Case 1 — local junk dir created: True
Case 1 — HF_MODEL_ID: Qwen/Qwen2.5-7B-Instruct
Case 2 — HF_MODEL_ID: Qwen/Qwen2.5-7B-Instruct

For comparison, the working lower-level deployment logs this within ~3s of
container start (weights served from S3, no HF download):

INFO text_generation_launcher: Args { model_id: "/opt/ml/model", ... }
INFO download: Files are already present on the host. Skipping download.
INFO download: Successfully downloaded weights for /opt/ml/model

System information

• SageMaker Python SDK version: 3.12.0 (also present in 3.3.x–3.13.x by source inspection)
• Framework name (eg. PyTorch) or algorithm (eg. KMeans): Hugging Face TGI (Text Generation Inference)
• Framework version: TGI DLC huggingface-pytorch-tgi-inference:2.4.0-tgi2.3.1-gpu-py311-cu124-ubuntu22.04
• Python version: 3.12
• CPU or GPU: GPU (ml.g5.2xlarge, ml.g5.12xlarge)
• Custom Docker image (Y/N): N (AWS-provided HF TGI DLC)

Additional context

Working workaround — bypass ModelBuilder and construct the resources
directly, attaching the S3 weights as an uncompressed ModelDataSource and
pointing HF_MODEL_ID at the mount path:

from sagemaker.core.resources import Endpoint, EndpointConfig, Model
from sagemaker.core.shapes.shapes import (
    ContainerDefinition, ModelDataSource, ProductionVariant, S3ModelDataSource,
)
from sagemaker.core.image_retriever.image_retriever import ImageRetriever

image_uri = ImageRetriever.retrieve(
    framework="huggingface-llm", region="us-east-1",
    version="2.3.1", image_scope="inference", instance_type="ml.g5.2xlarge",
)

primary = ContainerDefinition(
    image=image_uri,
    environment={
        "HF_MODEL_ID": "/opt/ml/model",   # mount path, not the HF repo id
        "HF_HUB_OFFLINE": "1",
        "SM_NUM_GPUS": "1",
        "MESSAGES_API_ENABLED": "true",
    },
    model_data_source=ModelDataSource(
        s3_data_source=S3ModelDataSource(
            s3_uri="s3://my-bucket/models/qwen2.5-7b-instruct/",
            s3_data_type="S3Prefix",
            compression_type="None",
        )
    ),
)
Model.create(model_name=..., primary_container=primary, execution_role_arn=ROLE, ...)
EndpointConfig.create(..., production_variants=[ProductionVariant(...)])
Endpoint.create(...)

This drops cold start from ~12 min to ~3–5 min for a 7B model.

Relationship to existing issues

	#5529	#5588	This issue
Model server	DJL Serving (LMI)	DJL Serving (LMI)	TGI
Mode	model + env_vars	passthrough (image_uri + env_vars, no model)	model + model_path/s3_model_data_url
Symptom	user HF_MODEL_ID overwritten	bad ModelDataUrl injected → ValidationException	S3 inputs silently ignored; junk local dir created
Does the engine accept S3 in HF_MODEL_ID?	Yes (LMI option.model_id)	Yes	No — TGI expects HF repo id or local path
Sufficient fix	setdefault on HF_MODEL_ID	clear s3_model_data_url in passthrough	setdefault is necessary but not sufficient; TGI also needs weights mounted at /opt/ml/model via ModelDataSource

Suggested fix (TGI path) — when an S3 URI is supplied:

1. Skip _create_dir_structure local-dir creation for S3 inputs.
2. Attach the S3 prefix as an uncompressed ModelDataSource on the container —
   the early branch in
   sagemaker/serve/model_server/tgi/server.py::SageMakerTgiServing._upload_tgi_artifacts
   already emits the correct shape when model_path is an S3 URI, but it's
   currently unreachable because _create_dir_structure runs first.
3. Set HF_MODEL_ID=/opt/ml/model instead of the HF repo id.

More broadly, consider unifying S3-as-weight-source handling across all
model_server branches so behavior is consistent (ties together #5529, #5588,
and this issue).