120 lines
5.8 KiB
Python
120 lines
5.8 KiB
Python
import os
|
|
from typing import Optional, Tuple, Any, List
|
|
import httpx
|
|
import psutil
|
|
import sentry_sdk
|
|
from loguru import logger
|
|
import functools
|
|
|
|
from BowongModalFunctions.models.web_model import WebhookNotify, WebhookMethodEnum, BaseFFMPEGTaskStatusResponse, \
|
|
TaskStatus, ErrorCode, FFMPEGResult
|
|
|
|
|
|
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):
|
|
logger.info(f"sentry-trace={sentry_trace_id}, baggage={sentry_baggage}")
|
|
if sentry_trace_id and sentry_baggage:
|
|
transaction = sentry_sdk.continue_trace(environ_or_headers={"sentry-trace": sentry_trace_id,
|
|
"baggage": sentry_baggage, },
|
|
op=op,
|
|
name=name)
|
|
else:
|
|
transaction = sentry_sdk.start_transaction(op=op, name=name)
|
|
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_tags({
|
|
'fn.id': fn_id,
|
|
'x_trace.id': sentry_trace_id,
|
|
'x_trace.baggage': sentry_baggage,
|
|
'cpu.count': cpu_count,
|
|
'cpu.frequency': f"{cpu_freq.current / 1000:.2f}",
|
|
'cpu.frequency.format': f"{cpu_freq.current / 1000:.2f} GHz",
|
|
'memory.available': total_mb,
|
|
'memory.available.format': f"{total_mb} Mb",
|
|
'modal.cloud.region': os.environ.get('MODAL_REGION', 'unknown'),
|
|
'modal.cloud.provider': os.environ.get('MODAL_CLOUD_PROVIDER', 'unknown'),
|
|
'modal.task.id': os.environ.get('MODAL_TASK_ID', 'unknown'),
|
|
'modal.identity.token': os.environ.get('MODAL_IDENTITY_TOKEN', 'unknown'),
|
|
'modal.image.id': os.environ.get('MODAL_IMAGE_ID', 'unknown'),
|
|
})
|
|
result = func(*args, **kwargs)
|
|
span.set_status('ok')
|
|
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
|
|
if not result:
|
|
status = TaskStatus.failed
|
|
error = "Internal Error"
|
|
code = ErrorCode.BUSINESS_ERROR.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:
|
|
logger.info(f"result = {result}")
|
|
if isinstance(result, str):
|
|
results = [result]
|
|
elif isinstance(result, FFMPEGResult):
|
|
results = [result]
|
|
elif isinstance(result, List):
|
|
results = result
|
|
else:
|
|
results = [result]
|
|
with httpx.Client() as client:
|
|
match webhook.method:
|
|
case WebhookMethodEnum.POST:
|
|
response = client.post(url=webhook.endpoint.__str__(),
|
|
headers=webhook.headers,
|
|
json=BaseFFMPEGTaskStatusResponse(
|
|
taskId=func_id, status=status,
|
|
error=error, code=code,
|
|
results=results,
|
|
).model_dump())
|
|
case WebhookMethodEnum.GET:
|
|
response = client.get(url=webhook.endpoint.__str__(),
|
|
headers=webhook.headers,
|
|
params=BaseFFMPEGTaskStatusResponse(
|
|
taskId=func_id, status=status, error=error,
|
|
code=code,
|
|
results=results,
|
|
).model_dump())
|
|
logger.info(f"webhook response = {response}")
|
|
response.raise_for_status()
|
|
return result
|
|
|
|
return async_wrapper
|
|
|
|
return decorator
|