feature: Trainium Neuron support for PyTorch #3423

knikure · 2022-10-18T21:48:25Z

Issue #, if available:

Description of changes:
Added the necessary changes to incorporate the v1.0-pt-1.11.0-tr-neuron-sdk2.3.0-py38 image into the PyTorch framework.

Testing done:

Locally tested the feature to extract image string with the following code, where the output matched with the released DLC image

import sagemaker, boto3

sagemaker.image_uris.retrieve(framework="pytorch",
    region=boto3.Session().region_name,
    instance_type="ml.trn1.xlarge")

'763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training-neuron:1.11.0-neuron-py38-sdk2.3.0-ubuntu20.04'

Merge Checklist

Put an x in the boxes that apply. You can also fill these out after creating the PR. If you're unsure about any of them, don't hesitate to ask. We're here to help! This is simply a reminder of what we are going to look for before merging your pull request.

General

I have read the CONTRIBUTING doc
I certify that the changes I am introducing will be backward compatible, and I have discussed concerns about this, if any, with the Python SDK team
I used the commit message format described in CONTRIBUTING
I have passed the region in to all S3 and STS clients that I've initialized as part of this change.
I have updated any necessary documentation, including READMEs and API docs (if appropriate)

Tests

I have added tests that prove my fix is effective or that my feature works (if appropriate)
I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes
I have checked that my tests are not configured for a specific region or account (if appropriate)
I have used unique_name_from_base to create resource names in integ tests (if appropriate)

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

sagemaker-bot · 2022-10-18T21:54:13Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-slow-tests
Commit ID: 464cb72
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-18T21:57:07Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 464cb72
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-18T22:33:15Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: 464cb72
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-18T22:53:45Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: 464cb72
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-18T23:03:46Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: ec0212c
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-18T23:40:09Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: ec0212c
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-18T23:50:50Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-pr
Commit ID: ec0212c
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-18T23:58:59Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: ec0212c
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-19T01:47:17Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-slow-tests
Commit ID: ec0212c
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

knikure

/bot run unit-tests

sagemaker-bot · 2022-10-19T16:53:17Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-slow-tests
Commit ID: ec0212c
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-19T17:32:06Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: c05b46b
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

codecov-commenter · 2022-10-19T17:34:04Z

Codecov Report

Merging #3423 (6371ba1) into master (1fa2377) will decrease coverage by 0.35%.
The diff coverage is 83.59%.

@@            Coverage Diff             @@
##           master    #3423      +/-   ##
==========================================
- Coverage   89.17%   88.82%   -0.36%     
==========================================
  Files         204      205       +1     
  Lines       18979    19458     +479     
==========================================
+ Hits        16924    17283     +359     
- Misses       2055     2175     +120

Impacted Files	Coverage Δ
src/sagemaker/sklearn/model.py	`92.98% <ø> (-1.76%)`	⬇️
...sagemaker/workflow/monitor_batch_transform_step.py	`0.00% <0.00%> (ø)`
src/sagemaker/xgboost/model.py	`96.36% <ø> (-1.82%)`	⬇️
...agemaker/model_monitor/clarify_model_monitoring.py	`89.84% <27.58%> (-5.09%)`	⬇️
src/sagemaker/session.py	`71.46% <50.00%> (-0.21%)`	⬇️
src/sagemaker/training_compiler/config.py	`80.43% <60.00%> (-2.90%)`	⬇️
src/sagemaker/pytorch/estimator.py	`94.73% <63.63%> (-5.27%)`	⬇️
src/sagemaker/sklearn/estimator.py	`93.93% <63.63%> (-6.07%)`	⬇️
src/sagemaker/model_monitor/dataset_format.py	`72.72% <70.00%> (-2.28%)`	⬇️
src/sagemaker/estimator.py	`89.46% <75.00%> (-0.05%)`	⬇️
... and 35 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

sagemaker-bot · 2022-10-19T17:46:32Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: c05b46b
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-19T18:02:58Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: c05b46b
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-19T18:21:46Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-slow-tests
Commit ID: c05b46b
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-19T20:06:38Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-pr
Commit ID: c05b46b
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-19T22:03:32Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 5e6372c
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-19T22:48:15Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: 5e6372c
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-19T23:04:47Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: 5e6372c
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-19T23:07:19Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 5e6372c
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-19T23:19:48Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-slow-tests
Commit ID: 5e6372c
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-20T02:17:28Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-pr
Commit ID: 5e6372c
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

mufaddal-rohawala · 2022-10-20T17:09:25Z

src/sagemaker/image_uris.py

@@ -30,6 +30,7 @@

 ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
 HUGGING_FACE_FRAMEWORK = "huggingface"
+TRAINIUM_ALLOWED_FRAMEWORKS = "pytorch"


We will need to grab this from json file. Can we add this action to the pending work post re-invent.

sagemaker-bot · 2022-10-24T23:38:12Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: bc5f8e1
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-25T00:19:11Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: bc5f8e1
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-26T00:35:08Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-slow-tests
Commit ID: fafc82c
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-26T00:54:53Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-pr
Commit ID: fafc82c
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-26T17:32:12Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 4625d13
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-26T17:43:27Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: 4625d13
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-26T17:50:44Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-pr
Commit ID: 4625d13
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-26T17:59:46Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: 4625d13
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-26T19:48:35Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-slow-tests
Commit ID: 4625d13
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

tejaschumbalkar

I see some references to sdk2.3.0, i will leave it up to you if you want to keep it as is

sagemaker-bot · 2022-10-27T20:29:49Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 6371ba1
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-27T20:38:48Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: 6371ba1
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

mchoi8739

please remove the documentation

sagemaker-bot · 2022-10-27T21:00:30Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-pr
Commit ID: 6371ba1
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-27T21:01:38Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: 6371ba1
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-27T22:46:16Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-slow-tests
Commit ID: 6371ba1
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

knikure

/bot run slow-tests

sagemaker-bot · 2022-10-27T23:03:42Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-slow-tests
Commit ID: 6371ba1
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2022-10-28T00:23:26Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-slow-tests
Commit ID: 6371ba1
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

Co-authored-by: Miyoung Choi <[email protected]>

knikure force-pushed the trainium-new branch from 464cb72 to ec0212c Compare October 18, 2022 22:55

knikure requested a review from mufaddal-rohawala October 18, 2022 23:00

knikure commented Oct 19, 2022

View reviewed changes

claytonparnell force-pushed the trainium-new branch from ec0212c to c05b46b Compare October 19, 2022 17:02

knikure force-pushed the trainium-new branch from 7baad15 to 5e6372c Compare October 19, 2022 21:59

mufaddal-rohawala approved these changes Oct 20, 2022

View reviewed changes

knikure force-pushed the trainium-new branch from fafc82c to 4625d13 Compare October 26, 2022 16:55

knikure and others added 3 commits October 27, 2022 19:03

Trainium Neuron support for PyTorch

a53fbc1

re-release the trn doc

a5e6753

Update Neuron SDK Version

6008732

knikure force-pushed the trainium-new branch from 4625d13 to 6008732 Compare October 27, 2022 19:41

tejaschumbalkar approved these changes Oct 27, 2022

View reviewed changes

tejaschumbalkar reviewed Oct 27, 2022

View reviewed changes

Fix unit test for sdk version

6371ba1

aws-vrnatham approved these changes Oct 27, 2022

View reviewed changes

mchoi8739 suggested changes Oct 27, 2022

View reviewed changes

knikure commented Oct 27, 2022

View reviewed changes

knikure merged commit ce7e4cd into aws:master Oct 27, 2022

knikure deleted the trainium-new branch October 27, 2022 23:17

mizanfiu pushed a commit to mizanfiu/sagemaker-python-sdk that referenced this pull request Dec 13, 2022

feature: Trainium Neuron support for PyTorch (aws#3423)

c56e1e1

Co-authored-by: Miyoung Choi <[email protected]>

JoseJuan98 pushed a commit to JoseJuan98/sagemaker-python-sdk that referenced this pull request Mar 4, 2023

feature: Trainium Neuron support for PyTorch (aws#3423)

4abd7ee

Co-authored-by: Miyoung Choi <[email protected]>

nmadan pushed a commit to nmadan/sagemaker-python-sdk that referenced this pull request Apr 18, 2023

feature: Trainium Neuron support for PyTorch (aws#3423)

9e8dd82

Co-authored-by: Miyoung Choi <[email protected]>

feature: Trainium Neuron support for PyTorch #3423

feature: Trainium Neuron support for PyTorch #3423

Conversation

knikure commented Oct 18, 2022

Merge Checklist

General

Tests

sagemaker-bot commented Oct 18, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 18, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 18, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 18, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 18, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 18, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 18, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 18, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 19, 2022

AWS CodeBuild CI Report

knikure left a comment • edited Loading

Choose a reason for hiding this comment

sagemaker-bot commented Oct 19, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 19, 2022

AWS CodeBuild CI Report

codecov-commenter commented Oct 19, 2022 • edited Loading

Codecov Report

sagemaker-bot commented Oct 19, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 19, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 19, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 19, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 19, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 19, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 19, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 19, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 19, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 20, 2022

AWS CodeBuild CI Report

mufaddal-rohawala Oct 20, 2022

Choose a reason for hiding this comment

sagemaker-bot commented Oct 24, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 25, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 26, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 26, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 26, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 26, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 26, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 26, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 26, 2022

AWS CodeBuild CI Report

tejaschumbalkar left a comment

Choose a reason for hiding this comment

sagemaker-bot commented Oct 27, 2022

AWS CodeBuild CI Report

sagemaker-bot commented Oct 27, 2022

AWS CodeBuild CI Report

mchoi8739 left a comment

knikure left a comment •

edited

Loading

codecov-commenter commented Oct 19, 2022 •

edited

Loading

knikure left a comment •

edited

Loading