Coverage for wrapper/wrapper_model.py: 88%
172 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Base class for model generation wrappers.
3"""
4import torch
5import logging
6import time
7import traceback
9import nvidia_smi
10from nvidia_smi import NVMLError
12from typing import Callable
13from typing import List
14from typing import Optional
15from typing import Dict
16from typing import Any
17from typing import Union
19from abc import ABC
20from abc import abstractmethod
22from datetime import datetime
24from console_utils import setup_logging
26from model_timing import LoadTimer
27from model_timing import GenTimer
30class GenerationInterruptedError(Exception):
31 """Exception raised when generation is interrupted."""
33 def __init__(self, message: str) -> None:
34 super().__init__(message)
35 self.message = message
38class ModelGeneration(ABC):
39 """Base class for model generation."""
41 def __init__(
42 self,
43 model_name: str,
44 torch_compile: bool = True,
45 ) -> None:
46 self.running: bool = False
47 self.interrupted: bool = False
48 self.status: str = "initializing"
49 self.model_name: str = model_name
51 # Parallelism
52 self.rank: int = 0
53 self.local_rank: int = 0
54 self.world_size: int = 1
55 self.device_id: Union[int, str] = 0
56 self.device: Optional[torch.device] = None
58 self.torch_compile: bool = torch_compile
60 # Timing
61 self.load_timer: LoadTimer = LoadTimer()
62 self.gen_timers: Dict[str, GenTimer] = {} # id -> GenTimer
64 self.gpu_setup: bool = True
66 def __del__(self) -> None:
67 if torch.cuda.is_available():
68 torch.cuda.empty_cache()
70 def interrupt(self) -> None:
71 """Interrupt the current generation process."""
72 if self.world_size > 1:
73 logging.warning(f"[{self.rank}] Interruption not supported in multi-GPU mode.")
75 if self.running:
76 logging.info(f"[{self.rank}] Interrupting.")
77 self.interrupted = True
78 else:
79 logging.warning(f"[{self.rank}] Not running, cannot interrupt.")
81 def is_interrupted(self) -> bool:
82 """Check and clear the interrupted flag. Returns True if generation was interrupted."""
83 if self.interrupted:
84 self.interrupted = False
85 return True
86 return False
88 def check_interrupted(self) -> None:
89 """Raise GenerationInterruptedError if generation has been interrupted."""
90 if self.is_interrupted():
91 raise GenerationInterruptedError("Generation interrupted.")
93 def init_logging(self) -> None:
94 """Initialize logging with colored output."""
95 setup_logging(
96 path="/tmp",
97 file_name="streamwise.log",
98 level=logging.INFO,
99 )
101 def init(self) -> None:
102 """Initialize the model, including loading and setting up parallelism."""
103 t0 = time.time()
105 try:
106 self.status = "initializing parallelism"
107 self.init_parallelism()
109 self.init_logging()
111 self.status = "loading model"
112 self.load_model()
114 self.status = "parallelizing model"
115 self.init_model_parallelism()
117 self.status = "compiling model"
118 self.model_compile()
120 self.status = "ok"
122 logging.info(f"[{self.rank}] Model loaded in {time.time() - t0:.3f} seconds.")
124 torch.cuda.empty_cache() # for tracking memory properly
125 mem_gb = torch.cuda.memory_allocated() / 1024 / 1024 ** 2
126 logging.info(f"[{self.rank}] Total memory allocated: {mem_gb:.2f} GB.")
127 except Exception as ex:
128 logging.error(f"Error during initialization: {ex}.")
129 logging.error(f"Trace: {traceback.format_exc()}.")
130 self.status = "failed"
131 raise ex
132 finally:
133 self.load_timer.end()
135 def init_parallelism(self) -> None:
136 """Initialize distributed parallelism if applicable."""
137 logging.info("Parallelism initialization should be implemented in subclass.")
139 def load_model(self) -> None:
140 """Load the model into memory."""
141 logging.info("Model initialization should be implemented in subclass.")
143 def init_model_parallelism(self) -> None:
144 """Set up model parallelism if applicable."""
145 logging.info("Model parallelism should be implemented in subclass.")
147 def model_compile(self) -> None:
148 """Compile the model using torch.compile if enabled."""
149 logging.info("Model parallelism should be implemented in subclass.")
151 def _assert_model_init(self) -> None:
152 """Assert that the model is initialized and ready."""
153 if self.status != "ok":
154 raise ValueError(f"Model not initialized. Current status: {self.status}.")
156 def _safe_nvml_call(self, fn: Callable[..., Any], *args: Any) -> Optional[Any]:
157 """Safely call an NVML function, returning None if not supported (e.g. on MIG instances)."""
158 try:
159 return fn(*args)
160 except NVMLError as ex:
161 logging.info("NVML call not supported: %s.", ex)
162 return None
164 def get_gpu_info(self) -> Optional[List[Dict[str, Any]]]:
165 """Get information about the GPUs on the system."""
166 if not self.gpu_setup:
167 return None
169 ret = []
170 try:
171 nvidia_smi.nvmlInit()
172 device_count = nvidia_smi.nvmlDeviceGetCount()
173 if device_count == 0:
174 logging.warning("No GPUs found.")
175 self.gpu_setup = False
176 return None
178 local_gpu_index = torch.cuda.current_device()
180 for gpu_index in range(device_count):
181 handle = nvidia_smi.nvmlDeviceGetHandleByIndex(gpu_index)
182 gpu_name_raw = nvidia_smi.nvmlDeviceGetName(handle)
183 gpu_name: str
184 if isinstance(gpu_name_raw, bytes):
185 gpu_name = gpu_name_raw.decode("utf-8")
186 else:
187 gpu_name = gpu_name_raw
188 gpu_mem_info = self._safe_nvml_call(nvidia_smi.nvmlDeviceGetMemoryInfo, handle)
190 # Some metrics are not supported on MIG instances; use _safe_nvml_call to
191 # return None for those fields rather than failing the entire GPU info query.
192 gpu_util = self._safe_nvml_call(nvidia_smi.nvmlDeviceGetUtilizationRates, handle)
193 temp = self._safe_nvml_call(
194 nvidia_smi.nvmlDeviceGetTemperature, handle, nvidia_smi.NVML_TEMPERATURE_GPU
195 )
196 power_draw_raw = self._safe_nvml_call(nvidia_smi.nvmlDeviceGetPowerUsage, handle)
197 power_limit_raw = self._safe_nvml_call(nvidia_smi.nvmlDeviceGetEnforcedPowerLimit, handle)
198 graphics_clock = self._safe_nvml_call(
199 nvidia_smi.nvmlDeviceGetClockInfo, handle, nvidia_smi.NVML_CLOCK_GRAPHICS
200 )
201 sm_clock = self._safe_nvml_call(nvidia_smi.nvmlDeviceGetClockInfo, handle, nvidia_smi.NVML_CLOCK_SM)
202 mem_clock = self._safe_nvml_call(nvidia_smi.nvmlDeviceGetClockInfo, handle, nvidia_smi.NVML_CLOCK_MEM)
204 ret.append({
205 "index": gpu_index,
206 "current": gpu_index == local_gpu_index,
207 "name": gpu_name,
208 "sm_util": gpu_util.gpu if gpu_util is not None else None,
209 "mem_util": gpu_util.memory if gpu_util is not None else None,
210 "mem_gib_used": gpu_mem_info.used / (1024 ** 3) if gpu_mem_info is not None else None,
211 "mem_gib_total": gpu_mem_info.total / (1024 ** 3) if gpu_mem_info is not None else None,
212 "temp": temp,
213 # Power in Watts (None if not supported by this device)
214 "power_draw_watts": power_draw_raw / 1000.0 if power_draw_raw is not None else None,
215 "power_limit_watts": power_limit_raw / 1000.0 if power_limit_raw is not None else None,
216 # Frequencies in MHz (None if not supported by this device)
217 "graphics_clock": graphics_clock,
218 "sm_clock": sm_clock,
219 "mem_clock": mem_clock,
220 })
222 return ret
223 except Exception as ex:
224 logging.warning(f"Error getting GPU info: {ex}.")
225 self.gpu_setup = False # Disable GPU info retrieval on error
226 return None
227 finally:
228 try:
229 nvidia_smi.nvmlShutdown()
230 except NVMLError as nvml_err:
231 logging.warning(f"Error shutting down NVML: {nvml_err}.")
233 def _new_gen_timer(
234 self,
235 job_id: Optional[str] = None
236 ) -> GenTimer:
237 if job_id is None:
238 job_id = datetime.now().strftime("%Y%m%dT%H%M%S%f")[:-3]
239 elif job_id in self.gen_timers:
240 new_job_id = job_id + "_" + datetime.now().strftime("%Y%m%dT%H%M%S%f")[:-3]
241 logging.info(f"[{self.rank}] Job '{job_id}' already exists, using '{new_job_id}'.")
242 job_id = new_job_id
243 gen_timer = GenTimer()
244 self.gen_timers[job_id] = gen_timer
245 return gen_timer
247 def get_health(self) -> Dict[str, Any]:
248 """Get health status of the model."""
249 ret = {
250 "model_name": self.model_name,
251 "status": self.status,
252 "running": self.running,
253 "load_timer": self.load_timer.to_dict() if self.load_timer else None,
254 "gen_timer": {
255 job_id: gen_timer.to_dict()
256 for job_id, gen_timer in self.gen_timers.items()
257 } if self.gen_timers else None,
258 }
260 gpu_info = self.get_gpu_info()
261 if gpu_info:
262 ret["gpu_info"] = gpu_info
264 return ret
266 def get_timestamps(self) -> List[dict]:
267 """Get timing timestamps for loading and generation."""
268 ret = []
269 if self.load_timer:
270 timestamps = self.load_timer.to_timestamps(
271 group=self.model_name,
272 subgroup="load")
273 if timestamps:
274 ret.extend(timestamps)
275 if self.gen_timers:
276 for gen_ix, (job_id, gen_timer) in enumerate(self.gen_timers.items()):
277 timestamps = gen_timer.to_timestamps(
278 group=self.model_name,
279 subgroup=job_id or f"gen_{gen_ix:04d}")
280 if timestamps:
281 ret.extend(timestamps)
282 return ret
284 @abstractmethod
285 async def generate(
286 self,
287 job_id: Optional[str] = None,
288 *args: Any,
289 **kwargs: Dict[str, Any]
290 ) -> Any:
291 """Generate output using the model."""
292 raise NotImplementedError("Method should be implemented in subclasses.")
294 @abstractmethod
295 async def warmup(self) -> None:
296 """Warmup the model with a sample generation."""
297 raise NotImplementedError("Method should be implemented in subclasses.")
299 @abstractmethod
300 async def get_rest_args(
301 self,
302 data_json: Dict[str, Union[str, int, float]]
303 ) -> Dict[str, Any]:
304 """Extract and validate arguments from the REST API request."""
305 if data_json is None or not isinstance(data_json, dict):
306 raise ValueError("Missing JSON body")
308 raise NotImplementedError("Method should be implemented in subclasses.")