Activity

1. humanzz

   changed the title [-][v3] FrameworkProcessor and ModelTrainer: 3 regressions from v2 migration[/-] [+][v3] FrameworkProcessor and ModelTrainer: 4 regressions from v2 migration[/+] on Apr 16, 2026
2. humanzz commented on Apr 16, 2026

   

   humanzz

   on Apr 16, 2026

   Last edited by humanzz

   ContributorAuthor

   More actions

   The following is a bit more context from an LLM-analysis I did on the codebase, along with historical context given I contributed the codeartifact support before, and proposal for solutions.

   CodeArtifact & SageMaker: How It Worked, What Broke, and How to Fix It

   ---

   1. The Problem CodeArtifact Solves
   2. How CodeArtifact Support Was Built (2023)
   3. The SDK v2 State: How Each Job Type Worked
   4. What Broke in SDK v3
   5. The Ideal Long-Term Solution (SDK + DLC Updates)
   6. The Stopgap Solution (SDK-Only Updates)
   7. Appendix: Container Inventory
   8. Appendix: The Three Existing Implementations

   ---

   The Problem CodeArtifact Solves

   SageMaker training, processing, and inference containers support installing additional Python dependencies from a requirements.txt file bundled with user code. By default, pip install resolves packages from the public PyPI index, which requires internet access.

   CodeArtifact solves this by providing a private PyPI mirror. With CodeArtifact:

   • Containers don't need internet access (improved security posture)
   • Only approved packages are installable (policy enforcement)

   The mechanism is simple: set the CA_REPOSITORY_ARN environment variable on the SageMaker job, and the container authenticates with CodeArtifact before running pip install.

   ---

   How CodeArtifact Support Was Built (2023)

   CodeArtifact support was contributed through a series of PRs across multiple repositories (see deep-learning-containers#2509 for the original feature request):

   1. sagemaker-training-toolkit#187 — Training containers read CA_REPOSITORY_ARN, use boto3 to get an auth token, build an authenticated pip index URL, pass it to pip install -i <url>
   2. sagemaker-inference-toolkit#130 — Same approach for inference containers
   3. sagemaker-pytorch-inference-toolkit#150 — Made PyTorch inference reuse the base toolkit's logic instead of duplicating it
   4. deep-learning-containers#3227 — Updated PyTorch inference container image to pick up the new toolkit
   5. sagemaker-python-sdk#4145 — Added codeartifact_repo_arn to SDK v2's FrameworkProcessor.run() for processing jobs

   The core authentication logic (parse ARN → boto3 get_authorization_token → build authenticated URL) was implemented identically in both sagemaker-training-toolkit and sagemaker-inference-toolkit. The new DLC ray inference containers (scripts/ray/sagemaker_serve.py) later independently reimplemented the same pattern.

   ---

   The SDK v2 State: How Each Job Type Worked

   The Clean Model: Container Owns the Lifecycle

   In SDK v2, the SDK's role was simple for training, inference, and tuning: upload user code, set environment variables, and let the container handle everything else. The container's native entrypoint invoked the appropriate toolkit, which handled:

   1. Downloading and extracting user code
   2. Reading CA_REPOSITORY_ARN from the environment
   3. Authenticating with CodeArtifact via boto3
   4. Installing requirements.txt with the authenticated pip index
   5. Running the user's entry script

   This worked because the toolkits were installed in the containers, boto3 was available, and the SDK didn't interfere with the container's entrypoint.

   The Exception: Processing Jobs

   Processing jobs were the one case where the SDK did override the container's behavior. FrameworkProcessor generated its own bash script (runproc.sh) that handled code extraction and dependency installation. Because this script bypassed the toolkit, it needed its own CodeArtifact implementation.

   Since runproc.sh is bash, the SDK v2 used AWS CLI (aws codeartifact login --tool pip) rather than boto3. This was added in PR #4145.

   SDK v2 Summary

   Job Type	Class	Who owns the lifecycle?	Entrypoint	CodeArtifact mechanism
   Training	PyTorch estimator	Container (sagemaker-training-toolkit)	Container's native	boto3 (pip install -i <url>)
   Inference	PyTorchModel	Container (sagemaker-inference-toolkit)	Container's native	boto3 (pip install -i <url>)
   Tuning	HyperparameterTuner	Container (sagemaker-training-toolkit)	Container's native	boto3 (pip install -i <url>)
   Processing	FrameworkProcessor	SDK (generated runproc.sh)	SDK overrides	AWS CLI (aws codeartifact login)

   The key insight: only processing jobs bypassed the container's toolkit, and that's the only place the SDK needed its own CodeArtifact implementation.

   ---

   What Broke in SDK v3

   Training: ModelTrainer Now Overrides the Entrypoint (Bug 4)

   The SDK v3 ModelTrainer follows the same pattern that FrameworkProcessor used — it overrides the container's entrypoint with its own SDK-generated script (sm_train.sh):

   # sagemaker-train/src/sagemaker/train/constants.py
   DEFAULT_CONTAINER_ENTRYPOINT = ["/bin/bash"]
   DEFAULT_CONTAINER_ARGUMENTS = ["-c", "chmod +x .../sm_train.sh && .../sm_train.sh"]

   This bypasses sagemaker-training-toolkit entirely. The SDK's sm_train.sh reimplements environment setup, requirements installation, and script execution — but without CodeArtifact support:

   # sagemaker-train/src/sagemaker/train/templates.py — INSTALL_REQUIREMENTS (no CodeArtifact)
   echo "Installing requirements"
   $SM_PIP_CMD install -r {requirements_file}

   Source: constants.py, templates.py

   The CA_REPOSITORY_ARN env var reaches the container but nothing reads it.

   Processing: codeartifact_repo_arn Not Ported (Bug 3)

   The SDK v3 FrameworkProcessor was rewritten in sagemaker-core but the codeartifact_repo_arn parameter from PR #4145 was not ported. The generated runproc.sh does a bare pip install without authentication.

   Tuning: Not Affected (Due to an API Limitation)

   The SDK v3 Tuner wraps ModelTrainer but cannot override the container entrypoint because the AWS CreateHyperParameterTuningJob API doesn't support it.

   The CreateHyperParameterTuningJob API has a TrainingJobDefinition field, which might suggest it supports the same options as CreateTrainingJob. However, the types differ despite confusingly similar naming:

   API	Field path	Actual type	Has ContainerEntrypoint?
   CreateTrainingJob	AlgorithmSpecification	AlgorithmSpecification	✅ Yes (8 fields)
   CreateHyperParameterTuningJob	TrainingJobDefinition.AlgorithmSpecification	HyperParameterAlgorithmSpecification	❌ No (4 fields)

   Both fields are named AlgorithmSpecification, but they map to different types. HyperParameterAlgorithmSpecification only has TrainingImage, TrainingInputMode, AlgorithmName, and MetricDefinitions. This applies to both TrainingJobDefinition (singular) and TrainingJobDefinitions (plural, for multi-algo tuning).

   Because of this, the Tuner falls back to the SDK v2-style approach: setting sagemaker_program and sagemaker_submit_directory as hyperparameters, letting the container's native toolkit handle everything — including CodeArtifact.

   Future risk: If AWS adds ContainerEntrypoint to HyperParameterAlgorithmSpecification, the SDK team may update the Tuner to match ModelTrainer's approach, breaking CodeArtifact for tuning too.

   Inference: Not Affected (Yet)

   The SDK v3 does not override inference container entrypoints. sagemaker-inference-toolkit still handles CA_REPOSITORY_ARN at model server startup. However, if a future SDK v3 inference path follows the ModelTrainer pattern, it would break here too.

   SDK v3 Summary

   Job Type	Class	Who owns the lifecycle?	Entrypoint	CodeArtifact
   Training	ModelTrainer	SDK (generated sm_train.sh)	SDK overrides	❌ Broken
   Inference	ModelBuilder	Container (toolkit)	Container's native	✅ Works
   Tuning	Tuner (wraps ModelTrainer)	Container (toolkit, due to API limitation)	Container's native	✅ Works (for now)
   Processing	FrameworkProcessor	SDK (generated runproc.sh)	SDK overrides	❌ Broken (param not ported)

   The pattern is clear: every time the SDK overrides the container entrypoint, CodeArtifact breaks.

   ---

   The Ideal Long-Term Solution (SDK + DLC Updates)

   The Root Problem

   CodeArtifact authentication logic has been copy-pasted across 4+ codebases:

   • sagemaker-training-toolkit (boto3)
   • sagemaker-inference-toolkit (boto3)
   • sagemaker-python-sdk SDK v2 FrameworkProcessor (AWS CLI)
   • deep-learning-containers ray sagemaker_serve.py (boto3)

   Each time the SDK takes over a container's entrypoint, this logic needs to be reimplemented again. This doesn't scale.

   Proposed Architecture

   A dedicated, shared Python package (or well-known module) that handles "configure pip for CodeArtifact from environment variables":

   sagemaker-pip-auth (or similar)
   ├── Depends on boto3 (declared in package dependencies)
   ├── Reads CA_REPOSITORY_ARN from environment
   ├── Uses boto3 to get auth token + repository endpoint
   ├── Returns authenticated pip index URL
   └── Provides a CLI entry point for bash callers
   

   Because the package declares boto3 as a dependency, installing it in a container automatically brings boto3 along — containers don't need to independently manage boto3.

   Requirements:

   1. This package is installed in all SageMaker containers (like the toolkits are today)
   2. The toolkits, the SDK's generated scripts, and any future entrypoint overrides all call this one package

   Usage from Python (toolkits, SDK driver scripts):

   from sagemaker_pip_auth import get_codeartifact_index
   index = get_codeartifact_index()  # Returns URL or None if CA_REPOSITORY_ARN not set
   if index:
       subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", "requirements.txt", "-i", index])

   Usage from bash (SDK-generated scripts like runproc.sh, sm_train.sh):

   # The package provides a CLI entry point
   if [ -n "$CA_REPOSITORY_ARN" ]; then
       CA_INDEX=$(python3 -m sagemaker_pip_auth)
       $SM_PIP_CMD install -r requirements.txt -i "$CA_INDEX"
   fi

   What This Requires

   Action	Owner	Dependency
   Create the shared package (with boto3 dependency)	SDK or DLC team	None
   Install the package in containers	DLC team	Package must exist first
   Update toolkits to use it	DLC team	Package must be in containers
   Update SDK scripts to use it	SDK team	Package must be in containers

   This requires cross-team coordination, but it's the right long-term answer.

   ---

   The Stopgap Solution (SDK-Only Updates)

   Given the current state — the SDK uploads custom entrypoint scripts that bypass the toolkit — we need a solution that works now, without waiting for container changes.

   Design: A Self-Contained Python Script

   A single Python script, uploaded by the SDK as part of its driver payload, that handles CodeArtifact authentication. It:

   1. Checks if CA_REPOSITORY_ARN is set — if not, does nothing
   2. Tries boto3 first (the established pattern used by all toolkits)
   3. Falls back to AWS CLI if boto3 is not available
   4. Hard fails if CA_REPOSITORY_ARN is set but neither boto3 nor AWS CLI is available — silent failure is worse than a clear error

   #!/usr/bin/env python3
   """Configure pip for CodeArtifact and install requirements.
   
   Reads CA_REPOSITORY_ARN from environment. Tries boto3 first (matching
   sagemaker-training-toolkit), falls back to AWS CLI, hard fails if neither
   is available when CA_REPOSITORY_ARN is set.
   
   Usage:
       python3 install_requirements.py requirements.txt
   """
   import os
   import re
   import subprocess
   import sys
   
   _ARN_RE = re.compile(
       r"arn:([^:]+):codeartifact:([^:]+):([^:]+):repository/([^/]+)/(.+)"
   )
   
   
   def _get_index_boto3(region, account, domain, repo):
       import boto3
       ca = boto3.client("codeartifact", region_name=region)
       token = ca.get_authorization_token(
           domain=domain, domainOwner=account
       )["authorizationToken"]
       endpoint = ca.get_repository_endpoint(
           domain=domain, domainOwner=account, repository=repo, format="pypi"
       )["repositoryEndpoint"]
       return re.sub(
           "https://", f"https://aws:{token}@",
           re.sub(f"{repo}/?$", f"{repo}/simple/", endpoint),
       )
   
   
   def _login_awscli(region, account, domain, repo):
       subprocess.check_call([
           "aws", "codeartifact", "login", "--tool", "pip",
           "--domain", domain, "--domain-owner", account,
           "--repository", repo, "--region", region,
       ])
   
   
   def configure_pip():
       """Returns an authenticated index URL, or None if CA not configured."""
       arn = os.environ.get("CA_REPOSITORY_ARN")
       if not arn:
           return None
   
       m = _ARN_RE.match(arn)
       if not m:
           raise ValueError(f"Invalid CA_REPOSITORY_ARN: {arn}")
       _, region, account, domain, repo = m.groups()
   
       # Try boto3 first (matches sagemaker-training-toolkit approach)
       try:
           return _get_index_boto3(region, account, domain, repo)
       except ImportError:
           pass
   
       # Fall back to AWS CLI
       try:
           _login_awscli(region, account, domain, repo)
           return None  # CLI modifies global pip config
       except FileNotFoundError:
           pass
   
       # Hard fail — CA is configured but we can't authenticate
       print(
           "ERROR: CA_REPOSITORY_ARN is set but neither boto3 nor AWS CLI "
           "is available to authenticate with CodeArtifact.",
           file=sys.stderr,
       )
       sys.exit(1)
   
   
   def main():
       req_file = sys.argv[1] if len(sys.argv) > 1 else "requirements.txt"
       pip_cmd = [sys.executable, "-m", "pip", "install", "-r", req_file]
       index = configure_pip()
       if index:
           pip_cmd.extend(["-i", index])
       subprocess.check_call(pip_cmd)
   
   
   if __name__ == "__main__":
       main()

   How It's Used

   For training (ModelTrainer's sm_train.sh):

   The script is uploaded as part of sm_drivers/scripts/. The INSTALL_REQUIREMENTS template becomes:

   echo "Installing requirements"
   $SM_PYTHON_CMD /opt/ml/input/data/sm_drivers/scripts/install_requirements.py {requirements_file}

   For processing (FrameworkProcessor's runproc.sh):

   The script is embedded in the source tarball or uploaded alongside it. The generated script becomes:

   if [[ -f 'requirements.txt' ]]; then
       python3 install_requirements.py requirements.txt
   fi

   Same script, both paths. No duplication.

   What This Doesn't Require

   • ❌ No container changes
   • ❌ No cross-team coordination
   • ❌ No new packages to install in containers
   • ❌ No changes to sagemaker-training-toolkit or sagemaker-inference-toolkit

   Behavior Matrix

   CA_REPOSITORY_ARN set?	boto3 available?	AWS CLI available?	Result
   No	—	—	Normal pip install (no CodeArtifact)
   Yes	Yes	—	pip install -i <authenticated_url> (boto3)
   Yes	No	Yes	aws codeartifact login then pip install (CLI)
   Yes	No	No	Hard fail with clear error message

   API Enhancement (Optional Layer)

   On top of the stopgap script, the SDK could expose codeartifact_repo_arn as a first-class parameter:

   # On SourceCode (for ModelTrainer)
   SourceCode(entry_script="train.py", source_dir="src",
              requirements="requirements.txt",
              codeartifact_repo_arn="arn:aws:codeartifact:...")
   
   # On FrameworkProcessor.run()
   processor.run(code="script.py", source_dir="src",
                 codeartifact_repo_arn="arn:aws:codeartifact:...")

   This would:

   1. Validate the ARN at submission time (fail fast, not at container runtime)
   2. Auto-set CA_REPOSITORY_ARN in the container environment
   3. Ensure the install_requirements.py script is included in the upload

   ---

   Appendix: Container Inventory

   Current Containers (DLC master branch)

   Container	boto3	awscli	sagemaker-training-toolkit	CodeArtifact?
   PyTorch Training 2.7–2.9	✅ explicit	✅ explicit	✅ explicit	✅ via toolkit
   PyTorch Inference (master)	✅ via toolkit	✅ explicit	— (inference toolkit)	✅ via toolkit

   New Containers (DLC main branch)

   Container	boto3	awscli	Toolkit	CodeArtifact?
   Base v1/v2	❌	❌	❌	❌
   Ray Inference	✅ (requirements.lock)	✅ (requirements.lock)	❌	✅ via sagemaker_serve.py
   vLLM / SGLang	❌	❌	❌	❌ (no requirements.txt support)

   The new DLC base images do NOT include boto3 or awscli. This is why the stopgap script's fallback chain (boto3 → CLI → hard fail) matters — and why the long-term solution uses a shared package that declares boto3 as its own dependency, so installing the package in a container automatically brings boto3 along.

   ---

   Appendix: The Three Existing Implementations

   All three use the same core logic — parse the ARN, get a token, build an authenticated URL — but implemented independently:

   1. sagemaker-training-toolkit (Python/boto3)

   # sagemaker_training/modules.py
   # https://github.com/aws/sagemaker-training-toolkit/blob/master/src/sagemaker_training/modules.py
   client = boto3.client("codeartifact", region_name=region)
   token = client.get_authorization_token(domain=domain, domainOwner=owner)["authorizationToken"]
   endpoint = client.get_repository_endpoint(domain=domain, domainOwner=owner, repository=repository, format="pypi")["repositoryEndpoint"]
   # → pip install -r requirements.txt -i https://aws:{token}@{endpoint}/{repo}/simple/

   2. SDK v2 FrameworkProcessor (Bash/AWS CLI)

   # Generated runproc.sh
   aws codeartifact login --tool pip --domain DOMAIN --domain-owner ACCOUNT --repository REPO --region REGION
   pip install -r requirements.txt

   3. DLC ray inference container (Python/boto3)

   # scripts/ray/sagemaker_serve.py
   # https://github.com/aws/deep-learning-containers/blob/main/scripts/ray/sagemaker_serve.py
   ca = boto3.client("codeartifact", region_name=region)
   token = ca.get_authorization_token(domain=domain, domainOwner=account, durationSeconds=43200)["authorizationToken"]
   endpoint = ca.get_repository_endpoint(domain=domain, domainOwner=account, repository=repository, format="pypi")["repositoryEndpoint"]
   # → pip install -r requirements.txt --index-url https://aws:{token}@{endpoint}/simple/

   The stopgap script consolidates these into one implementation with graceful fallback.
3. humanzz

   changed the title [-][v3] FrameworkProcessor and ModelTrainer: 4 regressions from v2 migration[/-] [+][v3] FrameworkProcessor and ModelTrainer: 4 regressions (including dropping CodeArtifact support) from v2 migration[/+] on Apr 16, 2026
4. humanzz

   added 2 commits that reference this issue on Apr 17, 2026

   feat: add CodeArtifact support for ModelTrainer and FrameworkProcesso…

   ...

   275ee1c

   fix: wire FrameworkProcessor code_location into code upload paths

   ...

   b3301d1
5. aviruthen

   added 2 commits that reference this issue on Apr 28, 2026

   fix: wire FrameworkProcessor code_location into code upload paths (#5773

   ...

   99d974d

   feat: add CodeArtifact support for ModelTrainer and FrameworkProcesso…

   ...

   9386fc0
6. d-vesely commented on May 8, 2026

   

   d-vesely

   on May 8, 2026

   More actions

   I would like to add that I described another regression of the FrameworkProcessor in Issue #5735, regarding the source_dir argument, which used to allow S3 URIs, but now requires a local path.