Patch WSL psutil bug (#6808)

python/ray/__init__.py

@@ -36,6 +36,7 @@ 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()
 
 # Expose ray ABI symbols which may be dependent by other shared

python/ray/compat.py

@@ -1,4 +1,6 @@
 import errno
+import io
+import platform
 import socket
 import sys
 
@@ -21,3 +23,37 @@ def patch_redis_empty_recv():
         return result
 
     redis.connection.recv = redis_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.
+        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":
+
+            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