-
Notifications
You must be signed in to change notification settings - Fork 118
/
service.py
67 lines (51 loc) · 2.06 KB
/
service.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
import atexit
import os
import bentoml
import numpy as np
from bentoml.io import NumpyNdarray, Text
import whylogs as why
# Set WhyLabs environment variables
os.environ["WHYLABS_API_KEY"] = "APIKEY"
os.environ["WHYLABS_DEFAULT_ORG_ID"] = "ORGID"
os.environ["WHYLABS_DEFAULT_DATASET_ID"] = "PROJECTID"
# Model Prediction Class Names
CLASS_NAMES = ["setosa", "versicolor", "virginica"]
# Load model from BentoML local model store & create a BentoML runner
iris_clf_runner = bentoml.sklearn.get("iris_knn:latest").to_runner()
svc = bentoml.Service("iris_classifier", runners=[iris_clf_runner])
# create a class to handle startup, shutdown, and prediction logic with Bentoml
class MyService:
def __init__(self):
# Initialize the whylogs rolling logger with a 5 minute interval
global logger
logger = why.logger(mode="rolling", interval=5, when="M", base_name="bentoml_predictions")
logger.append_writer("whylabs")
def on_exit_callback(self):
# Close the whylogs rolling logger when the service is shut down
logger.close()
# Create a BentoML API endpoint for model predictions & logging
@svc.api(
input=NumpyNdarray.from_sample(np.array([4.9, 3.0, 1.4, 0.2], dtype=np.double)),
output=Text(),
)
async def classify(features: np.ndarray):
results = await iris_clf_runner.predict.async_run([features])
probs = iris_clf_runner.predict_proba.run([features])
result = results[0]
category = CLASS_NAMES[result]
prob = max(probs[0])
# create a dict of data & prediction results w/ feature names
data = {
"sepal length": features[0],
"sepal width": features[1],
"petal length": features[2],
"petal width": features[3],
"class_output": category,
"proba_output": prob,
}
# Log data + model outputs to WhyLabs.ai
logger.log(data)
return category
my_service = MyService()
atexit.register(my_service.on_exit_callback)
# Run with: bentoml serve service:svc --reload