apache / apache/airflow

ONNX Model Inference Operator

Open
#41,702 8 comments 0 reactions 0 assignees View on GitHub
kind:feature kind:new provider request
Dominant language
Python
Stars
46.9k
Forks
17.8k
Avg merge
2d 10h
Merged PRs (30d)
483

Description

### Description

ONNX (Open Neural Network Exchange) provides cross-platform compatibility

An operator that can run inference using ONNX models, ideal for deploying machine learning models in a standardized format can provide us with direct model invocation.

this can be solved using a pythonOperator ofc as onnxruntime can be executed with pythonruntime, but this can also be built into airflow to minimize work, a simple onnx operator structure would be something like:

```

import onnxruntime as ort
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

def run_onnx_inference():
# Load the ONNX model
model_path = '/path/to/your/model.onnx'
session = ort.InferenceSession(model_path)

# Prepare input data
input_name = session.get_inputs()[0].name
input_data = {"your_input_key": your_input_data}

# Run inference
result = session.run(None, {input_name: input_data})
print(result)

# Define the DAG
with DAG(
dag_id='onnx_inference_dag',
start_date=datetime(2023, 1, 1),
schedule_interval='@once'
) as dag:

# Define the task
inference_task = PythonOperator(
task_id='onnx_inference_task',
python_callable=run_onnx_inference
)

```

Looking frwd to any suggestions.

### Use case/motivation

A direct support of onnx with Airflow's DAG-based orchestration can manage the entire lifecycle of data processing and model inference in one place, providing a more cohesive and manageable workflow.

### Related issues

_No response_

### Are you willing to submit a PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)

Contributor guide

Open the contributing guide

Research direction

Review the proposed PythonOperator approach and the onnxruntime example first. Define the operator's API, dependency and integration scope before implementation; the issue does not name repository files, tests, or a concrete completion criterion.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data-engineering, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.