_rust_bindings: expose Selection to python

Summary:
this diff provides a Python binding for the `Selection` type.

includes a class-method constructor `Selection.from_string(…)`. integrates into `monarch_extension` and adds a basic test.

Reviewed By: dulinriley

Differential Revision: D74971945

fbshipit-source-id: f63d6bdd9a9db23cef8f099f72a139bf046d079f

Author: shayne-fletcher

monarch_extension/src/lib.rs

@@ -47,6 +47,11 @@ pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> {
         "monarch_hyperactor.shape",
     )?)?;
 
+    monarch_hyperactor::selection::register_python_bindings(&get_or_add_new_module(
+        module,
+        "monarch_hyperactor.selection",
+    )?)?;
+
     client::register_python_bindings(&get_or_add_new_module(module, "monarch_extension.client")?)?;
     worker::register_python_bindings(&get_or_add_new_module(module, "monarch_extension.worker")?)?;
 

monarch_hyperactor/Cargo.toml

@@ -18,6 +18,7 @@ hyperactor_multiprocess = { version = "0.0.0", path = "../hyperactor_multiproces
 hyperactor_telemetry = { version = "0.0.0", path = "../hyperactor_telemetry" }
 monarch_types = { version = "0.0.0", path = "../monarch_types" }
 ndslice = { version = "0.0.0", path = "../ndslice" }
+nom = "7.1"
 pyo3 = { version = "0.22.6", features = ["anyhow"] }
 pyo3-async-runtimes = { git = "https://github.com/PyO3/pyo3-async-runtimes", rev = "f6bb9b471a5b7765dd770af36e83f26802459621", features = ["attributes", "tokio-runtime"] }
 serde = { version = "1.0.185", features = ["derive", "rc"] }

monarch_hyperactor/src/lib.rs

@@ -9,6 +9,7 @@ pub mod ndslice;
 pub mod proc;
 pub mod proc_mesh;
 pub mod runtime;
+pub mod selection;
 pub mod shape;
 
 use pyo3::Bound;
@@ -48,5 +49,7 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul
 
     hyperactor_mod.add_class::<shape::PyShape>()?;
 
+    hyperactor_mod.add_class::<selection::PySelection>()?;
+
     Ok(())
 }

monarch_hyperactor/src/selection.rs

@@ -0,0 +1,46 @@
+use ndslice::selection::Selection;
+use pyo3::PyResult;
+use pyo3::prelude::*;
+use pyo3::types::PyType;
+
+#[pyclass(name = "Selection", module = "monarch._monarch.selection", frozen)]
+pub struct PySelection {
+    inner: Selection,
+}
+
+#[pymethods]
+impl PySelection {
+    #[getter]
+    fn __repr__(&self) -> String {
+        format!("{:?}", self.inner)
+    }
+}
+
+impl From<Selection> for PySelection {
+    fn from(inner: Selection) -> Self {
+        Self { inner }
+    }
+}
+
+#[pymethods]
+impl PySelection {
+    #[classmethod]
+    #[pyo3(name = "from_string")]
+    pub fn parse(_cls: Bound<'_, PyType>, input: &str) -> PyResult<Self> {
+        // TODO: Make this a utility in ndslice.
+        use ndslice::selection::parse::expression;
+        use nom::combinator::all_consuming;
+
+        let input: String = input.chars().filter(|c| !c.is_whitespace()).collect();
+        let (_, selection) = all_consuming(expression)(&input).map_err(|err| {
+            pyo3::exceptions::PyValueError::new_err(format!("parse error: {err}"))
+        })?;
+
+        Ok(PySelection::from(selection))
+    }
+}
+
+pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    module.add_class::<PySelection>()?;
+    Ok(())
+}

python/monarch/__init__.py

@@ -39,6 +39,7 @@
     from monarch.common.opaque_ref import OpaqueRef
     from monarch.common.pipe import create_pipe, Pipe, remote_generator
     from monarch.common.remote import remote
+    from monarch.common.selection import Selection
     from monarch.common.shape import NDSlice, Shape
     from monarch.common.stream import get_active_stream, Stream
     from monarch.common.tensor import reduce, reduce_, Tensor

python/monarch/_monarch/hyperactor/__init__.py

@@ -28,6 +28,10 @@
     AllocSpec,
 )
 
+from monarch._rust_bindings.monarch_hyperactor.selection import (  # @manual=//monarch/monarch_extension:monarch_extension
+    Selection,
+)
+
 from monarch._rust_bindings.monarch_hyperactor.shape import (  # @manual=//monarch/monarch_extension:monarch_extension
     Shape,
 )
@@ -72,5 +76,6 @@ async def handle_cast(
     "PythonActorMesh",
     "ProcessAllocatorBase",
     "Shape",
+    "Selection",
     "LocalAllocatorBase",
 ]

python/monarch/_monarch/selection/__init__.py

@@ -0,0 +1,7 @@
+from monarch._rust_bindings.monarch_hyperactor.selection import (  # @manual=//monarch/monarch_extension:monarch_extension
+    Selection,
+)
+
+__all__ = [
+    "Selection",
+]

python/monarch/_rust_bindings/hyperactor.pyi

@@ -3,6 +3,7 @@
 from typing import Dict, final, List, Optional, Protocol, Type
 
 from monarch._rust_bindings.hyperactor_extension import Alloc, AllocSpec
+from monarch._rust_bindings.monarch_hyperactor.selection import Selection
 
 from monarch._rust_bindings.monarch_hyperactor.shape import Shape
 

python/monarch/_rust_bindings/monarch_hyperactor/selection.pyi

@@ -0,0 +1,22 @@
+from typing import final
+
+@final
+class Selection:
+    """Opaque representation of a selection expression used to represent
+    constraints over multidimensional shapes.
+
+    Construct via from_string()` and use with mesh APIs to filter,
+    evaluate, or route over structured topologies.
+    """
+    def __repr__(self) -> str: ...
+    @classmethod
+    def from_string(cls, s: str) -> Selection:
+        """Parse a selection expression from a string.
+
+        Accepts a compact string syntax such as `"(*, 0:4)"` or `"0 & (1 | 2)"`,
+            and returns a structured Selection object.
+
+        Raises:
+            ValueError: if the input string is not a valid selection expression.
+        """
+        ...

python/monarch/common/selection.py

@@ -0,0 +1,3 @@
+from monarch._monarch.selection import Selection
+
+__all__ = ["Selection"]

python/tests/_monarch/test_ndslice.py

@@ -6,6 +6,8 @@
 import random
 from unittest import TestCase
 
+from monarch._monarch.selection import Selection
+
 from monarch._monarch.shape import Shape, Slice
 
 
@@ -183,3 +185,17 @@ def test_shape_repr(self) -> None:
             repr(shape),
             'Shape { labels: ["label0", "label1"], slice: Slice { offset: 0, sizes: [2, 3], strides: [3, 1] } }',
         )
+
+
+class TestSelection(TestCase):
+    def test_parse_repr(self) -> None:
+        sel = Selection.from_string("(*, 1:3) & (0, *)")
+        self.assertIsInstance(sel, Selection)
+        self.assertEqual(
+            repr(sel),
+            "Intersection(All(Range(Range(1, Some(3), 1), True)), Range(Range(0, Some(1), 1), All(True)))",
+        )
+
+    def test_parse_invalid(self) -> None:
+        with self.assertRaises(ValueError):
+            Selection.from_string("this is not valid selection syntax")