├── .github └── workflows │ └── deploy.yaml ├── .gitignore ├── Dockerfile ├── LICENSE.md ├── Makefile ├── README.md ├── charts └── inferencedb │ ├── .helmignore │ ├── Chart.yaml │ ├── crds │ └── inferencelogger.yaml │ ├── templates │ ├── _helpers.tpl │ ├── deployment.yaml │ ├── role.yaml │ ├── rolebinding.yaml │ ├── service.yaml │ ├── serviceaccount.yaml │ └── servicemonitor.yaml │ └── values.yaml ├── examples └── kserve │ ├── kafka-broker │ ├── broker.yaml │ ├── inferencelogger.yaml │ └── inferenceservice.yaml │ └── multimodelserving │ ├── broker.yaml │ ├── inferenceloggers.yaml │ ├── inferenceservice.yaml │ └── trainedmodels.yaml ├── logo.svg ├── poetry.lock ├── pyproject.toml ├── skaffold.yaml ├── src └── inferencedb │ ├── __init__.py │ ├── app.py │ ├── config │ ├── __init__.py │ ├── component.py │ ├── config.py │ ├── factory.py │ └── providers │ │ ├── __init__.py │ │ ├── config_provider.py │ │ ├── file_config_provider.py │ │ └── kubernetes_config_provider.py │ ├── core │ ├── base_model.py │ ├── inference.py │ ├── inference_logger.py │ └── logging_utils.py │ ├── destinations │ ├── __init__.py │ ├── confluent_s3_destination.py │ ├── destination.py │ └── factory.py │ ├── event_processors │ ├── __init__.py │ ├── event_processor.py │ ├── factory.py │ ├── json_event_processor.py │ └── kserve_event_processor.py │ ├── main.py │ ├── py.typed │ ├── registry │ ├── __init__.py │ ├── decorators.py │ ├── factory.py │ ├── registered_object.py │ └── registry.py │ ├── schema_providers │ ├── __init__.py │ ├── avro_schema_provider.py │ ├── factory.py │ └── schema_provider.py │ ├── settings.py │ └── utils │ ├── __init__.py │ ├── asyncio_utils.py │ ├── kserve_utils.py │ └── pandas_utils.py └── tests ├── __init__.py ├── conftest.py └── test_kserve_utils.py /.github/workflows/deploy.yaml: -------------------------------------------------------------------------------- 1 | name: Deploy 2 | on: 3 | push: 4 | branches: 5 | - main 6 | env: 7 | REGISTRY: ghcr.io 8 | IMAGE_NAME: ${{ github.repository }} 9 | jobs: 10 | deploy: 11 | runs-on: ubuntu-20.04 12 | steps: 13 | - name: Checkout 14 | uses: actions/checkout@main 15 | with: 16 | ref: ${{ github.ref }} 17 | fetch-depth: 0 18 | token: ${{ secrets.GITHUB_TOKEN }} 19 | 20 | - name: Configure git 21 | run: | 22 | git config --global user.email "camparibot@aporia.com" 23 | git config --global user.name "camparibot" 24 | git config --global push.followTags true 25 | 26 | - name: Log in to the Container registry 27 | uses: docker/login-action@v1 28 | with: 29 | registry: ${{ env.REGISTRY }} 30 | username: ${{ github.actor }} 31 | password: ${{ secrets.GITHUB_TOKEN }} 32 | 33 | - name: Install dependencies 34 | run: make install-deps 35 | 36 | - name: Bump Version 37 | id: bump-version 38 | run: make bump-version 39 | env: 40 | IMAGE_NAME: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} 41 | CAMPARIBOT_TOKEN: ${{ secrets.CAMPARIBOT_TOKEN }} 42 | 43 | - name: Create check run 44 | id: create-check-run 45 | run: | 46 | CHECK_RUN_ID=`curl -X POST https://api.github.com/repos/${{ github.repository }}/check-runs \ 47 | -H "Authorization: token ${{ secrets.GITHUB_TOKEN }}" \ 48 | -H "Accept:application/vnd.github.antiope-preview+json" \ 49 | -d "{\"name\": \"Aporia / deploy (push)\", \"head_sha\": \"${{ steps.bump-version.outputs.bumped_version_commit_hash }}\", \"status\": \"in_progress\", \"details_url\": \"https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}\", \"output\": {\"title\": \"Versioned Commit\", \"summary\": \"This is a versioned commit. To see the full GitHub Action, [click here](https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}).\"}}" \ 50 | | jq .id`; 51 | echo "::set-output name=check_run_id::$CHECK_RUN_ID"; 52 | 53 | - name: Cache skaffold image builds & config 54 | id: cache-skaffold 55 | uses: actions/cache@v2 56 | with: 57 | path: ~/.skaffold/ 58 | key: fixed-${{ github.sha }} 59 | restore-keys: | 60 | fixed-${{ github.sha }} 61 | fixed- 62 | 63 | - name: Build 64 | run: make build 65 | env: 66 | NEW_VERSION: ${{ steps.bump-version.outputs.new-version }} 67 | CAMPARIBOT_TOKEN: ${{ secrets.CAMPARIBOT_TOKEN }} 68 | 69 | - name: Update check run to success 70 | run: | 71 | curl -X PATCH https://api.github.com/repos/${{ github.repository }}/check-runs/${{ steps.create-check-run.outputs.check_run_id }} \ 72 | -H "Authorization: token ${{ secrets.GITHUB_TOKEN }}" \ 73 | -H "Accept:application/vnd.github.antiope-preview+json" \ 74 | -d "{\"status\": \"completed\", \"conclusion\": \"success\"}"; 75 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | .DS_Store -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM python:3.9-slim 2 | WORKDIR /inferencedb 3 | STOPSIGNAL SIGINT 4 | 5 | # Install wget for readiness probe 6 | RUN apt update && apt install -y curl wget build-essential librocksdb-dev libsnappy-dev zlib1g-dev libbz2-dev liblz4-dev && rm -rf /var/lib/apt/lists/* 7 | 8 | # Install python dependencies with poetry 9 | COPY poetry.lock pyproject.toml ./ 10 | RUN pip3 install poetry 11 | RUN poetry config virtualenvs.create false 12 | RUN poetry install --no-interaction --no-ansi --no-dev 13 | 14 | COPY . . 15 | WORKDIR /inferencedb/src 16 | ENTRYPOINT [ "python", "-m", "inferencedb.main" ] 17 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | Elastic License 2.0 2 | 3 | URL: https://www.elastic.co/licensing/elastic-license 4 | 5 | ## Acceptance 6 | 7 | By using the software, you agree to all of the terms and conditions below. 8 | 9 | ## Copyright License 10 | 11 | The licensor grants you a non-exclusive, royalty-free, worldwide, 12 | non-sublicensable, non-transferable license to use, copy, distribute, make 13 | available, and prepare derivative works of the software, in each case subject to 14 | the limitations and conditions below. 15 | 16 | ## Limitations 17 | 18 | You may not provide the software to third parties as a hosted or managed 19 | service, where the service provides users with access to any substantial set of 20 | the features or functionality of the software. 21 | 22 | You may not move, change, disable, or circumvent the license key functionality 23 | in the software, and you may not remove or obscure any functionality in the 24 | software that is protected by the license key. 25 | 26 | You may not alter, remove, or obscure any licensing, copyright, or other notices 27 | of the licensor in the software. Any use of the licensor’s trademarks is subject 28 | to applicable law. 29 | 30 | ## Patents 31 | 32 | The licensor grants you a license, under any patent claims the licensor can 33 | license, or becomes able to license, to make, have made, use, sell, offer for 34 | sale, import and have imported the software, in each case subject to the 35 | limitations and conditions in this license. This license does not cover any 36 | patent claims that you cause to be infringed by modifications or additions to 37 | the software. If you or your company make any written claim that the software 38 | infringes or contributes to infringement of any patent, your patent license for 39 | the software granted under these terms ends immediately. If your company makes 40 | such a claim, your patent license ends immediately for work on behalf of your 41 | company. 42 | 43 | ## Notices 44 | 45 | You must ensure that anyone who gets a copy of any part of the software from you 46 | also gets a copy of these terms. 47 | 48 | If you modify the software, you must include in any modified copies of the 49 | software prominent notices stating that you have modified the software. 50 | 51 | ## No Other Rights 52 | 53 | These terms do not imply any licenses other than those expressly granted in 54 | these terms. 55 | 56 | ## Termination 57 | 58 | If you use the software in violation of these terms, such use is not licensed, 59 | and your licenses will automatically terminate. If the licensor provides you 60 | with a notice of your violation, and you cease all violation of this license no 61 | later than 30 days after you receive that notice, your licenses will be 62 | reinstated retroactively. However, if you violate these terms after such 63 | reinstatement, any additional violation of these terms will cause your licenses 64 | to terminate automatically and permanently. 65 | 66 | ## No Liability 67 | 68 | *As far as the law allows, the software comes as is, without any warranty or 69 | condition, and the licensor will not be liable to you for any damages arising 70 | out of these terms or the use or nature of the software, under any kind of 71 | legal claim.* 72 | 73 | ## Definitions 74 | 75 | The **licensor** is the entity offering these terms, and the **software** is the 76 | software the licensor makes available under these terms, including any portion 77 | of it. 78 | 79 | **you** refers to the individual or entity agreeing to these terms. 80 | 81 | **your company** is any legal entity, sole proprietorship, or other kind of 82 | organization that you work for, plus all organizations that have control over, 83 | are under the control of, or are under common control with that 84 | organization. **control** means ownership of substantially all the assets of an 85 | entity, or the power to direct its management and policies by vote, contract, or 86 | otherwise. Control can be direct or indirect. 87 | 88 | **your licenses** are all the licenses granted to you for the software under 89 | these terms. 90 | 91 | **use** means anything you do with the software requiring one of your licenses. 92 | 93 | **trademark** means trademarks, service marks, and similar rights. 94 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | SHELL := /bin/bash 2 | 3 | HELM_CHART=./charts/inferencedb 4 | DEFAULT_VERSION=1.0.0 5 | 6 | install-deps: 7 | @echo [!] Installing Semver 8 | @sudo wget https://raw.githubusercontent.com/fsaintjacques/semver-tool/master/src/semver -O /usr/bin/semver 9 | @sudo chmod +x /usr/bin/semver 10 | 11 | @echo [!] Installing yq 12 | @sudo wget https://github.com/mikefarah/yq/releases/download/v4.6.1/yq_linux_amd64 -O /usr/bin/yq && sudo chmod +x /usr/bin/yq 13 | 14 | @echo [!] Installing skaffold 15 | @curl -Lo skaffold https://storage.googleapis.com/skaffold/releases/latest/skaffold-linux-amd64 && sudo install skaffold /usr/local/bin/ 16 | 17 | build: 18 | skaffold build --tag=$(NEW_VERSION) 19 | 20 | bump-version: 21 | $(eval CURRENT_VERSION=$(shell git for-each-ref --sort=-v:refname --count=1 refs/tags/[0-9]*.[0-9]*.[0-9]* refs/tags/v[0-9]*.[0-9]*.[0-9]* | cut -d / -f 3-)) 22 | $(eval export NEW_VERSION=$(shell \ 23 | if [ -z $(CURRENT_VERSION) ]; then \ 24 | echo $(DEFAULT_VERSION); \ 25 | else \ 26 | semver bump patch $(CURRENT_VERSION); \ 27 | fi; \ 28 | )) 29 | 30 | @git log -1 --pretty="%B" > /tmp/commit-message 31 | 32 | @echo [!] Bumping version from $(CURRENT_VERSION) to $(NEW_VERSION) 33 | 34 | yq e '.version = "$(NEW_VERSION)"' -i $(HELM_CHART)/Chart.yaml 35 | yq e '.appVersion = "$(NEW_VERSION)"' -i $(HELM_CHART)/Chart.yaml 36 | git add $(HELM_CHART)/Chart.yaml 37 | 38 | git commit -F /tmp/commit-message --amend --no-edit 39 | 40 | git tag -a -m "Version v$(NEW_VERSION)" v$(NEW_VERSION) 41 | 42 | @BRANCH_PROTECTION=`curl https://api.github.com/repos/$(GITHUB_REPOSITORY)/branches/main/protection \ 43 | -H "Authorization: token $(CAMPARIBOT_TOKEN)" -H "Accept:application/vnd.github.luke-cage-preview+json" -X GET -s`; \ 44 | if [ "`echo $$BRANCH_PROTECTION | jq -r '.message'`" != "Branch not protected" ]; \ 45 | then \ 46 | echo [!] Disabling GitHub main branch protection; \ 47 | curl https://api.github.com/repos/$(GITHUB_REPOSITORY)/branches/main/protection \ 48 | -H "Authorization: token $(CAMPARIBOT_TOKEN)" -H "Accept:application/vnd.github.luke-cage-preview+json" -X DELETE; \ 49 | trap '\ 50 | echo [!] Re-enabling GitHub main branch protection; \ 51 | curl https://api.github.com/repos/$(GITHUB_REPOSITORY)/branches/main/protection -H "Authorization: token $(CAMPARIBOT_TOKEN)" \ 52 | -H "Accept:application/vnd.github.luke-cage-preview+json" -X PUT -d "{\"required_status_checks\":{\"strict\":false,\"contexts\":`echo $$BRANCH_PROTECTION | jq '.required_status_checks.contexts'`},\"restrictions\":{\"users\":[],\"teams\":[],\"apps\":[]},\"required_pull_request_reviews\":{\"dismiss_stale_reviews\":false,\"require_code_owner_reviews\":false},\"enforce_admins\":true,\"required_linear_history\":false,\"allow_force_pushes\":true,\"allow_deletions\":false}"; \ 53 | ' EXIT; \ 54 | fi; \ 55 | echo [!] Git Push; \ 56 | git push --force; 57 | 58 | echo "::set-output name=new-version::$(NEW_VERSION)"; 59 | echo "::set-output name=bumped_version_commit_hash::`git log --pretty=format:'%H' -n 1`"; 60 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |

2 | 3 |

4 | 5 | --- 6 | 7 | **InferenceDB** makes it easy to stream inferences of real-time ML models in production to a data lake, based on Kafka. This data can later be used for model retraining, data drift monitoring, performance degradation detection, AI incident investigation and more. 8 | 9 | ### Quickstart 10 | 11 | * [Flask](https://github.com/aporia-ai/inferencedb/wiki/Flask-Quickstart) 12 | * [FastAPI](https://github.com/aporia-ai/inferencedb/wiki/FastAPI-Quickstart) 13 | * [KServe](https://github.com/aporia-ai/inferencedb/wiki/KServe-Quickstart) 14 | 15 | 16 | ### Features 17 | 18 | * **Cloud Native** - Runs on top of Kubernetes and supports any cloud infrastructure 19 | * **Model Serving Integrations** - Connects to ML model serving tools like [KServe](https://kserve.github.io/website/) 20 | * **Extensible** - Add your own model serving frameworks and database destinations 21 | * **Horizontally Scalable** - Add more workers to support more models and more traffic 22 | * **Python Ecosystem** - Written in Python using [Faust](https://faust.readthedocs.io/en/latest/), so you can add your own data transformations using Numpy, Pandas, etc. 23 | 24 |

Made with :heart: by Aporia

25 | 26 | **WARNING:** InferenceDB is still experimental, use at your own risk! 💀 27 | 28 | ## Installation 29 | 30 | The only requirement to InferenceDB is a Kafka cluster, with [Schema Registry](https://docs.confluent.io/platform/current/schema-registry/index.html) and [Kafka Connect](https://docs.confluent.io/platform/current/connect/index.html). 31 | 32 | To install InferenceDB using Helm, run: 33 | 34 | ```sh 35 | helm install inferencedb inferencedb/inferencedb -n inferencedb --create-namespace \ 36 | --set kafka.broker=kafka:9092 \ 37 | --set kafka.schemaRegistryUrl=http://schema-registry:8081 \ 38 | --set kafka.connectUrl=http://kafka-connect:8083 39 | ``` 40 | 41 | ## Usage 42 | 43 | To start logging your model inferences, create an **InferenceLogger** Kubernetes resource. This is a [Kubernetes Custom Resource](https://kubernetes.io/docs/concepts/extend-kubernetes/api-extension/custom-resources/) that is defined and controlled by InferenceDB. 44 | 45 | **Example:** 46 | 47 | ```yaml 48 | apiVersion: inferencedb.aporia.com/v1alpha1 49 | kind: InferenceLogger 50 | metadata: 51 | name: my-model-inference-logger 52 | namespace: default 53 | spec: 54 | topic: my-model 55 | events: 56 | type: kserve 57 | config: {} 58 | destination: 59 | type: confluent-s3 60 | config: 61 | url: s3://my-bucket/inferencedb 62 | format: parquet 63 | awsRegion: us-east-2 64 | ``` 65 | 66 | This InferenceLogger will watch the `my-model` Kafka topic for events in KServe format, and log them to a Parquet file on S3. See the [KServe quickstart guide](https://github.com/aporia-ai/inferencedb/wiki/KServe-Quickstart) for more details. 67 | 68 | ## Development 69 | 70 | InferenceDB dev is done using [Skaffold](https://skaffold.dev/). 71 | 72 | Make sure you have a Kubernetes cluster with Kafka installed (can be local or remote), and edit [skaffold.yaml](skaffold.yaml) with the correct Kafka URLs and Docker image registry (for local, just use `local/inferencedb`). 73 | 74 | To start development, run: 75 | 76 | skaffold dev --trigger=manual 77 | 78 | This will build the Docker image, push it to the Docker registry you provided, and install the Helm chart on the cluster. Now, you can make changes to the code, click "Enter" on the Skaffold CLI and that would update the cluster. 79 | 80 | ## Roadmap 81 | 82 | ### Core 83 | 84 | * [ ] Add support for Spark Streaming in addition to Faust 85 | * [ ] Add more input validations on the Kafka URLs 86 | 87 | ### Event Processors 88 | 89 | * [x] JSON 90 | * [x] KServe 91 | * [ ] Seldon Core 92 | * [ ] BentoML 93 | * [ ] MLFlow Deployments 94 | 95 | ### Destinations 96 | 97 | * [x] Parquet on S3 98 | * [ ] HDF5 on S3 99 | * [ ] Azure Blob Storage 100 | * [ ] Google Cloud Storage 101 | * [ ] ADLS Gen2 102 | * [ ] AWS Glue 103 | * [ ] Delta Lake 104 | * [ ] PostgreSQL 105 | * [ ] Snowflake 106 | * [ ] Iceberg 107 | 108 | ### Documentation 109 | 110 | * [ ] How to set up Kafka using AWS / Azure / GCP managed services 111 | * [ ] API Reference for the CRDs 112 | -------------------------------------------------------------------------------- /charts/inferencedb/.helmignore: -------------------------------------------------------------------------------- 1 | # Patterns to ignore when building packages. 2 | # This supports shell glob matching, relative path matching, and 3 | # negation (prefixed with !). Only one pattern per line. 4 | .DS_Store 5 | # Common VCS dirs 6 | .git/ 7 | .gitignore 8 | .bzr/ 9 | .bzrignore 10 | .hg/ 11 | .hgignore 12 | .svn/ 13 | # Common backup files 14 | *.swp 15 | *.bak 16 | *.tmp 17 | *.orig 18 | *~ 19 | # Various IDEs 20 | .project 21 | .idea/ 22 | *.tmproj 23 | .vscode/ 24 | -------------------------------------------------------------------------------- /charts/inferencedb/Chart.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: v2 2 | name: inferencedb 3 | description: InferenceDB helm chart 4 | # A chart can be either an 'application' or a 'library' chart. 5 | # 6 | # Application charts are a collection of templates that can be packaged into versioned archives 7 | # to be deployed. 8 | # 9 | # Library charts provide useful utilities or functions for the chart developer. They're included as 10 | # a dependency of application charts to inject those utilities and functions into the rendering 11 | # pipeline. Library charts do not define any templates and therefore cannot be deployed. 12 | type: application 13 | # This is the chart version. This version number should be incremented each time you make changes 14 | # to the chart and its templates, including the app version. 15 | # Versions are expected to follow Semantic Versioning (https://semver.org/) 16 | version: 1.0.61 17 | # This is the version number of the application being deployed. This version number should be 18 | # incremented each time you make changes to the application. Versions are not expected to 19 | # follow Semantic Versioning. They should reflect the version the application is using. 20 | # It is recommended to use it with quotes. 21 | appVersion: 1.0.61 22 | -------------------------------------------------------------------------------- /charts/inferencedb/crds/inferencelogger.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: apiextensions.k8s.io/v1 2 | kind: CustomResourceDefinition 3 | metadata: 4 | name: inferenceloggers.inferencedb.aporia.com 5 | spec: 6 | scope: Namespaced 7 | group: inferencedb.aporia.com 8 | names: 9 | kind: InferenceLogger 10 | plural: inferenceloggers 11 | singular: inferencelogger 12 | shortNames: 13 | - inflog 14 | versions: 15 | - name: v1alpha1 16 | served: true 17 | storage: true 18 | schema: 19 | openAPIV3Schema: 20 | type: object 21 | properties: 22 | spec: 23 | type: object 24 | properties: 25 | topic: 26 | type: string 27 | schema: 28 | type: object 29 | properties: 30 | type: 31 | type: string 32 | config: 33 | type: object 34 | x-kubernetes-preserve-unknown-fields: true 35 | events: 36 | type: object 37 | properties: 38 | type: 39 | type: string 40 | config: 41 | type: object 42 | x-kubernetes-preserve-unknown-fields: true 43 | filters: 44 | type: object 45 | properties: 46 | modelName: 47 | type: string 48 | modelVersion: 49 | type: string 50 | destination: 51 | type: object 52 | properties: 53 | type: 54 | type: string 55 | config: 56 | type: object 57 | x-kubernetes-preserve-unknown-fields: true 58 | -------------------------------------------------------------------------------- /charts/inferencedb/templates/_helpers.tpl: -------------------------------------------------------------------------------- 1 | {{/* 2 | Expand the name of the chart. 3 | */}} 4 | {{- define "inferencedb.name" -}} 5 | {{- default .Chart.Name .Values.nameOverride | trunc 63 | trimSuffix "-" }} 6 | {{- end }} 7 | 8 | 9 | {{/* 10 | Create a default fully qualified app name. 11 | We truncate at 63 chars because some Kubernetes name fields are limited to this (by the DNS naming spec). 12 | If release name contains chart name it will be used as a full name. 13 | */}} 14 | {{- define "inferencedb.fullname" -}} 15 | {{- if .Values.fullnameOverride }} 16 | {{- .Values.fullnameOverride | trunc 63 | trimSuffix "-" }} 17 | {{- else }} 18 | {{- $name := default .Chart.Name .Values.nameOverride }} 19 | {{- if contains $name .Release.Name }} 20 | {{- .Release.Name | trunc 63 | trimSuffix "-" }} 21 | {{- else }} 22 | {{- printf "%s-%s" .Release.Name $name | trunc 63 | trimSuffix "-" }} 23 | {{- end }} 24 | {{- end }} 25 | {{- end }} 26 | 27 | {{/* 28 | Create chart name and version as used by the chart label. 29 | */}} 30 | {{- define "inferencedb.chart" -}} 31 | {{- printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" | trunc 63 | trimSuffix "-" }} 32 | {{- end }} 33 | 34 | {{/* 35 | Selector labels 36 | */}} 37 | {{- define "inferencedb.selectorLabels" -}} 38 | app.kubernetes.io/name: {{ include "inferencedb.name" . }} 39 | app.kubernetes.io/instance: {{ .Release.Name }} 40 | {{- end }} 41 | 42 | {{/* 43 | Common labels 44 | */}} 45 | {{- define "inferencedb.labels" -}} 46 | helm.sh/chart: {{ include "inferencedb.chart" . }} 47 | {{ include "inferencedb.selectorLabels" . }} 48 | {{- if .Chart.AppVersion }} 49 | app.kubernetes.io/version: {{ .Chart.AppVersion | quote }} 50 | {{- end }} 51 | app.kubernetes.io/managed-by: {{ .Release.Service }} 52 | {{- end }} 53 | 54 | {{/* 55 | Create the name of the service account to use 56 | */}} 57 | {{- define "inferencedb.serviceAccountName" -}} 58 | {{- if .Values.serviceAccount.create }} 59 | {{- default (include "inferencedb.fullname" .) .Values.serviceAccount.name }} 60 | {{- else }} 61 | {{- default "default" .Values.serviceAccount.name }} 62 | {{- end }} 63 | {{- end }} 64 | -------------------------------------------------------------------------------- /charts/inferencedb/templates/deployment.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: apps/v1 2 | kind: Deployment 3 | metadata: 4 | name: {{ include "inferencedb.fullname" . }} 5 | namespace: {{ .Release.Namespace }} 6 | labels: 7 | {{- include "inferencedb.labels" . | nindent 4 }} 8 | spec: 9 | replicas: {{ .Values.replicaCount }} 10 | selector: 11 | matchLabels: 12 | {{- include "inferencedb.selectorLabels" . | nindent 6 }} 13 | template: 14 | metadata: 15 | {{- with .Values.podAnnotations }} 16 | annotations: 17 | {{- toYaml . | nindent 8 }} 18 | {{- end }} 19 | labels: 20 | {{- include "inferencedb.selectorLabels" . | nindent 8 }} 21 | spec: 22 | {{- with .Values.imagePullSecrets }} 23 | imagePullSecrets: 24 | {{- toYaml . | nindent 8 }} 25 | {{- end }} 26 | serviceAccountName: {{ include "inferencedb.serviceAccountName" . }} 27 | securityContext: 28 | {{- toYaml .Values.podSecurityContext | nindent 8 }} 29 | containers: 30 | - name: {{ .Chart.Name }} 31 | securityContext: 32 | {{- toYaml .Values.securityContext | nindent 12 }} 33 | image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}" 34 | imagePullPolicy: {{ .Values.image.pullPolicy }} 35 | ports: 36 | - name: faust 37 | containerPort: 6066 38 | protocol: TCP 39 | env: 40 | - name: INFERENCEDB_KAFKA_BROKER 41 | value: {{ .Values.kafka.broker | quote }} 42 | - name: INFERENCEDB_KAFKA_SCHEMA_REGISTRY_URL 43 | value: {{ .Values.kafka.schemaRegistryUrl | quote }} 44 | - name: INFERENCEDB_KAFKA_CONNECT_URL 45 | value: {{ .Values.kafka.connectUrl | quote }} 46 | - name: INFERENCEDB_LOG_LEVEL 47 | value: {{ .Values.logLevel }} 48 | - name: INFERENCEDB_CONFIG_PROVIDER 49 | value: {{ .Values.configProvider.type | quote }} 50 | - name: INFERENCEDB_CONFIG_PROVIDER_ARGS 51 | value: {{ .Values.configProvider.args | toJson | quote }} 52 | resources: 53 | {{- toYaml .Values.resources | nindent 12 }} 54 | {{- if .Values.kafka.tls.secretName }} 55 | volumeMounts: 56 | - name: kafka-tls 57 | mountPath: /etc/kafka-tls 58 | readOnly: true 59 | volumes: 60 | - name: kafka-tls 61 | secret: 62 | secretName: {{ .Values.kafka.tls.secretName }} 63 | {{- end }} 64 | {{- with .Values.nodeSelector }} 65 | nodeSelector: 66 | {{- toYaml . | nindent 8 }} 67 | {{- end }} 68 | {{- with .Values.affinity }} 69 | affinity: 70 | {{- toYaml . | nindent 8 }} 71 | {{- end }} 72 | {{- with .Values.tolerations }} 73 | tolerations: 74 | {{- toYaml . | nindent 8 }} 75 | {{- end }} 76 | -------------------------------------------------------------------------------- /charts/inferencedb/templates/role.yaml: -------------------------------------------------------------------------------- 1 | {{- if .Values.rbac.create }} 2 | apiVersion: rbac.authorization.k8s.io/v1 3 | kind: {{ if .Values.rbac.namespaced }}Role{{ else }}ClusterRole{{ end }} 4 | metadata: 5 | name: {{ include "inferencedb.fullname" . }} 6 | {{ if .Values.rbac.namespaced }}namespace: {{ .Release.Namespace }}{{ end }} 7 | labels: 8 | {{- include "inferencedb.labels" . | nindent 4 }} 9 | rules: 10 | - apiGroups: 11 | - "inferencedb.aporia.com" # indicates the core API group 12 | resources: 13 | - "inferenceloggers" 14 | verbs: 15 | - "get" 16 | - "list" 17 | - "watch" 18 | - "create" 19 | - "update" 20 | - "patch" 21 | - "delete" 22 | {{- end }} 23 | -------------------------------------------------------------------------------- /charts/inferencedb/templates/rolebinding.yaml: -------------------------------------------------------------------------------- 1 | {{- if .Values.rbac.create }} 2 | apiVersion: rbac.authorization.k8s.io/v1 3 | kind: {{ if .Values.rbac.namespaced }}RoleBinding{{ else }}ClusterRoleBinding{{ end }} 4 | metadata: 5 | name: {{ include "inferencedb.fullname" . }} 6 | {{ if .Values.rbac.namespaced }}namespace: {{ .Release.Namespace }}{{ end }} 7 | labels: 8 | {{- include "inferencedb.labels" . | nindent 4 }} 9 | roleRef: 10 | apiGroup: rbac.authorization.k8s.io 11 | kind: {{ if .Values.rbac.namespaced }}Role{{ else }}ClusterRole{{ end }} 12 | name: {{ include "inferencedb.fullname" . }} 13 | subjects: 14 | - kind: ServiceAccount 15 | name: {{ include "inferencedb.serviceAccountName" . }} 16 | namespace: {{ .Release.Namespace }} 17 | {{- end }} 18 | -------------------------------------------------------------------------------- /charts/inferencedb/templates/service.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: v1 2 | kind: Service 3 | metadata: 4 | name: {{ include "inferencedb.fullname" . }} 5 | namespace: {{ .Release.Namespace }} 6 | labels: 7 | {{- include "inferencedb.labels" . | nindent 4 }} 8 | spec: 9 | type: {{ .Values.service.type }} 10 | ports: 11 | - port: {{ .Values.service.port }} 12 | targetPort: 6066 13 | protocol: TCP 14 | name: faust 15 | selector: 16 | {{- include "inferencedb.selectorLabels" . | nindent 4 }} 17 | -------------------------------------------------------------------------------- /charts/inferencedb/templates/serviceaccount.yaml: -------------------------------------------------------------------------------- 1 | {{- if .Values.serviceAccount.create -}} 2 | apiVersion: v1 3 | kind: ServiceAccount 4 | metadata: 5 | name: {{ include "inferencedb.serviceAccountName" . }} 6 | namespace: {{ .Release.Namespace }} 7 | labels: 8 | {{- include "inferencedb.labels" . | nindent 4 }} 9 | {{- with .Values.serviceAccount.annotations }} 10 | annotations: 11 | {{- toYaml . | nindent 4 }} 12 | {{- end }} 13 | {{- end }} 14 | -------------------------------------------------------------------------------- /charts/inferencedb/templates/servicemonitor.yaml: -------------------------------------------------------------------------------- 1 | {{- if .Values.serviceMonitor.enabled }} 2 | apiVersion: monitoring.coreos.com/v1 3 | kind: ServiceMonitor 4 | metadata: 5 | name: {{ include "inferencedb.fullname" . }} 6 | namespace: {{ .Release.Namespace }} 7 | labels: 8 | {{- include "inferencedb.labels" . | nindent 4 }} 9 | {{- if .Values.serviceMonitor.selector }} 10 | {{- toYaml .Values.serviceMonitor.selector | nindent 4 }} 11 | {{- end }} 12 | spec: 13 | selector: 14 | matchLabels: 15 | {{- include "inferencedb.selectorLabels" . | nindent 6 }} 16 | endpoints: 17 | - port: http 18 | path: {{ .Values.serviceMonitor.path }} 19 | interval: {{ .Values.serviceMonitor.interval }} 20 | {{- end }} 21 | -------------------------------------------------------------------------------- /charts/inferencedb/values.yaml: -------------------------------------------------------------------------------- 1 | # Default values for inferencedb. 2 | 3 | replicaCount: 3 4 | image: 5 | repository: ghcr.io/aporia-ai/inferencedb 6 | pullPolicy: IfNotPresent 7 | # Overrides the image tag whose default is the chart appVersion. 8 | tag: "" 9 | 10 | imagePullSecrets: [] 11 | nameOverride: "" 12 | fullnameOverride: "" 13 | 14 | service: 15 | type: ClusterIP 16 | port: 6066 17 | 18 | rbac: 19 | # Specifies whether RBAC resources should be created 20 | create: true 21 | # If this is false, the service account will be able to read Aporia's resources from any namespace. 22 | # NOTE: You also need to change this in the Kubernetes config provider's config! 23 | namespaced: false 24 | 25 | serviceAccount: 26 | # Specifies whether a service account should be created 27 | create: true 28 | # Annotations to add to the service account 29 | annotations: {} 30 | # The name of the service account to use. 31 | # If not set and create is true, a name is generated using the fullname template 32 | name: "" 33 | 34 | podAnnotations: {} 35 | 36 | podSecurityContext: {} 37 | # fsGroup: 2000 38 | 39 | securityContext: {} 40 | # capabilities: 41 | # drop: 42 | # - ALL 43 | # readOnlyRootFilesystem: true 44 | # runAsNonRoot: true 45 | # runAsUser: 1000 46 | 47 | resources: {} 48 | # We usually recommend not to specify default resources and to leave this as a conscious 49 | # choice for the user. This also increases chances charts run on environments with little 50 | # resources, such as Minikube. If you do want to specify resources, uncomment the following 51 | # lines, adjust them as necessary, and remove the curly braces after 'resources:'. 52 | # limits: 53 | # cpu: 100m 54 | # memory: 128Mi 55 | # requests: 56 | # cpu: 100m 57 | # memory: 128Mi 58 | 59 | nodeSelector: {} 60 | tolerations: [] 61 | affinity: {} 62 | 63 | serviceMonitor: 64 | ## if ServiceMonitor resources should be deployed for aporia-core 65 | enabled: false 66 | 67 | ## labels for ServiceMonitor, so that Prometheus can select it 68 | selector: 69 | prometheus: kube-prometheus 70 | 71 | ## the ServiceMonitor web endpoint path 72 | path: /metrics 73 | 74 | ## the ServiceMonitor web endpoint interval 75 | interval: "30s" 76 | 77 | logLevel: INFO 78 | 79 | kafka: 80 | broker: 81 | schemaRegistryUrl: 82 | connectUrl: 83 | tls: 84 | secretName: 85 | 86 | configProvider: 87 | type: kubernetes 88 | args: 89 | is_namespaced: false 90 | -------------------------------------------------------------------------------- /examples/kserve/kafka-broker/broker.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: eventing.knative.dev/v1 2 | kind: Broker 3 | metadata: 4 | name: sklearn-iris-broker 5 | namespace: default 6 | annotations: 7 | eventing.knative.dev/broker.class: Kafka 8 | spec: 9 | config: 10 | apiVersion: v1 11 | kind: ConfigMap 12 | name: inferencedb-kafka-broker-config 13 | namespace: knative-eventing 14 | --- 15 | apiVersion: v1 16 | kind: ConfigMap 17 | metadata: 18 | name: inferencedb-kafka-broker-config 19 | namespace: knative-eventing 20 | data: 21 | # Number of topic partitions 22 | default.topic.partitions: "8" 23 | # Replication factor of topic messages. 24 | default.topic.replication.factor: "1" 25 | # A comma separated list of bootstrap servers. (It can be in or out the k8s cluster) 26 | bootstrap.servers: "kafka-cp-kafka.default.svc.cluster.local:9092" 27 | -------------------------------------------------------------------------------- /examples/kserve/kafka-broker/inferencelogger.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: inferencedb.aporia.com/v1alpha1 2 | kind: InferenceLogger 3 | metadata: 4 | name: sklearn-iris 5 | namespace: default 6 | spec: 7 | # NOTE: The format is knative-broker-- 8 | topic: knative-broker-default-sklearn-iris-broker 9 | events: 10 | type: kserve 11 | config: {} 12 | destination: 13 | type: confluent-s3 14 | config: 15 | url: s3://aporia-data/inferencedb 16 | format: parquet 17 | awsRegion: us-east-2 18 | 19 | # Optional - Only if you want to override column names 20 | schema: 21 | type: avro 22 | config: 23 | columnNames: 24 | inputs: [sepal_width, petal_width, sepal_length, petal_length] 25 | outputs: [flower] 26 | -------------------------------------------------------------------------------- /examples/kserve/kafka-broker/inferenceservice.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: serving.kserve.io/v1beta1 2 | kind: InferenceService 3 | metadata: 4 | name: sklearn-iris 5 | spec: 6 | predictor: 7 | logger: 8 | mode: all 9 | url: http://kafka-broker-ingress.knative-eventing.svc.cluster.local/default/sklearn-iris-broker 10 | sklearn: 11 | protocolVersion: v2 12 | storageUri: gs://seldon-models/sklearn/iris 13 | -------------------------------------------------------------------------------- /examples/kserve/multimodelserving/broker.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: eventing.knative.dev/v1 2 | kind: Broker 3 | metadata: 4 | name: sklearn-mms-broker 5 | namespace: default 6 | annotations: 7 | eventing.knative.dev/broker.class: Kafka 8 | spec: 9 | config: 10 | apiVersion: v1 11 | kind: ConfigMap 12 | name: inferencedb-kafka-broker-config 13 | namespace: knative-eventing 14 | --- 15 | apiVersion: v1 16 | kind: ConfigMap 17 | metadata: 18 | name: inferencedb-kafka-broker-config 19 | namespace: knative-eventing 20 | data: 21 | # Number of topic partitions 22 | default.topic.partitions: "8" 23 | # Replication factor of topic messages. 24 | default.topic.replication.factor: "1" 25 | # A comma separated list of bootstrap servers. (It can be in or out the k8s cluster) 26 | bootstrap.servers: "kafka-cp-kafka.default.svc.cluster.local:9092" 27 | -------------------------------------------------------------------------------- /examples/kserve/multimodelserving/inferenceloggers.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: inferencedb.aporia.com/v1alpha1 2 | kind: InferenceLogger 3 | metadata: 4 | name: sklearn-mms-model-1 5 | namespace: default 6 | spec: 7 | # NOTE: The format is knative-broker-- 8 | topic: knative-broker-default-sklearn-mms-broker 9 | schema: 10 | type: avro 11 | config: 12 | columnNames: 13 | inputs: [sepal_width, petal_width, sepal_length, petal_length] 14 | outputs: [flower] 15 | events: 16 | type: kserve 17 | config: {} 18 | filters: 19 | modelName: model1 20 | # modelVersion: v1 21 | destination: 22 | type: confluent-s3 23 | config: 24 | url: s3://aporia-data/inferencedb 25 | format: parquet 26 | awsRegion: us-east-2 27 | --- 28 | apiVersion: inferencedb.aporia.com/v1alpha1 29 | kind: InferenceLogger 30 | metadata: 31 | name: sklearn-mms-model-2 32 | namespace: default 33 | spec: 34 | # NOTE: The format is knative-broker-- 35 | topic: knative-broker-default-sklearn-mms-broker 36 | schema: 37 | type: avro 38 | config: 39 | columnNames: 40 | inputs: [sepal_width, petal_width, sepal_length, petal_length] 41 | outputs: [flower] 42 | events: 43 | type: kserve 44 | config: {} 45 | filters: 46 | modelName: model2 47 | # modelVersion: v2 48 | destination: 49 | type: confluent-s3 50 | config: 51 | url: s3://aporia-data/inferencedb 52 | format: parquet 53 | awsRegion: us-east-2 54 | -------------------------------------------------------------------------------- /examples/kserve/multimodelserving/inferenceservice.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: serving.kserve.io/v1beta1 2 | kind: InferenceService 3 | metadata: 4 | name: sklearn-mms 5 | spec: 6 | predictor: 7 | minReplicas: 1 8 | logger: 9 | mode: all 10 | url: http://kafka-broker-ingress.knative-eventing.svc.cluster.local/default/sklearn-mms-broker 11 | sklearn: 12 | name: sklearn-mms 13 | protocolVersion: v2 14 | resources: 15 | limits: 16 | cpu: 500m 17 | memory: 1Gi 18 | requests: 19 | cpu: 500m 20 | memory: 1Gi 21 | -------------------------------------------------------------------------------- /examples/kserve/multimodelserving/trainedmodels.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: serving.kserve.io/v1alpha1 2 | kind: TrainedModel 3 | metadata: 4 | name: model1 5 | spec: 6 | inferenceService: sklearn-mms 7 | model: 8 | storageUri: gs://seldon-models/sklearn/mms/lr_model 9 | framework: sklearn 10 | memory: 512Mi 11 | --- 12 | apiVersion: serving.kserve.io/v1alpha1 13 | kind: TrainedModel 14 | metadata: 15 | name: model2 16 | spec: 17 | inferenceService: sklearn-mms 18 | model: 19 | storageUri: gs://seldon-models/sklearn/mms/lr_model 20 | framework: sklearn 21 | memory: 512Mi 22 | -------------------------------------------------------------------------------- /logo.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "inferencedb" 3 | version = "1.0.0" 4 | description = "" 5 | authors = ["Aporia "] 6 | 7 | [tool.poetry.dependencies] 8 | python = ">=3.9,<4.0" 9 | pydantic = "^1.9.0" 10 | PyYAML = "^6.0" 11 | python-schema-registry-client = "^2.2.2" 12 | pyhumps = "^3.5.3" 13 | numpy = "^1.22.2" 14 | pyarrow = "^7.0.0" 15 | pandavro = "^1.6.0" 16 | quickle = "^0.4.0" 17 | faust-streaming = {version = "^0.8.4", extras = ["rocksdb", "fast"]} 18 | python-dateutil = "^2.8.2" 19 | python-json-logger = "^2.0.2" 20 | 21 | [tool.poetry.dev-dependencies] 22 | pytest = "^5.2" 23 | black = "^21.12b0" 24 | flake8 = "^4.0.1" 25 | flake8-black = "^0.2.3" 26 | flake8-import-order = "^0.18.1" 27 | flake8-bugbear = "^21.11.29" 28 | flake8-bandit = "^2.1.2" 29 | mypy = "^0.910" 30 | flake8-annotations = "^2.7.0" 31 | flake8-docstrings = "^1.6.0" 32 | darglint = "^1.8.1" 33 | pytest-asyncio = "^0.16.0" 34 | isort = "^5.10.1" 35 | types-PyYAML = "^6.0.1" 36 | 37 | [tool.poetry.scripts] 38 | inferencedb = "inferencedb.main:main" 39 | 40 | [tool.black] 41 | line-length = 100 42 | 43 | [tool.isort] 44 | profile = "black" 45 | force_sort_within_sections = true 46 | lexicographical = true 47 | order_by_type = false 48 | group_by_package = true 49 | no_lines_before = ['LOCALFOLDER'] 50 | line_length = 100 51 | 52 | [tool.mypy] 53 | show_error_codes = true 54 | plugins = "pydantic.mypy" 55 | 56 | [[tool.mypy.overrides]] 57 | module = ["pandas", "pandas.*", "pytest", "apscheduler.*", "scipy.*", "uvicorn", "pyarrow.*", "sklearn.*", "parsedatetime"] 58 | ignore_missing_imports = true 59 | 60 | [build-system] 61 | requires = ["poetry-core>=1.0.0"] 62 | build-backend = "poetry.core.masonry.api" 63 | -------------------------------------------------------------------------------- /skaffold.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: skaffold/v2beta26 2 | kind: Config 3 | build: 4 | artifacts: 5 | - image: ghcr.io/aporia-ai/inferencedb 6 | deploy: 7 | helm: 8 | releases: 9 | - name: inferencedb 10 | chartPath: ./charts/inferencedb 11 | artifactOverrides: 12 | image: ghcr.io/aporia-ai/inferencedb 13 | imageStrategy: 14 | helm: {} 15 | setValues: 16 | image: 17 | pullPolicy: IfNotPresent 18 | devMode: true 19 | logLevel: DEBUG 20 | kafka: 21 | broker: kafka://kafka-cp-kafka:9092 22 | schemaRegistryUrl: http://kafka-cp-schema-registry:8081 23 | connectUrl: http://kafka-cp-kafka-connect:8083 24 | serviceMonitor: 25 | enabled: false 26 | selector: 27 | release: prometheus 28 | -------------------------------------------------------------------------------- /src/inferencedb/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aporia-ai/inferencedb/bee2cbe67f1e8ae3fa33948ac73f8656942f43cb/src/inferencedb/__init__.py -------------------------------------------------------------------------------- /src/inferencedb/app.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | from inferencedb.config.factory import create_config_provider, generate_config_from_dict 3 | from inferencedb.config.providers.kubernetes_config_provider import KubernetesConfigProvider 4 | from inferencedb.core.inference_logger import InferenceLogger 5 | from inferencedb.core.logging_utils import generate_logging_config, init_logging 6 | import faust 7 | import json 8 | import os 9 | import ssl 10 | import logging 11 | 12 | import aiohttp 13 | 14 | from inferencedb.utils.asyncio_utils import cancling_background_task 15 | from .settings import Settings 16 | from schema_registry.client import AsyncSchemaRegistryClient 17 | 18 | 19 | settings = Settings() 20 | 21 | # Initialize logging 22 | init_logging(log_level=settings.log_level) 23 | logging.info("Worker started.") 24 | 25 | 26 | # Load TLS certificates if necessary 27 | ssl_context = None 28 | if os.path.isdir("/etc/kafka-tls"): 29 | ssl_context = ssl.create_default_context( 30 | purpose=ssl.Purpose.SERVER_AUTH, 31 | cafile="/etc/kafka-tls/ca.crt", 32 | ) 33 | ssl_context.load_cert_chain("/etc/kafka-tls/user.crt", keyfile="/etc/kafka-tls/user.key") 34 | 35 | 36 | # Create Faust app 37 | app = faust.App( 38 | id="inferencedb", 39 | broker=settings.kafka_broker, 40 | store="rocksdb://", 41 | autodiscover=True, 42 | origin="inferencedb", 43 | broker_credentials=ssl_context, 44 | 45 | # Faust's internal logging level is INFO because DEBUG is just unreadable. 46 | # FUTURE: Maybe add here support for WARNING, ERROR, CRITICAL. 47 | logging_config=generate_logging_config(log_level="INFO"), 48 | ) 49 | 50 | 51 | config = generate_config_from_dict(json.loads(os.environ["CONFIG"])) 52 | 53 | # Connect to schema registry. 54 | schema_registry = AsyncSchemaRegistryClient(url=settings.kafka_schema_registry_url) 55 | 56 | # Remove old InferenceLogger Kafka connectors. 57 | @app.task(on_leader=True) 58 | async def remove_old_inferenceloggers(): 59 | config_provider = create_config_provider( 60 | name=settings.config_provider, 61 | params=settings.config_provider_args, 62 | ) 63 | 64 | async def remove_kafka_connector(inference_logger: dict): 65 | metadata = inference_logger['metadata'] 66 | 67 | # FUTURE: These should be in a separate utility function, which we can reuse in InferenceLogger and Destination. 68 | logger_name = f"{metadata['namespace']}-{metadata['name']}" 69 | connector_name = f"inferencedb-{logger_name}" 70 | 71 | for _ in range(10): 72 | async with aiohttp.ClientSession( 73 | base_url=settings.kafka_connect_url, 74 | timeout=aiohttp.ClientTimeout(total=5), 75 | ) as session: 76 | async with session.delete(url=f"/connectors/{connector_name}") as response: 77 | if response.status == 409: 78 | # Kafka connect rebalance 79 | await asyncio.sleep(3) 80 | continue 81 | elif response.status not in (204, 404): 82 | raise RuntimeError(f"Failed to delete Kafka Connector: {connector_name}") 83 | 84 | raise RuntimeError(f"Could not delete connector because of Kafka Connect rebalance.") 85 | 86 | await config_provider.manage_finalizers(remove_kafka_connector) 87 | 88 | 89 | # Create all inference loggers. 90 | for item in config.inference_loggers: 91 | inference_logger = InferenceLogger(app, schema_registry, item) 92 | inference_logger.register() 93 | -------------------------------------------------------------------------------- /src/inferencedb/config/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aporia-ai/inferencedb/bee2cbe67f1e8ae3fa33948ac73f8656942f43cb/src/inferencedb/config/__init__.py -------------------------------------------------------------------------------- /src/inferencedb/config/component.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, Optional 2 | 3 | from inferencedb.core.base_model import BaseModel 4 | 5 | 6 | class ComponentConfig(BaseModel): 7 | """Generic config for registered components.""" 8 | 9 | type: str 10 | config: Optional[Dict[str, Any]] = None 11 | -------------------------------------------------------------------------------- /src/inferencedb/config/config.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, List, Optional 2 | from inferencedb.registry.decorators import event_processor 3 | 4 | from pydantic import Field, Extra 5 | 6 | from inferencedb.core.base_model import BaseModel 7 | from .component import ComponentConfig 8 | 9 | 10 | class InferenceLoggerConfig(BaseModel): 11 | """A class representing the InferenceLogger config.""" 12 | 13 | class Config: 14 | """Configuration for InferenceLoggerConfig.""" 15 | 16 | extra = Extra.forbid 17 | 18 | name: str 19 | topic: str 20 | event_processor: ComponentConfig = Field(None, alias="events") 21 | schema_provider: Optional[ComponentConfig] = Field(None, alias="schema") 22 | filters: Optional[Dict[str, str]] 23 | destination: ComponentConfig 24 | 25 | 26 | class V1Alpha1ConfigHeader(BaseModel): 27 | """A class representing the header for the config.""" 28 | 29 | class Config: 30 | """Configuration for V1Alpha1ConfigHeader.""" 31 | 32 | # Ignore extra fields, so we can parse only the header out of the entire config 33 | extra = Extra.ignore 34 | 35 | api_version: str 36 | kind: str 37 | 38 | 39 | class V1Alpha1Config(V1Alpha1ConfigHeader): 40 | """A class representing the v1alpha1 config.""" 41 | 42 | class Config: 43 | """Configuration for V1Alpha1Config.""" 44 | 45 | extra = Extra.forbid 46 | 47 | inference_loggers: List[InferenceLoggerConfig] = [] 48 | -------------------------------------------------------------------------------- /src/inferencedb/config/factory.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, Optional 2 | 3 | from inferencedb.config.config import V1Alpha1Config, V1Alpha1ConfigHeader 4 | from inferencedb.config.providers.config_provider import ConfigProvider 5 | from inferencedb.registry.factory import create_registered_object 6 | from inferencedb.registry.registry import RegisteredObjectType 7 | 8 | HEADER_VERSIONS = {"v1alpha1": V1Alpha1ConfigHeader} 9 | CONFIG_VERSIONS = {"v1alpha1": V1Alpha1Config} 10 | 11 | 12 | def generate_config_header_from_dict(config_dict: Dict[str, Any]) -> V1Alpha1ConfigHeader: 13 | """Creates a config from a dictionary. 14 | 15 | Args: 16 | config_dict: The dictionary to create the config from. 17 | 18 | Returns: 19 | The config created from the dictionary. 20 | """ 21 | api_version = config_dict["api_version"].lower() 22 | if api_version not in HEADER_VERSIONS: 23 | raise ValueError( 24 | "Unsupported API version." f"Supported versions are {list(HEADER_VERSIONS.keys())}" 25 | ) 26 | 27 | header_class = HEADER_VERSIONS[api_version] 28 | return header_class(**config_dict) 29 | 30 | 31 | def generate_config_from_dict(config_dict: Dict[str, Any]) -> V1Alpha1Config: 32 | """Creates a config from a dictionary. 33 | 34 | Args: 35 | config_dict: The dictionary to create the config from. 36 | 37 | Returns: 38 | The config created from the dictionary. 39 | """ 40 | api_version = config_dict["api_version"].lower() 41 | if api_version not in CONFIG_VERSIONS: 42 | raise ValueError( 43 | "Unsupported API version." f"Supported versions are {list(CONFIG_VERSIONS.keys())}" 44 | ) 45 | 46 | config_class = CONFIG_VERSIONS[api_version] 47 | return config_class(**config_dict) 48 | 49 | 50 | def create_config_provider(name: str, params: Optional[Dict[str, Any]] = None) -> ConfigProvider: 51 | """Creates a config provider object. 52 | 53 | Args: 54 | name: Config provider name 55 | config: Config provider configuration 56 | 57 | Returns: 58 | ConfigProvider object 59 | """ 60 | return create_registered_object( 61 | object_type=RegisteredObjectType.CONFIG_PROVIDER, name=name, params=params 62 | ) 63 | -------------------------------------------------------------------------------- /src/inferencedb/config/providers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aporia-ai/inferencedb/bee2cbe67f1e8ae3fa33948ac73f8656942f43cb/src/inferencedb/config/providers/__init__.py -------------------------------------------------------------------------------- /src/inferencedb/config/providers/config_provider.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from asyncio import Event 3 | 4 | 5 | class ConfigProvider(ABC): 6 | """Abstract base class for config providers.""" 7 | 8 | def __init__(self): 9 | """Initializes a ConfigProvider.""" 10 | self._config = {} 11 | self._update_event = Event() 12 | 13 | def get_config(self) -> dict: 14 | """Returns the current configuration.""" 15 | return self._config 16 | 17 | def update_config(self, new_config: dict): 18 | """Updates the current configuration. 19 | 20 | Args: 21 | new_config: New config 22 | """ 23 | if new_config != self._config: 24 | self._config = new_config 25 | self._update_event.set() 26 | 27 | async def wait_for_update(self): 28 | """Waits for the current configuration to be updated.""" 29 | await self._update_event.wait() 30 | self._update_event.clear() 31 | 32 | @abstractmethod 33 | async def run(self): 34 | """Runs the config provider.""" 35 | ... 36 | -------------------------------------------------------------------------------- /src/inferencedb/config/providers/file_config_provider.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | import yaml 4 | 5 | from inferencedb.registry.decorators import config_provider 6 | from .config_provider import ConfigProvider 7 | 8 | 9 | @config_provider("file") 10 | class FileConfigProvider(ConfigProvider): 11 | """File config provider.""" 12 | 13 | def __init__(self, file_path: Path): 14 | """Initializes a FileConfigProvider. 15 | 16 | Args: 17 | file_path: Config file path. 18 | """ 19 | super().__init__() 20 | self.file_path = file_path 21 | 22 | async def run(self): 23 | """See base class.""" 24 | # The file config provider does not support dynamic updates 25 | with open(self.file_path, "r") as config_file: 26 | config = yaml.safe_load(config_file) 27 | 28 | self.update_config(config) 29 | -------------------------------------------------------------------------------- /src/inferencedb/config/providers/kubernetes_config_provider.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import logging 3 | import os 4 | from pathlib import Path 5 | import ssl 6 | from typing import Awaitable, Callable, List, Optional 7 | 8 | import aiohttp 9 | from inferencedb.config.config import InferenceLoggerConfig 10 | 11 | from inferencedb.registry.decorators import config_provider 12 | from .config_provider import ConfigProvider 13 | 14 | SERVICE_TOKEN_FILENAME = Path("/var/run/secrets/kubernetes.io/serviceaccount/token") 15 | SERVICE_CERT_FILENAME = Path("/var/run/secrets/kubernetes.io/serviceaccount/ca.crt") 16 | NAMESPACE_FILENAME = Path("/var/run/secrets/kubernetes.io/serviceaccount/namespace") 17 | 18 | FINALIZER_NAME = "inferencedb.aporia.com/finalizer" 19 | 20 | DEFAULT_HTTP_TIMEOUT_SEC = 5 21 | 22 | @config_provider("kubernetes") 23 | class KubernetesConfigProvider(ConfigProvider): 24 | """Kubernetes config provider.""" 25 | 26 | def __init__( 27 | self, 28 | is_namespaced: bool, 29 | polling_interval_sec: int = 5, 30 | ): 31 | """Initializes a KubernetesConfigProvider.""" 32 | super().__init__() 33 | self._is_namespaced = is_namespaced 34 | self._polling_interval_sec = polling_interval_sec 35 | 36 | # Read Kubernetes in-cluster token. 37 | with open(SERVICE_TOKEN_FILENAME) as token_file: 38 | token = token_file.read() 39 | 40 | # Read Kubernetes namespace. 41 | with open(NAMESPACE_FILENAME) as namespace_file: 42 | self._namespace = namespace_file.read() 43 | 44 | # Prepare aiohttp session. 45 | ssl_context = ssl.create_default_context(cafile=str(SERVICE_CERT_FILENAME)) 46 | connector = aiohttp.TCPConnector(ssl_context=ssl_context, verify_ssl=True) 47 | 48 | # URL Prefix - for some reason this doesn't work with aiohttp's base_url 49 | self._url_prefix = "/apis/inferencedb.aporia.com/v1alpha1" 50 | if is_namespaced: 51 | self._url_prefix += f"/namespaces/{self._namespace}" 52 | 53 | # The KUBERNETES_SERVICE_HOST and KUBERNETES_SERVICE_PORT environment variables are automatically 54 | # defined by Kubernetes to enable pods to easily access the Kubernetes API. 55 | host = os.environ["KUBERNETES_SERVICE_HOST"] # "kubernetes.default.svc" 56 | port = os.environ["KUBERNETES_SERVICE_PORT"] 57 | 58 | # HTTP session configuration 59 | self._session_config = { 60 | "base_url": f"https://{host}:{port}", 61 | "headers": {"Authorization": f"Bearer {token}"}, 62 | "connector": connector, 63 | "timeout": aiohttp.ClientTimeout(total=DEFAULT_HTTP_TIMEOUT_SEC) 64 | } 65 | 66 | async def run(self): 67 | """See base class.""" 68 | async with aiohttp.ClientSession(**self._session_config) as session: 69 | while True: 70 | config = await self._fetch_config(session) 71 | if config is not None: 72 | self.update_config(config) 73 | 74 | await asyncio.sleep(self._polling_interval_sec) 75 | 76 | async def _fetch_config(self, session: aiohttp.ClientSession) -> Optional[dict]: 77 | """Fetch configuration from Kubernetes resources. 78 | 79 | Args: 80 | session: AIOHttp session to use 81 | 82 | Returns: 83 | Configuration dict 84 | """ 85 | 86 | try: 87 | k8s_inference_loggers = await self._get_k8s_resources(session, "inferenceloggers") 88 | except RuntimeError: 89 | logging.error("Fetching K8s resources failed.", exc_info=True) 90 | return None 91 | 92 | # Iterate all InferenceLogger resources 93 | try: 94 | return { 95 | "api_version": "v1alpha1", 96 | "kind": "Config", 97 | "inferenceLoggers": [{ 98 | "name": f'{item["metadata"]["namespace"]}-{item["metadata"]["name"]}', 99 | "topic": item["spec"]["topic"], 100 | "events": item["spec"]["events"], 101 | "schema": item["spec"].get("schema"), 102 | "filters": item["spec"].get("filters"), 103 | "destination": item["spec"]["destination"], 104 | } for item in k8s_inference_loggers if item["metadata"].get("deletionTimestamp") is None], 105 | } 106 | except KeyError: 107 | logging.error("Invalid configuration format.", exc_info=True) 108 | return None 109 | 110 | async def _get_k8s_resources( 111 | self, session: aiohttp.ClientSession, plural_name: str 112 | ) -> List[dict]: 113 | """Fetch Kubernetes resources. 114 | 115 | Args: 116 | session: AIOHttp session to use 117 | plural_name: Plural lower case name of the resource (e.g pods) 118 | 119 | Returns: 120 | List of the resources. 121 | """ 122 | async with session.get(f"{self._url_prefix}/{plural_name}") as response: 123 | # All 2XX status codes (OK, CREATED, etc.) are considered success responses. 124 | # 3XX status codes are considered failures because we don't really handle redirects, 125 | # proxies, etc. All other statuses are obviously errors. 126 | if response.status < 200 or response.status >= 300: 127 | # This always raises an exception 128 | raise RuntimeError( 129 | f"Failed to fetch K8s {plural_name} resources (HTTP status: {response.status})" 130 | ) 131 | 132 | response_content = await response.json() 133 | 134 | if "items" not in response_content: 135 | raise RuntimeError("Invalid K8s API response (couldn't find 'items' in JSON)") 136 | 137 | return response_content["items"] 138 | 139 | 140 | async def manage_finalizers(self, finalizer: Callable[[dict], Awaitable[None]]): 141 | async with aiohttp.ClientSession(**self._session_config) as session: 142 | try: 143 | k8s_inference_loggers = await self._get_k8s_resources(session, "inferenceloggers") 144 | except RuntimeError: 145 | logging.error("Fetching K8s resources failed.", exc_info=True) 146 | return None 147 | 148 | for inference_logger in k8s_inference_loggers: 149 | metadata = inference_logger["metadata"] 150 | finalizers = metadata.get("finalizers", []) 151 | 152 | # If this resource doesn't have a finalizer and it wasn't marked for deletion, 153 | # add a finalizer. 154 | if FINALIZER_NAME not in finalizers and metadata.get("deletionTimestamp") is None: 155 | await self._patch_finalizers( 156 | session=session, 157 | kind_plural="inferenceloggers", 158 | name=metadata["name"], 159 | namespace=metadata["namespace"], 160 | finalizers=[*finalizers, FINALIZER_NAME], 161 | ) 162 | 163 | # If this resource has a finalizer and it _was_ marked for deletion, call the user-provided 164 | # finalizer function and remove the K8s finalizer from the resource, so K8s can actually delete it. 165 | elif FINALIZER_NAME in finalizers and metadata.get("deletionTimestamp") is not None: 166 | await finalizer(inference_logger) 167 | 168 | # Remove K8s finalizer. 169 | await self._patch_finalizers( 170 | session=session, 171 | kind_plural="inferenceloggers", 172 | name=metadata["name"], 173 | namespace=metadata["namespace"], 174 | finalizers=[item for item in finalizers if item != FINALIZER_NAME] 175 | ) 176 | 177 | async def _patch_finalizers(self, 178 | session: aiohttp.ClientSession, 179 | kind_plural: str, 180 | name: str, 181 | namespace: str, 182 | finalizers: List[str] 183 | ): 184 | patch = [{ 185 | "op": "add", 186 | "path": "/metadata/finalizers", 187 | "value": finalizers, 188 | }] 189 | 190 | # Build resource URL 191 | url_components = [self._url_prefix] 192 | 193 | if not self._is_namespaced: 194 | url_components = [*url_components, "namespaces", namespace] 195 | elif self._namespace != namespace: 196 | raise RuntimeError("Cannot patch finalizers for a resource in a different namespace when in namespaced mode.") 197 | 198 | url_components = [*url_components, kind_plural, name] 199 | 200 | # Patch! 201 | async with session.patch( 202 | url="/".join(url_components), 203 | json=patch, 204 | headers={"content-type": "application/json-patch+json"} 205 | ) as response: 206 | # All 2XX status codes (OK, CREATED, etc.) are considered success responses. 207 | # 3XX status codes are considered failures because we don't really handle redirects, 208 | # proxies, etc. All other statuses are obviously errors. 209 | if response.status < 200 or response.status >= 300: 210 | # This always raises an exception 211 | raise RuntimeError( 212 | f"Failed to patch K8s {kind_plural} resources (HTTP status: {response.status})" 213 | ) 214 | -------------------------------------------------------------------------------- /src/inferencedb/core/base_model.py: -------------------------------------------------------------------------------- 1 | from pydantic import BaseModel as PydanticBaseModel 2 | from pydantic.main import Extra 3 | from humps import camelize 4 | 5 | 6 | def to_camel(string): 7 | return camelize(string) 8 | 9 | 10 | class BaseModel(PydanticBaseModel): 11 | """Base class for all pydantic models. 12 | 13 | This class was added to change the pydantic default behavior globally. 14 | Reference: https://pydantic-docs.helpmanual.io/usage/model_config/#change-behaviour-globally 15 | """ 16 | 17 | class Config: 18 | """Configuration for BaseModel.""" 19 | 20 | extra = Extra.forbid 21 | arbitrary_types_allowed = True 22 | alias_generator = to_camel 23 | allow_population_by_field_name = True 24 | -------------------------------------------------------------------------------- /src/inferencedb/core/inference.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass 2 | from typing import Any, Dict, Optional, cast 3 | from datetime import datetime 4 | 5 | import pandas as pd 6 | import quickle 7 | from dateutil.parser import isoparse 8 | 9 | from inferencedb.utils.pandas_utils import deserialize_dataframe, serialize_dataframe 10 | 11 | _quickle_encoder = quickle.Encoder() 12 | _quickle_decoder = quickle.Decoder() 13 | 14 | 15 | @dataclass 16 | class Inference: 17 | id: Optional[str] = None 18 | model_name: Optional[str] = None 19 | model_version: Optional[str] = None 20 | inputs: Optional[pd.DataFrame] = None 21 | outputs: Optional[pd.DataFrame] = None 22 | occurred_at: Optional[datetime] = None 23 | 24 | def serialize(self): 25 | return _quickle_encoder.dumps({ 26 | "id": self.id, 27 | "model_name": self.model_name, 28 | "model_version": self.model_version, 29 | "inputs": serialize_dataframe(self.inputs), 30 | "outputs": serialize_dataframe(self.outputs), 31 | "occurred_at": self.occurred_at.isoformat() if self.occurred_at is not None else None, 32 | }) 33 | 34 | @staticmethod 35 | def deserialize(buf: bytes): 36 | if buf is None: 37 | return Inference() 38 | 39 | data = _quickle_decoder.loads(buf) 40 | return Inference( 41 | id=data.get("id"), 42 | model_name=data["model_name"], 43 | model_version=data["model_version"], 44 | inputs=deserialize_dataframe(data["inputs"]), 45 | outputs=deserialize_dataframe(data["outputs"]), 46 | occurred_at=isoparse(data["occurred_at"]) if data.get("occurred_at") is not None else None, 47 | ) 48 | -------------------------------------------------------------------------------- /src/inferencedb/core/inference_logger.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import faust 3 | from inferencedb.config.component import ComponentConfig 4 | 5 | from schema_registry.client import AsyncSchemaRegistryClient 6 | 7 | from inferencedb.config.config import InferenceLoggerConfig 8 | from inferencedb.event_processors.factory import create_event_processor 9 | from inferencedb.event_processors.kserve_event_processor import KServeEventProcessor 10 | from inferencedb.event_processors.json_event_processor import JSONEventProcessor 11 | from inferencedb.schema_providers.factory import create_schema_provider 12 | from inferencedb.schema_providers.avro_schema_provider import AvroSchemaProvider 13 | from inferencedb.destinations.factory import create_destination 14 | from inferencedb.destinations.confluent_s3_destination import ConfluentS3Destination 15 | 16 | DEFAULT_SCHEMA_PROVIDER_CONFIG = ComponentConfig(type="avro", config={}) 17 | 18 | 19 | class InferenceLogger: 20 | def __init__( 21 | self, 22 | app: faust.App, 23 | schema_registry: AsyncSchemaRegistryClient, 24 | config: InferenceLoggerConfig 25 | ): 26 | self._app = app 27 | self._config = config 28 | self._schema_registry = schema_registry 29 | 30 | # Create the target topic 31 | self._target_topic = app.topic(config.name) 32 | 33 | # Create schema provider 34 | schema_provider_config = config.schema_provider 35 | if schema_provider_config is None: 36 | schema_provider_config = DEFAULT_SCHEMA_PROVIDER_CONFIG 37 | 38 | self._schema_provider = create_schema_provider(schema_provider_config.type, { 39 | "logger_name": config.name, 40 | "subject": f"{self._target_topic.get_topic_name()}-value", 41 | "schema_registry": self._schema_registry, 42 | "config": schema_provider_config.config, 43 | }) 44 | 45 | # Create event processor 46 | self._source_topic = app.topic(self._config.topic) 47 | self._event_processor = create_event_processor(config.event_processor.type, { 48 | "logger_name": config.name, 49 | "app": app, 50 | "config": config.event_processor.config, 51 | }) 52 | 53 | # Create destination 54 | self._destination = create_destination(config.destination.type, { 55 | "logger_name": config.name, 56 | "topic": self._target_topic.get_topic_name(), 57 | "config": config.destination.config, 58 | }) 59 | 60 | def register(self): 61 | async def agent(stream): 62 | # Create the Kafka connector 63 | logging.debug("Creating Kafka connector") 64 | await self._destination.create_connector() 65 | 66 | # Process every inference event 67 | if hasattr(self._event_processor, "get_group_key"): 68 | stream = stream.group_by( 69 | name=self._config.name, 70 | key=self._event_processor.get_group_key, 71 | ) 72 | 73 | async for event in stream: 74 | inference = await self._event_processor.process_event(event) 75 | if inference is None: 76 | continue 77 | 78 | # Apply filters. 79 | # FUTURE: Do this in a more generic way. 80 | filters = self._config.filters 81 | if filters is not None: 82 | if filters.get("modelName") is not None and filters["modelName"] != inference.model_name: 83 | continue 84 | 85 | if filters.get("modelVersion") is not None and filters["modelVersion"] != inference.model_name: 86 | continue 87 | 88 | # Serialize each inference 89 | async for item in self._schema_provider.serialize(inference): 90 | yield item 91 | 92 | logging.debug("Starting agent") 93 | self._app.agent( 94 | name=self._config.name, 95 | channel=self._source_topic, 96 | sink=[self._target_topic] 97 | )(agent) 98 | -------------------------------------------------------------------------------- /src/inferencedb/core/logging_utils.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import logging.config 3 | 4 | 5 | def generate_logging_config(log_level: str) -> dict: 6 | """Generates logging config. 7 | 8 | Args: 9 | log_level: Log level 10 | 11 | Returns: 12 | log config dict. 13 | """ 14 | log_level = log_level.upper() 15 | log_handler = "simple" 16 | 17 | logger_config = { 18 | "level": log_level, 19 | "handlers": [log_handler], 20 | } 21 | 22 | logging_config = { 23 | "version": 1, 24 | "handlers": { 25 | "simple": { 26 | "class": "logging.StreamHandler", 27 | "formatter": "json", 28 | "stream": "ext://sys.stdout", 29 | }, 30 | }, 31 | "formatters": { 32 | "json": { 33 | "()": "pythonjsonlogger.jsonlogger.JsonFormatter", 34 | "format": "%(asctime)s:%(levelname)s:%(module)s:%(funcName)s:%(message)s", 35 | }, 36 | }, 37 | "loggers": { 38 | "sqlalchemy.engine": logger_config, 39 | "alembic": logger_config, 40 | }, 41 | "root": logger_config, 42 | "disable_existing_loggers": False, 43 | } 44 | 45 | return logging_config 46 | 47 | 48 | def init_logging(log_level: str): 49 | """Configures the logger. 50 | 51 | Args: 52 | log_level: Log level 53 | """ 54 | logging_config = generate_logging_config(log_level=log_level) 55 | logging.config.dictConfig(logging_config) 56 | -------------------------------------------------------------------------------- /src/inferencedb/destinations/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aporia-ai/inferencedb/bee2cbe67f1e8ae3fa33948ac73f8656942f43cb/src/inferencedb/destinations/__init__.py -------------------------------------------------------------------------------- /src/inferencedb/destinations/confluent_s3_destination.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import logging 3 | from typing import Any, Dict 4 | from urllib.parse import urlparse 5 | 6 | import aiohttp 7 | 8 | from inferencedb.destinations.destination import Destination 9 | from inferencedb.registry.decorators import destination 10 | from inferencedb.settings import Settings 11 | 12 | 13 | class ConnectorAlreadyExistsException(Exception): 14 | pass 15 | 16 | 17 | CONFLUENT_KAFKA_CONNECT_FORMATS = { 18 | "parquet": "io.confluent.connect.s3.format.parquet.ParquetFormat", 19 | } 20 | 21 | DEFAULT_HTTP_TIMEOUT_SEC = 5 22 | 23 | 24 | @destination("confluent-s3") 25 | class ConfluentS3Destination(Destination): 26 | def __init__(self, logger_name: str, topic: str, config: Dict[str, Any]): 27 | self._logger_name = logger_name 28 | self._topic = topic 29 | self._config = config 30 | 31 | async def create_connector(self): 32 | settings = Settings() 33 | url = urlparse(self._config["url"]) 34 | 35 | connector_name = f"inferencedb-{self._logger_name}" 36 | connector_config = { 37 | "connector.class": "io.confluent.connect.s3.S3SinkConnector", 38 | "storage.class": "io.confluent.connect.s3.storage.S3Storage", 39 | "s3.region": self._config["awsRegion"], 40 | "s3.bucket.name": url.netloc, 41 | "topics.dir": url.path.strip("/"), 42 | "flush.size": "2", 43 | "rotate.schedule.interval.ms": "20000", 44 | "auto.register.schemas": "false", 45 | "tasks.max": "1", 46 | "s3.part.size": "5242880", 47 | "timezone": "UTC", 48 | "parquet.codec": "snappy", 49 | "topics": self._topic, 50 | "s3.credentials.provider.class": "com.amazonaws.auth.DefaultAWSCredentialsProviderChain", 51 | "format.class": CONFLUENT_KAFKA_CONNECT_FORMATS[self._config["format"]], 52 | "value.converter": "io.confluent.connect.avro.AvroConverter", 53 | "key.converter": "org.apache.kafka.connect.storage.StringConverter", 54 | "schema.registry.url": settings.kafka_schema_registry_url, 55 | "value.converter.schema.registry.url": settings.kafka_schema_registry_url, 56 | **self._config.get("connector", {}) 57 | } 58 | 59 | async with aiohttp.ClientSession( 60 | base_url=settings.kafka_connect_url, 61 | timeout=aiohttp.ClientTimeout(total=DEFAULT_HTTP_TIMEOUT_SEC), 62 | ) as session: 63 | await self._upsert_connector(session, connector_name, connector_config) 64 | 65 | async def _upsert_connector(self, session: aiohttp.ClientSession, connector_name: str, connector_config: dict): 66 | for _ in range(10): 67 | async with session.put( 68 | url=f"/connectors/{connector_name}/config", 69 | json=connector_config, 70 | ) as response: 71 | if response.status in (201, 200): 72 | return 73 | elif response.status == 409: 74 | # Kafka connect rebalance 75 | await asyncio.sleep(3) 76 | continue 77 | else: 78 | raise RuntimeError(f"Could not create connector: {connector_name}, status: {response.status}") 79 | 80 | raise RuntimeError(f"Could not create connector because of Kafka Connect rebalance.") 81 | -------------------------------------------------------------------------------- /src/inferencedb/destinations/destination.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | 4 | class Destination(ABC): 5 | """Abstract base class for Destination objects.""" 6 | 7 | @abstractmethod 8 | async def create_connector(self): 9 | ... 10 | -------------------------------------------------------------------------------- /src/inferencedb/destinations/factory.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, Optional 2 | 3 | from inferencedb.registry.factory import create_registered_object 4 | from inferencedb.registry.registry import RegisteredObjectType 5 | from .destination import Destination 6 | 7 | 8 | def create_destination(type: str, params: Optional[Dict[str, Any]] = None) -> Destination: 9 | """Creates a Destination object. 10 | 11 | Args: 12 | type: Destination type 13 | params: Destination parameters 14 | 15 | Returns: 16 | Destination object 17 | """ 18 | return create_registered_object( 19 | object_type=RegisteredObjectType.DESTINATION, name=type, params=params 20 | ) 21 | -------------------------------------------------------------------------------- /src/inferencedb/event_processors/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aporia-ai/inferencedb/bee2cbe67f1e8ae3fa33948ac73f8656942f43cb/src/inferencedb/event_processors/__init__.py -------------------------------------------------------------------------------- /src/inferencedb/event_processors/event_processor.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | import pandas as pd 4 | 5 | from inferencedb.core.inference import Inference 6 | 7 | 8 | class EventProcessor(ABC): 9 | """Abstract base class for EventProcessor objects.""" 10 | 11 | @abstractmethod 12 | async def process_event(self, event) -> Inference: 13 | ... 14 | -------------------------------------------------------------------------------- /src/inferencedb/event_processors/factory.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, Optional 2 | 3 | from inferencedb.registry.factory import create_registered_object 4 | from inferencedb.registry.registry import RegisteredObjectType 5 | from .event_processor import EventProcessor 6 | 7 | 8 | def create_event_processor(type: str, params: Optional[Dict[str, Any]] = None) -> EventProcessor: 9 | """Creates an EventProcessor object. 10 | 11 | Args: 12 | type: EventProcessor type 13 | params: EventProcessor parameters 14 | 15 | Returns: 16 | EventProcessor object 17 | """ 18 | return create_registered_object( 19 | object_type=RegisteredObjectType.EVENT_PROCESSOR, name=type, params=params 20 | ) 21 | -------------------------------------------------------------------------------- /src/inferencedb/event_processors/json_event_processor.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict 2 | import pandas as pd 3 | from datetime import datetime 4 | 5 | import faust 6 | from dateutil.parser import isoparse 7 | 8 | from inferencedb.core.inference import Inference 9 | from inferencedb.registry.decorators import event_processor 10 | from .event_processor import EventProcessor 11 | 12 | 13 | @event_processor("json") 14 | class JSONEventProcessor(EventProcessor): 15 | def __init__(self, logger_name: str, app: faust.App, config: Dict[str, Any]): 16 | pass 17 | 18 | async def process_event(self, event) -> Inference: 19 | return Inference( 20 | id=event.get("id"), 21 | model_name=event.get("model_name"), 22 | model_version=event.get("model_version"), 23 | inputs=pd.DataFrame.from_records(event["inputs"]) if event.get("inputs") is not None else None, 24 | outputs=pd.DataFrame.from_records(event["outputs"]) if event.get("outputs") is not None else None, 25 | occurred_at=isoparse(event["occurred_at"]) if event.get("occurred_at") is not None else datetime.now(), 26 | ) 27 | -------------------------------------------------------------------------------- /src/inferencedb/event_processors/kserve_event_processor.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict 2 | 3 | import faust 4 | from dateutil.parser import isoparse 5 | 6 | from inferencedb.core.inference import Inference 7 | from inferencedb.registry.decorators import event_processor 8 | from inferencedb.utils.kserve_utils import parse_kserve_request, parse_kserve_response 9 | from .event_processor import EventProcessor 10 | 11 | 12 | INFERENCE_REQUEST_TYPE = "org.kubeflow.serving.inference.request" 13 | INFERENCE_RESPONSE_TYPE = "org.kubeflow.serving.inference.response" 14 | 15 | 16 | @event_processor("kserve") 17 | class KServeEventProcessor(EventProcessor): 18 | def __init__(self, 19 | logger_name: str, 20 | app: faust.App, 21 | config: Dict[str, Any]): 22 | # TODO: Validate config 23 | self._table = app.Table( 24 | name=f"kserve-event-processor-{logger_name}", 25 | value_type=bytes 26 | ).tumbling(10.0, expires=10.0) 27 | 28 | async def process_event(self, event) -> Inference: 29 | # Basic validations on the event version. 30 | event_id = faust.current_event().headers.get("ce_id", "").decode("utf-8") 31 | if not event_id: 32 | return 33 | 34 | if event is None: 35 | return 36 | 37 | # Add data to event, depending on the event type. 38 | event_type = faust.current_event().headers.get("ce_type", "").decode("utf-8") 39 | 40 | try: 41 | inference = Inference.deserialize(self._table[event_id].value()) 42 | except KeyError: 43 | inference = Inference() 44 | 45 | if event_type == INFERENCE_REQUEST_TYPE: 46 | inference.inputs = parse_kserve_request(event) 47 | elif event_type == INFERENCE_RESPONSE_TYPE: 48 | if "id" not in event: 49 | return 50 | 51 | inference.id = event["id"] 52 | inference.model_name = event.get("model_name") 53 | inference.model_version = event.get("model_version") 54 | inference.outputs = parse_kserve_response(event) 55 | inference.occurred_at = isoparse(faust.current_event().headers.get("ce_time", "").decode("utf-8")) 56 | else: 57 | return 58 | 59 | # Merge events. 60 | self._table[event_id] = inference.serialize() 61 | 62 | # Do we have all we need (request + response)? 63 | if ( 64 | inference.id is not None and 65 | inference.inputs is not None and 66 | inference.outputs is not None 67 | ): 68 | self._table.pop(event_id) 69 | return inference 70 | 71 | @staticmethod 72 | def get_group_key(event): 73 | return faust.current_event().headers.get("ce_id", "").decode("utf-8") 74 | -------------------------------------------------------------------------------- /src/inferencedb/main.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import json 3 | import os 4 | import logging 5 | import sys 6 | from signal import SIGTERM 7 | 8 | from inferencedb.config.factory import create_config_provider 9 | from inferencedb.config.providers.kubernetes_config_provider import KubernetesConfigProvider 10 | from inferencedb.utils.asyncio_utils import cancling_background_task, read_stream 11 | from inferencedb.settings import Settings 12 | from inferencedb.core.logging_utils import init_logging 13 | 14 | 15 | async def main(): 16 | settings = Settings() 17 | init_logging(settings.log_level) 18 | 19 | logging.info("InferenceDB started.") 20 | 21 | config_provider = create_config_provider( 22 | name=settings.config_provider, 23 | params=settings.config_provider_args, 24 | ) 25 | 26 | with cancling_background_task(config_provider.run()): 27 | # Wait for the first configuration. 28 | logging.debug("Waiting for a configuration update.") 29 | await config_provider.wait_for_update() 30 | 31 | while True: 32 | config = config_provider.get_config() 33 | 34 | # Create the worker process. 35 | logging.debug("Configuration update received - running worker.") 36 | process = await asyncio.create_subprocess_exec( 37 | "faust", "-A", "inferencedb.app", "worker", "-l", "info", 38 | stdout=asyncio.subprocess.PIPE, 39 | stderr=asyncio.subprocess.PIPE, 40 | env={ 41 | **os.environ, 42 | "CONFIG": json.dumps(config) 43 | }, 44 | ) 45 | 46 | # Wait for a configuration update and for the process to exit simultaneously. 47 | wait_for_process_task = asyncio.create_task(asyncio.wait([ 48 | read_stream(process.stdout, sys.stdout.buffer.write), 49 | read_stream(process.stderr, sys.stderr.buffer.write) 50 | ])) 51 | 52 | config_update_task = asyncio.create_task(config_provider.wait_for_update()) 53 | 54 | done, pending = await asyncio.wait( 55 | [wait_for_process_task, config_update_task], 56 | return_when=asyncio.FIRST_COMPLETED 57 | ) 58 | 59 | for task in pending: 60 | task.cancel() 61 | 62 | # If there is a configuration update, terminate the process and relaunch it. 63 | if config_update_task in done: 64 | # Warm shutdown, wait for tasks to complete. 65 | # https://faust.readthedocs.io/en/latest/userguide/workers.html 66 | process.send_signal(SIGTERM) 67 | await process.wait() 68 | 69 | 70 | if __name__ == "__main__": 71 | asyncio.run(main()) 72 | -------------------------------------------------------------------------------- /src/inferencedb/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aporia-ai/inferencedb/bee2cbe67f1e8ae3fa33948ac73f8656942f43cb/src/inferencedb/py.typed -------------------------------------------------------------------------------- /src/inferencedb/registry/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aporia-ai/inferencedb/bee2cbe67f1e8ae3fa33948ac73f8656942f43cb/src/inferencedb/registry/__init__.py -------------------------------------------------------------------------------- /src/inferencedb/registry/decorators.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, Union 2 | 3 | from .registry import get_registry, RegisteredObjectType 4 | 5 | 6 | def event_processor(name: str) -> Callable[[Callable], Callable]: 7 | """Decorator for registering event processors. 8 | 9 | Args: 10 | name: User friendly name for the event processor. 11 | 12 | Returns: 13 | Decorated class. 14 | """ 15 | 16 | def _decorator(callable: Callable) -> Callable: 17 | registry = get_registry() 18 | registry.register( 19 | object_type=RegisteredObjectType.EVENT_PROCESSOR, name=name, callable=callable 20 | ) 21 | return callable 22 | 23 | return _decorator 24 | 25 | 26 | def destination(name: str) -> Callable[[Callable], Callable]: 27 | """Decorator for registering destinations. 28 | 29 | Args: 30 | name: User friendly name for the destination. 31 | 32 | Returns: 33 | Decorated class. 34 | """ 35 | 36 | def _decorator(callable: Callable) -> Callable: 37 | registry = get_registry() 38 | registry.register(object_type=RegisteredObjectType.DESTINATION, name=name, callable=callable) 39 | return callable 40 | 41 | return _decorator 42 | 43 | 44 | def config_provider(name: str) -> Callable[[Callable], Callable]: 45 | """Decorator for registering config providers. 46 | 47 | Args: 48 | name: User friendly name for the config provider. 49 | 50 | Returns: 51 | Decorated class. 52 | """ 53 | 54 | def _decorator(callable: Callable) -> Callable: 55 | registry = get_registry() 56 | registry.register( 57 | object_type=RegisteredObjectType.CONFIG_PROVIDER, name=name, callable=callable 58 | ) 59 | return callable 60 | 61 | return _decorator 62 | 63 | 64 | def schema_provider(name: str) -> Callable[[Callable], Callable]: 65 | """Decorator for registering schema providers. 66 | 67 | Args: 68 | name: User friendly name for the schema provider. 69 | 70 | Returns: 71 | Decorated class. 72 | """ 73 | 74 | def _decorator(callable: Callable) -> Callable: 75 | registry = get_registry() 76 | registry.register( 77 | object_type=RegisteredObjectType.SCHEMA_PROVIDER, name=name, callable=callable 78 | ) 79 | return callable 80 | 81 | return _decorator 82 | -------------------------------------------------------------------------------- /src/inferencedb/registry/factory.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, Optional 2 | 3 | from .registry import get_registry, RegisteredObjectType 4 | 5 | 6 | def create_registered_object( 7 | object_type: RegisteredObjectType, name: str, params: Optional[Dict[str, Any]] = None 8 | ) -> Any: 9 | """Creates an object based on registry definitions. 10 | 11 | Args: 12 | object_type: Object type 13 | name: Object name 14 | config: Object config 15 | 16 | Returns: 17 | A new new object using the config. 18 | """ 19 | registry = get_registry() 20 | registered_object = registry.get(object_type=object_type, name=name) 21 | 22 | return registered_object.create_instance(params=params) 23 | -------------------------------------------------------------------------------- /src/inferencedb/registry/registered_object.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass 2 | from inspect import isclass, Signature, signature 3 | from typing import Any, Callable, Dict, Optional, Tuple, Type 4 | 5 | from pydantic.main import create_model 6 | 7 | from inferencedb.core.base_model import BaseModel 8 | 9 | 10 | class RegisteredObjectMetadata: 11 | """Base class for registered object metadata.""" 12 | 13 | pass 14 | 15 | 16 | @dataclass(frozen=True) 17 | class RegisteredObject: 18 | """Registered callable object.""" 19 | 20 | callable: Callable 21 | metadata: Optional[RegisteredObjectMetadata] = None 22 | 23 | @property 24 | def signature(self) -> Signature: 25 | """Returns the signature of the callable.""" 26 | return signature(self.callable) 27 | 28 | @property 29 | def parameters(self) -> Dict[str, Tuple[type, Any]]: 30 | """Returns the callable parameters as a name -> (type, default) mapping.""" 31 | params_dict = {} 32 | for param_name, param in self.signature.parameters.items(): 33 | param_default = ... 34 | if param.default is not param.empty: 35 | param_default = param.default 36 | 37 | params_dict[param_name] = (param.annotation, param_default) 38 | 39 | return params_dict 40 | 41 | @property 42 | def parameters_model(self) -> Type[BaseModel]: 43 | """Generates a pydantic model from the callable parameters.""" 44 | return create_model( 45 | f"{callable.__name__}_params", 46 | **self.parameters, # type: ignore 47 | __base__=BaseModel, 48 | ) 49 | 50 | def create_instance(self, params: Optional[dict] = None) -> Any: 51 | """Creates an instance of the registered object. 52 | 53 | Args: 54 | params: Constructor parameters. 55 | 56 | Returns: 57 | An instance of the registered object class. 58 | """ 59 | if not isclass(self.callable): 60 | raise TypeError(f"Cannot create instance of non-class object {self.callable}.") 61 | 62 | if params is None: 63 | params = {} 64 | 65 | validated_params = self.parameters_model(**params) # type: ignore 66 | 67 | return self.callable(**validated_params.dict()) 68 | -------------------------------------------------------------------------------- /src/inferencedb/registry/registry.py: -------------------------------------------------------------------------------- 1 | from enum import Enum 2 | from inspect import signature 3 | from typing import Callable, Dict, Optional 4 | 5 | from .registered_object import RegisteredObject, RegisteredObjectMetadata 6 | 7 | 8 | class RegisteredObjectType(Enum): 9 | """Registered object types.""" 10 | 11 | EVENT_PROCESSOR = "event_processor" 12 | DESTINATION = "destination" 13 | CONFIG_PROVIDER = "config_provider" 14 | SCHEMA_PROVIDER = "schema_provider" 15 | 16 | 17 | GENERIC_OBJECT_TYPES = [ 18 | RegisteredObjectType.EVENT_PROCESSOR, 19 | RegisteredObjectType.DESTINATION, 20 | RegisteredObjectType.CONFIG_PROVIDER, 21 | RegisteredObjectType.SCHEMA_PROVIDER, 22 | ] 23 | 24 | ObjectCollectionDict = Dict[RegisteredObjectType, Dict[str, RegisteredObject]] 25 | 26 | 27 | class Registry: 28 | """Generic callable obect registry. 29 | 30 | Callable code objects (classes and functions) can be registered to extend 31 | the functionality of aporia core. 32 | 33 | The registry supports the following objects: 34 | - sources 35 | - destinations 36 | - config providers 37 | - schema providers 38 | 39 | To register an object, use the appropriate decorator (e.g @source for sources). 40 | """ 41 | 42 | def __init__(self): 43 | """Initializes the Registry.""" 44 | self._generic_objects: ObjectCollectionDict = {} 45 | 46 | for object_type in GENERIC_OBJECT_TYPES: 47 | self._generic_objects[object_type] = {} 48 | 49 | def register( 50 | self, 51 | object_type: RegisteredObjectType, 52 | name: str, 53 | callable: Callable, 54 | metadata: Optional[RegisteredObjectMetadata] = None, 55 | ): 56 | """Registers an object. 57 | 58 | Args: 59 | object_type: Object type 60 | name: Object name 61 | callable: Class or function callable object 62 | metadata: Any other metadata to store with the object 63 | """ 64 | object_collection = self._generic_objects[object_type] 65 | 66 | if name in object_collection: 67 | raise ValueError(f"{object_type.value} {name} is already registered.") 68 | 69 | for param_name, param in signature(callable).parameters.items(): 70 | if param.annotation is param.empty: 71 | raise ValueError( 72 | f"Callable {callable} parameter {param_name} is missing a type annotation." 73 | ) 74 | 75 | object_collection[name] = RegisteredObject(callable=callable, metadata=metadata) 76 | 77 | def get(self, object_type: RegisteredObjectType, name: str) -> RegisteredObject: 78 | """Returns a registered object. 79 | 80 | Args: 81 | object_type: Object type. 82 | name: Object name. 83 | 84 | Returns: 85 | The registered object. 86 | """ 87 | object_collection = self._generic_objects[object_type] 88 | 89 | if name not in object_collection: 90 | raise KeyError(f"{object_type.value} {name} is not registered.") 91 | 92 | return object_collection[name] 93 | 94 | 95 | registry = Registry() 96 | 97 | 98 | def get_registry() -> Registry: 99 | """Returns the global registry.""" 100 | return registry 101 | -------------------------------------------------------------------------------- /src/inferencedb/schema_providers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aporia-ai/inferencedb/bee2cbe67f1e8ae3fa33948ac73f8656942f43cb/src/inferencedb/schema_providers/__init__.py -------------------------------------------------------------------------------- /src/inferencedb/schema_providers/avro_schema_provider.py: -------------------------------------------------------------------------------- 1 | from typing import Any, AsyncIterator, Dict, List 2 | 3 | import pandas as pd 4 | from pandavro import schema_infer 5 | from schema_registry.client import AsyncSchemaRegistryClient 6 | from schema_registry.serializers import AsyncAvroMessageSerializer 7 | from schema_registry.client.schema import AvroSchema 8 | 9 | from inferencedb.registry.decorators import schema_provider 10 | from inferencedb.schema_providers.schema_provider import SchemaProvider 11 | from inferencedb.core.inference import Inference 12 | 13 | AVRO_NAMESPACE = "com.aporia.inferencedb.v1alpha1" 14 | 15 | 16 | @schema_provider("avro") 17 | class AvroSchemaProvider(SchemaProvider): 18 | def __init__(self, 19 | logger_name: str, 20 | schema_registry: AsyncSchemaRegistryClient, 21 | subject: str, 22 | config: Dict[str, Any] 23 | ): 24 | self._logger_name = logger_name 25 | self._schema_registry = schema_registry 26 | self._subject = subject 27 | self._config = config 28 | self._serializer = AsyncAvroMessageSerializer(self._schema_registry) 29 | self._schema: AvroSchema = None 30 | self._is_columnar = config.get("columnar", True) 31 | 32 | self._input_column_names = None 33 | self._output_column_names = None 34 | 35 | if "columnNames" in config: 36 | self._input_column_names = config["columnNames"].get("inputs") 37 | self._output_column_names = config["columnNames"].get("outputs") 38 | 39 | async def fetch(self): 40 | # If a schema is already registered in the Schema Registry, use it. 41 | response = await self._schema_registry.get_schema(self._subject) 42 | if response is None: 43 | return 44 | 45 | self._schema = response.schema 46 | 47 | async def serialize(self, inference: Inference) -> AsyncIterator[bytes]: 48 | # Override input columns 49 | if self._input_column_names is not None: 50 | inference.inputs.columns = self._input_column_names 51 | elif ( 52 | isinstance(inference.inputs.columns, pd.RangeIndex) or 53 | list(inference.inputs.columns) == [str(i) for i in range(len(inference.inputs.columns))] 54 | ): 55 | inference.inputs.columns = [f"X{i}" for i in range(len(inference.inputs.columns))] 56 | 57 | # Override output columns 58 | if self._output_column_names is not None: 59 | inference.outputs.columns = self._output_column_names 60 | elif ( 61 | isinstance(inference.outputs.columns, pd.RangeIndex) or 62 | list(inference.outputs.columns) == [str(i) for i in range(len(inference.outputs.columns))] 63 | ): 64 | inference.outputs.columns = [f"Y{i}" for i in range(len(inference.outputs.columns))] 65 | 66 | if self._schema is None: 67 | await self._generate_schema_from_inference(inference) 68 | 69 | # TODO: Make sure the shape of every input & output is the same 70 | for i, ((_, inputs), (_, outputs)) in enumerate(zip(inference.inputs.iterrows(), inference.outputs.iterrows())): 71 | yield await self._serializer.encode_record_with_schema( 72 | subject=self._logger_name, 73 | schema=self._schema, 74 | record={ 75 | "id": f"{inference.id}_{i}", 76 | "occurred_at": inference.occurred_at.isoformat(), 77 | **inputs.to_dict(), 78 | **outputs.to_dict() 79 | }, 80 | ) 81 | 82 | async def _generate_schema_from_inference(self, inference: Inference): 83 | self._schema = AvroSchema({ 84 | "type": "record", 85 | "namespace": AVRO_NAMESPACE, 86 | "name": self._logger_name.replace("-", "_"), 87 | "fields": [ 88 | {"name": "id", "type": "string"}, 89 | {"name": "occurred_at", "type": "string"}, 90 | *schema_infer(inference.inputs)["fields"], 91 | *schema_infer(inference.outputs)["fields"], 92 | ], 93 | }) 94 | 95 | self._input_column_names = inference.inputs.columns 96 | self._output_column_names = inference.outputs.columns 97 | 98 | await self._schema_registry.register(self._subject, self._schema) 99 | -------------------------------------------------------------------------------- /src/inferencedb/schema_providers/factory.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, Optional 2 | 3 | from inferencedb.registry.factory import create_registered_object 4 | from inferencedb.registry.registry import RegisteredObjectType 5 | from .schema_provider import SchemaProvider 6 | 7 | 8 | def create_schema_provider(type: str, params: Optional[Dict[str, Any]] = None) -> SchemaProvider: 9 | """Creates a SchemaProvider object. 10 | 11 | Args: 12 | type: SchemaProvider type 13 | params: SchemaProvider parameters 14 | 15 | Returns: 16 | SchemaProvider object 17 | """ 18 | return create_registered_object( 19 | object_type=RegisteredObjectType.SCHEMA_PROVIDER, name=type, params=params 20 | ) 21 | -------------------------------------------------------------------------------- /src/inferencedb/schema_providers/schema_provider.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from typing import AsyncIterator 3 | 4 | from inferencedb.core.inference import Inference 5 | 6 | 7 | class SchemaProvider(ABC): 8 | """Abstract base class for SchemaProvider objects.""" 9 | 10 | @abstractmethod 11 | async def fetch(self): 12 | ... 13 | 14 | @abstractmethod 15 | async def serialize(self, inference: Inference) -> AsyncIterator[bytes]: 16 | ... 17 | -------------------------------------------------------------------------------- /src/inferencedb/settings.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict 2 | 3 | from pydantic import BaseSettings 4 | 5 | 6 | class Settings(BaseSettings): 7 | """Execution settings.""" 8 | 9 | kafka_broker: str 10 | kafka_schema_registry_url: str 11 | kafka_connect_url: str 12 | config_provider_args: Dict[str, Any] 13 | config_provider: str = "kubernetes" 14 | log_level: str = "WARNING" 15 | enable_pretty_logs: bool = False 16 | enable_dev_mode: bool = False 17 | 18 | class Config: 19 | """Execution settings config.""" 20 | 21 | env_prefix = "INFERENCEDB_" 22 | -------------------------------------------------------------------------------- /src/inferencedb/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aporia-ai/inferencedb/bee2cbe67f1e8ae3fa33948ac73f8656942f43cb/src/inferencedb/utils/__init__.py -------------------------------------------------------------------------------- /src/inferencedb/utils/asyncio_utils.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | from contextlib import contextmanager 3 | from typing import Coroutine, Dict, Generator, Optional, Union 4 | 5 | 6 | @contextmanager 7 | def cancling_background_task(coro: Coroutine) -> Generator[asyncio.Task, None, None]: 8 | """Schedules a coroutine to run in the background, canceling it when the context manager exits. 9 | 10 | Args: 11 | coro: Coroutine to schedule 12 | 13 | Yields: 14 | The scheduled task object. 15 | """ 16 | task = asyncio.create_task(coro) 17 | yield task 18 | task.cancel() 19 | 20 | 21 | async def read_stream(stream, cb): 22 | while True: 23 | line = await stream.readline() 24 | if line: 25 | cb(line) 26 | else: 27 | break 28 | -------------------------------------------------------------------------------- /src/inferencedb/utils/kserve_utils.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | import pandas as pd 4 | 5 | 6 | def parse_kserve_v2_tensor(columns: List[dict]) -> pd.DataFrame: 7 | # If there's one column with a 2D shape, then we'll assume the first dimension 8 | # is the amount of rows, and the second dimension is the amount of columns. 9 | if len(columns) == 1 and len(columns[0]["shape"]) == 2: 10 | return pd.DataFrame(data=columns[0]["data"]) 11 | 12 | return pd.DataFrame.from_dict({ 13 | col["name"]: col["data"] 14 | for col in columns 15 | }) 16 | 17 | 18 | 19 | def parse_kserve_request(request: dict) -> pd.DataFrame: 20 | if "instances" in request: 21 | return pd.DataFrame(data=request["instances"]) 22 | elif "inputs" in request: 23 | return parse_kserve_v2_tensor(request["inputs"]) 24 | else: 25 | raise ValueError(f"Invalid KServe request: no 'instances' or 'inputs' fields") 26 | 27 | 28 | 29 | def parse_kserve_response(response: dict) -> pd.DataFrame: 30 | if "predictions" in response: 31 | return pd.DataFrame(data=response["predictions"]) 32 | elif "outputs" in response: 33 | return parse_kserve_v2_tensor(response["outputs"]) 34 | else: 35 | raise ValueError(f"Invalid KServe response: no 'predictions' or 'outputs' fields") 36 | -------------------------------------------------------------------------------- /src/inferencedb/utils/pandas_utils.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Optional 2 | 3 | from faust.serializers import Codec 4 | import pandas as pd 5 | import pyarrow as pa 6 | from pyarrow.feather import read_feather, write_feather 7 | 8 | 9 | def serialize_dataframe(df: pd.DataFrame) -> Optional[bytes]: 10 | if df is None: 11 | return None 12 | 13 | with pa.BufferOutputStream() as output_stream: 14 | write_feather(df, output_stream) 15 | return output_stream.getvalue().to_pybytes() 16 | 17 | 18 | def deserialize_dataframe(data: Optional[bytes]) -> pd.DataFrame: 19 | if data is None: 20 | return None 21 | 22 | with pa.BufferReader(data) as reader: 23 | return read_feather(reader) 24 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aporia-ai/inferencedb/bee2cbe67f1e8ae3fa33948ac73f8656942f43cb/tests/__init__.py -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aporia-ai/inferencedb/bee2cbe67f1e8ae3fa33948ac73f8656942f43cb/tests/conftest.py -------------------------------------------------------------------------------- /tests/test_kserve_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pandas as pd 3 | 4 | from inferencedb.utils.kserve_utils import parse_kserve_request, parse_kserve_response, parse_kserve_v2_tensor 5 | 6 | 7 | def test_pd_single_col_single_datapoint(): 8 | result = parse_kserve_v2_tensor([ 9 | { 10 | "name": "col1", 11 | "shape": [1], 12 | "datatype": "FP32", 13 | "data": [7.4], 14 | }, 15 | ]) 16 | 17 | expected = pd.DataFrame(data=[[7.4]], columns=["col1"]) 18 | 19 | assert np.array_equal(result, expected) 20 | 21 | 22 | 23 | def test_pd_single_col_multiple_datapoint(): 24 | result = parse_kserve_v2_tensor([ 25 | { 26 | "name": "col1", 27 | "shape": [3], 28 | "datatype": "FP32", 29 | "data": [7.4, 3.5, 7.8], 30 | }, 31 | ]) 32 | 33 | expected = pd.DataFrame(data=[ 34 | [7.4], 35 | [3.5], 36 | [7.8], 37 | ], columns=["col1"]) 38 | 39 | assert np.array_equal(result, expected) 40 | 41 | 42 | 43 | def test_pd_multiple_col_single_datapoint(): 44 | result = parse_kserve_v2_tensor([ 45 | { 46 | "name": "col1", 47 | "shape": [1], 48 | "datatype": "FP32", 49 | "data": [3], 50 | }, 51 | { 52 | "name": "col2", 53 | "shape": [1], 54 | "datatype": "FP32", 55 | "data": [4], 56 | }, 57 | ]) 58 | 59 | expected = pd.DataFrame(data=[ 60 | [3, 4], 61 | ], columns=["col1", "col2"]) 62 | 63 | assert np.array_equal(result, expected) 64 | 65 | 66 | def test_pd_multiple_col_multiple_datapoints(): 67 | result = parse_kserve_v2_tensor([ 68 | { 69 | "name": "col1", 70 | "shape": [4], 71 | "datatype": "FP32", 72 | "data": [1, 2, 3, 4], 73 | }, 74 | { 75 | "name": "col2", 76 | "shape": [4], 77 | "datatype": "FP32", 78 | "data": [5, 6, 7, 8], 79 | }, 80 | { 81 | "name": "col3", 82 | "shape": [4], 83 | "datatype": "FP32", 84 | "data": [0, -1, -2, -3], 85 | }, 86 | ]) 87 | 88 | expected = pd.DataFrame(data=[ 89 | [1, 5, 0], 90 | [2, 6, -1], 91 | [3, 7, -2], 92 | [4, 8, -3], 93 | ], columns=["col1", "col2", "col3"]) 94 | 95 | assert np.array_equal(result, expected) 96 | 97 | 98 | def test_np_multiple_col_multiple_datapoints(): 99 | result = parse_kserve_v2_tensor([ 100 | { 101 | "name": "input-0", 102 | "shape": [2, 4], 103 | "datatype": "FP32", 104 | "data": [ 105 | [6.8, 2.8, 4.8, 1.4], 106 | [6.0, 3.4, 6.2, 1.8] 107 | ] 108 | } 109 | ]) 110 | 111 | expected = pd.DataFrame(data=[ 112 | [6.8, 2.8, 4.8, 1.4], 113 | [6.0, 3.4, 6.2, 1.8] 114 | ]) 115 | 116 | assert np.array_equal(result, expected) 117 | 118 | 119 | def test_pd_multiple_dtypes(): 120 | result = parse_kserve_v2_tensor([ 121 | { 122 | "name": "bool_column", 123 | "shape": [3], 124 | "datatype": "BOOL", 125 | "data": [True, True, False], 126 | }, 127 | { 128 | "name": "float_column", 129 | "shape": [3], 130 | "datatype": "FP32", 131 | "data": [0.5, -4.6, 103.4], 132 | }, 133 | { 134 | "name": "int_column", 135 | "shape": [3], 136 | "datatype": "INT32", 137 | "data": [32, 64, 128], 138 | }, 139 | { 140 | "name": "string_column", 141 | "shape": [3], 142 | "datatype": "BYTES", 143 | "data": ["hello", "world", "!"], 144 | }, 145 | ]) 146 | 147 | expected = pd.DataFrame(data=[ 148 | [True, 0.5, 32, "hello"], 149 | [True, -4.6, 64, "world"], 150 | [False, 103.4, 128, "!"], 151 | ], columns=["bool_column", "float_column", "int_column", "string_column"]) 152 | 153 | assert result.equals(expected) 154 | assert result["bool_column"].dtype == "bool" 155 | assert result["float_column"].dtype == "float64" 156 | assert result["int_column"].dtype == "int64" 157 | assert result["string_column"].dtype == "object" 158 | 159 | 160 | def test_text_and_embedding(): 161 | result = parse_kserve_v2_tensor([ 162 | { 163 | "name": "text", 164 | "shape": [2], 165 | "datatype": "BYTES", 166 | "data": ["Hello", "World"] 167 | }, 168 | { 169 | "name": "embedding", 170 | "shape": [2, 4], 171 | "datatype": "FP32", 172 | "data": [ 173 | [1, 2, 3, 4], 174 | [4, 5, 6, 7], 175 | ] 176 | } 177 | ]) 178 | 179 | expected = pd.DataFrame(data=[ 180 | ["Hello", [1, 2, 3, 4]], 181 | ["World", [4, 5, 6, 7]] 182 | ]) 183 | 184 | assert np.array_equal(result, expected) 185 | 186 | 187 | def test_kserve_v1_request(): 188 | result = parse_kserve_request({ 189 | "instances": [ 190 | [1, 2, 3, 4], 191 | [4, 5, 6, 7], 192 | ] 193 | }) 194 | 195 | expected = pd.DataFrame(data=[ 196 | [1, 2, 3, 4], 197 | [4, 5, 6, 7], 198 | ]) 199 | 200 | assert np.array_equal(result, expected) 201 | 202 | 203 | def test_kserve_v1_response(): 204 | result = parse_kserve_response({"predictions": [1, 1]}) 205 | 206 | expected = pd.DataFrame(data=[ 207 | [1], 208 | [1], 209 | ]) 210 | 211 | assert np.array_equal(result, expected) 212 | 213 | 214 | def test_kserve_v2_request(): 215 | result = parse_kserve_request({ 216 | "inputs": [ 217 | { 218 | "name": "text", 219 | "shape": [2], 220 | "datatype": "BYTES", 221 | "data": ["Hello", "World"] 222 | }, 223 | { 224 | "name": "embedding", 225 | "shape": [2, 4], 226 | "datatype": "FP32", 227 | "data": [ 228 | [1, 2, 3, 4], 229 | [4, 5, 6, 7], 230 | ] 231 | } 232 | ] 233 | }) 234 | 235 | expected = pd.DataFrame(data=[ 236 | ["Hello", [1, 2, 3, 4]], 237 | ["World", [4, 5, 6, 7]] 238 | ]) 239 | 240 | assert np.array_equal(result, expected) 241 | 242 | 243 | def test_kserve_v2_response(): 244 | result = parse_kserve_response({ 245 | "outputs": [ 246 | { 247 | "name": "predict", 248 | "shape": [3], 249 | "datatype": "FP32", 250 | "data": [2, 3, 4] 251 | } 252 | ] 253 | }) 254 | 255 | expected = pd.DataFrame(data=[ 256 | [2], [3], [4] 257 | ], columns=["predict"]) 258 | 259 | assert np.array_equal(result, expected) 260 | --------------------------------------------------------------------------------