rixa 0.0.2.dev0__tar.gz → 0.0.2rc2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rixa
-Version: 0.0.2.dev0
+Version: 0.0.2rc2
 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
@@ -14,9 +14,10 @@ Classifier: Operating System :: POSIX :: Linux
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Provides-Extra: pytorch
-Requires-Dist: torch>=2.7.0; extra == "pytorch"
-Requires-Dist: numpy>=1.26; extra == "pytorch"
+Provides-Extra: torch
+Requires-Dist: torch>=2.8.0; extra == "torch"
+Requires-Dist: numpy>=1.26; extra == "torch"
+Requires-Dist: ninja; extra == "torch"
 Provides-Extra: nvshmem
 Requires-Dist: nvshmem4py; extra == "nvshmem"
 Requires-Dist: numpy>=1.26; extra == "nvshmem"

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

@@ -1,6 +1,6 @@
 [project]
 name = "rixa"
-version = "0.0.2.dev"
+version = "0.0.2.rc2"
 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"]
+torch = ["torch>=2.8.0","numpy>=1.26","ninja"]
 nvshmem = ["nvshmem4py","numpy>=1.26"]
 
 [build-system]

{rixa-0.0.2.dev0 → rixa-0.0.2rc2}/src/rixa/bindings/core.c

@@ -153,8 +153,6 @@ static PyObject *get(PyObject *self, PyObject *args) {
 
 // 5. WATI
 static PyObject *wait_for_keys(PyObject *self, PyObject *args) {
-  float delta_T = 0.1;   // Fraction of second for every retry
-  float total_sleep = 0; // total amount of time spend on sleeping
   PyPMIx *self_pmix = (PyPMIx *)self;
   PyObject *keys_list;
   int timeout;

{rixa-0.0.2.dev0 → rixa-0.0.2rc2}/src/rixa/bindings/rixa_C10.cpp

@@ -2,14 +2,13 @@
 #include <cstdint>
 #include <pmix.h>
 #include <pybind11/pybind11.h>
-#include <pybind11/stl.h> // for std::vector, std::string conversions
+#include <pybind11/stl.h>
 #include <stdexcept>
 #include <string>
 #include <torch/csrc/distributed/c10d/Store.hpp>
 #include <torch/csrc/utils/pybind.h>
 
 namespace py = pybind11;
-// include your C header
 extern "C" {
 #include "rixa_pmix_store.h"
 }

{rixa-0.0.2.dev0 → rixa-0.0.2rc2}/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.2rc2/src/rixa.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rixa
-Version: 0.0.2.dev0
+Version: 0.0.2rc2
 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
@@ -14,9 +14,10 @@ Classifier: Operating System :: POSIX :: Linux
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Provides-Extra: pytorch
-Requires-Dist: torch>=2.7.0; extra == "pytorch"
-Requires-Dist: numpy>=1.26; extra == "pytorch"
+Provides-Extra: torch
+Requires-Dist: torch>=2.8.0; extra == "torch"
+Requires-Dist: numpy>=1.26; extra == "torch"
+Requires-Dist: ninja; extra == "torch"
 Provides-Extra: nvshmem
 Requires-Dist: nvshmem4py; extra == "nvshmem"
 Requires-Dist: numpy>=1.26; extra == "nvshmem"

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

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