fix(tracker.py): add base_trial_component_name parameter to create function

Author: alex6499cat

src/smexperiments/tracker.py

@@ -167,6 +167,7 @@ def load(
     @classmethod
     def create(
         cls,
+        base_trial_component_name="TrialComponent",
         display_name=None,
         artifact_bucket=None,
         artifact_prefix=None,
@@ -185,6 +186,7 @@ def create(
                 my_tracker = tracker.Tracker.create()
 
         Args:
+            base_trial_component_name: (str,optional). The name of the trial component resource that will be appended with a timestamp. Defaults to "TrialComponent".
             display_name: (str, optional). The display name of the trial component to track.
             artifact_bucket: (str, optional) The name of the S3 bucket to store artifacts to.
             artifact_prefix: (str, optional) The prefix to write artifacts to within ``artifact_bucket``
@@ -201,7 +203,7 @@ def create(
         sagemaker_boto_client = sagemaker_boto_client or _utils.sagemaker_client()
 
         tc = trial_component.TrialComponent.create(
-            trial_component_name=_utils.name("TrialComponent"),
+            trial_component_name=_utils.name(base_trial_component_name),
             display_name=display_name,
             sagemaker_boto_client=sagemaker_boto_client,
         )

tests/unit/test_tracker.py

@@ -149,13 +149,22 @@ def test_create(boto3_session, sagemaker_boto_client):
     trial_component_display_name = "foo-trial-component-display-name"
     sagemaker_boto_client.create_trial_component.return_value = {"TrialComponentName": trial_component_name}
     tracker_created = tracker.Tracker.create(
-        display_name=trial_component_display_name, sagemaker_boto_client=sagemaker_boto_client
+        base_trial_component_name="AlexName",
+        display_name=trial_component_display_name,
+        sagemaker_boto_client=sagemaker_boto_client,
     )
     assert trial_component_name == tracker_created.trial_component.trial_component_name
-
+    sagemaker_boto_client.create_trial_component.assert_called_with(
+        DisplayName="foo-trial-component-display-name", TrialComponentName=AnyStringWith("AlexName")
+    )
     assert tracker_created._metrics_writer is None
 
 
+class AnyStringWith(str):
+    def __eq__(self, other):
+        return self in other
+
+
 @pytest.fixture
 def trial_component_obj(sagemaker_boto_client):
     return trial_component.TrialComponent(sagemaker_boto_client)

tox.ini

@@ -59,6 +59,7 @@ deps =
     pytest
     pytest-cov
     docker
+    scikit-learn==0.24.2
 
 [testenv:flake8]
 basepython = python3