mirror of
https://github.com/langgenius/dify.git
synced 2026-11-11 05:41:09 +08:00
113 lines
4.1 KiB
Python
113 lines
4.1 KiB
Python
"""Celery entry points of the trace pipeline, routed by ``OPS_TRACE_QUEUE``.
|
|
|
|
Delivery is at most once per sink: an event is redelivered only while its
|
|
records are not persisted yet, or when the Agent collector could not decide
|
|
ownership (enterprise telemetry is then skipped, it already ran).
|
|
"""
|
|
|
|
import logging
|
|
import re
|
|
from datetime import UTC, date, datetime, timedelta
|
|
from typing import Any
|
|
|
|
from celery import shared_task
|
|
from dify_trace_sdk.events import TraceEvent
|
|
|
|
from extensions.ext_application_services import application_services
|
|
from extensions.ext_storage import storage
|
|
from services.tracing.adapters import PAYLOAD_PREFIX
|
|
from services.tracing.delivery import ExportJob
|
|
from services.tracing.errors import SourceNotFoundError, SourceNotReadyError, TraceCollectorUnavailableError
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
REDELIVERY_LIMIT = 10
|
|
PAYLOAD_RETENTION = timedelta(days=2)
|
|
_PAYLOAD_DAY = re.compile(rf"{PAYLOAD_PREFIX}/(\d{{4}}-\d{{2}}-\d{{2}})/")
|
|
|
|
|
|
def _expired(key: str, cutoff: date) -> bool:
|
|
match = _PAYLOAD_DAY.match(key)
|
|
if match is None:
|
|
return False
|
|
try:
|
|
return date.fromisoformat(match.group(1)) < cutoff
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def _redelivery_delay(retries: int) -> int:
|
|
return min(60, 5 * 2**retries)
|
|
|
|
|
|
def _delete_quietly(key: str) -> None:
|
|
try:
|
|
storage.delete(key)
|
|
except Exception:
|
|
logger.warning("Trace payload cleanup failed: key=%s", key, exc_info=True)
|
|
|
|
|
|
@shared_task(bind=True, max_retries=REDELIVERY_LIMIT)
|
|
def deliver_trace_event(self, envelope: dict[str, Any]) -> None:
|
|
payload_key = envelope.get("payload_key")
|
|
redelivering = False
|
|
try:
|
|
raw = storage.load(payload_key) if isinstance(payload_key, str) else envelope["event"]
|
|
event = TraceEvent.model_validate_json(raw)
|
|
except Exception:
|
|
logger.warning("Trace event unreadable; dropped", exc_info=True)
|
|
if isinstance(payload_key, str):
|
|
_delete_quietly(payload_key)
|
|
return
|
|
try:
|
|
application_services().tracing.delivery.deliver(
|
|
event, enterprise_delivered=envelope.get("enterprise_delivered") is True
|
|
)
|
|
except SourceNotFoundError:
|
|
logger.info("Trace event source not found; dropped: kind=%s app_id=%s", event.payload.kind, event.app_id)
|
|
except SourceNotReadyError:
|
|
if self.request.retries < REDELIVERY_LIMIT:
|
|
redelivering = True
|
|
raise self.retry(args=(envelope,), countdown=_redelivery_delay(self.request.retries)) from None
|
|
logger.warning("Trace event source never became ready; dropped: kind=%s", event.payload.kind)
|
|
except TraceCollectorUnavailableError:
|
|
if self.request.retries < REDELIVERY_LIMIT:
|
|
redelivering = True
|
|
raise self.retry(
|
|
args=({**envelope, "enterprise_delivered": True},),
|
|
countdown=_redelivery_delay(self.request.retries),
|
|
) from None
|
|
logger.warning("Agent trace collector unavailable; dropped: kind=%s", event.payload.kind)
|
|
finally:
|
|
if isinstance(payload_key, str) and not redelivering:
|
|
_delete_quietly(payload_key)
|
|
|
|
|
|
@shared_task
|
|
def retry_trace_export(job_key: str) -> None:
|
|
try:
|
|
job = ExportJob.model_validate_json(storage.load(job_key))
|
|
except Exception:
|
|
logger.warning("Trace export job unreadable; dropped: key=%s", job_key, exc_info=True)
|
|
_delete_quietly(job_key)
|
|
return
|
|
try:
|
|
application_services().tracing.delivery.retry(job)
|
|
finally:
|
|
_delete_quietly(job_key)
|
|
|
|
|
|
@shared_task
|
|
def sweep_trace_payloads() -> None:
|
|
"""Delete expired new-pipeline payloads; legacy workers still own their payloads during an upgrade."""
|
|
cutoff = (datetime.now(UTC) - PAYLOAD_RETENTION).date()
|
|
# Directory listings need the trailing slash on some storage backends.
|
|
prefix = f"{PAYLOAD_PREFIX}/"
|
|
try:
|
|
keys = storage.scan(prefix, files=True)
|
|
except (FileNotFoundError, NotImplementedError):
|
|
return
|
|
for key in keys:
|
|
if key.startswith(prefix) and _expired(key, cutoff):
|
|
_delete_quietly(key)
|