@@ -0,0 +1,243 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user