aws / aws/sagemaker-python-sdk
[v3] FrameworkProcessor and ModelTrainer: 4 regressions (including dropping CodeArtifact support) from v2 migration
- 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)
**Additional context**
The following four issues were found during a migration from `sagemaker==2.257.1` to `sagemaker==3.7.1`, moving training jobs from `sagemaker.pytorch.PyTorch` to `sagemaker.train.ModelTrainer` and processing jobs from `sagemaker.pytorch.processing.PyTorchProcessor` to `sagemaker.core.processing.FrameworkProcessor`.
---
**Describe the bug**
Four regressions found when migrating from v2 to v3, specifically around `FrameworkProcessor` (processing jobs) and `ModelTrainer` (training jobs). All four worked correctly in v2.
**System information**
- **SageMaker Python SDK version**: 3.7.1
- **Framework name**: PyTorch
- **Framework version**: 2.10
- **Python version**: 3.13
- **CPU or GPU**: Both
- **Custom Docker image**: N
---
### Bug 1: `wait=True` does not respect sagemaker session
**Affects:** `ModelTrainer.train(wait=True)` and `FrameworkProcessor.run(wait=True)`
`ProcessingJob.refresh()` and `TrainingJob.refresh()` use `Base.get_sagemaker_client()` — a global/default client — instead of the `sagemaker_session` passed to the processor/trainer. This fails with `NoCredentialsError` when using assumed-role sessions (via STS).
**To reproduce:**
```python
import boto3
from sagemaker.core.helper.session_helper import Session
from sagemaker.core.processing import FrameworkProcessor
from sagemaker.train import ModelTrainer
from sagemaker.core.training.configs import Compute, SourceCode
# Assumed-role session
sts = boto3.client("sts")
assumed = sts.assume_role(RoleArn="arn:aws:iam::123456789:role/MyRole", RoleSessionName="test")
creds = assumed["Credentials"]
assumed_session = boto3.Session(
aws_access_key_id=creds["AccessKeyId"],
aws_secret_access_key=creds["SecretAccessKey"],
aws_session_token=creds["SessionToken"],
region_name="us-west-2",
)
sm_session = Session(boto_session=assumed_session)
# FrameworkProcessor — job created OK, wait fails
processor = FrameworkProcessor(
image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-cpu-py313",
command=["python3"],
role="arn:aws:iam::123456789:role/MyRole",
instance_count=1,
instance_type="ml.m5.xlarge",
sagemaker_session=sm_session,
)
processor.run(code="my_script.py", source_dir="src", wait=True)
# → NoCredentialsError
# ModelTrainer — same issue
trainer = ModelTrainer(
training_image="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-cpu-py313",
role="arn:aws:iam::123456789:role/MyRole",
source_code=SourceCode(entry_script="train.py", source_dir="src"),
compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
sagemaker_session=sm_session,
)
trainer.train(wait=True)
# → NoCredentialsError
```
**Root cause** (`sagemaker/core/resources.py`):
```python
# ProcessingJob.refresh() / TrainingJob.refresh()
client = Base.get_sagemaker_client() # ← ignores the session
response = client.describe_processing_job(**operation_input_args)
```
**v2 behaviour** — `ProcessingJob.wait()` used the session directly:
```python
def wait(self, logs=True):
if logs:
self.sagemaker_session.logs_for_processing_job(self.job_name, wait=True)
else:
self.sagemaker_session.wait_for_processing_job(self.job_name)
```
---
### Bug 2: `FrameworkProcessor.code_location` is accepted but ignored
`FrameworkProcessor.__init__` accepts `code_location` and stores it as `self.code_location`. The docstring states it controls where code is uploaded. However, `_package_code` ignores it and always uploads to `self.sagemaker_session.default_bucket()`.
**To reproduce:**
```python
processor = FrameworkProcessor(
image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-cpu-py313",
command=["python3"],
role="arn:aws:iam::123456789:role/MyRole",
instance_count=1,
instance_type="ml.m5.xlarge",
code_location="s3://my-custom-bucket", # ← ignored
)
processor.run(code="my_script.py", source_dir="src", wait=False)
# Code uploads to s3://sagemaker-us-west-2-123456789/... instead of s3://my-custom-bucket/...
```
**Root cause** (`sagemaker-core/src/sagemaker/core/processing.py`, `_package_code`):
```python
s3_uri = s3.s3_path_join(
"s3://",
self.sagemaker_session.default_bucket(), # ← always uses default bucket
self.sagemaker_session.default_bucket_prefix or "",
job_name, "source", "sourcedir.tar.gz",
)
```
`self.code_location` is never referenced.
**v2 behaviour** — `FrameworkProcessor` delegated to an estimator that honored `code_location`:
```python
# v2 FrameworkProcessor._create_estimator
return self.estimator_cls(
...
code_location=self.code_location,
...
)
```
**Note:** `ModelTrainer` does not offer `code_location` at all — it always uses `session.default_bucket()`. Suggested fix: either remove `code_location` from `FrameworkProcessor` to align with `ModelTrainer`, or update `_package_code` to use it when set.
---
Bugs 3 and 4 are both about losing the ability to install `requirements.txt` dependencies from a CodeArtifact repository.
1. PyTorch training and inference containers supported this via the `CA_REPOSITORY_ARN` environment variable (see https://github.com/aws/deep-learning-containers/issues/2509 for details)
2. https://github.com/aws/sagemaker-python-sdk/pull/4145 extended that support to processing jobs by exposing `codeartifact_repo_arn` on `FrameworkProcessor.run()`
For context, the new ray-based inference containers are also adding codeartifact support via `CA_REPOSITORY_ARN` environment variable as can be seen in https://github.com/aws/deep-learning-containers/blob/0fc07f317a4db68ff728274070fbe332dde1ca26/scripts/ray/sagemaker_serve.py#L68
---
### Bug 3: CodeArtifact support missing from `FrameworkProcessor`
In v2, `FrameworkProcessor.run()` accepted `codeartifact_repo_arn` (added in [PR #4145](https://github.com/aws/sagemaker-python-sdk/pull/4145)). This configured pip inside the container to authenticate with CodeArtifact before installing `requirements.txt`.
In v3, `FrameworkProcessor.run()` does not accept `codeartifact_repo_arn`, and `_generate_framework_script` has no CodeArtifact support. The generated `runproc.sh` runs `pip install -r requirements.txt` without authentication, failing for packages hosted on private CodeArtifact repositories.
**To reproduce:**
```python
processor = FrameworkProcessor(
image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-cpu-py313",
command=["python3"],
role="arn:aws:iam::123456789:role/MyRole",
instance_count=1,
instance_type="ml.m5.xlarge",
)
# v2 supported: processor.run(..., codeartifact_repo_arn="arn:aws:codeartifact:us-west-2:123:repository/domain/repo")
# v3 does not — parameter doesn't exist
processor.run(code="my_script.py", source_dir="src", wait=False)
# Container fails: pip install -r requirements.txt → package not found (private CodeArtifact repo)
```
**v2 behaviour** — `_generate_framework_script` injected CodeArtifact login into `runproc.sh`:
```bash
if [[ -f 'requirements.txt' ]]; then
if ! hash aws 2>/dev/null; then
echo "AWS CLI is not installed. Skipping CodeArtifact login."
else
aws codeartifact login --tool pip --domain {domain} --domain-owner {owner} --repository {repository} --region {region}
fi
pip install -r requirements.txt
fi
```
**Suggested fix:** Port `codeartifact_repo_arn` and `_get_codeartifact_command` from v2's [PR #4145](https://github.com/aws/sagemaker-python-sdk/pull/4145/files) into v3's `FrameworkProcessor`.
---
### Bug 4: `ModelTrainer` bypasses `sagemaker-training-toolkit`, losing CodeArtifact support for `requirements.txt`
**Affects:** `ModelTrainer.train()` with `SourceCode(requirements="requirements.txt")`
`ModelTrainer` overrides the container's `ENTRYPOINT` with its own `sm_train.sh` driver script. This bypasses the `sagemaker-training-toolkit` installed in the container, which handled `requirements.txt` installation with CodeArtifact support (via the `CA_REPOSITORY_ARN` environment variable, added in [sagemaker-training-toolkit#187](https://github.com/aws/sagemaker-training-toolkit/pull/187)).
The v3 `sm_train.sh` template for requirements installation is a bare `pip install`:
```bash
# from sagemaker.train.templates.INSTALL_REQUIREMENTS
echo "Installing requirements"
$SM_PIP_CMD install -r {requirements_file}
```
It does not check for `CA_REPOSITORY_ARN` or configure pip to use CodeArtifact. The env var is passed to the container but nothing reads it.
**To reproduce:**
```python
from sagemaker.train import ModelTrainer
from sagemaker.core.training.configs import Compute, SourceCode
trainer = ModelTrainer(
training_image="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-cpu-py313",
role="arn:aws:iam::123456789:role/MyRole",
source_code=SourceCode(
entry_script="train.py",
source_dir="src",
requirements="requirements.txt", # ← installed without CodeArtifact
),
compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
environment={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:123:repository/domain/repo"},
)
trainer.train(wait=False)
# Container runs: pip install -r requirements.txt (using public PyPI, not CodeArtifact)
# Fails in VPC-isolated environments where PyPI is unreachable
```
**v2 behaviour** — the `PyTorch` estimator used the container's native entrypoint, which invoked `sagemaker-training-toolkit`. The toolkit checked `CA_REPOSITORY_ARN`, ran `aws codeartifact login --tool pip`, then installed requirements. This was added in [sagemaker-training-toolkit v4.7.0](https://github.com/aws/sagemaker-training-toolkit/releases/tag/v4.7.0) and tracked in [deep-learning-containers#2509](https://github.com/aws/deep-learning-containers/issues/2509).
**Suggested fix:** The `INSTALL_REQUIREMENTS` template in `sagemaker.train.templates` should check for `CA_REPOSITORY_ARN` and configure pip accordingly before installing, matching the behaviour of `sagemaker-training-toolkit`.
Contributor guide
Assessment
This issue has not been assessed yet.