set num_processes=1 in MAB test (facebookresearch#602)

Summary: Pull Request resolved: facebookresearch#602 Attempt to fix the test timing out on CircleCI Differential Revision: D33900623 fbshipit-source-id: ef9bb3a44ea726df9ba3e82a94fa8467704abe0a
alexnikulkov · Jan 31, 2022 · 3739292 · 3739292
1 parent 09e36a9
commit 3739292
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 7 deletions.
diff --git a/reagent/mab/simulation.py b/reagent/mab/simulation.py
@@ -188,15 +188,25 @@ def multiple_evaluations_bandit_algo(
         )
         for _ in range(n_bandits)
     )
-    with Pool(num_processes) as pool:
-        pseudo_regrets = pool.starmap(
-            partial(
-                single_evaluation_bandit_algo,
+    if num_processes == 1:
+        pseudo_regrets = [
+            single_evaluation_bandit_algo(
+                *a,
                 update_every=update_every,
                 freeze_scores_btw_updates=freeze_scores_btw_updates,
-            ),
-            arguments,
-        )
+            )
+            for a in arguments
+        ]
+    else:
+        with Pool(num_processes) as pool:
+            pseudo_regrets = pool.starmap(
+                partial(
+                    single_evaluation_bandit_algo,
+                    update_every=update_every,
+                    freeze_scores_btw_updates=freeze_scores_btw_updates,
+                ),
+                arguments,
+            )
     return np.stack(pseudo_regrets).mean(0)
 
 

diff --git a/reagent/test/mab/test_mab.py b/reagent/test/mab/test_mab.py
@@ -374,6 +374,7 @@ def test_multiple_evaluations_bandit_algo(self):
             max_steps=max_steps,
             algo_kwargs={"n_arms": 2},
             bandit_kwargs={"probs": torch.Tensor([0.3, 0.5])},
+            num_processes=1,
         )
 
         self.assertIsInstance(regret_trajectory, np.ndarray)