Fix pylint, flake8 and inte test failures

Author: Chia-Eng

src/sagemaker/estimator.py

@@ -123,7 +123,7 @@ def __init__(
         enable_network_isolation=False,
         profiler_config=None,
         disable_profiler=False,
-        environment = None,
+        environment=None,
         **kwargs,
     ):
         """Initialize an ``EstimatorBase`` instance.
@@ -267,7 +267,8 @@ def __init__(
                 ``disable_profiler`` parameter to ``True``.
             disable_profiler (bool): Specifies whether Debugger monitoring and profiling
                 will be disabled (default: ``False``).
-            environment (dict[str, str]) : Environment variables to be set for use during training job (default: ``None``)
+            environment (dict[str, str]) : Environment variables to be set for
+                use during training job (default: ``None``)
 
         """
         instance_count = renamed_kwargs(
@@ -1664,7 +1665,7 @@ def __init__(
         enable_sagemaker_metrics=None,
         profiler_config=None,
         disable_profiler=False,
-        environment = None,
+        environment=None,
         **kwargs,
     ):
         """Initialize an ``Estimator`` instance.
@@ -1813,7 +1814,8 @@ def __init__(
                 ``disable_profiler`` parameter to ``True``.
             disable_profiler (bool): Specifies whether Debugger monitoring and profiling
                 will be disabled (default: ``False``).
-            environment (dict[str, str]) : Environment variables to be set for use during training job (default: ``None``)
+            environment (dict[str, str]) : Environment variables to be set for
+                use during training job (default: ``None``)
         """
         self.image_uri = image_uri
         self.hyperparam_dict = hyperparameters.copy() if hyperparameters else {}

src/sagemaker/session.py

@@ -523,10 +523,12 @@ def train(  # noqa: C901
                 Series. For more information see:
                 https://docs.aws.amazon.com/sagemaker/latest/dg/API_AlgorithmSpecification.html#SageMaker-Type-AlgorithmSpecification-EnableSageMakerMetricsTimeSeries
                 (default: ``None``).
-            profiler_rule_configs (list[dict]): A list of profiler rule configurations.
+            profiler_rule_configs (list[dict]): A list of profiler rule
+                configurations.src/sagemaker/lineage/artifact.py:285
             profiler_config (dict): Configuration for how profiling information is emitted
                 with SageMaker Profiler. (default: ``None``).
-            environment (dict[str, str]) : Environment variables to be set for use during training job (default: ``None``)
+            environment (dict[str, str]) : Environment variables to be set for
+                use during training job (default: ``None``)
 
         Returns:
             str: ARN of the training job, if it is created.
@@ -661,7 +663,8 @@ def _get_train_request(  # noqa: C901
             profiler_rule_configs (list[dict]): A list of profiler rule configurations.
             profiler_config(dict): Configuration for how profiling information is emitted with
                 SageMaker Profiler. (default: ``None``).
-            environment (dict[str, str]) : Environment variables to be set for use during training job (default: ``None``)
+            environment (dict[str, str]) : Environment variables to be set for
+                use during training job (default: ``None``)
 
         Returns:
             Dict: a training request dict

tests/integ/test_tf.py

@@ -36,7 +36,7 @@
 PARAMETER_SERVER_DISTRIBUTION = {"parameter_server": {"enabled": True}}
 MPI_DISTRIBUTION = {"mpi": {"enabled": True}}
 TAGS = [{"Key": "some-key", "Value": "some-value"}]
-ENV_INPUT= {'env_key1': 'env_val1', 'env_key2': 'env_val2', 'env_key3': 'env_val3'}
+ENV_INPUT = {'env_key1': 'env_val1', 'env_key2': 'env_val2', 'env_key3': 'env_val3'}
 
 
 def test_mnist_with_checkpoint_config(
@@ -84,8 +84,11 @@ def test_mnist_with_checkpoint_config(
     actual_training_checkpoint_config = sagemaker_session.sagemaker_client.describe_training_job(
         TrainingJobName=training_job_name
     )["CheckpointConfig"]
+    actual_training_environment_variable_config = sagemaker_session.sagemaker_client.describe_training_job(
+        TrainingJobName=training_job_name
+    )["Environment"]
     assert actual_training_checkpoint_config == expected_training_checkpoint_config
-    assert actual_training_checkpoint_config['Environment'] == ENV_INPUT
+    assert actual_training_environment_variable_config == ENV_INPUT
 
 
 def test_server_side_encryption(sagemaker_session, tf_full_version, tf_full_py_version):

tests/unit/sagemaker/tensorflow/test_estimator_init.py

@@ -20,7 +20,8 @@
 
 REGION = "us-west-2"
 
-ENV_INPUT= {'env_key1': 'env_val1', 'env_key2': 'env_val2', 'env_key3': 'env_val3'}
+ENV_INPUT = {'env_key1': 'env_val1', 'env_key2': 'env_val2', 'env_key3': 'env_val3'}
+
 
 @pytest.fixture()
 def sagemaker_session():
@@ -68,6 +69,7 @@ def test_framework_name(sagemaker_session):
     tf = _build_tf(sagemaker_session, framework_version="1.15.2", py_version="py3")
     assert tf._framework_name == "tensorflow"
 
+
 def test_tf_add_environment_variables(sagemaker_session):
     tf = _build_tf(
         sagemaker_session,
@@ -77,6 +79,7 @@ def test_tf_add_environment_variables(sagemaker_session):
     )
     assert tf.environment == ENV_INPUT
 
+
 def test_tf_miss_environment_variables(sagemaker_session):
     tf = _build_tf(
         sagemaker_session,
@@ -86,6 +89,7 @@ def test_tf_miss_environment_variables(sagemaker_session):
     )
     assert not tf.environment
 
+
 def test_enable_sm_metrics(sagemaker_session):
     tf = _build_tf(
         sagemaker_session,

tests/unit/test_estimator.py

@@ -71,7 +71,7 @@
 CODECOMMIT_REPO_SSH = "ssh://git-codecommit.us-west-2.amazonaws.com/v1/repos/test-repo/"
 CODECOMMIT_BRANCH = "master"
 REPO_DIR = "/tmp/repo_dir"
-ENV_INPUT = {'env_key1':'env_val1', 'env_key2':'env_val2', 'env_key3':'env_val3'}
+ENV_INPUT = {'env_key1': 'env_val1', 'env_key2': 'env_val2', 'env_key3': 'env_val3'}
 
 DESCRIBE_TRAINING_JOB_RESULT = {"ModelArtifacts": {"S3ModelArtifacts": MODEL_DATA}}
 
@@ -2682,6 +2682,7 @@ def test_generic_to_fit_with_sagemaker_metrics_missing(sagemaker_session):
     args = sagemaker_session.train.call_args[1]
     assert "enable_sagemaker_metrics" not in args
 
+
 def test_add_environment_variables_to_train_args(sagemaker_session):
     e = Estimator(
         IMAGE_URI,
@@ -2699,6 +2700,7 @@ def test_add_environment_variables_to_train_args(sagemaker_session):
     args = sagemaker_session.train.call_args[1]
     assert args["environment"] == ENV_INPUT
 
+
 def test_generic_to_fit_with_sagemaker_metrics_enabled(sagemaker_session):
     e = Estimator(
         IMAGE_URI,

tests/unit/test_mxnet.py

@@ -65,7 +65,7 @@
 
 MODEL_PKG_RESPONSE = {"ModelPackageArn": "arn:model-pkg-arn"}
 
-ENV_INPUT= {'env_key1': 'env_val1', 'env_key2': 'env_val2', 'env_key3': 'env_val3'}
+ENV_INPUT = {'env_key1': 'env_val1', 'env_key2': 'env_val2', 'env_key3': 'env_val3'}
 
 
 @pytest.fixture()
@@ -961,6 +961,7 @@ def test_create_model_with_custom_hosting_image(sagemaker_session):
 
     assert model.image_uri == custom_hosting_image
 
+
 def test_mx_add_environment_variables(sagemaker_session, mxnet_training_version, mxnet_training_py_version):
     mx = MXNet(
         entry_point=SCRIPT_PATH,
@@ -974,6 +975,7 @@ def test_mx_add_environment_variables(sagemaker_session, mxnet_training_version,
     )
     assert mx.environment == ENV_INPUT
 
+
 def test_mx_missing_environment_variables(sagemaker_session, mxnet_training_version, mxnet_training_py_version):
     mx = MXNet(
         entry_point=SCRIPT_PATH,
@@ -987,6 +989,7 @@ def test_mx_missing_environment_variables(sagemaker_session, mxnet_training_vers
     )
     assert not mx.environment
 
+
 def test_mx_enable_sm_metrics(sagemaker_session, mxnet_training_version, mxnet_training_py_version):
     mx = MXNet(
         entry_point=SCRIPT_PATH,

tests/unit/test_pytorch.py

@@ -46,7 +46,7 @@
 
 ENDPOINT_CONFIG_DESC = {"ProductionVariants": [{"ModelName": "model-1"}, {"ModelName": "model-2"}]}
 
-ENV_INPUT= {'env_key1': 'env_val1', 'env_key2': 'env_val2', 'env_key3': 'env_val3'}
+ENV_INPUT = {'env_key1': 'env_val1', 'env_key2': 'env_val2', 'env_key3': 'env_val3'}
 
 LIST_TAGS_RESULT = {"Tags": [{"Key": "TagtestKey", "Value": "TagtestValue"}]}
 
@@ -629,7 +629,7 @@ def test_pt_enable_sm_metrics(
 
 
 def test_pt_disable_sm_metrics(
-    sagemaker_session, pytorch_training_version, pytorch_training_py_version
+        sagemaker_session, pytorch_training_version, pytorch_training_py_version
 ):
     pytorch = _pytorch_estimator(
         sagemaker_session,
@@ -639,9 +639,10 @@ def test_pt_disable_sm_metrics(
     )
     assert not pytorch.enable_sagemaker_metrics
 
+
 def test_pt_add_environment_variables(
-            sagemaker_session, pytorch_training_version, pytorch_training_py_version
-    ):
+        sagemaker_session, pytorch_training_version, pytorch_training_py_version
+):
     pytorch = _pytorch_estimator(
         sagemaker_session,
         framework_version=pytorch_training_version,
@@ -650,9 +651,10 @@ def test_pt_add_environment_variables(
     )
     assert pytorch.environment
 
+
 def test_pt_miss_environment_variables(
-            sagemaker_session, pytorch_training_version, pytorch_training_py_version
-    ):
+        sagemaker_session, pytorch_training_version, pytorch_training_py_version
+):
     pytorch = _pytorch_estimator(
         sagemaker_session,
         framework_version=pytorch_training_version,

tests/unit/test_session.py

@@ -36,7 +36,7 @@
 
 SAMPLE_PARAM_RANGES = [{"Name": "mini_batch_size", "MinValue": "10", "MaxValue": "100"}]
 
-ENV_INPUT= {'env_key1': 'env_val1', 'env_key2': 'env_val2', 'env_key3': 'env_val3'}
+ENV_INPUT = {'env_key1': 'env_val1', 'env_key2': 'env_val2', 'env_key3': 'env_val3'}
 
 REGION = "us-west-2"
 STS_ENDPOINT = "sts.us-west-2.amazonaws.com"