Add wheel install for remote function integ test to use local code

Author: mollyheamazon

sagemaker-core/tests/integ/remote_function/conftest.py

@@ -0,0 +1,70 @@
+"""Shared fixtures for remote function integration tests."""
+
+import glob
+import os
+import subprocess
+import tempfile
+
+import cloudpickle
+import pytest
+
+from sagemaker.core.helper.session_helper import Session
+from sagemaker.core.s3 import S3Uploader
+
+
+def _get_repo_root():
+    return os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", ".."))
+
+
+def _build_and_upload_core_wheel(sagemaker_session):
+    """Build sagemaker-core wheel and upload to S3. Returns (s3_prefix, wheel_basename)."""
+    repo_root = _get_repo_root()
+    dist_dir = tempfile.mkdtemp(prefix="sagemaker_core_wheel_")
+
+    subprocess.run(
+        f"python -m build --wheel --outdir {dist_dir}",
+        shell=True,
+        cwd=os.path.join(repo_root, "sagemaker-core"),
+        check=True,
+    )
+
+    matches = glob.glob(os.path.join(dist_dir, "sagemaker_core-*.whl"))
+    if not matches:
+        raise FileNotFoundError(f"No sagemaker-core wheel found in {dist_dir}")
+    wheel_path = matches[0]
+
+    s3_prefix = f"s3://{sagemaker_session.default_bucket()}/remote-function-test/wheels"
+    S3Uploader.upload(wheel_path, s3_prefix, sagemaker_session=sagemaker_session)
+
+    return s3_prefix, os.path.basename(wheel_path)
+
+
+@pytest.fixture(scope="module")
+def sagemaker_session():
+    import boto3
+    return Session(boto3.Session())
+
+
+@pytest.fixture(scope="module")
+def role(sagemaker_session):
+    import boto3
+    account_id = boto3.client("sts").get_caller_identity()["Account"]
+    return f"arn:aws:iam::{account_id}:role/Admin"
+
+
+@pytest.fixture(scope="module")
+def image_uri(sagemaker_session):
+    region = sagemaker_session.boto_region_name
+    return f"763104351884.dkr.ecr.{region}.amazonaws.com/pytorch-training:2.0.0-cpu-py310"
+
+
+@pytest.fixture(scope="module")
+def dev_sdk_pre_execution_commands(sagemaker_session):
+    """Build dev sagemaker-core wheel, upload to S3, and return pre_execution_commands."""
+    s3_prefix, wheel_name = _build_and_upload_core_wheel(sagemaker_session)
+    cp_version = cloudpickle.__version__
+    return [
+        f"pip install cloudpickle=={cp_version}",
+        f"aws s3 cp {s3_prefix}/{wheel_name} /tmp/{wheel_name}",
+        f"pip install /tmp/{wheel_name}",
+    ]

sagemaker-core/tests/integ/remote_function/test_sagemaker_dependency_injection.py

@@ -9,12 +9,6 @@
 import tempfile
 import pytest
 
-# Skip decorator for AWS configuration
-# skip_if_no_aws_region = pytest.mark.skipif(
-#     not os.environ.get('AWS_DEFAULT_REGION'),
-#     reason="AWS credentials not configured"
-# )
-
 # Add src to path
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../../../src'))
 
@@ -25,58 +19,47 @@ class TestRemoteFunctionDependencyInjection:
     """Integration tests for dependency injection in remote functions."""
 
     @pytest.mark.integ
-    # @skip_if_no_aws_region
-    def test_remote_function_without_dependencies(self):
-        """Test remote function execution without explicit dependencies.
-        
-        This test verifies that when no dependencies are provided, the remote
-        function still executes successfully because sagemaker>=3.2.0 is
-        automatically injected.
-        """
+    def test_remote_function_without_dependencies(
+        self, dev_sdk_pre_execution_commands, role, image_uri, sagemaker_session
+    ):
+        """Test remote function execution without explicit dependencies."""
         @remote(
             instance_type="ml.m5.large",
-            # No dependencies specified - sagemaker should be injected automatically
+            role=role,
+            image_uri=image_uri,
+            sagemaker_session=sagemaker_session,
+            pre_execution_commands=dev_sdk_pre_execution_commands,
         )
         def simple_add(x, y):
-            """Simple function that adds two numbers."""
             return x + y
-        
-        # Execute the function
+
         result = simple_add(5, 3)
-        
-        # Verify result
         assert result == 8, f"Expected 8, got {result}"
-        print("✓ Remote function without dependencies executed successfully")
 
     @pytest.mark.integ
-    # @skip_if_no_aws_region
-    def test_remote_function_with_user_dependencies_no_sagemaker(self):
-        """Test remote function with user dependencies but no sagemaker.
-        
-        This test verifies that when user provides dependencies without sagemaker,
-        sagemaker>=3.2.0 is automatically appended.
-        """
-        # Create a temporary requirements.txt without sagemaker
+    def test_remote_function_with_user_dependencies_no_sagemaker(
+        self, dev_sdk_pre_execution_commands, role, image_uri, sagemaker_session
+    ):
+        """Test remote function with user dependencies but no sagemaker."""
         with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
             f.write("numpy>=1.20.0\npandas>=1.3.0\n")
             req_file = f.name
-        
+
         try:
             @remote(
                 instance_type="ml.m5.large",
+                role=role,
+                image_uri=image_uri,
+                sagemaker_session=sagemaker_session,
                 dependencies=req_file,
+                pre_execution_commands=dev_sdk_pre_execution_commands,
             )
             def compute_with_numpy(x):
-                """Function that uses numpy."""
                 import numpy as np
                 return np.array([x, x*2, x*3]).sum()
-            
-            # Execute the function
+
             result = compute_with_numpy(5)
-            
-            # Verify result (5 + 10 + 15 = 30)
             assert result == 30, f"Expected 30, got {result}"
-            print("✓ Remote function with user dependencies executed successfully")
         finally:
             os.remove(req_file)
 
@@ -85,52 +68,55 @@ class TestRemoteFunctionVersionCompatibility:
     """Tests for version compatibility between local and remote environments."""
 
     @pytest.mark.integ
-    # @skip_if_no_aws_region
-    def test_deserialization_with_injected_sagemaker(self):
-        """Test that deserialization works with injected sagemaker dependency.
-        
-        This test verifies that the remote environment can properly deserialize
-        functions when sagemaker>=3.2.0 is available.
-        """
+    def test_deserialization_with_injected_sagemaker(
+        self, dev_sdk_pre_execution_commands, role, image_uri, sagemaker_session
+    ):
+        """Test that deserialization works with injected sagemaker dependency."""
         @remote(
             instance_type="ml.m5.large",
+            role=role,
+            image_uri=image_uri,
+            sagemaker_session=sagemaker_session,
+            pre_execution_commands=dev_sdk_pre_execution_commands,
         )
         def complex_computation(data):
-            """Function that performs complex computation."""
             result = sum(data) * len(data)
             return result
-        
-        # Execute with various data types
+
         test_data = [1, 2, 3, 4, 5]
         result = complex_computation(test_data)
-        
-        # Verify result (sum=15, len=5, 15*5=75)
         assert result == 75, f"Expected 75, got {result}"
-        print("✓ Deserialization with injected sagemaker works correctly")
 
     @pytest.mark.integ
-    # @skip_if_no_aws_region
-    def test_multiple_remote_functions_with_dependencies(self):
-        """Test multiple remote functions with different dependency configurations.
-        
-        This test verifies that the dependency injection works correctly
-        when multiple remote functions are defined and executed.
-        """
-        @remote(instance_type="ml.m5.large")
+    def test_multiple_remote_functions_with_dependencies(
+        self, dev_sdk_pre_execution_commands, role, image_uri, sagemaker_session
+    ):
+        """Test multiple remote functions with different dependency configurations."""
+        @remote(
+            instance_type="ml.m5.large",
+            role=role,
+            image_uri=image_uri,
+            sagemaker_session=sagemaker_session,
+            pre_execution_commands=dev_sdk_pre_execution_commands,
+        )
         def func1(x):
             return x + 1
-        
-        @remote(instance_type="ml.m5.large")
+
+        @remote(
+            instance_type="ml.m5.large",
+            role=role,
+            image_uri=image_uri,
+            sagemaker_session=sagemaker_session,
+            pre_execution_commands=dev_sdk_pre_execution_commands,
+        )
         def func2(x):
             return x * 2
-        
-        # Execute both functions
+
         result1 = func1(5)
         result2 = func2(5)
-        
+
         assert result1 == 6, f"func1: Expected 6, got {result1}"
         assert result2 == 10, f"func2: Expected 10, got {result2}"
-        print("✓ Multiple remote functions with dependencies executed successfully")
 
 
 if __name__ == "__main__":