refactor: Improve C API support

pyproject.toml

@@ -56,7 +56,7 @@ ignore = [
 convention = "google"
 
 [tool.ruff.lint.per-file-ignores]
-"tests/**.py" = ["ANN001", "ANN201", "D", "INP001", "PLR2004", "S101", "S311"]
+"tests/**.py" = ["ANN001", "ANN201", "D", "INP001", "PLR2004", "S101", "S311", "S603"]
 "examples/**.py" = ["ANN001", "ANN201", "D", "INP001", "T201"]
 "run-examples.py" = ["ANN001", "ANN201", "D", "INP001", "S603", "T201"]
 "build_ffi.py" = ["ANN001", "ANN201", "S603"]

tests/conftest.py

@@ -2,13 +2,17 @@
 
 import importlib.util
 import os
+import subprocess
+import sys
 from pathlib import Path
 
 import pkgconfig
 import pytest
 
 from build_ffi import compile_ffi
 
+_REEXEC_MARKER = "_VACCEL_PYTEST_REEXEC_DONE"
+
 
 def pytest_configure():
     # Build ffi lib if package is not installed
@@ -22,6 +26,23 @@ def pytest_configure():
     os.environ["VACCEL_LOG_LEVEL"] = "4"
 
 
+def pytest_cmdline_main(config):
+    _ = config
+
+    # Set library path from pkgconfig
+    if _REEXEC_MARKER not in os.environ:
+        lib_path = os.environ.get("LD_LIBRARY_PATH", "")
+        vaccel_lib_path = pkgconfig.variables("vaccel")["libdir"]
+
+        env = os.environ.copy()
+        env["LD_LIBRARY_PATH"] = (
+            f"{vaccel_lib_path}:{lib_path}" if lib_path else vaccel_lib_path
+        )
+        env[_REEXEC_MARKER] = "1"
+
+        sys.exit(subprocess.call([sys.executable, *sys.argv], env=env))
+
+
 @pytest.fixture(scope="session")
 def vaccel_paths():
     variables = pkgconfig.variables("vaccel")

tests/test_general.py

@@ -1,27 +1,23 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from pathlib import Path
-
-import pytest
-
-from vaccel import Resource, ResourceType, Session
-
-
-@pytest.fixture
-def test_lib(vaccel_paths) -> Path:
-    return vaccel_paths["lib"] / "libmytestlib.so"
+from vaccel import PluginType, Session
 
 
 def test_session():
-    ses_a = Session(flags=0)
+    ses_a = Session()
+    assert ses_a.id > 0
+    assert ses_a.flags == 0
+    assert ses_a.is_remote == 0
+
     ses_b = Session(flags=1)
     assert ses_b.id == ses_a.id + 1
+    assert ses_b.flags == 1
+    assert ses_b.is_remote == 0
 
-
-def test_resource(test_lib):
-    res_a = Resource(test_lib, ResourceType.LIB)
-    res_b = Resource([test_lib, test_lib], ResourceType.LIB)
-    assert res_b.id == res_a.id + 1
+    ses_c = Session(flags=PluginType.GENERIC | PluginType.DEBUG)
+    assert ses_c.id == ses_b.id + 1
+    assert ses_c.flags == PluginType.GENERIC | PluginType.DEBUG
+    assert ses_c.is_remote == 0
 
 
 def test_noop():

tests/test_genop.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from vaccel import Arg, OpType, Session
+from vaccel import Arg, ArgType, OpType, Session
 
 
 @pytest.fixture
@@ -41,9 +41,12 @@ def test_data():
 
 
 def test_exec(test_lib, test_args):
-    arg_read = [OpType.EXEC, test_lib["path"], test_lib["symbol"]]
-    arg_read.extend(test_args["read"])
-    g_arg_read = [Arg(arg) for arg in arg_read]
+    g_arg_read = [
+        Arg(OpType.EXEC, ArgType.UINT8),
+        Arg(test_lib["path"], ArgType.STRING),
+        Arg(test_lib["symbol"], ArgType.STRING),
+    ]
+    g_arg_read += [Arg(arg) for arg in test_args["read"]]
     g_arg_write = [Arg(arg) for arg in test_args["write"]]
 
     session = Session()
@@ -55,20 +58,20 @@ def test_exec(test_lib, test_args):
 
 def test_sgemm(test_data):
     arg_read = [
-        Arg(OpType.BLAS_SGEMM),
-        Arg(test_data["m"]),
-        Arg(test_data["n"]),
-        Arg(test_data["k"]),
-        Arg(test_data["alpha"]),
-        Arg(test_data["a"]),
-        Arg(test_data["lda"]),
-        Arg(test_data["b"]),
-        Arg(test_data["ldb"]),
-        Arg(test_data["beta"]),
-        Arg(test_data["ldc"]),
+        Arg(OpType.BLAS_SGEMM, ArgType.UINT8),
+        Arg(test_data["m"], ArgType.INT64),
+        Arg(test_data["n"], ArgType.INT64),
+        Arg(test_data["k"], ArgType.INT64),
+        Arg(test_data["alpha"], ArgType.FLOAT32),
+        Arg(test_data["a"], ArgType.FLOAT32_ARRAY),
+        Arg(test_data["lda"], ArgType.INT64),
+        Arg(test_data["b"], ArgType.FLOAT32_ARRAY),
+        Arg(test_data["ldb"], ArgType.INT64),
+        Arg(test_data["beta"], ArgType.FLOAT32),
+        Arg(test_data["ldc"], ArgType.INT64),
     ]
     c = [float(0)] * test_data["m"] * test_data["n"]
-    arg_write = [Arg(c)]
+    arg_write = [Arg(c, ArgType.FLOAT32_ARRAY)]
 
     session = Session()
     session.genop(arg_read, arg_write)

tests/test_resource.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from vaccel import Resource, ResourceType, Session
+from vaccel._c_types import CBytes
+
+
+@pytest.fixture
+def test_lib(vaccel_paths) -> Path:
+    return vaccel_paths["lib"] / "libmytestlib.so"
+
+
+@pytest.fixture
+def test_buffer() -> bytes:
+    data = [1.0] * 30
+    data_np = np.array(data, dtype=np.float32)
+    return {
+        "data": data,
+        "data_np": data_np,
+        "data_bytes": np.ascontiguousarray(data_np).tobytes(),
+    }
+
+
+def test_resource(test_lib):
+    res_a = Resource(test_lib, ResourceType.LIB)
+    assert res_a.id > 0
+    assert res_a.remote_id == 0
+
+    res_b = Resource([test_lib, test_lib], ResourceType.LIB)
+    assert res_b.id == res_a.id + 1
+    assert res_b.remote_id == 0
+
+
+def test_resource_from_buffer(test_buffer):
+    res = Resource.from_buffer(test_buffer["data_bytes"], ResourceType.DATA)
+    res_data = res.value.blobs[0].data
+    res_size = res.value.blobs[0].size
+    assert res.id > 0
+    assert res.remote_id == 0
+    assert (
+        CBytes.from_c_obj(res_data, res_size).value == test_buffer["data_bytes"]
+    )
+
+
+def test_resource_from_numpy(test_buffer):
+    res = Resource.from_numpy(test_buffer["data_np"])
+    res_data = res.value.blobs[0].data
+    res_size = res.value.blobs[0].size
+    assert res.id > 0
+    assert res.remote_id == 0
+    assert (
+        CBytes.from_c_obj(res_data, res_size).value == test_buffer["data_bytes"]
+    )
+
+
+def test_resource_register(test_lib):
+    res = Resource(test_lib, ResourceType.LIB)
+    ses = Session()
+
+    res.register(ses)
+    assert res.is_registered(ses)
+
+
+def test_resource_unregister(test_buffer):
+    res = Resource.from_buffer(test_buffer["data_bytes"], ResourceType.DATA)
+    ses = Session()
+
+    res.register(ses)
+    res.unregister(ses)
+    assert not res.is_registered(ses)
+
+
+def test_resource_register_unregister_multi(test_lib, test_buffer):
+    res_a = Resource.from_buffer(test_buffer["data_bytes"], ResourceType.DATA)
+    res_b = Resource(test_lib, ResourceType.LIB)
+    ses_a = Session()
+    ses_b = Session()
+
+    res_a.register(ses_a)
+    assert res_a.is_registered(ses_a)
+    res_b.register(ses_a)
+    assert res_b.is_registered(ses_a)
+
+    res_a.register(ses_b)
+    assert res_a.is_registered(ses_b)
+    res_b.register(ses_b)
+    assert res_b.is_registered(ses_b)
+
+    res_a.sync(ses_a)
+    res_a.sync(ses_b)
+
+    res_a.unregister(ses_a)
+    assert not res_a.is_registered(ses_a)
+    res_b.unregister(ses_a)
+    assert not res_b.is_registered(ses_a)
+    res_a.unregister(ses_b)
+    assert not res_a.is_registered(ses_b)
+    res_b.unregister(ses_b)
+    assert not res_b.is_registered(ses_b)

vaccel/__init__.py

@@ -3,17 +3,20 @@
 """Python API for vAccel."""
 
 from ._version import __version__
-from .arg import Arg
+from .arg import Arg, ArgType
 from .config import Config
 from .op import OpType
+from .plugin import PluginType
 from .resource import Resource, ResourceType
 from .session import Session
 from .vaccel import bootstrap, cleanup
 
 __all__ = [
     "Arg",
+    "ArgType",
     "Config",
     "OpType",
+    "PluginType",
     "Resource",
     "ResourceType",
     "Session",

vaccel/_c_types/types.py

@@ -75,13 +75,18 @@ class CAny(CType):
         _wrapped (CType): The wrapped C object.
     """
 
-    def __init__(self, obj: Any):
+    def __init__(self, obj: Any, *, precision: str | None = None):
         """Initializes a new `CAny` object.
 
         Args:
             obj: The python object to be wrapped as a C type.
+            precision: The C type that will be used to represent the python
+                object type, if the type is numeric.
         """
-        self._wrapped = to_ctype(obj)
+        if precision is not None:
+            self._wrapped = to_ctype(obj, precision=precision)
+        else:
+            self._wrapped = to_ctype(obj)
 
     def _init_c_obj(self):
         msg = "CAny is a generic adapter, not meant for initialization."
@@ -111,11 +116,13 @@ def __repr__(self):
 
 
 @singledispatch
-def to_ctype(value: Any):
+def to_ctype(value: Any, *, precision: str | None = None):
+    _ = precision
     msg = f"No CType wrapper registered for {type(value)}"
     raise TypeError(msg)
 
 
 @to_ctype.register
-def _(value: CType):
+def _(value: CType, *, precision: str | None = None):
+    _ = precision
     return value

vaccel/_c_types/utils.py

@@ -2,7 +2,7 @@
 
 """Utilities for C type conversions."""
 
-from enum import IntEnum
+from enum import IntEnum, IntFlag
 from typing import Any
 
 
@@ -26,36 +26,46 @@ def __init__(self, lib: Any):
         self.lib = lib
         self._cache = {}
 
-    def from_prefix(self, enum_name: str, prefix: str) -> IntEnum:
+    def from_prefix(
+        self,
+        enum_name: str,
+        prefix: str,
+        enum_type: type[IntEnum] | type[IntFlag] = IntEnum,
+    ) -> IntEnum | IntFlag:
         """Generates a Python enum from a C enum prefix.
 
         Dynamically create a Python `IntEnum` from C enum constants that share a
         prefix.
 
         Args:
-            enum_name: The name to give the `IntEnum`.
+            enum_name: The name to give the generated Enum.
             prefix: The prefix of the C constants (e.g., "VACCEL_").
+            enum_type: The Enum base class to use.
 
         Returns:
             A Python IntEnum with values mapped from the C library.
         """
-        if enum_name in self._cache:
-            return self._cache[enum_name]
+        cache_key = (enum_name, enum_type)
+        if cache_key in self._cache:
+            return self._cache[cache_key]
 
         members = {
             attr[len(prefix) :]: getattr(self.lib, attr)
             for attr in dir(self.lib)
             if attr.startswith(prefix)
         }
 
+        if not members:
+            msg = f"No constants found with prefix '{prefix}'"
+            raise ValueError(msg)
+
         # Build docstring
         docstring = self._build_enum_docstring(enum_name, members)
 
-        # Generate enum
-        enum_cls = IntEnum(enum_name, members)
+        enum_cls = enum_type(enum_name, members)
         enum_cls.__doc__ = docstring
 
-        self._cache[enum_name] = enum_cls
+        self._cache[cache_key] = enum_cls
         return enum_cls
 
     def _build_enum_docstring(

vaccel/_c_types/wrappers/cbytes.py

@@ -85,6 +85,8 @@ def __eq__(self, other: "CBytes | bytes | bytearray | memoryview"):
             return self._data == other
         return NotImplemented
 
+    __hash__ = None
+
     def __bytes__(self):
         return bytes(self._data)
 
@@ -102,10 +104,12 @@ def __repr__(self):
 
 
 @to_ctype.register
-def _(value: bytes):
+def _(value: bytes, *, precision: str | None = None):
+    _ = precision
     return CBytes(value)
 
 
 @to_ctype.register
-def _(value: bytearray):
+def _(value: bytearray, *, precision: str | None = None):
+    _ = precision
     return CBytes(value)