feat: Support get runner type in context by get_runner_type

Signed-off-by: plotor <[email protected]>

Author: plotor

daft/context.py

@@ -47,6 +47,12 @@ def __init__(self, ctx: PyDaftContext | None = None):
         else:
             self._ctx = PyDaftContext()
 
+    def get_runner_type(self) -> str:
+        if self._ctx._runner is not None:
+            return self._ctx._runner.name
+
+        return self._ctx.get_runner_type()
+
     def get_or_create_runner(self) -> Runner[PartitionT]:
         return self._ctx.get_or_create_runner()
 

daft/daft/__init__.pyi

@@ -1955,6 +1955,7 @@ class PyDaftContext:
     def __init__(self) -> None: ...
     _runner: Runner[Any]
     def get_or_create_runner(self) -> Runner[PartitionT]: ...
+    def get_runner_type(self) -> str: ...
     _daft_execution_config: PyDaftExecutionConfig
     _daft_planning_config: PyDaftPlanningConfig
     @property

daft/utils.py

@@ -126,7 +126,7 @@ def column_inputs_to_expressions(columns: ManyColumnsInputType) -> list[Expressi
     return [col(c) if isinstance(c, str) else c for c in column_iter]
 
 
-def detect_ray_state() -> bool:
+def detect_ray_state() -> tuple[bool, bool]:
     ray_is_initialized = False
     ray_is_in_job = False
     in_ray_worker = False
@@ -145,4 +145,4 @@ def detect_ray_state() -> bool:
     except ImportError:
         pass
 
-    return not in_ray_worker and (ray_is_initialized or ray_is_in_job)
+    return ray_is_initialized or ray_is_in_job, in_ray_worker

src/daft-context/src/lib.rs

@@ -244,6 +244,14 @@ pub fn set_runner_ray(
 ) -> DaftResult<DaftContext> {
     let ctx = get_context();
 
+    let runner_type = get_runner_type_from_env();
+    if !runner_type.is_empty() && runner_type != RayRunner::NAME {
+        log::warn!(
+            "Ignore inconsistent $DAFT_RUNNER='{}' env when setting runner as ray",
+            runner_type
+        );
+    }
+
     let runner = Runner::Ray(RayRunner::try_new(
         address,
         max_task_backlog,
@@ -268,6 +276,14 @@ pub fn set_runner_ray(
 pub fn set_runner_native(num_threads: Option<usize>) -> DaftResult<DaftContext> {
     let ctx = get_context();
 
+    let runner_type = get_runner_type_from_env();
+    if !runner_type.is_empty() && runner_type != NativeRunner::NAME {
+        log::warn!(
+            "Ignore inconsistent $DAFT_RUNNER='{}' env when setting runner as native",
+            runner_type
+        );
+    }
+
     let runner = Runner::Native(NativeRunner::try_new(num_threads)?);
     let runner = Arc::new(runner);
 
@@ -322,30 +338,45 @@ fn get_ray_runner_config_from_env() -> RunnerConfig {
 
 /// Helper function to automatically detect whether to use the ray runner.
 #[cfg(feature = "python")]
-fn detect_ray_state() -> bool {
+fn detect_ray_state() -> (bool, bool) {
     Python::with_gil(|py| {
         py.import(pyo3::intern!(py, "daft.utils"))
             .and_then(|m| m.getattr(pyo3::intern!(py, "detect_ray_state")))
             .and_then(|m| m.call0())
             .and_then(|m| m.extract())
-            .unwrap_or(false)
+            .unwrap_or((false, false))
     })
 }
 
 #[cfg(feature = "python")]
-fn get_runner_config_from_env() -> DaftResult<RunnerConfig> {
+fn get_runner_type_from_env() -> String {
     const DAFT_RUNNER: &str = "DAFT_RUNNER";
 
-    let runner_from_envvar = std::env::var(DAFT_RUNNER)
+    std::env::var(DAFT_RUNNER)
         .unwrap_or_default()
-        .to_lowercase();
-
-    match runner_from_envvar.as_str() {
-        "native" => Ok(RunnerConfig::Native { num_threads: None }),
-        "ray" => Ok(get_ray_runner_config_from_env()),
-        "py" => Err(DaftError::ValueError("The PyRunner was removed from Daft from v0.5.0 onwards. Please set the env to `DAFT_RUNNER=native` instead.".to_string())),
-        "" => Ok(if detect_ray_state() { get_ray_runner_config_from_env() } else { RunnerConfig::Native { num_threads: None }}),
-        other => Err(DaftError::ValueError(format!("Invalid runner type `DAFT_RUNNER={other}` specified through the env. Please use either `native` or `ray` instead.")))
+        .to_lowercase()
+}
+
+#[cfg(feature = "python")]
+fn get_runner_config_from_env() -> DaftResult<RunnerConfig> {
+    match get_runner_type_from_env().as_str() {
+        NativeRunner::NAME => Ok(RunnerConfig::Native { num_threads: None }),
+        RayRunner::NAME => Ok(get_ray_runner_config_from_env()),
+        "py" => Err(DaftError::ValueError(
+            "The PyRunner was removed from Daft from v0.5.0 onwards. \
+            Please set the env to `DAFT_RUNNER=native` instead."
+                .to_string(),
+        )),
+        "" => Ok(if detect_ray_state() == (true, false) {
+            // on ray but not in ray worker
+            get_ray_runner_config_from_env()
+        } else {
+            RunnerConfig::Native { num_threads: None }
+        }),
+        other => Err(DaftError::ValueError(format!(
+            "Invalid runner type `DAFT_RUNNER={other}` specified through the env. \
+            Please use either `native` or `ray` instead."
+        ))),
     }
 }
 
@@ -366,7 +397,7 @@ pub fn reset_runner() {
 }
 
 #[cfg(feature = "python")]
-pub fn register_modules(parent: &Bound<PyModule>) -> pyo3::PyResult<()> {
+pub fn register_modules(parent: &Bound<PyModule>) -> PyResult<()> {
     parent.add_function(wrap_pyfunction!(
         python::get_runner_config_from_env,
         parent

src/daft-context/src/python.rs

@@ -2,9 +2,10 @@ use std::sync::Arc;
 
 use common_daft_config::{PyDaftExecutionConfig, PyDaftPlanningConfig};
 use common_error::DaftError;
-use pyo3::prelude::*;
+use daft_py_runners::{NativeRunner, RayRunner};
+use pyo3::{prelude::*, IntoPyObjectExt};
 
-use crate::{DaftContext, Runner, RunnerConfig};
+use crate::{detect_ray_state, DaftContext, Runner, RunnerConfig};
 
 #[pyclass]
 pub struct PyRunnerConfig {
@@ -13,7 +14,7 @@ pub struct PyRunnerConfig {
 
 #[pyclass]
 pub struct PyDaftContext {
-    inner: crate::DaftContext,
+    inner: DaftContext,
 }
 
 impl Default for PyDaftContext {
@@ -45,6 +46,24 @@ impl PyDaftContext {
             }
         }
     }
+
+    pub fn get_runner_type(&self, py: Python) -> PyResult<PyObject> {
+        let runner_type = self.inner.runner().map_or_else(
+            || {
+                if detect_ray_state().0 {
+                    RayRunner::NAME
+                } else {
+                    NativeRunner::NAME
+                }
+            },
+            |runner| match runner.as_ref() {
+                Runner::Ray(_) => RayRunner::NAME,
+                Runner::Native(_) => NativeRunner::NAME,
+            },
+        );
+        runner_type.into_py_any(py)
+    }
+
     #[getter(_daft_execution_config)]
     pub fn get_daft_execution_config(&self, py: Python) -> PyResult<PyDaftExecutionConfig> {
         let config = py.allow_threads(|| self.inner.execution_config());

src/daft-py-runners/src/lib.rs

@@ -22,6 +22,8 @@ pub struct RayRunner {
 
 #[cfg(feature = "python")]
 impl RayRunner {
+    pub const NAME: &'static str = "ray";
+
     pub fn try_new(
         address: Option<String>,
         max_task_backlog: Option<usize>,
@@ -53,6 +55,8 @@ pub struct NativeRunner {
 
 #[cfg(feature = "python")]
 impl NativeRunner {
+    pub const NAME: &'static str = "native";
+
     pub fn try_new(num_threads: Option<usize>) -> DaftResult<Self> {
         Python::with_gil(|py| {
             let native_runner_module = py.import(intern!(py, "daft.runners.native_runner"))?;
@@ -84,13 +88,13 @@ impl Runner {
         Python::with_gil(|py| {
             let name = obj.getattr(py, "name")?.extract::<String>(py)?;
             match name.as_ref() {
-                "ray" => {
+                RayRunner::NAME => {
                     let ray_runner = RayRunner {
                         pyobj: Arc::new(obj),
                     };
                     Ok(Self::Ray(ray_runner))
                 }
-                "native" => {
+                NativeRunner::NAME => {
                     let native_runner = NativeRunner {
                         pyobj: Arc::new(obj),
                     };