Python bindings

.github/workflows/ci.yml

@@ -24,7 +24,7 @@ jobs:
         if: runner.os == 'Linux'
         run: |
           sudo apt-get update
-          sudo apt-get install -y clang libc++-dev libc++abi-dev
+          sudo apt-get install -y clang libc++-dev libc++abi-dev libomp-18-dev
 
       - name: Install OpenMP (MacOS)
         if: runner.os == 'macOS'
@@ -98,3 +98,51 @@ jobs:
           docker run --rm \
           -v "$(pwd):/workspace:Z" \
             tgn-dev /bin/bash -c "make test"
+
+  python-tests:
+    needs: build
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+        python-version: ["3.10"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Download Build Artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: build-${{ matrix.os }}-Release
+          path: build/
+
+      - name: Setup uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      - name: Install Clang Dependencies
+        if: runner.os == 'Linux'
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y clang libc++-dev libc++abi-dev libomp-18-dev
+
+      - name: Install OpenMP (MacOS)
+        if: runner.os == 'macOS'
+        run: brew install libomp
+
+      - name: Install & Test Python Bindings
+        run: make test-python
+
+  python-tests-on-container:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Run Python Suite in Container
+        run: |
+          docker build -t tgn-dev .
+          docker run --rm \
+            -v "$(pwd):/workspace:Z" \
+            tgn-dev /bin/bash -c "make test-python"

.github/workflows/integration.yml

@@ -23,7 +23,7 @@ jobs:
         if: runner.os == 'Linux'
         run: |
           sudo apt-get update
-          sudo apt-get install -y clang libc++-dev libc++abi-dev
+          sudo apt-get install -y clang libc++-dev libc++abi-dev libomp-18-dev
 
       - name: Install OpenMP (MacOS)
         if: runner.os == 'macOS'

CMakeLists.txt

@@ -9,6 +9,7 @@ project(
 option(TGN_BUILD_TESTS "Build test binaries" OFF)
 option(TGN_BUILD_EXAMPLES "Build example binaries" OFF)
 option(TGN_BUILD_TOOLS "Build tool binaries" OFF)
+option(TGN_BUILD_PYTHON "Build Python Bindings" OFF)
 
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 set(CMAKE_CXX_STANDARD 20)
@@ -90,6 +91,10 @@ if (TGN_BUILD_TOOLS)
     add_subdirectory(tools/tguf_cli)
 endif()
 
+if (TGN_BUILD_PYTHON)
+    add_subdirectory(python)
+endif()
+
 
 if (TGN_BUILD_TESTS)
     enable_testing()

Makefile

@@ -42,6 +42,11 @@ help:
 	@echo ""
 	@echo "Data Targets:"
 	@echo "  make download-<ds>       - Download TGB dataset (e.g., make download-tgbl-wiki)"
+	@echo ""
+	@echo "Python Targets:"
+	@echo "  make python                - Build and install Python bindings"
+	@echo "  make test-python           - Run Python-specific tests"
+	@echo "  make clean-python          - Run Python-specific build artifacts"
 	@echo "========================================================================"
 
 $(BUILD_DIR)/CMakeCache.txt:
@@ -124,6 +129,21 @@ perf-node-%: profile-build data/%.tguf
 download-%: data/%.tguf
 	@echo "Dataset $* is up to date."
 
+.PHONY: python
+python:
+	@(cd python && \
+		uv sync --group dev --no-install-project && \
+		SKBUILD_CMAKE_ARGS="-DTGN_BUILD_PYTHON=ON" \
+		uv pip install -e . --no-build-isolation)
+
+.PHONY: test-python
+test-python: python
+	@(cd python && uv run pytest test/)
+
+.PHONY: clean-python
+clean-python:
+	rm -rf python/build
+
 .PHONY: clean
 clean:
 	rm -rf $(BUILD_DIR)

README.md

@@ -7,7 +7,7 @@ C++ Port of "Temporal Graph Networks for Deep Learning on Dynamic Graphs".
 
 ### Prerequisites
 
-You can just use the provided [Dockerfile](./Dockerfile). If you prefer to install dependencies manually, see below.
+You should just use the [Dockerfile](./Dockerfile), but if you prefer to install dependencies manually, see below.
 
 ##### Linux
 
@@ -43,3 +43,7 @@ make run-link-tgbl-wiki
 # Download `tgbn-trade` data, convert to `.tguf` and run examples/node_pred.cpp
 make run-node-tgbn-trade
 ```
+
+### Python Bindings
+
+We expose [Python bindings](./python) for the `TGUFBuilder` via [nanobind](https://nanobind.readthedocs.io/en/latest/).

python/.python-version

@@ -0,0 +1 @@
+3.10

python/CMakeLists.txt

@@ -0,0 +1,35 @@
+cmake_minimum_required(VERSION 3.16)
+project(
+    tgn_python
+    VERSION 0.1.0
+    DESCRIPTION "Python bindings for tgn.cpp"
+    LANGUAGES CXX
+)
+
+if (CMAKE_VERSION VERSION_LESS 3.18)
+  set(DEV_MODULE Development)
+else()
+  set(DEV_MODULE Development.Module)
+endif()
+
+find_package(Python 3.10 COMPONENTS Interpreter ${DEV_MODULE} REQUIRED)
+
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
+  OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE nanobind_ROOT)
+find_package(nanobind CONFIG REQUIRED)
+
+set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
+set(CMAKE_INSTALL_RPATH "$ORIGIN;${CMAKE_BINARY_DIR}")
+
+nanobind_add_module(
+    _core
+    NB_STATIC
+    nanobind/main.cpp
+)
+
+target_link_libraries(_core PRIVATE tgn)
+
+install(TARGETS _core
+        LIBRARY DESTINATION tguf
+        COMPONENT python_bindings)

python/README.md

@@ -0,0 +1,50 @@
+### tguf (Temporal Graph Universal Format)
+
+Python bindings for the `tgn.cpp` TGUF builder.
+
+### Prerequisites
+
+Requires [uv](https://docs.astral.sh/uv/) and C++20 toolchain.
+
+```sh
+# Build an editable install with bindings from the project root
+make python
+```
+
+### Quickstart
+
+Our bindings use zero-copy semantics to ingest any Python object supporting the DLPack or Buffer protocol (e.g., numpy.ndarray, torch.Tensor).
+
+```python
+import tguf
+import numpy as np
+import torch
+
+# Define your tguf schema
+schema = tguf.TGUFSchema(
+    path="my_data.tguf",
+    edge_capacity=10,
+    node_feat_capacity=5,
+    node_feat_dim=4,
+    msg_dim=8,
+)
+builder = tguf.TGUFBuilder(schema)
+
+# Stream some edge events
+edge_batch = tguf.Batch(
+    src=np.array([0, 1], dtype=np.int64),
+    dst=np.array([2, 3], dtype=np.int64),
+    time=np.array([0, 1], dtype=np.int64),
+    msg=np.random.randn(2, 8), dtype=np.float32),
+)
+builder.append_edges(batch)
+
+# Stream through some node features
+builder.append_node_feats(
+    n_id=torch.arange(5, dtype=torch.int64) # tensor's work too
+    node_feat = torch.randn(5, 4, dtype=torch.float32)
+)
+
+# Finalize the .tguf file and flush to disk
+builder.finalize()
+```

python/nanobind/main.cpp

@@ -0,0 +1,130 @@
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+#include <nanobind/stl/optional.h>
+#include <nanobind/stl/string.h>
+#include <torch/torch.h>
+
+#include <cstdint>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tgn.h"
+
+namespace nb = nanobind;
+
+// This takes any Python object supporting DLPack/Buffer Protocol
+torch::Tensor tensor_view(nb::ndarray<> array, torch::ScalarType type) {
+  std::vector<std::int64_t> shape;
+  for (auto i = 0; i < array.ndim(); ++i) {
+    shape.push_back(array.shape(i));
+  }
+  return torch::from_blob(array.data(), shape,
+                          torch::TensorOptions().dtype(type));
+}
+
+NB_MODULE(_core, m) {
+  nb::class_<tgn::TGUFSchema>(m, "TGUFSchema")
+      .def(
+          "__init__",
+          [](tgn::TGUFSchema *self, std::string path,
+             std::optional<std::size_t> edge_capacity,
+             std::optional<std::size_t> msg_dim,
+             std::optional<std::size_t> label_dim,
+             std::optional<std::size_t> node_feat_capacity,
+             std::optional<std::size_t> node_feat_dim,
+             std::optional<std::size_t> label_capacity,
+             std::optional<std::size_t> negatives_start_e_id,
+             std::optional<std::size_t> negatives_per_edge,
+             std::optional<std::size_t> val_start,
+             std::optional<std::size_t> test_start) {
+            new (self) tgn::TGUFSchema();
+            self->path = std::move(path);
+            self->edge_capacity = edge_capacity.value_or(0);
+            self->msg_dim = msg_dim.value_or(0);
+            self->node_feat_capacity = node_feat_capacity.value_or(0);
+            self->node_feat_dim = node_feat_dim.value_or(0);
+            self->label_capacity = label_capacity.value_or(0);
+            self->label_dim = label_dim.value_or(0);
+            self->negatives_start_e_id = negatives_start_e_id.value_or(0);
+            self->negatives_per_edge = negatives_per_edge.value_or(0);
+            self->val_start = val_start;
+            self->test_start = test_start;
+          },
+          nb::arg("path"), nb::arg("edge_capacity") = nb::none(),
+          nb::arg("msg_dim") = nb::none(), nb::arg("label_dim") = nb::none(),
+          nb::arg("node_feat_capacity") = nb::none(),
+          nb::arg("node_feat_dim") = nb::none(),
+          nb::arg("label_capacity") = nb::none(),
+          nb::arg("negatives_start_e_id") = nb::none(),
+          nb::arg("negatives_per_edge") = nb::none(),
+          nb::arg("val_start") = nb::none(), nb::arg("test_start") = nb::none())
+
+      .def_rw("path", &tgn::TGUFSchema::path)
+      .def_rw("edge_capacity", &tgn::TGUFSchema::edge_capacity)
+      .def_rw("msg_dim", &tgn::TGUFSchema::msg_dim)
+      .def_rw("label_dim", &tgn::TGUFSchema::label_dim)
+      .def_rw("node_feat_capacity", &tgn::TGUFSchema::node_feat_capacity)
+      .def_rw("node_feat_dim", &tgn::TGUFSchema::node_feat_dim)
+      .def_rw("label_capacity", &tgn::TGUFSchema::label_capacity)
+      .def_rw("negatives_start_e_id", &tgn::TGUFSchema::negatives_start_e_id)
+      .def_rw("negatives_per_edge", &tgn::TGUFSchema::negatives_per_edge)
+      .def_rw("val_start", &tgn::TGUFSchema::val_start)
+      .def_rw("test_start", &tgn::TGUFSchema::test_start);
+
+  nb::class_<tgn::Batch>(m, "Batch")
+      .def(
+          "__init__",
+          [](tgn::Batch *self, nb::ndarray<> src, nb::ndarray<> dst,
+             nb::ndarray<> time, nb::ndarray<> msg,
+             std::optional<nb::ndarray<>> neg_dst) {
+            new (self) tgn::Batch{
+                tensor_view(src, torch::kLong), tensor_view(dst, torch::kLong),
+                tensor_view(time, torch::kLong),
+                tensor_view(msg, torch::kFloat),
+                neg_dst
+                    ? std::make_optional(tensor_view(*neg_dst, torch::kLong))
+                    : std::nullopt};
+          },
+          nb::arg("src"), nb::arg("dst"), nb::arg("time"), nb::arg("msg"),
+          nb::arg("neg_dst") = nb::none());
+
+  nb::class_<tgn::TGUFBuilder>(m, "TGUFBuilder")
+      .def(nb::init<const tgn::TGUFSchema &>(), nb::arg("schema"))
+
+      .def(
+          "append_edges",
+          [](const tgn::TGUFBuilder &self, const tgn::Batch &batch) {
+            nb::gil_scoped_release release;
+            self.append_edges(batch);
+          },
+          nb::arg("batch"))
+
+      .def(
+          "append_labels",
+          [](const tgn::TGUFBuilder &self, nb::ndarray<> n_id,
+             nb::ndarray<> time, nb::ndarray<> target) {
+            nb::gil_scoped_release release;
+
+            self.append_labels(tensor_view(n_id, torch::kLong),
+                               tensor_view(time, torch::kLong),
+                               tensor_view(target, torch::kFloat));
+          },
+          nb::arg("n_id"), nb::arg("time"), nb::arg("target"))
+
+      .def(
+          "append_node_feats",
+          [](const tgn::TGUFBuilder &self, nb::ndarray<> n_id,
+             nb::ndarray<> node_feat) {
+            nb::gil_scoped_release release;
+
+            self.append_node_feats(tensor_view(n_id, torch::kLong),
+                                   tensor_view(node_feat, torch::kFloat));
+          },
+          nb::arg("n_id"), nb::arg("node_feat"))
+
+      .def("finalize", [](tgn::TGUFBuilder &self) {
+        nb::gil_scoped_release release;
+        self.finalize();
+      });
+}