|
6 | 6 |
|
7 | 7 | import argparse |
8 | 8 | import contextlib |
| 9 | +import dataclasses |
9 | 10 | import datetime |
| 11 | +import inspect |
10 | 12 | import logging |
11 | 13 | import os |
12 | 14 | import sys |
13 | 15 | import time |
14 | 16 | from contextlib import contextmanager |
| 17 | +from enum import Enum |
15 | 18 | from functools import lru_cache, partial |
16 | 19 | from logging import Logger |
17 | 20 | from types import MethodType |
@@ -281,6 +284,107 @@ def _method( |
281 | 284 | logger = init_logger(__name__) |
282 | 285 |
|
283 | 286 |
|
| 287 | +def _is_torch_tensor(obj: Any) -> tuple[bool, Any]: |
| 288 | + """Return (is_tensor, torch_module_or_None) without importing torch at module import time.""" |
| 289 | + try: |
| 290 | + import torch # type: ignore |
| 291 | + |
| 292 | + return isinstance(obj, torch.Tensor), torch |
| 293 | + except Exception: |
| 294 | + return False, None |
| 295 | + |
| 296 | + |
| 297 | +def _sanitize_for_logging(obj: Any, key_hint: str | None = None) -> Any: |
| 298 | + """Recursively convert objects to JSON-serializable forms for concise logging. |
| 299 | +
|
| 300 | + Rules: |
| 301 | + - Drop any field/dict key named 'param_names_mapping'. |
| 302 | + - Render Enums using their value. |
| 303 | + - Render torch.Tensor as a compact summary; if key name is 'scaling_factor', include stats. |
| 304 | + - Dataclasses are expanded to dicts and sanitized recursively. |
| 305 | + - Callables/functions are rendered as their qualified name. |
| 306 | + - Redact sensitive fields like 'prompt' and 'negative_prompt' (only show length). |
| 307 | + - Fallback to str(...) for unknown types. |
| 308 | + """ |
| 309 | + if obj is None or isinstance(obj, (str, int, float, bool)): |
| 310 | + if key_hint in ("prompt", "negative_prompt"): |
| 311 | + if isinstance(obj, str): |
| 312 | + return f"<redacted, len={len(obj)}>" |
| 313 | + return obj |
| 314 | + |
| 315 | + if isinstance(obj, Enum): |
| 316 | + return obj.value |
| 317 | + |
| 318 | + is_tensor, torch_mod = _is_torch_tensor(obj) |
| 319 | + if is_tensor: |
| 320 | + try: |
| 321 | + ten = obj.detach().cpu() |
| 322 | + if key_hint == "scaling_factor": |
| 323 | + stats = { |
| 324 | + "shape": list(ten.shape), |
| 325 | + "dtype": str(ten.dtype), |
| 326 | + } |
| 327 | + try: |
| 328 | + stats["min"] = float(ten.min().item()) |
| 329 | + except Exception: |
| 330 | + pass |
| 331 | + try: |
| 332 | + stats["max"] = float(ten.max().item()) |
| 333 | + except Exception: |
| 334 | + pass |
| 335 | + try: |
| 336 | + stats["mean"] = float(ten.float().mean().item()) |
| 337 | + except Exception: |
| 338 | + pass |
| 339 | + return {"tensor": "scaling_factor", **stats} |
| 340 | + return {"tensor": True, "shape": list(ten.shape), "dtype": str(ten.dtype)} |
| 341 | + except Exception: |
| 342 | + return "<tensor>" |
| 343 | + |
| 344 | + if dataclasses.is_dataclass(obj): |
| 345 | + result: dict[str, Any] = {} |
| 346 | + for f in dataclasses.fields(obj): |
| 347 | + if not f.repr: |
| 348 | + continue |
| 349 | + name = f.name |
| 350 | + if "names_mapping" in name: |
| 351 | + continue |
| 352 | + try: |
| 353 | + value = getattr(obj, name) |
| 354 | + except Exception: |
| 355 | + continue |
| 356 | + result[name] = _sanitize_for_logging(value, key_hint=name) |
| 357 | + return result |
| 358 | + |
| 359 | + if isinstance(obj, dict): |
| 360 | + result_dict: dict[str, Any] = {} |
| 361 | + for k, v in obj.items(): |
| 362 | + try: |
| 363 | + key_str = str(k) |
| 364 | + except Exception: |
| 365 | + key_str = "<key>" |
| 366 | + if key_str == "param_names_mapping": |
| 367 | + continue |
| 368 | + result_dict[key_str] = _sanitize_for_logging(v, key_hint=key_str) |
| 369 | + return result_dict |
| 370 | + |
| 371 | + if isinstance(obj, (list, tuple, set)): |
| 372 | + return [_sanitize_for_logging(x, key_hint=key_hint) for x in obj] |
| 373 | + |
| 374 | + try: |
| 375 | + if inspect.isroutine(obj) or inspect.isclass(obj): |
| 376 | + module = getattr(obj, "__module__", "") |
| 377 | + qn = getattr(obj, "__qualname__", getattr(obj, "__name__", "<callable>")) |
| 378 | + return f"{module}.{qn}" if module else qn |
| 379 | + except Exception: |
| 380 | + pass |
| 381 | + |
| 382 | + try: |
| 383 | + return str(obj) |
| 384 | + except Exception: |
| 385 | + return "<unserializable>" |
| 386 | + |
| 387 | + |
284 | 388 | def _trace_calls(log_path, root_dir, frame, event, arg=None): |
285 | 389 | if event in ["call", "return"]: |
286 | 390 | # Extract the filename, line number, function name, and the code object |
@@ -403,6 +507,8 @@ def globally_suppress_loggers(): |
403 | 507 | "python_multipart.multipart", |
404 | 508 | "filelock", |
405 | 509 | "urllib3", |
| 510 | + "httpx", |
| 511 | + "httpcore", |
406 | 512 | ] |
407 | 513 |
|
408 | 514 | for name in target_names: |
@@ -457,7 +563,7 @@ def log_generation_timer( |
457 | 563 | "Processing prompt %d/%d: %s", |
458 | 564 | request_idx, |
459 | 565 | total_requests, |
460 | | - prompt[:100], |
| 566 | + _sanitize_for_logging(prompt, key_hint="prompt"), |
461 | 567 | ) |
462 | 568 |
|
463 | 569 | timer = GenerationTimer() |
|
0 commit comments