├── .python-version ├── examples ├── test_model │ ├── examples.yaml │ ├── model │ │ ├── __init__.py │ │ └── model.py │ ├── data │ │ └── model │ │ │ └── model.joblib │ └── config.yaml ├── xgboost │ ├── examples.yaml │ ├── model │ │ ├── __init__.py │ │ └── model.py │ ├── data │ │ └── model │ │ │ └── model.ubj │ └── config.yaml ├── bert_uncased │ ├── examples.yaml │ ├── model │ │ ├── __init__.py │ │ └── model.py │ └── config.yaml ├── mlflow_sklearn │ ├── examples.yaml │ ├── model │ │ ├── __init__.py │ │ └── model.py │ ├── data │ │ └── model │ │ │ └── model │ │ │ ├── model.pkl │ │ │ ├── python_env.yaml │ │ │ ├── requirements.txt │ │ │ ├── conda.yaml │ │ │ └── MLmodel │ └── config.yaml ├── test_model_1 │ ├── examples.yaml │ ├── model │ │ ├── __init__.py │ │ └── model.py │ ├── data │ │ └── model │ │ │ └── model.joblib │ └── config.yaml └── mlruns │ └── 0 │ ├── 3a98a956c3114efcbceb6cce54407227 │ ├── tags │ │ ├── mlflow.user │ │ ├── mlflow.source.type │ │ ├── mlflow.runName │ │ ├── mlflow.source.name │ │ └── mlflow.log-model.history │ ├── artifacts │ │ └── model │ │ │ ├── model.pkl │ │ │ ├── python_env.yaml │ │ │ ├── requirements.txt │ │ │ ├── conda.yaml │ │ │ └── MLmodel │ └── meta.yaml │ ├── a33bc3b0ca6e4ab6990673faa2c44aa0 │ ├── tags │ │ ├── mlflow.user │ │ ├── mlflow.source.type │ │ ├── mlflow.runName │ │ ├── mlflow.source.name │ │ └── mlflow.log-model.history │ ├── artifacts │ │ └── model │ │ │ ├── model.pkl │ │ │ ├── python_env.yaml │ │ │ ├── requirements.txt │ │ │ ├── conda.yaml │ │ │ └── MLmodel │ └── meta.yaml │ ├── ae2db1d887d84e73a40da9ad90a4c516 │ ├── tags │ │ ├── mlflow.user │ │ ├── mlflow.source.type │ │ ├── mlflow.runName │ │ ├── mlflow.source.name │ │ └── mlflow.log-model.history │ ├── artifacts │ │ └── model │ │ │ ├── model.pkl │ │ │ ├── python_env.yaml │ │ │ ├── requirements.txt │ │ │ ├── conda.yaml │ │ │ └── MLmodel │ └── meta.yaml │ └── meta.yaml ├── model_to_docker ├── __init__.py ├── truss_saver │ ├── .dockerignore │ ├── truss │ │ ├── tests │ │ │ ├── __init__.py │ │ │ ├── local │ │ │ │ ├── __init__.py │ │ │ │ └── test_local_config_handler.py │ │ │ ├── patch │ │ │ │ ├── test_types.py │ │ │ │ ├── test_signature.py │ │ │ │ └── test_dir_signature.py │ │ │ ├── templates │ │ │ │ ├── control │ │ │ │ │ └── control │ │ │ │ │ │ └── helpers │ │ │ │ │ │ ├── test_context_managers.py │ │ │ │ │ │ └── test_patch_applier.py │ │ │ │ ├── core │ │ │ │ │ └── server │ │ │ │ │ │ ├── common │ │ │ │ │ │ └── test_util.py │ │ │ │ │ │ └── test_secrets_resolver.py │ │ │ │ └── server │ │ │ │ │ ├── common │ │ │ │ │ └── test_retry.py │ │ │ │ │ └── test_model_wrapper.py │ │ │ ├── test_truss_util.py │ │ │ ├── test_docker.py │ │ │ ├── test_notebooks.py │ │ │ ├── contexts │ │ │ │ ├── local_loader │ │ │ │ │ └── test_load_local.py │ │ │ │ └── image_builder │ │ │ │ │ └── test_serving_image_builder.py │ │ │ ├── environments_inference │ │ │ │ └── test_requirements_inference.py │ │ │ ├── test_truss_gatherer.py │ │ │ ├── test_context_builder_image.py │ │ │ ├── test_testing_utilities_for_other_tests.py │ │ │ ├── test_backward.py │ │ │ ├── test_validation.py │ │ │ └── test_config.py │ │ ├── templates │ │ │ ├── __init__.py │ │ │ ├── server │ │ │ │ ├── __init__.py │ │ │ │ ├── common │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── retry.py │ │ │ │ │ ├── logging.py │ │ │ │ │ ├── patches.py │ │ │ │ │ ├── util.py │ │ │ │ │ ├── patches │ │ │ │ │ │ └── whisper │ │ │ │ │ │ │ └── patch.py │ │ │ │ │ ├── errors.py │ │ │ │ │ └── serialization.py │ │ │ │ ├── requirements.txt │ │ │ │ └── inference_server.py │ │ │ ├── shared │ │ │ │ ├── __init__.py │ │ │ │ ├── README.md │ │ │ │ └── secrets_resolver.py │ │ │ ├── custom │ │ │ │ ├── model │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── model.py │ │ │ │ ├── examples.yaml │ │ │ │ └── train │ │ │ │ │ └── train.py │ │ │ ├── keras │ │ │ │ └── model │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── model.py │ │ │ ├── mlflow │ │ │ │ └── model │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── model.py │ │ │ ├── lightgbm │ │ │ │ └── model │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── model.py │ │ │ ├── pipeline │ │ │ │ └── model │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── model.py │ │ │ ├── pytorch │ │ │ │ └── model │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── model.py │ │ │ ├── sklearn │ │ │ │ └── model │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── model.py │ │ │ ├── xgboost │ │ │ │ └── model │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── model.py │ │ │ ├── huggingface_transformer │ │ │ │ ├── model │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── model.py │ │ │ │ └── requirements.txt │ │ │ ├── training │ │ │ │ ├── requirements.txt │ │ │ │ └── job.py │ │ │ ├── control │ │ │ │ ├── requirements.txt │ │ │ │ └── control │ │ │ │ │ ├── helpers │ │ │ │ │ ├── context_managers.py │ │ │ │ │ ├── errors.py │ │ │ │ │ ├── inference_server_starter.py │ │ │ │ │ ├── inference_server_process_controller.py │ │ │ │ │ └── types.py │ │ │ │ │ ├── server.py │ │ │ │ │ └── application.py │ │ │ ├── training.Dockerfile.jinja │ │ │ ├── base.Dockerfile.jinja │ │ │ ├── README.md.jinja │ │ │ └── server.Dockerfile.jinja │ │ ├── pytest.ini │ │ ├── util │ │ │ ├── data_structures.py │ │ │ ├── jinja.py │ │ │ ├── gpu.py │ │ │ └── path.py │ │ ├── blob │ │ │ ├── blob_backend.py │ │ │ ├── http_public_blob_backend.py │ │ │ └── blob_backend_registry.py │ │ ├── __init__.py │ │ ├── decorators.py │ │ ├── contexts │ │ │ ├── truss_context.py │ │ │ ├── local_loader │ │ │ │ ├── utils.py │ │ │ │ ├── load_model_local.py │ │ │ │ ├── docker_build_emulator.py │ │ │ │ └── train_local.py │ │ │ └── image_builder │ │ │ │ ├── image_builder.py │ │ │ │ └── util.py │ │ ├── patch │ │ │ ├── signature.py │ │ │ ├── dir_signature.py │ │ │ ├── types.py │ │ │ └── hash.py │ │ ├── notebook.py │ │ ├── readme_generator.py │ │ ├── errors.py │ │ ├── local │ │ │ └── local_config.py │ │ ├── model_frameworks │ │ │ ├── keras.py │ │ │ ├── xgboost.py │ │ │ ├── sklearn.py │ │ │ ├── lightgbm.py │ │ │ ├── __init__.py │ │ │ ├── huggingface_transformer.py │ │ │ ├── mlflow.py │ │ │ └── pytorch.py │ │ ├── types.py │ │ ├── model_framework.py │ │ ├── validation.py │ │ ├── truss_gatherer.py │ │ ├── constants.py │ │ ├── docker.py │ │ └── environment_inference │ │ │ └── requirements_inference.py │ ├── .python-version │ ├── .tool-versions │ ├── .gitbook.yaml │ ├── __init__.py │ ├── bin │ │ └── codespace_post_create.sh │ ├── .gitattributes │ ├── Dockerfile │ ├── .flake8 │ ├── docs.py │ ├── context_builder.Dockerfile │ ├── .gitignore │ ├── .pre-commit-config.yaml │ ├── docker │ │ └── base_images │ │ │ └── base_image.Dockerfile.jinja │ └── pyproject.toml └── api.py ├── requires-build.txt ├── requires-ci.txt ├── .ipynb_checkpoints └── truss_basics-checkpoint.ipynb ├── docs └── CHANGELOG.md ├── CONTRIBUTING.md ├── .gitignore ├── setup.py ├── README.md └── requires-install.txt /.python-version: -------------------------------------------------------------------------------- 1 | 3.9 2 | -------------------------------------------------------------------------------- /examples/test_model/examples.yaml: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/xgboost/examples.yaml: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/bert_uncased/examples.yaml: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/mlflow_sklearn/examples.yaml: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/test_model/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/test_model_1/examples.yaml: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/xgboost/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/bert_uncased/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/mlflow_sklearn/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/test_model_1/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model_to_docker/__init__.py: -------------------------------------------------------------------------------- 1 | from .api import * -------------------------------------------------------------------------------- /model_to_docker/truss_saver/.dockerignore: -------------------------------------------------------------------------------- 1 | **/.git 2 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/.python-version: -------------------------------------------------------------------------------- 1 | 3.9.11 2 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/local/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /requires-build.txt: -------------------------------------------------------------------------------- 1 | build==0.10.0 2 | twine==4.0.2 3 | 4 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/server/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/shared/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/custom/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/keras/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/mlflow/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /requires-ci.txt: -------------------------------------------------------------------------------- 1 | pylint==2.17.2 2 | flake8==6.0.0 3 | black==23.3.0 -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/lightgbm/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/pipeline/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/pytorch/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/sklearn/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/xgboost/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/mlruns/0/3a98a956c3114efcbceb6cce54407227/tags/mlflow.user: -------------------------------------------------------------------------------- 1 | faizank -------------------------------------------------------------------------------- /examples/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/tags/mlflow.user: -------------------------------------------------------------------------------- 1 | faizank -------------------------------------------------------------------------------- /examples/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/tags/mlflow.user: -------------------------------------------------------------------------------- 1 | faizank -------------------------------------------------------------------------------- /model_to_docker/truss_saver/.tool-versions: -------------------------------------------------------------------------------- 1 | python 3.9.11 2 | poetry 1.4.2 3 | -------------------------------------------------------------------------------- /examples/mlruns/0/3a98a956c3114efcbceb6cce54407227/tags/mlflow.source.type: -------------------------------------------------------------------------------- 1 | LOCAL -------------------------------------------------------------------------------- /examples/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/tags/mlflow.source.type: -------------------------------------------------------------------------------- 1 | LOCAL -------------------------------------------------------------------------------- /examples/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/tags/mlflow.source.type: -------------------------------------------------------------------------------- 1 | LOCAL -------------------------------------------------------------------------------- /examples/mlruns/0/3a98a956c3114efcbceb6cce54407227/tags/mlflow.runName: -------------------------------------------------------------------------------- 1 | masked-gnu-148 -------------------------------------------------------------------------------- /examples/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/tags/mlflow.runName: -------------------------------------------------------------------------------- 1 | salty-steed-823 -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/huggingface_transformer/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/training/requirements.txt: -------------------------------------------------------------------------------- 1 | PyYAML==6.0 2 | -------------------------------------------------------------------------------- /examples/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/tags/mlflow.runName: -------------------------------------------------------------------------------- 1 | thoughtful-foal-616 -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/custom/examples.yaml: -------------------------------------------------------------------------------- 1 | - input: 2 | inputs: 3 | - - 0 4 | name: example1 5 | -------------------------------------------------------------------------------- /examples/xgboost/data/model/model.ubj: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/slashml/slash_docker/HEAD/examples/xgboost/data/model/model.ubj -------------------------------------------------------------------------------- /model_to_docker/truss_saver/.gitbook.yaml: -------------------------------------------------------------------------------- 1 | root: ./docs 2 | 3 | ​structure: 4 | readme: ../README.md 5 | summary: SUMMARY.md​ 6 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/__init__.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | 4 | dirname = os.path.dirname(__file__) 5 | sys.path.append(dirname) -------------------------------------------------------------------------------- /examples/test_model/data/model/model.joblib: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/slashml/slash_docker/HEAD/examples/test_model/data/model/model.joblib -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/huggingface_transformer/requirements.txt: -------------------------------------------------------------------------------- 1 | -i https://pypi.org/simple 2 | 3 | transformers 4 | torch -------------------------------------------------------------------------------- /.ipynb_checkpoints/truss_basics-checkpoint.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [], 3 | "metadata": {}, 4 | "nbformat": 4, 5 | "nbformat_minor": 5 6 | } 7 | -------------------------------------------------------------------------------- /examples/test_model_1/data/model/model.joblib: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/slashml/slash_docker/HEAD/examples/test_model_1/data/model/model.joblib -------------------------------------------------------------------------------- /model_to_docker/truss_saver/bin/codespace_post_create.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | poetry install 3 | poetry run pre-commit install 4 | git lfs install 5 | -------------------------------------------------------------------------------- /examples/mlflow_sklearn/data/model/model/model.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/slashml/slash_docker/HEAD/examples/mlflow_sklearn/data/model/model/model.pkl -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/server/common/__init__.py: -------------------------------------------------------------------------------- 1 | """This module is designated to exist inside of Truss environments and containers""" 2 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/.gitattributes: -------------------------------------------------------------------------------- 1 | *.joblib filter=lfs diff=lfs merge=lfs -text 2 | examples/gpu-sharing/examples.yaml filter=lfs diff=lfs merge=lfs -text 3 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/control/requirements.txt: -------------------------------------------------------------------------------- 1 | dataclasses-json==0.5.7 2 | Flask==2.0.3 3 | waitress==2.1.2 4 | truss==0.4.8rc7 5 | tenacity==8.1.0 6 | -------------------------------------------------------------------------------- /examples/mlruns/0/3a98a956c3114efcbceb6cce54407227/tags/mlflow.source.name: -------------------------------------------------------------------------------- 1 | /Users/faizank/workspace/experiments/model_deployment_experiments/venv/lib/python3.10/site-packages/ipykernel_launcher.py -------------------------------------------------------------------------------- /examples/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/tags/mlflow.source.name: -------------------------------------------------------------------------------- 1 | /Users/faizank/workspace/experiments/model_deployment_experiments/venv/lib/python3.10/site-packages/ipykernel_launcher.py -------------------------------------------------------------------------------- /examples/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/tags/mlflow.source.name: -------------------------------------------------------------------------------- 1 | /Users/faizank/workspace/experiments/model_deployment_experiments/venv/lib/python3.10/site-packages/ipykernel_launcher.py -------------------------------------------------------------------------------- /examples/mlflow_sklearn/data/model/model/python_env.yaml: -------------------------------------------------------------------------------- 1 | python: 3.10.7 2 | build_dependencies: 3 | - pip==23.3.1 4 | - setuptools==63.2.0 5 | - wheel==0.41.2 6 | dependencies: 7 | - -r requirements.txt 8 | -------------------------------------------------------------------------------- /examples/mlruns/0/3a98a956c3114efcbceb6cce54407227/artifacts/model/model.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/slashml/slash_docker/HEAD/examples/mlruns/0/3a98a956c3114efcbceb6cce54407227/artifacts/model/model.pkl -------------------------------------------------------------------------------- /examples/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/artifacts/model/model.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/slashml/slash_docker/HEAD/examples/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/artifacts/model/model.pkl -------------------------------------------------------------------------------- /examples/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/artifacts/model/model.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/slashml/slash_docker/HEAD/examples/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/artifacts/model/model.pkl -------------------------------------------------------------------------------- /examples/mlflow_sklearn/data/model/model/requirements.txt: -------------------------------------------------------------------------------- 1 | mlflow==2.7.1 2 | cloudpickle==2.2.1 3 | numpy==1.23.5 4 | packaging==20.9 5 | psutil==5.9.6 6 | pyyaml==6.0.1 7 | scikit-learn==1.3.1 8 | scipy==1.11.3 -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/shared/README.md: -------------------------------------------------------------------------------- 1 | # Shared code between training and serving images 2 | 3 | Code in this directory is common to both training and serving and is copied into them. 4 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM python:3.9 2 | 3 | RUN curl -sSL https://install.python-poetry.org | python - 4 | 5 | ENV PATH="/root/.local/bin:${PATH}" 6 | COPY . . 7 | 8 | RUN poetry install --only main 9 | -------------------------------------------------------------------------------- /examples/mlruns/0/3a98a956c3114efcbceb6cce54407227/artifacts/model/python_env.yaml: -------------------------------------------------------------------------------- 1 | python: 3.10.7 2 | build_dependencies: 3 | - pip==23.3.1 4 | - setuptools==63.2.0 5 | - wheel==0.41.2 6 | dependencies: 7 | - -r requirements.txt 8 | -------------------------------------------------------------------------------- /examples/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/artifacts/model/python_env.yaml: -------------------------------------------------------------------------------- 1 | python: 3.10.7 2 | build_dependencies: 3 | - pip==23.3.1 4 | - setuptools==63.2.0 5 | - wheel==0.41.2 6 | dependencies: 7 | - -r requirements.txt 8 | -------------------------------------------------------------------------------- /examples/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/artifacts/model/python_env.yaml: -------------------------------------------------------------------------------- 1 | python: 3.10.7 2 | build_dependencies: 3 | - pip==23.3.1 4 | - setuptools==63.2.0 5 | - wheel==0.41.2 6 | dependencies: 7 | - -r requirements.txt 8 | -------------------------------------------------------------------------------- /examples/mlruns/0/3a98a956c3114efcbceb6cce54407227/artifacts/model/requirements.txt: -------------------------------------------------------------------------------- 1 | mlflow==2.7.1 2 | cloudpickle==2.2.1 3 | numpy==1.23.5 4 | packaging==20.9 5 | psutil==5.9.6 6 | pyyaml==6.0.1 7 | scikit-learn==1.3.1 8 | scipy==1.11.3 -------------------------------------------------------------------------------- /examples/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/artifacts/model/requirements.txt: -------------------------------------------------------------------------------- 1 | mlflow==2.7.1 2 | cloudpickle==2.2.1 3 | numpy==1.23.5 4 | packaging==20.9 5 | psutil==5.9.6 6 | pyyaml==6.0.1 7 | scikit-learn==1.3.1 8 | scipy==1.11.3 -------------------------------------------------------------------------------- /examples/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/artifacts/model/requirements.txt: -------------------------------------------------------------------------------- 1 | mlflow==2.7.1 2 | cloudpickle==2.2.1 3 | numpy==1.23.5 4 | packaging==20.9 5 | psutil==5.9.6 6 | pyyaml==6.0.1 7 | scikit-learn==1.3.1 8 | scipy==1.11.3 -------------------------------------------------------------------------------- /examples/mlruns/0/meta.yaml: -------------------------------------------------------------------------------- 1 | artifact_location: file:///Users/faizank/workspace/experiments/model_deployment_experiments/slashdocker/mlruns/0 2 | creation_time: 1699402176987 3 | experiment_id: '0' 4 | last_update_time: 1699402176987 5 | lifecycle_stage: active 6 | name: Default 7 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | addopts = --capture=no 3 | filterwarnings = 4 | ignore::DeprecationWarning 5 | ignore::PendingDeprecationWarning 6 | markers = 7 | integration: mark tests for use in local integration testing only (deselect with '-m "not integration"') 8 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/util/data_structures.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, Optional, TypeVar 2 | 3 | X = TypeVar("X") 4 | Y = TypeVar("Y") 5 | 6 | 7 | def transform_optional(x: Optional[X], fn: Callable[[X], Optional[Y]]) -> Optional[Y]: 8 | if x is None: 9 | return None 10 | 11 | return fn(x) 12 | -------------------------------------------------------------------------------- /docs/CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | 3 | Release notes for new versions of Slash_Docker, in reverse chronological order. 4 | 5 | ### Version 0.0.1 (initial release) 6 | 7 | This release introduces Slash_Docker, with support for the following frameworks 8 | 9 | * Scikit-learn 10 | * XGBoost 11 | * PyTorch 12 | * TensorFlow 13 | * MLFLow 14 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/blob/blob_backend.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from pathlib import Path 3 | 4 | 5 | class BlobBackend(ABC): 6 | """A blob backend downloads large remote files.""" 7 | 8 | @abstractmethod 9 | def download(self, url: str, download_to: Path): 10 | raise NotImplementedError() 11 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/patch/test_types.py: -------------------------------------------------------------------------------- 1 | from truss.patch.signature import calc_truss_signature 2 | from truss.patch.types import TrussSignature 3 | 4 | 5 | def test_truss_signature_type(custom_model_truss_dir): 6 | sign = calc_truss_signature(custom_model_truss_dir) 7 | assert TrussSignature.from_dict(sign.to_dict()) == sign 8 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/server/requirements.txt: -------------------------------------------------------------------------------- 1 | -i https://pypi.org/simple 2 | 3 | argparse==1.4.0 4 | aiocontextvars==0.2.2 5 | cython==0.29.23 6 | msgpack-numpy==0.4.8 7 | msgpack==1.0.2 8 | python-json-logger==2.0.2 9 | pyyaml==6.0.0 10 | fastapi==0.95.0 11 | uvicorn==0.21.1 12 | psutil==5.9.4 13 | joblib==1.2.0 14 | requests==2.31.0 15 | -------------------------------------------------------------------------------- /examples/mlflow_sklearn/data/model/model/conda.yaml: -------------------------------------------------------------------------------- 1 | channels: 2 | - conda-forge 3 | dependencies: 4 | - python=3.10.7 5 | - pip<=23.3.1 6 | - pip: 7 | - mlflow==2.7.1 8 | - cloudpickle==2.2.1 9 | - numpy==1.23.5 10 | - packaging==20.9 11 | - psutil==5.9.6 12 | - pyyaml==6.0.1 13 | - scikit-learn==1.3.1 14 | - scipy==1.11.3 15 | name: mlflow-env 16 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | import-order-style=google 3 | # Note: this forces all google imports to be in the third group. See 4 | # https://github.com/PyCQA/flake8-import-order/issues/111 5 | ignore = E203, E266, W503 6 | max-line-length = 120 7 | exclude = 8 | __pycache__, 9 | .git, 10 | *.pyc, 11 | conf.py, 12 | **tests/samples.py 13 | -------------------------------------------------------------------------------- /examples/mlruns/0/3a98a956c3114efcbceb6cce54407227/artifacts/model/conda.yaml: -------------------------------------------------------------------------------- 1 | channels: 2 | - conda-forge 3 | dependencies: 4 | - python=3.10.7 5 | - pip<=23.3.1 6 | - pip: 7 | - mlflow==2.7.1 8 | - cloudpickle==2.2.1 9 | - numpy==1.23.5 10 | - packaging==20.9 11 | - psutil==5.9.6 12 | - pyyaml==6.0.1 13 | - scikit-learn==1.3.1 14 | - scipy==1.11.3 15 | name: mlflow-env 16 | -------------------------------------------------------------------------------- /examples/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/artifacts/model/conda.yaml: -------------------------------------------------------------------------------- 1 | channels: 2 | - conda-forge 3 | dependencies: 4 | - python=3.10.7 5 | - pip<=23.3.1 6 | - pip: 7 | - mlflow==2.7.1 8 | - cloudpickle==2.2.1 9 | - numpy==1.23.5 10 | - packaging==20.9 11 | - psutil==5.9.6 12 | - pyyaml==6.0.1 13 | - scikit-learn==1.3.1 14 | - scipy==1.11.3 15 | name: mlflow-env 16 | -------------------------------------------------------------------------------- /examples/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/artifacts/model/conda.yaml: -------------------------------------------------------------------------------- 1 | channels: 2 | - conda-forge 3 | dependencies: 4 | - python=3.10.7 5 | - pip<=23.3.1 6 | - pip: 7 | - mlflow==2.7.1 8 | - cloudpickle==2.2.1 9 | - numpy==1.23.5 10 | - packaging==20.9 11 | - psutil==5.9.6 12 | - pyyaml==6.0.1 13 | - scikit-learn==1.3.1 14 | - scipy==1.11.3 15 | name: mlflow-env 16 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/util/jinja.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | from jinja2 import Environment, FileSystemLoader, Template 4 | 5 | 6 | def read_template_from_fs(base_dir: Path, template_file_name: str) -> Template: 7 | template_loader = FileSystemLoader(str(base_dir)) 8 | template_env = Environment(loader=template_loader) 9 | return template_env.get_template(template_file_name) 10 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/docs.py: -------------------------------------------------------------------------------- 1 | """ 2 | Auto-generate CLI and Python client docstrings from 3 | source code using Sphinx. Format as Markdown for GitBook. 4 | """ 5 | 6 | 7 | # Generate Markdown files from docstrings 8 | def run_sphinx(): 9 | pass 10 | 11 | 12 | # Put desired files in proper location 13 | def move_docs(): 14 | pass 15 | 16 | 17 | # Delete unneeded files 18 | def cleanup(): 19 | pass 20 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/templates/control/control/helpers/test_context_managers.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from truss.templates.control.control.helpers.context_managers import current_directory 4 | 5 | 6 | def test_current_directory(tmp_path): 7 | orig_cwd = os.getcwd() 8 | with current_directory(tmp_path): 9 | assert os.getcwd() == str(tmp_path) 10 | 11 | assert os.getcwd() == orig_cwd 12 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa F401 2 | 3 | from pathlib import Path 4 | 5 | from single_source import get_version 6 | 7 | __version__ = get_version(__name__, Path(__file__).parent.parent) 8 | 9 | 10 | def version(): 11 | return __version__ 12 | 13 | 14 | from truss.build import ( 15 | create, 16 | create_from_mlflow_uri, 17 | from_directory, 18 | init, 19 | kill_all, 20 | load, 21 | mk_truss, 22 | ) 23 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/decorators.py: -------------------------------------------------------------------------------- 1 | def proxy_to_shadow_if_scattered(func): 2 | def wrapper(*args, **kwargs): 3 | from truss.truss_handle import TrussHandle 4 | 5 | truss_handle = args[0] 6 | if not truss_handle.is_scattered(): 7 | return func(*args, **kwargs) 8 | 9 | gathered_truss_handle = TrussHandle(truss_handle.gather()) 10 | return func(gathered_truss_handle, *args[1:], **kwargs) 11 | 12 | return wrapper 13 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/patch/test_signature.py: -------------------------------------------------------------------------------- 1 | from truss.patch.signature import calc_truss_signature 2 | 3 | 4 | def test_calc_truss_signature(custom_model_truss_dir): 5 | sign = calc_truss_signature(custom_model_truss_dir) 6 | assert len(sign.content_hashes_by_path) > 0 7 | assert "config.yaml" in sign.content_hashes_by_path 8 | with (custom_model_truss_dir / "config.yaml").open() as config_file: 9 | assert config_file.read() == sign.config 10 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/contexts/truss_context.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from pathlib import Path 3 | 4 | 5 | class TrussContext(ABC): 6 | """Marker class for Truss context. 7 | 8 | A model is represented in a standard form called Truss. Truss contexts 9 | perform a certain operation on this Truss. Some examples are: running the 10 | model directly, building an image and, deploying to baseten. 11 | """ 12 | 13 | @staticmethod 14 | @abstractmethod 15 | def run(truss_dir: Path): 16 | pass 17 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/pipeline/model/model.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | from typing import Any 3 | 4 | 5 | class Model: 6 | def __init__(self, **kwargs) -> None: 7 | self._data_dir = kwargs["data_dir"] 8 | self._config = kwargs["config"] 9 | self._pipeline = None 10 | 11 | def load(self): 12 | with open(self._data_dir / "pipeline.cpick", "rb") as f: 13 | self._pipeline = pickle.load(f) 14 | 15 | def predict(self, model_input: Any) -> Any: 16 | return self._pipeline(model_input) 17 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/test_truss_util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | 4 | from truss.util import path 5 | 6 | 7 | def test_max_modified(): 8 | epoch_time = int(time.time()) 9 | with path.given_or_temporary_dir() as dir: 10 | time.sleep(0.1) 11 | t1 = path.get_max_modified_time_of_dir(dir) 12 | assert t1 > epoch_time 13 | time.sleep(0.1) 14 | os.makedirs(os.path.join(dir, "test")) 15 | t2 = path.get_max_modified_time_of_dir(dir) 16 | assert t2 > t1 17 | 18 | 19 | test_max_modified() 20 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/patch/signature.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | from truss.patch.dir_signature import directory_content_signature 4 | from truss.patch.types import TrussSignature 5 | 6 | 7 | def calc_truss_signature(truss_dir: Path) -> TrussSignature: 8 | content_signature = directory_content_signature(truss_dir) 9 | with (truss_dir / "config.yaml").open("r") as config_file: 10 | config = config_file.read() 11 | return TrussSignature( 12 | content_hashes_by_path=content_signature, 13 | config=config, 14 | ) 15 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/control/control/helpers/context_managers.py: -------------------------------------------------------------------------------- 1 | import os 2 | from contextlib import contextmanager 3 | 4 | 5 | @contextmanager 6 | def current_directory(directory: str): 7 | """ 8 | Execute with using given directory as current working directory. 9 | 10 | This can be problematic in a multi-threaded scenario. It's assumed that 11 | the caller will take locks as needed. 12 | """ 13 | cwd = os.getcwd() 14 | os.chdir(directory) 15 | try: 16 | yield 17 | finally: 18 | os.chdir(cwd) 19 | -------------------------------------------------------------------------------- /examples/mlruns/0/3a98a956c3114efcbceb6cce54407227/meta.yaml: -------------------------------------------------------------------------------- 1 | artifact_uri: file:///Users/faizank/workspace/experiments/model_deployment_experiments/slashdocker/mlruns/0/3a98a956c3114efcbceb6cce54407227/artifacts 2 | end_time: 1699404188700 3 | entry_point_name: '' 4 | experiment_id: '0' 5 | lifecycle_stage: active 6 | run_id: 3a98a956c3114efcbceb6cce54407227 7 | run_name: masked-gnu-148 8 | run_uuid: 3a98a956c3114efcbceb6cce54407227 9 | source_name: '' 10 | source_type: 4 11 | source_version: '' 12 | start_time: 1699404186342 13 | status: 3 14 | tags: [] 15 | user_id: faizank 16 | -------------------------------------------------------------------------------- /examples/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/meta.yaml: -------------------------------------------------------------------------------- 1 | artifact_uri: file:///Users/faizank/workspace/experiments/model_deployment_experiments/slashdocker/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/artifacts 2 | end_time: 1699403589533 3 | entry_point_name: '' 4 | experiment_id: '0' 5 | lifecycle_stage: active 6 | run_id: ae2db1d887d84e73a40da9ad90a4c516 7 | run_name: salty-steed-823 8 | run_uuid: ae2db1d887d84e73a40da9ad90a4c516 9 | source_name: '' 10 | source_type: 4 11 | source_version: '' 12 | start_time: 1699403587361 13 | status: 3 14 | tags: [] 15 | user_id: faizank 16 | -------------------------------------------------------------------------------- /examples/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/meta.yaml: -------------------------------------------------------------------------------- 1 | artifact_uri: file:///Users/faizank/workspace/experiments/model_deployment_experiments/slashdocker/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/artifacts 2 | end_time: 1699402179860 3 | entry_point_name: '' 4 | experiment_id: '0' 5 | lifecycle_stage: active 6 | run_id: a33bc3b0ca6e4ab6990673faa2c44aa0 7 | run_name: thoughtful-foal-616 8 | run_uuid: a33bc3b0ca6e4ab6990673faa2c44aa0 9 | source_name: '' 10 | source_type: 4 11 | source_version: '' 12 | start_time: 1699402177165 13 | status: 3 14 | tags: [] 15 | user_id: faizank 16 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/notebook.py: -------------------------------------------------------------------------------- 1 | def is_notebook_or_ipython() -> bool: 2 | """Based on https://stackoverflow.com/a/39662359""" 3 | try: 4 | shell = get_ipython().__class__.__name__ # type: ignore 5 | if shell == "ZMQInteractiveShell": 6 | return True # Jupyter notebook or qtconsole 7 | elif shell == "TerminalInteractiveShell": 8 | return True # Terminal running IPython 9 | else: 10 | return False # Other type (?) 11 | except NameError: 12 | return False # Probably standard Python interpreter 13 | -------------------------------------------------------------------------------- /examples/mlruns/0/3a98a956c3114efcbceb6cce54407227/tags/mlflow.log-model.history: -------------------------------------------------------------------------------- 1 | [{"run_id": "3a98a956c3114efcbceb6cce54407227", "artifact_path": "model", "utc_time_created": "2023-11-08 00:43:06.352779", "flavors": {"python_function": {"model_path": "model.pkl", "predict_fn": "predict", "loader_module": "mlflow.sklearn", "python_version": "3.10.7", "env": {"conda": "conda.yaml", "virtualenv": "python_env.yaml"}}, "sklearn": {"pickled_model": "model.pkl", "sklearn_version": "1.3.1", "serialization_format": "cloudpickle", "code": null}}, "model_uuid": "202c5d4b4e6b4acb8a85d76d0a5147ec", "mlflow_version": "2.7.1"}] -------------------------------------------------------------------------------- /examples/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/tags/mlflow.log-model.history: -------------------------------------------------------------------------------- 1 | [{"run_id": "a33bc3b0ca6e4ab6990673faa2c44aa0", "artifact_path": "model", "utc_time_created": "2023-11-08 00:09:37.192688", "flavors": {"python_function": {"model_path": "model.pkl", "predict_fn": "predict", "loader_module": "mlflow.sklearn", "python_version": "3.10.7", "env": {"conda": "conda.yaml", "virtualenv": "python_env.yaml"}}, "sklearn": {"pickled_model": "model.pkl", "sklearn_version": "1.3.1", "serialization_format": "cloudpickle", "code": null}}, "model_uuid": "f715c46aeb0b430c9e488c462018ed0d", "mlflow_version": "2.7.1"}] -------------------------------------------------------------------------------- /examples/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/tags/mlflow.log-model.history: -------------------------------------------------------------------------------- 1 | [{"run_id": "ae2db1d887d84e73a40da9ad90a4c516", "artifact_path": "model", "utc_time_created": "2023-11-08 00:33:07.372597", "flavors": {"python_function": {"model_path": "model.pkl", "predict_fn": "predict", "loader_module": "mlflow.sklearn", "python_version": "3.10.7", "env": {"conda": "conda.yaml", "virtualenv": "python_env.yaml"}}, "sklearn": {"pickled_model": "model.pkl", "sklearn_version": "1.3.1", "serialization_format": "cloudpickle", "code": null}}, "model_uuid": "f5a5cc0f71a4488aa2f21ab62a8f56e4", "mlflow_version": "2.7.1"}] -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/custom/train/train.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from typing import Dict 3 | 4 | 5 | class Train: 6 | def __init__( 7 | self, 8 | config, 9 | output_dir: Path, 10 | variables: Dict, 11 | secrets, 12 | ): 13 | self._config = config 14 | self._output_dir = output_dir 15 | self._variables = variables 16 | self._secrets = secrets 17 | 18 | def train(self): 19 | # Write your training code here, populating generated artifacts in 20 | # self._output_dir 21 | pass 22 | -------------------------------------------------------------------------------- /examples/mlflow_sklearn/data/model/model/MLmodel: -------------------------------------------------------------------------------- 1 | artifact_path: model 2 | flavors: 3 | python_function: 4 | env: 5 | conda: conda.yaml 6 | virtualenv: python_env.yaml 7 | loader_module: mlflow.sklearn 8 | model_path: model.pkl 9 | predict_fn: predict 10 | python_version: 3.10.7 11 | sklearn: 12 | code: null 13 | pickled_model: model.pkl 14 | serialization_format: cloudpickle 15 | sklearn_version: 1.3.1 16 | mlflow_version: 2.7.1 17 | model_uuid: 202c5d4b4e6b4acb8a85d76d0a5147ec 18 | run_id: 3a98a956c3114efcbceb6cce54407227 19 | utc_time_created: '2023-11-08 00:43:06.352779' 20 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/training.Dockerfile.jinja: -------------------------------------------------------------------------------- 1 | {% extends "base.Dockerfile.jinja" %} 2 | 3 | {% block install_requirements %} 4 | COPY ./training_requirements.txt training_requirements.txt 5 | RUN pip install -r training_requirements.txt --no-cache-dir && rm -rf /root/.cache/pip 6 | {{ super() }} 7 | {% endblock %} 8 | 9 | {% block app_copy %} 10 | COPY ./training /app 11 | # todo: put this behind a flag 12 | # COPY ./{{ config.train.train_module_dir }} /app/train 13 | COPY ./config.yaml /app/config.yaml 14 | {% endblock %} 15 | 16 | 17 | {% block run %} 18 | CMD exec python3 /app/job.py 19 | {% endblock %} 20 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/patch/test_dir_signature.py: -------------------------------------------------------------------------------- 1 | from truss.patch.dir_signature import directory_content_signature 2 | 3 | 4 | def test_directory_content_signature(tmp_path): 5 | root = tmp_path / "root" 6 | root.mkdir() 7 | (root / "file1").touch() 8 | (root / "file2").touch() 9 | subdir = root / "dir" 10 | subdir.mkdir() 11 | (subdir / "file3").touch() 12 | 13 | content_sign = directory_content_signature(root) 14 | print(content_sign) 15 | 16 | assert content_sign.keys() == { 17 | "dir", 18 | "dir/file3", 19 | "file1", 20 | "file2", 21 | } 22 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/test_docker.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from python_on_whales import docker 3 | from truss.docker import get_urls_from_container 4 | 5 | 6 | @pytest.fixture 7 | def docker_container(): 8 | container = docker.container.create("nginx", publish=[[19051, 19051]]) 9 | try: 10 | container.start() 11 | yield container 12 | finally: 13 | container.stop() 14 | 15 | 16 | @pytest.mark.integration 17 | def test_get_urls_from_container(docker_container): 18 | resp = get_urls_from_container(docker_container) 19 | assert resp == {19051: ["http://0.0.0.0:19051", "http://:::19051"]} 20 | -------------------------------------------------------------------------------- /examples/mlruns/0/3a98a956c3114efcbceb6cce54407227/artifacts/model/MLmodel: -------------------------------------------------------------------------------- 1 | artifact_path: model 2 | flavors: 3 | python_function: 4 | env: 5 | conda: conda.yaml 6 | virtualenv: python_env.yaml 7 | loader_module: mlflow.sklearn 8 | model_path: model.pkl 9 | predict_fn: predict 10 | python_version: 3.10.7 11 | sklearn: 12 | code: null 13 | pickled_model: model.pkl 14 | serialization_format: cloudpickle 15 | sklearn_version: 1.3.1 16 | mlflow_version: 2.7.1 17 | model_uuid: 202c5d4b4e6b4acb8a85d76d0a5147ec 18 | run_id: 3a98a956c3114efcbceb6cce54407227 19 | utc_time_created: '2023-11-08 00:43:06.352779' 20 | -------------------------------------------------------------------------------- /examples/mlruns/0/a33bc3b0ca6e4ab6990673faa2c44aa0/artifacts/model/MLmodel: -------------------------------------------------------------------------------- 1 | artifact_path: model 2 | flavors: 3 | python_function: 4 | env: 5 | conda: conda.yaml 6 | virtualenv: python_env.yaml 7 | loader_module: mlflow.sklearn 8 | model_path: model.pkl 9 | predict_fn: predict 10 | python_version: 3.10.7 11 | sklearn: 12 | code: null 13 | pickled_model: model.pkl 14 | serialization_format: cloudpickle 15 | sklearn_version: 1.3.1 16 | mlflow_version: 2.7.1 17 | model_uuid: f715c46aeb0b430c9e488c462018ed0d 18 | run_id: a33bc3b0ca6e4ab6990673faa2c44aa0 19 | utc_time_created: '2023-11-08 00:09:37.192688' 20 | -------------------------------------------------------------------------------- /examples/mlruns/0/ae2db1d887d84e73a40da9ad90a4c516/artifacts/model/MLmodel: -------------------------------------------------------------------------------- 1 | artifact_path: model 2 | flavors: 3 | python_function: 4 | env: 5 | conda: conda.yaml 6 | virtualenv: python_env.yaml 7 | loader_module: mlflow.sklearn 8 | model_path: model.pkl 9 | predict_fn: predict 10 | python_version: 3.10.7 11 | sklearn: 12 | code: null 13 | pickled_model: model.pkl 14 | serialization_format: cloudpickle 15 | sklearn_version: 1.3.1 16 | mlflow_version: 2.7.1 17 | model_uuid: f5a5cc0f71a4488aa2f21ab62a8f56e4 18 | run_id: ae2db1d887d84e73a40da9ad90a4c516 19 | utc_time_created: '2023-11-08 00:33:07.372597' 20 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/util/gpu.py: -------------------------------------------------------------------------------- 1 | import subprocess as sp 2 | from typing import Optional 3 | 4 | 5 | def get_gpu_memory() -> Optional[int]: 6 | # https://stackoverflow.com/questions/59567226/how-to-programmatically-determine-available-gpu-memory-with-tensorflow 7 | try: 8 | command = "nvidia-smi --query-gpu=memory.used --format=csv" 9 | memory_free_info = ( 10 | sp.check_output(command.split()).decode("ascii").split("\n")[1] 11 | ) 12 | memory_free_values = int(memory_free_info.split()[0]) 13 | return memory_free_values 14 | except FileNotFoundError: 15 | return None 16 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/custom/model/model.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | 4 | class Model: 5 | def __init__(self, **kwargs) -> None: 6 | self._data_dir = kwargs["data_dir"] 7 | self._config = kwargs["config"] 8 | self._secrets = kwargs["secrets"] 9 | self._model = None 10 | 11 | def load(self): 12 | # Load model here and assign to self._model. 13 | pass 14 | 15 | def predict(self, model_input: Any) -> Any: 16 | model_output = {} 17 | # Invoke model on model_input and calculate predictions here. 18 | model_output["predictions"] = [] 19 | return model_output 20 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/server/common/retry.py: -------------------------------------------------------------------------------- 1 | import time 2 | from typing import Callable 3 | 4 | 5 | def retry( 6 | fn: Callable, 7 | count: int, 8 | logging_fn: Callable, 9 | base_message: str, 10 | gap_seconds: float = 0.0, 11 | ): 12 | i = 0 13 | while i <= count: 14 | try: 15 | fn() 16 | return 17 | except Exception as exc: 18 | msg = base_message 19 | if i >= count: 20 | raise exc 21 | 22 | if i == 0: 23 | msg = f"{msg} Retrying..." 24 | else: 25 | msg = f"{msg} Retrying. Retry count: {i}" 26 | logging_fn(msg) 27 | i += 1 28 | time.sleep(gap_seconds) 29 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/test_notebooks.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import nbformat 4 | import pytest 5 | from nbconvert.preprocessors import ExecutePreprocessor 6 | 7 | 8 | @pytest.mark.parametrize("notebook", ["happy.ipynb"]) 9 | def test_notebook_exec(notebook): 10 | """ 11 | Test for Jupyter notebooks to establish some base exercising. Future tests can add new notebooks similar to 12 | happy.ipynb in the test_data directory, then add that file to the above parameterization 13 | """ 14 | with open( 15 | f"{os.path.dirname(os.path.realpath(__file__))}/../test_data/{notebook}" 16 | ) as f: 17 | nb = nbformat.read(f, as_version=4) 18 | ep = ExecutePreprocessor(timeout=600, kernel_name="python3") 19 | ep.preprocess(nb) 20 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/blob/http_public_blob_backend.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | from pathlib import Path 3 | 4 | import requests 5 | from truss.blob.blob_backend import BlobBackend 6 | 7 | BLOB_DOWNLOAD_TIMEOUT_SECS = 600 # 10 minutes 8 | 9 | 10 | class HttpPublic(BlobBackend): 11 | """Downloads without auth, files must be publicly available.""" 12 | 13 | def download(self, URL: str, download_to: Path): 14 | # Streaming download to keep memory usage low 15 | resp = requests.get( 16 | URL, 17 | allow_redirects=True, 18 | stream=True, 19 | timeout=BLOB_DOWNLOAD_TIMEOUT_SECS, 20 | ) 21 | resp.raise_for_status() 22 | with download_to.open("wb") as file: 23 | shutil.copyfileobj(resp.raw, file) 24 | -------------------------------------------------------------------------------- /model_to_docker/api.py: -------------------------------------------------------------------------------- 1 | # import sys 2 | # import os 3 | # sys.path.append(f'{os.getcwd()}/truss_saver') 4 | 5 | from .truss_saver import truss 6 | from .truss_saver.truss.docker import Docker 7 | from .truss_saver.truss import cli 8 | 9 | 10 | def save_model(model, model_name, python_dependencies=None, system_dependencies=None): 11 | tr = truss.create(model, model_name) 12 | from pprint import pprint as pp 13 | return pp({'config': tr.__dict__['_spec'].__dict__['_config'].__dict__}) 14 | 15 | def run_model_server(model_path, port=8080): 16 | container = cli.run_image(model_path, port=port) 17 | return container 18 | 19 | def stop_model_server(container_id): 20 | Docker.client().stop(containers=container_id) 21 | 22 | 23 | def active_model_server(): 24 | Docker.client().containers.list() -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/readme_generator.py: -------------------------------------------------------------------------------- 1 | from jinja2 import Template 2 | from truss.constants import README_TEMPLATE_NAME, TEMPLATES_DIR 3 | from truss.truss_spec import TrussSpec 4 | 5 | 6 | def generate_readme(_spec: TrussSpec) -> str: 7 | readme_template_path = TEMPLATES_DIR / README_TEMPLATE_NAME 8 | with readme_template_path.open() as readme_template_file: 9 | readme_template = Template(readme_template_file.read()) 10 | # examples.yaml may not exist 11 | # if examples.yaml does exist, but it's empty, examples_raw is None 12 | examples_raw = _spec.examples if _spec.examples_path.exists() else None 13 | readme_contents = readme_template.render( 14 | config=_spec.config, examples=examples_raw 15 | ) 16 | return readme_contents 17 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | ## Contributing 2 | 3 | Slash_Docker is created at [slashml](https://slashml.co), but as an open and living project eagerly accepts contributions of all kinds from the broader developer community. Please note that all participation with slash_docker falls under our [code of conduct](CODE_OF_CONDUCT.md). 4 | 5 | 6 | You can perform one of the following tasks: 7 | 8 | * For bugs and feature requests, file an issue. 9 | * For changes and updates, create a pull request. 10 | 11 | ## Local development 12 | 13 | To get started contributing to the library, all you have to do is clone this repository! 14 | 15 | ### Setup 16 | 17 | Git clone this repo `git clone git@github.com:slashml/slash_docker.git`. Then `cd` into `slash_docker`. 18 | 19 | 20 | Then install the dependencies with: 21 | ``` 22 | pip install -r requirements.txt 23 | ``` -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/errors.py: -------------------------------------------------------------------------------- 1 | class Error(Exception): 2 | """Base Baseten Error""" 3 | 4 | def __init__(self, message): 5 | super().__init__(message) 6 | self.message = message 7 | 8 | 9 | class FrameworkNotSupportedError(Error): 10 | """Raised in places where the user attempts to use Baseten with an unsupported framework""" 11 | 12 | pass 13 | 14 | 15 | class ModelFilesMissingError(Error): 16 | pass 17 | 18 | 19 | class ModelClassImplementationError(Error): 20 | pass 21 | 22 | 23 | class InvalidConfigurationError(Error): 24 | pass 25 | 26 | 27 | class ValidationError(Error): 28 | pass 29 | 30 | 31 | class ContainerIsDownError(Error): 32 | pass 33 | 34 | 35 | class ContainerNotFoundError(Error): 36 | pass 37 | 38 | 39 | class ContainerAPINoResponseError(Error): 40 | pass 41 | -------------------------------------------------------------------------------- /examples/mlflow_sklearn/model/model.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | import mlflow 4 | import numpy as np 5 | 6 | 7 | class Model: 8 | def __init__(self, **kwargs) -> None: 9 | self._data_dir = kwargs["data_dir"] 10 | config = kwargs["config"] 11 | model_metadata = config["model_metadata"] 12 | self._model_binary_dir = model_metadata["model_binary_dir"] 13 | self._model = None 14 | 15 | def load(self): 16 | model_binary_dir_path = self._data_dir / self._model_binary_dir 17 | self._model = mlflow.pyfunc.load_model(model_binary_dir_path / "model") 18 | 19 | def predict(self, model_input: Any) -> Any: 20 | model_output = {} 21 | inputs = np.array(model_input) 22 | result = self._model.predict(inputs).tolist() 23 | model_output["predictions"] = result 24 | return model_output 25 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/blob/blob_backend_registry.py: -------------------------------------------------------------------------------- 1 | from typing import Dict 2 | 3 | from truss.blob.blob_backend import BlobBackend 4 | from truss.blob.http_public_blob_backend import HttpPublic 5 | from truss.constants import HTTP_PUBLIC_BLOB_BACKEND 6 | 7 | 8 | class _BlobBackendRegistry: 9 | def __init__(self) -> None: 10 | self._backends: Dict[str, BlobBackend] = {} 11 | # Register default backend 12 | self._backends[HTTP_PUBLIC_BLOB_BACKEND] = HttpPublic() 13 | 14 | def register_backend(self, name: str, backend: BlobBackend): 15 | self._backends[name] = backend 16 | 17 | def get_backend(self, name: str): 18 | if name not in self._backends: 19 | raise ValueError(f"Backend {name} is not registered.") 20 | return self._backends[name] 21 | 22 | 23 | BLOB_BACKEND_REGISTRY = _BlobBackendRegistry() 24 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/keras/model/model.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | import numpy as np 4 | from tensorflow import keras 5 | 6 | 7 | class Model: 8 | def __init__(self, **kwargs) -> None: 9 | self._data_dir = kwargs["data_dir"] 10 | config = kwargs["config"] 11 | model_metadata = config["model_metadata"] 12 | self._model_binary_dir = model_metadata["model_binary_dir"] 13 | self._model = None 14 | 15 | def load(self): 16 | self._model = keras.models.load_model( 17 | str(self._data_dir / self._model_binary_dir) 18 | ) 19 | 20 | def predict(self, model_input: Any) -> Any: 21 | model_output = {} 22 | inputs = np.array(model_input) 23 | result = self._model.predict(inputs).tolist() 24 | model_output["predictions"] = result 25 | return model_output 26 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/templates/core/server/common/test_util.py: -------------------------------------------------------------------------------- 1 | from unittest import mock 2 | 3 | 4 | def model_supports_predict_proba(): 5 | mock_not_predict_proba = mock.Mock(name="mock_not_predict_proba") 6 | mock_not_predict_proba.predict_proba.return_value = False 7 | 8 | mock_check_proba = mock.Mock(name="mock_check_proba") 9 | mock_check_proba.predict_proba.return_value = True 10 | mock_check_proba._check_proba.return_value = True 11 | 12 | mock_not_check_proba = mock.Mock(name="mock_not_check_proba") 13 | mock_not_check_proba.predict_proba.return_value = True 14 | mock_not_check_proba._check_proba.side_effect = AttributeError 15 | 16 | assert not model_supports_predict_proba(mock_not_predict_proba) 17 | assert model_supports_predict_proba(mock_check_proba) 18 | assert not model_supports_predict_proba(mock_not_check_proba) 19 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/mlflow/model/model.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | import mlflow 4 | import numpy as np 5 | 6 | 7 | class Model: 8 | def __init__(self, **kwargs) -> None: 9 | self._data_dir = kwargs["data_dir"] 10 | config = kwargs["config"] 11 | model_metadata = config["model_metadata"] 12 | self._model_binary_dir = model_metadata["model_binary_dir"] 13 | self._model = None 14 | 15 | def load(self): 16 | model_binary_dir_path = self._data_dir / self._model_binary_dir 17 | self._model = mlflow.pyfunc.load_model(model_binary_dir_path / "model") 18 | 19 | def predict(self, model_input: Any) -> Any: 20 | model_output = {} 21 | inputs = np.array(model_input) 22 | result = self._model.predict(inputs).tolist() 23 | model_output["predictions"] = result 24 | return model_output 25 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/server/inference_server.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import Dict 3 | 4 | import yaml 5 | from common.logging import setup_logging 6 | from common.truss_server import TrussServer # noqa: E402 7 | 8 | CONFIG_FILE = "config.yaml" 9 | 10 | setup_logging() 11 | 12 | 13 | class ConfiguredTrussServer: 14 | _config: Dict 15 | _port: int 16 | 17 | def __init__(self, config_path: str, port: int): 18 | self._port = port 19 | with open(config_path, encoding="utf-8") as config_file: 20 | self._config = yaml.safe_load(config_file) 21 | 22 | def start(self): 23 | server = TrussServer(http_port=self._port, config=self._config) 24 | server.start() 25 | 26 | 27 | if __name__ == "__main__": 28 | env_port = int(os.environ.get("INFERENCE_SERVER_PORT", "8080")) 29 | ConfiguredTrussServer(CONFIG_FILE, env_port).start() 30 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/patch/dir_signature.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from typing import Dict 3 | 4 | from truss.patch.hash import file_content_hash_str 5 | 6 | 7 | def directory_content_signature(root: Path) -> Dict[str, str]: 8 | """Calculate content signature of a filesystem directory. 9 | 10 | Sort all files by path, store file path with content hash. 11 | Signatures are meant to track changes in directory. e.g. 12 | A previous signature and the directory can be combined to create 13 | a patch from the previous state. 14 | 15 | Hash of directories is marked None. 16 | """ 17 | paths = list(root.glob("**/*")) 18 | paths.sort(key=lambda p: p.relative_to(root)) 19 | 20 | def path_hash(pth: Path): 21 | if pth.is_file(): 22 | return file_content_hash_str(pth) 23 | 24 | return {str(path.relative_to(root)): path_hash(path) for path in paths} 25 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/contexts/local_loader/test_load_local.py: -------------------------------------------------------------------------------- 1 | from truss.contexts.local_loader.utils import prepare_secrets 2 | from truss.local.local_config_handler import LocalConfigHandler 3 | from truss.truss_handle import TrussHandle 4 | from truss.truss_spec import TrussSpec 5 | 6 | 7 | def test_prepare_secrets(custom_model_truss_dir, tmp_path): 8 | orig_truss_config_dir = LocalConfigHandler.TRUSS_CONFIG_DIR 9 | LocalConfigHandler.TRUSS_CONFIG_DIR = tmp_path 10 | try: 11 | LocalConfigHandler.set_secret("secret_name", "secret_value") 12 | handle = TrussHandle(custom_model_truss_dir) 13 | handle.add_secret("secret_name") 14 | spec = TrussSpec(custom_model_truss_dir) 15 | secrets = prepare_secrets(spec) 16 | assert secrets["secret_name"] == "secret_value" 17 | finally: 18 | LocalConfigHandler.TRUSS_CONFIG_DIR = orig_truss_config_dir 19 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/context_builder.Dockerfile: -------------------------------------------------------------------------------- 1 | # Builds baseten/truss-context-builder, a light-weight image that can be used 2 | # for creating docker build context out of a Truss. 3 | # Build that image as: 4 | # docker buildx build . -f context_builder.Dockerfile --platform=linux/amd64 -t baseten/truss-context-builder 5 | FROM python:3.9-slim 6 | 7 | RUN apt-get update \ 8 | && apt-get install --yes --no-install-recommends curl tar \ 9 | && apt-get autoremove -y \ 10 | && apt-get clean -y \ 11 | && rm -rf /var/lib/apt/lists/* /tmp/library-scripts/ 12 | 13 | RUN curl -sSL https://install.python-poetry.org | python - 14 | 15 | ENV PATH="/root/.local/bin:${PATH}" 16 | COPY . . 17 | 18 | # https://python-poetry.org/docs/configuration/#virtualenvsin-project 19 | # to write to project root .venv file to be used for context builder test 20 | RUN poetry config virtualenvs.in-project true \ 21 | && poetry install --only builder 22 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | *.egg-info/ 3 | dict 4 | venv* 5 | 6 | *.csv 7 | *.log 8 | .DS_Store 9 | 10 | # Environments 11 | .env 12 | .venv 13 | env/ 14 | venv/ 15 | ENV/ 16 | env.bak/ 17 | venv.bak/ 18 | vv 19 | 20 | # IDE 21 | .idea/* 22 | .vscode/* 23 | 24 | # Byte-compiled / optimized / DLL files 25 | __pycache__/ 26 | *.py[cod] 27 | *$py.class 28 | .mypy_cache/ 29 | 30 | # Unit test / coverage reports 31 | htmlcov/ 32 | .tox/ 33 | .nox/ 34 | .coverage 35 | .coverage.* 36 | .cache 37 | nosetests.xml 38 | coverage.xml 39 | *.cover 40 | .hypothesis/ 41 | .pytest_cache/ 42 | 43 | # Distribution / packaging 44 | .Python 45 | build/ 46 | develop-eggs/ 47 | dist/ 48 | downloads/ 49 | eggs/ 50 | .eggs/ 51 | lib/ 52 | lib64/ 53 | parts/ 54 | sdist/ 55 | var/ 56 | wheels/ 57 | pip-wheel-metadata/ 58 | share/python-wheels/ 59 | *.egg-info/ 60 | .installed.cfg 61 | *.egg 62 | MANIFEST 63 | node_modules/ 64 | .npm 65 | npm-debug* 66 | 67 | # vim 68 | *.swp 69 | 70 | docs/_build/ -------------------------------------------------------------------------------- /examples/xgboost/config.yaml: -------------------------------------------------------------------------------- 1 | apply_library_patches: true 2 | bundled_packages_dir: packages 3 | data_dir: data 4 | description: null 5 | environment_variables: {} 6 | examples_filename: examples.yaml 7 | external_package_dirs: [] 8 | input_type: Any 9 | live_reload: false 10 | model_class_filename: model.py 11 | model_class_name: Model 12 | model_framework: xgboost 13 | model_metadata: 14 | model_binary_dir: model 15 | supports_predict_proba: false 16 | model_module_dir: model 17 | model_name: null 18 | model_type: Model 19 | python_version: py310 20 | requirements: 21 | - xgboost==2.0.1 22 | resources: 23 | accelerator: null 24 | cpu: 500m 25 | memory: 512Mi 26 | use_gpu: false 27 | secrets: {} 28 | spec_version: '2.0' 29 | system_packages: [] 30 | train: 31 | resources: 32 | accelerator: null 33 | cpu: 500m 34 | memory: 512Mi 35 | use_gpu: false 36 | training_class_filename: train.py 37 | training_class_name: Train 38 | training_module_dir: train 39 | variables: {} 40 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/contexts/local_loader/utils.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import inspect 3 | from typing import Dict 4 | 5 | from truss.local.local_config_handler import LocalConfigHandler 6 | from truss.truss_spec import TrussSpec 7 | 8 | 9 | def prepare_secrets(spec: TrussSpec) -> Dict[str, str]: 10 | secrets = copy.deepcopy(spec.secrets) 11 | local_secerts = LocalConfigHandler.get_config().secrets 12 | for secret_name in secrets: 13 | if secret_name in local_secerts: 14 | secrets[secret_name] = local_secerts[secret_name] 15 | return secrets 16 | 17 | 18 | def signature_accepts_keyword_arg(signature: inspect.Signature, kwarg: str) -> bool: 19 | return kwarg in signature.parameters or _signature_accepts_kwargs(signature) 20 | 21 | 22 | def _signature_accepts_kwargs(signature: inspect.Signature) -> bool: 23 | for param in signature.parameters.values(): 24 | if param.kind == inspect.Parameter.VAR_KEYWORD: 25 | return True 26 | return False 27 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/local/local_config.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass, field 2 | from pathlib import Path 3 | from typing import Dict 4 | 5 | import yaml 6 | 7 | 8 | @dataclass 9 | class LocalConfig: 10 | secrets: Dict[str, str] = field(default_factory=dict) 11 | use_sudo: bool = False 12 | 13 | @staticmethod 14 | def from_dict(d): 15 | return LocalConfig( 16 | secrets=d.get("secrets", {}), 17 | use_sudo=d.get("use_sudo", False), 18 | ) 19 | 20 | @staticmethod 21 | def from_yaml(yaml_path: Path): 22 | with yaml_path.open() as yaml_file: 23 | return LocalConfig.from_dict(yaml.safe_load(yaml_file)) 24 | 25 | def to_dict(self): 26 | return { 27 | "secrets": self.secrets, 28 | "use_sudo": self.use_sudo, 29 | } 30 | 31 | def write_to_yaml_file(self, path: Path): 32 | with path.open("w") as config_file: 33 | yaml.dump(self.to_dict(), config_file) 34 | -------------------------------------------------------------------------------- /examples/bert_uncased/config.yaml: -------------------------------------------------------------------------------- 1 | apply_library_patches: true 2 | bundled_packages_dir: packages 3 | data_dir: data 4 | description: null 5 | environment_variables: {} 6 | examples_filename: examples.yaml 7 | external_package_dirs: [] 8 | input_type: Any 9 | live_reload: false 10 | model_class_filename: model.py 11 | model_class_name: Model 12 | model_framework: huggingface_transformer 13 | model_metadata: 14 | has_hybrid_args: false 15 | has_named_args: false 16 | transformer_config: 17 | model: bert-base-uncased 18 | model_module_dir: model 19 | model_name: null 20 | model_type: fill-mask 21 | python_version: py310 22 | requirements: 23 | - transformers==4.34.1 24 | resources: 25 | accelerator: null 26 | cpu: 500m 27 | memory: 512Mi 28 | use_gpu: false 29 | secrets: {} 30 | spec_version: '2.0' 31 | system_packages: [] 32 | train: 33 | resources: 34 | accelerator: null 35 | cpu: 500m 36 | memory: 512Mi 37 | use_gpu: false 38 | training_class_filename: train.py 39 | training_class_name: Train 40 | training_module_dir: train 41 | variables: {} 42 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/model_frameworks/keras.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from typing import Dict, Set 3 | 4 | from truss.constants import TENSORFLOW_REQ_MODULE_NAMES 5 | from truss.model_framework import ModelFramework 6 | from truss.types import ModelFrameworkType 7 | 8 | 9 | class Keras(ModelFramework): 10 | def typ(self) -> ModelFrameworkType: 11 | return ModelFrameworkType.KERAS 12 | 13 | def required_python_depedencies(self) -> Set[str]: 14 | return TENSORFLOW_REQ_MODULE_NAMES 15 | 16 | def serialize_model_to_directory(self, model, target_directory: Path): 17 | model.save(target_directory) 18 | 19 | def model_metadata(self, model) -> Dict[str, str]: 20 | return { 21 | "model_binary_dir": "model", 22 | } 23 | 24 | def supports_model_class(self, model_class) -> bool: 25 | model_framework, _, _ = model_class.__module__.partition(".") 26 | return model_framework in [ 27 | ModelFrameworkType.KERAS.value, 28 | ModelFrameworkType.TENSORFLOW.value, 29 | ] 30 | -------------------------------------------------------------------------------- /examples/test_model/config.yaml: -------------------------------------------------------------------------------- 1 | apply_library_patches: true 2 | bundled_packages_dir: packages 3 | data_dir: data 4 | description: null 5 | environment_variables: {} 6 | examples_filename: examples.yaml 7 | external_package_dirs: [] 8 | input_type: Any 9 | live_reload: false 10 | model_class_filename: model.py 11 | model_class_name: Model 12 | model_framework: sklearn 13 | model_metadata: 14 | model_binary_dir: model 15 | supports_predict_proba: false 16 | model_module_dir: model 17 | model_name: null 18 | model_type: Model 19 | python_version: py310 20 | requirements: 21 | - numpy==1.23.5 22 | - scipy==1.11.3 23 | - joblib==1.3.2 24 | - scikit-learn==1.3.1 25 | - threadpoolctl==3.2.0 26 | resources: 27 | accelerator: null 28 | cpu: 500m 29 | memory: 512Mi 30 | use_gpu: false 31 | secrets: {} 32 | spec_version: '2.0' 33 | system_packages: [] 34 | train: 35 | resources: 36 | accelerator: null 37 | cpu: 500m 38 | memory: 512Mi 39 | use_gpu: false 40 | training_class_filename: train.py 41 | training_class_name: Train 42 | training_module_dir: train 43 | variables: {} 44 | -------------------------------------------------------------------------------- /examples/test_model_1/config.yaml: -------------------------------------------------------------------------------- 1 | apply_library_patches: true 2 | bundled_packages_dir: packages 3 | data_dir: data 4 | description: null 5 | environment_variables: {} 6 | examples_filename: examples.yaml 7 | external_package_dirs: [] 8 | input_type: Any 9 | live_reload: false 10 | model_class_filename: model.py 11 | model_class_name: Model 12 | model_framework: sklearn 13 | model_metadata: 14 | model_binary_dir: model 15 | supports_predict_proba: false 16 | model_module_dir: model 17 | model_name: null 18 | model_type: Model 19 | python_version: py310 20 | requirements: 21 | - threadpoolctl==3.2.0 22 | - scipy==1.11.3 23 | - joblib==1.3.2 24 | - numpy==1.23.5 25 | - scikit-learn==1.3.1 26 | resources: 27 | accelerator: null 28 | cpu: 500m 29 | memory: 512Mi 30 | use_gpu: false 31 | secrets: {} 32 | spec_version: '2.0' 33 | system_packages: [] 34 | train: 35 | resources: 36 | accelerator: null 37 | cpu: 500m 38 | memory: 512Mi 39 | use_gpu: false 40 | training_class_filename: train.py 41 | training_class_name: Train 42 | training_module_dir: train 43 | variables: {} 44 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/patch/types.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass 2 | from typing import Dict 3 | 4 | 5 | @dataclass 6 | class TrussSignature: 7 | """Truss signature stores information for calculating patches for future 8 | changes to Truss. 9 | 10 | Currently, it stores hashes of all of the paths in the truss directory, and 11 | the truss config contents. Path hashes allow calculating added/updated/removes 12 | paths in future trusses compared to this. Config contents allow calculating 13 | config changes, such as add/update/remove of python requirements etc. 14 | """ 15 | 16 | content_hashes_by_path: Dict[str, str] 17 | config: str 18 | 19 | def to_dict(self) -> dict: 20 | return { 21 | "content_hashes_by_path": self.content_hashes_by_path, 22 | "config": self.config, 23 | } 24 | 25 | @staticmethod 26 | def from_dict(d) -> "TrussSignature": 27 | return TrussSignature( 28 | content_hashes_by_path=d["content_hashes_by_path"], 29 | config=d["config"], 30 | ) 31 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/control/control/helpers/errors.py: -------------------------------------------------------------------------------- 1 | class Error(Exception): 2 | """Base Truss Error""" 3 | 4 | def __init__(self, message: str): 5 | super(Error, self).__init__(message) 6 | self.message = message 7 | self.type = type 8 | 9 | 10 | class PatchApplicatonError(Error): 11 | pass 12 | 13 | 14 | class UnsupportedPatch(PatchApplicatonError): 15 | """Patch unsupported by this truss""" 16 | 17 | pass 18 | 19 | 20 | class PatchFailedRecoverable(PatchApplicatonError): 21 | """Patch admissible but failed to apply. Recoverable via further patching.""" 22 | 23 | pass 24 | 25 | 26 | class PatchFailedUnrecoverable(PatchApplicatonError): 27 | """Patch admissible but failed to apply, leaving truss in unrecoverable state. 28 | Full deploy is required.""" 29 | 30 | pass 31 | 32 | 33 | class InadmissiblePatch(PatchApplicatonError): 34 | """Patch does not apply to current state of Truss.""" 35 | 36 | pass 37 | 38 | 39 | class ModelNotReady(Error): 40 | """Model has started running, but not ready yet.""" 41 | 42 | pass 43 | -------------------------------------------------------------------------------- /examples/mlflow_sklearn/config.yaml: -------------------------------------------------------------------------------- 1 | apply_library_patches: true 2 | bundled_packages_dir: packages 3 | data_dir: data 4 | description: null 5 | environment_variables: {} 6 | examples_filename: examples.yaml 7 | external_package_dirs: [] 8 | input_type: Any 9 | live_reload: false 10 | model_class_filename: model.py 11 | model_class_name: Model 12 | model_framework: mlflow 13 | model_metadata: 14 | model_binary_dir: model 15 | supports_predict_proba: false 16 | model_module_dir: model 17 | model_name: null 18 | model_type: Model 19 | python_version: py310 20 | requirements: 21 | - mlflow==2.7.1 22 | - cloudpickle==2.2.1 23 | - numpy==1.23.5 24 | - packaging==20.9 25 | - psutil==5.9.6 26 | - pyyaml==6.0.1 27 | - scikit-learn==1.3.1 28 | - scipy==1.11.3 29 | resources: 30 | accelerator: null 31 | cpu: 500m 32 | memory: 512Mi 33 | use_gpu: false 34 | secrets: {} 35 | spec_version: '2.0' 36 | system_packages: [] 37 | train: 38 | resources: 39 | accelerator: null 40 | cpu: 500m 41 | memory: 512Mi 42 | use_gpu: false 43 | training_class_filename: train.py 44 | training_class_name: Train 45 | training_module_dir: train 46 | variables: {} 47 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.pth 3 | *.code-workspace 4 | backend/*.sqlite3 5 | .DS_Store 6 | .env 7 | .env.* 8 | backend/baseten/settings/local.py 9 | /staticfiles/* 10 | /mediafiles/* 11 | __pycache__/ 12 | /.vscode/* 13 | !/.vscode/settings.json 14 | 15 | # coverage result 16 | .coverage 17 | /coverage/ 18 | 19 | # pycharm 20 | .idea 21 | 22 | # data 23 | *.dump 24 | 25 | # npm 26 | node_modules/ 27 | npm-debug.log 28 | 29 | #Folders 30 | .DS_Store 31 | 32 | # Webpack 33 | /frontend/bundles/* 34 | /frontend/webpack_bundles/* 35 | /webpack-stats.json 36 | /webpack-stats.dev.json 37 | 38 | # Cypress 39 | /cypress/screenshots/ 40 | /cypress/videos/ 41 | __diff_output__ 42 | 43 | # Sass 44 | .sass-cache 45 | *.map 46 | 47 | # General 48 | /venv/ 49 | /env/ 50 | /output/ 51 | /cache/ 52 | /dist/ 53 | 54 | # Spritesmith 55 | spritesmith-generated/ 56 | spritesmith.scss 57 | 58 | # templated email 59 | tmp_email/ 60 | 61 | # process ids 62 | *.pid 63 | backend/celerybeat-schedule 64 | 65 | backend/staticfiles/**/* 66 | 67 | backend/cov_html/* 68 | k8s/on-prem/deps/* 69 | 70 | # yarn lock file 71 | yarn.lock 72 | 73 | .eslintcache 74 | /.pytest_cache/ 75 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/environments_inference/test_requirements_inference.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | import pandas as pd # noqa 4 | from truss.build import create 5 | 6 | 7 | def test_infer_deps_through_create_local_import(sklearn_rfc_model, tmp_path): 8 | dir_path = tmp_path / "truss" 9 | import requests # noqa 10 | 11 | tr = create( 12 | sklearn_rfc_model, 13 | target_directory=dir_path, 14 | ) 15 | spec = tr.spec 16 | _validate_that_package_is_in_requirements(spec.requirements, "requests") 17 | 18 | 19 | def test_infer_deps_through_create_top_of_the_file_import(sklearn_rfc_model, tmp_path): 20 | dir_path = tmp_path / "truss" 21 | tr = create( 22 | sklearn_rfc_model, 23 | target_directory=dir_path, 24 | ) 25 | spec = tr.spec 26 | _validate_that_package_is_in_requirements(spec.requirements, "pandas") 27 | 28 | 29 | def _validate_that_package_is_in_requirements( 30 | requirements_list: List[str], 31 | package_name: str, 32 | ): 33 | requirement_entries = [ 34 | req for req in requirements_list if req.startswith(package_name) 35 | ] 36 | assert len(requirement_entries) == 1 37 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/model_frameworks/xgboost.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from typing import Any, Dict, Set 3 | 4 | from truss.constants import XGBOOST_REQ_MODULE_NAMES 5 | from truss.model_framework import ModelFramework 6 | from truss.types import ModelFrameworkType 7 | 8 | MODEL_FILENAME = "model.ubj" 9 | 10 | 11 | class XGBoost(ModelFramework): 12 | def typ(self) -> ModelFrameworkType: 13 | return ModelFrameworkType.XGBOOST 14 | 15 | def required_python_depedencies(self) -> Set[str]: 16 | return XGBOOST_REQ_MODULE_NAMES 17 | 18 | def serialize_model_to_directory(self, model, target_directory: Path): 19 | model_filename = MODEL_FILENAME 20 | model_filepath = target_directory / model_filename 21 | model.save_model(model_filepath) 22 | 23 | def model_metadata(self, model) -> Dict[str, Any]: 24 | return { 25 | "model_binary_dir": "model", 26 | "supports_predict_proba": False, 27 | } 28 | 29 | def supports_model_class(self, model_class) -> bool: 30 | model_framework, _, _ = model_class.__module__.partition(".") 31 | return model_framework == ModelFrameworkType.XGBOOST.value 32 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | import os 4 | try: 5 | BUILD_DIR = os.environ['BUILD_DIR'] 6 | print('using the env var', BUILD_DIR) 7 | except: 8 | print('using the current directory as root for build') 9 | BUILD_DIR = f'{os.getcwd()}' 10 | 11 | def read_req_file(req_type): 12 | with open(f"{BUILD_DIR}/requires-{req_type}.txt") as fp: # pylint: disable=W1514 13 | requires = (line.strip() for line in fp) 14 | return [req for req in requires if req and not req.startswith("#")] 15 | 16 | # Read the contents of the README file 17 | with open("README.md", "r") as fh: 18 | long_description = fh.read() 19 | 20 | 21 | 22 | setup( 23 | name="model_to_docker", 24 | version="0.0.1", 25 | author="eff-kay", 26 | author_email="faiizan14@gmail.com", 27 | description=( 28 | "A unified interface to convert pre-trained ML models to Docker images" "Developed by SlashML." 29 | ), 30 | long_description=long_description, # Use the contents of the README file 31 | long_description_content_type="text/markdown", # Set the type of the README file 32 | packages=find_packages("."), 33 | package_dir={"": "."}, 34 | install_requires=read_req_file("install"), 35 | license="GNU GPLv3", 36 | python_requires=">=3.9", 37 | ) 38 | -------------------------------------------------------------------------------- /examples/test_model/model/model.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | MODEL_BASENAME = "model" 4 | MODEL_EXTENSIONS = [".joblib", ".pkl", ".pickle"] 5 | 6 | 7 | class Model: 8 | def __init__(self, **kwargs) -> None: 9 | self._data_dir = kwargs["data_dir"] 10 | config = kwargs["config"] 11 | model_metadata = config["model_metadata"] 12 | self._supports_predict_proba = model_metadata["supports_predict_proba"] 13 | self._model_binary_dir = model_metadata["model_binary_dir"] 14 | self._model = None 15 | 16 | def load(self): 17 | import joblib 18 | 19 | model_binary_dir_path = self._data_dir / self._model_binary_dir 20 | paths = [ 21 | (model_binary_dir_path / MODEL_BASENAME).with_suffix(model_extension) 22 | for model_extension in MODEL_EXTENSIONS 23 | ] 24 | model_file_path = next(path for path in paths if path.exists()) 25 | self._model = joblib.load(model_file_path) 26 | 27 | def predict(self, model_input: Any) -> Any: 28 | model_output = {} 29 | result = self._model.predict(model_input) 30 | model_output["predictions"] = result 31 | if self._supports_predict_proba: 32 | model_output["probabilities"] = self._model.predict_proba( 33 | model_input 34 | ).tolist() 35 | return model_output 36 | -------------------------------------------------------------------------------- /examples/test_model_1/model/model.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | MODEL_BASENAME = "model" 4 | MODEL_EXTENSIONS = [".joblib", ".pkl", ".pickle"] 5 | 6 | 7 | class Model: 8 | def __init__(self, **kwargs) -> None: 9 | self._data_dir = kwargs["data_dir"] 10 | config = kwargs["config"] 11 | model_metadata = config["model_metadata"] 12 | self._supports_predict_proba = model_metadata["supports_predict_proba"] 13 | self._model_binary_dir = model_metadata["model_binary_dir"] 14 | self._model = None 15 | 16 | def load(self): 17 | import joblib 18 | 19 | model_binary_dir_path = self._data_dir / self._model_binary_dir 20 | paths = [ 21 | (model_binary_dir_path / MODEL_BASENAME).with_suffix(model_extension) 22 | for model_extension in MODEL_EXTENSIONS 23 | ] 24 | model_file_path = next(path for path in paths if path.exists()) 25 | self._model = joblib.load(model_file_path) 26 | 27 | def predict(self, model_input: Any) -> Any: 28 | model_output = {} 29 | result = self._model.predict(model_input) 30 | model_output["predictions"] = result 31 | if self._supports_predict_proba: 32 | model_output["probabilities"] = self._model.predict_proba( 33 | model_input 34 | ).tolist() 35 | return model_output 36 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/sklearn/model/model.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | MODEL_BASENAME = "model" 4 | MODEL_EXTENSIONS = [".joblib", ".pkl", ".pickle"] 5 | 6 | 7 | class Model: 8 | def __init__(self, **kwargs) -> None: 9 | self._data_dir = kwargs["data_dir"] 10 | config = kwargs["config"] 11 | model_metadata = config["model_metadata"] 12 | self._supports_predict_proba = model_metadata["supports_predict_proba"] 13 | self._model_binary_dir = model_metadata["model_binary_dir"] 14 | self._model = None 15 | 16 | def load(self): 17 | import joblib 18 | 19 | model_binary_dir_path = self._data_dir / self._model_binary_dir 20 | paths = [ 21 | (model_binary_dir_path / MODEL_BASENAME).with_suffix(model_extension) 22 | for model_extension in MODEL_EXTENSIONS 23 | ] 24 | model_file_path = next(path for path in paths if path.exists()) 25 | self._model = joblib.load(model_file_path) 26 | 27 | def predict(self, model_input: Any) -> Any: 28 | model_output = {} 29 | result = self._model.predict(model_input) 30 | model_output["predictions"] = result 31 | if self._supports_predict_proba: 32 | model_output["probabilities"] = self._model.predict_proba( 33 | model_input 34 | ).tolist() 35 | return model_output 36 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/model_frameworks/sklearn.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from typing import Any, Dict, Set 3 | 4 | from truss.constants import SKLEARN_REQ_MODULE_NAMES 5 | from truss.model_framework import ModelFramework 6 | from truss.templates.server.common.util import model_supports_predict_proba 7 | from truss.types import ModelFrameworkType 8 | 9 | MODEL_FILENAME = "model.joblib" 10 | 11 | 12 | class SKLearn(ModelFramework): 13 | def typ(self) -> ModelFrameworkType: 14 | return ModelFrameworkType.SKLEARN 15 | 16 | def required_python_depedencies(self) -> Set[str]: 17 | return SKLEARN_REQ_MODULE_NAMES 18 | 19 | def serialize_model_to_directory(self, model, target_directory: Path): 20 | import joblib 21 | 22 | model_filename = MODEL_FILENAME 23 | model_filepath = target_directory / model_filename 24 | joblib.dump(model, model_filepath, compress=True) 25 | 26 | def model_metadata(self, model) -> Dict[str, Any]: 27 | supports_predict_proba = model_supports_predict_proba(model) 28 | return { 29 | "model_binary_dir": "model", 30 | "supports_predict_proba": supports_predict_proba, 31 | } 32 | 33 | def supports_model_class(self, model_class) -> bool: 34 | model_framework, _, _ = model_class.__module__.partition(".") 35 | return model_framework == ModelFrameworkType.SKLEARN.value 36 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/model_frameworks/lightgbm.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from typing import Any, Dict, Set 3 | 4 | from truss.constants import LIGHTGBM_REQ_MODULE_NAMES 5 | from truss.model_framework import ModelFramework 6 | from truss.templates.server.common.util import model_supports_predict_proba 7 | from truss.types import ModelFrameworkType 8 | 9 | MODEL_FILENAME = "model.joblib" 10 | 11 | 12 | class LightGBM(ModelFramework): 13 | def typ(self) -> ModelFrameworkType: 14 | return ModelFrameworkType.LIGHTGBM 15 | 16 | def required_python_depedencies(self) -> Set[str]: 17 | return LIGHTGBM_REQ_MODULE_NAMES 18 | 19 | def serialize_model_to_directory(self, model, target_directory: Path): 20 | import joblib 21 | 22 | model_filename = MODEL_FILENAME 23 | model_filepath = target_directory / model_filename 24 | joblib.dump(model, model_filepath, compress=True) 25 | 26 | def model_metadata(self, model) -> Dict[str, Any]: 27 | supports_predict_proba = model_supports_predict_proba(model) 28 | return { 29 | "model_binary_dir": "model", 30 | "supports_predict_proba": supports_predict_proba, 31 | } 32 | 33 | def supports_model_class(self, model_class) -> bool: 34 | model_framework, _, _ = model_class.__module__.partition(".") 35 | return model_framework == ModelFrameworkType.LIGHTGBM.value 36 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/lightgbm/model/model.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | MODEL_BASENAME = "model" 4 | MODEL_EXTENSIONS = [".joblib", ".pkl", ".pickle"] 5 | 6 | 7 | class Model: 8 | def __init__(self, **kwargs) -> None: 9 | self._data_dir = kwargs["data_dir"] 10 | config = kwargs["config"] 11 | model_metadata = config["model_metadata"] 12 | # LightGBM does not implement a `predict_proba` function 13 | self._supports_predict_proba = False 14 | self._model_binary_dir = model_metadata["model_binary_dir"] 15 | self._model = None 16 | 17 | def load(self): 18 | import joblib 19 | 20 | model_binary_dir_path = self._data_dir / self._model_binary_dir 21 | paths = [ 22 | (model_binary_dir_path / MODEL_BASENAME).with_suffix(model_extension) 23 | for model_extension in MODEL_EXTENSIONS 24 | ] 25 | model_file_path = next(path for path in paths if path.exists()) 26 | self._model = joblib.load(model_file_path) 27 | 28 | def predict(self, model_input: Any) -> Any: 29 | model_output = {} 30 | result = self._model.predict(model_input) 31 | model_output["predictions"] = result 32 | if self._supports_predict_proba: 33 | model_output["probabilities"] = self._model.predict_proba( 34 | model_input 35 | ).tolist() 36 | return model_output 37 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/test_truss_gatherer.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from typing import List 3 | 4 | from truss.patch.dir_signature import directory_content_signature 5 | from truss.truss_gatherer import gather 6 | 7 | 8 | def test_gather(custom_model_with_external_package): 9 | gathered_truss_path = gather(custom_model_with_external_package) 10 | subdir = gathered_truss_path / "packages" / "subdir" 11 | sub_module = gathered_truss_path / "packages" / "subdir" / "sub_module.py" 12 | ext_pkg_top_module = gathered_truss_path / "packages" / "top_module.py" 13 | ext_pkg_top_module2 = gathered_truss_path / "packages" / "top_module2.py" 14 | assert subdir.exists() 15 | assert ext_pkg_top_module.exists() 16 | assert ext_pkg_top_module2.exists() 17 | assert sub_module.exists() 18 | sub_module.unlink() 19 | subdir.rmdir() 20 | ext_pkg_top_module.unlink() 21 | ext_pkg_top_module2.unlink() 22 | 23 | assert _same_dir_content( 24 | custom_model_with_external_package, 25 | gathered_truss_path, 26 | ["config.yaml"], 27 | ) 28 | 29 | 30 | def _same_dir_content(dir1: Path, dir2: Path, ignore_paths: List[str]) -> bool: 31 | sig1 = directory_content_signature(dir1) 32 | for path in ignore_paths: 33 | del sig1[path] 34 | sig2 = directory_content_signature(dir2) 35 | for path in ignore_paths: 36 | del sig2[path] 37 | return sig1 == sig2 38 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/pytorch/model/model.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | import torch 4 | from torch import package 5 | 6 | 7 | class Model: 8 | def __init__(self, **kwargs) -> None: 9 | self._data_dir = kwargs["data_dir"] 10 | config = kwargs["config"] 11 | self._model_metadata = config["model_metadata"] 12 | self._model = None 13 | self._model_dtype = None 14 | self._device = torch.device( 15 | "cuda:0" 16 | if torch.cuda.is_available() and config["resources"]["use_gpu"] 17 | else "cpu" 18 | ) 19 | 20 | def load(self): 21 | imp = package.PackageImporter( 22 | self._data_dir / "model" / self._model_metadata["torch_package_file"] 23 | ) 24 | package_name = self._model_metadata["torch_model_package_name"] 25 | model_pickle_filename = self._model_metadata["torch_model_pickle_filename"] 26 | self._model = imp.load_pickle(package_name, model_pickle_filename) 27 | self._model_dtype = list(self._model.parameters())[0].dtype 28 | 29 | def predict(self, model_input: Any) -> Any: 30 | model_output = {} 31 | with torch.no_grad(): 32 | inputs = torch.tensor( 33 | model_input, dtype=self._model_dtype, device=self._device 34 | ) 35 | model_output["predictions"] = self._model(inputs).tolist() 36 | return model_output 37 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/contexts/image_builder/test_serving_image_builder.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from tempfile import TemporaryDirectory 3 | 4 | from truss.contexts.image_builder.serving_image_builder import ( 5 | ServingImageBuilderContext, 6 | ) 7 | from truss.truss_handle import TrussHandle 8 | 9 | BASE_DIR = Path(__file__).parent 10 | 11 | 12 | def test_serving_image_dockerfile_from_user_base_image(custom_model_truss_dir): 13 | th = TrussHandle(custom_model_truss_dir) 14 | th.set_base_image("baseten/truss-server-base:3.9-v0.4.3", "/usr/local/bin/python3") 15 | builder_context = ServingImageBuilderContext 16 | image_builder = builder_context.run(th.spec.truss_dir) 17 | with TemporaryDirectory() as tmp_dir: 18 | tmp_path = Path(tmp_dir) 19 | image_builder.prepare_image_build_dir(tmp_path) 20 | with open(tmp_path / "Dockerfile", "r") as f: 21 | gen_docker_lines = f.readlines() 22 | with open( 23 | f"{BASE_DIR}/../../../test_data/server.Dockerfile", 24 | "r", 25 | ) as f: 26 | server_docker_lines = f.readlines() 27 | 28 | def filter_empty_lines(lines): 29 | return list(filter(lambda x: x and x != "\n", lines)) 30 | 31 | gen_docker_lines = filter_empty_lines(gen_docker_lines) 32 | server_docker_lines = filter_empty_lines(server_docker_lines) 33 | assert gen_docker_lines == server_docker_lines 34 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/test_context_builder_image.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | from pathlib import Path 3 | 4 | import pytest 5 | 6 | 7 | @pytest.mark.integration 8 | def test_build_docker_image(): 9 | root_path = Path(__file__).parent.parent.parent 10 | root = str(root_path) 11 | context_builder_image_test_dir = str( 12 | root_path / "truss" / "test_data" / "context_builder_image_test" 13 | ) 14 | 15 | subprocess.run( 16 | [ 17 | "docker", 18 | "buildx", 19 | "build", 20 | ".", 21 | "-f", 22 | "context_builder.Dockerfile", 23 | "--platform=linux/amd64", 24 | "-t", 25 | "baseten/truss-context-builder:test", 26 | ], 27 | check=True, 28 | cwd=root, 29 | ) 30 | 31 | subprocess.run( 32 | [ 33 | "docker", 34 | "buildx", 35 | "build", 36 | context_builder_image_test_dir, 37 | "--platform=linux/amd64", 38 | "-t", 39 | "baseten/truss-context-builder-test", 40 | ], 41 | check=True, 42 | cwd=root, 43 | ) 44 | 45 | # This will throw if building docker build context fails 46 | subprocess.run( 47 | [ 48 | "docker", 49 | "run", 50 | "baseten/truss-context-builder-test", 51 | ], 52 | check=True, 53 | cwd=root, 54 | ) 55 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/contexts/image_builder/image_builder.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from pathlib import Path 3 | from typing import Optional 4 | 5 | from truss.docker import Docker 6 | from truss.util.path import given_or_temporary_dir 7 | 8 | 9 | class ImageBuilder(ABC): 10 | def build_image( 11 | self, 12 | build_dir: Optional[Path] = None, 13 | tag: Optional[str] = None, 14 | labels: Optional[dict] = None, 15 | cache: bool = True, 16 | ): 17 | """Build image. 18 | 19 | Arguments: 20 | build_dir(Path): Directory to use for building the docker image. If None 21 | then a temporary directory is used. 22 | tag(str): A tag to assign to the docker image. 23 | """ 24 | 25 | with given_or_temporary_dir(build_dir) as build_dir_path: 26 | self.prepare_image_build_dir(build_dir_path) 27 | return Docker.client().build( 28 | str(build_dir_path), 29 | labels=labels if labels else {}, 30 | tags=tag or self.default_tag, 31 | cache=cache, 32 | ) 33 | 34 | @property 35 | @abstractmethod 36 | def default_tag(self): 37 | pass 38 | 39 | @abstractmethod 40 | def prepare_image_build_dir(self, build_dir: Optional[Path] = None): 41 | pass 42 | 43 | def docker_build_command(self, build_dir) -> str: 44 | return f"docker build {build_dir} -t {self.default_tag}" 45 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/server/common/logging.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import sys 3 | 4 | from pythonjsonlogger import jsonlogger 5 | 6 | LEVEL: int = logging.INFO 7 | 8 | JSON_LOG_HANDLER = logging.StreamHandler(stream=sys.stdout) 9 | JSON_LOG_HANDLER.set_name("json_logger_handler") 10 | JSON_LOG_HANDLER.setLevel(LEVEL) 11 | JSON_LOG_HANDLER.setFormatter( 12 | jsonlogger.JsonFormatter("%(asctime)s %(levelname)s %(message)s") 13 | ) 14 | 15 | 16 | class HealthCheckFilter(logging.Filter): 17 | def filter(self, record: logging.LogRecord) -> bool: 18 | # for any health check endpoints, lets skip logging 19 | return ( 20 | record.getMessage().find("GET / ") == -1 21 | and record.getMessage().find("GET /v1/models/model ") == -1 22 | ) 23 | 24 | 25 | def setup_logging() -> None: 26 | loggers = [logging.getLogger()] + [ 27 | logging.getLogger(name) for name in logging.root.manager.loggerDict 28 | ] 29 | 30 | for logger in loggers: 31 | logger.setLevel(LEVEL) 32 | logger.propagate = False 33 | 34 | setup = False 35 | 36 | # let's not thrash the handlers unnecessarily 37 | for handler in logger.handlers: 38 | if handler.name == JSON_LOG_HANDLER.name: 39 | setup = True 40 | 41 | if not setup: 42 | logger.handlers.clear() 43 | logger.addHandler(JSON_LOG_HANDLER) 44 | 45 | # some special handling for request logging 46 | if logger.name == "uvicorn.access": 47 | logger.addFilter(HealthCheckFilter()) 48 | -------------------------------------------------------------------------------- /examples/xgboost/model/model.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | import xgboost as xgb 4 | 5 | MODEL_BASENAME = "model" 6 | MODEL_EXTENSIONS = [".ubj"] 7 | 8 | 9 | class Model: 10 | def __init__(self, **kwargs) -> None: 11 | self._data_dir = kwargs["data_dir"] 12 | config = kwargs["config"] 13 | model_metadata = config["model_metadata"] 14 | # XGBoost models, saved and loaded via the native save/load 15 | # in XGBoost do not support predicting probabilities unless 16 | # they are a multi-class classification problem. 17 | # TODO: Integrate model_metadata field to determine model 18 | # objective function to determine if an XGBoost model 19 | # supports predicting probabilities. 20 | self._supports_predict_proba = False 21 | self._model_binary_dir = model_metadata["model_binary_dir"] 22 | self._model = None 23 | 24 | def load(self): 25 | model_binary_dir_path = self._data_dir / self._model_binary_dir 26 | paths = [ 27 | (model_binary_dir_path / MODEL_BASENAME).with_suffix(model_extension) 28 | for model_extension in MODEL_EXTENSIONS 29 | ] 30 | model_file_path = next(path for path in paths if path.exists()) 31 | self._model = xgb.Booster() 32 | self._model.load_model(model_file_path) 33 | 34 | def predict(self, model_input: Any) -> Any: 35 | model_output = {} 36 | dmatrix_inputs = xgb.DMatrix(model_input) 37 | result = self._model.predict(dmatrix_inputs) 38 | model_output["predictions"] = result 39 | return model_output 40 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/server/common/patches.py: -------------------------------------------------------------------------------- 1 | import importlib.util 2 | import logging 3 | from pathlib import Path 4 | 5 | # Set up logging 6 | logging.basicConfig(level=logging.INFO) 7 | logger = logging.getLogger(__name__) 8 | 9 | 10 | def apply_patches(enabled: bool, requirements: list): 11 | """ 12 | Apply patches to certain functions. The patches are contained in each patch module under 'patches' directory. 13 | If a patch cannot be applied, it logs the name of the function and the exception details. 14 | """ 15 | PATCHES_DIR = Path(__file__).parent / "patches" 16 | if not enabled: 17 | return 18 | for requirement in requirements: 19 | for patch_name in PATCHES_DIR.iterdir(): 20 | if patch_name.name in requirement: 21 | try: 22 | patch_file = PATCHES_DIR / patch_name / "patch.py" 23 | if patch_file.exists(): 24 | spec = importlib.util.spec_from_file_location( 25 | f"{patch_name}_patch", patch_file 26 | ) 27 | patch_module = importlib.util.module_from_spec(spec) # type: ignore 28 | spec.loader.exec_module(patch_module) # type: ignore 29 | patch_module.patch() # Apply the patch 30 | logger.info(f"{patch_name} patch applied successfully") 31 | except Exception as e: 32 | logger.debug( 33 | f"{patch_name} patch could not be applied. Exception: {str(e)}" 34 | ) 35 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/xgboost/model/model.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | import xgboost as xgb 4 | 5 | MODEL_BASENAME = "model" 6 | MODEL_EXTENSIONS = [".ubj"] 7 | 8 | 9 | class Model: 10 | def __init__(self, **kwargs) -> None: 11 | self._data_dir = kwargs["data_dir"] 12 | config = kwargs["config"] 13 | model_metadata = config["model_metadata"] 14 | # XGBoost models, saved and loaded via the native save/load 15 | # in XGBoost do not support predicting probabilities unless 16 | # they are a multi-class classification problem. 17 | # TODO: Integrate model_metadata field to determine model 18 | # objective function to determine if an XGBoost model 19 | # supports predicting probabilities. 20 | self._supports_predict_proba = False 21 | self._model_binary_dir = model_metadata["model_binary_dir"] 22 | self._model = None 23 | 24 | def load(self): 25 | model_binary_dir_path = self._data_dir / self._model_binary_dir 26 | paths = [ 27 | (model_binary_dir_path / MODEL_BASENAME).with_suffix(model_extension) 28 | for model_extension in MODEL_EXTENSIONS 29 | ] 30 | model_file_path = next(path for path in paths if path.exists()) 31 | self._model = xgb.Booster() 32 | self._model.load_model(model_file_path) 33 | 34 | def predict(self, model_input: Any) -> Any: 35 | model_output = {} 36 | dmatrix_inputs = xgb.DMatrix(model_input) 37 | result = self._model.predict(dmatrix_inputs) 38 | model_output["predictions"] = result 39 | return model_output 40 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/shared/secrets_resolver.py: -------------------------------------------------------------------------------- 1 | import os 2 | from collections.abc import Mapping 3 | from pathlib import Path 4 | from typing import Dict 5 | 6 | 7 | class SecretsResolver: 8 | SECRETS_MOUNT_DIR = "/secrets" 9 | SECRET_ENV_VAR_PREFIX = "TRUSS_SECRET_" 10 | 11 | @staticmethod 12 | def get_secrets(config: Dict): 13 | return Secrets(config.get("secrets", {})) 14 | 15 | @staticmethod 16 | def _resolve_secret(secret_name: str, default_value: str): 17 | secret_value = default_value 18 | secret_env_var_name = SecretsResolver.SECRET_ENV_VAR_PREFIX + secret_name 19 | if secret_env_var_name in os.environ: 20 | secret_value = os.environ[secret_env_var_name] 21 | secret_path = SecretsResolver._secrets_mount_dir_path() / secret_name 22 | if secret_path.exists() and secret_path.is_file(): 23 | with secret_path.open() as secret_file: 24 | secret_value = secret_file.read() 25 | return secret_value 26 | 27 | @staticmethod 28 | def _secrets_mount_dir_path(): 29 | return Path(SecretsResolver.SECRETS_MOUNT_DIR) 30 | 31 | 32 | class Secrets(Mapping): 33 | def __init__(self, base_secrets: Dict[str, str]): 34 | self._base_secrets = base_secrets 35 | 36 | def __getitem__(self, key: str) -> str: 37 | return SecretsResolver._resolve_secret(key, self._base_secrets[key]) 38 | 39 | def __iter__(self): 40 | raise NotImplementedError( 41 | "Secrets are meant for lookup and can't be iterated on" 42 | ) 43 | 44 | def __len__(self): 45 | return len(self._base_secrets) 46 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/server/common/util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | import psutil 5 | 6 | 7 | def model_supports_predict_proba(model: object) -> bool: 8 | if not hasattr(model, "predict_proba"): 9 | return False 10 | if hasattr( 11 | model, "_check_proba" 12 | ): # noqa eg Support Vector Machines *can* predict proba if they made certain choices while training 13 | try: 14 | model._check_proba() 15 | return True 16 | except AttributeError: 17 | return False 18 | return True 19 | 20 | 21 | def cpu_count(): 22 | """Get the available CPU count for this system. 23 | Takes the minimum value from the following locations: 24 | - Total system cpus available on the host. 25 | - CPU Affinity (if set) 26 | - Cgroups limit (if set) 27 | """ 28 | count = os.cpu_count() 29 | 30 | # Check CPU affinity if available 31 | try: 32 | affinity_count = len(psutil.Process().cpu_affinity()) 33 | if affinity_count > 0: 34 | count = min(count, affinity_count) 35 | except Exception: 36 | pass 37 | 38 | # Check cgroups if available 39 | if sys.platform == "linux": 40 | try: 41 | with open("/sys/fs/cgroup/cpu,cpuacct/cpu.cfs_quota_us") as f: 42 | quota = int(f.read()) 43 | with open("/sys/fs/cgroup/cpu,cpuacct/cpu.cfs_period_us") as f: 44 | period = int(f.read()) 45 | cgroups_count = int(quota / period) 46 | if cgroups_count > 0: 47 | count = min(count, cgroups_count) 48 | except Exception: 49 | pass 50 | 51 | return count 52 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | fail_fast: true 2 | repos: 3 | - repo: https://github.com/pre-commit/pre-commit-hooks 4 | rev: v4.4.0 5 | hooks: 6 | - id: check-added-large-files 7 | args: ["--maxkb=500"] 8 | - id: trailing-whitespace 9 | - id: end-of-file-fixer 10 | - id: check-yaml 11 | - id: fix-byte-order-marker 12 | - id: check-case-conflict 13 | - id: check-merge-conflict 14 | - id: check-symlinks 15 | - id: debug-statements 16 | - repo: local 17 | hooks: 18 | - id: isort 19 | name: isort-local 20 | entry: poetry run isort 21 | language: system 22 | types: [python] 23 | pass_filenames: true 24 | - id: flake8 25 | name: flake8-local 26 | entry: poetry run flake8 27 | language: system 28 | types: [python] 29 | pass_filenames: true 30 | - id: mypy 31 | name: mypy-local 32 | entry: poetry run mypy 33 | args: ["--install-types", "--non-interactive"] 34 | language: python 35 | types: [python] 36 | exclude: ^examples/|^truss/test.+/|model.py$ 37 | pass_filenames: true 38 | - repo: https://github.com/tcort/markdown-link-check 39 | rev: "v3.10.3" 40 | hooks: 41 | - id: markdown-link-check 42 | - repo: https://github.com/psf/black 43 | rev: 22.10.0 44 | hooks: 45 | - id: black 46 | # It is recommended to specify the latest version of Python 47 | # supported by your project here, or alternatively use 48 | # pre-commit's default_language_version, see 49 | # https://pre-commit.com/#top_level-default_language_version 50 | language_version: python3.9 51 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/test_testing_utilities_for_other_tests.py: -------------------------------------------------------------------------------- 1 | # This file contains shared code to be used in other tests 2 | # TODO(pankaj): Using a tests file for shared code is not ideal, we should 3 | # move it to a regular file. This is a short term hack. 4 | 5 | import shutil 6 | import subprocess 7 | import time 8 | from contextlib import contextmanager 9 | 10 | from truss.build import kill_all 11 | from truss.constants import TRUSS 12 | from truss.docker import get_containers 13 | 14 | DISK_SPACE_LOW_PERCENTAGE = 20 15 | 16 | 17 | @contextmanager 18 | def ensure_kill_all(): 19 | try: 20 | yield 21 | finally: 22 | kill_all_with_retries() 23 | ensure_free_disk_space() 24 | 25 | 26 | def kill_all_with_retries(num_retries: int = 10): 27 | kill_all() 28 | attempts = 0 29 | while attempts < num_retries: 30 | containers = get_containers({TRUSS: True}) 31 | if len(containers) == 0: 32 | return 33 | attempts += 1 34 | time.sleep(1) 35 | 36 | 37 | def ensure_free_disk_space(): 38 | """Check if disk space is low.""" 39 | if is_disk_space_low(): 40 | clear_disk_space() 41 | 42 | 43 | def is_disk_space_low() -> bool: 44 | disk_usage = shutil.disk_usage("/") 45 | disk_free_percent = 100 * float(disk_usage.free) / disk_usage.total 46 | return disk_free_percent <= DISK_SPACE_LOW_PERCENTAGE 47 | 48 | 49 | def clear_disk_space(): 50 | docker_ps_output = subprocess.check_output( 51 | ["docker", "ps", "-a", "-f", "status=exited", "-q"] 52 | ).decode("utf-8") 53 | docker_containers = docker_ps_output.split("\n")[:-1] 54 | subprocess.run(["docker", "rm", *docker_containers]) 55 | subprocess.run(["docker", "system", "prune", "-a", "-f"]) 56 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/templates/core/server/test_secrets_resolver.py: -------------------------------------------------------------------------------- 1 | import os 2 | from contextlib import contextmanager 3 | from pathlib import Path 4 | 5 | from truss.templates.shared.secrets_resolver import SecretsResolver 6 | 7 | CONFIG = {"secrets": {"secret_key": "default_secret_value"}} 8 | 9 | 10 | def test_resolve_default_value(): 11 | secrets = SecretsResolver.get_secrets(CONFIG) 12 | assert secrets["secret_key"] == "default_secret_value" 13 | 14 | 15 | def test_resolve_env_var(): 16 | secrets = SecretsResolver.get_secrets(CONFIG) 17 | with _override_env_var("TRUSS_SECRET_secret_key", "secret_value_from_env"): 18 | assert secrets["secret_key"] == "secret_value_from_env" 19 | 20 | 21 | def test_resolve_mounted_secrets(tmp_path): 22 | secrets = SecretsResolver.get_secrets(CONFIG) 23 | with (tmp_path / "secret_key").open("w") as f: 24 | f.write("secret_value_from_mounted_secrets") 25 | with _secrets_mount_dir(tmp_path), _override_env_var( 26 | "TRUSS_SECRET_secret_key", "secret_value_from_env" 27 | ): 28 | assert secrets["secret_key"] == "secret_value_from_mounted_secrets" 29 | 30 | 31 | @contextmanager 32 | def _secrets_mount_dir(path: Path): 33 | orig_secrets_mount_dir = SecretsResolver.SECRETS_MOUNT_DIR 34 | SecretsResolver.SECRETS_MOUNT_DIR = str(path) 35 | try: 36 | yield 37 | finally: 38 | SecretsResolver.SECRETS_MOUNT_DIR = orig_secrets_mount_dir 39 | 40 | 41 | @contextmanager 42 | def _override_env_var(key: str, value: str): 43 | orig_value = os.environ.get(key, None) 44 | try: 45 | os.environ[key] = value 46 | yield 47 | finally: 48 | if orig_value is not None: 49 | os.environ[key] = orig_value 50 | else: 51 | del os.environ[key] 52 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/model_frameworks/__init__.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | from truss.errors import FrameworkNotSupportedError 4 | from truss.model_framework import ModelFramework 5 | from truss.model_frameworks.huggingface_transformer import HuggingfaceTransformer 6 | from truss.model_frameworks.keras import Keras 7 | from truss.model_frameworks.lightgbm import LightGBM 8 | from truss.model_frameworks.mlflow import Mlflow 9 | from truss.model_frameworks.pytorch import PyTorch 10 | from truss.model_frameworks.sklearn import SKLearn 11 | from truss.model_frameworks.xgboost import XGBoost 12 | from truss.types import ModelFrameworkType 13 | 14 | MODEL_FRAMEWORKS_BY_TYPE = { 15 | ModelFrameworkType.SKLEARN: SKLearn(), 16 | ModelFrameworkType.KERAS: Keras(), 17 | ModelFrameworkType.HUGGINGFACE_TRANSFORMER: HuggingfaceTransformer(), 18 | ModelFrameworkType.PYTORCH: PyTorch(), 19 | ModelFrameworkType.XGBOOST: XGBoost(), 20 | ModelFrameworkType.LIGHTGBM: LightGBM(), 21 | ModelFrameworkType.MLFLOW: Mlflow(), 22 | } 23 | 24 | 25 | SUPPORTED_MODEL_FRAMEWORKS = [ 26 | ModelFrameworkType.SKLEARN, 27 | ModelFrameworkType.KERAS, 28 | ModelFrameworkType.TENSORFLOW, 29 | ModelFrameworkType.HUGGINGFACE_TRANSFORMER, 30 | ModelFrameworkType.XGBOOST, 31 | ModelFrameworkType.PYTORCH, 32 | ModelFrameworkType.LIGHTGBM, 33 | ] 34 | 35 | 36 | def model_framework_from_model(model: Any) -> ModelFramework: 37 | return model_framework_from_model_class(model.__class__) 38 | 39 | 40 | def model_framework_from_model_class(model_class) -> ModelFramework: 41 | for model_framework in MODEL_FRAMEWORKS_BY_TYPE.values(): 42 | if model_framework.supports_model_class(model_class): 43 | return model_framework 44 | 45 | raise FrameworkNotSupportedError( 46 | "Model must be one of " 47 | + "/".join([t.value for t in SUPPORTED_MODEL_FRAMEWORKS]) 48 | ) 49 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/contexts/local_loader/load_model_local.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | from pathlib import Path 3 | 4 | from truss.contexts.local_loader.truss_module_loader import truss_module_loaded 5 | from truss.contexts.local_loader.utils import ( 6 | prepare_secrets, 7 | signature_accepts_keyword_arg, 8 | ) 9 | from truss.contexts.truss_context import TrussContext 10 | from truss.templates.server.common.patches import apply_patches 11 | from truss.truss_spec import TrussSpec 12 | 13 | 14 | class LoadModelLocal(TrussContext): 15 | """Loads a Truss model locally. 16 | 17 | The loaded model can be used to make predictions for quick testing. 18 | Runs in the current pip environment directly. Assumes all requirements and 19 | system packages are already installed. 20 | """ 21 | 22 | @staticmethod 23 | def run(truss_dir: Path): 24 | spec = TrussSpec(truss_dir) 25 | with truss_module_loaded( 26 | str(truss_dir), 27 | spec.model_module_fullname, 28 | spec.bundled_packages_dir.name, 29 | [str(path.resolve()) for path in spec.external_package_dirs_paths], 30 | ) as module: 31 | model_class = getattr(module, spec.model_class_name) 32 | model_class_signature = inspect.signature(model_class) 33 | model_init_params = {} 34 | if signature_accepts_keyword_arg(model_class_signature, "config"): 35 | model_init_params["config"] = spec.config.to_dict() 36 | if signature_accepts_keyword_arg(model_class_signature, "data_dir"): 37 | model_init_params["data_dir"] = truss_dir / "data" 38 | if signature_accepts_keyword_arg(model_class_signature, "secrets"): 39 | model_init_params["secrets"] = prepare_secrets(spec) 40 | apply_patches(spec.apply_library_patches, spec.requirements) 41 | model = model_class(**model_init_params) 42 | if hasattr(model, "load"): 43 | model.load() 44 | return model 45 | -------------------------------------------------------------------------------- /examples/bert_uncased/model/model.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import traceback 3 | from typing import Any 4 | 5 | import torch 6 | from transformers import pipeline 7 | 8 | 9 | class Model: 10 | def __init__(self, **kwargs) -> None: 11 | self._data_dir = kwargs["data_dir"] 12 | config = kwargs["config"] 13 | self._config = config 14 | model_metadata = config["model_metadata"] 15 | self._transformer_config = model_metadata["transformer_config"] 16 | self._has_named_args = model_metadata["has_named_args"] 17 | self._has_hybrid_args = model_metadata["has_hybrid_args"] 18 | self._model = None 19 | 20 | def load(self): 21 | transformer_config = self._transformer_config.copy() 22 | if torch.cuda.is_available(): 23 | transformer_config["device"] = 0 24 | 25 | self._model = pipeline( 26 | task=self._config["model_type"], 27 | **transformer_config, 28 | ) 29 | 30 | def predict(self, model_input: Any) -> Any: 31 | model_output = {} 32 | instances = model_input 33 | 34 | with torch.no_grad(): 35 | if self._has_named_args: 36 | result = [self._model(**instance) for instance in instances] 37 | elif self._has_hybrid_args: 38 | try: 39 | result = [] 40 | for instance in instances: 41 | prompt = instance.pop("prompt") 42 | result.append(self._model(prompt, **instance)) 43 | except (KeyError, AttributeError): 44 | logging.error(traceback.format_exc()) 45 | model_output["error"] = { 46 | "traceback": f'Expected request as an object with text in "prompt"\n{traceback.format_exc()}' 47 | } 48 | return model_output 49 | else: 50 | result = self._model(instances) 51 | model_output["predictions"] = result 52 | return model_output 53 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/huggingface_transformer/model/model.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import traceback 3 | from typing import Any 4 | 5 | import torch 6 | from transformers import pipeline 7 | 8 | 9 | class Model: 10 | def __init__(self, **kwargs) -> None: 11 | self._data_dir = kwargs["data_dir"] 12 | config = kwargs["config"] 13 | self._config = config 14 | model_metadata = config["model_metadata"] 15 | self._transformer_config = model_metadata["transformer_config"] 16 | self._has_named_args = model_metadata["has_named_args"] 17 | self._has_hybrid_args = model_metadata["has_hybrid_args"] 18 | self._model = None 19 | 20 | def load(self): 21 | transformer_config = self._transformer_config.copy() 22 | if torch.cuda.is_available(): 23 | transformer_config["device"] = 0 24 | 25 | self._model = pipeline( 26 | task=self._config["model_type"], 27 | **transformer_config, 28 | ) 29 | 30 | def predict(self, model_input: Any) -> Any: 31 | model_output = {} 32 | instances = model_input 33 | 34 | with torch.no_grad(): 35 | if self._has_named_args: 36 | result = [self._model(**instance) for instance in instances] 37 | elif self._has_hybrid_args: 38 | try: 39 | result = [] 40 | for instance in instances: 41 | prompt = instance.pop("prompt") 42 | result.append(self._model(prompt, **instance)) 43 | except (KeyError, AttributeError): 44 | logging.error(traceback.format_exc()) 45 | model_output["error"] = { 46 | "traceback": f'Expected request as an object with text in "prompt"\n{traceback.format_exc()}' 47 | } 48 | return model_output 49 | else: 50 | result = self._model(instances) 51 | model_output["predictions"] = result 52 | return model_output 53 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/control/control/server.py: -------------------------------------------------------------------------------- 1 | import os 2 | from pathlib import Path 3 | from threading import Thread 4 | from typing import Union 5 | 6 | from application import create_app 7 | from helpers.inference_server_starter import inference_server_startup_flow 8 | from waitress.server import BaseWSGIServer, MultiSocketServer 9 | 10 | CONTROL_SERVER_PORT = int(os.environ.get("CONTROL_SERVER_PORT", "8080")) 11 | INFERENCE_SERVER_PORT = int(os.environ.get("INFERENCE_SERVER_PORT", "8090")) 12 | PYTHON_EXECUTABLE_LOOKUP_PATHS = [ 13 | "/usr/local/bin/python", 14 | "/usr/local/bin/python3", 15 | "/usr/bin/python", 16 | "/usr/bin/python3", 17 | ] 18 | 19 | 20 | def _identify_python_executable_path() -> str: 21 | for path in PYTHON_EXECUTABLE_LOOKUP_PATHS: 22 | if Path(path).exists(): 23 | return path 24 | 25 | raise RuntimeError("Unable to find python, make sure it's installed.") 26 | 27 | 28 | if __name__ == "__main__": 29 | from waitress import create_server 30 | 31 | inf_serv_home: str = os.environ["APP_HOME"] 32 | python_executable_path: str = _identify_python_executable_path() 33 | application = create_app( 34 | { 35 | "inference_server_home": inf_serv_home, 36 | "inference_server_process_args": [ 37 | python_executable_path, 38 | f"{inf_serv_home}/inference_server.py", 39 | ], 40 | "control_server_host": "*", 41 | "control_server_port": CONTROL_SERVER_PORT, 42 | "inference_server_port": INFERENCE_SERVER_PORT, 43 | } 44 | ) 45 | 46 | # Perform inference server startup flow in background 47 | Thread(target=inference_server_startup_flow, args=(application,)).start() 48 | 49 | application.logger.info( 50 | f"Starting live reload server on port {CONTROL_SERVER_PORT}" 51 | ) 52 | server: Union[BaseWSGIServer, MultiSocketServer] = create_server( 53 | application, 54 | host=application.config["control_server_host"], 55 | port=application.config["control_server_port"], 56 | ) 57 | server.run() 58 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/test_backward.py: -------------------------------------------------------------------------------- 1 | from truss.build import from_directory, mk_truss 2 | 3 | 4 | def test_mk_truss_passthrough(sklearn_rfc_model, tmp_path): 5 | dir_path = tmp_path / "truss" 6 | data_file_path = tmp_path / "data.txt" 7 | with data_file_path.open("w") as data_file: 8 | data_file.write("test") 9 | req_file_path = tmp_path / "requirements.txt" 10 | requirements = [ 11 | "tensorflow==2.3.1", 12 | "uvicorn==0.12.2", 13 | ] 14 | with req_file_path.open("w") as req_file: 15 | for req in requirements: 16 | req_file.write(f"{req}\n") 17 | scaf = mk_truss( 18 | sklearn_rfc_model, 19 | target_directory=dir_path, 20 | data_files=[str(data_file_path)], 21 | requirements_file=str(req_file_path), 22 | ) 23 | spec = scaf.spec 24 | assert spec.model_module_dir.exists() 25 | assert spec.truss_dir == dir_path 26 | assert spec.config_path.exists() 27 | assert spec.data_dir.exists() 28 | assert (spec.data_dir / "data.txt").exists() 29 | assert spec.requirements == requirements 30 | 31 | 32 | def test_from_directory_passthrough(sklearn_rfc_model, tmp_path): 33 | dir_path = tmp_path / "truss" 34 | data_file_path = tmp_path / "data.txt" 35 | with data_file_path.open("w") as data_file: 36 | data_file.write("test") 37 | req_file_path = tmp_path / "requirements.txt" 38 | requirements = [ 39 | "tensorflow==2.3.1", 40 | "uvicorn==0.12.2", 41 | ] 42 | with req_file_path.open("w") as req_file: 43 | for req in requirements: 44 | req_file.write(f"{req}\n") 45 | mk_truss( 46 | sklearn_rfc_model, 47 | target_directory=dir_path, 48 | data_files=[str(data_file_path)], 49 | requirements_file=str(req_file_path), 50 | ) 51 | scaf = from_directory(dir_path) 52 | spec = scaf.spec 53 | assert spec.model_module_dir.exists() 54 | assert spec.truss_dir == dir_path 55 | assert spec.config_path.exists() 56 | assert spec.data_dir.exists() 57 | assert (spec.data_dir / "data.txt").exists() 58 | assert spec.requirements == requirements 59 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/model_frameworks/huggingface_transformer.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from pathlib import Path 3 | from typing import Any, Dict, Optional, Set 4 | 5 | from truss.constants import HUGGINGFACE_TRANSFORMER_MODULE_NAME 6 | from truss.model_framework import ModelFramework 7 | from truss.types import ModelFrameworkType 8 | 9 | logger: logging.Logger = logging.getLogger(__name__) 10 | 11 | 12 | class HuggingfaceTransformer(ModelFramework): 13 | def typ(self) -> ModelFrameworkType: 14 | return ModelFrameworkType.HUGGINGFACE_TRANSFORMER 15 | 16 | def required_python_depedencies(self) -> Set[str]: 17 | return HUGGINGFACE_TRANSFORMER_MODULE_NAME 18 | 19 | def serialize_model_to_directory(self, model, target_directory: Path): 20 | # For Huggingface models, all the important details are in metadata. 21 | pass 22 | 23 | def model_metadata(self, model) -> Dict[str, Any]: 24 | hf_task = self.model_type(model) 25 | return { 26 | "transformer_config": { 27 | "model": _hf_model_name(model), 28 | }, 29 | "has_hybrid_args": hf_task in {"text-generation"}, 30 | "has_named_args": hf_task in {"zero-shot-classification"}, 31 | } 32 | 33 | def model_type(self, model) -> str: 34 | return _infer_hf_task(model) 35 | 36 | def supports_model_class(self, model_class) -> bool: 37 | model_framework, _, _ = model_class.__module__.partition(".") 38 | return model_framework == "transformers" 39 | 40 | 41 | def _hf_model_name(model) -> Optional[str]: 42 | try: 43 | return model.model.config._name_or_path # type: ignore 44 | except AttributeError: 45 | logger.info( 46 | "Unable to infer a HuggingFace model on this task pipeline; transformer library will default" 47 | ) 48 | return None 49 | 50 | 51 | def _infer_hf_task(model) -> str: 52 | try: 53 | return model.task 54 | except AttributeError as error: 55 | logger.exception( 56 | "Unable to find a HuggingFace task on this object, did you call with a `Pipeline` object" 57 | ) 58 | raise error 59 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/contexts/image_builder/util.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from typing import Optional 3 | 4 | from truss import __version__ 5 | 6 | # This needs to be updated whenever new truss base images are published. This 7 | # manual nature is by design; at this point we want any new base image releases 8 | # to be reviewed via code change before landing. This value should be typically 9 | # set to the latest version of truss library. 10 | # 11 | # [IMPORTANT] Make sure all images for this version are published to dockerhub 12 | # before change to this value lands. This value is used to look for base images 13 | # when building docker image for a truss. 14 | TRUSS_BASE_IMAGE_VERSION_TAG = "v0.4.8" 15 | 16 | 17 | def file_is_empty(path: Path, ignore_hash_style_comments: bool = True) -> bool: 18 | if not path.exists(): 19 | return True 20 | 21 | with path.open() as file: 22 | for line in file.readlines(): 23 | if ignore_hash_style_comments and _is_hash_style_comment(line): 24 | continue 25 | if line.strip() != "": 26 | return False 27 | 28 | return True 29 | 30 | 31 | def file_is_not_empty(path: Path, ignore_hash_style_comments: bool = True) -> bool: 32 | return not file_is_empty(path, ignore_hash_style_comments) 33 | 34 | 35 | def _is_hash_style_comment(line: str) -> bool: 36 | return line.lstrip().startswith("#") 37 | 38 | 39 | def truss_base_image_name(job_type: str) -> str: 40 | return f"baseten/truss-{job_type}-base" 41 | 42 | 43 | def truss_base_image_tag( 44 | python_version: str, 45 | use_gpu: bool, 46 | version_tag: Optional[str] = None, 47 | ) -> str: 48 | if version_tag is None: 49 | version_tag = f"v{__version__}" 50 | 51 | base_tag = python_version 52 | if use_gpu: 53 | base_tag = f"{base_tag}-gpu" 54 | return f"{base_tag}-{version_tag}" 55 | 56 | 57 | def to_dotted_python_version(truss_python_version: str) -> str: 58 | """Converts python version string using in truss config to the conventional dotted form. 59 | 60 | e.g. py39 to 3.9 61 | """ 62 | return f"{truss_python_version[2]}.{truss_python_version[3:]}" 63 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/test_validation.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from truss.errors import ValidationError 3 | from truss.validation import ( 4 | validate_cpu_spec, 5 | validate_memory_spec, 6 | validate_secret_name, 7 | ) 8 | 9 | 10 | @pytest.mark.parametrize( 11 | "secret_name, should_error", 12 | [ 13 | (None, True), 14 | (1, True), 15 | ("", True), 16 | (".", True), 17 | ("..", True), 18 | ("a" * 253, False), 19 | ("a" * 254, True), 20 | ("-", False), 21 | ("-.", False), 22 | ("a-.", False), 23 | ("-.a", False), 24 | ("a-foo", False), 25 | ("a.foo", False), 26 | (".foo", False), 27 | ("x\\", True), 28 | ("a_b", False), 29 | ("_a", False), 30 | ("a_", False), 31 | ], 32 | ) 33 | def test_validate_secret_name(secret_name, should_error): 34 | does_error = False 35 | try: 36 | validate_secret_name(secret_name) 37 | except: # noqa 38 | does_error = True 39 | 40 | assert does_error == should_error 41 | 42 | 43 | @pytest.mark.parametrize( 44 | "cpu_spec, expected_valid", 45 | [ 46 | (None, False), 47 | ("", False), 48 | ("1", True), 49 | ("1.5", True), 50 | ("1.5m", True), 51 | (1, False), 52 | ("1m", True), 53 | ("1M", False), 54 | ("M", False), 55 | ("M1", False), 56 | ], 57 | ) 58 | def test_validate_cpu_spec(cpu_spec, expected_valid): 59 | if not expected_valid: 60 | with pytest.raises(ValidationError): 61 | validate_cpu_spec(cpu_spec) 62 | else: 63 | validate_cpu_spec(cpu_spec) 64 | 65 | 66 | @pytest.mark.parametrize( 67 | "mem_spec, expected_valid", 68 | [ 69 | (None, False), 70 | (1, False), 71 | ("1m", False), 72 | ("1k", True), 73 | ("512k", True), 74 | ("512M", True), 75 | ("1.5Gi", True), 76 | ("abc", False), 77 | ], 78 | ) 79 | def test_validate_mem_spec(mem_spec, expected_valid): 80 | if not expected_valid: 81 | with pytest.raises(ValidationError): 82 | validate_memory_spec(mem_spec) 83 | else: 84 | validate_memory_spec(mem_spec) 85 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/base.Dockerfile.jinja: -------------------------------------------------------------------------------- 1 | ARG PYVERSION={{config.python_version}} 2 | FROM --platform=linux/x86_64 {{base_image_name_and_tag}} 3 | 4 | ENV PYTHON_EXECUTABLE {{ config.base_image.python_executable_path or 'python3' }} 5 | 6 | {% block fail_fast %} 7 | RUN grep -w 'ID=debian\|ID_LIKE=debian' /etc/os-release || { echo "ERROR: Supplied base image is not a debian image"; exit 1; } 8 | RUN $PYTHON_EXECUTABLE -c "import sys; sys.exit(0) if sys.version_info.major == 3 and sys.version_info.minor >=8 and sys.version_info.minor <=11 else sys.exit(1)" \ 9 | || { echo "ERROR: Supplied base image does not have 3.8 <= python <= 3.11"; exit 1; } 10 | {% endblock %} 11 | 12 | RUN pip install --upgrade pip --no-cache-dir \ 13 | && rm -rf /root/.cache/pip 14 | 15 | {% block base_image_patch %} 16 | {% endblock %} 17 | 18 | {% if config.model_framework.value == 'huggingface_transformer' %} 19 | {% if config.resources.use_gpu %} 20 | # HuggingFace pytorch gpu support needs mkl 21 | RUN pip install mkl 22 | {% endif %} 23 | {% endif%} 24 | 25 | {% block post_base %} 26 | {% endblock %} 27 | 28 | {% block install_system_requirements %} 29 | {%- if should_install_system_requirements %} 30 | COPY ./system_packages.txt system_packages.txt 31 | RUN apt-get update && apt-get install --yes --no-install-recommends $(cat system_packages.txt) \ 32 | && apt-get autoremove -y \ 33 | && apt-get clean -y \ 34 | && rm -rf /var/lib/apt/lists/* 35 | {%- endif %} 36 | {% endblock %} 37 | 38 | 39 | {% block install_requirements %} 40 | {%- if should_install_requirements %} 41 | COPY ./requirements.txt requirements.txt 42 | RUN pip install -r requirements.txt --no-cache-dir && rm -rf /root/.cache/pip 43 | {%- endif %} 44 | {% endblock %} 45 | 46 | 47 | ENV APP_HOME /app 48 | WORKDIR $APP_HOME 49 | 50 | 51 | {% block app_copy %} 52 | {% endblock %} 53 | 54 | {% block bundled_packages_copy %} 55 | {%- if bundled_packages_dir_exists %} 56 | COPY ./{{config.bundled_packages_dir}} /packages 57 | {%- endif %} 58 | {% endblock %} 59 | 60 | 61 | {% for env_var_name, env_var_value in config.environment_variables.items() %} 62 | ENV {{ env_var_name }} {{ env_var_value }} 63 | {% endfor %} 64 | 65 | 66 | {% block run %} 67 | {% endblock %} 68 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/util/path.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import string 4 | import tempfile 5 | from contextlib import contextmanager 6 | from distutils.dir_util import copy_tree, remove_tree 7 | from distutils.file_util import copy_file 8 | from pathlib import Path 9 | from typing import List, Optional, Tuple, Union 10 | 11 | 12 | def copy_tree_path(src: Path, dest: Path) -> List[str]: 13 | return copy_tree(str(src), str(dest)) 14 | 15 | 16 | def copy_file_path(src: Path, dest: Path) -> Tuple[str, str]: 17 | return copy_file(str(src), str(dest)) 18 | 19 | 20 | def copy_tree_or_file(src: Path, dest: Path) -> Union[List[str], Tuple[str, str]]: 21 | if src.is_file(): 22 | return copy_file_path(src, dest) 23 | 24 | return copy_tree_path(src, dest) 25 | 26 | 27 | def remove_tree_path(target: Path) -> None: 28 | return remove_tree(str(target)) 29 | 30 | 31 | def get_max_modified_time_of_dir(path: Path) -> float: 32 | max_modified_time = os.path.getmtime(path) 33 | for root, dirs, files in os.walk(path): 34 | if os.path.islink(root): 35 | raise ValueError(f"Symlinks not allowed in Truss: {root}") 36 | files = [f for f in files if not f.startswith(".")] 37 | dirs[:] = [d for d in dirs if not d.startswith(".")] 38 | max_modified_time = max(max_modified_time, os.path.getmtime(root)) 39 | for file in files: 40 | max_modified_time = max( 41 | max_modified_time, os.path.getmtime(os.path.join(root, file)) 42 | ) 43 | return max_modified_time 44 | 45 | 46 | @contextmanager 47 | def given_or_temporary_dir(given_dir: Optional[Path] = None): 48 | if given_dir is not None: 49 | yield given_dir 50 | else: 51 | with tempfile.TemporaryDirectory() as temp_dir: 52 | yield Path(temp_dir) 53 | 54 | 55 | def build_truss_target_directory(stub: str) -> Path: 56 | """Builds a directory under ~/.truss/models for the purpose of creating a Truss at.""" 57 | rand_suffix = "".join(random.choices(string.ascii_uppercase + string.digits, k=6)) 58 | target_directory_path = Path( 59 | Path.home(), ".truss", "models", f"{stub}-{rand_suffix}" 60 | ) 61 | target_directory_path.mkdir(parents=True) 62 | return target_directory_path 63 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/docker/base_images/base_image.Dockerfile.jinja: -------------------------------------------------------------------------------- 1 | {% if use_gpu %} 2 | FROM nvidia/cuda:11.2.1-base-ubuntu20.04 3 | ENV CUDNN_VERSION=8.1.0.77 4 | ENV CUDA=11.2 5 | ENV LD_LIBRARY_PATH /usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH 6 | 7 | RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub && \ 8 | apt-get update && apt-get install -y --no-install-recommends \ 9 | ca-certificates \ 10 | cuda-command-line-tools-11-2 \ 11 | libcublas-11-2 \ 12 | libcublas-dev-11-2 \ 13 | libcufft-11-2 \ 14 | libcurand-11-2 \ 15 | libcusolver-11-2 \ 16 | libcusparse-11-2 \ 17 | libcudnn8=${CUDNN_VERSION}-1+cuda${CUDA} \ 18 | libgomp1 \ 19 | && \ 20 | apt-get clean && \ 21 | rm -rf /var/lib/apt/lists/* 22 | 23 | # Allow statements and log messages to immediately appear in the Knative logs 24 | ENV PYTHONUNBUFFERED True 25 | ENV DEBIAN_FRONTEND=noninteractive 26 | 27 | RUN apt update && \ 28 | apt install -y bash \ 29 | build-essential \ 30 | git \ 31 | curl \ 32 | ca-certificates \ 33 | software-properties-common && \ 34 | add-apt-repository -y ppa:deadsnakes/ppa && \ 35 | apt update -y && \ 36 | apt install -y python{{python_version}} && \ 37 | apt install -y python{{python_version}}-distutils && \ 38 | apt install -y python{{python_version}}-venv && \ 39 | apt install -y python{{python_version}}-dev && \ 40 | rm -rf /var/lib/apt/lists 41 | 42 | RUN ln -sf /usr/bin/python{{python_version}} /usr/bin/python3 && \ 43 | curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py && \ 44 | python3 get-pip.py 45 | 46 | RUN python3 -m pip install --no-cache-dir --upgrade pip 47 | 48 | {% else %} 49 | FROM python:{{python_version}} 50 | RUN apt update && apt install -y 51 | 52 | # Allow statements and log messages to immediately appear in the Knative logs 53 | ENV PYTHONUNBUFFERED True 54 | {% endif %} 55 | 56 | 57 | RUN pip install --no-cache-dir --upgrade pip \ 58 | && rm -rf /root/.cache/pip 59 | 60 | COPY ./requirements.txt requirements.txt 61 | RUN pip install --no-cache-dir -r requirements.txt \ 62 | && rm -rf /root/.cache/pip 63 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/templates/control/control/helpers/test_patch_applier.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from pathlib import Path 3 | from unittest.mock import Mock 4 | 5 | import pytest 6 | 7 | # Needed to simulate the set up on the model docker container 8 | sys.path.append( 9 | str( 10 | Path(__file__).parent.parent.parent.parent.parent.parent 11 | / "templates" 12 | / "control" 13 | / "control" 14 | ) 15 | ) 16 | 17 | # Have to use imports in this form, otherwise isinstance checks fail on helper classes 18 | from helpers.patch_applier import PatchApplier # noqa 19 | from helpers.types import Action, ModelCodePatch, Patch, PatchType # noqa 20 | 21 | 22 | @pytest.fixture 23 | def patch_applier(truss_container_fs): 24 | return PatchApplier(truss_container_fs / "app", Mock()) 25 | 26 | 27 | def test_patch_applier_add(patch_applier: PatchApplier, truss_container_fs): 28 | patch = Patch( 29 | type=PatchType.MODEL_CODE, 30 | body=ModelCodePatch( 31 | action=Action.ADD, 32 | path="dummy", 33 | content="", 34 | ), 35 | ) 36 | patch_applier.apply_patch(patch) 37 | assert (truss_container_fs / "app" / "model" / "dummy").exists() 38 | 39 | 40 | def test_patch_applier_remove(patch_applier: PatchApplier, truss_container_fs): 41 | patch = Patch( 42 | type=PatchType.MODEL_CODE, 43 | body=ModelCodePatch( 44 | action=Action.REMOVE, 45 | path="model.py", 46 | ), 47 | ) 48 | assert (truss_container_fs / "app" / "model" / "model.py").exists() 49 | patch_applier.apply_patch(patch) 50 | assert not (truss_container_fs / "app" / "model" / "model.py").exists() 51 | 52 | 53 | def test_patch_applier_update(patch_applier: PatchApplier, truss_container_fs): 54 | new_model_file_content = """ 55 | class Model: 56 | pass 57 | """ 58 | patch = Patch( 59 | type=PatchType.MODEL_CODE, 60 | body=ModelCodePatch( 61 | action=Action.UPDATE, 62 | path="model.py", 63 | content=new_model_file_content, 64 | ), 65 | ) 66 | patch_applier.apply_patch(patch) 67 | with (truss_container_fs / "app" / "model" / "model.py").open() as model_file: 68 | assert model_file.read() == new_model_file_content 69 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/templates/server/common/test_retry.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | from unittest.mock import Mock 3 | 4 | import pytest 5 | from truss.templates.server.common.retry import retry 6 | 7 | 8 | class FailForCallCount: 9 | def __init__(self, count: int) -> None: 10 | self._fail_for_call_count = count 11 | self._call_count = 0 12 | 13 | def __call__(self, *args: Any, **kwds: Any) -> Any: 14 | self._call_count += 1 15 | if self._call_count <= self._fail_for_call_count: 16 | raise RuntimeError() 17 | 18 | @property 19 | def call_count(self) -> int: 20 | return self._call_count 21 | 22 | 23 | def fail_for_call_count(count: int) -> callable: 24 | call_count = 0 25 | 26 | def inner(): 27 | nonlocal call_count 28 | 29 | return inner 30 | 31 | 32 | def test_no_fail_no_retry(): 33 | mock = Mock() 34 | mock_logging_fn = Mock() 35 | retry(mock, 3, mock_logging_fn, "") 36 | mock.assert_called_once() 37 | mock_logging_fn.assert_not_called() 38 | 39 | 40 | def test_retry_fail_once(): 41 | mock_logging_fn = Mock() 42 | fail_mock = FailForCallCount(1) 43 | retry(fail_mock, 3, mock_logging_fn, "Failed") 44 | assert fail_mock.call_count == 2 45 | assert mock_logging_fn.call_count == 1 46 | mock_logging_fn.assert_called_once_with("Failed Retrying...") 47 | 48 | 49 | def test_retry_fail_twice(): 50 | mock_logging_fn = Mock() 51 | fail_mock = FailForCallCount(2) 52 | retry(fail_mock, 3, mock_logging_fn, "Failed") 53 | assert fail_mock.call_count == 3 54 | assert mock_logging_fn.call_count == 2 55 | mock_logging_fn.assert_called_with("Failed Retrying. Retry count: 1") 56 | 57 | 58 | def test_retry_fail_twice_gap(): 59 | mock_logging_fn = Mock() 60 | fail_mock = FailForCallCount(2) 61 | retry(fail_mock, 3, mock_logging_fn, "Failed", gap_seconds=0.1) 62 | assert fail_mock.call_count == 3 63 | assert mock_logging_fn.call_count == 2 64 | mock_logging_fn.assert_called_with("Failed Retrying. Retry count: 1") 65 | 66 | 67 | def test_retry_fail_more_than_limit(): 68 | mock_logging_fn = Mock() 69 | fail_mock = FailForCallCount(3) 70 | with pytest.raises(RuntimeError): 71 | retry(fail_mock, 1, mock_logging_fn, "Failed") 72 | assert fail_mock.call_count == 2 73 | assert mock_logging_fn.call_count == 1 74 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/contexts/local_loader/docker_build_emulator.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass, field 2 | from pathlib import Path 3 | from typing import Dict, List 4 | 5 | import dockerfile 6 | from truss.util.path import copy_tree_or_file 7 | 8 | 9 | @dataclass 10 | class DockerBuildEmulatorResult: 11 | workdir: Path = field(default_factory=lambda: Path("/")) 12 | env: Dict = field(default_factory=dict) 13 | entrypoint: List = field(default_factory=list) 14 | 15 | 16 | class DockerBuildEmulator: 17 | """Emulates Docker Builds 18 | 19 | As running docker builds is expensive, this class emulates the docker build 20 | by parsing the docker file and applying certain commands to create an 21 | appropriate enviroment in a directory to simulate the root of the file system. 22 | 23 | Support COPY, ENV, ENTRYPOINT, WORKDIR commands. All other commands are ignored. 24 | """ 25 | 26 | def __init__( 27 | self, 28 | dockerfile_path: Path, 29 | context_dir: Path, 30 | ) -> None: 31 | self._commands = dockerfile.parse_file(str(dockerfile_path)) 32 | self._context_dir = context_dir 33 | 34 | def run(self, fs_root_dir: Path) -> DockerBuildEmulatorResult: 35 | def _resolve_env(key: str) -> str: 36 | if key.startswith("$"): 37 | key = key.replace("$", "", 1) 38 | v = result.env[key] 39 | return v 40 | return key 41 | 42 | def _resolve_values(keys: List[str]) -> List[str]: 43 | return list(map(_resolve_env, keys)) 44 | 45 | result = DockerBuildEmulatorResult() 46 | for cmd in self._commands: 47 | if cmd.cmd not in ["ENV", "ENTRYPOINT", "COPY", "WORKDIR"]: 48 | continue 49 | values = _resolve_values(cmd.value) 50 | if cmd.cmd == "ENV": 51 | result.env[values[0]] = values[1] 52 | if cmd.cmd == "ENTRYPOINT": 53 | result.entrypoint = list(values) 54 | if cmd.cmd == "COPY": 55 | src, dst = values 56 | src = src.replace("./", "", 1) 57 | dst = dst.replace("/", "", 1) 58 | copy_tree_or_file(self._context_dir / src, fs_root_dir / dst) 59 | if cmd.cmd == "WORKDIR": 60 | result.workdir = result.workdir / values[0] 61 | return result 62 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Model to Docker 2 | A unified way to package pre-trained models into docker files 3 | 4 | 5 | # Model to Docker in action, or how it works 6 | 7 | https://youtu.be/rc6ylq01D0c 8 | 9 | # Installation guide 10 | 11 | ``` 12 | git clone git@github.com:slashml/slash_docker.git 13 | cd slash_docker 14 | pip install -e . 15 | ``` 16 | 17 | # Quickstart 18 | 19 | 20 | ```python 21 | from sklearn.linear_model import LinearRegression 22 | 23 | lm = LinearRegression() 24 | lm.fit([[2], [3], [4]], [4,6,8]) 25 | 26 | #test the model locally 27 | lm.predict([[1024]]) 28 | 29 | ``` 30 | 31 | ```python 32 | # serialize and save the model 33 | from model_to_docker import save_model 34 | save_model(lm, 'sklearn_linear_regression') 35 | 36 | ``` 37 | 38 | ```python 39 | # start a docker container with the previously saved model 40 | # this will create a docker image, and start a docker container 41 | from model_to_docker import run_model_server 42 | 43 | # this will take a few seconds if its the first time 44 | container = run_model_server('sklearn_linear_regression', port=5000) 45 | 46 | ``` 47 | 48 | ```python 49 | # perform inference on the locally running docker container 50 | import requests 51 | 52 | resp = requests.post('http://127.0.0.1:8080/v1/models/model:predict', '[[1024]]') 53 | 54 | print(resp.json()) 55 | ``` 56 | 57 | ```bash 58 | # you can also use the curl command in terminal to perform inference 59 | curl http://127.0.0.1:8080/v1/models/model:predict -d '[[1025]]' 60 | ``` 61 | 62 | ```python 63 | # stop the docker container 64 | from model_to_docker import stop_model_server 65 | stop_model_server(container.id) 66 | ``` 67 | 68 | 69 | For other examples look inside the examples folder 70 | 71 | 72 | ### Supported frameworks 73 | 74 | * Scikit-learn 75 | * XGBoost 76 | * PyTorch 77 | * Mlflow 78 | * TensorFlow 79 | * MLFLow 80 | * FastAI 81 | * HuggingFace Transformers 82 | * Keras 83 | * LightGBM 84 | * ONNX 85 | * PyCaret 86 | * SegmentAnything 87 | 88 | ### Contact 89 | If you run into any issues, or need a new feature, or any custom help, please feel free to reach out to us at faizank@slashml.com, support@slashml.com 90 | 91 | ### Contributing 92 | We are actively looking for contributors. Please look at the [CONTRIBUTING.md](https://github.com/slashml/slash_docker/blob/main/CONTRIBUTING.md) guide for more details. We have a growing list of good_first_issues, you can reach us at support@slashml.com, to get started. -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/types.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass 2 | from enum import Enum 3 | from typing import Any, Dict, List 4 | 5 | from truss.patch.types import TrussSignature 6 | from truss.templates.control.control.helpers.types import Patch 7 | 8 | 9 | class ModelFrameworkType(Enum): 10 | SKLEARN = "sklearn" 11 | TENSORFLOW = "tensorflow" 12 | KERAS = "keras" 13 | PYTORCH = "pytorch" 14 | HUGGINGFACE_TRANSFORMER = "huggingface_transformer" 15 | XGBOOST = "xgboost" 16 | LIGHTGBM = "lightgbm" 17 | MLFLOW = "mlflow" 18 | CUSTOM = "custom" 19 | 20 | 21 | @dataclass 22 | class Example: 23 | name: str 24 | input: Any 25 | 26 | @staticmethod 27 | def from_dict(example_dict): 28 | return Example( 29 | name=example_dict["name"], 30 | input=example_dict["input"], 31 | ) 32 | 33 | def to_dict(self) -> dict: 34 | return { 35 | "name": self.name, 36 | "input": self.input, 37 | } 38 | 39 | 40 | @dataclass 41 | class PatchDetails: 42 | prev_hash: str 43 | prev_signature: TrussSignature 44 | next_hash: str 45 | next_signature: TrussSignature 46 | patch_ops: List[Patch] 47 | 48 | def to_dict(self): 49 | return { 50 | "prev_hash": self.prev_hash, 51 | "prev_signature": self.prev_signature.to_dict(), 52 | "next_hash": self.next_hash, 53 | "next_signature": self.next_signature.to_dict(), 54 | "patch_ops": [patch_op.to_dict() for patch_op in self.patch_ops], 55 | } 56 | 57 | def is_empty(self) -> bool: 58 | # It's possible for prev_hash and next_hash to be different and yet 59 | # patch_ops to be empty, because certain parts of truss may be ignored. 60 | # e.g. training code is ignored for serving. 61 | return len(self.patch_ops) == 0 62 | 63 | @staticmethod 64 | def from_dict(patch_details: Dict) -> "PatchDetails": 65 | return PatchDetails( 66 | prev_hash=patch_details["prev_hash"], 67 | prev_signature=TrussSignature.from_dict(patch_details["prev_signature"]), 68 | next_hash=patch_details["next_hash"], 69 | next_signature=TrussSignature.from_dict(patch_details["next_signature"]), 70 | patch_ops=[ 71 | Patch.from_dict(patch_op) for patch_op in patch_details["patch_ops"] 72 | ], 73 | ) 74 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/contexts/local_loader/train_local.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | from pathlib import Path 3 | from typing import Dict, Optional 4 | 5 | from truss.contexts.local_loader.truss_module_loader import truss_module_loaded 6 | from truss.contexts.local_loader.utils import ( 7 | prepare_secrets, 8 | signature_accepts_keyword_arg, 9 | ) 10 | from truss.contexts.truss_context import TrussContext 11 | from truss.truss_spec import TrussSpec 12 | 13 | 14 | class LocalTrainer(TrussContext): 15 | """Allows training a truss locally.""" 16 | 17 | @staticmethod 18 | def run(truss_dir: Path): 19 | def train(variables: Optional[Dict] = None): 20 | spec = TrussSpec(truss_dir) 21 | with truss_module_loaded( 22 | str(truss_dir), 23 | spec.train_module_fullname, 24 | spec.bundled_packages_dir.name, 25 | [str(path.resolve()) for path in spec.external_package_dirs_paths], 26 | ) as module: 27 | train_class = getattr(module, spec.train_class_name) 28 | train_class_signature = inspect.signature(train_class) 29 | train_init_params = {} 30 | config = spec.config 31 | if signature_accepts_keyword_arg(train_class_signature, "config"): 32 | train_init_params["config"] = config.to_dict() 33 | if signature_accepts_keyword_arg(train_class_signature, "output_dir"): 34 | train_init_params["output_dir"] = truss_dir / "data" 35 | if signature_accepts_keyword_arg(train_class_signature, "secrets"): 36 | train_init_params["secrets"] = prepare_secrets(spec) 37 | 38 | # Wire up variables 39 | if signature_accepts_keyword_arg(train_class_signature, "variables"): 40 | runtime_variables = {} 41 | runtime_variables.update(config.train.variables) 42 | if variables is not None: 43 | runtime_variables.update(variables) 44 | train_init_params["variables"] = runtime_variables 45 | trainer = train_class(**train_init_params) 46 | 47 | if hasattr(trainer, "pre_train"): 48 | trainer.pre_train() 49 | 50 | trainer.train() 51 | 52 | if hasattr(trainer, "post_train"): 53 | trainer.post_train() 54 | 55 | return train 56 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/control/control/application.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import re 3 | from pathlib import Path 4 | from typing import Dict 5 | 6 | from endpoints import control_app 7 | from flask import Flask 8 | from helpers.errors import PatchApplicatonError 9 | from helpers.inference_server_controller import InferenceServerController 10 | from helpers.inference_server_process_controller import InferenceServerProcessController 11 | from helpers.patch_applier import PatchApplier 12 | from werkzeug.exceptions import HTTPException 13 | 14 | 15 | def create_app(base_config: Dict): 16 | app = Flask(__name__) 17 | # TODO(pankaj): change this back to info once things are stable 18 | app.logger.setLevel(logging.DEBUG) 19 | app.config.update(base_config) 20 | app.config[ 21 | "inference_server_process_controller" 22 | ] = InferenceServerProcessController( 23 | app.config["inference_server_home"], 24 | app.config["inference_server_process_args"], 25 | app.config["inference_server_port"], 26 | app_logger=app.logger, 27 | ) 28 | patch_applier = PatchApplier( 29 | Path(app.config["inference_server_home"]), 30 | app.logger, 31 | app.config.get("pip_path"), 32 | ) 33 | app.config["inference_server_controller"] = InferenceServerController( 34 | app.config["inference_server_process_controller"], 35 | patch_applier, 36 | app.logger, 37 | app.config.get("oversee_inference_server", True), 38 | ) 39 | app.register_blueprint(control_app) 40 | 41 | def handle_error(exc): 42 | try: 43 | raise exc 44 | except HTTPException: 45 | return exc 46 | except PatchApplicatonError: 47 | app.logger.exception(exc) 48 | error_type = _camel_to_snake_case(type(exc).__name__) 49 | return { 50 | "error": { 51 | "type": error_type, 52 | "msg": str(exc), 53 | } 54 | } 55 | except Exception: 56 | app.logger.exception(exc) 57 | return { 58 | "error": { 59 | "type": "unknown", 60 | "msg": f"{type(exc)}: {exc}", 61 | } 62 | } 63 | 64 | app.register_error_handler(Exception, handle_error) 65 | return app 66 | 67 | 68 | def _camel_to_snake_case(camel_cased: str) -> str: 69 | return re.sub(r"(? ModelFrameworkType: 15 | return ModelFrameworkType.MLFLOW 16 | 17 | def required_python_depedencies(self) -> Set[str]: 18 | return MLFLOW_REQ_MODULE_NAMES 19 | 20 | def serialize_model_to_directory(self, model, target_directory: Path): 21 | import mlflow 22 | 23 | if isinstance(model, str): 24 | self._download_model_from_uri(uri=model, target_directory=target_directory) 25 | elif isinstance(model, mlflow.pyfunc.PyFuncModel): 26 | self._download_model_from_pyfunc( 27 | model=model, target_directory=target_directory 28 | ) 29 | 30 | def model_metadata(self, model) -> Dict[str, Any]: 31 | supports_predict_proba = model_supports_predict_proba(model) 32 | return { 33 | "model_binary_dir": "model", 34 | "supports_predict_proba": supports_predict_proba, 35 | } 36 | 37 | def supports_model_class(self, model_class) -> bool: 38 | model_framework, _, _ = model_class.__module__.partition(".") 39 | return model_framework == ModelFrameworkType.MLFLOW.value 40 | 41 | def to_truss(self, model, target_directory: Path) -> None: 42 | super().to_truss(model, target_directory) 43 | self._add_mlflow_requirements(str(target_directory)) 44 | 45 | def _download_model_from_uri(self, uri: str, target_directory: Path): 46 | from mlflow.artifacts import download_artifacts 47 | 48 | download_artifacts(artifact_uri=uri, dst_path=target_directory) 49 | 50 | def _download_model_from_pyfunc(self, model, target_directory: Path): 51 | from mlflow.artifacts import download_artifacts 52 | 53 | run_id = model._model_meta.run_id 54 | download_artifacts(run_id=run_id, dst_path=target_directory) 55 | 56 | def _add_mlflow_requirements(self, target_directory: str): 57 | truss = TrussHandle(truss_dir=Path(target_directory)) 58 | requirements_file = ( 59 | truss._spec.data_dir / "model" / "model" / "requirements.txt" 60 | ) 61 | if not requirements_file.exists(): 62 | return 63 | truss.update_requirements_from_file(str(requirements_file)) 64 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/control/control/helpers/inference_server_starter.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import requests 4 | from tenacity import Retrying, stop_after_attempt, wait_exponential 5 | 6 | 7 | def inference_server_startup_flow(application) -> None: 8 | """ 9 | Perform the inference server startup flow 10 | 11 | Inference server startup flow supports checking for patches. If a patch ping 12 | url is provided then we hit that url to start the sync mechanism. The ping 13 | calls with current truss hash. The patch ping endpoint should return a 14 | response indicating, either that the supplied hash is current or that the 15 | request has been accepted. Acceptance of request means that a patch will be 16 | supplied soon to the truss (by calling of /control/patch endpoint). 17 | 18 | If we find that our hash is current, we start the inference server 19 | immediately. Otherwise, we delay the start to when the patch is supplied. 20 | 21 | The goal is to start the inference server as soon as we have the latest 22 | code, but not before. 23 | Example responses: 24 | {"is_current": true} 25 | {"accepted": true} 26 | """ 27 | inference_server_controller = application.config["inference_server_controller"] 28 | patch_ping_url = os.environ.get("PATCH_PING_URL_TRUSS") 29 | if patch_ping_url is None: 30 | inference_server_controller.start() 31 | return 32 | 33 | truss_hash = inference_server_controller.truss_hash() 34 | payload = {"truss_hash": truss_hash} 35 | 36 | for attempt in Retrying( 37 | stop=stop_after_attempt(15), 38 | wait=wait_exponential(multiplier=2, min=1, max=4), 39 | ): 40 | with attempt: 41 | try: 42 | application.logger.info( 43 | f"Pinging {patch_ping_url} for patch with hash {truss_hash}" 44 | ) 45 | resp = requests.post(patch_ping_url, json=payload) 46 | resp.raise_for_status() 47 | resp_body = resp.json() 48 | 49 | # If hash is current start inference server, otherwise delay that 50 | # for when patch is applied. 51 | if "is_current" in resp_body and resp_body["is_current"] is True: 52 | application.logger.info( 53 | "Hash is current, starting inference server" 54 | ) 55 | inference_server_controller.start() 56 | except Exception as exc: # noqa 57 | application.logger.warning( 58 | f"Patch ping attempt failed with error {exc}" 59 | ) 60 | raise exc 61 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/server/common/patches/whisper/patch.py: -------------------------------------------------------------------------------- 1 | """ 2 | Patch for OpenAi/Whisper: 3 | 4 | Whisper currently uses tqdm for the progress bar during the download of Whisper weights from Azure. 5 | However, when Transfer-Encoding=Chunked, the Content-Length is not set, causing the loading process 6 | to fail. This patch makes the download more defensive, addressing this issue. tqdm progress bar still 7 | works when the Content-Length is None. 8 | Open PR (https://github.com/openai/whisper/pull/1366) - once merged we can remove the patch. 9 | """ 10 | import os 11 | from typing import Union 12 | 13 | 14 | def _download(url: str, root: str, in_memory: bool) -> Union[bytes, str]: 15 | import hashlib 16 | import urllib.request 17 | import warnings 18 | 19 | from tqdm import tqdm 20 | 21 | os.makedirs(root, exist_ok=True) 22 | expected_sha256 = url.split("/")[-2] 23 | download_target = os.path.join(root, os.path.basename(url)) 24 | 25 | if os.path.exists(download_target) and not os.path.isfile(download_target): 26 | raise RuntimeError(f"{download_target} exists and is not a regular file") 27 | 28 | if os.path.isfile(download_target): 29 | with open(download_target, "rb") as f: 30 | model_bytes = f.read() 31 | if hashlib.sha256(model_bytes).hexdigest() == expected_sha256: 32 | return model_bytes if in_memory else download_target 33 | else: 34 | warnings.warn( 35 | f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file" 36 | ) 37 | 38 | with urllib.request.urlopen(url) as source, open(download_target, "wb") as output: 39 | with tqdm( 40 | total=int(source.info().get("Content-Length")) 41 | if source.info().get("Content-Length") is not None 42 | else None, 43 | ncols=80, 44 | unit="iB", 45 | unit_scale=True, 46 | unit_divisor=1024, 47 | ) as loop: 48 | while True: 49 | buffer = source.read(8192) 50 | if not buffer: 51 | break 52 | 53 | output.write(buffer) 54 | loop.update(len(buffer)) 55 | 56 | model_bytes = open(download_target, "rb").read() 57 | if hashlib.sha256(model_bytes).hexdigest() != expected_sha256: 58 | raise RuntimeError( 59 | "Model has been downloaded but the SHA256 checksum does not not match. Please retry loading the model." 60 | ) 61 | 62 | return model_bytes if in_memory else download_target 63 | 64 | 65 | def patch(): 66 | import whisper 67 | 68 | whisper._download = _download 69 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "truss" 3 | version = "0.4.8" 4 | description = "A seamless bridge from model development to model delivery" 5 | license = "MIT" 6 | readme = "README.md" 7 | authors = ["Pankaj Gupta ", "Phil Howes "] 8 | include = ["*.txt", "*.Dockerfile", "*.md"] 9 | repository = "https://github.com/basetenlabs/truss" 10 | keywords = ["MLOps", "AI", "Model Serving", "Model Deployment", "Machine Learning"] 11 | 12 | [tool.poetry.urls] 13 | "Homepage" = "https://truss.baseten.co" 14 | "Bug Reports" = "https://github.com/basetenlabs/truss/issues" 15 | "Documentation" = "https://truss.baseten.co" 16 | "Baseten" = "https://baseten.co" 17 | 18 | [tool.poetry.dependencies] 19 | python = ">=3.8,<3.12" 20 | numpy = "1.23.5" 21 | msgpack = ">=1.0.2" 22 | msgpack-numpy = ">=0.4.7.1" 23 | packaging = "^20.9" 24 | python-json-logger = ">=2.0.2" 25 | PyYAML = "^6.0" 26 | Jinja2 = "^3.1.2" 27 | python-on-whales = "^0.46.0" 28 | tenacity = "^8.0.1" 29 | single-source = "^0.3.0" 30 | cloudpickle = "^2.2.0" 31 | blake3 = "^0.3.3" 32 | fastapi = "^0.95.0" 33 | uvicorn = "^0.21.1" 34 | psutil = "^5.9.4" 35 | joblib = "^1.2.0" 36 | dockerfile = "^3.2.0" 37 | 38 | [tool.poetry.group.builder.dependencies] 39 | python = ">=3.8,<3.12" 40 | packaging = "^20.9" 41 | python-json-logger = ">=2.0.2" 42 | PyYAML = "^6.0" 43 | Jinja2 = "^3.1.2" 44 | tenacity = "^8.0.1" 45 | cloudpickle = "^2.2.0" 46 | single-source = "^0.3.0" 47 | click = "^8.0.3" 48 | requests = "^2.28.1" 49 | blake3 = "^0.3.3" 50 | fastapi = "^0.95.0" 51 | uvicorn = "^0.21.1" 52 | psutil = "^5.9.4" 53 | 54 | [tool.poetry.dev-dependencies] 55 | torch = "^1.9.0" 56 | ipython = "^7.16" 57 | pytest = "7.2.0" 58 | tensorflow = { version = "^2.4.4", markers = "sys_platform == 'linux'" } 59 | tensorflow-macos = { version = "^2.4.4", markers = "sys_platform == 'darwin'" } 60 | pre-commit = "^2.18.1" 61 | scikit-learn = "1.0.2" 62 | pandas = "1.5.2" 63 | tensorflow-hub = "^0.12.0" 64 | isort = "^5.10.1" 65 | flake8 = "^4.0.1" 66 | ipdb = "^0.13.9" 67 | coverage = "^6.4.1" 68 | pytest-cov = "^3.0.0" 69 | xgboost = "^1.6.1" 70 | lightgbm = "^3.3.2" 71 | transformers = "^4.20.1" 72 | black = "^22.6.0" 73 | Flask = "^2.2.2" 74 | waitress = "^2.1.2" 75 | nbconvert = "^7.2.1" 76 | ipykernel = "^6.16.0" 77 | 78 | [tool.poetry.scripts] 79 | truss = 'truss.cli:cli_group' 80 | 81 | [tool.poetry.group.dev.dependencies] 82 | mlflow = "^1.29.0" 83 | mypy = "^1.0.0" 84 | pytest-split = "^0.8.1" 85 | 86 | [build-system] 87 | requires = ["poetry-core>=1.2.1"] 88 | build-backend = "poetry.core.masonry.api" 89 | 90 | [tool.isort] 91 | profile = "black" 92 | src_paths = ["isort", "test"] 93 | 94 | [tool.mypy] 95 | python_version = 3.8 96 | ignore_missing_imports = true 97 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/patch/hash.py: -------------------------------------------------------------------------------- 1 | import fnmatch 2 | from pathlib import Path 3 | from typing import Any, List, Optional 4 | 5 | from blake3 import blake3 6 | 7 | 8 | def directory_content_hash( 9 | root: Path, 10 | ignore_patterns: Optional[List[str]] = None, 11 | ) -> str: 12 | """Calculate content based hash of a filesystem directory. 13 | 14 | Rough algo: Sort all files by path, then take hash of a content stream, where 15 | we write path hash to the stream followed by hash of content if path is a file. 16 | Note the hash of hash aspect. 17 | 18 | Also, note that name of the root directory is not taken into account, only the contents 19 | underneath. The (root) Directory will have the same hash, even if renamed. 20 | """ 21 | hasher = blake3() 22 | paths = [ 23 | path 24 | for path in root.glob("**/*") 25 | if not _path_matches_any_pattern(path.relative_to(root), ignore_patterns) 26 | ] 27 | paths.sort(key=lambda p: p.relative_to(root)) 28 | for path in paths: 29 | hasher.update(str_hash(str(path.relative_to(root)))) 30 | if path.is_file(): 31 | hasher.update(file_content_hash(path)) 32 | return hasher.hexdigest() 33 | 34 | 35 | def file_content_hash(file: Path) -> bytes: 36 | """Calculate blake3 hash of file content. 37 | Returns: binary hash of content 38 | """ 39 | return _file_content_hash_loaded_hasher(file).digest() 40 | 41 | 42 | def file_content_hash_str(file: Path) -> str: 43 | """Calculate blake3 hash of file content. 44 | 45 | Returns: string hash of content 46 | """ 47 | return _file_content_hash_loaded_hasher(file).hexdigest() 48 | 49 | 50 | def _file_content_hash_loaded_hasher(file: Path) -> Any: 51 | hasher = blake3() 52 | buffer = bytearray(128 * 1024) 53 | mem_view = memoryview(buffer) 54 | with file.open("rb") as f: 55 | done = False 56 | while not done: 57 | n = f.readinto(mem_view) 58 | if n > 0: 59 | hasher.update(mem_view[:n]) 60 | else: 61 | done = True 62 | return hasher 63 | 64 | 65 | def str_hash(content: str) -> bytes: 66 | hasher = blake3() 67 | hasher.update(content.encode("utf-8")) 68 | return hasher.digest() 69 | 70 | 71 | def str_hash_str(content: str) -> str: 72 | hasher = blake3() 73 | hasher.update(content.encode("utf-8")) 74 | return hasher.hexdigest() 75 | 76 | 77 | def _path_matches_any_pattern(path: Path, patterns: Optional[List[str]]) -> bool: 78 | if patterns is None: 79 | return False 80 | 81 | for pattern in patterns: 82 | if fnmatch.fnmatch(str(path), pattern): 83 | return True 84 | 85 | return False 86 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/model_frameworks/pytorch.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from typing import Dict, Set 3 | 4 | from truss.constants import PYTORCH_REQ_MODULE_NAMES 5 | from truss.model_framework import ModelFramework 6 | from truss.types import ModelFrameworkType 7 | 8 | TORCH_PACKAGE_FILE = "model_package.pt" 9 | TORCH_MODEL_PICKLE_FILENAME = "model.pkl" 10 | TORCH_MODEL_PACKAGE_NAME = "torch_model" 11 | 12 | 13 | class PyTorch(ModelFramework): 14 | def typ(self) -> ModelFrameworkType: 15 | return ModelFrameworkType.PYTORCH 16 | 17 | def required_python_depedencies(self) -> Set[str]: 18 | return PYTORCH_REQ_MODULE_NAMES 19 | 20 | def serialize_model_to_directory(self, model, target_directory: Path): 21 | from torch import package 22 | 23 | try: 24 | _torch_package(model, target_directory / TORCH_PACKAGE_FILE, []) 25 | except package.package_exporter.PackagingError as pkg_err: 26 | # Make sure previous, potentially partially written, package is gone 27 | (target_directory / TORCH_PACKAGE_FILE).unlink() 28 | modules_to_extern = _broken_modules(pkg_err) 29 | _torch_package( 30 | model, target_directory / TORCH_PACKAGE_FILE, modules_to_extern 31 | ) 32 | 33 | def model_metadata(self, model) -> Dict[str, str]: 34 | return { 35 | "torch_package_file": TORCH_PACKAGE_FILE, 36 | "torch_model_pickle_filename": TORCH_MODEL_PICKLE_FILENAME, 37 | "torch_model_package_name": TORCH_MODEL_PACKAGE_NAME, 38 | "model_binary_dir": "model", 39 | } 40 | 41 | def supports_model_class(self, model_class) -> bool: 42 | try: 43 | import torch 44 | 45 | return issubclass(model_class, torch.nn.Module) 46 | except ModuleNotFoundError: 47 | return False 48 | 49 | 50 | def _torch_package(model, path: Path, extern_modules) -> None: 51 | from torch import package 52 | 53 | with package.PackageExporter(path) as exp: 54 | exp.intern(f"{model.__class__.__module__}.**") 55 | for extern_module in extern_modules: 56 | exp.extern(f"{extern_module}.**") 57 | exp.save_pickle(TORCH_MODEL_PACKAGE_NAME, TORCH_MODEL_PICKLE_FILENAME, model) 58 | 59 | 60 | def _broken_modules(pkg_error): 61 | """Extract the broken modules from the torch package error. 62 | 63 | We would extern these modules. 64 | Args: 65 | pkg_error (package.package_exporter.PackagingError): Error to extract broken modules from. 66 | """ 67 | return [ 68 | module_name 69 | for module_name, attrs in pkg_error.dependency_graph.nodes.items() 70 | if attrs.get("error") is not None 71 | ] 72 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/README.md.jinja: -------------------------------------------------------------------------------- 1 | {% if config.model_name == None %} 2 | # Model how-to 3 | {% else %} 4 | # {{ config.model_name }} how-to 5 | {% endif %} 6 | Welcome to your {{ config.model_framework.value }} model's Truss. Below are some useful commands to work with this Truss. For all the commands below, make sure you're in this directory when you run them. You can find the docs for Truss [here](https://truss.baseten.co). 7 | 8 | ### Build a Docker image from your Truss 9 | ``` 10 | truss build-image 11 | ``` 12 | 13 | ### Run the Docker image from your Truss 14 | ``` 15 | # We assume you've built the image first. 16 | truss run-image 17 | ``` 18 | 19 | ### Run inference on your Truss 20 | There are two ways to run inference on your Truss model. 21 | #### Via Truss CLI 22 | ``` 23 | {%- if examples == None %} 24 | truss predict --target_directory ./ --request 'YOUR_INPUT' 25 | {% else %} 26 | {%- for example in examples %} 27 | {%- for second_key, second_value in example.input.items() %} 28 | truss predict --target_directory ./ --request '{{ second_value[0] }}' 29 | {%- endfor %} 30 | {%- endfor %} 31 | {%- endif %} 32 | ``` 33 | 34 | #### Via CURL 35 | In order to run inference via CURL, we assume you've built and run the Docker image generated from your Truss. Refer above for more instructions on how to do this. 36 | ``` 37 | {%- if examples == None %} 38 | curl -H 'Content-Type: application/json' \ 39 | -d 'YOUR_INPUT' \ 40 | -X POST http://localhost:8080/v1/models/model:predict 41 | {% else %} 42 | {%- for example in examples %} 43 | {%- for second_key, second_value in example.input.items() %} 44 | curl -H 'Content-Type: application/json' \ 45 | -d '{{ second_value[0] }}' \ 46 | -X POST http://localhost:8080/v1/models/model:predict 47 | {%- endfor %} 48 | {%- endfor %} 49 | {%- endif %} 50 | 51 | ``` 52 | 53 | ### Running all of the example inputs on your Truss 54 | ``` 55 | truss run-example 56 | ``` 57 | 58 | ### Run a specific example input on your Truss 59 | ``` 60 | {%- if examples == None %} 61 | truss run-example --name EXAMPLE_NAME 62 | {% else %} 63 | {%- for example in examples %} 64 | truss run-example --name {{ example.name }} 65 | {%- endfor %} 66 | {%- endif %} 67 | ``` 68 | 69 | ### Adding pre/postprocessing code 70 | You may find that you'd like to preprocess the inputs to your model or postprocess the outputs of your model. 71 | 1. Navigate to `model/{{ config.model_class_filename }}`. This is where your preprocessing/postprocessing logic will live. 72 | 2. Define a `preprocess` or `postprocess` function as such. 73 | ``` 74 | def preprocess(self, request: Dict) -> Dict: 75 | # Code that runs before inference 76 | return request 77 | 78 | def postprocess(self, request: Dict) -> Dict: 79 | # Code that runs after inference 80 | return request 81 | ``` 82 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/server/common/errors.py: -------------------------------------------------------------------------------- 1 | from http import HTTPStatus 2 | from typing import Optional 3 | 4 | from fastapi.responses import JSONResponse 5 | 6 | 7 | class ModelMissingError(Exception): 8 | def __init__(self, path): 9 | self.path = path 10 | 11 | def __str__(self): 12 | return self.path 13 | 14 | 15 | class InferenceError(RuntimeError): 16 | def __init__(self, reason): 17 | self.reason = reason 18 | 19 | def __str__(self): 20 | return self.reason 21 | 22 | 23 | class InvalidInput(ValueError): 24 | """ 25 | Exception class indicating invalid input arguments. 26 | HTTP Servers should return HTTP_400 (Bad Request). 27 | """ 28 | 29 | def __init__(self, reason): 30 | self.reason = reason 31 | 32 | def __str__(self): 33 | return self.reason 34 | 35 | 36 | class ModelNotFound(Exception): 37 | """ 38 | Exception class indicating requested model does not exist. 39 | HTTP Servers should return HTTP_404 (Not Found). 40 | """ 41 | 42 | def __init__(self, model_name=None): 43 | self.reason = f"Model with name {model_name} does not exist." 44 | 45 | def __str__(self): 46 | return self.reason 47 | 48 | 49 | class ModelNotReady(RuntimeError): 50 | def __init__(self, model_name: str, detail: Optional[str] = None): 51 | self.model_name = model_name 52 | self.error_msg = f"Model with name {self.model_name} is not ready." 53 | if detail: 54 | self.error_msg = self.error_msg + " " + detail 55 | 56 | def __str__(self): 57 | return self.error_msg 58 | 59 | 60 | async def exception_handler(_, exc): 61 | return JSONResponse( 62 | status_code=HTTPStatus.INTERNAL_SERVER_ERROR, content={"error": str(exc)} 63 | ) 64 | 65 | 66 | async def invalid_input_handler(_, exc): 67 | return JSONResponse(status_code=HTTPStatus.BAD_REQUEST, content={"error": str(exc)}) 68 | 69 | 70 | async def inference_error_handler(_, exc): 71 | return JSONResponse( 72 | status_code=HTTPStatus.INTERNAL_SERVER_ERROR, content={"error": str(exc)} 73 | ) 74 | 75 | 76 | async def generic_exception_handler(_, exc): 77 | return JSONResponse( 78 | status_code=HTTPStatus.INTERNAL_SERVER_ERROR, 79 | content={"error": f"{type(exc).__name__} : {str(exc)}"}, 80 | ) 81 | 82 | 83 | async def model_not_found_handler(_, exc): 84 | return JSONResponse(status_code=HTTPStatus.NOT_FOUND, content={"error": str(exc)}) 85 | 86 | 87 | async def model_not_ready_handler(_, exc): 88 | return JSONResponse( 89 | status_code=HTTPStatus.SERVICE_UNAVAILABLE, content={"error": str(exc)} 90 | ) 91 | 92 | 93 | async def not_implemented_error_handler(_, exc): 94 | return JSONResponse( 95 | status_code=HTTPStatus.NOT_IMPLEMENTED, content={"error": str(exc)} 96 | ) 97 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/control/control/helpers/inference_server_process_controller.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import signal 4 | import subprocess 5 | from typing import List, Optional 6 | 7 | from helpers.context_managers import current_directory 8 | 9 | 10 | class InferenceServerProcessController: 11 | 12 | _inference_server_process: Optional[subprocess.Popen] = None 13 | _inference_server_port: int 14 | _inference_server_home: str 15 | _app_logger: logging.Logger 16 | _inference_server_process_args: List[str] 17 | 18 | def __init__( 19 | self, 20 | inference_server_home: str, 21 | inference_server_process_args: List[str], 22 | inference_server_port: int, 23 | app_logger: logging.Logger, 24 | ) -> None: 25 | self._inference_server_home = inference_server_home 26 | self._inference_server_process_args = inference_server_process_args 27 | self._inference_server_port = inference_server_port 28 | self._inference_server_started = False 29 | self._inference_server_ever_started = False 30 | self._app_logger = app_logger 31 | 32 | def start(self): 33 | with current_directory(self._inference_server_home): 34 | inf_env = os.environ.copy() 35 | inf_env["INFERENCE_SERVER_PORT"] = str(self._inference_server_port) 36 | self._inference_server_process = subprocess.Popen( 37 | self._inference_server_process_args, 38 | env=inf_env, 39 | ) 40 | 41 | self._inference_server_started = True 42 | self._inference_server_ever_started = True 43 | 44 | def stop(self): 45 | if self._inference_server_process is not None: 46 | name = " ".join(self._inference_server_process_args) 47 | 48 | for line in os.popen("ps ax | grep '" + name + "' | grep -v grep"): 49 | pid = line.split()[0] 50 | os.kill(int(pid), signal.SIGKILL) 51 | 52 | self._inference_server_started = False 53 | 54 | def inference_server_started(self) -> bool: 55 | return self._inference_server_started 56 | 57 | def inference_server_ever_started(self) -> bool: 58 | return self._inference_server_ever_started 59 | 60 | def is_inference_server_running(self) -> bool: 61 | # Explicitly check if inference server process is up, this is a bit expensive. 62 | if not self._inference_server_started: 63 | return False 64 | 65 | if self._inference_server_process is None: 66 | return False 67 | 68 | return self._inference_server_process.poll() is None 69 | 70 | def check_and_recover_inference_server(self): 71 | if self.inference_server_started() and not self.is_inference_server_running(): 72 | self._app_logger.warning( 73 | "Inference server seems to have crashed, restarting" 74 | ) 75 | self.start() 76 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/model_framework.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from pathlib import Path 3 | from typing import Dict, List, Optional, Set 4 | 5 | import yaml 6 | from truss.constants import CONFIG_FILE, TEMPLATES_DIR 7 | from truss.environment_inference.requirements_inference import infer_deps 8 | from truss.model_inference import infer_python_version, map_to_supported_python_version 9 | from truss.truss_config import DEFAULT_EXAMPLES_FILENAME, TrussConfig 10 | from truss.types import ModelFrameworkType 11 | from truss.util.path import copy_file_path, copy_tree_path 12 | 13 | 14 | class ModelFramework(ABC): 15 | @abstractmethod 16 | def typ(self) -> ModelFrameworkType: 17 | pass 18 | 19 | @abstractmethod 20 | def required_python_depedencies(self) -> Set[str]: 21 | """Returns a set of packages required by this framework. 22 | 23 | e.g. {'tensorflow'} 24 | """ 25 | pass 26 | 27 | def requirements_txt(self) -> List[str]: 28 | 29 | return list(infer_deps(must_include_deps=self.required_python_depedencies())) 30 | 31 | @abstractmethod 32 | def serialize_model_to_directory(self, model, target_directory: Path): 33 | pass 34 | 35 | @abstractmethod 36 | def model_metadata(self, model) -> Dict[str, str]: 37 | pass 38 | 39 | def model_type(self, model) -> str: 40 | return "Model" 41 | 42 | def model_name(self, model) -> Optional[str]: 43 | return None 44 | 45 | def to_truss(self, model, target_directory: Path) -> None: 46 | """Exports in-memory model to a Truss, in a target directory.""" 47 | model_binary_dir = target_directory / "data" / "model" 48 | model_binary_dir.mkdir(parents=True, exist_ok=True) 49 | 50 | # Serialize model and write it 51 | self.serialize_model_to_directory(model, model_binary_dir) 52 | template_path = TEMPLATES_DIR / self.typ().value 53 | copy_tree_path(template_path / "model", target_directory / "model") 54 | examples_path = template_path / DEFAULT_EXAMPLES_FILENAME 55 | target_examples_path = target_directory / DEFAULT_EXAMPLES_FILENAME 56 | if examples_path.exists(): 57 | copy_file_path(examples_path, target_examples_path) 58 | else: 59 | target_examples_path.touch() 60 | 61 | python_version = map_to_supported_python_version(infer_python_version()) 62 | 63 | # Create config 64 | config = TrussConfig( 65 | model_name=self.model_name(model), 66 | model_type=self.model_type(model), 67 | model_framework=self.typ(), 68 | model_metadata=self.model_metadata(model), 69 | requirements=self.requirements_txt(), 70 | python_version=python_version, 71 | ) 72 | with (target_directory / CONFIG_FILE).open("w") as config_file: 73 | yaml.dump(config.to_dict(), config_file) 74 | 75 | @abstractmethod 76 | def supports_model_class(self, model_class) -> bool: 77 | pass 78 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/validation.py: -------------------------------------------------------------------------------- 1 | import re 2 | from pathlib import Path 3 | from typing import Pattern, Set 4 | 5 | from truss.errors import ValidationError 6 | 7 | SECRET_NAME_MATCH_REGEX: Pattern[str] = re.compile(r"^[-._a-zA-Z0-9]+$") 8 | MILLI_CPU_REGEX: Pattern[str] = re.compile(r"^[0-9.]*m$") 9 | MEMORY_REGEX: Pattern[str] = re.compile(r"^[0-9.]*(\w*)$") 10 | MEMORY_UNITS: Set[str] = set( 11 | ["k", "M", "G", "T", "P", "E", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei"] 12 | ) 13 | 14 | 15 | def validate_secret_name(secret_name: str) -> None: 16 | if secret_name is None or not isinstance(secret_name, str) or secret_name == "": 17 | raise ValueError(f"Invalid secret name `{secret_name}`") 18 | 19 | def constraint_violation_msg(): 20 | return f"Constraint violation for {secret_name}" 21 | 22 | if len(secret_name) > 253: 23 | raise ValueError( 24 | f"Secret name `{secret_name}` is longer than max allowed 253 chars." 25 | ) 26 | 27 | if not SECRET_NAME_MATCH_REGEX.match(secret_name): 28 | raise ValueError( 29 | constraint_violation_msg() + ", invalid characters found in secret name." 30 | ) 31 | 32 | if secret_name == ".": 33 | raise ValueError(constraint_violation_msg() + ", secret name cannot be `.`") 34 | 35 | if secret_name == "..": 36 | raise ValueError(constraint_violation_msg() + ", secret name cannot be `..`") 37 | 38 | 39 | def validate_cpu_spec(cpu_spec: str) -> None: 40 | if not isinstance(cpu_spec, str): 41 | raise ValidationError( 42 | f"{cpu_spec} needs to be a string, but is {type(cpu_spec)}" 43 | ) 44 | 45 | if _is_numeric(cpu_spec): 46 | return 47 | 48 | is_milli_cpu_format = MILLI_CPU_REGEX.search(cpu_spec) is not None 49 | if not is_milli_cpu_format: 50 | raise ValidationError(f"Invalid cpu specification {cpu_spec}") 51 | 52 | 53 | def validate_memory_spec(mem_spec: str) -> None: 54 | if not isinstance(mem_spec, str): 55 | raise ValidationError( 56 | f"{mem_spec} needs to be a string, but is {type(mem_spec)}" 57 | ) 58 | if _is_numeric(mem_spec): 59 | return 60 | 61 | match = MEMORY_REGEX.search(mem_spec) 62 | if match is None: 63 | raise ValidationError(f"Invalid memory specification {mem_spec}") 64 | 65 | unit = match.group(1) 66 | if unit not in MEMORY_UNITS: 67 | raise ValidationError(f"Invalid memory unit {unit} in {mem_spec}") 68 | 69 | 70 | def _is_numeric(number_like: str) -> bool: 71 | try: 72 | float(number_like) 73 | return True 74 | except ValueError: 75 | return False 76 | 77 | 78 | def validate_python_executable_path(path: str) -> None: 79 | """ 80 | This python executable path determines the python executable 81 | used to run the inference server - check to see that it is an absolute path 82 | """ 83 | if path and not Path(path).is_absolute(): 84 | raise ValidationError( 85 | f"Invalid relative python executable path {path}. Provide an absolute path" 86 | ) 87 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/server.Dockerfile.jinja: -------------------------------------------------------------------------------- 1 | {% extends "base.Dockerfile.jinja" %} 2 | 3 | {% block base_image_patch %} 4 | # If user base image is supplied in config, apply build commands from truss base image 5 | {% if config.base_image %} 6 | ENV PYTHONUNBUFFERED True 7 | ENV DEBIAN_FRONTEND=noninteractive 8 | 9 | RUN apt update && \ 10 | apt install -y bash \ 11 | build-essential \ 12 | git \ 13 | curl \ 14 | ca-certificates \ 15 | software-properties-common \ 16 | && apt-get autoremove -y \ 17 | && apt-get clean -y \ 18 | && rm -rf /var/lib/apt/lists/* 19 | 20 | COPY ./base_server_requirements.txt base_server_requirements.txt 21 | RUN pip install -r base_server_requirements.txt --no-cache-dir && rm -rf /root/.cache/pip 22 | 23 | {%- if config.live_reload %} 24 | RUN $PYTHON_EXECUTABLE -m venv -h >/dev/null \ 25 | || { pythonVersion=$(echo $($PYTHON_EXECUTABLE --version) | cut -d" " -f2 | cut -d"." -f1,2) \ 26 | && add-apt-repository -y ppa:deadsnakes/ppa \ 27 | && apt update -y && apt install -y --no-install-recommends python$pythonVersion-venv \ 28 | && apt-get autoremove -y \ 29 | && apt-get clean -y \ 30 | && rm -rf /var/lib/apt/lists/*; } 31 | # Create symlink for control server to start inference server process with correct python executable 32 | RUN readlink {{config.base_image.python_executable_path}} &>/dev/null \ 33 | && echo "WARNING: Overwriting existing link at /usr/local/bin/python" 34 | RUN ln -sf {{config.base_image.python_executable_path}} /usr/local/bin/python 35 | {%- endif %} 36 | {% endif %} 37 | 38 | {% endblock %} 39 | 40 | {% block install_requirements %} 41 | {%- if should_install_server_requirements %} 42 | COPY ./server_requirements.txt server_requirements.txt 43 | RUN pip install -r server_requirements.txt --no-cache-dir && rm -rf /root/.cache/pip 44 | {%- endif %} 45 | {{ super() }} 46 | {% endblock %} 47 | 48 | {% block app_copy %} 49 | # Copy data before code for better caching 50 | {%- if data_dir_exists %} 51 | COPY ./{{config.data_dir}} /app/data 52 | {%- endif %} 53 | 54 | COPY ./server /app 55 | COPY ./{{ config.model_module_dir }} /app/model 56 | COPY ./config.yaml /app/config.yaml 57 | {%- if config.live_reload %} 58 | COPY ./control /control 59 | RUN python3 -m venv /control/.env \ 60 | && /control/.env/bin/pip3 install -r /control/requirements.txt 61 | {%- endif %} 62 | {% endblock %} 63 | 64 | 65 | {% block run %} 66 | {%- if config.live_reload %} 67 | ENV HASH_TRUSS {{truss_hash}} 68 | ENV CONTROL_SERVER_PORT 8080 69 | ENV INFERENCE_SERVER_PORT 8090 70 | ENV SERVER_START_CMD="/control/.env/bin/python3 /control/control/server.py" 71 | ENTRYPOINT ["/control/.env/bin/python3", "/control/control/server.py"] 72 | {%- else %} 73 | ENV INFERENCE_SERVER_PORT 8080 74 | ENV SERVER_START_CMD="{{(config.base_image.python_executable_path or "python3") ~ " /app/inference_server.py"}}" 75 | ENTRYPOINT ["{{config.base_image.python_executable_path or "python3"}}", "/app/inference_server.py"] 76 | {%- endif %} 77 | {% endblock %} 78 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/truss_gatherer.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | import yaml 4 | from truss.local.local_config_handler import LocalConfigHandler 5 | from truss.patch.hash import str_hash_str 6 | from truss.truss_handle import TrussHandle 7 | from truss.util.path import copy_file_path, copy_tree_path, remove_tree_path 8 | 9 | 10 | def gather(truss_path: Path) -> Path: 11 | handle = TrussHandle(truss_path) 12 | shadow_truss_dir_name = _calc_shadow_truss_dirname(truss_path) 13 | shadow_truss_metdata_file_path = ( 14 | LocalConfigHandler.shadow_trusses_dir_path() 15 | / f"{shadow_truss_dir_name}.metadata.yaml" 16 | ) 17 | shadow_truss_path = ( 18 | LocalConfigHandler.shadow_trusses_dir_path() / shadow_truss_dir_name 19 | ) 20 | if shadow_truss_metdata_file_path.exists(): 21 | with shadow_truss_metdata_file_path.open() as fp: 22 | metadata = yaml.safe_load(fp) 23 | max_mod_time = metadata["max_mod_time"] 24 | if max_mod_time == handle.max_modified_time: 25 | return shadow_truss_path 26 | 27 | # Shadow truss is out of sync, clear it 28 | shadow_truss_metdata_file_path.unlink() 29 | remove_tree_path(shadow_truss_path) 30 | 31 | copy_tree_path(truss_path, shadow_truss_path) 32 | packages_dir_path_in_shadow = ( 33 | shadow_truss_path / handle.spec.config.bundled_packages_dir 34 | ) 35 | packages_dir_path_in_shadow.mkdir(exist_ok=True) 36 | for path in handle.spec.external_package_dirs_paths: 37 | if not path.is_dir(): 38 | raise ValueError( 39 | f"External packages directory at {path} is not a directory" 40 | ) 41 | # We copy over contents of the external package directory, not the 42 | # directory itself. This mimics the local load behavior and is meant to 43 | # replicate adding external package directory to sys.path which doesn't 44 | # make the directory available as a package to python but the contents 45 | # inside. 46 | # 47 | # Note that this operation can fail if there are conflicts. Onus is on 48 | # the creator of truss to make sure that there are no conflicts. 49 | for sub_path in path.iterdir(): 50 | if sub_path.is_dir(): 51 | copy_tree_path(sub_path, packages_dir_path_in_shadow / sub_path.name) 52 | if sub_path.is_file(): 53 | copy_file_path(sub_path, packages_dir_path_in_shadow / sub_path.name) 54 | 55 | # Don't run validation because they will fail until we clear external 56 | # packages. We do it after. 57 | shadow_handle = TrussHandle(shadow_truss_path, validate=False) 58 | shadow_handle.clear_external_packages() 59 | shadow_handle.validate() 60 | with shadow_truss_metdata_file_path.open("w") as fp: 61 | yaml.safe_dump({"max_mod_time": handle.max_modified_time}, fp) 62 | return shadow_truss_path 63 | 64 | 65 | def _calc_shadow_truss_dirname(truss_path: Path) -> str: 66 | resolved_path_str = str(truss_path.resolve()) 67 | return str_hash_str(resolved_path_str) 68 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/constants.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pathlib 3 | from typing import Set 4 | 5 | SKLEARN = "sklearn" 6 | TENSORFLOW = "tensorflow" 7 | KERAS = "keras" 8 | XGBOOST = "xgboost" 9 | PYTORCH = "pytorch" 10 | CUSTOM = "custom" 11 | HUGGINGFACE_TRANSFORMER = "huggingface_transformer" 12 | LIGHTGBM = "lightgbm" 13 | 14 | BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 15 | CODE_DIR = pathlib.Path(BASE_DIR, "truss") 16 | 17 | TEMPLATES_DIR = pathlib.Path(CODE_DIR, "templates") 18 | SERVER_CODE_DIR: pathlib.Path = TEMPLATES_DIR / "server" 19 | TRAINING_JOB_WRAPPER_CODE_DIR_NAME = "training" 20 | TRAINING_JOB_WRAPPER_CODE_DIR: pathlib.Path = ( 21 | TEMPLATES_DIR / TRAINING_JOB_WRAPPER_CODE_DIR_NAME 22 | ) 23 | SHARED_SERVING_AND_TRAINING_CODE_DIR_NAME = "shared" 24 | SHARED_SERVING_AND_TRAINING_CODE_DIR: pathlib.Path = ( 25 | TEMPLATES_DIR / SHARED_SERVING_AND_TRAINING_CODE_DIR_NAME 26 | ) 27 | CONTROL_SERVER_CODE_DIR: pathlib.Path = TEMPLATES_DIR / "control" 28 | 29 | 30 | # Alias for TEMPLATES_DIR 31 | SERVING_DIR: pathlib.Path = TEMPLATES_DIR 32 | 33 | REQUIREMENTS_TXT_FILENAME = "requirements.txt" 34 | BASE_SERVER_REQUIREMENTS_TXT_FILENAME = "base_server_requirements.txt" 35 | SERVER_REQUIREMENTS_TXT_FILENAME = "server_requirements.txt" 36 | TRAINING_REQUIREMENTS_TXT_FILENAME = "training_requirements.txt" 37 | SYSTEM_PACKAGES_TXT_FILENAME = "system_packages.txt" 38 | 39 | SERVER_DOCKERFILE_TEMPLATE_NAME = "server.Dockerfile.jinja" 40 | TRAINING_DOCKERFILE_TEMPLATE_NAME = "training.Dockerfile.jinja" 41 | MODEL_DOCKERFILE_NAME = "Dockerfile" 42 | TRAINING_DOCKERFILE_NAME = "Dockerfile" 43 | 44 | README_TEMPLATE_NAME = "README.md.jinja" 45 | MODEL_README_NAME = "README.md" 46 | 47 | CONFIG_FILE = "config.yaml" 48 | DOCKERFILE = "Dockerfile" 49 | # Used to indicate whether to associate a container with Truss 50 | TRUSS = "truss" 51 | # Used to create unique identifier based on last time truss was updated 52 | TRUSS_MODIFIED_TIME = "truss_modified_time" 53 | # Path of the Truss used to identify which Truss is being referred 54 | TRUSS_DIR = "truss_dir" 55 | TRUSS_HASH = "truss_hash" 56 | TRAINING_TRUSS_HASH = "training_truss_hash" 57 | TRAINING_LABEL = "training" 58 | 59 | HUGGINGFACE_TRANSFORMER_MODULE_NAME: Set[str] = set({}) 60 | 61 | # list from https://scikit-learn.org/stable/developers/advanced_installation.html 62 | SKLEARN_REQ_MODULE_NAMES: Set[str] = { 63 | "numpy", 64 | "scipy", 65 | "joblib", 66 | "scikit-learn", 67 | "threadpoolctl", 68 | } 69 | 70 | XGBOOST_REQ_MODULE_NAMES: Set[str] = {"xgboost"} 71 | 72 | # list from https://www.tensorflow.org/install/pip 73 | # if problematic, lets look to https://www.tensorflow.org/install/source 74 | TENSORFLOW_REQ_MODULE_NAMES: Set[str] = { 75 | "tensorflow", 76 | } 77 | 78 | LIGHTGBM_REQ_MODULE_NAMES: Set[str] = { 79 | "lightgbm", 80 | } 81 | 82 | # list from https://pytorch.org/get-started/locally/ 83 | PYTORCH_REQ_MODULE_NAMES: Set[str] = { 84 | "torch", 85 | "torchvision", 86 | "torchaudio", 87 | } 88 | 89 | MLFLOW_REQ_MODULE_NAMES: Set[str] = {"mlflow"} 90 | 91 | INFERENCE_SERVER_PORT = 8080 92 | 93 | TRAINING_VARIABLES_FILENAME = "variables.yaml" 94 | 95 | HTTP_PUBLIC_BLOB_BACKEND = "http_public" 96 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/control/control/helpers/types.py: -------------------------------------------------------------------------------- 1 | from abc import abstractmethod 2 | from dataclasses import dataclass 3 | from enum import Enum 4 | from typing import Dict, Optional, Type, Union 5 | 6 | 7 | class PatchType(Enum): 8 | """Types of console requests sent to Django and passed along to pynode.""" 9 | 10 | MODEL_CODE = "model_code" 11 | PYTHON_REQUIREMENT = "python_requirement" 12 | SYSTEM_PACKAGE = "system_package" 13 | 14 | 15 | class Action(Enum): 16 | """Types of console requests sent to Django and passed along to pynode.""" 17 | 18 | ADD = "ADD" 19 | UPDATE = "UPDATE" 20 | REMOVE = "REMOVE" 21 | 22 | 23 | @dataclass 24 | class PatchBody: 25 | """Marker class""" 26 | 27 | action: Action 28 | 29 | @abstractmethod 30 | def to_dict(self): 31 | pass 32 | 33 | 34 | @dataclass 35 | class ModelCodePatch(PatchBody): 36 | path: str # Relative to model module directory 37 | content: Optional[str] = None 38 | 39 | def to_dict(self): 40 | return { 41 | "action": self.action.value, 42 | "path": self.path, 43 | "content": self.content, 44 | } 45 | 46 | @staticmethod 47 | def from_dict(patch_dict: Dict): 48 | action_str = patch_dict["action"] 49 | return ModelCodePatch( 50 | action=Action[action_str], 51 | path=patch_dict["path"], 52 | content=patch_dict["content"], 53 | ) 54 | 55 | 56 | @dataclass 57 | class PythonRequirementPatch(PatchBody): 58 | # For uninstall this should just be the name of the package, but for update 59 | # should be the full line in the requirements.txt format. 60 | requirement: str 61 | 62 | def to_dict(self): 63 | return { 64 | "action": self.action.value, 65 | "requirement": self.requirement, 66 | } 67 | 68 | @staticmethod 69 | def from_dict(patch_dict: Dict): 70 | action_str = patch_dict["action"] 71 | return PythonRequirementPatch( 72 | action=Action[action_str], 73 | requirement=patch_dict["requirement"], 74 | ) 75 | 76 | 77 | @dataclass 78 | class SystemPackagePatch(PatchBody): 79 | # For uninstall this should just be the name of the package, but for update 80 | # should be the full line in the requirements.txt format. 81 | package: str 82 | 83 | def to_dict(self): 84 | return { 85 | "action": self.action.value, 86 | "package": self.package, 87 | } 88 | 89 | @staticmethod 90 | def from_dict(patch_dict: Dict): 91 | action_str = patch_dict["action"] 92 | return SystemPackagePatch( 93 | action=Action[action_str], 94 | package=patch_dict["package"], 95 | ) 96 | 97 | 98 | PATCH_BODY_BY_TYPE: Dict[ 99 | PatchType, 100 | Type[Union[ModelCodePatch, PythonRequirementPatch, SystemPackagePatch]], 101 | ] = { 102 | PatchType.MODEL_CODE: ModelCodePatch, 103 | PatchType.PYTHON_REQUIREMENT: PythonRequirementPatch, 104 | PatchType.SYSTEM_PACKAGE: SystemPackagePatch, 105 | } 106 | 107 | 108 | @dataclass 109 | class Patch: 110 | """Request to execute code on console.""" 111 | 112 | type: PatchType 113 | body: PatchBody 114 | 115 | def to_dict(self): 116 | return { 117 | "type": self.type.value, 118 | "body": self.body.to_dict(), 119 | } 120 | 121 | @staticmethod 122 | def from_dict(patch_dict: Dict): 123 | typ = PatchType(patch_dict["type"]) 124 | body = PATCH_BODY_BY_TYPE[typ].from_dict(patch_dict["body"]) 125 | return Patch(typ, body) 126 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/tests/test_config.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from truss.truss_config import ( 3 | DEFAULT_CPU, 4 | DEFAULT_MEMORY, 5 | DEFAULT_USE_GPU, 6 | Accelerator, 7 | AcceleratorSpec, 8 | BaseImage, 9 | Resources, 10 | ) 11 | 12 | 13 | @pytest.mark.parametrize( 14 | "input_dict, expect_resources, output_dict", 15 | [ 16 | ( 17 | {}, 18 | Resources(), 19 | { 20 | "cpu": DEFAULT_CPU, 21 | "memory": DEFAULT_MEMORY, 22 | "use_gpu": DEFAULT_USE_GPU, 23 | "accelerator": None, 24 | }, 25 | ), 26 | ( 27 | {"accelerator": None}, 28 | Resources(), 29 | { 30 | "cpu": DEFAULT_CPU, 31 | "memory": DEFAULT_MEMORY, 32 | "use_gpu": DEFAULT_USE_GPU, 33 | "accelerator": None, 34 | }, 35 | ), 36 | ( 37 | {"accelerator": "V100"}, 38 | Resources(accelerator=AcceleratorSpec(Accelerator.V100, 1), use_gpu=True), 39 | { 40 | "cpu": DEFAULT_CPU, 41 | "memory": DEFAULT_MEMORY, 42 | "use_gpu": True, 43 | "accelerator": "V100", 44 | }, 45 | ), 46 | ( 47 | {"accelerator": "T4:1"}, 48 | Resources(accelerator=AcceleratorSpec(Accelerator.T4, 1), use_gpu=True), 49 | { 50 | "cpu": DEFAULT_CPU, 51 | "memory": DEFAULT_MEMORY, 52 | "use_gpu": True, 53 | "accelerator": "T4", 54 | }, 55 | ), 56 | ( 57 | {"accelerator": "A10G:4"}, 58 | Resources(accelerator=AcceleratorSpec(Accelerator.A10G, 4), use_gpu=True), 59 | { 60 | "cpu": DEFAULT_CPU, 61 | "memory": DEFAULT_MEMORY, 62 | "use_gpu": True, 63 | "accelerator": "A10G:4", 64 | }, 65 | ), 66 | ], 67 | ) 68 | def test_parse_resources(input_dict, expect_resources, output_dict): 69 | parsed_result = Resources.from_dict(input_dict) 70 | assert parsed_result == expect_resources 71 | assert parsed_result.to_dict() == output_dict 72 | 73 | 74 | @pytest.mark.parametrize( 75 | "input_str, expected_acc", 76 | [ 77 | (None, AcceleratorSpec(None, 0)), 78 | ("T4", AcceleratorSpec(Accelerator.T4, 1)), 79 | ("A10G:4", AcceleratorSpec(Accelerator.A10G, 4)), 80 | ("A100:8", AcceleratorSpec(Accelerator.A100, 8)), 81 | ], 82 | ) 83 | def test_acc_spec_from_str(input_str, expected_acc): 84 | assert AcceleratorSpec.from_str(input_str) == expected_acc 85 | 86 | 87 | @pytest.mark.parametrize( 88 | "input_dict, expect_base_image, output_dict", 89 | [ 90 | ( 91 | {}, 92 | BaseImage(), 93 | { 94 | "image": "", 95 | "python_executable_path": "", 96 | }, 97 | ), 98 | ( 99 | {"image": "custom_base_image", "python_executable_path": "/path/python"}, 100 | BaseImage(image="custom_base_image", python_executable_path="/path/python"), 101 | { 102 | "image": "custom_base_image", 103 | "python_executable_path": "/path/python", 104 | }, 105 | ), 106 | ], 107 | ) 108 | def test_parse_base_image(input_dict, expect_base_image, output_dict): 109 | parsed_result = BaseImage.from_dict(input_dict) 110 | assert parsed_result == expect_base_image 111 | assert parsed_result.to_dict() == output_dict 112 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/server/common/serialization.py: -------------------------------------------------------------------------------- 1 | import json 2 | import uuid 3 | from datetime import date, datetime, time, timedelta 4 | from decimal import Decimal 5 | from typing import Any, Callable, Dict, Optional, Union 6 | 7 | 8 | # mostly cribbed from django.core.serializer.DjangoJSONEncoder 9 | def truss_msgpack_encoder( 10 | obj: Union[Decimal, date, time, timedelta, uuid.UUID, Dict], 11 | chain: Optional[Callable] = None, 12 | ) -> Dict: 13 | if isinstance(obj, datetime): 14 | r = obj.isoformat() 15 | if r.endswith("+00:00"): 16 | r = r[:-6] + "Z" 17 | return {b"__dt_datetime_iso__": True, b"data": r} 18 | elif isinstance(obj, date): 19 | r = obj.isoformat() 20 | return {b"__dt_date_iso__": True, b"data": r} 21 | elif isinstance(obj, time): 22 | if obj.utcoffset() is not None: 23 | raise ValueError("Cannot represent timezone-aware times.") 24 | r = obj.isoformat() 25 | return {b"__dt_time_iso__": True, b"data": r} 26 | elif isinstance(obj, timedelta): 27 | return { 28 | b"__dt_timedelta__": True, 29 | b"data": (obj.days, obj.seconds, obj.microseconds), 30 | } 31 | elif isinstance(obj, Decimal): 32 | return {b"__decimal__": True, b"data": str(obj)} 33 | elif isinstance(obj, uuid.UUID): 34 | return {b"__uuid__": True, b"data": str(obj)} 35 | else: 36 | return obj if chain is None else chain(obj) 37 | 38 | 39 | def truss_msgpack_decoder(obj: Any, chain=None): 40 | try: 41 | if b"__dt_datetime_iso__" in obj: 42 | return datetime.fromisoformat(obj[b"data"]) 43 | elif b"__dt_date_iso__" in obj: 44 | return date.fromisoformat(obj[b"data"]) 45 | elif b"__dt_time_iso__" in obj: 46 | return time.fromisoformat(obj[b"data"]) 47 | elif b"__dt_timedelta__" in obj: 48 | days, seconds, microseconds = obj[b"data"] 49 | return timedelta(days=days, seconds=seconds, microseconds=microseconds) 50 | elif b"__decimal__" in obj: 51 | return Decimal(obj[b"data"]) 52 | elif b"__uuid__" in obj: 53 | return uuid.UUID(obj[b"data"]) 54 | else: 55 | return obj if chain is None else chain(obj) 56 | except KeyError: 57 | return obj if chain is None else chain(obj) 58 | 59 | 60 | # this json object is JSONType + np.array + datetime 61 | def is_truss_serializable(obj) -> bool: 62 | import numpy as np 63 | 64 | # basic JSON types 65 | if isinstance(obj, (str, int, float, bool, type(None), dict, list)): 66 | return True 67 | elif isinstance(obj, (datetime, date, time, timedelta)): 68 | return True 69 | elif isinstance(obj, np.ndarray): 70 | return True 71 | else: 72 | return False 73 | 74 | 75 | def truss_msgpack_serialize(obj): 76 | import msgpack 77 | import msgpack_numpy as mp_np 78 | 79 | return msgpack.packb( 80 | obj, default=lambda x: truss_msgpack_encoder(x, chain=mp_np.encode) 81 | ) 82 | 83 | 84 | def truss_msgpack_deserialize(obj): 85 | import msgpack 86 | import msgpack_numpy as mp_np 87 | 88 | return msgpack.unpackb( 89 | obj, object_hook=lambda x: truss_msgpack_decoder(x, chain=mp_np.decode) 90 | ) 91 | 92 | 93 | class DeepNumpyEncoder(json.JSONEncoder): 94 | def default(self, obj): 95 | import numpy as np 96 | 97 | if isinstance(obj, np.integer): 98 | return int(obj) 99 | elif isinstance(obj, np.floating): 100 | return float(obj) 101 | elif isinstance(obj, np.ndarray): 102 | return obj.tolist() 103 | else: 104 | return super(DeepNumpyEncoder, self).default(obj) 105 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/templates/training/job.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | import inspect 3 | import os 4 | import sys 5 | from pathlib import Path 6 | 7 | import yaml 8 | from shared.secrets_resolver import SecretsResolver 9 | 10 | CONFIG_FILE = "config.yaml" 11 | 12 | # This is where user training module will be mounted 13 | TRAINING_CODE_PATH = "/train" 14 | OUTPUT_PATH = "/output" 15 | VARIABLES_PATH = "/variables" 16 | VARIABLES_FILE = "variables.yaml" 17 | 18 | 19 | def _signature_accepts_keyword_arg(signature: inspect.Signature, kwarg: str) -> bool: 20 | return kwarg in signature.parameters or _signature_accepts_kwargs(signature) 21 | 22 | 23 | # todo: avoid duplication with model_wrapper 24 | def _signature_accepts_kwargs(signature: inspect.Signature) -> bool: 25 | for param in signature.parameters.values(): 26 | if param.kind == inspect.Parameter.VAR_KEYWORD: 27 | return True 28 | return False 29 | 30 | 31 | def _add_bundled_packages_to_path(config) -> None: 32 | if "bundled_packages_dir" in config: 33 | bundled_packages_path = Path("/packages") 34 | if bundled_packages_path.exists(): 35 | sys.path.append(str(bundled_packages_path)) 36 | 37 | 38 | def _load_train_class(config): 39 | if not Path(TRAINING_CODE_PATH).exists(): 40 | raise ValueError("Training code not mounted") 41 | 42 | # This should be created via mounting, but if not then create it. 43 | if not Path(OUTPUT_PATH).exists(): 44 | Path(OUTPUT_PATH).mkdir() 45 | 46 | sys.path.append(TRAINING_CODE_PATH) 47 | train_config = config["train"] 48 | training_module_name = str( 49 | Path(train_config["training_class_filename"]).with_suffix("") 50 | ) 51 | train_module = importlib.import_module(f"{training_module_name}") 52 | return getattr(train_module, train_config["training_class_name"]) 53 | 54 | 55 | def _create_trainer(config): 56 | train_class = _load_train_class(config) 57 | train_class_signature = inspect.signature(train_class) 58 | train_init_params = {} 59 | if _signature_accepts_keyword_arg(train_class_signature, "config"): 60 | train_init_params["config"] = config 61 | if _signature_accepts_keyword_arg(train_class_signature, "output_dir"): 62 | train_init_params["output_dir"] = Path(OUTPUT_PATH) 63 | 64 | # Wire up secrets 65 | if _signature_accepts_keyword_arg(train_class_signature, "secrets"): 66 | train_init_params["secrets"] = SecretsResolver.get_secrets(config) 67 | 68 | # Wire up variables 69 | if _signature_accepts_keyword_arg(train_class_signature, "variables"): 70 | default_variables = {} 71 | if "train" in config and "variables" in config["train"]: 72 | default_variables = config["train"]["variables"] 73 | 74 | runtime_variables = {} 75 | vars_path = Path(VARIABLES_PATH) / VARIABLES_FILE 76 | if vars_path.exists(): 77 | with vars_path.open() as vars_file: 78 | runtime_variables = yaml.safe_load(vars_file) 79 | 80 | variables = { 81 | **default_variables, 82 | **runtime_variables, 83 | } 84 | train_init_params["variables"] = variables 85 | return train_class(**train_init_params) 86 | 87 | 88 | if __name__ == "__main__": 89 | with open(CONFIG_FILE, encoding="utf-8") as config_file: 90 | truss_config = yaml.safe_load(config_file) 91 | _add_bundled_packages_to_path(truss_config) 92 | sys.path.append(os.environ["APP_HOME"]) 93 | trainer = _create_trainer(truss_config) 94 | if hasattr(trainer, "pre_train"): 95 | trainer.pre_train() 96 | 97 | # train is a required method, so no check 98 | trainer.train() 99 | 100 | if hasattr(trainer, "post_train"): 101 | trainer.post_train() 102 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/docker.py: -------------------------------------------------------------------------------- 1 | import enum 2 | import logging 3 | from typing import TYPE_CHECKING, Any, Dict, List 4 | 5 | if TYPE_CHECKING: 6 | from python_on_whales.components.container.cli_wrapper import Container 7 | 8 | from truss.constants import TRUSS_DIR 9 | from truss.local.local_config_handler import LocalConfigHandler 10 | 11 | 12 | class Docker: 13 | _client = None 14 | 15 | @staticmethod 16 | def client(): 17 | if Docker._client is None: 18 | from python_on_whales import DockerClient, docker 19 | 20 | if LocalConfigHandler.get_config().use_sudo: 21 | Docker._client = DockerClient(client_call=["sudo", "docker"]) 22 | else: 23 | Docker._client = docker 24 | return Docker._client 25 | 26 | 27 | def get_containers(labels: Dict, all: bool = False): 28 | """Gets containers given labels.""" 29 | return Docker.client().container.list( 30 | filters=_create_label_filters(labels), all=all 31 | ) 32 | 33 | 34 | def get_images(labels: Dict): 35 | """Gets images given labels.""" 36 | return Docker.client().image.list(filters=_create_label_filters(labels)) 37 | 38 | 39 | def get_urls_from_container(container_details) -> Dict[int, List[str]]: 40 | """Gets url where docker container is hosted.""" 41 | if ( 42 | container_details.network_settings is None 43 | or container_details.network_settings.ports is None 44 | ): 45 | return {} 46 | ports = container_details.network_settings.ports 47 | 48 | def parse_port(port_protocol) -> int: 49 | return int(port_protocol.split("/")[0]) 50 | 51 | def url_from_port_protocol_value(port_protocol_value: Dict[str, str]) -> str: 52 | return ( 53 | "http://" 54 | + port_protocol_value["HostIp"] 55 | + ":" 56 | + port_protocol_value["HostPort"] 57 | ) 58 | 59 | def urls_from_port_protocol_values( 60 | port_protocol_values: List[Dict[str, str]] 61 | ) -> List[str]: 62 | return [url_from_port_protocol_value(v) for v in port_protocol_values] 63 | 64 | return { 65 | parse_port(port_protocol): urls_from_port_protocol_values(value) 66 | for port_protocol, value in ports.items() 67 | if value is not None 68 | } 69 | 70 | 71 | def kill_containers(labels: Dict[str, bool]) -> None: 72 | from python_on_whales.exceptions import DockerException 73 | 74 | containers = get_containers(labels) 75 | for container in containers: 76 | container_labels = container.config.labels 77 | if TRUSS_DIR in container_labels: 78 | truss_dir = container_labels[TRUSS_DIR] 79 | logging.info(f"Killing Container: {container.id} for {truss_dir}") 80 | try: 81 | Docker.client().container.kill(containers) 82 | except DockerException: 83 | # The container may have stopped running by this point, this path 84 | # is for catching that. Unfortunately, there's no separate exception 85 | # for this scenario, so we catch the general one. Specific exceptions 86 | # such as NoSuchContainer are still allowed to error out. 87 | pass 88 | 89 | 90 | def get_container_logs(container, follow, stream): 91 | return Docker.client().container.logs(container, follow=follow, stream=stream) 92 | 93 | 94 | class DockerStates(enum.Enum): 95 | CREATED = "created" 96 | RUNNING = "running" 97 | PAUSED = "paused" 98 | RESTARTING = "restarting" 99 | OOMKILLED = "oomkilled" 100 | DEAD = "dead" 101 | EXITED = "exited" 102 | 103 | 104 | def inspect_container(container) -> "Container": 105 | """Inspects truss container""" 106 | return Docker.client().container.inspect(container) 107 | 108 | 109 | def get_container_state(container) -> DockerStates: 110 | "Get state of the container" 111 | return DockerStates(inspect_container(container).state.status) 112 | 113 | 114 | def _create_label_filters(labels: Dict) -> Dict[str, Any]: 115 | return { 116 | f"label={label_key}": label_value for label_key, label_value in labels.items() 117 | } 118 | -------------------------------------------------------------------------------- /model_to_docker/truss_saver/truss/environment_inference/requirements_inference.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | import types 3 | from itertools import dropwhile 4 | from typing import Optional, Set 5 | 6 | import pkg_resources # type: ignore 7 | from pkg_resources import WorkingSet # type: ignore 8 | 9 | # Some packages are weird and have different 10 | # imported names vs. system/pip names. Unfortunately, 11 | # there is no systematic way to get pip names from 12 | # a package's imported name. You'll have to add 13 | # exceptions to this list manually! 14 | POORLY_NAMED_PACKAGES = {"PIL": "Pillow", "sklearn": "scikit-learn"} 15 | 16 | # We don't want a few foundation packages 17 | IGNORED_PACKAGES = {"baseten", "pip", "truss", "pluggy", "pytest", "py"} 18 | 19 | TOP_LEVEL_NAMESPACES_TO_DROP_FOR_INFERENCE = ["truss", "baseten"] 20 | 21 | 22 | def infer_deps(must_include_deps: Optional[Set[str]] = None) -> Set[str]: 23 | """Infers the depedencies based on imports into the global namespace 24 | 25 | Args: 26 | must_include_deps (Set, optional): The set of package names that 27 | must necessarily be imported. 28 | Defaults to None. 29 | 30 | Returns: 31 | Set[str]: set of required python requirements, including versions. E.g. `{"xgboost==1.6.1"}` 32 | """ 33 | 34 | # Find the stack frame that likely has the relevant global inputs 35 | stack = inspect.stack() 36 | try: 37 | relevant_stack = _filter_truss_frames(stack) 38 | except StopIteration: 39 | return set() 40 | 41 | if not must_include_deps: 42 | must_include_deps = set() 43 | 44 | pkg_candidates = _extract_packages_from_frame(relevant_stack[0].frame) 45 | imports = must_include_deps.union(pkg_candidates) 46 | requirements = set([]) 47 | 48 | # Must refresh working set manually to get latest installed 49 | pkg_resources.working_set = ( 50 | WorkingSet._build_master() # type: ignore # pylint: disable=protected-access 51 | ) 52 | 53 | # Cross-check the names of installed packages vs. imported packages to get versions 54 | for pkg_in_frame in pkg_resources.working_set: 55 | if ( 56 | pkg_in_frame.project_name in imports 57 | and pkg_in_frame.project_name not in IGNORED_PACKAGES 58 | ): 59 | requirements.add(f"{pkg_in_frame.project_name}=={pkg_in_frame.version}") 60 | # Remove the package from imports as it was added into requirements 61 | imports.remove(pkg_in_frame.project_name) 62 | 63 | # Add the must include deps not found in frame to requirements 64 | deps_not_found_in_frame = imports.intersection(must_include_deps) 65 | requirements = requirements.union(deps_not_found_in_frame) 66 | 67 | return requirements 68 | 69 | 70 | def _filter_truss_frames(stack_frames): 71 | def is_truss_invocation_frame(stack_frame): 72 | module = inspect.getmodule(stack_frame.frame) 73 | if not module: 74 | return False 75 | 76 | module_name = module.__name__ 77 | for namespace in TOP_LEVEL_NAMESPACES_TO_DROP_FOR_INFERENCE: 78 | if module_name.startswith(f"{namespace}."): 79 | return True 80 | return False 81 | 82 | return list(dropwhile(is_truss_invocation_frame, stack_frames)) 83 | 84 | 85 | def _extract_packages_from_frame(frame) -> Set[str]: 86 | candidate_symbols = {**frame.f_globals, **frame.f_locals} 87 | 88 | pkg_names = set() 89 | for name, val in candidate_symbols.items(): 90 | if name.startswith("__"): 91 | continue 92 | 93 | if isinstance(val, types.ModuleType) and val.__name__ is not None: 94 | # Split ensures you get root package, 95 | # not just imported function 96 | pkg_name = val.__name__.split(".")[0] 97 | elif hasattr(val, "__module__") and val.__module__ is not None: 98 | pkg_name = val.__module__.split(".")[0] 99 | else: 100 | continue 101 | 102 | if pkg_name in POORLY_NAMED_PACKAGES: 103 | pkg_name = POORLY_NAMED_PACKAGES[pkg_name] 104 | 105 | pkg_names.add(pkg_name) 106 | 107 | return pkg_names 108 | -------------------------------------------------------------------------------- /requires-install.txt: -------------------------------------------------------------------------------- 1 | absl-py==1.4.0 2 | alembic==1.11.1 3 | anyio==3.7.0 4 | appnope==0.1.3 5 | astunparse==1.6.3 6 | attrs==23.1.0 7 | backcall==0.2.0 8 | beautifulsoup4==4.12.2 9 | black==22.12.0 10 | blake3==0.3.3 11 | bleach==6.0.0 12 | blinker==1.6.2 13 | build==1.0.3 14 | CacheControl==0.13.1 15 | cachetools==5.3.1 16 | certifi==2023.5.7 17 | cffi==1.16.0 18 | cfgv==3.3.1 19 | charset-normalizer==3.1.0 20 | cleo==2.1.0 21 | click==8.1.3 22 | cloudpickle==2.2.1 23 | comm==0.1.3 24 | coverage==6.5.0 25 | crashtest==0.4.1 26 | databricks-cli==0.17.6 27 | debugpy==1.6.7 28 | decorator==5.1.1 29 | defusedxml==0.7.1 30 | distlib==0.3.6 31 | docker==6.1.2 32 | dockerfile==3.2.0 33 | dulwich==0.21.6 34 | entrypoints==0.4 35 | exceptiongroup==1.1.1 36 | fastapi==0.95.2 37 | fastjsonschema 38 | filelock==3.12.0 39 | flake8==4.0.1 40 | Flask==2.3.2 41 | flatbuffers==23.5.26 42 | fsspec==2023.5.0 43 | gast==0.4.0 44 | gitdb==4.0.10 45 | GitPython==3.1.31 46 | google-auth==2.17.3 47 | google-auth-oauthlib==1.0.0 48 | google-pasta==0.2.0 49 | greenlet==2.0.2 50 | grpcio==1.54.2 51 | gunicorn==20.1.0 52 | h11==0.14.0 53 | h5py==3.8.0 54 | huggingface-hub==0.14.1 55 | identify==2.5.24 56 | idna==3.4 57 | importlib-metadata==5.2.0 58 | iniconfig==2.0.0 59 | installer==0.7.0 60 | ipdb==0.13.13 61 | ipykernel==6.23.1 62 | ipython==7.34.0 63 | isort==5.12.0 64 | itsdangerous==2.1.2 65 | jaraco.classes==3.3.0 66 | jax==0.4.10 67 | jedi==0.18.2 68 | Jinja2==3.1.2 69 | joblib==1.2.0 70 | jsonschema==4.17.3 71 | jupyter_client==8.2.0 72 | jupyter_core==5.3.0 73 | jupyterlab-pygments==0.2.2 74 | keras==2.12.0 75 | keyring==24.2.0 76 | libclang==16.0.0 77 | lightgbm==3.3.5 78 | Mako==1.2.4 79 | Markdown==3.4.3 80 | MarkupSafe==2.1.2 81 | matplotlib-inline==0.1.6 82 | mccabe==0.6.1 83 | mistune==2.0.5 84 | ml-dtypes==0.1.0 85 | mlflow==1.30.1 86 | more-itertools==10.1.0 87 | msgpack==1.0.5 88 | msgpack-numpy==0.4.8 89 | mypy==1.3.0 90 | mypy-extensions==1.0.0 91 | nbclient==0.8.0 92 | nbconvert==7.4.0 93 | nbformat==5.8.0 94 | nest-asyncio==1.5.6 95 | nodeenv==1.8.0 96 | numpy==1.23.5 97 | oauthlib==3.2.2 98 | opt-einsum==3.3.0 99 | packaging==20.9 100 | pandas==1.5.2 101 | pandocfilters==1.5.0 102 | parso==0.8.3 103 | pathspec==0.11.1 104 | pexpect==4.8.0 105 | pickleshare==0.7.5 106 | pkginfo==1.9.6 107 | platformdirs==3.5.1 108 | pluggy==1.0.0 109 | poetry 110 | poetry-core 111 | poetry-plugin-export 112 | pre-commit==2.21.0 113 | prometheus-client==0.17.0 114 | prometheus-flask-exporter==0.22.4 115 | prompt-toolkit==3.0.38 116 | protobuf==4.23.2 117 | psutil==5.9.5 118 | ptyprocess==0.7.0 119 | pyasn1==0.5.0 120 | pyasn1-modules==0.3.0 121 | pycodestyle==2.8.0 122 | pycparser==2.21 123 | pydantic==1.10.8 124 | pyflakes==2.4.0 125 | Pygments==2.15.1 126 | PyJWT==2.7.0 127 | pyparsing==3.0.9 128 | pyproject_hooks==1.0.0 129 | pyrsistent==0.19.3 130 | pytest==7.2.0 131 | pytest-cov==3.0.0 132 | pytest-split==0.8.1 133 | python-dateutil==2.8.2 134 | python-json-logger==2.0.7 135 | python-on-whales==0.46.0 136 | pytz==2022.7.1 137 | PyYAML==6.0 138 | pyzmq==25.1.0 139 | querystring-parser==1.2.4 140 | rapidfuzz==3.5.2 141 | regex==2023.5.5 142 | requests==2.31.0 143 | requests-oauthlib==1.3.1 144 | requests-toolbelt==1.0.0 145 | rsa==4.9 146 | scikit-learn==1.0.2 147 | scipy==1.10.1 148 | shellingham==1.5.4 149 | single-source==0.3.0 150 | six==1.16.0 151 | smmap==5.0.0 152 | sniffio==1.3.0 153 | soupsieve==2.4.1 154 | SQLAlchemy==1.4.48 155 | sqlparse==0.4.4 156 | starlette==0.27.0 157 | tabulate==0.9.0 158 | tenacity==8.2.2 159 | tensorboard==2.12.3 160 | tensorboard-data-server==0.7.0 161 | tensorflow-estimator==2.12.0 162 | tensorflow-hub==0.12.0 163 | tensorflow-io-gcs-filesystem==0.32.0 164 | tensorflow-macos==2.12.0 165 | termcolor==2.3.0 166 | threadpoolctl==3.1.0 167 | tinycss2==1.2.1 168 | tokenizers==0.13.3 169 | tomli==2.0.1 170 | tomlkit==0.12.2 171 | torch==1.13.1 172 | tornado==6.3.2 173 | tqdm==4.65.0 174 | traitlets==5.9.0 175 | transformers==4.29.2 176 | trove-classifiers==2023.11.7 177 | typer==0.9.0 178 | typing_extensions==4.6.2 179 | urllib3==2.0.2 180 | uvicorn==0.21.1 181 | virtualenv==20.23.0 182 | waitress==2.1.2 183 | wcwidth==0.2.6 184 | webencodings==0.5.1 185 | websocket-client==1.5.2 186 | Werkzeug==2.3.4 187 | wrapt==1.14.1 188 | xattr==0.10.1 189 | xgboost==1.7.5 190 | zipp==3.15.0 191 | --------------------------------------------------------------------------------