Spaces:
Running
Running
File size: 4,332 Bytes
a55ea57 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import time
from typing import Any
from celery import signals
from src.lib.tracing.tracker import tracer
_task_starts: dict[str, float] = {}
_task_context_tokens: dict[str, Any] = {}
def _task_id(task_id: str | None, task=None) -> str:
if task_id:
return task_id
request = getattr(task, "request", None)
return getattr(request, "id", None) or "unknown"
def _task_name(task=None, sender=None) -> str:
return getattr(task, "name", None) or getattr(sender, "name", None) or "unknown"
def _tenant_id(args=None, kwargs=None) -> str | None:
args = args or ()
kwargs = kwargs or {}
candidates = list(args)
candidates.extend(kwargs.values())
for candidate in candidates:
if isinstance(candidate, dict):
brain_id = candidate.get("brain_id")
if brain_id:
return str(brain_id)
brain_id = kwargs.get("brain_id")
return str(brain_id) if brain_id else None
def _trace_id(task_id: str, kwargs=None) -> str:
kwargs = kwargs or {}
trace_id = kwargs.get("trace_id")
return str(trace_id) if trace_id else task_id
def _metadata(task_id: str, task_name: str, args=None, kwargs=None) -> dict[str, Any]:
return {
"task_id": task_id,
"task_name": task_name,
"args_count": len(args or ()),
"kwargs_keys": sorted((kwargs or {}).keys()),
}
def install_celery_tracing(_app=None, service_name: str = "brainapi-worker") -> None:
signals.task_prerun.connect(
_trace_task_prerun,
weak=False,
dispatch_uid=f"{service_name}.trace.task_prerun",
)
signals.task_postrun.connect(
_trace_task_postrun,
weak=False,
dispatch_uid=f"{service_name}.trace.task_postrun",
)
signals.task_failure.connect(
_trace_task_failure,
weak=False,
dispatch_uid=f"{service_name}.trace.task_failure",
)
def _trace_task_prerun(sender=None, task_id=None, task=None, args=None, kwargs=None, **_):
resolved_task_id = _task_id(task_id, task)
resolved_task_name = _task_name(task, sender)
_task_starts[resolved_task_id] = time.perf_counter()
_task_context_tokens[resolved_task_id] = tracer.set_context(
trace_id=_trace_id(resolved_task_id, kwargs),
tenant_id=_tenant_id(args, kwargs),
)
tracer.publish(
"latency",
f"{resolved_task_name}.started",
service="brainapi-worker",
operation=resolved_task_name,
metadata=_metadata(resolved_task_id, resolved_task_name, args, kwargs),
)
def _trace_task_postrun(sender=None, task_id=None, task=None, args=None, kwargs=None, retval=None, state=None, **_):
resolved_task_id = _task_id(task_id, task)
resolved_task_name = _task_name(task, sender)
started_at = _task_starts.pop(resolved_task_id, None)
duration_ms = (time.perf_counter() - started_at) * 1000 if started_at else None
metadata = {
**_metadata(resolved_task_id, resolved_task_name, args, kwargs),
"state": state,
"returned": retval is not None,
}
if duration_ms is not None and duration_ms >= tracer.default_sla_ms:
tracer.publish(
"sla_breach",
f"{resolved_task_name}.sla_breach",
service="brainapi-worker",
operation=resolved_task_name,
duration_ms=duration_ms,
threshold_ms=tracer.default_sla_ms,
metadata=metadata,
)
token = _task_context_tokens.pop(resolved_task_id, None)
if token is not None:
tracer.reset_context(token)
def _trace_task_failure(sender=None, task_id=None, exception=None, args=None, kwargs=None, einfo=None, **_):
resolved_task_id = _task_id(task_id, sender)
resolved_task_name = _task_name(sender=sender)
started_at = _task_starts.get(resolved_task_id)
duration_ms = (time.perf_counter() - started_at) * 1000 if started_at else None
metadata = _metadata(resolved_task_id, resolved_task_name, args, kwargs)
if einfo is not None:
metadata["traceback"] = str(einfo)
tracer.exception(
f"{resolved_task_name}.failed",
exception or RuntimeError("Celery task failed"),
service="brainapi-worker",
operation=resolved_task_name,
duration_ms=duration_ms,
metadata=metadata,
)
|