Coverage for apps/client.py: 64%
309 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Client for managing requests to LMM services.
3"""
5import sys
6import time
7import logging
8import asyncio
9import json
11from urllib.parse import urlparse
13from pydantic import BaseModel
14from pydantic import Field
15from pydantic import model_validator
17from http import HTTPStatus
19from aiohttp import TCPConnector
20from aiohttp import ClientTimeout
21from aiohttp import ClientSession
22from aiohttp import ClientOSError
23from aiohttp import ClientError
24from asyncio import TimeoutError
26from enum import Enum
28from typing import List
29from typing import Optional
30from typing import Dict
31from typing import Tuple
32from typing import Any
33from typing import Union
35from lmm_service_manager import LMMServiceManager
37from client_timeout import SERVICE_TIMEOUT
38from client_timeout import SERVICE_LONG_TIMEOUT
40from client_headers import JSON_HEADERS
41from client_headers import BINARY_HEADERS
43sys.path.append("..") # noqa: E402
45from console_utils import setup_logging
46from console_utils import bytes_to_human
48import k8s_utils
49from k8s_utils import NoActiveContainerError
50from k8s_utils import NoRunnableContainerError
51from k8s_utils import ServiceNotFoundError
54class RequestStatus(str, Enum):
55 """Status of a service request, JSON-serializable as strings."""
56 CREATED = "CREATED"
57 PENDING = "PENDING"
58 RETRYING = "RETRYING"
59 RUNNING = "RUNNING"
60 COMPLETED = "COMPLETED"
61 FAILED = "FAILED"
62 CANCELLED = "CANCELLED"
63 EXPIRED = "EXPIRED"
64 UNKNOWN = "UNKNOWN"
67class ServiceRequest(BaseModel):
68 """A request to a service with its associated metadata and status."""
70 request_id: str
71 service_name: str
72 path: Optional[str] = None
73 base_url: Optional[str] = None
74 url: Optional[str] = None
75 deadline: Optional[float] = None
76 status: RequestStatus = RequestStatus.CREATED
77 retries: int = 0
78 times: List[Tuple[RequestStatus, float]] = Field(
79 default_factory=lambda: [(RequestStatus.CREATED, time.time())])
81 # Non-serializable, exclude from JSON
82 payload_json: Optional[Dict] = Field(default=None, exclude=True)
83 payload_bytes: Optional[bytes] = Field(default=None, exclude=True)
84 timeout: Optional[float] = Field(default_factory=lambda: SERVICE_TIMEOUT.total, exclude=True)
85 future: Optional[asyncio.Future] = Field(default=None, exclude=True)
86 exception: Optional[Exception] = Field(default=None, exclude=True)
87 tasks: List[asyncio.Task] = Field(default_factory=list, exclude=True)
89 model_config = {
90 "arbitrary_types_allowed": True
91 }
93 @model_validator(mode="before")
94 def validate_payload_and_convert(
95 cls: Any,
96 values: Dict[str, Any]
97 ) -> Dict[str, Any]:
98 """Ensure payload exists and do conversions before model creation."""
99 if not values.get("payload_json") and not values.get("payload_bytes"):
100 raise ValueError("Either payload_json or payload_bytes must be provided.")
101 timeout = values.get("timeout")
102 if isinstance(timeout, ClientTimeout): # ClientTimeout object
103 values["timeout"] = timeout.total
104 return values
106 @property
107 def client_timeout(self) -> ClientTimeout:
108 """Return an aiohttp ClientTimeout object for this request."""
109 return ClientTimeout(total=self.timeout)
111 def set_status(self, status: RequestStatus) -> None:
112 """Set the status of the request and record the time."""
113 self.status = status
114 self.times.append((status, time.time()))
116 def set_failure(self, ex: Exception) -> None:
117 """Set the request status to FAILED and record the exception."""
118 self.set_status(RequestStatus.FAILED)
119 if self.future:
120 self.future.set_exception(ex)
121 self.exception = ex
123 def done(self) -> bool:
124 """Check if the request is in a terminal state (COMPLETED, FAILED, CANCELLED, EXPIRED)."""
125 return self.status in {
126 RequestStatus.COMPLETED,
127 RequestStatus.FAILED,
128 RequestStatus.CANCELLED,
129 RequestStatus.EXPIRED,
130 }
132 def is_running(self) -> bool:
133 """Check if the request is currently running."""
134 return self.status == RequestStatus.RUNNING
136 def get_base_request_url(self) -> Optional[str]:
137 """
138 Get the base request URL (scheme, hostname, port) for the service request.
139 For example: http://10.244.0.97:8080/hunyuanframepackf1 -> http://10.244.0.97:8080
140 """
141 if self.url is None:
142 return None
143 parsed_url = urlparse(self.url)
144 base_request_url = f"{parsed_url.scheme}://{parsed_url.hostname}"
145 if parsed_url.port is not None:
146 base_request_url += f":{parsed_url.port}"
147 return base_request_url
149 def set_retry(self) -> None:
150 """Set the request status to RETRYING and increment the retry count."""
151 self.set_status(RequestStatus.RETRYING)
152 self.retries += 1
154 def dict(self, **kwargs: Any) -> Dict[str, Any]:
155 """Custom dict for JSON serialization."""
156 d = super().dict(**kwargs)
157 d["status"] = self.status.value
158 d["times"] = [(s.value, t) for s, t in self.times]
159 return d
161 def json(self, **kwargs: Any) -> str:
162 """Serialize to JSON string."""
163 return json.dumps(self.dict(**kwargs), **kwargs)
165 @classmethod
166 def parse_json(cls, data: str) -> "ServiceRequest":
167 """Parse a JSON string back into a ServiceRequest."""
168 obj = json.loads(data)
169 if "status" in obj:
170 obj["status"] = RequestStatus(obj["status"])
171 if "times" in obj:
172 obj["times"] = [(RequestStatus(s), t) for s, t in obj["times"]]
173 return cls(**obj)
175 def get_payload_len(self) -> int:
176 """Get the length of the payload in bytes, whether it's JSON or binary data."""
177 if self.payload_bytes:
178 return len(self.payload_bytes)
179 if self.payload_json:
180 return len(json.dumps(self.payload_json).encode("utf-8"))
181 return 0
184class ServiceError(Exception):
185 """Exception for service errors."""
187 def __init__(
188 self,
189 service_name: Optional[str] = None,
190 job_id: Optional[str] = None,
191 request_id: Optional[str] = None,
192 message: Optional[str] = None,
193 url: Optional[str] = None,
194 status: Optional[Union[int, HTTPStatus, str]] = None,
195 response_body: Optional[str] = None,
196 content_type: Optional[str] = None
197 ) -> None:
198 """Initialize the ServiceError exception."""
199 super().__init__(message)
200 self.job_id = job_id
201 self.request_id = request_id
202 self.message = message
203 self.service_name = service_name
204 self.url = url
205 self.status: int | HTTPStatus | str = status or HTTPStatus.INTERNAL_SERVER_ERROR
206 self.response_body = response_body
207 self.content_type = content_type
209 def __str__(self) -> str:
210 """String representation of the ServiceError."""
211 ret = f"{self.message} (Service:{self.service_name}"
212 if self.job_id:
213 ret += f" Job:{self.job_id}"
214 if self.url:
215 ret += f" URL:{self.url}"
216 if self.status:
217 ret += f" Status:{self.status}"
218 if self.response_body:
219 ret += f" Response:{self.response_body}"
220 if self.content_type:
221 ret += f" Type:{self.content_type}"
222 ret += ")"
223 return ret
226class ServiceRequestWorker:
227 """A worker that processes service requests asynchronously and manages their execution."""
229 def __init__(
230 self,
231 app_name: str,
232 service_manager: LMMServiceManager
233 ) -> None:
234 """Initialize the service request worker."""
235 self.app_name = app_name
236 self.service_manager = service_manager
237 self.running = True
239 connector = TCPConnector(
240 limit=100,
241 limit_per_host=10,
242 use_dns_cache=True,
243 force_close=True,
244 ssl=k8s_utils.VERIFY_SSL)
245 self.session: Optional[ClientSession] = ClientSession(
246 connector=connector,
247 timeout=SERVICE_LONG_TIMEOUT)
249 self.queues: Dict[str, List[ServiceRequest]] = {}
250 self.requests: Dict[str, ServiceRequest] = {}
252 self.logger = self._get_logger()
254 def _get_logger(self) -> logging.Logger:
255 """Get the logger for the service worker."""
256 logger = setup_logging(
257 path=f"/tmp/{self.app_name}/logs",
258 file_name="service_worker.log",
259 level=logging.INFO)
260 return logger
262 """
263 def _get_logger(self) -> logging.Logger:
264 log_dir = "/tmp/streamwise/logs"
265 os.makedirs(log_dir, exist_ok=True)
266 log_file = os.path.join(log_dir, "service_manager.log")
267 logger = logging.getLogger("service_manager")
268 logger.setLevel(logging.INFO)
269 if not logger.handlers:
270 file_handler = logging.FileHandler(log_file)
271 # formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
272 formatter = ColoredFormatter(
273 fmt="[%(asctime)s] %(log_color)s%(levelname)s: %(message)s",
274 datefmt="%Y-%m-%d %H:%M:%S",
275 log_colors={
276 "DEBUG": "cyan",
277 "INFO": "white",
278 "WARNING": "yellow",
279 "ERROR": 'red',
280 "CRITICAL": 'bold_red',
281 }
282 )
283 file_handler.setFormatter(formatter)
284 logger.addHandler(file_handler)
285 return logger
286 """
288 def get_queued_requests(self) -> List[str]:
289 """Get a list of request ids that are currently queued for processing."""
290 ret = []
291 for queue in self.queues.values():
292 for request in queue:
293 ret.append(request.request_id)
294 return ret
296 def get_requests(self) -> Dict[str, ServiceRequest]:
297 """Get a dictionary of all requests currently being managed by the worker."""
298 return self.requests
300 async def start(self) -> None:
301 """
302 Start the worker to process requests asynchronously.
303 It waits if nothing available.
304 """
305 MAX_RETRY_SECONDS = 0.5
306 INIT_RETRY_SECONDS = 0.05
307 retry_seconds = INIT_RETRY_SECONDS
308 while self.running:
309 processed_requests = self._process_requests()
310 if processed_requests == 0:
311 await asyncio.sleep(retry_seconds)
312 # Exponential backoff up to max
313 retry_seconds = min(retry_seconds * 2, MAX_RETRY_SECONDS)
314 else:
315 retry_seconds = INIT_RETRY_SECONDS # Reset backoff
317 def _process_requests(self) -> int:
318 """Process queued requests for all services."""
319 processed_requests = 0
320 for queue in self.queues.values():
321 processed_requests += self._process_requests_queue(queue)
322 return processed_requests
324 def _process_requests_queue(
325 self,
326 queue: List[ServiceRequest]
327 ) -> int:
328 """Process requests in a specific service queue."""
329 if not queue:
330 return 0
332 next_request = None
333 for queued_request in queue:
334 if not next_request:
335 next_request = queued_request
336 elif not next_request.deadline and queued_request.deadline:
337 # Take the one with a deadline over one without
338 next_request = queued_request
339 elif (queued_request.deadline is not None
340 and next_request.deadline is not None
341 and queued_request.deadline < next_request.deadline):
342 # Take requests with earlier deadlines first
343 next_request = queued_request
345 if not next_request:
346 return 0
348 # Take from the queue and send it
349 queue.remove(next_request)
350 task = asyncio.create_task(self._http_request(next_request))
351 next_request.tasks.append(task)
352 return 1
354 async def stop(self) -> None:
355 """Stop the service request worker."""
356 self.running = False
357 if self.session:
358 await self.session.close()
359 self.session = None
361 async def submit_request(
362 self,
363 request: ServiceRequest
364 ) -> asyncio.Future:
365 """Submit a request to the service and return an asyncio future."""
366 self.logger.debug(f"Submitting request {request.request_id} to service {request.service_name}.")
367 if request.service_name not in self.service_manager.services:
368 raise ServiceNotFoundError(request.service_name)
370 if request.future is not None:
371 if request.future.done():
372 raise ValueError(f"Request {request.request_id} already has a result, cannot submit again.")
373 else:
374 async_loop = asyncio.get_running_loop()
375 request.future = async_loop.create_future()
377 if request.service_name not in self.queues:
378 queue: List[ServiceRequest] = []
379 self.queues[request.service_name] = queue
380 else:
381 queue = self.queues[request.service_name]
382 queue.append(request)
383 self.requests[request.request_id] = request
385 return request.future
387 async def _http_request(
388 self,
389 request: ServiceRequest
390 ) -> None:
391 """Perform the HTTP request to the service and handle the response."""
392 t0 = time.time()
393 try:
394 if request.base_url is not None:
395 base_url = request.base_url
396 else:
397 base_urls = self.service_manager.get_service_urls(request.service_name)
398 base_url = base_urls[0]
400 path = request.service_name
401 if request.path is not None:
402 path = request.path
403 request.url = f"{base_url}/{path}"
404 except NoRunnableContainerError:
405 # Wait until a container becomes available, queue it again to retry later
406 self.logger.debug(f"Service {request.service_name} has no runnable containers, retrying...")
407 request.set_retry()
408 task = asyncio.create_task(self.submit_request(request))
409 request.tasks.append(task)
410 return
411 except NoActiveContainerError as nac_ex:
412 # Sometimes the update of the status gets out of sync so we retry
413 if not nac_ex.containers:
414 request.set_status(RequestStatus.FAILED)
415 ex = ServiceError(
416 message=f"No active containers for {request.service_name}: {nac_ex.containers}.",
417 request_id=request.request_id,
418 service_name=request.service_name)
419 request.set_failure(ex)
420 return
421 else:
422 # Wait until a container becomes available, queue it again to retry later
423 self.logger.debug(f"Service {request.service_name} has no active containers, retrying...")
424 request.set_retry()
425 task = asyncio.create_task(self.submit_request(request))
426 request.tasks.append(task)
427 return
429 try:
430 request.set_status(RequestStatus.RUNNING)
431 post_args: Dict[str, Any] = {
432 "headers": {},
433 }
434 if request.payload_json is not None:
435 post_args = {
436 "json": request.payload_json,
437 "headers": JSON_HEADERS,
438 }
439 elif request.payload_bytes is not None:
440 post_args = {
441 "data": request.payload_bytes,
442 "headers": BINARY_HEADERS,
443 }
444 post_args["timeout"] = request.timeout
446 assert self.session is not None
447 async with self.session.post(request.url, **post_args) as response:
448 content_type = response.headers.get("Content-Type", "")
449 if response.status == HTTPStatus.OK:
450 response_binary = await response.read()
452 if request.future is not None:
453 request.future.set_result((
454 content_type,
455 response_binary
456 ))
457 request.set_status(RequestStatus.COMPLETED)
458 self.logger.info(
459 f"Request {request.request_id} with {bytes_to_human(request.get_payload_len())} "
460 f"to {request.service_name}@{request.url} completed with "
461 f"{bytes_to_human(len(response_binary))} and type {content_type}.")
462 return
464 # Handle error response
465 err_msg = ""
466 if content_type == "application/json":
467 response_json = await response.json()
468 err_msg = response_json.get("error")
469 else:
470 err_msg = await response.text()
472 if "generation in progress" in err_msg.lower() or "no runnable containers" in err_msg.lower():
473 # Queue it again to retry later
474 request.set_retry()
475 task = asyncio.create_task(self.submit_request(request))
476 request.tasks.append(task)
477 return
478 if "request entity too large" in err_msg.lower():
479 payload_size = request.get_payload_len()
480 self.logger.error(
481 f"Request too long for '{request.service_name}': {err_msg}. Size: {payload_size} bytes.")
483 self.logger.error(
484 f"Request {request.request_id} to {request.service_name}@{request.url} failed: {err_msg}.")
485 ex = ServiceError(
486 message="Service request failed",
487 service_name=request.service_name,
488 request_id=request.request_id,
489 url=request.url,
490 status=response.status,
491 response_body=err_msg)
492 request.set_failure(ex)
493 except TimeoutError: # usually empty
494 msg = f"Timeout for {request.request_id} to {request.service_name}@{request.url}: "
495 msg += f"{time.time() - t0:.3f} > {request.timeout}"
496 self.logger.error(f"{msg}.")
497 # TODO should we retry?
498 request.set_failure(TimeoutError(msg))
499 except ClientOSError as client_os_err:
500 payload_len = len(json.dumps(post_args["data"]).encode("utf-8"))
501 self.logger.error(
502 f"Client OS error for {request.request_id} "
503 f"with {bytes_to_human(payload_len)} "
504 f"to {request.service_name}@{request.url}: {client_os_err}.")
505 if "broken pipe" in str(client_os_err).lower():
506 # Usually happens when querying the same multiple times
507 # [Errno 32] Broken pipe
508 request.set_retry()
509 task = asyncio.create_task(self.submit_request(request))
510 request.tasks.append(task)
511 else:
512 request.set_failure(client_os_err)
513 except ClientError as client_err:
514 self.logger.error(
515 f"Client error for {request.request_id} to {request.service_name}@{request.url}: {client_err}.")
516 if "server disconnected" in str(client_err).lower():
517 request.set_retry()
518 task = asyncio.create_task(self.submit_request(request))
519 request.tasks.append(task)
520 else:
521 request.set_failure(client_err)
522 except AssertionError as assert_err:
523 self.logger.error(
524 f"Assertion error for {request.request_id} to {request.service_name}@{request.url}: {assert_err}.")
525 request.set_failure(assert_err)
526 except ValueError as value_err:
527 self.logger.error(
528 f"Value error for {request.request_id} to {request.service_name}@{request.url}: {value_err}.")
529 request.set_failure(value_err)
530 except Exception as ex:
531 ex_name = type(ex).__name__
532 self.logger.error(
533 f"Error for {request.request_id} to {request.service_name}@{request.url}: {ex} [type:{ex_name}].")
534 request.set_failure(ex)