wehub-resource-sync
587 lines
18 KiB
Python
587 lines
18 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
import multiprocessing as mp
|
|
import threading
|
|
import time
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.v1.multiprocess.futures import CUDAMessagingFuture, MessagingFuture
|
|
|
|
# ==============================================================================
|
|
# Helper Functions for CUDAMessagingFuture Tests
|
|
# ==============================================================================
|
|
|
|
|
|
def _create_cuda_event_in_process(event_queue: mp.Queue, delay: float = 0.0):
|
|
"""Helper process that creates a CUDA event and sends the IPC handle."""
|
|
torch.cuda.init()
|
|
if delay > 0:
|
|
time.sleep(delay)
|
|
|
|
# Create and record a CUDA event with interprocess flag
|
|
event = torch.cuda.Event(interprocess=True)
|
|
event.record()
|
|
event_bytes = event.ipc_handle()
|
|
|
|
# Send the event handle to the main process
|
|
event_queue.put(event_bytes)
|
|
|
|
|
|
def test_messaging_future_basic_usage():
|
|
"""Test basic usage of MessagingFuture: set result and retrieve it."""
|
|
future = MessagingFuture[int]()
|
|
|
|
# Initially, future should not be done
|
|
assert not future.query(), "Future should not be done initially"
|
|
|
|
# Set result
|
|
future.set_result(42)
|
|
|
|
# Future should now be done
|
|
assert future.query(), "Future should be done after setting result"
|
|
|
|
# Get result (should be immediate)
|
|
result = future.result(timeout=1)
|
|
assert result == 42, f"Expected result 42, got {result}"
|
|
|
|
|
|
def test_messaging_future_with_thread():
|
|
"""Test MessagingFuture with result set from another thread."""
|
|
future = MessagingFuture[str]()
|
|
|
|
def set_future_result():
|
|
time.sleep(0.5)
|
|
future.set_result("Hello from thread")
|
|
|
|
# Start thread that will set the result
|
|
thread = threading.Thread(target=set_future_result)
|
|
thread.start()
|
|
|
|
# Initially should not be done
|
|
assert not future.query(), "Future should not be done before thread sets result"
|
|
|
|
# Wait for result
|
|
result = future.result(timeout=2)
|
|
assert result == "Hello from thread", f"Expected 'Hello from thread', got {result}"
|
|
|
|
# Should be done now
|
|
assert future.query(), "Future should be done after getting result"
|
|
|
|
thread.join()
|
|
|
|
|
|
def test_messaging_future_wait_success():
|
|
"""Test wait method when result becomes available."""
|
|
future = MessagingFuture[int]()
|
|
|
|
def set_future_result():
|
|
time.sleep(0.3)
|
|
future.set_result(100)
|
|
|
|
thread = threading.Thread(target=set_future_result)
|
|
thread.start()
|
|
|
|
# Wait should return True when result is set
|
|
success = future.wait(timeout=1)
|
|
assert success, "Wait should return True when result is available"
|
|
assert future.query(), "Future should be done after wait returns True"
|
|
|
|
thread.join()
|
|
|
|
|
|
def test_messaging_future_wait_timeout():
|
|
"""Test wait method when timeout is reached."""
|
|
future = MessagingFuture[int]()
|
|
|
|
# Wait with short timeout (result never set)
|
|
start_time = time.time()
|
|
success = future.wait(timeout=0.2)
|
|
elapsed = time.time() - start_time
|
|
|
|
assert not success, "Wait should return False on timeout"
|
|
assert not future.query(), "Future should not be done after timeout"
|
|
assert 0.15 < elapsed < 0.3, f"Wait should respect timeout, elapsed: {elapsed}"
|
|
|
|
|
|
def test_messaging_future_result_timeout():
|
|
"""Test result method raises TimeoutError when timeout is reached."""
|
|
future = MessagingFuture[int]()
|
|
|
|
# Try to get result with timeout (result never set)
|
|
with pytest.raises(
|
|
TimeoutError, match="Future result not available within timeout"
|
|
):
|
|
future.result(timeout=0.2)
|
|
|
|
assert not future.query(), "Future should not be done after timeout"
|
|
|
|
|
|
def test_messaging_future_wait_no_timeout():
|
|
"""Test wait method without timeout (waits indefinitely until result is set)."""
|
|
future = MessagingFuture[float]()
|
|
|
|
def set_future_result():
|
|
time.sleep(0.3)
|
|
future.set_result(3.14)
|
|
|
|
thread = threading.Thread(target=set_future_result)
|
|
thread.start()
|
|
|
|
# Wait without timeout should wait until result is available
|
|
success = future.wait() # No timeout parameter
|
|
assert success, "Wait should return True when result is set"
|
|
assert future.result() == 3.14, "Result should be accessible after wait"
|
|
|
|
thread.join()
|
|
|
|
|
|
def test_messaging_future_multiple_result_calls():
|
|
"""Test that result can be retrieved multiple times after being set."""
|
|
future = MessagingFuture[str]()
|
|
future.set_result("persistent value")
|
|
|
|
# Get result multiple times
|
|
result1 = future.result(timeout=0.1)
|
|
result2 = future.result(timeout=0.1)
|
|
result3 = future.result(timeout=0.1)
|
|
|
|
assert result1 == result2 == result3 == "persistent value", (
|
|
"Result should be retrievable multiple times"
|
|
)
|
|
|
|
|
|
def test_messaging_future_complex_type():
|
|
"""Test MessagingFuture with complex types like lists and dicts."""
|
|
future = MessagingFuture[dict]()
|
|
|
|
complex_data = {"key1": [1, 2, 3], "key2": {"nested": "value"}, "key3": 42}
|
|
|
|
def set_future_result():
|
|
time.sleep(0.2)
|
|
future.set_result(complex_data)
|
|
|
|
thread = threading.Thread(target=set_future_result)
|
|
thread.start()
|
|
|
|
result = future.result(timeout=1)
|
|
assert result == complex_data, "Complex types should be preserved"
|
|
|
|
thread.join()
|
|
|
|
|
|
# ==============================================================================
|
|
# CUDAMessagingFuture Tests
|
|
# ==============================================================================
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="CUDA is required for CUDAMessagingFuture tests",
|
|
)
|
|
def test_cuda_messaging_future_basic_usage():
|
|
"""Test basic usage of CUDAMessagingFuture: create, wait, and get result."""
|
|
torch.cuda.init()
|
|
|
|
# Create CUDA event in a separate process
|
|
ctx = mp.get_context("spawn")
|
|
event_queue = ctx.Queue()
|
|
process = ctx.Process(target=_create_cuda_event_in_process, args=(event_queue,))
|
|
process.start()
|
|
|
|
# Get event bytes from the process
|
|
event_bytes = event_queue.get(timeout=30)
|
|
process.join(timeout=2)
|
|
|
|
# Create the raw future that will return (event_bytes, result_value)
|
|
raw_future = MessagingFuture[tuple[bytes, int]]()
|
|
|
|
# Create CUDAMessagingFuture from raw future
|
|
cuda_future = CUDAMessagingFuture.FromMessagingFuture(raw_future)
|
|
|
|
# Initially, future should not be done
|
|
assert not cuda_future.query(), "CUDAMessagingFuture should not be done initially"
|
|
|
|
# Set result in raw future
|
|
raw_future.set_result((event_bytes, 42))
|
|
|
|
# Wait for CUDA future to complete
|
|
success = cuda_future.wait()
|
|
assert success, "Wait should return True when result is available"
|
|
|
|
# Get result
|
|
result = cuda_future.result()
|
|
assert result == 42, f"Expected result 42, got {result}"
|
|
|
|
# Query should return True after completion
|
|
assert cuda_future.query(), "CUDAMessagingFuture should be done after wait"
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="CUDA is required for CUDAMessagingFuture tests",
|
|
)
|
|
def test_cuda_messaging_future_with_thread():
|
|
"""Test CUDAMessagingFuture with result set from another thread."""
|
|
torch.cuda.init()
|
|
|
|
# Create CUDA event in a separate process
|
|
ctx = mp.get_context("spawn")
|
|
event_queue = ctx.Queue()
|
|
process = ctx.Process(target=_create_cuda_event_in_process, args=(event_queue,))
|
|
process.start()
|
|
|
|
# Get event bytes from the process
|
|
event_bytes = event_queue.get(timeout=30)
|
|
process.join(timeout=2)
|
|
|
|
raw_future = MessagingFuture[tuple[bytes, str]]()
|
|
cuda_future = CUDAMessagingFuture.FromMessagingFuture(raw_future)
|
|
|
|
def set_future_result():
|
|
time.sleep(0.5)
|
|
raw_future.set_result((event_bytes, "Hello CUDA"))
|
|
|
|
# Start thread that will set the result
|
|
thread = threading.Thread(target=set_future_result)
|
|
thread.start()
|
|
|
|
# Initially should not be done
|
|
assert not cuda_future.query(), (
|
|
"Future should not be done before thread sets result"
|
|
)
|
|
|
|
# Wait for result
|
|
result = cuda_future.result()
|
|
assert result == "Hello CUDA", f"Expected 'Hello CUDA', got {result}"
|
|
|
|
# Should be done now
|
|
assert cuda_future.query(), "Future should be done after getting result"
|
|
|
|
thread.join()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="CUDA is required for CUDAMessagingFuture tests",
|
|
)
|
|
def test_cuda_messaging_future_wait_no_timeout():
|
|
"""Test wait method without timeout (waits indefinitely
|
|
until result is set).
|
|
"""
|
|
torch.cuda.init()
|
|
|
|
# Create CUDA event in a separate process
|
|
ctx = mp.get_context("spawn")
|
|
event_queue = ctx.Queue()
|
|
process = ctx.Process(target=_create_cuda_event_in_process, args=(event_queue,))
|
|
process.start()
|
|
|
|
# Get event bytes from the process
|
|
event_bytes = event_queue.get(timeout=30)
|
|
process.join(timeout=2)
|
|
|
|
raw_future = MessagingFuture[tuple[bytes, float]]()
|
|
cuda_future = CUDAMessagingFuture.FromMessagingFuture(raw_future)
|
|
|
|
def set_future_result():
|
|
time.sleep(0.3)
|
|
raw_future.set_result((event_bytes, 3.14))
|
|
|
|
thread = threading.Thread(target=set_future_result)
|
|
thread.start()
|
|
|
|
# Wait without timeout should wait until result is available
|
|
success = cuda_future.wait() # No timeout parameter
|
|
assert success, "Wait should return True when result is set"
|
|
assert cuda_future.result() == 3.14, "Result should be accessible after wait"
|
|
|
|
thread.join()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="CUDA is required for CUDAMessagingFuture tests",
|
|
)
|
|
def test_cuda_messaging_future_wait_with_timeout_success():
|
|
"""Test that wait method works correctly with timeout when result is available."""
|
|
torch.cuda.init()
|
|
|
|
# Create CUDA event in a separate process
|
|
ctx = mp.get_context("spawn")
|
|
event_queue = ctx.Queue()
|
|
process = ctx.Process(target=_create_cuda_event_in_process, args=(event_queue,))
|
|
process.start()
|
|
|
|
# Get event bytes from the process
|
|
event_bytes = event_queue.get(timeout=30)
|
|
process.join(timeout=2)
|
|
|
|
raw_future = MessagingFuture[tuple[bytes, int]]()
|
|
cuda_future = CUDAMessagingFuture.FromMessagingFuture(raw_future)
|
|
|
|
def set_future_result():
|
|
time.sleep(0.3)
|
|
raw_future.set_result((event_bytes, 123))
|
|
|
|
thread = threading.Thread(target=set_future_result)
|
|
thread.start()
|
|
|
|
# Wait with timeout should return True when result is available
|
|
success = cuda_future.wait(timeout=2.0)
|
|
assert success, "Wait with timeout should return True when result is available"
|
|
assert cuda_future.result() == 123, "Result should be accessible after wait"
|
|
|
|
thread.join()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="CUDA is required for CUDAMessagingFuture tests",
|
|
)
|
|
def test_cuda_messaging_future_wait_timeout_reached():
|
|
"""Test that wait method returns False when timeout is reached."""
|
|
torch.cuda.init()
|
|
|
|
raw_future = MessagingFuture[tuple[bytes, int]]()
|
|
cuda_future = CUDAMessagingFuture.FromMessagingFuture(raw_future)
|
|
|
|
# Wait with short timeout (result never set)
|
|
start_time = time.time()
|
|
success = cuda_future.wait(timeout=0.2)
|
|
elapsed = time.time() - start_time
|
|
|
|
assert not success, "Wait should return False on timeout"
|
|
assert not cuda_future.query(), "Future should not be done after timeout"
|
|
assert 0.15 < elapsed < 0.4, f"Wait should respect timeout, elapsed: {elapsed}"
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="CUDA is required for CUDAMessagingFuture tests",
|
|
)
|
|
def test_cuda_messaging_future_result_with_timeout_success():
|
|
"""Test that result method works correctly with timeout when result is available."""
|
|
torch.cuda.init()
|
|
|
|
# Create CUDA event in a separate process
|
|
ctx = mp.get_context("spawn")
|
|
event_queue = ctx.Queue()
|
|
process = ctx.Process(target=_create_cuda_event_in_process, args=(event_queue,))
|
|
process.start()
|
|
|
|
# Get event bytes from the process
|
|
event_bytes = event_queue.get(timeout=30)
|
|
process.join(timeout=2)
|
|
|
|
raw_future = MessagingFuture[tuple[bytes, int]]()
|
|
cuda_future = CUDAMessagingFuture.FromMessagingFuture(raw_future)
|
|
|
|
def set_future_result():
|
|
time.sleep(0.3)
|
|
raw_future.set_result((event_bytes, 456))
|
|
|
|
thread = threading.Thread(target=set_future_result)
|
|
thread.start()
|
|
|
|
# Get result with timeout should succeed when result is available
|
|
result = cuda_future.result(timeout=2.0)
|
|
assert result == 456, f"Expected result 456, got {result}"
|
|
|
|
thread.join()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="CUDA is required for CUDAMessagingFuture tests",
|
|
)
|
|
def test_cuda_messaging_future_result_timeout_reached():
|
|
"""Test that result method raises TimeoutError when timeout is reached."""
|
|
torch.cuda.init()
|
|
|
|
raw_future = MessagingFuture[tuple[bytes, int]]()
|
|
cuda_future = CUDAMessagingFuture.FromMessagingFuture(raw_future)
|
|
|
|
# Try to get result with timeout (result never set)
|
|
with pytest.raises(
|
|
TimeoutError, match="CUDAMessagingFuture result not available within timeout"
|
|
):
|
|
cuda_future.result(timeout=0.2)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="CUDA is required for CUDAMessagingFuture tests",
|
|
)
|
|
def test_cuda_messaging_future_multiple_result_calls():
|
|
"""Test that result can be retrieved multiple times after being set."""
|
|
torch.cuda.init()
|
|
|
|
# Create CUDA event in a separate process
|
|
ctx = mp.get_context("spawn")
|
|
event_queue = ctx.Queue()
|
|
process = ctx.Process(target=_create_cuda_event_in_process, args=(event_queue,))
|
|
process.start()
|
|
|
|
# Get event bytes from the process
|
|
event_bytes = event_queue.get(timeout=30)
|
|
process.join(timeout=2)
|
|
|
|
raw_future = MessagingFuture[tuple[bytes, str]]()
|
|
cuda_future = CUDAMessagingFuture.FromMessagingFuture(raw_future)
|
|
|
|
raw_future.set_result((event_bytes, "persistent cuda value"))
|
|
|
|
# Get result multiple times
|
|
result1 = cuda_future.result()
|
|
result2 = cuda_future.result()
|
|
result3 = cuda_future.result()
|
|
|
|
assert result1 == result2 == result3 == "persistent cuda value", (
|
|
"Result should be retrievable multiple times"
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="CUDA is required for CUDAMessagingFuture tests",
|
|
)
|
|
def test_cuda_messaging_future_query_before_and_after():
|
|
"""Test query method returns False before completion and True after."""
|
|
torch.cuda.init()
|
|
|
|
# Create CUDA event in a separate process
|
|
ctx = mp.get_context("spawn")
|
|
event_queue = ctx.Queue()
|
|
process = ctx.Process(target=_create_cuda_event_in_process, args=(event_queue,))
|
|
process.start()
|
|
|
|
# Get event bytes from the process
|
|
event_bytes = event_queue.get(timeout=30)
|
|
process.join(timeout=2)
|
|
|
|
raw_future = MessagingFuture[tuple[bytes, int]]()
|
|
cuda_future = CUDAMessagingFuture.FromMessagingFuture(raw_future)
|
|
|
|
# Query before setting result
|
|
assert not cuda_future.query(), "Query should return False before result is set"
|
|
|
|
# Set result
|
|
raw_future.set_result((event_bytes, 100))
|
|
|
|
# Wait for completion
|
|
cuda_future.wait()
|
|
|
|
# Query after setting result
|
|
assert cuda_future.query(), (
|
|
"Query should return True after result is set and waited"
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="CUDA is required for CUDAMessagingFuture tests",
|
|
)
|
|
def test_cuda_messaging_future_complex_type():
|
|
"""Test CUDAMessagingFuture with complex types like lists and dicts."""
|
|
torch.cuda.init()
|
|
|
|
# Create CUDA event in a separate process
|
|
ctx = mp.get_context("spawn")
|
|
event_queue = ctx.Queue()
|
|
process = ctx.Process(target=_create_cuda_event_in_process, args=(event_queue,))
|
|
process.start()
|
|
|
|
# Get event bytes from the process
|
|
event_bytes = event_queue.get(timeout=30)
|
|
process.join(timeout=2)
|
|
|
|
complex_data = {"key1": [1, 2, 3], "key2": {"nested": "value"}, "key3": 42}
|
|
|
|
raw_future = MessagingFuture[tuple[bytes, dict]]()
|
|
cuda_future = CUDAMessagingFuture.FromMessagingFuture(raw_future)
|
|
|
|
def set_future_result():
|
|
time.sleep(0.2)
|
|
raw_future.set_result((event_bytes, complex_data))
|
|
|
|
thread = threading.Thread(target=set_future_result)
|
|
thread.start()
|
|
|
|
result = cuda_future.result()
|
|
assert result == complex_data, "Complex types should be preserved"
|
|
|
|
thread.join()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="CUDA is required for CUDAMessagingFuture tests",
|
|
)
|
|
def test_messaging_future_to_cuda_future():
|
|
"""Test converting MessagingFuture to CUDAMessagingFuture
|
|
using to_cuda_future method.
|
|
"""
|
|
torch.cuda.init()
|
|
|
|
# Create CUDA event in a separate process
|
|
ctx = mp.get_context("spawn")
|
|
event_queue = ctx.Queue()
|
|
process = ctx.Process(target=_create_cuda_event_in_process, args=(event_queue,))
|
|
process.start()
|
|
|
|
# Get event bytes from the process
|
|
event_bytes = event_queue.get(timeout=30)
|
|
process.join(timeout=2)
|
|
|
|
raw_future = MessagingFuture[tuple[bytes, int]]()
|
|
|
|
# Convert to CUDA future
|
|
cuda_future = raw_future.to_cuda_future()
|
|
|
|
# Verify it's a CUDAMessagingFuture instance
|
|
assert isinstance(cuda_future, CUDAMessagingFuture), (
|
|
"to_cuda_future should return CUDAMessagingFuture instance"
|
|
)
|
|
|
|
# Set result and verify it works
|
|
raw_future.set_result((event_bytes, 999))
|
|
|
|
result = cuda_future.result()
|
|
assert result == 999, f"Expected result 999, got {result}"
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="CUDA is required for CUDAMessagingFuture tests",
|
|
)
|
|
def test_cuda_messaging_future_with_explicit_device():
|
|
"""Test CUDAMessagingFuture with explicit device parameter."""
|
|
torch.cuda.init()
|
|
|
|
# Create CUDA event in a separate process
|
|
ctx = mp.get_context("spawn")
|
|
event_queue = ctx.Queue()
|
|
process = ctx.Process(target=_create_cuda_event_in_process, args=(event_queue,))
|
|
process.start()
|
|
|
|
# Get event bytes from the process
|
|
event_bytes = event_queue.get(timeout=30)
|
|
process.join(timeout=2)
|
|
|
|
device = torch.cuda.current_device()
|
|
raw_future = MessagingFuture[tuple[bytes, str]]()
|
|
|
|
# Create CUDA future with explicit device
|
|
cuda_future = CUDAMessagingFuture.FromMessagingFuture(raw_future, device=device)
|
|
|
|
# Set result
|
|
raw_future.set_result((event_bytes, "explicit device"))
|
|
|
|
# Get result
|
|
result = cuda_future.result()
|
|
assert result == "explicit device", f"Expected 'explicit device', got {result}"
|