Repository: mosaicml/composer

Issue:
WandB error due to multiple artifacts with the same ID
When running a baseline resnet50 model on imagenet, I encountered this error:
```
wandb: ERROR Error while calling W&B API: Error 1062: Duplicate entry '6394579-1' for key 'unique_artifact_collection_membership_version' (<Response [409]>)
Exception in thread Thread-7:
Traceback (most recent call last):
  File "/usr/lib/python3.8/threading.py", line 932, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.8/threading.py", line 870, in run
    self._target(*self._args, **self._kwargs)
  File "/usr/local/lib/python3.8/dist-packages/wandb/filesync/step_upload.py", line 50, in _thread_body
    self._handle_event(event)
  File "/usr/local/lib/python3.8/dist-packages/wandb/filesync/step_upload.py", line 79, in _handle_event
    self._maybe_commit_artifact(job.artifact_id)
  File "/usr/local/lib/python3.8/dist-packages/wandb/filesync/step_upload.py", line 161, in _maybe_commit_artifact
    self._api.commit_artifact(artifact_id)
  File "/usr/local/lib/python3.8/dist-packages/wandb/sdk/internal/internal_api.py", line 2235, in commit_artifact
    response = self.gql(mutation, variable_values={"artifactID": artifact_id})
  File "/usr/local/lib/python3.8/dist-packages/wandb/sdk/lib/retry.py", line 102, in __call__
    result = self._call_fn(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/wandb/sdk/internal/internal_api.py", line 147, in execute
    six.reraise(*sys.exc_info())
  File "/usr/local/lib/python3.8/dist-packages/six.py", line 719, in reraise
    raise value
  File "/usr/local/lib/python3.8/dist-packages/wandb/sdk/internal/internal_api.py", line 141, in execute
    return self.client.execute(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/wandb/vendor/gql-0.2.0/gql/client.py", line 52, in execute
    result = self._get_result(document, *args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/wandb/vendor/gql-0.2.0/gql/client.py", line 60, in _get_result
    return self.transport.execute(document, *args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/wandb/vendor/gql-0.2.0/gql/transport/requests.py", line 39, in execute
    request.raise_for_status()
  File "/usr/local/lib/python3.8/dist-packages/requests/models.py", line 953, in raise_for_status
    raise HTTPError(http_error_msg, response=self)
requests.exceptions.HTTPError: 409 Client Error: Conflict for url: https://api.wandb.ai/graphql
```

I've asked the WandB folks and they think it's from an attempted upload of an artifact with the same ID as another. The recent addition of artifact uploading from `run_directory` seems to be causing this, so PR #89 will disable it by default, but we need to verify that artifact uploads are working as expected.


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
@@ -20,7 +20,7 @@
 from composer.core.logging.logger import LogLevel
 from composer.core.state import State
 from composer.utils import ddp
-from composer.utils.run_directory import get_run_directory
+from composer.utils.run_directory import get_modified_files, get_run_directory
 
 log = logging.getLogger(__name__)
 
@@ -206,47 +206,43 @@ def _trigger_upload(self, logger: Optional[Logger], log_level: Optional[LogLevel
         # Assuming only the main thread on each rank writes to the run directory, then the barrier here will ensure
         # that the run directory is not being modified after we pass this barrier
         ddp.barrier()
-        if not ddp.get_local_rank() == 0:
-            return
-        run_directory = get_run_directory()
-        assert run_directory is not None, "invariant error"
-        # the disk time can differ from system time, so going to touch a file and then read the timestamp from it to get the real time
-        python_time = time.time()
-        touch_file = (pathlib.Path(run_directory) / f".{python_time}")
-        touch_file.touch()
-        new_last_uploaded_timestamp = os.path.getmtime(str(touch_file))
-
-        # Now, for each file that was modified since self._last_upload_timestamp, copy it to the temporary directory
-        files_to_be_uploaded = []
-
-        # check if any upload threads have crashed. if so, then shutdown the training process
-        for worker in self._workers:
-            if not worker.is_alive():
-                assert self._finished is not None, "invariant error"
-                self._finished.set()
-                raise RuntimeError("Upload worker crashed unexpectedly")
-        for root, dirs, files in os.walk(run_directory):
-            del dirs  # unused
-            for file in files:
-                if any(x.startswith(".") for x in file.split(os.path.sep)):
-                    # skip hidden files and folders
-                    continue
-                filepath = os.path.join(root, file)
+        if ddp.get_local_rank() == 0:
+            run_directory = get_run_directory()
+            assert run_directory is not None, "invariant error"
+            # the disk time can differ from system time, so going to touch a file and then read the timestamp from it to get the real time
+            python_time = time.time()
+            touch_file = (pathlib.Path(run_directory) / f".{python_time}")
+            touch_file.touch()
+            new_last_uploaded_timestamp = os.path.getmtime(str(touch_file))
+
+            # Now, for each file that was modified since self._last_upload_timestamp, copy it to the temporary directory
+            files_to_be_uploaded = []
+
+            # check if any upload threads have crashed. if so, then shutdown the training process
+            for worker in self._workers:
+                if not worker.is_alive():
+                    assert self._finished is not None, "invariant error"
+                    self._finished.set()
+                    raise RuntimeError("Upload worker crashed unexpectedly")
+            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
-                modified_time = os.path.getmtime(filepath)
-                if modified_time > self._last_upload_timestamp:
-                    copied_path = os.path.join(self._upload_staging_folder, str(new_last_uploaded_timestamp), relpath)
-                    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._last_upload_timestamp = new_last_uploaded_timestamp
-        if logger is not None and log_level is not None:
-            # now log which files are being uploaded. OK to do, since we're done reading the directory,
-            # and any logfiles will now have their last modified timestamp
-            # incremented past self._last_upload_timestamp
-            logger.metric(log_level, {"run_directory/uploaded_files": files_to_be_uploaded})
+                copied_path = os.path.join(self._upload_staging_folder, str(new_last_uploaded_timestamp), relpath)
+                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._last_upload_timestamp = new_last_uploaded_timestamp
+            if logger is not None and log_level is not None:
+                # now log which files are being uploaded. OK to do, since we're done reading the directory,
+                # and any logfiles will now have their last modified timestamp
+                # incremented past self._last_upload_timestamp
+                logger.metric(log_level, {"run_directory/uploaded_files": files_to_be_uploaded})
+
+        # Ensure that other callbacks do not start writing to the run directory until we copying everything
+        ddp.barrier()
 
 
 def _validate_credentials(
diff --git a/composer/loggers/wandb_logger.py b/composer/loggers/wandb_logger.py
--- a/composer/loggers/wandb_logger.py
+++ b/composer/loggers/wandb_logger.py
@@ -6,14 +6,14 @@
 import sys
 from typing import Any, Dict, Optional
 
-from composer.core.logging import LogLevel, RankZeroLoggerBackend, TLogData
+from composer.core.logging import BaseLoggerBackend, LogLevel, TLogData
 from composer.core.types import Logger, State, StateDict
-from composer.utils.run_directory import get_run_directory
+from composer.utils import ddp, run_directory
 
 import wandb  # isort: skip
 
 
-class WandBLoggerBackend(RankZeroLoggerBackend):
+class WandBLoggerBackend(BaseLoggerBackend):
     """Log to Weights and Biases (https://wandb.ai/)
 
     Args:
@@ -33,21 +33,26 @@ def __init__(self,
         super().__init__()
         self._log_artifacts = log_artifacts
         self._log_artifacts_every_n_batches = log_artifacts_every_n_batches
+        self._last_upload_timestamp = 0.0
         if init_params is None:
             init_params = {}
         self._init_params = init_params
 
     def _log_metric(self, epoch: int, step: int, log_level: LogLevel, data: TLogData):
         del epoch, log_level  # unused
-        wandb.log(data, step=step)
+        if ddp.get_local_rank() == 0:
+            wandb.log(data, step=step)
 
     def state_dict(self) -> StateDict:
         # Storing these fields in the state dict to support run resuming in the future.
+        if ddp.get_local_rank() != 0:
+            raise RuntimeError("WandB can only be checkpointed on rank 0")
         return {"name": wandb.run.name, "project": wandb.run.project, "entity": wandb.run.entity, "id": wandb.run.id}
 
     def init(self, state: State, logger: Logger) -> None:
         del state, logger  # unused
-        wandb.init(**self._init_params)
+        if ddp.get_local_rank() == 0:
+            wandb.init(**self._init_params)
 
     def batch_end(self, state: State, logger: Logger) -> None:
         del logger  # unused
@@ -72,16 +77,25 @@ def _upload_artifacts(self):
         # so uploads will not block the training loop
         # slowdown is likely from extra I/O of scanning the directory and/or
         # scheduling uploads
-        run_directory = get_run_directory()
-        if run_directory is not None:
-            for subfile in os.listdir(run_directory):
-                artifact = wandb.Artifact(name=subfile, type=subfile)
-                full_path = os.path.join(run_directory, subfile)
-                if os.path.isdir(full_path):
-                    artifact.add_dir(full_path)
-                else:
-                    artifact.add_file(full_path)
+        run_directory_name = run_directory.get_run_directory()
+        if run_directory_name is None:
+            return
+        # barrier that every process has reached this point and that
+        # previous callbacks are finished writing to the run directory
+        ddp.barrier()
+        if ddp.get_local_rank() == 0:
+            modified_files = run_directory.get_modified_files(self._last_upload_timestamp)
+            for modified_file in modified_files:
+                file_type = modified_file.split(".")[-1]
+                relpath = os.path.relpath(modified_file, run_directory_name)
+                relpath = relpath.replace("/", "-")
+                artifact = wandb.Artifact(name=relpath, type=file_type)
+                artifact.add_file(os.path.abspath(modified_file))
                 wandb.log_artifact(artifact)
+            self._last_upload_timestamp = run_directory.get_run_directory_timestamp()
+        # barrier to ensure that other processes do not continue to other callbacks
+        # that could start writing to the run directory
+        ddp.barrier()
 
     def post_close(self) -> None:
         # Cleaning up on post_close so all artifacts are uploaded
@@ -90,8 +104,9 @@ def post_close(self) -> None:
 
         exc_tpe, exc_info, tb = sys.exc_info()
 
-        if (exc_tpe, exc_info, tb) == (None, None, None):
-            wandb.finish(0)
-        else:
-            # record there was an error
-            wandb.finish(1)
+        if ddp.get_local_rank() == 0:
+            if (exc_tpe, exc_info, tb) == (None, None, None):
+                wandb.finish(0)
+            else:
+                # record there was an error
+                wandb.finish(1)
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
@@ -2,6 +2,8 @@
 
 import logging
 import os
+import pathlib
+import time
 
 log = logging.getLogger(__name__)
 
@@ -12,8 +14,56 @@ def get_run_directory():
     return os.environ.get(_RUN_DIRECTORY_KEY)
 
 
-def get_relative_to_run_directory(path: str, base: str = ".") -> str:
+def get_relative_to_run_directory(*path: str, base: str = ".") -> str:
     run_directory = get_run_directory()
     if run_directory is None:
-        return os.path.join(base, path)
-    return os.path.join(run_directory, path)
+        raise ValueError("shouldn't be triggered")
+        return os.path.join(base, *path)
+    return os.path.join(run_directory, *path)
+
+
+def get_modified_files(modified_since_timestamp: float, *, ignore_hidden: bool = True):
+    """Returns a list of files (recursively) in the run directory that have been modified since
+    ``modified_since_timestamp``.
+
+    Args:
+        modified_since_timestamp (float): Minimum last modified timestamp(in seconds since EPOCH)
+            of files to include.
+        ignore_hidden (bool, optional): Whether to ignore hidden files and folders (default: ``True``)
+    Returns:
+        List[str]: List of filepaths that have been modified since ``modified_since_timestamp``
+    """
+    modified_files = []
+    run_directory = get_run_directory()
+    if run_directory is None:
+        raise RuntimeError("Run directory is not defined")
+    for root, dirs, files in os.walk(run_directory):
+        del dirs  # unused
+        for file in files:
+            if ignore_hidden and any(x.startswith(".") for x in file.split(os.path.sep)):
+                # skip hidden files and folders
+                continue
+            filepath = os.path.join(root, file)
+            modified_time = os.path.getmtime(filepath)
+            if modified_time >= modified_since_timestamp:
+                modified_files.append(filepath)
+    return modified_files
+
+
+def get_run_directory_timestamp() -> float:
+    """Returns the current timestamp on the run directory filesystem.
+    Note that the disk time can differ from system time (e.g. when using
+    network filesystems).
+
+    Returns:
+        float: the current timestamp on the run directory filesystem.
+    """
+    run_directory = get_run_directory()
+    if run_directory is None:
+        raise RuntimeError("Run directory is not defined")
+    python_time = time.time()
+    touch_file = (pathlib.Path(run_directory) / f".{python_time}")
+    touch_file.touch()
+    new_last_uploaded_timestamp = os.path.getmtime(str(touch_file))
+    os.remove(str(touch_file))
+    return new_last_uploaded_timestamp


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -13,16 +13,16 @@
 import _pytest.mark
 import pytest
 import torch.distributed
+from _pytest.monkeypatch import MonkeyPatch
 
 import composer
-from composer.utils.run_directory import get_relative_to_run_directory
+from composer.utils import run_directory
 
 # Allowed options for pytest.mark.world_size()
 # Important: when updating this list, make sure to also up scripts/test.sh
 # so tests of all world sizes will be executed
 WORLD_SIZE_OPTIONS = (1, 2)
 
-PYTEST_DDP_LOCKFILE_DIR = get_relative_to_run_directory(os.path.join("pytest_lockfiles"))
 DDP_TIMEOUT = datetime.timedelta(seconds=5)
 
 # Add the path of any pytest fixture files you want to make global
@@ -112,12 +112,13 @@ def set_loglevels():
     logging.getLogger(composer.__name__).setLevel(logging.DEBUG)
 
 
-@pytest.fixture
-def ddp_tmpdir(tmpdir: pathlib.Path) -> str:
+@pytest.fixture(autouse=True)
+def subfolder_run_directory(tmpdir: pathlib.Path, monkeypatch: MonkeyPatch) -> None:
     tmpdir_test_folder_name = os.path.basename(os.path.normpath(str(tmpdir)))
-    test_folder_tmpdir = get_relative_to_run_directory(tmpdir_test_folder_name, base=os.path.join(str(tmpdir), ".."))
+    test_folder_tmpdir = run_directory.get_relative_to_run_directory(tmpdir_test_folder_name,
+                                                                     base=os.path.join(str(tmpdir), ".."))
+    monkeypatch.setenv(run_directory._RUN_DIRECTORY_KEY, test_folder_tmpdir)
     os.makedirs(test_folder_tmpdir, exist_ok=True)
-    return os.path.abspath(test_folder_tmpdir)
 
 
 @pytest.fixture(autouse=True)
@@ -139,24 +140,24 @@ def configure_ddp(request: pytest.FixtureRequest):
 
 
 @pytest.fixture(autouse=True)
-def wait_for_all_procs(ddp_tmpdir: str):
+def wait_for_all_procs(subfolder_run_directory: None):
     yield
     if not 'RANK' in os.environ:
         # Not running in a DDP environment
         return
     global_rank = int(os.environ['RANK'])
     world_size = int(os.environ['WORLD_SIZE'])
-    proc_lockfile = os.path.join(ddp_tmpdir, f"{global_rank}_finished")
+    proc_lockfile = run_directory.get_relative_to_run_directory(f"{global_rank}_finished")
     pathlib.Path(proc_lockfile).touch(exist_ok=False)
     # other processes shouldn't be (too) far behind the current one
     end_time = datetime.datetime.now() + datetime.timedelta(seconds=15)
     for rank in range(world_size):
-        if not os.path.exists(os.path.join(ddp_tmpdir, f"{rank}_finished")):
+        if not os.path.exists(run_directory.get_relative_to_run_directory(f"{rank}_finished")):
             # sleep for the other procs to write their finished file
             if datetime.datetime.now() < end_time:
                 time.sleep(0.1)
             else:
-                test_name = os.path.basename(os.path.normpath(ddp_tmpdir))
+                test_name = os.path.basename(os.path.normpath(str(run_directory.get_run_directory())))
                 raise RuntimeError(f"Rank {rank} did not finish test {test_name}")
 
 
diff --git a/tests/test_amp.py b/tests/test_amp.py
--- a/tests/test_amp.py
+++ b/tests/test_amp.py
@@ -32,7 +32,7 @@ def run_and_measure_memory(precision: Precision) -> int:
 
 @pytest.mark.timeout(60)
 @pytest.mark.gpu
-def test_fp16_mixed(ddp_tmpdir: str):
+def test_fp16_mixed():
     memory_full = run_and_measure_memory(Precision.FP32)
     memory_amp = run_and_measure_memory(Precision.AMP)
     assert memory_amp < 0.7 * memory_full
diff --git a/tests/test_logger.py b/tests/test_logger.py
--- a/tests/test_logger.py
+++ b/tests/test_logger.py
@@ -13,7 +13,8 @@
 from composer.core.logging import Logger, LogLevel
 from composer.core.state import State
 from composer.loggers.file_logger import FileLoggerBackend
-from composer.loggers.logger_hparams import FileLoggerBackendHparams, TQDMLoggerBackendHparams
+from composer.loggers.logger_hparams import (FileLoggerBackendHparams, TQDMLoggerBackendHparams,
+                                             WandBLoggerBackendHparams)
 from composer.trainer.trainer_hparams import TrainerHparams
 
 
@@ -90,3 +91,15 @@ def get_mock_tqdm(position: int, *args, **kwargs):
     for mock_tqdm in is_train_to_mock_tqdms[False]:
         assert mock_tqdm.update.call_count == trainer._eval_subset_num_batches
         mock_tqdm.close.assert_called_once()
+
+
+@pytest.mark.parametrize("world_size", [
+    pytest.param(1),
+    pytest.param(2, marks=pytest.mark.world_size(2)),
+])
+@pytest.mark.timeout(10)
+def test_wandb_logger(mosaic_trainer_hparams: TrainerHparams, world_size: int, monkeypatch: MonkeyPatch):
+    del world_size  # unused. Set via launcher script
+    mosaic_trainer_hparams.loggers = [WandBLoggerBackendHparams(log_artifacts=True, log_artifacts_every_n_batches=1)]
+    trainer = mosaic_trainer_hparams.initialize_object()
+    trainer.fit()
diff --git a/tests/trainer/test_checkpoint.py b/tests/trainer/test_checkpoint.py
--- a/tests/trainer/test_checkpoint.py
+++ b/tests/trainer/test_checkpoint.py
@@ -18,7 +18,7 @@
 from composer.trainer.devices import CPUDeviceHparams, DeviceHparams, GPUDeviceHparams
 from composer.trainer.trainer import Trainer
 from composer.trainer.trainer_hparams import TrainerHparams, callback_registry
-from composer.utils import ddp
+from composer.utils import ddp, run_directory
 from tests.test_state import assert_state_equivalent
 from tests.utils.deep_compare import deep_compare
 
@@ -139,7 +139,6 @@ def test_checkpoint(
     checkpoint_filename: str,
     validate_every_n_batches: int,
     validate_every_n_epochs: int,
-    ddp_tmpdir: str,
 ):
     """strategy:
     - train two epochs. capture checkpoints after `checkpoint_interval` and ep2.
@@ -147,32 +146,28 @@ def test_checkpoint(
     - assert that the checkpoint from the new trainer at the end is the same as the checkpoint from the first trainer at the end.
     """
     del world_size  # unused. Read via env variable
-    checkpoint_a_folder = os.path.join(ddp_tmpdir, "first")
-    checkpointing_trainer_hparams.checkpoint_folder = checkpoint_a_folder
 
+    checkpointing_trainer_hparams.device = device_hparams
+
+    checkpoint_a_folder = "first"
+    checkpointing_trainer_hparams.checkpoint_folder = checkpoint_a_folder
     checkpointing_trainer_hparams.checkpoint_interval_unit = "ep" if checkpoint_filename.startswith("ep") else "it"
     checkpointing_trainer_hparams.validate_every_n_batches = validate_every_n_batches
     checkpointing_trainer_hparams.validate_every_n_epochs = validate_every_n_epochs
     final_checkpoint = "ep2.pt" if checkpoint_filename.startswith("ep") else "it8.pt"
-
-    checkpointing_trainer_hparams.device = device_hparams
-
     _test_checkpoint_trainer(checkpointing_trainer_hparams)
     checkpoint_a_file_path = os.path.join(checkpoint_a_folder, f"{checkpoint_filename}.pt")
-    checkpoint_b_file_path = os.path.join(checkpoint_a_folder, final_checkpoint)
-    trainer_1_hparams_filepath = os.path.join(checkpoint_a_folder, "hparams.yaml")
+    checkpoint_b_file_path = run_directory.get_relative_to_run_directory(checkpoint_a_folder, final_checkpoint)
+    trainer_1_hparams_filepath = run_directory.get_relative_to_run_directory(checkpoint_a_folder, "hparams.yaml")
 
     second_trainer_hparams = TrainerHparams.create(trainer_1_hparams_filepath, cli_args=False)
-    checkpoint_b_folder = os.path.join(ddp_tmpdir, "second")
+    checkpoint_b_folder = "second"
     second_trainer_hparams.checkpoint_folder = checkpoint_b_folder
-    second_trainer_hparams.checkpoint_filepath = checkpoint_a_file_path
+    second_trainer_hparams.checkpoint_filepath = run_directory.get_relative_to_run_directory(checkpoint_a_file_path)
     _test_checkpoint_trainer(second_trainer_hparams)
 
-    checkpoint_c_file_path = os.path.join(ddp_tmpdir, "checkpoint_c.pt")
-    trainer_2_hparams_filepath = os.path.join(ddp_tmpdir, "trainer_2_hparams.yaml")
-
-    checkpoint_c_file_path = os.path.join(checkpoint_b_folder, final_checkpoint)
-    trainer_2_hparams_filepath = os.path.join(checkpoint_b_folder, "hparams.yaml")
+    checkpoint_c_file_path = run_directory.get_relative_to_run_directory(checkpoint_b_folder, final_checkpoint)
+    trainer_2_hparams_filepath = run_directory.get_relative_to_run_directory(checkpoint_b_folder, "hparams.yaml")
 
     if ddp.get_global_rank() == 0:
 
diff --git a/tests/trainer/test_ddp.py b/tests/trainer/test_ddp.py
--- a/tests/trainer/test_ddp.py
+++ b/tests/trainer/test_ddp.py
@@ -2,9 +2,8 @@
 
 import collections.abc
 import os
-import pathlib
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Type, Union
+from typing import Dict, List, Optional
 
 import pytest
 import torch
@@ -20,22 +19,25 @@
 from composer.core.state import State
 from composer.datasets import DataloaderHparams, SyntheticBatchPairDataset, SyntheticHparamsMixin
 from composer.datasets.hparams import DatasetHparams
-from composer.models import ModelHparams
 from composer.trainer.deepspeed import DeepSpeedHparams
 from composer.trainer.devices import CPUDeviceHparams, DeviceHparams, GPUDeviceHparams
 from composer.trainer.trainer_hparams import TrainerHparams, callback_registry, dataset_registry
-from composer.utils import ddp
+from composer.utils import ddp, run_directory
 from tests.fixtures.models import SimpleBatchPairModel
 
 
-def get_file_path(tmpdir: Union[str, pathlib.Path], *, rank: int, is_train: bool) -> str:
+def get_file_path(*, rank: int, is_train: bool) -> str:
     train_str = "train" if is_train else "val"
-    return os.path.join(tmpdir, f"{train_str}_rank_{rank}_num_accesses")
+    rundir = run_directory.get_run_directory()
+    assert rundir is not None
+    return os.path.join(rundir, f"{train_str}_rank_{rank}_num_accesses")
 
 
-def get_batch_file_path(tmpdir: Union[str, pathlib.Path], *, rank: int, epoch: int, is_train: bool) -> str:
+def get_batch_file_path(*, rank: int, epoch: int, is_train: bool) -> str:
     train_str = "train" if is_train else "val"
-    return os.path.join(tmpdir, f"{train_str}-rank-{rank}-epoch-{epoch}-batch0.pt")
+    rundir = run_directory.get_run_directory()
+    assert rundir is not None
+    return os.path.join(rundir, f"{train_str}-rank-{rank}-epoch-{epoch}-batch0.pt")
 
 
 class TrackedDataset(types.Dataset):
@@ -45,15 +47,14 @@ class TrackedDataset(types.Dataset):
     Because of atomic file writes, it is slow and should not be used in any performance measurements.
     """
 
-    def __init__(self, is_train: bool, tmpdir: str, synthetic_dataset: SyntheticBatchPairDataset):
+    def __init__(self, is_train: bool, synthetic_dataset: SyntheticBatchPairDataset):
         self.dataset = synthetic_dataset
         self.is_train = is_train
-        self.tmpdir = tmpdir
         self.counter = 0
 
     def __getitem__(self, idx: int):
         self.counter += 1
-        with open(get_file_path(self.tmpdir, rank=ddp.get_global_rank(), is_train=self.is_train), "w+") as f:
+        with open(get_file_path(rank=ddp.get_global_rank(), is_train=self.is_train), "w+") as f:
             f.write(str(self.counter))
         return self.dataset[idx]
 
@@ -63,12 +64,10 @@ def __len__(self):
 
 @dataclass
 class TrackedDatasetHparams(DatasetHparams, SyntheticHparamsMixin):
-    tmpdir: Optional[str] = hp.optional("tmpdir", default=None)
     num_classes: Optional[int] = hp.optional("num_classes", default=None)
     data_shape: Optional[List[int]] = hp.optional("data_shape", default=None)
 
     def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHparams) -> types.DataLoader:
-        assert self.tmpdir is not None
         assert self.num_classes is not None
         assert self.data_shape is not None
         synthetic_dataset = SyntheticBatchPairDataset(
@@ -78,9 +77,7 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
             num_classes=self.num_classes,
         )
         drop_last = False
-        tracked_dataset = TrackedDataset(is_train=self.is_train,
-                                         tmpdir=self.tmpdir,
-                                         synthetic_dataset=synthetic_dataset)
+        tracked_dataset = TrackedDataset(is_train=self.is_train, synthetic_dataset=synthetic_dataset)
         sampler = ddp.get_sampler(tracked_dataset, drop_last=drop_last, shuffle=True)
         return dataloader_hparams.initialize_object(
             dataset=tracked_dataset,
@@ -92,16 +89,12 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
 
 class CheckBatch0(Callback):
 
-    def __init__(self, tmpdir: str):
+    def __init__(self):
         super().__init__()
-        self.tmpdir = tmpdir
 
     def _run_event(self, event: Event, state: State, logger: Logger) -> None:
         if event in (Event.BEFORE_FORWARD, Event.EVAL_BEFORE_FORWARD):
-            filepath = get_batch_file_path(self.tmpdir,
-                                           rank=ddp.get_global_rank(),
-                                           epoch=state.epoch,
-                                           is_train=state.model.training)
+            filepath = get_batch_file_path(rank=ddp.get_global_rank(), epoch=state.epoch, is_train=state.model.training)
             if os.path.exists(filepath):
                 return
             last_input, last_target = state.batch_pair
@@ -114,10 +107,9 @@ def _run_event(self, event: Event, state: State, logger: Logger) -> None:
 
 @dataclass
 class CheckBatch0Hparams(CallbackHparams):
-    tmpdir: str = hp.required("tmpdir")
 
     def initialize_object(self) -> Callback:
-        return CheckBatch0(self.tmpdir)
+        return CheckBatch0()
 
 
 @pytest.fixture(autouse=True)
@@ -136,8 +128,7 @@ def patch_registries(monkeypatch: MonkeyPatch):
     pytest.param(1),
     pytest.param(2, marks=pytest.mark.world_size(2)),
 ])
-def test_ddp(device: DeviceHparams, world_size: int, ddp_tmpdir: str, mosaic_trainer_hparams: TrainerHparams,
-             SimpleBatchPairModelHparams: Type[ModelHparams], deepspeed: bool) -> None:
+def test_ddp(device: DeviceHparams, world_size: int, mosaic_trainer_hparams: TrainerHparams, deepspeed: bool) -> None:
     """
     test strategy for ddp:
     1) Train a dummy model on two gps, for two epochs, using the tracked dataset.
@@ -166,7 +157,6 @@ def test_ddp(device: DeviceHparams, world_size: int, ddp_tmpdir: str, mosaic_tra
     hparams.train_dataset = TrackedDatasetHparams(
         synthetic_num_unique_samples=hparams.train_batch_size * hparams.train_subset_num_batches,
         is_train=True,
-        tmpdir=ddp_tmpdir,
         data_shape=list(model.in_shape),
         num_classes=model.num_classes,
     )
@@ -176,7 +166,6 @@ def test_ddp(device: DeviceHparams, world_size: int, ddp_tmpdir: str, mosaic_tra
     hparams.val_dataset = TrackedDatasetHparams(
         synthetic_num_unique_samples=hparams.eval_batch_size * hparams.eval_subset_num_batches,
         is_train=False,
-        tmpdir=ddp_tmpdir,
         data_shape=list(model.in_shape),
         num_classes=model.num_classes,
     )
@@ -193,7 +182,7 @@ def test_ddp(device: DeviceHparams, world_size: int, ddp_tmpdir: str, mosaic_tra
     hparams.loggers = []
     hparams.validate_every_n_batches = 0
     hparams.validate_every_n_epochs = 1
-    hparams.callbacks.append(CheckBatch0Hparams(tmpdir=ddp_tmpdir))
+    hparams.callbacks.append(CheckBatch0Hparams())
     if deepspeed:
         hparams.deepspeed = DeepSpeedHparams(enabled=True)
     trainer = hparams.initialize_object()
@@ -211,9 +200,9 @@ def test_ddp(device: DeviceHparams, world_size: int, ddp_tmpdir: str, mosaic_tra
     actual_val_num_loads = 0
 
     for i in range(ddp.get_world_size()):
-        with open(get_file_path(ddp_tmpdir, is_train=True, rank=i), "r") as f:
+        with open(get_file_path(is_train=True, rank=i), "r") as f:
             actual_train_num_loads += int(f.read())
-        with open(get_file_path(ddp_tmpdir, is_train=False, rank=i), "r") as f:
+        with open(get_file_path(is_train=False, rank=i), "r") as f:
             actual_val_num_loads += int(f.read())
     assert actual_train_num_loads == expected_train_num_loads, f"actual_train_num_loads({actual_train_num_loads}) != expected_train_num_loads({expected_train_num_loads})"
     assert actual_val_num_loads == expected_val_num_loads, f"actual_val_num_loads({actual_val_num_loads}) != expected_val_num_loads({expected_val_num_loads})"
@@ -224,7 +213,7 @@ def test_ddp(device: DeviceHparams, world_size: int, ddp_tmpdir: str, mosaic_tra
         for local_rank in range(ddp.get_local_world_size()):
             for is_train in (True, False):
                 data: Dict[str, types.Tensor] = torch.load(  # type: ignore
-                    get_batch_file_path(ddp_tmpdir, rank=local_rank, epoch=epoch, is_train=is_train),
+                    get_batch_file_path(rank=local_rank, epoch=epoch, is_train=is_train),
                     map_location='cpu',
                 )
                 for pickle in is_train_to_pickles[is_train]: