Merge branch 'master' into command-injection

Author: pintaoz-aws

requirements/extras/test_requirements.txt

@@ -4,8 +4,11 @@ pytest-xdist
 mock
 pydantic==2.11.9
 pydantic_core==2.33.2
-pandas
+pandas>=2.3.0
+numpy>=2.0.0, <3.0
+scikit-learn==1.6.1
 scipy
 omegaconf
 graphene
-typing_extensions>=4.9.0
+typing_extensions>=4.9.0
+tensorflow>=2.16.2,<=2.19.0

sagemaker-core/src/sagemaker/core/common_utils.py

@@ -59,6 +59,7 @@
     "us-isob-east-1": "sc2s.sgov.gov",
     "us-isof-south-1": "csp.hci.ic.gov",
     "us-isof-east-1": "csp.hci.ic.gov",
+    "eu-isoe-west-1": "cloud.adc-e.uk",
 }
 
 ECR_URI_PATTERN = r"^(\d+)(\.)dkr(\.)ecr(\.)(.+)(\.)(.*)(/)(.*:.*)$"
@@ -74,6 +75,20 @@
 WAITING_DOT_NUMBER = 10
 MAX_ITEMS = 100
 PAGE_SIZE = 10
+_MAX_BUFFER_SIZE = 100 * 1024 * 1024  # 100 MB - Maximum buffer size for streaming iterators
+
+_SENSITIVE_SYSTEM_PATHS = [
+    abspath(os.path.expanduser("~/.aws")),
+    abspath(os.path.expanduser("~/.ssh")),
+    abspath(os.path.expanduser("~/.kube")),
+    abspath(os.path.expanduser("~/.docker")),
+    abspath(os.path.expanduser("~/.config")),
+    abspath(os.path.expanduser("~/.credentials")),
+    "/etc",
+    "/root",
+    "/var/lib",
+    "/opt/ml/metadata",
+]
 
 logger = logging.getLogger(__name__)
 
@@ -607,11 +622,73 @@ def _save_model(repacked_model_uri, tmp_model_path, sagemaker_session, kms_key):
         shutil.move(tmp_model_path, repacked_model_uri.replace("file://", ""))
 
 
+def _validate_source_directory(source_directory):
+    """Validate that source_directory is safe to use.
+
+    Ensures the source directory path does not access restricted system locations.
+
+    Args:
+        source_directory (str): The source directory path to validate.
+
+    Raises:
+        ValueError: If the path is not allowed.
+    """
+    if not source_directory or source_directory.lower().startswith("s3://"):
+        # S3 paths and None are safe
+        return
+
+    # Resolve symlinks to get the actual path
+    abs_source = abspath(realpath(source_directory))
+
+    # Check if the source path is under any sensitive directory
+    for sensitive_path in _SENSITIVE_SYSTEM_PATHS:
+        if abs_source != "/" and abs_source.startswith(sensitive_path):
+            raise ValueError(
+                f"source_directory cannot access sensitive system paths. "
+                f"Got: {source_directory} (resolved to {abs_source})"
+            )
+
+
+def _validate_dependency_path(dependency):
+    """Validate that a dependency path is safe to use.
+
+    Ensures the dependency path does not access restricted system locations.
+
+    Args:
+        dependency (str): The dependency path to validate.
+
+    Raises:
+        ValueError: If the path is not allowed.
+    """
+    if not dependency:
+        return
+
+    # Resolve symlinks to get the actual path
+    abs_dependency = abspath(realpath(dependency))
+
+    # Check if the dependency path is under any sensitive directory
+    for sensitive_path in _SENSITIVE_SYSTEM_PATHS:
+        if abs_dependency != "/" and abs_dependency.startswith(sensitive_path):
+            raise ValueError(
+                f"dependency path cannot access sensitive system paths. "
+                f"Got: {dependency} (resolved to {abs_dependency})"
+            )
+
+
 def _create_or_update_code_dir(
     model_dir, inference_script, source_directory, dependencies, sagemaker_session, tmp
 ):
     """Placeholder docstring"""
     code_dir = os.path.join(model_dir, "code")
+    resolved_code_dir = _get_resolved_path(code_dir)
+    
+    # Validate that code_dir does not resolve to a sensitive system path
+    for sensitive_path in _SENSITIVE_SYSTEM_PATHS:
+        if resolved_code_dir != "/" and resolved_code_dir.startswith(sensitive_path):
+            raise ValueError(
+                f"Invalid code_dir path: {code_dir} resolves to sensitive system path {resolved_code_dir}"
+            )
+
     if source_directory and source_directory.lower().startswith("s3://"):
         local_code_path = os.path.join(tmp, "local_code.tar.gz")
         download_file_from_url(source_directory, local_code_path, sagemaker_session)
@@ -620,6 +697,8 @@ def _create_or_update_code_dir(
             custom_extractall_tarfile(t, code_dir)
 
     elif source_directory:
+        # Validate source_directory for security
+        _validate_source_directory(source_directory)
         if os.path.exists(code_dir):
             shutil.rmtree(code_dir)
         shutil.copytree(source_directory, code_dir)
@@ -635,6 +714,8 @@ def _create_or_update_code_dir(
                 raise
 
     for dependency in dependencies:
+        # Validate dependency path for security
+        _validate_dependency_path(dependency)
         lib_dir = os.path.join(code_dir, "lib")
         if os.path.isdir(dependency):
             shutil.copytree(dependency, os.path.join(lib_dir, os.path.basename(dependency)))
@@ -1555,7 +1636,7 @@ def get_instance_type_family(instance_type: str) -> str:
     """
     instance_type_family = ""
     if isinstance(instance_type, str):
-        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        match = re.match(r"^ml[\._]([a-z\d\-]+)\.?\w*$", instance_type)
         if match is not None:
             instance_type_family = match[1]
     return instance_type_family
@@ -1646,6 +1727,38 @@ def _get_safe_members(members):
             yield file_info
 
 
+def _validate_extracted_paths(extract_path):
+    """Validate that extracted paths remain within the expected directory.
+
+    Performs post-extraction validation to ensure all extracted files and directories
+    are within the intended extraction path.
+
+    Args:
+        extract_path (str): The path where files were extracted.
+
+    Raises:
+        ValueError: If any extracted file is outside the expected extraction path.
+    """
+    base = _get_resolved_path(extract_path)
+
+    for root, dirs, files in os.walk(extract_path):
+        # Check directories
+        for dir_name in dirs:
+            dir_path = os.path.join(root, dir_name)
+            resolved = _get_resolved_path(dir_path)
+            if not resolved.startswith(base):
+                logger.error("Extracted directory escaped extraction path: %s", dir_path)
+                raise ValueError(f"Extracted path outside expected directory: {dir_path}")
+
+        # Check files
+        for file_name in files:
+            file_path = os.path.join(root, file_name)
+            resolved = _get_resolved_path(file_path)
+            if not resolved.startswith(base):
+                logger.error("Extracted file escaped extraction path: %s", file_path)
+                raise ValueError(f"Extracted path outside expected directory: {file_path}")
+
+
 def custom_extractall_tarfile(tar, extract_path):
     """Extract a tarfile, optionally using data_filter if available.
 
@@ -1666,6 +1779,8 @@ def custom_extractall_tarfile(tar, extract_path):
         tar.extractall(path=extract_path, filter="data")
     else:
         tar.extractall(path=extract_path, members=_get_safe_members(tar))
+        # Re-validate extracted paths to catch symlink race conditions
+        _validate_extracted_paths(extract_path)
 
 
 def can_model_package_source_uri_autopopulate(source_uri: str):

sagemaker-core/src/sagemaker/core/fw_utils.py

@@ -25,13 +25,13 @@
 
 from packaging import version
 
-import sagemaker.core.common_utils as sagemaker_utils
-from sagemaker.core.deprecations import deprecation_warn_base, renamed_kwargs
+import sagemaker.core.common_utils as utils
+from sagemaker.core.deprecations import deprecation_warn_base, renamed_kwargs, renamed_warning
 from sagemaker.core.instance_group import InstanceGroup
-from sagemaker.core.s3 import s3_path_join
+from sagemaker.core.s3.utils import s3_path_join
 from sagemaker.core.session_settings import SessionSettings
 from sagemaker.core.workflow import is_pipeline_variable
-from sagemaker.core.helper.pipeline_variable import PipelineVariable
+from sagemaker.core.workflow.entities import PipelineVariable
 
 logger = logging.getLogger(__name__)
 
@@ -155,6 +155,9 @@
     "2.3.1",
     "2.4.1",
     "2.5.1",
+    "2.6.0",
+    "2.7.1",
+    "2.8.0",
 ]
 
 TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES = ["torch_distributed"]
@@ -455,7 +458,7 @@ def tar_and_upload_dir(
 
     try:
         source_files = _list_files_to_compress(script, directory) + dependencies
-        tar_file = sagemaker_utils.create_tar_file(
+        tar_file = utils.create_tar_file(
             source_files, os.path.join(tmp, _TAR_SOURCE_FILENAME)
         )
 
@@ -516,7 +519,7 @@ def framework_name_from_image(image_uri):
             - str: The image tag
             - str: If the TensorFlow image is script mode
     """
-    sagemaker_pattern = re.compile(sagemaker_utils.ECR_URI_PATTERN)
+    sagemaker_pattern = re.compile(utils.ECR_URI_PATTERN)
     sagemaker_match = sagemaker_pattern.match(image_uri)
     if sagemaker_match is None:
         return None, None, None, None
@@ -595,7 +598,7 @@ def model_code_key_prefix(code_location_key_prefix, model_name, image):
     """
     name_from_image = f"/model_code/{int(time.time())}"
     if not is_pipeline_variable(image):
-        name_from_image = sagemaker_utils.name_from_image(image)
+        name_from_image = utils.name_from_image(image)
     return s3_path_join(code_location_key_prefix, model_name or name_from_image)
 
 
@@ -961,7 +964,7 @@ def validate_distribution_for_instance_type(instance_type, distribution):
     """
     err_msg = ""
     if isinstance(instance_type, str):
-        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        match = re.match(r"^ml[\._]([a-z\d\-]+)\.?\w*$", instance_type)
         if match and match[1].startswith("trn"):
             keys = list(distribution.keys())
             if len(keys) == 0:
@@ -1062,7 +1065,7 @@ def validate_torch_distributed_distribution(
             )
 
     # Check entry point type
-    if not entry_point.endswith(".py"):
+    if entry_point is not None and not entry_point.endswith(".py"):
         err_msg += (
             "Unsupported entry point type for the distribution torch_distributed.\n"
             "Only python programs (*.py) are supported."
@@ -1082,7 +1085,7 @@ def _is_gpu_instance(instance_type):
         bool: Whether or not the instance_type supports GPU
     """
     if isinstance(instance_type, str):
-        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        match = re.match(r"^ml[\._]([a-z\d\-]+)\.?\w*$", instance_type)
         if match:
             if match[1].startswith("p") or match[1].startswith("g"):
                 return True
@@ -1101,7 +1104,7 @@ def _is_trainium_instance(instance_type):
         bool: Whether or not the instance_type is a Trainium instance
     """
     if isinstance(instance_type, str):
-        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        match = re.match(r"^ml[\._]([a-z\d\-]+)\.?\w*$", instance_type)
         if match and match[1].startswith("trn"):
             return True
     return False
@@ -1148,7 +1151,7 @@ def _instance_type_supports_profiler(instance_type):
         bool: Whether or not the region supports Amazon SageMaker Debugger profiling feature.
     """
     if isinstance(instance_type, str):
-        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        match = re.match(r"^ml[\._]([a-z\d\-]+)\.?\w*$", instance_type)
         if match and match[1].startswith("trn"):
             return True
     return False
@@ -1174,3 +1177,44 @@ def validate_version_or_image_args(framework_version, py_version, image_uri):
             "framework_version or py_version was None, yet image_uri was also None. "
             "Either specify both framework_version and py_version, or specify image_uri."
         )
+
+
+def create_image_uri(
+    region,
+    framework,
+    instance_type,
+    framework_version,
+    py_version=None,
+    account=None,  # pylint: disable=W0613
+    accelerator_type=None,
+    optimized_families=None,  # pylint: disable=W0613
+):
+    """Deprecated method. Please use sagemaker.image_uris.retrieve().
+
+    Args:
+        region (str): AWS region where the image is uploaded.
+        framework (str): framework used by the image.
+        instance_type (str): SageMaker instance type. Used to determine device
+            type (cpu/gpu/family-specific optimized).
+        framework_version (str): The version of the framework.
+        py_version (str): Optional. Python version Ex: `py38, py39, py310, py311`.
+            If not specified, image uri will not include a python component.
+        account (str): AWS account that contains the image. (default:
+            '520713654638')
+        accelerator_type (str): SageMaker Elastic Inference accelerator type.
+        optimized_families (str): Deprecated. A no-op argument.
+
+    Returns:
+        the image uri
+    """
+    from sagemaker.core import image_uris
+    
+    renamed_warning("The method create_image_uri")
+    return image_uris.retrieve(
+        framework=framework,
+        region=region,
+        version=framework_version,
+        py_version=py_version,
+        instance_type=instance_type,
+        accelerator_type=accelerator_type,
+    )