diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py index d4f2e0cc37..e1e9692fa4 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py @@ -543,18 +543,20 @@ 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. + self.env_vars.setdefault("HF_HOME", "/tmp") + self.env_vars.setdefault("HUGGINGFACE_HUB_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..910a0f199d 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 @@ -146,3 +146,59 @@ 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 HF_HOME=/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"