wehub-resource-sync
279 lines
10 KiB
Python
279 lines
10 KiB
Python
import argparse
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
import traceback
|
|
from typing import Optional
|
|
from urllib.parse import quote
|
|
|
|
import aiohttp
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_PROMETHEUS_HOST = "http://localhost:9090"
|
|
PROMETHEUS_HOST_ENV_VAR = "RAY_PROMETHEUS_HOST"
|
|
RETRIES = 3
|
|
|
|
|
|
class PrometheusQueryError(Exception):
|
|
def __init__(self, status, message):
|
|
self.message = (
|
|
"Error fetching data from prometheus. "
|
|
f"status: {status}, message: {message}"
|
|
)
|
|
super().__init__(self.message)
|
|
|
|
|
|
class PrometheusClient:
|
|
def __init__(self) -> None:
|
|
self.http_session = aiohttp.ClientSession()
|
|
self.prometheus_host = os.environ.get(
|
|
PROMETHEUS_HOST_ENV_VAR, DEFAULT_PROMETHEUS_HOST
|
|
)
|
|
|
|
async def query_prometheus(self, query_type, **kwargs):
|
|
url = f"{self.prometheus_host}/api/v1/{query_type}?" + "&".join(
|
|
[f"{k}={quote(str(v), safe='')}" for k, v in kwargs.items()]
|
|
)
|
|
query_str = kwargs.get("query", url)
|
|
logger.debug(f"Running Prometheus query {url}")
|
|
last_error = None
|
|
for attempt in range(RETRIES):
|
|
try:
|
|
async with self.http_session.get(url) as resp:
|
|
if resp.status == 200:
|
|
prom_data = await resp.json()
|
|
return prom_data["data"]["result"]
|
|
body = (await resp.text())[:500]
|
|
last_error = f"non-200 status {resp.status}: {body}"
|
|
logger.warning(
|
|
f"Prometheus query returned non-200 status {resp.status} "
|
|
f"(attempt {attempt + 1}/{RETRIES}). Query: {query_str!r}. "
|
|
f"Body: {body}"
|
|
)
|
|
except asyncio.TimeoutError:
|
|
last_error = "request timed out"
|
|
logger.warning(
|
|
f"Prometheus query timed out "
|
|
f"(attempt {attempt + 1}/{RETRIES}). Query: {query_str!r}."
|
|
)
|
|
except aiohttp.ClientError as e:
|
|
last_error = f"connection error: {e}"
|
|
logger.warning(
|
|
f"Prometheus query connection error "
|
|
f"(attempt {attempt + 1}/{RETRIES}). Query: {query_str!r}. "
|
|
f"Error: {e}"
|
|
)
|
|
if attempt < RETRIES - 1:
|
|
await asyncio.sleep(1)
|
|
logger.error(
|
|
f"Prometheus query failed after {RETRIES} attempts and returned no data. "
|
|
f"Query: {query_str!r}. Last error: {last_error}. "
|
|
"This is a metrics-collection failure (Prometheus unreachable/erroring), "
|
|
"NOT an empty result for a healthy metric."
|
|
)
|
|
return None
|
|
|
|
async def close(self):
|
|
await self.http_session.close()
|
|
|
|
|
|
# Metrics here mirror what we have in Grafana.
|
|
async def _get_prometheus_metrics(
|
|
start_time: float, end_time: float, session_name: Optional[str] = None
|
|
) -> dict:
|
|
client = PrometheusClient()
|
|
kwargs = {
|
|
"query_type": "query_range",
|
|
"start": int(start_time),
|
|
"end": int(end_time),
|
|
"step": 15,
|
|
}
|
|
sf = f'{{SessionName="{session_name}"}}' if session_name else ""
|
|
sf_spilled = (
|
|
f'{{SessionName="{session_name}",State="Spilled"}}'
|
|
if session_name
|
|
else '{State="Spilled"}'
|
|
)
|
|
metrics = {
|
|
"cpu_utilization": client.query_prometheus(
|
|
query=f"ray_node_cpu_utilization{sf} * ray_node_cpu_count{sf} / 100",
|
|
**kwargs,
|
|
),
|
|
"cpu_count": client.query_prometheus(query=f"ray_node_cpu_count{sf}", **kwargs),
|
|
"gpu_utilization": client.query_prometheus(
|
|
query=f"ray_node_gpus_utilization{sf} / 100", **kwargs
|
|
),
|
|
"gpu_count": client.query_prometheus(
|
|
query=f"ray_node_gpus_available{sf}", **kwargs
|
|
),
|
|
"disk_usage": client.query_prometheus(
|
|
query=f"ray_node_disk_usage{sf}", **kwargs
|
|
),
|
|
"disk_space": client.query_prometheus(
|
|
query=f"sum(ray_node_disk_free{sf}) + sum(ray_node_disk_usage{sf})",
|
|
**kwargs,
|
|
),
|
|
"memory_usage": client.query_prometheus(
|
|
query=f"ray_node_mem_used{sf}", **kwargs
|
|
),
|
|
"total_memory": client.query_prometheus(
|
|
query=f"ray_node_mem_total{sf}", **kwargs
|
|
),
|
|
"memory_usage_host": client.query_prometheus(
|
|
query=f"ray_node_mem_used_host{sf}", **kwargs
|
|
),
|
|
"total_memory_host": client.query_prometheus(
|
|
query=f"ray_node_mem_total_host{sf}", **kwargs
|
|
),
|
|
"memory_usage_cgroup": client.query_prometheus(
|
|
query=f"ray_node_cgroup_mem_used{sf}", **kwargs
|
|
),
|
|
"total_memory_cgroup": client.query_prometheus(
|
|
query=f"ray_node_cgroup_mem_total{sf}", **kwargs
|
|
),
|
|
"gpu_memory_usage": client.query_prometheus(
|
|
query=f"ray_node_gram_used{sf} * 1024 * 1024", **kwargs
|
|
),
|
|
"gpu_total_memory": client.query_prometheus(
|
|
query=(
|
|
f"(sum(ray_node_gram_available{sf}) + sum(ray_node_gram_used{sf}))"
|
|
" * 1024 * 1024"
|
|
),
|
|
**kwargs,
|
|
),
|
|
"network_receive_speed": client.query_prometheus(
|
|
query=f"ray_node_network_receive_speed{sf}", **kwargs
|
|
),
|
|
"network_send_speed": client.query_prometheus(
|
|
query=f"ray_node_network_send_speed{sf}", **kwargs
|
|
),
|
|
"cluster_active_nodes": client.query_prometheus(
|
|
query=f"ray_cluster_active_nodes{sf}", **kwargs
|
|
),
|
|
"cluster_failed_nodes": client.query_prometheus(
|
|
query=f"ray_cluster_failed_nodes{sf}", **kwargs
|
|
),
|
|
"cluster_pending_nodes": client.query_prometheus(
|
|
query=f"ray_cluster_pending_nodes{sf}", **kwargs
|
|
),
|
|
"worker_oom_kills": client.query_prometheus(
|
|
query=(
|
|
f"sum(ray_memory_manager_worker_eviction_total{sf}) by (Type, Name)"
|
|
),
|
|
**kwargs,
|
|
),
|
|
"unexpected_worker_failures": client.query_prometheus(
|
|
query=f"sum(ray_node_manager_unexpected_worker_failure_total{sf}) by (Type, Name)",
|
|
**kwargs,
|
|
),
|
|
# `State="Spilled"` is the cumulative-bytes counter (the other States
|
|
# are point-in-time / transient); `> 0` drops always-emitted 0 points.
|
|
"spilled_bytes": client.query_prometheus(
|
|
query=f"sum(ray_spill_manager_objects_bytes{sf_spilled}) > 0",
|
|
**kwargs,
|
|
),
|
|
}
|
|
metrics = {k: await v for k, v in metrics.items()}
|
|
await client.close()
|
|
|
|
# Summarise the outcome so a glance at the logs tells whether the metrics
|
|
# are trustworthy. `None` => the query failed to collect (infra/timeout);
|
|
# `[]` => collected fine but no matching series (e.g. no OOMs happened);
|
|
# truthy => collected data.
|
|
failed = sorted(k for k, v in metrics.items() if v is None)
|
|
empty = sorted(k for k, v in metrics.items() if v == [])
|
|
with_data = sorted(k for k, v in metrics.items() if v)
|
|
logger.info(
|
|
f"Prometheus collection summary: {len(with_data)} metric(s) with data, "
|
|
f"{len(empty)} empty, {len(failed)} failed to collect "
|
|
f"(out of {len(metrics)} total)."
|
|
)
|
|
if failed:
|
|
logger.error(
|
|
f"{len(failed)} metric(s) FAILED to collect and will be null in the "
|
|
f"output: {failed}. See the per-query warnings above for the cause "
|
|
"(timeout / connection error / non-200). This indicates a "
|
|
"metrics-collection/infra problem, not a real metric signal."
|
|
)
|
|
return metrics
|
|
|
|
|
|
def get_prometheus_metrics(start_time: float, end_time: float) -> dict:
|
|
session_name = None
|
|
try:
|
|
import ray
|
|
|
|
if not ray.is_initialized():
|
|
ray.init("auto")
|
|
session_name = ray.get_runtime_context().get_session_name()
|
|
except Exception:
|
|
logger.warning(
|
|
"Couldn't obtain Ray session name for Prometheus query filtering. "
|
|
f"Exception below:\n{traceback.format_exc()}"
|
|
)
|
|
try:
|
|
return asyncio.run(_get_prometheus_metrics(start_time, end_time, session_name))
|
|
except Exception:
|
|
logger.error(
|
|
"Couldn't obtain Prometheus metrics. "
|
|
f"Exception below:\n{traceback.format_exc()}"
|
|
)
|
|
return {}
|
|
|
|
|
|
def save_prometheus_metrics(
|
|
start_time: float,
|
|
end_time: Optional[float] = None,
|
|
path: Optional[str] = None,
|
|
use_ray: bool = False,
|
|
) -> bool:
|
|
path = path or os.environ.get("METRICS_OUTPUT_JSON", None)
|
|
if path:
|
|
if not end_time:
|
|
end_time = time.time()
|
|
if use_ray:
|
|
import ray
|
|
from ray.air.util.node import _force_on_current_node
|
|
|
|
addr = os.environ.get("RAY_ADDRESS", None)
|
|
ray.init(addr)
|
|
|
|
@ray.remote(num_cpus=0)
|
|
def get_metrics():
|
|
end_time = time.time()
|
|
return get_prometheus_metrics(start_time, end_time)
|
|
|
|
remote_run = _force_on_current_node(get_metrics)
|
|
ref = remote_run.remote()
|
|
metrics = ray.get(ref, timeout=900)
|
|
else:
|
|
metrics = get_prometheus_metrics(start_time, end_time)
|
|
with open(path, "w") as metrics_output_file:
|
|
json.dump(metrics, metrics_output_file)
|
|
return path
|
|
return None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="[%(levelname)s %(asctime)s] %(filename)s: %(lineno)d %(message)s",
|
|
)
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("start_time", type=float, help="Start time")
|
|
parser.add_argument(
|
|
"--path", default="", type=str, help="Where to save the metrics json"
|
|
)
|
|
parser.add_argument(
|
|
"--use_ray",
|
|
default=False,
|
|
action="store_true",
|
|
help="Whether to run this script in a ray.remote call (for Ray Client)",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
save_prometheus_metrics(args.start_time, path=args.path, use_ray=args.use_ray)
|