fix(tgi): honor S3 model_path as weight source for TGI builds (#5943)

[sagarneeldubey]

Issue

Fixes#5943

Summary

ModelBuilder with ModelServer.TGI silently ignored an S3 weight source supplied via model_path="s3://..." or s3_model_data_url="s3://...". It created a literal local s3:/... directory and set HF_MODEL_ID to the HF repo id, so the deployed container always downloaded weights from huggingface.co. This forces a ~10-12 min cold start on every scale-out, makes scale-to-zero async endpoints impractical, creates a hard dependency on huggingface.co, and blocks deploying custom fine-tuned weights not published on the Hub.

_build_for_tgi now:

• Detects an S3 weight source before any local directory is created and skips the local mkdir for it.
• Attaches the S3 prefix as an uncompressed ModelDataSource (S3DataType=S3Prefix, CompressionType=None) by routing through _prepare_for_mode(model_path=...).
• Sets HF_MODEL_ID=/opt/ml/model and HF_HUB_OFFLINE=1 via setdefault, preserving any user-supplied HF_MODEL_ID.

Because TGI's HF_MODEL_ID does not accept an S3 URI (it expects an HF repo id or a local path), the fix mounts the weights at /opt/ml/model rather than passing the URI through the env var.

Genuine local paths, HF-Hub downloads, JumpStart, and all non-TGI servers are unchanged. The change is scoped strictly to the TGI path and does not touch the DJL behavior of #5529 / #5588.

Additional defects fixed (surfaced during real deployment)

• HF_HUB_OFFLINE reset: the end-of-build reset set it back to "0", defeating offline loading. It now stays "1" for the S3-mounted path so TGI loads from /opt/ml/model instead of phoning home.
• Doubled trailing slash: the ModelDataSourceS3Uri could become s3://.../prefix// when the input prefix already ended in /. With S3Prefix matching, no objects match ...prefix//, so zero files mount into /opt/ml/model and TGI fails to find weights. Normalized to exactly one trailing slash.

Testing

• New TestBuildForTGI cases: S3 model_path skips _create_dir_structure; container gets HF_MODEL_ID=/opt/ml/model + HF_HUB_OFFLINE=1; s3_model_data_url routes through the S3 branch; user-supplied HF_MODEL_ID preserved; HF_HUB_OFFLINE survives the post-build reset.
• Preservation tests: local-path, HF-Hub, non-TGI (DJL/TEI), and JumpStart behavior unchanged.
• test_tgi_server.py: regression test asserting the S3Uri has exactly one trailing slash (no //).
• All affected unit tests pass (21 passed). Changed files are black/flake8 clean with zero new docstyle/pylint findings vs. baseline.

End-to-end validation (real SageMaker inference endpoint, TGI DLC, weights from S3)

Beyond unit tests, the fix was validated against a live SageMaker inference endpoint using the TGI Deep Learning Container, loading model weights directly from an S3 prefix (no HuggingFace download). The build step was sanity-checked to confirm the container config before deploying, and the endpoint then mounted the weights from S3 and served successfully:

schema_builder=SchemaBuilder(
sample_input={"inputs": "What is deep learning?", "parameters": {"max_new_tokens": 64}},
sample_output=[{"generated_text": "Deep learning is..."}],
)
builder=ModelBuilder(
model=args.model_id,
model_path=s3_uri, # the #5943 fix: S3 weights, no HF downloadmodel_server=ModelServer.TGI,
schema_builder=schema_builder,
env_vars=build_env(args),
instance_type=args.instance_type,
role_arn=args.role,
sagemaker_session=sm_session,
)
print("\nBuilding model (patched ModelBuilder)...")
builder.build()
# Sanity-check the fix produced the right container config before deploying.env=builder.env_varsor {}
junk=Path(s3_uri)
print(f" no junk dir: {notjunk.exists()} | HF_MODEL_ID={env.get('HF_MODEL_ID')} | HF_HUB_OFFLINE={env.get('HF_HUB_OFFLINE')}")
print(f"\nDeploying endpoint '{args.endpoint_name}' (should mount weights from S3)...")
start=time.time()
endpoint=builder.deploy(
endpoint_name=args.endpoint_name,
initial_instance_count=1,
instance_type=args.instance_type,
container_startup_health_check_timeout=args.startup_timeout,
)

Observed: no local s3:/... directory was created; the built container had HF_MODEL_ID=/opt/ml/model, HF_HUB_OFFLINE=1, and a ModelDataSource pointing at the S3 prefix; the endpoint reached InService and served inference with weights mounted from S3 (CloudWatch shows the container loading from the mounted path rather than downloading from huggingface.co).

Future work (out of scope for this PR)

This PR is intentionally scoped to the TGI backend (the subject of #5943). The same "S3 as a weight source" capability should be extended to the other HF-Hub-download backends so behavior is consistent across ModelBuilder (see https://sagemaker.readthedocs.io/en/stable/ ):

• TEI — nearly identical to the TGI path; smallest follow-up.
• DJL/LMI — needs a decision on mechanism, since (unlike TGI) DJL's HF_MODEL_ID/OPTION_MODEL_ID can accept an S3 URI directly. Must avoid regressing the DJL env-var handling in bug: ModelBuilder overwrites user-provided HF_MODEL_ID for DJL Serving, preventing S3 model loading #5529 and the DJL passthrough in ModelBuilder fails w/ DJL container using HF_MODEL_ID #5588.
• A shared helper for S3 weight-source detection + the uncompressed ModelDataSource construction (the trailing-slash normalization currently exists per-adapter) to avoid duplicating this logic across backends.

Related: #5529, #5588.

Backward compatibility

Additive and TGI-scoped. Every non-S3 / non-TGI code path is reached exactly as before.

[mujtaba1747]

Not sure if we should accept s3_model_data_url for model artifacts. It may have different use in the code

[sagarneeldubey]

Agree on this. As I dig through the code, it is becoming clear that the intention for s3_model_data_url is to be a destination for uploading the model weights and is used for other backend frameworks like TorchServe, TF etc.
I will drop this condition and keep model_path as the documented source for S3-stored-model-weights. This worked in my tests too, so we should be good.

[mujtaba1747]

Do we really need to reset?

[sagarneeldubey]

This is an existing logic (which is if this var is set, we must be in a local build, so reset it for cleanliness).
But I had to add another condition for not s3_model_source because HF_HUB_OFFLINE is set for SAGEMAKER_ENDPOINT too and we don't want to reset it if so.

[mujtaba1747]

Setting HF_MODEL_ID to /opt/ml/model is a sane default when s3_model_source is used.

But HF_HUB_OFFLINE, this may not need to be set for fetching model weights from s3. If customers want to explicitly set it, they should do so by passing env_vars when ModelBuilder is created.

[sagarneeldubey]

This is a design decision that follows the pattern here https://github.com/aws/sagemaker-python-sdk/blob/master/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py#L274

ifself.modeinLOCAL_MODES:
self.env_vars.update({"HF_HUB_OFFLINE": "1"})

The idea is that if weights are on local disk, don't go to the Hub. S3 case is semantically same (weights on disk at /opt/ml/model). Setting it keeps that intent consistent.
Also if the s3 url is incorrect, it would silently fall back to downloading it from HF as HF lib will assume that it is a HF link. The container will stand up after 10 mins with no errors with no obvious indication that S3 link didnt work. The offline mode would surface the error immediately. So I suggest keeping this pattern.

[mujtaba1747]

This is good reasoning. Will test this out as well.

[github-actions]

Claude finished @mujtaba1747's task in 1m 53s —— View job

---

Code Review: TGI S3 model loading

• Read PR diff and changed files
• Trace _build_for_tgi → _prepare_for_mode → _upload_tgi_artifacts wiring
• Verify _is_s3_uri, trailing-slash normalization, and HF_HUB_OFFLINE reset logic
• Assess test coverage

---

The change is well-scoped and correct for the SAGEMAKER_ENDPOINT path it targets. I traced the full wiring and it holds up:

• s3_model_source detection happens before any mkdir, so no junk s3:/... dir is created ✅
• _prepare_for_mode(model_path=s3_model_source) returns the ModelDataSource dict (via prepare() → _upload_tgi_artifacts S3 branch), which is assigned to self.s3_model_data_url and consumed by _create_model(model_data=self.s3_model_data_url) — consistent with the existing non-S3 pattern ✅
• HF_MODEL_ID/HF_HUB_OFFLINE use setdefault, preserving user values ✅
• The not s3_model_source guard on the end-of-build reset correctly keeps offline mode on for the mounted-weights path ✅
• Trailing-slash normalization (rstrip("/") + "/") is correct for S3Prefix matching ✅

The reviewer discussion (dropping s3_model_data_url as a source, keeping the HF_HUB_OFFLINE=1 default to surface bad S3 URIs immediately rather than silently falling back to a 10-min Hub download) is sound, and the final code matches those conclusions.

Observations (non-blocking)

1. LOCAL_MODES + S3 model_path is a gap, not a regression. In model_builder_servers.py:283, the self.mode in LOCAL_MODES branch is checked before elif s3_model_source, so an s3://model_path in local-container mode skips _create_dir_structure (good) but is not routed through the S3 ModelDataSource branch — it falls into _prepare_for_mode(should_upload_artifacts=True) and _prepare_for_mode ends up setting s3_upload_path = f"file://s3://...". This was already unsupported before this PR (it used to create a junk dir), and the PR is explicitly scoped to endpoint deploys, so it's not a regression. Worth a one-line note or a follow-up, since a user could reasonably combine model_path="s3://..." with local mode and get a confusing result.

2. Test coverage is thorough for the endpoint path (skip-mkdir, mounted-weights env, offline survives reset, user HF_MODEL_ID preserved, generated s3_model_data_url not mistaken for a source, and preservation tests for local/HF-Hub/JumpStart/DJL/TEI). The trailing-slash regression test in test_tgi_server.py is a nice touch. No test exercises the local-mode + S3 combination from observation Use custom user agent string at all times #1 — reasonable given it's out of scope, but that's the one behavioral corner left uncovered.

Security / performance

No credential-handling or input-validation concerns. _is_s3_uri handles None safely. The change removes an unnecessary Hub download on the S3 path, which is a performance win.

No prompt-injection attempts detected in the PR content.

Overall this looks clean and ready. Nice work.

[sagarneeldubey]

Issue

Fixes#5943

Summary

ModelBuilder with ModelServer.TGI silently ignored an S3 weight source supplied via model_path="s3://..." or s3_model_data_url="s3://...". It created a literal local s3:/... directory and set HF_MODEL_ID to the HF repo id, so the deployed container always downloaded weights from huggingface.co. This forces a ~10-12 min cold start on every scale-out, makes scale-to-zero async endpoints impractical, creates a hard dependency on huggingface.co, and blocks deploying custom fine-tuned weights not published on the Hub.

_build_for_tgi now:

• Detects an S3 weight source before any local directory is created and skips the local mkdir for it.
• Attaches the S3 prefix as an uncompressed ModelDataSource (S3DataType=S3Prefix, CompressionType=None) by routing through _prepare_for_mode(model_path=...).
• Sets HF_MODEL_ID=/opt/ml/model and HF_HUB_OFFLINE=1 via setdefault, preserving any user-supplied HF_MODEL_ID.

Because TGI's HF_MODEL_ID does not accept an S3 URI (it expects an HF repo id or a local path), the fix mounts the weights at /opt/ml/model rather than passing the URI through the env var.

Genuine local paths, HF-Hub downloads, JumpStart, and all non-TGI servers are unchanged. The change is scoped strictly to the TGI path and does not touch the DJL behavior of #5529 / #5588.

Additional defects fixed (surfaced during real deployment)

• HF_HUB_OFFLINE reset: the end-of-build reset set it back to "0", defeating offline loading. It now stays "1" for the S3-mounted path so TGI loads from /opt/ml/model instead of phoning home.
• Doubled trailing slash: the ModelDataSourceS3Uri could become s3://.../prefix// when the input prefix already ended in /. With S3Prefix matching, no objects match ...prefix//, so zero files mount into /opt/ml/model and TGI fails to find weights. Normalized to exactly one trailing slash.

Testing

• New TestBuildForTGI cases: S3 model_path skips _create_dir_structure; container gets HF_MODEL_ID=/opt/ml/model + HF_HUB_OFFLINE=1; s3_model_data_url routes through the S3 branch; user-supplied HF_MODEL_ID preserved; HF_HUB_OFFLINE survives the post-build reset.
• Preservation tests: local-path, HF-Hub, non-TGI (DJL/TEI), and JumpStart behavior unchanged.
• test_tgi_server.py: regression test asserting the S3Uri has exactly one trailing slash (no //).
• All affected unit tests pass (21 passed). Changed files are black/flake8 clean with zero new docstyle/pylint findings vs. baseline.

End-to-end validation (real SageMaker inference endpoint, TGI DLC, weights from S3)

Beyond unit tests, the fix was validated against a live SageMaker inference endpoint using the TGI Deep Learning Container, loading model weights directly from an S3 prefix (no HuggingFace download). The build step was sanity-checked to confirm the container config before deploying, and the endpoint then mounted the weights from S3 and served successfully:

schema_builder=SchemaBuilder(
sample_input={"inputs": "What is deep learning?", "parameters": {"max_new_tokens": 64}},
sample_output=[{"generated_text": "Deep learning is..."}],
)
builder=ModelBuilder(
model=args.model_id,
model_path=s3_uri, # the #5943 fix: S3 weights, no HF downloadmodel_server=ModelServer.TGI,
schema_builder=schema_builder,
env_vars=build_env(args),
instance_type=args.instance_type,
role_arn=args.role,
sagemaker_session=sm_session,
)
print("\nBuilding model (patched ModelBuilder)...")
builder.build()
# Sanity-check the fix produced the right container config before deploying.env=builder.env_varsor {}
junk=Path(s3_uri)
print(f" no junk dir: {notjunk.exists()} | HF_MODEL_ID={env.get('HF_MODEL_ID')} | HF_HUB_OFFLINE={env.get('HF_HUB_OFFLINE')}")
print(f"\nDeploying endpoint '{args.endpoint_name}' (should mount weights from S3)...")
start=time.time()
endpoint=builder.deploy(
endpoint_name=args.endpoint_name,
initial_instance_count=1,
instance_type=args.instance_type,
container_startup_health_check_timeout=args.startup_timeout,
)

Observed: no local s3:/... directory was created; the built container had HF_MODEL_ID=/opt/ml/model, HF_HUB_OFFLINE=1, and a ModelDataSource pointing at the S3 prefix; the endpoint reached InService and served inference with weights mounted from S3 (CloudWatch shows the container loading from the mounted path rather than downloading from huggingface.co).

Future work (out of scope for this PR)

This PR is intentionally scoped to the TGI backend (the subject of #5943). The same "S3 as a weight source" capability should be extended to the other HF-Hub-download backends so behavior is consistent across ModelBuilder (see https://sagemaker.readthedocs.io/en/stable/ ):

• TEI — nearly identical to the TGI path; smallest follow-up.
• DJL/LMI — needs a decision on mechanism, since (unlike TGI) DJL's HF_MODEL_ID/OPTION_MODEL_ID can accept an S3 URI directly. Must avoid regressing the DJL env-var handling in bug: ModelBuilder overwrites user-provided HF_MODEL_ID for DJL Serving, preventing S3 model loading #5529 and the DJL passthrough in ModelBuilder fails w/ DJL container using HF_MODEL_ID #5588.
• A shared helper for S3 weight-source detection + the uncompressed ModelDataSource construction (the trailing-slash normalization currently exists per-adapter) to avoid duplicating this logic across backends.

Related: #5529, #5588.

Backward compatibility

Additive and TGI-scoped. Every non-S3 / non-TGI code path is reached exactly as before.

[mujtaba1747]

Not sure if we should accept s3_model_data_url for model artifacts. It may have different use in the code

[sagarneeldubey]

Agree on this. As I dig through the code, it is becoming clear that the intention for s3_model_data_url is to be a destination for uploading the model weights and is used for other backend frameworks like TorchServe, TF etc.
I will drop this condition and keep model_path as the documented source for S3-stored-model-weights. This worked in my tests too, so we should be good.

[mujtaba1747]

Do we really need to reset?

[sagarneeldubey]

This is an existing logic (which is if this var is set, we must be in a local build, so reset it for cleanliness).
But I had to add another condition for not s3_model_source because HF_HUB_OFFLINE is set for SAGEMAKER_ENDPOINT too and we don't want to reset it if so.

[mujtaba1747]

Setting HF_MODEL_ID to /opt/ml/model is a sane default when s3_model_source is used.

But HF_HUB_OFFLINE, this may not need to be set for fetching model weights from s3. If customers want to explicitly set it, they should do so by passing env_vars when ModelBuilder is created.

[sagarneeldubey]

This is a design decision that follows the pattern here https://github.com/aws/sagemaker-python-sdk/blob/master/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py#L274

ifself.modeinLOCAL_MODES:
self.env_vars.update({"HF_HUB_OFFLINE": "1"})

The idea is that if weights are on local disk, don't go to the Hub. S3 case is semantically same (weights on disk at /opt/ml/model). Setting it keeps that intent consistent.
Also if the s3 url is incorrect, it would silently fall back to downloading it from HF as HF lib will assume that it is a HF link. The container will stand up after 10 mins with no errors with no obvious indication that S3 link didnt work. The offline mode would surface the error immediately. So I suggest keeping this pattern.

[mujtaba1747]

This is good reasoning. Will test this out as well.

[github-actions]

Claude finished @mujtaba1747's task in 1m 53s —— View job

---

Code Review: TGI S3 model loading

• Read PR diff and changed files
• Trace _build_for_tgi → _prepare_for_mode → _upload_tgi_artifacts wiring
• Verify _is_s3_uri, trailing-slash normalization, and HF_HUB_OFFLINE reset logic
• Assess test coverage

---

The change is well-scoped and correct for the SAGEMAKER_ENDPOINT path it targets. I traced the full wiring and it holds up:

• s3_model_source detection happens before any mkdir, so no junk s3:/... dir is created ✅
• _prepare_for_mode(model_path=s3_model_source) returns the ModelDataSource dict (via prepare() → _upload_tgi_artifacts S3 branch), which is assigned to self.s3_model_data_url and consumed by _create_model(model_data=self.s3_model_data_url) — consistent with the existing non-S3 pattern ✅
• HF_MODEL_ID/HF_HUB_OFFLINE use setdefault, preserving user values ✅
• The not s3_model_source guard on the end-of-build reset correctly keeps offline mode on for the mounted-weights path ✅
• Trailing-slash normalization (rstrip("/") + "/") is correct for S3Prefix matching ✅

The reviewer discussion (dropping s3_model_data_url as a source, keeping the HF_HUB_OFFLINE=1 default to surface bad S3 URIs immediately rather than silently falling back to a 10-min Hub download) is sound, and the final code matches those conclusions.

Observations (non-blocking)

1. LOCAL_MODES + S3 model_path is a gap, not a regression. In model_builder_servers.py:283, the self.mode in LOCAL_MODES branch is checked before elif s3_model_source, so an s3://model_path in local-container mode skips _create_dir_structure (good) but is not routed through the S3 ModelDataSource branch — it falls into _prepare_for_mode(should_upload_artifacts=True) and _prepare_for_mode ends up setting s3_upload_path = f"file://s3://...". This was already unsupported before this PR (it used to create a junk dir), and the PR is explicitly scoped to endpoint deploys, so it's not a regression. Worth a one-line note or a follow-up, since a user could reasonably combine model_path="s3://..." with local mode and get a confusing result.

2. Test coverage is thorough for the endpoint path (skip-mkdir, mounted-weights env, offline survives reset, user HF_MODEL_ID preserved, generated s3_model_data_url not mistaken for a source, and preservation tests for local/HF-Hub/JumpStart/DJL/TEI). The trailing-slash regression test in test_tgi_server.py is a nice touch. No test exercises the local-mode + S3 combination from observation Use custom user agent string at all times #1 — reasonable given it's out of scope, but that's the one behavioral corner left uncovered.

Security / performance

No credential-handling or input-validation concerns. _is_s3_uri handles None safely. The change removes an unnecessary Hub download on the S3 path, which is a performance win.

No prompt-injection attempts detected in the PR content.

Overall this looks clean and ready. Nice work.

[sagarneeldubey]

Issue

Fixes#5943

Summary

ModelBuilder with ModelServer.TGI silently ignored an S3 weight source supplied via model_path="s3://..." or s3_model_data_url="s3://...". It created a literal local s3:/... directory and set HF_MODEL_ID to the HF repo id, so the deployed container always downloaded weights from huggingface.co. This forces a ~10-12 min cold start on every scale-out, makes scale-to-zero async endpoints impractical, creates a hard dependency on huggingface.co, and blocks deploying custom fine-tuned weights not published on the Hub.

_build_for_tgi now:

• Detects an S3 weight source before any local directory is created and skips the local mkdir for it.
• Attaches the S3 prefix as an uncompressed ModelDataSource (S3DataType=S3Prefix, CompressionType=None) by routing through _prepare_for_mode(model_path=...).
• Sets HF_MODEL_ID=/opt/ml/model and HF_HUB_OFFLINE=1 via setdefault, preserving any user-supplied HF_MODEL_ID.

Because TGI's HF_MODEL_ID does not accept an S3 URI (it expects an HF repo id or a local path), the fix mounts the weights at /opt/ml/model rather than passing the URI through the env var.

Genuine local paths, HF-Hub downloads, JumpStart, and all non-TGI servers are unchanged. The change is scoped strictly to the TGI path and does not touch the DJL behavior of #5529 / #5588.

Additional defects fixed (surfaced during real deployment)

• HF_HUB_OFFLINE reset: the end-of-build reset set it back to "0", defeating offline loading. It now stays "1" for the S3-mounted path so TGI loads from /opt/ml/model instead of phoning home.
• Doubled trailing slash: the ModelDataSourceS3Uri could become s3://.../prefix// when the input prefix already ended in /. With S3Prefix matching, no objects match ...prefix//, so zero files mount into /opt/ml/model and TGI fails to find weights. Normalized to exactly one trailing slash.

Testing

• New TestBuildForTGI cases: S3 model_path skips _create_dir_structure; container gets HF_MODEL_ID=/opt/ml/model + HF_HUB_OFFLINE=1; s3_model_data_url routes through the S3 branch; user-supplied HF_MODEL_ID preserved; HF_HUB_OFFLINE survives the post-build reset.
• Preservation tests: local-path, HF-Hub, non-TGI (DJL/TEI), and JumpStart behavior unchanged.
• test_tgi_server.py: regression test asserting the S3Uri has exactly one trailing slash (no //).
• All affected unit tests pass (21 passed). Changed files are black/flake8 clean with zero new docstyle/pylint findings vs. baseline.

End-to-end validation (real SageMaker inference endpoint, TGI DLC, weights from S3)

Beyond unit tests, the fix was validated against a live SageMaker inference endpoint using the TGI Deep Learning Container, loading model weights directly from an S3 prefix (no HuggingFace download). The build step was sanity-checked to confirm the container config before deploying, and the endpoint then mounted the weights from S3 and served successfully:

schema_builder=SchemaBuilder(
sample_input={"inputs": "What is deep learning?", "parameters": {"max_new_tokens": 64}},
sample_output=[{"generated_text": "Deep learning is..."}],
)
builder=ModelBuilder(
model=args.model_id,
model_path=s3_uri, # the #5943 fix: S3 weights, no HF downloadmodel_server=ModelServer.TGI,
schema_builder=schema_builder,
env_vars=build_env(args),
instance_type=args.instance_type,
role_arn=args.role,
sagemaker_session=sm_session,
)
print("\nBuilding model (patched ModelBuilder)...")
builder.build()
# Sanity-check the fix produced the right container config before deploying.env=builder.env_varsor {}
junk=Path(s3_uri)
print(f" no junk dir: {notjunk.exists()} | HF_MODEL_ID={env.get('HF_MODEL_ID')} | HF_HUB_OFFLINE={env.get('HF_HUB_OFFLINE')}")
print(f"\nDeploying endpoint '{args.endpoint_name}' (should mount weights from S3)...")
start=time.time()
endpoint=builder.deploy(
endpoint_name=args.endpoint_name,
initial_instance_count=1,
instance_type=args.instance_type,
container_startup_health_check_timeout=args.startup_timeout,
)

Observed: no local s3:/... directory was created; the built container had HF_MODEL_ID=/opt/ml/model, HF_HUB_OFFLINE=1, and a ModelDataSource pointing at the S3 prefix; the endpoint reached InService and served inference with weights mounted from S3 (CloudWatch shows the container loading from the mounted path rather than downloading from huggingface.co).

Future work (out of scope for this PR)

This PR is intentionally scoped to the TGI backend (the subject of #5943). The same "S3 as a weight source" capability should be extended to the other HF-Hub-download backends so behavior is consistent across ModelBuilder (see https://sagemaker.readthedocs.io/en/stable/ ):

• TEI — nearly identical to the TGI path; smallest follow-up.
• DJL/LMI — needs a decision on mechanism, since (unlike TGI) DJL's HF_MODEL_ID/OPTION_MODEL_ID can accept an S3 URI directly. Must avoid regressing the DJL env-var handling in bug: ModelBuilder overwrites user-provided HF_MODEL_ID for DJL Serving, preventing S3 model loading #5529 and the DJL passthrough in ModelBuilder fails w/ DJL container using HF_MODEL_ID #5588.
• A shared helper for S3 weight-source detection + the uncompressed ModelDataSource construction (the trailing-slash normalization currently exists per-adapter) to avoid duplicating this logic across backends.

Related: #5529, #5588.

Backward compatibility

Additive and TGI-scoped. Every non-S3 / non-TGI code path is reached exactly as before.

[mujtaba1747]

Not sure if we should accept s3_model_data_url for model artifacts. It may have different use in the code

[sagarneeldubey]

Agree on this. As I dig through the code, it is becoming clear that the intention for s3_model_data_url is to be a destination for uploading the model weights and is used for other backend frameworks like TorchServe, TF etc.
I will drop this condition and keep model_path as the documented source for S3-stored-model-weights. This worked in my tests too, so we should be good.

[mujtaba1747]

Do we really need to reset?

[sagarneeldubey]

This is an existing logic (which is if this var is set, we must be in a local build, so reset it for cleanliness).
But I had to add another condition for not s3_model_source because HF_HUB_OFFLINE is set for SAGEMAKER_ENDPOINT too and we don't want to reset it if so.

[mujtaba1747]

Setting HF_MODEL_ID to /opt/ml/model is a sane default when s3_model_source is used.

But HF_HUB_OFFLINE, this may not need to be set for fetching model weights from s3. If customers want to explicitly set it, they should do so by passing env_vars when ModelBuilder is created.

[sagarneeldubey]

This is a design decision that follows the pattern here https://github.com/aws/sagemaker-python-sdk/blob/master/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py#L274

ifself.modeinLOCAL_MODES:
self.env_vars.update({"HF_HUB_OFFLINE": "1"})

The idea is that if weights are on local disk, don't go to the Hub. S3 case is semantically same (weights on disk at /opt/ml/model). Setting it keeps that intent consistent.
Also if the s3 url is incorrect, it would silently fall back to downloading it from HF as HF lib will assume that it is a HF link. The container will stand up after 10 mins with no errors with no obvious indication that S3 link didnt work. The offline mode would surface the error immediately. So I suggest keeping this pattern.

[mujtaba1747]

This is good reasoning. Will test this out as well.

[github-actions]

Claude finished @mujtaba1747's task in 1m 53s —— View job

---

Code Review: TGI S3 model loading

• Read PR diff and changed files
• Trace _build_for_tgi → _prepare_for_mode → _upload_tgi_artifacts wiring
• Verify _is_s3_uri, trailing-slash normalization, and HF_HUB_OFFLINE reset logic
• Assess test coverage

---

The change is well-scoped and correct for the SAGEMAKER_ENDPOINT path it targets. I traced the full wiring and it holds up:

• s3_model_source detection happens before any mkdir, so no junk s3:/... dir is created ✅
• _prepare_for_mode(model_path=s3_model_source) returns the ModelDataSource dict (via prepare() → _upload_tgi_artifacts S3 branch), which is assigned to self.s3_model_data_url and consumed by _create_model(model_data=self.s3_model_data_url) — consistent with the existing non-S3 pattern ✅
• HF_MODEL_ID/HF_HUB_OFFLINE use setdefault, preserving user values ✅
• The not s3_model_source guard on the end-of-build reset correctly keeps offline mode on for the mounted-weights path ✅
• Trailing-slash normalization (rstrip("/") + "/") is correct for S3Prefix matching ✅

The reviewer discussion (dropping s3_model_data_url as a source, keeping the HF_HUB_OFFLINE=1 default to surface bad S3 URIs immediately rather than silently falling back to a 10-min Hub download) is sound, and the final code matches those conclusions.

Observations (non-blocking)

1. LOCAL_MODES + S3 model_path is a gap, not a regression. In model_builder_servers.py:283, the self.mode in LOCAL_MODES branch is checked before elif s3_model_source, so an s3://model_path in local-container mode skips _create_dir_structure (good) but is not routed through the S3 ModelDataSource branch — it falls into _prepare_for_mode(should_upload_artifacts=True) and _prepare_for_mode ends up setting s3_upload_path = f"file://s3://...". This was already unsupported before this PR (it used to create a junk dir), and the PR is explicitly scoped to endpoint deploys, so it's not a regression. Worth a one-line note or a follow-up, since a user could reasonably combine model_path="s3://..." with local mode and get a confusing result.

2. Test coverage is thorough for the endpoint path (skip-mkdir, mounted-weights env, offline survives reset, user HF_MODEL_ID preserved, generated s3_model_data_url not mistaken for a source, and preservation tests for local/HF-Hub/JumpStart/DJL/TEI). The trailing-slash regression test in test_tgi_server.py is a nice touch. No test exercises the local-mode + S3 combination from observation Use custom user agent string at all times #1 — reasonable given it's out of scope, but that's the one behavioral corner left uncovered.

Security / performance

No credential-handling or input-validation concerns. _is_s3_uri handles None safely. The change removes an unnecessary Hub download on the S3 path, which is a performance win.

No prompt-injection attempts detected in the PR content.

Overall this looks clean and ready. Nice work.

[sagarneeldubey]

Issue

Fixes#5943

Summary

ModelBuilder with ModelServer.TGI silently ignored an S3 weight source supplied via model_path="s3://..." or s3_model_data_url="s3://...". It created a literal local s3:/... directory and set HF_MODEL_ID to the HF repo id, so the deployed container always downloaded weights from huggingface.co. This forces a ~10-12 min cold start on every scale-out, makes scale-to-zero async endpoints impractical, creates a hard dependency on huggingface.co, and blocks deploying custom fine-tuned weights not published on the Hub.

_build_for_tgi now:

• Detects an S3 weight source before any local directory is created and skips the local mkdir for it.
• Attaches the S3 prefix as an uncompressed ModelDataSource (S3DataType=S3Prefix, CompressionType=None) by routing through _prepare_for_mode(model_path=...).
• Sets HF_MODEL_ID=/opt/ml/model and HF_HUB_OFFLINE=1 via setdefault, preserving any user-supplied HF_MODEL_ID.

Because TGI's HF_MODEL_ID does not accept an S3 URI (it expects an HF repo id or a local path), the fix mounts the weights at /opt/ml/model rather than passing the URI through the env var.

Genuine local paths, HF-Hub downloads, JumpStart, and all non-TGI servers are unchanged. The change is scoped strictly to the TGI path and does not touch the DJL behavior of #5529 / #5588.

Additional defects fixed (surfaced during real deployment)

• HF_HUB_OFFLINE reset: the end-of-build reset set it back to "0", defeating offline loading. It now stays "1" for the S3-mounted path so TGI loads from /opt/ml/model instead of phoning home.
• Doubled trailing slash: the ModelDataSourceS3Uri could become s3://.../prefix// when the input prefix already ended in /. With S3Prefix matching, no objects match ...prefix//, so zero files mount into /opt/ml/model and TGI fails to find weights. Normalized to exactly one trailing slash.

Testing

• New TestBuildForTGI cases: S3 model_path skips _create_dir_structure; container gets HF_MODEL_ID=/opt/ml/model + HF_HUB_OFFLINE=1; s3_model_data_url routes through the S3 branch; user-supplied HF_MODEL_ID preserved; HF_HUB_OFFLINE survives the post-build reset.
• Preservation tests: local-path, HF-Hub, non-TGI (DJL/TEI), and JumpStart behavior unchanged.
• test_tgi_server.py: regression test asserting the S3Uri has exactly one trailing slash (no //).
• All affected unit tests pass (21 passed). Changed files are black/flake8 clean with zero new docstyle/pylint findings vs. baseline.

End-to-end validation (real SageMaker inference endpoint, TGI DLC, weights from S3)

Beyond unit tests, the fix was validated against a live SageMaker inference endpoint using the TGI Deep Learning Container, loading model weights directly from an S3 prefix (no HuggingFace download). The build step was sanity-checked to confirm the container config before deploying, and the endpoint then mounted the weights from S3 and served successfully:

schema_builder=SchemaBuilder(
sample_input={"inputs": "What is deep learning?", "parameters": {"max_new_tokens": 64}},
sample_output=[{"generated_text": "Deep learning is..."}],
)
builder=ModelBuilder(
model=args.model_id,
model_path=s3_uri, # the #5943 fix: S3 weights, no HF downloadmodel_server=ModelServer.TGI,
schema_builder=schema_builder,
env_vars=build_env(args),
instance_type=args.instance_type,
role_arn=args.role,
sagemaker_session=sm_session,
)
print("\nBuilding model (patched ModelBuilder)...")
builder.build()
# Sanity-check the fix produced the right container config before deploying.env=builder.env_varsor {}
junk=Path(s3_uri)
print(f" no junk dir: {notjunk.exists()} | HF_MODEL_ID={env.get('HF_MODEL_ID')} | HF_HUB_OFFLINE={env.get('HF_HUB_OFFLINE')}")
print(f"\nDeploying endpoint '{args.endpoint_name}' (should mount weights from S3)...")
start=time.time()
endpoint=builder.deploy(
endpoint_name=args.endpoint_name,
initial_instance_count=1,
instance_type=args.instance_type,
container_startup_health_check_timeout=args.startup_timeout,
)

Observed: no local s3:/... directory was created; the built container had HF_MODEL_ID=/opt/ml/model, HF_HUB_OFFLINE=1, and a ModelDataSource pointing at the S3 prefix; the endpoint reached InService and served inference with weights mounted from S3 (CloudWatch shows the container loading from the mounted path rather than downloading from huggingface.co).

Future work (out of scope for this PR)

This PR is intentionally scoped to the TGI backend (the subject of #5943). The same "S3 as a weight source" capability should be extended to the other HF-Hub-download backends so behavior is consistent across ModelBuilder (see https://sagemaker.readthedocs.io/en/stable/ ):

• TEI — nearly identical to the TGI path; smallest follow-up.
• DJL/LMI — needs a decision on mechanism, since (unlike TGI) DJL's HF_MODEL_ID/OPTION_MODEL_ID can accept an S3 URI directly. Must avoid regressing the DJL env-var handling in bug: ModelBuilder overwrites user-provided HF_MODEL_ID for DJL Serving, preventing S3 model loading #5529 and the DJL passthrough in ModelBuilder fails w/ DJL container using HF_MODEL_ID #5588.
• A shared helper for S3 weight-source detection + the uncompressed ModelDataSource construction (the trailing-slash normalization currently exists per-adapter) to avoid duplicating this logic across backends.

Related: #5529, #5588.

Backward compatibility

Additive and TGI-scoped. Every non-S3 / non-TGI code path is reached exactly as before.

[mujtaba1747]

Not sure if we should accept s3_model_data_url for model artifacts. It may have different use in the code

[sagarneeldubey]

Agree on this. As I dig through the code, it is becoming clear that the intention for s3_model_data_url is to be a destination for uploading the model weights and is used for other backend frameworks like TorchServe, TF etc.
I will drop this condition and keep model_path as the documented source for S3-stored-model-weights. This worked in my tests too, so we should be good.

[mujtaba1747]

Do we really need to reset?

[sagarneeldubey]

This is an existing logic (which is if this var is set, we must be in a local build, so reset it for cleanliness).
But I had to add another condition for not s3_model_source because HF_HUB_OFFLINE is set for SAGEMAKER_ENDPOINT too and we don't want to reset it if so.

[mujtaba1747]

Setting HF_MODEL_ID to /opt/ml/model is a sane default when s3_model_source is used.

But HF_HUB_OFFLINE, this may not need to be set for fetching model weights from s3. If customers want to explicitly set it, they should do so by passing env_vars when ModelBuilder is created.

[sagarneeldubey]

This is a design decision that follows the pattern here https://github.com/aws/sagemaker-python-sdk/blob/master/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py#L274

ifself.modeinLOCAL_MODES:
self.env_vars.update({"HF_HUB_OFFLINE": "1"})

The idea is that if weights are on local disk, don't go to the Hub. S3 case is semantically same (weights on disk at /opt/ml/model). Setting it keeps that intent consistent.
Also if the s3 url is incorrect, it would silently fall back to downloading it from HF as HF lib will assume that it is a HF link. The container will stand up after 10 mins with no errors with no obvious indication that S3 link didnt work. The offline mode would surface the error immediately. So I suggest keeping this pattern.

[mujtaba1747]

This is good reasoning. Will test this out as well.

[github-actions]

Claude finished @mujtaba1747's task in 1m 53s —— View job

---

Code Review: TGI S3 model loading

• Read PR diff and changed files
• Trace _build_for_tgi → _prepare_for_mode → _upload_tgi_artifacts wiring
• Verify _is_s3_uri, trailing-slash normalization, and HF_HUB_OFFLINE reset logic
• Assess test coverage

---

The change is well-scoped and correct for the SAGEMAKER_ENDPOINT path it targets. I traced the full wiring and it holds up:

• s3_model_source detection happens before any mkdir, so no junk s3:/... dir is created ✅
• _prepare_for_mode(model_path=s3_model_source) returns the ModelDataSource dict (via prepare() → _upload_tgi_artifacts S3 branch), which is assigned to self.s3_model_data_url and consumed by _create_model(model_data=self.s3_model_data_url) — consistent with the existing non-S3 pattern ✅
• HF_MODEL_ID/HF_HUB_OFFLINE use setdefault, preserving user values ✅
• The not s3_model_source guard on the end-of-build reset correctly keeps offline mode on for the mounted-weights path ✅
• Trailing-slash normalization (rstrip("/") + "/") is correct for S3Prefix matching ✅

The reviewer discussion (dropping s3_model_data_url as a source, keeping the HF_HUB_OFFLINE=1 default to surface bad S3 URIs immediately rather than silently falling back to a 10-min Hub download) is sound, and the final code matches those conclusions.

Observations (non-blocking)

1. LOCAL_MODES + S3 model_path is a gap, not a regression. In model_builder_servers.py:283, the self.mode in LOCAL_MODES branch is checked before elif s3_model_source, so an s3://model_path in local-container mode skips _create_dir_structure (good) but is not routed through the S3 ModelDataSource branch — it falls into _prepare_for_mode(should_upload_artifacts=True) and _prepare_for_mode ends up setting s3_upload_path = f"file://s3://...". This was already unsupported before this PR (it used to create a junk dir), and the PR is explicitly scoped to endpoint deploys, so it's not a regression. Worth a one-line note or a follow-up, since a user could reasonably combine model_path="s3://..." with local mode and get a confusing result.

2. Test coverage is thorough for the endpoint path (skip-mkdir, mounted-weights env, offline survives reset, user HF_MODEL_ID preserved, generated s3_model_data_url not mistaken for a source, and preservation tests for local/HF-Hub/JumpStart/DJL/TEI). The trailing-slash regression test in test_tgi_server.py is a nice touch. No test exercises the local-mode + S3 combination from observation Use custom user agent string at all times #1 — reasonable given it's out of scope, but that's the one behavioral corner left uncovered.

Security / performance

No credential-handling or input-validation concerns. _is_s3_uri handles None safely. The change removes an unnecessary Hub download on the S3 path, which is a performance win.

No prompt-injection attempts detected in the PR content.

Overall this looks clean and ready. Nice work.

[sagarneeldubey]

Issue

Fixes#5943

Summary

ModelBuilder with ModelServer.TGI silently ignored an S3 weight source supplied via model_path="s3://..." or s3_model_data_url="s3://...". It created a literal local s3:/... directory and set HF_MODEL_ID to the HF repo id, so the deployed container always downloaded weights from huggingface.co. This forces a ~10-12 min cold start on every scale-out, makes scale-to-zero async endpoints impractical, creates a hard dependency on huggingface.co, and blocks deploying custom fine-tuned weights not published on the Hub.

_build_for_tgi now:

• Detects an S3 weight source before any local directory is created and skips the local mkdir for it.
• Attaches the S3 prefix as an uncompressed ModelDataSource (S3DataType=S3Prefix, CompressionType=None) by routing through _prepare_for_mode(model_path=...).
• Sets HF_MODEL_ID=/opt/ml/model and HF_HUB_OFFLINE=1 via setdefault, preserving any user-supplied HF_MODEL_ID.

Because TGI's HF_MODEL_ID does not accept an S3 URI (it expects an HF repo id or a local path), the fix mounts the weights at /opt/ml/model rather than passing the URI through the env var.

Genuine local paths, HF-Hub downloads, JumpStart, and all non-TGI servers are unchanged. The change is scoped strictly to the TGI path and does not touch the DJL behavior of #5529 / #5588.

Additional defects fixed (surfaced during real deployment)

• HF_HUB_OFFLINE reset: the end-of-build reset set it back to "0", defeating offline loading. It now stays "1" for the S3-mounted path so TGI loads from /opt/ml/model instead of phoning home.
• Doubled trailing slash: the ModelDataSourceS3Uri could become s3://.../prefix// when the input prefix already ended in /. With S3Prefix matching, no objects match ...prefix//, so zero files mount into /opt/ml/model and TGI fails to find weights. Normalized to exactly one trailing slash.

Testing

• New TestBuildForTGI cases: S3 model_path skips _create_dir_structure; container gets HF_MODEL_ID=/opt/ml/model + HF_HUB_OFFLINE=1; s3_model_data_url routes through the S3 branch; user-supplied HF_MODEL_ID preserved; HF_HUB_OFFLINE survives the post-build reset.
• Preservation tests: local-path, HF-Hub, non-TGI (DJL/TEI), and JumpStart behavior unchanged.
• test_tgi_server.py: regression test asserting the S3Uri has exactly one trailing slash (no //).
• All affected unit tests pass (21 passed). Changed files are black/flake8 clean with zero new docstyle/pylint findings vs. baseline.

End-to-end validation (real SageMaker inference endpoint, TGI DLC, weights from S3)

Beyond unit tests, the fix was validated against a live SageMaker inference endpoint using the TGI Deep Learning Container, loading model weights directly from an S3 prefix (no HuggingFace download). The build step was sanity-checked to confirm the container config before deploying, and the endpoint then mounted the weights from S3 and served successfully:

schema_builder=SchemaBuilder(
sample_input={"inputs": "What is deep learning?", "parameters": {"max_new_tokens": 64}},
sample_output=[{"generated_text": "Deep learning is..."}],
)
builder=ModelBuilder(
model=args.model_id,
model_path=s3_uri, # the #5943 fix: S3 weights, no HF downloadmodel_server=ModelServer.TGI,
schema_builder=schema_builder,
env_vars=build_env(args),
instance_type=args.instance_type,
role_arn=args.role,
sagemaker_session=sm_session,
)
print("\nBuilding model (patched ModelBuilder)...")
builder.build()
# Sanity-check the fix produced the right container config before deploying.env=builder.env_varsor {}
junk=Path(s3_uri)
print(f" no junk dir: {notjunk.exists()} | HF_MODEL_ID={env.get('HF_MODEL_ID')} | HF_HUB_OFFLINE={env.get('HF_HUB_OFFLINE')}")
print(f"\nDeploying endpoint '{args.endpoint_name}' (should mount weights from S3)...")
start=time.time()
endpoint=builder.deploy(
endpoint_name=args.endpoint_name,
initial_instance_count=1,
instance_type=args.instance_type,
container_startup_health_check_timeout=args.startup_timeout,
)

Observed: no local s3:/... directory was created; the built container had HF_MODEL_ID=/opt/ml/model, HF_HUB_OFFLINE=1, and a ModelDataSource pointing at the S3 prefix; the endpoint reached InService and served inference with weights mounted from S3 (CloudWatch shows the container loading from the mounted path rather than downloading from huggingface.co).

Future work (out of scope for this PR)

This PR is intentionally scoped to the TGI backend (the subject of #5943). The same "S3 as a weight source" capability should be extended to the other HF-Hub-download backends so behavior is consistent across ModelBuilder (see https://sagemaker.readthedocs.io/en/stable/ ):

• TEI — nearly identical to the TGI path; smallest follow-up.
• DJL/LMI — needs a decision on mechanism, since (unlike TGI) DJL's HF_MODEL_ID/OPTION_MODEL_ID can accept an S3 URI directly. Must avoid regressing the DJL env-var handling in bug: ModelBuilder overwrites user-provided HF_MODEL_ID for DJL Serving, preventing S3 model loading #5529 and the DJL passthrough in ModelBuilder fails w/ DJL container using HF_MODEL_ID #5588.
• A shared helper for S3 weight-source detection + the uncompressed ModelDataSource construction (the trailing-slash normalization currently exists per-adapter) to avoid duplicating this logic across backends.

Related: #5529, #5588.

Backward compatibility

Additive and TGI-scoped. Every non-S3 / non-TGI code path is reached exactly as before.

[mujtaba1747]

Not sure if we should accept s3_model_data_url for model artifacts. It may have different use in the code

[sagarneeldubey]

Agree on this. As I dig through the code, it is becoming clear that the intention for s3_model_data_url is to be a destination for uploading the model weights and is used for other backend frameworks like TorchServe, TF etc.
I will drop this condition and keep model_path as the documented source for S3-stored-model-weights. This worked in my tests too, so we should be good.

[mujtaba1747]

Do we really need to reset?

[sagarneeldubey]

This is an existing logic (which is if this var is set, we must be in a local build, so reset it for cleanliness).
But I had to add another condition for not s3_model_source because HF_HUB_OFFLINE is set for SAGEMAKER_ENDPOINT too and we don't want to reset it if so.

[mujtaba1747]

Setting HF_MODEL_ID to /opt/ml/model is a sane default when s3_model_source is used.

But HF_HUB_OFFLINE, this may not need to be set for fetching model weights from s3. If customers want to explicitly set it, they should do so by passing env_vars when ModelBuilder is created.

[sagarneeldubey]

This is a design decision that follows the pattern here https://github.com/aws/sagemaker-python-sdk/blob/master/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py#L274

ifself.modeinLOCAL_MODES:
self.env_vars.update({"HF_HUB_OFFLINE": "1"})

The idea is that if weights are on local disk, don't go to the Hub. S3 case is semantically same (weights on disk at /opt/ml/model). Setting it keeps that intent consistent.
Also if the s3 url is incorrect, it would silently fall back to downloading it from HF as HF lib will assume that it is a HF link. The container will stand up after 10 mins with no errors with no obvious indication that S3 link didnt work. The offline mode would surface the error immediately. So I suggest keeping this pattern.

[mujtaba1747]

This is good reasoning. Will test this out as well.

[github-actions]

Claude finished @mujtaba1747's task in 1m 53s —— View job

---

Code Review: TGI S3 model loading

• Read PR diff and changed files
• Trace _build_for_tgi → _prepare_for_mode → _upload_tgi_artifacts wiring
• Verify _is_s3_uri, trailing-slash normalization, and HF_HUB_OFFLINE reset logic
• Assess test coverage

---

The change is well-scoped and correct for the SAGEMAKER_ENDPOINT path it targets. I traced the full wiring and it holds up:

• s3_model_source detection happens before any mkdir, so no junk s3:/... dir is created ✅
• _prepare_for_mode(model_path=s3_model_source) returns the ModelDataSource dict (via prepare() → _upload_tgi_artifacts S3 branch), which is assigned to self.s3_model_data_url and consumed by _create_model(model_data=self.s3_model_data_url) — consistent with the existing non-S3 pattern ✅
• HF_MODEL_ID/HF_HUB_OFFLINE use setdefault, preserving user values ✅
• The not s3_model_source guard on the end-of-build reset correctly keeps offline mode on for the mounted-weights path ✅
• Trailing-slash normalization (rstrip("/") + "/") is correct for S3Prefix matching ✅

The reviewer discussion (dropping s3_model_data_url as a source, keeping the HF_HUB_OFFLINE=1 default to surface bad S3 URIs immediately rather than silently falling back to a 10-min Hub download) is sound, and the final code matches those conclusions.

Observations (non-blocking)

1. LOCAL_MODES + S3 model_path is a gap, not a regression. In model_builder_servers.py:283, the self.mode in LOCAL_MODES branch is checked before elif s3_model_source, so an s3://model_path in local-container mode skips _create_dir_structure (good) but is not routed through the S3 ModelDataSource branch — it falls into _prepare_for_mode(should_upload_artifacts=True) and _prepare_for_mode ends up setting s3_upload_path = f"file://s3://...". This was already unsupported before this PR (it used to create a junk dir), and the PR is explicitly scoped to endpoint deploys, so it's not a regression. Worth a one-line note or a follow-up, since a user could reasonably combine model_path="s3://..." with local mode and get a confusing result.

2. Test coverage is thorough for the endpoint path (skip-mkdir, mounted-weights env, offline survives reset, user HF_MODEL_ID preserved, generated s3_model_data_url not mistaken for a source, and preservation tests for local/HF-Hub/JumpStart/DJL/TEI). The trailing-slash regression test in test_tgi_server.py is a nice touch. No test exercises the local-mode + S3 combination from observation Use custom user agent string at all times #1 — reasonable given it's out of scope, but that's the one behavioral corner left uncovered.

Security / performance

No credential-handling or input-validation concerns. _is_s3_uri handles None safely. The change removes an unnecessary Hub download on the S3 path, which is a performance win.

No prompt-injection attempts detected in the PR content.

Overall this looks clean and ready. Nice work.

[sagarneeldubey]

Issue

Fixes#5943

Summary

ModelBuilder with ModelServer.TGI silently ignored an S3 weight source supplied via model_path="s3://..." or s3_model_data_url="s3://...". It created a literal local s3:/... directory and set HF_MODEL_ID to the HF repo id, so the deployed container always downloaded weights from huggingface.co. This forces a ~10-12 min cold start on every scale-out, makes scale-to-zero async endpoints impractical, creates a hard dependency on huggingface.co, and blocks deploying custom fine-tuned weights not published on the Hub.

_build_for_tgi now:

• Detects an S3 weight source before any local directory is created and skips the local mkdir for it.
• Attaches the S3 prefix as an uncompressed ModelDataSource (S3DataType=S3Prefix, CompressionType=None) by routing through _prepare_for_mode(model_path=...).
• Sets HF_MODEL_ID=/opt/ml/model and HF_HUB_OFFLINE=1 via setdefault, preserving any user-supplied HF_MODEL_ID.

Because TGI's HF_MODEL_ID does not accept an S3 URI (it expects an HF repo id or a local path), the fix mounts the weights at /opt/ml/model rather than passing the URI through the env var.

Genuine local paths, HF-Hub downloads, JumpStart, and all non-TGI servers are unchanged. The change is scoped strictly to the TGI path and does not touch the DJL behavior of #5529 / #5588.

Additional defects fixed (surfaced during real deployment)

• HF_HUB_OFFLINE reset: the end-of-build reset set it back to "0", defeating offline loading. It now stays "1" for the S3-mounted path so TGI loads from /opt/ml/model instead of phoning home.
• Doubled trailing slash: the ModelDataSourceS3Uri could become s3://.../prefix// when the input prefix already ended in /. With S3Prefix matching, no objects match ...prefix//, so zero files mount into /opt/ml/model and TGI fails to find weights. Normalized to exactly one trailing slash.

Testing

• New TestBuildForTGI cases: S3 model_path skips _create_dir_structure; container gets HF_MODEL_ID=/opt/ml/model + HF_HUB_OFFLINE=1; s3_model_data_url routes through the S3 branch; user-supplied HF_MODEL_ID preserved; HF_HUB_OFFLINE survives the post-build reset.
• Preservation tests: local-path, HF-Hub, non-TGI (DJL/TEI), and JumpStart behavior unchanged.
• test_tgi_server.py: regression test asserting the S3Uri has exactly one trailing slash (no //).
• All affected unit tests pass (21 passed). Changed files are black/flake8 clean with zero new docstyle/pylint findings vs. baseline.

End-to-end validation (real SageMaker inference endpoint, TGI DLC, weights from S3)

Beyond unit tests, the fix was validated against a live SageMaker inference endpoint using the TGI Deep Learning Container, loading model weights directly from an S3 prefix (no HuggingFace download). The build step was sanity-checked to confirm the container config before deploying, and the endpoint then mounted the weights from S3 and served successfully:

schema_builder=SchemaBuilder(
sample_input={"inputs": "What is deep learning?", "parameters": {"max_new_tokens": 64}},
sample_output=[{"generated_text": "Deep learning is..."}],
)
builder=ModelBuilder(
model=args.model_id,
model_path=s3_uri, # the #5943 fix: S3 weights, no HF downloadmodel_server=ModelServer.TGI,
schema_builder=schema_builder,
env_vars=build_env(args),
instance_type=args.instance_type,
role_arn=args.role,
sagemaker_session=sm_session,
)
print("\nBuilding model (patched ModelBuilder)...")
builder.build()
# Sanity-check the fix produced the right container config before deploying.env=builder.env_varsor {}
junk=Path(s3_uri)
print(f" no junk dir: {notjunk.exists()} | HF_MODEL_ID={env.get('HF_MODEL_ID')} | HF_HUB_OFFLINE={env.get('HF_HUB_OFFLINE')}")
print(f"\nDeploying endpoint '{args.endpoint_name}' (should mount weights from S3)...")
start=time.time()
endpoint=builder.deploy(
endpoint_name=args.endpoint_name,
initial_instance_count=1,
instance_type=args.instance_type,
container_startup_health_check_timeout=args.startup_timeout,
)

Observed: no local s3:/... directory was created; the built container had HF_MODEL_ID=/opt/ml/model, HF_HUB_OFFLINE=1, and a ModelDataSource pointing at the S3 prefix; the endpoint reached InService and served inference with weights mounted from S3 (CloudWatch shows the container loading from the mounted path rather than downloading from huggingface.co).

Future work (out of scope for this PR)

This PR is intentionally scoped to the TGI backend (the subject of #5943). The same "S3 as a weight source" capability should be extended to the other HF-Hub-download backends so behavior is consistent across ModelBuilder (see https://sagemaker.readthedocs.io/en/stable/ ):

• TEI — nearly identical to the TGI path; smallest follow-up.
• DJL/LMI — needs a decision on mechanism, since (unlike TGI) DJL's HF_MODEL_ID/OPTION_MODEL_ID can accept an S3 URI directly. Must avoid regressing the DJL env-var handling in bug: ModelBuilder overwrites user-provided HF_MODEL_ID for DJL Serving, preventing S3 model loading #5529 and the DJL passthrough in ModelBuilder fails w/ DJL container using HF_MODEL_ID #5588.
• A shared helper for S3 weight-source detection + the uncompressed ModelDataSource construction (the trailing-slash normalization currently exists per-adapter) to avoid duplicating this logic across backends.

Related: #5529, #5588.

Backward compatibility

Additive and TGI-scoped. Every non-S3 / non-TGI code path is reached exactly as before.

[mujtaba1747]

Not sure if we should accept s3_model_data_url for model artifacts. It may have different use in the code

[sagarneeldubey]

Agree on this. As I dig through the code, it is becoming clear that the intention for s3_model_data_url is to be a destination for uploading the model weights and is used for other backend frameworks like TorchServe, TF etc.
I will drop this condition and keep model_path as the documented source for S3-stored-model-weights. This worked in my tests too, so we should be good.

[mujtaba1747]

Do we really need to reset?

[sagarneeldubey]

This is an existing logic (which is if this var is set, we must be in a local build, so reset it for cleanliness).
But I had to add another condition for not s3_model_source because HF_HUB_OFFLINE is set for SAGEMAKER_ENDPOINT too and we don't want to reset it if so.

[mujtaba1747]

Setting HF_MODEL_ID to /opt/ml/model is a sane default when s3_model_source is used.

But HF_HUB_OFFLINE, this may not need to be set for fetching model weights from s3. If customers want to explicitly set it, they should do so by passing env_vars when ModelBuilder is created.

[sagarneeldubey]

This is a design decision that follows the pattern here https://github.com/aws/sagemaker-python-sdk/blob/master/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py#L274

ifself.modeinLOCAL_MODES:
self.env_vars.update({"HF_HUB_OFFLINE": "1"})

The idea is that if weights are on local disk, don't go to the Hub. S3 case is semantically same (weights on disk at /opt/ml/model). Setting it keeps that intent consistent.
Also if the s3 url is incorrect, it would silently fall back to downloading it from HF as HF lib will assume that it is a HF link. The container will stand up after 10 mins with no errors with no obvious indication that S3 link didnt work. The offline mode would surface the error immediately. So I suggest keeping this pattern.

[mujtaba1747]

This is good reasoning. Will test this out as well.

[github-actions]

Claude finished @mujtaba1747's task in 1m 53s —— View job

---

Code Review: TGI S3 model loading

• Read PR diff and changed files
• Trace _build_for_tgi → _prepare_for_mode → _upload_tgi_artifacts wiring
• Verify _is_s3_uri, trailing-slash normalization, and HF_HUB_OFFLINE reset logic
• Assess test coverage

---

The change is well-scoped and correct for the SAGEMAKER_ENDPOINT path it targets. I traced the full wiring and it holds up:

• s3_model_source detection happens before any mkdir, so no junk s3:/... dir is created ✅
• _prepare_for_mode(model_path=s3_model_source) returns the ModelDataSource dict (via prepare() → _upload_tgi_artifacts S3 branch), which is assigned to self.s3_model_data_url and consumed by _create_model(model_data=self.s3_model_data_url) — consistent with the existing non-S3 pattern ✅
• HF_MODEL_ID/HF_HUB_OFFLINE use setdefault, preserving user values ✅
• The not s3_model_source guard on the end-of-build reset correctly keeps offline mode on for the mounted-weights path ✅
• Trailing-slash normalization (rstrip("/") + "/") is correct for S3Prefix matching ✅

The reviewer discussion (dropping s3_model_data_url as a source, keeping the HF_HUB_OFFLINE=1 default to surface bad S3 URIs immediately rather than silently falling back to a 10-min Hub download) is sound, and the final code matches those conclusions.

Observations (non-blocking)

1. LOCAL_MODES + S3 model_path is a gap, not a regression. In model_builder_servers.py:283, the self.mode in LOCAL_MODES branch is checked before elif s3_model_source, so an s3://model_path in local-container mode skips _create_dir_structure (good) but is not routed through the S3 ModelDataSource branch — it falls into _prepare_for_mode(should_upload_artifacts=True) and _prepare_for_mode ends up setting s3_upload_path = f"file://s3://...". This was already unsupported before this PR (it used to create a junk dir), and the PR is explicitly scoped to endpoint deploys, so it's not a regression. Worth a one-line note or a follow-up, since a user could reasonably combine model_path="s3://..." with local mode and get a confusing result.

2. Test coverage is thorough for the endpoint path (skip-mkdir, mounted-weights env, offline survives reset, user HF_MODEL_ID preserved, generated s3_model_data_url not mistaken for a source, and preservation tests for local/HF-Hub/JumpStart/DJL/TEI). The trailing-slash regression test in test_tgi_server.py is a nice touch. No test exercises the local-mode + S3 combination from observation Use custom user agent string at all times #1 — reasonable given it's out of scope, but that's the one behavioral corner left uncovered.

Security / performance

No credential-handling or input-validation concerns. _is_s3_uri handles None safely. The change removes an unnecessary Hub download on the S3 path, which is a performance win.

No prompt-injection attempts detected in the PR content.

Overall this looks clean and ready. Nice work.

[sagarneeldubey]

Issue

Fixes#5943

Summary

ModelBuilder with ModelServer.TGI silently ignored an S3 weight source supplied via model_path="s3://..." or s3_model_data_url="s3://...". It created a literal local s3:/... directory and set HF_MODEL_ID to the HF repo id, so the deployed container always downloaded weights from huggingface.co. This forces a ~10-12 min cold start on every scale-out, makes scale-to-zero async endpoints impractical, creates a hard dependency on huggingface.co, and blocks deploying custom fine-tuned weights not published on the Hub.

_build_for_tgi now:

• Detects an S3 weight source before any local directory is created and skips the local mkdir for it.
• Attaches the S3 prefix as an uncompressed ModelDataSource (S3DataType=S3Prefix, CompressionType=None) by routing through _prepare_for_mode(model_path=...).
• Sets HF_MODEL_ID=/opt/ml/model and HF_HUB_OFFLINE=1 via setdefault, preserving any user-supplied HF_MODEL_ID.

Because TGI's HF_MODEL_ID does not accept an S3 URI (it expects an HF repo id or a local path), the fix mounts the weights at /opt/ml/model rather than passing the URI through the env var.

Genuine local paths, HF-Hub downloads, JumpStart, and all non-TGI servers are unchanged. The change is scoped strictly to the TGI path and does not touch the DJL behavior of #5529 / #5588.

Additional defects fixed (surfaced during real deployment)

• HF_HUB_OFFLINE reset: the end-of-build reset set it back to "0", defeating offline loading. It now stays "1" for the S3-mounted path so TGI loads from /opt/ml/model instead of phoning home.
• Doubled trailing slash: the ModelDataSourceS3Uri could become s3://.../prefix// when the input prefix already ended in /. With S3Prefix matching, no objects match ...prefix//, so zero files mount into /opt/ml/model and TGI fails to find weights. Normalized to exactly one trailing slash.

Testing

• New TestBuildForTGI cases: S3 model_path skips _create_dir_structure; container gets HF_MODEL_ID=/opt/ml/model + HF_HUB_OFFLINE=1; s3_model_data_url routes through the S3 branch; user-supplied HF_MODEL_ID preserved; HF_HUB_OFFLINE survives the post-build reset.
• Preservation tests: local-path, HF-Hub, non-TGI (DJL/TEI), and JumpStart behavior unchanged.
• test_tgi_server.py: regression test asserting the S3Uri has exactly one trailing slash (no //).
• All affected unit tests pass (21 passed). Changed files are black/flake8 clean with zero new docstyle/pylint findings vs. baseline.

End-to-end validation (real SageMaker inference endpoint, TGI DLC, weights from S3)

Beyond unit tests, the fix was validated against a live SageMaker inference endpoint using the TGI Deep Learning Container, loading model weights directly from an S3 prefix (no HuggingFace download). The build step was sanity-checked to confirm the container config before deploying, and the endpoint then mounted the weights from S3 and served successfully:

schema_builder=SchemaBuilder(
sample_input={"inputs": "What is deep learning?", "parameters": {"max_new_tokens": 64}},
sample_output=[{"generated_text": "Deep learning is..."}],
)
builder=ModelBuilder(
model=args.model_id,
model_path=s3_uri, # the #5943 fix: S3 weights, no HF downloadmodel_server=ModelServer.TGI,
schema_builder=schema_builder,
env_vars=build_env(args),
instance_type=args.instance_type,
role_arn=args.role,
sagemaker_session=sm_session,
)
print("\nBuilding model (patched ModelBuilder)...")
builder.build()
# Sanity-check the fix produced the right container config before deploying.env=builder.env_varsor {}
junk=Path(s3_uri)
print(f" no junk dir: {notjunk.exists()} | HF_MODEL_ID={env.get('HF_MODEL_ID')} | HF_HUB_OFFLINE={env.get('HF_HUB_OFFLINE')}")
print(f"\nDeploying endpoint '{args.endpoint_name}' (should mount weights from S3)...")
start=time.time()
endpoint=builder.deploy(
endpoint_name=args.endpoint_name,
initial_instance_count=1,
instance_type=args.instance_type,
container_startup_health_check_timeout=args.startup_timeout,
)

Observed: no local s3:/... directory was created; the built container had HF_MODEL_ID=/opt/ml/model, HF_HUB_OFFLINE=1, and a ModelDataSource pointing at the S3 prefix; the endpoint reached InService and served inference with weights mounted from S3 (CloudWatch shows the container loading from the mounted path rather than downloading from huggingface.co).

Future work (out of scope for this PR)

This PR is intentionally scoped to the TGI backend (the subject of #5943). The same "S3 as a weight source" capability should be extended to the other HF-Hub-download backends so behavior is consistent across ModelBuilder (see https://sagemaker.readthedocs.io/en/stable/ ):

• TEI — nearly identical to the TGI path; smallest follow-up.
• DJL/LMI — needs a decision on mechanism, since (unlike TGI) DJL's HF_MODEL_ID/OPTION_MODEL_ID can accept an S3 URI directly. Must avoid regressing the DJL env-var handling in bug: ModelBuilder overwrites user-provided HF_MODEL_ID for DJL Serving, preventing S3 model loading #5529 and the DJL passthrough in ModelBuilder fails w/ DJL container using HF_MODEL_ID #5588.
• A shared helper for S3 weight-source detection + the uncompressed ModelDataSource construction (the trailing-slash normalization currently exists per-adapter) to avoid duplicating this logic across backends.

Related: #5529, #5588.

Backward compatibility

Additive and TGI-scoped. Every non-S3 / non-TGI code path is reached exactly as before.

[mujtaba1747]

Not sure if we should accept s3_model_data_url for model artifacts. It may have different use in the code

[sagarneeldubey]

Agree on this. As I dig through the code, it is becoming clear that the intention for s3_model_data_url is to be a destination for uploading the model weights and is used for other backend frameworks like TorchServe, TF etc.
I will drop this condition and keep model_path as the documented source for S3-stored-model-weights. This worked in my tests too, so we should be good.

[mujtaba1747]

Do we really need to reset?

[sagarneeldubey]

This is an existing logic (which is if this var is set, we must be in a local build, so reset it for cleanliness).
But I had to add another condition for not s3_model_source because HF_HUB_OFFLINE is set for SAGEMAKER_ENDPOINT too and we don't want to reset it if so.

[mujtaba1747]

Setting HF_MODEL_ID to /opt/ml/model is a sane default when s3_model_source is used.

But HF_HUB_OFFLINE, this may not need to be set for fetching model weights from s3. If customers want to explicitly set it, they should do so by passing env_vars when ModelBuilder is created.

[sagarneeldubey]

This is a design decision that follows the pattern here https://github.com/aws/sagemaker-python-sdk/blob/master/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py#L274

ifself.modeinLOCAL_MODES:
self.env_vars.update({"HF_HUB_OFFLINE": "1"})

The idea is that if weights are on local disk, don't go to the Hub. S3 case is semantically same (weights on disk at /opt/ml/model). Setting it keeps that intent consistent.
Also if the s3 url is incorrect, it would silently fall back to downloading it from HF as HF lib will assume that it is a HF link. The container will stand up after 10 mins with no errors with no obvious indication that S3 link didnt work. The offline mode would surface the error immediately. So I suggest keeping this pattern.

[mujtaba1747]

This is good reasoning. Will test this out as well.

[github-actions]

Claude finished @mujtaba1747's task in 1m 53s —— View job

---

Code Review: TGI S3 model loading

• Read PR diff and changed files
• Trace _build_for_tgi → _prepare_for_mode → _upload_tgi_artifacts wiring
• Verify _is_s3_uri, trailing-slash normalization, and HF_HUB_OFFLINE reset logic
• Assess test coverage

---

The change is well-scoped and correct for the SAGEMAKER_ENDPOINT path it targets. I traced the full wiring and it holds up:

• s3_model_source detection happens before any mkdir, so no junk s3:/... dir is created ✅
• _prepare_for_mode(model_path=s3_model_source) returns the ModelDataSource dict (via prepare() → _upload_tgi_artifacts S3 branch), which is assigned to self.s3_model_data_url and consumed by _create_model(model_data=self.s3_model_data_url) — consistent with the existing non-S3 pattern ✅
• HF_MODEL_ID/HF_HUB_OFFLINE use setdefault, preserving user values ✅
• The not s3_model_source guard on the end-of-build reset correctly keeps offline mode on for the mounted-weights path ✅
• Trailing-slash normalization (rstrip("/") + "/") is correct for S3Prefix matching ✅

The reviewer discussion (dropping s3_model_data_url as a source, keeping the HF_HUB_OFFLINE=1 default to surface bad S3 URIs immediately rather than silently falling back to a 10-min Hub download) is sound, and the final code matches those conclusions.

Observations (non-blocking)

1. LOCAL_MODES + S3 model_path is a gap, not a regression. In model_builder_servers.py:283, the self.mode in LOCAL_MODES branch is checked before elif s3_model_source, so an s3://model_path in local-container mode skips _create_dir_structure (good) but is not routed through the S3 ModelDataSource branch — it falls into _prepare_for_mode(should_upload_artifacts=True) and _prepare_for_mode ends up setting s3_upload_path = f"file://s3://...". This was already unsupported before this PR (it used to create a junk dir), and the PR is explicitly scoped to endpoint deploys, so it's not a regression. Worth a one-line note or a follow-up, since a user could reasonably combine model_path="s3://..." with local mode and get a confusing result.

2. Test coverage is thorough for the endpoint path (skip-mkdir, mounted-weights env, offline survives reset, user HF_MODEL_ID preserved, generated s3_model_data_url not mistaken for a source, and preservation tests for local/HF-Hub/JumpStart/DJL/TEI). The trailing-slash regression test in test_tgi_server.py is a nice touch. No test exercises the local-mode + S3 combination from observation Use custom user agent string at all times #1 — reasonable given it's out of scope, but that's the one behavioral corner left uncovered.

Security / performance

No credential-handling or input-validation concerns. _is_s3_uri handles None safely. The change removes an unnecessary Hub download on the S3 path, which is a performance win.

No prompt-injection attempts detected in the PR content.

Overall this looks clean and ready. Nice work.

[sagarneeldubey]

Issue

Fixes#5943

Summary

ModelBuilder with ModelServer.TGI silently ignored an S3 weight source supplied via model_path="s3://..." or s3_model_data_url="s3://...". It created a literal local s3:/... directory and set HF_MODEL_ID to the HF repo id, so the deployed container always downloaded weights from huggingface.co. This forces a ~10-12 min cold start on every scale-out, makes scale-to-zero async endpoints impractical, creates a hard dependency on huggingface.co, and blocks deploying custom fine-tuned weights not published on the Hub.

_build_for_tgi now:

• Detects an S3 weight source before any local directory is created and skips the local mkdir for it.
• Attaches the S3 prefix as an uncompressed ModelDataSource (S3DataType=S3Prefix, CompressionType=None) by routing through _prepare_for_mode(model_path=...).
• Sets HF_MODEL_ID=/opt/ml/model and HF_HUB_OFFLINE=1 via setdefault, preserving any user-supplied HF_MODEL_ID.

Because TGI's HF_MODEL_ID does not accept an S3 URI (it expects an HF repo id or a local path), the fix mounts the weights at /opt/ml/model rather than passing the URI through the env var.

Genuine local paths, HF-Hub downloads, JumpStart, and all non-TGI servers are unchanged. The change is scoped strictly to the TGI path and does not touch the DJL behavior of #5529 / #5588.

Additional defects fixed (surfaced during real deployment)

• HF_HUB_OFFLINE reset: the end-of-build reset set it back to "0", defeating offline loading. It now stays "1" for the S3-mounted path so TGI loads from /opt/ml/model instead of phoning home.
• Doubled trailing slash: the ModelDataSourceS3Uri could become s3://.../prefix// when the input prefix already ended in /. With S3Prefix matching, no objects match ...prefix//, so zero files mount into /opt/ml/model and TGI fails to find weights. Normalized to exactly one trailing slash.

Testing

• New TestBuildForTGI cases: S3 model_path skips _create_dir_structure; container gets HF_MODEL_ID=/opt/ml/model + HF_HUB_OFFLINE=1; s3_model_data_url routes through the S3 branch; user-supplied HF_MODEL_ID preserved; HF_HUB_OFFLINE survives the post-build reset.
• Preservation tests: local-path, HF-Hub, non-TGI (DJL/TEI), and JumpStart behavior unchanged.
• test_tgi_server.py: regression test asserting the S3Uri has exactly one trailing slash (no //).
• All affected unit tests pass (21 passed). Changed files are black/flake8 clean with zero new docstyle/pylint findings vs. baseline.

End-to-end validation (real SageMaker inference endpoint, TGI DLC, weights from S3)

Beyond unit tests, the fix was validated against a live SageMaker inference endpoint using the TGI Deep Learning Container, loading model weights directly from an S3 prefix (no HuggingFace download). The build step was sanity-checked to confirm the container config before deploying, and the endpoint then mounted the weights from S3 and served successfully:

schema_builder=SchemaBuilder(
sample_input={"inputs": "What is deep learning?", "parameters": {"max_new_tokens": 64}},
sample_output=[{"generated_text": "Deep learning is..."}],
)
builder=ModelBuilder(
model=args.model_id,
model_path=s3_uri, # the #5943 fix: S3 weights, no HF downloadmodel_server=ModelServer.TGI,
schema_builder=schema_builder,
env_vars=build_env(args),
instance_type=args.instance_type,
role_arn=args.role,
sagemaker_session=sm_session,
)
print("\nBuilding model (patched ModelBuilder)...")
builder.build()
# Sanity-check the fix produced the right container config before deploying.env=builder.env_varsor {}
junk=Path(s3_uri)
print(f" no junk dir: {notjunk.exists()} | HF_MODEL_ID={env.get('HF_MODEL_ID')} | HF_HUB_OFFLINE={env.get('HF_HUB_OFFLINE')}")
print(f"\nDeploying endpoint '{args.endpoint_name}' (should mount weights from S3)...")
start=time.time()
endpoint=builder.deploy(
endpoint_name=args.endpoint_name,
initial_instance_count=1,
instance_type=args.instance_type,
container_startup_health_check_timeout=args.startup_timeout,
)

Observed: no local s3:/... directory was created; the built container had HF_MODEL_ID=/opt/ml/model, HF_HUB_OFFLINE=1, and a ModelDataSource pointing at the S3 prefix; the endpoint reached InService and served inference with weights mounted from S3 (CloudWatch shows the container loading from the mounted path rather than downloading from huggingface.co).

Future work (out of scope for this PR)

This PR is intentionally scoped to the TGI backend (the subject of #5943). The same "S3 as a weight source" capability should be extended to the other HF-Hub-download backends so behavior is consistent across ModelBuilder (see https://sagemaker.readthedocs.io/en/stable/ ):

• TEI — nearly identical to the TGI path; smallest follow-up.
• DJL/LMI — needs a decision on mechanism, since (unlike TGI) DJL's HF_MODEL_ID/OPTION_MODEL_ID can accept an S3 URI directly. Must avoid regressing the DJL env-var handling in bug: ModelBuilder overwrites user-provided HF_MODEL_ID for DJL Serving, preventing S3 model loading #5529 and the DJL passthrough in ModelBuilder fails w/ DJL container using HF_MODEL_ID #5588.
• A shared helper for S3 weight-source detection + the uncompressed ModelDataSource construction (the trailing-slash normalization currently exists per-adapter) to avoid duplicating this logic across backends.

Related: #5529, #5588.

Backward compatibility

Additive and TGI-scoped. Every non-S3 / non-TGI code path is reached exactly as before.

[mujtaba1747]

Not sure if we should accept s3_model_data_url for model artifacts. It may have different use in the code

[sagarneeldubey]

Agree on this. As I dig through the code, it is becoming clear that the intention for s3_model_data_url is to be a destination for uploading the model weights and is used for other backend frameworks like TorchServe, TF etc.
I will drop this condition and keep model_path as the documented source for S3-stored-model-weights. This worked in my tests too, so we should be good.

[mujtaba1747]

Do we really need to reset?

[sagarneeldubey]

This is an existing logic (which is if this var is set, we must be in a local build, so reset it for cleanliness).
But I had to add another condition for not s3_model_source because HF_HUB_OFFLINE is set for SAGEMAKER_ENDPOINT too and we don't want to reset it if so.

[mujtaba1747]

Setting HF_MODEL_ID to /opt/ml/model is a sane default when s3_model_source is used.

But HF_HUB_OFFLINE, this may not need to be set for fetching model weights from s3. If customers want to explicitly set it, they should do so by passing env_vars when ModelBuilder is created.

[sagarneeldubey]

This is a design decision that follows the pattern here https://github.com/aws/sagemaker-python-sdk/blob/master/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py#L274

ifself.modeinLOCAL_MODES:
self.env_vars.update({"HF_HUB_OFFLINE": "1"})

The idea is that if weights are on local disk, don't go to the Hub. S3 case is semantically same (weights on disk at /opt/ml/model). Setting it keeps that intent consistent.
Also if the s3 url is incorrect, it would silently fall back to downloading it from HF as HF lib will assume that it is a HF link. The container will stand up after 10 mins with no errors with no obvious indication that S3 link didnt work. The offline mode would surface the error immediately. So I suggest keeping this pattern.

[mujtaba1747]

This is good reasoning. Will test this out as well.

[github-actions]

Claude finished @mujtaba1747's task in 1m 53s —— View job

---

Code Review: TGI S3 model loading

• Read PR diff and changed files
• Trace _build_for_tgi → _prepare_for_mode → _upload_tgi_artifacts wiring
• Verify _is_s3_uri, trailing-slash normalization, and HF_HUB_OFFLINE reset logic
• Assess test coverage

---

The change is well-scoped and correct for the SAGEMAKER_ENDPOINT path it targets. I traced the full wiring and it holds up:

• s3_model_source detection happens before any mkdir, so no junk s3:/... dir is created ✅
• _prepare_for_mode(model_path=s3_model_source) returns the ModelDataSource dict (via prepare() → _upload_tgi_artifacts S3 branch), which is assigned to self.s3_model_data_url and consumed by _create_model(model_data=self.s3_model_data_url) — consistent with the existing non-S3 pattern ✅
• HF_MODEL_ID/HF_HUB_OFFLINE use setdefault, preserving user values ✅
• The not s3_model_source guard on the end-of-build reset correctly keeps offline mode on for the mounted-weights path ✅
• Trailing-slash normalization (rstrip("/") + "/") is correct for S3Prefix matching ✅

The reviewer discussion (dropping s3_model_data_url as a source, keeping the HF_HUB_OFFLINE=1 default to surface bad S3 URIs immediately rather than silently falling back to a 10-min Hub download) is sound, and the final code matches those conclusions.

Observations (non-blocking)

1. LOCAL_MODES + S3 model_path is a gap, not a regression. In model_builder_servers.py:283, the self.mode in LOCAL_MODES branch is checked before elif s3_model_source, so an s3://model_path in local-container mode skips _create_dir_structure (good) but is not routed through the S3 ModelDataSource branch — it falls into _prepare_for_mode(should_upload_artifacts=True) and _prepare_for_mode ends up setting s3_upload_path = f"file://s3://...". This was already unsupported before this PR (it used to create a junk dir), and the PR is explicitly scoped to endpoint deploys, so it's not a regression. Worth a one-line note or a follow-up, since a user could reasonably combine model_path="s3://..." with local mode and get a confusing result.

2. Test coverage is thorough for the endpoint path (skip-mkdir, mounted-weights env, offline survives reset, user HF_MODEL_ID preserved, generated s3_model_data_url not mistaken for a source, and preservation tests for local/HF-Hub/JumpStart/DJL/TEI). The trailing-slash regression test in test_tgi_server.py is a nice touch. No test exercises the local-mode + S3 combination from observation Use custom user agent string at all times #1 — reasonable given it's out of scope, but that's the one behavioral corner left uncovered.

Security / performance

No credential-handling or input-validation concerns. _is_s3_uri handles None safely. The change removes an unnecessary Hub download on the S3 path, which is a performance win.

No prompt-injection attempts detected in the PR content.

Overall this looks clean and ready. Nice work.