Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
28 commits
Select commit Hold shift + click to select a range
752ab1a
Add python profiling utilities
Eden-D-Zhang Oct 14, 2025
29436b6
Add profiling to query functions
Eden-D-Zhang Oct 14, 2025
4d19c81
Remove `CLP_ENABLE_PROFILING` environment variable check
Eden-D-Zhang Oct 14, 2025
b5bdcf1
Lint
Eden-D-Zhang Oct 14, 2025
02b29fb
Merge branch 'main' of https://github.com/y-scope/clp into pyinstrume…
Eden-D-Zhang Oct 14, 2025
6f65035
Clean up docstrings
Eden-D-Zhang Oct 14, 2025
c2be1f6
Remove __all__
Eden-D-Zhang Oct 15, 2025
a677177
Merge branch 'main' of https://github.com/Eden-D-Zhang/clp into pyins…
Eden-D-Zhang Oct 15, 2025
cc85f93
Remove unnecessary function
Eden-D-Zhang Oct 15, 2025
9810d31
Address review
Eden-D-Zhang Oct 16, 2025
b013429
Merge branch 'main' into pyinstrument_profile
Eden-D-Zhang Oct 16, 2025
a33a9c1
Address review
Eden-D-Zhang Oct 16, 2025
99d134b
Lint
Eden-D-Zhang Oct 16, 2025
61e8332
Merge branch 'main' of https://github.com/y-scope/clp into pyinstrume…
Eden-D-Zhang Oct 21, 2025
ab6ee4f
Upate lock file
Eden-D-Zhang Oct 21, 2025
cbb32de
Update dependencies and lock file
Eden-D-Zhang Oct 21, 2025
8447f12
Merge branch 'main' of https://github.com/y-scope/clp into pyinstrume…
Eden-D-Zhang Oct 21, 2025
159e645
Address review
Eden-D-Zhang Oct 21, 2025
4080df6
Delete file
Eden-D-Zhang Oct 21, 2025
7df8e09
Merge branch 'main' into pyinstrument_profile
Eden-D-Zhang Oct 21, 2025
4089eb0
Change constant name
Eden-D-Zhang Oct 21, 2025
11641e2
Merge branch 'pyinstrument_profile' of https://github.com/Eden-D-Zhan…
Eden-D-Zhang Oct 21, 2025
9437246
Change typing imports according to coderabbit
Eden-D-Zhang Oct 21, 2025
ada4073
Lint
Eden-D-Zhang Oct 21, 2025
e94d2cc
Add clp_logging logger, type annotations
Eden-D-Zhang Oct 23, 2025
8beeda2
Merge branch 'main' of https://github.com/y-scope/clp into pyinstrume…
Eden-D-Zhang Oct 23, 2025
e6dfbea
Lint
Eden-D-Zhang Oct 23, 2025
022a7ab
Merge branch 'main' into pyinstrument_profile
Eden-D-Zhang Oct 23, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
258 changes: 258 additions & 0 deletions components/clp-py-utils/clp_py_utils/profiling_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,258 @@
"""
Profiling utilities for CLP query execution performance analysis.

This module provides lightweight profiling decorators using pyinstrument.

Profile outputs include:
- HTML files with interactive flame graphs and call trees.
- Text summaries showing call hierarchy and timing.
"""

import datetime
import functools
import inspect
import os
from collections.abc import Callable
from pathlib import Path
from typing import Any, TypeVar

from pyinstrument import Profiler

from clp_py_utils.clp_logging import get_logger

logger = get_logger("profiler")

F = TypeVar("F", bound=Callable[..., Any])

PROFILING_INTERVAL_SECONDS = 0.001


def profile(
section_name: str | None = None,
job_id_param: str = "job_id",
task_id_param: str = "task_id",
) -> Callable[[F], F]:
"""
Profiles function execution as decorator with automatic context extraction.

Output files are written to $CLP_LOGS_DIR/profiles/ (e.g., clp-package/var/log/query_worker/
profiles/).

:param section_name: Override for profile section name. If None, uses function name.
:param job_id_param: Parameter name to extract job_id from (default: "job_id").
Can use dot notation for attributes, e.g., "job.id".
:param task_id_param: Parameter name to extract task_id from (default: "task_id").
Can use dot notation for attributes, e.g., "task.id".
:return: Decorated function with profiling capabilities.
"""

def decorator(func: F) -> F:
name = section_name or func.__name__
is_async = inspect.iscoroutinefunction(func)

if is_async:

@functools.wraps(func)
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
if not _is_profiling_enabled():
return await func(*args, **kwargs)

# Profiling enabled: extract context and profile execution
job_id, task_id = _extract_context_from_args(
func, args, kwargs, job_id_param, task_id_param
)

profiler = Profiler(interval=PROFILING_INTERVAL_SECONDS)
try:
profiler.start()
except RuntimeError as e:
# Skip profiling this function to avoid conflicts
if "already a profiler running" in str(e):
logger.debug(
f"Skipping nested profiling for {name} "
f"(parent profiler already active)"
)
return await func(*args, **kwargs)
raise

try:
result = await func(*args, **kwargs)
return result
finally:
profiler.stop()
_save_profile(profiler, name, job_id, task_id)

return async_wrapper # type: ignore

@functools.wraps(func)
def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
if not _is_profiling_enabled():
return func(*args, **kwargs)

# Profiling enabled: extract context and profile execution
job_id, task_id = _extract_context_from_args(
func, args, kwargs, job_id_param, task_id_param
)

profiler = Profiler(interval=PROFILING_INTERVAL_SECONDS)
try:
profiler.start()
except RuntimeError as e:
# Skip profiling this function to avoid conflicts
if "already a profiler running" in str(e):
logger.debug(
f"Skipping nested profiling for {name} (parent profiler already active)"
)
return func(*args, **kwargs)
raise

try:
result = func(*args, **kwargs)
return result
finally:
profiler.stop()
_save_profile(profiler, name, job_id, task_id)

return sync_wrapper # type: ignore

return decorator


def _extract_context_from_args(
func: Callable,
args: tuple,
kwargs: dict,
job_id_param: str = "job_id",
task_id_param: str = "task_id",
) -> tuple[str, str]:
"""
Extracts job_id and task_id from function arguments.

:param func: The function being profiled.
:param args: Positional arguments passed to the function.
:param kwargs: Keyword arguments passed to the function.
:param job_id_param: Name/path of the parameter containing job_id (default: "job_id").
:param task_id_param: Name/path of the parameter containing task_id (default: "task_id").
:return: tuple of (job_id, task_id) as strings. Empty strings if not found.
"""
job_id = ""
task_id = ""

try:
# Get function signature
sig = inspect.signature(func)
param_names = list(sig.parameters.keys())

def extract_value(param_spec: str) -> str:
"""Extract value from parameter, supporting dot notation for attributes."""
if not param_spec:
return ""

# Split on '.' to handle attribute access
parts = param_spec.split(".")
param_name = parts[0]

# Find the parameter value
value = None
if param_name in kwargs:
value = kwargs[param_name]
elif param_name in param_names:
idx = param_names.index(param_name)
if idx < len(args):
value = args[idx]

if value is None:
return ""

# Navigate through attributes if dot notation was used
for attr_name in parts[1:]:
if hasattr(value, attr_name):
value = getattr(value, attr_name)
else:
return ""

return str(value)

# Extract job_id and task_id
job_id = extract_value(job_id_param)
task_id = extract_value(task_id_param)

except Exception as e:
logger.debug(f"Failed to extract context from {func.__name__}: {e}")

return job_id, task_id


def _is_profiling_enabled() -> bool:
"""
Checks if profiling is enabled.
TODO: Add `CLPConfig` mechanism to enable/disable profiling for each component.

:return: Whether the profiler is enabled.
"""
return False


def _save_profile(
profiler: Profiler, section_name: str, job_id: str = "", task_id: str = ""
) -> None:
"""
Saves profiler output to HTML and text formats. Generates .html and .txt files.

:param profiler: The pyinstrument Profiler object containing profiling data.
:param section_name: Name identifying this profiling section.
:param job_id: Optional job identifier for filename.
:param task_id: Optional task identifier for filename.
"""
try:
# Get the session for logging
session = profiler.last_session
if not session:
logger.debug(f"No profiling session for {section_name}")
return

duration = session.duration
sample_count = session.sample_count

timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
filename_parts = [section_name]

if job_id:
filename_parts.append(f"job{job_id}")
if task_id:
filename_parts.append(f"task{task_id}")

filename_parts.append(timestamp)
base_filename = "_".join(filename_parts)

output_dir = Path(os.getenv("CLP_LOGS_DIR")) / "profiles"
output_dir.mkdir(exist_ok=True, parents=True)

# Save HTML with interactive visualization
html_path = output_dir / f"{base_filename}.html"
with open(html_path, "w", encoding="utf-8") as f:
f.write(profiler.output_html())

# Save human-readable text summary with call hierarchy
txt_path = output_dir / f"{base_filename}.txt"
with open(txt_path, "w", encoding="utf-8") as f:
# Header
f.write("=" * 80 + "\n")
f.write(f"CLP Query Profiling Report (pyinstrument)\n")
f.write(f"Section: {section_name}\n")
if job_id:
f.write(f"Job ID: {job_id}\n")
if task_id:
f.write(f"Task ID: {task_id}\n")
f.write(f"Timestamp: {timestamp}\n")
f.write("=" * 80 + "\n\n")
f.write(profiler.output_text(unicode=True, color=False))

logger.info(
f"Profile saved: {section_name} "
f"(duration={duration:.6f}s, samples={sample_count}) "
f"HTML={html_path}, TXT={txt_path}"
)

except Exception as e:
logger.error(f"Failed to save profile for {section_name}: {e}", exc_info=True)
1 change: 1 addition & 0 deletions components/clp-py-utils/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ dependencies = [
"mariadb>=1.0.11,<1.1.dev0",
"mysql-connector-python>=9.4.0",
"pydantic>=2.12.3",
"pyinstrument>=5.1.1",
"python-Levenshtein>=0.27.1",
"PyYAML>=6.0.3",
"result>=0.17.0",
Expand Down
Loading
Loading