Description
Hello team,
I have recently worked with several SageMaker teams to deliver on a feature I requested a while back - to support installing requirements.txt
dependencies from a specified CodeArtifact repository in sagemaker training jobs and deployed endpoints/models. You can see the Requests/PRs/Releases at
The GitHub feature request above was meant to cover all (or at least as many) SageMaker images (used in training jobs/inference), but I've prioritized working on delivering this capability in PyTorch 2.0.1 training/inference containers and that has already been delivered.
A while back, a blog post about leveraging CodeArtifact in SageMaker notebooks was published at https://aws.amazon.com/blogs/machine-learning/secure-aws-codeartifact-access-for-isolated-amazon-sagemaker-notebook-instances/ which, in addition to the feature requests above, provides good context on why users might want to install their requirements.txt
dependencies from CodeArtifact.
This is a request to update the PyTorch requirements.txt
support documentation to add details about using CodeArtifact with PyTorch 2.0.1+ e.g. at https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#using-third-party-libraries
Below, I list out the instructions to leverage CodeArtifact
Steps to leverage CodeArtifact in PyTorch 2.0.1
- Set the relevant CodeArtifact environment variable in Training jobs and in Models
- Update the IAM permissions of the SageMaker execution role, and CodeArtifact repository to allow the training job/model
Set the relevant CodeArtifact environment variable in Training jobs and in Models
- The environment variable to set is
CA_REPOSITORY_ARN
and the value is the CodeArtifact Repository ARN - Where the environment variable needs to be set is
- Training Job (Environment): https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateTrainingJob.html#API_CreateTrainingJob_RequestParameters
- Create Model (Environment): https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateModel.html#API_CreateModel_RequestParameters
- Using the SageMaker Python SDK
PyTorch
estimator'senvironment
argument to use in training jobPyTorch
estimator'sfit()
'senv
argument to use in creating the model
Update the IAM permissions of the SageMaker execution role, and CodeArtifact repository to allow the training job/model
-
The SageMaker execution role needs to have a policy that allows access and retrieval of packages from CodeArtifact repository. Examples for this are
AWSCodeArtifactReadOnlyAccess
from https://docs.aws.amazon.com/codeartifact/latest/ug/security_iam_id-based-policy-examples.html
-
The CodeArtifact repository's resource policy needs to allow the SageMaker execution role to execute the necessary actions
Happy to provide any additional context/details if needed.