feature: Adding Jumpstart retrieval functions #2789

evakravi · 2021-12-07T21:48:16Z

Issue #, if available:

Description of changes:
This PR implements retrieval functions for JumpStart model specs, allowing training jobs and inference endpoints to be launched with pre-trained models in script mode. More specifically, this PR modifies sagemaker.model_uris.retrieve function so that container images compatible with the requested JumpStrart version, model ids, and script (training or inference) can be used. In addition, this PR adds functions sagemaker.script_uris.retrieve and sagemaker.model_uris.retrieve to allow users to retrieve s3 uris for model scripts and tarballs (respectively).

These utility functions rely on the the JumpStart model cache implemented in #2756, with an additional utility developed in this cache sagemaker.jumpstart.accessors to create a singleton for the model cache.

This PR also removes the semantic-version library from the package dependencies. To deal with versions, JumpStart will instead use packaging.versions module.
Testing done:
This PR includes unit tests for the retrieve functions, verifying that given a mocked spec file, the functions return the expected value. In addition, integration tests were written that check that the actual retrieve functions, using live jumpstart-hosted content, can successfully launch sagemaker inference and training jobs.

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 backword 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.

codecov-commenter · 2021-12-07T22:01:00Z

Codecov Report

Merging #2789 (8d95d30) into master-jumpstart (b09793a) will increase coverage by 0.13%.
The diff coverage is 95.29%.

@@                 Coverage Diff                  @@
##           master-jumpstart    #2789      +/-   ##
====================================================
+ Coverage             88.76%   88.90%   +0.13%     
====================================================
  Files                   179      181       +2     
  Lines                 15734    15855     +121     
====================================================
+ Hits                  13967    14096     +129     
+ Misses                 1767     1759       -8

Impacted Files	Coverage Δ
src/sagemaker/sklearn/estimator.py	`100.00% <ø> (ø)`
src/sagemaker/xgboost/estimator.py	`98.18% <ø> (ø)`
src/sagemaker/jumpstart/artifacts.py	`89.61% <89.61%> (ø)`
src/sagemaker/script_uris.py	`92.30% <92.30%> (ø)`
src/sagemaker/model_uris.py	`92.85% <92.85%> (ø)`
src/sagemaker/image_uris.py	`98.81% <95.23%> (-0.51%)`	⬇️
src/sagemaker/estimator.py	`90.53% <100.00%> (-0.06%)`	⬇️
src/sagemaker/jumpstart/accessors.py	`100.00% <100.00%> (ø)`
src/sagemaker/jumpstart/cache.py	`96.49% <100.00%> (+0.61%)`	⬆️
src/sagemaker/jumpstart/types.py	`92.30% <100.00%> (+2.51%)`	⬆️
... and 5 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update b09793a...8d95d30. Read the comment docs.

sagemaker-bot · 2021-12-07T22:03:41Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2021-12-07T22:44:59Z

AWS CodeBuild CI Report

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

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

src/sagemaker/image_uris.py

src/sagemaker/jumpstart/accessors.py

tests/unit/sagemaker/jumpstart/test_cache.py

tests/unit/sagemaker/test_model_uris.py

tests/unit/sagemaker/test_script_uris.py

tests/unit/sagemaker/jumpstart/utils.py

…and code

sagemaker-bot · 2021-12-13T21:12:56Z

AWS CodeBuild CI Report

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

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

…e uris

sagemaker-bot · 2021-12-16T14:18:01Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2021-12-29T00:04:14Z

AWS CodeBuild CI Report

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

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

JGuinegagne

suggestion 1: you are running all the unit tests in us-west-2. Since you are testing all frameworks, you could also vary the regions to ensure the bucket mapping is correct.

suggestion 2: it is important to unit test the retrieve image URI utility for all frameworks because of the framework-specific subtleties, but I don't think it's necessary for the script and model_artifact utilities. For these two utils, maybe it's enough to test one framework and reduce the size of the PR. WDYT?

src/sagemaker/estimator.py

src/sagemaker/image_uris.py

JGuinegagne · 2021-12-29T12:40:26Z

src/sagemaker/jumpstart/accessors.py

+class SageMakerSettings(object):
+    """Static class for storing the SageMaker settings."""
+
+    _PARSED_SAGEMAKER_VERSION = ""


non-blocking: technically, this static attribute isn't a constant, so lower case might be more appropriate.

src/sagemaker/jumpstart/accessors.py

tests/integ/sagemaker/jumpstart/retrieve_uri/test_inference.py

tests/integ/sagemaker/jumpstart/retrieve_uri/constants.py

tests/integ/sagemaker/jumpstart/retrieve_uri/utils.py

sagemaker-bot · 2021-12-30T00:36:03Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2021-12-30T01:22:54Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2021-12-30T01:31:40Z

AWS CodeBuild CI Report

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

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

src/sagemaker/jumpstart/artifacts.py

src/sagemaker/model_uris.py

src/sagemaker/script_uris.py

tests/integ/sagemaker/jumpstart/retrieve_uri/training.py

tests/unit/sagemaker/jumpstart/utils.py

JGuinegagne

A few clean-up comments, otherwise LGTM

src/sagemaker/jumpstart/artifacts.py

sagemaker-bot · 2021-12-30T15:22:17Z

AWS CodeBuild CI Report

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

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

src/sagemaker/jumpstart/accessors.py

shreyapandit · 2022-01-05T22:56:06Z

src/sagemaker/jumpstart/accessors.py

+        if JumpStartModelsCache._cache is None or region != JumpStartModelsCache._curr_region:
+            JumpStartModelsCache._cache = cache.JumpStartModelsCache(region=region, **cache_kwargs)
+            JumpStartModelsCache._curr_region = region


nit: @evakravi This section is same as lines 80-82 above, can we refactor this to avoid duplicacy?

…Start cache

sagemaker-bot · 2022-01-06T15:34:29Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-06T16:26:20Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-06T16:47:24Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-07T19:45:13Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-07T20:57:06Z

AWS CodeBuild CI Report

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

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

src/sagemaker/sklearn/estimator.py

sagemaker-bot · 2022-01-07T22:44:48Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-07T23:04:38Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-07T23:46:38Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-07T23:57:42Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-08T00:41:46Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-10T14:48:15Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-10T15:40:10Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-10T15:50:32Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-11T19:56:31Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-11T20:34:02Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-11T20:37:17Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-12T00:12:43Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-12T00:53:39Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2022-01-12T01:48:02Z

AWS CodeBuild CI Report

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

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

evakravi added 4 commits December 7, 2021 15:38

feat: jumpstart retrieve functions (wip)

fa79a5d

change: use packaging library for jumpstart versions

84bf597

fix: linting, mypy, logical issues for jumpstart models

86cf577

fix: tox.ini

fea1020

change: retrieve script uri argument name, comments

f85b7f0

JGuinegagne requested changes Dec 8, 2021

View reviewed changes

change: improve jumpstart retrieve fx impl, cleanup tests, comments, …

ed7b772

…and code

change: improve jumpstart retrieve uri unit tests, fix logic for imag…

a427d4b

…e uris

shreyapandit changed the title ~~Feat/jumpstart retrieve functions~~ feature: jumpstart retrieve functions Dec 16, 2021

feat: integration tests for jumpstart sdk retrieve functions

f6ade25

evakravi marked this pull request as ready for review December 29, 2021 01:02

JGuinegagne requested changes Dec 29, 2021

View reviewed changes

change: cleanup code, remove redundant tests

5b98f42

JGuinegagne reviewed Dec 30, 2021

View reviewed changes

JGuinegagne requested changes Dec 30, 2021

View reviewed changes

src/sagemaker/jumpstart/artifacts.py Show resolved Hide resolved

change: minor updates to jumpstart retrieve functions

64b51d6

JGuinegagne approved these changes Dec 30, 2021

View reviewed changes

shreyapandit reviewed Jan 5, 2022

View reviewed changes

src/sagemaker/jumpstart/accessors.py Show resolved Hide resolved

shreyapandit reviewed Jan 5, 2022

View reviewed changes

change: create helper function for setting kwargs and region for Jump…

4e4590c

…Start cache

mufaddal-rohawala approved these changes Jan 7, 2022

View reviewed changes

shreyapandit approved these changes Jan 7, 2022

View reviewed changes

mufaddal-rohawala reviewed Jan 7, 2022

View reviewed changes

src/sagemaker/sklearn/estimator.py Outdated Show resolved Hide resolved

fix: image_uri_region argument formatting

42743aa

fix: jumpstart works without region in aws config

8d95d30

shreyapandit approved these changes Jan 10, 2022

View reviewed changes

Merge branch 'master-jumpstart' into feat/jumpstart-retrieve-functions

c3d158c

fix: reduce MaxRuntimeInSeconds for jumpstart integ test training jobs

9a396de

mufaddal-rohawala approved these changes Jan 12, 2022

View reviewed changes

shreyapandit changed the title ~~feature: jumpstart retrieve functions~~ feature: Adding Jumpstart retrieval functions Jan 12, 2022

shreyapandit merged commit b691d3d into aws:master-jumpstart Jan 12, 2022

feature: Adding Jumpstart retrieval functions #2789

feature: Adding Jumpstart retrieval functions #2789

Uh oh!

Conversation

evakravi commented Dec 7, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge Checklist

General

Tests

Uh oh!

codecov-commenter commented Dec 7, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

sagemaker-bot commented Dec 7, 2021

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Dec 7, 2021

AWS CodeBuild CI Report

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sagemaker-bot commented Dec 13, 2021

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Dec 16, 2021

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Dec 29, 2021

AWS CodeBuild CI Report

Uh oh!

JGuinegagne left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

JGuinegagne Dec 29, 2021

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sagemaker-bot commented Dec 30, 2021

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Dec 30, 2021

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Dec 30, 2021

AWS CodeBuild CI Report

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

JGuinegagne left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sagemaker-bot commented Dec 30, 2021

AWS CodeBuild CI Report

Uh oh!

Uh oh!

shreyapandit Jan 5, 2022

Choose a reason for hiding this comment

Uh oh!

evakravi commented Dec 7, 2021 •

edited

Loading

codecov-commenter commented Dec 7, 2021 •

edited

Loading