Merge branch 'master' into feature/aws-rft-sdk-integration

Author: mollyheamazon

sagemaker-core/src/sagemaker/core/jumpstart/search.py

@@ -188,15 +188,6 @@ def _parse_primary_expression(self, tokens: List[str], pos: int) -> tuple[_Expre
             # Pattern token
             return _PatternNode(token), pos + 1
 
-    def _convert_expression(self, expr: str) -> str:
-        """
-        Legacy method for backward compatibility.
-        This method is no longer used but kept to avoid breaking changes.
-        """
-        # This method is deprecated and should not be used
-        # It's kept only for backward compatibility
-        return expr
-
 
 def _list_all_hub_models(hub_name: str, sm_client: Session) -> Iterator[HubContent]:
     """

sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py

@@ -17,6 +17,7 @@
 import os
 import re
 import shutil
+import stat
 import subprocess
 from tempfile import TemporaryDirectory
 from typing import Any, Dict, List, Optional
@@ -57,6 +58,17 @@
 SM_STUDIO_LOCAL_MODE = "SM_STUDIO_LOCAL_MODE"
 
 
+def _rmtree(path):
+    """Remove a directory tree, handling root-owned files from Docker containers."""
+    def _onerror(func, path, exc_info):
+        if isinstance(exc_info[1], PermissionError):
+            os.chmod(path, stat.S_IRWXU)
+            func(path)
+        else:
+            raise exc_info[1]
+    shutil.rmtree(path, onerror=_onerror)
+
+
 class _LocalContainer(BaseModel):
     """A local training job class for local mode model trainer.
 
@@ -209,12 +221,12 @@ def train(
         # Print our Job Complete line
         logger.info("Local training job completed, output artifacts saved to %s", artifacts)
 
-        shutil.rmtree(os.path.join(self.container_root, "input"))
-        shutil.rmtree(os.path.join(self.container_root, "shared"))
+        _rmtree(os.path.join(self.container_root, "input"))
+        _rmtree(os.path.join(self.container_root, "shared"))
         for host in self.hosts:
-            shutil.rmtree(os.path.join(self.container_root, host))
+            _rmtree(os.path.join(self.container_root, host))
         for folder in self._temporary_folders:
-            shutil.rmtree(os.path.join(self.container_root, folder))
+            _rmtree(os.path.join(self.container_root, folder))
         return artifacts
 
     def retrieve_artifacts(

sagemaker-core/src/sagemaker/core/processing.py

@@ -487,6 +487,9 @@ def _normalize_outputs(self, outputs=None):
                 # If the output's s3_uri is not an s3_uri, create one.
                 parse_result = urlparse(output.s3_output.s3_uri)
                 if parse_result.scheme != "s3":
+                    if getattr(self.sagemaker_session, "local_mode", False) and parse_result.scheme == "file":
+                        normalized_outputs.append(output)
+                        continue
                     if _pipeline_config:
                         s3_uri = Join(
                             on="/",

sagemaker-core/tests/unit/test_processing.py

@@ -238,6 +238,261 @@ def test_normalize_outputs_invalid_type(self, mock_session):
             processor._normalize_outputs(["invalid"])
 
 
+
+
+class TestBugConditionFileUriReplacedInLocalMode:
+    """Bug condition exploration test: file:// URIs should be preserved in local mode.
+
+    **Validates: Requirements 1.1, 1.2, 2.1, 2.2**
+
+    EXPECTED TO FAIL on unfixed code — failure confirms the bug exists.
+    The bug is that _normalize_outputs() replaces file:// URIs with s3:// paths
+    even when the session is a LocalSession (local_mode=True).
+    """
+
+    @pytest.fixture
+    def local_mock_session(self):
+        session = Mock()
+        session.boto_session = Mock()
+        session.boto_session.region_name = "us-west-2"
+        session.sagemaker_client = Mock()
+        session.default_bucket = Mock(return_value="default-bucket")
+        session.default_bucket_prefix = "prefix"
+        session.expand_role = Mock(side_effect=lambda x: x)
+        session.sagemaker_config = {}
+        session.local_mode = True
+        return session
+
+    @pytest.mark.parametrize(
+        "file_uri",
+        [
+            "file:///tmp/output",
+            "file:///home/user/results",
+            "file:///data/processed",
+        ],
+    )
+    def test_normalize_outputs_preserves_file_uri_in_local_mode(self, local_mock_session, file_uri):
+        """file:// URIs must be preserved when local_mode=True.
+
+        On unfixed code, _normalize_outputs replaces file:// URIs with
+        s3://default-bucket/prefix/job-name/output/output-1, which is the bug.
+        """
+        processor = Processor(
+            role="arn:aws:iam::123456789012:role/SageMakerRole",
+            image_uri="test-image:latest",
+            instance_count=1,
+            instance_type="ml.m5.xlarge",
+            sagemaker_session=local_mock_session,
+        )
+        processor._current_job_name = "test-job"
+
+        s3_output = ProcessingS3Output(
+            s3_uri=file_uri,
+            local_path="/opt/ml/processing/output",
+            s3_upload_mode="EndOfJob",
+        )
+        outputs = [ProcessingOutput(output_name="my-output", s3_output=s3_output)]
+
+        with patch("sagemaker.core.workflow.utilities._pipeline_config", None):
+            result = processor._normalize_outputs(outputs)
+
+        assert len(result) == 1
+        assert result[0].s3_output.s3_uri == file_uri, (
+            f"Expected file:// URI to be preserved as '{file_uri}' in local mode, "
+            f"but got '{result[0].s3_output.s3_uri}'"
+        )
+
+
+class TestPreservationNonLocalFileBehavior:
+    """Preservation property tests: Non-local-file behavior must remain unchanged.
+
+    **Validates: Requirements 3.1, 3.2, 3.3, 3.4**
+
+    These tests capture baseline behavior on UNFIXED code. They MUST PASS on both
+    unfixed and fixed code, confirming no regressions are introduced by the fix.
+    """
+
+    @pytest.fixture
+    def session_local_mode_true(self):
+        session = Mock()
+        session.boto_session = Mock()
+        session.boto_session.region_name = "us-west-2"
+        session.sagemaker_client = Mock()
+        session.default_bucket = Mock(return_value="default-bucket")
+        session.default_bucket_prefix = "prefix"
+        session.expand_role = Mock(side_effect=lambda x: x)
+        session.sagemaker_config = {}
+        session.local_mode = True
+        return session
+
+    @pytest.fixture
+    def session_local_mode_false(self):
+        session = Mock()
+        session.boto_session = Mock()
+        session.boto_session.region_name = "us-west-2"
+        session.sagemaker_client = Mock()
+        session.default_bucket = Mock(return_value="default-bucket")
+        session.default_bucket_prefix = "prefix"
+        session.expand_role = Mock(side_effect=lambda x: x)
+        session.sagemaker_config = {}
+        session.local_mode = False
+        return session
+
+    def _make_processor(self, session):
+        processor = Processor(
+            role="arn:aws:iam::123456789012:role/SageMakerRole",
+            image_uri="test-image:latest",
+            instance_count=1,
+            instance_type="ml.m5.xlarge",
+            sagemaker_session=session,
+        )
+        processor._current_job_name = "test-job"
+        return processor
+
+    # --- Requirement 3.1: S3 URIs pass through unchanged regardless of local_mode ---
+
+    @pytest.mark.parametrize(
+        "s3_uri,local_mode_fixture",
+        [
+            ("s3://my-bucket/path", "session_local_mode_true"),
+            ("s3://my-bucket/path", "session_local_mode_false"),
+            ("s3://another-bucket/deep/nested/path", "session_local_mode_true"),
+            ("s3://another-bucket/deep/nested/path", "session_local_mode_false"),
+        ],
+    )
+    def test_s3_uri_preserved_regardless_of_local_mode(self, s3_uri, local_mode_fixture, request):
+        """S3 URIs must pass through unchanged regardless of local_mode setting.
+
+        **Validates: Requirements 3.1**
+        """
+        session = request.getfixturevalue(local_mode_fixture)
+        processor = self._make_processor(session)
+
+        s3_output = ProcessingS3Output(
+            s3_uri=s3_uri,
+            local_path="/opt/ml/processing/output",
+            s3_upload_mode="EndOfJob",
+        )
+        outputs = [ProcessingOutput(output_name="my-output", s3_output=s3_output)]
+
+        with patch("sagemaker.core.workflow.utilities._pipeline_config", None):
+            result = processor._normalize_outputs(outputs)
+
+        assert len(result) == 1
+        assert result[0].s3_output.s3_uri == s3_uri
+
+    # --- Requirement 3.2: Non-S3 URIs with local_mode=False replaced with S3 paths ---
+
+    @pytest.mark.parametrize(
+        "non_s3_uri",
+        [
+            "/local/output/path",
+            "http://example.com/output",
+            "ftp://server/output",
+        ],
+    )
+    def test_non_s3_uri_replaced_when_not_local_mode(self, non_s3_uri, session_local_mode_false):
+        """Non-S3 URIs in non-local sessions are replaced with auto-generated S3 paths.
+
+        **Validates: Requirements 3.2**
+        """
+        processor = self._make_processor(session_local_mode_false)
+
+        s3_output = ProcessingS3Output(
+            s3_uri=non_s3_uri,
+            local_path="/opt/ml/processing/output",
+            s3_upload_mode="EndOfJob",
+        )
+        outputs = [ProcessingOutput(output_name="output-1", s3_output=s3_output)]
+
+        with patch("sagemaker.core.workflow.utilities._pipeline_config", None):
+            result = processor._normalize_outputs(outputs)
+
+        assert len(result) == 1
+        assert result[0].s3_output.s3_uri.startswith("s3://default-bucket/")
+
+    # --- Requirement 3.3: Pipeline variable URIs skip normalization ---
+
+    def test_pipeline_variable_uri_skips_normalization(self, session_local_mode_false):
+        """Pipeline variable URIs skip normalization entirely.
+
+        **Validates: Requirements 3.3**
+        """
+        processor = self._make_processor(session_local_mode_false)
+
+        s3_output = ProcessingS3Output(
+            s3_uri="s3://bucket/output",
+            local_path="/opt/ml/processing/output",
+            s3_upload_mode="EndOfJob",
+        )
+        outputs = [ProcessingOutput(output_name="output-1", s3_output=s3_output)]
+
+        with patch("sagemaker.core.processing.is_pipeline_variable", return_value=True):
+            result = processor._normalize_outputs(outputs)
+
+        assert len(result) == 1
+        # Pipeline variable outputs are appended as-is without URI modification
+        assert result[0].s3_output.s3_uri == "s3://bucket/output"
+
+    # --- Requirement 3.4: Non-ProcessingOutput objects raise TypeError ---
+
+    @pytest.mark.parametrize(
+        "invalid_output",
+        [
+            ["a string"],
+            [42],
+            [{"key": "value"}],
+        ],
+    )
+    def test_non_processing_output_raises_type_error(self, invalid_output, session_local_mode_false):
+        """Non-ProcessingOutput objects must raise TypeError.
+
+        **Validates: Requirements 3.4**
+        """
+        processor = self._make_processor(session_local_mode_false)
+
+        with pytest.raises(TypeError, match="must be provided as ProcessingOutput objects"):
+            processor._normalize_outputs(invalid_output)
+
+    # --- Output name auto-generation ---
+
+    def test_multiple_outputs_with_s3_uris_preserved(self, session_local_mode_false):
+        """Multiple outputs with S3 URIs are all preserved unchanged.
+
+        **Validates: Requirements 3.1, 3.2**
+        """
+        processor = self._make_processor(session_local_mode_false)
+
+        outputs = [
+            ProcessingOutput(
+                output_name="first-output",
+                s3_output=ProcessingS3Output(
+                    s3_uri="s3://my-bucket/first",
+                    local_path="/opt/ml/processing/output1",
+                    s3_upload_mode="EndOfJob",
+                ),
+            ),
+            ProcessingOutput(
+                output_name="second-output",
+                s3_output=ProcessingS3Output(
+                    s3_uri="s3://my-bucket/second",
+                    local_path="/opt/ml/processing/output2",
+                    s3_upload_mode="EndOfJob",
+                ),
+            ),
+        ]
+
+        with patch("sagemaker.core.workflow.utilities._pipeline_config", None):
+            result = processor._normalize_outputs(outputs)
+
+        assert len(result) == 2
+        assert result[0].output_name == "first-output"
+        assert result[1].output_name == "second-output"
+        # S3 URIs should be preserved since they already have s3:// scheme
+        assert result[0].s3_output.s3_uri == "s3://my-bucket/first"
+        assert result[1].s3_output.s3_uri == "s3://my-bucket/second"
+
+
 class TestProcessorStartNew:
     def test_start_new_with_pipeline_session(self, mock_session):
         from sagemaker.core.workflow.pipeline_context import PipelineSession

sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py

@@ -3,13 +3,16 @@
 from __future__ import absolute_import
 from pathlib import Path
 import logging
+import os
 from datetime import datetime, timedelta
 from typing import Dict, Type
 import base64
 import time
 import subprocess
 import docker
 
+from sagemaker.core.local.utils import check_for_studio
+
 from sagemaker.serve.model_server.tensorflow_serving.server import LocalTensorflowServing
 from sagemaker.serve.spec.inference_spec import InferenceSpec
 from sagemaker.serve.builder.schema_builder import SchemaBuilder
@@ -33,6 +36,25 @@
     + "Please increase container_timeout_seconds or review your inference code."
 )
 
+STUDIO_DOCKER_SOCKET_PATHS = [
+    "/docker/proxy/docker.sock",
+    "/var/run/docker.sock",
+]
+
+
+def _get_docker_client():
+    """Get a Docker client, handling SageMaker Studio's non-standard socket path."""
+    if os.environ.get("DOCKER_HOST"):
+        return docker.from_env()
+    try:
+        if check_for_studio():
+            for socket_path in STUDIO_DOCKER_SOCKET_PATHS:
+                if os.path.exists(socket_path):
+                    return docker.DockerClient(base_url=f"unix://{socket_path}")
+    except (NotImplementedError, Exception):
+        pass
+    return docker.from_env()
+
 
 class LocalContainerMode(
     LocalTorchServe,
@@ -212,7 +234,7 @@ def _pull_image(self, image: str):
 
         # Check if Docker is available first
         try:
-            self.client = docker.from_env()
+            self.client = _get_docker_client()
             self.client.ping()  # Test Docker connection
         except Exception as e:
             raise RuntimeError(