| # | |
| # Licensed to the Apache Software Foundation (ASF) under one or more | |
| # contributor license agreements. See the NOTICE file distributed with | |
| # this work for additional information regarding copyright ownership. | |
| # The ASF licenses this file to You under the Apache License, Version 2.0 | |
| # (the "License"); you may not use this file except in compliance with | |
| # the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| """Module for managing ML models in Apache Beam pipelines. | |
| This module provides classes and functions to efficiently manage multiple | |
| machine learning models within Apache Beam pipelines. It includes functionality | |
| for loading, caching, and updating models using multi-process shared memory, | |
| ensuring that models are reused across different workers to optimize resource | |
| usage and performance. | |
| """ | |
| import gc | |
| import heapq | |
| import itertools | |
| import logging | |
| import subprocess | |
| import threading | |
| import time | |
| from collections import Counter | |
| from collections import OrderedDict | |
| from collections import defaultdict | |
| from collections import deque | |
| from typing import Any | |
| from typing import Callable | |
| from typing import Optional | |
| import numpy as np | |
| import torch | |
| from scipy.optimize import nnls | |
| from apache_beam.metrics.metric import Metrics | |
| from apache_beam.utils import multi_process_shared | |
| logger = logging.getLogger(__name__) | |
| _MODEL_MANAGER_METRICS_NAMESPACE = "BeamML_ModelManager" | |
| class GPUMonitor: | |
| """Monitors GPU memory usage in a separate thread using nvidia-smi. | |
| This class continuously polls GPU memory statistics to track current usage | |
| and peak usage over a sliding time window. It serves as the source of truth | |
| for the ModelManager's resource decisions. | |
| Attributes: | |
| fallback_memory_mb: Default total memory if hardware detection fails. | |
| poll_interval: Seconds between memory checks. | |
| peak_window_seconds: Duration to track peak memory usage. | |
| """ | |
| def __init__( | |
| self, | |
| fallback_memory_mb: float = 16000.0, | |
| poll_interval: float = 0.5, | |
| peak_window_seconds: float = 30.0): | |
| self._current_usage = 0.0 | |
| self._peak_usage = 0.0 | |
| self._total_memory = fallback_memory_mb | |
| self._poll_interval = poll_interval | |
| self._peak_window_seconds = peak_window_seconds | |
| self._memory_history = deque() | |
| self._running = False | |
| self._thread = None | |
| self._lock = threading.Lock() | |
| def _detect_hardware(self): | |
| try: | |
| cmd = [ | |
| "nvidia-smi", | |
| "--query-gpu=memory.total", | |
| "--format=csv,noheader,nounits" | |
| ] | |
| output = subprocess.check_output(cmd, text=True).strip() | |
| self._total_memory = float(output) | |
| return True | |
| except (FileNotFoundError, subprocess.CalledProcessError): | |
| logger.warning( | |
| "nvidia-smi not found or failed. Defaulting total memory to %s MB", | |
| self._total_memory) | |
| return False | |
| except Exception as e: | |
| logger.warning( | |
| "Error parsing nvidia-smi output: %s. " | |
| "Defaulting total memory to %s MB", | |
| e, | |
| self._total_memory) | |
| return False | |
| def start(self): | |
| self._gpu_available = self._detect_hardware() | |
| if self._running or not self._gpu_available: | |
| return | |
| self._running = True | |
| self._thread = threading.Thread(target=self._poll_loop, daemon=True) | |
| self._thread.start() | |
| def stop(self): | |
| self._running = False | |
| if self._thread: | |
| self._thread.join() | |
| def reset_peak(self): | |
| with self._lock: | |
| now = time.time() | |
| self._memory_history.clear() | |
| self._memory_history.append((now, self._current_usage)) | |
| self._peak_usage = self._current_usage | |
| def get_stats(self) -> tuple[float, float, float]: | |
| with self._lock: | |
| return self._current_usage, self._peak_usage, self._total_memory | |
| def refresh(self): | |
| """Forces an immediate poll of the GPU.""" | |
| usage = self._get_nvidia_smi_used() | |
| now = time.time() | |
| with self._lock: | |
| self._current_usage = usage | |
| self._memory_history.append((now, usage)) | |
| # Recalculate peak immediately | |
| while self._memory_history and (now - self._memory_history[0][0] | |
| > self._peak_window_seconds): | |
| self._memory_history.popleft() | |
| self._peak_usage = ( | |
| max(m for _, m in self._memory_history) | |
| if self._memory_history else usage) | |
| def _get_nvidia_smi_used(self) -> float: | |
| try: | |
| cmd = [ | |
| "nvidia-smi", | |
| "--query-gpu=memory.free", | |
| "--format=csv,noheader,nounits" | |
| ] | |
| output = subprocess.check_output(cmd, text=True).strip() | |
| free_memory = float(output) | |
| return self._total_memory - free_memory | |
| except Exception as e: | |
| logger.warning('Failed to get GPU memory usage: %s', e) | |
| return 0.0 | |
| def _poll_loop(self): | |
| while self._running: | |
| usage = self._get_nvidia_smi_used() | |
| now = time.time() | |
| with self._lock: | |
| self._current_usage = usage | |
| self._memory_history.append((now, usage)) | |
| while self._memory_history and (now - self._memory_history[0][0] | |
| > self._peak_window_seconds): | |
| self._memory_history.popleft() | |
| self._peak_usage = ( | |
| max(m for _, m in self._memory_history) | |
| if self._memory_history else usage) | |
| time.sleep(self._poll_interval) | |
| class ResourceEstimator: | |
| """Estimates individual model memory usage using statistical observation. | |
| Uses Non-Negative Least Squares (NNLS) to deduce the memory footprint of | |
| individual models based on aggregate system memory readings and the | |
| configuration of active models at that time. | |
| """ | |
| def __init__( | |
| self, | |
| smoothing_factor: float = 0.2, | |
| min_data_points: int = 5, | |
| verbose_logging: bool = False): | |
| self.smoothing_factor = smoothing_factor | |
| self.min_data_points = min_data_points | |
| self.verbose_logging = verbose_logging | |
| self.estimates: dict[str, float] = {} | |
| self.history = defaultdict(lambda: deque(maxlen=20)) | |
| self.known_models = set() | |
| self._lock = threading.Lock() | |
| def logging_info(self, message: str, *args): | |
| if self.verbose_logging: | |
| logger.info(message, *args) | |
| def is_unknown(self, model_tag: str) -> bool: | |
| with self._lock: | |
| return model_tag not in self.estimates | |
| def get_estimate(self, model_tag: str, default_mb: float = 4000.0) -> float: | |
| with self._lock: | |
| return self.estimates.get(model_tag, default_mb) | |
| def set_initial_estimate(self, model_tag: str, cost: float): | |
| with self._lock: | |
| self.estimates[model_tag] = cost | |
| self.known_models.add(model_tag) | |
| Metrics.distribution( | |
| _MODEL_MANAGER_METRICS_NAMESPACE, | |
| f"memory_estimate_mb_{model_tag}", | |
| process_wide=True).update(int(cost)) | |
| self.logging_info("Initial Profile for %s: %s MB", model_tag, cost) | |
| def add_observation( | |
| self, active_snapshot: dict[str, int], peak_memory: float): | |
| if active_snapshot: | |
| model_list = "\n".join( | |
| f"\t- {model}: {count}" | |
| for model, count in sorted(active_snapshot.items())) | |
| else: | |
| model_list = "\t- None" | |
| self.logging_info( | |
| "Adding Observation:\n PeakMemory: %.1f MB\n Instances:\n%s", | |
| peak_memory, | |
| model_list) | |
| if not active_snapshot: | |
| return | |
| with self._lock: | |
| config_key = tuple(sorted(active_snapshot.items())) | |
| self.history[config_key].append(peak_memory) | |
| for tag in active_snapshot: | |
| self.known_models.add(tag) | |
| self._solve() | |
| def _solve(self): | |
| """ | |
| Solves Ax=b using raw readings (no pre-averaging) and NNLS. | |
| This creates a 'tall' matrix A where every memory reading is | |
| a separate equation. | |
| """ | |
| unique = sorted(list(self.known_models)) | |
| # We need to build the matrix first to know if we have enough data points | |
| A, b = [], [] | |
| for config_key, mem_values in self.history.items(): | |
| if not mem_values: | |
| continue | |
| # 1. Create the feature row for this configuration ONCE | |
| # (It represents the model counts + bias) | |
| counts = dict(config_key) | |
| feature_row = [counts.get(model, 0) for model in unique] | |
| feature_row.append(1) # Bias column | |
| # 2. Add a separate row to the matrix for EVERY individual reading | |
| # Instead of averaging, we flatten the history into the matrix | |
| for reading in mem_values: | |
| A.append(feature_row) # The inputs (models) stay the same | |
| b.append(reading) # The output (memory) varies due to noise | |
| # Convert to numpy for SciPy | |
| A = np.array(A) | |
| b = np.array(b) | |
| if len( | |
| self.history.keys()) < len(unique) + 1 or len(A) < self.min_data_points: | |
| # Not enough data to solve yet | |
| return | |
| self.logging_info( | |
| "Solving with %s total observations for %s models.", | |
| len(A), | |
| len(unique)) | |
| try: | |
| # Solve using Non-Negative Least Squares | |
| # x will be >= 0 | |
| x, _ = nnls(A, b) | |
| weights = x[:-1] | |
| bias = x[-1] | |
| for i, model in enumerate(unique): | |
| calculated_cost = weights[i] | |
| if model in self.estimates: | |
| old = self.estimates[model] | |
| new = (old * (1 - self.smoothing_factor)) + ( | |
| calculated_cost * self.smoothing_factor) | |
| self.estimates[model] = new | |
| else: | |
| self.estimates[model] = calculated_cost | |
| self.logging_info( | |
| "Updated Estimate for %s: %.1f MB", model, self.estimates[model]) | |
| Metrics.distribution( | |
| _MODEL_MANAGER_METRICS_NAMESPACE, | |
| f"memory_estimate_mb_{model}", | |
| process_wide=True).update(int(self.estimates[model])) | |
| self.logging_info("System Bias: %s MB", bias) | |
| except Exception as e: | |
| logger.error("Solver failed: %s", e) | |
| class QueueTicket: | |
| def __init__(self, priority, ticket_num, tag): | |
| self.priority = priority | |
| self.ticket_num = ticket_num | |
| self.tag = tag | |
| self.wake_event = threading.Event() | |
| def __lt__(self, other): | |
| return (self.priority, self.ticket_num) < (other.priority, other.ticket_num) | |
| class ModelManager: | |
| """Manages model lifecycles, caching, and resource arbitration. | |
| This class acts as the central controller for acquiring model instances. | |
| 1. LRU Caching of idle models. | |
| 2. Resource estimation and admission control (preventing OOM). | |
| 3. Dynamic eviction of low-priority models, determined by count of | |
| pending requests, when space is needed. | |
| 4. 'Isolation Mode' for safely profiling unknown models. | |
| """ | |
| def __init__( | |
| self, | |
| monitor: Optional['GPUMonitor'] = None, | |
| slack_percentage: float = 0.10, | |
| poll_interval: float = 0.5, | |
| peak_window_seconds: float = 30.0, | |
| min_data_points: int = 5, | |
| smoothing_factor: float = 0.2, | |
| eviction_cooldown_seconds: float = 10.0, | |
| min_model_copies: int = 1, | |
| wait_timeout_seconds: float = 300.0, | |
| lock_timeout_seconds: float = 60.0, | |
| verbose_logging: bool = False): | |
| self._estimator = ResourceEstimator( | |
| min_data_points=min_data_points, | |
| smoothing_factor=smoothing_factor, | |
| verbose_logging=verbose_logging) | |
| self._monitor = monitor if monitor else GPUMonitor( | |
| poll_interval=poll_interval, peak_window_seconds=peak_window_seconds) | |
| self._slack_percentage = slack_percentage | |
| self._eviction_cooldown = eviction_cooldown_seconds | |
| self._min_model_copies = min_model_copies | |
| self._wait_timeout_seconds = wait_timeout_seconds | |
| self._lock_timeout_seconds = lock_timeout_seconds | |
| self._verbose_logging = verbose_logging | |
| # Resource State | |
| self._models = defaultdict(list) | |
| # Idle LRU used to track released models that | |
| # can be freed or reused upon request. | |
| self._idle_lru = OrderedDict() | |
| self._active_counts = Counter() | |
| self._total_active_jobs = 0 | |
| self._pending_reservations = 0.0 | |
| # Isolation state used to profile unknown models, | |
| # ensuring they run alone to get accurate readings. | |
| # isolation_baseline represents the GPU usage before | |
| # loading the unknown model. | |
| self._isolation_mode = False | |
| self._isolation_baseline = 0.0 | |
| # Waiting Queue and Ticketing to make sure we have fair ordering | |
| # and also priority for unknown models. | |
| self._wait_queue = [] | |
| self._ticket_counter = itertools.count() | |
| self._cancelled_tickets = set() | |
| # TODO: Consider making the wait to be smarter, i.e. | |
| # splitting read/write etc. to avoid potential contention. | |
| self._cv = threading.Condition() | |
| self._monitor.start() | |
| def _update_model_count_metric(self): | |
| for tag, instances in self._models.items(): | |
| Metrics.distribution( | |
| _MODEL_MANAGER_METRICS_NAMESPACE, | |
| f"num_loaded_models_{tag}", | |
| process_wide=True).update(len(instances)) | |
| def _clear_all_model_metrics(self): | |
| for tag in self._models: | |
| Metrics.distribution( | |
| _MODEL_MANAGER_METRICS_NAMESPACE, | |
| f"num_loaded_models_{tag}", | |
| process_wide=True).update(0) | |
| def logging_info(self, message: str, *args): | |
| if self._verbose_logging: | |
| logger.info(message, *args) | |
| def all_models(self, tag) -> list[Any]: | |
| return self._models[tag] | |
| # Should hold _cv lock when calling | |
| def try_enter_isolation_mode(self, tag: str, ticket_num: int) -> bool: | |
| if self._total_active_jobs > 0: | |
| self.logging_info( | |
| "Waiting to enter isolation: tag=%s ticket num=%s", tag, ticket_num) | |
| self._cv.wait(timeout=self._lock_timeout_seconds) | |
| # return False since we have waited and need to re-evaluate | |
| # in caller to make sure our priority is still valid. | |
| return False | |
| self.logging_info("Unknown model %s detected. Flushing GPU.", tag) | |
| self._delete_all_models() | |
| self._isolation_mode = True | |
| self._total_active_jobs += 1 | |
| self._isolation_baseline, _, _ = self._monitor.get_stats() | |
| self._monitor.reset_peak() | |
| return True | |
| # Should hold _cv lock when calling | |
| def should_spawn_model(self, tag: str, ticket_num: int) -> bool: | |
| curr, _, total = self._monitor.get_stats() | |
| est_cost = self._estimator.get_estimate(tag) | |
| limit = total * (1 - self._slack_percentage) | |
| # Use current usage for capacity check (ignore old spikes) | |
| if (curr + self._pending_reservations + est_cost) <= limit: | |
| self._pending_reservations += est_cost | |
| self._total_active_jobs += 1 | |
| self._active_counts[tag] += 1 | |
| return True | |
| # Evict to make space (passing tag to check demand/existence) | |
| if self._evict_to_make_space(limit, est_cost, requesting_tag=tag): | |
| return True | |
| # Manually log status for debugging if we are going to wait | |
| idle_count = 0 | |
| other_idle_count = 0 | |
| for item in self._idle_lru.items(): | |
| if item[1][0] == tag: | |
| idle_count += 1 | |
| else: | |
| other_idle_count += 1 | |
| total_model_count = 0 | |
| for _, instances in self._models.items(): | |
| total_model_count += len(instances) | |
| curr, _, _ = self._monitor.get_stats() | |
| self.logging_info( | |
| "Waiting for resources to free up: " | |
| "tag=%s ticket num%s model count=%s " | |
| "idle count=%s resource usage=%.1f MB " | |
| "total models count=%s other idle=%s", | |
| tag, | |
| ticket_num, | |
| len(self._models[tag]), | |
| idle_count, | |
| curr, | |
| total_model_count, | |
| other_idle_count) | |
| # Wait since we couldn't make space and | |
| # added timeout to avoid missed notify call. | |
| self._cv.wait(timeout=self._lock_timeout_seconds) | |
| return False | |
| def _wake_next_in_queue(self): | |
| if self._wait_queue: | |
| # Clean up cancelled tickets at head of queue | |
| while self._wait_queue and self._wait_queue[ | |
| 0].ticket_num in self._cancelled_tickets: | |
| self._cancelled_tickets.remove(self._wait_queue[0].ticket_num) | |
| heapq.heappop(self._wait_queue) | |
| next_inline = self._wait_queue[0] | |
| next_inline.wake_event.set() | |
| def _wait_in_queue(self, ticket: QueueTicket): | |
| self._cv.release() | |
| try: | |
| ticket.wake_event.wait(timeout=self._lock_timeout_seconds) | |
| ticket.wake_event.clear() | |
| finally: | |
| self._cv.acquire() | |
| def acquire_model(self, tag: str, loader_func: Callable[[], Any]) -> Any: | |
| current_priority = 0 if self._estimator.is_unknown(tag) else 1 | |
| ticket_num = next(self._ticket_counter) | |
| my_ticket = QueueTicket(current_priority, ticket_num, tag) | |
| with self._cv: | |
| # FAST PATH: Grab from idle LRU if available | |
| if not self._isolation_mode: | |
| cached_instance = self._try_grab_from_lru(tag) | |
| if cached_instance: | |
| return cached_instance | |
| # SLOW PATH: Enqueue and wait for turn to acquire model, | |
| # with unknown models having priority and order enforced | |
| # by ticket number as FIFO. | |
| self.logging_info( | |
| "Acquire Queued: tag=%s, priority=%d " | |
| "total models count=%s ticket num=%s", | |
| tag, | |
| current_priority, | |
| len(self._models[tag]), | |
| ticket_num) | |
| heapq.heappush(self._wait_queue, my_ticket) | |
| est_cost = 0.0 | |
| is_unknown = False | |
| wait_time_start = time.time() | |
| try: | |
| while True: | |
| wait_time_elapsed = time.time() - wait_time_start | |
| if wait_time_elapsed > self._wait_timeout_seconds: | |
| raise RuntimeError( | |
| f"Timeout waiting to acquire model: {tag} " | |
| f"after {wait_time_elapsed:.1f} seconds.") | |
| if not self._wait_queue or self._wait_queue[ | |
| 0].ticket_num != ticket_num: | |
| self.logging_info( | |
| "Waiting for its turn: tag=%s ticket num=%s", tag, ticket_num) | |
| self._wait_in_queue(my_ticket) | |
| continue | |
| # Re-evaluate priority in case model became known during wait | |
| is_unknown = self._estimator.is_unknown(tag) | |
| real_priority = 0 if is_unknown else 1 | |
| # If priority changed, reinsert into queue and wait | |
| if current_priority != real_priority: | |
| heapq.heappop(self._wait_queue) | |
| current_priority = real_priority | |
| my_ticket = QueueTicket(current_priority, ticket_num, tag) | |
| heapq.heappush(self._wait_queue, my_ticket) | |
| self._wake_next_in_queue() | |
| continue | |
| # Try grab from LRU again in case model was released during wait | |
| cached_instance = self._try_grab_from_lru(tag) | |
| if cached_instance: | |
| return cached_instance | |
| # Path A: Isolation | |
| if is_unknown: | |
| if self.try_enter_isolation_mode(tag, ticket_num): | |
| # We got isolation, can proceed to spawn | |
| break | |
| else: | |
| # We waited, need to re-evaluate our turn | |
| # because priority may have changed during the wait | |
| continue | |
| # Path B: Concurrent | |
| else: | |
| if self._isolation_mode: | |
| self.logging_info( | |
| "Waiting due to isolation in progress: tag=%s ticket num%s", | |
| tag, | |
| ticket_num) | |
| self._wait_in_queue(my_ticket) | |
| continue | |
| if self.should_spawn_model(tag, ticket_num): | |
| est_cost = self._estimator.get_estimate(tag) | |
| # We can proceed to spawn since we have resources | |
| break | |
| else: | |
| # We waited, need to re-evaluate our turn | |
| # because priority may have changed during the wait | |
| continue | |
| finally: | |
| # Remove self from wait queue once done | |
| if self._wait_queue and self._wait_queue[0].ticket_num == ticket_num: | |
| heapq.heappop(self._wait_queue) | |
| else: | |
| # Marked as cancelled so that we skip when we reach head later | |
| self._cancelled_tickets.add(ticket_num) | |
| self._wake_next_in_queue() | |
| return self._spawn_new_model(tag, loader_func, is_unknown, est_cost) | |
| def release_model(self, tag: str, instance: Any): | |
| with self._cv: | |
| try: | |
| self._total_active_jobs -= 1 | |
| if self._active_counts[tag] > 0: | |
| self._active_counts[tag] -= 1 | |
| self._idle_lru[id(instance)] = (tag, instance, time.time()) | |
| # Update estimator with latest stats | |
| _, peak_during_job, _ = self._monitor.get_stats() | |
| if self._isolation_mode and self._active_counts[tag] == 0: | |
| # For isolation mode, we directly set the initial estimate | |
| # so that we can quickly learn the model cost. | |
| cost = max(0, peak_during_job - self._isolation_baseline) | |
| self._estimator.set_initial_estimate(tag, cost) | |
| self._isolation_mode = False | |
| self._isolation_baseline = 0.0 | |
| else: | |
| # Regular update for known models | |
| snapshot = { | |
| t: len(instances) | |
| for t, instances in self._models.items() if len(instances) > 0 | |
| } | |
| if snapshot: | |
| self._estimator.add_observation(snapshot, peak_during_job) | |
| finally: | |
| self._wake_next_in_queue() | |
| self._cv.notify_all() | |
| def _try_grab_from_lru(self, tag: str) -> Any: | |
| target_key = None | |
| target_instance = None | |
| for key, (t, instance, _) in reversed(self._idle_lru.items()): | |
| if t == tag: | |
| target_key = key | |
| target_instance = instance | |
| break | |
| if target_instance: | |
| # Found an idle model, remove from LRU and return | |
| del self._idle_lru[target_key] | |
| self._active_counts[tag] += 1 | |
| self._total_active_jobs += 1 | |
| return target_instance | |
| self.logging_info("No idle model found for tag: %s", tag) | |
| return None | |
| def _evict_to_make_space( | |
| self, limit: float, est_cost: float, requesting_tag: str) -> bool: | |
| """ | |
| Evicts models based on Demand Magnitude + Tiers. | |
| Crucially: If we have 0 active copies of 'requesting_tag', we FORCE eviction | |
| of the lowest-demand candidate to avoid starvation. | |
| Returns True if space was made, False otherwise. | |
| """ | |
| curr, _, _ = self._monitor.get_stats() | |
| projected_usage = curr + self._pending_reservations + est_cost | |
| if projected_usage <= limit: | |
| # Memory usage changed and we are already under limit | |
| return True | |
| now = time.time() | |
| # Calculate the demand from the wait queue | |
| # TODO: Also factor in the active counts to avoid thrashing | |
| demand_map = Counter() | |
| for item in self._wait_queue: | |
| demand_map[item.tag] += 1 | |
| my_demand = demand_map[requesting_tag] | |
| am_i_starving = len(self._models[requesting_tag]) == 0 | |
| candidates = [] | |
| for key, (tag, instance, release_time) in self._idle_lru.items(): | |
| candidate_demand = demand_map[tag] | |
| # TODO: Try to avoid churn if demand is similar | |
| if not am_i_starving and candidate_demand >= my_demand: | |
| continue | |
| # Attempts to score candidates based on hotness and manually | |
| # specified minimum copies. Demand is weighted heavily to | |
| # ensure we evict low-demand models first. | |
| age = now - release_time | |
| is_cold = age >= self._eviction_cooldown | |
| total_copies = len(self._models[tag]) | |
| is_surplus = total_copies > self._min_model_copies | |
| if is_cold and is_surplus: tier = 0 | |
| elif not is_cold and is_surplus: tier = 1 | |
| elif is_cold and not is_surplus: tier = 2 | |
| else: tier = 3 | |
| score = (candidate_demand * 10) + tier | |
| candidates.append((score, release_time, key, tag, instance)) | |
| candidates.sort(key=lambda x: (x[0], x[1])) | |
| # Evict candidates until we are under limit | |
| for score, _, key, tag, instance in candidates: | |
| if projected_usage <= limit: | |
| break | |
| if key not in self._idle_lru: continue | |
| self._perform_eviction(key, tag, instance, score) | |
| curr, _, _ = self._monitor.get_stats() | |
| projected_usage = curr + self._pending_reservations + est_cost | |
| return projected_usage <= limit | |
| def _delete_instance(self, instance: Any): | |
| if isinstance(instance, str): | |
| # If the instance is a string, it's a uuid used | |
| # to retrieve the model from MultiProcessShared | |
| try: | |
| multi_process_shared.MultiProcessShared( | |
| lambda: "N/A", tag=instance).unsafe_hard_delete() | |
| except (EOFError, OSError, BrokenPipeError): | |
| # This can happen even in normal operation. | |
| pass | |
| if hasattr(instance, 'mock_model_unsafe_hard_delete'): | |
| # Call the mock unsafe hard delete method for testing | |
| instance.mock_model_unsafe_hard_delete() | |
| del instance | |
| def _perform_eviction(self, key: str, tag: str, instance: Any, score: int): | |
| self.logging_info("Evicting Model: %s (Score %d)", tag, score) | |
| curr, _, _ = self._monitor.get_stats() | |
| self.logging_info("Resource Usage Before Eviction: %.1f MB", curr) | |
| if key in self._idle_lru: | |
| del self._idle_lru[key] | |
| for i, inst in enumerate(self._models[tag]): | |
| if instance == inst: | |
| del self._models[tag][i] | |
| break | |
| self._delete_instance(instance) | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| self._monitor.refresh() | |
| self._monitor.reset_peak() | |
| curr, _, _ = self._monitor.get_stats() | |
| self.logging_info("Resource Usage After Eviction: %.1f MB", curr) | |
| self._update_model_count_metric() | |
| def _spawn_new_model( | |
| self, | |
| tag: str, | |
| loader_func: Callable[[], Any], | |
| is_unknown: bool, | |
| est_cost: float) -> Any: | |
| try: | |
| with self._cv: | |
| self.logging_info("Loading Model: %s (Unknown: %s)", tag, is_unknown) | |
| baseline_snap, _, _ = self._monitor.get_stats() | |
| instance = loader_func() | |
| _, peak_during_load, _ = self._monitor.get_stats() | |
| snapshot = {tag: 1} | |
| self._estimator.add_observation( | |
| snapshot, peak_during_load - baseline_snap) | |
| if not is_unknown: | |
| self._pending_reservations = max( | |
| 0.0, self._pending_reservations - est_cost) | |
| self._models[tag].append(instance) | |
| self._update_model_count_metric() | |
| return instance | |
| except Exception as e: | |
| logger.error("Load Failed: %s. Error: %s", tag, e) | |
| with self._cv: | |
| self._total_active_jobs -= 1 | |
| if is_unknown: | |
| self._isolation_mode = False | |
| self._isolation_baseline = 0.0 | |
| else: | |
| self._pending_reservations = max( | |
| 0.0, self._pending_reservations - est_cost) | |
| self._active_counts[tag] -= 1 | |
| self._cv.notify_all() | |
| raise e | |
| def _delete_all_models(self): | |
| self._clear_all_model_metrics() | |
| self._idle_lru.clear() | |
| for _, instances in self._models.items(): | |
| for instance in instances: | |
| self._delete_instance(instance) | |
| self._models.clear() | |
| self._active_counts.clear() | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| self._monitor.refresh() | |
| self._monitor.reset_peak() | |
| def _force_reset(self): | |
| logger.warning("Force Reset Triggered") | |
| self._delete_all_models() | |
| self._models = defaultdict(list) | |
| self._idle_lru = OrderedDict() | |
| self._active_counts = Counter() | |
| self._wait_queue = [] | |
| self._total_active_jobs = 0 | |
| self._pending_reservations = 0.0 | |
| self._isolation_mode = False | |
| self._isolation_baseline = 0.0 | |
| def shutdown(self): | |
| self._delete_all_models() | |
| self._monitor.stop() | |
| def __del__(self): | |
| self.shutdown() | |
| def __exit__(self, exc_type, exc_value, traceback): | |
| self.shutdown() |