Merge branch 'master' into fix/bug-modelbuilder-overwrites-user-provided-hf-5529

Author: aviruthen

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

@@ -326,36 +326,36 @@ def _build_for_djl(self) -> Model:
             self.hf_model_config = _get_model_config_properties_from_hf(
                 self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
             )
-            
+
             # Apply DJL-specific configurations
             default_djl_configurations, _default_max_new_tokens = _get_default_djl_configurations(
                 self.model, self.hf_model_config, self.schema_builder
             )
             self.env_vars.update(default_djl_configurations)
-            
+
             # Configure schema builder for text generation
             if "parameters" not in self.schema_builder.sample_input:
                 self.schema_builder.sample_input["parameters"] = {}
             self.schema_builder.sample_input["parameters"]["max_new_tokens"] = _default_max_new_tokens
-            
-            # Set DJL serving defaults
+
+            # Set DJL serving defaults (only if not already set by user)
             djl_env_vars = {
                 "OPTION_ENGINE": "Python",
                 "SERVING_MIN_WORKERS": "1",
-                "SERVING_MAX_WORKERS": "1", 
+                "SERVING_MAX_WORKERS": "1",
                 "OPTION_MODEL_LOADING_TIMEOUT": "240",
                 "OPTION_PREDICT_TIMEOUT": "60",
-                "TENSOR_PARALLEL_DEGREE": "1"  # Default, will be overridden below
+                "TENSOR_PARALLEL_DEGREE": "1",
             }
-            
+
             # Add HuggingFace authentication
             if self.env_vars.get("HUGGING_FACE_HUB_TOKEN"):
                 djl_env_vars["HF_TOKEN"] = self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
-            
+
             # Update with defaults only if not already set
             for key, value in djl_env_vars.items():
                 self.env_vars.setdefault(key, value)
-            
+
             # DJL downloads models directly from HuggingFace Hub
             self.s3_upload_path = None
 
@@ -367,6 +367,12 @@ def _build_for_djl(self) -> Model:
             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"})

sagemaker-serve/tests/unit/servers/__init__.py

@@ -0,0 +1,151 @@
+"""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
+- 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
+"""
+
+import pytest
+from unittest.mock import Mock, patch
+
+from sagemaker.serve.model_builder import ModelBuilder
+from sagemaker.serve.utils.types import ModelServer
+from sagemaker.serve.mode.function_pointers import Mode
+from sagemaker.core.resources import Model
+
+
+MOCK_ROLE_ARN = "arn:aws:iam::000000000000:role/SageMakerRole"
+MOCK_IMAGE_URI = "000000000000.dkr.ecr.us-east-1.amazonaws.com/djl-inference:latest"
+MOCK_HF_MODEL_CONFIG = {"model_type": "gpt2", "architectures": ["GPT2LMHeadModel"]}
+
+
+# Common patches needed for _build_for_djl
+_DJL_PATCHES = [
+    "sagemaker.serve.model_builder_servers._get_nb_instance",
+    "sagemaker.serve.model_builder_servers._get_default_djl_configurations",
+    "sagemaker.serve.model_builder_servers._get_model_config_properties_from_hf",
+    "sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id",
+    "sagemaker.serve.model_builder.ModelBuilder._validate_djl_serving_sample_data",
+    "sagemaker.serve.model_builder.ModelBuilder._auto_detect_image_uri",
+    "sagemaker.serve.model_builder.ModelBuilder._prepare_for_mode",
+    "sagemaker.serve.model_builder.ModelBuilder._create_model",
+    "sagemaker.serve.model_builder_servers._get_default_tensor_parallel_degree",
+    "sagemaker.serve.model_builder_servers._get_gpu_info",
+]
+
+
+def _mock_sagemaker_session():
+    """Create a mock SageMaker session."""
+    session = Mock()
+    session.boto_region_name = "us-east-1"
+    session.sagemaker_config = {}
+    session.default_bucket.return_value = "mock-bucket"
+    session.upload_data.return_value = "s3://mock-bucket/model.tar.gz"
+    return session
+
+
+def _create_djl_builder(tmp_path, env_vars=None, mode=Mode.SAGEMAKER_ENDPOINT):
+    """Create a ModelBuilder configured for DJL serving tests."""
+    builder = ModelBuilder(
+        model="test-org/test-model",
+        role_arn=MOCK_ROLE_ARN,
+        sagemaker_session=_mock_sagemaker_session(),
+        model_path=str(tmp_path),
+        mode=mode,
+        image_uri=MOCK_IMAGE_URI,
+        model_server=ModelServer.DJL_SERVING,
+        instance_type="ml.g6e.12xlarge",
+        env_vars=env_vars or {},
+    )
+    builder.schema_builder = Mock()
+    builder.schema_builder.sample_input = {"inputs": "Hello"}
+    builder._optimizing = False
+    builder.hf_model_config = MOCK_HF_MODEL_CONFIG
+    return builder
+
+
+def _setup_mocks(mocks):
+    """Configure common mock return values for DJL build."""
+    # mocks are in reverse order of _DJL_PATCHES
+    mock_gpu_info = mocks[-1]
+    mock_tp_degree = mocks[-2]
+    mock_create = mocks[-3]
+    mock_prepare = mocks[-4]
+    # mock_auto_detect = mocks[-5]  # no setup needed
+    # mock_validate = mocks[-6]  # no setup needed
+    mock_is_js = mocks[-7]
+    mock_hf_config = mocks[-8]
+    mock_djl_config = mocks[-9]
+    mock_nb = mocks[-10]
+
+    mock_nb.return_value = None
+    mock_djl_config.return_value = ({}, 256)
+    mock_hf_config.return_value = MOCK_HF_MODEL_CONFIG
+    mock_is_js.return_value = False
+    mock_prepare.return_value = ("s3://bucket/model", None)
+    mock_create.return_value = Mock(spec=Model)
+    mock_tp_degree.return_value = 4
+    mock_gpu_info.return_value = 4
+
+
+class TestDjlHfCacheAndModelId:
+    """Tests for DJL builder HF cache env vars and HF_MODEL_ID handling."""
+
+    @pytest.fixture(autouse=True)
+    def _patch_djl(self):
+        """Apply all DJL-related patches for each test."""
+        patchers = [patch(p) for p in _DJL_PATCHES]
+        self._mocks = [p.start() for p in patchers]
+        _setup_mocks(self._mocks)
+        yield
+        for p in patchers:
+            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."""
+        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"
+
+    def test_preserves_user_provided_hf_model_id(self, tmp_path):
+        """User-provided HF_MODEL_ID must NOT be overridden by model param."""
+        builder = _create_djl_builder(
+            tmp_path, env_vars={"HF_MODEL_ID": "/opt/ml/model"}
+        )
+        builder._build_for_djl()
+
+        assert builder.env_vars["HF_MODEL_ID"] == "/opt/ml/model"
+
+    def test_sets_hf_model_id_from_model_param_when_not_provided(self, tmp_path):
+        """When no user-provided HF_MODEL_ID, it should come from model param."""
+        builder = _create_djl_builder(tmp_path)
+        builder._build_for_djl()
+
+        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."""
+        builder = _create_djl_builder(
+            tmp_path,
+            env_vars={
+                "HF_HOME": "/my/custom/cache",
+                "HUGGINGFACE_HUB_CACHE": "/my/custom/hub",
+            },
+        )
+        builder._build_for_djl()
+
+        assert builder.env_vars["HF_HOME"] == "/my/custom/cache"
+        assert builder.env_vars["HUGGINGFACE_HUB_CACHE"] == "/my/custom/hub"
+
+    def test_local_mode_sets_hf_hub_offline(self, tmp_path):
+        """HF_HUB_OFFLINE=1 should be set in LOCAL_CONTAINER mode."""
+        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
+        builder._build_for_djl()
+
+        assert builder.env_vars["HF_HUB_OFFLINE"] == "1"

sagemaker-serve/tests/unit/servers/test_djl_hf_cache_env.py

@@ -1504,7 +1504,13 @@ def _build_training_job_definition(self, inputs):
                     model_trainer.stopping_condition.max_wait_time_in_seconds
                 )
 
-        definition = HyperParameterTrainingJobDefinition(
+        # Propagate environment variables from ModelTrainer.
+        # Only include when it's a dict (even empty); omit otherwise so the
+        # Pydantic field stays Unassigned and is excluded during serialization.
+        env = model_trainer.environment
+
+        # Build base kwargs for the definition
+        definition_kwargs = dict(
             algorithm_specification=algorithm_spec,
             role_arn=model_trainer.role,
             input_data_config=input_data_config if input_data_config else None,
@@ -1515,10 +1521,11 @@ def _build_training_job_definition(self, inputs):
             enable_managed_spot_training=model_trainer.compute.enable_managed_spot_training,
         )
 
-        # Pass through environment variables from model_trainer
-        env = getattr(model_trainer, "environment", None)
-        if env and isinstance(env, dict):
-            definition.environment = env
+        # Include environment only when it's a dict (including empty).
+        if isinstance(env, dict):
+            definition_kwargs["environment"] = env
+
+        definition = HyperParameterTrainingJobDefinition(**definition_kwargs)
 
         # Pass through VPC config from model_trainer
         networking = getattr(model_trainer, "networking", None)

sagemaker-train/src/sagemaker/train/tuner.py

@@ -596,3 +596,73 @@ def test_build_training_job_definition_includes_spot_params(self):
         assert isinstance(
             definition.stopping_condition.max_wait_time_in_seconds, int
         ), "Max wait time should be set"
+
+    def test_build_training_job_definition_includes_environment_variables(self):
+        """Test that _build_training_job_definition includes environment variables.
+
+        This test verifies the fix for GitHub issue #5613 where tuning jobs were
+        missing environment variables that were set on the ModelTrainer.
+        """
+        mock_trainer = _create_mock_model_trainer()
+        mock_trainer.environment = {
+            "FOO": "bar",
+            "RANDOM_STATE": "42",
+        }
+
+        tuner = HyperparameterTuner(
+            model_trainer=mock_trainer,
+            objective_metric_name="accuracy",
+            hyperparameter_ranges=_create_single_hp_range(),
+        )
+
+        definition = tuner._build_training_job_definition(None)
+
+        assert definition.environment is not None, "Environment should not be None"
+        assert definition.environment == {
+            "FOO": "bar",
+            "RANDOM_STATE": "42",
+        }, "Environment variables should match those set on ModelTrainer"
+
+    def test_build_training_job_definition_with_none_environment(self):
+        """Test that _build_training_job_definition handles None environment gracefully.
+
+        When environment is None, it should not be passed to the Pydantic constructor,
+        so the field stays as Unassigned (excluded from serialization).
+        """
+        from sagemaker.core.utils.utils import Unassigned
+
+        mock_trainer = _create_mock_model_trainer()
+        mock_trainer.environment = None
+
+        tuner = HyperparameterTuner(
+            model_trainer=mock_trainer,
+            objective_metric_name="accuracy",
+            hyperparameter_ranges=_create_single_hp_range(),
+        )
+
+        definition = tuner._build_training_job_definition(None)
+
+        assert isinstance(definition.environment, Unassigned), (
+            "Environment should be Unassigned when model_trainer.environment is None"
+        )
+
+    def test_build_training_job_definition_with_empty_environment(self):
+        """Test that _build_training_job_definition passes through empty environment.
+
+        An empty dict is valid for the SageMaker API, so we pass it through as-is
+        rather than silently converting it to None.
+        """
+        mock_trainer = _create_mock_model_trainer()
+        mock_trainer.environment = {}
+
+        tuner = HyperparameterTuner(
+            model_trainer=mock_trainer,
+            objective_metric_name="accuracy",
+            hyperparameter_ranges=_create_single_hp_range(),
+        )
+
+        definition = tuner._build_training_job_definition(None)
+
+        assert definition.environment == {}, (
+            "Empty dict environment should be passed through as-is"
+        )

sagemaker-train/tests/unit/train/test_tuner.py

@@ -405,8 +405,31 @@ def test_passes_environment_variables(self):
         definition = tuner._build_training_job_definition(inputs=None)
         assert definition.environment == {"MY_VAR": "value", "OTHER": "123"}
 
+    def test_passes_empty_environment(self):
+        """Should pass through empty dict environment as-is.
+
+        An empty dict is valid for the SageMaker API, so we pass it through
+        rather than silently converting it to None/Unassigned.
+        """
+        trainer = _mock_model_trainer(environment={})
+
+        tuner = HyperparameterTuner(
+            model_trainer=trainer,
+            objective_metric_name="accuracy",
+            hyperparameter_ranges=_hp_ranges(),
+        )
+
+        definition = tuner._build_training_job_definition(inputs=None)
+        assert definition.environment == {}, (
+            "Empty dict environment should be passed through as-is"
+        )
+
     def test_skips_environment_when_none(self):
-        """Should not set environment when model_trainer.environment is None."""
+        """Should not set environment when model_trainer.environment is None.
+
+        When environment is None, it is not passed to the Pydantic constructor,
+        so the field stays as Unassigned (excluded from serialization).
+        """
         trainer = _mock_model_trainer(environment=None)
 
         tuner = HyperparameterTuner(
@@ -416,10 +439,16 @@ def test_skips_environment_when_none(self):
         )
 
         definition = tuner._build_training_job_definition(inputs=None)
-        assert _is_unassigned(definition.environment)
+        assert _is_unassigned(definition.environment), (
+            "Environment should be Unassigned when model_trainer.environment is None"
+        )
 
     def test_skips_environment_when_not_dict(self):
-        """Should not set environment when it's not a dict (e.g. MagicMock)."""
+        """Should not set environment when it's not a dict (e.g. MagicMock).
+
+        Non-dict values are not passed to the Pydantic constructor to avoid
+        validation errors. The field stays as Unassigned.
+        """
         trainer = _mock_model_trainer(environment=MagicMock())
 
         tuner = HyperparameterTuner(
@@ -429,7 +458,9 @@ def test_skips_environment_when_not_dict(self):
         )
 
         definition = tuner._build_training_job_definition(inputs=None)
-        assert _is_unassigned(definition.environment)
+        assert _is_unassigned(definition.environment), (
+            "Environment should be Unassigned when model_trainer.environment is not a dict"
+        )
 
     def test_passes_vpc_config(self):
         """Should set definition.vpc_config from model_trainer.networking._to_vpc_config()."""