feat: JumpStart alternative config parsing #4566

Captainia · 2024-04-08T21:46:51Z

Issue #, if available:

Description of changes:

Initial round of reading and parsing the JumpStart alternative configs.

Class structures that goes from bottom to top:

                       JumpStartMetadataBaseFields 
             (inherits) /                       \(inherits)
       JumpStartModelSpecs             JumpStartConfigComponent   
          (includes)/                            \ (consists)
   JumpStartMetadataConfigs                    JumpStartMetadataConfig
                                                         \ (consists)
                                                JumpStartMetadataConfigs

Testing done:

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.

src/sagemaker/jumpstart/types.py

evakravi · 2024-04-08T23:11:16Z

src/sagemaker/jumpstart/types.py

+        return json_obj
+
+    @property
+    def resolved_config(self) -> Dict[str, Any]:


this logic won't work for dicts inside dicts. can you instead do:

def deep_override_dict(dict1: dict, dict2: dict) -> dict: """Overrides any overlapping contents of dict1 with the contents of dict2.""" flattened_dict = flatten(dict1) flattened_dict.update(flatten(dict2)) return unflatten(flattened_dict) if flattened_dict else {}

I remembered I used something similar before to this but right now I don't see a native Python supported method flatten. There's something similar but doesn't seem it trivially support unflatten: https://pandas.pydata.org/docs/reference/api/pandas.json_normalize.html

src/sagemaker/jumpstart/types.py

evakravi · 2024-04-10T23:34:44Z

src/sagemaker/jumpstart/types.py

+class JumpStartPresetComponent(JumpStartMetadataBaseFields):
+    """Data class of JumpStart config component."""
+
+    slots = ["component_name", "component_override_fields"]


what is component_override_fields?

i guess this is fine, gives us flexibility in the future

It is the fields that this particular component will override to the JumpStartMetadataBaseFields. Each JumpStartPresetComponent inherits the base class so contains all the base fields, but since we only want to override to these specific fields so I kept a separate list only for them.

src/sagemaker/jumpstart/types.py

src/sagemaker/jumpstart/utils.py

src/sagemaker/jumpstart/types.py

src/sagemaker/jumpstart/utils.py

whittech1 · 2024-04-16T13:22:01Z

src/sagemaker/jumpstart/types.py

+                continue
+            return self.configs[config_name]
+
+        return None


The return type should be Optional right?

Yes, I'll address this in my next PR, thanks for catching

akozd · 2024-04-16T13:44:17Z

src/sagemaker/jumpstart/utils.py

+    model_type: enums.JumpStartModelType = enums.JumpStartModelType.OPEN_WEIGHTS,
+) -> Dict[str, List[JumpStartMetadataConfig]]:
+    """Returns metadata configs for the given model ID and region."""
+    model_specs = verify_model_region_and_return_specs(


nit: this block of code appears 3 times in get_jumpstart_configs, get_benchmark_stats, and get_config_names

Yea I think it's required because we use verify_model_region_and_return_specs as the entry point to get specs in all these utils, and each function can be called standalone. But I think we can probably choose a better naming of verify_model_region_and_return_specs and move outside of utils as it's used heavily everywhere.

src/sagemaker/jumpstart/types.py

evakravi · 2024-04-16T14:11:22Z

src/sagemaker/jumpstart/types.py

@@ -1028,22 +1038,23 @@ class JumpStartPresetComponent(JumpStartMetadataBaseFields):

    __slots__ = slots + JumpStartMetadataBaseFields.__slots__

-    def __init__(
+    def __init__(  # pylint: disable=super-init-not-called


why aren't we calling the superclass constructor? This isn't good OOP design IMO, it violates Liskov Substution Principle

I'll update, I left the component to only have its specific fields but not construct all the parent class fields which are optional, but it should work with superclass constructor. Good callout

src/sagemaker/jumpstart/types.py

JGuinegagne · 2024-04-17T11:23:55Z

src/sagemaker/jumpstart/types.py

        """
-        self.from_json(spec)
+        self.from_json(fields)


how does this work, I don't mypy should accept this signature given that from_json accepts a Dict[str, Any] and get passed an Optional[Dict[str, Any]]. This could cause null pointer errors.

Sorry this is confusing typing, the fields object is Dict[str, Any] but not optional. I think mypy is not checking these typings, updated in my next PR

JGuinegagne · 2024-04-17T11:24:27Z

src/sagemaker/jumpstart/types.py

        self.incremental_training_supported: bool = bool(
-            json_obj.get("incremental_training_supported", False)
+            json_obj.get("incremental_training_supported")
        )


all these types are de facto Optional[*] now?

src/sagemaker/jumpstart/types.py

src/sagemaker/jumpstart/utils.py

* add alternative config parsing * rename a few functions * refactor based on new mock * remove commented files * add interfaces and utils * various fixes and add some tests * add more tests * address comments * swap presets with configs * fix: docstyle * fix: doc * fix: docstyle * updates * typing

Captainia added 2 commits April 8, 2024 21:45

add alternative config parsing

35f235a

rename a few functions

d62a8fc

evakravi reviewed Apr 8, 2024

View reviewed changes

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

evakravi reviewed Apr 8, 2024

View reviewed changes

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

evakravi reviewed Apr 8, 2024

View reviewed changes

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

evakravi reviewed Apr 8, 2024

View reviewed changes

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

evakravi reviewed Apr 8, 2024

View reviewed changes

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

whittech1 reviewed Apr 9, 2024

View reviewed changes

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

refactor based on new mock

750b9f0

Captainia temporarily deployed to auto-approve April 10, 2024 22:10 — with GitHub Actions Inactive

remove commented files

6db438e

Captainia temporarily deployed to auto-approve April 10, 2024 22:12 — with GitHub Actions Inactive

evakravi reviewed Apr 10, 2024

View reviewed changes

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

evakravi reviewed Apr 10, 2024

View reviewed changes

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

evakravi reviewed Apr 10, 2024

View reviewed changes

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

evakravi reviewed Apr 10, 2024

View reviewed changes

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

evakravi reviewed Apr 10, 2024

View reviewed changes

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

add interfaces and utils

21a0e60

Captainia temporarily deployed to auto-approve April 11, 2024 21:54 — with GitHub Actions Inactive

various fixes and add some tests

76fee33

Captainia temporarily deployed to auto-approve April 12, 2024 12:59 — with GitHub Actions Inactive

add more tests

cd27c33

Captainia temporarily deployed to auto-approve April 12, 2024 15:51 — with GitHub Actions Inactive

Merge branch 'master' into alternate-config

183a0c9

Captainia marked this pull request as ready for review April 12, 2024 15:54

Captainia requested a review from a team as a code owner April 12, 2024 15:54

Captainia requested review from nargokul and removed request for a team April 12, 2024 15:54

Merge branch 'master' into alternate-config

8dade8a

Captainia temporarily deployed to auto-approve April 12, 2024 23:10 — with GitHub Actions Inactive

fix: docstyle

1481ba8

Captainia temporarily deployed to auto-approve April 15, 2024 14:21 — with GitHub Actions Inactive

fix: doc

4d7a31a

Captainia temporarily deployed to auto-approve April 15, 2024 14:40 — with GitHub Actions Inactive

fix: docstyle

decc22e

Captainia temporarily deployed to auto-approve April 15, 2024 15:04 — with GitHub Actions Inactive

makungaj1 approved these changes Apr 15, 2024

View reviewed changes

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

updates

c19a28a

Captainia temporarily deployed to auto-approve April 15, 2024 19:12 — with GitHub Actions Inactive

typing

866956c

Captainia temporarily deployed to auto-approve April 15, 2024 20:20 — with GitHub Actions Inactive

qiyunz reviewed Apr 15, 2024

View reviewed changes

src/sagemaker/jumpstart/types.py Outdated Show resolved Hide resolved

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

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

Merge branch 'master' into alternate-config

4850e4f

Captainia temporarily deployed to auto-approve April 16, 2024 13:17 — with GitHub Actions Inactive

whittech1 approved these changes Apr 16, 2024

View reviewed changes

Merge branch 'master' into alternate-config

1a21b82

Captainia temporarily deployed to auto-approve April 16, 2024 13:38 — with GitHub Actions Inactive

akozd reviewed Apr 16, 2024

View reviewed changes

akozd approved these changes Apr 16, 2024

View reviewed changes

knikure self-assigned this Apr 16, 2024

evakravi reviewed Apr 16, 2024

View reviewed changes

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

evakravi reviewed Apr 16, 2024

View reviewed changes

evakravi approved these changes Apr 16, 2024

View reviewed changes

knikure approved these changes Apr 16, 2024

View reviewed changes

knikure merged commit 8bb857f into aws:master Apr 16, 2024
10 checks passed

JGuinegagne reviewed Apr 17, 2024

View reviewed changes

feat: JumpStart alternative config parsing #4566

feat: JumpStart alternative config parsing #4566

Uh oh!

Conversation

Captainia commented Apr 8, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge Checklist

General

Tests

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Captainia commented Apr 8, 2024 •

edited

Loading