diff --git a/python/ray/__init__.py b/python/ray/__init__.py index 2474893b3..81cc2b528 100644 --- a/python/ray/__init__.py +++ b/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 diff --git a/python/ray/compat.py b/python/ray/compat.py index 3ddb3bf8d..7443f082b 100644 --- a/python/ray/compat.py +++ b/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