aws / aws/sagemaker-python-sdk

@remote decorator forces sagemaker upgrade via pip install -U, causing DeserializationError

Open
#5,872 0 comments 0 reactions 0 assignees View on GitHub
remote-function
Dominant language
Python
Stars
2.3k
Forks
1.3k
Avg merge
1d 22h
Merged PRs (30d)
35

Description

PySDK Version
- [ ] PySDK V2 (2.x)
- [x] PySDK V3 (3.x)

**Describe the bug**

When using the `@remote` decorator with dependencies=None, the SDK creates a temporary requirements.txt containing sagemaker>=3.2.0,<4.0.0 and installs it in the container with pip install -r ... -U. The -U (upgrade) flag forces pip to upgrade sagemaker to the latest version in the range, even if a compatible version is already installed. This creates a version mismatch between the client (which serialized the function with 3.5.0) and the container (which deserializes it with 3.11.0), resulting in a DeserializationError.

Additionally, when a dependencies file is provided but does not contain sagemaker, ensure_sagemaker_dependency appends sagemaker>=3.2.0,<4.0.0 to it. Combined with the -U flag, this also forces an upgrade even if the container already has a compatible version installed.

The root cause is two-fold:

1. _ensure_sagemaker_dependency (in sagemaker/core/remote_function/job.py) creates a temp requirements file with a loose range (sagemaker>=3.2.0,<4.0.0) instead of pinning the client's exact version.

Source: _ensure_sagemaker_dependency

```python

def _ensure_sagemaker_dependency(local_dependencies_path: str) -> str:
"""Ensure sagemaker>=3.2.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 < 3.2.0 use HMAC-based
integrity checks which require the REMOTE_FUNCTION_SECRET_KEY environment variable.
Versions >= 3.2.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>=3.2.0,<4.0.0"

if local_dependencies_path is None:
fd, req_file = tempfile.mkstemp(suffix=".txt", prefix="sagemaker_requirements_")
os.close(fd)

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 local_dependencies_path.endswith(".txt"):
with open(local_dependencies_path, "r") as f:
content = f.read()

if "sagemaker" in content.lower():
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
```

2. _install_requirements_txt (in sagemaker/core/remote_function/runtime_environment/runtime_environment_manager.py) hardcodes the -U flag:
```python
cmd = [python_executable, "-m", "pip", "install", "-r", validated_path, "-U"]
```

Source: _install_requirements_txt

```python
def _install_requirements_txt(self, local_path, python_executable):
"""Install requirements.txt file"""
# Validate path to prevent command injection
validated_path = self._validate_path(local_path)
cmd = [python_executable, "-m", "pip", "install", "-r", validated_path, "-U"]
logger.info("Running command: '%s' in the dir: '%s' ", " ".join(cmd), os.getcwd())
_run_shell_cmd(cmd)
logger.info("Command %s ran successfully", " ".join(cmd))
```

**To reproduce**

```python
from sagemaker.core.remote_function import remote
from sagemaker.core.helper.session_helper import Session, get_execution_role

BUCKET = "my-bucket"

@remote(
role=get_execution_role(),
instance_type="ml.m5.large",
sagemaker_session=Session(default_bucket=BUCKET),
s3_root_uri=f"s3://{BUCKET}/remote",
image_uri="763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.5-cpu-py311",
dependencies=None,
)
def hello():
return "hello world"

hello()
```

Even explicitly installing the correct version via pre_execution_commands doesn't help, since bootstrap runs after and overwrites it:

```python
@remote(
role=get_execution_role(),
instance_type="ml.m5.large",
sagemaker_session=Session(default_bucket=BUCKET),
s3_root_uri=f"s3://{BUCKET}/remote",
image_uri="763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.5-cpu-py311",
dependencies=None,
pre_execution_commands=["pip install sagemaker==3.5.0"],
)
def hello():
return "hello world"

hello()
```

**Expected behavior**

The container should use the same sagemaker version as the client that serialized the function. A version mismatch in the serialization layer should never be introduced silently by the SDK itself.

Suggested fixes:
1. Pin the exact client version in the generated requirements (sagemaker== — the version is already known via --client_sagemaker_pysdk_version)
2. Remove the -U flag from _install_requirements_txt
3. Both

**Screenshots or logs**

```markdown
sagemaker.remote_function INFO Running command: '/opt/conda/bin/python -m pip install -r /sagemaker_remote_function_workspace/sagemaker_requirements_4x1lpm7l.txt -U'
sagemaker.remote_function INFO Requirement already satisfied: sagemaker<4.0.0,>=3.2.0 in /opt/conda/lib/python3.11/site-packages (from -r /sagemaker_remote_function_workspace/sagemaker_requirements_4x1lpm7l.txt (line 1)) (3.5.0)
sagemaker.remote_function INFO Collecting sagemaker<4.0.0,>=3.2.0 (from -r /sagemaker_remote_function_workspace/sagemaker_requirements_4x1lpm7l.txt (line 1))
sagemaker.remote_function INFO Downloading sagemaker-3.11.0-py3-none-any.whl.metadata (20 kB)
sagemaker.remote_function INFO Found existing installation: sagemaker 3.5.0
sagemaker.remote_function INFO Uninstalling sagemaker-3.5.0:
sagemaker.remote_function INFO Successfully uninstalled sagemaker-3.5.0
sagemaker.remote_function INFO Successfully installed sagemaker-3.11.0
```

Followed by:
```markdown
DeserializationError: Integrity check for the serialized function or data failed. The payload metadata does not contain an asymmetric signature. Please upgrade your SageMaker Python SDK version.
```

**System information**
- **SageMaker Python SDK version**: 3.5.0
- **Framework name (eg. PyTorch) or algorithm (eg. KMeans)**: Any
- **Framework version**: -
- **Python version**: 3.11
- **CPU or GPU**: CPU
- **Custom Docker image (Y/N)**: N (AWS DLC pytorch-training:2.5-cpu-py311)

**Additional context**

- This affects both `@remote` and `@step` (SageMaker Pipelines), which share the same remote_function/job.py code path.
- **Workaround**: Provide a dependencies file containing the pinned sagemaker version (e.g. sagemaker==3.5.0). This prevents _ensure_sagemaker_dependency from injecting the loose range.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.