rixa 0.0.2.dev0__tar.gz → 0.0.2rc1__tar.gz

{rixa-0.0.2.dev0/src/rixa.egg-info → rixa-0.0.2rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rixa
-Version: 0.0.2.dev0
+Version: 0.0.2rc1
 Summary: PMIx bootstrap method for modern AI/ML applications
 Author-email: Mateusz Kapusta <mr.kapusta@student.uw.edu.pl>
 Keywords: PMIx,HPC,distributed-computing,machine-learning,slurm
@@ -17,6 +17,7 @@ License-File: LICENSE
 Provides-Extra: pytorch
 Requires-Dist: torch>=2.7.0; extra == "pytorch"
 Requires-Dist: numpy>=1.26; extra == "pytorch"
+Requires-Dist: ninja; extra == "pytorch"
 Provides-Extra: nvshmem
 Requires-Dist: nvshmem4py; extra == "nvshmem"
 Requires-Dist: numpy>=1.26; extra == "nvshmem"

{rixa-0.0.2.dev0 → rixa-0.0.2rc1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rixa"
-version = "0.0.2.dev"
+version = "0.0.2.rc1"
 description = "PMIx bootstrap method for modern AI/ML applications"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -22,7 +22,7 @@ classifiers = [
 keywords = ["PMIx", "HPC", "distributed-computing", "machine-learning", "slurm"]
 
 [project.optional-dependencies]
-pytorch = ["torch>=2.7.0","numpy>=1.26"]
+pytorch = ["torch>=2.7.0","numpy>=1.26","ninja"]
 nvshmem = ["nvshmem4py","numpy>=1.26"]
 
 [build-system]

{rixa-0.0.2.dev0 → rixa-0.0.2rc1}/src/rixa/pytorch.py

@@ -4,10 +4,22 @@ from rixa.PMIx_core import PMIxStore
 from typing import overload
 import os
 import importlib.util
+from pathlib import Path
+import ctypes
 
 _ext = None
 
 
+def load_pmix():
+    libdir = Path(__file__).resolve().parent.parent / "rixa.libs"
+
+    candidates = list(libdir.glob("libpmix*.so*"))
+    if not candidates:
+        raise RuntimeError("PMIx not found")
+    lib = sorted(candidates)[0]  # or better selection logic
+    ctypes.CDLL(str(lib), mode=ctypes.RTLD_GLOBAL)
+
+
 def _get_ext():
     global _ext
     if _ext is not None:
@@ -28,6 +40,7 @@ def _get_ext():
         )
 
     src_dir = os.path.join(os.path.dirname(__file__), "bindings")
+    load_pmix()
 
     _ext = cpp_ext.load(
         name="_rixa_torch",  # no dots in JIT name
@@ -35,8 +48,11 @@ def _get_ext():
             os.path.join(src_dir, "rixa_pmix_store.c"),
             os.path.join(src_dir, "rixa_C10.cpp"),
         ],
-        extra_cflags=["-O3"],
-        extra_ldflags=["-lpmix"],
+        extra_cflags=[
+            "-O3",
+        ],
+        extra_ldflags=["-Wl,-rpath,$ORIGIN/../rixa.libs"],
+        # extra_ldflags=["-lpmix"],
         extra_include_paths=[
             src_dir,
         ],

{rixa-0.0.2.dev0 → rixa-0.0.2rc1/src/rixa.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rixa
-Version: 0.0.2.dev0
+Version: 0.0.2rc1
 Summary: PMIx bootstrap method for modern AI/ML applications
 Author-email: Mateusz Kapusta <mr.kapusta@student.uw.edu.pl>
 Keywords: PMIx,HPC,distributed-computing,machine-learning,slurm
@@ -17,6 +17,7 @@ License-File: LICENSE
 Provides-Extra: pytorch
 Requires-Dist: torch>=2.7.0; extra == "pytorch"
 Requires-Dist: numpy>=1.26; extra == "pytorch"
+Requires-Dist: ninja; extra == "pytorch"
 Provides-Extra: nvshmem
 Requires-Dist: nvshmem4py; extra == "nvshmem"
 Requires-Dist: numpy>=1.26; extra == "nvshmem"

{rixa-0.0.2.dev0 → rixa-0.0.2rc1}/src/rixa.egg-info/requires.txt

@@ -6,3 +6,4 @@ numpy>=1.26
 [pytorch]
 torch>=2.7.0
 numpy>=1.26
+ninja