Log all files in a directory

Author: yzhu0

src/smexperiments/tracker.py

@@ -20,6 +20,7 @@
 import botocore
 import json
 from smexperiments._utils import get_module
+from os.path import join
 
 import dateutil
 
@@ -264,7 +265,7 @@ def log_output(self, name, value, media_type=None):
         Examples
             .. code-block:: python
 
-                # log input dataset s3 location
+                # log output dataset s3 location
                 my_tracker.log_output(name='prediction', value='s3://outputs/path')
 
         Args:
@@ -276,6 +277,26 @@ def log_output(self, name, value, media_type=None):
             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_artifacts(self, directory, media_type=None):
+        """Upload all the files under the directory to s3 and store it as artifacts in this trial component. The file
+        name is used as the artifact name
+
+        Examples
+            .. code-block:: python
+
+                # log local artifact
+                my_tracker.log_artifact(directory='/local/path)
+
+        Args:
+            directory (str): The directory of the local files to upload.
+            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``.
+        """
+        for dir_file in os.listdir(directory):
+            file_path = join(directory, dir_file)
+            artifact_name = os.path.splitext(dir_file)[0]
+            self.log_artifact(file_path=file_path, name=artifact_name, media_type=media_type)
+
     def log_artifact(self, file_path, name=None, media_type=None):
         """Legacy overload method to prevent breaking existing code.
 

tests/integ/test_tracker.py

@@ -95,6 +95,32 @@ def test_log_artifact(trial_component_obj, bucket, tempdir, sagemaker_boto_clien
     assert prefix in loaded.output_artifacts[artifact_name].value
 
 
+def test_log_artifacts(trial_component_obj, bucket, tempdir, sagemaker_boto_client):
+    prefix = name()
+    file_contents = "happy monkey monkey"
+    file_path = os.path.join(tempdir, "foo.txt")
+    file_path1 = os.path.join(tempdir, "bar.txt")
+    with open(file_path, "w") as foo_file:
+        foo_file.write(file_contents)
+    with open(file_path1, "w") as bar_file:
+        bar_file.write(file_contents)
+
+    with tracker.Tracker.load(
+        trial_component_obj.trial_component_name,
+        artifact_bucket=bucket,
+        artifact_prefix=prefix,
+        sagemaker_boto_client=sagemaker_boto_client,
+    ) as tracker_obj:
+        tracker_obj.log_artifacts(tempdir)
+    loaded = trial_component.TrialComponent.load(
+        trial_component_name=trial_component_obj.trial_component_name, sagemaker_boto_client=sagemaker_boto_client
+    )
+    assert "text/plain" == loaded.output_artifacts["foo"].media_type
+    assert prefix in loaded.output_artifacts["foo"].value
+    assert "text/plain" == loaded.output_artifacts["bar"].media_type
+    assert prefix in loaded.output_artifacts["bar"].value
+
+
 def test_create_default_bucket(boto3_session):
     bucket_name_prefix = _utils.name("sm-test")
     bucket = _utils.get_or_create_default_bucket(boto3_session, default_bucket_prefix=bucket_name_prefix)