Merge pull request #1056 from PReithofer/main

fix race condition in generate_units
facebookresearch · Aug 25, 2023 · 9bcd52b · 9bcd52b
2 parents c15e0d3 + aca08d9
commit 9bcd52b
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/mephisto/operations/task_launcher.py b/mephisto/operations/task_launcher.py
@@ -158,15 +158,15 @@ def generate_units(self):
             )
 
             units_id_to_remove = []
-            for i, item in enumerate(self.unlaunched_units.items()):
-                db_id, unit = item
-                if i < num_avail_units:
-                    self.launched_units[unit.db_id] = unit
-                    units_id_to_remove.append(db_id)
-                    yield unit
-                else:
-                    break
             with self.unlaunched_units_access_condition:
+                for i, item in enumerate(self.unlaunched_units.items()):
+                    db_id, unit = item
+                    if i < num_avail_units:
+                        self.launched_units[unit.db_id] = unit
+                        units_id_to_remove.append(db_id)
+                        yield unit
+                    else:
+                        break
                 for db_id in units_id_to_remove:
                     self.unlaunched_units.pop(db_id)