Use ray backend by default when we're running in a Ray worker. (#1959)

Author: rjpower

lib/zephyr/src/zephyr/backend_factory.py

@@ -22,6 +22,7 @@
 from typing import Literal
 
 import humanfriendly
+import ray
 
 from zephyr.backends import Backend, BackendConfig, RayBackend, SyncBackend, ThreadPoolBackend
 
@@ -145,14 +146,22 @@ def flow_backend(
     )
     if not has_params:
         if current is None:
-            logger.warning("No backend configured in context, using ThreadPoolBackend as default.")
-            return create_backend("threadpool")
+            # Default to Ray backend if Ray is already initialized
+            if ray.is_initialized():
+                logger.info("Ray is initialized, using RayBackend as default.")
+                return create_backend("ray")
+            else:
+                logger.warning("No backend configured in context, using ThreadPoolBackend as default.")
+                return create_backend("threadpool")
         return current
 
     # Parameters provided: create new backend with merged config
     if current is None:
-        # No current backend, create default threadpool
-        current = create_backend("threadpool")
+        # No current backend, create default based on Ray availability
+        if ray.is_initialized():
+            current = create_backend("ray")
+        else:
+            current = create_backend("threadpool")
 
     # Build override dict, only including non-None values
     overrides = {}

lib/zephyr/tests/test_backends.py

@@ -14,7 +14,10 @@
 
 """Tests for backend implementations."""
 
-from zephyr.backends import format_shard_path
+import ray
+
+from zephyr.backend_factory import flow_backend
+from zephyr.backends import RayBackend, format_shard_path
 
 
 def test_format_shard_path_basic():
@@ -109,3 +112,16 @@ def test_write_jsonl_no_compression_without_gz_extension(tmp_path):
         assert len(lines) == 2
         assert json.loads(lines[0]) == {"id": 1, "text": "hello"}
         assert json.loads(lines[1]) == {"id": 2, "text": "world"}
+
+
+def test_flow_backend_defaults_to_ray_when_initialized():
+    """Test that flow_backend returns RayBackend when Ray is initialized."""
+    if ray.is_initialized():
+        ray.shutdown()
+
+    try:
+        ray.init(ignore_reinit_error=True)
+        backend = flow_backend()
+        assert isinstance(backend, RayBackend)
+    finally:
+        ray.shutdown()