divide out tensor engine tests (#649)

Summary:
Pull Request resolved: #649

Main changes:
- Pulls out tests that require GPUs into `tensor_engine` folder
- Changes `test_cpu.py` step of OSS to run all non tensor engine tests
- Changes `test_gpu.py` step of OSS to run all tensor engine tests.
- Changes wheels.yml to run `test_python_actors` for a minimal set of tests to validate wheel correctness.

Differential Revision: D78989844

Author: allenwang28

.github/workflows/test-cpu.yml

@@ -25,6 +25,9 @@ jobs:
         # Source common setup functions
         source scripts/common-setup.sh
 
+        # Setup test environment
+        setup_test_environment
+
         # Setup test environment
         setup_conda_environment
 
@@ -34,7 +37,4 @@ jobs:
         # Install the built wheel from artifact
         install_wheel_from_artifact
 
-        # Currently a no-op.
-        # Tests requiring tensor engine / GPU need to be identified and flagged to skip.
-        # We will just ensure monarch can be imported successfully.
-        python -c "import monarch; print('Monarch imported successfully')"
+        LC_ALL=C pytest python/tests/ -s -v -m "not oss_skip" --ignore=python/tests/tensor_engine -n 4

.github/workflows/test-cuda.yml

@@ -55,5 +55,5 @@ jobs:
         pyright python/tests/test_python_actors.py
 
         # Run CUDA tests
-        LC_ALL=C pytest python/tests/ -s -v -m "not oss_skip"
-        python python/tests/test_mock_cuda.py
+        LC_ALL=C pytest python/tests/tensor_engine -s -v -m "not oss_skip" -n 4
+        python python/tests/tensor_engine/test_mock_cuda.py

.github/workflows/wheels.yml

@@ -36,6 +36,9 @@ jobs:
         source scripts/common-setup.sh
         setup_build_environment ${{ matrix.python-version }}
 
+        # Setup test environment
+        setup_test_environment
+
         # Setup Tensor Engine dependencies
         setup_tensor_engine
 
@@ -54,7 +57,8 @@ jobs:
         # Run tests
         install_python_test_dependencies
         pip install dist/*.whl
-        python -c "import monarch"
+        LC_ALL=C pytest python/tests/test_python_actors.py -s -v -m "not oss_skip"
+
   publish:
     name: Publish to PyPI
     needs: build

python/monarch/_testing.py

@@ -10,6 +10,7 @@
 import tempfile
 import time
 from contextlib import contextmanager, ExitStack
+from enum import Enum
 from typing import Any, Callable, Dict, Generator, Literal, Optional
 
 import monarch_supervisor
@@ -225,7 +226,7 @@ def exit(
     return dm
 
 
-class BackendType:
+class BackendType(Enum):
     PY = "py"
     RS = "rs"
     MESH = "mesh"

python/tests/test_coalescing.py renamed to python/tests/tensor_engine/test_coalescing.py

@@ -78,7 +78,7 @@ def local_device_mesh(
             num_hosts,
             gpu_per_host,
             activate,
-            backend=str(backend_type),
+            backend=backend_type.value,
         )
 
     @property

python/tests/test_controller.py renamed to python/tests/tensor_engine/test_controller.py

@@ -96,7 +96,9 @@ def local_rust_device_mesh(
     torch.cuda.device_count() < 2,
     reason="Not enough GPUs, this test requires at least 2 GPUs",
 )
-@pytest.mark.parametrize("backend_type", [BackendType.PY, BackendType.RS, "mesh"])
+@pytest.mark.parametrize(
+    "backend_type", [BackendType.PY, BackendType.RS, BackendType.MESH]
+)
 # Set global timeout--sandcastle's timeout is 600s. A test that sandcastle times
 # out is not counted as a failure, so we set a more restrictive timeout to
 # ensure we see a hard failure in CI.
@@ -114,7 +116,7 @@ def local_device_mesh(
             N,
             gpu_per_host,
             activate,
-            backend=str(backend_type),
+            backend=backend_type.value,
         )
 
     def test_errors(self, backend_type):
@@ -176,7 +178,7 @@ def test_sub_mesh_use_only_one(self, backend_type):
             local_x = local_x.result(timeout=20)
             assert torch.equal(local_x, torch.ones(3, 4))
 
-    def test_sub_mesh_process_grop(self, backend_type):
+    def test_sub_mesh_process_group(self, backend_type):
         with self.local_device_mesh(2, 2, backend_type, activate=False) as device_mesh:
             h0 = device_mesh.slice(host=0)
             pg0 = h0.process_group(("gpu",))
@@ -603,7 +605,7 @@ def test_to_mesh_pytree(self, backend_type):
         assert torch.equal(moved_tensor_b, torch.tensor([2.0]))
 
     def test_hanging_error(self, backend_type):
-        if backend_type != "mesh":
+        if backend_type != BackendType.MESH:
             pytest.skip("only relevant for mesh backend")
         with self.local_device_mesh(2, 2, backend_type) as device_mesh:
             remote(lambda: torch.rand(3) + torch.rand(4), propagate=lambda: None)()