Coverage for quart_utils.py: 79%

1"""

2Utils for Quart/Flask applications, including filters and helper functions.

3"""

4import os

5import re

6import logging

7import json

8import aiofiles

9import aiofiles.os

10import mimetypes

12from http import HTTPStatus

14from datetime import datetime

15from datetime import timedelta

17from typing import Union

18from typing import Tuple

19from typing import Any

20from typing import List

21from typing import Dict

22from typing import Optional

24from quart import Response

27QuartReturn = Union[

28 str,

29 Response,

30 Tuple[Response, int],

31 Dict[str, Any],

32 Tuple[str, HTTPStatus],

33 Tuple[Dict[str, Any], HTTPStatus],

34]

37def json_pretty_filter(

38 value: str,

39 max_len: int = 128

40) -> str:

41 def truncate(obj: Any) -> Union[str, List, Dict]:

42 if isinstance(obj, str) and len(obj) > max_len:

43 return obj[:max_len] + f"... [truncated, {format_bytes(len(obj))}]"

44 if isinstance(obj, list):

45 return [truncate(item) for item in obj]

46 if isinstance(obj, dict):

47 return {k: truncate(v) for k, v in obj.items()}

48 return obj

50 try:

51 obj = json.loads(value) if isinstance(value, str) else value

52 truncated = truncate(obj)

53 return json.dumps(truncated, indent=2, ensure_ascii=False)

54 except Exception:

55 return value

58def format_datetime(value: int) -> str:

59 try:

60 dt = datetime.fromtimestamp(value)

61 return dt.strftime('%Y-%m-%d %H:%M:%S')

62 except (ValueError, OSError, TypeError):

63 return "Invalid date"

66def format_bytes(memory: int) -> str:

67 if memory == 0:

68 return '-'

69 if memory < 1024:

70 return f"{memory} B"

71 if memory < 1024 ** 2:

72 return f"{memory / 1024:.1f} KiB"

73 if memory < 1024 ** 3:

74 return f"{memory / 1024 / 1024:.1f} MiB"

75 if memory < 1024 ** 4:

76 return f"{memory / 1024 / 1024 / 1024:.1f} GiB"

77 return f"{memory / 1024 / 1024 / 1024 / 1024:.1f} TiB"

80def format_string(in_string: Optional[str]) -> Optional[str]:

81 """Format a string to be more readable."""

82 if not isinstance(in_string, str):

83 return in_string

84 ret = in_string.replace("_", " ").title()

86 # Make some words uppercase

87 UPPER_WORDS = ["gpu", "cpu", "api", "sm", "http"]

88 # Fix the capitalization of some words

89 REPLACE_WORDS = {

90 "gib": "GiB",

91 "mib": "MiB",

92 "kib": "KiB",

93 "mbps": "Mbps",

94 "gbps": "Gbps",

95 "tbps": "Tbps",

96 "vcpu": "vCPU",

97 }

99 words = ret.split(" ")

100 for i in range(len(words)):

101 word_lower = words[i].lower()

102 if word_lower in UPPER_WORDS:

103 words[i] = word_lower.upper()

104 elif word_lower in REPLACE_WORDS:

105 words[i] = REPLACE_WORDS[word_lower]

106 ret = " ".join(words)

107

108 return ret

109

110

111def format_duration(input_date: Optional[Union[timedelta, float]]) -> str:

112 """Format a duration (timedelta) into a human-readable string."""

113 if not input_date:

114 return "0"

115 if isinstance(input_date, float):

116 total_seconds = int(input_date)

117 elif isinstance(input_date, timedelta):

118 total_seconds = int(input_date.total_seconds())

119 else:

120 raise TypeError("input_date must be a timedelta or float")

121 if total_seconds < 0:

122 return "?"

123

124 hours, remainder = divmod(total_seconds, 60 * 60)

125 minutes, seconds = divmod(remainder, 60)

126

127 if hours > 1:

128 return f"{hours} hours {minutes} minutes {seconds} seconds"

129 if hours > 0:

130 return f"{hours} hour {minutes} minutes {seconds} seconds"

131 if minutes > 1:

132 return f"{minutes} minutes {seconds} seconds"

133 if minutes > 0:

134 return f"{minutes} minute {seconds} seconds"

135 return f"{seconds} seconds"

136

137

138def format_duration_short(input_date: Optional[Union[timedelta, float]]) -> str:

139 """Format a duration (timedelta) into a human-readable string."""

140 if not input_date:

141 return "0"

142 if isinstance(input_date, float):

143 total_seconds = int(input_date)

144 elif isinstance(input_date, timedelta):

145 total_seconds = int(input_date.total_seconds())

146 else:

147 raise TypeError("input_date must be a timedelta or float")

148 hours, remainder = divmod(total_seconds, 60 * 60)

149 minutes, seconds = divmod(remainder, 60)

150 return f"{hours:02d}:{minutes:02d}:{seconds:02d}"

151

152

153def format_url(url: Optional[str]) -> Optional[str]:

154 """Format a URL to remove the scheme (http/https)."""

155 if not url:

156 return url

157 url = url.strip()

158 if url.startswith("http://"):

159 return url[7:]

160 if url.startswith("https://"):

161 return url[8:]

162 return url

163

164

165_AZURE_VM_SKU_RE = re.compile(r'^Standard_(?:NC|ND|NV)\d+[a-zA-Z]*_([A-Za-z0-9]+)_v\d+$', re.IGNORECASE)

166

167# Matches a trailing MIG profile like "MIG 1g.10gb" or "MIG 3g.40gb"

168_MIG_PROFILE_RE = re.compile(r'\bMIG\s+(\d+g\.\d+gb)\b', re.IGNORECASE)

169

170_AZURE_VM_GPU_MAP: Dict[str, str] = {

171 "A100": "A100 80GB",

172 "GB200": "GB200",

173 "GB300": "GB300",

174 "H100": "H100",

175 "H200": "H200",

176 "MI300X": "MI300X",

177 "T4": "T4",

178}

179

180# Ordered list of (pattern, display_name) for raw GPU model strings.

181# Patterns are matched against the model string with hyphens replaced by spaces.

182# More specific patterns (e.g. "A100 40GB") must appear before broader ones ("A100 80GB").

183_GPU_MODEL_PATTERNS: List[Tuple[str, str]] = [

184 (r'\bA100\b.*\b40\s*GB\b', "A100 40GB"),

185 (r'\bA100\b', "A100 80GB"),

186 (r'\bGB300\b', "GB300"),

187 (r'\bGB200\b', "GB200"),

188 (r'\bH100\b.*\bNVL\b', "H100 NVL"),

189 (r'\bNVL\b.*\bH100\b', "H100 NVL"),

190 (r'\bH100\b', "H100"),

191 (r'\bH200\b', "H200"),

192 (r'\bV100\b.*\b(16\s*GB|PCIE)\b', "V100 16GB"),

193 (r'\b(16\s*GB|PCIE)\b.*\bV100\b', "V100 16GB"),

194 (r'\bV100\b.*\b(32\s*GB|SXM)\b', "V100 32GB"),

195 (r'\b(32\s*GB|SXM)\b.*\bV100\b', "V100 32GB"),

196 (r'\bV100\b', "V100"),

197 (r'\bT4\b', "T4"),

198]

199

200

201def format_gpu_model(

202 gpu_model: Optional[str]

203) -> Optional[str]:

204 """Format GPU model names to be more user-friendly.

205

206 Handles Azure VM SKU names (e.g. Standard_ND96ams_A100_v4) and raw GPU

207 model strings reported by nvidia-smi (e.g. NVIDIA A100-SXM4-80GB).

208 MIG instance names (e.g. NVIDIA A100-SXM4-80GB MIG 1g.10gb) are formatted

209 without the MIG profile (e.g. A100 80GB).

210

211 Azure GPU VM sizes reference:

212 https://learn.microsoft.com/en-us/azure/virtual-machines/sizes/gpu-accelerated/nd-family

213 """

214 if not gpu_model or gpu_model == "N/A":

215 return gpu_model

216 if not isinstance(gpu_model, str):

217 return gpu_model

218 gpu_model = gpu_model.strip()

219

220 # Azure VM SKU: Standard_<series><size>_<GPU>_v<version>

221 # e.g. Standard_NC96ads_A100_v4, Standard_ND96ams_A100_v4, Standard_ND96isrf_H100_v5

222 azure_match = _AZURE_VM_SKU_RE.match(gpu_model)

223 if azure_match:

224 gpu_part = azure_match.group(1).upper()

225 return _AZURE_VM_GPU_MAP.get(gpu_part, gpu_part)

226

227 # Extract and strip the MIG profile before matching the base GPU model,

228 # then re-append it so the profile is preserved in the display name.

229 # e.g. "NVIDIA A100-SXM4-80GB MIG 1g.10gb" -> "A100 80GB"

230 mig_match = _MIG_PROFILE_RE.search(gpu_model)

231 base_model = _MIG_PROFILE_RE.sub("", gpu_model).strip() if mig_match else gpu_model

232

233 # Normalize hyphens to spaces for pattern matching

234 normalized = base_model.replace("-", " ")

235

236 for pattern, display_name in _GPU_MODEL_PATTERNS:

237 if re.search(pattern, normalized, re.IGNORECASE):

238 return display_name # + mig_suffix

239

240 return gpu_model

241

242

243def get_gpu_mem(gpu_model: Optional[str]) -> int:

244 """Get GPU memory in GB based on the GPU model string."""

245 if gpu_model is None:

246 return 0

247 if "16GB" in gpu_model:

248 return 16

249 if "16GB" in gpu_model:

250 return 32

251 if "40GB" in gpu_model:

252 return 40

253 if "A100" in gpu_model:

254 return 80

255 if "H100" in gpu_model:

256 return 80

257 if "H200" in gpu_model:

258 return 141

259 if "GB200" in gpu_model:

260 return 196

261 if "GB300" in gpu_model:

262 return 288

263 if "MI300" in gpu_model:

264 return 192

265 return 80

266

267

268def format_gpu_model_mig(

269 gpu_model: Optional[str],

270 mig_profile: Optional[str],

271) -> str:

272 if gpu_model is None:

273 return ""

274 mem_gb = get_gpu_mem(gpu_model)

275 if mig_profile is None:

276 return gpu_model

277 if mem_gb == 80:

278 if mig_profile == "1g.10gb":

279 return f"⅛ {gpu_model}"

280 if mig_profile == "2g.20gb":

281 return f"¼ {gpu_model}"

282 if mig_profile == "3g.40gb":

283 return f"½ {gpu_model}"

284 if mig_profile == "4g.40gb":

285 return f"½ {gpu_model}"

286 if mem_gb == 40:

287 if mig_profile == "1g.5gb":

288 return f"⅛ {gpu_model}"

289 if mig_profile == "2g.10gb":

290 return f"¼ {gpu_model}"

291 if mig_profile == "3g.20gb":

292 return f"½ {gpu_model}"

293 if mig_profile == "4g.20gb":

294 return f"½ {gpu_model}"

295 return f"{gpu_model} {mig_profile}"

296

297

298def get_aspect_ratio(ratio: float) -> str:

299 if abs(ratio - 1) < 0.01:

300 return "1:1"

301 if abs(ratio - (16. / 9.0)) < 0.01:

302 return "16:9"

303 if abs(ratio - (16. / 10.0)) < 0.01:

304 return "16:10"

305 if abs(ratio - (4. / 3.0)) < 0.01:

306 return "4:3"

307 if abs(ratio - (5. / 4.0)) < 0.01:

308 return "5:4"

309 if abs(ratio - (3. / 2.0)) < 0.01:

310 return "3:2"

311 if abs(ratio - (2. / 1.0)) < 0.01:

312 return "2:1"

313 return f"{ratio:.2f}:1"

314

315

316async def get_service_json_filename() -> Optional[str]:

317 if await aiofiles.os.path.exists("../services.json"):

318 return "../services.json"

319 if await aiofiles.os.path.exists("services.json"):

320 return "services.json"

321 logging.warning("services.json file not found")

322 return None

323

324

325def get_k8s_service_emoji(container_name: Optional[str]) -> str:

326 if container_name is None:

327 return ""

328 if container_name.startswith("nvidia-device-plugin-"):

329 return "⚙️"

330 if container_name.startswith("node-driver-"):

331 return "🧩"

332 if container_name.startswith("azure"):

333 return "☁️"

334 if container_name.startswith("metrics-server"):

335 return "📊"

336 if container_name.startswith("liveness-"):

337 return "💓"

338 if container_name.startswith("kube-proxy"):

339 return "🔀"

340 if container_name.startswith("gatekeeper"):

341 return "🛡️"

342 if container_name.startswith("cns-container"):

343 return "🧩"

344 if container_name.startswith("coredns"):

345 return "🌐"

346 if container_name.startswith("autoscaler"):

347 return "📈"

348 if container_name.startswith("cloud-node-manager"):

349 return "☁️"

350 if "security" in container_name:

351 return "🛡️"

352 if "azsec" in container_name:

353 return "🛡️"

354 if "konnect" in container_name:

355 return "🌐"

356 if "log" in container_name:

357 return "📄"

358 if "auoms" in container_name:

359 return "📄"

360 if "mdsdmgr" in container_name:

361 return "📊"

362 if "debug" in container_name:

363 return "🐞"

364 return f"❓"

365

366

367async def get_class_emoji(container_name: Optional[str]) -> str:

368 if container_name is None:

369 return ""

370 services_file_name = await get_service_json_filename()

371 if services_file_name is None:

372 return ""

373 async with aiofiles.open(services_file_name, mode="r") as file:

374 data_str = await file.read()

375 data_json = json.loads(data_str)

376

377 if container_name not in data_json:

378 return get_k8s_service_emoji(container_name)

379

380 data_service_json = data_json[container_name]

381 if "class" not in data_service_json:

382 return f"❓"

383 service_class = data_service_json["class"]

384

385 if service_class == "text2audio" or service_class == "text2speech":

386 return "📄→🔉"

387 if service_class == "video2audio":

388 return "🎬→🔉"

389 if service_class == "text2image":

390 return "📄→🖼️"

391 if service_class == "text2video":

392 return "📄→🎬"

393 if service_class == "image2video":

394 return "📄→🎬"

395 if service_class == "textimageaudio2video":

396 return "📄🖼️🔉→🎬"

397 if service_class == "image2image":

398 return "🖼️→🖼️"

399 if service_class == "text2text":

400 return "📄→📄"

401 if service_class == "textimage2video":

402 return "📄🖼️→🎬"

403 if service_class == "manager":

404 return "⚒️"

405 if service_class == "doc2video":

406 return "📄→🎬"

407 if service_class == "video2video":

408 return "🎬→🎬"

409 if service_class == "audio2audio":

410 return "🔉→🔉"

411 if service_class == "audio2text":

412 return "🔉→📄"

413

414 # VAE

415 if service_class == "latent2image":

416 return "🔢→🖼️"

417 if service_class == "latent2video":

418 return "🔢→🎬"

419 if service_class == "image2latent":

420 return "🖼️→🔢"

421 if service_class == "video2latent":

422 return "🎬→🔢"

423

424 # Unknown class

425 return f"❓"

426

427

428def get_file_type_emoji(file_type: Optional[str]) -> str:

429 if file_type is None:

430 return "❓"

431 if file_type == "directory":

432 return "📁"

433 if file_type == "image":

434 return "🖼️"

435 if file_type == "audio":

436 return "🎵"

437 if file_type == "video":

438 return "🎥"

439 if file_type == "archive":

440 return "📦"

441 if file_type in ("file", "text", "json", "jsonl", "x-ndjson", "pdf"):

442 return "📄"

443 if file_type == "tensor":

444 return "📊"

445 if file_type == "kernel":

446 return "📊"

447 if file_type == "presentation":

448 return "📊"

449 if file_type == "base64":

450 return "📄"

451 return f"❓ {file_type}"

452

453

454def get_file_type(file_name: str) -> str:

455 """Determine the file type based on its extension."""

456 file_name = file_name.lower()

457 if file_name.endswith((".mp4", ".avi", ".mkv")):

458 return "video"

459 if file_name.endswith((".wav", ".mp3", ".aac", ".flac", ".ogg", ".m4a")):

460 return "audio"

461 if file_name.endswith((".png", ".jpg", ".jpeg")):

462 return "image"

463 if file_name.endswith((".log", ".txt")):

464 return "text"

465 if file_name.endswith(".pt"):

466 return "tensor"

467 if file_name.endswith(".ptx"):

468 return "kernel"

469 if file_name.endswith(".json"):

470 return "json"

471 if file_name.endswith(".jsonl"):

472 return "jsonl"

473 if file_name.endswith(".pdf"):

474 return "pdf"

475 if file_name.endswith((".pptx", ".ppt")):

476 return "presentation"

477 if file_name.endswith(".base64"):

478 return "base64"

479 if file_name.endswith((".zip", ".tar", ".gz", ".bz2", ".7z")):

480 return "archive"

481 return "unknown"

482

483

484def get_mime_type(file_name: str) -> str:

485 mimetype, _ = mimetypes.guess_type(file_name)

486 if mimetype is None:

487 mimetype = "application/octet-stream"

488 if file_name.endswith(".log"):

489 mimetype = "text/plain"

490 if file_name.endswith(".jsonl"):

491 mimetype = "application/x-ndjson"

492 return mimetype

493

494

495def get_content_type_emoji(content_type: str) -> str:

496 if content_type is None:

497 return "❓"

498 if content_type.startswith("text/"):

499 return "📄"

500 if content_type.startswith("image/"):

501 return "🖼️"

502 if content_type.startswith("audio/"):

503 return "🎵"

504 if content_type.startswith("video/"):

505 return "🎥"

506 if content_type == "application/json":

507 return "📄"

508 if content_type == "application/x-ndjson":

509 return "📄"

510 if content_type == "application/pdf":

511 return "📄"

512 if content_type == "application/octet-stream":

513 return "📦"

514 return f"❓ {content_type}"

515

516

517def get_friendly_region_name(region: Optional[str]) -> str:

518 if not region or region == "N/A":

519 return "N/A"

520 region_lower = region.lower()

521 if region_lower.startswith("eastus2"):

522 return "East US 2"

523 if region_lower.startswith("eastus"):

524 return "East US"

525 if region_lower.startswith("westus3"):

526 return "West US 3"

527 if region_lower.startswith("westus2"):

528 return "West US 2"

529 if region_lower.startswith("westus"):

530 return "West US"

531 if region_lower.startswith("centralus"):

532 return "Central US"

533 if region_lower.startswith("northcentralus"):

534 return "North Central US"

535 if region_lower.startswith("southcentralus"):

536 return "South Central US"

537 if region_lower.startswith("westeurope"):

538 return "West Europe"

539 if region_lower.startswith("eastasia"):

540 return "East Asia"

541 if region_lower.startswith("southeastasia"):

542 return "Southeast Asia"

543 if region_lower.startswith("swedencentral"):

544 return "Sweden Central"

545 return region.capitalize()

546

547

548async def get_friendly_container_name(container_name: str) -> str:

549 """

550 Get a friendly name for a container from services.json.

551 For example,

552 "imageresize" -> "Image Resize"

553 "fantasytalking" -> "Fantasy Talking"

554 """

555 services_file_name = await get_service_json_filename()

556 if services_file_name is None:

557 return container_name

558 async with aiofiles.open(services_file_name) as file:

559 data_str = await file.read()

560 data_json = json.loads(data_str)

561 if container_name not in data_json:

562 return container_name

563 data_service_json = data_json[container_name]

564 friendly_name = data_service_json["friendlyName"]

565 return friendly_name

566

567

568async def get_friendly_pod_name(pod_name: str) -> str:

569 """Get a friendly name for a pod from services.json."""

570 container_name = pod_name.split("-")[0]

571 return await get_friendly_container_name(container_name)

572

573

574async def get_friendly_model_name(model_name: str) -> str:

575 """Get a friendly name for a service from services.json."""

576 return await get_friendly_container_name(model_name)

577

578

579async def is_rtgen_container(container_name: str) -> bool:

580 services_file_name = await get_service_json_filename()

581 if services_file_name is None:

582 return False

583 async with aiofiles.open(services_file_name) as file:

584 data_str = await file.read()

585 data_json = json.loads(data_str)

586 return container_name in data_json.keys()

587

588

589async def get_docker_image(

590 container_name: str,

591 tag: Optional[str] = None

592) -> Optional[str]:

593 """

594 Get the docker image for a container from services.json.

595 If tag is provided, it overrides the tag defined in services.json.

596 """

597 services_file_name = await get_service_json_filename()

598 if services_file_name is None:

599 return ""

600 async with aiofiles.open(services_file_name, "r") as file:

601 data_str = await file.read()

602 data_json = json.loads(data_str)

603 if container_name not in data_json:

604 return None

605 data_service_json = data_json[container_name]

606 docker_image = data_service_json["dockerImage"]

607 docker_repo = docker_image.get("repository", os.environ.get("DOCKER_REPO", ""))

608 docker_name = docker_image["name"]

609 docker_tag = tag if tag else docker_image["tag"]

610 return f"{docker_repo}/{docker_name}:{docker_tag}"

611

612

613def parse_request_id(request_id: str) -> Dict[str, str]:

614 """

615 Parse request ids like:

616 20250904T010335296_flux

617 20250904T010335296_007_kokoro

618 20250904T010335296_006_001_fantasytalking

619 20260105T194652416_main_image_flux

620 """

621 match = re.match(

622 r"^(\d{8}T\d{9})" # job_id

623 r"(?:_(\d+))?" # scene_id (optional)

624 r"(?:_(\d+))?" # sub_scene_id (optional)

625 r"(?:_(.*?))?" # task_id (optional, lazy)

626 r"_([^_]+)$", # service_name (last part)

627 request_id

628 )

629 if not match:

630 return {}

631 job_id = match.group(1) or ""

632 scene_id = match.group(2) or ""

633 sub_scene_id = match.group(3) or ""

634 task_id = match.group(4) or ""

635 service_name = match.group(5) or ""

636 return {

637 "job_id": job_id,

638 "scene_id": scene_id,

639 "sub_scene_id": sub_scene_id,

640 "task_id": task_id,

641 "service_name": service_name,

642 }

643

644

645async def list_files(

646 folder_path: str

647) -> List[Dict[str, Any]]:

648 """List files in a folder with their metadata."""

649 files = []

650 file_names = await aiofiles.os.listdir(folder_path)

651 for file_name in file_names:

652 try:

653 file_path = os.path.join(folder_path, file_name)

654 file_date = await aiofiles.os.path.getmtime(file_path)

655 mimetype, _ = mimetypes.guess_type(file_name)

656 if await aiofiles.os.path.isfile(file_path):

657 file_size = await aiofiles.os.path.getsize(file_path)

658 file_type = get_file_type(file_name)

659 files.append({

660 "name": file_name,

661 "size": file_size,

662 "date": file_date,

663 "type": file_type,

664 "mimetype": mimetype,

665 })

666 elif await aiofiles.os.path.isdir(file_path):

667 files.append({

668 "name": file_name,

669 "size": 0,

670 "date": file_date,

671 "type": "directory",

672 "mimetype": "inode/directory",

673 })

674 except PermissionError:

675 files.append({

676 "name": file_name,

677 "size": 0,

678 "date": 0,

679 "type": "unknown",

680 "mimetype": "unknown",

681 "error": "Permission denied"

682 })

683 except Exception as ex:

684 logging.error(f"Error accessing file {file_name}: {ex}.")

685 files.append({

686 "name": file_name,

687 "size": 0,

688 "date": 0,

689 "type": "unknown",

690 "mimetype": "unknown",

691 "error": str(ex)

692 })

693 return files