sagemaker-core 2.24.0__tar.gz → 2.26.0__tar.gz

{sagemaker_core-2.24.0/src/sagemaker_core.egg-info → sagemaker_core-2.26.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sagemaker-core
-Version: 2.24.0
+Version: 2.26.0
 Summary: An python package for sagemaker core functionalities
 Author-email: AWS <sagemaker-interests@amazon.com>
 Project-URL: Homepage, https://github.com/aws/sagemaker-python-sdk

sagemaker_core-2.26.0/VERSION

@@ -0,0 +1 @@
+2.26.0

{sagemaker_core-2.24.0 → sagemaker_core-2.26.0}/src/sagemaker/core/config_schema.py

@@ -501,7 +501,10 @@ SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA = {
                                 },
                                 "MlflowApp": {
                                     "type": "object",
-                                    "properties": {"role_arn": {"type": "string"}},
+                                    "properties": {
+                                        "role_arn": {"type": "string"},
+                                        "kms_key_id": {"type": "string"},
+                                    },
                                 },
                                 "MlflowTrackingServer": {
                                     "type": "object",
@@ -873,6 +876,27 @@ SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA = {
                                         "profiler_config": {"s3_output_path": {"type": "string"}},
                                     },
                                 },
+                                "TrainingSession": {
+                                    "type": "object",
+                                    "properties": {
+                                        "training_session_detail": {
+                                            "vpc_config": {
+                                                "security_group_ids": {
+                                                    "type": "array",
+                                                    "items": {"type": "string"},
+                                                },
+                                                "subnets": {
+                                                    "type": "array",
+                                                    "items": {"type": "string"},
+                                                },
+                                            },
+                                            "role_arn": {"type": "string"},
+                                            "output_config": {
+                                                "s3_location": {"kms_key_arn": {"type": "string"}}
+                                            },
+                                        }
+                                    },
+                                },
                                 "TransformJob": {
                                     "type": "object",
                                     "properties": {

{sagemaker_core-2.24.0 → sagemaker_core-2.26.0}/src/sagemaker/core/helper/iam_role_resolver.py

@@ -16,6 +16,7 @@ from __future__ import absolute_import
 
 import json
 import logging
+import os
 import time
 from typing import List, Optional, Set, Tuple, Union
 from urllib.parse import unquote
@@ -369,13 +370,47 @@ def _resolve_config_default_role(role_type: str, sagemaker_session=None) -> Opti
 # ---------------------------------------------------------------------------
 # Read-only permission / trust validation
 # ---------------------------------------------------------------------------
-def _simulate_denied_actions(iam_client, role_arn: str, actions: List[str]) -> List[str]:
-    """Return the subset of ``actions`` that ``role_arn`` is NOT allowed to perform.
+def _is_unverifiable_denial(result: dict) -> bool:
+    """Return True if an evaluation result's non-allowed decision is unverifiable.
+
+    Because iam:SimulatePrincipalPolicy does not evaluate condition-based SCPs or
+    condition-based permissions boundaries, an `implicitDeny` caused by (or
+    masked by) `AllowedByOrganizations == false` or
+    `AllowedByPermissionsBoundary == false` cannot be definitively verified by
+    the client-side policy simulator.
+    """
+    if result.get("EvalDecision") == "explicitDeny":
+        return False
+    org_detail = result.get("OrganizationsDecisionDetail", {})
+    org_allowed = org_detail.get("AllowedByOrganizations") if isinstance(org_detail, dict) else None
+    pb_detail = result.get("PermissionsBoundaryDecisionDetail", {})
+    pb_allowed = (
+        pb_detail.get("AllowedByPermissionsBoundary") if isinstance(pb_detail, dict) else None
+    )
+
+    # Handle boolean False, string "false", "False", etc.
+    if org_allowed in (False, "false", "False") or pb_allowed in (False, "false", "False"):
+        return True
+    return False
+
+
+def _simulate_denied_actions(
+    iam_client, role_arn: str, actions: List[str]
+) -> Tuple[List[str], List[str]]:
+    """Return lists of (denied_actions, unverifiable_actions) for ``role_arn``.
 
     Wraps the paginated ``iam:SimulatePrincipalPolicy`` call so both the role-
     validation path and the HyperPod caller-side check share one implementation.
-    An empty list means every action is allowed. The pagination matters: a
-    truncated first page must not produce a false "all allowed" verdict.
+    The pagination matters: a truncated first page must not produce a false
+    "all allowed" verdict.
+
+    Returns:
+        (denied_actions, unverifiable_actions):
+            denied_actions: actions definitively denied (explicitDeny, or
+                implicitDeny without org/permissions-boundary conditions).
+            unverifiable_actions: actions returning implicitDeny with
+                AllowedByOrganizations=False or AllowedByPermissionsBoundary=False,
+                which the simulator cannot evaluate due to conditions.
 
     Raises ClientError on failures the caller must interpret (e.g. AccessDenied
     when the principal can't self-simulate, NoSuchEntity when the role is gone).
@@ -385,11 +420,16 @@ def _simulate_denied_actions(iam_client, role_arn: str, actions: List[str]) -> L
     for page in paginator.paginate(PolicySourceArn=role_arn, ActionNames=actions):
         evaluation_results.extend(page.get("EvaluationResults", []))
 
-    return [
-        result["EvalActionName"]
-        for result in evaluation_results
-        if result["EvalDecision"] != "allowed"
-    ]
+    denied = []
+    unverifiable = []
+    for result in evaluation_results:
+        if result.get("EvalDecision") != "allowed":
+            action_name = result["EvalActionName"]
+            if _is_unverifiable_denial(result):
+                unverifiable.append(action_name)
+            else:
+                denied.append(action_name)
+    return denied, unverifiable
 
 
 def _evaluate_permissions(
@@ -405,14 +445,15 @@ def _evaluate_permissions(
         verdict True  — all gated actions are allowed (denied_actions empty).
         verdict False — at least one gated action is denied (denied_actions lists them).
         verdict None  — could not be determined, e.g. the caller lacks
-                        iam:SimulatePrincipalPolicy (denied_actions empty).
+                        iam:SimulatePrincipalPolicy or conditional SCPs prevent
+                        evaluation (denied_actions empty).
     """
     required_actions = _get_smoke_test_actions(role_type)
     if not required_actions:
         return True, []
 
     try:
-        denied = _simulate_denied_actions(iam_client, role_arn, required_actions)
+        denied, unverifiable = _simulate_denied_actions(iam_client, role_arn, required_actions)
         if denied:
             logger.info(
                 "Role '%s' is missing permissions for: %s",
@@ -420,6 +461,16 @@ def _evaluate_permissions(
                 ", ".join(denied[:5]) + ("..." if len(denied) > 5 else ""),
             )
             return False, denied
+        if unverifiable:
+            logger.info(
+                "Cannot definitively verify permissions for role '%s' due to "
+                "Organizations SCPs or permissions boundaries (%s returned implicitDeny "
+                "with AllowedByOrganizations/AllowedByPermissionsBoundary=false); "
+                "permission verdict unknown.",
+                role_arn,
+                ", ".join(unverifiable[:5]) + ("..." if len(unverifiable) > 5 else ""),
+            )
+            return None, []
         return True, []
 
     except ClientError as e:
@@ -580,9 +631,12 @@ def _build_validation_error_message(
 
 
 def resolve_and_validate_role(
-    provided_role: Optional[str],
-    role_type: str,
+    provided_role: Optional[str] = None,
+    role_type: str = "training",
     sagemaker_session=None,
+    validate_role: bool = True,
+    *,
+    role_arn: Optional[str] = None,
 ) -> str:
     """Resolve the role to use and validate it (read-only; does not mutate IAM).
 
@@ -593,17 +647,21 @@ def resolve_and_validate_role(
         3. Otherwise → resolve the caller's own identity role.
 
     The resolved role is then VALIDATED (read-only, via iam:SimulatePrincipalPolicy
-    + trust inspection):
+    + trust inspection) unless ``validate_role=False`` or environment variable
+    ``SAGEMAKER_VALIDATE_ROLE=false`` is set:
         * permissions allowed AND trusted → return the ARN.
         * a required permission is definitively denied → raise RoleValidationError.
         * the trust policy definitively excludes the service → raise RoleValidationError.
         * permissions cannot be verified (caller lacks iam:SimulatePrincipalPolicy,
-          the common Studio/notebook case) → return the ARN with a WARNING.
+          or conditional SCPs/permissions boundaries prevent evaluation) → return
+          the ARN with a WARNING.
 
     Args:
         provided_role: User-supplied role name or ARN. If set, used directly.
         role_type: One of ROLE_TYPES.
         sagemaker_session: SageMaker session (used to get the boto session).
+        validate_role: If False, skips client-side permission/trust validation.
+        role_arn: Keyword alias for ``provided_role``.
 
     Returns:
         IAM role ARN.
@@ -616,6 +674,7 @@ def resolve_and_validate_role(
     if role_type not in ROLE_TYPES:
         raise ValueError(f"Invalid role_type '{role_type}'. Must be one of: {ROLE_TYPES}")
 
+    provided_role = provided_role or role_arn
     boto_session = _get_boto_session(sagemaker_session)
     iam_client = boto_session.client("iam")
 
@@ -640,6 +699,14 @@ def resolve_and_validate_role(
             if not role_arn:
                 raise RoleValidationError(_build_validation_error_message(None, role_type))
 
+    if not validate_role or os.getenv("SAGEMAKER_VALIDATE_ROLE", "true").lower() in (
+        "false",
+        "0",
+        "no",
+    ):
+        logger.info("Skipping IAM role validation for '%s' (%s).", role_arn, role_type)
+        return role_arn
+
     # Permission check (definitive denial blocks; unverifiable warns).
     verdict, denied = _evaluate_permissions(iam_client, role_arn, role_type)
     if verdict is False:
@@ -657,7 +724,8 @@ def resolve_and_validate_role(
     if verdict is None:
         logger.warning(
             "Could not verify permissions for role '%s' (caller lacks "
-            "iam:SimulatePrincipalPolicy). Proceeding with it. If the operation "
+            "iam:SimulatePrincipalPolicy or conditional SCPs/permissions boundaries "
+            "prevent evaluation). Proceeding with it. If the operation "
             "later fails with an access-denied error, ensure the role has the "
             "required permissions for '%s' (see "
             "IamRoleResolver().get_required_actions('%s')) or create a dedicated "
@@ -716,7 +784,7 @@ def verify_hyperpod_connect_permissions(
         return None
 
     try:
-        denied = _simulate_denied_actions(
+        denied, unverifiable = _simulate_denied_actions(
             iam_client, caller_role_arn, list(HYPERPOD_CLI_CONNECT_ACTIONS)
         )
     except ClientError as e:
@@ -744,6 +812,15 @@ def verify_hyperpod_connect_permissions(
         )
         return False
 
+    if unverifiable:
+        logger.info(
+            "Cannot definitively verify HyperPod connect permissions for '%s' due to "
+            "Organizations SCPs or permissions boundaries; the HyperPod CLI will "
+            "validate access at submit time.",
+            caller_role_arn,
+        )
+        return None
+
     logger.info("Caller '%s' has the HyperPod CLI connect permissions.", caller_role_arn)
     return True
 
@@ -765,7 +842,9 @@ def verify_evaluation_caller_permissions(
 
     Returns:
         True  — all evaluation caller actions are allowed.
-        None  — could not be determined (caller is not a role, or cannot simulate).
+        None  — could not be determined (caller is not a role, cannot simulate, or
+            the decision is unverifiable because of condition-based SCPs or
+            permissions boundaries).
 
     Raises:
         RoleValidationError: If permissions are definitively denied.
@@ -788,7 +867,7 @@ def verify_evaluation_caller_permissions(
         return None
 
     try:
-        denied = _simulate_denied_actions(
+        denied, unverifiable = _simulate_denied_actions(
             iam_client, caller_role_arn, list(EVALUATION_CALLER_ACTIONS)
         )
     except ClientError as e:
@@ -816,6 +895,15 @@ def verify_evaluation_caller_permissions(
         )
         raise RoleValidationError(message)
 
+    if unverifiable:
+        logger.info(
+            "Cannot definitively verify evaluation pipeline permissions for '%s' due to "
+            "Organizations SCPs or permissions boundaries; errors will surface at "
+            "pipeline creation time.",
+            caller_role_arn,
+        )
+        return None
+
     logger.info(
         "Caller '%s' has the evaluation pipeline orchestration permissions.",
         caller_role_arn,
@@ -878,6 +966,22 @@ class IamRoleResolver:
         self._sts_client = self._boto_session.client("sts")
 
     # -- public API ---------------------------------------------------------
+    def resolve_and_validate_role(
+        self,
+        provided_role: Optional[str] = None,
+        role_type: str = "training",
+        *,
+        role_arn: Optional[str] = None,
+        validate_role: bool = True,
+    ) -> str:
+        """Resolve and validate an IAM role (read-only; does not mutate IAM)."""
+        return resolve_and_validate_role(
+            provided_role=provided_role or role_arn,
+            role_type=role_type,
+            sagemaker_session=self._sagemaker_session,
+            validate_role=validate_role,
+        )
+
     def get_required_actions(self, role_type: str) -> List[str]:
         """Return the IAM actions a role of ``role_type`` needs (read-only preview)."""
         self._validate_role_type(role_type)

{sagemaker_core-2.24.0 → sagemaker_core-2.26.0}/src/sagemaker/core/image_uri_config/djl-neuronx.json

@@ -23,6 +23,7 @@
                 "eu-west-3": "763104351884",
                 "mx-central-1": "637423239942",
                 "sa-east-1": "763104351884",
+                "sa-west-1": "434457397001",
                 "us-east-1": "763104351884",
                 "us-east-2": "763104351884",
                 "us-gov-east-1": "446045086412",
@@ -59,6 +60,7 @@
                 "eu-west-3": "763104351884",
                 "mx-central-1": "637423239942",
                 "sa-east-1": "763104351884",
+                "sa-west-1": "434457397001",
                 "us-east-1": "763104351884",
                 "us-east-2": "763104351884",
                 "us-gov-east-1": "446045086412",
@@ -91,6 +93,7 @@
                 "eu-west-3": "763104351884",
                 "mx-central-1": "637423239942",
                 "sa-east-1": "763104351884",
+                "sa-west-1": "434457397001",
                 "us-east-1": "763104351884",
                 "us-east-2": "763104351884",
                 "us-gov-east-1": "446045086412",
@@ -123,6 +126,7 @@
                 "eu-west-3": "763104351884",
                 "mx-central-1": "637423239942",
                 "sa-east-1": "763104351884",
+                "sa-west-1": "434457397001",
                 "us-east-1": "763104351884",
                 "us-east-2": "763104351884",
                 "us-gov-east-1": "446045086412",
@@ -155,6 +159,7 @@
                 "eu-west-3": "763104351884",
                 "mx-central-1": "637423239942",
                 "sa-east-1": "763104351884",
+                "sa-west-1": "434457397001",
                 "us-east-1": "763104351884",
                 "us-east-2": "763104351884",
                 "us-gov-east-1": "446045086412",
@@ -185,6 +190,7 @@
                 "eu-west-3": "763104351884",
                 "mx-central-1": "637423239942",
                 "sa-east-1": "763104351884",
+                "sa-west-1": "434457397001",
                 "us-east-1": "763104351884",
                 "us-east-2": "763104351884",
                 "us-gov-east-1": "446045086412",
@@ -215,6 +221,7 @@
                 "eu-west-3": "763104351884",
                 "mx-central-1": "637423239942",
                 "sa-east-1": "763104351884",
+                "sa-west-1": "434457397001",
                 "us-east-1": "763104351884",
                 "us-east-2": "763104351884",
                 "us-gov-east-1": "446045086412",
@@ -245,6 +252,7 @@
                 "eu-west-3": "763104351884",
                 "mx-central-1": "637423239942",
                 "sa-east-1": "763104351884",
+                "sa-west-1": "434457397001",
                 "us-east-1": "763104351884",
                 "us-east-2": "763104351884",
                 "us-gov-east-1": "446045086412",

{sagemaker_core-2.24.0 → sagemaker_core-2.26.0}/src/sagemaker/core/image_uris.py

@@ -258,7 +258,7 @@ def retrieve(
         pt_or_tf_version = (
             re.compile("^(pytorch|tensorflow)(.*)$").match(base_framework_version).group(2)
         )
-        _version = original_version
+        _version = _version_for_config(version, config)
 
         if repo in [
             "huggingface-pytorch-trcomp-training",

{sagemaker_core-2.24.0 → sagemaker_core-2.26.0}/src/sagemaker/core/jumpstart/models.py

@@ -14,12 +14,15 @@
 
 from __future__ import absolute_import, annotations
 
+import logging
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
-from pydantic import Field, ConfigDict
+from pydantic import Field, ConfigDict, field_validator
 
 from sagemaker.core.jumpstart.configs import BaseConfig
 
+logger = logging.getLogger(__name__)
+
 
 class StrEnum(str, Enum):
     """A string-valued enumeration."""
@@ -406,6 +409,31 @@ class CapabilityEnum(StrEnum):
     FINE_TUNING = "FINE_TUNING"
     VALIDATION = "VALIDATION"
     INCREMENTAL_TRAINING = "INCREMENTAL_TRAINING"
+    CUSTOMIZATION = "CUSTOMIZATION"
+
+
+def _coerce_capabilities(value: Any) -> Any:
+    """Coerce ``Capabilities`` values, tolerating unknown entries.
+
+    Known values are converted to ``CapabilityEnum`` members. Unknown values
+    are kept as plain strings with a warning, so a new service-side capability
+    cannot fail validation of the entire ``HubContentDocument`` for a pinned
+    SDK version.
+    """
+    if not isinstance(value, list):
+        return value
+    coerced = []
+    for item in value:
+        try:
+            coerced.append(CapabilityEnum(item))
+        except ValueError:
+            logger.warning(
+                "Unknown capability %r is not a member of CapabilityEnum; "
+                "keeping it as a string.",
+                item,
+            )
+            coerced.append(item)
+    return coerced
 
 
 class DemoNotebookModel(BaseConfig):
@@ -444,7 +472,14 @@ class HostingComponentsModel(BaseConfig):
     DynamicContainerDeploymentSupported: Optional[bool] = None
     HostingAdditionalDataSources: Optional[Dict[str, List[HostingAdditionalDataSourceModel]]] = None
     BedrockIOMappingId: Optional[str] = None
-    Capabilities: Optional[List[CapabilityEnum]] = None
+    Capabilities: Optional[List[Union[CapabilityEnum, str]]] = None
+
+    @field_validator("Capabilities", mode="before")
+    @classmethod
+    def _validate_capabilities(cls, value: Any) -> Any:
+        """Keep unknown capability values as strings instead of failing."""
+        return _coerce_capabilities(value)
+
     HostingEcrUri: Optional[str] = None
     HostingArtifactS3DataType: Optional[HostingArtifactS3DataTypeEnum] = None
     HostingArtifactCompressionType: Optional[HostingArtifactCompressionTypeEnum] = None
@@ -506,4 +541,10 @@ class HubContentDocument(HostingComponentsModel, TrainingComponentsModel):
     TrainingConfigRankings: Optional[Dict[str, RankingSchemaModel]] = None
     Resources: Optional[List[ResourceModel]] = None
     Highlights: Optional[List[str]] = None
-    Capabilities: List[CapabilityEnum] = None
+    Capabilities: List[Union[CapabilityEnum, str]] = None
+
+    @field_validator("Capabilities", mode="before")
+    @classmethod
+    def _validate_capabilities(cls, value: Any) -> Any:
+        """Keep unknown capability values as strings instead of failing."""
+        return _coerce_capabilities(value)

{sagemaker_core-2.24.0 → sagemaker_core-2.26.0}/src/sagemaker/core/jumpstart/region_config.json

@@ -139,6 +139,10 @@
         "content_bucket": "jumpstart-cache-prod-me-south-1",
         "gated_content_bucket": "jumpstart-private-cache-prod-me-south-1"
     },
+    "me-west-1": {
+        "content_bucket": "jumpstart-cache-prod-me-west-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-me-west-1"
+    },
     "mx-central-1": {
         "content_bucket": "jumpstart-cache-prod-mx-central-1",
         "gated_content_bucket": "jumpstart-private-cache-prod-mx-central-1"
@@ -148,6 +152,10 @@
         "gated_content_bucket": "jumpstart-private-cache-prod-sa-east-1",
         "neo_content_bucket": "sagemaker-sd-models-prod-sa-east-1"
     },
+    "sa-west-1": {
+        "content_bucket": "jumpstart-cache-prod-sa-west-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-sa-west-1"
+    },
     "us-east-1": {
         "content_bucket": "jumpstart-cache-prod-us-east-1",
         "gated_content_bucket": "jumpstart-private-cache-prod-us-east-1",

{sagemaker_core-2.24.0 → sagemaker_core-2.26.0}/src/sagemaker/core/local/data.py

@@ -357,15 +357,19 @@ class MultiRecordStrategy(BatchStrategy):
         Returns:
             generator of records
         """
-        buffer = ""
+        buffer = None
         for element in self.splitter.split(file):
+            if buffer is None:
+                # Match the buffer type to the record type so binary inputs
+                # concatenate correctly instead of raising a TypeError.
+                buffer = b"" if isinstance(element, bytes) else ""
             if _payload_size_within_limit(buffer + element, size):
                 buffer += element
             else:
                 tmp = buffer
                 buffer = element
                 yield tmp
-        if _validate_payload_size(buffer, size):
+        if buffer is not None and _validate_payload_size(buffer, size):
             yield buffer
 
 

{sagemaker_core-2.24.0 → sagemaker_core-2.26.0}/src/sagemaker/core/local/entities.py

@@ -338,7 +338,11 @@ class _LocalTransformJob(object):
         self.container.serve(self.primary_container["ModelDataUrl"], environment)
 
         serving_port = get_config_value("local.serving_port", self.local_session.config) or 8080
-        _wait_for_serving_container(serving_port)
+        health_check_timeout = (
+            get_config_value("local.health_check_timeout", self.local_session.config)
+            or HEALTH_CHECK_TIMEOUT_LIMIT
+        )
+        _wait_for_serving_container(serving_port, timeout=health_check_timeout)
 
         # Get capabilities from Container if needed
         endpoint_url = "http://%s:%d/execution-parameters" % (get_docker_host(), serving_port)
@@ -632,7 +636,11 @@ class _LocalEndpoint(object):
         )
 
         serving_port = get_config_value("local.serving_port", self.local_session.config) or 8080
-        _wait_for_serving_container(serving_port)
+        health_check_timeout = (
+            get_config_value("local.health_check_timeout", self.local_session.config)
+            or HEALTH_CHECK_TIMEOUT_LIMIT
+        )
+        _wait_for_serving_container(serving_port, timeout=health_check_timeout)
         # the container is running and it passed the healthcheck status is now InService
         self.state = _LocalEndpoint._IN_SERVICE
 
@@ -655,15 +663,21 @@ class _LocalEndpoint(object):
         return response
 
 
-def _wait_for_serving_container(serving_port):
-    """Placeholder docstring."""
+def _wait_for_serving_container(serving_port, timeout=HEALTH_CHECK_TIMEOUT_LIMIT):
+    """Wait for the serving container to become healthy.
+
+    Args:
+        serving_port (int): The port the serving container is listening on.
+        timeout (int): Maximum number of seconds to wait for the container to become healthy.
+            Defaults to ``HEALTH_CHECK_TIMEOUT_LIMIT``.
+    """
     i = 0
     http = urllib3.PoolManager()
 
     endpoint_url = "http://%s:%d/ping" % (get_docker_host(), serving_port)
     while True:
         i += 5
-        if i >= HEALTH_CHECK_TIMEOUT_LIMIT:
+        if i >= timeout:
             raise RuntimeError("Giving up, endpoint didn't launch correctly")
 
         logger.info("Checking if serving container is up, attempt: %s", i)

{sagemaker_core-2.24.0 → sagemaker_core-2.26.0}/src/sagemaker/core/local/image.py

@@ -39,6 +39,7 @@ import sagemaker
 from sagemaker.core.config.config_schema import CONTAINER_CONFIG, LOCAL
 import sagemaker.core
 from sagemaker.core.common_utils import custom_extractall_tarfile
+from sagemaker.core.constants import DIR_PARAM_NAME, SAGEMAKER_OUTPUT_LOCATION
 
 CONTAINER_PREFIX = "algo"
 STUDIO_HOST_NAME = "sagemaker-local"
@@ -277,9 +278,7 @@ class _SageMakerContainer(object):
             data_dir, input_data_config, output_data_config, hyperparameters
         )
         # If local, source directory needs to be updated to mounted /opt/ml/code path
-        hyperparameters = self._update_local_src_path(
-            hyperparameters, key=sagemaker.serve.model_builder.DIR_PARAM_NAME
-        )
+        hyperparameters = self._update_local_src_path(hyperparameters, key=DIR_PARAM_NAME)
 
         # Create the configuration files for each container that we will create
         # Each container will map the additional local volumes (if any).
@@ -344,15 +343,15 @@ class _SageMakerContainer(object):
         volumes = self._prepare_serving_volumes(model_dir)
 
         # If the user script was passed as a file:// mount it to the container.
-        if sagemaker.serve.model_builder.DIR_PARAM_NAME.upper() in environment:
-            script_dir = environment[sagemaker.serve.model_builder.DIR_PARAM_NAME.upper()]
+        if DIR_PARAM_NAME.upper() in environment:
+            script_dir = environment[DIR_PARAM_NAME.upper()]
             parsed_uri = urlparse(script_dir)
             if parsed_uri.scheme == "file":
                 host_dir = os.path.abspath(parsed_uri.netloc + parsed_uri.path)
                 volumes.append(_Volume(host_dir, "/opt/ml/code"))
                 # Update path to mount location
                 environment = environment.copy()
-                environment[sagemaker.serve.model_builder.DIR_PARAM_NAME.upper()] = "/opt/ml/code"
+                environment[DIR_PARAM_NAME.upper()] = "/opt/ml/code"
 
         if _ecr_login_if_needed(self.sagemaker_session.boto_session, self.image):
             _pull_image(self.image)
@@ -583,8 +582,8 @@ class _SageMakerContainer(object):
 
         # If there is a training script directory and it is a local directory,
         # mount it to the container.
-        if sagemaker.serve.model_builder.DIR_PARAM_NAME in hyperparameters:
-            training_dir = json.loads(hyperparameters[sagemaker.serve.model_builder.DIR_PARAM_NAME])
+        if DIR_PARAM_NAME in hyperparameters:
+            training_dir = json.loads(hyperparameters[DIR_PARAM_NAME])
             parsed_uri = urlparse(training_dir)
             if parsed_uri.scheme == "file":
                 host_dir = os.path.abspath(parsed_uri.netloc + parsed_uri.path)
@@ -593,10 +592,7 @@ class _SageMakerContainer(object):
                 volumes.append(_Volume(shared_dir, "/opt/ml/shared"))
 
         parsed_uri = urlparse(output_data_config["S3OutputPath"])
-        if (
-            parsed_uri.scheme == "file"
-            and sagemaker.serve.model_builder.SAGEMAKER_OUTPUT_LOCATION in hyperparameters
-        ):
+        if parsed_uri.scheme == "file" and SAGEMAKER_OUTPUT_LOCATION in hyperparameters:
             dir_path = os.path.abspath(parsed_uri.netloc + parsed_uri.path)
             intermediate_dir = os.path.join(dir_path, "output", "intermediate")
             if not os.path.exists(intermediate_dir):