From 0ae4fe020dd00a81fc3303b6a0d98085a8e56a5e Mon Sep 17 00:00:00 2001
From: Eric Liang <ekhliang@gmail.com>
Date: Sun, 23 Feb 2020 21:26:49 -0800
Subject: [PATCH] revert omp threads fix (#7288)

---
 python/ray/__init__.py         | 6 ++++++
 python/ray/tests/test_basic.py | 7 +++++++
 2 files changed, 13 insertions(+)

diff --git a/python/ray/__init__.py b/python/ray/__init__.py
index ecd967738..9874d36c2 100644
--- a/python/ray/__init__.py
+++ b/python/ray/__init__.py
@@ -13,6 +13,12 @@ if "pickle5" in sys.modules:
                       "requires a specific version of pickle5 (which is "
                       "packaged along with Ray).")
 
+if "OMP_NUM_THREADS" not in os.environ:
+    logger.debug("[ray] Forcing OMP_NUM_THREADS=1 to avoid performance "
+                 "degradation with many workers (issue #6998). You can "
+                 "override this by explicitly setting OMP_NUM_THREADS.")
+    os.environ["OMP_NUM_THREADS"] = "1"
+
 # Add the directory containing pickle5 to the Python path so that we find the
 # pickle5 version packaged with ray and not a pre-existing pickle5.
 pickle5_path = os.path.join(
diff --git a/python/ray/tests/test_basic.py b/python/ray/tests/test_basic.py
index 19557a9ac..7140f393c 100644
--- a/python/ray/tests/test_basic.py
+++ b/python/ray/tests/test_basic.py
@@ -38,6 +38,13 @@ def test_ignore_http_proxy(shutdown_only):
     assert ray.get(f.remote()) == 1
 
 
+# https://github.com/ray-project/ray/issues/7287
+def test_omp_threads_set(shutdown_only):
+    ray.init(num_cpus=1)
+    # Should have been auto set by ray init.
+    assert os.environ["OMP_NUM_THREADS"] == "1"
+
+
 def test_simple_serialization(ray_start_regular):
     primitive_objects = [
         # Various primitive types.