implement __len__ for MeshTrait (#821)

Summary:
Pull Request resolved: #821

straightforward enough
ghstack-source-id: 302243074
exported-using-ghexport

Reviewed By: mariusae

Differential Revision: D79483634

fbshipit-source-id: 785ab0cec0494bf5276179e87d27898d9caefba6

Author: suo

python/monarch/_src/actor/actor_mesh.py

@@ -624,9 +624,6 @@ def items(self) -> Iterable[Tuple[Point, R]]:
     def __iter__(self) -> Iterator[Tuple[Point, R]]:
         return iter(self.items())
 
-    def __len__(self) -> int:
-        return len(self._shape)
-
     def __repr__(self) -> str:
         return f"ValueMesh({self._shape})"
 

python/monarch/_src/actor/shape.py

@@ -224,5 +224,8 @@ def size(self, dim: Union[None, str, Sequence[str]] = None) -> int:
     def sizes(self) -> dict[str, int]:
         return dict(zip(self._labels, self._ndslice.sizes))
 
+    def __len__(self) -> int:
+        return len(self._ndslice)
+
 
 __all__ = ["NDSlice", "Shape", "MeshTrait"]

python/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)

python/tests/test_python_actors.py

@@ -1004,3 +1004,9 @@ def s(t):
     b = PythonTask.spawn_blocking(lambda: s(0))
     r = PythonTask.select_one([a.task(), b.task()]).block_on()
     assert r == (0, 1)
+
+
+def test_mesh_len():
+    proc_mesh = local_proc_mesh(gpus=12).get()
+    s = proc_mesh.spawn("sync_actor", SyncActor).get()
+    assert 12 == len(s)