link launch and sync conda/workspace locations

Summary:
Make sure the conda/workspace locations during launch map with the
locations when we sync.

Differential Revision: D79516268

Author: highker

python/monarch/_src/actor/proc_mesh.py

@@ -13,6 +13,7 @@
 import threading
 import warnings
 from contextlib import AbstractContextManager
+from pathlib import Path
 
 from typing import (
     Any,
@@ -69,6 +70,7 @@
 from monarch._src.actor.endpoint import endpoint
 from monarch._src.actor.future import DeprecatedNotAFuture, Future
 from monarch._src.actor.shape import MeshTrait
+from monarch.tools.config import Config
 
 HAS_TENSOR_ENGINE = False
 try:
@@ -356,7 +358,11 @@ def rank_tensor(self, dim: str | Sequence[str]) -> "Tensor":
     def rank_tensors(self) -> Dict[str, "Tensor"]:
         return self._device_mesh.ranks
 
-    async def sync_workspace(self, conda: bool = False, auto_reload: bool = False) -> None:
+    # TODO(kiuk): once we have the Workspace dataclass in config,
+    # we should pass in Workspace here instead of the full Config
+    async def sync_workspace(
+        self, config: Config, conda: bool = False, auto_reload: bool = False
+    ) -> None:
         if self._code_sync_client is None:
             self._code_sync_client = CodeSyncMeshClient.spawn_blocking(
                 proc_mesh=await self._proc_mesh_for_asyncio_fixme,
@@ -368,27 +374,27 @@ async def sync_workspace(self, conda: bool = False, auto_reload: bool = False) -
         # The workspace shape (i.e. only perform one rsync per host).
         assert set(self._shape.labels).issubset({"gpus", "hosts"})
 
-        # TODO(agallagher): Is there a better way to infer/set the local
-        # workspace dir, rather than use PWD?
-        workspaces = [
-            WorkspaceConfig(
-                local=os.getcwd(),
-                remote=RemoteWorkspace(
-                    location=WorkspaceLocation.FromEnvVar("WORKSPACE_DIR"),
-                    shape=WorkspaceShape.shared("gpus"),
+        workspaces = []
+        if config.workspace is not None:
+            workspaces.append(
+                WorkspaceConfig(
+                    local=Path(config.workspace),
+                    remote=RemoteWorkspace(
+                        location=WorkspaceLocation.FromEnvVar("WORKSPACE_DIR"),
+                        shape=WorkspaceShape.shared("gpus"),
+                    ),
+                    method=CodeSyncMethod.Rsync,
                 ),
-                method=CodeSyncMethod.Rsync,
-            ),
-        ]
+            )
 
         # If `conda` is set, also sync the currently activated conda env.
         conda_prefix = os.environ.get("CONDA_PREFIX")
         if conda and conda_prefix is not None:
             workspaces.append(
                 WorkspaceConfig(
-                    local=conda_prefix,
+                    local=Path(conda_prefix),
                     remote=RemoteWorkspace(
-                        location=WorkspaceLocation.FromEnvVar("CONDA_PREFIX"),
+                        location=WorkspaceLocation.FromEnvVar("CONDA_DIR"),
                         shape=WorkspaceShape.shared("gpus"),
                     ),
                     method=CodeSyncMethod.CondaSync,

python/monarch/tools/config/__init__.py

@@ -32,6 +32,7 @@ class Config:
 
     scheduler: str = NOT_SET
     scheduler_args: dict[str, Any] = field(default_factory=dict)
+    # TODO: make workspace a list
     workspace: Optional[str] = None
     dryrun: bool = False
     appdef: UnnamedAppDef = field(default_factory=UnnamedAppDef)

python/tests/test_python_actors.py

@@ -36,6 +36,7 @@
     local_proc_mesh,
     proc_mesh,
 )
+from monarch.tools.config import defaults
 from typing_extensions import assert_type
 
 
@@ -769,6 +770,49 @@ async def test_same_actor_twice() -> None:
     await pm.stop()
 
 
+class LsActor(Actor):
+    def __init__(self, workspace: str):
+        self.workspace = workspace
+
+    @endpoint
+    async def ls(self) -> list[str]:
+        return os.listdir(self.workspace)
+
+
+async def test_sync_workspace() -> None:
+    pm = await proc_mesh(gpus=1)
+
+    # create two workspaces: one for local and one for remote
+    with tempfile.TemporaryDirectory() as workspace_src, tempfile.TemporaryDirectory() as workspace_dst:
+        try:
+            os.environ["WORKSPACE_DIR"] = workspace_dst
+            config = defaults.config("slurm", workspace_src)
+            await pm.sync_workspace(config, conda=False, auto_reload=True)
+
+            # now file in remote workspace initially
+            am = await pm.spawn("ls", LsActor, workspace_dst)
+            for item in list(am.ls.call().get()):
+                assert len(item[1]) == 0
+
+            # write a file to local workspace
+            file_path = os.path.join(workspace_src, "new_file")
+            with open(file_path, "w") as f:
+                f.write("hello world")
+                f.flush()
+
+            # force a sync and it should populate on the dst workspace
+            await pm.sync_workspace(config, conda=False, auto_reload=True)
+            for item in list(am.ls.call().get()):
+                assert len(item[1]) == 1
+                assert item[1][0] == "new_file"
+                file_path = os.path.join(workspace_dst, item[1][0])
+                with open(file_path, "r") as f:
+                    assert f.readline() == "hello world"
+
+        finally:
+            os.environ.pop("WORKSPACE_DIR", None)  # Remove if it didn't
+
+
 class TestActorMeshStop(unittest.IsolatedAsyncioTestCase):
     async def test_actor_mesh_stop(self) -> None:
         pm = proc_mesh(gpus=2)