Implement mohu-py Crate for High-Performance Python Bindings - #259
Implement mohu-py Crate for High-Performance Python Bindings#259Aryan0819 wants to merge 11 commits into
Conversation
📝 WalkthroughWalkthroughThis PR introduces mohu-py, a PyO3-based Python extension crate that exposes mohu tensors to Python. It includes tensor construction and introspection, zero-copy NumPy buffer protocol access, DLPack interoperability for ML framework integration, and Array API Standard compliance with NumPy ufunc support. ChangesPython Bindings Implementation
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related issues
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 12
🧹 Nitpick comments (2)
docs/src/tensor.rs (1)
77-79: 💤 Low valueMissing
// SAFETY:comment on unsafe function.Document why it's safe to receive and ignore the
_viewpointer here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/src/tensor.rs` around lines 77 - 79, Add a // SAFETY: comment above the unsafe fn __releasebuffer__ explaining why it is safe to accept and ignore the _view: state that the pointer originates from the Python buffer protocol and is only an observer (not owned by Rust), that no dereferencing or mutation of the pointer occurs, and that all actual resource/ownership cleanup is handled by Rust's drop of the backing object (so ignoring the pointer cannot violate memory safety); reference the unsafe function name __releasebuffer__ and parameter _view in the comment so future readers can quickly locate the justification.docs/src/dlpack.rs (1)
1-38: ⚖️ Poor tradeoffDuplicate implementation with same issues as
bindings/mohu-py/src/dlpack.rs.This file appears to be a copy of the DLPack implementation. The same issues apply:
- Memory leak if
PyCapsule_Newfails (line 13 allocation not freed at line 19)- Missing
// SAFETY:comments on the unsafe block (lines 15-22) and unsafe extern function (lines 31-38)Consider whether this duplication is intentional (documentation examples) or if these should share a single implementation.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/src/dlpack.rs` around lines 1 - 38, The duplicated DLPack implementation in docs/src/dlpack.rs repeats the same bugs: __dlpack__ allocates raw_dl_struct via Box::into_raw but does not free it if PyCapsule_New fails, and both the unsafe block in __dlpack__ and the unsafe extern "C" fn dlpack_capsule_deleter lack SAFETY doc comments; either deduplicate this implementation with bindings/mohu-py/src/dlpack.rs or fix here by (1) after calling PyCapsule_New, if capsule.is_null() convert the raw pointer back into a Box to drop it before returning the PyRuntimeError, and (2) add concise // SAFETY: comments above the unsafe block in __dlpack__ and above dlpack_capsule_deleter explaining why the raw pointer usage, casting, and FFI calls are safe.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@bindings/mohu-py/pyproject.toml`:
- Around line 5-14: The pyproject's [project] table is missing a required
version field; to fix, add either a static version entry (version = "X.Y.Z") or
enable dynamic versioning by adding dynamic = ["version"] under [project] and
configure tool.maturin.version to point at Cargo.toml (e.g., set
tool.maturin.version = {from = "cargo"} or appropriate key), ensuring the
project metadata includes a resolvable version for wheel builds; update the
pyproject.toml entries (referencing the [project] table, dynamic = ["version"],
and tool.maturin.version) accordingly.
In `@bindings/mohu-py/python/mohu/__init__.py`:
- Line 1: Remove the unused import of the sys module in the module
initialization file: delete the top-level "import sys" statement from
mohu.__init__.py (the unused import symbol "sys" is not referenced anywhere in
this file), leaving only necessary imports/exports.
In `@bindings/mohu-py/python/mohu/array_api.py`:
- Around line 9-19: The current stubbed implementations of abs and add in
functions abs(x: Tensor, /) -> Tensor and add(x1: Tensor, x2: Tensor, /) ->
Tensor return inputs instead of performing operations; change both to raise
NotImplementedError (with a clear message like "abs not implemented" and "add
not implemented") after validating argument types so callers fail loudly until
the real Rust-backed ops are wired up, keeping the existing isinstance checks in
place for Tensor.
- Line 5: Remove the unused top-level import "import mohu as mh" from this
module (the unused symbol is mh) to eliminate dead code; locate the import
statement in bindings/mohu-py/python/mohu/array_api.py and delete it, ensuring
no other code in functions or classes references mh before committing.
In `@bindings/mohu-py/src/dlpack.rs`:
- Around line 31-38: Add a SAFETY comment above the unsafe extern "C" fn
dlpack_capsule_deleter documenting the contract: this function is invoked only
by Python's capsule destructor with a non-null capsule pointer whose payload was
created with Box::into_raw for a mohu_core::DLManagedTensor and associated with
the "dltensor" name via PyCapsule_SetPointer; the function therefore may safely
call PyCapsule_GetPointer and convert the returned pointer into Box::from_raw to
drop/restore ownership exactly once (ensuring no double-free or use-after-free),
and callers must ensure the capsule name and allocation provenance are correct
and the pointer is properly aligned for DLManagedTensor.
- Around line 15-22: Add a // SAFETY: comment immediately above the unsafe block
that documents the required invariants for safety: that raw_dl_struct was
created via Box::into_raw (so points to a valid heap allocation), the capsule
name (b"dltensor\0") is a null-terminated static C string, the call to
PyCapsule_New transfers ownership of raw_dl_struct to the Python capsule on
success, and that the dlpack_capsule_deleter will free that allocation; also
note that the subsequent null check on capsule ensures we only relinquish
ownership to Python when PyCapsule_New succeeded (otherwise ownership remains
with Rust). Ensure the comment references raw_dl_struct, PyCapsule_New,
dlpack_capsule_deleter, and Bound::from_owned_ptr.
- Around line 13-20: The code leaks the Box allocated by Box::into_raw
(raw_dl_struct) if PyCapsule_New returns null; after creating raw_dl_struct from
self.inner.to_dlpack_struct(), ensure you reclaim/free it on failure by
converting back with Box::from_raw(raw_dl_struct) (dropping it) before returning
the PyRuntimeError. Update the unsafe block around PyCapsule_New so that on
capsule.is_null() you call Box::from_raw(raw_dl_struct) to free the allocation
and then return the error; leave dlpack_capsule_deleter and successful capsule
path unchanged.
In `@bindings/mohu-py/src/tensor.rs`:
- Around line 27-40: The functions dtype_str and shape use invalid Python-style
decorators `@getter`; replace those with the PyO3 attribute macro `#[getter]`
(i.e., annotate the Rust methods with #[getter] so PyO3 exposes them as Python
getters for the struct that holds self.inner), leaving the match on
self.inner.dtype() in dtype_str and the shape().to_vec() logic in shape
unchanged.
- Around line 43-75: The __getbuffer__ implementation is unsafe and violates
buffer-protocol invariants: add a SAFETY comment above unsafe fn __getbuffer__
documenting required invariants (caller guarantees, pointer validity, lifetime
expectations); set (*view).obj to a new reference to the exporting PyObject
(e.g., an owned/new reference to self or Py::new wrapper) instead of null and
implement a matching releasebuffer function that decrements that reference
exactly once to preserve Python lifetime tracking; ensure shape and strides are
stored in exporter-owned arrays of the correct element type (Py_ssize_t/isize)
rather than casting from usize-based slices, expose pointers to those
exporter-owned buffers for (*view).shape and (*view).strides and guarantee they
remain valid for the lifetime of the buffer.
In `@docs/src/tensor.rs`:
- Around line 27-28: The attribute syntax on the PyO3 getters is wrong: replace
the invalid `@getter` occurrences with the correct Rust attribute `#[getter]` on
the getter functions (e.g., change the attribute above fn dtype_str(&self) ->
&'static str to `#[getter]`), and do the same for the other getter function(s)
referenced around lines 37–38 (apply `#[getter]` to those function definitions
as well) so the code compiles with PyO3.
- Around line 55-56: The buffer view currently sets (*view).obj = null_mut(),
which violates PEP 3118 and can cause use-after-free; change the __getbuffer__
implementation to accept slf: &pyo3::PyClassInitializer/Bound<'py, Self> (i.e.,
replace &self with slf: &Bound<'py, Self>) so you can store a live Python-side
reference in view->obj, and assign the exporting object's pointer (from slf) to
(*view).obj so the exporter is kept alive for the lifetime of the buffer view;
update the __getbuffer__ signature and any callers to use the Bound<'py, Self>
pattern and ensure you convert slf to the raw PyObject pointer for view->obj.
- Around line 43-47: Add a SAFETY comment immediately above the unsafe fn
__getbuffer__ documenting the invariants that make the raw pointer writes and
casts safe: state that `view` is assumed non-null and points to a valid mutable
`Py_buffer`, the GIL (or equivalent) is held if required by the Python C-API,
and that all fields written (e.g. buf, len, itemsize, readonly, ndim, format,
shape, strides, suboffsets, internal) are initialized appropriately; explain
that pointers derived from `self` (e.g. the internal data pointer
referenced/cast into `buf`) are properly aligned, non-null, valid for the
lifetime required by the consumer, and will not be mutated/freed concurrently
(no aliasing/mutability violations), and note any invariants about ownership
transfer or required release behavior (e.g. who must call release). Ensure the
comment explicitly names __getbuffer__, Py_buffer, and the self data pointer so
reviewers can locate and verify the invariants.
---
Nitpick comments:
In `@docs/src/dlpack.rs`:
- Around line 1-38: The duplicated DLPack implementation in docs/src/dlpack.rs
repeats the same bugs: __dlpack__ allocates raw_dl_struct via Box::into_raw but
does not free it if PyCapsule_New fails, and both the unsafe block in __dlpack__
and the unsafe extern "C" fn dlpack_capsule_deleter lack SAFETY doc comments;
either deduplicate this implementation with bindings/mohu-py/src/dlpack.rs or
fix here by (1) after calling PyCapsule_New, if capsule.is_null() convert the
raw pointer back into a Box to drop it before returning the PyRuntimeError, and
(2) add concise // SAFETY: comments above the unsafe block in __dlpack__ and
above dlpack_capsule_deleter explaining why the raw pointer usage, casting, and
FFI calls are safe.
In `@docs/src/tensor.rs`:
- Around line 77-79: Add a // SAFETY: comment above the unsafe fn
__releasebuffer__ explaining why it is safe to accept and ignore the _view:
state that the pointer originates from the Python buffer protocol and is only an
observer (not owned by Rust), that no dereferencing or mutation of the pointer
occurs, and that all actual resource/ownership cleanup is handled by Rust's drop
of the backing object (so ignoring the pointer cannot violate memory safety);
reference the unsafe function name __releasebuffer__ and parameter _view in the
comment so future readers can quickly locate the justification.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 0ef77a32-c6ce-4ac9-bfaf-78eb9d00394b
📒 Files selected for processing (9)
Cargo.tomlbindings/mohu-py/pyproject.tomlbindings/mohu-py/python/mohu/__init__.pybindings/mohu-py/python/mohu/array_api.pybindings/mohu-py/src/dlpack.rsbindings/mohu-py/src/lib.rsbindings/mohu-py/src/tensor.rsdocs/src/dlpack.rsdocs/src/tensor.rs
|
@manishworkss @mugiwaraluffy56 @Bbn08 please review it |
2 similar comments
|
@manishworkss @mugiwaraluffy56 @Bbn08 please review it |
|
@manishworkss @mugiwaraluffy56 @Bbn08 please review it |
What
Why
How
Checklist
cargo test --workspacepassescargo clippy --workspace -- -D warningspassescargo fmt --allappliedCHANGELOG.mdupdated (if user-facing change)Close #227
under gssoc 2026
Summary by CodeRabbit
Release Notes