allow SAGEMAKER_HUB_NAME env var override for HUB_NAME constant

Author: mollyheamazon

sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py

@@ -15,6 +15,7 @@
 from sagemaker.core.shapes import ServerlessJobConfig, Channel, DataSource, ModelPackageConfig, MlflowConfig
 from sagemaker.train.configs import InputData, OutputDataConfig
 from sagemaker.train.defaults import TrainDefaults
+from sagemaker.train.constants import HUB_NAME
 
 logger = logging.getLogger(__name__)
 
@@ -317,7 +318,7 @@ def _resolve_model_package_arn(model_package) -> Optional[str]:
 
 
 def _get_fine_tuning_options_and_model_arn(model_name: str, customization_technique: str, training_type, sagemaker_session,
-                                         hub_name: str = "SageMakerPublicHub") -> tuple:
+                                         hub_name: str = HUB_NAME) -> tuple:
     """Get fine-tuning options and model ARN for given customization technique.
     
     Returns:

sagemaker-train/src/sagemaker/train/common_utils/model_resolution.py

@@ -14,6 +14,7 @@
 from enum import Enum
 import re
 from sagemaker.train.base_trainer import BaseTrainer
+from sagemaker.train.constants import HUB_NAME
 from sagemaker.core.utils.utils import Unassigned
 
 
@@ -52,7 +53,7 @@ class _ModelResolver:
     and fine-tuned ModelPackage objects/ARNs.
     """
 
-    DEFAULT_HUB_NAME = "SageMakerPublicHub"
+    DEFAULT_HUB_NAME = HUB_NAME
 
     def __init__(self, sagemaker_session=None):
         """

sagemaker-train/src/sagemaker/train/constants.py

@@ -40,7 +40,7 @@
     + f"&& {SM_DRIVERS_CONTAINER_PATH}/{TRAIN_SCRIPT}",
 ]
 
-HUB_NAME = "SageMakerPublicHub"
+HUB_NAME = os.environ.get("SAGEMAKER_HUB_NAME", "SageMakerPublicHub")
 
 # Allowed reward model IDs for RLAIF trainer with region restrictions
 _ALLOWED_REWARD_MODEL_IDS = {

sagemaker-train/src/sagemaker/train/evaluate/benchmark_evaluator.py

@@ -21,6 +21,7 @@
 from .execution import EvaluationPipelineExecution
 from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter
 from sagemaker.core.telemetry.constants import Feature
+from sagemaker.train.constants import HUB_NAME
 
 _logger = logging.getLogger(__name__)
 
@@ -466,7 +467,7 @@ def hyperparameters(self):
 
             override_params = _get_evaluation_override_params(
                 hub_content_name=hub_content_name,
-                hub_name="SageMakerPublicHub",
+                hub_name=HUB_NAME,
                 evaluation_type=evaluation_type,
                 region=region,
                 session=boto_session

sagemaker-train/src/sagemaker/train/evaluate/custom_scorer_evaluator.py

@@ -16,6 +16,7 @@
 from .execution import EvaluationPipelineExecution
 from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter
 from sagemaker.core.telemetry.constants import Feature
+from sagemaker.train.constants import HUB_NAME
 
 _logger = logging.getLogger(__name__)
 
@@ -240,7 +241,7 @@ def hyperparameters(self):
 
             override_params = _get_evaluation_override_params(
                 hub_content_name=hub_content_name,
-                hub_name="SageMakerPublicHub",
+                hub_name=HUB_NAME,
                 evaluation_type="DeterministicEvaluation",
                 region=region,
                 session=boto_session
@@ -365,7 +366,7 @@ def _get_inference_params_from_hub(self, region: str) -> dict:
             _logger.info(f"Fetching evaluation recipe override parameters from hub for model: {hub_content_name}")
             override_params = _get_evaluation_override_params(
                 hub_content_name=hub_content_name,
-                hub_name="SageMakerPublicHub",
+                hub_name=HUB_NAME,
                 evaluation_type="DeterministicEvaluation",
                 region=region,
                 session=session

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

@@ -0,0 +1,26 @@
+"""Tests for SAGEMAKER_HUB_NAME env-var override of the HUB_NAME constant."""
+from __future__ import absolute_import
+
+import importlib
+import os
+from unittest.mock import patch
+
+
+def _reload_hub_name():
+    """Reload the constants module under the current env and return HUB_NAME."""
+    from sagemaker.train import constants
+    importlib.reload(constants)
+    return constants.HUB_NAME
+
+
+def test_hub_name_defaults_to_public_hub():
+    """When SAGEMAKER_HUB_NAME is unset, HUB_NAME is SageMakerPublicHub."""
+    env = {k: v for k, v in os.environ.items() if k != "SAGEMAKER_HUB_NAME"}
+    with patch.dict(os.environ, env, clear=True):
+        assert _reload_hub_name() == "SageMakerPublicHub"
+
+
+def test_hub_name_overridden_by_env_var():
+    """When SAGEMAKER_HUB_NAME is set, HUB_NAME reflects the override."""
+    with patch.dict(os.environ, {"SAGEMAKER_HUB_NAME": "MyPrivateHub"}):
+        assert _reload_hub_name() == "MyPrivateHub"