handle max input/output artifacts on TC

Author: yzhu0

src/smexperiments/tracker.py

@@ -252,6 +252,8 @@ def log_input(self, name, value, media_type=None):
             value (str): The value.
             media_type (str, optional): The MediaType (MIME type) of the value
         """
+        if len(self.trial_component.input_artifacts) >= 30:
+            raise ValueError("Cannot add more than 30 input_artifacts under tracker trial_component.")
         self.trial_component.input_artifacts[name] = api_types.TrialComponentArtifact(value, media_type=media_type)
 
     def log_output(self, name, value, media_type=None):
@@ -270,6 +272,8 @@ def log_output(self, name, value, media_type=None):
             value (str): The value.
             media_type (str, optional): The MediaType (MIME type) of the value.
         """
+        if len(self.trial_component.output_artifacts) >= 30:
+            raise ValueError("Cannot add more than 30 output_artifacts under tracker trial_component")
         self.trial_component.output_artifacts[name] = api_types.TrialComponentArtifact(value, media_type=media_type)
 
     def log_artifact(self, file_path, name=None, media_type=None):
@@ -303,6 +307,8 @@ def log_output_artifact(self, file_path, name=None, media_type=None):
             media_type (str, optional): The MediaType (MIME type) of the file. If not specified, this library
                 will attempt to infer the media type from the file extension of ``file_path``.
         """
+        if len(self.trial_component.output_artifacts) >= 30:
+            raise ValueError("Cannot add more than 30 output_artifacts under tracker trial_component")
         media_type = media_type or _guess_media_type(file_path)
         name = name or _resolve_artifact_name(file_path)
         s3_uri, etag = self._artifact_uploader.upload_artifact(file_path)
@@ -326,6 +332,8 @@ def log_input_artifact(self, file_path, name=None, media_type=None):
             media_type (str, optional): The MediaType (MIME type) of the file. If not specified, this library
                 will attempt to infer the media type from the file extension of ``file_path``.
         """
+        if len(self.trial_component.input_artifacts) >= 30:
+            raise ValueError("Cannot add more than 30 input_artifacts under tracker trial_component.")
         media_type = media_type or _guess_media_type(file_path)
         name = name or _resolve_artifact_name(file_path)
         s3_uri, etag = self._artifact_uploader.upload_artifact(file_path)

tests/unit/test_tracker.py

@@ -292,6 +292,56 @@ def test_log_input_artifact(under_test):
     assert "text/plain" == under_test.trial_component.input_artifacts["foo.txt"].media_type
 
 
+def test_log_inputs_error(under_test):
+    for index in range(0, 30):
+        file_path = "foo" + str(index) + ".txt"
+        under_test.trial_component.input_artifacts[file_path] = {
+            "foo": api_types.TrialComponentArtifact(value="baz" + str(index), media_type="text/text")
+        }
+    with pytest.raises(ValueError):
+        under_test.log_input("foo.txt", "name", "whizz/bang")
+
+
+def test_log_outputs(under_test):
+    for index in range(0, 30):
+        file_path = "foo" + str(index) + ".txt"
+        under_test.trial_component.output_artifacts[file_path] = {
+            "foo": api_types.TrialComponentArtifact(value="baz" + str(index), media_type="text/text")
+        }
+    with pytest.raises(ValueError):
+        under_test.log_output("foo.txt", "name", "whizz/bang")
+
+
+def test_log_multiple_input_artifact(under_test):
+    for index in range(0, 30):
+        file_path = "foo" + str(index) + ".txt"
+        under_test._artifact_uploader.upload_artifact.return_value = (
+            "s3uri_value" + str(index),
+            "etag_value" + str(index),
+        )
+        under_test.log_input_artifact(file_path, "name" + str(index), "whizz/bang" + str(index))
+        under_test._artifact_uploader.upload_artifact.assert_called_with(file_path)
+
+    under_test._artifact_uploader.upload_artifact.return_value = ("s3uri_value", "etag_value")
+    with pytest.raises(ValueError):
+        under_test.log_input_artifact("foo.txt", "name", "whizz/bang")
+
+
+def test_log_multiple_output_artifact(under_test):
+    for index in range(0, 30):
+        file_path = "foo" + str(index) + ".txt"
+        under_test._artifact_uploader.upload_artifact.return_value = (
+            "s3uri_value" + str(index),
+            "etag_value" + str(index),
+        )
+        under_test.log_output_artifact(file_path, "name" + str(index), "whizz/bang" + str(index))
+        under_test._artifact_uploader.upload_artifact.assert_called_with(file_path)
+
+    under_test._artifact_uploader.upload_artifact.return_value = ("s3uri_value", "etag_value")
+    with pytest.raises(ValueError):
+        under_test.log_output_artifact("foo.txt", "name", "whizz/bang")
+
+
 def test_log_pr_curve(under_test):
 
     y_true = [0, 0, 1, 1]