wehub-resource-sync
53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
"""Admit operation handler"""
|
|
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# Standard
|
|
from typing import TYPE_CHECKING, Any
|
|
import random
|
|
|
|
# First Party
|
|
from lmcache.v1.cache_controller.message import (
|
|
BatchedKVOperationMsg,
|
|
KVOpEvent,
|
|
OpType,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
# Local
|
|
from ..benchmark import TestData, ZMQControllerBenchmark
|
|
|
|
# Local
|
|
from .base import OperationHandler
|
|
|
|
|
|
class AdmitHandler(OperationHandler):
|
|
"""Handler for admit operations"""
|
|
|
|
@property
|
|
def operation_name(self) -> str:
|
|
return "admit"
|
|
|
|
def create_message(
|
|
self, benchmark: "ZMQControllerBenchmark", test_data: "TestData"
|
|
) -> Any:
|
|
instance = random.choice(test_data.instances)
|
|
worker = random.choice(test_data.workers)
|
|
location = random.choice(test_data.locations)
|
|
|
|
operations = []
|
|
for _ in range(benchmark.config.batch_size):
|
|
key = random.choice(test_data.keys)
|
|
seq_num = benchmark.get_next_sequence_number(instance, worker, location)
|
|
operations.append(KVOpEvent(op_type=OpType.ADMIT, key=key, seq_num=seq_num))
|
|
|
|
return BatchedKVOperationMsg(
|
|
instance_id=instance,
|
|
worker_id=worker,
|
|
location=location,
|
|
operations=operations,
|
|
)
|
|
|
|
def get_message_count(self, benchmark: "ZMQControllerBenchmark") -> int:
|
|
return benchmark.config.batch_size
|