Added support for SageMaker defaults

Author: ivan-khvostishkov

tests/test_batch_inference.py

@@ -1,4 +1,5 @@
 import logging
+from datetime import timedelta
 
 import sagemaker
 from sagemaker.pytorch import PyTorch
@@ -8,19 +9,18 @@
 import test_util
 
 
-def test_clean_batch_inference(request):
+def test_clean_batch_inference():
     # noinspection DuplicatedCode
     sagemaker_session = sagemaker.Session()
     bucket = sagemaker_session.default_bucket()
 
     estimator = PyTorch(entry_point='train_clean.py',
                         source_dir='source_dir/training_clean/',
-                        role=request.config.getini('sagemaker_role'),
                         framework_version='1.9.1',
                         py_version='py38',
                         instance_count=1,
                         instance_type='ml.m5.xlarge',
-                        max_run=60 * 60 * 3,
+                        max_run=int(timedelta(minutes=15).total_seconds()),
                         keep_alive_period_in_seconds=1800,
                         container_log_level=logging.INFO)
     estimator.fit()
@@ -51,24 +51,23 @@ def test_clean_batch_inference(request):
                                     key_prefix='batch-transform/output')
 
 
-def test_batch_ssh(request):
+def test_batch_ssh():
     # noinspection DuplicatedCode
     sagemaker_session = sagemaker.Session()
     bucket = sagemaker_session.default_bucket()
 
     estimator = PyTorch(entry_point='train_clean.py',
                         source_dir='source_dir/training_clean/',
-                        role=request.config.getini('sagemaker_role'),
                         framework_version='1.9.1',
                         py_version='py38',
                         instance_count=1,
                         instance_type='ml.m5.xlarge',
-                        max_run=60 * 60 * 3,
+                        max_run=int(timedelta(minutes=15).total_seconds()),
                         keep_alive_period_in_seconds=1800,
                         container_log_level=logging.INFO)
     estimator.fit()
 
-    model = estimator.create_model(entry_point='inference.py',
+    model = estimator.create_model(entry_point='inference_ssh.py',
                                    source_dir='source_dir/inference/',
                                    dependencies=[SSHModelWrapper.dependency_dir()])
 

tests/test_clean_core.py

@@ -1,5 +1,7 @@
 import logging
 import os
+from datetime import timedelta
+from typing import Optional
 
 import pytest
 import sagemaker
@@ -21,7 +23,7 @@ def test_clean_train_warm_pool():
                         py_version='py38',
                         instance_count=1,
                         instance_type='ml.m5.xlarge',
-                        max_run=60 * 60 * 3,
+                        max_run=int(timedelta(minutes=15).total_seconds()),
                         keep_alive_period_in_seconds=1800,
                         container_log_level=logging.INFO)
     estimator.fit()
@@ -32,15 +34,14 @@ def test_clean_train_warm_pool():
 
 
 # noinspection DuplicatedCode
-def test_clean_inference(request):
+def test_clean_inference():
     estimator = PyTorch(entry_point='train_clean.py',
                         source_dir='source_dir/training_clean/',
-                        role=request.config.getini('sagemaker_role'),
                         framework_version='1.9.1',
                         py_version='py38',
                         instance_count=1,
                         instance_type='ml.m5.xlarge',
-                        max_run=60 * 60 * 3,
+                        max_run=int(timedelta(minutes=15).total_seconds()),
                         keep_alive_period_in_seconds=1800,
                         container_log_level=logging.INFO)
     estimator.fit()
@@ -65,41 +66,30 @@ def test_clean_inference(request):
 
 # noinspection DuplicatedCode
 @pytest.mark.parametrize("instance_type", ["ml.m5.xlarge"])
-def test_clean_inference_mms(request, instance_type):
+def test_clean_inference_mms(instance_type):
     estimator = PyTorch(entry_point='train_clean.py',
                         source_dir='source_dir/training_clean/',
-                        role=request.config.getini('sagemaker_role'),
                         framework_version='1.9.1',
                         py_version='py38',
                         instance_count=1,
                         instance_type=instance_type,
-                        max_run=60 * 60 * 3,
+                        max_run=int(timedelta(minutes=15).total_seconds()),
                         keep_alive_period_in_seconds=1800,
                         container_log_level=logging.INFO)
     estimator.fit()
 
     model_1 = estimator.create_model(entry_point='inference_clean.py',
                                      source_dir='source_dir/inference_clean/')
 
-    # we need a temp endpoint to produce 'repacked_model_data'
-    temp_endpoint_name = name_from_base('temp-inference-mms')
-    temp_predictor: Predictor = model_1.deploy(initial_instance_count=1,
-                                               instance_type='ml.m5.xlarge',
-                                               endpoint_name=temp_endpoint_name,
-                                               wait=True)
+    _ = model_1.prepare_container_def(instance_type='ml.m5.xlarge')
     repacked_model_data_1 = model_1.repacked_model_data
-    temp_predictor.delete_endpoint()
 
     # MUST have the same entry point file name as for the model_1
     model_2 = estimator.create_model(entry_point='inference_clean.py',
                                      source_dir='source_dir/inference_clean_model2/')
-    temp_endpoint_name = name_from_base('temp-inference-mms')
-    temp_predictor: Predictor = model_2.deploy(initial_instance_count=1,
-                                               instance_type='ml.m5.xlarge',
-                                               endpoint_name=temp_endpoint_name,
-                                               wait=True)
+
+    _ = model_2.prepare_container_def(instance_type='ml.m5.xlarge')
     repacked_model_data_2 = model_2.repacked_model_data
-    temp_predictor.delete_endpoint()
 
     bucket = sagemaker.Session().default_bucket()
     job_name = estimator.latest_training_job.name
@@ -115,12 +105,15 @@ def test_clean_inference_mms(request, instance_type):
 
     endpoint_name = name_from_base('inference-mms')
 
-    predictor: Predictor = mdm.deploy(initial_instance_count=1,
-                                      instance_type='ml.m5.xlarge',
-                                      endpoint_name=endpoint_name,
-                                      wait=True)
-
+    predictor: Optional[Predictor] = None
     try:
+        predictor = mdm.deploy(
+            initial_instance_count=1,
+            instance_type='ml.m5.xlarge',
+            endpoint_name=endpoint_name,
+            wait=True
+        )
+
         # Note: we need a repacked model data here, not an estimator data
         mdm.add_model(model_data_source=repacked_model_data_1, model_data_path='model_1.tar.gz')
         mdm.add_model(model_data_source=repacked_model_data_2, model_data_path='model_2.tar.gz')
@@ -143,49 +136,32 @@ def test_clean_inference_mms(request, instance_type):
 
 # noinspection DuplicatedCode
 @pytest.mark.parametrize("instance_type", ["ml.m5.xlarge"])
-def test_clean_inference_mms_without_model(request, instance_type):
+def test_clean_inference_mms_without_model(instance_type):
     estimator = PyTorch(entry_point='train_clean.py',
                         source_dir='source_dir/training_clean/',
-                        role=request.config.getini('sagemaker_role'),
                         framework_version='1.9.1',
                         py_version='py38',
                         instance_count=1,
                         instance_type=instance_type,
-                        max_run=60 * 60 * 3,
+                        max_run=int(timedelta(minutes=15).total_seconds()),
                         keep_alive_period_in_seconds=1800,
                         container_log_level=logging.INFO)
     estimator.fit()
 
     model_1 = estimator.create_model(entry_point='inference_clean.py',
                                      source_dir='source_dir/inference_clean/')
 
-    # we need a temp endpoint to produce 'repacked_model_data'
-    temp_endpoint_name = name_from_base('temp-inference-mms')
-    temp_predictor: Predictor = model_1.deploy(initial_instance_count=1,
-                                               instance_type='ml.m5.xlarge',
-                                               endpoint_name=temp_endpoint_name,
-                                               wait=True)
+    model_1_description = model_1.prepare_container_def(instance_type='ml.m5.xlarge')
     repacked_model_data_1 = model_1.repacked_model_data
-    temp_predictor.delete_endpoint()
-
-    # But we still don't have access to the deployed container URI from Model object, so still need to use boto3.
-    # Re-fetch container and model data location from Container 1 of the model:
-    model_1_description = model_1.sagemaker_session.describe_model(model_1.name)
-    container_uri = model_1_description['PrimaryContainer']['Image']
-    # Also re-fetch deploy environment:
-    deploy_env = model_1_description['PrimaryContainer']['Environment']
+    container_uri = model_1_description['Image']
+    deploy_env = model_1_description['Environment']
 
     # MUST have the same entry point file name as for the model_1
     model_2 = estimator.create_model(entry_point='inference_clean.py',
                                      source_dir='source_dir/inference_clean_model2/')
 
-    temp_endpoint_name = name_from_base('temp-inference-mms')
-    temp_predictor: Predictor = model_2.deploy(initial_instance_count=1,
-                                               instance_type='ml.m5.xlarge',
-                                               endpoint_name=temp_endpoint_name,
-                                               wait=True)
+    _ = model_2.prepare_container_def(instance_type='ml.m5.xlarge')
     repacked_model_data_2 = model_2.repacked_model_data
-    temp_predictor.delete_endpoint()
 
     bucket = sagemaker.Session().default_bucket()
     job_name = estimator.latest_training_job.name
@@ -196,8 +172,8 @@ def test_clean_inference_mms_without_model(request, instance_type):
     mdm = MultiDataModel(
         name=mdm_name,
         model_data_prefix=model_data_prefix,
-        role=model_1.role,
         image_uri=container_uri,
+        # entry_point=model_1.entry_point,  # NOTE: entry point ignored
         env=deploy_env,  # will copy 'SAGEMAKER_PROGRAM' env variable with entry point file name
         predictor_cls=PyTorchPredictor
     )

tests/test_distributed.py

@@ -1,5 +1,6 @@
 import logging
 import os
+from datetime import timedelta
 
 import pytest
 from sagemaker.pytorch import PyTorch
@@ -20,19 +21,18 @@ def test_node_rank_from_env_json_non_existing_rc():
     assert node_rank == 0
 
 
-def test_distributed_training_with_default_instance_count(request):
+def test_distributed_training_with_default_instance_count():
     instance_count = 3
     default_ssh_instance_count = 2
     estimator = PyTorch(entry_point='train.py',
                         source_dir='source_dir/training/',
                         dependencies=[SSHEstimatorWrapper.dependency_dir()],
                         base_job_name='ssh-training',
-                        role=request.config.getini('sagemaker_role'),
                         framework_version='1.9.1',
                         py_version='py38',
                         instance_count=instance_count,
                         instance_type='ml.m5.xlarge',
-                        max_run=60 * 60 * 3,
+                        max_run=int(timedelta(minutes=15).total_seconds()),
                         keep_alive_period_in_seconds=1800,
                         container_log_level=logging.INFO)
 
@@ -45,18 +45,17 @@ def test_distributed_training_with_default_instance_count(request):
 
 
 @pytest.mark.parametrize("ssh_instance_count", [3, 1])
-def test_distributed_training_with_changed_instance_count(request, ssh_instance_count):
+def test_distributed_training_with_changed_instance_count(ssh_instance_count):
     instance_count = 3
     estimator = PyTorch(entry_point='train.py',
                         source_dir='source_dir/training/',
                         dependencies=[SSHEstimatorWrapper.dependency_dir()],
                         base_job_name='ssh-training',
-                        role=request.config.getini('sagemaker_role'),
                         framework_version='1.9.1',
                         py_version='py38',
                         instance_count=instance_count,
                         instance_type='ml.m5.xlarge',
-                        max_run=60 * 60 * 3,
+                        max_run=int(timedelta(minutes=15).total_seconds()),
                         keep_alive_period_in_seconds=1800,
                         container_log_level=logging.INFO)
 
@@ -69,18 +68,17 @@ def test_distributed_training_with_changed_instance_count(request, ssh_instance_
     assert len(mi_ids) == ssh_instance_count
 
 
-def test_distributed_training_mpi_single_node(request):
+def test_distributed_training_mpi_single_node():
     instance_count = 1
     estimator = PyTorch(entry_point='train.py',
                         source_dir='source_dir/training/',
                         dependencies=[SSHEstimatorWrapper.dependency_dir()],
                         base_job_name='ssh-training',
-                        role=request.config.getini('sagemaker_role'),
                         framework_version='1.9.1',
                         py_version='py38',
                         instance_count=instance_count,
                         instance_type='ml.g4dn.xlarge',
-                        max_run=60 * 60 * 3,
+                        max_run=int(timedelta(minutes=15).total_seconds()),
                         keep_alive_period_in_seconds=1800,
                         container_log_level=logging.INFO,
                         distribution={