Merge origin/main into dev/vllm-align

Resolved conflicts:
- vllm_omni/inputs/preprocess.py: took origin/main's elif mm_processor_kwargs branch, fixed undefined mm_uuids reference
- vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py: took origin/main's NPU support

Made-with: Cursor

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/distributed/omni_connectors/test_adapter_and_flow.py

@@ -1,19 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from unittest.mock import MagicMock
-
 import pytest
+from pytest_mock import MockerFixture
 
 from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector, try_send_via_connector
 from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
 from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec, OmniTransferConfig
 from vllm_omni.distributed.omni_connectors.utils.initialization import get_connectors_config_for_stage
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 @pytest.fixture
-def mock_objects():
-    return {"connector": MagicMock(), "metrics": MagicMock(), "queue_fn": MagicMock()}
+def mock_objects(mocker: MockerFixture):
+    return {"connector": mocker.MagicMock(), "metrics": mocker.MagicMock(), "queue_fn": mocker.MagicMock()}
 
 
 def test_send_success(mock_objects):
@@ -138,7 +139,7 @@ def test_recv_no_connector():
     assert inputs is None
 
 
-def test_shm_connector_flow():
+def test_shm_connector_flow(mocker: MockerFixture):
     """
     Verify the full flow: Send -> Adapter -> Connector -> Adapter -> Recv.
     Using real SharedMemoryConnector (inline mode for simplicity).
@@ -161,7 +162,7 @@ def test_shm_connector_flow():
     def mock_submit(payload):
         queue_capture.append(payload)
 
-    mock_metrics = MagicMock()
+    mock_metrics = mocker.MagicMock()
 
     # 3. Send
     success = try_send_via_connector(
@@ -217,7 +218,7 @@ def test_get_connectors_for_stage():
     assert stage_2_config["from_stage_1"]["spec"]["name"] == "C2"
 
 
-def test_recv_with_missing_metadata():
+def test_recv_with_missing_metadata(mocker: MockerFixture):
     """Test recv when queue payload is malformed (missing metadata)."""
     # Connector expects metadata but task doesn't have it
     task = {
@@ -226,7 +227,7 @@ def test_recv_with_missing_metadata():
         "from_stage": "0",
         # Missing "connector_metadata"
     }
-    mock_conn = MagicMock()
+    mock_conn = mocker.MagicMock()
     # If get is called with None metadata, connector usually handles it or adapter handles exception
     mock_conn.get.side_effect = Exception("Get failed")
 

tests/distributed/omni_connectors/test_basic_connectors.py

@@ -1,14 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from unittest.mock import MagicMock
 
 import pytest
+from pytest_mock import MockerFixture
 
 from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
 from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
 from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
 
+# pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 def test_basic_serialization():
     """Test basic msgpack serialization."""
@@ -78,19 +80,19 @@ def test_put_get_inline(shm_connector):
     assert size == ret_size
 
 
-def test_put_get_shm(shm_connector, monkeypatch):
+def test_put_get_shm(mocker: MockerFixture, shm_connector, monkeypatch: pytest.MonkeyPatch):
     """Test SHM transfer logic for large data (Mocked)."""
     # Create data larger than 100 bytes
     data = {"large": "x" * 200}
 
     # Mock SHM return values
     mock_handle = {"name": "test_shm", "size": 200}
-    mock_write = MagicMock(return_value=mock_handle)
+    mock_write = mocker.MagicMock(return_value=mock_handle)
     monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_write_bytes", mock_write)
 
     # When reading, return the serialized bytes of the data
     serialized_data = shm_connector.serialize_obj(data)
-    mock_read = MagicMock(return_value=serialized_data)
+    mock_read = mocker.MagicMock(return_value=serialized_data)
     monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_read_bytes", mock_read)
 
     # Put

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -4,10 +4,10 @@
 import threading
 from collections import deque
 from types import SimpleNamespace
-from unittest.mock import MagicMock
 
 import pytest
 import torch
+from pytest_mock import MockerFixture
 from vllm.v1.request import RequestStatus
 
 from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
@@ -35,25 +35,25 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
         prompt_token_ids=[],
         num_computed_tokens=0,
         additional_information=None,
+        is_finished=lambda: status == RequestStatus.FINISHED_STOPPED,
     )
 
 
 @pytest.fixture
-def build_adapter(monkeypatch):
+def build_adapter(monkeypatch, mocker: MockerFixture):
     def _build(*, stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2):
-        connector = MagicMock()
+        connector = mocker.MagicMock()
         connector.stage_id = stage_id
         connector.get.return_value = None
         connector.put.return_value = (True, 1, {})
 
         def _fake_base_init(self, config):
             self.config = config
-            self._pending_load_reqs = {}
+            self._pending_load_reqs = deque()
             self._finished_load_reqs = set()
-            self._pending_save_reqs = {}
+            self._pending_save_reqs = deque()
             self._finished_save_reqs = set()
             self.stop_event = threading.Event()
-            self.lock = threading.Lock()
 
         monkeypatch.setattr(OmniTransferAdapterBase, "__init__", _fake_base_init)
         monkeypatch.setattr(
@@ -108,9 +108,8 @@ def test_load_poll(build_adapter):
     adapter.load_async(request)
     payload = {"code_predictor_codes": [[1]], "hidden_states": torch.tensor([[2.0]]), "finished": True}
     connector.get.return_value = (payload, 16)
-    adapter._poll_single_request("req-1")
+    adapter._poll_single_request(request)
 
-    connector.get.assert_called_once_with("1", "2", "external-1_1_0")
     assert request.additional_information == payload
     assert adapter.get_req_chunk["req-1"] == 1
     assert "req-1" in adapter._finished_load_reqs
@@ -120,17 +119,15 @@ def test_load_poll(build_adapter):
 
 def test_save_async(build_adapter):
     adapter, _ = build_adapter(stage_id=1)
-    request = SimpleNamespace(external_req_id="external-1")
+    request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
 
     adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": False}
     adapter.save_async(pooling_output=None, request=request)
     adapter.custom_process_next_stage_input_func = lambda **kwargs: {}
     adapter.save_async(pooling_output=None, request=request)
 
-    assert adapter.put_req_chunk["external-1"] == 1
-    queued = adapter._pending_save_reqs["external-1"]
-    assert len(queued) == 1
-    assert queued[0]["put_key"] == "external-1_1_0"
+    task = adapter._pending_save_reqs.popleft()
+    assert task["is_finished"] is False
 
 
 def test_update_request_payload(build_adapter):
@@ -180,3 +177,232 @@ def test_postprocess_scheduler_output(build_adapter):
     assert cached_info["cached-ready"] == {"k": "v"}
     assert cached_info["missing"] is None
     assert adapter.requests_with_ready_chunks == {"leftover"}
+
+
+# ---------------------------------------------------------------
+# Cleanup tests
+# ---------------------------------------------------------------
+
+
+def _populate_adapter_state(adapter, req_id="req-1", ext_id="ext-1"):
+    """Fill every per-request structure so cleanup can be verified."""
+    adapter.finished_requests.add(req_id)
+    adapter.get_req_chunk[req_id] = 3
+    adapter.requests_with_ready_chunks.add(req_id)
+    adapter.request_ids_mapping[req_id] = ext_id
+    adapter._pending_load_reqs.append(SimpleNamespace(request_id=req_id))
+    adapter._finished_load_reqs.add(req_id)
+
+    adapter.put_req_chunk[ext_id] = 5
+    adapter.request_payload[ext_id] = {"hidden": [1, 2]}
+    adapter.code_prompt_token_ids[ext_id] = [[10, 20]]
+
+
+def test_cleanup_clears_all_state(build_adapter):
+    """After cleanup, no per-request key should remain in any dict/set."""
+    adapter, _ = build_adapter(stage_id=1)
+    req_id, ext_id = "req-1", "ext-1"
+    _populate_adapter_state(adapter, req_id, ext_id)
+
+    adapter.cleanup(req_id, ext_id)
+
+    assert req_id not in adapter.finished_requests
+    assert req_id not in adapter.get_req_chunk
+    assert req_id not in adapter.requests_with_ready_chunks
+    assert req_id not in adapter.request_ids_mapping
+    assert all(getattr(r, "request_id", None) != req_id for r in adapter._pending_load_reqs)
+    assert req_id not in adapter._finished_load_reqs
+
+    assert ext_id not in adapter.put_req_chunk
+    assert ext_id not in adapter.request_payload
+    assert ext_id not in adapter.code_prompt_token_ids
+
+
+def test_cleanup_infers_external_id(build_adapter):
+    """When external_req_id is None, cleanup should look it up from the mapping."""
+    adapter, _ = build_adapter(stage_id=1)
+    req_id, ext_id = "req-2", "ext-2"
+    _populate_adapter_state(adapter, req_id, ext_id)
+
+    adapter.cleanup(req_id)
+
+    assert ext_id not in adapter.put_req_chunk
+    assert ext_id not in adapter.request_payload
+
+
+def test_cleanup_idempotent(build_adapter):
+    """Calling cleanup multiple times for the same (or nonexistent) request must not raise."""
+    adapter, _ = build_adapter(stage_id=1)
+
+    try:
+        adapter.cleanup("nonexistent")
+        adapter.cleanup("nonexistent")
+    except Exception as e:
+        pytest.fail(f"cleanup should be idempotent: {e}")
+
+    req_id, ext_id = "req-3", "ext-3"
+    _populate_adapter_state(adapter, req_id, ext_id)
+    adapter.cleanup(req_id, ext_id)
+
+    try:
+        adapter.cleanup(req_id, ext_id)
+    except Exception as e:
+        pytest.fail(f"second cleanup should be idempotent: {e}")
+
+
+def test_cleanup_request_id_reuse_not_polluted(build_adapter):
+    """After cleanup, reusing the same request_id must not be treated as finished."""
+    adapter, _ = build_adapter(stage_id=1)
+    req_id, ext_id = "req-reuse", "ext-reuse"
+    _populate_adapter_state(adapter, req_id, ext_id)
+
+    adapter.cleanup(req_id, ext_id)
+
+    assert req_id not in adapter.finished_requests
+    assert req_id not in adapter.get_req_chunk
+
+
+def test_cleanup_preserves_pending_save(build_adapter):
+    """Cleanup must NOT remove _pending_save_reqs to avoid losing unsent chunks."""
+    adapter, _ = build_adapter(stage_id=1)
+    req_id, ext_id = "req-4", "ext-4"
+    _populate_adapter_state(adapter, req_id, ext_id)
+
+    pending_task = {"put_key": f"{ext_id}_1_0", "data": {"x": 1}}
+    adapter._pending_save_reqs.append(pending_task)
+
+    adapter.cleanup(req_id, ext_id)
+
+    assert len(adapter._pending_save_reqs) == 1
+
+
+def test_cleanup_only_affects_target_request(build_adapter):
+    """Cleanup for one request must not affect another request's state."""
+    adapter, _ = build_adapter(stage_id=1)
+    _populate_adapter_state(adapter, "req-a", "ext-a")
+    _populate_adapter_state(adapter, "req-b", "ext-b")
+
+    adapter.cleanup("req-a", "ext-a")
+
+    assert "req-b" in adapter.finished_requests
+    assert "req-b" in adapter.get_req_chunk
+    assert "ext-b" in adapter.put_req_chunk
+    assert "ext-b" in adapter.request_payload
+    assert "ext-b" in adapter.code_prompt_token_ids
+    assert "req-b" in adapter.request_ids_mapping
+
+
+def test_cleanup_after_poll_flow(build_adapter):
+    """Simulate full load_async -> poll -> finished -> cleanup cycle."""
+    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
+    request = _req("req-flow", RequestStatus.WAITING, external_req_id="ext-flow")
+
+    adapter.load_async(request)
+
+    adapter.request_ids_mapping["req-flow"] = "ext-flow"
+    payload = {"hidden_states": torch.tensor([[1.0]]), "finished": True}
+    connector.get.return_value = (payload, 8)
+    adapter._poll_single_request(request)
+
+    assert "req-flow" in adapter.finished_requests
+    assert adapter.get_req_chunk["req-flow"] == 1
+    assert "req-flow" in adapter.request_ids_mapping
+
+    adapter.cleanup("req-flow", "ext-flow")
+
+    assert "req-flow" not in adapter.finished_requests
+    assert "req-flow" not in adapter.get_req_chunk
+    assert "req-flow" not in adapter.request_ids_mapping
+    assert "ext-flow" not in adapter.request_payload
+
+
+# ---------------------------------------------------------------
+# Scheduler trigger tests
+# ---------------------------------------------------------------
+
+
+class _HashableRequest(SimpleNamespace):
+    """SimpleNamespace that can be added to a set (needed by scheduler internals)."""
+
+    def __hash__(self):
+        return hash(self.request_id)
+
+    def __eq__(self, other):
+        return getattr(other, "request_id", None) == self.request_id
+
+
+def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: MockerFixture):
+    """OmniGenerationScheduler must call adapter.cleanup when request finishes."""
+    cleanup_calls = []
+
+    adapter_mock = mocker.MagicMock()
+    adapter_mock.finished_requests = {"req-s1"}
+    adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
+
+    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+    scheduler = mocker.MagicMock()
+    scheduler.chunk_transfer_adapter = adapter_mock
+    scheduler.connector = None
+    scheduler.ec_connector = None
+    scheduler.perf_metrics = None
+    scheduler.log_stats = False
+    scheduler.recompute_kv_load_failures = False
+    scheduler.structured_output_manager = mocker.MagicMock()
+    scheduler.structured_output_manager.should_advance.return_value = False
+    scheduler.finished_req_ids_dict = {}
+    scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_event_publisher = mocker.MagicMock()
+
+    request = _HashableRequest(
+        request_id="req-s1",
+        external_req_id="ext-s1",
+        status=RequestStatus.RUNNING,
+        is_finished=lambda: False,
+        num_computed_tokens=10,
+        num_prompt_tokens=10,
+        prompt_token_ids=list(range(10)),
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        num_cached_tokens=0,
+        num_external_computed_tokens=0,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler.requests = {"req-s1": request}
+
+    scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
+    scheduler._free_request = mocker.MagicMock(return_value=None)
+    scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
+    scheduler.running = [request]
+    scheduler.waiting = mocker.MagicMock()
+    scheduler.waiting.remove_requests = mocker.MagicMock()
+    scheduler.make_stats = mocker.MagicMock(return_value=None)
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"req-s1": 10},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    model_runner_output = SimpleNamespace(
+        sampled_token_ids=None,
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=None,
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        req_id_to_index={"req-s1": 0},
+    )
+
+    OmniGenerationScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
+
+    assert len(cleanup_calls) == 1
+    args, _ = cleanup_calls[0]
+    assert args[0] == "req-s1"
+    assert args[1] == "ext-s1"

tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py

@@ -7,7 +7,6 @@ These tests do NOT require Mooncake or RDMA environment.
 """
 
 import threading
-import unittest
 
 import pytest
 import torch
@@ -22,7 +21,7 @@ pytestmark = [pytest.mark.cpu, pytest.mark.parallel]
 
 
 @pytest.mark.core_model
-class TestBufferAllocator(unittest.TestCase):
+class TestBufferAllocator:
     """Unit tests for BufferAllocator."""
 
     def test_basic_alloc_free(self):
@@ -30,15 +29,15 @@ class TestBufferAllocator(unittest.TestCase):
         allocator = BufferAllocator(total_size=4096, alignment=64)
 
         offset1 = allocator.alloc(512)
-        self.assertEqual(offset1, 0)
+        assert offset1 == 0
 
         offset2 = allocator.alloc(512)
-        self.assertGreater(offset2, 0)
+        assert offset2 > 0
 
         # Free first block, should be reusable
         allocator.free(offset1, 512)
         offset3 = allocator.alloc(512)
-        self.assertEqual(offset3, 0)
+        assert offset3 == 0
 
     def test_alignment(self):
         """Verify allocation respects alignment."""
@@ -47,22 +46,22 @@ class TestBufferAllocator(unittest.TestCase):
         _offset1 = allocator.alloc(100)
         offset2 = allocator.alloc(100)
 
-        self.assertEqual(offset2 % 128, 0)
-        self.assertEqual(offset2, 128)
+        assert offset2 % 128 == 0
+        assert offset2 == 128
 
     def test_exhaustion_and_recovery(self):
         """Test that full allocation fails, then succeeds after free."""
         allocator = BufferAllocator(total_size=1024, alignment=64)
 
         offset = allocator.alloc(1024)
-        self.assertEqual(offset, 0)
+        assert offset == 0
 
-        with self.assertRaises(MemoryError):
+        with pytest.raises(MemoryError):
             allocator.alloc(64)
 
         allocator.free(offset, 1024)
         offset2 = allocator.alloc(1024)
-        self.assertEqual(offset2, 0)
+        assert offset2 == 0
 
     def test_thread_safety(self):
         """Verify allocator is thread-safe under concurrent access."""
@@ -84,10 +83,10 @@ class TestBufferAllocator(unittest.TestCase):
         for t in threads:
             t.join()
 
-        self.assertEqual(len(errors), 0, f"Thread safety errors: {errors}")
+        assert len(errors) == 0, f"Thread safety errors: {errors}"
 
 
-class TestAllocatorInvariants(unittest.TestCase):
+class TestAllocatorInvariants:
     """
     Defensive invariant tests for BufferAllocator: double-free, partial
     overlap corruption, and adjacent-block merging.
@@ -106,7 +105,7 @@ class TestAllocatorInvariants(unittest.TestCase):
         allocator.free(offset, 256)
         # Pool should still be consistent: allocate full size back
         offset2 = allocator.alloc(4096)
-        self.assertEqual(offset2, 0)
+        assert offset2 == 0
 
     @pytest.mark.slow
     def test_double_free_after_merge_is_safe(self):
@@ -123,7 +122,7 @@ class TestAllocatorInvariants(unittest.TestCase):
         allocator.free(a, 64)  # should not raise
         # Pool should still be fully usable
         offset = allocator.alloc(4096)
-        self.assertEqual(offset, 0)
+        assert offset == 0
 
     @pytest.mark.slow
     def test_partial_overlap_raises_corruption(self):
@@ -133,7 +132,7 @@ class TestAllocatorInvariants(unittest.TestCase):
         b = allocator.alloc(128)
         allocator.free(a, 128)  # [0, 128) is now free
         # Try to free a region that starts inside [0,128) but extends beyond
-        with self.assertRaises(RuntimeError):
+        with pytest.raises(RuntimeError):
             allocator.free(64, 128)  # [64, 192) overlaps with free [0, 128)
 
         # b is still allocated; freeing b should be fine
@@ -154,7 +153,7 @@ class TestAllocatorInvariants(unittest.TestCase):
         # covering at least 192 bytes (3 * 64).
         # Verify by allocating a contiguous block of 192 bytes.
         offset = allocator.alloc(192)
-        self.assertEqual(offset, 0, "Adjacent blocks were not merged properly")
+        assert offset == 0, "Adjacent blocks were not merged properly"
 
     @pytest.mark.slow
     def test_fragmentation_and_defrag(self):
@@ -174,7 +173,7 @@ class TestAllocatorInvariants(unittest.TestCase):
         allocator.free(d, 64)  # free blocks: [64, 128) and [192, 256)
 
         # Pool has two 64-byte holes; contiguous 128 is unavailable
-        with self.assertRaises(MemoryError):
+        with pytest.raises(MemoryError):
             allocator.alloc(128)
 
         allocator.free(a, 64)
@@ -182,14 +181,16 @@ class TestAllocatorInvariants(unittest.TestCase):
 
         # After freeing everything, full pool should be available
         offset = allocator.alloc(256)
-        self.assertEqual(offset, 0)
+        assert offset == 0
 
 
 @pytest.mark.core_model
-class TestManagedBuffer(unittest.TestCase):
+class TestManagedBuffer:
     """Unit tests for ManagedBuffer."""
 
-    def setUp(self):
+    @pytest.fixture(autouse=True)
+    def _setup(self):
+        # automatically invoked for every test method in the class
         self.allocator = BufferAllocator(total_size=4096, alignment=64)
         self.pool = torch.zeros(4096, dtype=torch.uint8)
 
@@ -203,12 +204,12 @@ class TestManagedBuffer(unittest.TestCase):
         self.pool[offset : offset + 64] = src.view(torch.uint8)
 
         # Raw uint8 view
-        self.assertEqual(buf.tensor.shape[0], 64)
+        assert buf.tensor.shape[0] == 64
 
         # Typed view
         typed = buf.as_tensor(dtype=torch.float32, shape=(4, 4))
-        self.assertEqual(typed.shape, (4, 4))
-        self.assertTrue(torch.equal(typed.flatten(), src))
+        assert typed.shape == (4, 4)
+        assert torch.equal(typed.flatten(), src)
 
         buf.release()
 
@@ -217,14 +218,10 @@ class TestManagedBuffer(unittest.TestCase):
         offset = self.allocator.alloc(128)
 
         with ManagedBuffer(self.allocator, offset, 128, self.pool) as buf:
-            self.assertFalse(buf._released)
+            assert not buf._released
 
-        self.assertTrue(buf._released)
+        assert buf._released
 
         # Space should be reusable
         new_offset = self.allocator.alloc(128)
-        self.assertEqual(new_offset, offset)
-
-
-if __name__ == "__main__":
-    unittest.main()
+        assert new_offset == offset

tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py

@@ -13,7 +13,6 @@ import os
 import subprocess
 import threading
 import time
-import unittest
 
 import pytest
 import torch
@@ -195,43 +194,43 @@ def _md5(tensor: torch.Tensor) -> str:
 # ---------------------------------------------------------------------------
 
 
-@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
-class TestBasicConnector(unittest.TestCase):
+@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+class TestBasicConnector:
     """Verify connector initialization, put, cleanup, and health check."""
 
     def test_initialization(self):
         port = _free_port()
         with MooncakeTransferEngineConnector(_connector_config(port, pool_size=1024 * 1024)) as c:
-            self.assertNotEqual(c.rpc_port, 0)
-            self.assertEqual(c.pool_size, 1024 * 1024)
-            self.assertTrue(c.pool.is_pinned())
+            assert c.rpc_port != 0
+            assert c.pool_size == 1024 * 1024
+            assert c.pool.is_pinned()
             health = c.health()
-            self.assertEqual(health["status"], "healthy")
+            assert health["status"] == "healthy"
 
     def test_put_tensor_bytes_object(self):
         """Put tensor / bytes / dict and verify metadata."""
         port = _free_port()
         with MooncakeTransferEngineConnector(_connector_config(port)) as c:
             ok, sz, meta = c.put("s0", "s1", "t", torch.randn(100))
-            self.assertTrue(ok)
-            self.assertTrue(meta["is_fast_path"])
+            assert ok
+            assert meta["is_fast_path"]
 
             ok, sz, meta = c.put("s0", "s1", "b", b"hello" * 100)
-            self.assertTrue(ok)
-            self.assertTrue(meta["is_fast_path"])
+            assert ok
+            assert meta["is_fast_path"]
 
             ok, sz, meta = c.put("s0", "s1", "d", {"k": [1, 2, 3]})
-            self.assertTrue(ok)
-            self.assertFalse(meta["is_fast_path"])
+            assert ok
+            assert not meta["is_fast_path"]
 
     def test_cleanup_releases_buffer(self):
         port = _free_port()
         with MooncakeTransferEngineConnector(_connector_config(port)) as c:
             c.put("s0", "s1", "r1", torch.randn(100))
             key = MooncakeTransferEngineConnector._make_key("r1", "s0", "s1")
-            self.assertIn(key, c._local_buffers)
+            assert key in c._local_buffers
             c.cleanup("r1", from_stage="s0", to_stage="s1")
-            self.assertNotIn(key, c._local_buffers)
+            assert key not in c._local_buffers
 
     def test_pool_exhaustion_and_recovery(self):
         """Fill pool, verify failure, free, verify recovery."""
@@ -247,7 +246,7 @@ class TestBasicConnector(unittest.TestCase):
             for rid in ids:
                 c.cleanup(rid, from_stage="s", to_stage="s")
             ok, _, _ = c.put("s", "s", "recovery", torch.randn(1000))
-            self.assertTrue(ok, "Pool recovery failed after cleanup")
+            assert ok, "Pool recovery failed after cleanup"
 
 
 # ---------------------------------------------------------------------------
@@ -255,8 +254,8 @@ class TestBasicConnector(unittest.TestCase):
 # ---------------------------------------------------------------------------
 
 
-@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
-class TestEndToEnd(unittest.TestCase):
+@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+class TestEndToEnd:
     """E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload, mixed types."""
 
     def _pair(self, pool_size=16 * 1024 * 1024):
@@ -269,14 +268,14 @@ class TestEndToEnd(unittest.TestCase):
         try:
             orig = torch.randn(1024, 1024, dtype=torch.float32)
             ok, sz, meta = p.put("s0", "s1", "t1", orig)
-            self.assertTrue(ok)
+            assert ok
             time.sleep(0.5)
             result = c.get("s0", "s1", "t1", meta)
-            self.assertIsNotNone(result)
+            assert result is not None
             buf, _ = result
-            self.assertIsInstance(buf, ManagedBuffer)
+            assert isinstance(buf, ManagedBuffer)
             recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
-            self.assertTrue(torch.equal(recon, orig))
+            assert torch.equal(recon, orig)
             buf.release()
         finally:
             p.close()
@@ -287,15 +286,15 @@ class TestEndToEnd(unittest.TestCase):
         try:
             orig = b"Hello RDMA! " * 1000
             ok, _, meta = p.put("s0", "s1", "b1", orig)
-            self.assertTrue(ok)
+            assert ok
             time.sleep(0.5)
             result = c.get("s0", "s1", "b1", meta)
-            self.assertIsNotNone(result)
+            assert result is not None
             buf, _ = result
             recv = buf.to_bytes() if hasattr(buf, "to_bytes") else buf
             if hasattr(buf, "release"):
                 buf.release()
-            self.assertEqual(recv, orig)
+            assert recv == orig
         finally:
             p.close()
             c.close()
@@ -305,13 +304,13 @@ class TestEndToEnd(unittest.TestCase):
         try:
             orig = {"msg": "hello", "nums": [1, 2, 3], "nested": {"k": "v"}}
             ok, _, meta = p.put("s0", "s1", "o1", orig)
-            self.assertTrue(ok)
-            self.assertFalse(meta["is_fast_path"])
+            assert ok
+            assert not meta["is_fast_path"]
             time.sleep(1.0)
             result = c.get("s0", "s1", "o1", meta)
-            self.assertIsNotNone(result)
+            assert result is not None
             recv, _ = result
-            self.assertEqual(recv, orig)
+            assert recv == orig
         finally:
             p.close()
             c.close()
@@ -327,15 +326,15 @@ class TestEndToEnd(unittest.TestCase):
             sbuf.as_tensor(dtype=torch.float32, shape=shape).copy_(ref)
 
             ok, sz, meta = p.put("s0", "s1", "zc", sbuf)
-            self.assertTrue(ok)
-            self.assertTrue(meta.get("is_fast_path"))
+            assert ok
+            assert meta.get("is_fast_path")
             time.sleep(0.5)
 
             result = c.get("s0", "s1", "zc", meta)
-            self.assertIsNotNone(result)
+            assert result is not None
             rbuf, _ = result
             recon = rbuf.as_tensor(dtype=torch.float32, shape=shape)
-            self.assertTrue(torch.equal(recon, ref))
+            assert torch.equal(recon, ref)
             rbuf.release()
         finally:
             p.close()
@@ -348,13 +347,13 @@ class TestEndToEnd(unittest.TestCase):
             orig = torch.randn(5000, 5000, dtype=torch.float32)
             md5_send = _md5(orig)
             ok, _, meta = p.put("s0", "s1", "lg", orig)
-            self.assertTrue(ok)
+            assert ok
             time.sleep(1.0)
             result = c.get("s0", "s1", "lg", meta)
-            self.assertIsNotNone(result)
+            assert result is not None
             buf, _ = result
             recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
-            self.assertEqual(_md5(recon), md5_send)
+            assert _md5(recon) == md5_send
             buf.release()
         finally:
             p.close()
@@ -374,23 +373,23 @@ class TestEndToEnd(unittest.TestCase):
             time.sleep(0.5)
             for name, data in cases:
                 ok, _, meta = p.put("s0", "s1", name, data)
-                self.assertTrue(ok, f"put failed: {name}")
+                assert ok, f"put failed: {name}"
                 time.sleep(0.3)
                 result = c.get("s0", "s1", name, meta)
-                self.assertIsNotNone(result, f"get failed: {name}")
+                assert result is not None, f"get failed: {name}"
                 recv, _ = result
                 if isinstance(data, torch.Tensor):
-                    self.assertIsInstance(recv, ManagedBuffer)
+                    assert isinstance(recv, ManagedBuffer)
                     recon = recv.as_tensor(dtype=data.dtype, shape=data.shape)
-                    self.assertTrue(torch.equal(recon, data), f"mismatch: {name}")
+                    assert torch.equal(recon, data), f"mismatch: {name}"
                     recv.release()
                 elif isinstance(data, bytes):
                     got = recv.to_bytes() if hasattr(recv, "to_bytes") else recv
                     if hasattr(recv, "release"):
                         recv.release()
-                    self.assertEqual(got, data)
+                    assert got == data
                 else:
-                    self.assertEqual(recv, data)
+                    assert recv == data
         finally:
             p.close()
             c.close()
@@ -420,24 +419,24 @@ class TestEndToEnd(unittest.TestCase):
                 t.start()
             for t in threads:
                 t.join()
-            self.assertEqual(len(errors), 0, f"errors: {errors}")
+            assert len(errors) == 0, f"errors: {errors}"
 
     def test_auto_cleanup(self):
         """Producer buffer should be released after consumer get."""
         p, c = self._pair()
         try:
             ok, _, meta = p.put("s0", "s1", "ac", torch.randn(100))
-            self.assertTrue(ok)
+            assert ok
             key = MooncakeTransferEngineConnector._make_key("ac", "s0", "s1")
-            self.assertIn(key, p._local_buffers)
+            assert key in p._local_buffers
             time.sleep(0.5)
             result = c.get("s0", "s1", "ac", meta)
-            self.assertIsNotNone(result)
+            assert result is not None
             buf, _ = result
             if isinstance(buf, ManagedBuffer):
                 buf.release()
             time.sleep(0.3)
-            self.assertNotIn(key, p._local_buffers)
+            assert key not in p._local_buffers
         finally:
             p.close()
             c.close()
@@ -448,22 +447,22 @@ class TestEndToEnd(unittest.TestCase):
 # ---------------------------------------------------------------------------
 
 
-@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
-class TestLifecycle(unittest.TestCase):
+@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+class TestLifecycle:
     """Close, context manager, double-close safety."""
 
     def test_close_releases_resources(self):
         c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size=1024 * 1024))
         c.put("s0", "s1", "x", torch.randn(100))
         c.close()
-        self.assertTrue(c._stop_event.is_set())
-        self.assertEqual(len(c._local_buffers), 0)
+        assert c._stop_event.is_set()
+        assert len(c._local_buffers) == 0
 
     def test_context_manager(self):
         with MooncakeTransferEngineConnector(_connector_config(_free_port())) as c:
             ok, _, _ = c.put("s0", "s1", "ctx", torch.randn(50))
-            self.assertTrue(ok)
-        self.assertTrue(c._stop_event.is_set())
+            assert ok
+        assert c._stop_event.is_set()
 
     def test_double_close_safe(self):
         c = MooncakeTransferEngineConnector(_connector_config(_free_port()))
@@ -477,9 +476,9 @@ class TestLifecycle(unittest.TestCase):
 
 
 @pytest.mark.cuda
-@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
-@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
-class TestGPUPool(unittest.TestCase):
+@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestGPUPool:
     """GPU memory pool: initialization, put (CPU/GPU tensor), E2E transfer."""
 
     @staticmethod
@@ -488,18 +487,18 @@ class TestGPUPool(unittest.TestCase):
 
     def test_gpu_pool_init(self):
         with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c:
-            self.assertEqual(c.pool_device, "cuda:0")
-            self.assertTrue(c.pool.is_cuda)
+            assert c.pool_device == "cuda:0"
+            assert c.pool.is_cuda
 
     def test_gpu_pool_put_cpu_and_gpu_tensor(self):
         with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c:
             ok, _, meta = c.put("s0", "s1", "h2d", torch.randn(256, 256))
-            self.assertTrue(ok)
-            self.assertTrue(meta["is_fast_path"])
+            assert ok
+            assert meta["is_fast_path"]
 
             ok, _, meta = c.put("s0", "s1", "d2d", torch.randn(256, 256, device="cuda:0"))
-            self.assertTrue(ok)
-            self.assertTrue(meta["is_fast_path"])
+            assert ok
+            assert meta["is_fast_path"]
 
     def test_gpu_e2e_transfer(self):
         p = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port()))
@@ -507,14 +506,14 @@ class TestGPUPool(unittest.TestCase):
         try:
             orig = torch.randn(512, 512, dtype=torch.float32, device="cuda:0")
             ok, _, meta = p.put("s0", "s1", "ge", orig)
-            self.assertTrue(ok)
+            assert ok
             time.sleep(0.5)
             result = c.get("s0", "s1", "ge", meta)
-            self.assertIsNotNone(result)
+            assert result is not None
             buf, _ = result
             recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
-            self.assertTrue(recon.is_cuda)
-            self.assertTrue(torch.equal(recon.cpu(), orig.cpu()))
+            assert recon.is_cuda
+            assert torch.equal(recon.cpu(), orig.cpu())
             buf.release()
         finally:
             p.close()
@@ -527,8 +526,8 @@ class TestGPUPool(unittest.TestCase):
 
 
 @pytest.mark.slow
-@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
-class TestStressCorrectness(unittest.TestCase):
+@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+class TestStressCorrectness:
     """
     Slow but high-value regression tests: concurrent put+get with data
     integrity, large payload, edge cases, and sustained stress.
@@ -581,20 +580,20 @@ class TestStressCorrectness(unittest.TestCase):
             for t in threads:
                 t.join()
 
-            self.assertEqual(len(errors), 0, f"Producer errors: {errors}")
-            self.assertEqual(len(meta_map), num_workers)
+            assert len(errors) == 0, f"Producer errors: {errors}"
+            assert len(meta_map) == num_workers
 
             time.sleep(1.0)
 
             # Consumer: get each tensor and verify integrity
             for rid, meta in meta_map.items():
                 result = c.get("s0", "s1", rid, meta)
-                self.assertIsNotNone(result, f"get returned None for {rid}")
+                assert result is not None, f"get returned None for {rid}"
                 buf, _ = result
-                self.assertIsInstance(buf, ManagedBuffer, f"{rid} not ManagedBuffer")
+                assert isinstance(buf, ManagedBuffer), f"{rid} not ManagedBuffer"
                 recon = buf.as_tensor(dtype=torch.float32, shape=(512, 512))
                 recv_md5 = _md5(recon)
-                self.assertEqual(recv_md5, md5_sent[rid], f"MD5 mismatch for {rid}")
+                assert recv_md5 == md5_sent[rid], f"MD5 mismatch for {rid}"
                 buf.release()
         finally:
             p.close()
@@ -683,12 +682,8 @@ class TestStressCorrectness(unittest.TestCase):
             for t in threads:
                 t.join()
             ct.join(timeout=60)
-            self.assertEqual(len(errors), 0, f"errors: {errors}")
-            self.assertEqual(
-                consumed_count[0],
-                num_items,
-                f"Consumer only processed {consumed_count[0]}/{num_items} items",
-            )
+            assert len(errors) == 0, f"errors: {errors}"
+            assert consumed_count[0] == num_items, f"Consumer only processed {consumed_count[0]}/{num_items} items"
         finally:
             p.close()
             c.close()
@@ -701,13 +696,13 @@ class TestStressCorrectness(unittest.TestCase):
         try:
             orig = torch.tensor([3.14], dtype=torch.float32)
             ok, _, meta = p.put("s0", "s1", "tiny", orig)
-            self.assertTrue(ok)
+            assert ok
             time.sleep(0.5)
             result = c.get("s0", "s1", "tiny", meta)
-            self.assertIsNotNone(result)
+            assert result is not None
             buf, _ = result
             recon = buf.as_tensor(dtype=torch.float32, shape=(1,))
-            self.assertTrue(torch.equal(recon, orig))
+            assert torch.equal(recon, orig)
             buf.release()
         finally:
             p.close()
@@ -718,7 +713,7 @@ class TestStressCorrectness(unittest.TestCase):
         port = _free_port()
         with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c:
             ok, sz, meta = c.put("s0", "s1", "empty_b", b"")
-            self.assertFalse(ok, "Empty bytes should be rejected by connector")
+            assert not ok, "Empty bytes should be rejected by connector"
 
     # -- Large payload stress (500 MB) --
 
@@ -731,13 +726,13 @@ class TestStressCorrectness(unittest.TestCase):
             orig = torch.randn(side, side, dtype=torch.float32)
             md5_send = _md5(orig)
             ok, _, meta = p.put("s0", "s1", "500mb", orig)
-            self.assertTrue(ok)
+            assert ok
             time.sleep(2.0)
             result = c.get("s0", "s1", "500mb", meta)
-            self.assertIsNotNone(result)
+            assert result is not None
             buf, _ = result
             recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
-            self.assertEqual(_md5(recon), md5_send)
+            assert _md5(recon) == md5_send
             buf.release()
         finally:
             p.close()
@@ -752,12 +747,8 @@ class TestStressCorrectness(unittest.TestCase):
             for i in range(50):
                 rid = f"cycle_{i}"
                 ok, _, _ = c.put("s0", "s1", rid, torch.randn(256, 256))
-                self.assertTrue(ok, f"Put failed at iteration {i}")
+                assert ok, f"Put failed at iteration {i}"
                 c.cleanup(rid, from_stage="s0", to_stage="s1")
             # After all cycles, pool should be fully recovered
             ok, _, _ = c.put("s0", "s1", "final", torch.randn(1024, 1024))
-            self.assertTrue(ok, "Pool not fully recovered after rapid cycles")
-
-
-if __name__ == "__main__":
-    unittest.main()
+            assert ok, "Pool not fully recovered after rapid cycles"