wehub-resource-sync
132 lines
3.8 KiB
Python
132 lines
3.8 KiB
Python
import multiprocessing
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from unittest import mock
|
|
|
|
from mlflow.tracing.export.async_export_queue import AsyncTraceExportQueue, Task
|
|
|
|
from tests.tracing.helper import skip_when_testing_trace_sdk
|
|
|
|
|
|
def test_async_queue_handle_tasks():
|
|
queue = AsyncTraceExportQueue()
|
|
|
|
counter = 0
|
|
|
|
def increment(delta):
|
|
nonlocal counter
|
|
counter += delta
|
|
|
|
for _ in range(10):
|
|
task = Task(handler=increment, args=(1,))
|
|
queue.put(task)
|
|
|
|
queue.flush(terminate=True)
|
|
assert counter == 10
|
|
|
|
|
|
def exporter_process(counter):
|
|
# This process exits before waiting for the tasks to finish
|
|
|
|
queue = AsyncTraceExportQueue()
|
|
|
|
def increment(counter):
|
|
time.sleep(1)
|
|
with counter.get_lock():
|
|
counter.value += 1
|
|
|
|
for _ in range(10):
|
|
task = Task(handler=increment, args=(counter,))
|
|
queue.put(task)
|
|
|
|
|
|
@skip_when_testing_trace_sdk
|
|
def test_async_queue_complete_task_process_finished():
|
|
multiprocessing.set_start_method("spawn", force=True)
|
|
counter = multiprocessing.Value("i", 0)
|
|
process = multiprocessing.Process(target=exporter_process, args=(counter,))
|
|
process.start()
|
|
process.join(timeout=15)
|
|
|
|
assert counter.value == 10
|
|
|
|
|
|
def test_async_queue_activate_thread_safe():
|
|
with mock.patch("atexit.register") as mock_atexit:
|
|
queue = AsyncTraceExportQueue()
|
|
|
|
def count_threads():
|
|
main_thread = threading.main_thread()
|
|
return sum(
|
|
t.is_alive()
|
|
for t in threading.enumerate()
|
|
if t is not main_thread and t.name.startswith("MLflowTraceLogging")
|
|
)
|
|
|
|
# 1. Validate activation
|
|
with ThreadPoolExecutor(
|
|
max_workers=10, thread_name_prefix="test-async-export-queue-activate"
|
|
) as executor:
|
|
for _ in range(10):
|
|
executor.submit(queue.activate)
|
|
assert count_threads() > 0 # Logging thread + max 5 worker threads
|
|
mock_atexit.assert_called_once()
|
|
mock_atexit.reset_mock()
|
|
|
|
# 2. Validate flush (continue)
|
|
queue.flush(terminate=False)
|
|
assert queue.is_active()
|
|
assert count_threads() > 0 # New threads should be created
|
|
mock_atexit.assert_not_called() # Exit callback should not be registered again
|
|
|
|
# 3. Validate flush with termination
|
|
with ThreadPoolExecutor(
|
|
max_workers=10, thread_name_prefix="test-async-export-queue-flush"
|
|
) as executor:
|
|
for _ in range(10):
|
|
executor.submit(queue.flush(terminate=True))
|
|
assert count_threads() == 0
|
|
|
|
|
|
def test_put_after_terminate_executes_synchronously():
|
|
queue = AsyncTraceExportQueue()
|
|
|
|
calls = []
|
|
queue.put(Task(handler=calls.append, args=(1,)))
|
|
queue.flush(terminate=True)
|
|
|
|
assert not queue.is_active()
|
|
assert queue._stop_event.is_set()
|
|
|
|
# Calling put() after termination must not deadlock; task must run synchronously.
|
|
queue.put(Task(handler=calls.append, args=(2,)))
|
|
|
|
assert calls == [1, 2]
|
|
|
|
|
|
def test_async_queue_drop_task_when_full(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_ASYNC_TRACE_LOGGING_MAX_QUEUE_SIZE", "3")
|
|
monkeypatch.setenv("MLFLOW_ASYNC_TRACE_LOGGING_MAX_WORKERS", "1")
|
|
|
|
queue = AsyncTraceExportQueue()
|
|
|
|
processed_tasks = 0
|
|
|
|
# Create a slow handler to keep tasks in the queue
|
|
def slow_handler():
|
|
time.sleep(0.5)
|
|
|
|
nonlocal processed_tasks
|
|
processed_tasks += 1
|
|
|
|
for _ in range(10):
|
|
task = Task(handler=slow_handler, args=())
|
|
queue.put(task)
|
|
|
|
queue.flush(terminate=True)
|
|
|
|
# One more task than the queue size might be processed, because the first task
|
|
# can be drained from the queue immediately, which creates a slot for another task
|
|
assert processed_tasks <= 4
|