Refactor workspace location enum (#404)

Summary:
Pull Request resolved: #404

This does some cleanup of the workspace location enum we use to
represent the remote side location, to generalize for use in other
locations (right now it's overly tied to the rsync module).

Reviewed By: suo

Differential Revision: D77456893

fbshipit-source-id: 7ef89c7ccdcc8fe3d10feded31a23d2a14d7532b

Author: andrewjcg

hyperactor_mesh/src/code_sync.rs

@@ -7,3 +7,6 @@
  */
 
 pub mod rsync;
+mod workspace;
+
+pub use workspace::WorkspaceLocation;

hyperactor_mesh/src/code_sync/rsync.rs

@@ -37,6 +37,7 @@ use tokio::process::Child;
 use tokio::process::Command;
 
 use crate::actor_mesh::ActorMesh;
+use crate::code_sync::WorkspaceLocation;
 use crate::connect::Connect;
 use crate::connect::accept;
 use crate::connect::connect_mesh;
@@ -155,15 +156,9 @@ impl RsyncDaemon {
     }
 }
 
-#[derive(Clone, Debug, Named, Serialize, Deserialize)]
-pub enum Workspace {
-    Constant(PathBuf),
-    FromEnvVar(String),
-}
-
 #[derive(Debug, Named, Serialize, Deserialize)]
 pub struct RsyncParams {
-    pub workspace: Workspace,
+    pub workspace: WorkspaceLocation,
 }
 
 #[derive(Debug)]
@@ -180,10 +175,7 @@ impl Actor for RsyncActor {
     type Params = RsyncParams;
 
     async fn new(RsyncParams { workspace }: Self::Params) -> Result<Self> {
-        let workspace = match workspace {
-            Workspace::Constant(p) => p,
-            Workspace::FromEnvVar(v) => PathBuf::from(std::env::var(v)?),
-        };
+        let workspace = workspace.resolve()?;
         let daemon = RsyncDaemon::spawn(TcpListener::bind(("::1", 0)).await?, &workspace).await?;
         Ok(Self { daemon })
     }

hyperactor_mesh/src/code_sync/workspace.rs

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use std::path::PathBuf;
+
+use anyhow::Result;
+use serde::Deserialize;
+use serde::Serialize;
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub enum WorkspaceLocation {
+    Constant(PathBuf),
+    FromEnvVar(String),
+}
+
+impl WorkspaceLocation {
+    pub fn resolve(&self) -> Result<PathBuf> {
+        Ok(match self {
+            WorkspaceLocation::Constant(p) => p.clone(),
+            WorkspaceLocation::FromEnvVar(v) => PathBuf::from(std::env::var(v)?),
+        })
+    }
+}

monarch_extension/Cargo.toml

@@ -15,6 +15,7 @@ crate-type = ["cdylib"]
 
 [dependencies]
 anyhow = "1.0.98"
+bincode = "1.3.3"
 clap = { version = "4.5.38", features = ["derive", "env", "string", "unicode", "wrap_help"] }
 controller = { version = "0.0.0", path = "../controller", optional = true }
 hyperactor = { version = "0.0.0", path = "../hyperactor" }
@@ -31,6 +32,7 @@ nccl-sys = { path = "../nccl-sys", optional = true }
 ndslice = { version = "0.0.0", path = "../ndslice" }
 pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods"] }
 pyo3-async-runtimes = { git = "https://github.com/PyO3/pyo3-async-runtimes", rev = "c5a3746f110b4d246556b0f6c29f5f555919eee3", features = ["attributes", "tokio-runtime"] }
+serde = { version = "1.0.185", features = ["derive", "rc"] }
 tokio = { version = "1.45.0", features = ["full", "test-util", "tracing"] }
 torch-sys = { version = "0.0.0", path = "../torch-sys", optional = true }
 torch-sys-cuda = { version = "0.0.0", path = "../torch-sys-cuda", optional = true }

monarch_extension/src/code_sync.rs

@@ -13,32 +13,66 @@ use std::sync::Arc;
 
 use hyperactor_mesh::RootActorMesh;
 use hyperactor_mesh::SlicedActorMesh;
+use hyperactor_mesh::code_sync::WorkspaceLocation;
 use hyperactor_mesh::code_sync::rsync;
 use hyperactor_mesh::proc_mesh::SharedSpawnable;
 use hyperactor_mesh::shape::Shape;
 use monarch_hyperactor::proc_mesh::PyProcMesh;
 use monarch_hyperactor::runtime::signal_safe_block_on;
 use monarch_hyperactor::shape::PyShape;
 use pyo3::Bound;
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
+use pyo3::types::PyBytes;
 use pyo3::types::PyModule;
+use serde::Deserialize;
+use serde::Serialize;
 
-#[pyclass(frozen, module = "monarch._rust_bindings.monarch_extension.code_sync")]
-#[derive(Clone, Debug)]
-enum RemoteWorkspace {
+#[pyclass(
+    frozen,
+    name = "WorkspaceLocation",
+    module = "monarch._rust_bindings.monarch_extension.code_sync"
+)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
+enum PyWorkspaceLocation {
     Constant(PathBuf),
     FromEnvVar(String),
 }
 
-impl From<RemoteWorkspace> for rsync::Workspace {
-    fn from(workspace: RemoteWorkspace) -> rsync::Workspace {
+impl From<PyWorkspaceLocation> for WorkspaceLocation {
+    fn from(workspace: PyWorkspaceLocation) -> WorkspaceLocation {
         match workspace {
-            RemoteWorkspace::Constant(v) => rsync::Workspace::Constant(v),
-            RemoteWorkspace::FromEnvVar(v) => rsync::Workspace::FromEnvVar(v),
+            PyWorkspaceLocation::Constant(v) => WorkspaceLocation::Constant(v),
+            PyWorkspaceLocation::FromEnvVar(v) => WorkspaceLocation::FromEnvVar(v),
         }
     }
 }
 
+#[pymethods]
+impl PyWorkspaceLocation {
+    #[staticmethod]
+    fn from_bytes(bytes: &Bound<'_, PyBytes>) -> PyResult<Self> {
+        bincode::deserialize(bytes.as_bytes())
+            .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))
+    }
+
+    fn __reduce__<'py>(
+        slf: &Bound<'py, Self>,
+    ) -> PyResult<(Bound<'py, PyAny>, (Bound<'py, PyBytes>,))> {
+        let bytes = bincode::serialize(&*slf.borrow())
+            .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))?;
+        let py_bytes = PyBytes::new(slf.py(), &bytes);
+        Ok((slf.as_any().getattr("from_bytes")?, (py_bytes,)))
+    }
+
+    fn resolve(&self) -> PyResult<PathBuf> {
+        let loc: WorkspaceLocation = self.clone().into();
+        loc.resolve()
+            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))
+    }
+}
+
 #[pyclass(
     frozen,
     name = "RsyncMeshClient",
@@ -59,7 +93,7 @@ impl RsyncMeshClient {
         proc_mesh: &PyProcMesh,
         shape: &PyShape,
         local_workspace: PathBuf,
-        remote_workspace: RemoteWorkspace,
+        remote_workspace: PyWorkspaceLocation,
     ) -> PyResult<Self> {
         let proc_mesh = Arc::clone(&proc_mesh.inner);
         let shape = shape.get_inner().clone();
@@ -92,7 +126,7 @@ impl RsyncMeshClient {
 }
 
 pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
-    module.add_class::<RemoteWorkspace>()?;
+    module.add_class::<PyWorkspaceLocation>()?;
     module.add_class::<RsyncMeshClient>()?;
     Ok(())
 }

python/monarch/_rust_bindings/monarch_extension/code_sync.pyi

@@ -4,24 +4,31 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from pathlib import Path
 from typing import final
 
 from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcMesh
 
 from monarch._rust_bindings.monarch_hyperactor.shape import Shape
 
-class RemoteWorkspace:
+class WorkspaceLocation:
     """
-    Python binding for the Rust RemoteWorkspace enum.
+    Python binding for the Rust WorkspaceLocation enum.
     """
     @final
-    class Constant(RemoteWorkspace):
+    class Constant(WorkspaceLocation):
         def __init__(self, path) -> None: ...
 
     @final
-    class FromEnvVar(RemoteWorkspace):
+    class FromEnvVar(WorkspaceLocation):
         def __init__(self, var) -> None: ...
 
+    def resolve(self) -> Path:
+        """
+        Resolve the workspace location to a Path.
+        """
+        ...
+
 @final
 class RsyncMeshClient:
     """
@@ -32,6 +39,6 @@ class RsyncMeshClient:
         proc_mesh: ProcMesh,
         shape: Shape,
         local_workspace: str,
-        remote_workspace: RemoteWorkspace,
+        remote_workspace: WorkspaceLocation,
     ) -> RsyncMeshClient: ...
     async def sync_workspace(self) -> None: ...

python/monarch/code_sync/__init__.py

@@ -5,6 +5,6 @@
 # LICENSE file in the root directory of this source tree.
 
 from monarch._rust_bindings.monarch_extension.code_sync import (  # noqa: F401
-    RemoteWorkspace,
     RsyncMeshClient,
+    WorkspaceLocation,
 )

python/monarch/proc_mesh.py

@@ -45,7 +45,7 @@
 from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice
 from monarch.actor_mesh import _Actor, _ActorMeshRefImpl, Actor, ActorMeshRef
 
-from monarch.code_sync import RemoteWorkspace, RsyncMeshClient
+from monarch.code_sync import RsyncMeshClient, WorkspaceLocation
 from monarch.common._device_utils import _local_device_count
 from monarch.common.shape import MeshTrait
 from monarch.rdma import RDMAManager
@@ -231,7 +231,7 @@ async def sync_workspace(self) -> None:
                 # TODO(agallagher): Is there a better way to infer/set the local
                 # workspace dir, rather than use PWD?
                 local_workspace=os.getcwd(),
-                remote_workspace=RemoteWorkspace.FromEnvVar("WORKSPACE_DIR"),
+                remote_workspace=WorkspaceLocation.FromEnvVar("WORKSPACE_DIR"),
             )
         await self._rsync_mesh_client.sync_workspace()