diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py index d4f2e0cc37..d18ef7585e 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py @@ -579,7 +579,7 @@ def _build_for_djl(self) -> Model: model = self._create_model() if "HF_HUB_OFFLINE" in self.env_vars: - self.env_vars.update({"TRANSFORMERS_OFFLINE": "0"}) + self.env_vars.update({"HF_HUB_OFFLINE": "0"}) return model diff --git a/sagemaker-serve/tests/unit/servers/test_djl_hf_cache_env.py b/sagemaker-serve/tests/unit/servers/test_djl_hf_cache_env.py index 40d3a2a1a9..0a6b842d22 100644 --- a/sagemaker-serve/tests/unit/servers/test_djl_hf_cache_env.py +++ b/sagemaker-serve/tests/unit/servers/test_djl_hf_cache_env.py @@ -143,6 +143,13 @@ def test_local_mode_sets_hf_hub_offline(self, tmp_path): builder = _create_djl_builder(tmp_path, mode=Mode.LOCAL_CONTAINER) # Local mode doesn't need GPU info mocks for instance_type validation builder.instance_type = None + envs_at_create = [] + mock_create = self._mocks[ + _DJL_PATCHES.index("sagemaker.serve.model_builder.ModelBuilder._create_model") + ] + mock_create.side_effect = lambda: envs_at_create.append(dict(builder.env_vars)) builder._build_for_djl() - assert builder.env_vars["HF_HUB_OFFLINE"] == "1" + # The model (and so the local container) is created with HF_HUB_OFFLINE=1; the + # in-memory flag is reset afterwards so it does not leak into a later build. + assert envs_at_create[-1]["HF_HUB_OFFLINE"] == "1" diff --git a/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py b/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py index e4509c6af8..dea27e5c52 100644 --- a/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py +++ b/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py @@ -992,6 +992,52 @@ def test_build_sagemaker_endpoint_tensor_parallel( self.assertEqual(self.builder.env_vars["TENSOR_PARALLEL_DEGREE"], "4") mock_create.assert_called_once() + @patch("sagemaker.serve.model_builder_servers._get_gpu_info") + @patch("sagemaker.serve.model_builder_servers._get_default_tensor_parallel_degree") + @patch("sagemaker.serve.model_builder_servers._get_model_config_properties_from_hf") + @patch("sagemaker.serve.model_builder_servers._get_default_djl_configurations") + @patch("sagemaker.serve.model_builder_servers._get_nb_instance") + @patch("sagemaker.serve.model_server.djl_serving.prepare._create_dir_structure") + @patch.object(MockModelBuilderServers, "_validate_djl_serving_sample_data") + @patch.object(MockModelBuilderServers, "_is_jumpstart_model_id") + @patch.object(MockModelBuilderServers, "_auto_detect_image_uri") + @patch.object(MockModelBuilderServers, "_prepare_for_mode") + @patch.object(MockModelBuilderServers, "_create_model") + def test_local_build_does_not_leave_hf_hub_offline_for_endpoint_build( + self, + mock_create, + mock_prepare, + mock_detect, + mock_js, + mock_validate, + mock_dir, + mock_nb, + mock_djl_config, + mock_hf_config, + mock_tp, + mock_gpu, + ): + """HF_HUB_OFFLINE=1 set for a local build must be reset, as TGI/TEI do.""" + mock_js.return_value = False + mock_nb.return_value = None + mock_hf_config.return_value = {"model_type": "gpt2"} + mock_djl_config.return_value = ({}, 512) + mock_prepare.return_value = ("s3://bucket/model.tar.gz", None) + mock_tp.return_value = 1 + mock_gpu.return_value = 1 + envs_at_create = [] + mock_create.side_effect = lambda: envs_at_create.append(dict(self.builder.env_vars)) + self.builder.model = "gpt2" + self.builder.instance_type = "ml.g5.2xlarge" + + self.builder.mode = Mode.LOCAL_CONTAINER + self.builder._build_for_djl() + self.builder.mode = Mode.SAGEMAKER_ENDPOINT + self.builder._build_for_djl() + + self.assertEqual(envs_at_create[0]["HF_HUB_OFFLINE"], "1") + self.assertNotEqual(envs_at_create[1].get("HF_HUB_OFFLINE"), "1") + # ------------------------------------------------------------------ # Preservation property test (tgi-s3-model-loading bugfix spec) #