pg_onnx: ONNX Runtime integrated with PostgreSQL. Perform ML inference with data in your database.
graph LR
subgraph P[PostgreSQL]
direction LR
PM((postmaster))
PS1[postgres]
PS2[postgres]
PS3[postgres]
TABLE[(onnx model data)]
subgraph ONNX[pg_onnx Background Worker]
direction TB
OS[onnxruntime-server]
ONNX1([onnxruntime\nsession:\nmodel1])
ONNX2([onnxruntime\nsession:\nmodel2])
OS -. " create/execute session " .-o ONNX1 & ONNX2
end
PM -. " fork " .-o PS1 & PS2 & PS3
PM -. " dynamic background worker " .-o ONNX
PS3 <-- " import model " --> TABLE
PS1 & PS2 & PS3 <-- " ONNX operation " --> ONNX
end
C[Client trying to use pg_onnx] ==> PS3
download-onnxruntime-linux.sh
script
/usr/local/onnxruntime
./usr/local/onnxruntime/lib
to /etc/ld.so.conf.d/onnxruntime.conf
and run ldconfig
.brew install onnxruntime
sudo apt install cmake libboost-all-dev libpq-dev postgresql-server-dev-all
# optional, for Nvidia GPU support
sudo apt install nvidia-cuda-toolkit nvidia-cudnn
brew install cmake boost postgresql
git submodule update --init --recursive
to update submodules.git clone --recursive https://github.com/kibae/pg_onnx.git
cmake -B build -S . -DCMAKE_BUILD_TYPE=Release
cmake --build build --target pg_onnx --parallel
sudo cmake --install build/pg_onnx
CREATE EXTENSION IF NOT EXISTS pg_onnx;
SELECT pg_onnx_import_model(
'sample_model', --------------- model name
'v20230101', ------------------ model version
PG_READ_BINARY_FILE('/your_model_path/model.onnx')::bytea, -- model binary data
'{"cuda": true}'::jsonb, ------ options
'sample model' ---------------- description
);
SELECT pg_onnx_execute_session(
'sample_model', -- model name
'v20230101', ----- model version
'{
"x": [[1], [2], [3]],
"y": [[3], [4], [5]],
"z": [[5], [6], [7]]
}' --------------- inputs
);
pg_onnx_execute
--------------------------------------------------------------------------------
{"output": [[0.7488641738891602], [0.8607008457183838], [0.9725375175476074]]}
-- Create a test table
CREATE TABLE trigger_test
(
id SERIAL PRIMARY KEY,
value1 INT,
value2 INT,
value3 INT,
prediction FLOAT
);
-- Create a trigger function
CREATE OR REPLACE FUNCTION trigger_test_insert()
RETURNS TRIGGER AS
$$
DECLARE
result jsonb;
BEGIN
result := pg_onnx_execute_session(
'sample_model', 'v20230101',
JSONB_BUILD_OBJECT(
'x', ARRAY [[NEW.value1]],
'y', ARRAY [[NEW.value2]],
'z', ARRAY [[NEW.value3]]));
-- output shape: float[-1,1]
-- eg: {"output": [[0.6492120623588562]]}
NEW.prediction := result -> 'output' -> 0 -> 0;
RETURN NEW;
END;
$$
LANGUAGE plpgsql;
-- Create a trigger
CREATE TRIGGER trigger_test_insert
BEFORE INSERT
ON trigger_test
FOR EACH ROW
EXECUTE PROCEDURE trigger_test_insert();