GitXplorerGitXplorer
d

rust-tf-pluggabledevice

public
0 stars
0 forks
0 issues

Commits

List of commits on branch main.
Unverified
23a0dbdc895114d4ec3507316e0610aa96ae7e52

Make crate available for external usage

ddsseng committed 2 years ago
Unverified
dd448e633cf8e0ac59494c274b11e98f19b4877a

Clarify requirements

ddsseng committed 2 years ago
Unverified
1965ad90271cf1766170c6465c8c9afe1904a146

Remove extra includes

ddsseng committed 2 years ago
Unverified
10bccb95f00005037cb8cca1089a58fe5223ada7

Fix clippy

ddsseng committed 2 years ago
Unverified
a571464bb3ad9ddacbe0063a428b5da9a8e8049f

Remove the need to have a build script in plugin

ddsseng committed 2 years ago
Unverified
22ffbe8a2ac44352ffc85abd96698ac0cb932048

Strip release library

ddsseng committed 2 years ago

README

The README file for this repository.

Rust TensorFlow PluggableDevice

Ported over from https://github.com/tensorflow/community/pull/352

The tfp-bindings crate mostly provides unsafe bindings to all TensorFlow C APIs required by plug-ins via libtensorflow_framework.so.2 library (yes, and currently in only supports Linux and Python 3.10). It also has some should-be-safe bindings to main functions required by compute kernels and their registration. Its safety mostly relies on validating data, checking result codes and asserts in debug mode to prevent common plug-in issues like unterminated strings.

Plugin itself (tfp-plugin) should link to that library via its build.rs script and implement SE_InitPlugin and TF_InitKernel functions with proper types (these are excluded from bindgen). It may also include optimizer, but it is not yet implemented in this repository.

Requirements

  • Linux, tested on Fedora 36, any should work fine.
  • MSRV not yet specified, I use nightly. No strict limits are being set.
  • Python 3 for test script. pip3 and venv to install TensorFlow
  • clang for binding generation via bindgen

Usage in other packages

You can import the bindings crate by using such a dependency string. There are environment variables you can use to point the paths to TensorFlow library:

  • TF_INCLUDE_PATH
  • TF_LIBRARY_PATH

They could be set, for example, by Cargo project config

Cargo.toml example. Replace x by needed commit hash.

tfp-bindings = { git = "https://github.com/sh7dm/rust-tf-pluggabledevice", rev = "x" }

Try it out

python3.10 -m venv venv
source venv/bin/activate
pip3 install --no-cache-dir tf-nightly-cpu # or tensorflow-cpu
# https://github.com/tensorflow/tensorflow/issues/55497
mv venv/lib64 venv/_lib64
cargo build
mkdir venv/lib/python3.10/site-packages/tensorflow-plugins/
cp target/debug/libtfp.so venv/lib/python3.10/site-packages/tensorflow-plugins/
python3 test.py
import tensorflow as tf
tf.config.list_physical_devices()

Running tests

Specifying LD_LIBRARY_PATH manually is necessary as of https://github.com/rust-lang/cargo/issues/4044

LD_LIBRARY_PATH=$(pwd)/venv/lib/python3.10/site-packages/tensorflow cargo test