monarch_hyperactor/
selection.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use hyperactor_mesh::sel;
10use ndslice::selection::Selection;
11use pyo3::PyResult;
12use pyo3::prelude::*;
13use pyo3::types::PyType;
14
15#[pyclass(
16    name = "Selection",
17    module = "monarch._rust_bindings.monarch_hyperactor.selection",
18    frozen
19)]
20pub struct PySelection {
21    inner: Selection,
22}
23
24impl PySelection {
25    #[allow(dead_code)]
26    pub(crate) fn inner(&self) -> &Selection {
27        &self.inner
28    }
29}
30
31impl From<Selection> for PySelection {
32    fn from(inner: Selection) -> Self {
33        Self { inner }
34    }
35}
36
37#[pymethods]
38impl PySelection {
39    #[getter]
40    fn __repr__(&self) -> String {
41        format!("{:?}", self.inner)
42    }
43
44    /// Parses a selection expression from a string.
45    ///
46    /// This allows you to construct a `PySelection` using the
47    /// selection algebra surface syntax, such as `"(*, 0:4, ?)"`.
48    ///
49    /// Raises:
50    ///     ValueError: If the input string is not a valid selection
51    ///     expression.
52    ///
53    /// Example:
54    ///     PySelection.from_string("(*, 1:3, ?)") # subset of a mesh
55    #[classmethod]
56    #[pyo3(name = "from_string")]
57    pub fn parse(_cls: Bound<'_, PyType>, input: &str) -> PyResult<Self> {
58        let selection = ndslice::selection::parse::parse(input).map_err(|err| {
59            pyo3::exceptions::PyValueError::new_err(format!("parse error: {err}"))
60        })?;
61
62        Ok(PySelection::from(selection))
63    }
64
65    /// Selects all elements in the mesh — use this to mean "route to
66    /// all nodes".
67    ///
68    /// The '*' expression is automatically expanded to match the
69    /// dimensionality of the slice. For example, in a 3D slice, the
70    /// selection becomes `*, *, *`.
71    #[classmethod]
72    pub fn all(_cls: Bound<'_, PyType>) -> Self {
73        PySelection::from(sel!(*))
74    }
75
76    /// Selects one element nondeterministically — use this to mean
77    /// "route to a single random node".
78    ///
79    /// The '?' expression is automatically expanded to match the
80    /// dimensionality of the slice. For example, in a 3D slice, the
81    /// selection becomes `?, ?, ?`.
82    #[classmethod]
83    pub fn any(_cls: Bound<'_, PyType>) -> Self {
84        PySelection::from(sel!(?))
85    }
86}
87
88pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
89    module.add_class::<PySelection>()?;
90    Ok(())
91}
monarch_hyperactor/selection.rs

monarch_hyperactor/
selection.rs
