Add actor to sync conda environments (#860)

Summary:
Pull Request resolved: #860

Adds support to coda sync to handle syncing conda environments, via:

```
await proc_mesh.sync_workspace(conda=True)
```

Reviewed By: LucasLLC, highker

Differential Revision: D78458897

fbshipit-source-id: 173f9f986a40592872bdbff24223829a5374924f

Author: andrewjcg

monarch_extension/src/code_sync.rs

@@ -12,13 +12,13 @@ use std::path::PathBuf;
 
 use anyhow::Result;
 use futures::TryFutureExt;
-use futures::future::try_join_all;
 use hyperactor_mesh::Mesh;
 use hyperactor_mesh::RootActorMesh;
 use hyperactor_mesh::shared_cell::SharedCell;
 use monarch_hyperactor::code_sync::WorkspaceLocation;
 use monarch_hyperactor::code_sync::manager::CodeSyncManager;
 use monarch_hyperactor::code_sync::manager::CodeSyncManagerParams;
+use monarch_hyperactor::code_sync::manager::CodeSyncMethod;
 use monarch_hyperactor::code_sync::manager::WorkspaceConfig;
 use monarch_hyperactor::code_sync::manager::WorkspaceShape;
 use monarch_hyperactor::code_sync::manager::code_sync_mesh;
@@ -118,6 +118,44 @@ impl RemoteWorkspace {
     }
 }
 
+#[pyclass(
+    frozen,
+    name = "CodeSyncMethod",
+    module = "monarch._rust_bindings.monarch_extension.code_sync"
+)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
+enum PyCodeSyncMethod {
+    Rsync,
+    CondaSync,
+}
+
+impl From<PyCodeSyncMethod> for CodeSyncMethod {
+    fn from(method: PyCodeSyncMethod) -> CodeSyncMethod {
+        match method {
+            PyCodeSyncMethod::Rsync => CodeSyncMethod::Rsync,
+            PyCodeSyncMethod::CondaSync => CodeSyncMethod::CondaSync,
+        }
+    }
+}
+
+#[pymethods]
+impl PyCodeSyncMethod {
+    #[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,)))
+    }
+}
+
 #[pyclass(
     frozen,
     name = "WorkspaceConfig",
@@ -127,14 +165,19 @@ impl RemoteWorkspace {
 struct PyWorkspaceConfig {
     local: PathBuf,
     remote: RemoteWorkspace,
+    method: PyCodeSyncMethod,
 }
 
 #[pymethods]
 impl PyWorkspaceConfig {
     #[new]
-    #[pyo3(signature = (*, local, remote))]
-    fn new(local: PathBuf, remote: RemoteWorkspace) -> Self {
-        Self { local, remote }
+    #[pyo3(signature = (*, local, remote, method = PyCodeSyncMethod::Rsync))]
+    fn new(local: PathBuf, remote: RemoteWorkspace, method: PyCodeSyncMethod) -> Self {
+        Self {
+            local,
+            remote,
+            method,
+        }
     }
 }
 
@@ -152,6 +195,7 @@ impl CodeSyncMeshClient {
         actor_mesh: SharedCell<RootActorMesh<'static, CodeSyncManager>>,
         local: PathBuf,
         remote: RemoteWorkspace,
+        method: CodeSyncMethod,
         auto_reload: bool,
     ) -> Result<()> {
         let actor_mesh = actor_mesh.borrow()?;
@@ -164,7 +208,9 @@ impl CodeSyncMeshClient {
             location: remote.location.into(),
             shape,
         };
-        code_sync_mesh(&actor_mesh, local, remote, auto_reload).await?;
+        code_sync_mesh(&actor_mesh, local, remote, method, auto_reload)
+            .await
+            .map_err(|err| PyRuntimeError::new_err(format!("{:#?}", err)))?;
         Ok(())
     }
 }
@@ -183,12 +229,13 @@ impl CodeSyncMeshClient {
         })?
     }
 
-    #[pyo3(signature = (*, local, remote, auto_reload = false))]
+    #[pyo3(signature = (*, local, remote, method = PyCodeSyncMethod::Rsync, auto_reload = false))]
     fn sync_workspace<'py>(
         &self,
         py: Python<'py>,
         local: PathBuf,
         remote: RemoteWorkspace,
+        method: PyCodeSyncMethod,
         auto_reload: bool,
     ) -> PyResult<Bound<'py, PyAny>> {
         monarch_hyperactor::runtime::future_into_py(
@@ -197,6 +244,7 @@ impl CodeSyncMeshClient {
                 self.actor_mesh.clone(),
                 local,
                 remote,
+                method.into(),
                 auto_reload,
             )
             .err_into(),
@@ -213,21 +261,27 @@ impl CodeSyncMeshClient {
         let actor_mesh = self.actor_mesh.clone();
         monarch_hyperactor::runtime::future_into_py(
             py,
-            try_join_all(workspaces.into_iter().map(|workspace| {
-                CodeSyncMeshClient::sync_workspace_(
-                    actor_mesh.clone(),
-                    workspace.local,
-                    workspace.remote,
-                    auto_reload,
-                )
-            }))
+            async move {
+                for workspace in workspaces.into_iter() {
+                    CodeSyncMeshClient::sync_workspace_(
+                        actor_mesh.clone(),
+                        workspace.local,
+                        workspace.remote,
+                        workspace.method.into(),
+                        auto_reload,
+                    )
+                    .await?
+                }
+                anyhow::Ok(())
+            }
             .err_into(),
         )
     }
 }
 
 pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
     module.add_class::<CodeSyncMeshClient>()?;
+    module.add_class::<PyCodeSyncMethod>()?;
     module.add_class::<PyWorkspaceConfig>()?;
     module.add_class::<PyWorkspaceLocation>()?;
     module.add_class::<PyWorkspaceShape>()?;

monarch_hyperactor/Cargo.toml

@@ -25,8 +25,10 @@ hyperactor_mesh = { version = "0.0.0", path = "../hyperactor_mesh" }
 hyperactor_multiprocess = { version = "0.0.0", path = "../hyperactor_multiprocess" }
 hyperactor_telemetry = { version = "0.0.0", path = "../hyperactor_telemetry" }
 inventory = "0.3.8"
+lazy_errors = "0.10.1"
 lazy_static = "1.5"
 libc = "0.2.139"
+monarch_conda = { version = "0.0.0", path = "../monarch_conda" }
 monarch_types = { version = "0.0.0", path = "../monarch_types" }
 ndslice = { version = "0.0.0", path = "../ndslice" }
 nix = { version = "0.30.1", features = ["dir", "event", "hostname", "inotify", "ioctl", "mman", "mount", "net", "poll", "ptrace", "reboot", "resource", "sched", "signal", "term", "time", "user", "zerocopy"] }

monarch_hyperactor/src/code_sync.rs

@@ -7,6 +7,7 @@
  */
 
 pub mod auto_reload;
+pub mod conda_sync;
 pub mod manager;
 pub mod rsync;
 mod workspace;

monarch_hyperactor/src/code_sync/conda_sync.rs

@@ -0,0 +1,179 @@
+/*
+ * 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::collections::HashMap;
+use std::path::PathBuf;
+
+use anyhow::Result;
+use async_trait::async_trait;
+use futures::FutureExt;
+use futures::StreamExt;
+use futures::TryStreamExt;
+use hyperactor::Actor;
+use hyperactor::Bind;
+use hyperactor::Handler;
+use hyperactor::Named;
+use hyperactor::PortRef;
+use hyperactor::Unbind;
+use hyperactor_mesh::actor_mesh::ActorMesh;
+use hyperactor_mesh::connect::Connect;
+use hyperactor_mesh::connect::accept;
+use hyperactor_mesh::sel;
+use lazy_errors::ErrorStash;
+use lazy_errors::OrStash;
+use lazy_errors::StashedResult;
+use lazy_errors::TryCollectOrStash;
+use monarch_conda::sync::Action;
+use monarch_conda::sync::receiver;
+use monarch_conda::sync::sender;
+use ndslice::Selection;
+use serde::Deserialize;
+use serde::Serialize;
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
+
+use crate::code_sync::WorkspaceLocation;
+
+/// Represents the result of an conda sync operation with details about what was transferred
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Named)]
+pub struct CondaSyncResult {
+    /// All changes that occurred during the sync operation
+    pub changes: HashMap<PathBuf, Action>,
+}
+
+#[derive(Debug, Clone, Named, Serialize, Deserialize, Bind, Unbind)]
+pub struct CondaSyncMessage {
+    /// The connect message to create a duplex bytestream with the client.
+    pub connect: PortRef<Connect>,
+    /// A port to send back the result or any errors.
+    pub result: PortRef<Result<CondaSyncResult, String>>,
+    /// The location of the workspace to sync.
+    pub workspace: WorkspaceLocation,
+}
+
+#[derive(Debug, Named, Serialize, Deserialize)]
+pub struct CondaSyncParams {}
+
+#[derive(Debug)]
+#[hyperactor::export(spawn = true, handlers = [CondaSyncMessage { cast = true }])]
+pub struct CondaSyncActor {}
+
+#[async_trait]
+impl Actor for CondaSyncActor {
+    type Params = CondaSyncParams;
+
+    async fn new(CondaSyncParams {}: Self::Params) -> Result<Self> {
+        Ok(Self {})
+    }
+}
+
+#[async_trait]
+impl Handler<CondaSyncMessage> for CondaSyncActor {
+    async fn handle(
+        &mut self,
+        cx: &hyperactor::Context<Self>,
+        CondaSyncMessage {
+            workspace,
+            connect,
+            result,
+        }: CondaSyncMessage,
+    ) -> Result<(), anyhow::Error> {
+        let res = async {
+            let workspace = workspace.resolve()?;
+            let (connect_msg, completer) = Connect::allocate(cx.self_id().clone(), cx);
+            connect.send(cx, connect_msg)?;
+            let (mut read, mut write) = completer.complete().await?.into_split();
+            let changes_result = receiver(&workspace, &mut read, &mut write).await;
+
+            // Shutdown our end, then read from the other end till exhaustion to avoid undeliverable
+            // message spam.
+            write.shutdown().await?;
+            let mut buf = vec![];
+            read.read_to_end(&mut buf).await?;
+
+            anyhow::Ok(CondaSyncResult {
+                changes: changes_result?,
+            })
+        }
+        .await;
+        result.send(cx, res.map_err(|e| format!("{:#?}", e)))?;
+        Ok(())
+    }
+}
+
+pub async fn conda_sync_mesh<M>(
+    actor_mesh: &M,
+    local_workspace: PathBuf,
+    remote_workspace: WorkspaceLocation,
+) -> Result<Vec<CondaSyncResult>>
+where
+    M: ActorMesh<Actor = CondaSyncActor>,
+{
+    let mailbox = actor_mesh.proc_mesh().client();
+    let (conns_tx, conns_rx) = mailbox.open_port::<Connect>();
+
+    let (res1, res2) = futures::future::join(
+        conns_rx
+            .take(actor_mesh.shape().slice().len())
+            .err_into::<anyhow::Error>()
+            .try_for_each_concurrent(None, |connect| async {
+                let (mut read, mut write) = accept(mailbox, mailbox.actor_id().clone(), connect)
+                    .await?
+                    .into_split();
+                let res = sender(&local_workspace, &mut read, &mut write).await;
+
+                // Shutdown our end, then read from the other end till exhaustion to avoid undeliverable
+                // message spam.
+                write.shutdown().await?;
+                let mut buf = vec![];
+                read.read_to_end(&mut buf).await?;
+
+                res
+            })
+            .boxed(),
+        async move {
+            let (result_tx, result_rx) = mailbox.open_port::<Result<CondaSyncResult, String>>();
+            actor_mesh.cast(
+                mailbox,
+                sel!(*),
+                CondaSyncMessage {
+                    connect: conns_tx.bind(),
+                    result: result_tx.bind(),
+                    workspace: remote_workspace,
+                },
+            )?;
+
+            // Wait for all actors to report result.
+            let results = result_rx
+                .take(actor_mesh.shape().slice().len())
+                .try_collect::<Vec<_>>()
+                .await?;
+
+            // Combine all errors into one.
+            let mut errs = ErrorStash::<_, _, anyhow::Error>::new(|| "remote failures");
+            match results
+                .into_iter()
+                .map(|res| res.map_err(anyhow::Error::msg))
+                .try_collect_or_stash::<Vec<_>>(&mut errs)
+            {
+                StashedResult::Ok(results) => anyhow::Ok(results),
+                StashedResult::Err(_) => Err(errs.into_result().unwrap_err().into()),
+            }
+        },
+    )
+    .await;
+
+    // Combine code sync handler and cast errors into one.
+    let mut errs = ErrorStash::<_, _, anyhow::Error>::new(|| "code sync failed");
+    res1.or_stash(&mut errs);
+    if let StashedResult::Ok(results) = res2.or_stash(&mut errs) {
+        errs.into_result()?;
+        return Ok(results);
+    }
+    Err(errs.into_result().unwrap_err().into())
+}