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

Research direction

Start by reading _ensure_sagemaker_dependency in sagemaker/core/remote_function/job.py and _install_requirements_txt in sagemaker/core/remote_function/runtime_environment/runtime_environment_manager.py. Trace how the client SDK version reaches dependency generation and how requirements are installed. The work is done when remote environments retain the compatible client version without an unintended upgrade or serialization mismatch.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, python
Domain
cloud, machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.