wehub-resource-sync
98 lines
3.0 KiB
Python
98 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Payload construction and HTTP transport for usage telemetry."""
|
|
|
|
# Future
|
|
from __future__ import annotations
|
|
|
|
# Standard
|
|
from urllib.parse import urljoin
|
|
import json
|
|
import os
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.connections import global_http_connection
|
|
from lmcache.logging import init_logger
|
|
from lmcache.usage_telemetry.identity import UsageIdentity
|
|
from lmcache.usage_telemetry.messages import (
|
|
USAGE_SCHEMA_VERSION,
|
|
DeploymentMode,
|
|
UsageMessage,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def usage_server_url(endpoint: str) -> str:
|
|
"""Build the stats-server URL for *endpoint*.
|
|
|
|
Args:
|
|
endpoint: Path component on the stats server, e.g. ``"context"``.
|
|
|
|
Returns:
|
|
The full URL, honoring the ``LMCACHE_USAGE_TRACK_URL`` override.
|
|
"""
|
|
base = os.getenv("LMCACHE_USAGE_TRACK_URL", "http://stats.lmcache.ai:8080")
|
|
return urljoin(base, endpoint)
|
|
|
|
|
|
class UsageMessageSender:
|
|
"""Default HTTP transport for usage messages.
|
|
|
|
Posts JSON payloads with a short timeout and swallows every failure
|
|
(telemetry must never disturb serving). Inject a stub instance into the
|
|
usage contexts to capture payloads in tests.
|
|
"""
|
|
|
|
def send(self, url: str, payload: dict[str, object]) -> None:
|
|
"""POST one JSON payload to *url*; never raises.
|
|
|
|
Args:
|
|
url: Full endpoint URL on the stats server.
|
|
payload: Flat JSON-serializable payload.
|
|
"""
|
|
try:
|
|
client = global_http_connection.get_sync_client()
|
|
client.post(
|
|
url,
|
|
data=json.dumps(payload),
|
|
headers={"Content-Type": "application/json"},
|
|
timeout=5,
|
|
)
|
|
except Exception as e:
|
|
logger.debug("Unable to send lmcache usage message: %s", e)
|
|
|
|
|
|
DEFAULT_SENDER = UsageMessageSender()
|
|
"""Shared default sender; stateless, safe to reuse across contexts."""
|
|
|
|
|
|
def build_usage_payload(
|
|
message: UsageMessage, identity: UsageIdentity, mode: DeploymentMode
|
|
) -> dict[str, object]:
|
|
"""Flatten *message* into a payload stamped with identity and schema info.
|
|
|
|
The wire ``message_type`` is the message's class name, per the schema
|
|
contract in :mod:`lmcache.usage_telemetry.messages`.
|
|
|
|
Args:
|
|
message: The message to serialize; its fields become payload keys.
|
|
identity: Identifiers stamped on the payload.
|
|
mode: Deployment mode stamped as the ``deployment_mode`` key.
|
|
|
|
Returns:
|
|
A flat JSON-serializable dict (``torch.dtype`` values stringified).
|
|
"""
|
|
payload: dict[str, object] = {
|
|
"message_type": type(message).__name__,
|
|
"schema_version": USAGE_SCHEMA_VERSION,
|
|
"session_id": identity.session_id,
|
|
"machine_id": identity.machine_id,
|
|
"deployment_mode": mode.value,
|
|
}
|
|
for key, value in vars(message).items():
|
|
payload[key] = str(value) if isinstance(value, torch.dtype) else value
|
|
return payload
|