Merge branch 'master' into framework-versioning

Author: metrizable

setup.py

@@ -66,6 +66,7 @@ def read_version():
         "pytest<6.1.0",
         "pytest-cov",
         "pytest-rerunfailures",
+        "pytest-timeout",
         "pytest-xdist",
         "mock",
         "contextlib2",

src/sagemaker/image_uris.py

@@ -98,6 +98,8 @@ def retrieve(
             "mxnet-1.8.0-gpu-py37": "cu110-ubuntu16.04",
             "pytorch-1.6-gpu-py36": "cu110-ubuntu18.04-v3",
             "pytorch-1.6.0-gpu-py36": "cu110-ubuntu18.04",
+            "pytorch-1.6-gpu-py3": "cu110-ubuntu18.04-v3",
+            "pytorch-1.6.0-gpu-py3": "cu110-ubuntu18.04",
         }
         key = "-".join([framework, tag])
         if key in container_versions:

tests/integ/sagemaker/lineage/conftest.py

@@ -155,7 +155,8 @@ def artifact_obj_with_association(sagemaker_session, artifact_obj):
 @pytest.fixture
 def trial_component_obj(sagemaker_session):
     trial_component_obj = trial_component.TrialComponent.create(
-        trial_component_name=name(), sagemaker_boto_client=sagemaker_session.sagemaker_client
+        trial_component_name=name(),
+        sagemaker_boto_client=sagemaker_session.sagemaker_client,
     )
     yield trial_component_obj
     time.sleep(0.5)

tests/integ/sagemaker/lineage/test_action.py

@@ -15,6 +15,9 @@
 
 import datetime
 import logging
+import time
+
+import pytest
 
 from sagemaker.lineage import action
 
@@ -80,6 +83,7 @@ def test_list(action_objs, sagemaker_session):
     assert action_names
 
 
+@pytest.mark.timeout(30)
 def test_tag(action_obj, sagemaker_session):
     tag = {"Key": "foo", "Value": "bar"}
     action_obj.set_tag(tag)
@@ -90,12 +94,14 @@ def test_tag(action_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0
     assert actual_tags[0] == tag
 
 
+@pytest.mark.timeout(30)
 def test_tags(action_obj, sagemaker_session):
     tags = [{"Key": "foo1", "Value": "bar1"}]
     action_obj.set_tags(tags)
@@ -106,6 +112,7 @@ def test_tags(action_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0

tests/integ/sagemaker/lineage/test_artifact.py

@@ -17,6 +17,8 @@
 import logging
 import time
 
+import pytest
+
 from sagemaker.lineage import artifact
 
 
@@ -111,6 +113,7 @@ def test_downstream_trials(trial_associated_artifact, trial_obj, sagemaker_sessi
     assert trial_obj.trial_name in trials
 
 
+@pytest.mark.timeout(30)
 def test_tag(artifact_obj, sagemaker_session):
     tag = {"Key": "foo", "Value": "bar"}
     artifact_obj.set_tag(tag)
@@ -121,12 +124,14 @@ def test_tag(artifact_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0
     assert actual_tags[0] == tag
 
 
+@pytest.mark.timeout(30)
 def test_tags(artifact_obj, sagemaker_session):
     tags = [{"Key": "foo1", "Value": "bar1"}]
     artifact_obj.set_tags(tags)
@@ -137,6 +142,7 @@ def test_tags(artifact_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0

tests/integ/sagemaker/lineage/test_association.py

@@ -16,6 +16,8 @@
 import datetime
 import time
 
+import pytest
+
 from sagemaker.lineage import association
 
 
@@ -55,6 +57,7 @@ def test_list(association_objs, sagemaker_session):
     assert association_keys_listed
 
 
+@pytest.mark.timeout(30)
 def test_set_tag(association_obj, sagemaker_session):
     tag = {"Key": "foo", "Value": "bar"}
     association_obj.set_tag(tag)
@@ -65,13 +68,14 @@ def test_set_tag(association_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
-        time.sleep(1)
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0
     assert actual_tags[0] == tag
 
 
+@pytest.mark.timeout(30)
 def test_tags(association_obj, sagemaker_session):
     tags = [{"Key": "foo1", "Value": "bar1"}]
     association_obj.set_tags(tags)
@@ -82,7 +86,7 @@ def test_tags(association_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
-        time.sleep(1)
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0

tests/integ/sagemaker/lineage/test_context.py

@@ -15,6 +15,9 @@
 
 import datetime
 import logging
+import time
+
+import pytest
 
 from sagemaker.lineage import context
 
@@ -78,6 +81,7 @@ def test_list(context_objs, sagemaker_session):
     assert context_names
 
 
+@pytest.mark.timeout(30)
 def test_tag(context_obj, sagemaker_session):
     tag = {"Key": "foo", "Value": "bar"}
     context_obj.set_tag(tag)
@@ -88,12 +92,14 @@ def test_tag(context_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0
     assert actual_tags[0] == tag
 
 
+@pytest.mark.timeout(30)
 def test_tags(context_obj, sagemaker_session):
     tags = [{"Key": "foo1", "Value": "bar1"}]
     context_obj.set_tags(tags)
@@ -104,6 +110,7 @@ def test_tags(context_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
+        time.sleep(5)
     # When sagemaker-client-config endpoint-url is passed as argument to hit some endpoints,
     # length of actual tags will be greater than 1
     assert len(actual_tags) > 0

tests/unit/sagemaker/image_uris/test_retrieve.py

@@ -553,6 +553,21 @@ def test_retrieve_auto_selected_container_version():
     )
 
 
+def test_retrieve_pytorch_container_version():
+    uri = image_uris.retrieve(
+        framework="pytorch",
+        region="us-west-2",
+        version="1.6",
+        py_version="py3",
+        instance_type="ml.p4d.24xlarge",
+        image_scope="training",
+    )
+    assert (
+        "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.6-gpu-py3-cu110-ubuntu18.04-v3"
+        == uri
+    )
+
+
 @patch("sagemaker.image_uris.config_for_framework", return_value=BASE_CONFIG)
 def test_retrieve_unsupported_processor_type(config_for_framework):
     with pytest.raises(ValueError) as e:

tox.ini

@@ -55,6 +55,7 @@ markers =
     canary_quick
     cron
     local_mode
+    timeout: mark a test as a timeout.
 
 [testenv]
 passenv =