Acquire locks before mp.Value updates

learning-at-home · Nov 4, 2024 · 0590954 · 0590954
1 parent aa9caf4
commit 0590954
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/tests/test_moe.py b/tests/test_moe.py
@@ -319,7 +319,8 @@ def test_client_anomaly_detection():
 def _measure_coro_running_time(n_coros, elapsed_fut, counter):
     async def coro():
         await asyncio.sleep(0.1)
-        counter.value += 1
+        with counter.get_lock():
+            counter.value += 1
 
     try:
         start_time = time.perf_counter()

diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
@@ -414,8 +414,8 @@ def run_trainer(batch_size: int, batch_time: float, client_mode: bool):
             loss.backward()
 
             optimizer.step()
-
-            total_samples_accumulated.value += batch_size
+            with total_samples_accumulated.get_lock():
+                total_samples_accumulated.value += batch_size
 
             if not reuse_grad_buffers:
                 optimizer.zero_grad()