torchmonarch-nightly 2025.8.2__cp311-cp311-manylinux2014_x86_64.whl → 2025.9.3__cp311-cp311-manylinux2014_x86_64.whl

tests/test_env_before_cuda.py

@@ -75,9 +75,9 @@ class TestEnvBeforeCuda(unittest.IsolatedAsyncioTestCase):
 
         spec = AllocSpec(AllocConstraints(), gpus=1, hosts=1)
         allocator = LocalAllocator()
-        alloc = await allocator.allocate(spec)
+        alloc = allocator.allocate(spec)
 
-        proc_mesh = await ProcMesh.from_alloc(alloc, setup=setup_cuda_env)
+        proc_mesh = ProcMesh.from_alloc(alloc, setup=setup_cuda_env)
 
         try:
             actor = await proc_mesh.spawn("cuda_init", CudaInitTestActor)
@@ -110,7 +110,7 @@ class TestEnvBeforeCuda(unittest.IsolatedAsyncioTestCase):
             for name, value in cuda_env_vars.items():
                 os.environ[name] = value
 
-        proc_mesh_instance = await proc_mesh(gpus=1, hosts=1, setup=setup_cuda_env)
+        proc_mesh_instance = proc_mesh(gpus=1, hosts=1, setup=setup_cuda_env)
 
         try:
             actor = await proc_mesh_instance.spawn("cuda_init", CudaInitTestActor)
@@ -136,7 +136,7 @@ class TestEnvBeforeCuda(unittest.IsolatedAsyncioTestCase):
             "CUDA_DEVICE_MAX_CONNECTIONS": "1",
         }
 
-        proc_mesh_instance = await proc_mesh(gpus=1, hosts=1, env=cuda_env_vars)
+        proc_mesh_instance = proc_mesh(gpus=1, hosts=1, env=cuda_env_vars)
 
         try:
             actor = await proc_mesh_instance.spawn("cuda_init", CudaInitTestActor)

tests/test_mesh_trait.py

@@ -0,0 +1,38 @@
+# 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.
+
+from typing import Iterable
+
+from monarch._src.actor.shape import MeshTrait, NDSlice, Shape, Slice
+
+
+class Mesh(MeshTrait):
+    """
+    A simple implementor of MeshTrait.
+    """
+
+    def __init__(self, shape: Shape, values: list[int]) -> None:
+        self._shape = shape
+        self._values = values
+
+    def _new_with_shape(self, shape: Shape) -> "Mesh":
+        return Mesh(shape, self._values)
+
+    @property
+    def _ndslice(self) -> NDSlice:
+        return self._shape.ndslice
+
+    @property
+    def _labels(self) -> Iterable[str]:
+        return self._shape.labels
+
+
+def test_len() -> None:
+    s = Slice(offset=0, sizes=[2, 3], strides=[3, 1])
+    shape = Shape(["label0", "label1"], s)
+
+    mesh = Mesh(shape, [1, 2, 3, 4, 5, 6])
+    assert 6 == len(mesh)