wehub-resource-sync
67 lines
2.1 KiB
Python
67 lines
2.1 KiB
Python
import random
|
|
import time
|
|
from threading import Thread
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.pyfunc.context import (
|
|
Context,
|
|
get_prediction_context,
|
|
set_prediction_context,
|
|
)
|
|
|
|
|
|
def test_prediction_context_thread_safe():
|
|
def set_context(context):
|
|
with set_prediction_context(context):
|
|
time.sleep(0.2 * random.random())
|
|
assert get_prediction_context() == context
|
|
context.update(is_evaluate=not context.is_evaluate)
|
|
assert get_prediction_context() == context
|
|
|
|
threads = []
|
|
for i in range(10):
|
|
context = Context(request_id=f"request_{i}", is_evaluate=random.choice([True, False]))
|
|
thread = Thread(name=f"test-pyfunc-context-{i}", target=set_context, args=(context,))
|
|
thread.start()
|
|
threads.append(thread)
|
|
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
assert get_prediction_context() is None
|
|
|
|
|
|
def test_set_prediction_context_raise_on_invalid_context():
|
|
with pytest.raises(TypeError, match="Expected context to be an instance of Context"):
|
|
with set_prediction_context("invalid"):
|
|
pass
|
|
|
|
|
|
def test_prediction_context_pyfunc_predict():
|
|
class MyModel(mlflow.pyfunc.PythonModel):
|
|
def predict(self, model_inputs):
|
|
context = get_prediction_context()
|
|
return context.request_id
|
|
|
|
def predict_stream(self, model_inputs):
|
|
for _ in range(3):
|
|
context = get_prediction_context()
|
|
yield context.request_id
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pyfunc.log_model(name="model", python_model=MyModel())
|
|
pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
|
|
with set_prediction_context(Context(request_id="request_1")):
|
|
assert pyfunc_model.predict(None) == "request_1"
|
|
|
|
with set_prediction_context(Context(request_id="request_2")):
|
|
generator = pyfunc_model.predict_stream(None)
|
|
|
|
# When a prediction context is set for predict_stream call, it should also
|
|
# be effective during the iteration of the generator.
|
|
for output in generator:
|
|
assert output == "request_2"
|