We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Configure Jenkins
Enable github actions for:

- [x] pytest CPU runner
- [x] formatting and type checking (`yapf`, `pyright`)
- [ ] docker builds
- [x] docs builds

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `composer/callbacks/profiler.py`
Content:
```
1 from __future__ import annotations
2 
3 import atexit
4 from dataclasses import asdict, dataclass
5 from typing import TYPE_CHECKING, Optional
6 
7 import torch.profiler
8 import yahp as hp
9 from torch.profiler.profiler import ProfilerAction
10 
11 from composer import Callback
12 from composer.callbacks.callback_hparams import CallbackHparams
13 from composer.core.types import StateDict
14 from composer.utils.ddp import get_global_rank
15 
16 if TYPE_CHECKING:
17     from composer.core import Logger, State
18 
19 _PROFILE_MISSING_ERROR = "The profiler has not been setup. Please call profiler.training_start() before training starts."
20 
21 
22 @dataclass
23 class TorchProfilerHparams(CallbackHparams):
24     tensorboard_trace_handler_dir: str = hp.required("directory to store trace results")
25     tensorboard_use_gzip: bool = hp.optional("use gzip for trace", default=False)
26     record_shapes: bool = hp.optional(doc="record_shapes", default=True)
27     profile_memory: bool = hp.optional(doc="track tensor memory allocations and frees", default=False)
28     with_stack: bool = hp.optional(doc="record stack info", default=True)
29     with_flops: bool = hp.optional(doc="estimate flops for operators", default=True)
30 
31     skip: int = hp.optional("Number of batches to skip at epoch start", default=0)
32     warmup: int = hp.optional("Number of warmup batches in a cycle", default=1)
33     active: int = hp.optional("Number of batches to profile in a cycle", default=5)
34     wait: int = hp.optional("Number of batches to skip at the end of each cycle", default=0)
35 
36     def initialize_object(self) -> TorchProfiler:
37         return TorchProfiler(**asdict(self))
38 
39 
40 @dataclass
41 class TorchProfilerState:
42     batch_in_epoch: int = 0
43     batches_per_epoch: int = 0
44 
45 
46 class TorchProfiler(Callback):
47 
48     def __init__(
49         self,
50         *,
51         skip: int,
52         warmup: int,
53         active: int,
54         wait: int,
55         tensorboard_trace_handler_dir: str,
56         tensorboard_use_gzip: bool,
57         record_shapes: bool,
58         profile_memory: bool,
59         with_stack: bool,
60         with_flops: bool,
61     ) -> None:
62         super().__init__()
63         self.hparams = TorchProfilerHparams(
64             skip=skip,
65             warmup=warmup,
66             active=active,
67             wait=wait,
68             tensorboard_trace_handler_dir=tensorboard_trace_handler_dir,
69             tensorboard_use_gzip=tensorboard_use_gzip,
70             record_shapes=record_shapes,
71             profile_memory=profile_memory,
72             with_stack=with_stack,
73             with_flops=with_flops,
74         )
75         self.profiler: Optional[torch.profiler.profile] = None
76         self.profiler_state: TorchProfilerState = TorchProfilerState()
77         self._torch_profiler_scheduler = torch.profiler.profiler.schedule(
78             wait=self.hparams.wait,
79             warmup=self.hparams.warmup,
80             active=self.hparams.active,
81             skip_first=self.hparams.skip,
82         )
83 
84     def state_dict(self) -> StateDict:
85         return asdict(self.profiler_state)
86 
87     def load_state_dict(self, state: StateDict) -> None:
88         self.profiler_state = TorchProfilerState(**state)
89 
90     def scheduler_fn(self, profiler_step: int) -> ProfilerAction:
91         # Invoked on every batch, at the batch end
92         # But, it's called one batch in advance.
93         # Wrapping the default scheduling function to deal with epoch boundaries
94         # Giving the torch scheduler the batch in the epoch, not the global step
95 
96         # adding 1 since this is called before the step is incremented
97         next_batch_in_epoch = self.profiler_state.batch_in_epoch + 1
98         if profiler_step == 0:
99             next_batch_in_epoch = 0
100         torch_scheduler_action = self._torch_profiler_scheduler(next_batch_in_epoch)
101         if next_batch_in_epoch == self.profiler_state.batches_per_epoch:
102             if torch_scheduler_action == ProfilerAction.RECORD:
103                 # force saving at epoch boundaries
104                 torch_scheduler_action = ProfilerAction.RECORD_AND_SAVE
105         return torch_scheduler_action
106 
107     def training_start(self, state: State, logger: Logger) -> None:
108         assert self.profiler is None, _PROFILE_MISSING_ERROR
109         self.profiler = torch.profiler.profile(
110             schedule=self.scheduler_fn,
111             on_trace_ready=torch.profiler.tensorboard_trace_handler(
112                 self.hparams.tensorboard_trace_handler_dir,
113                 worker_name=str(get_global_rank()),
114                 use_gzip=self.hparams.tensorboard_use_gzip,
115             ),
116             activities=None,  # auto-set
117             record_shapes=self.hparams.record_shapes,
118             profile_memory=self.hparams.profile_memory,
119             with_stack=self.hparams.with_stack,
120             with_flops=self.hparams.with_flops,
121         )
122         self.profiler.__enter__()
123         atexit.register(self._close_profiler)
124 
125     def batch_end(self, state: State, logger: Logger) -> None:
126         assert self.profiler is not None, _PROFILE_MISSING_ERROR
127         self.profiler.step()
128 
129     def epoch_start(self, state: State, logger: Logger) -> None:
130         self.profiler_state.batches_per_epoch = state.steps_per_epoch
131 
132     def batch_start(self, state: State, logger: Logger) -> None:
133         self.profiler_state.batch_in_epoch = state.batch_idx
134         assert self.profiler is not None, _PROFILE_MISSING_ERROR
135         logger.metric_batch({"profiler/state": self.profiler.current_action.name})
136 
137     def _close_profiler(self) -> None:
138         assert self.profiler is not None
139         self.profiler.__exit__(None, None, None)
140
```

Path: `composer/callbacks/timing_monitor.py`
Content:
```
1 from __future__ import annotations
2 
3 import logging
4 import time
5 from dataclasses import dataclass
6 from typing import List
7 
8 import yahp as hp
9 
10 from composer import Logger, State
11 from composer.callbacks.callback_hparams import CallbackHparams
12 from composer.core.callback import Callback
13 from composer.core.types import BreakEpochException
14 
15 log = logging.getLogger(__name__)
16 
17 
18 class TimingMonitor(Callback):
19 
20     def __init__(self, min_steps: int, epoch_list: List[int], step_list: List[int], all_epochs: bool):
21         super().__init__()
22         self.hparams = TimingMonitorHparams(min_steps=min_steps,
23                                             epoch_list=epoch_list,
24                                             step_list=step_list,
25                                             all_epochs=all_epochs)
26 
27         if not all_epochs:
28             assert len(epoch_list) > 0, "'epoch_list'  must be non-empty."
29             assert 0 in epoch_list, \
30                 "'epoch_list' must contain 0, otherwise first K epochs have unknown speed"
31         assert len(step_list) > 0, "'step_list'  must be non-empty."
32         assert 0 in step_list, \
33             "'step_list' must contain 0 because `EPOCH_START` requires batch_idx of 0"
34 
35         # Sort lists so that time only moves forward
36         epoch_list = list(sorted(epoch_list))
37         step_list = list(sorted(step_list))
38 
39         self.current_time = None
40         self.profile_examples = 0
41         self.profile_steps = 0
42         self.profile_time = 0
43         self.wall_clock_train = 0
44 
45         self.min_steps = min_steps
46 
47         self.all_epochs = all_epochs
48         self.epoch_list = epoch_list
49         self.epoch_ix = 0
50         self.step_list = step_list
51         self.step_ix = 0
52 
53         # initialized at training_start
54         self.original_max_epochs = -1
55         self.wct_dict = {}
56 
57     def _compute_elapsed_wct(self, epoch_wct_dict, steps_per_epoch, n_epochs):
58         wct = 0.0
59         wct_per_step = 0
60         assert 0 in epoch_wct_dict, "epoch_wct_dict must contain 0"
61         for step in range(steps_per_epoch):
62             if step in epoch_wct_dict:
63                 wct_per_step = epoch_wct_dict[step]
64             wct += wct_per_step
65         return wct * n_epochs
66 
67     def training_start(self, state: State, logger: Logger):
68         self.wall_clock_train = 0.0
69         self.original_max_epochs = state.max_epochs
70         # maybe override epoch_list
71         if self.all_epochs:
72             self.epoch_list = list(range(state.max_epochs))
73             log.info(f"all_epochs=True, overriding epoch_list to be every epoch from 0 to {state.max_epochs}")
74         self.wct_dict = {e: {s: -1.0 for s in self.step_list} for e in self.epoch_list}
75         state.max_epochs = len(self.epoch_list)
76 
77     def epoch_end(self, state: State, logger: Logger):
78         prev_epoch = self.epoch_list[self.epoch_ix]
79         epoch_wct_dict = self.wct_dict[prev_epoch]
80         self.epoch_ix += 1
81         if self.epoch_ix < len(self.epoch_list):
82             next_epoch = self.epoch_list[self.epoch_ix]
83         else:
84             next_epoch = self.original_max_epochs
85 
86         state.epoch = next_epoch - 1
87         state.step = next_epoch * state.steps_per_epoch
88         n_epochs = next_epoch - prev_epoch
89 
90         self.wall_clock_train += self._compute_elapsed_wct(epoch_wct_dict, state.steps_per_epoch, n_epochs)
91         logger.metric_epoch({'wall_clock_train': self.wall_clock_train})
92 
93     def batch_start(self, state: State, logger: Logger):
94         if self.current_time is None:
95             self.current_time = time.time()
96             self.profile_examples = 0
97             self.profile_steps = 0
98             self.profile_time = 0.0
99 
100     def batch_end(self, state: State, logger: Logger):
101         if self.current_time is not None:
102             now = time.time()
103             elapsed = now - self.current_time
104             self.current_time = now
105             self.profile_examples += state.last_batch_size * state.world_size
106             self.profile_steps += 1
107             self.profile_time += elapsed
108 
109             if self.profile_steps >= self.min_steps:
110                 avg_throughput = self.profile_examples / self.profile_time
111                 avg_time_per_step = self.profile_time / self.profile_steps
112                 profile_epoch = self.epoch_list[self.epoch_ix]
113                 profile_step = self.step_list[self.step_ix]
114                 self.wct_dict[profile_epoch][profile_step] = avg_time_per_step
115                 logger.metric_batch({'throughput/step': avg_throughput})
116 
117                 self.current_time = None
118                 self.step_ix += 1
119                 if self.step_ix == len(self.step_list):
120                     self.step_ix = 0
121                     raise BreakEpochException
122                 else:
123                     state.step = state.epoch * state.steps_per_epoch + self.step_list[self.step_ix]
124 
125 
126 @dataclass
127 class TimingMonitorHparams(CallbackHparams):
128     min_steps: int = hp.optional(
129         doc="minimum number of steps to use for measuring throughput",
130         default=50,
131     )
132     epoch_list: List[int] = hp.optional(
133         doc="list of epochs at which to measure throughput",
134         default_factory=lambda: [0, 1],
135     )
136     step_list: List[int] = hp.optional(
137         doc="list of steps at which to measure throughput",
138         default_factory=lambda: [0, 50],
139     )
140     all_epochs: bool = hp.optional(
141         doc="if true, override epoch_list and profile at all epochs.",
142         default=False,
143     )
144 
145     def initialize_object(self) -> TimingMonitor:
146         return TimingMonitor(
147             min_steps=self.min_steps,
148             epoch_list=self.epoch_list,
149             step_list=self.step_list,
150             all_epochs=self.all_epochs,
151         )
152
```

Path: `composer/trainer/entrypoint.py`
Content:
```
1 import logging
2 
3 import composer
4 from composer.trainer.trainer import Trainer
5 from composer.trainer.trainer_hparams import TrainerHparams
6 
7 
8 def main() -> None:
9     logging.basicConfig()
10     logging.captureWarnings(True)
11 
12     hparams = TrainerHparams.create()
13     logging.getLogger(composer.__name__).setLevel(hparams.log_level)
14     trainer = Trainer.create_from_hparams(hparams=hparams)
15     trainer.fit()
16 
17 
18 if __name__ == "__main__":
19     main()
20
```
--- 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/callbacks/profiler.py b/composer/callbacks/profiler.py
--- a/composer/callbacks/profiler.py
+++ b/composer/callbacks/profiler.py
@@ -1,3 +1,5 @@
+# Copyright 2021 MosaicML. All Rights Reserved.
+
 from __future__ import annotations
 
 import atexit
diff --git a/composer/callbacks/timing_monitor.py b/composer/callbacks/timing_monitor.py
--- a/composer/callbacks/timing_monitor.py
+++ b/composer/callbacks/timing_monitor.py
@@ -1,3 +1,5 @@
+# Copyright 2021 MosaicML. All Rights Reserved.
+
 from __future__ import annotations
 
 import logging
diff --git a/composer/trainer/entrypoint.py b/composer/trainer/entrypoint.py
--- a/composer/trainer/entrypoint.py
+++ b/composer/trainer/entrypoint.py
@@ -1,3 +1,5 @@
+# Copyright 2021 MosaicML. All Rights Reserved.
+
 import logging
 
 import composer

{"golden_diff": "diff --git a/composer/callbacks/profiler.py b/composer/callbacks/profiler.py\n--- a/composer/callbacks/profiler.py\n+++ b/composer/callbacks/profiler.py\n@@ -1,3 +1,5 @@\n+# Copyright 2021 MosaicML. All Rights Reserved.\n+\n from __future__ import annotations\n \n import atexit\ndiff --git a/composer/callbacks/timing_monitor.py b/composer/callbacks/timing_monitor.py\n--- a/composer/callbacks/timing_monitor.py\n+++ b/composer/callbacks/timing_monitor.py\n@@ -1,3 +1,5 @@\n+# Copyright 2021 MosaicML. All Rights Reserved.\n+\n from __future__ import annotations\n \n import logging\ndiff --git a/composer/trainer/entrypoint.py b/composer/trainer/entrypoint.py\n--- a/composer/trainer/entrypoint.py\n+++ b/composer/trainer/entrypoint.py\n@@ -1,3 +1,5 @@\n+# Copyright 2021 MosaicML. All Rights Reserved.\n+\n import logging\n \n import composer\n", "issue": "Configure Jenkins\nEnable github actions for:\r\n\r\n- [x] pytest CPU runner\r\n- [x] formatting and type checking (`yapf`, `pyright`)\r\n- [ ] docker builds\r\n- [x] docs builds\n", "before_files": [{"content": "from __future__ import annotations\n\nimport atexit\nfrom dataclasses import asdict, dataclass\nfrom typing import TYPE_CHECKING, Optional\n\nimport torch.profiler\nimport yahp as hp\nfrom torch.profiler.profiler import ProfilerAction\n\nfrom composer import Callback\nfrom composer.callbacks.callback_hparams import CallbackHparams\nfrom composer.core.types import StateDict\nfrom composer.utils.ddp import get_global_rank\n\nif TYPE_CHECKING:\n    from composer.core import Logger, State\n\n_PROFILE_MISSING_ERROR = \"The profiler has not been setup. Please call profiler.training_start() before training starts.\"\n\n\n@dataclass\nclass TorchProfilerHparams(CallbackHparams):\n    tensorboard_trace_handler_dir: str = hp.required(\"directory to store trace results\")\n    tensorboard_use_gzip: bool = hp.optional(\"use gzip for trace\", default=False)\n    record_shapes: bool = hp.optional(doc=\"record_shapes\", default=True)\n    profile_memory: bool = hp.optional(doc=\"track tensor memory allocations and frees\", default=False)\n    with_stack: bool = hp.optional(doc=\"record stack info\", default=True)\n    with_flops: bool = hp.optional(doc=\"estimate flops for operators\", default=True)\n\n    skip: int = hp.optional(\"Number of batches to skip at epoch start\", default=0)\n    warmup: int = hp.optional(\"Number of warmup batches in a cycle\", default=1)\n    active: int = hp.optional(\"Number of batches to profile in a cycle\", default=5)\n    wait: int = hp.optional(\"Number of batches to skip at the end of each cycle\", default=0)\n\n    def initialize_object(self) -> TorchProfiler:\n        return TorchProfiler(**asdict(self))\n\n\n@dataclass\nclass TorchProfilerState:\n    batch_in_epoch: int = 0\n    batches_per_epoch: int = 0\n\n\nclass TorchProfiler(Callback):\n\n    def __init__(\n        self,\n        *,\n        skip: int,\n        warmup: int,\n        active: int,\n        wait: int,\n        tensorboard_trace_handler_dir: str,\n        tensorboard_use_gzip: bool,\n        record_shapes: bool,\n        profile_memory: bool,\n        with_stack: bool,\n        with_flops: bool,\n    ) -> None:\n        super().__init__()\n        self.hparams = TorchProfilerHparams(\n            skip=skip,\n            warmup=warmup,\n            active=active,\n            wait=wait,\n            tensorboard_trace_handler_dir=tensorboard_trace_handler_dir,\n            tensorboard_use_gzip=tensorboard_use_gzip,\n            record_shapes=record_shapes,\n            profile_memory=profile_memory,\n            with_stack=with_stack,\n            with_flops=with_flops,\n        )\n        self.profiler: Optional[torch.profiler.profile] = None\n        self.profiler_state: TorchProfilerState = TorchProfilerState()\n        self._torch_profiler_scheduler = torch.profiler.profiler.schedule(\n            wait=self.hparams.wait,\n            warmup=self.hparams.warmup,\n            active=self.hparams.active,\n            skip_first=self.hparams.skip,\n        )\n\n    def state_dict(self) -> StateDict:\n        return asdict(self.profiler_state)\n\n    def load_state_dict(self, state: StateDict) -> None:\n        self.profiler_state = TorchProfilerState(**state)\n\n    def scheduler_fn(self, profiler_step: int) -> ProfilerAction:\n        # Invoked on every batch, at the batch end\n        # But, it's called one batch in advance.\n        # Wrapping the default scheduling function to deal with epoch boundaries\n        # Giving the torch scheduler the batch in the epoch, not the global step\n\n        # adding 1 since this is called before the step is incremented\n        next_batch_in_epoch = self.profiler_state.batch_in_epoch + 1\n        if profiler_step == 0:\n            next_batch_in_epoch = 0\n        torch_scheduler_action = self._torch_profiler_scheduler(next_batch_in_epoch)\n        if next_batch_in_epoch == self.profiler_state.batches_per_epoch:\n            if torch_scheduler_action == ProfilerAction.RECORD:\n                # force saving at epoch boundaries\n                torch_scheduler_action = ProfilerAction.RECORD_AND_SAVE\n        return torch_scheduler_action\n\n    def training_start(self, state: State, logger: Logger) -> None:\n        assert self.profiler is None, _PROFILE_MISSING_ERROR\n        self.profiler = torch.profiler.profile(\n            schedule=self.scheduler_fn,\n            on_trace_ready=torch.profiler.tensorboard_trace_handler(\n                self.hparams.tensorboard_trace_handler_dir,\n                worker_name=str(get_global_rank()),\n                use_gzip=self.hparams.tensorboard_use_gzip,\n            ),\n            activities=None,  # auto-set\n            record_shapes=self.hparams.record_shapes,\n            profile_memory=self.hparams.profile_memory,\n            with_stack=self.hparams.with_stack,\n            with_flops=self.hparams.with_flops,\n        )\n        self.profiler.__enter__()\n        atexit.register(self._close_profiler)\n\n    def batch_end(self, state: State, logger: Logger) -> None:\n        assert self.profiler is not None, _PROFILE_MISSING_ERROR\n        self.profiler.step()\n\n    def epoch_start(self, state: State, logger: Logger) -> None:\n        self.profiler_state.batches_per_epoch = state.steps_per_epoch\n\n    def batch_start(self, state: State, logger: Logger) -> None:\n        self.profiler_state.batch_in_epoch = state.batch_idx\n        assert self.profiler is not None, _PROFILE_MISSING_ERROR\n        logger.metric_batch({\"profiler/state\": self.profiler.current_action.name})\n\n    def _close_profiler(self) -> None:\n        assert self.profiler is not None\n        self.profiler.__exit__(None, None, None)\n", "path": "composer/callbacks/profiler.py"}, {"content": "from __future__ import annotations\n\nimport logging\nimport time\nfrom dataclasses import dataclass\nfrom typing import List\n\nimport yahp as hp\n\nfrom composer import Logger, State\nfrom composer.callbacks.callback_hparams import CallbackHparams\nfrom composer.core.callback import Callback\nfrom composer.core.types import BreakEpochException\n\nlog = logging.getLogger(__name__)\n\n\nclass TimingMonitor(Callback):\n\n    def __init__(self, min_steps: int, epoch_list: List[int], step_list: List[int], all_epochs: bool):\n        super().__init__()\n        self.hparams = TimingMonitorHparams(min_steps=min_steps,\n                                            epoch_list=epoch_list,\n                                            step_list=step_list,\n                                            all_epochs=all_epochs)\n\n        if not all_epochs:\n            assert len(epoch_list) > 0, \"'epoch_list'  must be non-empty.\"\n            assert 0 in epoch_list, \\\n                \"'epoch_list' must contain 0, otherwise first K epochs have unknown speed\"\n        assert len(step_list) > 0, \"'step_list'  must be non-empty.\"\n        assert 0 in step_list, \\\n            \"'step_list' must contain 0 because `EPOCH_START` requires batch_idx of 0\"\n\n        # Sort lists so that time only moves forward\n        epoch_list = list(sorted(epoch_list))\n        step_list = list(sorted(step_list))\n\n        self.current_time = None\n        self.profile_examples = 0\n        self.profile_steps = 0\n        self.profile_time = 0\n        self.wall_clock_train = 0\n\n        self.min_steps = min_steps\n\n        self.all_epochs = all_epochs\n        self.epoch_list = epoch_list\n        self.epoch_ix = 0\n        self.step_list = step_list\n        self.step_ix = 0\n\n        # initialized at training_start\n        self.original_max_epochs = -1\n        self.wct_dict = {}\n\n    def _compute_elapsed_wct(self, epoch_wct_dict, steps_per_epoch, n_epochs):\n        wct = 0.0\n        wct_per_step = 0\n        assert 0 in epoch_wct_dict, \"epoch_wct_dict must contain 0\"\n        for step in range(steps_per_epoch):\n            if step in epoch_wct_dict:\n                wct_per_step = epoch_wct_dict[step]\n            wct += wct_per_step\n        return wct * n_epochs\n\n    def training_start(self, state: State, logger: Logger):\n        self.wall_clock_train = 0.0\n        self.original_max_epochs = state.max_epochs\n        # maybe override epoch_list\n        if self.all_epochs:\n            self.epoch_list = list(range(state.max_epochs))\n            log.info(f\"all_epochs=True, overriding epoch_list to be every epoch from 0 to {state.max_epochs}\")\n        self.wct_dict = {e: {s: -1.0 for s in self.step_list} for e in self.epoch_list}\n        state.max_epochs = len(self.epoch_list)\n\n    def epoch_end(self, state: State, logger: Logger):\n        prev_epoch = self.epoch_list[self.epoch_ix]\n        epoch_wct_dict = self.wct_dict[prev_epoch]\n        self.epoch_ix += 1\n        if self.epoch_ix < len(self.epoch_list):\n            next_epoch = self.epoch_list[self.epoch_ix]\n        else:\n            next_epoch = self.original_max_epochs\n\n        state.epoch = next_epoch - 1\n        state.step = next_epoch * state.steps_per_epoch\n        n_epochs = next_epoch - prev_epoch\n\n        self.wall_clock_train += self._compute_elapsed_wct(epoch_wct_dict, state.steps_per_epoch, n_epochs)\n        logger.metric_epoch({'wall_clock_train': self.wall_clock_train})\n\n    def batch_start(self, state: State, logger: Logger):\n        if self.current_time is None:\n            self.current_time = time.time()\n            self.profile_examples = 0\n            self.profile_steps = 0\n            self.profile_time = 0.0\n\n    def batch_end(self, state: State, logger: Logger):\n        if self.current_time is not None:\n            now = time.time()\n            elapsed = now - self.current_time\n            self.current_time = now\n            self.profile_examples += state.last_batch_size * state.world_size\n            self.profile_steps += 1\n            self.profile_time += elapsed\n\n            if self.profile_steps >= self.min_steps:\n                avg_throughput = self.profile_examples / self.profile_time\n                avg_time_per_step = self.profile_time / self.profile_steps\n                profile_epoch = self.epoch_list[self.epoch_ix]\n                profile_step = self.step_list[self.step_ix]\n                self.wct_dict[profile_epoch][profile_step] = avg_time_per_step\n                logger.metric_batch({'throughput/step': avg_throughput})\n\n                self.current_time = None\n                self.step_ix += 1\n                if self.step_ix == len(self.step_list):\n                    self.step_ix = 0\n                    raise BreakEpochException\n                else:\n                    state.step = state.epoch * state.steps_per_epoch + self.step_list[self.step_ix]\n\n\n@dataclass\nclass TimingMonitorHparams(CallbackHparams):\n    min_steps: int = hp.optional(\n        doc=\"minimum number of steps to use for measuring throughput\",\n        default=50,\n    )\n    epoch_list: List[int] = hp.optional(\n        doc=\"list of epochs at which to measure throughput\",\n        default_factory=lambda: [0, 1],\n    )\n    step_list: List[int] = hp.optional(\n        doc=\"list of steps at which to measure throughput\",\n        default_factory=lambda: [0, 50],\n    )\n    all_epochs: bool = hp.optional(\n        doc=\"if true, override epoch_list and profile at all epochs.\",\n        default=False,\n    )\n\n    def initialize_object(self) -> TimingMonitor:\n        return TimingMonitor(\n            min_steps=self.min_steps,\n            epoch_list=self.epoch_list,\n            step_list=self.step_list,\n            all_epochs=self.all_epochs,\n        )\n", "path": "composer/callbacks/timing_monitor.py"}, {"content": "import logging\n\nimport composer\nfrom composer.trainer.trainer import Trainer\nfrom composer.trainer.trainer_hparams import TrainerHparams\n\n\ndef main() -> None:\n    logging.basicConfig()\n    logging.captureWarnings(True)\n\n    hparams = TrainerHparams.create()\n    logging.getLogger(composer.__name__).setLevel(hparams.log_level)\n    trainer = Trainer.create_from_hparams(hparams=hparams)\n    trainer.fit()\n\n\nif __name__ == \"__main__\":\n    main()\n", "path": "composer/trainer/entrypoint.py"}], "after_files": [{"content": "# Copyright 2021 MosaicML. All Rights Reserved.\n\nfrom __future__ import annotations\n\nimport atexit\nfrom dataclasses import asdict, dataclass\nfrom typing import TYPE_CHECKING, Optional\n\nimport torch.profiler\nimport yahp as hp\nfrom torch.profiler.profiler import ProfilerAction\n\nfrom composer import Callback\nfrom composer.callbacks.callback_hparams import CallbackHparams\nfrom composer.core.types import StateDict\nfrom composer.utils.ddp import get_global_rank\n\nif TYPE_CHECKING:\n    from composer.core import Logger, State\n\n_PROFILE_MISSING_ERROR = \"The profiler has not been setup. Please call profiler.training_start() before training starts.\"\n\n\n@dataclass\nclass TorchProfilerHparams(CallbackHparams):\n    tensorboard_trace_handler_dir: str = hp.required(\"directory to store trace results\")\n    tensorboard_use_gzip: bool = hp.optional(\"use gzip for trace\", default=False)\n    record_shapes: bool = hp.optional(doc=\"record_shapes\", default=True)\n    profile_memory: bool = hp.optional(doc=\"track tensor memory allocations and frees\", default=False)\n    with_stack: bool = hp.optional(doc=\"record stack info\", default=True)\n    with_flops: bool = hp.optional(doc=\"estimate flops for operators\", default=True)\n\n    skip: int = hp.optional(\"Number of batches to skip at epoch start\", default=0)\n    warmup: int = hp.optional(\"Number of warmup batches in a cycle\", default=1)\n    active: int = hp.optional(\"Number of batches to profile in a cycle\", default=5)\n    wait: int = hp.optional(\"Number of batches to skip at the end of each cycle\", default=0)\n\n    def initialize_object(self) -> TorchProfiler:\n        return TorchProfiler(**asdict(self))\n\n\n@dataclass\nclass TorchProfilerState:\n    batch_in_epoch: int = 0\n    batches_per_epoch: int = 0\n\n\nclass TorchProfiler(Callback):\n\n    def __init__(\n        self,\n        *,\n        skip: int,\n        warmup: int,\n        active: int,\n        wait: int,\n        tensorboard_trace_handler_dir: str,\n        tensorboard_use_gzip: bool,\n        record_shapes: bool,\n        profile_memory: bool,\n        with_stack: bool,\n        with_flops: bool,\n    ) -> None:\n        super().__init__()\n        self.hparams = TorchProfilerHparams(\n            skip=skip,\n            warmup=warmup,\n            active=active,\n            wait=wait,\n            tensorboard_trace_handler_dir=tensorboard_trace_handler_dir,\n            tensorboard_use_gzip=tensorboard_use_gzip,\n            record_shapes=record_shapes,\n            profile_memory=profile_memory,\n            with_stack=with_stack,\n            with_flops=with_flops,\n        )\n        self.profiler: Optional[torch.profiler.profile] = None\n        self.profiler_state: TorchProfilerState = TorchProfilerState()\n        self._torch_profiler_scheduler = torch.profiler.profiler.schedule(\n            wait=self.hparams.wait,\n            warmup=self.hparams.warmup,\n            active=self.hparams.active,\n            skip_first=self.hparams.skip,\n        )\n\n    def state_dict(self) -> StateDict:\n        return asdict(self.profiler_state)\n\n    def load_state_dict(self, state: StateDict) -> None:\n        self.profiler_state = TorchProfilerState(**state)\n\n    def scheduler_fn(self, profiler_step: int) -> ProfilerAction:\n        # Invoked on every batch, at the batch end\n        # But, it's called one batch in advance.\n        # Wrapping the default scheduling function to deal with epoch boundaries\n        # Giving the torch scheduler the batch in the epoch, not the global step\n\n        # adding 1 since this is called before the step is incremented\n        next_batch_in_epoch = self.profiler_state.batch_in_epoch + 1\n        if profiler_step == 0:\n            next_batch_in_epoch = 0\n        torch_scheduler_action = self._torch_profiler_scheduler(next_batch_in_epoch)\n        if next_batch_in_epoch == self.profiler_state.batches_per_epoch:\n            if torch_scheduler_action == ProfilerAction.RECORD:\n                # force saving at epoch boundaries\n                torch_scheduler_action = ProfilerAction.RECORD_AND_SAVE\n        return torch_scheduler_action\n\n    def training_start(self, state: State, logger: Logger) -> None:\n        assert self.profiler is None, _PROFILE_MISSING_ERROR\n        self.profiler = torch.profiler.profile(\n            schedule=self.scheduler_fn,\n            on_trace_ready=torch.profiler.tensorboard_trace_handler(\n                self.hparams.tensorboard_trace_handler_dir,\n                worker_name=str(get_global_rank()),\n                use_gzip=self.hparams.tensorboard_use_gzip,\n            ),\n            activities=None,  # auto-set\n            record_shapes=self.hparams.record_shapes,\n            profile_memory=self.hparams.profile_memory,\n            with_stack=self.hparams.with_stack,\n            with_flops=self.hparams.with_flops,\n        )\n        self.profiler.__enter__()\n        atexit.register(self._close_profiler)\n\n    def batch_end(self, state: State, logger: Logger) -> None:\n        assert self.profiler is not None, _PROFILE_MISSING_ERROR\n        self.profiler.step()\n\n    def epoch_start(self, state: State, logger: Logger) -> None:\n        self.profiler_state.batches_per_epoch = state.steps_per_epoch\n\n    def batch_start(self, state: State, logger: Logger) -> None:\n        self.profiler_state.batch_in_epoch = state.batch_idx\n        assert self.profiler is not None, _PROFILE_MISSING_ERROR\n        logger.metric_batch({\"profiler/state\": self.profiler.current_action.name})\n\n    def _close_profiler(self) -> None:\n        assert self.profiler is not None\n        self.profiler.__exit__(None, None, None)\n", "path": "composer/callbacks/profiler.py"}, {"content": "# Copyright 2021 MosaicML. All Rights Reserved.\n\nfrom __future__ import annotations\n\nimport logging\nimport time\nfrom dataclasses import dataclass\nfrom typing import List\n\nimport yahp as hp\n\nfrom composer import Logger, State\nfrom composer.callbacks.callback_hparams import CallbackHparams\nfrom composer.core.callback import Callback\nfrom composer.core.types import BreakEpochException\n\nlog = logging.getLogger(__name__)\n\n\nclass TimingMonitor(Callback):\n\n    def __init__(self, min_steps: int, epoch_list: List[int], step_list: List[int], all_epochs: bool):\n        super().__init__()\n        self.hparams = TimingMonitorHparams(min_steps=min_steps,\n                                            epoch_list=epoch_list,\n                                            step_list=step_list,\n                                            all_epochs=all_epochs)\n\n        if not all_epochs:\n            assert len(epoch_list) > 0, \"'epoch_list'  must be non-empty.\"\n            assert 0 in epoch_list, \\\n                \"'epoch_list' must contain 0, otherwise first K epochs have unknown speed\"\n        assert len(step_list) > 0, \"'step_list'  must be non-empty.\"\n        assert 0 in step_list, \\\n            \"'step_list' must contain 0 because `EPOCH_START` requires batch_idx of 0\"\n\n        # Sort lists so that time only moves forward\n        epoch_list = list(sorted(epoch_list))\n        step_list = list(sorted(step_list))\n\n        self.current_time = None\n        self.profile_examples = 0\n        self.profile_steps = 0\n        self.profile_time = 0\n        self.wall_clock_train = 0\n\n        self.min_steps = min_steps\n\n        self.all_epochs = all_epochs\n        self.epoch_list = epoch_list\n        self.epoch_ix = 0\n        self.step_list = step_list\n        self.step_ix = 0\n\n        # initialized at training_start\n        self.original_max_epochs = -1\n        self.wct_dict = {}\n\n    def _compute_elapsed_wct(self, epoch_wct_dict, steps_per_epoch, n_epochs):\n        wct = 0.0\n        wct_per_step = 0\n        assert 0 in epoch_wct_dict, \"epoch_wct_dict must contain 0\"\n        for step in range(steps_per_epoch):\n            if step in epoch_wct_dict:\n                wct_per_step = epoch_wct_dict[step]\n            wct += wct_per_step\n        return wct * n_epochs\n\n    def training_start(self, state: State, logger: Logger):\n        self.wall_clock_train = 0.0\n        self.original_max_epochs = state.max_epochs\n        # maybe override epoch_list\n        if self.all_epochs:\n            self.epoch_list = list(range(state.max_epochs))\n            log.info(f\"all_epochs=True, overriding epoch_list to be every epoch from 0 to {state.max_epochs}\")\n        self.wct_dict = {e: {s: -1.0 for s in self.step_list} for e in self.epoch_list}\n        state.max_epochs = len(self.epoch_list)\n\n    def epoch_end(self, state: State, logger: Logger):\n        prev_epoch = self.epoch_list[self.epoch_ix]\n        epoch_wct_dict = self.wct_dict[prev_epoch]\n        self.epoch_ix += 1\n        if self.epoch_ix < len(self.epoch_list):\n            next_epoch = self.epoch_list[self.epoch_ix]\n        else:\n            next_epoch = self.original_max_epochs\n\n        state.epoch = next_epoch - 1\n        state.step = next_epoch * state.steps_per_epoch\n        n_epochs = next_epoch - prev_epoch\n\n        self.wall_clock_train += self._compute_elapsed_wct(epoch_wct_dict, state.steps_per_epoch, n_epochs)\n        logger.metric_epoch({'wall_clock_train': self.wall_clock_train})\n\n    def batch_start(self, state: State, logger: Logger):\n        if self.current_time is None:\n            self.current_time = time.time()\n            self.profile_examples = 0\n            self.profile_steps = 0\n            self.profile_time = 0.0\n\n    def batch_end(self, state: State, logger: Logger):\n        if self.current_time is not None:\n            now = time.time()\n            elapsed = now - self.current_time\n            self.current_time = now\n            self.profile_examples += state.last_batch_size * state.world_size\n            self.profile_steps += 1\n            self.profile_time += elapsed\n\n            if self.profile_steps >= self.min_steps:\n                avg_throughput = self.profile_examples / self.profile_time\n                avg_time_per_step = self.profile_time / self.profile_steps\n                profile_epoch = self.epoch_list[self.epoch_ix]\n                profile_step = self.step_list[self.step_ix]\n                self.wct_dict[profile_epoch][profile_step] = avg_time_per_step\n                logger.metric_batch({'throughput/step': avg_throughput})\n\n                self.current_time = None\n                self.step_ix += 1\n                if self.step_ix == len(self.step_list):\n                    self.step_ix = 0\n                    raise BreakEpochException\n                else:\n                    state.step = state.epoch * state.steps_per_epoch + self.step_list[self.step_ix]\n\n\n@dataclass\nclass TimingMonitorHparams(CallbackHparams):\n    min_steps: int = hp.optional(\n        doc=\"minimum number of steps to use for measuring throughput\",\n        default=50,\n    )\n    epoch_list: List[int] = hp.optional(\n        doc=\"list of epochs at which to measure throughput\",\n        default_factory=lambda: [0, 1],\n    )\n    step_list: List[int] = hp.optional(\n        doc=\"list of steps at which to measure throughput\",\n        default_factory=lambda: [0, 50],\n    )\n    all_epochs: bool = hp.optional(\n        doc=\"if true, override epoch_list and profile at all epochs.\",\n        default=False,\n    )\n\n    def initialize_object(self) -> TimingMonitor:\n        return TimingMonitor(\n            min_steps=self.min_steps,\n            epoch_list=self.epoch_list,\n            step_list=self.step_list,\n            all_epochs=self.all_epochs,\n        )\n", "path": "composer/callbacks/timing_monitor.py"}, {"content": "# Copyright 2021 MosaicML. All Rights Reserved.\n\nimport logging\n\nimport composer\nfrom composer.trainer.trainer import Trainer\nfrom composer.trainer.trainer_hparams import TrainerHparams\n\n\ndef main() -> None:\n    logging.basicConfig()\n    logging.captureWarnings(True)\n\n    hparams = TrainerHparams.create()\n    logging.getLogger(composer.__name__).setLevel(hparams.log_level)\n    trainer = Trainer.create_from_hparams(hparams=hparams)\n    trainer.fit()\n\n\nif __name__ == \"__main__\":\n    main()\n", "path": "composer/trainer/entrypoint.py"}]}