diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py index d4f2e0cc37..8d2162508c 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py @@ -543,18 +543,24 @@ def _build_for_djl(self) -> Model: self._auto_detect_image_uri() + # Set HF cache env vars to writable location (unconditionally, using setdefault + # to preserve user-provided values). This is needed because /opt/ml/model/ may be + # read-only when source_code artifacts are mounted there. This must happen before + # _prepare_for_mode, which merges the DJL server defaults (HF_HOME=/opt/ml/model/) + # into env_vars with setdefault. + # TRANSFORMERS_CACHE is included because transformers uses it as the default + # explicit cache_dir for from_pretrained, which overrides HF_HUB_CACHE: leaving it + # at the DJL default of /opt/ml/model/ keeps the read-only writes going. + self.env_vars.setdefault("HF_HOME", "/tmp") + self.env_vars.setdefault("HUGGINGFACE_HUB_CACHE", "/tmp") + self.env_vars.setdefault("TRANSFORMERS_CACHE", "/tmp") + if not self._optimizing: if self.mode in LOCAL_MODES: self._prepare_for_mode(should_upload_artifacts=True) else: self.s3_model_data_url, _ = self._prepare_for_mode() - # Set HF cache env vars to writable location (unconditionally, using setdefault - # to preserve user-provided values). This is needed because /opt/ml/model/ may be - # read-only when source_code artifacts are mounted there. - self.env_vars.setdefault("HF_HOME", "/tmp") - self.env_vars.setdefault("HUGGINGFACE_HUB_CACHE", "/tmp") - # Cache management based on mode if self.mode in LOCAL_MODES: self.env_vars.update({"HF_HUB_OFFLINE": "1"}) 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..e9d65ea147 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 @@ -1,10 +1,10 @@ """Tests for DJL builder HF cache environment variables and HF_MODEL_ID handling. Verifies that _build_for_djl() correctly: -- Sets HF_HOME and HUGGINGFACE_HUB_CACHE to /tmp for writable cache +- Sets HF_HOME, HUGGINGFACE_HUB_CACHE and TRANSFORMERS_CACHE to /tmp for writable cache - Preserves user-provided HF_MODEL_ID values (uses setdefault) - Sets HF_MODEL_ID from model param when not provided by user -- Preserves user-provided HF_HOME and HUGGINGFACE_HUB_CACHE values +- Preserves user-provided HF cache directory values """ import pytest @@ -103,12 +103,13 @@ def _patch_djl(self): p.stop() def test_sets_hf_cache_env_vars_to_tmp(self, tmp_path): - """HF_HOME and HUGGINGFACE_HUB_CACHE should be /tmp in endpoint mode.""" + """The HF cache env vars should all be /tmp in endpoint mode.""" builder = _create_djl_builder(tmp_path) builder._build_for_djl() assert builder.env_vars["HF_HOME"] == "/tmp" assert builder.env_vars["HUGGINGFACE_HUB_CACHE"] == "/tmp" + assert builder.env_vars["TRANSFORMERS_CACHE"] == "/tmp" def test_preserves_user_provided_hf_model_id(self, tmp_path): """User-provided HF_MODEL_ID must NOT be overridden by model param.""" @@ -125,18 +126,20 @@ def test_sets_hf_model_id_from_model_param_when_not_provided(self, tmp_path): assert builder.env_vars["HF_MODEL_ID"] == "test-org/test-model" def test_preserves_user_provided_hf_cache_dirs(self, tmp_path): - """User-provided HF_HOME and HUGGINGFACE_HUB_CACHE should be preserved.""" + """User-provided HF cache dirs should be preserved.""" builder = _create_djl_builder( tmp_path, env_vars={ "HF_HOME": "/my/custom/cache", "HUGGINGFACE_HUB_CACHE": "/my/custom/hub", + "TRANSFORMERS_CACHE": "/my/custom/transformers", }, ) builder._build_for_djl() assert builder.env_vars["HF_HOME"] == "/my/custom/cache" assert builder.env_vars["HUGGINGFACE_HUB_CACHE"] == "/my/custom/hub" + assert builder.env_vars["TRANSFORMERS_CACHE"] == "/my/custom/transformers" def test_local_mode_sets_hf_hub_offline(self, tmp_path): """HF_HUB_OFFLINE=1 should be set in LOCAL_CONTAINER mode.""" @@ -146,3 +149,60 @@ def test_local_mode_sets_hf_hub_offline(self, tmp_path): builder._build_for_djl() assert builder.env_vars["HF_HUB_OFFLINE"] == "1" + + +def test_endpoint_build_hf_cache_is_tmp_with_real_prepare_for_mode(tmp_path): + """The DJL server defaults merged by _prepare_for_mode must not win over /tmp. + + The tests above mock _prepare_for_mode; here it runs for real (only the S3 upload is + mocked), since it merges the DJL server's /opt/ml/model/ defaults into env_vars. + """ + builder = ModelBuilder( + model="test-org/test-model", + role_arn=MOCK_ROLE_ARN, + sagemaker_session=_mock_sagemaker_session(), + model_path=str(tmp_path), + mode=Mode.SAGEMAKER_ENDPOINT, + image_uri=MOCK_IMAGE_URI, + model_server=ModelServer.DJL_SERVING, + instance_type="ml.g5.2xlarge", + ) + builder.schema_builder = Mock( + sample_input={"inputs": "Hello", "parameters": {}}, + sample_output=[{"generated_text": "Hi"}], + ) + builder._optimizing = False + # Normally initialized by build() + builder.modes = {} + builder.serve_settings = Mock(s3_model_data_url=None) + envs_at_create = [] + + with ( + patch("sagemaker.serve.model_builder_servers._get_nb_instance", return_value=None), + patch( + "sagemaker.serve.model_builder_servers._get_model_config_properties_from_hf", + return_value={"num_attention_heads": 32}, + ), + patch("sagemaker.serve.model_builder_servers._get_gpu_info", return_value=1), + patch.object(ModelBuilder, "_is_jumpstart_model_id", return_value=False), + patch.object(ModelBuilder, "_validate_djl_serving_sample_data"), + patch.object(ModelBuilder, "_auto_detect_image_uri"), + patch.object( + ModelBuilder, + "_create_model", + side_effect=lambda: envs_at_create.append(dict(builder.env_vars)), + ), + patch( + "sagemaker.serve.model_server.djl_serving.server.determine_bucket_and_prefix", + return_value=("mock-bucket", "prefix"), + ), + patch( + "sagemaker.serve.model_server.djl_serving.server.S3Uploader.upload", + return_value="s3://mock-bucket/prefix/code", + ), + ): + builder._build_for_djl() + + assert envs_at_create[-1]["HF_HOME"] == "/tmp" + assert envs_at_create[-1]["HUGGINGFACE_HUB_CACHE"] == "/tmp" + assert envs_at_create[-1]["TRANSFORMERS_CACHE"] == "/tmp"