Repository: mosaicml/composer

Issue:
Run directory upload test is flaky
The `tests/callbacks/test_run_directory_uploader.py::test_run_directory_uploader` test is flaky.

** To reproduce

Running the test one GitHub actions gives the following output:

```
tests/callbacks/test_run_directory_uploader.py::test_run_directory_uploader[False]
197
  /opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/_pytest/threadexception.py:75: PytestUnhandledThreadExceptionWarning: Exception in thread Thread-3
198
  
199
  Traceback (most recent call last):
200
    File "/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/threading.py", line 932, in _bootstrap_inner
201
      self.run()
202
    File "/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/threading.py", line 870, in run
203
      self._target(*self._args, **self._kwargs)
204
    File "/home/runner/work/composer/composer/composer/callbacks/run_directory_uploader.py", line 313, in _upload_worker
205
      provider.upload_object(
206
    File "/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/libcloud/storage/drivers/local.py", line 521, in upload_object
207
      shutil.copy(file_path, obj_path)
208
    File "/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/libcloud/storage/drivers/local.py", line 106, in __exit__
209
      raise value
210
    File "/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/libcloud/storage/drivers/local.py", line 521, in upload_object
211
      shutil.copy(file_path, obj_path)
212
    File "/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/shutil.py", line 418, in copy
213
      copyfile(src, dst, follow_symlinks=follow_symlinks)
214
    File "/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/shutil.py", line 264, in copyfile
215
      with open(src, 'rb') as fsrc, open(dst, 'wb') as fdst:
216
  FileNotFoundError: [Errno 2] No such file or directory: '/tmp/tmpxw_0yxo6/1639786439.0068386/dummy_file'
217
  
```

## Expected behavior

The test should not be flaky.




Discussion:

Resolving patch:
diff --git a/composer/callbacks/run_directory_uploader.py b/composer/callbacks/run_directory_uploader.py
--- a/composer/callbacks/run_directory_uploader.py
+++ b/composer/callbacks/run_directory_uploader.py
@@ -12,8 +12,9 @@
 import tempfile
 import threading
 import time
+import uuid
 import warnings
-from typing import Any, Callable, Dict, Optional, Type, Union
+from typing import Any, Callable, Dict, Optional, Tuple, Type, Union
 
 from composer.core.callback import Callback
 from composer.core.logging import Logger
@@ -136,8 +137,8 @@ def __init__(
 
         if use_procs:
             mp_ctx = multiprocessing.get_context('spawn')
-            self._file_upload_queue: Union[queue.Queue[str],
-                                           multiprocessing.JoinableQueue[str]] = mp_ctx.JoinableQueue()
+            self._file_upload_queue: Union[queue.Queue[Tuple[str, str]],
+                                           multiprocessing.JoinableQueue[Tuple[str, str]]] = mp_ctx.JoinableQueue()
             self._finished_cls: Union[Callable[[], multiprocessing._EventType], Type[threading.Event]] = mp_ctx.Event
             self._proc_class = mp_ctx.Process
         else:
@@ -163,7 +164,6 @@ def init(self, state: State, logger: Logger) -> None:
                              kwargs={
                                  "file_queue": self._file_upload_queue,
                                  "is_finished": self._finished,
-                                 "upload_staging_dir": self._upload_staging_folder,
                                  "provider_name": self._provider,
                                  "container_name": self._container,
                                  "object_name_prefix": self._object_name_prefix,
@@ -200,6 +200,8 @@ def post_close(self):
             self._finished.set()
         for worker in self._workers:
             worker.join()
+        if self._tempdir is not None:
+            self._tempdir.cleanup()
 
     def _trigger_upload(self, logger: Optional[Logger], log_level: Optional[LogLevel]) -> None:
         # Ensure that every rank is at this point
@@ -227,12 +229,12 @@ def _trigger_upload(self, logger: Optional[Logger], log_level: Optional[LogLevel
             modified_files = get_modified_files(self._last_upload_timestamp)
             for filepath in modified_files:
                 relpath = os.path.relpath(filepath, run_directory)  # chop off the run directory
-                copied_path = os.path.join(self._upload_staging_folder, str(new_last_uploaded_timestamp), relpath)
+                copied_path = os.path.join(self._upload_staging_folder, str(uuid.uuid4()))
                 files_to_be_uploaded.append(relpath)
                 copied_path_dirname = os.path.dirname(copied_path)
                 os.makedirs(copied_path_dirname, exist_ok=True)
                 shutil.copy2(filepath, copied_path)
-                self._file_upload_queue.put_nowait(copied_path)
+                self._file_upload_queue.put_nowait((copied_path, relpath))
 
             self._last_upload_timestamp = new_last_uploaded_timestamp
             if logger is not None and log_level is not None:
@@ -266,7 +268,6 @@ def _validate_credentials(
 def _upload_worker(
     file_queue: Union[queue.Queue[str], multiprocessing.JoinableQueue[str]],
     is_finished: Union[multiprocessing._EventType, threading.Event],
-    upload_staging_dir: str,
     provider_name: str,
     container_name: str,
     object_name_prefix: str,
@@ -281,7 +282,6 @@ def _upload_worker(
             set when training is finished and no new files will be added to the queue.
             The worker will continue to upload existing files that are in the queue.
             When the queue is empty, the worker will exit.
-        upload_staging_dir (str): The upload staging directory.
         provider_name (str): The cloud provider name.
         container_name (str): The container name (e.g. s3 bucket) for the provider
             where files will be uploaded.
@@ -297,14 +297,12 @@ def _upload_worker(
     container = provider.get_container(container_name)
     while True:
         try:
-            file_path_to_upload = file_queue.get(block=True, timeout=0.5)
+            file_path_to_upload, obj_name = file_queue.get(block=True, timeout=0.5)
         except queue.Empty:
             if is_finished.is_set():
                 break
             else:
                 continue
-        obj_name = os.path.sep.join(os.path.relpath(file_path_to_upload, upload_staging_dir).split(
-            os.path.sep)[1:])  # the first folder is the upload timestamp. Chop that off.
         log.info("Uploading file %s to %s://%s/%s%s", file_path_to_upload, provider_name, container_name,
                  object_name_prefix, obj_name)
         retry_counter = 0
diff --git a/composer/utils/run_directory.py b/composer/utils/run_directory.py
--- a/composer/utils/run_directory.py
+++ b/composer/utils/run_directory.py
@@ -17,7 +17,6 @@ def get_run_directory():
 def get_relative_to_run_directory(*path: str, base: str = ".") -> str:
     run_directory = get_run_directory()
     if run_directory is None:
-        raise ValueError("shouldn't be triggered")
         return os.path.join(base, *path)
     return os.path.join(run_directory, *path)
 


Test patch:
diff --git a/tests/callbacks/test_run_directory_uploader.py b/tests/callbacks/test_run_directory_uploader.py
--- a/tests/callbacks/test_run_directory_uploader.py
+++ b/tests/callbacks/test_run_directory_uploader.py
@@ -12,12 +12,8 @@
 from composer.utils.run_directory import get_run_directory
 
 
-# This test is flaky see: https://github.com/mosaicml/composer/issues/176
 @pytest.mark.parametrize("use_procs", [False, True])
-# TODO(ravi) -- remove the pytest.in #110. The TRAINING_END event is likely slow as it has to copy many
-# files created by the ddp test. #110 grately reduces the number of files from the DDP test.
 @pytest.mark.timeout(15)
-@pytest.mark.xfail
 def test_run_directory_uploader(tmpdir: pathlib.Path, use_procs: bool, dummy_state: State, dummy_logger: Logger):
     dummy_state.epoch = 0
     dummy_state.step = 0