Repository: allenai/tango (language: python)

Issue:
Remote workspace
This is about a workspace implementation that stores the cache and assorted run information remotely. It's an implementation of the `Workspace` class that keeps the data in the cloud. We want this because we want to see the same workspace from servers at Emerald Landing, Cirrascale, GCE, and also from our work laptops.

I think the two leading candidates for storing step results are wandb and GCS. In my mind Google is going to be a more stable, conservative choice, but wandb has a ton of cool integrations and visualizations. If wandb is equally easy to set up, equally fast, and equally scalable (1TB for a model checkpoint?), I'm not sure we still need something that saves the data in GCS.

In `LocalWorkspace` I went out of my way to store run information and step information directly in the file system. It's nice to have the same system for both step results and metadata, because that's one less system. For a remote workspace, it might not scale to access a ton of tiny files in GCS though, so maybe it's time to bite the bullet and put it into a DB. If you choose to go that route, please give me a chance to convince you of the merits of writing SQL queries directly in Python code before you start using a persistence framework.

The remote workspace will need a local cache for step results, both for performance, and also because we want to memory map some of those files. This should be easy, since all the step results are immutable.
Deprecation warnings not being logged

Problems running the Euler example
I'm a bit confused, because when I run the following:

```
tango run euler_general.jsonnet -i complex_arithmetic -d workspace
```

I would expect my work to be in the folder workspace.  But when I ls this folder it's empty and instead I see my run directly in runs.

```
$ ls workspace/runs/
(empty)
$ ls runs
hot-marlin/   viable-goose/
```

I also don't seem to have step_result_for_run that the CLI output suggested.

```
>>> from tango import Workspace
>>> workspace = Workspace.from_url('.')
>>> workspace.step_result_for_run
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'LocalWorkspace' object has no attribute 'step_result_for_run'
```
Problems running the Euler example
I'm a bit confused, because when I run the following:

```
tango run euler_general.jsonnet -i complex_arithmetic -d workspace
```

I would expect my work to be in the folder workspace.  But when I ls this folder it's empty and instead I see my run directly in runs.

```
$ ls workspace/runs/
(empty)
$ ls runs
hot-marlin/   viable-goose/
```

I also don't seem to have step_result_for_run that the CLI output suggested.

```
>>> from tango import Workspace
>>> workspace = Workspace.from_url('.')
>>> workspace.step_result_for_run
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'LocalWorkspace' object has no attribute 'step_result_for_run'
```


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index a42e86ff5..c40dae852 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -18,6 +18,7 @@ env:
   CACHE_PREFIX: v3  # Change this to invalidate existing cache.
   PYTHON_PATH: ./
   DEFAULT_PYTHON: 3.7
+  WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
   BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
   BEAKER_WORKSPACE: ai2/tango-testing
   BEAKER_DEFAULT_CLUSTER: ai2/tango-gpu-tests
@@ -70,7 +71,7 @@ jobs:
             extras: dev
             requires_torch: false
             run: |
-              pytest -v --color=yes --doctest-modules --ignore=tests/integrations --ignore=tango/integrations tests/ tango/
+              pytest -v --durations=10 --color=yes --doctest-modules --ignore=tests/integrations --ignore=tango/integrations tests/ tango/
 
           - name: Datasets integration
             extras: dev,datasets
@@ -102,6 +103,12 @@ jobs:
             run: |
               pytest -v --color=yes --doctest-modules tango/integrations/fairscale tests/integrations/fairscale
 
+          - name: W&B integration
+            extras: dev,torch,wandb
+            requires_torch: true
+            run: |
+              pytest -v --color=yes --doctest-modules tango/integrations/wandb tests/integrations/wandb
+
           - name: Example - train_lm
             extras: dev,all
             requires_torch: true
@@ -116,7 +123,7 @@ jobs:
               extras: dev
               requires_torch: false
               run: |
-                pytest -v --color=yes --doctest-modules --ignore=tests/integrations --ignore=tango/integrations tests/ tango/
+                pytest -v --durations=10 --color=yes --doctest-modules --ignore=tests/integrations --ignore=tango/integrations tests/ tango/
             python: '3.8'
 
           - task:
@@ -124,7 +131,7 @@ jobs:
               extras: dev
               requires_torch: false
               run: |
-                pytest -v --color=yes --doctest-modules --ignore=tests/integrations --ignore=tango/integrations tests/ tango/
+                pytest -v --durations=10 --color=yes --doctest-modules --ignore=tests/integrations --ignore=tango/integrations tests/ tango/
             python: '3.9'
 
           - task:
@@ -132,7 +139,7 @@ jobs:
               extras: dev
               requires_torch: false
               run: |
-                pytest -v --color=yes --doctest-modules --ignore=tests/integrations --ignore=tango/integrations tests/ tango/
+                pytest -v --durations=10 --color=yes --doctest-modules --ignore=tests/integrations --ignore=tango/integrations tests/ tango/
             python: '3.10'
 
     steps:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c80372676..b06965cac 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,8 +7,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- Added a Weights & Baises remote `Workspace` implementation: `WandbWorkspace`, registered as "wandb".
+  This can be instantiated from a workspace URL in the form "wandb://entity/project".
+- Added a method `Workspace.step_result_for_run` which gives the result of a step given the run name and step name within that run.
+- Added property `Workspace.url`, which returns a URL for the workspace that can be used to instantiate the exact same workspace using `Workspace.from_url()`. Subclasses must implement this.
+
+### Changed
+
+- `StepInfo` start and end times will be always be in UTC now.
+- `WandbTrainCallback` now logs system metrics from each worker process in distributed training.
+- `StepCache.__contains__()` and `StepCache.__getitem__()` now take accept either a `Step` or `StepInfo` as an argument (`Union[Step, StepInfo]`).
+- Refactored `tango.step_graph.StepGraph` to allow initialization from a `Dict[str, Step]`.
+- `Executor.execute_step_graph()` now attempts to execute all steps and summarizes success/failures.
+
 ### Fixed
 
+- Fixed bug with `LocalWorkspace.from_parsed_url()` ([#278](https://github.com/allenai/tango/issues/278)).
+- Deprecation warnings will now be logged from `tango` CLI.
 - Fixed the text format in the case of serializing an iterator of string.
 - Added missing default value of `None` to `TangoGlobalSettings.find_or_default()`.
 - Mypy has become incompatible with transformers and datasets, so we have to disable the checks in some places.
diff --git a/dev-requirements.txt b/dev-requirements.txt
index 035e34255..7adfda49a 100644
--- a/dev-requirements.txt
+++ b/dev-requirements.txt
@@ -9,6 +9,7 @@ flake8
 mypy==0.950
 types-PyYAML
 types-setuptools
+types-pytz
 
 # Automatic code formatting
 black==22.3.0
diff --git a/docs/source/api/components/index.rst b/docs/source/api/components/index.rst
index c059b8b43..06f9f7787 100644
--- a/docs/source/api/components/index.rst
+++ b/docs/source/api/components/index.rst
@@ -8,6 +8,7 @@ The core components of **AI2 Tango**.
    :caption: Components
 
    step
+   step_info
    step_graph
    workspace
    step_cache
diff --git a/docs/source/api/components/step_info.rst b/docs/source/api/components/step_info.rst
new file mode 100644
index 000000000..aa2abcd9a
--- /dev/null
+++ b/docs/source/api/components/step_info.rst
@@ -0,0 +1,8 @@
+StepInfo
+========
+
+.. autoclass:: tango.step_info.StepInfo
+   :member-order: bysource
+   :members:
+
+.. autoclass:: tango.step_info.StepState
diff --git a/docs/source/api/components/workspace.rst b/docs/source/api/components/workspace.rst
index bb99f09e9..4d9d107b6 100644
--- a/docs/source/api/components/workspace.rst
+++ b/docs/source/api/components/workspace.rst
@@ -17,12 +17,6 @@ Implementations
 Metadata
 --------
 
-.. autoclass:: tango.workspace.StepInfo
-   :members:
-
-.. autoclass:: tango.workspace.StepState
-   :members:
-
 .. autoclass:: tango.workspace.Run
    :members:
 
diff --git a/docs/source/api/integrations/torch.rst b/docs/source/api/integrations/torch.rst
index 629c828a5..b54ea62ff 100644
--- a/docs/source/api/integrations/torch.rst
+++ b/docs/source/api/integrations/torch.rst
@@ -70,9 +70,11 @@ Callbacks
 
 .. autoclass:: tango.integrations.torch.TrainCallback
    :members:
+   :member-order: bysource
 
 .. autoclass:: tango.integrations.torch.EvalCallback
    :members:
+   :member-order: bysource
 
 .. autoclass:: tango.integrations.torch.StopEarlyCallback
 
diff --git a/docs/source/api/integrations/wandb.rst b/docs/source/api/integrations/wandb.rst
index 6d6091b76..680fe3762 100644
--- a/docs/source/api/integrations/wandb.rst
+++ b/docs/source/api/integrations/wandb.rst
@@ -2,10 +2,12 @@ Weights & Biases
 ================
 
 .. automodule:: tango.integrations.wandb
-   :members:
-   :exclude-members: WandbTrainCallback
 
 Reference
 ---------
 
+.. autoclass:: tango.integrations.wandb.WandbWorkspace
+
+.. autoclass:: tango.integrations.wandb.WandbStepCache
+
 .. autoclass:: tango.integrations.wandb.WandbTrainCallback
diff --git a/docs/source/conf.py b/docs/source/conf.py
index eb82d5b65..5683419c4 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -26,7 +26,6 @@
 version = VERSION_SHORT
 release = VERSION
 
-
 # -- General configuration ---------------------------------------------------
 
 # Add any Sphinx extension module names here, as strings. They can be
@@ -43,9 +42,6 @@
     "sphinx_autodoc_typehints",
 ]
 
-# Tell myst-parser to assign header anchors for h1-h3.
-myst_heading_anchors = 3
-
 suppress_warnings = ["myst.header"]
 
 # Add any paths that contain templates here, relative to this directory.
@@ -58,6 +54,8 @@
 
 source_suffix = [".rst", ".md"]
 
+# -- Extension configuration -------------------------------------------------
+
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3", None),
     "rich": ("https://rich.readthedocs.io/en/latest", None),
@@ -70,6 +68,16 @@
     "transformers": ("https://huggingface.co/docs/transformers/master/en", None),
 }
 
+# Tell myst-parser to assign header anchors for h1-h3.
+myst_heading_anchors = 3
+
+# By default, sort documented members by type within classes and modules.
+autodoc_member_order = "groupwise"
+
+python_use_unqualified_type_names = True
+
+# Include default values when documenting parameter types.
+typehints_defaults = "comma"
 
 # -- Options for HTML output -------------------------------------------------
 
@@ -110,5 +118,3 @@
         },
     ],
 }
-
-python_use_unqualified_type_names = True
diff --git a/docs/source/examples/euler.md b/docs/source/examples/euler.md
index 6215b09a1..e874f08ed 100644
--- a/docs/source/examples/euler.md
+++ b/docs/source/examples/euler.md
@@ -3,14 +3,16 @@
 
 ## Running the experiment
 
-Try running the experiment with:
+If you haven't already, clone the [tango repository](https://github.com/allenai/tango) and then
+change directories into `examples/euler`.
+
+You can then run the experiment with:
 
 ```bash
 tango run euler_general.jsonnet -i complex_arithmetic -d workspace
 ```
 
-This will leave its results in a directory called `workspace/`.
-
+This will leave its results in a subdirectory of `workspace/runs/` corresponding to the name of the run.
 The output it prints should look something like this:
 ```
 Starting new run fluent-jay
@@ -49,6 +51,8 @@ Server started at http://localhost:8080/run/fluent-jay
 Finished run fluent-jay
 ```
 
+In this case the assigned run name is "fluent-jay," so the output will be in `workspace/runs/fluent-jay`.
+
 ## Try the workspace UI
 
 You can visualize the workspace by running
diff --git a/requirements.txt b/requirements.txt
index 00e1bc19a..ba0bf341d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -16,6 +16,7 @@ more-itertools>=8.0,<9.0
 sqlitedict
 glob2>=0.7
 petname>=2.6,<3.0
+pytz
 
 ##################################################
 ###### Extra dependencies for integrations #######
diff --git a/tango/__main__.py b/tango/__main__.py
index 98c00a3f8..d4b4b48f1 100644
--- a/tango/__main__.py
+++ b/tango/__main__.py
@@ -167,6 +167,9 @@ def main(
     start_method: Optional[str] = None,
     called_by_executor: bool = False,
 ):
+    if not sys.warnoptions:
+        warnings.simplefilter("default", category=DeprecationWarning)
+
     settings: TangoGlobalSettings = (
         TangoGlobalSettings.from_file(settings)
         if settings is not None
@@ -825,7 +828,7 @@ def _display_run_results(
             ),
             "python",
         )
-        cli_logger.info("Use your workspace to get the cached result of a step, .e.g.")
+        cli_logger.info("Use your workspace to get the cached result of a step, e.g.")
         cli_logger.info(example)
 
 
diff --git a/tango/common/util.py b/tango/common/util.py
index 7d5c8a7b3..f7196c255 100644
--- a/tango/common/util.py
+++ b/tango/common/util.py
@@ -5,13 +5,26 @@
 import sys
 import traceback
 from contextlib import contextmanager
+from datetime import datetime, tzinfo
 from pathlib import Path
 from typing import Iterable, Optional, Set, Tuple, Union
 
+import pytz
+
 from .aliases import PathOrStr
 from .exceptions import SigTermReceived
 
 
+def tango_cache_dir() -> Path:
+    """
+    Returns a directory suitable for caching things from Tango, defaulting
+    to ``$HOME/.cache/tango``.
+    """
+    cache_dir = Path.home() / ".cache" / "tango"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
 def _handle_sigterm(sig, frame):
     raise SigTermReceived
 
@@ -243,3 +256,11 @@ def exception_to_string(e: BaseException) -> str:
     else:
         formatted = traceback.format_exception(etype=type(e), value=e, tb=e.__traceback__)
     return "".join(formatted)
+
+
+def utc_now_datetime() -> datetime:
+    return datetime.utcnow().replace(tzinfo=pytz.utc)
+
+
+def local_timezone() -> Optional[tzinfo]:
+    return datetime.now().astimezone().tzinfo
diff --git a/tango/executors/multicore_executor.py b/tango/executors/multicore_executor.py
index 53a36c11a..0e0453146 100644
--- a/tango/executors/multicore_executor.py
+++ b/tango/executors/multicore_executor.py
@@ -8,7 +8,8 @@
 from tango.executor import Executor, ExecutorOutput
 from tango.step import Step
 from tango.step_graph import StepGraph
-from tango.workspace import StepState, Workspace
+from tango.step_info import StepState
+from tango.workspace import Workspace
 
 logger = logging.getLogger(__name__)
 
@@ -129,6 +130,10 @@ def _update_running_steps(step_states: Dict[str, StepState]) -> List[str]:
             for step_name, process in _running.items():
                 poll_status = process.poll()
                 if poll_status is not None:
+                    # The step may have finished since we synced step states.
+                    if step_states[step_name] == StepState.RUNNING:
+                        step_states[step_name] = self._get_state(step_graph[step_name])
+
                     # We check for uncacheable leaf step too.
                     if step_states[step_name] in [StepState.COMPLETED, StepState.UNCACHEABLE]:
                         done.append(step_name)
@@ -141,7 +146,7 @@ def _update_running_steps(step_states: Dict[str, StepState]) -> List[str]:
                         # it thinks that the process is not running.
                         errors.append(step_name)
                     else:
-                        logger.warning(
+                        raise RuntimeError(
                             f"Step '{step_name}' has the state {step_states[step_name]}, "
                             "but the corresponding process has ended!"
                         )
@@ -220,10 +225,10 @@ def _try_to_execute_next_step(config_path: str, run_name: Optional[str] = None)
                     config_path,
                     "-s",
                     step_name,
+                    "-w",
+                    self.workspace.url,
                     "--no-server",
                 ]
-                if hasattr(self.workspace, "dir"):
-                    command += ["-w", self.workspace.dir]  # type: ignore
                 if self.include_package is not None:
                     for package in self.include_package:
                         command += ["-i", package]
diff --git a/tango/integrations/torch/eval_callback.py b/tango/integrations/torch/eval_callback.py
index 7e318bf33..6aef730cf 100644
--- a/tango/integrations/torch/eval_callback.py
+++ b/tango/integrations/torch/eval_callback.py
@@ -3,6 +3,7 @@
 
 from tango.common.dataset_dict import DatasetDictBase
 from tango.common.registrable import Registrable
+from tango.workspace import Workspace
 
 from .data import DataLoader
 from .model import Model
@@ -19,24 +20,24 @@ class EvalCallback(Registrable):
         All of the parameters to this base class will be automatically set within
         the training loop, so you shouldn't include them in your config for your callbacks.
 
-    Attributes
-    ----------
-    step_id : :class:`str`
-    work_dir : :class:`~pathlib.Path`
-    model : :class:`Model`
-    dataset_dict : :class:`tango.common.DatasetDictBase`
-    dataloader : :class:`DataLoader`
-
+    :ivar Workspace workspace: The tango workspace being used.
+    :ivar str step_id: The unique ID of the step.
+    :ivar pathlib.Path work_dir: The working directory of the step
+    :ivar Model model: The model being evaluated.
+    :ivar DatasetDictBase dataset_dict: The dataset dict containing the evaluation split.
+    :ivar DataLoader dataloader: The data loader used to load the evaluation split data.
     """
 
     def __init__(
         self,
+        workspace: Workspace,
         step_id: str,
         work_dir: Path,
         model: Model,
         dataset_dict: DatasetDictBase,
         dataloader: DataLoader,
     ) -> None:
+        self.workspace = workspace
         self.step_id = step_id
         self.work_dir = work_dir
         self.model = model
diff --git a/tango/integrations/torch/train.py b/tango/integrations/torch/train.py
index 8dcc99480..faede8561 100644
--- a/tango/integrations/torch/train.py
+++ b/tango/integrations/torch/train.py
@@ -17,6 +17,7 @@
 from tango.common.util import get_extra_imported_modules, import_extra_module
 from tango.format import Format
 from tango.step import Step
+from tango.workspace import Workspace
 
 from .data import DataLoader
 from .exceptions import StopEarly
@@ -266,6 +267,7 @@ def _train(
         config = TrainConfig(
             self.unique_id,
             self.work_dir,
+            step_name=self.name,
             train_split=train_split,
             validation_split=validation_split,
             seed=seed,
@@ -293,6 +295,7 @@ def _train(
             mp.spawn(
                 _train,
                 args=(
+                    self.workspace,
                     config,
                     model,
                     training_engine,
@@ -309,6 +312,7 @@ def _train(
         else:
             final_model = _train(  # type: ignore[assignment]
                 0,
+                self.workspace,
                 config,
                 model,
                 training_engine,
@@ -334,6 +338,7 @@ def _train(
 
 def _train(
     worker_id: int,
+    workspace: Workspace,
     config: TrainConfig,
     model: Lazy[Model],
     training_engine: Lazy[TrainingEngine],
@@ -426,6 +431,7 @@ def _train(
     # Initialize callbacks.
     callbacks: List[TrainCallback] = [
         callback.construct(
+            workspace=workspace,
             train_config=config,
             training_engine=training_engine,
             dataset_dict=dataset_dict,
diff --git a/tango/integrations/torch/train_callback.py b/tango/integrations/torch/train_callback.py
index 5153f79cf..8c92a6e98 100644
--- a/tango/integrations/torch/train_callback.py
+++ b/tango/integrations/torch/train_callback.py
@@ -4,6 +4,7 @@
 
 from tango.common.dataset_dict import DatasetDictBase
 from tango.common.registrable import Registrable
+from tango.workspace import Workspace
 
 from .data import DataLoader
 from .exceptions import StopEarly
@@ -14,9 +15,9 @@
 
 class TrainCallback(Registrable):
     """
-    A ``TrainCallback`` is a :class:`~tango.common.Registrable` class
+    A :class:`TrainCallback` is a :class:`~tango.common.Registrable` class
     that can be used within :class:`TorchTrainStep` to customize behavior in the training
-    loop.
+    loop. You can set the training callbacks with the ``callbacks`` parameter to :class:`TorchTrainStep`.
 
     .. tip::
         All of the parameters to this base class will be automatically set within
@@ -33,26 +34,25 @@ class TrainCallback(Registrable):
         See :class:`~tango.integrations.wandb.WandbTrainCallback` for an example
         implementation.
 
-    Attributes
-    ----------
-    train_config : :class:`TrainConfig`
-    training_engine : :class:`TrainingEngine`
-    optimizer : :class:`Optimizer`
-    dataset_dict : :class:`tango.common.DatasetDictBase`
-    train_dataloader : :class:`DataLoader`
-    validation_dataloader : :class:`DataLoader`, optional
-    lr_scheduler : :class:`LRScheduler`, optional
-
+    :ivar Workspace workspace: The tango workspace being used.
+    :ivar TrainConfig train_config: The training config.
+    :ivar TrainingEngine training_engine: The engine used to train the model.
+    :ivar tango.common.DatasetDictBase dataset_dict: The dataset dict containing train and
+        optional validation splits.
+    :ivar DataLoader train_dataloader: The dataloader used for the training split.
+    :ivar DataLoader validation_dataloader: Optional dataloader used for the validation split.
     """
 
     def __init__(
         self,
+        workspace: Workspace,
         train_config: TrainConfig,
         training_engine: TrainingEngine,
         dataset_dict: DatasetDictBase,
         train_dataloader: DataLoader,
         validation_dataloader: Optional[DataLoader] = None,
     ) -> None:
+        self.workspace = workspace
         self.train_config = train_config
         self.training_engine = training_engine
         self.dataset_dict = dataset_dict
@@ -67,6 +67,13 @@ def step_id(self) -> str:
         """
         return self.train_config.step_id
 
+    @property
+    def step_name(self) -> Optional[str]:
+        """
+        The name of the current :class:`~tango.Step`.
+        """
+        return self.train_config.step_name
+
     @property
     def work_dir(self) -> Path:
         """
diff --git a/tango/integrations/torch/train_config.py b/tango/integrations/torch/train_config.py
index bc5053f02..8aa2ecb01 100644
--- a/tango/integrations/torch/train_config.py
+++ b/tango/integrations/torch/train_config.py
@@ -22,6 +22,14 @@ class TrainConfig:
     The working directory for the training run.
     """
 
+    step_name: Optional[str] = None
+    """
+    The name of the current step.
+
+    .. note::
+        The same step can be run under different names.
+    """
+
     worker_id: int = 0
     """
     The ID of the distributed worker.
diff --git a/tango/integrations/torch/training_engine.py b/tango/integrations/torch/training_engine.py
index 66b7e8255..4dc447a9f 100644
--- a/tango/integrations/torch/training_engine.py
+++ b/tango/integrations/torch/training_engine.py
@@ -21,12 +21,10 @@ class TrainingEngine(Registrable):
     A :class:`TrainingEngine` defines and drives the strategy for training a model
     in :class:`TorchTrainStep`.
 
-    Attributes
-    ----------
-    train_config : :class:`TrainConfig`
-    model : :class:`Model`
-    optimizer : :class:`Optimizer`
-    lr_scheduler : :class:`LRScheduler`, optional
+    :ivar TrainConfig train_config: The training config.
+    :ivar Model model: The model being trained.
+    :ivar Optimizer optimizer: The optimizer being used to train the model.
+    :ivar LRScheduler lr_scheduler: The optional learning rate scheduler.
     """
 
     default_implementation = "torch"
diff --git a/tango/integrations/wandb/__init__.py b/tango/integrations/wandb/__init__.py
index d2be78019..bccb000f4 100644
--- a/tango/integrations/wandb/__init__.py
+++ b/tango/integrations/wandb/__init__.py
@@ -6,9 +6,25 @@
 
 Components for Tango integration with `Weights & Biases <https://wandb.ai/>`_.
 
+Overview
+--------
+
+The main components provided by this integration are the :class:`WandbWorkspace` and
+the :class:`WandbTrainCallback`.
+
+The :class:`WandbWorkspace` is a :class:`~tango.workspace.Workspace` implementation that is
+great for collaboration. It tracks Tango runs and steps in the W&B project of your choosing
+and uses W&B Artifacts to cache step results in the cloud so that they're accessible anywhere.
+
+And if you're training PyTorch models via the :class:`~tango.integrations.torch.TorchTrainStep`,
+you can use the :class:`WandbTrainCallback` to track metrics throughout the run.
+
 """
 
-__all__ = []
+__all__ = ["WandbWorkspace", "WandbStepCache"]
+
+from .step_cache import WandbStepCache
+from .workspace import WandbWorkspace
 
 try:
     from .torch_train_callback import WandbTrainCallback
diff --git a/tango/integrations/wandb/step_cache.py b/tango/integrations/wandb/step_cache.py
new file mode 100644
index 000000000..246282e2b
--- /dev/null
+++ b/tango/integrations/wandb/step_cache.py
@@ -0,0 +1,223 @@
+import logging
+import os
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import wandb
+from wandb.errors import Error as WandbError
+
+from tango.common.aliases import PathOrStr
+from tango.common.file_lock import FileLock
+from tango.common.params import Params
+from tango.common.util import tango_cache_dir
+from tango.step import Step
+from tango.step_cache import CacheMetadata, StepCache
+from tango.step_caches.local_step_cache import LocalStepCache
+from tango.step_info import StepInfo
+
+from .util import ArtifactKind, check_environment, is_missing_artifact_error
+
+logger = logging.getLogger(__name__)
+
+
+# This class inherits from `LocalStepCache` to benefit from its in-memory "weak cache" and "strong cache",
+# but it handles saving artifacts to disk a little differently.
+@StepCache.register("wandb")
+class WandbStepCache(LocalStepCache):
+    """
+    This is a :class:`~tango.step_cache.StepCache` that's used by :class:`WandbWorkspace`.
+    It stores the results of steps on W&B as Artifacts.
+
+    It also keeps a limited in-memory cache as well as a local backup on disk, so fetching a
+    step's resulting subsequent times should be fast.
+
+    :param project: The W&B project to use.
+    :param entity: The W&B entity (user or organization account) to use.
+
+    .. tip::
+        Registered as :class:`~tango.step_cache.StepCache` under the name "wandb".
+    """
+
+    def __init__(self, project: str, entity: str):
+        check_environment()
+        super().__init__(tango_cache_dir() / "wandb_cache")
+        self.project = project
+        self.entity = entity
+
+    @property
+    def wandb_client(self) -> wandb.Api:
+        return wandb.Api(overrides={"entity": self.entity, "project": self.project})
+
+    def _acquire_step_lock_file(self, step: Union[Step, StepInfo], read_only_ok: bool = False):
+        return FileLock(
+            self.step_dir(step).with_suffix(".lock"), read_only_ok=read_only_ok
+        ).acquire_with_updates(desc=f"acquiring step cache lock for '{step.unique_id}'")
+
+    def _step_artifact_name(self, step: Union[Step, StepInfo]) -> str:
+        if isinstance(step, Step):
+            return step.__class__.__name__
+        else:
+            return step.step_class_name
+
+    def get_step_result_artifact(
+        self, step: Union[Step, StepInfo]
+    ) -> Optional[wandb.apis.public.Artifact]:
+        try:
+            return self.wandb_client.artifact(
+                f"{self.entity}/{self.project}/{self._step_artifact_name(step)}:{step.unique_id}",
+                type=ArtifactKind.STEP_RESULT.value,
+            )
+        except WandbError as exc:
+            if is_missing_artifact_error(exc):
+                return None
+            else:
+                raise
+
+    def create_step_result_artifact(
+        self, step: Step, objects_dir: Optional[PathOrStr] = None
+    ) -> None:
+        """
+        Create an artifact for the result of a step.
+        """
+        artifact = wandb.Artifact(
+            self._step_artifact_name(step), type=ArtifactKind.STEP_RESULT.value
+        )
+
+        # Add files
+        if objects_dir is not None:
+            artifact.add_dir(str(objects_dir))
+
+        # Log/persist the artifact to W&B.
+        artifact.save()
+        artifact.wait()
+
+        # Add an alias for the step's unique ID.
+        # Only after we've logged the artifact can we add an alias.
+        artifact.aliases.append(step.unique_id)
+        artifact.save()
+        artifact.wait()
+
+    def use_step_result_artifact(self, step: Union[Step, StepInfo]) -> None:
+        """
+        "Use" the artifact corresponding to the result of a step.
+        """
+        if wandb.run is None:
+            raise RuntimeError("This can only be called from within a W&B run")
+        wandb.run.use_artifact(
+            f"{self.entity}/{self.project}/{self._step_artifact_name(step)}:{step.unique_id}"
+        )
+
+    def __contains__(self, step: Any) -> bool:
+        if isinstance(step, (Step, StepInfo)):
+            cacheable = step.cache_results if isinstance(step, Step) else step.cacheable
+            if not cacheable:
+                return False
+
+            key = step.unique_id
+
+            # First check if we have a copy in memory.
+            if key in self.strong_cache:
+                return True
+            if key in self.weak_cache:
+                return True
+
+            # Then check if we have a copy on disk in our cache directory.
+            with self._acquire_step_lock_file(step, read_only_ok=True):
+                if self.step_dir(step).is_dir():
+                    return True
+
+            # If not, check W&B for the corresponding artifact.
+            return self.get_step_result_artifact(step) is not None
+        else:
+            return False
+
+    def __getitem__(self, step: Union[Step, StepInfo]) -> Any:
+        key = step.unique_id
+
+        if wandb.run is not None:
+            # Mark that the current run uses the step's result artifact.
+            try:
+                self.use_step_result_artifact(step)
+            except WandbError as exc:
+                if is_missing_artifact_error(exc):
+                    raise KeyError(step)
+                else:
+                    raise
+
+        # Try getting the result from our in-memory caches first.
+        result = self._get_from_cache(key)
+        if result is not None:
+            return result
+
+        def load_and_return():
+            metadata = CacheMetadata.from_params(Params.from_file(self._metadata_path(step)))
+            result = metadata.format.read(self.step_dir(step))
+            self._add_to_cache(key, result)
+            return result
+
+        # Next check our local on-disk cache.
+        with self._acquire_step_lock_file(step, read_only_ok=True):
+            if self.step_dir(step).is_dir():
+                return load_and_return()
+
+        # Finally, check W&B for the corresponding artifact.
+        with self._acquire_step_lock_file(step):
+            # Make sure the step wasn't cached since the last time we checked (above).
+            if self.step_dir(step).is_dir():
+                return load_and_return()
+
+            artifact = self.get_step_result_artifact(step)
+            if artifact is None:
+                raise KeyError(step)
+
+            # We'll download the artifact to a temporary directory first, in case something goes wrong.
+            temp_dir = tempfile.mkdtemp(dir=self.dir, prefix=key)
+            try:
+                artifact.download(root=temp_dir, recursive=True)
+                # Download and extraction was successful, rename temp directory to final step result directory.
+                os.replace(temp_dir, self.step_dir(step))
+            finally:
+                shutil.rmtree(temp_dir, ignore_errors=True)
+
+            return load_and_return()
+
+    def __setitem__(self, step: Step, value: Any) -> None:
+        if not step.cache_results:
+            logger.warning("Tried to cache step %s despite being marked as uncacheable.", step.name)
+            return
+
+        if wandb.run is None:
+            raise RuntimeError("Can only add results to the WandbStepCache within a W&B run")
+
+        with self._acquire_step_lock_file(step):
+            # We'll write the step's results to temporary directory first, and try to upload to W&B
+            # from there in case anything goes wrong.
+            temp_dir = Path(tempfile.mkdtemp(dir=self.dir, prefix=step.unique_id))
+            try:
+                step.format.write(value, temp_dir)
+                metadata = CacheMetadata(step=step.unique_id, format=step.format)
+                metadata.to_params().to_file(temp_dir / self.METADATA_FILE_NAME)
+                # Create the artifact and upload serialized result to it.
+                self.create_step_result_artifact(step, temp_dir)
+                # Upload successful, rename temp directory to the final step result directory.
+                if self.step_dir(step).is_dir():
+                    shutil.rmtree(self.step_dir(step), ignore_errors=True)
+                os.replace(temp_dir, self.step_dir(step))
+            finally:
+                shutil.rmtree(temp_dir, ignore_errors=True)
+
+        # Finally, add to in-memory caches.
+        self._add_to_cache(step.unique_id, value)
+
+    def __len__(self) -> int:
+        completed_cacheable_step_runs = self.wandb_client.runs(
+            f"{self.entity}/{self.project}",
+            filters={
+                "config.job_type": "step",
+                "config.cacheable": True,
+                "state": "finished",
+            },
+        )
+        return len(list(completed_cacheable_step_runs))
diff --git a/tango/integrations/wandb/torch_train_callback.py b/tango/integrations/wandb/torch_train_callback.py
index 730130f95..926c06f5a 100644
--- a/tango/integrations/wandb/torch_train_callback.py
+++ b/tango/integrations/wandb/torch_train_callback.py
@@ -1,12 +1,15 @@
-import os
-import warnings
 from typing import Any, Dict, List, Optional
 
 import torch
+import wandb
 
+from tango.common.exceptions import ConfigurationError
 from tango.integrations.torch.train_callback import TrainCallback
 from tango.integrations.torch.util import peak_gpu_memory
 
+from .util import check_environment
+from .workspace import WandbWorkspace
+
 
 @TrainCallback.register("wandb::log")
 class WandbTrainCallback(TrainCallback):
@@ -15,36 +18,47 @@ class WandbTrainCallback(TrainCallback):
     the :class:`~tango.integrations.torch.TorchTrainStep` that logs training and
     validation metrics to W&B.
 
+    This can be used with any :class:`~tango.workspace.Workspace` implementation,
+    including :class:`WandbWorkspace`.
+
     .. tip::
 
         Registered as a :class:`~tango.integrations.torch.TrainCallback`
         under the name "wandb::log".
 
     .. important::
-        This class also requires the :mod:`~tango.integrations.torch` integration.
+
+        When this callback is used with the :class:`WandbWorkspace` it will log metrics
+        to the same W&B project that the workspace uses. The ``group`` and ``name``
+        parameters will also automatically be set, so a :class:`~tango.common.exceptions.ConfigurationError`
+        will be raised if any of ``project``, ``entity``, ``group``, or ``name`` are set in this callback.
 
     :param project:
         W&B project to associated this run with.
+
     :param entity:
         W&B entity (user or organization) to associated this run with.
+
     :param group:
         W&B group to associated this run with.
+
     :param name:
-        Set the name of the run in W&B.
+        Set the name of the run in W&B. If not set, the default will be the name of the step.
+
     :param notes:
         Arbitrary notes to add in W&B to this run.
+
     :param tags:
         Arbitrary tags to add in W&B to this run.
+
     :param watch_model:
         If ``True``, ``wandb.watch()`` is called to collect gradients and other information
         about the model throughout training.
         See `docs.wandb.ai/ref/python/watch <https://docs.wandb.ai/ref/python/watch>`_.
+
     :param wandb_config:
         Arbitrary configuration fields to set in W&B for this run.
         See `docs.wandb.ai/guides/track/config <https://docs.wandb.ai/guides/track/config>`_.
-    :param wandb_kwargs:
-        Additional keyword arguments that will be passed to ``wandb.init()``.
-        See `docs.wandb.ai/ref/python/init <https://docs.wandb.ai/ref/python/init>`_.
 
     """
 
@@ -59,81 +73,107 @@ def __init__(
         tags: Optional[List[str]] = None,
         watch_model: bool = False,
         wandb_config: Optional[Dict[str, Any]] = None,
-        wandb_kwargs: Optional[Dict[str, Any]] = None,
         **kwargs,
     ) -> None:
         super().__init__(*args, **kwargs)
 
-        if self.is_local_main_process and "WANDB_API_KEY" not in os.environ:
-            warnings.warn(
-                "Missing environment variable 'WANDB_API_KEY' required to authenticate to Weights & Biases.",
-                UserWarning,
-            )
+        if self.is_local_main_process:
+            check_environment()
+
+        if isinstance(self.workspace, WandbWorkspace) or wandb.run is not None:
+            err_msg_template = "Cannot set '{var_name}' in WandbTrainCallback "
+            if isinstance(self.workspace, WandbWorkspace):
+                err_msg_template += "since it has already been set from the WandbWorkspace."
+            else:
+                err_msg_template += "since a W&B run has already been initialized."
+            for var, var_name in [
+                (project, "project"),
+                (entity, "entity"),
+                (group, "group"),
+                (name, "name"),
+            ]:
+                if var is not None:
+                    raise ConfigurationError(err_msg_template.format(var_name=var_name))
+
+        self.project = (
+            project if not isinstance(self.workspace, WandbWorkspace) else self.workspace.project
+        )
+        self.entity = (
+            entity if not isinstance(self.workspace, WandbWorkspace) else self.workspace.entity
+        )
+        self.group = group or self.step_id
+        self.notes = notes or self._get_default_notes()
+        self.tags = tags
+        self.watch_model = watch_model
 
-        _wandb_config = self.train_config.as_dict()
-        del _wandb_config["worker_id"]
-        _wandb_config["work_dir"] = str(self.train_config.work_dir.resolve())
+        self.wandb_config = self.train_config.as_dict()
+        del self.wandb_config["worker_id"]
         if wandb_config is not None:
-            _wandb_config.update(wandb_config)
-
-        self._watch_model = watch_model
-        self._run_id: Optional[str] = None
-        self._wandb_kwargs: Dict[str, Any] = dict(
-            dir=str(self.work_dir),
-            project=project,
-            entity=entity,
-            group=group,
-            name=name or self._get_run_name_from_step(),
-            notes=notes,
-            config=_wandb_config,
-            tags=tags,
-            anonymous="allow",
-            **(wandb_kwargs or {}),
+            self.wandb_config.update(wandb_config)
+        if wandb.run is None:
+            self.wandb_config["job_type"] = "train_metrics"
+
+        self.run_name: str = name or self.step_name or "train"
+        if self.train_config.is_distributed:
+            self.run_name += f" (rank {self.train_config.worker_id})"
+
+        self.run_id: str = (
+            wandb.run.id  # type: ignore[attr-defined]
+            if wandb.run is not None
+            else self.step_id + f"-rank{self.train_config.worker_id}"
         )
-
-    def _get_run_name_from_step(self) -> str:
-        """
-        Generate a run name for W&B based on the step unique ID.
-        """
-        step_name, step_hash = self.step_id.split("-", maxsplit=1)
-        return f"{step_name}-{step_hash[:6]}"
+        self.resume: Optional[str] = None
+        self.should_finalize_run: bool = (
+            wandb.run is None
+        )  # if we have to start out own W&B run, we need to finish it
 
     def state_dict(self) -> Dict[str, Any]:
-        if self.is_local_main_process:
-            return {
-                "run_id": self._run_id,
-            }
-        else:
-            return {}
+        return {}
 
     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
-        if self.is_local_main_process:
-            self._wandb_kwargs["resume"] = "allow"
-            self._run_id = state_dict["run_id"]
+        self.resume = "allow"
 
     def pre_train_loop(self) -> None:
+        if wandb.run is None:
+            if self.run_id is None:
+                self.run_id = self.step_id + f"-rank{self.train_config.worker_id}"
+            # Initialize a new W&B run.
+            wandb.init(
+                id=self.run_id,
+                dir=str(self.work_dir),
+                project=self.project,
+                entity=self.entity,
+                group=self.group,
+                name=self.run_name,
+                notes=self.notes,
+                config=self.wandb_config,
+                tags=self.tags,
+                job_type="train_metrics",
+            )
+        else:
+            # We are already running inside of a W&B run, possibly because
+            # we're using the WandbWorkspace.
+            wandb.config.update(self.wandb_config)
+            if self.tags:
+                wandb.run.tags = (wandb.run.tags or tuple()) + tuple(self.tags)
+            if self.notes:
+                wandb.run.notes = self.notes
+
+        if self.watch_model:
+            wandb.watch(self.training_engine.model)
+
+        # Log GPU memory statistics.
         if torch.cuda.is_available():
             torch.cuda.reset_peak_memory_stats()
         peak_gpu_mbs = peak_gpu_memory()
         if self.is_local_main_process:
-            import wandb
-
-            self.wandb = wandb
-
-            if self._run_id is None:
-                self._run_id = self.wandb.util.generate_id()
-
-            self.wandb.init(id=self._run_id, **self._wandb_kwargs)
-
-            if self._watch_model:
-                self.wandb.watch(self.training_engine.model)
             metrics = {f"sys/worker{rank}_peak_gpu_mem": mbs for rank, mbs in peak_gpu_mbs.items()}
             metrics["epoch"] = 0
-            self.wandb.log(metrics, step=0)
+            wandb.log(metrics, step=0)
 
     def post_train_loop(self, step: int, epoch: int) -> None:
-        if self.is_local_main_process:
-            self.wandb.finish()
+        if self.should_finalize_run:
+            wandb.finish()
 
     def log_batch(self, step: int, epoch: int, batch_loss: float) -> None:
         peak_gpu_mbs = peak_gpu_memory()
@@ -146,7 +186,7 @@ def log_batch(self, step: int, epoch: int, batch_loss: float) -> None:
             metrics.update(
                 {f"sys/worker{rank}_peak_gpu_mem": mbs for rank, mbs in peak_gpu_mbs.items()}
             )
-            self.wandb.log(
+            wandb.log(
                 metrics,
                 step=step + 1,
             )
@@ -155,7 +195,7 @@ def post_val_loop(
         self, step: int, epoch: int, val_metric: float, best_val_metric: float
     ) -> None:
         if self.is_local_main_process:
-            self.wandb.log(
+            wandb.log(
                 {
                     f"val/{self.train_config.val_metric_name}": val_metric,
                     f"val/best_{self.train_config.val_metric_name}": best_val_metric,
@@ -163,3 +203,11 @@ def post_val_loop(
                 },
                 step=step + 1,
             )
+
+    def _get_default_notes(self) -> str:
+        notes = (
+            f'Metrics for Tango step "{self.step_name}" from worker {self.train_config.worker_id}.'
+        )
+        if isinstance(self.workspace, WandbWorkspace):
+            notes += f"\nMain run for step: {self.workspace.wandb_project_url}/runs/{self.step_id}/overview"
+        return notes
diff --git a/tango/integrations/wandb/util.py b/tango/integrations/wandb/util.py
new file mode 100644
index 000000000..29d5ae644
--- /dev/null
+++ b/tango/integrations/wandb/util.py
@@ -0,0 +1,42 @@
+import os
+import warnings
+from enum import Enum
+
+from wandb.errors import Error as WandbError
+
+_API_KEY_WARNING_ISSUED = False
+_SILENCE_WARNING_ISSUED = False
+
+
+def is_missing_artifact_error(err: WandbError):
+    """
+    Check if a specific W&B error is caused by a 404 on the artifact we're looking for.
+    """
+    # This is brittle, but at least we have a test for it.
+    return "does not contain artifact" in err.message
+
+
+def check_environment():
+    global _API_KEY_WARNING_ISSUED, _SILENCE_WARNING_ISSUED
+    if "WANDB_API_KEY" not in os.environ and not _API_KEY_WARNING_ISSUED:
+        warnings.warn(
+            "Missing environment variable 'WANDB_API_KEY' required to authenticate to Weights & Biases.",
+            UserWarning,
+        )
+        _API_KEY_WARNING_ISSUED = True
+    if "WANDB_SILENT" not in os.environ and not _SILENCE_WARNING_ISSUED:
+        warnings.warn(
+            "The Weights & Biases client may produce a lot of log messages. "
+            "You can silence these by setting the environment variable 'WANDB_SILENT=true'",
+            UserWarning,
+        )
+        _SILENCE_WARNING_ISSUED = True
+
+
+class RunKind(Enum):
+    STEP = "step"
+    TANGO_RUN = "tango_run"
+
+
+class ArtifactKind(Enum):
+    STEP_RESULT = "step_result"
diff --git a/tango/integrations/wandb/workspace.py b/tango/integrations/wandb/workspace.py
new file mode 100644
index 000000000..4228e61d0
--- /dev/null
+++ b/tango/integrations/wandb/workspace.py
@@ -0,0 +1,426 @@
+import logging
+import tempfile
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Iterable, Iterator, Optional, TypeVar, Union
+from urllib.parse import ParseResult
+
+import pytz
+import wandb
+
+from tango.common.file_lock import FileLock
+from tango.common.util import exception_to_string, tango_cache_dir, utc_now_datetime
+from tango.step import Step
+from tango.step_cache import StepCache
+from tango.step_info import StepInfo, StepState
+from tango.workspace import Run, Workspace
+
+from .step_cache import WandbStepCache
+from .util import ArtifactKind, RunKind, check_environment
+
+T = TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+
+@Workspace.register("wandb")
+class WandbWorkspace(Workspace):
+    """
+    This is a :class:`~tango.workspace.Workspace` that tracks Tango runs in a W&B project.
+    It also stores step results as W&B Artifacts via :class:`WandbStepCache`.
+
+    Each Tango run with this workspace will generate multiple runs in your W&B project.
+    There will always be a W&B run corresponding to each Tango run with the same name,
+    which will contain some metadata about the Tango run. Then there will be one W&B run
+    for each cacheable step that runs with a name corresponding to the name of the step.
+    So if your Tango run includes 3 cacheable steps, that will result in a total of 4 new runs in W&B.
+
+    :param project: The W&B project to use for the workspace.
+    :param entity: The W&B entity (user or organization account) to use for the workspace.
+
+    .. tip::
+        Registered as a :class:`~tango.workspace.Workspace` under the name "wandb".
+    """
+
+    def __init__(self, project: str, entity: Optional[str] = None):
+        check_environment()
+        super().__init__()
+        self.project = project
+        self._entity = entity
+        self.cache = WandbStepCache(project=self.project, entity=self.entity)
+        self.steps_dir = tango_cache_dir() / "wandb_workspace"
+        self.locks: Dict[Step, FileLock] = {}
+        self._running_step_info: Dict[str, StepInfo] = {}
+
+    def __getstate__(self):
+        """
+        We override `__getstate__()` to customize how instances of this class are pickled
+        since we don't want to persist certain attributes.
+        """
+        out = super().__getstate__()
+        out["locks"] = {}
+        return out
+
+    @property
+    def wandb_client(self) -> wandb.Api:
+        overrides = {"project": self.project}
+        if self._entity is not None:
+            overrides["entity"] = self._entity
+        return wandb.Api(overrides=overrides)
+
+    @property
+    def entity(self) -> str:
+        return self._entity or self.wandb_client.default_entity
+
+    @property
+    def url(self) -> str:
+        return f"wandb://{self.entity}/{self.project}"
+
+    @classmethod
+    def from_parsed_url(cls, parsed_url: ParseResult) -> Workspace:
+        entity = parsed_url.netloc
+        project = parsed_url.path
+        if project:
+            project = project.strip("/")
+        return cls(project=project, entity=entity)
+
+    @property
+    def step_cache(self) -> StepCache:
+        return self.cache
+
+    @property
+    def wandb_project_url(self) -> str:
+        """
+        The URL of the W&B project this workspace uses.
+        """
+        app_url = self.wandb_client.client.app_url
+        app_url = app_url.rstrip("/")
+        return f"{app_url}/{self.entity}/{self.project}"
+
+    def _get_unique_id(self, step_or_unique_id: Union[Step, str]) -> str:
+        if isinstance(step_or_unique_id, Step):
+            unique_id = step_or_unique_id.unique_id
+        else:
+            unique_id = step_or_unique_id
+        return unique_id
+
+    def step_dir(self, step_or_unique_id: Union[Step, str]) -> Path:
+        unique_id = self._get_unique_id(step_or_unique_id)
+        path = self.steps_dir / unique_id
+        path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    def work_dir(self, step: Step) -> Path:
+        path = self.step_dir(step) / "work"
+        path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    def step_info(self, step_or_unique_id: Union[Step, str]) -> StepInfo:
+        unique_id = self._get_unique_id(step_or_unique_id)
+        if unique_id in self._running_step_info:
+            return self._running_step_info[unique_id]
+        step_info = self._get_updated_step_info(
+            unique_id,
+            step_name=step_or_unique_id.name if isinstance(step_or_unique_id, Step) else None,
+        )
+        if step_info is None:
+            raise KeyError(step_or_unique_id)
+        else:
+            return step_info
+
+    def step_starting(self, step: Step) -> None:
+        if wandb.run is not None:
+            raise RuntimeError(
+                "There is already a W&B run initialized, cannot initialize another one."
+            )
+
+        work_dir = self.work_dir(step)
+
+        lock_path = self.step_dir(step) / "lock"
+        lock = FileLock(lock_path, read_only_ok=True)
+        lock.acquire_with_updates(desc=f"acquiring lock for '{step.name}'")
+        self.locks[step] = lock
+
+        step_info = self._get_updated_step_info(step.unique_id) or StepInfo.new_from_step(step)
+        if step_info.state not in {StepState.INCOMPLETE, StepState.FAILED, StepState.UNCACHEABLE}:
+            raise RuntimeError(
+                f"Step '{step.name}' is trying to start, but it is already {step_info.state}. "
+                "If you are certain the step is not running somewhere else, delete the lock "
+                f"file at {lock_path}."
+            )
+
+        try:
+            # Initialize W&B run for the step.
+            wandb.init(
+                name=step_info.step_name,
+                job_type=RunKind.STEP.value,
+                group=step.unique_id,
+                dir=str(work_dir),
+                entity=self.entity,
+                project=self.project,
+                # For cacheable steps we can just use the step's unique ID as the W&B run ID,
+                # but not for uncacheable steps since those might be ran more than once, and
+                # and will need a unique W&B run ID each time.
+                id=step.unique_id if step.cache_results else None,
+                resume="allow" if step.cache_results else None,
+                notes="\n".join(
+                    [
+                        f'Tango step "{step.name}"',
+                        f"\N{bullet} type: {step_info.step_class_name}",
+                        f"\N{bullet} ID: {step.unique_id}",
+                    ]
+                ),
+                config={
+                    "job_type": RunKind.STEP.value,
+                    "_run_suite_id": self._generate_run_suite_id(),  # used for testing only
+                },
+            )
+
+            assert wandb.run is not None
+            logger.info(
+                "Tracking '%s' step on Weights and Biases: %s/runs/%s/overview",
+                step.name,
+                self.wandb_project_url,
+                wandb.run.id,
+            )
+
+            # "Use" all of the result artifacts for this step's dependencies in order to declare
+            # those dependencies to W&B.
+            for dependency in step.dependencies:
+                self.cache.use_step_result_artifact(dependency)
+
+            # Update StepInfo to mark as running.
+            step_info.start_time = utc_now_datetime()
+            step_info.end_time = None
+            step_info.error = None
+            step_info.result_location = None
+            wandb.run.config.update({"step_info": step_info.to_json_dict()}, allow_val_change=True)
+            self._running_step_info[step.unique_id] = step_info
+        except:  # noqa: E722
+            lock.release()
+            del self.locks[step]
+            raise
+
+    def step_finished(self, step: Step, result: T) -> T:
+        if wandb.run is None:
+            raise RuntimeError(
+                f"{self.__class__.__name__}.step_finished() called outside of a W&B run. "
+                f"Did you forget to call {self.__class__.__name__}.step_starting() first?"
+            )
+
+        step_info = self._running_step_info.pop(step.unique_id)
+
+        try:
+            if step.cache_results:
+                self.step_cache[step] = result
+                if hasattr(result, "__next__"):
+                    assert isinstance(result, Iterator)
+                    # Caching the iterator will consume it, so we write it to the
+                    # cache and then read from the cache for the return value.
+                    result = self.step_cache[step]
+                result_artifact = self.cache.get_step_result_artifact(step)
+                if result_artifact is None:
+                    raise RuntimeError(f"Failed to find step result artifact for {step.unique_id}")
+                step_info.result_location = (
+                    f"{self.wandb_project_url}/artifacts/{ArtifactKind.STEP_RESULT.value}"
+                    f"/{result_artifact._sequence_name}/{result_artifact.commit_hash}"
+                )
+            else:
+                # Create an empty artifact in order to build the DAG in W&B.
+                self.cache.create_step_result_artifact(step)
+
+            step_info.end_time = utc_now_datetime()
+            wandb.run.config.update({"step_info": step_info.to_json_dict()}, allow_val_change=True)
+
+            # Finalize the step's W&B run.
+            wandb.finish()
+        finally:
+            self.locks[step].release()
+            del self.locks[step]
+
+        return result
+
+    def step_failed(self, step: Step, e: BaseException) -> None:
+        if wandb.run is None:
+            raise RuntimeError(
+                f"{self.__class__.__name__}.step_failed() called outside of a W&B run. "
+                f"Did you forget to call {self.__class__.__name__}.step_starting() first?"
+            )
+
+        step_info = self._running_step_info.pop(step.unique_id)
+
+        try:
+            # Update StepInfo, marking the step as failed.
+            if step_info.state != StepState.RUNNING:
+                raise RuntimeError(f"Step '{step.name}' is failing, but it never started.")
+            step_info.end_time = utc_now_datetime()
+            step_info.error = exception_to_string(e)
+            wandb.run.config.update({"step_info": step_info.to_json_dict()}, allow_val_change=True)
+
+            # Finalize the step's W&B run.
+            wandb.finish(exit_code=1)
+        finally:
+            self.locks[step].release()
+            del self.locks[step]
+
+    def register_run(self, targets: Iterable[Step], name: Optional[str] = None) -> Run:
+        all_steps = set(targets)
+        for step in targets:
+            all_steps |= step.recursive_dependencies
+
+        wandb_run_id: str
+        wandb_run_name: str
+        with tempfile.TemporaryDirectory() as temp_dir_name:
+            with wandb.init(  # type: ignore[union-attr]
+                job_type=RunKind.TANGO_RUN.value,
+                entity=self.entity,
+                project=self.project,
+                name=name,
+                dir=temp_dir_name,
+                config={
+                    "job_type": RunKind.TANGO_RUN.value,  # need this in the config so we can filter runs by this
+                    "_run_suite_id": self._generate_run_suite_id(),  # used for testing only
+                },
+            ) as wandb_run:
+                wandb_run_id = wandb_run.id
+                wandb_run_name = wandb_run.name  # type: ignore[assignment]
+                logger.info("Registering run %s with Weights and Biases", wandb_run.name)
+                logger.info(
+                    "View run at: %s/runs/%s/overview", self.wandb_project_url, wandb_run_id
+                )
+
+                # Collect step info for all steps.
+                step_ids: Dict[str, bool] = {}
+                step_name_to_info: Dict[str, Dict[str, Any]] = {}
+                for step in all_steps:
+                    step_info = StepInfo.new_from_step(step)
+                    step_name_to_info[step.name] = {
+                        k: v for k, v in step_info.to_json_dict().items() if v is not None
+                    }
+                    step_ids[step.unique_id] = True
+
+                # Update config with step info.
+                wandb_run.config.update({"steps": step_name_to_info, "_step_ids": step_ids})
+
+                # Update notes.
+                notes = "Tango run\n--------------"
+                cacheable_steps = {step for step in all_steps if step.cache_results}
+                if cacheable_steps:
+                    notes += "\nCacheable steps:\n"
+                    for step in sorted(cacheable_steps, key=lambda step: step.name):
+                        notes += f"\N{bullet} {step.name}"
+                        dependencies = step.dependencies
+                        if dependencies:
+                            notes += ", depends on: " + ", ".join(
+                                sorted(
+                                    [f"'{dep.name}'" for dep in dependencies],
+                                )
+                            )
+                        notes += "\n  \N{rightwards arrow with hook} "
+                        notes += f"{self.wandb_project_url}/runs/{step.unique_id}/overview\n"
+                wandb_run.notes = notes
+
+        return self.registered_run(wandb_run_name)
+
+    def _generate_run_suite_id(self) -> str:
+        return wandb.util.generate_id()
+
+    def registered_runs(self) -> Dict[str, Run]:
+        runs: Dict[str, Run] = {}
+        matching_runs = list(
+            self.wandb_client.runs(
+                f"{self.entity}/{self.project}",
+                filters={"config.job_type": RunKind.TANGO_RUN.value},
+            )
+        )
+        for wandb_run in matching_runs:
+            runs[wandb_run.name] = self._get_run_from_wandb_run(wandb_run)
+        return runs
+
+    def registered_run(self, name: str) -> Run:
+        matching_runs = list(
+            self.wandb_client.runs(
+                f"{self.entity}/{self.project}",
+                filters={"display_name": name, "config.job_type": RunKind.TANGO_RUN.value},
+            )
+        )
+        if not matching_runs:
+            raise KeyError(f"Run '{name}' not found in workspace")
+        elif len(matching_runs) > 1:
+            raise ValueError(f"Found more than one run named '{name}' in W&B project")
+        return self._get_run_from_wandb_run(matching_runs[0])
+
+    def _get_run_from_wandb_run(
+        self,
+        wandb_run: wandb.apis.public.Run,
+    ) -> Run:
+        step_name_to_info = {}
+        for step_name, step_info_dict in wandb_run.config["steps"].items():
+            step_info = StepInfo.from_json_dict(step_info_dict)
+            if step_info.cacheable:
+                updated_step_info = self._get_updated_step_info(
+                    step_info.unique_id, step_name=step_name
+                )
+                if updated_step_info is not None:
+                    step_info = updated_step_info
+            step_name_to_info[step_name] = step_info
+        return Run(
+            name=wandb_run.name,
+            steps=step_name_to_info,
+            start_date=datetime.strptime(wandb_run.created_at, "%Y-%m-%dT%H:%M:%S").replace(
+                tzinfo=pytz.utc
+            ),
+        )
+
+    def _get_updated_step_info(
+        self, step_id: str, step_name: Optional[str] = None
+    ) -> Optional[StepInfo]:
+        # First try to find the W&B run corresponding to the step. This will only
+        # work if the step execution was started already.
+        filters = {
+            "config.job_type": RunKind.STEP.value,
+            "config.step_info.unique_id": step_id,
+        }
+        if step_name is not None:
+            filters["display_name"] = step_name
+        for wandb_run in self.wandb_client.runs(
+            f"{self.entity}/{self.project}",
+            filters=filters,
+        ):
+            step_info = StepInfo.from_json_dict(wandb_run.config["step_info"])
+            # Might need to fix the step info the step failed and we failed to update the config.
+            if step_info.start_time is None:
+                step_info.start_time = datetime.strptime(
+                    wandb_run.created_at, "%Y-%m-%dT%H:%M:%S"
+                ).replace(tzinfo=pytz.utc)
+            if wandb_run.state in {"failed", "finished"}:
+                if step_info.end_time is None:
+                    step_info.end_time = datetime.strptime(
+                        wandb_run.heartbeatAt, "%Y-%m-%dT%H:%M:%S"
+                    ).replace(tzinfo=pytz.utc)
+                if wandb_run.state == "failed" and step_info.error is None:
+                    step_info.error = "Exception"
+            return step_info
+
+        # If the step hasn't been started yet, we'll have to pull the step info from the
+        # registered run.
+        filters = {
+            "config.job_type": RunKind.TANGO_RUN.value,
+            f"config._step_ids.{step_id}": True,
+        }
+        if step_name is not None:
+            filters[f"config.steps.{step_name}.unique_id"] = step_id
+        for wandb_run in self.wandb_client.runs(
+            f"{self.entity}/{self.project}",
+            filters=filters,
+        ):
+            if step_name is not None:
+                step_info_data = wandb_run.config["steps"][step_name]
+            else:
+                step_info_data = next(
+                    d for d in wandb_run.config["steps"].values() if d["unique_id"] == step_id
+                )
+            step_info = StepInfo.from_json_dict(step_info_data)
+            return step_info
+
+        return None
diff --git a/tango/server/workspace_server.py b/tango/server/workspace_server.py
index ca737c59b..c7d308695 100644
--- a/tango/server/workspace_server.py
+++ b/tango/server/workspace_server.py
@@ -49,8 +49,10 @@ def _serialize_step_info(cls, step_info: StepInfo) -> Dict[str, Any]:
             "step_class_name": step_info.step_class_name,
             "version": step_info.version,
             "dependencies": list(step_info.dependencies),
-            "start_time": step_info.start_time.isoformat() if step_info.start_time else None,
-            "end_time": step_info.end_time.isoformat() if step_info.end_time else None,
+            "start_time": step_info.start_time_local.isoformat()
+            if step_info.start_time_local
+            else None,
+            "end_time": step_info.end_time_local.isoformat() if step_info.end_time_local else None,
             "error": error,
             "result_location": result_location,
             "state": step_info.state.value,
diff --git a/tango/step.py b/tango/step.py
index 7d795f3a5..683451838 100644
--- a/tango/step.py
+++ b/tango/step.py
@@ -23,17 +23,6 @@
     cast,
 )
 
-try:
-    from typing import get_args, get_origin  # type: ignore
-except ImportError:
-
-    def get_origin(tp):  # type: ignore
-        return getattr(tp, "__origin__", None)
-
-    def get_args(tp):  # type: ignore
-        return getattr(tp, "__args__", ())
-
-
 from tango.common.det_hash import CustomDetHash, det_hash
 from tango.common.exceptions import ConfigurationError
 from tango.common.from_params import (
@@ -47,6 +36,17 @@ def get_args(tp):  # type: ignore
 from tango.common.registrable import Registrable
 from tango.format import DillFormat, Format
 
+try:
+    from typing import get_args, get_origin  # type: ignore
+except ImportError:
+
+    def get_origin(tp):  # type: ignore
+        return getattr(tp, "__origin__", None)
+
+    def get_args(tp):  # type: ignore
+        return getattr(tp, "__args__", ())
+
+
 if TYPE_CHECKING:
     from tango.workspace import Workspace
 
@@ -117,6 +117,11 @@ class Step(Registrable, Generic[T]):
     the model output, not about how many outputs you can produce at the same time.
     """
 
+    _UNIQUE_ID_SUFFIX: Optional[str] = None
+    """
+    Used internally for testing.
+    """
+
     def __init__(
         self,
         step_name: Optional[str] = None,
@@ -337,7 +342,7 @@ def run(self, **kwargs) -> T:
         """
         raise NotImplementedError()
 
-    def _run_with_work_dir(self, workspace: "Workspace", **kwargs) -> T:
+    def _run_with_work_dir(self, workspace: "Workspace", needed_by: Optional["Step"] = None) -> T:
         if self.work_dir_for_run is not None:
             raise RuntimeError("You can only run a Step's run() method once at a time.")
 
@@ -354,18 +359,29 @@ def _run_with_work_dir(self, workspace: "Workspace", **kwargs) -> T:
             self.work_dir_for_run = Path(dir_for_cleanup.name)
 
         try:
-            if self.cache_results:
-                workspace.step_starting(self)
+            kwargs = self._replace_steps_with_results(self.kwargs, workspace)
+
+            if needed_by:
+                cli_logger.info(
+                    '[blue]\N{black circle} Starting step [bold]"%s"[/] (needed by "%s")...[/]',
+                    self.name,
+                    needed_by.name,
+                )
+            else:
+                cli_logger.info(
+                    '[blue]\N{black circle} Starting step [bold]"%s"[/]...[/]',
+                    self.name,
+                )
+
+            workspace.step_starting(self)
+
             try:
                 result = self.run(**kwargs)
-                if self.cache_results:
-                    result = workspace.step_finished(self, result)
+                result = workspace.step_finished(self, result)
+                cli_logger.info(f'[green]\N{check mark} Finished step [bold]"{self.name}"[/][/]')
                 return result
             except BaseException as e:
-                # TODO (epwalsh): do we want to handle KeyboardInterrupts differently?
-                # Maybe have a `workspace.step_interrupted()` method?
-                if self.cache_results:
-                    workspace.step_failed(self, e)
+                workspace.step_failed(self, e)
                 raise
         finally:
             self._workspace = None
@@ -448,6 +464,8 @@ def unique_id(self) -> str:
                 self.unique_id_cache += det_hash(
                     _random_for_step_names.getrandbits((58**32).bit_length())
                 )[:32]
+            if self._UNIQUE_ID_SUFFIX is not None:
+                self.unique_id_cache += f"-{self._UNIQUE_ID_SUFFIX}"
 
         return self.unique_id_cache
 
@@ -520,23 +538,8 @@ def result(
                     self.name,
                 )
             return workspace.step_cache[self]
-
-        kwargs = self._replace_steps_with_results(self.kwargs, workspace)
-
-        if needed_by:
-            cli_logger.info(
-                '[blue]\N{black circle} Starting step [bold]"%s"[/] (needed by "%s")...[/]',
-                self.name,
-                needed_by.name,
-            )
         else:
-            cli_logger.info(
-                '[blue]\N{black circle} Starting step [bold]"%s"[/]...[/]',
-                self.name,
-            )
-        result = self._run_with_work_dir(workspace, **kwargs)
-        cli_logger.info(f'[green]\N{check mark} Finished step [bold]"{self.name}"[/][/]')
-        return result
+            return self._run_with_work_dir(workspace, needed_by=needed_by)
 
     def ensure_result(
         self,
diff --git a/tango/step_cache.py b/tango/step_cache.py
index d7531689a..29ef70b35 100644
--- a/tango/step_cache.py
+++ b/tango/step_cache.py
@@ -1,11 +1,13 @@
 import logging
 from abc import abstractmethod
 from dataclasses import dataclass
-from typing import Any, TypeVar
+from typing import Any, TypeVar, Union
 
-from tango.common.from_params import FromParams
-from tango.common.registrable import Registrable
-from tango.step import Step
+from .common.from_params import FromParams
+from .common.registrable import Registrable
+from .format import Format
+from .step import Step
+from .step_info import StepInfo
 
 logger = logging.getLogger(__name__)
 
@@ -28,7 +30,7 @@ class StepCache(Registrable):
     def __contains__(self, step: Any) -> bool:
         """This is a generic implementation of ``__contains__``. If you are writing your own
         ``StepCache``, you might want to write a faster one yourself."""
-        if not isinstance(step, Step):
+        if not isinstance(step, (Step, StepInfo)):
             return False
         try:
             self.__getitem__(step)
@@ -37,7 +39,7 @@ def __contains__(self, step: Any) -> bool:
             return False
 
     @abstractmethod
-    def __getitem__(self, step: Step) -> Any:
+    def __getitem__(self, step: Union[Step, StepInfo]) -> Any:
         """Returns the results for the given step."""
         raise NotImplementedError()
 
@@ -58,3 +60,8 @@ class CacheMetadata(FromParams):
     """
     The step name.
     """
+
+    format: Format
+    """
+    The format used to serialize the step's result.
+    """
diff --git a/tango/step_caches/local_step_cache.py b/tango/step_caches/local_step_cache.py
index 92f369a8c..7f4d3272b 100644
--- a/tango/step_caches/local_step_cache.py
+++ b/tango/step_caches/local_step_cache.py
@@ -1,12 +1,15 @@
 import collections
 import logging
+import warnings
 import weakref
 from pathlib import Path
-from typing import Any, MutableMapping, Optional, OrderedDict, Union
+from typing import Any, MutableMapping, Optional, OrderedDict, Union, cast
 
 from tango.common.aliases import PathOrStr
+from tango.common.params import Params
 from tango.step import Step
 from tango.step_cache import CacheMetadata, StepCache
+from tango.step_info import StepInfo
 
 logger = logging.getLogger(__name__)
 
@@ -30,6 +33,7 @@ class LocalStepCache(StepCache):
     """
 
     LRU_CACHE_MAX_SIZE = 8
+    METADATA_FILE_NAME = "cache-metadata.json"
 
     def __init__(self, dir: PathOrStr):
         self.dir = Path(dir)
@@ -37,11 +41,28 @@ def __init__(self, dir: PathOrStr):
 
         # We keep an in-memory cache as well so we don't have to de-serialize stuff
         # we happen to have in memory already.
-        self.weak_cache: MutableMapping[str, Any] = weakref.WeakValueDictionary()
-
+        self.weak_cache: MutableMapping[str, Any]
         # Not all Python objects can be referenced weakly, and even if they can they
         # might get removed too quickly, so we also keep an LRU cache.
-        self.strong_cache: OrderedDict[str, Any] = collections.OrderedDict()
+        self.strong_cache: OrderedDict[str, Any]
+        self._init_mem_caches()
+
+    def _init_mem_caches(self):
+        self.weak_cache = weakref.WeakValueDictionary()
+        self.strong_cache = collections.OrderedDict()
+
+    def __getstate__(self):
+        """
+        We override `__getstate__()` to customize how instances of this class are pickled
+        since we don't want to persist values in the weak and strong in-memory caches
+        during pickling. And `WeakValueDictionary` can't be pickled anyway.
+        """
+        return {"dir": self.dir}
+
+    def __setstate__(self, state):
+        for k, v in state.items():
+            setattr(self, k, v)
+        self._init_mem_caches()
 
     def _add_to_cache(self, key: str, o: Any) -> None:
         if hasattr(o, "__next__"):
@@ -68,44 +89,54 @@ def _get_from_cache(self, key: str) -> Optional[Any]:
         except KeyError:
             return None
 
+    def _metadata_path(self, step_or_unique_id: Union[Step, StepInfo, str]) -> Path:
+        return self.step_dir(step_or_unique_id) / self.METADATA_FILE_NAME
+
     def __contains__(self, step: object) -> bool:
-        if isinstance(step, Step) and step.cache_results:
+        if (isinstance(step, Step) and step.cache_results) or (
+            isinstance(step, StepInfo) and step.cacheable
+        ):
             key = step.unique_id
             if key in self.strong_cache:
                 return True
             if key in self.weak_cache:
                 return True
-            metadata_file = self.step_dir(step) / "cache-metadata.json"
-            return metadata_file.exists()
+            return self._metadata_path(
+                cast(Union[Step, StepInfo], step)  # cast is for mypy :/
+            ).exists()
         else:
             return False
 
-    def __getitem__(self, step: Step) -> Any:
+    def __getitem__(self, step: Union[Step, StepInfo]) -> Any:
         key = step.unique_id
         result = self._get_from_cache(key)
         if result is None:
             if step not in self:
                 raise KeyError(step)
-            result = step.format.read(self.step_dir(step))
+            metadata = CacheMetadata.from_params(Params.from_file(self._metadata_path(step)))
+            result = metadata.format.read(self.step_dir(step))
             self._add_to_cache(key, result)
         return result
 
     def __setitem__(self, step: Step, value: Any) -> None:
         if not step.cache_results:
-            logger.warning("Tried to cache step %s despite being marked as uncacheable.", step.name)
+            warnings.warn(
+                f"Tried to cache step '{step.name}' despite being marked as uncacheable",
+                UserWarning,
+            )
             return
 
         location = self.step_dir(step)
         location.mkdir(parents=True, exist_ok=True)
 
-        metadata_location = location / "cache-metadata.json"
+        metadata_location = self._metadata_path(step)
         if metadata_location.exists():
             raise ValueError(f"{metadata_location} already exists! Will not overwrite.")
         temp_metadata_location = metadata_location.with_suffix(".temp")
 
         try:
             step.format.write(value, location)
-            metadata = CacheMetadata(step=step.unique_id)
+            metadata = CacheMetadata(step=step.unique_id, format=step.format)
             metadata.to_params().to_file(temp_metadata_location)
             self._add_to_cache(step.unique_id, value)
             temp_metadata_location.rename(metadata_location)
@@ -117,17 +148,27 @@ def __setitem__(self, step: Step, value: Any) -> None:
             raise
 
     def __len__(self) -> int:
-        return sum(1 for _ in self.dir.glob("*/cache-metadata.json"))
+        return sum(1 for _ in self.dir.glob(f"*/{self.METADATA_FILE_NAME}"))
 
-    def step_dir(self, step_or_unique_id: Union[Step, str]) -> Path:
+    def step_dir(self, step_or_unique_id: Union[Step, StepInfo, str]) -> Path:
         """Returns the directory that contains the results of the step.
 
         You can use this even for a step that's not cached yet. In that case it will return the directory where
         the results will be written."""
-        if isinstance(step_or_unique_id, Step):
-            if not step_or_unique_id.cache_results:
+        if isinstance(step_or_unique_id, (Step, StepInfo)):
+            cacheable = (
+                step_or_unique_id.cache_results
+                if isinstance(step_or_unique_id, Step)
+                else step_or_unique_id.cacheable
+            )
+            if not cacheable:
+                class_name = (
+                    step_or_unique_id.__class__.__name__
+                    if isinstance(step_or_unique_id, Step)
+                    else step_or_unique_id.step_class_name
+                )
                 raise RuntimeError(
-                    f"Uncacheable steps (like '{step_or_unique_id.name}') don't have step directories."
+                    f"Uncacheable steps (like '{class_name}') don't have step directories."
                 )
             unique_id = step_or_unique_id.unique_id
         else:
diff --git a/tango/step_caches/memory_step_cache.py b/tango/step_caches/memory_step_cache.py
index a4351273e..e57751c70 100644
--- a/tango/step_caches/memory_step_cache.py
+++ b/tango/step_caches/memory_step_cache.py
@@ -1,8 +1,10 @@
 import logging
-from typing import Any, Dict
+import warnings
+from typing import Any, Dict, Union
 
 from tango.step import Step
 from tango.step_cache import StepCache
+from tango.step_info import StepInfo
 
 logger = logging.getLogger(__name__)
 
@@ -19,7 +21,7 @@ class MemoryStepCache(StepCache):
     def __init__(self):
         self.cache: Dict[str, Any] = {}
 
-    def __getitem__(self, step: Step) -> Any:
+    def __getitem__(self, step: Union[Step, StepInfo]) -> Any:
         return self.cache[step.unique_id]
 
     def __setitem__(self, step: Step, value: Any) -> None:
@@ -28,10 +30,13 @@ def __setitem__(self, step: Step, value: Any) -> None:
         if step.cache_results:
             self.cache[step.unique_id] = value
         else:
-            logger.warning("Tried to cache step %s despite being marked as uncacheable.", step.name)
+            warnings.warn(
+                f"Tried to cache step '{step.name}' despite being marked as uncacheable.",
+                UserWarning,
+            )
 
-    def __contains__(self, step: object):
-        if isinstance(step, Step):
+    def __contains__(self, step: object) -> bool:
+        if isinstance(step, (Step, StepInfo)):
             return step.unique_id in self.cache
         else:
             return False
diff --git a/tango/step_info.py b/tango/step_info.py
new file mode 100644
index 000000000..86216020f
--- /dev/null
+++ b/tango/step_info.py
@@ -0,0 +1,193 @@
+from dataclasses import asdict, dataclass
+from datetime import datetime, timedelta
+from enum import Enum
+from typing import Any, Dict, Optional, Set
+
+import pytz
+
+from .common.util import local_timezone
+from .step import Step
+
+
+class StepState(Enum):
+    """Describes the possible state a step can be in."""
+
+    INCOMPLETE = "incomplete"
+    """The step has not run yet."""
+
+    RUNNING = "running"
+    """The step is running right now."""
+
+    COMPLETED = "completed"
+    """The step finished running successfully."""
+
+    FAILED = "failed"
+    """The step ran, but failed."""
+
+    UNCACHEABLE = "uncacheable"
+    """The step is uncacheable. It will be executed as many times as the results are needed,
+    so we don't keep track of the state."""
+
+
+@dataclass
+class StepInfo:
+    """Stores step information without being the :class:`.Step` itself.
+
+    It's not always possible to get a :class:`.Step` object, because :class:`.Step` objects can't be serialized.
+    But you can always serialize a :class:`.StepInfo` object.
+    """
+
+    unique_id: str
+    """
+    The unique ID of the step
+    """
+
+    step_class_name: str
+    """
+    The name of the :class:`.Step` class
+    """
+
+    dependencies: Set[str]
+    """
+    The unique ids of all the steps that this step depends on
+    """
+
+    cacheable: bool
+    """
+    Whether or not the step is cacheable.
+    """
+
+    step_name: Optional[str] = None
+    """
+    The name of the step, if it has one. Anonymous steps are identified only by their unique ID.
+
+    The same step can have different names in different runs. The last run wins, so don't rely
+    on this property in your code. It is just here to aid readability.
+    """
+
+    version: Optional[str] = None
+    """
+    The version string of the :class:`.Step`, if it has one.
+    """
+
+    start_time: Optional[datetime] = None
+    """
+    The time (in UTC) that this step started running.
+
+    .. seealso::
+        :meth:`start_time_local()`.
+    """
+
+    end_time: Optional[datetime] = None
+    """
+    The time (in UTC) that this step stopped running. This will be set whether the step succeeded or failed.
+
+    .. seealso::
+        :meth:`end_time_local()`.
+    """
+
+    error: Optional[str] = None
+    """
+    If the step failed, this is where the error goes.
+
+    .. note::
+        Some ``Workspace`` implementations need to serialize ``StepInfo`` (using pickle or dill, for example),
+        but some exceptions can't be pickled. In those cases ``error`` will just be a string representation
+        of the exception.
+    """
+
+    result_location: Optional[str] = None
+    """
+    Location of the result. This could be a path or a URL.
+    """
+
+    @property
+    def start_time_local(self) -> Optional[datetime]:
+        """
+        The time the step started running with respect to the local timezone, if the timezone
+        can be determined.
+        """
+        return None if self.start_time is None else self.start_time.astimezone(local_timezone())
+
+    @property
+    def end_time_local(self) -> Optional[datetime]:
+        """
+        The time the step stopped running with respect to the local timezone, if the timezone
+        can be determined.
+        """
+        return None if self.end_time is None else self.end_time.astimezone(local_timezone())
+
+    @property
+    def duration(self) -> Optional[timedelta]:
+        """
+        The time it took to run this step.
+        """
+        if self.start_time is not None and self.end_time is not None:
+            return self.end_time - self.start_time
+        else:
+            return None
+
+    @property
+    def state(self) -> StepState:
+        """
+        Returns the state of the step
+        """
+        if self.cacheable:
+            if self.start_time is None and self.end_time is None and self.error is None:
+                return StepState.INCOMPLETE
+            if self.start_time is not None and self.end_time is None and self.error is None:
+                return StepState.RUNNING
+            if self.start_time is not None and self.end_time is not None and self.error is None:
+                return StepState.COMPLETED
+            if self.start_time is not None and self.end_time is not None and self.error is not None:
+                return StepState.FAILED
+        else:
+            return StepState.UNCACHEABLE
+        raise RuntimeError(f"{self.__class__.__name__} is in an invalid state.")
+
+    def to_json_dict(self) -> Dict[str, Any]:
+        """
+        Generates a JSON-safe, human-readable, dictionary representation of this dataclass.
+        """
+        return {
+            k: (
+                v.strftime("%Y-%m-%dT%H:%M:%S")
+                if isinstance(v, datetime)
+                else list(v)
+                if isinstance(v, set)
+                else v
+            )
+            for k, v in asdict(self).items()
+        }
+
+    @classmethod
+    def from_json_dict(cls, json_dict: Dict[str, Any]) -> "StepInfo":
+        """
+        The inverse of :meth:`to_json_dict()`.
+
+        :param json_dict: A dictionary representation, such as the one produced by :meth:`to_json_dict()`.
+        """
+        return cls(
+            **{
+                k: (
+                    datetime.strptime(v, "%Y-%m-%dT%H:%M:%S").replace(tzinfo=pytz.utc)
+                    if k in {"start_time", "end_time"} and v is not None
+                    else set(v)
+                    if k == "dependencies"
+                    else v
+                )
+                for k, v in json_dict.items()
+            }
+        )
+
+    @classmethod
+    def new_from_step(cls, step: Step, **kwargs) -> "StepInfo":
+        return cls(
+            unique_id=step.unique_id,
+            step_name=step.name,
+            step_class_name=step.__class__.__name__,
+            version=step.VERSION,
+            dependencies={dep.unique_id for dep in step.dependencies},
+            cacheable=step.cache_results,
+            **kwargs,
+        )
diff --git a/tango/workspace.py b/tango/workspace.py
index 7b8269e2e..fc9277b1b 100644
--- a/tango/workspace.py
+++ b/tango/workspace.py
@@ -7,9 +7,8 @@
 import time
 from abc import abstractmethod
 from contextlib import contextmanager
-from dataclasses import dataclass, field, replace
-from datetime import datetime, timedelta
-from enum import Enum
+from dataclasses import dataclass, field
+from datetime import datetime
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import (
@@ -20,165 +19,23 @@
     Iterable,
     List,
     Optional,
-    Set,
     TypeVar,
     Union,
     cast,
 )
 from urllib.parse import ParseResult, urlparse
 
-import dill
-
-from tango.common import FromParams, Registrable
-from tango.step import Step
-from tango.step_cache import StepCache
-from tango.version import VERSION
+from .common import FromParams, Registrable
+from .step import Step
+from .step_cache import StepCache
+from .step_info import StepInfo
+from .version import VERSION
 
 logger = logging.getLogger(__name__)
 
 T = TypeVar("T")
 
 
-class StepState(Enum):
-    """Describes the possible state a step can be in."""
-
-    INCOMPLETE = "incomplete"
-    """The step has not run yet."""
-
-    RUNNING = "running"
-    """The step is running right now."""
-
-    COMPLETED = "completed"
-    """The step finished running successfully."""
-
-    FAILED = "failed"
-    """The step ran, but failed."""
-
-    UNCACHEABLE = "uncacheable"
-    """The step is uncacheable. It will be executed as many times as the results are needed,
-    so we don't keep track of the state."""
-
-
-@dataclass
-class StepInfo:
-    """Stores step information without being the :class:`.Step` itself.
-
-    It's not always possible to get a :class:`.Step` object, because :class:`.Step` objects can't be serialized.
-    But you can always serialize a :class:`.StepInfo` object.
-    """
-
-    unique_id: str
-    """
-    The unique ID of the step
-    """
-
-    step_name: Optional[str]
-    """
-    The name of the step, if it has one. Anonymous steps are identified only by their unique ID.
-
-    The same step can have different names in different runs. The last run wins, so don't rely
-    on this property in your code. It is just here to aid readability.
-    """
-
-    step_class_name: str
-    """
-    The name of the :class:`.Step` class
-    """
-
-    version: Optional[str]
-    """
-    The version string of the :class:`.Step`, if it has one
-    """
-
-    dependencies: Set[str]
-    """
-    The unique ids of all the steps that this step depends on
-    """
-
-    cacheable: bool
-    """
-    Whether or not the step is cacheable.
-    """
-
-    start_time: Optional[datetime] = None
-    """
-    The time this step started running
-    """
-
-    end_time: Optional[datetime] = None
-    """
-    The time this step stopped running. This will be set whether the step succeeded or failed.
-    """
-
-    error: Optional[str] = None
-    """
-    If the step failed, this is where the error goes.
-
-    .. note::
-        Some ``Workspace`` implementations need to serialize ``StepInfo`` (using pickle or dill, for example),
-        but some exceptions can't be pickled. In those cases ``error`` will just be a string representation
-        of the exception.
-    """
-
-    result_location: Optional[str] = None
-    """
-    Location of the result. This could be a path or a URL.
-    """
-
-    @property
-    def duration(self) -> Optional[timedelta]:
-        """
-        The time it took to run this step.
-        """
-        if self.start_time is not None and self.end_time is not None:
-            return self.end_time - self.start_time
-        else:
-            return None
-
-    @property
-    def state(self) -> StepState:
-        """
-        Returns the state of the step
-        """
-        if self.cacheable:
-            if self.start_time is None and self.end_time is None and self.error is None:
-                return StepState.INCOMPLETE
-            if self.start_time is not None and self.end_time is None and self.error is None:
-                return StepState.RUNNING
-            if self.start_time is not None and self.end_time is not None and self.error is None:
-                return StepState.COMPLETED
-            if self.start_time is not None and self.end_time is not None and self.error is not None:
-                return StepState.FAILED
-        else:
-            return StepState.UNCACHEABLE
-        raise RuntimeError(f"{self.__class__.__name__} is in an invalid state.")
-
-    def serialize(self) -> bytes:
-        """
-        Returns a serialized form of the ``StepInfo``.
-        """
-        instance_to_dump = self
-        if isinstance(self.error, BaseException):
-            # See if we can pickle and unpickle the exception.
-            # When we can't, we'll fallback to storing the exception as a string
-            # representation of it.
-            dump = dill.dumps(self.error)
-            try:
-                dill.loads(dump)
-            except TypeError:
-                # Fails with TypeError for some exceptions that take multiple positional
-                # arguments.
-                instance_to_dump = replace(self, error=repr(self.error))
-        return dill.dumps(instance_to_dump)
-
-    @classmethod
-    def deserialize(cls, data: bytes) -> "StepInfo":
-        """
-        Deserialize the result of :meth:`serialize()` into a ``StepInfo`` instance.
-        """
-        return dill.loads(data)
-
-
 @dataclass
 class Run:
     """
@@ -192,7 +49,7 @@ class Run:
 
     steps: Dict[str, StepInfo]
     """
-    A mapping from step names to :class:`StepInfo`, for all the target steps in the run.
+    A mapping from step names to :class:`~tango.step_info.StepInfo`, for all the target steps in the run.
 
     This only contains the targets of a run. Usually, that means it contains all named steps.
     Un-named dependencies (or dependencies that are not targets) are not contained in ``steps``.
@@ -228,12 +85,32 @@ class Workspace(Registrable):
     def __init__(self):
         self._delayed_cleanup_temp_dirs: List[TemporaryDirectory] = []
 
+    def __getstate__(self):
+        """
+        We override `__getstate__()` to customize how instances of this class are pickled
+        since we don't want to persist certain attributes.
+        """
+        out = {k: v for k, v in self.__dict__.items() if k not in {"_delayed_cleanup_temp_dirs"}}
+        out["_delayed_cleanup_temp_dirs"] = []
+        return out
+
+    @property
+    def url(self) -> str:
+        """
+        Get a URL for the workspace that can be used to instantiate the same workspace
+        using :meth:`.from_url()`.
+        """
+        raise NotImplementedError
+
     @classmethod
     def from_url(cls, url: str) -> "Workspace":
         """
         Initialize a :class:`Workspace` from a workspace URL or path, e.g. ``local:///tmp/workspace``
         would give you a :class:`~tango.workspaces.LocalWorkspace` in the directory ``/tmp/workspace``.
 
+        For :class:`~tango.workspaces.LocalWorkspace`, you can also just pass in a plain
+        path, e.g. ``/tmp/workspace``.
+
         .. tip::
             Registered as a workspace constructor under the name "from_url".
 
@@ -278,7 +155,9 @@ def work_dir(self, step: Step) -> Path:
     @abstractmethod
     def step_info(self, step_or_unique_id: Union[Step, str]) -> StepInfo:
         """
-        Returns a :class:`.StepInfo` for a given step
+        Returns a :class:`~tango.step_info.StepInfo` for a given step.
+
+        :raises KeyError: If the step has not been registered as part of a run yet.
         """
         raise NotImplementedError()
 
@@ -343,14 +222,20 @@ def registered_run(self, name: str) -> Run:
 
         :return: A :class:`Run` object representing the named run
 
-        Note that the :class:`Run` object only contains the targets of a run. Usually, that means it
-        contains all named steps. Un-named dependencies (or dependencies that are not targets)
-        are not contained in the result.
-
-        This method throws ``KeyError`` if there is no run with the given name.
+        :raises KeyError: If there is no run with the given name.
         """
         raise NotImplementedError()
 
+    def step_result_for_run(self, run_name: str, step_name: str) -> Any:
+        """
+        Get the result of a step from a run.
+
+        :raises KeyError: If there is no run or step with the given name.
+        """
+        run = self.registered_run(run_name)
+        step_info = run.steps[step_name]
+        return self.step_cache[step_info]
+
     def capture_logs_for_run(self, name: str) -> ContextManager[None]:
         """
         Should return a context manager that can be used to capture the logs for a run.
@@ -452,10 +337,15 @@ def check_for_repo(cls) -> Optional["GitMetadata"]:
                 .strip()
                 .split("\n")
             ):
-                if "(fetch)" in line:
-                    _, line = line.split("\t")
-                    remote = line.split(" ")[0]
-                    break
+                remotes: Dict[str, str] = {}
+                if line.endswith("(fetch)"):
+                    name, info = line.split("\t")
+                    url = info.split(" ")[0]
+                    remotes[name] = url
+                if "origin" in remotes:
+                    remote = remotes["origin"]
+                elif remotes:
+                    remote = list(remotes.values())[0]
             return cls(commit=commit, remote=remote)
         except (subprocess.CalledProcessError, FileNotFoundError):
             return None
diff --git a/tango/workspaces/local_workspace.py b/tango/workspaces/local_workspace.py
index f32f31a5d..0499de4c1 100644
--- a/tango/workspaces/local_workspace.py
+++ b/tango/workspaces/local_workspace.py
@@ -6,17 +6,19 @@
 from typing import Any, Dict, Iterable, Iterator, Optional, Set, TypeVar, Union
 from urllib.parse import ParseResult
 
+import dill
 import petname
 from sqlitedict import SqliteDict
 
 from tango.common import PathOrStr
 from tango.common.file_lock import FileLock
 from tango.common.logging import file_handler
-from tango.common.util import exception_to_string
+from tango.common.util import exception_to_string, utc_now_datetime
 from tango.step import Step
 from tango.step_cache import StepCache
 from tango.step_caches import LocalStepCache
-from tango.workspace import Run, StepExecutionMetadata, StepInfo, StepState, Workspace
+from tango.step_info import StepInfo, StepState
+from tango.workspace import Run, StepExecutionMetadata, Workspace
 
 logger = logging.getLogger(__name__)
 
@@ -72,7 +74,7 @@ def __init__(self, dir: PathOrStr):
             with SqliteDict(self.step_info_file) as d:
                 for stepinfo_file in self.cache.dir.glob("*/stepinfo.dill"):
                     with stepinfo_file.open("rb") as f:
-                        stepinfo = StepInfo.deserialize(f.read())
+                        stepinfo = dill.load(f)
 
                     # The `StepInfo` class changed from one version to the next. The deserialized version
                     # ends up being a `StepInfo` instance that is missing the `cacheable` member. This
@@ -90,9 +92,31 @@ def __init__(self, dir: PathOrStr):
             with open(self.dir / "settings.json", "w") as settings_file:
                 json.dump(settings, settings_file)
 
+    def __getstate__(self):
+        """
+        We override `__getstate__()` to customize how instances of this class are pickled
+        since we don't want to persist certain attributes.
+        """
+        out = super().__getstate__()
+        out["locks"] = {}
+        return out
+
+    @property
+    def url(self) -> str:
+        return "local://" + str(self.dir)
+
     @classmethod
     def from_parsed_url(cls, parsed_url: ParseResult) -> "Workspace":
-        return cls(parsed_url.path)
+        workspace_dir: Path
+        if parsed_url.netloc:
+            workspace_dir = Path(parsed_url.netloc)
+            if parsed_url.path:
+                workspace_dir = workspace_dir / parsed_url.path.lstrip("/")
+        elif parsed_url.path:
+            workspace_dir = Path(parsed_url.path)
+        else:
+            workspace_dir = Path(".")
+        return cls(workspace_dir.resolve())
 
     def step_dir(self, step_or_unique_id: Union[Step, str]) -> Path:
         return self.cache.step_dir(step_or_unique_id)
@@ -173,14 +197,7 @@ def find_or_add_step_info(step_or_unique_id: Union[Step, str]) -> StepInfo:
                     for dep in step.dependencies:
                         find_or_add_step_info(dep)
 
-                    step_info = StepInfo(
-                        step.unique_id,
-                        step.name if step.name != step.unique_id else None,
-                        step.__class__.__name__,
-                        step.VERSION,
-                        {dep.unique_id for dep in step.dependencies},
-                        step.cache_results,
-                    )
+                    step_info = StepInfo.new_from_step(step)
                     d[unique_id] = step_info
                     del step
 
@@ -237,7 +254,7 @@ def step_starting(self, step: Step) -> None:
         self.locks[step] = lock
 
         try:
-            step_info.start_time = datetime.now()
+            step_info.start_time = utc_now_datetime()
             step_info.end_time = None
             step_info.error = None
             step_info.result_location = None
@@ -290,7 +307,7 @@ def replace_steps_with_unique_id(o: Any):
             metadata.save(self.step_dir(step))
 
         # Mark the step as finished
-        step_info.end_time = datetime.now()
+        step_info.end_time = utc_now_datetime()
         step_info.result_location = str(self.step_dir(step).absolute())
         with SqliteDict(self.step_info_file) as d:
             d[step.unique_id] = step_info
@@ -312,7 +329,7 @@ def step_failed(self, step: Step, e: BaseException) -> None:
             step_info = self.step_info(step)
             if step_info.state != StepState.RUNNING:
                 raise RuntimeError(f"Step '{step.name}' is failing, but it never started.")
-            step_info.end_time = datetime.now()
+            step_info.end_time = utc_now_datetime()
             step_info.error = exception_to_string(e)
             with SqliteDict(self.step_info_file) as d:
                 d[step.unique_id] = step_info
@@ -382,14 +399,7 @@ def _save_registered_run(self, name: str, all_steps: Iterable[Step]) -> None:
                     step_info.name = step.name
                     d[step.unique_id] = step_info
                 except KeyError:
-                    d[step.unique_id] = StepInfo(
-                        step.unique_id,
-                        step.name,
-                        step.__class__.__name__,
-                        step.VERSION,
-                        {dep.unique_id for dep in step.dependencies},
-                        step.cache_results,
-                    )
+                    d[step.unique_id] = StepInfo.new_from_step(step)
                 step_unique_ids[step.name] = step.unique_id
 
             d.commit()
diff --git a/tango/workspaces/memory_workspace.py b/tango/workspaces/memory_workspace.py
index 2b22d8b37..452f30907 100644
--- a/tango/workspaces/memory_workspace.py
+++ b/tango/workspaces/memory_workspace.py
@@ -1,15 +1,15 @@
 import copy
-from datetime import datetime
 from typing import Dict, Iterable, Iterator, Optional, TypeVar, Union
 from urllib.parse import ParseResult
 
 import petname
 
-from tango.common.util import exception_to_string
+from tango.common.util import exception_to_string, utc_now_datetime
 from tango.step import Step
 from tango.step_cache import StepCache
 from tango.step_caches import default_step_cache
-from tango.workspace import Run, StepInfo, StepState, Workspace
+from tango.step_info import StepInfo, StepState
+from tango.workspace import Run, Workspace
 
 T = TypeVar("T")
 
@@ -30,6 +30,10 @@ def __init__(self):
         self.unique_id_to_info: Dict[str, StepInfo] = {}
         self.runs: Dict[str, Run] = {}
 
+    @property
+    def url(self) -> str:
+        return "memory://"
+
     @classmethod
     def from_parsed_url(cls, parsed_url: ParseResult) -> "Workspace":
         return cls()
@@ -49,14 +53,7 @@ def step_info(self, step_or_unique_id: Union[Step, str]) -> StepInfo:
         except KeyError:
             if isinstance(step_or_unique_id, Step):
                 step = step_or_unique_id
-                return StepInfo(
-                    step.unique_id,
-                    step.name if step.name != step.unique_id else None,
-                    step.__class__.__name__,
-                    step.VERSION,
-                    {dep.unique_id for dep in step.dependencies},
-                    step.cache_results,
-                )
+                return StepInfo.new_from_step(step)
             else:
                 raise KeyError()
 
@@ -65,14 +62,8 @@ def step_starting(self, step: Step) -> None:
         if not step.cache_results:
             return
 
-        self.unique_id_to_info[step.unique_id] = StepInfo(
-            step.unique_id,
-            step.name if step.name != step.unique_id else None,
-            step.__class__.__name__,
-            step.VERSION,
-            {dep.unique_id for dep in step.dependencies},
-            step.cache_results,
-            datetime.now(),
+        self.unique_id_to_info[step.unique_id] = StepInfo.new_from_step(
+            step, start_time=utc_now_datetime()
         )
 
     def step_finished(self, step: Step, result: T) -> T:
@@ -83,7 +74,7 @@ def step_finished(self, step: Step, result: T) -> T:
         existing_step_info = self.unique_id_to_info[step.unique_id]
         if existing_step_info.state != StepState.RUNNING:
             raise RuntimeError(f"Step {step.name} is ending, but it never started.")
-        existing_step_info.end_time = datetime.now()
+        existing_step_info.end_time = utc_now_datetime()
 
         if step.cache_results:
             self.step_cache[step] = result
@@ -103,7 +94,7 @@ def step_failed(self, step: Step, e: BaseException) -> None:
         existing_step_info = self.unique_id_to_info[step.unique_id]
         if existing_step_info.state != StepState.RUNNING:
             raise RuntimeError(f"Step {step.name} is failing, but it never started.")
-        existing_step_info.end_time = datetime.now()
+        existing_step_info.end_time = utc_now_datetime()
         existing_step_info.error = exception_to_string(e)
 
     def register_run(self, targets: Iterable[Step], name: Optional[str] = None) -> Run:
@@ -111,17 +102,10 @@ def register_run(self, targets: Iterable[Step], name: Optional[str] = None) -> R
             name = petname.generate()
         steps: Dict[str, StepInfo] = {}
         for step in targets:
-            step_info = StepInfo(
-                step.unique_id,
-                step.name if step.name != step.unique_id else None,
-                step.__class__.__name__,
-                step.VERSION,
-                {dep.unique_id for dep in step.dependencies},
-                step.cache_results,
-            )
+            step_info = StepInfo.new_from_step(step)
             self.unique_id_to_info[step.unique_id] = step_info
             steps[step.unique_id] = step_info
-        run = Run(name, steps, datetime.now())
+        run = Run(name, steps, utc_now_datetime())
         self.runs[name] = run
         return run
 


Test patch:
diff --git a/pytest.ini b/pytest.ini
index c6c219aa3..d9e0a5038 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -8,4 +8,5 @@ markers =
 filterwarnings =
     ignore:.*Consider increasing the value of the `num_workers` argument.*:UserWarning:pytorch_lightning\.trainer\.data_loading
     ignore:.*you defined a validation_step but have no val_dataloader.*:UserWarning:pytorch_lightning\.trainer\.configuration_validator
+    ignore::UserWarning:tango\.*
 doctest_optionflags = NORMALIZE_WHITESPACE
diff --git a/tango/common/testing.py b/tango/common/testing.py
index 7463eca35..314a62704 100644
--- a/tango/common/testing.py
+++ b/tango/common/testing.py
@@ -67,6 +67,7 @@ def run(
         config: Union[PathOrStr, Dict[str, Any], Params],
         overrides: Optional[Union[Dict[str, Any], str]] = None,
         include_package: Optional[List[str]] = None,
+        workspace_url: Optional[str] = None,
         step_name: Optional[str] = None,
         parallelism: int = 1,
         multicore: bool = False,
@@ -87,7 +88,7 @@ def run(
         run_name = _run(
             TangoGlobalSettings(),
             str(config),
-            workspace_url="local://" + str(self.TEST_DIR / "workspace"),
+            workspace_url=workspace_url or "local://" + str(self.TEST_DIR / "workspace"),
             overrides=overrides,
             include_package=include_package,
             start_server=False,
diff --git a/test_fixtures/integrations/torch/__init__.py b/test_fixtures/integrations/torch/__init__.py
index 5608905c4..4caf0f803 100644
--- a/test_fixtures/integrations/torch/__init__.py
+++ b/test_fixtures/integrations/torch/__init__.py
@@ -1,9 +1,6 @@
-import torch
 import torch.nn as nn
 
-from tango.common.dataset_dict import DatasetDict
 from tango.integrations.torch import Model
-from tango.step import Step
 
 
 @Model.register("basic_regression")
diff --git a/tests/integrations/torch/train_test.py b/tests/integrations/torch/train_test.py
index 9696f0699..d14be6747 100644
--- a/tests/integrations/torch/train_test.py
+++ b/tests/integrations/torch/train_test.py
@@ -21,7 +21,7 @@ def teardown_method(self):
     @pytest.mark.parametrize("with_validation", [True, False])
     def test_basic_train(self, with_validation: bool):
         result_dir = self.run(
-            self.FIXTURES_ROOT / "integrations/torch/train.jsonnet",
+            self.FIXTURES_ROOT / "integrations" / "torch" / "train.jsonnet",
             include_package=[
                 "test_fixtures.integrations.common",
                 "test_fixtures.integrations.torch",
@@ -46,7 +46,7 @@ def test_basic_train(self, with_validation: bool):
 
     def test_basic_train_with_epochs(self):
         result_dir = self.run(
-            self.FIXTURES_ROOT / "integrations/torch/train.jsonnet",
+            self.FIXTURES_ROOT / "integrations" / "torch" / "train.jsonnet",
             include_package=[
                 "test_fixtures.integrations.common",
                 "test_fixtures.integrations.torch",
@@ -63,7 +63,7 @@ def test_basic_train_with_epochs(self):
 
     def test_basic_train_with_streaming_data(self):
         result_dir = self.run(
-            self.FIXTURES_ROOT / "integrations/torch/train.jsonnet",
+            self.FIXTURES_ROOT / "integrations" / "torch" / "train.jsonnet",
             include_package=[
                 "test_fixtures.integrations.common",
                 "test_fixtures.integrations.torch",
@@ -73,7 +73,7 @@ def test_basic_train_with_streaming_data(self):
 
     def test_train_distributed(self):
         result_dir = self.run(
-            self.FIXTURES_ROOT / "integrations/torch/train_dist.jsonnet",
+            self.FIXTURES_ROOT / "integrations" / "torch" / "train_dist.jsonnet",
             include_package=[
                 "test_fixtures.integrations.common",
                 "test_fixtures.integrations.torch",
diff --git a/tests/integrations/wandb/__init__.py b/tests/integrations/wandb/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/integrations/wandb/step_cache_test.py b/tests/integrations/wandb/step_cache_test.py
new file mode 100644
index 000000000..eb82761f6
--- /dev/null
+++ b/tests/integrations/wandb/step_cache_test.py
@@ -0,0 +1,43 @@
+import os
+import pickle
+import sys
+
+import pytest
+
+from tango import Step
+from tango.integrations.wandb import WandbStepCache
+
+WANDB_ENTITY = os.environ.get("WANDB_ENTITY", "allennlp")
+WANDB_PROJECT = "tango-workspace-testing"
+
+
+class SomeFakeStep(Step):
+    DETERMINISTIC = True
+    CACHEABLE = True
+
+    def run(self) -> int:  # type: ignore
+        return 1
+
+
+def test_step_cache_artifact_not_found():
+    step = SomeFakeStep(step_name="hi there")
+    step_cache = WandbStepCache(project=WANDB_PROJECT, entity=WANDB_ENTITY)
+    assert step not in step_cache
+
+
+@pytest.mark.parametrize(
+    "protocol",
+    [pytest.param(protocol, id=f"protocol={protocol}") for protocol in range(4)]
+    + [
+        pytest.param(
+            5,
+            id="protocol=5",
+            marks=pytest.mark.skipif(
+                sys.version_info < (3, 8), reason="Protocol 5 requires Python 3.8 or newer"
+            ),
+        ),
+    ],
+)
+def test_pickling(protocol: int):
+    step_cache = WandbStepCache(project=WANDB_PROJECT, entity=WANDB_ENTITY)
+    pickle.loads(pickle.dumps(step_cache, protocol=protocol))
diff --git a/tests/integrations/wandb/workspace_test.py b/tests/integrations/wandb/workspace_test.py
new file mode 100644
index 000000000..23ae8b9ed
--- /dev/null
+++ b/tests/integrations/wandb/workspace_test.py
@@ -0,0 +1,164 @@
+import json
+import os
+import pickle
+import shutil
+import sys
+import uuid
+
+import pytest
+import wandb
+
+from tango import Step, StepGraph, Workspace
+from tango.common import Params, util
+from tango.common.logging import initialize_logging, teardown_logging
+from tango.common.testing import TangoTestCase
+from tango.integrations.wandb import WandbWorkspace
+from tango.step_info import StepState
+from test_fixtures.package.steps import *  # noqa: F403,F401
+
+WANDB_ENTITY = os.environ.get("WANDB_ENTITY", "allennlp")
+WANDB_PROJECT = "tango-workspace-testing"
+
+
+class TestWandbWorkspace(TangoTestCase):
+    # Need to define the `setup_method()` as fixture so we can use other fixtures within it.
+    @pytest.fixture(autouse=True)
+    def setup_method(self, monkeypatch):
+        super().setup_method()
+        # Patch tango_cache_dir()
+        monkeypatch.setattr(util, "tango_cache_dir", lambda: self.TEST_DIR)
+
+    @pytest.mark.parametrize(
+        "protocol",
+        [pytest.param(protocol, id=f"protocol={protocol}") for protocol in range(4)]
+        + [
+            pytest.param(
+                5,
+                id="protocol=5",
+                marks=pytest.mark.skipif(
+                    sys.version_info < (3, 8), reason="Protocol 5 requires Python 3.8 or newer"
+                ),
+            ),
+        ],
+    )
+    def test_pickle_workspace(self, protocol):
+        workspace = WandbWorkspace(project=WANDB_PROJECT, entity=WANDB_ENTITY)
+        unpickled_workspace = pickle.loads(pickle.dumps(workspace, protocol=protocol))
+        assert unpickled_workspace.wandb_client is not None
+        assert unpickled_workspace.project == workspace.project
+        assert unpickled_workspace.entity == workspace.entity
+        assert unpickled_workspace.steps_dir == workspace.steps_dir
+
+    def test_from_url(self):
+        workspace = Workspace.from_url(f"wandb://{WANDB_ENTITY}/{WANDB_PROJECT}")
+        assert isinstance(workspace, WandbWorkspace)
+        assert workspace.entity == WANDB_ENTITY
+        assert workspace.project == WANDB_PROJECT
+
+
+class TestWandbWorkspaceUsage(TangoTestCase):
+    # Need to define the `setup_method()` as fixture so we can use other fixtures within it.
+    @pytest.fixture(autouse=True)
+    def setup_method(self, monkeypatch):
+        super().setup_method()
+        self.UNIQUE_ID_SUFFIX = os.environ.get("GITHUB_SHA", "")[:6] + "-" + str(uuid.uuid1())[:6]
+        # Patch tango_cache_dir()
+        monkeypatch.setattr(util, "tango_cache_dir", lambda: self.TEST_DIR)
+        # Patch Step unique IDs and W&B run IDs.
+        monkeypatch.setattr(Step, "_UNIQUE_ID_SUFFIX", self.UNIQUE_ID_SUFFIX)
+        monkeypatch.setattr(
+            WandbWorkspace,
+            "_generate_run_suite_id",
+            lambda workspace: wandb.util.generate_id() + "-" + self.UNIQUE_ID_SUFFIX,
+        )
+
+        self.workspace = WandbWorkspace(project=WANDB_PROJECT, entity=WANDB_ENTITY)
+
+        initialize_logging(enable_cli_logs=True)
+
+    def teardown_method(self):
+        super().teardown_method()
+
+        # Delete W&B runs and their artifacts produced by the test.
+        for wandb_run in self.workspace.wandb_client.runs(
+            f"{WANDB_ENTITY}/{WANDB_PROJECT}",
+        ):
+            if (
+                self.UNIQUE_ID_SUFFIX in wandb_run.id
+                or self.UNIQUE_ID_SUFFIX in wandb_run.config.get("_run_suite_id", "")
+            ):
+                wandb_run.delete(delete_artifacts=True)
+
+        teardown_logging()
+
+    def test_direct_usage(self):
+        params = Params.from_file(self.FIXTURES_ROOT / "experiment" / "hello_world.jsonnet")
+        step_graph = StepGraph.from_params(params.pop("steps", keep_as_dict=True))
+        tango_run = self.workspace.register_run(step for step in step_graph.values())
+
+        # Test 'registered_run()' and 'registered_runs()' methods.
+        assert self.workspace.registered_run(tango_run.name) == tango_run
+        assert self.workspace.registered_runs()[tango_run.name] == tango_run
+
+        hello_step = step_graph["hello"]
+        hello_world_step = step_graph["hello_world"]
+
+        # Test getting step info.
+        step_info = self.workspace.step_info(hello_step)
+        assert step_info.unique_id.endswith(self.UNIQUE_ID_SUFFIX)
+        assert step_info.step_name == "hello"
+        assert step_info.state == StepState.INCOMPLETE
+
+        # Mark the "hello" step as starting.
+        self.workspace.step_starting(hello_step)
+        assert self.workspace.step_info(hello_step).state == StepState.RUNNING
+
+        # Mark the "hello" step as finished.
+        self.workspace.step_finished(hello_step, "hello")
+        assert self.workspace.step_info(hello_step).state == StepState.COMPLETED
+
+        # Make sure the result is in the cache, exists locally, and on W&B.
+        cache = self.workspace.cache
+        assert hello_step in cache
+        assert cache.step_dir(hello_step).is_dir()
+        assert cache.get_step_result_artifact(hello_step) is not None
+
+        # Now make sure we can fetch the item from the cache, even if it's not in memory
+        # or in the cache directory.
+        if hello_step.unique_id in cache.weak_cache:
+            del cache.weak_cache[hello_step.unique_id]
+        if hello_step.unique_id in cache.strong_cache:
+            del cache.strong_cache[hello_step.unique_id]
+        shutil.rmtree(cache.step_dir(hello_step))
+        assert hello_step in cache
+        assert cache[hello_step] == "hello"
+
+        # Now start the "hello_world" step and then mark it as failed.
+        self.workspace.step_starting(hello_world_step)
+        self.workspace.step_failed(hello_world_step, ValueError("oh no!"))
+        assert self.workspace.step_info(hello_world_step).state == StepState.FAILED
+
+    @pytest.mark.parametrize(
+        "multicore", [pytest.param(True, id="multicore"), pytest.param(False, id="singe-core")]
+    )
+    @pytest.mark.parametrize(
+        "distributed",
+        [
+            pytest.param(True, id="distributed"),
+            pytest.param(False, id="single-device"),
+        ],
+    )
+    def test_with_wandb_train_callback(self, multicore: bool, distributed: bool):
+        self.run(
+            self.FIXTURES_ROOT
+            / "integrations"
+            / "torch"
+            / ("train.jsonnet" if not distributed else "train_dist.jsonnet"),
+            include_package=[
+                "test_fixtures.integrations.common",
+                "test_fixtures.integrations.torch",
+            ],
+            overrides=json.dumps({"steps.train.callbacks": [{"type": "wandb::log"}]}),
+            workspace_url=f"wandb://{WANDB_ENTITY}/{WANDB_PROJECT}",
+            multicore=multicore,
+        )
diff --git a/tests/step_caches/__init__.py b/tests/step_caches/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/step_caches/local_step_cache_test.py b/tests/step_caches/local_step_cache_test.py
new file mode 100644
index 000000000..76ec44c17
--- /dev/null
+++ b/tests/step_caches/local_step_cache_test.py
@@ -0,0 +1,39 @@
+import pickle
+import sys
+
+import pytest
+
+from tango.common.testing import TangoTestCase
+from tango.step import Step
+from tango.step_caches.local_step_cache import LocalStepCache
+
+
+class DummyStep(Step):
+    def run(self, x: int) -> int:  # type: ignore[override]
+        return x
+
+
+class TestLocalStepCache(TangoTestCase):
+    @pytest.mark.parametrize(
+        "protocol",
+        [pytest.param(protocol, id=f"protocol={protocol}") for protocol in range(4)]
+        + [
+            pytest.param(
+                5,
+                id="protocol=5",
+                marks=pytest.mark.skipif(
+                    sys.version_info < (3, 8), reason="Protocol 5 requires Python 3.8 or newer"
+                ),
+            ),
+        ],
+    )
+    def test_pickling(self, protocol: int):
+        step = DummyStep(step_name="dummy", x=1)
+        step_cache = LocalStepCache(self.TEST_DIR)
+        step_cache[step] = 1
+        assert step in step_cache
+        assert step.unique_id in step_cache.strong_cache
+        pickled_step_cache = pickle.dumps(step_cache, protocol=protocol)
+        unpickled_step_cache = pickle.loads(pickled_step_cache)
+        assert step.unique_id not in unpickled_step_cache.strong_cache
+        assert step in unpickled_step_cache
diff --git a/tests/workspace_test.py b/tests/workspace_test.py
deleted file mode 100644
index efe0a8fad..000000000
--- a/tests/workspace_test.py
+++ /dev/null
@@ -1,29 +0,0 @@
-from tango.workspace import StepInfo
-
-
-class CustomError(Exception):
-    def __init__(self, msg, foo, bar):
-        super().__init__(msg)
-        self.foo = foo
-        self.bar = bar
-
-
-def test_step_info_serde_with_custom_exception():
-    """
-    Tests that we can serialize and deserialize a ``StepInfo`` instance with
-    a custom exception that has positional arguments.
-    """
-    step_info = StepInfo(
-        unique_id="id-123",
-        step_name=None,
-        step_class_name=None,
-        version=None,
-        dependencies=set(),
-        cacheable=True,
-        error=CustomError("Oh no!", "foo", "bar"),
-    )
-    dump = step_info.serialize()
-    # Since 'CustomError' can't be pickled, we'll just have the string representation
-    # of it.
-    result = StepInfo.deserialize(dump)
-    assert result.error == "CustomError('Oh no!')"
diff --git a/tests/workspaces/local_workspace_test.py b/tests/workspaces/local_workspace_test.py
index 40a1f536f..b987a87c3 100644
--- a/tests/workspaces/local_workspace_test.py
+++ b/tests/workspaces/local_workspace_test.py
@@ -4,7 +4,7 @@
 
 from tango import Step
 from tango.common.testing import TangoTestCase
-from tango.workspace import StepState
+from tango.step_info import StepState
 from tango.workspaces import LocalWorkspace