feat(zero): add --profile support for num_workers=1

Author: justindujardin

libraries/mathy_python/mathy/agents/zero/config.py

@@ -9,3 +9,7 @@ class SelfPlayConfig(BaseConfig):
     self_play_problems: int = 64
     training_iterations: int = 100
     cpuct: float = 1.0
+
+
+    # When profile is true and workers == 1 the main thread will output worker_0.profile on exit
+    profile: bool = False

libraries/mathy_python/mathy/agents/zero/practice_runner.py

@@ -1,3 +1,4 @@
+import os
 import queue
 import time
 from multiprocessing import Array, Pool, Process, Queue, cpu_count
@@ -141,25 +142,44 @@ def execute_episodes(
         """
         examples = []
         results: List[EpisodeSummary] = []
+        if self.config.profile:
+            import cProfile
 
-        game = self.get_game()
-        predictor = self.get_predictor(game)
-        for i, args in enumerate(episode_args_list):
-            start = time.time()
-            (episode_examples, episode_reward, is_win, problem,) = self.execute_episode(
-                i, game, predictor, **args
-            )
-            duration = time.time() - start
-            examples.append(episode_examples)
-            episode_summary = EpisodeSummary(
-                solved=bool(is_win),
-                text=problem.text,
-                complexity=problem.complexity,
-                reward=episode_reward,
-                duration=duration,
-            )
-            results.append(episode_summary)
-            self.episode_complete(i, episode_summary)
+            pr = cProfile.Profile()
+            pr.enable()
+
+        try:
+            game = self.get_game()
+            predictor = self.get_predictor(game)
+            for i, args in enumerate(episode_args_list):
+                start = time.time()
+                (
+                    episode_examples,
+                    episode_reward,
+                    is_win,
+                    problem,
+                ) = self.execute_episode(i, game, predictor, **args)
+                duration = time.time() - start
+                examples.append(episode_examples)
+                episode_summary = EpisodeSummary(
+                    solved=bool(is_win),
+                    text=problem.text,
+                    complexity=problem.complexity,
+                    reward=episode_reward,
+                    duration=duration,
+                )
+                results.append(episode_summary)
+                self.episode_complete(i, episode_summary)
+        except KeyboardInterrupt:
+            print("Interrupt received. Exiting.")
+
+        if self.config.profile:
+            profile_name = f"worker_0.profile"
+            profile_path = os.path.join(self.config.model_dir, profile_name)
+            pr.disable()
+            pr.dump_stats(profile_path)
+            if self.config.verbose:
+                print(f"PROFILER: saved {profile_path}")
         return examples, results
 
     def execute_episode(
@@ -231,7 +251,7 @@ class ParallelPracticeRunner(PracticeRunner):
     def execute_episodes(
         self, episode_args_list
     ) -> Tuple[List[EpisodeHistory], List[EpisodeSummary]]:
-        def worker(work_queue, result_queue):
+        def worker(work_queue: Queue, result_queue: Queue):
             """Pull items out of the work queue and execute episodes until there are
             no items left"""
             game = self.get_game()

libraries/mathy_python/mathy/agents/zero/self_play_runner.py

@@ -20,6 +20,8 @@ def self_play_runner(config: SelfPlayConfig):
     BaseEpisodeRunner = (
         PracticeRunner if config.num_workers < 2 else ParallelPracticeRunner
     )
+    if config.profile and config.num_workers > 1:
+        raise NotImplementedError("zero agent does not support multiprocess profiling")
     lesson_name = "mathy-poly-easy-v0"
     if config.verbose:
         print(config.json(indent=2))

libraries/mathy_python/mathy/cli.py

@@ -258,7 +258,7 @@ def cli_train(
         instance = A3CAgent(args)
         instance.train()
     elif agent == "zero":
-        setup_tf_env(use_mp=True)
+        setup_tf_env(use_mp=workers > 1)
         from .agents.zero import SelfPlayConfig, self_play_runner
 
         self_play_cfg = SelfPlayConfig(
@@ -275,6 +275,7 @@ def cli_train(
             training_iterations=training_iterations,
             self_play_problems=self_play_problems,
             print_training=show,
+            profile=profile,
         )
         if episodes is not None:
             self_play_cfg.max_eps = episodes

libraries/mathy_python/tests/test_self_play.py

@@ -0,0 +1,12 @@
+import pytest
+
+from ..mathy.agents.zero import SelfPlayConfig, self_play_runner
+
+
+def test_self_play_runner_errors():
+    """Throws error for bad profiling configuration of zero agent.
+    
+    This is only because it wasn't clear how to profile multiple processes that
+    stop/start overtime during training/evaluation loops"""
+    with pytest.raises(NotImplementedError):
+        self_play_runner(SelfPlayConfig(profile=True, num_workers=2))

libraries/website/docs/ml/a3c.md

@@ -93,29 +93,7 @@ The CLI A3C agent accepts a `--profile` option, or a config option when the API
 {!./snippets/ml/a3c_profiling.py!}
 ```
 
-The output files can be visualized using **[Snakeviz](https://jiffyclub.github.io/snakeviz/){target=\_blank}**:
-
-You can install Snakeviz if it's not already present:
-
-```bash
-pip install snakeviz
-```
-
-You can view the output from the previous run:
-
-```bash
-snakeviz /the_folder_my_model_is_in/worker_0.profile
-```
-
-Running the above command will launch a webpage on your local system.
-
-The webpage has performance information about the training run that just finished.
-
-Clicking on the various functions will expand them further.
-
-You can use this "drilling in" to find spots of code that may be using lots of time when they should not be.
-
-<img mathy-logo src="/img/snakeviz_profile.png" alt="View agent performance profile in Skakeviz">
+Learn about how to view output profiles on the [debugging page](/ml/debugging)
 
 ## Multiprocess A3C
 

libraries/website/docs/ml/zero.md

@@ -27,3 +27,18 @@ In this mode you can set breakpoints in the debugger to help diagnose errors.
 ```python
 {!./snippets/ml/zero_debugging.py!}
 ```
+
+### Performance Profiling
+
+The CLI Zero agent accepts `--profile` option along with `--num-workers=1` or the same config options when the API is used.
+
+!!! warning "Limited to num_workers=1"
+
+    It wasn't immediately clear how to profile multiple processes that start and stop over time, so to profile zero you must
+    specify a single worker process.
+
+```python
+{!./snippets/ml/a3c_profiling.py!}
+```
+
+Learn about how to view output profiles on the [debugging page](/ml/debugging)