Support s3 log_dir in Tensorboard logger.

Author: hengoren-exp

composer/loggers/tensorboard_logger.py

@@ -5,6 +5,7 @@
 
 from pathlib import Path
 from typing import Any, Optional, Sequence, Union
+from urllib.parse import urlparse
 
 import numpy as np
 import torch
@@ -108,9 +109,16 @@ def _initialize_summary_writer(self):
 
         assert self.run_name is not None
         assert self.log_dir is not None
-        # We name the child directory after the run_name to ensure the run_name shows up
-        # in the Tensorboard GUI.
-        summary_writer_log_dir = Path(self.log_dir) / self.run_name
+
+        parsed = urlparse(self.log_dir)
+        # TODO: Handle other remote storage schemes
+        if parsed.scheme == 's3':
+            scheme, bucket, prefix, _, _, _ = parsed
+            summary_writer_log_dir = f"{scheme}://{bucket}/{prefix.strip('/')}/{self.run_name}"
+        else:
+            # We name the child directory after the run_name to ensure the run_name shows up
+            # in the Tensorboard GUI.
+            summary_writer_log_dir = str(Path(self.log_dir) / self.run_name)
 
         # Disable SummaryWriter's internal flushing to avoid file corruption while
         # file staged for upload to an ObjectStore.

tests/loggers/test_tensorboard_logger.py

@@ -3,6 +3,8 @@
 
 from typing import Sequence
 
+import boto3
+import moto
 import pytest
 import torch
 
@@ -54,8 +56,12 @@ def test_tensorboard_log_image(test_tensorboard_logger, dummy_state):
     # Tensorboard images are stored inline, so we can't check them automatically.
 
 
+@moto.mock_aws
 def test_tensorboard_logger_s3_log_dir(dummy_state):
     bucket_name = 'test-tensorboard-bucket'
+    s3 = boto3.client('s3')
+    s3.create_bucket(Bucket=bucket_name)
+
     test_s3_log_dir = f's3://{bucket_name}/log_prefix'
 
     dummy_state.run_name = 'tensorboard-test-log-s3'
@@ -64,4 +70,4 @@ def test_tensorboard_logger_s3_log_dir(dummy_state):
     tensorboard_logger.init(dummy_state, logger)
     assert tensorboard_logger.writer is not None
     expected_log_dir = f'{test_s3_log_dir}/{dummy_state.run_name}'
-    assert str(tensorboard_logger.writer.log_dir) == expected_log_dir
+    assert tensorboard_logger.writer.log_dir == expected_log_dir