|
| 1 | +"""Prometheus metrics for Bastion runners.""" |
| 2 | + |
| 3 | +import logging |
| 4 | +import threading |
| 5 | +from http.server import BaseHTTPRequestHandler, HTTPServer |
| 6 | + |
| 7 | +from prometheus_client import CollectorRegistry, Gauge, Histogram, generate_latest, multiprocess |
| 8 | + |
| 9 | +BASTION_DOWNLOAD_JOB_ALL_SECONDS = Histogram( |
| 10 | + "bastion_download_job_all_seconds", |
| 11 | + "Time in seconds consumed in each scheduling loop to download all the job specs", |
| 12 | + [], |
| 13 | + buckets=(1, 5, 10, 30, 60, 120, 300, 600, 1800, 3600), # 1s to 1hr |
| 14 | +) |
| 15 | + |
| 16 | +JOB_TIME_TO_RUNNING_SECONDS = Histogram( |
| 17 | + "bastion_job_time_to_running_seconds", |
| 18 | + "Time in seconds from runner starts a job until until it reaches Running state", |
| 19 | + ["job_name"], |
| 20 | + buckets=(1, 5, 10, 30, 60, 120, 300, 600, 1800, 3600), # 1s to 1hr |
| 21 | +) |
| 22 | + |
| 23 | +JOB_RUN_LATENCY_GAUGE = Gauge( |
| 24 | + "bastion_job_run_latency_seconds", |
| 25 | + "Time in seconds from runner starts a job until until it reaches Running state", |
| 26 | + ["job_name"], |
| 27 | +) |
| 28 | + |
| 29 | +JOB_WAIT_FOR_BUILD_SECONDS = Gauge( |
| 30 | + "bastion_job_wait_for_build_seconds", |
| 31 | + "time in seconds until the job is built successfully", |
| 32 | + ["job_name"], |
| 33 | +) |
| 34 | + |
| 35 | +JOB_SUSPENDED_DURATION_SECONDS = Histogram( |
| 36 | + "bastion_job_suspended_duration_seconds", |
| 37 | + "Time in seconds the jobset stayed in SUSPENDED state", |
| 38 | + ["job_name"], |
| 39 | + buckets=(1, 5, 10, 30, 60, 120, 300, 600, 1800, 3600), # 1s to 1hr |
| 40 | +) |
| 41 | + |
| 42 | + |
| 43 | +class _MetricsHandler(BaseHTTPRequestHandler): |
| 44 | + """HTTP handler that serves aggregated metrics from multiprocess mode.""" |
| 45 | + |
| 46 | + def do_GET(self): # pylint: disable=invalid-name |
| 47 | + if self.path == "/metrics": |
| 48 | + registry = CollectorRegistry() |
| 49 | + multiprocess.MultiProcessCollector(registry) |
| 50 | + metrics = generate_latest(registry) |
| 51 | + |
| 52 | + self.send_response(200) |
| 53 | + self.send_header("Content-Type", "text/plain; charset=utf-8") |
| 54 | + self.end_headers() |
| 55 | + self.wfile.write(metrics) |
| 56 | + else: |
| 57 | + self.send_error(404) |
| 58 | + |
| 59 | + def log_message(self, format, *args): # pylint: disable=redefined-builtin |
| 60 | + # Suppress default logging |
| 61 | + pass |
| 62 | + |
| 63 | + |
| 64 | +class MetricsServer: |
| 65 | + """Singleton class to manage Prometheus metrics server. |
| 66 | +
|
| 67 | + Supports multiprocess mode where metrics from parent and child processes |
| 68 | + are aggregated and exposed via a single endpoint. |
| 69 | + """ |
| 70 | + |
| 71 | + _instance = None |
| 72 | + _lock = threading.Lock() |
| 73 | + _server_started = False |
| 74 | + _multiprocess_dir = None |
| 75 | + |
| 76 | + def __new__(cls): |
| 77 | + if cls._instance is None: |
| 78 | + with cls._lock: |
| 79 | + if cls._instance is None: |
| 80 | + cls._instance = super().__new__(cls) |
| 81 | + return cls._instance |
| 82 | + |
| 83 | + @classmethod |
| 84 | + def start(cls, port: int = 8000): |
| 85 | + """Start Prometheus metrics HTTP server on the specified port. |
| 86 | +
|
| 87 | + Args: |
| 88 | + port: Port number to start the metrics server on. Defaults to 8000. |
| 89 | + use_multiprocess: If True, use multiprocess mode which aggregates metrics |
| 90 | + from child processes. Defaults to False for backward compatibility. |
| 91 | + """ |
| 92 | + with cls._lock: |
| 93 | + if not cls._server_started: |
| 94 | + try: |
| 95 | + # Use custom handler for multiprocess aggregation |
| 96 | + server = HTTPServer(("", port), _MetricsHandler) |
| 97 | + thread = threading.Thread(target=server.serve_forever, daemon=True) |
| 98 | + thread.start() |
| 99 | + logging.info( |
| 100 | + "Prometheus metrics server (multiprocess mode) started on port %s", port |
| 101 | + ) |
| 102 | + cls._server_started = True |
| 103 | + except OSError as e: |
| 104 | + # Port might already be in use |
| 105 | + logging.warning("Failed to start metrics server on port %s: %s", port, e) |
| 106 | + |
| 107 | + @classmethod |
| 108 | + def is_started(cls) -> bool: |
| 109 | + """Check if the metrics server has been started.""" |
| 110 | + return cls._server_started |
| 111 | + |
| 112 | + @classmethod |
| 113 | + def get_multiprocess_dir(cls) -> str: |
| 114 | + """Get the multiprocess directory path if configured.""" |
| 115 | + return cls._multiprocess_dir |
| 116 | + |
| 117 | + |
| 118 | +def record_job_time_to_running(job_name: str, time_seconds: float): |
| 119 | + """Record the time it took for a job to reach RUNNING state. |
| 120 | +
|
| 121 | + Args: |
| 122 | + job_name: Name of the job |
| 123 | + time_seconds: Time in seconds from first seen until RUNNING state |
| 124 | + """ |
| 125 | + JOB_TIME_TO_RUNNING_SECONDS.labels(job_name=job_name).observe(time_seconds) |
| 126 | + logging.info("Metrics: job %s took %.2f seconds to reach RUNNING state", job_name, time_seconds) |
| 127 | + |
| 128 | + |
| 129 | +def record_job_run_latency(job_name: str, time_seconds: float): |
| 130 | + JOB_RUN_LATENCY_GAUGE.labels(job_name).set(time_seconds) |
| 131 | + |
| 132 | + |
| 133 | +def record_job_wait_for_build(job_name: str, time_seconds: float): |
| 134 | + """Record the time it took to wait for the cloud build to finish |
| 135 | +
|
| 136 | + This does not necessarily mean the whole cloud build time for a job, |
| 137 | + as the build starts when the client submits the job. This only measures |
| 138 | + how long the runner waits for the build to complete before launching |
| 139 | + the job. |
| 140 | + """ |
| 141 | + JOB_WAIT_FOR_BUILD_SECONDS.labels(job_name).set(time_seconds) |
| 142 | + |
| 143 | + |
| 144 | +def record_job_suspended_duration(job_name: str, time_seconds: float): |
| 145 | + """Record the time a job spent in SUSPENDED state. |
| 146 | +
|
| 147 | + Args: |
| 148 | + job_name: Name of the job |
| 149 | + time_seconds: Time in seconds the job was in SUSPENDED state |
| 150 | + """ |
| 151 | + JOB_SUSPENDED_DURATION_SECONDS.labels(job_name=job_name).observe(time_seconds) |
| 152 | + logging.info("Metrics: job %s was in SUSPENDED state for %.2f seconds", job_name, time_seconds) |
| 153 | + |
| 154 | + |
| 155 | +def cleanup(pid=None): |
| 156 | + if pid: |
| 157 | + logging.info("Cleaning up prometheus metrics DB file for %s", pid) |
| 158 | + multiprocess.mark_process_dead(pid) |
0 commit comments