Remove manual monarch installation from CI (#72)

* Remove manual monarch installation from CI

torchmonarch is already declared in pyproject.toml dependencies,
so manual wheel installation is redundant. Let pip install it
automatically like the forge repo does.

* Remove manual torch installation from CI

torch==2.9.0 is already declared in pyproject.toml dependencies,
so manual installation is redundant.

* Adapt CI workflow to match torchforge pattern

- Use conda-incubator/setup-miniconda for environment setup
- Use uv package manager for installation
- Add uv configuration with PyTorch index to pyproject.toml
- Add concurrency control and shell defaults

This matches the torchforge workflow and should resolve linking issues.

* Add repository owner check to unit test workflow

* Add workflow_dispatch to allow manual workflow runs

* Use GPU runner to match forge workflow

* Comment out _proc_mesh.stop() calls due to monarch bug

- Comment out all _proc_mesh.stop() calls in test files
- Add TODO comments to investigate monarch bug
- Tests are failing due to these calls
- Files updated: test_store.py, test_tensor_slice.py, test_state_dict.py, test_large_tensors.py, test_resharding_basic.py

Author: casteryh

.github/workflows/unit_test.yaml

@@ -2,34 +2,38 @@ name: Unit Test
 
 on:
   pull_request:
+  workflow_dispatch:
 
+concurrency:
+  group: unit-test-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
+  cancel-in-progress: true
+
+defaults:
+  run:
+    shell: bash -l -eo pipefail {0}
 
 jobs:
   unit_tests:
-    runs-on: ubuntu-latest
+    if: github.repository_owner == 'meta-pytorch'
+    runs-on: linux.g5.12xlarge.nvidia.gpu
     timeout-minutes: 60
     strategy:
       matrix:
         python-version: ['3.10']
     steps:
       - name: Check out repo
         uses: actions/checkout@v4
-      - name: Setup Python
-        uses: actions/setup-python@v4
+      - name: Setup conda env
+        uses: conda-incubator/setup-miniconda@v2
         with:
+          auto-update-conda: true
+          miniconda-version: "latest"
+          activate-environment: test
           python-version: ${{ matrix.python-version }}
-      - name: Verify Python version
-        run: |
-          python --version
-          python -c "import sys; print(f'Python version: {sys.version_info.major}.{sys.version_info.minor}')"
       - name: Update pip
         run: python -m pip install --upgrade pip
-      - name: Install torch
-        run: python -m pip install torch
-      - name: Install dependencies
-        run: python -m pip install -e ".[dev]"
-      - name: Install monarch from local wheel
-        run: python -m pip install assets/ci/monarch-0.0.1-cp310-cp310-linux_x86_64.whl --force-reinstall
+      - name: Install torchstore
+        run: pip install uv && uv pip install . && uv pip install .[dev]
       - name: Run slice tests (test_slice.py) with coverage
         run: |
           TORCHSTORE_RDMA_ENABLED=0 \

assets/ci/monarch-0.0.1-cp310-cp310-linux_x86_64.whl

@@ -46,3 +46,15 @@ dev = [
 [tool.setuptools.packages.find]
 include = ["torchstore*"]
 exclude = ["assets*", "tests*", "docs*"]
+
+# pytorch
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url  = "https://download.pytorch.org/whl/cu128"
+
+[tool.uv.sources]
+torch = { index = "pytorch-cu128" }
+
+[tool.uv]
+index-strategy = "unsafe-best-match"
+prerelease = "allow"

pyproject.toml

@@ -97,7 +97,8 @@ async def get(self):
         await actor.get.call_one()
         # TODO: assert equal tensors from put/get
     finally:
-        await actor._proc_mesh.stop()
+        # TODO: Investigate monarch bug with proc_mesh.stop()
+        # await actor._proc_mesh.stop()
         await ts.shutdown()
 
 

tests/test_large_tensors.py

@@ -247,9 +247,11 @@ async def _test_resharding(
 
         # teardown distributed or the next test will complain
         await put_mesh.destroy_process_group.call()
-        await put_mesh._proc_mesh.stop()
+        # TODO: Investigate monarch bug with proc_mesh.stop()
+        # await put_mesh._proc_mesh.stop()
         await get_mesh.destroy_process_group.call()
-        await get_mesh._proc_mesh.stop()
+        # TODO: Investigate monarch bug with proc_mesh.stop()
+        # await get_mesh._proc_mesh.stop()
         await ts.shutdown()
 
 

tests/test_resharding_basic.py

@@ -265,8 +265,9 @@ async def test_dcp_sharding_parity(strategy_params, use_rdma):
                             f"Assertion failed on rank {coord.rank} ({save_mesh_shape=} {get_mesh_shape=}): {e}"
                         ) from e
         finally:
-            await save_world._proc_mesh.stop()
-            await get_world._proc_mesh.stop()
+            # TODO: Investigate monarch bug with proc_mesh.stop()
+            # await save_world._proc_mesh.stop()
+            # await get_world._proc_mesh.stop()
             await ts.shutdown()
 
 

tests/test_state_dict.py

@@ -79,8 +79,9 @@ async def get(self, rank_offset=0):
             expected = torch.tensor([other_rank + 1] * 10)
             assert torch.equal(expected, val), f"{expected} != {val}"
     finally:
-        await actor_mesh_0._proc_mesh.stop()
-        await actor_mesh_1._proc_mesh.stop()
+        # TODO: Investigate monarch bug with proc_mesh.stop()
+        # await actor_mesh_0._proc_mesh.stop()
+        # await actor_mesh_1._proc_mesh.stop()
         await ts.shutdown()
 
 
@@ -142,8 +143,9 @@ def __eq__(self, other: object) -> bool:
                 assert expected == val, f"{expected.val} != {val.val}"
 
     finally:
-        await actor_mesh_0._proc_mesh.stop()
-        await actor_mesh_1._proc_mesh.stop()
+        # TODO: Investigate monarch bug with proc_mesh.stop()
+        # await actor_mesh_0._proc_mesh.stop()
+        # await actor_mesh_1._proc_mesh.stop()
         await ts.shutdown()
 
 
@@ -211,7 +213,8 @@ async def exists(self, key):
                 assert exists_result
 
     finally:
-        await actor_mesh._proc_mesh.stop()
+        # TODO: Investigate monarch bug with proc_mesh.stop()
+        # await actor_mesh._proc_mesh.stop()
         await ts.shutdown()
 
 
@@ -292,7 +295,8 @@ async def get(self, key):
             await actor_mesh.get.call("tensor_key_0")
 
     finally:
-        await actor_mesh._proc_mesh.stop()
+        # TODO: Investigate monarch bug with proc_mesh.stop()
+        # await actor_mesh._proc_mesh.stop()
         await ts.shutdown()
 
 

tests/test_store.py

@@ -81,7 +81,8 @@ async def put(self, key, tensor):
         assert tensor_slice.shape == (5, 10)
 
     finally:
-        await put_actor_mesh._proc_mesh.stop()
+        # TODO: Investigate monarch bug with proc_mesh.stop()
+        # await put_actor_mesh._proc_mesh.stop()
         await ts.shutdown()
 
 
@@ -147,7 +148,8 @@ async def test_put_dtensor_get_full_tensor():
         finally:
             # Clean up process groups
             await put_mesh.destroy_process_group.call()
-            await put_mesh._proc_mesh.stop()
+            # TODO: Investigate monarch bug with proc_mesh.stop()
+            # await put_mesh._proc_mesh.stop()
             await ts.shutdown()
 
 
@@ -232,7 +234,8 @@ async def test_dtensor_fetch_slice():
         finally:
             if put_mesh is not None:
                 await put_mesh.destroy_process_group.call()
-                await put_mesh._proc_mesh.stop()
+                # TODO: Investigate monarch bug with proc_mesh.stop()
+                # await put_mesh._proc_mesh.stop()
             await ts.shutdown()
 
 
@@ -280,7 +283,8 @@ async def test_partial_put():
         finally:
             # Clean up process groups
             await put_mesh.destroy_process_group.call()
-            await put_mesh._proc_mesh.stop()
+            # TODO: Investigate monarch bug with proc_mesh.stop()
+            # await put_mesh._proc_mesh.stop()
             await ts.shutdown()