Repository: mosaicml/composer

Issue:
Docstrings cleanup! `composer.core`
The `composer.core` needs docstrings. As you work through adding docstrings, please:

## Things to pay attention to
1. Ensure that each public module has a module-level docstring.
   i. _Need to make a module private? Rename it so the file (or folder) begins with an underscore._ 
2. On each public module (e.g. `__init__.py`, as well as individual public files), ensure that `__all__` is set to the members that should be re-exported. This is necessary so our API docs will properly list the members of that file, and not all imported members (such as from python's stdlib imports).
2. Ensure that each public class has a class-level docstring
  i. _Need to make a class private? Rename it so the class name begins with an underscore._
3. Ensure that each public function (except for `__init__`, which is handled by the class-level docstring) has a docstring.
  i. _Need to make a function private? Rename it so the function name begins with an underscore._
4. Double-check cross references. Double-tick marks (e.g. ``` ``code`` ```) should ONLY be used for values or function parameter names; never for attributes, functions, classes, or modules. We have [intersphinx](https://www.sphinx-doc.org/en/master/usage/extensions/intersphinx.html) configured, so this applies to imports as well. Instead, for a cross reference, use the following format:
  i. `` :attr:`MyClass.attribute` `` for attributes
  i. `` :meth:`MyClass.my_method` `` for methods
  i `` :class:`MyClass` `` for classes
  i: `` :mod:`composer.utils` `` for modules
5. Any public functions that we expect a user to invoke should have an example. Please use [doctest](https://www.sphinx-doc.org/en/master/usage/extensions/doctest.html) to structure the example codes. Doctest ensures that the examples actually run and works and do not get out-of-date.

## Previewing docs locally

```bash
cd docs
python3 -m http.server --directory _build/html/ &  # run a web server on http://localhost:8000
rm -rf source/api_reference && make clean && make html  # clean the api reference docstrings, and rebuild the docs folder
```

Then, go to http://localhost:8000/api_reference.html to preview the api reference


Discussion:
Also see https://github.com/mosaicml/composer/pull/516 -- that addresses some of the files.

Resolving patch:
diff --git a/composer/core/state.py b/composer/core/state.py
--- a/composer/core/state.py
+++ b/composer/core/state.py
@@ -58,14 +58,8 @@ class State(Serializable):
 
     .. note::
 
-        * An instance of this class is automatically constructed by the :class:`~.Trainer` constructor. A user need
-          not instantiate this class.
-
-        * To support multi-GPU training, ``State.model`` may be wrapped in :class:`torch.nn.parallel.DistributedDataParallel`,
-          and the dataloaders may be wrapped in a device-specific dataloader that handles moving tensors to device.
-
-        * :attr:`~.types.PyTorchScheduler` are wrapped in :attr:`~.ComposerScheduler`, which handles stepping either
-          stepwise or epochwise, and also properly sets up learning rate warmups.
+        An instance of this class is automatically constructed by the :class:`~.Trainer` constructor. A user need
+        not instantiate this class.
 
     Args:
         model (:attr:`~.types.Model`): The model, typically as a subclass of :class:`~.ComposerModel`.
@@ -96,9 +90,9 @@ class State(Serializable):
         loss (:attr:`~.types.Tensors`): The most recently computed loss.
         outputs (:attr:`~.types.Tensors`): The most recently computed output from the model's forward pass.
         timer (Timer): The timer that tracks training loop progress.
-        serialized_attributes (List[str]): The state of attributes in this list will be serialized in a checkpoint.
+        serialized_attributes (List[str]): The names of the attribute which are serialized in a checkpoint.
 
-            This list contains the following attributes:
+            By default, the following attributes are serialized:
 
             +-----------------------+-------------------------------------------------------------+
             | Attribute             | Description                                                 |
diff --git a/composer/core/time.py b/composer/core/time.py
--- a/composer/core/time.py
+++ b/composer/core/time.py
@@ -11,6 +11,8 @@
 
 The :class:`~.time.Time` class represents static durations of training time or points in the training process in terms
 of a specific :class:`~.time.TimeUnit` enum. This class supports comparisons, arithmetic, and conversions.
+
+See the :doc:`Time Guide </trainer/time>` for more details on tracking time during training.
 """
 from __future__ import annotations
 
@@ -60,6 +62,8 @@ class Time(Generic[TValue]):
     """Time represents static durations of training time or points in the training process in terms of a
     :class:`TimeUnit` enum (epochs, batches, samples, tokens, or duration).
 
+    See the :doc:`Time Guide </trainer/time>` for more details on tracking time during training.
+
     To construct an instance of :class:`Time`, you can either:
 
     #. Use a value followed by a :class:`TimeUnit` enum or string. For example,
@@ -354,7 +358,15 @@ def from_timestring(cls, timestring: str) -> Time:
 
 
 class Timer(Serializable):
-    """Timer tracks the current training progress, in terms of epochs, batches, samples, and tokens."""
+    """Timer tracks the current training progress, in terms of epochs, batches, samples, and tokens.
+
+    See the :doc:`Time Guide </trainer/time>` for more details on tracking time during training.
+
+    .. note::
+
+        An instance of this class is automatically constructed by the :class:`~composer.core.state.State` constructor.
+        A user need not instantiate this class.
+    """
 
     def __init__(self):
         self._epoch = Time(0, TimeUnit.EPOCH)
@@ -387,37 +399,37 @@ def load_state_dict(self, state: StateDict) -> None:
 
     @property
     def epoch(self) -> Time[int]:
-        """The current epoch."""
+        """The total epoch count."""
         return self._epoch
 
     @property
     def batch(self) -> Time[int]:
-        """The current batch."""
+        """The total batch count."""
         return self._batch
 
     @property
     def sample(self) -> Time[int]:
-        """The current sample."""
+        """The total sample count."""
         return self._sample
 
     @property
     def token(self) -> Time[int]:
-        """The current token."""
+        """The total token count."""
         return self._token
 
     @property
     def batch_in_epoch(self) -> Time[int]:
-        """The number of batches seen in the current epoch (resets at 0 at the beginning of every epoch)."""
+        """The batch count in the current epoch (resets at 0 at the beginning of every epoch)."""
         return self._batch_in_epoch
 
     @property
     def sample_in_epoch(self) -> Time[int]:
-        """The number of samples seen in the current epoch (resets at 0 at the beginning of every epoch)."""
+        """The sample count in the current epoch (resets at 0 at the beginning of every epoch)."""
         return self._sample_in_epoch
 
     @property
     def token_in_epoch(self) -> Time[int]:
-        """The number of tokens seen in the current epoch (resets at 0 at the beginning of every epoch)."""
+        """The token count in the current epoch (resets at 0 at the beginning of every epoch)."""
         return self._token_in_epoch
 
     def get(self, unit: Union[str, TimeUnit]) -> Time[int]:
@@ -554,13 +566,13 @@ class Timestamp(NamedTuple):
         :class:`Timestamp` should not be instantiated directly; instead use :meth:`Timer.get_timestamp`.
 
     Attributes:
-        epoch (Time[int]): epoch at a snapshot.
-        batch (Time[int]): batch at a snapshot.
-        batch_in_epoch (Time[int]): The number of batches seen in the epoch at a snapshot.
-        sample (Time[int]): sample at a snapshot.
-        sample_in_epoch (Time[int]): The number of samples seen in the epoch at a snapshot.
-        token (Time[int]): token at a snapshot.
-        token_in_epoch (Time[int]): The number of tokens seen in the epoch at a snapshot.
+        epoch (Time[int]): The total epoch count when the :class`Timestamp` was generated.
+        batch (Time[int]): The total batch count when the :class`Timestamp` was generated.
+        batch_in_epoch (Time[int]): The batch count in the epoch when the :class`Timestamp` was generated.
+        sample (Time[int]): The total sample count when the :class`Timestamp` was generated.
+        sample_in_epoch (Time[int]): The sample count in the epoch when the :class`Timestamp` was generated.
+        token (Time[int]): The total token count when the :class`Timestamp` was generated.
+        token_in_epoch (Time[int]): The token count in the epoch when the :class`Timestamp` was generated.
     """
     epoch: Time[int]
     batch: Time[int]


Test patch: