import math from collections import namedtuple from typing import AnyStr ServerInfo = namedtuple("ServerInfo", ["instance_type", "instance_idx"]) class Task: """Request task (carries request length info).""" def __init__(self, task_id, task_length, task_load): self.id = task_id self.length = task_length self.bucket_idx = -1 self.load = task_load self.server_info: ServerInfo = ServerInfo("Unknown", -1) def __repr__(self): return ( f"Task(id={self.id}, length={self.length}, load={self.load}, " f"instance_type={self.server_info.instance_type}, instance_idx={self.server_info.instance_idx})" ) class Bucket: """A bucket.""" def __init__(self, bucket_ranges: tuple[int, int]): self.min_length = bucket_ranges[0] self.max_length = bucket_ranges[1] self.task_count = 0 self.total_load = 0.0 class DynamicBucketLoadBalancer: """ Statically buckets requests by length first, then dynamically adjusts the assignment of new requests based on bucket load and length affinity to achieve load balancing. """ def __init__( self, buckets: list[tuple[int, int]], sensitivity=1.0, affinity_strength=0.1, log_func=print, all_neighbor=False ): """ Initialize the load balancer. :param buckets: length range of each bucket :param sensitivity: sensitivity to the load gap (higher = more sensitive) :param affinity_strength: strength of length affinity (higher = a request stays more in its standard bucket) :param log_func: logging function :param all_neighbor: if False, only the left/right buckets are neighbors; if True, all buckets are neighbors """ self.num_buckets = len(buckets) self.sensitivity = sensitivity self.affinity_strength = affinity_strength self.log_func = log_func self.all_neighbor = all_neighbor self.buckets = {idx: Bucket(bucket_ranges) for idx, bucket_ranges in enumerate(buckets)} bucket_boundaries = ", ".join( f"bucket {idx}: [{bucket.min_length}, {bucket.max_length})" for idx, bucket in self.buckets.items() ) self._log_info(f"Initialized {self.num_buckets} buckets: {bucket_boundaries}") # Redirect only when the redirect probability exceeds this threshold self.base_probability_threshold = 0.12 self._log_info(f"Load Balance base_probability_threshold: {self.base_probability_threshold:.2f} ") # Track request tasks self.tasks: dict[AnyStr, Task] = {} # type: ignore # Statistics self.redirected_tasks = 0 self.total_tasks = 0 def _log_info(self, msg, *args, **kwargs): if self.log_func: self.log_func(msg, *args, **kwargs) def _get_standard_bucket_index(self, task_length): """Return the standard bucket index for the given request length.""" for bucket_idx, bucket in self.buckets.items(): if bucket.min_length <= task_length < bucket.max_length: return bucket_idx # Fall back to the last bucket if the length is outside every range return self.num_buckets - 1 def _get_neighbor_indices(self, bucket_idx): """Return the left and right neighbor indices of the given bucket.""" if self.all_neighbor: return list(range(self.num_buckets)) neighbors = [] if bucket_idx > 0: neighbors.append(bucket_idx - 1) if bucket_idx < self.num_buckets - 1: neighbors.append(bucket_idx + 1) return neighbors def _calculate_length_affinity(self, task_length, neighbor_bucket_idx): """ Compute the affinity factor between the task length and the neighbor bucket (0.0 to 1.0). 1.0 means right next to the neighbor bucket, 0.0 means far away from it. """ neighbor_bucket = self.buckets[neighbor_bucket_idx] neighbor_bucket_min = neighbor_bucket.min_length neighbor_bucket_max = neighbor_bucket.max_length if neighbor_bucket_min < task_length < neighbor_bucket_max: raise RuntimeError("task_length must be outside the neighbor bucket range") neighbor_bucket_center = (neighbor_bucket_min + neighbor_bucket_max) / 2.0 neighbor_bucket_half_width = (neighbor_bucket_max - neighbor_bucket_min) / 2.0 distance_to_center = abs(task_length - neighbor_bucket_center) # The closer to the neighbor bucket boundary, the closer the affinity to 1 if neighbor_bucket_half_width > 0: # Relative distance from the neighbor bucket half-width normalized_distance = (distance_to_center - neighbor_bucket_half_width) / neighbor_bucket_half_width # Exponential decay, e.g. normalized_distance=0.1, affinity_strength=1.0 -> 0.9 neighbor_affinity = math.exp(-self.affinity_strength * normalized_distance) else: neighbor_affinity = 1.0 # Safeguard; unreachable in practice # Clamp to [0, 1] return max(0.0, min(neighbor_affinity, 1.0)) def _calculate_redirect_probability(self, task_length, standard_bucket_idx, neighbor_bucket_idx): """ Compute the probability of redirecting to the neighbor bucket based on the load gap and length affinity. """ standard_load = self.buckets[standard_bucket_idx].total_load neighbor_load = self.buckets[neighbor_bucket_idx].total_load # --- 1. Base probability from the load gap --- if standard_load <= 0: load_probability = 0.0 # No load in the standard bucket -> no redirect else: load_ratio = neighbor_load / max(standard_load, 1e-9) # Guard against division by zero # The smaller the neighbor load relative to the standard load, the # higher the redirect probability # e.g. neighbor/standard = 5/6, sensitivity=1.0 -> 1/6 ≈ 0.1667 load_probability = 1 - load_ratio**self.sensitivity load_probability = max(0.0, min(load_probability, 1)) # --- 2. Length affinity factor --- affinity_factor = self._calculate_length_affinity(task_length, neighbor_bucket_idx) # --- 3. Final probability: load gap * length affinity --- # A large load gap is suppressed when the request length is far from the # neighbor bucket; a modest gap can still redirect when it is close. final_probability = load_probability * affinity_factor return final_probability def dispatch_single_task(self, task_id: AnyStr, task_length: int, task_load): return self.dispatch_task(Task(task_id, task_length, task_load)) def dispatch_task(self, cur_task): """ Assign a bucket to a new request, considering dynamic load balancing and length affinity. """ self.total_tasks += 1 standard_bucket_idx = self._get_standard_bucket_index(cur_task.length) neighbor_indices = self._get_neighbor_indices(standard_bucket_idx) best_neighbor_idx = None best_redirect_prob = 0.0 # Pick the neighbor with the highest redirect probability for neighbor_idx in neighbor_indices: # Only consider neighbors with lower load if self.buckets[neighbor_idx].total_load < self.buckets[standard_bucket_idx].total_load: prob = self._calculate_redirect_probability(cur_task.length, standard_bucket_idx, neighbor_idx) if prob > best_redirect_prob: best_redirect_prob = prob best_neighbor_idx = neighbor_idx # Decide the final bucket final_bucket_idx = standard_bucket_idx if best_neighbor_idx is not None and best_redirect_prob > 0: if self.base_probability_threshold < best_redirect_prob: final_bucket_idx = best_neighbor_idx self.redirected_tasks += 1 self._log_info( f"{cur_task} redirected from bucket {standard_bucket_idx} to {final_bucket_idx}" f"(prob={best_redirect_prob:.4f})" ) # Bookkeeping on the chosen bucket self.buckets[final_bucket_idx].task_count += 1 self.buckets[final_bucket_idx].total_load += cur_task.load cur_task.bucket_idx = final_bucket_idx if cur_task.id in self.tasks: raise RuntimeError(f"Task {cur_task.id} is existed!") else: self.tasks[cur_task.id] = cur_task return final_bucket_idx, cur_task def release_task(self, task_id): """Release the load of a request.""" if task_id in self.tasks: found_task = self.tasks.pop(task_id) if 0 <= found_task.bucket_idx < self.num_buckets: self.buckets[found_task.bucket_idx].task_count -= 1 self.buckets[found_task.bucket_idx].total_load -= found_task.load return True else: raise RuntimeError(f"Bucket {found_task.bucket_idx} not found") else: raise RuntimeError(f"Task {task_id} not found") def release_all_tasks(self): for bucket in self.buckets.values(): bucket.task_count = 0 bucket.total_load = 0 self.tasks.clear() class NoStandardBucketLoadBalancer(DynamicBucketLoadBalancer): """Dispatch requests by load only, with no standard bucket.""" def __init__(self, num_buckets: int, max_length: int, log_func=print): bucket_range = math.ceil(max_length / num_buckets) start_length = 0 buckets = [] for _ in range(num_buckets): end_length = start_length + bucket_range if end_length > max_length: end_length = max_length buckets.append((start_length, end_length)) start_length += bucket_range super().__init__(buckets=buckets, log_func=log_func, sensitivity=100, affinity_strength=0, all_neighbor=True)