Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 18 additions & 8 deletions tests/v1/kv_connector/unit/test_moriio_connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import importlib.util
import os
import subprocess
import uuid
from unittest.mock import MagicMock, patch

import msgspec
Expand Down Expand Up @@ -99,6 +100,11 @@ def _setup_kv_transfer_request(
"remote_engine_id": "test_engine",
}
)
zmq_addr = f"host:{remote_host},handshake:{fake_port},notify:{fake_port}"
fake_uuid = uuid.uuid4().hex
request.request_id = (
f"___prefill_addr_{zmq_addr}___decode_addr_{zmq_addr}_{fake_uuid}"
)
return request


Expand Down Expand Up @@ -254,13 +260,14 @@ def test_write_mode_saves_local_block_ids():
do_remote_decode=True,
do_remote_prefill=False,
)

# Setup KV transfer params and embed ZMQ addrs in request_id before
# adding to scheduler so the ID is consistent everywhere.
request = _setup_kv_transfer_request(request)
request_id = request.request_id

scheduler.add_request(request)

# Fake Config
request = _setup_kv_transfer_request(request)

# Remote Prefill, triggers MoRIIOConnectorMetadata.
scheduler_output = scheduler.schedule()
kv_connector_metadata = scheduler_output.kv_connector_metadata
Expand Down Expand Up @@ -312,13 +319,14 @@ def test_write_mode_with_chunked_prefill_saves_local_block_ids():
do_remote_decode=True,
do_remote_prefill=False,
)

# Setup KV transfer params and embed ZMQ addrs in request_id before
# adding to scheduler so the ID is consistent everywhere.
request = _setup_kv_transfer_request(request)
request_id = request.request_id

scheduler.add_request(request)

# Fake Config
request = _setup_kv_transfer_request(request)

# Remote Prefill with chunked prefill, triggers multiple schedules.
expected_counts = [(0, 0, 0), (0, 0, 0), (1, 0, 0)]
kv_connector_metadata = None
Expand Down Expand Up @@ -363,15 +371,17 @@ def test_read_mode_loads_remote_block_ids(moriio_read_mode):
do_remote_decode=False,
do_remote_prefill=True,
)

# Setup KV transfer params and embed ZMQ addrs in request_id before
# adding to scheduler so the ID is consistent everywhere.
request = _setup_kv_transfer_request(request)
request_id = request.request_id

scheduler.add_request(request)
block_list = scheduler.kv_cache_manager.coordinator.single_type_managers[
0
].req_to_blocks[request_id]

request = _setup_kv_transfer_request(request)

# Set remote block ids to be fetched.
request.kv_transfer_params["remote_block_ids"] = block_list

Expand Down
Loading