Add sagemaker dependency for remote function by default

Author: aviruthen

src/sagemaker/remote_function/job.py

@@ -1235,6 +1235,11 @@ def _generate_input_data_config(job_settings: _JobSettings, s3_base_uri: str):
 
     local_dependencies_path = RuntimeEnvironmentManager().snapshot(job_settings.dependencies)
 
+    # Ensure sagemaker dependency is included to prevent version mismatch issues
+    # Resolves issue where computing hash for integrity check changed in 2.256.0
+    local_dependencies_path = _ensure_sagemaker_dependency(local_dependencies_path)
+    job_settings.dependencies = local_dependencies_path
+
     if step_compilation_context:
         with _tmpdir() as tmp_dir:
             script_and_dependencies_s3uri = _prepare_dependencies_and_pre_execution_scripts(
@@ -1291,6 +1296,113 @@ def _generate_input_data_config(job_settings: _JobSettings, s3_base_uri: str):
     return input_data_config
 
 
+def _check_sagemaker_version_compatibility(sagemaker_requirement: str) -> None:
+    """Check if the sagemaker version requirement uses incompatible hashing.
+
+    Raises ValueError if the requirement would install a version that uses HMAC hashing
+    (which is incompatible with the current SHA256-based integrity checks).
+
+    Args:
+        sagemaker_requirement: The sagemaker requirement string (e.g., "sagemaker>=2.200.0")
+
+    Raises:
+        ValueError: If the requirement would install a version using HMAC hashing
+    """
+    import re
+    from packaging.specifiers import SpecifierSet
+    from packaging import version as pkg_version
+
+    match = re.search(r'sagemaker\s*(.+)$', sagemaker_requirement.strip(), re.IGNORECASE)
+    if not match:
+        return
+
+    specifier_str = match.group(1).strip()
+    
+    try:
+        specifier_set = SpecifierSet(specifier_str)
+    except Exception:
+        return
+
+    # Test if any HMAC version would satisfy the specifier
+    # V2 HMAC versions: < 2.256.0
+    v2_hmac_test_versions = ["2.0.0", "2.100.0", "2.200.0", "2.255.0", "2.255.1", "2.255.99"]
+    for test_version in v2_hmac_test_versions:
+        if test_version in specifier_set:
+            raise ValueError(
+                f"The sagemaker version specified in requirements.txt ({sagemaker_requirement}) "
+                f"could install a version using HMAC-based integrity checks which are incompatible "
+                f"with the current SHA256-based integrity checks. Please update to "
+                f"sagemaker>=2.256.0,<3.0.0 (for V2) or sagemaker>=3.2.0,<4.0.0 (for V3)."
+            )
+    
+    # V3 HMAC versions: < 3.2.0
+    v3_hmac_test_versions = ["3.0.0", "3.0.1", "3.1.0", "3.1.99"]
+    for test_version in v3_hmac_test_versions:
+        if test_version in specifier_set:
+            raise ValueError(
+                f"The sagemaker version specified in requirements.txt ({sagemaker_requirement}) "
+                f"could install a version using HMAC-based integrity checks which are incompatible "
+                f"with the current SHA256-based integrity checks. Please update to "
+                f"sagemaker>=2.256.0,<3.0.0 (for V2) or sagemaker>=3.2.0,<4.0.0 (for V3)."
+            )
+
+
+def _ensure_sagemaker_dependency(local_dependencies_path: str) -> str:
+    """Ensure sagemaker>=2.256.0 is in the dependencies.
+
+    This function ensures that the remote environment has a compatible version of sagemaker
+    that includes the fix for the HMAC key security issue. Versions < 2.256.0 use HMAC-based
+    integrity checks which require the REMOTE_FUNCTION_SECRET_KEY environment variable.
+    Versions >= 2.256.0 use SHA256-based integrity checks which are secure and don't require
+    the secret key.
+
+    If no dependencies are provided, creates a temporary requirements.txt with sagemaker.
+    If dependencies are provided, appends sagemaker if not already present.
+
+    Args:
+        local_dependencies_path: Path to user's dependencies file or None
+
+    Returns:
+        Path to the dependencies file (created or modified)
+
+    Raises:
+        ValueError: If user has pinned sagemaker to a version using HMAC hashing
+    """
+    import tempfile
+
+    SAGEMAKER_MIN_VERSION = "sagemaker>=2.256.0,<3.0.0"
+
+    if local_dependencies_path is None:
+        # Create a temporary requirements.txt in the system temp directory
+        # This avoids overwriting any user files in their working directory
+        fd, req_file = tempfile.mkstemp(suffix=".txt", prefix="sagemaker_requirements_")
+        os.close(fd)  # Close the file descriptor, we'll write to it ourselves
+
+        with open(req_file, "w") as f:
+            f.write(f"{SAGEMAKER_MIN_VERSION}\n")
+        logger.info("Created temporary requirements.txt at %s with %s", req_file, SAGEMAKER_MIN_VERSION)
+        return req_file
+
+    # If dependencies provided, ensure sagemaker is included
+    if local_dependencies_path.endswith(".txt"):
+        with open(local_dependencies_path, "r") as f:
+            content = f.read()
+
+        # Check if sagemaker is already specified
+        if "sagemaker" in content.lower():
+            # Extract the sagemaker requirement line for compatibility check
+            for line in content.split('\n'):
+                if 'sagemaker' in line.lower():
+                    _check_sagemaker_version_compatibility(line.strip())
+                    break
+        else:
+            with open(local_dependencies_path, "a") as f:
+                f.write(f"\n{SAGEMAKER_MIN_VERSION}\n")
+            logger.info("Appended %s to requirements.txt", SAGEMAKER_MIN_VERSION)
+
+    return local_dependencies_path
+
+
 def _prepare_dependencies_and_pre_execution_scripts(
     local_dependencies_path: str,
     pre_execution_commands: List[str],

tests/integ/sagemaker/remote_function/test_sagemaker_dependency_injection.py

@@ -0,0 +1,127 @@
+"""Integration tests for sagemaker dependency injection in remote functions.
+
+These tests verify that the sagemaker>=2.256.0 dependency is properly injected
+into remote function jobs, preventing version mismatch issues.
+"""
+
+import os
+import sys
+import tempfile
+import pytest
+
+# Add src to path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../../../src'))
+
+from sagemaker.remote_function import remote
+
+
+class TestRemoteFunctionDependencyInjection:
+    """Integration tests for dependency injection in remote functions."""
+
+    @pytest.mark.integ
+    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>=2.256.0 is
+        automatically injected.
+        """
+        @remote(
+            instance_type="ml.m5.large",
+            # No dependencies specified - sagemaker should be injected automatically
+        )
+        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
+    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>=2.256.0 is automatically appended.
+        """
+        # Create a temporary requirements.txt without 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",
+                dependencies=req_file,
+            )
+            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)
+
+
+class TestRemoteFunctionVersionCompatibility:
+    """Tests for version compatibility between local and remote environments."""
+
+    @pytest.mark.integ
+    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>=2.256.0 is available.
+        """
+        @remote(
+            instance_type="ml.m5.large",
+        )
+        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
+    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 func1(x):
+            return x + 1
+        
+        @remote(instance_type="ml.m5.large")
+        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__":
+    pytest.main([__file__, "-v", "-m", "integ"])