Repository: mosaicml/composer

Issue:
Consistency in docs [State, time, callbacks]
I am re-writing a callback that was working on composer 0.5 and now it does not.
I see that there is a new class `time` and that `state.epoch` is not there anymore, the problem is that the docs still point to that. Like [here](https://docs.mosaicml.com/en/v0.7.1/trainer/logging.html#logging-additional-data) or that `Logger` cannot be imported from `composer` (`from composer.loggers import Logger`)
I am having a hard time creating a custom callback, maybe a tutorial with the new API?

The callback is for logging images to wandb:
```python
import wandb
from composer import Callback, State, Logger

class LogPredictions(Callback):
    def __init__(self, num_samples=100, seed=1234):
        super().__init__()
        self.num_samples = num_samples
        self.data = []
        
    def eval_batch_end(self, state: State, logger: Logger):
        """Compute predictions per batch and stores them on self.data"""
        if state.epoch == state.max_duration: #on last val epoch
            if len(self.data) < self.num_samples:
                n = self.num_samples
                x, y = state.batch_pair
                outputs = state.outputs.argmax(-1)
                data = [[wandb.Image(x_i), y_i, y_pred] for x_i, y_i, y_pred in list(zip(x[:n], y[:n], outputs[:n]))]
                self.data += data
            
    def eval_end(self, state: State, logger: Logger):
        "Create a wandb.Table and logs it"
        columns = ['image', 'ground truth', 'prediction']
        table = wandb.Table(columns=columns, data=self.data[:self.num_samples])
        wandb.log({'sample_table':table}, step=int(state.timer.batch))
```



Discussion:
Hi there!

We removed `state.epoch`, `state.step`, and `state.timer`. Instead, for the epoch count, could you try using `int(state.timestamp.epoch)`, and for the batch count, `int(state.timestamp.batch)`? See the [Timestamp API Reference](https://docs.mosaicml.com/en/v0.7.1/api_reference/composer.core.time.html#composer.core.time.Timestamp) for more info.

We'll add the re-export of the `Logger` to composer to the next release; but for now, yes, could you use `from composer.loggers import Logger`.

Sorry for the confusion but thanks for letting us know -- we'll get it fixed!


Resolving patch:
diff --git a/composer/__init__.py b/composer/__init__.py
--- a/composer/__init__.py
+++ b/composer/__init__.py
@@ -3,28 +3,25 @@
 
 """Composer."""
 
-from composer import algorithms as algorithms
-from composer import callbacks as callbacks
-from composer import datasets as datasets
-from composer import functional as functional
-from composer import loggers as loggers
-from composer import models as models
-from composer import optim as optim
-from composer import profiler as profiler
-from composer import trainer as trainer
-from composer import utils as utils
-from composer.core import Algorithm as Algorithm
-from composer.core import Callback as Callback
-from composer.core import DataSpec as DataSpec
-from composer.core import Engine as Engine
-from composer.core import Evaluator as Evaluator
-from composer.core import Event as Event
-from composer.core import State as State
-from composer.core import Time as Time
-from composer.core import Timestamp as Timestamp
-from composer.core import TimeUnit as TimeUnit
-from composer.core import types as types
-from composer.models import ComposerModel as ComposerModel
-from composer.trainer import Trainer as Trainer
+from composer.core import Algorithm, Callback, DataSpec, Engine, Evaluator, Event, State, Time, Timestamp, TimeUnit
+from composer.loggers import Logger
+from composer.models import ComposerModel
+from composer.trainer import Trainer
+
+__all__ = [
+    'Algorithm',
+    'Callback',
+    'DataSpec',
+    'Engine',
+    'Evaluator',
+    'Event',
+    'State',
+    'Time',
+    'Timestamp',
+    'TimeUnit',
+    'Logger',
+    'ComposerModel',
+    'Trainer',
+]
 
 __version__ = '0.7.0'
diff --git a/docs/source/doctest_cleanup.py b/docs/source/doctest_cleanup.py
--- a/docs/source/doctest_cleanup.py
+++ b/docs/source/doctest_cleanup.py
@@ -4,6 +4,7 @@
 """Cleanup script that is executed at the end of each doctest."""
 
 import os
+import shutil
 
 # variables are defined in doctest_fixtures.py
 # pyright: reportUndefinedVariable=none
@@ -12,4 +13,7 @@
 
 os.chdir(cwd)
 
-tmpdir.cleanup()
+try:
+    shutil.rmtree(tmpdir)
+except OSError:
+    pass
diff --git a/docs/source/doctest_fixtures.py b/docs/source/doctest_fixtures.py
--- a/docs/source/doctest_fixtures.py
+++ b/docs/source/doctest_fixtures.py
@@ -63,9 +63,9 @@
 from tests.common import SimpleModel
 
 # Change the cwd to be the tempfile, so we don't pollute the documentation source folder
-tmpdir = tempfile.TemporaryDirectory()
+tmpdir = tempfile.mkdtemp()
 cwd = os.path.abspath('.')
-os.chdir(tmpdir.name)
+os.chdir(tmpdir)
 
 num_channels = 3
 num_classes = 10


Test patch:
diff --git a/tests/algorithms/algorithm_settings.py b/tests/algorithms/algorithm_settings.py
--- a/tests/algorithms/algorithm_settings.py
+++ b/tests/algorithms/algorithm_settings.py
@@ -14,6 +14,7 @@
 from torch.utils.data import Dataset
 
 import composer
+import composer.algorithms
 from composer import Algorithm
 from composer.algorithms import (EMA, SAM, SWA, Alibi, AugMix, BlurPool, ChannelsLast, ColOut, CutMix, CutOut,
                                  Factorize, GhostBatchNorm, GradientClipping, LabelSmoothing, LayerFreezing, MixUp,