tests: Implement integration tests covering JumpStart PrivateHub workflows #4883

malav-shastri · 2024-10-03T07:34:10Z

Issue #, if available:

Description of changes:
Adding 6 new integ test cases for PrivateHub functionalities/workflows

Added tests:
test_private_hub
test_hub_model_reference
test_jumpstart_hub_model
test_jumpstart_hub_gated_model
test_jumpstart_gated_model_inference_component_enabled
test_instatiating_model

Testing done:

all the newly added tests are passing
integ tests
unit tests
black -l 100 .
flake 8

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.

…flows

Captainia

There are some errors in integ test hook Hub with name PySDK-HubTest-7d6019c6-e043-484e-83c4-fa08706ec4f2 does not exist., could be the tests are ran in parallel

Captainia · 2024-10-03T13:41:53Z

tests/integ/sagemaker/jumpstart/conftest.py

+            sagemaker_session.delete_hub(hub["HubName"])
+
+
+def _delete_hub_contents(sagemaker_session, test_hub_name):


nit: this deletes only the model references (not models or other contents), recommend to take in a optional arg of which content type to delete

Captainia · 2024-10-03T13:45:33Z

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

+    JUMPSTART_LOGGER.info("starting test")
+    JUMPSTART_LOGGER.info(f"get identity {get_sm_session().get_caller_identity_arn()}")
+
+    model_id = "catboost-classification-model"


can we make TEST_MODEL_IDS a enum? so you don't need to repeat

Captainia · 2024-10-03T13:47:21Z

tests/integ/sagemaker/jumpstart/private_hub/test_hub.py

+    # Createhub
+    create_hub_response = hub_instance.create(
+        description="This is a Test Private Hub.",
+        display_name="malavhs Test hub",


Let's not commit this to mainline

Captainia · 2024-10-03T13:49:08Z

tests/integ/sagemaker/jumpstart/utils.py

+        )
+        models.extend(response["hub_content_summaries"])
+
+    return models[0]["hub_content_arn"]


If we only care about the first model arn, do we need to paginate all the responses?

you are right we don't need to paginate here, infact this function only returns a single model because I am filtering on the exact model id. But the output is list so I am accessing the first element.

can you move this to the unit tests? integ tests are only for when we need to create new resources.

@evakravi sorry I didn't get that. this is the utility function to get the Public hub content arn which in turn we use for creating the model reference in the Hub, why do we need to move it to unit tests?

Captainia · 2024-10-03T13:52:29Z

tests/integ/sagemaker/jumpstart/conftest.py

+    list_hub_response = sagemaker_session.list_hubs(name_contains=HUB_NAME_PREFIX)
+
+    for hub in list_hub_response["HubSummaries"]:
+        if hub["HubName"] != SM_JUMPSTART_PUBLIC_HUB_NAME:


While we don't delete public hub, should we also restrict to only delete the hubs that are associated with a pysdk integ test? This can delete the hubs we don't want to in the account that runs the integ test

I am already filtering the hubs related only to PySDK integ test on line 142
sagemaker_session.list_hubs(name_contains=HUB_NAME_PREFIX)

Concurrent runs of the test will cause an issue with this clean-up strategy. Should it clean-up the hub created by this specific test?

yeah i actually had to change it later on, I have updated the commit

isn't the suite id unique to each run? That would prevent overlapping test sessions on the same account/region from interfering with each other.

@evakravi yes but I wasn't deleting the unique Hub created by a specific Pytest session. I was rather deleting all the hubs created by PySDK test regardless of whether it was created by a specific pytest session or not. PySDK runs this command to run integ tests pytest tests/integ -m 'not local_mode and not cron and not slow_test' -n 240 which means there are 240 parallel workers and each would corresponds to a separate pytest session. Because I am not cleaning up specific test session created hub, this will delete the ones which are being used by other pytest sessions and tests in parallel.

AWS-pratab

We should add tests for the fine-tuning/training paths as well.

AWS-pratab · 2024-10-03T18:10:39Z

tests/integ/sagemaker/jumpstart/conftest.py

+    list_hub_response = sagemaker_session.list_hubs(name_contains=HUB_NAME_PREFIX)
+
+    for hub in list_hub_response["HubSummaries"]:
+        if hub["HubName"] != SM_JUMPSTART_PUBLIC_HUB_NAME:


Concurrent runs of the test will cause an issue with this clean-up strategy. Should it clean-up the hub created by this specific test?

AWS-pratab · 2024-10-03T18:14:02Z

tests/integ/sagemaker/jumpstart/constants.py

 JUMPSTART_TAG = "JumpStart-SDK-Integ-Test-Suite-Id"

+SM_JUMPSTART_PUBLIC_HUB_NAME = "SageMakerPublicHub"


nit: Likely this is already defined somewhere in src

AWS-pratab · 2024-10-03T18:15:00Z

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

+TEST_MODEL_IDS = {
+    "catboost-classification-model",
+    "huggingface-txt2img-conflictx-complex-lineart",
+    "meta-textgeneration-llama-2-7b",
+    "meta-textgeneration-llama-3-2-1b",
+    "catboost-regression-model",
+}


I think the integ test runs in PDX. Double check these are available in PDX region.

these should be, I have chosen these models from the existing Jumpstart hub integ tests

AWS-pratab · 2024-10-03T18:20:59Z

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

+        model_id=model_id,
+        role=get_sm_session().get_caller_identity_arn(),
+        sagemaker_session=get_sm_session(),
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],


Can we test via the HubArn path as well, since we noticed an issue around that once.

sorry didn't get it, this is the HubArn path right? we ask customers to provide hub_name in the jumpstart model parameter, but we convert it into HubArn right after it. Sure customers can provide arn directly to model class but in that case we just leave it as it is and that gets passed to the rest of the code.

AWS-pratab · 2024-10-03T18:23:21Z

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

+    )
+
+    # uses ml.m5.4xlarge instance
+    model.deploy(


We should assert the success status of the endpoint, by adding a wait step to poll on status.

+1, we may even want to consider sending it one request using the default payload, although that's optional.

I think I have missed that default payload request here. Similar to other inference tests let me add it here

malav-shastri · 2024-10-03T22:27:27Z

We should add tests for the fine-tuning/training paths as well.

I have purposely skipped it because training isn't supported for model references. And training for jumpstart model is already being covered in the test suit.

evakravi · 2024-10-03T23:01:48Z

src/sagemaker/jumpstart/model.py

@@ -1036,13 +1036,15 @@ def _get_deployment_configs(
                image_uri=image_uri,
                region=self.region,
                model_version=self.model_version,
+                hub_arn=self.hub_arn,


can we add a unit test for this? seems like the current coverage didn't cover this bug

synced with @evakravi offline, this needs to be covered through unit tests and I'll add it as a fast follow.

evakravi · 2024-10-03T23:02:15Z

tests/integ/sagemaker/jumpstart/conftest.py

 )

 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME


 def _setup():
    print("Setting up...")
-    os.environ.update({ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID: get_test_suite_id()})
+    test_suit_id = get_test_suite_id()


nit: test_suite_id

evakravi · 2024-10-03T23:03:20Z

tests/integ/sagemaker/jumpstart/conftest.py

+    hub = Hub(
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session()
+    )
+    hub.create(description=test_hub_description)


should we necessarily create a Hub every time a JS integ test is run? Does this bring any problems?

I don't think of any problems tbh with this strategy, we're cleaning it up in the end. do you think we should be approaching it differently?

evakravi · 2024-10-03T23:04:53Z

tests/integ/sagemaker/jumpstart/conftest.py

+def _delete_hubs(sagemaker_session):
+    # list Hubs created by PySDK integration tests
+    list_hub_response = sagemaker_session.list_hubs(
+        name_contains=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME]


can we create a utility to get the hub name from the env var?

just to confirm you mean a function to get os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME] ?

question: would there ever be several hubs here?
The _setup() method seems to create only one hub, and since the name is set there can only be one, can it not?

@JGuinegagne yeah there'll be only one, I changed this recently previously I was deleting all the hubs starting with a specific prefix but that's messing up with the concurrent pytest executions. Let me no use list_hubs here, thanks.

evakravi · 2024-10-03T23:08:48Z

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

+    assert model.inference_component_name == predictor.component_name
+
+
+def test_instatiating_model(setup, add_models):


i wonder if we can run all the tests currently for JumpStart but for PrivateHub.

RIsk is that this test coverage drifts from the non-private-hub tests as new models/features are added.

I wonder if we can run all the tests currently for JumpStart but for PrivateHub.

I am not sure, are you thinking about reusing the test code of Jumpstart public hub specific tests and integrate PrivateHub workflow triggers there? won't it be too complicated? keeping separate increases the readability? also not all features available for content type Model is available for content type ModelReferences?

RIsk is that this test coverage drifts from the non-private-hub tests as new models/features are added.

sorry I can't understand what you mean by this.

I think Evan is suggesting defining a single set of model tests, and systematically test them against the public hub and a private hub. That would be substantial rework from your current PR though.

fix typo please: test_instantiating_model

I think Evan is suggesting defining a single set of model tests, and systematically test them against the public hub and a private hub. That would be substantial rework from your current PR though.

+1 That would require more efforts and rework. I can take it as an improvement but my question is I want to understand what's the benefit of it? is it better design or just a personal preference to implement it?

JGuinegagne · 2024-10-04T13:09:41Z

tests/integ/sagemaker/jumpstart/conftest.py

+def _delete_hubs(sagemaker_session):
+    # list Hubs created by PySDK integration tests
+    list_hub_response = sagemaker_session.list_hubs(
+        name_contains=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME]


question: would there ever be several hubs here?
The _setup() method seems to create only one hub, and since the name is set there can only be one, can it not?

JGuinegagne · 2024-10-04T13:11:35Z

tests/integ/sagemaker/jumpstart/conftest.py

+        if hub["HubName"] != SM_JUMPSTART_PUBLIC_HUB_NAME:
+            # delete all hub contents first
+            _delete_hub_contents(sagemaker_session, hub["HubName"])
+            sagemaker_session.delete_hub(hub["HubName"])


careful, if there are more than one hub to delete, you might get throttled by the TPS limit of 1.

there'll be only 1 hub at a time

JGuinegagne · 2024-10-04T13:12:05Z

tests/integ/sagemaker/jumpstart/conftest.py

+            hub_name=test_hub_name,
+            hub_content_type=HubContentType.MODEL_REFERENCE.value,
+            hub_content_name=models["HubContentName"],
+        )


careful with throttling here.

JGuinegagne · 2024-10-04T13:13:43Z

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

+
+
+@pytest.fixture(scope="session")
+def add_models():


nit: consider renaming add_model_references

sure, thanks

JGuinegagne · 2024-10-04T13:15:29Z

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

+    )
+
+    # uses ml.m5.4xlarge instance
+    model.deploy(


+1, we may even want to consider sending it one request using the default payload, although that's optional.

JGuinegagne · 2024-10-04T13:21:24Z

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

+    assert model.inference_component_name == predictor.component_name
+
+
+def test_instatiating_model(setup, add_models):


I think Evan is suggesting defining a single set of model tests, and systematically test them against the public hub and a private hub. That would be substantial rework from your current PR though.

JGuinegagne · 2024-10-04T13:21:54Z

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

+    assert model.inference_component_name == predictor.component_name
+
+
+def test_instatiating_model(setup, add_models):


fix typo please: test_instantiating_model

JGuinegagne · 2024-10-04T13:23:21Z

tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py

+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session()
+    )
+
+    # Create Model Reference


nit: unnecessary comment

JGuinegagne · 2024-10-04T13:25:42Z

tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py

+    # Describe Model
+    describe_model_response = hub_instance.describe_model(model_name=model_id)
+    assert describe_model_response is not None
+    assert type(describe_model_response) == DescribeHubContentResponse


optional: consider asserting that

the HubContentName corresponds to the model_id

the HubContentVersion corresponds to the latest version in the public hub

the HubContentType corresponds is a ModelReference

+1, could you address this if you have time? I think what we are currently checking is pretty shallow, no harm to assert these in the tests.

JGuinegagne · 2024-10-04T13:26:35Z

tests/integ/sagemaker/jumpstart/utils.py

@@ -115,6 +116,20 @@ def download_file(local_download_path, s3_bucket, s3_key, s3_client) -> None:
    s3_client.download_file(s3_bucket, s3_key, local_download_path)


+def get_public_hub_model_arn(hub: Hub, model_id: str) -> str:
+    filter_value = f"model_id == {model_id}"
+    response = hub.list_sagemaker_public_hub_models(filter=filter_value)


question: why not use a describe_ method?

Or do you intend to test discovery through list? In that case, please rename the method accordingly.

this is a util function for these tests not a test itself. use of describe method would get us a public hub content arn which would consist model version in it at the end. Create Model Reference don't accept arn with model version in it. On the other side we have implemented this list method in a way that it gets us model name and a public hub arn which can be accepted by create_model_reference api call.

Captainia · 2024-10-07T23:35:12Z

tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py

+    # Describe Model
+    describe_model_response = hub_instance.describe_model(model_name=model_id)
+    assert describe_model_response is not None
+    assert type(describe_model_response) == DescribeHubContentResponse


+1, could you address this if you have time? I think what we are currently checking is pretty shallow, no harm to assert these in the tests.

malavhs added 3 commits October 3, 2024 07:15

tests: Implement integration tests covering JumpStart PrivateHub work…

8910f50

…flows

linting

bfeb2c0

formating

705ceb9

malav-shastri requested a review from a team as a code owner October 3, 2024 07:34

malav-shastri requested a review from liujiaorr October 3, 2024 07:34

malav-shastri temporarily deployed to auto-approve October 3, 2024 07:34 — with GitHub Actions Inactive

malav-shastri changed the title ~~Ch integ tests~~ tests: Implement integration tests covering JumpStart PrivateHub workflows Oct 3, 2024

Merge branch 'master' into ch_integ_tests

386d836

malav-shastri temporarily deployed to auto-approve October 3, 2024 07:43 — with GitHub Actions Inactive

Captainia reviewed Oct 3, 2024

View reviewed changes

AWS-pratab reviewed Oct 3, 2024

View reviewed changes

Only delete the pytest session specific test

fa7e47c

malav-shastri temporarily deployed to auto-approve October 3, 2024 20:02 — with GitHub Actions Inactive

change scope to session

7c50ee8

malav-shastri temporarily deployed to auto-approve October 3, 2024 22:12 — with GitHub Actions Inactive

address nits

cb5f1c7

malav-shastri temporarily deployed to auto-approve October 3, 2024 22:15 — with GitHub Actions Inactive

evakravi reviewed Oct 3, 2024

View reviewed changes

Merge branch 'master' into ch_integ_tests

52991a0

malav-shastri temporarily deployed to auto-approve October 4, 2024 01:53 — with GitHub Actions Inactive

Address test failures

4bd94ec

malav-shastri temporarily deployed to auto-approve October 4, 2024 02:27 — with GitHub Actions Inactive

address typo

3fed9f4

JGuinegagne reviewed Oct 4, 2024

View reviewed changes

address comments

6456883

malav-shastri temporarily deployed to auto-approve October 7, 2024 00:21 — with GitHub Actions Inactive

resolve flake8 errors

bac00dd

malav-shastri temporarily deployed to auto-approve October 7, 2024 03:26 — with GitHub Actions Inactive

implement throttle handling

556d120

malav-shastri temporarily deployed to auto-approve October 7, 2024 05:03 — with GitHub Actions Inactive

flake8

8ff04d3

malav-shastri temporarily deployed to auto-approve October 7, 2024 05:07 — with GitHub Actions Inactive

flake8

d79b8a3

malav-shastri temporarily deployed to auto-approve October 7, 2024 05:24 — with GitHub Actions Inactive

Captainia approved these changes Oct 7, 2024

View reviewed changes

Merge branch 'master' into ch_integ_tests

9e29524

malav-shastri temporarily deployed to auto-approve October 7, 2024 23:49 — with GitHub Actions Inactive

Adding more assertions

e0b8467

malav-shastri deployed to auto-approve October 8, 2024 00:42 — with GitHub Actions Active

Captainia approved these changes Oct 8, 2024

View reviewed changes

chad119 approved these changes Oct 8, 2024

View reviewed changes

benieric approved these changes Oct 8, 2024

View reviewed changes

pintaoz-aws merged commit d18e41b into aws:master Oct 8, 2024
14 checks passed

		sagemaker_session.delete_hub(hub["HubName"])


		def _delete_hub_contents(sagemaker_session, test_hub_name):

		JUMPSTART_TAG = "JumpStart-SDK-Integ-Test-Suite-Id"

		SM_JUMPSTART_PUBLIC_HUB_NAME = "SageMakerPublicHub"

		assert model.inference_component_name == predictor.component_name


		def test_instatiating_model(setup, add_models):

tests: Implement integration tests covering JumpStart PrivateHub workflows #4883

tests: Implement integration tests covering JumpStart PrivateHub workflows #4883

Conversation

malav-shastri commented Oct 3, 2024 • edited Loading

Merge Checklist

General

Tests

Captainia left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

AWS-pratab left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

malav-shastri Oct 3, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

malav-shastri commented Oct 3, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

malav-shastri Oct 4, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

malav-shastri Oct 4, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

malav-shastri Oct 4, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

malav-shastri commented Oct 3, 2024 •

edited

Loading

malav-shastri Oct 3, 2024 •

edited

Loading

malav-shastri Oct 4, 2024 •

edited

Loading

malav-shastri Oct 4, 2024 •

edited

Loading

malav-shastri Oct 4, 2024 •

edited

Loading