Refactor helper methods for new scheduler integration (#6354)

python/ray/tests/BUILD

@@ -94,6 +94,14 @@ py_test(
     deps = ["//:ray_lib"],
 )
 
+py_test(
+    name = "test_stress_direct",
+    size = "medium",
+    srcs = ["test_stress_direct.py", "test_stress.py"],
+    tags = ["exclusive"],
+    deps = ["//:ray_lib"],
+)
+
 py_test(
     name = "test_stress_sharded",
     size = "medium",
@@ -102,6 +110,14 @@ py_test(
     deps = ["//:ray_lib"],
 )
 
+py_test(
+    name = "test_stress_sharded_direct",
+    size = "medium",
+    srcs = ["test_stress_sharded_direct.py", "test_stress_sharded.py"],
+    tags = ["exclusive"],
+    deps = ["//:ray_lib"],
+)
+
 py_test(
     name = "test_stress_failure",
     size = "large",

python/ray/tests/test_stress_direct.py

@@ -0,0 +1,16 @@
+"""Wrapper script that sets RAY_FORCE_DIRECT."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import pytest
+import sys
+import os
+
+if __name__ == "__main__":
+    os.environ["RAY_FORCE_DIRECT"] = "1"
+    sys.exit(
+        pytest.main(
+            ["-v",
+             os.path.join(os.path.dirname(__file__), "test_stress.py")]))

python/ray/tests/test_stress_sharded_direct.py

@@ -0,0 +1,17 @@
+"""Wrapper script that sets RAY_FORCE_DIRECT."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import pytest
+import sys
+import os
+
+if __name__ == "__main__":
+    os.environ["RAY_FORCE_DIRECT"] = "1"
+    sys.exit(
+        pytest.main([
+            "-v",
+            os.path.join(os.path.dirname(__file__), "test_stress_sharded.py")
+        ]))