84 lines
3.7 KiB
Python
84 lines
3.7 KiB
Python
from typing import Optional, Tuple, Any
|
|
import httpx
|
|
import psutil
|
|
import sentry_sdk
|
|
from loguru import logger
|
|
import functools
|
|
|
|
from BowongModalFunctions.models.web_model import WebhookNotify, WebhookMethodEnum, BaseFFMPEGTaskStatusResponse, \
|
|
TaskStatus, ErrorCode
|
|
|
|
|
|
class SentryUtils:
|
|
@staticmethod
|
|
def capture_hardware_info() -> Tuple[Any, int, Any]:
|
|
cpu_freq = psutil.cpu_freq()
|
|
cpu_count = psutil.cpu_count()
|
|
mem = psutil.virtual_memory()
|
|
|
|
return cpu_freq, cpu_count, mem
|
|
|
|
@staticmethod
|
|
def sentry_tracker(sentry_trace_id: Optional[str], sentry_baggage: Optional[str],
|
|
op: Optional[str], name: Optional[str], fn_id: str):
|
|
def decorator(func):
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
if sentry_trace_id and sentry_baggage:
|
|
transaction = sentry_sdk.continue_trace(environ_or_headers={"sentry-trace": sentry_trace_id,
|
|
"baggage": sentry_baggage, })
|
|
else:
|
|
transaction = sentry_sdk.start_transaction(op='modal.function', name='Modal Function直接调用')
|
|
with transaction:
|
|
with transaction.start_child(op=op, name=name) as span:
|
|
cpu_freq, cpu_count, mem = SentryUtils.capture_hardware_info()
|
|
total_mb = mem.total >> 20
|
|
# available_mb = mem.available >> 20
|
|
sentry_sdk.set_tag('fn.id', fn_id)
|
|
span.set_data('cpu.count', cpu_count)
|
|
span.set_data('cpu.frequency', f"{cpu_freq.current / 1000:.2f} GHz")
|
|
span.set_data('memory.available', f"{total_mb} Mb")
|
|
result = func(*args, **kwargs)
|
|
return result
|
|
|
|
return wrapper
|
|
|
|
return decorator
|
|
|
|
@staticmethod
|
|
def webhook_handler(webhook: WebhookNotify, func_id: str):
|
|
def decorator(func):
|
|
@functools.wraps(func)
|
|
async def async_wrapper(*args, **kwargs):
|
|
try:
|
|
result = await func(*args, **kwargs)
|
|
status = TaskStatus.success
|
|
error = None
|
|
code = ErrorCode.SUCCESS.value
|
|
except Exception as e:
|
|
logger.exception(e)
|
|
result = None
|
|
status = TaskStatus.failed
|
|
error = e.message if hasattr(e, 'message') else str(e)
|
|
code = ErrorCode.SYSTEM_ERROR.value
|
|
logger.info(f"webhook = {webhook}")
|
|
if webhook:
|
|
with httpx.Client() as client:
|
|
match webhook.method:
|
|
case WebhookMethodEnum.POST:
|
|
response = client.post(url=webhook.endpoint.__str__(),
|
|
json=BaseFFMPEGTaskStatusResponse(
|
|
taskId=func_id, status=status, error=error,
|
|
code=code,
|
|
results=[result] if isinstance(result, str) else result,
|
|
).model_dump())
|
|
logger.info(f"webhook response = {response}")
|
|
case WebhookMethodEnum.GET:
|
|
response = client.post(url=webhook.endpoint.__str__())
|
|
response.raise_for_status()
|
|
return result
|
|
|
|
return async_wrapper
|
|
|
|
return decorator
|