From 882f60012fca805cad30f0e9e5d281117e10eae4 Mon Sep 17 00:00:00 2001
From: mehrdadn <mehrdadn@users.noreply.github.com>
Date: Mon, 29 Jun 2020 19:00:00 -0700
Subject: [PATCH] Fix WSL patch (it doesn't fall under "win32") (#9195)

---
 python/ray/__init__.py |  7 ++++-
 python/ray/compat.py   | 58 ++++++++++++++++++++++--------------------
 2 files changed, 36 insertions(+), 29 deletions(-)

diff --git a/python/ray/__init__.py b/python/ray/__init__.py
index 81cc2b528..178d6523b 100644
--- a/python/ray/__init__.py
+++ b/python/ray/__init__.py
@@ -1,6 +1,7 @@
 import os
 import logging
 from os.path import dirname
+import platform
 import sys
 
 logger = logging.getLogger(__name__)
@@ -36,9 +37,13 @@ sys.path.insert(0, thirdparty_files)
 
 if sys.platform == "win32":
     import ray.compat  # noqa: E402
-    ray.compat.patch_psutil()
     ray.compat.patch_redis_empty_recv()
 
+if (platform.system() == "Linux"
+        and "Microsoft".lower() in platform.release().lower()):
+    import ray.compat  # noqa: E402
+    ray.compat.patch_psutil()
+
 # Expose ray ABI symbols which may be dependent by other shared
 # libraries such as _streaming.so. See BUILD.bazel:_raylet
 python_shared_lib_suffix = ".so" if sys.platform != "win32" else ".pyd"
diff --git a/python/ray/compat.py b/python/ray/compat.py
index 7443f082b..eeac3240b 100644
--- a/python/ray/compat.py
+++ b/python/ray/compat.py
@@ -26,34 +26,36 @@ def patch_redis_empty_recv():
 
 
 def patch_psutil():
-    if (platform.system() == "Linux"
-            and "Microsoft".lower() in platform.release().lower()):
-        # WSL's /proc/meminfo has an inconsistency where it
-        # nondeterministically omits a space after colons (after "SwapFree:"
-        # in my case).
-        # psutil then splits on spaces and then parses the wrong field,
-        # crashing on the 'int(fields[1])' expression in
-        # psutil._pslinux.virtual_memory().
-        # Workaround: We ensure there is a space following each colon.
+    """WSL's /proc/meminfo has an inconsistency where it
+    nondeterministically omits a space after colons (after "SwapFree:"
+    in my case).
+    psutil then splits on spaces and then parses the wrong field,
+    crashing on the 'int(fields[1])' expression in
+    psutil._pslinux.virtual_memory().
+    Workaround: We ensure there is a space following each colon.
+    """
+    assert (platform.system() == "Linux"
+            and "Microsoft".lower() in platform.release().lower())
+
+    try:
+        import psutil._pslinux
+    except ImportError:
+        psutil = None
+    psutil_open_binary = None
+    if psutil:
         try:
-            import psutil._pslinux
-        except ImportError:
-            psutil = None
-        psutil_open_binary = None
-        if psutil:
-            try:
-                psutil_open_binary = psutil._pslinux.open_binary
-            except AttributeError:
-                pass
-        # Only patch it if it doesn't seem to have been patched already
-        if psutil_open_binary and psutil_open_binary.__name__ == "open_binary":
+            psutil_open_binary = psutil._pslinux.open_binary
+        except AttributeError:
+            pass
+    # Only patch it if it doesn't seem to have been patched already
+    if psutil_open_binary and psutil_open_binary.__name__ == "open_binary":
 
-            def psutil_open_binary_patched(fname, *args, **kwargs):
-                f = psutil_open_binary(fname, *args, **kwargs)
-                if fname == "/proc/meminfo":
-                    with f:
-                        # Make sure there's a space after colons
-                        return io.BytesIO(f.read().replace(b":", b": "))
-                return f
+        def psutil_open_binary_patched(fname, *args, **kwargs):
+            f = psutil_open_binary(fname, *args, **kwargs)
+            if fname == "/proc/meminfo":
+                with f:
+                    # Make sure there's a space after colons
+                    return io.BytesIO(f.read().replace(b":", b": "))
+            return f
 
-            psutil._pslinux.open_binary = psutil_open_binary_patched
+        psutil._pslinux.open_binary = psutil_open_binary_patched