wehub-resource-sync
45 lines
1.7 KiB
Python
45 lines
1.7 KiB
Python
import asyncio
|
|
import inspect
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
from mlflow.utils.autologging_utils.safety import update_wrapper_extended
|
|
|
|
|
|
def asyncify(is_async):
|
|
"""
|
|
Decorator that converts a function to an async function if `is_async` is True. This is useful
|
|
for testing purposes, where we want to test both synchronous and asynchronous code paths.
|
|
"""
|
|
|
|
def decorator(fn):
|
|
if is_async:
|
|
|
|
async def async_fn(*args, **kwargs):
|
|
if args and inspect.iscoroutinefunction(args[0]):
|
|
original = args[0]
|
|
|
|
def wrapped_original(*og_args, **og_kwargs):
|
|
# Run the original async function in a separate thread. This is a workaround
|
|
# for the fact that we cannot use asyncio.run here because an event loop is
|
|
# already running in the main thread.
|
|
with ThreadPoolExecutor(thread_name_prefix="async-helper") as executor:
|
|
future = executor.submit(asyncio.run, original(*og_args, **og_kwargs))
|
|
return future.result()
|
|
|
|
args = (update_wrapper_extended(wrapped_original, original), *args[1:])
|
|
return fn(*args, **kwargs)
|
|
|
|
return update_wrapper_extended(async_fn, fn)
|
|
else:
|
|
return fn
|
|
|
|
return decorator
|
|
|
|
|
|
def run_sync_or_async(fn, *args, **kwargs):
|
|
"""Convenience function that runs a function synchronously regardless of whether it is async."""
|
|
if inspect.iscoroutinefunction(fn):
|
|
return asyncio.run(fn(*args, **kwargs))
|
|
else:
|
|
return fn(*args, **kwargs)
|