Fix unit tests

Author: Ignacio Quintero

tests/unit/test_chainer.py

@@ -66,15 +66,14 @@ def _get_full_gpu_image_uri(version):
 
 
 def _chainer_estimator(sagemaker_session, framework_version=defaults.CHAINER_VERSION, train_instance_type=None,
-                       enable_cloudwatch_metrics=False, base_job_name=None, use_mpi=None, num_processes=None,
+                       base_job_name=None, use_mpi=None, num_processes=None,
                        process_slots_per_host=None, additional_mpi_options=None, **kwargs):
     return Chainer(entry_point=SCRIPT_PATH,
                    framework_version=framework_version,
                    role=ROLE,
                    sagemaker_session=sagemaker_session,
                    train_instance_count=INSTANCE_COUNT,
                    train_instance_type=train_instance_type if train_instance_type else INSTANCE_TYPE,
-                   enable_cloudwatch_metrics=enable_cloudwatch_metrics,
                    base_job_name=base_job_name,
                    use_mpi=use_mpi,
                    num_processes=num_processes,
@@ -152,7 +151,6 @@ def _create_train_job_with_additional_hyperparameters(version):
         },
         'hyperparameters': {
             'sagemaker_program': json.dumps('dummy_script.py'),
-            'sagemaker_enable_cloudwatch_metrics': 'false',
             'sagemaker_container_log_level': str(logging.INFO),
             'sagemaker_job_name': json.dumps(JOB_NAME),
             'sagemaker_submit_directory':
@@ -225,12 +223,10 @@ def test_attach_with_additional_hyperparameters(sagemaker_session, chainer_versi
 def test_create_model(sagemaker_session, chainer_version):
     container_log_level = '"logging.INFO"'
     source_dir = 's3://mybucket/source'
-    enable_cloudwatch_metrics = 'true'
     chainer = Chainer(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                       train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,
                       framework_version=chainer_version, container_log_level=container_log_level,
-                      py_version=PYTHON_VERSION, base_job_name='job', source_dir=source_dir,
-                      enable_cloudwatch_metrics=enable_cloudwatch_metrics)
+                      py_version=PYTHON_VERSION, base_job_name='job', source_dir=source_dir)
 
     job_name = 'new_name'
     chainer.fit(inputs='s3://mybucket/train', job_name='new_name')
@@ -244,19 +240,16 @@ def test_create_model(sagemaker_session, chainer_version):
     assert model.name == job_name
     assert model.container_log_level == container_log_level
     assert model.source_dir == source_dir
-    assert model.enable_cloudwatch_metrics == enable_cloudwatch_metrics
 
 
 def test_create_model_with_custom_image(sagemaker_session):
     container_log_level = '"logging.INFO"'
     source_dir = 's3://mybucket/source'
-    enable_cloudwatch_metrics = 'true'
     custom_image = 'ubuntu:latest'
     chainer = Chainer(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                       train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,
                       image_name=custom_image, container_log_level=container_log_level,
-                      py_version=PYTHON_VERSION, base_job_name='job', source_dir=source_dir,
-                      enable_cloudwatch_metrics=enable_cloudwatch_metrics)
+                      py_version=PYTHON_VERSION, base_job_name='job', source_dir=source_dir)
 
     chainer.fit(inputs='s3://mybucket/train', job_name='new_name')
     model = chainer.create_model()

tests/unit/test_mxnet.py

@@ -101,11 +101,10 @@ def _create_train_job(version):
 def test_create_model(sagemaker_session, mxnet_version):
     container_log_level = '"logging.INFO"'
     source_dir = 's3://mybucket/source'
-    enable_cloudwatch_metrics = 'true'
     mx = MXNet(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,
                framework_version=mxnet_version, container_log_level=container_log_level,
-               base_job_name='job', source_dir=source_dir, enable_cloudwatch_metrics=enable_cloudwatch_metrics)
+               base_job_name='job', source_dir=source_dir)
 
     job_name = 'new_name'
     mx.fit(inputs='s3://mybucket/train', job_name='new_name')
@@ -119,18 +118,16 @@ def test_create_model(sagemaker_session, mxnet_version):
     assert model.name == job_name
     assert model.container_log_level == container_log_level
     assert model.source_dir == source_dir
-    assert model.enable_cloudwatch_metrics == enable_cloudwatch_metrics
 
 
 def test_create_model_with_custom_image(sagemaker_session):
     container_log_level = '"logging.INFO"'
     source_dir = 's3://mybucket/source'
-    enable_cloudwatch_metrics = 'true'
     custom_image = 'mxnet:2.0'
     mx = MXNet(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,
                image_name=custom_image, container_log_level=container_log_level,
-               base_job_name='job', source_dir=source_dir, enable_cloudwatch_metrics=enable_cloudwatch_metrics)
+               base_job_name='job', source_dir=source_dir)
 
     job_name = 'new_name'
     mx.fit(inputs='s3://mybucket/train', job_name='new_name')
@@ -143,7 +140,6 @@ def test_create_model_with_custom_image(sagemaker_session):
     assert model.name == job_name
     assert model.container_log_level == container_log_level
     assert model.source_dir == source_dir
-    assert model.enable_cloudwatch_metrics == enable_cloudwatch_metrics
 
 
 @patch('time.strftime', return_value=TIMESTAMP)

tests/unit/test_pytorch.py

@@ -64,15 +64,14 @@ def _get_full_gpu_image_uri(version, py_version=PYTHON_VERSION):
 
 
 def _pytorch_estimator(sagemaker_session, framework_version=defaults.PYTORCH_VERSION, train_instance_type=None,
-                       enable_cloudwatch_metrics=False, base_job_name=None, **kwargs):
+                       base_job_name=None, **kwargs):
     return PyTorch(entry_point=SCRIPT_PATH,
                    framework_version=framework_version,
                    py_version=PYTHON_VERSION,
                    role=ROLE,
                    sagemaker_session=sagemaker_session,
                    train_instance_count=INSTANCE_COUNT,
                    train_instance_type=train_instance_type if train_instance_type else INSTANCE_TYPE,
-                   enable_cloudwatch_metrics=enable_cloudwatch_metrics,
                    base_job_name=base_job_name,
                    **kwargs)
 
@@ -119,11 +118,10 @@ def _create_train_job(version):
 def test_create_model(sagemaker_session, pytorch_version):
     container_log_level = '"logging.INFO"'
     source_dir = 's3://mybucket/source'
-    enable_cloudwatch_metrics = 'true'
     pytorch = PyTorch(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                       train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,
                       framework_version=pytorch_version, container_log_level=container_log_level,
-                      base_job_name='job', source_dir=source_dir, enable_cloudwatch_metrics=enable_cloudwatch_metrics)
+                      base_job_name='job', source_dir=source_dir)
 
     job_name = 'new_name'
     pytorch.fit(inputs='s3://mybucket/train', job_name='new_name')
@@ -137,18 +135,16 @@ def test_create_model(sagemaker_session, pytorch_version):
     assert model.name == job_name
     assert model.container_log_level == container_log_level
     assert model.source_dir == source_dir
-    assert model.enable_cloudwatch_metrics == enable_cloudwatch_metrics
 
 
 def test_create_model_with_custom_image(sagemaker_session):
     container_log_level = '"logging.INFO"'
     source_dir = 's3://mybucket/source'
-    enable_cloudwatch_metrics = 'true'
     image = 'pytorch:9000'
     pytorch = PyTorch(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                       train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,
                       container_log_level=container_log_level, image_name=image,
-                      base_job_name='job', source_dir=source_dir, enable_cloudwatch_metrics=enable_cloudwatch_metrics)
+                      base_job_name='job', source_dir=source_dir)
 
     job_name = 'new_name'
     pytorch.fit(inputs='s3://mybucket/train', job_name='new_name')
@@ -161,7 +157,6 @@ def test_create_model_with_custom_image(sagemaker_session):
     assert model.name == job_name
     assert model.container_log_level == container_log_level
     assert model.source_dir == source_dir
-    assert model.enable_cloudwatch_metrics == enable_cloudwatch_metrics
 
 
 @patch('time.strftime', return_value=TIMESTAMP)

tests/unit/test_tf_estimator.py

@@ -107,7 +107,7 @@ def _create_train_job(tf_version):
 
 
 def _build_tf(sagemaker_session, framework_version=defaults.TF_VERSION, train_instance_type=None,
-              checkpoint_path=None, enable_cloudwatch_metrics=False, base_job_name=None,
+              checkpoint_path=None, base_job_name=None,
               training_steps=None, evaluation_steps=None, **kwargs):
     return TensorFlow(entry_point=SCRIPT_PATH,
                       training_steps=training_steps,
@@ -118,7 +118,6 @@ def _build_tf(sagemaker_session, framework_version=defaults.TF_VERSION, train_in
                       train_instance_count=INSTANCE_COUNT,
                       train_instance_type=train_instance_type if train_instance_type else INSTANCE_TYPE,
                       checkpoint_path=checkpoint_path,
-                      enable_cloudwatch_metrics=enable_cloudwatch_metrics,
                       base_job_name=base_job_name,
                       **kwargs)
 
@@ -183,12 +182,11 @@ def test_tf_nonexistent_requirements_path(sagemaker_session):
 def test_create_model(sagemaker_session, tf_version):
     container_log_level = '"logging.INFO"'
     source_dir = 's3://mybucket/source'
-    enable_cloudwatch_metrics = 'true'
     tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                     training_steps=1000, evaluation_steps=10, train_instance_count=INSTANCE_COUNT,
                     train_instance_type=INSTANCE_TYPE, framework_version=tf_version,
                     container_log_level=container_log_level, base_job_name='job',
-                    source_dir=source_dir, enable_cloudwatch_metrics=enable_cloudwatch_metrics)
+                    source_dir=source_dir)
 
     job_name = 'doing something'
     tf.fit(inputs='s3://mybucket/train', job_name=job_name)
@@ -202,19 +200,17 @@ def test_create_model(sagemaker_session, tf_version):
     assert model.name == job_name
     assert model.container_log_level == container_log_level
     assert model.source_dir == source_dir
-    assert model.enable_cloudwatch_metrics == enable_cloudwatch_metrics
 
 
 def test_create_model_with_custom_image(sagemaker_session):
     container_log_level = '"logging.INFO"'
     source_dir = 's3://mybucket/source'
-    enable_cloudwatch_metrics = 'true'
     custom_image = 'tensorflow:1.0'
     tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                     training_steps=1000, evaluation_steps=10, train_instance_count=INSTANCE_COUNT,
                     train_instance_type=INSTANCE_TYPE, image_name=custom_image,
                     container_log_level=container_log_level, base_job_name='job',
-                    source_dir=source_dir, enable_cloudwatch_metrics=enable_cloudwatch_metrics)
+                    source_dir=source_dir)
 
     job_name = 'doing something'
     tf.fit(inputs='s3://mybucket/train', job_name=job_name)