torch.set_num_threads(1) in run_server.py if OpenMP not available

src/petals/cli/run_server.py

@@ -1,6 +1,5 @@
 import argparse
 import logging
-import platform
 
 import configargparse
 import torch
@@ -212,7 +211,8 @@ def main():
 
     validate_version()
 
-    if platform.platform() == "Darwin":
+    if not torch.backends.openmp.is_available():
+        # Necessary to prevent the server from freezing after forks
         torch.set_num_threads(1)
 
     server = Server(

tests/test_priority_pool.py

@@ -9,8 +9,6 @@ from petals.server.task_pool import PrioritizedTaskPool
 
 
 def _submit_tasks(pools, results_valid):
-    torch.set_num_threads(1)
-
     futures = []
     futures.append(pools[0].submit_task(torch.tensor([0]), priority=1))
     futures.append(pools[0].submit_task(torch.tensor([1]), priority=1))