File size: 771 Bytes
baac5bb |
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 |
from threading import Lock
import execution
CPACK_LAST_ID = None
CPACK_OUTPUT_CACHE = {}
_lock = Lock()
def store_bentoml_value(func):
def wrapped(
inputs, class_def, unique_id, outputs=None, dynprompt=None, extra_data={}
):
global CPACK_LAST_ID
if getattr(class_def, "CPACK_NODE", False):
with _lock:
CPACK_LAST_ID = unique_id
if outputs is None:
outputs = CPACK_OUTPUT_CACHE
return func(inputs, class_def, unique_id, outputs, dynprompt, extra_data)
return wrapped
execution.get_input_data = store_bentoml_value(execution.get_input_data)
def set_bentoml_output(output):
with _lock:
CPACK_OUTPUT_CACHE[CPACK_LAST_ID] = output
|