From bbe56d02d9479e5469d2223c6e1748f82d2394be Mon Sep 17 00:00:00 2001
From: Marius Eriksen <meriksen@meta.com>
Date: Thu, 9 Oct 2025 13:35:40 -0700
Subject: [PATCH] [monarch] extension: make CodeSyncMeshClient ActorMesh
 version polymorphic

This uses shimming to accept both v0 and v1 proc meshes when constructing a codesync client.

Differential Revision: [D84284336](https://our.internmc.facebook.com/intern/diff/D84284336/)

**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D84284336/)!

[ghstack-poisoned]
---
 monarch_extension/src/code_sync.rs            | 41 ++++++++++++++-----
 .../monarch_extension/code_sync.pyi           |  7 +++-
 2 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/monarch_extension/src/code_sync.rs b/monarch_extension/src/code_sync.rs
index 11a556d3d..4ef276ce4 100644
--- a/monarch_extension/src/code_sync.rs
+++ b/monarch_extension/src/code_sync.rs
@@ -27,7 +27,9 @@ use monarch_hyperactor::context::PyInstance;
 use monarch_hyperactor::instance_dispatch;
 use monarch_hyperactor::proc_mesh::PyProcMesh;
 use monarch_hyperactor::runtime::signal_safe_block_on;
+use monarch_hyperactor::v1::proc_mesh::PyProcMesh as PyProcMeshV1;
 use pyo3::Bound;
+use pyo3::exceptions::PyException;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
@@ -264,16 +266,35 @@ impl CodeSyncMeshClient {
 impl CodeSyncMeshClient {
     #[staticmethod]
     #[pyo3(signature = (*, client, proc_mesh))]
-    fn spawn_blocking(py: Python, client: PyInstance, proc_mesh: &PyProcMesh) -> PyResult<Self> {
-        let proc_mesh = proc_mesh.try_inner()?;
-        signal_safe_block_on(py, async move {
-            let actor_mesh = instance_dispatch!(client, |cx| {
-                proc_mesh
-                    .spawn(cx, "code_sync_manager", &CodeSyncManagerParams {})
-                    .await?
-            });
-            Ok(Self { actor_mesh })
-        })?
+    fn spawn_blocking(
+        py: Python,
+        client: PyInstance,
+        proc_mesh: &Bound<'_, PyAny>,
+    ) -> PyResult<Self> {
+        if let Ok(v0) = proc_mesh.downcast::<PyProcMesh>() {
+            let proc_mesh = v0.borrow().try_inner()?;
+            signal_safe_block_on(py, async move {
+                let actor_mesh = instance_dispatch!(client, |cx| {
+                    proc_mesh
+                        .spawn(cx, "code_sync_manager", &CodeSyncManagerParams {})
+                        .await?
+                });
+                Ok(Self { actor_mesh })
+            })?
+        } else {
+            let proc_mesh = proc_mesh.downcast::<PyProcMeshV1>()?.borrow().mesh_ref()?;
+            signal_safe_block_on(py, async move {
+                let actor_mesh = instance_dispatch!(client, |cx| {
+                    proc_mesh
+                        .spawn_service(cx, "code_sync_manager", &CodeSyncManagerParams {})
+                        .await
+                        .map_err(|err| PyException::new_err(err.to_string()))?
+                });
+                Ok(Self {
+                    actor_mesh: SharedCell::from(RootActorMesh::from(actor_mesh)),
+                })
+            })?
+        }
     }
 
     #[pyo3(signature = (*, local, remote, method = PyCodeSyncMethod::Rsync {}, auto_reload = false))]
diff --git a/python/monarch/_rust_bindings/monarch_extension/code_sync.pyi b/python/monarch/_rust_bindings/monarch_extension/code_sync.pyi
index 4b91e3492..e93a7c961 100644
--- a/python/monarch/_rust_bindings/monarch_extension/code_sync.pyi
+++ b/python/monarch/_rust_bindings/monarch_extension/code_sync.pyi
@@ -9,9 +9,12 @@
 from pathlib import Path
 from typing import Any, Dict, final
 
-from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcMesh
+from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcMesh as ProcMeshV0
 
 from monarch._rust_bindings.monarch_hyperactor.shape import Shape
+from monarch._rust_bindings.monarch_hyperactor.v1.proc_mesh import (
+    ProcMesh as ProcMeshV1,
+)
 
 class WorkspaceLocation:
     """
@@ -84,7 +87,7 @@ class CodeSyncMeshClient:
     @staticmethod
     def spawn_blocking(
         client: Any,
-        proc_mesh: ProcMesh,
+        proc_mesh: ProcMeshV0 | ProcMeshV1,
     ) -> CodeSyncMeshClient: ...
     async def sync_workspace(
         self,