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

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

ModelBuilder with ModelServer.TGI silently ignored an S3 weight source
(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 container always downloaded weights from huggingface.co.

_build_for_tgi now detects an S3 weight source, skips the local mkdir for
it, attaches the S3 prefix as an uncompressed ModelDataSource, and sets
HF_MODEL_ID=/opt/ml/model and HF_HUB_OFFLINE=1 (via setdefault, preserving
any user-supplied HF_MODEL_ID). Genuine local paths, HF-Hub downloads,
JumpStart, and all non-TGI servers are unchanged.

Also fixes two defects surfaced during real deployment:
- HF_HUB_OFFLINE was reset to "0" at the end of the build; it now stays
  "1" for the S3-mounted path so TGI loads from /opt/ml/model.
- The uncompressed ModelDataSource S3Uri could end in "//" when the input
  prefix already had a trailing slash; normalized to exactly one slash so
  S3Prefix matching finds the weight objects.

Adds unit and regression tests for all of the above.

* fix(tgi): distinguish S3 model source from upload destination

---------

Co-authored-by: Sagar Dubey <dubeysag@amazon.com>
Co-authored-by: Syed Mujtaba <42322958+mujtaba1747@users.noreply.github.com>

Author: 3 people

sagemaker-serve/src/sagemaker/serve/model_builder_servers.py

@@ -26,6 +26,7 @@
 # SageMaker core imports
 from sagemaker.core.resources import Model, Endpoint
 from sagemaker.core.utils.utils import logger
+from sagemaker.core.common_utils import _is_s3_uri
 
 
 # SageMaker serve imports
@@ -220,11 +221,25 @@ def _build_for_tgi(self) -> Model:
 
         from sagemaker.serve.model_server.tgi.prepare import _create_dir_structure
 
-        _create_dir_structure(self.model_path)
+        # Detect an S3 weight source from model_path before any local directory is
+        # created. TGI's HF_MODEL_ID does not accept an S3 URI, so an S3 source is
+        # attached as an uncompressed ModelDataSource (mounted at /opt/ml/model)
+        # instead of being downloaded from the HuggingFace Hub.
+        s3_model_source = self.model_path if _is_s3_uri(self.model_path) else None
+
+        # Skip the local mkdir for an S3 source so we do not create a literal
+        # local "s3:/..." directory tree; only create it for genuine local paths.
+        if not s3_model_source:
+            _create_dir_structure(self.model_path)
 
         if isinstance(self.model, str) and not self._is_jumpstart_model_id():
             # Configure HuggingFace model for TGI
-            self.env_vars.setdefault("HF_MODEL_ID", self.model)
+            if s3_model_source:
+                # Weights are mounted at /opt/ml/model; do not download from the Hub.
+                self.env_vars.setdefault("HF_MODEL_ID", "/opt/ml/model")
+                self.env_vars.setdefault("HF_HUB_OFFLINE", "1")
+            else:
+                self.env_vars.setdefault("HF_MODEL_ID", self.model)
 
             self.hf_model_config = _get_model_config_properties_from_hf(
                 self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
@@ -267,6 +282,11 @@ def _build_for_tgi(self) -> Model:
         if not self._optimizing:
             if self.mode in LOCAL_MODES:
                 self._prepare_for_mode(should_upload_artifacts=True)
+            elif s3_model_source:
+                # Route the S3 weight source through _prepare_for_mode so the
+                # _upload_tgi_artifacts S3 branch builds the uncompressed
+                # ModelDataSource (CompressionType="None", S3DataType="S3Prefix").
+                self.s3_model_data_url, _ = self._prepare_for_mode(model_path=s3_model_source)
             else:
                 self.s3_model_data_url, _ = self._prepare_for_mode()
 
@@ -299,7 +319,10 @@ def _build_for_tgi(self) -> Model:
 
         model = self._create_model()
 
-        if "HF_HUB_OFFLINE" in self.env_vars:
+        # Reset the in-memory HF_HUB_OFFLINE flag after the container is built,
+        # EXCEPT when weights are mounted from S3: those must stay offline so TGI
+        # loads from /opt/ml/model instead of phoning home to the HuggingFace Hub.
+        if "HF_HUB_OFFLINE" in self.env_vars and not s3_model_source:
             self.env_vars.update({"HF_HUB_OFFLINE": "0"})
 
         return model

sagemaker-serve/src/sagemaker/serve/model_server/tgi/server.py

@@ -131,7 +131,11 @@ def _upload_tgi_artifacts(
                 "S3DataSource": {
                     "CompressionType": "None",
                     "S3DataType": "S3Prefix",
-                    "S3Uri": model_data_url + "/",
+                    # Normalize trailing slashes so an S3Prefix URI has exactly one.
+                    # A user-supplied weight prefix may already end in "/", and a
+                    # doubled "s3://.../prefix//" would not match the actual object
+                    # keys under "s3://.../prefix/" with S3Prefix matching.
+                    "S3Uri": model_data_url.rstrip("/") + "/",
                 }
             }
             if model_data_url

sagemaker-serve/tests/unit/model_server/test_tgi_server.py

@@ -127,6 +127,33 @@ def test_upload_tgi_artifacts_with_s3_path(self, mock_is_s3):
         self.assertIsNotNone(model_data)
         self.assertEqual(model_data["S3DataSource"]["S3Uri"], "s3://bucket/model/")
 
+    @patch("sagemaker.serve.model_server.tgi.server._is_s3_uri")
+    def test_upload_tgi_artifacts_s3_path_trailing_slash_normalized(self, mock_is_s3):
+        """An S3 model_path that already ends in "/" must not yield a doubled "//".
+
+        With S3DataType="S3Prefix", SageMaker does literal prefix matching, so a
+        doubled "s3://bucket/model//" would fail to match the actual object keys
+        under "s3://bucket/model/" and nothing would mount into /opt/ml/model.
+        Regression test for the trailing-slash deploy-breaker.
+        """
+        from sagemaker.serve.model_server.tgi.server import SageMakerTgiServing
+
+        server = SageMakerTgiServing()
+        mock_is_s3.return_value = True
+        mock_session = Mock()
+
+        model_data, _ = server._upload_tgi_artifacts(
+            model_path="s3://bucket/model/",
+            sagemaker_session=mock_session,
+            jumpstart=False,
+            should_upload_artifacts=False,
+        )
+
+        self.assertIsNotNone(model_data)
+        # Exactly one trailing slash, regardless of the input's trailing slash.
+        self.assertEqual(model_data["S3DataSource"]["S3Uri"], "s3://bucket/model/")
+        self.assertNotIn("//", model_data["S3DataSource"]["S3Uri"].split("://", 1)[1])
+
     @patch("sagemaker.serve.model_server.tgi.server._is_s3_uri")
     def test_upload_tgi_artifacts_jumpstart(self, mock_is_s3):
         """Test _upload_tgi_artifacts with jumpstart=True."""