Make more tests compatible with Windows (#9303)

python/ray/test_utils.py

@@ -1,5 +1,6 @@
 import asyncio
 import errno
+import io
 import json
 import fnmatch
 import os
@@ -9,8 +10,12 @@ import time
 import socket
 import math
 
+from contextlib import redirect_stdout, redirect_stderr
+
 import ray
 import ray.services
+import ray.utils
+from ray.scripts.scripts import main as ray_main
 
 import psutil  # We must import psutil after ray because we bundle it with ray.
 
@@ -52,6 +57,63 @@ def _pid_alive(pid):
     return alive
 
 
+def check_call_module(main, argv, capture_stdout=False, capture_stderr=False):
+    # We use this function instead of calling the "ray" command to work around
+    # some deadlocks that occur when piping ray's output on Windows
+    stream = io.TextIOWrapper(io.BytesIO(), encoding=sys.stdout.encoding)
+    old_argv = sys.argv[:]
+    try:
+        sys.argv = argv[:]
+        try:
+            with redirect_stderr(stream if capture_stderr else sys.stderr):
+                with redirect_stdout(stream if capture_stdout else sys.stdout):
+                    main()
+        finally:
+            stream.flush()
+    except SystemExit as ex:
+        if ex.code:
+            output = stream.buffer.getvalue()
+            raise subprocess.CalledProcessError(ex.code, argv, output)
+    except Exception as ex:
+        output = stream.buffer.getvalue()
+        raise subprocess.CalledProcessError(1, argv, output, ex.args[0])
+    finally:
+        sys.argv = old_argv
+        if capture_stdout:
+            sys.stdout.buffer.write(stream.buffer.getvalue())
+        elif capture_stderr:
+            sys.stderr.buffer.write(stream.buffer.getvalue())
+    return stream.buffer.getvalue()
+
+
+def check_call_ray(args, capture_stdout=False, capture_stderr=False):
+    # We use this function instead of calling the "ray" command to work around
+    # some deadlocks that occur when piping ray's output on Windows
+    argv = ["ray"] + args
+    if sys.platform == "win32":
+        result = check_call_module(
+            ray_main,
+            argv,
+            capture_stdout=capture_stdout,
+            capture_stderr=capture_stderr)
+    else:
+        stdout_redir = None
+        stderr_redir = None
+        if capture_stdout:
+            stdout_redir = subprocess.PIPE
+        if capture_stderr and capture_stdout:
+            stderr_redir = subprocess.STDOUT
+        elif capture_stderr:
+            stderr_redir = subprocess.PIPE
+        proc = subprocess.Popen(argv, stdout=stdout_redir, stderr=stderr_redir)
+        (stdout, stderr) = proc.communicate()
+        if proc.returncode:
+            raise subprocess.CalledProcessError(proc.returncode, argv, stdout,
+                                                stderr)
+        result = b"".join([s for s in [stdout, stderr] if s is not None])
+    return result
+
+
 def wait_for_pid_to_exit(pid, timeout=20):
     start_time = time.time()
     while time.time() - start_time < timeout: