feat: jumpstart default payloads #4149

evakravi · 2023-10-02T18:40:24Z

Issue #, if available:

Description of changes:
Adds default payload functionality for JumpStart Models. The following workflow is now possible:

from sagemaker.jumpstart.model import JumpStartModel

model = JumpStartModel(model_id=<model_id>)
default_payload = model.retrieve_default_payload()
predictor = model.deploy()
response = predictor.predict(default_payload)

The JumpStartModel class now exposes a function retrieve_default_payload to give customers a starting point for interacting with the model.

Standalone utilities has also been created in sagemaker.payloads module:

retrieve_all_examples
retrieve_example

Testing done:
Integration tests run locally against staged metadata.

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 · 2023-10-02T18:42:45Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-02T18:45:13Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-02T18:52:26Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-02T19:42:16Z

AWS CodeBuild CI Report

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

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

whittech1 · 2023-10-02T20:51:29Z

src/sagemaker/jumpstart/types.py

+            self.accept = accept
+
+    def to_json(self) -> Dict[str, Any]:
+        """Returns json representation of JumpStartDefaultPayloads object."""


nit: supposed to be JumpStartSerializedPayload?

whittech1 · 2023-10-02T20:55:25Z

src/sagemaker/payloads.py

+logger = logging.getLogger(__name__)
+
+
+def retrieve_samples(


nit: reasoning for switching to samples noun?

Should we use defaults? I thought it would be confusing to use default in the context of returning both a single and multiple payloads

haven't we used retrieve_options in other accessors?

whittech1 · 2023-10-02T20:56:29Z

src/sagemaker/payloads.py

+    """
+    if not jumpstart_utils.is_jumpstart_model_input(model_id, model_version):
+        raise ValueError(
+            "Must specify JumpStart `model_id` and `model_version` when retrieving model URIs."


nit: when retrieving payloads?

whittech1 · 2023-10-02T20:59:10Z

src/sagemaker/jumpstart/payload_utils.py

+        self.region = region
+        self.s3_client = s3_client
+
+    def get_bytes_payload_with_s3_references(


I think MH has similar logic atm in the other PR, is the plan to reuse the functions in this PR post-merge?

Yes that's the idea. I was working on this PR when I realized a lot of code is duplicated. So MH will eventually call these utilities.

whittech1 · 2023-10-02T21:00:44Z

src/sagemaker/jumpstart/model.py

+
+        if sample_payloads is None or len(sample_payloads) == 0:
+            raise NotImplementedError(
+                f"No default payload supported for model ID '{self.model_id}'."


nit: error vs returning None?

Not sure, i'm cool with both

are you now? That's something you should definitely have an opinion on Evan.

What I mean is that both experiences can be justified from a UX perspective. I guess, to maintain consistency, we should return None, since we don't typically use NotImplementedError

sagemaker-bot · 2023-10-02T21:28:26Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-02T21:37:23Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-02T22:08:00Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-02T22:21:40Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-02T23:19:28Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-03T22:13:06Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-03T22:21:24Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-03T22:33:39Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-03T22:53:53Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-03T23:08:16Z

AWS CodeBuild CI Report

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

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

JGuinegagne · 2023-10-03T23:24:14Z

src/sagemaker/base_predictor.py

@@ -201,20 +204,42 @@ def _create_request_args(
        custom_attributes=None,
    ):
        """Placeholder docstring"""
+
+        js_accept = None


typing please

also consider spelling out for better readability: js_ -> jumpstart_

JGuinegagne · 2023-10-03T23:29:00Z

src/sagemaker/jumpstart/accessors.py

+
+        Requests are cached so that the same s3 request is never made more
+        than once, unless a different region or client is used.
+        """


hum, not sure about this, you are effectively caching objects in memory aren't you?

Could you:
(a) determine the max memory you are willing to use for such cache?
(b) add a head object with a size limits for such objects
(c) derive the max number of items in the @lru_cache(max_items) from round((a)/(b)) ?

The maximum memory is very system dependent and the payloads will come in all different sizes. How about we expose a function that clears the cache JumpStartS3Accessor.clear_cache()? This can call JumpStartS3Accessor.get_object_cached.cache_clear() under the hood. See: https://stackoverflow.com/questions/37653784/how-do-i-use-cache-clear-on-python-functools-lru-cache

JGuinegagne · 2023-10-03T23:29:56Z

src/sagemaker/jumpstart/accessors.py

+        if s3_client is None:
+            s3_client = JumpStartS3Accessor._get_default_s3_client(region)
+
+        return s3_client.get_object(Bucket=bucket, Key=key)["Body"].read()


this line will buffer the whole object in memory? Is that acceptable and shouldn't you build in safeguards?

What kind of safeguards are you referring to?

for posterity as you addressed above: inference file size mainly.

JGuinegagne · 2023-10-03T23:31:57Z

src/sagemaker/jumpstart/artifacts/payloads.py

+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This module contains functions for obtaining JumpStart payloads."""


review with @judyheflin please. Maybe obtaining example payloads for JumpStart models.

JGuinegagne · 2023-10-03T23:33:23Z

src/sagemaker/jumpstart/model.py

+    def retrieve_default_payload(self) -> JumpStartSerializablePayload:
+        """Returns default payload associated with the model.
+
+        Payload can be directly used with the `sagemaker.predictor.Predictor.predict(...)` function.


this is customer facing: please add a Raises section.

if payloads.retrieve_options can throw, please port over the appropriate item of its Raises section.

Related question: do you need a way to propagate the tolerate_vulnerable and tolerate_deprecated flags?

JGuinegagne · 2023-10-03T23:46:23Z

src/sagemaker/payloads.py

+logger = logging.getLogger(__name__)
+
+
+def retrieve_samples(


haven't we used retrieve_options in other accessors?

JGuinegagne · 2023-10-03T23:47:21Z

src/sagemaker/payloads.py

+        Dict[str, JumpStartSerializablePayload]
+    ] = artifacts._retrieve_default_payloads(
+        model_id,
+        model_version,  # type: ignore


question: why do you need to silence the linter?

JGuinegagne · 2023-10-03T23:49:01Z

src/sagemaker/payloads.py

+
+    unserialized_payloads: List[JumpStartSerializablePayload] = list(
+        unserialized_payload_dict.values()
+    )


here aren't we buffering a list of content-embedded payloads in memory?

Let's sync up offline on this point please.

By default we embed the content in the payload, however we give the option not to.

JGuinegagne · 2023-10-03T23:49:38Z

tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py

+def test_default_payload_jumpstart_model(setup):
+
+    # DO NOT COMMIT THIS LINE
+    os.environ.update({"AWS_JUMPSTART_CONTENT_BUCKET_OVERRIDE": "jumpstart-cache-alpha-us-west-2"})


please don't expose this bucket name in publicly available code.

This won't get committed, but yes I will remove

JGuinegagne · 2023-10-03T23:50:25Z

tests/unit/sagemaker/jumpstart/constants.py

+            "huggingface_transformers_version": "4.17",
+        },
+        "hosting_artifact_key": "stabilityai-infer/infer-model-depth2img-st"
+        "able-diffusion-v1-5-controlnet-v1-1-fp16.tar.gz",


sanity-check: did you intend to jump a line here?

this was just to reduce the line length

…loads

sagemaker-bot · 2023-10-04T16:36:08Z

AWS CodeBuild CI Report

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

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

evakravi

/bot run slow-tests

sagemaker-bot · 2023-10-09T23:31:24Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-09T23:40:24Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-10T00:34:59Z

AWS CodeBuild CI Report

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

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

JGuinegagne · 2023-10-10T01:17:35Z

src/sagemaker/jumpstart/model.py

@@ -312,6 +314,27 @@ def _is_valid_model_id_hook():

        super(JumpStartModel, self).__init__(**model_init_kwargs.to_kwargs_dict())

+    def retrieve_default_payload(self) -> JumpStartSerializablePayload:


sorry I missed this: this instance method needs be renamed retrieve_example_payload()

Please also add a retrieve_all_example_payloads() instance method.

sagemaker-bot · 2023-10-10T01:43:10Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-10T01:45:03Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-10T02:25:09Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-10T02:36:10Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-10T04:19:34Z

AWS CodeBuild CI Report

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

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

evakravi

/bot run notebook-tests, pr, unit-tests

sagemaker-bot · 2023-10-10T12:30:43Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-10T12:35:15Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-10T13:19:42Z

AWS CodeBuild CI Report

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

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

problems addressed

sagemaker-bot · 2023-10-10T17:49:28Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-10T18:03:21Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-10T18:03:35Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-10T18:28:12Z

AWS CodeBuild CI Report

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

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

sagemaker-bot · 2023-10-10T18:44:55Z

AWS CodeBuild CI Report

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

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

feat: jumpstart default payloads

cba9034

evakravi requested a review from a team as a code owner October 2, 2023 18:40

evakravi requested review from jgoyani1 and removed request for a team October 2, 2023 18:40

chore: use lru cache for python compatibility

6082e59

evakravi marked this pull request as draft October 2, 2023 18:43

fix: lru cache annotation

bf13065

evakravi added 2 commits October 2, 2023 19:13

fix: switch order of annotations

9750dc5

fix: docstring, add integ test

5dc4aac

fix: flake8

ca20d87

whittech1 reviewed Oct 2, 2023

View reviewed changes

fix: incorrect words

b2a6374

JGuinegagne requested changes Oct 3, 2023

View reviewed changes

evakravi added 2 commits October 4, 2023 15:42

chore: address PR comments

214e458

Merge remote-tracking branch 'origin' into feat/jumpstart-default-pay…

0a61992

…loads

evakravi commented Oct 9, 2023

View reviewed changes

Merge branch 'master' into feat/jumpstart-default-payloads

2aecf36

JGuinegagne previously requested changes Oct 10, 2023

View reviewed changes

evakravi commented Oct 10, 2023

View reviewed changes

evakravi added the do-not-merge Use when PR needs to be marked as do not merge label Oct 10, 2023

evakravi marked this pull request as draft October 10, 2023 16:28

jgoyani1 approved these changes Oct 10, 2023

View reviewed changes

chore: address comments and fix typing/docstrings

e5a7058

evakravi marked this pull request as ready for review October 10, 2023 16:56

JGuinegagne approved these changes Oct 10, 2023

View reviewed changes

evakravi removed the do-not-merge Use when PR needs to be marked as do not merge label Oct 10, 2023

jgoyani1 approved these changes Oct 10, 2023

View reviewed changes

jgoyani1 merged commit 7f6f3f9 into aws:master Oct 10, 2023

		@@ -312,6 +314,27 @@ def _is_valid_model_id_hook():

		super(JumpStartModel, self).__init__(**model_init_kwargs.to_kwargs_dict())

		def retrieve_default_payload(self) -> JumpStartSerializablePayload:

feat: jumpstart default payloads #4149

feat: jumpstart default payloads #4149

Uh oh!

Conversation

evakravi commented Oct 2, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge Checklist

General

Tests

Uh oh!

sagemaker-bot commented Oct 2, 2023

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Oct 2, 2023

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Oct 2, 2023

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Oct 2, 2023

AWS CodeBuild CI Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

evakravi Oct 4, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sagemaker-bot commented Oct 2, 2023

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Oct 2, 2023

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Oct 2, 2023

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Oct 2, 2023

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Oct 2, 2023

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Oct 3, 2023

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Oct 3, 2023

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Oct 3, 2023

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Oct 3, 2023

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Oct 3, 2023

AWS CodeBuild CI Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

evakravi commented Oct 2, 2023 •

edited

Loading

evakravi Oct 4, 2023 •

edited

Loading

evakravi Oct 4, 2023 •

edited

Loading