[monarch] Make sure context is propagated through PythonTask.spawn_blocking

monarch_hyperactor/src/pytokio.rs

@@ -413,14 +413,30 @@ impl PyPythonTask {
     }
 
     #[staticmethod]
-    fn spawn_blocking(f: PyObject) -> PyResult<PyShared> {
+    fn spawn_blocking(py: Python<'_>, f: PyObject) -> PyResult<PyShared> {
         let (tx, rx) = watch::channel(None);
         let traceback = current_traceback()?;
         let traceback1 = traceback
             .as_ref()
             .map_or_else(|| None, |t| Python::with_gil(|py| Some(t.clone_ref(py))));
+        let monarch_context = py
+            .import("monarch._src.actor.actor_mesh")?
+            .call_method0("context")?
+            .unbind();
+        // The `_context` contextvar needs to be propagated through to the thread that
+        // runs the blocking tokio task. Upon completion, the original value of `_context`
+        // is restored.
         let handle = get_tokio_runtime().spawn_blocking(move || {
-            let result = Python::with_gil(|py| f.call0(py));
+            let result = Python::with_gil(|py| {
+                let _context = py
+                    .import("monarch._src.actor.actor_mesh")?
+                    .getattr("_context")?;
+                let old_context = _context.call_method1("get", (PyNone::get(py),))?;
+                _context.call_method1("set", (monarch_context.clone_ref(py),))?;
+                let result = f.call0(py);
+                _context.call_method1("set", (old_context,))?;
+                result
+            });
             send_result(tx, result, traceback1);
         });
         Ok(PyShared {

python/tests/test_python_actors.py

@@ -1686,3 +1686,35 @@ def test_login_job():
             assert v == "hello!"
 
         j.kill()
+
+
+class TestPytokioActor(Actor):
+    @endpoint
+    def context_propagated_through_spawn(self) -> None:
+        cx = context()
+
+        async def task():
+            assert cx is context()
+
+        PythonTask.from_coroutine(coro=task()).spawn().block_on()
+
+    @endpoint
+    def context_propagated_through_spawn_blocking(self) -> None:
+        cx = context()
+
+        def task():
+            assert cx is context()
+
+        PythonTask.spawn_blocking(task).block_on()
+
+
+def test_context_propagated_through_python_task_spawn():
+    p = this_host().spawn_procs()
+    a = p.spawn("test_pytokio_actor", TestPytokioActor)
+    a.context_propagated_through_spawn.call().get()
+
+
+def test_context_propagated_through_python_task_spawn_blocking():
+    p = this_host().spawn_procs()
+    a = p.spawn("test_pytokio_actor", TestPytokioActor)
+    a.context_propagated_through_spawn_blocking.call().get()