feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation - #5772

Merged
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix
Apr 28, 2026
Merged

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation#5772
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix

Conversation

@humanzz

@humanzzhumanzz commented Apr 17, 2026

Copy link
Copy Markdown
Contributor

SDK v3's ModelTrainer and FrameworkProcessor override the container entrypoint with SDK-generated scripts (sm_train.sh, runproc.sh), bypassing the container's entrypoint which involved sagemaker-training-toolkit handling
CA_REPOSITORY_ARN-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in #5765.

This is the stopgap solution proposed in this comment: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.

  • Add install_requirements.py in sagemaker-core — reads CA_REPOSITORY_ARN from container environment; no-op if unset
  • Try boto3 first (matching sagemaker-training-toolkit), fall back to AWS CLI, hard-fail if neither is available
  • Wire into ModelTrainer: copy script into sm_drivers/scripts/, update INSTALL_REQUIREMENTS templates to call it instead of bare pip install
  • Wire into FrameworkProcessor: upload script as sibling file alongside runproc.sh, update generated script to call it

Issue #, if available:

#5765

Description of changes:

When the SDK overrides a container's entrypoint — as ModelTrainer does for training jobs (Bug 4) and FrameworkProcessor does for processing jobs (Bug 3) — the container's native sagemaker-training-toolkit is bypassed. This toolkit handled CA_REPOSITORY_ARN-based CodeArtifact authentication for requirements.txt installation via boto3. Without it, pip install -r requirements.txt runs against public PyPI, failing in VPC-isolated environments or when packages are only available in a private CodeArtifact repository.

See #5765 and the detailed analysis comment for full context.

Solution: Stopgap install_requirements.py

A self-contained Python script in sagemaker-core that handles CodeArtifact authentication before installing requirements. It:

  1. Reads CA_REPOSITORY_ARN from the container environment — if not set, does a normal pip install
  2. Tries boto3 first (matching sagemaker-training-toolkit's approach) to build an authenticated pip index URL
  3. Falls back to AWS CLI (aws codeartifact login --tool pip) if boto3 is unavailable
  4. Hard-fails with a clear error if CA_REPOSITORY_ARN is set but neither boto3 nor AWS CLI is available

The script can be used as:

  • A standalone script: python install_requirements.py requirements.txt (used by bash-based entrypoints)
  • An importable module: from sagemaker.core.utils.install_requirements import configure_pip, install_requirements (for Python-native callers like @remote or ModelBuilder)

Changes

FileChange
sagemaker-core/.../utils/install_requirements.pyNew module with configure_pip(), install_requirements(), main(), and CodeArtifactAuthMethod enum
sagemaker-core/tests/unit/test_install_requirements.py22 unit tests covering all auth methods, fallback chains, error propagation
sagemaker-train/.../templates.pyINSTALL_REQUIREMENTS and INSTALL_AUTO_REQUIREMENTS now call install_requirements.py instead of bare pip install
sagemaker-train/.../model_trainer.pyCopy install_requirements.py from sagemaker-core into sm_drivers/scripts/ at runtime
sagemaker-core/.../processing.pyUpload install_requirements.py as sibling file alongside runproc.sh and sourcedir.tar.gz; update generated script to call it
sagemaker-core/tests/unit/test_processing.pyVerify install_requirements.py is uploaded and referenced in generated script

What this covers

Job TypeClassCodeArtifact with this PR
TrainingModelTrainer✅ Fixed — install_requirements.py in sm_drivers/scripts/
ProcessingFrameworkProcessor✅ Fixed — install_requirements.py uploaded as sibling file
TuningTuner✅ Already works — Tuner uses container's native toolkit (not affected by this PR)
InferenceModelBuilder✅ Already works — SDK doesn't override inference entrypoints

What this does NOT cover

PathStatusNotes
@remote function (runtime_environment_manager.py)❌ Not wiredHas its own _install_requirements_txt() that does bare pip install. Could use configure_pip() via import.
sagemaker-serve (requirements_manager.py)❌ Not wiredSame — bare pip install in-process. Could import configure_pip().
sagemaker-core/modules (templates.py)❌ Not wiredDuplicate of sagemaker-train/templates.py without INSTALL_REQUIREMENTS. Lower priority.

These are follow-up opportunities — the module is available for them to import.

Known risks

  1. Tuning jobs depend on the container's toolkit — The CreateHyperParameterTuningJob API uses HyperParameterAlgorithmSpecification which lacks ContainerEntrypoint, so the Tuner cannot use sm_train.sh. If future containers drop sagemaker-training-toolkit, tuning jobs will lose CodeArtifact support with no SDK-side fix possible until the API adds entrypoint support.

  2. boto3 availability in future containers — Current PyTorch training containers (2.7–2.9) include boto3. New DLC base images on the main branch do not. The script's fallback to AWS CLI mitigates this, but if neither is available, the script hard-fails. The long-term solution is a shared package with boto3 as a declared dependency (see analysis).

Long-term solution

This PR is a stopgap that works within the SDK alone. The long-term solution requires coordination between the SDK and DLC to ensure that both the container's default entrypoint and any SDK-overridden entrypoint have access to the same CodeArtifact-aware installer — ideally a shared package with boto3 as a declared dependency, installed in all SageMaker containers. See the proposed ideal solution for details.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

…r requirements.txt installation
SDK v3's `ModelTrainer` and `FrameworkProcessor` override the container entrypoint with SDK-generated scripts (`sm_train.sh`, `runproc.sh`), bypassing the container's entrypoint which involved `sagemaker-training-toolkit` handling
`CA_REPOSITORY_ARN`-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in aws#5765.
This is the stopgap solution proposed in this comment[aws#5765 (comment)]: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.
- Add `install_requirements.py` in sagemaker-core — reads `CA_REPOSITORY_ARN` from container environment; no-op if unset
- Try `boto3` first (matching sagemaker-training-toolkit), fall back to `AWS CLI`, hard-fail if neither is available
- Wire into `ModelTrainer`: copy script into `sm_drivers/scripts/`, update `INSTALL_REQUIREMENTS` templates to call it instead of bare `pip install`
- Wire into `FrameworkProcessor`: upload script as sibling file alongside `runproc.sh`, update generated script to call it
@humanzz

Copy link
Copy Markdown
ContributorAuthor

I've also tested this with my code (as an integration test) to verify the behaviours

Training Job (ModelTrainer)

fromsagemaker.trainimportModelTrainerfromsagemaker.core.training.configsimportCompute, SourceCodetrainer=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"
),
compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
sagemaker_session=sm_session,
environment={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
trainer.train(input_data_config=inputs, wait=False)

CloudWatch logsinstall_requirements.py ran from sm_drivers/scripts/, authenticated via boto3, pip resolved from CodeArtifact:

Installing requirements
++ /usr/local/bin/python3 /opt/ml/input/data/sm_drivers/scripts/install_requirements.py requirements.txt
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Training job completed successfully. ✅

Processing Job (FrameworkProcessor)

fromsagemaker.core.processingimportFrameworkProcessorprocessor=FrameworkProcessor(
image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-gpu-py313",
command=["python3"],
role="arn:aws:iam::123456789:role/MyRole",
instance_count=1,
instance_type="ml.g6.4xlarge",
sagemaker_session=sm_session,
env={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
processor.run(code="my_script.py", source_dir="src", wait=False)

CloudWatch logsinstall_requirements.py uploaded as sibling file, authenticated via boto3:

Files in /opt/ml/processing/input/code/ before extraction:
-rw-r--r-- 1 root root 6652 Apr 17 10:22 install_requirements.py
-rw-r--r-- 1 root root 685 Apr 17 10:22 runproc.sh
-rw-r--r-- 1 root root 81582 Apr 17 10:22 sourcedir.tar.gz
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Processing job completed successfully. ✅

@aviruthen
aviruthen merged commit 9386fc0 into aws:masterApr 28, 2026
37 of 53 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants

@humanzz@zhaoqizqwang@aviruthen@mollyheamazon
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Skip to content

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation - #5772

Merged
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix
Apr 28, 2026
Merged

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation#5772
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix

Conversation

@humanzz

@humanzzhumanzz commented Apr 17, 2026

Copy link
Copy Markdown
Contributor

SDK v3's ModelTrainer and FrameworkProcessor override the container entrypoint with SDK-generated scripts (sm_train.sh, runproc.sh), bypassing the container's entrypoint which involved sagemaker-training-toolkit handling
CA_REPOSITORY_ARN-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in #5765.

This is the stopgap solution proposed in this comment: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.

  • Add install_requirements.py in sagemaker-core — reads CA_REPOSITORY_ARN from container environment; no-op if unset
  • Try boto3 first (matching sagemaker-training-toolkit), fall back to AWS CLI, hard-fail if neither is available
  • Wire into ModelTrainer: copy script into sm_drivers/scripts/, update INSTALL_REQUIREMENTS templates to call it instead of bare pip install
  • Wire into FrameworkProcessor: upload script as sibling file alongside runproc.sh, update generated script to call it

Issue #, if available:

#5765

Description of changes:

When the SDK overrides a container's entrypoint — as ModelTrainer does for training jobs (Bug 4) and FrameworkProcessor does for processing jobs (Bug 3) — the container's native sagemaker-training-toolkit is bypassed. This toolkit handled CA_REPOSITORY_ARN-based CodeArtifact authentication for requirements.txt installation via boto3. Without it, pip install -r requirements.txt runs against public PyPI, failing in VPC-isolated environments or when packages are only available in a private CodeArtifact repository.

See #5765 and the detailed analysis comment for full context.

Solution: Stopgap install_requirements.py

A self-contained Python script in sagemaker-core that handles CodeArtifact authentication before installing requirements. It:

  1. Reads CA_REPOSITORY_ARN from the container environment — if not set, does a normal pip install
  2. Tries boto3 first (matching sagemaker-training-toolkit's approach) to build an authenticated pip index URL
  3. Falls back to AWS CLI (aws codeartifact login --tool pip) if boto3 is unavailable
  4. Hard-fails with a clear error if CA_REPOSITORY_ARN is set but neither boto3 nor AWS CLI is available

The script can be used as:

  • A standalone script: python install_requirements.py requirements.txt (used by bash-based entrypoints)
  • An importable module: from sagemaker.core.utils.install_requirements import configure_pip, install_requirements (for Python-native callers like @remote or ModelBuilder)

Changes

FileChange
sagemaker-core/.../utils/install_requirements.pyNew module with configure_pip(), install_requirements(), main(), and CodeArtifactAuthMethod enum
sagemaker-core/tests/unit/test_install_requirements.py22 unit tests covering all auth methods, fallback chains, error propagation
sagemaker-train/.../templates.pyINSTALL_REQUIREMENTS and INSTALL_AUTO_REQUIREMENTS now call install_requirements.py instead of bare pip install
sagemaker-train/.../model_trainer.pyCopy install_requirements.py from sagemaker-core into sm_drivers/scripts/ at runtime
sagemaker-core/.../processing.pyUpload install_requirements.py as sibling file alongside runproc.sh and sourcedir.tar.gz; update generated script to call it
sagemaker-core/tests/unit/test_processing.pyVerify install_requirements.py is uploaded and referenced in generated script

What this covers

Job TypeClassCodeArtifact with this PR
TrainingModelTrainer✅ Fixed — install_requirements.py in sm_drivers/scripts/
ProcessingFrameworkProcessor✅ Fixed — install_requirements.py uploaded as sibling file
TuningTuner✅ Already works — Tuner uses container's native toolkit (not affected by this PR)
InferenceModelBuilder✅ Already works — SDK doesn't override inference entrypoints

What this does NOT cover

PathStatusNotes
@remote function (runtime_environment_manager.py)❌ Not wiredHas its own _install_requirements_txt() that does bare pip install. Could use configure_pip() via import.
sagemaker-serve (requirements_manager.py)❌ Not wiredSame — bare pip install in-process. Could import configure_pip().
sagemaker-core/modules (templates.py)❌ Not wiredDuplicate of sagemaker-train/templates.py without INSTALL_REQUIREMENTS. Lower priority.

These are follow-up opportunities — the module is available for them to import.

Known risks

  1. Tuning jobs depend on the container's toolkit — The CreateHyperParameterTuningJob API uses HyperParameterAlgorithmSpecification which lacks ContainerEntrypoint, so the Tuner cannot use sm_train.sh. If future containers drop sagemaker-training-toolkit, tuning jobs will lose CodeArtifact support with no SDK-side fix possible until the API adds entrypoint support.

  2. boto3 availability in future containers — Current PyTorch training containers (2.7–2.9) include boto3. New DLC base images on the main branch do not. The script's fallback to AWS CLI mitigates this, but if neither is available, the script hard-fails. The long-term solution is a shared package with boto3 as a declared dependency (see analysis).

Long-term solution

This PR is a stopgap that works within the SDK alone. The long-term solution requires coordination between the SDK and DLC to ensure that both the container's default entrypoint and any SDK-overridden entrypoint have access to the same CodeArtifact-aware installer — ideally a shared package with boto3 as a declared dependency, installed in all SageMaker containers. See the proposed ideal solution for details.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

…r requirements.txt installation
SDK v3's `ModelTrainer` and `FrameworkProcessor` override the container entrypoint with SDK-generated scripts (`sm_train.sh`, `runproc.sh`), bypassing the container's entrypoint which involved `sagemaker-training-toolkit` handling
`CA_REPOSITORY_ARN`-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in aws#5765.
This is the stopgap solution proposed in this comment[aws#5765 (comment)]: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.
- Add `install_requirements.py` in sagemaker-core — reads `CA_REPOSITORY_ARN` from container environment; no-op if unset
- Try `boto3` first (matching sagemaker-training-toolkit), fall back to `AWS CLI`, hard-fail if neither is available
- Wire into `ModelTrainer`: copy script into `sm_drivers/scripts/`, update `INSTALL_REQUIREMENTS` templates to call it instead of bare `pip install`
- Wire into `FrameworkProcessor`: upload script as sibling file alongside `runproc.sh`, update generated script to call it
@humanzz

Copy link
Copy Markdown
ContributorAuthor

I've also tested this with my code (as an integration test) to verify the behaviours

Training Job (ModelTrainer)

fromsagemaker.trainimportModelTrainerfromsagemaker.core.training.configsimportCompute, SourceCodetrainer=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"
),
compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
sagemaker_session=sm_session,
environment={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
trainer.train(input_data_config=inputs, wait=False)

CloudWatch logsinstall_requirements.py ran from sm_drivers/scripts/, authenticated via boto3, pip resolved from CodeArtifact:

Installing requirements
++ /usr/local/bin/python3 /opt/ml/input/data/sm_drivers/scripts/install_requirements.py requirements.txt
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Training job completed successfully. ✅

Processing Job (FrameworkProcessor)

fromsagemaker.core.processingimportFrameworkProcessorprocessor=FrameworkProcessor(
image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-gpu-py313",
command=["python3"],
role="arn:aws:iam::123456789:role/MyRole",
instance_count=1,
instance_type="ml.g6.4xlarge",
sagemaker_session=sm_session,
env={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
processor.run(code="my_script.py", source_dir="src", wait=False)

CloudWatch logsinstall_requirements.py uploaded as sibling file, authenticated via boto3:

Files in /opt/ml/processing/input/code/ before extraction:
-rw-r--r-- 1 root root 6652 Apr 17 10:22 install_requirements.py
-rw-r--r-- 1 root root 685 Apr 17 10:22 runproc.sh
-rw-r--r-- 1 root root 81582 Apr 17 10:22 sourcedir.tar.gz
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Processing job completed successfully. ✅

@aviruthen
aviruthen merged commit 9386fc0 into aws:masterApr 28, 2026
37 of 53 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants

@humanzz@zhaoqizqwang@aviruthen@mollyheamazon
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix"); } } catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation - #5772

Merged
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix
Apr 28, 2026
Merged

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation#5772
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix

Conversation

@humanzz

@humanzzhumanzz commented Apr 17, 2026

Copy link
Copy Markdown
Contributor

SDK v3's ModelTrainer and FrameworkProcessor override the container entrypoint with SDK-generated scripts (sm_train.sh, runproc.sh), bypassing the container's entrypoint which involved sagemaker-training-toolkit handling
CA_REPOSITORY_ARN-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in #5765.

This is the stopgap solution proposed in this comment: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.

  • Add install_requirements.py in sagemaker-core — reads CA_REPOSITORY_ARN from container environment; no-op if unset
  • Try boto3 first (matching sagemaker-training-toolkit), fall back to AWS CLI, hard-fail if neither is available
  • Wire into ModelTrainer: copy script into sm_drivers/scripts/, update INSTALL_REQUIREMENTS templates to call it instead of bare pip install
  • Wire into FrameworkProcessor: upload script as sibling file alongside runproc.sh, update generated script to call it

Issue #, if available:

#5765

Description of changes:

When the SDK overrides a container's entrypoint — as ModelTrainer does for training jobs (Bug 4) and FrameworkProcessor does for processing jobs (Bug 3) — the container's native sagemaker-training-toolkit is bypassed. This toolkit handled CA_REPOSITORY_ARN-based CodeArtifact authentication for requirements.txt installation via boto3. Without it, pip install -r requirements.txt runs against public PyPI, failing in VPC-isolated environments or when packages are only available in a private CodeArtifact repository.

See #5765 and the detailed analysis comment for full context.

Solution: Stopgap install_requirements.py

A self-contained Python script in sagemaker-core that handles CodeArtifact authentication before installing requirements. It:

  1. Reads CA_REPOSITORY_ARN from the container environment — if not set, does a normal pip install
  2. Tries boto3 first (matching sagemaker-training-toolkit's approach) to build an authenticated pip index URL
  3. Falls back to AWS CLI (aws codeartifact login --tool pip) if boto3 is unavailable
  4. Hard-fails with a clear error if CA_REPOSITORY_ARN is set but neither boto3 nor AWS CLI is available

The script can be used as:

  • A standalone script: python install_requirements.py requirements.txt (used by bash-based entrypoints)
  • An importable module: from sagemaker.core.utils.install_requirements import configure_pip, install_requirements (for Python-native callers like @remote or ModelBuilder)

Changes

FileChange
sagemaker-core/.../utils/install_requirements.pyNew module with configure_pip(), install_requirements(), main(), and CodeArtifactAuthMethod enum
sagemaker-core/tests/unit/test_install_requirements.py22 unit tests covering all auth methods, fallback chains, error propagation
sagemaker-train/.../templates.pyINSTALL_REQUIREMENTS and INSTALL_AUTO_REQUIREMENTS now call install_requirements.py instead of bare pip install
sagemaker-train/.../model_trainer.pyCopy install_requirements.py from sagemaker-core into sm_drivers/scripts/ at runtime
sagemaker-core/.../processing.pyUpload install_requirements.py as sibling file alongside runproc.sh and sourcedir.tar.gz; update generated script to call it
sagemaker-core/tests/unit/test_processing.pyVerify install_requirements.py is uploaded and referenced in generated script

What this covers

Job TypeClassCodeArtifact with this PR
TrainingModelTrainer✅ Fixed — install_requirements.py in sm_drivers/scripts/
ProcessingFrameworkProcessor✅ Fixed — install_requirements.py uploaded as sibling file
TuningTuner✅ Already works — Tuner uses container's native toolkit (not affected by this PR)
InferenceModelBuilder✅ Already works — SDK doesn't override inference entrypoints

What this does NOT cover

PathStatusNotes
@remote function (runtime_environment_manager.py)❌ Not wiredHas its own _install_requirements_txt() that does bare pip install. Could use configure_pip() via import.
sagemaker-serve (requirements_manager.py)❌ Not wiredSame — bare pip install in-process. Could import configure_pip().
sagemaker-core/modules (templates.py)❌ Not wiredDuplicate of sagemaker-train/templates.py without INSTALL_REQUIREMENTS. Lower priority.

These are follow-up opportunities — the module is available for them to import.

Known risks

  1. Tuning jobs depend on the container's toolkit — The CreateHyperParameterTuningJob API uses HyperParameterAlgorithmSpecification which lacks ContainerEntrypoint, so the Tuner cannot use sm_train.sh. If future containers drop sagemaker-training-toolkit, tuning jobs will lose CodeArtifact support with no SDK-side fix possible until the API adds entrypoint support.

  2. boto3 availability in future containers — Current PyTorch training containers (2.7–2.9) include boto3. New DLC base images on the main branch do not. The script's fallback to AWS CLI mitigates this, but if neither is available, the script hard-fails. The long-term solution is a shared package with boto3 as a declared dependency (see analysis).

Long-term solution

This PR is a stopgap that works within the SDK alone. The long-term solution requires coordination between the SDK and DLC to ensure that both the container's default entrypoint and any SDK-overridden entrypoint have access to the same CodeArtifact-aware installer — ideally a shared package with boto3 as a declared dependency, installed in all SageMaker containers. See the proposed ideal solution for details.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

…r requirements.txt installation
SDK v3's `ModelTrainer` and `FrameworkProcessor` override the container entrypoint with SDK-generated scripts (`sm_train.sh`, `runproc.sh`), bypassing the container's entrypoint which involved `sagemaker-training-toolkit` handling
`CA_REPOSITORY_ARN`-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in aws#5765.
This is the stopgap solution proposed in this comment[aws#5765 (comment)]: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.
- Add `install_requirements.py` in sagemaker-core — reads `CA_REPOSITORY_ARN` from container environment; no-op if unset
- Try `boto3` first (matching sagemaker-training-toolkit), fall back to `AWS CLI`, hard-fail if neither is available
- Wire into `ModelTrainer`: copy script into `sm_drivers/scripts/`, update `INSTALL_REQUIREMENTS` templates to call it instead of bare `pip install`
- Wire into `FrameworkProcessor`: upload script as sibling file alongside `runproc.sh`, update generated script to call it
@humanzz

Copy link
Copy Markdown
ContributorAuthor

I've also tested this with my code (as an integration test) to verify the behaviours

Training Job (ModelTrainer)

fromsagemaker.trainimportModelTrainerfromsagemaker.core.training.configsimportCompute, SourceCodetrainer=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"
),
compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
sagemaker_session=sm_session,
environment={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
trainer.train(input_data_config=inputs, wait=False)

CloudWatch logsinstall_requirements.py ran from sm_drivers/scripts/, authenticated via boto3, pip resolved from CodeArtifact:

Installing requirements
++ /usr/local/bin/python3 /opt/ml/input/data/sm_drivers/scripts/install_requirements.py requirements.txt
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Training job completed successfully. ✅

Processing Job (FrameworkProcessor)

fromsagemaker.core.processingimportFrameworkProcessorprocessor=FrameworkProcessor(
image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-gpu-py313",
command=["python3"],
role="arn:aws:iam::123456789:role/MyRole",
instance_count=1,
instance_type="ml.g6.4xlarge",
sagemaker_session=sm_session,
env={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
processor.run(code="my_script.py", source_dir="src", wait=False)

CloudWatch logsinstall_requirements.py uploaded as sibling file, authenticated via boto3:

Files in /opt/ml/processing/input/code/ before extraction:
-rw-r--r-- 1 root root 6652 Apr 17 10:22 install_requirements.py
-rw-r--r-- 1 root root 685 Apr 17 10:22 runproc.sh
-rw-r--r-- 1 root root 81582 Apr 17 10:22 sourcedir.tar.gz
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Processing job completed successfully. ✅

@aviruthen
aviruthen merged commit 9386fc0 into aws:masterApr 28, 2026
37 of 53 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants

@humanzz@zhaoqizqwang@aviruthen@mollyheamazon
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms"); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation - #5772

Merged
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix
Apr 28, 2026
Merged

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation#5772
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix

Conversation

@humanzz

@humanzzhumanzz commented Apr 17, 2026

Copy link
Copy Markdown
Contributor

SDK v3's ModelTrainer and FrameworkProcessor override the container entrypoint with SDK-generated scripts (sm_train.sh, runproc.sh), bypassing the container's entrypoint which involved sagemaker-training-toolkit handling
CA_REPOSITORY_ARN-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in #5765.

This is the stopgap solution proposed in this comment: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.

  • Add install_requirements.py in sagemaker-core — reads CA_REPOSITORY_ARN from container environment; no-op if unset
  • Try boto3 first (matching sagemaker-training-toolkit), fall back to AWS CLI, hard-fail if neither is available
  • Wire into ModelTrainer: copy script into sm_drivers/scripts/, update INSTALL_REQUIREMENTS templates to call it instead of bare pip install
  • Wire into FrameworkProcessor: upload script as sibling file alongside runproc.sh, update generated script to call it

Issue #, if available:

#5765

Description of changes:

When the SDK overrides a container's entrypoint — as ModelTrainer does for training jobs (Bug 4) and FrameworkProcessor does for processing jobs (Bug 3) — the container's native sagemaker-training-toolkit is bypassed. This toolkit handled CA_REPOSITORY_ARN-based CodeArtifact authentication for requirements.txt installation via boto3. Without it, pip install -r requirements.txt runs against public PyPI, failing in VPC-isolated environments or when packages are only available in a private CodeArtifact repository.

See #5765 and the detailed analysis comment for full context.

Solution: Stopgap install_requirements.py

A self-contained Python script in sagemaker-core that handles CodeArtifact authentication before installing requirements. It:

  1. Reads CA_REPOSITORY_ARN from the container environment — if not set, does a normal pip install
  2. Tries boto3 first (matching sagemaker-training-toolkit's approach) to build an authenticated pip index URL
  3. Falls back to AWS CLI (aws codeartifact login --tool pip) if boto3 is unavailable
  4. Hard-fails with a clear error if CA_REPOSITORY_ARN is set but neither boto3 nor AWS CLI is available

The script can be used as:

  • A standalone script: python install_requirements.py requirements.txt (used by bash-based entrypoints)
  • An importable module: from sagemaker.core.utils.install_requirements import configure_pip, install_requirements (for Python-native callers like @remote or ModelBuilder)

Changes

FileChange
sagemaker-core/.../utils/install_requirements.pyNew module with configure_pip(), install_requirements(), main(), and CodeArtifactAuthMethod enum
sagemaker-core/tests/unit/test_install_requirements.py22 unit tests covering all auth methods, fallback chains, error propagation
sagemaker-train/.../templates.pyINSTALL_REQUIREMENTS and INSTALL_AUTO_REQUIREMENTS now call install_requirements.py instead of bare pip install
sagemaker-train/.../model_trainer.pyCopy install_requirements.py from sagemaker-core into sm_drivers/scripts/ at runtime
sagemaker-core/.../processing.pyUpload install_requirements.py as sibling file alongside runproc.sh and sourcedir.tar.gz; update generated script to call it
sagemaker-core/tests/unit/test_processing.pyVerify install_requirements.py is uploaded and referenced in generated script

What this covers

Job TypeClassCodeArtifact with this PR
TrainingModelTrainer✅ Fixed — install_requirements.py in sm_drivers/scripts/
ProcessingFrameworkProcessor✅ Fixed — install_requirements.py uploaded as sibling file
TuningTuner✅ Already works — Tuner uses container's native toolkit (not affected by this PR)
InferenceModelBuilder✅ Already works — SDK doesn't override inference entrypoints

What this does NOT cover

PathStatusNotes
@remote function (runtime_environment_manager.py)❌ Not wiredHas its own _install_requirements_txt() that does bare pip install. Could use configure_pip() via import.
sagemaker-serve (requirements_manager.py)❌ Not wiredSame — bare pip install in-process. Could import configure_pip().
sagemaker-core/modules (templates.py)❌ Not wiredDuplicate of sagemaker-train/templates.py without INSTALL_REQUIREMENTS. Lower priority.

These are follow-up opportunities — the module is available for them to import.

Known risks

  1. Tuning jobs depend on the container's toolkit — The CreateHyperParameterTuningJob API uses HyperParameterAlgorithmSpecification which lacks ContainerEntrypoint, so the Tuner cannot use sm_train.sh. If future containers drop sagemaker-training-toolkit, tuning jobs will lose CodeArtifact support with no SDK-side fix possible until the API adds entrypoint support.

  2. boto3 availability in future containers — Current PyTorch training containers (2.7–2.9) include boto3. New DLC base images on the main branch do not. The script's fallback to AWS CLI mitigates this, but if neither is available, the script hard-fails. The long-term solution is a shared package with boto3 as a declared dependency (see analysis).

Long-term solution

This PR is a stopgap that works within the SDK alone. The long-term solution requires coordination between the SDK and DLC to ensure that both the container's default entrypoint and any SDK-overridden entrypoint have access to the same CodeArtifact-aware installer — ideally a shared package with boto3 as a declared dependency, installed in all SageMaker containers. See the proposed ideal solution for details.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

…r requirements.txt installation
SDK v3's `ModelTrainer` and `FrameworkProcessor` override the container entrypoint with SDK-generated scripts (`sm_train.sh`, `runproc.sh`), bypassing the container's entrypoint which involved `sagemaker-training-toolkit` handling
`CA_REPOSITORY_ARN`-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in aws#5765.
This is the stopgap solution proposed in this comment[aws#5765 (comment)]: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.
- Add `install_requirements.py` in sagemaker-core — reads `CA_REPOSITORY_ARN` from container environment; no-op if unset
- Try `boto3` first (matching sagemaker-training-toolkit), fall back to `AWS CLI`, hard-fail if neither is available
- Wire into `ModelTrainer`: copy script into `sm_drivers/scripts/`, update `INSTALL_REQUIREMENTS` templates to call it instead of bare `pip install`
- Wire into `FrameworkProcessor`: upload script as sibling file alongside `runproc.sh`, update generated script to call it
@humanzz

Copy link
Copy Markdown
ContributorAuthor

I've also tested this with my code (as an integration test) to verify the behaviours

Training Job (ModelTrainer)

fromsagemaker.trainimportModelTrainerfromsagemaker.core.training.configsimportCompute, SourceCodetrainer=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"
),
compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
sagemaker_session=sm_session,
environment={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
trainer.train(input_data_config=inputs, wait=False)

CloudWatch logsinstall_requirements.py ran from sm_drivers/scripts/, authenticated via boto3, pip resolved from CodeArtifact:

Installing requirements
++ /usr/local/bin/python3 /opt/ml/input/data/sm_drivers/scripts/install_requirements.py requirements.txt
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Training job completed successfully. ✅

Processing Job (FrameworkProcessor)

fromsagemaker.core.processingimportFrameworkProcessorprocessor=FrameworkProcessor(
image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-gpu-py313",
command=["python3"],
role="arn:aws:iam::123456789:role/MyRole",
instance_count=1,
instance_type="ml.g6.4xlarge",
sagemaker_session=sm_session,
env={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
processor.run(code="my_script.py", source_dir="src", wait=False)

CloudWatch logsinstall_requirements.py uploaded as sibling file, authenticated via boto3:

Files in /opt/ml/processing/input/code/ before extraction:
-rw-r--r-- 1 root root 6652 Apr 17 10:22 install_requirements.py
-rw-r--r-- 1 root root 685 Apr 17 10:22 runproc.sh
-rw-r--r-- 1 root root 81582 Apr 17 10:22 sourcedir.tar.gz
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Processing job completed successfully. ✅

@aviruthen
aviruthen merged commit 9386fc0 into aws:masterApr 28, 2026
37 of 53 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants

@humanzz@zhaoqizqwang@aviruthen@mollyheamazon
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links"); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "youtube.com"; var __re = new RegExp('^' + "youtube\\.com" + '
Skip to content

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation - #5772

Merged
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix
Apr 28, 2026
Merged

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation#5772
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix

Conversation

@humanzz

@humanzzhumanzz commented Apr 17, 2026

Copy link
Copy Markdown
Contributor

SDK v3's ModelTrainer and FrameworkProcessor override the container entrypoint with SDK-generated scripts (sm_train.sh, runproc.sh), bypassing the container's entrypoint which involved sagemaker-training-toolkit handling
CA_REPOSITORY_ARN-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in #5765.

This is the stopgap solution proposed in this comment: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.

  • Add install_requirements.py in sagemaker-core — reads CA_REPOSITORY_ARN from container environment; no-op if unset
  • Try boto3 first (matching sagemaker-training-toolkit), fall back to AWS CLI, hard-fail if neither is available
  • Wire into ModelTrainer: copy script into sm_drivers/scripts/, update INSTALL_REQUIREMENTS templates to call it instead of bare pip install
  • Wire into FrameworkProcessor: upload script as sibling file alongside runproc.sh, update generated script to call it

Issue #, if available:

#5765

Description of changes:

When the SDK overrides a container's entrypoint — as ModelTrainer does for training jobs (Bug 4) and FrameworkProcessor does for processing jobs (Bug 3) — the container's native sagemaker-training-toolkit is bypassed. This toolkit handled CA_REPOSITORY_ARN-based CodeArtifact authentication for requirements.txt installation via boto3. Without it, pip install -r requirements.txt runs against public PyPI, failing in VPC-isolated environments or when packages are only available in a private CodeArtifact repository.

See #5765 and the detailed analysis comment for full context.

Solution: Stopgap install_requirements.py

A self-contained Python script in sagemaker-core that handles CodeArtifact authentication before installing requirements. It:

  1. Reads CA_REPOSITORY_ARN from the container environment — if not set, does a normal pip install
  2. Tries boto3 first (matching sagemaker-training-toolkit's approach) to build an authenticated pip index URL
  3. Falls back to AWS CLI (aws codeartifact login --tool pip) if boto3 is unavailable
  4. Hard-fails with a clear error if CA_REPOSITORY_ARN is set but neither boto3 nor AWS CLI is available

The script can be used as:

  • A standalone script: python install_requirements.py requirements.txt (used by bash-based entrypoints)
  • An importable module: from sagemaker.core.utils.install_requirements import configure_pip, install_requirements (for Python-native callers like @remote or ModelBuilder)

Changes

FileChange
sagemaker-core/.../utils/install_requirements.pyNew module with configure_pip(), install_requirements(), main(), and CodeArtifactAuthMethod enum
sagemaker-core/tests/unit/test_install_requirements.py22 unit tests covering all auth methods, fallback chains, error propagation
sagemaker-train/.../templates.pyINSTALL_REQUIREMENTS and INSTALL_AUTO_REQUIREMENTS now call install_requirements.py instead of bare pip install
sagemaker-train/.../model_trainer.pyCopy install_requirements.py from sagemaker-core into sm_drivers/scripts/ at runtime
sagemaker-core/.../processing.pyUpload install_requirements.py as sibling file alongside runproc.sh and sourcedir.tar.gz; update generated script to call it
sagemaker-core/tests/unit/test_processing.pyVerify install_requirements.py is uploaded and referenced in generated script

What this covers

Job TypeClassCodeArtifact with this PR
TrainingModelTrainer✅ Fixed — install_requirements.py in sm_drivers/scripts/
ProcessingFrameworkProcessor✅ Fixed — install_requirements.py uploaded as sibling file
TuningTuner✅ Already works — Tuner uses container's native toolkit (not affected by this PR)
InferenceModelBuilder✅ Already works — SDK doesn't override inference entrypoints

What this does NOT cover

PathStatusNotes
@remote function (runtime_environment_manager.py)❌ Not wiredHas its own _install_requirements_txt() that does bare pip install. Could use configure_pip() via import.
sagemaker-serve (requirements_manager.py)❌ Not wiredSame — bare pip install in-process. Could import configure_pip().
sagemaker-core/modules (templates.py)❌ Not wiredDuplicate of sagemaker-train/templates.py without INSTALL_REQUIREMENTS. Lower priority.

These are follow-up opportunities — the module is available for them to import.

Known risks

  1. Tuning jobs depend on the container's toolkit — The CreateHyperParameterTuningJob API uses HyperParameterAlgorithmSpecification which lacks ContainerEntrypoint, so the Tuner cannot use sm_train.sh. If future containers drop sagemaker-training-toolkit, tuning jobs will lose CodeArtifact support with no SDK-side fix possible until the API adds entrypoint support.

  2. boto3 availability in future containers — Current PyTorch training containers (2.7–2.9) include boto3. New DLC base images on the main branch do not. The script's fallback to AWS CLI mitigates this, but if neither is available, the script hard-fails. The long-term solution is a shared package with boto3 as a declared dependency (see analysis).

Long-term solution

This PR is a stopgap that works within the SDK alone. The long-term solution requires coordination between the SDK and DLC to ensure that both the container's default entrypoint and any SDK-overridden entrypoint have access to the same CodeArtifact-aware installer — ideally a shared package with boto3 as a declared dependency, installed in all SageMaker containers. See the proposed ideal solution for details.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

…r requirements.txt installation
SDK v3's `ModelTrainer` and `FrameworkProcessor` override the container entrypoint with SDK-generated scripts (`sm_train.sh`, `runproc.sh`), bypassing the container's entrypoint which involved `sagemaker-training-toolkit` handling
`CA_REPOSITORY_ARN`-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in aws#5765.
This is the stopgap solution proposed in this comment[aws#5765 (comment)]: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.
- Add `install_requirements.py` in sagemaker-core — reads `CA_REPOSITORY_ARN` from container environment; no-op if unset
- Try `boto3` first (matching sagemaker-training-toolkit), fall back to `AWS CLI`, hard-fail if neither is available
- Wire into `ModelTrainer`: copy script into `sm_drivers/scripts/`, update `INSTALL_REQUIREMENTS` templates to call it instead of bare `pip install`
- Wire into `FrameworkProcessor`: upload script as sibling file alongside `runproc.sh`, update generated script to call it
@humanzz

Copy link
Copy Markdown
ContributorAuthor

I've also tested this with my code (as an integration test) to verify the behaviours

Training Job (ModelTrainer)

fromsagemaker.trainimportModelTrainerfromsagemaker.core.training.configsimportCompute, SourceCodetrainer=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"
),
compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
sagemaker_session=sm_session,
environment={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
trainer.train(input_data_config=inputs, wait=False)

CloudWatch logsinstall_requirements.py ran from sm_drivers/scripts/, authenticated via boto3, pip resolved from CodeArtifact:

Installing requirements
++ /usr/local/bin/python3 /opt/ml/input/data/sm_drivers/scripts/install_requirements.py requirements.txt
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Training job completed successfully. ✅

Processing Job (FrameworkProcessor)

fromsagemaker.core.processingimportFrameworkProcessorprocessor=FrameworkProcessor(
image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-gpu-py313",
command=["python3"],
role="arn:aws:iam::123456789:role/MyRole",
instance_count=1,
instance_type="ml.g6.4xlarge",
sagemaker_session=sm_session,
env={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
processor.run(code="my_script.py", source_dir="src", wait=False)

CloudWatch logsinstall_requirements.py uploaded as sibling file, authenticated via boto3:

Files in /opt/ml/processing/input/code/ before extraction:
-rw-r--r-- 1 root root 6652 Apr 17 10:22 install_requirements.py
-rw-r--r-- 1 root root 685 Apr 17 10:22 runproc.sh
-rw-r--r-- 1 root root 81582 Apr 17 10:22 sourcedir.tar.gz
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Processing job completed successfully. ✅

@aviruthen
aviruthen merged commit 9386fc0 into aws:masterApr 28, 2026
37 of 53 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants

@humanzz@zhaoqizqwang@aviruthen@mollyheamazon
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default"); } } catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation - #5772

Merged
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix
Apr 28, 2026
Merged

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation#5772
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix

Conversation

@humanzz

@humanzzhumanzz commented Apr 17, 2026

Copy link
Copy Markdown
Contributor

SDK v3's ModelTrainer and FrameworkProcessor override the container entrypoint with SDK-generated scripts (sm_train.sh, runproc.sh), bypassing the container's entrypoint which involved sagemaker-training-toolkit handling
CA_REPOSITORY_ARN-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in #5765.

This is the stopgap solution proposed in this comment: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.

  • Add install_requirements.py in sagemaker-core — reads CA_REPOSITORY_ARN from container environment; no-op if unset
  • Try boto3 first (matching sagemaker-training-toolkit), fall back to AWS CLI, hard-fail if neither is available
  • Wire into ModelTrainer: copy script into sm_drivers/scripts/, update INSTALL_REQUIREMENTS templates to call it instead of bare pip install
  • Wire into FrameworkProcessor: upload script as sibling file alongside runproc.sh, update generated script to call it

Issue #, if available:

#5765

Description of changes:

When the SDK overrides a container's entrypoint — as ModelTrainer does for training jobs (Bug 4) and FrameworkProcessor does for processing jobs (Bug 3) — the container's native sagemaker-training-toolkit is bypassed. This toolkit handled CA_REPOSITORY_ARN-based CodeArtifact authentication for requirements.txt installation via boto3. Without it, pip install -r requirements.txt runs against public PyPI, failing in VPC-isolated environments or when packages are only available in a private CodeArtifact repository.

See #5765 and the detailed analysis comment for full context.

Solution: Stopgap install_requirements.py

A self-contained Python script in sagemaker-core that handles CodeArtifact authentication before installing requirements. It:

  1. Reads CA_REPOSITORY_ARN from the container environment — if not set, does a normal pip install
  2. Tries boto3 first (matching sagemaker-training-toolkit's approach) to build an authenticated pip index URL
  3. Falls back to AWS CLI (aws codeartifact login --tool pip) if boto3 is unavailable
  4. Hard-fails with a clear error if CA_REPOSITORY_ARN is set but neither boto3 nor AWS CLI is available

The script can be used as:

  • A standalone script: python install_requirements.py requirements.txt (used by bash-based entrypoints)
  • An importable module: from sagemaker.core.utils.install_requirements import configure_pip, install_requirements (for Python-native callers like @remote or ModelBuilder)

Changes

FileChange
sagemaker-core/.../utils/install_requirements.pyNew module with configure_pip(), install_requirements(), main(), and CodeArtifactAuthMethod enum
sagemaker-core/tests/unit/test_install_requirements.py22 unit tests covering all auth methods, fallback chains, error propagation
sagemaker-train/.../templates.pyINSTALL_REQUIREMENTS and INSTALL_AUTO_REQUIREMENTS now call install_requirements.py instead of bare pip install
sagemaker-train/.../model_trainer.pyCopy install_requirements.py from sagemaker-core into sm_drivers/scripts/ at runtime
sagemaker-core/.../processing.pyUpload install_requirements.py as sibling file alongside runproc.sh and sourcedir.tar.gz; update generated script to call it
sagemaker-core/tests/unit/test_processing.pyVerify install_requirements.py is uploaded and referenced in generated script

What this covers

Job TypeClassCodeArtifact with this PR
TrainingModelTrainer✅ Fixed — install_requirements.py in sm_drivers/scripts/
ProcessingFrameworkProcessor✅ Fixed — install_requirements.py uploaded as sibling file
TuningTuner✅ Already works — Tuner uses container's native toolkit (not affected by this PR)
InferenceModelBuilder✅ Already works — SDK doesn't override inference entrypoints

What this does NOT cover

PathStatusNotes
@remote function (runtime_environment_manager.py)❌ Not wiredHas its own _install_requirements_txt() that does bare pip install. Could use configure_pip() via import.
sagemaker-serve (requirements_manager.py)❌ Not wiredSame — bare pip install in-process. Could import configure_pip().
sagemaker-core/modules (templates.py)❌ Not wiredDuplicate of sagemaker-train/templates.py without INSTALL_REQUIREMENTS. Lower priority.

These are follow-up opportunities — the module is available for them to import.

Known risks

  1. Tuning jobs depend on the container's toolkit — The CreateHyperParameterTuningJob API uses HyperParameterAlgorithmSpecification which lacks ContainerEntrypoint, so the Tuner cannot use sm_train.sh. If future containers drop sagemaker-training-toolkit, tuning jobs will lose CodeArtifact support with no SDK-side fix possible until the API adds entrypoint support.

  2. boto3 availability in future containers — Current PyTorch training containers (2.7–2.9) include boto3. New DLC base images on the main branch do not. The script's fallback to AWS CLI mitigates this, but if neither is available, the script hard-fails. The long-term solution is a shared package with boto3 as a declared dependency (see analysis).

Long-term solution

This PR is a stopgap that works within the SDK alone. The long-term solution requires coordination between the SDK and DLC to ensure that both the container's default entrypoint and any SDK-overridden entrypoint have access to the same CodeArtifact-aware installer — ideally a shared package with boto3 as a declared dependency, installed in all SageMaker containers. See the proposed ideal solution for details.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

…r requirements.txt installation
SDK v3's `ModelTrainer` and `FrameworkProcessor` override the container entrypoint with SDK-generated scripts (`sm_train.sh`, `runproc.sh`), bypassing the container's entrypoint which involved `sagemaker-training-toolkit` handling
`CA_REPOSITORY_ARN`-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in aws#5765.
This is the stopgap solution proposed in this comment[aws#5765 (comment)]: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.
- Add `install_requirements.py` in sagemaker-core — reads `CA_REPOSITORY_ARN` from container environment; no-op if unset
- Try `boto3` first (matching sagemaker-training-toolkit), fall back to `AWS CLI`, hard-fail if neither is available
- Wire into `ModelTrainer`: copy script into `sm_drivers/scripts/`, update `INSTALL_REQUIREMENTS` templates to call it instead of bare `pip install`
- Wire into `FrameworkProcessor`: upload script as sibling file alongside `runproc.sh`, update generated script to call it
@humanzz

Copy link
Copy Markdown
ContributorAuthor

I've also tested this with my code (as an integration test) to verify the behaviours

Training Job (ModelTrainer)

fromsagemaker.trainimportModelTrainerfromsagemaker.core.training.configsimportCompute, SourceCodetrainer=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"
),
compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
sagemaker_session=sm_session,
environment={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
trainer.train(input_data_config=inputs, wait=False)

CloudWatch logsinstall_requirements.py ran from sm_drivers/scripts/, authenticated via boto3, pip resolved from CodeArtifact:

Installing requirements
++ /usr/local/bin/python3 /opt/ml/input/data/sm_drivers/scripts/install_requirements.py requirements.txt
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Training job completed successfully. ✅

Processing Job (FrameworkProcessor)

fromsagemaker.core.processingimportFrameworkProcessorprocessor=FrameworkProcessor(
image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-gpu-py313",
command=["python3"],
role="arn:aws:iam::123456789:role/MyRole",
instance_count=1,
instance_type="ml.g6.4xlarge",
sagemaker_session=sm_session,
env={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
processor.run(code="my_script.py", source_dir="src", wait=False)

CloudWatch logsinstall_requirements.py uploaded as sibling file, authenticated via boto3:

Files in /opt/ml/processing/input/code/ before extraction:
-rw-r--r-- 1 root root 6652 Apr 17 10:22 install_requirements.py
-rw-r--r-- 1 root root 685 Apr 17 10:22 runproc.sh
-rw-r--r-- 1 root root 81582 Apr 17 10:22 sourcedir.tar.gz
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Processing job completed successfully. ✅

@aviruthen
aviruthen merged commit 9386fc0 into aws:masterApr 28, 2026
37 of 53 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants

@humanzz@zhaoqizqwang@aviruthen@mollyheamazon
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation - #5772

Merged
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix
Apr 28, 2026
Merged

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation#5772
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix

Conversation

@humanzz

@humanzzhumanzz commented Apr 17, 2026

Copy link
Copy Markdown
Contributor

SDK v3's ModelTrainer and FrameworkProcessor override the container entrypoint with SDK-generated scripts (sm_train.sh, runproc.sh), bypassing the container's entrypoint which involved sagemaker-training-toolkit handling
CA_REPOSITORY_ARN-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in #5765.

This is the stopgap solution proposed in this comment: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.

  • Add install_requirements.py in sagemaker-core — reads CA_REPOSITORY_ARN from container environment; no-op if unset
  • Try boto3 first (matching sagemaker-training-toolkit), fall back to AWS CLI, hard-fail if neither is available
  • Wire into ModelTrainer: copy script into sm_drivers/scripts/, update INSTALL_REQUIREMENTS templates to call it instead of bare pip install
  • Wire into FrameworkProcessor: upload script as sibling file alongside runproc.sh, update generated script to call it

Issue #, if available:

#5765

Description of changes:

When the SDK overrides a container's entrypoint — as ModelTrainer does for training jobs (Bug 4) and FrameworkProcessor does for processing jobs (Bug 3) — the container's native sagemaker-training-toolkit is bypassed. This toolkit handled CA_REPOSITORY_ARN-based CodeArtifact authentication for requirements.txt installation via boto3. Without it, pip install -r requirements.txt runs against public PyPI, failing in VPC-isolated environments or when packages are only available in a private CodeArtifact repository.

See #5765 and the detailed analysis comment for full context.

Solution: Stopgap install_requirements.py

A self-contained Python script in sagemaker-core that handles CodeArtifact authentication before installing requirements. It:

  1. Reads CA_REPOSITORY_ARN from the container environment — if not set, does a normal pip install
  2. Tries boto3 first (matching sagemaker-training-toolkit's approach) to build an authenticated pip index URL
  3. Falls back to AWS CLI (aws codeartifact login --tool pip) if boto3 is unavailable
  4. Hard-fails with a clear error if CA_REPOSITORY_ARN is set but neither boto3 nor AWS CLI is available

The script can be used as:

  • A standalone script: python install_requirements.py requirements.txt (used by bash-based entrypoints)
  • An importable module: from sagemaker.core.utils.install_requirements import configure_pip, install_requirements (for Python-native callers like @remote or ModelBuilder)

Changes

FileChange
sagemaker-core/.../utils/install_requirements.pyNew module with configure_pip(), install_requirements(), main(), and CodeArtifactAuthMethod enum
sagemaker-core/tests/unit/test_install_requirements.py22 unit tests covering all auth methods, fallback chains, error propagation
sagemaker-train/.../templates.pyINSTALL_REQUIREMENTS and INSTALL_AUTO_REQUIREMENTS now call install_requirements.py instead of bare pip install
sagemaker-train/.../model_trainer.pyCopy install_requirements.py from sagemaker-core into sm_drivers/scripts/ at runtime
sagemaker-core/.../processing.pyUpload install_requirements.py as sibling file alongside runproc.sh and sourcedir.tar.gz; update generated script to call it
sagemaker-core/tests/unit/test_processing.pyVerify install_requirements.py is uploaded and referenced in generated script

What this covers

Job TypeClassCodeArtifact with this PR
TrainingModelTrainer✅ Fixed — install_requirements.py in sm_drivers/scripts/
ProcessingFrameworkProcessor✅ Fixed — install_requirements.py uploaded as sibling file
TuningTuner✅ Already works — Tuner uses container's native toolkit (not affected by this PR)
InferenceModelBuilder✅ Already works — SDK doesn't override inference entrypoints

What this does NOT cover

PathStatusNotes
@remote function (runtime_environment_manager.py)❌ Not wiredHas its own _install_requirements_txt() that does bare pip install. Could use configure_pip() via import.
sagemaker-serve (requirements_manager.py)❌ Not wiredSame — bare pip install in-process. Could import configure_pip().
sagemaker-core/modules (templates.py)❌ Not wiredDuplicate of sagemaker-train/templates.py without INSTALL_REQUIREMENTS. Lower priority.

These are follow-up opportunities — the module is available for them to import.

Known risks

  1. Tuning jobs depend on the container's toolkit — The CreateHyperParameterTuningJob API uses HyperParameterAlgorithmSpecification which lacks ContainerEntrypoint, so the Tuner cannot use sm_train.sh. If future containers drop sagemaker-training-toolkit, tuning jobs will lose CodeArtifact support with no SDK-side fix possible until the API adds entrypoint support.

  2. boto3 availability in future containers — Current PyTorch training containers (2.7–2.9) include boto3. New DLC base images on the main branch do not. The script's fallback to AWS CLI mitigates this, but if neither is available, the script hard-fails. The long-term solution is a shared package with boto3 as a declared dependency (see analysis).

Long-term solution

This PR is a stopgap that works within the SDK alone. The long-term solution requires coordination between the SDK and DLC to ensure that both the container's default entrypoint and any SDK-overridden entrypoint have access to the same CodeArtifact-aware installer — ideally a shared package with boto3 as a declared dependency, installed in all SageMaker containers. See the proposed ideal solution for details.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

…r requirements.txt installation
SDK v3's `ModelTrainer` and `FrameworkProcessor` override the container entrypoint with SDK-generated scripts (`sm_train.sh`, `runproc.sh`), bypassing the container's entrypoint which involved `sagemaker-training-toolkit` handling
`CA_REPOSITORY_ARN`-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in aws#5765.
This is the stopgap solution proposed in this comment[aws#5765 (comment)]: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.
- Add `install_requirements.py` in sagemaker-core — reads `CA_REPOSITORY_ARN` from container environment; no-op if unset
- Try `boto3` first (matching sagemaker-training-toolkit), fall back to `AWS CLI`, hard-fail if neither is available
- Wire into `ModelTrainer`: copy script into `sm_drivers/scripts/`, update `INSTALL_REQUIREMENTS` templates to call it instead of bare `pip install`
- Wire into `FrameworkProcessor`: upload script as sibling file alongside `runproc.sh`, update generated script to call it
@humanzz

Copy link
Copy Markdown
ContributorAuthor

I've also tested this with my code (as an integration test) to verify the behaviours

Training Job (ModelTrainer)

fromsagemaker.trainimportModelTrainerfromsagemaker.core.training.configsimportCompute, SourceCodetrainer=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"
),
compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
sagemaker_session=sm_session,
environment={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
trainer.train(input_data_config=inputs, wait=False)

CloudWatch logsinstall_requirements.py ran from sm_drivers/scripts/, authenticated via boto3, pip resolved from CodeArtifact:

Installing requirements
++ /usr/local/bin/python3 /opt/ml/input/data/sm_drivers/scripts/install_requirements.py requirements.txt
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Training job completed successfully. ✅

Processing Job (FrameworkProcessor)

fromsagemaker.core.processingimportFrameworkProcessorprocessor=FrameworkProcessor(
image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-gpu-py313",
command=["python3"],
role="arn:aws:iam::123456789:role/MyRole",
instance_count=1,
instance_type="ml.g6.4xlarge",
sagemaker_session=sm_session,
env={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
processor.run(code="my_script.py", source_dir="src", wait=False)

CloudWatch logsinstall_requirements.py uploaded as sibling file, authenticated via boto3:

Files in /opt/ml/processing/input/code/ before extraction:
-rw-r--r-- 1 root root 6652 Apr 17 10:22 install_requirements.py
-rw-r--r-- 1 root root 685 Apr 17 10:22 runproc.sh
-rw-r--r-- 1 root root 81582 Apr 17 10:22 sourcedir.tar.gz
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Processing job completed successfully. ✅

@aviruthen
aviruthen merged commit 9386fc0 into aws:masterApr 28, 2026
37 of 53 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants

@humanzz@zhaoqizqwang@aviruthen@mollyheamazon
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode"); } } catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); } })(); })();
Skip to content

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation - #5772

Merged
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix
Apr 28, 2026
Merged

feat: add CodeArtifact support for ModelTrainer and FrameworkProcessor requirements.txt installation#5772
aviruthen merged 1 commit into
aws:masterfrom
humanzz:codeartifact-fix

Conversation

@humanzz

@humanzzhumanzz commented Apr 17, 2026

Copy link
Copy Markdown
Contributor

SDK v3's ModelTrainer and FrameworkProcessor override the container entrypoint with SDK-generated scripts (sm_train.sh, runproc.sh), bypassing the container's entrypoint which involved sagemaker-training-toolkit handling
CA_REPOSITORY_ARN-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in #5765.

This is the stopgap solution proposed in this comment: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.

  • Add install_requirements.py in sagemaker-core — reads CA_REPOSITORY_ARN from container environment; no-op if unset
  • Try boto3 first (matching sagemaker-training-toolkit), fall back to AWS CLI, hard-fail if neither is available
  • Wire into ModelTrainer: copy script into sm_drivers/scripts/, update INSTALL_REQUIREMENTS templates to call it instead of bare pip install
  • Wire into FrameworkProcessor: upload script as sibling file alongside runproc.sh, update generated script to call it

Issue #, if available:

#5765

Description of changes:

When the SDK overrides a container's entrypoint — as ModelTrainer does for training jobs (Bug 4) and FrameworkProcessor does for processing jobs (Bug 3) — the container's native sagemaker-training-toolkit is bypassed. This toolkit handled CA_REPOSITORY_ARN-based CodeArtifact authentication for requirements.txt installation via boto3. Without it, pip install -r requirements.txt runs against public PyPI, failing in VPC-isolated environments or when packages are only available in a private CodeArtifact repository.

See #5765 and the detailed analysis comment for full context.

Solution: Stopgap install_requirements.py

A self-contained Python script in sagemaker-core that handles CodeArtifact authentication before installing requirements. It:

  1. Reads CA_REPOSITORY_ARN from the container environment — if not set, does a normal pip install
  2. Tries boto3 first (matching sagemaker-training-toolkit's approach) to build an authenticated pip index URL
  3. Falls back to AWS CLI (aws codeartifact login --tool pip) if boto3 is unavailable
  4. Hard-fails with a clear error if CA_REPOSITORY_ARN is set but neither boto3 nor AWS CLI is available

The script can be used as:

  • A standalone script: python install_requirements.py requirements.txt (used by bash-based entrypoints)
  • An importable module: from sagemaker.core.utils.install_requirements import configure_pip, install_requirements (for Python-native callers like @remote or ModelBuilder)

Changes

FileChange
sagemaker-core/.../utils/install_requirements.pyNew module with configure_pip(), install_requirements(), main(), and CodeArtifactAuthMethod enum
sagemaker-core/tests/unit/test_install_requirements.py22 unit tests covering all auth methods, fallback chains, error propagation
sagemaker-train/.../templates.pyINSTALL_REQUIREMENTS and INSTALL_AUTO_REQUIREMENTS now call install_requirements.py instead of bare pip install
sagemaker-train/.../model_trainer.pyCopy install_requirements.py from sagemaker-core into sm_drivers/scripts/ at runtime
sagemaker-core/.../processing.pyUpload install_requirements.py as sibling file alongside runproc.sh and sourcedir.tar.gz; update generated script to call it
sagemaker-core/tests/unit/test_processing.pyVerify install_requirements.py is uploaded and referenced in generated script

What this covers

Job TypeClassCodeArtifact with this PR
TrainingModelTrainer✅ Fixed — install_requirements.py in sm_drivers/scripts/
ProcessingFrameworkProcessor✅ Fixed — install_requirements.py uploaded as sibling file
TuningTuner✅ Already works — Tuner uses container's native toolkit (not affected by this PR)
InferenceModelBuilder✅ Already works — SDK doesn't override inference entrypoints

What this does NOT cover

PathStatusNotes
@remote function (runtime_environment_manager.py)❌ Not wiredHas its own _install_requirements_txt() that does bare pip install. Could use configure_pip() via import.
sagemaker-serve (requirements_manager.py)❌ Not wiredSame — bare pip install in-process. Could import configure_pip().
sagemaker-core/modules (templates.py)❌ Not wiredDuplicate of sagemaker-train/templates.py without INSTALL_REQUIREMENTS. Lower priority.

These are follow-up opportunities — the module is available for them to import.

Known risks

  1. Tuning jobs depend on the container's toolkit — The CreateHyperParameterTuningJob API uses HyperParameterAlgorithmSpecification which lacks ContainerEntrypoint, so the Tuner cannot use sm_train.sh. If future containers drop sagemaker-training-toolkit, tuning jobs will lose CodeArtifact support with no SDK-side fix possible until the API adds entrypoint support.

  2. boto3 availability in future containers — Current PyTorch training containers (2.7–2.9) include boto3. New DLC base images on the main branch do not. The script's fallback to AWS CLI mitigates this, but if neither is available, the script hard-fails. The long-term solution is a shared package with boto3 as a declared dependency (see analysis).

Long-term solution

This PR is a stopgap that works within the SDK alone. The long-term solution requires coordination between the SDK and DLC to ensure that both the container's default entrypoint and any SDK-overridden entrypoint have access to the same CodeArtifact-aware installer — ideally a shared package with boto3 as a declared dependency, installed in all SageMaker containers. See the proposed ideal solution for details.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

…r requirements.txt installation
SDK v3's `ModelTrainer` and `FrameworkProcessor` override the container entrypoint with SDK-generated scripts (`sm_train.sh`, `runproc.sh`), bypassing the container's entrypoint which involved `sagemaker-training-toolkit` handling
`CA_REPOSITORY_ARN`-based CodeArtifact authentication. This broke CodeArtifact support for both training (Bug 4) and processing (Bug 3) reported in aws#5765.
This is the stopgap solution proposed in this comment[aws#5765 (comment)]: a self-contained install_requirements.py script that the SDK uploads to the container alongside its generated entrypoint scripts.
- Add `install_requirements.py` in sagemaker-core — reads `CA_REPOSITORY_ARN` from container environment; no-op if unset
- Try `boto3` first (matching sagemaker-training-toolkit), fall back to `AWS CLI`, hard-fail if neither is available
- Wire into `ModelTrainer`: copy script into `sm_drivers/scripts/`, update `INSTALL_REQUIREMENTS` templates to call it instead of bare `pip install`
- Wire into `FrameworkProcessor`: upload script as sibling file alongside `runproc.sh`, update generated script to call it
@humanzz

Copy link
Copy Markdown
ContributorAuthor

I've also tested this with my code (as an integration test) to verify the behaviours

Training Job (ModelTrainer)

fromsagemaker.trainimportModelTrainerfromsagemaker.core.training.configsimportCompute, SourceCodetrainer=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"
),
compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
sagemaker_session=sm_session,
environment={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
trainer.train(input_data_config=inputs, wait=False)

CloudWatch logsinstall_requirements.py ran from sm_drivers/scripts/, authenticated via boto3, pip resolved from CodeArtifact:

Installing requirements
++ /usr/local/bin/python3 /opt/ml/input/data/sm_drivers/scripts/install_requirements.py requirements.txt
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Training job completed successfully. ✅

Processing Job (FrameworkProcessor)

fromsagemaker.core.processingimportFrameworkProcessorprocessor=FrameworkProcessor(
image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.10-gpu-py313",
command=["python3"],
role="arn:aws:iam::123456789:role/MyRole",
instance_count=1,
instance_type="ml.g6.4xlarge",
sagemaker_session=sm_session,
env={"CA_REPOSITORY_ARN": "arn:aws:codeartifact:us-west-2:ACCOUNT:repository/DOMAIN/REPO"},
)
processor.run(code="my_script.py", source_dir="src", wait=False)

CloudWatch logsinstall_requirements.py uploaded as sibling file, authenticated via boto3:

Files in /opt/ml/processing/input/code/ before extraction:
-rw-r--r-- 1 root root 6652 Apr 17 10:22 install_requirements.py
-rw-r--r-- 1 root root 685 Apr 17 10:22 runproc.sh
-rw-r--r-- 1 root root 81582 Apr 17 10:22 sourcedir.tar.gz
Looking in indexes: https://aws:****@amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/pyarrow/20.0.0/pyarrow-20.0.0-cp313-cp313-manylinux_2_28_x86_64.whl (42.3 MB)
Downloading https://amazon-ACCOUNT.d.codeartifact.us-west-2.amazonaws.com/pypi/REPO/simple/sentence-transformers/5.4.1/sentence_transformers-5.4.1-py3-none-any.whl (571 kB)

Processing job completed successfully. ✅

@aviruthen
aviruthen merged commit 9386fc0 into aws:masterApr 28, 2026
37 of 53 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants

@humanzz@zhaoqizqwang@aviruthen@mollyheamazon