We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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))
```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `docs/source/doctest_cleanup.py`
Content:
```
1 # Copyright 2022 MosaicML Composer authors
2 # SPDX-License-Identifier: Apache-2.0
3 
4 """Cleanup script that is executed at the end of each doctest."""
5 
6 import os
7 
8 # variables are defined in doctest_fixtures.py
9 # pyright: reportUndefinedVariable=none
10 
11 # tmpdir and cwd were defined in doctest_fixtures.py
12 
13 os.chdir(cwd)
14 
15 tmpdir.cleanup()
16
```

Path: `composer/__init__.py`
Content:
```
1 # Copyright 2022 MosaicML Composer authors
2 # SPDX-License-Identifier: Apache-2.0
3 
4 """Composer."""
5 
6 from composer import algorithms as algorithms
7 from composer import callbacks as callbacks
8 from composer import datasets as datasets
9 from composer import functional as functional
10 from composer import loggers as loggers
11 from composer import models as models
12 from composer import optim as optim
13 from composer import profiler as profiler
14 from composer import trainer as trainer
15 from composer import utils as utils
16 from composer.core import Algorithm as Algorithm
17 from composer.core import Callback as Callback
18 from composer.core import DataSpec as DataSpec
19 from composer.core import Engine as Engine
20 from composer.core import Evaluator as Evaluator
21 from composer.core import Event as Event
22 from composer.core import State as State
23 from composer.core import Time as Time
24 from composer.core import Timestamp as Timestamp
25 from composer.core import TimeUnit as TimeUnit
26 from composer.core import types as types
27 from composer.models import ComposerModel as ComposerModel
28 from composer.trainer import Trainer as Trainer
29 
30 __version__ = '0.7.0'
31
```

Path: `docs/source/doctest_fixtures.py`
Content:
```
1 # Copyright 2022 MosaicML Composer authors
2 # SPDX-License-Identifier: Apache-2.0
3 
4 # disabling general type issues because of monkeypatching
5 #yright: reportGeneralTypeIssues=none
6 
7 """Fixtures available in doctests.
8 
9 The script is run before any doctests are executed,
10 so all imports and variables are available in any doctest.
11 The output of this setup script does not show up in the documentation.
12 """
13 import os
14 import sys
15 import tempfile
16 from typing import Any
17 from typing import Callable as Callable
18 
19 import numpy as np
20 import torch
21 import torch.optim
22 import torch.utils.data
23 from PIL import Image
24 from torch.optim.lr_scheduler import CosineAnnealingLR
25 
26 import composer
27 import composer.loggers
28 import composer.loggers.object_store_logger
29 import composer.trainer
30 import composer.trainer.trainer
31 import composer.utils
32 import composer.utils.checkpoint
33 import composer.utils.file_helpers
34 from composer import Trainer
35 from composer.core import Algorithm as Algorithm
36 from composer.core import Callback as Callback
37 from composer.core import DataSpec as DataSpec
38 from composer.core import Engine as Engine
39 from composer.core import Evaluator as Evaluator
40 from composer.core import Event as Event
41 from composer.core import State as State
42 from composer.core import Time as Time
43 from composer.core import Timestamp as Timestamp
44 from composer.core import TimeUnit as TimeUnit
45 from composer.core import types as types
46 from composer.datasets.synthetic import SyntheticBatchPairDataset
47 from composer.loggers import InMemoryLogger as InMemoryLogger
48 from composer.loggers import Logger as Logger
49 from composer.loggers import LogLevel as LogLevel
50 from composer.loggers import ObjectStoreLogger
51 from composer.models import ComposerModel as ComposerModel
52 from composer.optim.scheduler import ConstantScheduler
53 from composer.utils import LibcloudObjectStore
54 from composer.utils import ensure_tuple as ensure_tuple
55 
56 # Need to insert the repo root at the beginning of the path, since there may be other modules named `tests`
57 # Assuming that docs generation is running from the `docs` directory
58 _docs_dir = os.path.abspath('.')
59 _repo_root = os.path.dirname(_docs_dir)
60 if sys.path[0] != _repo_root:
61     sys.path.insert(0, _repo_root)
62 
63 from tests.common import SimpleModel
64 
65 # Change the cwd to be the tempfile, so we don't pollute the documentation source folder
66 tmpdir = tempfile.TemporaryDirectory()
67 cwd = os.path.abspath('.')
68 os.chdir(tmpdir.name)
69 
70 num_channels = 3
71 num_classes = 10
72 data_shape = (num_channels, 5, 5)
73 
74 Model = SimpleModel
75 
76 model = SimpleModel(num_channels, num_classes)
77 
78 optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
79 
80 scheduler = CosineAnnealingLR(optimizer, T_max=1)
81 
82 dataset = SyntheticBatchPairDataset(
83     total_dataset_size=100,
84     data_shape=data_shape,
85     num_classes=num_classes,
86     num_unique_samples_to_create=10,
87 )
88 
89 train_dataset = dataset
90 eval_dataset = dataset
91 
92 batch_size = 10
93 
94 train_dataloader = torch.utils.data.DataLoader(
95     train_dataset,
96     batch_size=batch_size,
97     num_workers=0,
98     pin_memory=False,
99     drop_last=True,
100 )
101 
102 eval_dataloader = torch.utils.data.DataLoader(
103     eval_dataset,
104     batch_size=batch_size,
105     num_workers=0,
106     pin_memory=False,
107     drop_last=False,
108 )
109 
110 state = State(
111     rank_zero_seed=0,
112     model=model,
113     run_name='run_name',
114     optimizers=optimizer,
115     grad_accum=1,
116     dataloader=train_dataloader,
117     dataloader_label='train',
118     max_duration='1ep',
119     precision='fp32',
120 )
121 
122 logger = Logger(state)
123 
124 engine = Engine(state, logger)
125 
126 image = Image.fromarray(np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8))
127 
128 # error: "randn" is not a known member of module (reportGeneralTypeIssues)
129 X_example = torch.randn(batch_size, num_channels, 32, 32)  # type: ignore
130 # error: "randn" is not a known member of module (reportGeneralTypeIssues)
131 logits = torch.randn(batch_size, num_classes)  # type: ignore
132 # error: "randint" is not a known member of module (reportGeneralTypeIssues)
133 y_example = torch.randint(num_classes, (batch_size,))  # type: ignore
134 
135 
136 def loss_fun(output, target, reduction='none'):
137     """Dummy loss function."""
138     return torch.ones_like(target)
139 
140 
141 # Patch Trainer __init__ function to replace arguments while preserving type
142 _original_trainer_init = Trainer.__init__
143 
144 
145 def _new_trainer_init(self, fake_ellipses: None = None, **kwargs: Any):
146     if 'model' not in kwargs:
147         kwargs['model'] = model
148     if 'optimizers' not in kwargs:
149         kwargs['optimizers'] = torch.optim.SGD(kwargs['model'].parameters(), lr=0.01)
150     if 'schedulers' not in kwargs:
151         kwargs['schedulers'] = ConstantScheduler()
152     if 'max_duration' not in kwargs:
153         kwargs['max_duration'] = '1ep'
154     if 'train_dataloader' not in kwargs:
155         kwargs['train_dataloader'] = train_dataloader
156     if 'eval_dataloader' not in kwargs:
157         kwargs['eval_dataloader'] = eval_dataloader
158     if 'progress_bar' not in kwargs:
159         kwargs['progress_bar'] = False  # hide tqdm logging
160     if 'log_to_console' not in kwargs:
161         kwargs['log_to_console'] = False  # hide console logging
162     _original_trainer_init(self, **kwargs)
163 
164 
165 Trainer.__init__ = _new_trainer_init
166 
167 
168 # Do not attempt to validate cloud credentials
169 def _do_not_validate(*args, **kwargs) -> None:
170     pass
171 
172 
173 composer.loggers.object_store_logger._validate_credentials = _do_not_validate  # type: ignore
174 
175 # Patch ObjectStoreLogger __init__ function to replace arguments while preserving type
176 _original_objectStoreLogger_init = ObjectStoreLogger.__init__
177 
178 
179 def _new_objectStoreLogger_init(self, fake_ellipses: None = None, **kwargs: Any):
180     os.makedirs('./object_store', exist_ok=True)
181     kwargs.update(use_procs=False,
182                   num_concurrent_uploads=1,
183                   object_store_cls=LibcloudObjectStore,
184                   object_store_kwargs={
185                       'provider': 'local',
186                       'container': '.',
187                       'provider_kwargs': {
188                           'key': os.path.abspath('./object_store'),
189                       },
190                   })
191     _original_objectStoreLogger_init(self, **kwargs)
192 
193 
194 ObjectStoreLogger.__init__ = _new_objectStoreLogger_init  # type: ignore
195 
196 # Patch ObjectStore __init__ function to replace arguments while preserving type
197 _original_libcloudObjectStore_init = LibcloudObjectStore.__init__
198 
199 
200 def _new_libcloudObjectStore_init(self, fake_ellipses: None = None, **kwargs: Any):
201     os.makedirs('./object_store', exist_ok=True)
202     kwargs.update(
203         provider='local',
204         container='.',
205         provider_kwargs={
206             'key': os.path.abspath('./object_store'),
207         },
208     )
209     _original_libcloudObjectStore_init(self, **kwargs)
210 
211 
212 LibcloudObjectStore.__init__ = _new_libcloudObjectStore_init  # type: ignore
213
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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

{"golden_diff": "diff --git a/composer/__init__.py b/composer/__init__.py\n--- a/composer/__init__.py\n+++ b/composer/__init__.py\n@@ -3,28 +3,25 @@\n \n \"\"\"Composer.\"\"\"\n \n-from composer import algorithms as algorithms\n-from composer import callbacks as callbacks\n-from composer import datasets as datasets\n-from composer import functional as functional\n-from composer import loggers as loggers\n-from composer import models as models\n-from composer import optim as optim\n-from composer import profiler as profiler\n-from composer import trainer as trainer\n-from composer import utils as utils\n-from composer.core import Algorithm as Algorithm\n-from composer.core import Callback as Callback\n-from composer.core import DataSpec as DataSpec\n-from composer.core import Engine as Engine\n-from composer.core import Evaluator as Evaluator\n-from composer.core import Event as Event\n-from composer.core import State as State\n-from composer.core import Time as Time\n-from composer.core import Timestamp as Timestamp\n-from composer.core import TimeUnit as TimeUnit\n-from composer.core import types as types\n-from composer.models import ComposerModel as ComposerModel\n-from composer.trainer import Trainer as Trainer\n+from composer.core import Algorithm, Callback, DataSpec, Engine, Evaluator, Event, State, Time, Timestamp, TimeUnit\n+from composer.loggers import Logger\n+from composer.models import ComposerModel\n+from composer.trainer import Trainer\n+\n+__all__ = [\n+    'Algorithm',\n+    'Callback',\n+    'DataSpec',\n+    'Engine',\n+    'Evaluator',\n+    'Event',\n+    'State',\n+    'Time',\n+    'Timestamp',\n+    'TimeUnit',\n+    'Logger',\n+    'ComposerModel',\n+    'Trainer',\n+]\n \n __version__ = '0.7.0'\ndiff --git a/docs/source/doctest_cleanup.py b/docs/source/doctest_cleanup.py\n--- a/docs/source/doctest_cleanup.py\n+++ b/docs/source/doctest_cleanup.py\n@@ -4,6 +4,7 @@\n \"\"\"Cleanup script that is executed at the end of each doctest.\"\"\"\n \n import os\n+import shutil\n \n # variables are defined in doctest_fixtures.py\n # pyright: reportUndefinedVariable=none\n@@ -12,4 +13,7 @@\n \n os.chdir(cwd)\n \n-tmpdir.cleanup()\n+try:\n+    shutil.rmtree(tmpdir)\n+except OSError:\n+    pass\ndiff --git a/docs/source/doctest_fixtures.py b/docs/source/doctest_fixtures.py\n--- a/docs/source/doctest_fixtures.py\n+++ b/docs/source/doctest_fixtures.py\n@@ -63,9 +63,9 @@\n from tests.common import SimpleModel\n \n # Change the cwd to be the tempfile, so we don't pollute the documentation source folder\n-tmpdir = tempfile.TemporaryDirectory()\n+tmpdir = tempfile.mkdtemp()\n cwd = os.path.abspath('.')\n-os.chdir(tmpdir.name)\n+os.chdir(tmpdir)\n \n num_channels = 3\n num_classes = 10\n", "issue": "Consistency in docs [State, time, callbacks]\nI am re-writing a callback that was working on composer 0.5 and now it does not.\r\nI 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`)\r\nI am having a hard time creating a custom callback, maybe a tutorial with the new API?\r\n\r\nThe callback is for logging images to wandb:\r\n```python\r\nimport wandb\r\nfrom composer import Callback, State, Logger\r\n\r\nclass LogPredictions(Callback):\r\n    def __init__(self, num_samples=100, seed=1234):\r\n        super().__init__()\r\n        self.num_samples = num_samples\r\n        self.data = []\r\n        \r\n    def eval_batch_end(self, state: State, logger: Logger):\r\n        \"\"\"Compute predictions per batch and stores them on self.data\"\"\"\r\n        if state.epoch == state.max_duration: #on last val epoch\r\n            if len(self.data) < self.num_samples:\r\n                n = self.num_samples\r\n                x, y = state.batch_pair\r\n                outputs = state.outputs.argmax(-1)\r\n                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]))]\r\n                self.data += data\r\n            \r\n    def eval_end(self, state: State, logger: Logger):\r\n        \"Create a wandb.Table and logs it\"\r\n        columns = ['image', 'ground truth', 'prediction']\r\n        table = wandb.Table(columns=columns, data=self.data[:self.num_samples])\r\n        wandb.log({'sample_table':table}, step=int(state.timer.batch))\r\n```\r\n\n", "before_files": [{"content": "# Copyright 2022 MosaicML Composer authors\n# SPDX-License-Identifier: Apache-2.0\n\n\"\"\"Cleanup script that is executed at the end of each doctest.\"\"\"\n\nimport os\n\n# variables are defined in doctest_fixtures.py\n# pyright: reportUndefinedVariable=none\n\n# tmpdir and cwd were defined in doctest_fixtures.py\n\nos.chdir(cwd)\n\ntmpdir.cleanup()\n", "path": "docs/source/doctest_cleanup.py"}, {"content": "# Copyright 2022 MosaicML Composer authors\n# SPDX-License-Identifier: Apache-2.0\n\n\"\"\"Composer.\"\"\"\n\nfrom composer import algorithms as algorithms\nfrom composer import callbacks as callbacks\nfrom composer import datasets as datasets\nfrom composer import functional as functional\nfrom composer import loggers as loggers\nfrom composer import models as models\nfrom composer import optim as optim\nfrom composer import profiler as profiler\nfrom composer import trainer as trainer\nfrom composer import utils as utils\nfrom composer.core import Algorithm as Algorithm\nfrom composer.core import Callback as Callback\nfrom composer.core import DataSpec as DataSpec\nfrom composer.core import Engine as Engine\nfrom composer.core import Evaluator as Evaluator\nfrom composer.core import Event as Event\nfrom composer.core import State as State\nfrom composer.core import Time as Time\nfrom composer.core import Timestamp as Timestamp\nfrom composer.core import TimeUnit as TimeUnit\nfrom composer.core import types as types\nfrom composer.models import ComposerModel as ComposerModel\nfrom composer.trainer import Trainer as Trainer\n\n__version__ = '0.7.0'\n", "path": "composer/__init__.py"}, {"content": "# Copyright 2022 MosaicML Composer authors\n# SPDX-License-Identifier: Apache-2.0\n\n# disabling general type issues because of monkeypatching\n#yright: reportGeneralTypeIssues=none\n\n\"\"\"Fixtures available in doctests.\n\nThe script is run before any doctests are executed,\nso all imports and variables are available in any doctest.\nThe output of this setup script does not show up in the documentation.\n\"\"\"\nimport os\nimport sys\nimport tempfile\nfrom typing import Any\nfrom typing import Callable as Callable\n\nimport numpy as np\nimport torch\nimport torch.optim\nimport torch.utils.data\nfrom PIL import Image\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\n\nimport composer\nimport composer.loggers\nimport composer.loggers.object_store_logger\nimport composer.trainer\nimport composer.trainer.trainer\nimport composer.utils\nimport composer.utils.checkpoint\nimport composer.utils.file_helpers\nfrom composer import Trainer\nfrom composer.core import Algorithm as Algorithm\nfrom composer.core import Callback as Callback\nfrom composer.core import DataSpec as DataSpec\nfrom composer.core import Engine as Engine\nfrom composer.core import Evaluator as Evaluator\nfrom composer.core import Event as Event\nfrom composer.core import State as State\nfrom composer.core import Time as Time\nfrom composer.core import Timestamp as Timestamp\nfrom composer.core import TimeUnit as TimeUnit\nfrom composer.core import types as types\nfrom composer.datasets.synthetic import SyntheticBatchPairDataset\nfrom composer.loggers import InMemoryLogger as InMemoryLogger\nfrom composer.loggers import Logger as Logger\nfrom composer.loggers import LogLevel as LogLevel\nfrom composer.loggers import ObjectStoreLogger\nfrom composer.models import ComposerModel as ComposerModel\nfrom composer.optim.scheduler import ConstantScheduler\nfrom composer.utils import LibcloudObjectStore\nfrom composer.utils import ensure_tuple as ensure_tuple\n\n# Need to insert the repo root at the beginning of the path, since there may be other modules named `tests`\n# Assuming that docs generation is running from the `docs` directory\n_docs_dir = os.path.abspath('.')\n_repo_root = os.path.dirname(_docs_dir)\nif sys.path[0] != _repo_root:\n    sys.path.insert(0, _repo_root)\n\nfrom tests.common import SimpleModel\n\n# Change the cwd to be the tempfile, so we don't pollute the documentation source folder\ntmpdir = tempfile.TemporaryDirectory()\ncwd = os.path.abspath('.')\nos.chdir(tmpdir.name)\n\nnum_channels = 3\nnum_classes = 10\ndata_shape = (num_channels, 5, 5)\n\nModel = SimpleModel\n\nmodel = SimpleModel(num_channels, num_classes)\n\noptimizer = torch.optim.SGD(model.parameters(), lr=0.001)\n\nscheduler = CosineAnnealingLR(optimizer, T_max=1)\n\ndataset = SyntheticBatchPairDataset(\n    total_dataset_size=100,\n    data_shape=data_shape,\n    num_classes=num_classes,\n    num_unique_samples_to_create=10,\n)\n\ntrain_dataset = dataset\neval_dataset = dataset\n\nbatch_size = 10\n\ntrain_dataloader = torch.utils.data.DataLoader(\n    train_dataset,\n    batch_size=batch_size,\n    num_workers=0,\n    pin_memory=False,\n    drop_last=True,\n)\n\neval_dataloader = torch.utils.data.DataLoader(\n    eval_dataset,\n    batch_size=batch_size,\n    num_workers=0,\n    pin_memory=False,\n    drop_last=False,\n)\n\nstate = State(\n    rank_zero_seed=0,\n    model=model,\n    run_name='run_name',\n    optimizers=optimizer,\n    grad_accum=1,\n    dataloader=train_dataloader,\n    dataloader_label='train',\n    max_duration='1ep',\n    precision='fp32',\n)\n\nlogger = Logger(state)\n\nengine = Engine(state, logger)\n\nimage = Image.fromarray(np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8))\n\n# error: \"randn\" is not a known member of module (reportGeneralTypeIssues)\nX_example = torch.randn(batch_size, num_channels, 32, 32)  # type: ignore\n# error: \"randn\" is not a known member of module (reportGeneralTypeIssues)\nlogits = torch.randn(batch_size, num_classes)  # type: ignore\n# error: \"randint\" is not a known member of module (reportGeneralTypeIssues)\ny_example = torch.randint(num_classes, (batch_size,))  # type: ignore\n\n\ndef loss_fun(output, target, reduction='none'):\n    \"\"\"Dummy loss function.\"\"\"\n    return torch.ones_like(target)\n\n\n# Patch Trainer __init__ function to replace arguments while preserving type\n_original_trainer_init = Trainer.__init__\n\n\ndef _new_trainer_init(self, fake_ellipses: None = None, **kwargs: Any):\n    if 'model' not in kwargs:\n        kwargs['model'] = model\n    if 'optimizers' not in kwargs:\n        kwargs['optimizers'] = torch.optim.SGD(kwargs['model'].parameters(), lr=0.01)\n    if 'schedulers' not in kwargs:\n        kwargs['schedulers'] = ConstantScheduler()\n    if 'max_duration' not in kwargs:\n        kwargs['max_duration'] = '1ep'\n    if 'train_dataloader' not in kwargs:\n        kwargs['train_dataloader'] = train_dataloader\n    if 'eval_dataloader' not in kwargs:\n        kwargs['eval_dataloader'] = eval_dataloader\n    if 'progress_bar' not in kwargs:\n        kwargs['progress_bar'] = False  # hide tqdm logging\n    if 'log_to_console' not in kwargs:\n        kwargs['log_to_console'] = False  # hide console logging\n    _original_trainer_init(self, **kwargs)\n\n\nTrainer.__init__ = _new_trainer_init\n\n\n# Do not attempt to validate cloud credentials\ndef _do_not_validate(*args, **kwargs) -> None:\n    pass\n\n\ncomposer.loggers.object_store_logger._validate_credentials = _do_not_validate  # type: ignore\n\n# Patch ObjectStoreLogger __init__ function to replace arguments while preserving type\n_original_objectStoreLogger_init = ObjectStoreLogger.__init__\n\n\ndef _new_objectStoreLogger_init(self, fake_ellipses: None = None, **kwargs: Any):\n    os.makedirs('./object_store', exist_ok=True)\n    kwargs.update(use_procs=False,\n                  num_concurrent_uploads=1,\n                  object_store_cls=LibcloudObjectStore,\n                  object_store_kwargs={\n                      'provider': 'local',\n                      'container': '.',\n                      'provider_kwargs': {\n                          'key': os.path.abspath('./object_store'),\n                      },\n                  })\n    _original_objectStoreLogger_init(self, **kwargs)\n\n\nObjectStoreLogger.__init__ = _new_objectStoreLogger_init  # type: ignore\n\n# Patch ObjectStore __init__ function to replace arguments while preserving type\n_original_libcloudObjectStore_init = LibcloudObjectStore.__init__\n\n\ndef _new_libcloudObjectStore_init(self, fake_ellipses: None = None, **kwargs: Any):\n    os.makedirs('./object_store', exist_ok=True)\n    kwargs.update(\n        provider='local',\n        container='.',\n        provider_kwargs={\n            'key': os.path.abspath('./object_store'),\n        },\n    )\n    _original_libcloudObjectStore_init(self, **kwargs)\n\n\nLibcloudObjectStore.__init__ = _new_libcloudObjectStore_init  # type: ignore\n", "path": "docs/source/doctest_fixtures.py"}], "after_files": [{"content": "# Copyright 2022 MosaicML Composer authors\n# SPDX-License-Identifier: Apache-2.0\n\n\"\"\"Cleanup script that is executed at the end of each doctest.\"\"\"\n\nimport os\nimport shutil\n\n# variables are defined in doctest_fixtures.py\n# pyright: reportUndefinedVariable=none\n\n# tmpdir and cwd were defined in doctest_fixtures.py\n\nos.chdir(cwd)\n\ntry:\n    shutil.rmtree(tmpdir)\nexcept OSError:\n    pass\n", "path": "docs/source/doctest_cleanup.py"}, {"content": "# Copyright 2022 MosaicML Composer authors\n# SPDX-License-Identifier: Apache-2.0\n\n\"\"\"Composer.\"\"\"\n\nfrom composer.core import Algorithm, Callback, DataSpec, Engine, Evaluator, Event, State, Time, Timestamp, TimeUnit\nfrom composer.loggers import Logger\nfrom composer.models import ComposerModel\nfrom composer.trainer import Trainer\n\n__all__ = [\n    'Algorithm',\n    'Callback',\n    'DataSpec',\n    'Engine',\n    'Evaluator',\n    'Event',\n    'State',\n    'Time',\n    'Timestamp',\n    'TimeUnit',\n    'Logger',\n    'ComposerModel',\n    'Trainer',\n]\n\n__version__ = '0.7.0'\n", "path": "composer/__init__.py"}, {"content": "# Copyright 2022 MosaicML Composer authors\n# SPDX-License-Identifier: Apache-2.0\n\n# disabling general type issues because of monkeypatching\n#yright: reportGeneralTypeIssues=none\n\n\"\"\"Fixtures available in doctests.\n\nThe script is run before any doctests are executed,\nso all imports and variables are available in any doctest.\nThe output of this setup script does not show up in the documentation.\n\"\"\"\nimport os\nimport sys\nimport tempfile\nfrom typing import Any\nfrom typing import Callable as Callable\n\nimport numpy as np\nimport torch\nimport torch.optim\nimport torch.utils.data\nfrom PIL import Image\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\n\nimport composer\nimport composer.loggers\nimport composer.loggers.object_store_logger\nimport composer.trainer\nimport composer.trainer.trainer\nimport composer.utils\nimport composer.utils.checkpoint\nimport composer.utils.file_helpers\nfrom composer import Trainer\nfrom composer.core import Algorithm as Algorithm\nfrom composer.core import Callback as Callback\nfrom composer.core import DataSpec as DataSpec\nfrom composer.core import Engine as Engine\nfrom composer.core import Evaluator as Evaluator\nfrom composer.core import Event as Event\nfrom composer.core import State as State\nfrom composer.core import Time as Time\nfrom composer.core import Timestamp as Timestamp\nfrom composer.core import TimeUnit as TimeUnit\nfrom composer.core import types as types\nfrom composer.datasets.synthetic import SyntheticBatchPairDataset\nfrom composer.loggers import InMemoryLogger as InMemoryLogger\nfrom composer.loggers import Logger as Logger\nfrom composer.loggers import LogLevel as LogLevel\nfrom composer.loggers import ObjectStoreLogger\nfrom composer.models import ComposerModel as ComposerModel\nfrom composer.optim.scheduler import ConstantScheduler\nfrom composer.utils import LibcloudObjectStore\nfrom composer.utils import ensure_tuple as ensure_tuple\n\n# Need to insert the repo root at the beginning of the path, since there may be other modules named `tests`\n# Assuming that docs generation is running from the `docs` directory\n_docs_dir = os.path.abspath('.')\n_repo_root = os.path.dirname(_docs_dir)\nif sys.path[0] != _repo_root:\n    sys.path.insert(0, _repo_root)\n\nfrom tests.common import SimpleModel\n\n# Change the cwd to be the tempfile, so we don't pollute the documentation source folder\ntmpdir = tempfile.mkdtemp()\ncwd = os.path.abspath('.')\nos.chdir(tmpdir)\n\nnum_channels = 3\nnum_classes = 10\ndata_shape = (num_channels, 5, 5)\n\nModel = SimpleModel\n\nmodel = SimpleModel(num_channels, num_classes)\n\noptimizer = torch.optim.SGD(model.parameters(), lr=0.001)\n\nscheduler = CosineAnnealingLR(optimizer, T_max=1)\n\ndataset = SyntheticBatchPairDataset(\n    total_dataset_size=100,\n    data_shape=data_shape,\n    num_classes=num_classes,\n    num_unique_samples_to_create=10,\n)\n\ntrain_dataset = dataset\neval_dataset = dataset\n\nbatch_size = 10\n\ntrain_dataloader = torch.utils.data.DataLoader(\n    train_dataset,\n    batch_size=batch_size,\n    num_workers=0,\n    pin_memory=False,\n    drop_last=True,\n)\n\neval_dataloader = torch.utils.data.DataLoader(\n    eval_dataset,\n    batch_size=batch_size,\n    num_workers=0,\n    pin_memory=False,\n    drop_last=False,\n)\n\nstate = State(\n    rank_zero_seed=0,\n    model=model,\n    run_name='run_name',\n    optimizers=optimizer,\n    grad_accum=1,\n    dataloader=train_dataloader,\n    dataloader_label='train',\n    max_duration='1ep',\n    precision='fp32',\n)\n\nlogger = Logger(state)\n\nengine = Engine(state, logger)\n\nimage = Image.fromarray(np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8))\n\n# error: \"randn\" is not a known member of module (reportGeneralTypeIssues)\nX_example = torch.randn(batch_size, num_channels, 32, 32)  # type: ignore\n# error: \"randn\" is not a known member of module (reportGeneralTypeIssues)\nlogits = torch.randn(batch_size, num_classes)  # type: ignore\n# error: \"randint\" is not a known member of module (reportGeneralTypeIssues)\ny_example = torch.randint(num_classes, (batch_size,))  # type: ignore\n\n\ndef loss_fun(output, target, reduction='none'):\n    \"\"\"Dummy loss function.\"\"\"\n    return torch.ones_like(target)\n\n\n# Patch Trainer __init__ function to replace arguments while preserving type\n_original_trainer_init = Trainer.__init__\n\n\ndef _new_trainer_init(self, fake_ellipses: None = None, **kwargs: Any):\n    if 'model' not in kwargs:\n        kwargs['model'] = model\n    if 'optimizers' not in kwargs:\n        kwargs['optimizers'] = torch.optim.SGD(kwargs['model'].parameters(), lr=0.01)\n    if 'schedulers' not in kwargs:\n        kwargs['schedulers'] = ConstantScheduler()\n    if 'max_duration' not in kwargs:\n        kwargs['max_duration'] = '1ep'\n    if 'train_dataloader' not in kwargs:\n        kwargs['train_dataloader'] = train_dataloader\n    if 'eval_dataloader' not in kwargs:\n        kwargs['eval_dataloader'] = eval_dataloader\n    if 'progress_bar' not in kwargs:\n        kwargs['progress_bar'] = False  # hide tqdm logging\n    if 'log_to_console' not in kwargs:\n        kwargs['log_to_console'] = False  # hide console logging\n    _original_trainer_init(self, **kwargs)\n\n\nTrainer.__init__ = _new_trainer_init\n\n\n# Do not attempt to validate cloud credentials\ndef _do_not_validate(*args, **kwargs) -> None:\n    pass\n\n\ncomposer.loggers.object_store_logger._validate_credentials = _do_not_validate  # type: ignore\n\n# Patch ObjectStoreLogger __init__ function to replace arguments while preserving type\n_original_objectStoreLogger_init = ObjectStoreLogger.__init__\n\n\ndef _new_objectStoreLogger_init(self, fake_ellipses: None = None, **kwargs: Any):\n    os.makedirs('./object_store', exist_ok=True)\n    kwargs.update(use_procs=False,\n                  num_concurrent_uploads=1,\n                  object_store_cls=LibcloudObjectStore,\n                  object_store_kwargs={\n                      'provider': 'local',\n                      'container': '.',\n                      'provider_kwargs': {\n                          'key': os.path.abspath('./object_store'),\n                      },\n                  })\n    _original_objectStoreLogger_init(self, **kwargs)\n\n\nObjectStoreLogger.__init__ = _new_objectStoreLogger_init  # type: ignore\n\n# Patch ObjectStore __init__ function to replace arguments while preserving type\n_original_libcloudObjectStore_init = LibcloudObjectStore.__init__\n\n\ndef _new_libcloudObjectStore_init(self, fake_ellipses: None = None, **kwargs: Any):\n    os.makedirs('./object_store', exist_ok=True)\n    kwargs.update(\n        provider='local',\n        container='.',\n        provider_kwargs={\n            'key': os.path.abspath('./object_store'),\n        },\n    )\n    _original_libcloudObjectStore_init(self, **kwargs)\n\n\nLibcloudObjectStore.__init__ = _new_libcloudObjectStore_init  # type: ignore\n", "path": "docs/source/doctest_fixtures.py"}]}