Formatted files

Author: ananth102

src/sagemaker/experiments/run.py

@@ -342,9 +342,7 @@ def log_precision_recall(
         if positive_label is not None:
             kwargs["pos_label"] = positive_label
 
-        precision, recall, _ = precision_recall_curve(
-            y_true, predicted_probabilities, **kwargs
-        )
+        precision, recall, _ = precision_recall_curve(y_true, predicted_probabilities, **kwargs)
 
         kwargs["average"] = "micro"
         ap = average_precision_score(y_true, predicted_probabilities, **kwargs)
@@ -564,9 +562,7 @@ def _is_input_valid(input_type, field_name, field_value) -> bool:
             field_name (str): The name of the field to be checked.
             field_value (str or int or float): The value of the field to be checked.
         """
-        if isinstance(field_value, Number) and (
-            isnan(field_value) or isinf(field_value)
-        ):
+        if isinstance(field_value, Number) and (isnan(field_value) or isinf(field_value)):
             logger.warning(
                 "Failed to log %s %s. Received invalid value: %s.",
                 input_type,
@@ -628,14 +624,10 @@ def _verify_trial_component_artifacts_length(self, is_output):
         err_msg_template = "Cannot add more than {} {}_artifacts under run"
         if is_output:
             if len(self._trial_component.output_artifacts) >= MAX_RUN_TC_ARTIFACTS_LEN:
-                raise ValueError(
-                    err_msg_template.format(MAX_RUN_TC_ARTIFACTS_LEN, "output")
-                )
+                raise ValueError(err_msg_template.format(MAX_RUN_TC_ARTIFACTS_LEN, "output"))
         else:
             if len(self._trial_component.input_artifacts) >= MAX_RUN_TC_ARTIFACTS_LEN:
-                raise ValueError(
-                    err_msg_template.format(MAX_RUN_TC_ARTIFACTS_LEN, "input")
-                )
+                raise ValueError(err_msg_template.format(MAX_RUN_TC_ARTIFACTS_LEN, "input"))
 
     @staticmethod
     def _generate_trial_component_name(run_name: str, experiment_name: str) -> str:
@@ -656,28 +648,20 @@ def _generate_trial_component_name(run_name: str, experiment_name: str) -> str:
         """
         buffer = 1  # leave length buffers for delimiters
         max_len = int(MAX_NAME_LEN_IN_BACKEND / 2) - buffer
-        err_msg_template = (
-            "The {} (length: {}) must have length less than or equal to {}"
-        )
+        err_msg_template = "The {} (length: {}) must have length less than or equal to {}"
         if len(run_name) > max_len:
-            raise ValueError(
-                err_msg_template.format("run_name", len(run_name), max_len)
-            )
+            raise ValueError(err_msg_template.format("run_name", len(run_name), max_len))
         if len(experiment_name) > max_len:
             raise ValueError(
-                err_msg_template.format(
-                    "experiment_name", len(experiment_name), max_len
-                )
+                err_msg_template.format("experiment_name", len(experiment_name), max_len)
             )
         trial_component_name = "{}{}{}".format(experiment_name, DELIMITER, run_name)
         # due to mixed-case concerns on the backend
         trial_component_name = trial_component_name.lower()
         return trial_component_name
 
     @staticmethod
-    def _extract_run_name_from_tc_name(
-        trial_component_name: str, experiment_name: str
-    ) -> str:
+    def _extract_run_name_from_tc_name(trial_component_name: str, experiment_name: str) -> str:
         """Extract the user supplied run name from a trial component name.
 
         Args:
@@ -694,9 +678,7 @@ def _extract_run_name_from_tc_name(
         )
 
     @staticmethod
-    def _append_run_tc_label_to_tags(
-        tags: Optional[List[Dict[str, str]]] = None
-    ) -> list:
+    def _append_run_tc_label_to_tags(tags: Optional[List[Dict[str, str]]] = None) -> list:
         """Append the run trial component label to tags used to create a trial component.
 
         Args:

tests/unit/sagemaker/experiments/test_run.py

@@ -150,9 +150,8 @@ def test_run_init_name_length_exceed_limit(sagemaker_session):
             sagemaker_session=sagemaker_session,
         )
 
-    assert (
-        f"The run_name (length: {MAX_NAME_LEN_IN_BACKEND}) must have length less than"
-        in str(err)
+    assert f"The run_name (length: {MAX_NAME_LEN_IN_BACKEND}) must have length less than" in str(
+        err
     )
 
 
@@ -224,9 +223,7 @@ def test_run_load_no_run_name_and_in_train_job(
             }
         ]
     }
-    expmock = MagicMock(
-        return_value=Experiment(experiment_name=TEST_EXP_NAME, tags=expected_tags)
-    )
+    expmock = MagicMock(return_value=Experiment(experiment_name=TEST_EXP_NAME, tags=expected_tags))
     with patch("sagemaker.experiments.run.Experiment._load_or_create", expmock):
         with load_run(sagemaker_session=sagemaker_session, **kwargs) as run_obj:
             assert run_obj._in_load
@@ -239,12 +236,8 @@ def test_run_load_no_run_name_and_in_train_job(
             assert run_obj.experiment_name == TEST_EXP_NAME
             assert run_obj._experiment
             assert run_obj.experiment_config == exp_config
-            assert (
-                run_obj._artifact_uploader.artifact_bucket == expected_artifact_bucket
-            )
-            assert (
-                run_obj._artifact_uploader.artifact_prefix == expected_artifact_prefix
-            )
+            assert run_obj._artifact_uploader.artifact_bucket == expected_artifact_bucket
+            assert run_obj._artifact_uploader.artifact_prefix == expected_artifact_prefix
             assert run_obj._experiment.tags == expected_tags
 
     client.describe_training_job.assert_called_once_with(TrainingJobName=job_name)
@@ -269,9 +262,7 @@ def test_run_load_no_run_name_and_in_train_job_but_fail_to_get_exp_cfg(
         with load_run(sagemaker_session=sagemaker_session):
             pass
 
-    assert "Not able to fetch RunName in ExperimentConfig of the sagemaker job" in str(
-        err
-    )
+    assert "Not able to fetch RunName in ExperimentConfig of the sagemaker job" in str(err)
 
 
 def test_run_load_no_run_name_and_not_in_train_job(run_obj, sagemaker_session):
@@ -291,9 +282,7 @@ def test_run_load_no_run_name_and_not_in_train_job_but_no_obj_in_context(
 
     # experiment_name is given but is not supplied along with the run_name so it's ignored.
     with pytest.raises(RuntimeError) as err:
-        with load_run(
-            experiment_name=TEST_EXP_NAME, sagemaker_session=sagemaker_session
-        ):
+        with load_run(experiment_name=TEST_EXP_NAME, sagemaker_session=sagemaker_session):
             pass
 
     assert "Failed to load a Run object" in str(err)
@@ -621,17 +610,12 @@ def test_log_output_artifact(run_obj):
     with run_obj:
         run_obj.log_file("foo.txt", "name", "whizz/bang")
         run_obj._artifact_uploader.upload_artifact.assert_called_with("foo.txt")
-        assert (
-            "whizz/bang" == run_obj._trial_component.output_artifacts["name"].media_type
-        )
+        assert "whizz/bang" == run_obj._trial_component.output_artifacts["name"].media_type
 
         run_obj.log_file("foo.txt")
         run_obj._artifact_uploader.upload_artifact.assert_called_with("foo.txt")
         assert "foo.txt" in run_obj._trial_component.output_artifacts
-        assert (
-            "text/plain"
-            == run_obj._trial_component.output_artifacts["foo.txt"].media_type
-        )
+        assert "text/plain" == run_obj._trial_component.output_artifacts["foo.txt"].media_type
 
 
 def test_log_input_artifact_outside_run_context(run_obj):
@@ -648,51 +632,36 @@ def test_log_input_artifact(run_obj):
     with run_obj:
         run_obj.log_file("foo.txt", "name", "whizz/bang", is_output=False)
         run_obj._artifact_uploader.upload_artifact.assert_called_with("foo.txt")
-        assert (
-            "whizz/bang" == run_obj._trial_component.input_artifacts["name"].media_type
-        )
+        assert "whizz/bang" == run_obj._trial_component.input_artifacts["name"].media_type
 
         run_obj.log_file("foo.txt", is_output=False)
         run_obj._artifact_uploader.upload_artifact.assert_called_with("foo.txt")
         assert "foo.txt" in run_obj._trial_component.input_artifacts
-        assert (
-            "text/plain"
-            == run_obj._trial_component.input_artifacts["foo.txt"].media_type
-        )
+        assert "text/plain" == run_obj._trial_component.input_artifacts["foo.txt"].media_type
 
 
 def test_log_multiple_inputs(run_obj):
     with run_obj:
         for index in range(0, MAX_RUN_TC_ARTIFACTS_LEN):
             file_path = "foo" + str(index) + ".txt"
             run_obj._trial_component.input_artifacts[file_path] = {
-                "foo": TrialComponentArtifact(
-                    value="baz" + str(index), media_type="text/text"
-                )
+                "foo": TrialComponentArtifact(value="baz" + str(index), media_type="text/text")
             }
         with pytest.raises(ValueError) as error:
             run_obj.log_artifact("foo.txt", "name", "whizz/bang", False)
-        assert (
-            f"Cannot add more than {MAX_RUN_TC_ARTIFACTS_LEN} input_artifacts"
-            in str(error)
-        )
+        assert f"Cannot add more than {MAX_RUN_TC_ARTIFACTS_LEN} input_artifacts" in str(error)
 
 
 def test_log_multiple_outputs(run_obj):
     with run_obj:
         for index in range(0, MAX_RUN_TC_ARTIFACTS_LEN):
             file_path = "foo" + str(index) + ".txt"
             run_obj._trial_component.output_artifacts[file_path] = {
-                "foo": TrialComponentArtifact(
-                    value="baz" + str(index), media_type="text/text"
-                )
+                "foo": TrialComponentArtifact(value="baz" + str(index), media_type="text/text")
             }
         with pytest.raises(ValueError) as error:
             run_obj.log_artifact("foo.txt", "name", "whizz/bang")
-        assert (
-            f"Cannot add more than {MAX_RUN_TC_ARTIFACTS_LEN} output_artifacts"
-            in str(error)
-        )
+        assert f"Cannot add more than {MAX_RUN_TC_ARTIFACTS_LEN} output_artifacts" in str(error)
 
 
 def test_log_multiple_input_artifacts(run_obj):
@@ -722,10 +691,7 @@ def test_log_multiple_input_artifacts(run_obj):
         # log an extra input artifact, should raise exception
         with pytest.raises(ValueError) as error:
             run_obj.log_file("foo.txt", "name", "whizz/bang", is_output=False)
-        assert (
-            f"Cannot add more than {MAX_RUN_TC_ARTIFACTS_LEN} input_artifacts"
-            in str(error)
-        )
+        assert f"Cannot add more than {MAX_RUN_TC_ARTIFACTS_LEN} input_artifacts" in str(error)
 
 
 def test_log_multiple_output_artifacts(run_obj):
@@ -750,10 +716,7 @@ def test_log_multiple_output_artifacts(run_obj):
         # log an extra output artifact, should raise exception
         with pytest.raises(ValueError) as error:
             run_obj.log_file("foo.txt", "name", "whizz/bang")
-        assert (
-            f"Cannot add more than {MAX_RUN_TC_ARTIFACTS_LEN} output_artifacts"
-            in str(error)
-        )
+        assert f"Cannot add more than {MAX_RUN_TC_ARTIFACTS_LEN} output_artifacts" in str(error)
 
 
 def test_log_precision_recall_outside_run_context(run_obj):
@@ -820,10 +783,7 @@ def test_log_precision_recall_invalid_input(run_obj):
                 no_skill=no_skill,
                 is_output=False,
             )
-        assert (
-            "Lengths mismatch between true labels and predicted probabilities"
-            in str(error)
-        )
+        assert "Lengths mismatch between true labels and predicted probabilities" in str(error)
 
 
 def test_log_confusion_matrix_outside_run_context(run_obj):
@@ -921,9 +881,7 @@ def test_log_roc_curve_invalid_input(run_obj):
 
     with run_obj:
         with pytest.raises(ValueError) as error:
-            run_obj.log_roc_curve(
-                y_true, y_scores, title="TestROCCurve", is_output=False
-            )
+            run_obj.log_roc_curve(y_true, y_scores, title="TestROCCurve", is_output=False)
         assert "Lengths mismatch between true labels and predicted scores" in str(error)
 
 
@@ -940,18 +898,10 @@ def test_log_roc_curve_invalid_input(run_obj):
 @patch("sagemaker.experiments.run._TrialComponent.list")
 @patch("sagemaker.experiments.run._TrialComponent.search")
 def test_list(mock_tc_search, mock_tc_list, mock_tc_load, run_obj, sagemaker_session):
-    start_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
-        hours=1
-    )
-    end_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
-        hours=2
-    )
-    creation_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
-        hours=3
-    )
-    last_modified_time = datetime.datetime.now(
-        datetime.timezone.utc
-    ) + datetime.timedelta(hours=4)
+    start_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=1)
+    end_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=2)
+    creation_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=3)
+    last_modified_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=4)
     tc_list_len = 20
     tc_list_len_half = int(tc_list_len / 2)
     mock_tc_search.side_effect = [
@@ -1039,9 +989,8 @@ def test_list(mock_tc_search, mock_tc_list, mock_tc_load, run_obj, sagemaker_ses
         assert run._experiment
         assert run._trial
         assert isinstance(run._trial_component, _TrialComponent)
-        assert (
-            run._trial_component.trial_component_name
-            == Run._generate_trial_component_name("a" + str(i), TEST_EXP_NAME)
+        assert run._trial_component.trial_component_name == Run._generate_trial_component_name(
+            "a" + str(i), TEST_EXP_NAME
         )
         assert run._in_load is False
         assert run._inside_load_context is False
@@ -1054,9 +1003,7 @@ def test_list(mock_tc_search, mock_tc_list, mock_tc_load, run_obj, sagemaker_ses
 @patch("sagemaker.experiments.run._TrialComponent.list")
 def test_list_empty(mock_tc_list, sagemaker_session):
     mock_tc_list.return_value = []
-    assert [] == list_runs(
-        experiment_name=TEST_EXP_NAME, sagemaker_session=sagemaker_session
-    )
+    assert [] == list_runs(experiment_name=TEST_EXP_NAME, sagemaker_session=sagemaker_session)
 
 
 @patch(
@@ -1122,10 +1069,7 @@ def test_exit_fail(sagemaker_session, run_obj):
     except ValueError:
         pass
 
-    assert (
-        run_obj._trial_component.status.primary_status
-        == _TrialComponentStatusType.Failed.value
-    )
+    assert run_obj._trial_component.status.primary_status == _TrialComponentStatusType.Failed.value
     assert run_obj._trial_component.status.message
     assert isinstance(run_obj._trial_component.end_time, datetime.datetime)
 
@@ -1182,9 +1126,7 @@ def _verify_tc_status_before_enter_init(trial_component):
     assert not trial_component.status
 
 
-def _verify_tc_status_when_entering(
-    trial_component, init_start_time=None, has_completed=False
-):
+def _verify_tc_status_when_entering(trial_component, init_start_time=None, has_completed=False):
     if not init_start_time:
         assert isinstance(trial_component.start_time, datetime.datetime)
         now = datetime.datetime.now(dateutil.tz.tzlocal())
@@ -1194,17 +1136,11 @@ def _verify_tc_status_when_entering(
 
     if not has_completed:
         assert not trial_component.end_time
-    assert (
-        trial_component.status.primary_status
-        == _TrialComponentStatusType.InProgress.value
-    )
+    assert trial_component.status.primary_status == _TrialComponentStatusType.InProgress.value
 
 
 def _verify_tc_status_when_successfully_exit(trial_component, old_end_time=None):
-    assert (
-        trial_component.status.primary_status
-        == _TrialComponentStatusType.Completed.value
-    )
+    assert trial_component.status.primary_status == _TrialComponentStatusType.Completed.value
     assert isinstance(trial_component.start_time, datetime.datetime)
     assert isinstance(trial_component.end_time, datetime.datetime)
     if old_end_time: