Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 57 additions & 5 deletions python/packages/jumpstarter-cli/jumpstarter_cli/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,15 @@ def _reap_zombie_processes(capture_child=None):
logger.warning(f"PARENT: Error during zombie reaping: {e}")


def _handle_child(config, parsed_bind=None, tls_insecure=False, tls_cert=None, tls_key=None, passphrase=None): # noqa: C901
def _handle_child( # noqa: C901
config,
parsed_bind=None,
tls_insecure=False,
tls_cert=None,
tls_key=None,
passphrase=None,
metrics_bind_address=":8080",
):
"""Handle child process with graceful shutdown."""
async def serve_with_graceful_shutdown(): # noqa: C901
received_signal = 0
Expand Down Expand Up @@ -99,6 +107,12 @@ async def signal_handler():
# Start signal handler immediately
signal_tg.start_soon(signal_handler)

from jumpstarter.metrics import start_metrics_server

listen_addr = start_metrics_server(metrics_bind_address)
if listen_addr:
logger.info("Serving metrics server at http://%s/metrics", listen_addr)

if parsed_bind is not None:
host, port = parsed_bind
tls_credentials = None
Expand Down Expand Up @@ -204,7 +218,13 @@ def parent_signal_handler(signum, _):


def _serve_with_exc_handling(
config, parsed_bind=None, tls_insecure=False, tls_cert=None, tls_key=None, passphrase=None
config,
parsed_bind=None,
tls_insecure=False,
tls_cert=None,
tls_key=None,
passphrase=None,
metrics_bind_address=":8080",
):
max_rapid_failures = config.failure_detection.max_rapid_failures
rapid_failure_window = config.failure_detection.rapid_failure_window
Expand Down Expand Up @@ -253,7 +273,15 @@ def _serve_with_exc_handling(
rapid_failure_count = 0
else:
os.setsid() # Become group leader so all spawned subprocesses are reached by parent's signals
_handle_child(config, parsed_bind, tls_insecure, tls_cert, tls_key, passphrase)
_handle_child(
config,
parsed_bind,
tls_insecure,
tls_cert,
tls_key,
passphrase,
metrics_bind_address,
)
sys.exit(1) # should never happen


Expand Down Expand Up @@ -294,8 +322,24 @@ def _serve_with_exc_handling(
default=False,
help="Exit after the current lease ends instead of waiting for a new one.",
)
@click.option(
"--metrics-bind-address",
"metrics_bind_address",
default=":8080",
show_default=True,
help="Address for HTTP GET /metrics (Prometheus/OpenMetrics). Use 0 to disable.",
)
@handle_exceptions
def run(config, listener_bind, tls_insecure, tls_cert, tls_key, passphrase, exit_on_lease_end):
def run(
config,
listener_bind,
tls_insecure,
tls_cert,
tls_key,
passphrase,
exit_on_lease_end,
metrics_bind_address,
):
"""Run an exporter locally."""
if listener_bind is not None and config is None:
raise click.UsageError("--exporter-config (or --exporter) is required when using --tls-grpc-listener")
Expand All @@ -313,4 +357,12 @@ def run(config, listener_bind, tls_insecure, tls_cert, tls_key, passphrase, exit
if exit_on_lease_end:
config.exit_on_lease_end = True
parsed_bind = _parse_listener_bind(listener_bind) if listener_bind is not None else None
return _serve_with_exc_handling(config, parsed_bind, tls_insecure, tls_cert, tls_key, passphrase)
return _serve_with_exc_handling(
config,
parsed_bind,
tls_insecure,
tls_cert,
tls_key,
passphrase,
metrics_bind_address,
)
109 changes: 109 additions & 0 deletions python/packages/jumpstarter/jumpstarter/driver/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

import logging
import os
import time
from abc import ABCMeta, abstractmethod
from contextlib import asynccontextmanager
from dataclasses import field
Expand Down Expand Up @@ -113,11 +114,36 @@ def client(cls) -> str:
def extra_labels(self) -> dict[str, str]:
return {}

def _record_operation_metrics(
self,
*,
operation: str,
result: str,
duration_seconds: float,
error_type: str | None = None,
) -> None:
from jumpstarter.metrics.registry import (
exemplars_from_log_context,
exporter_from_log_context,
get_registry,
)

get_registry().record_operation(
exporter=exporter_from_log_context(default=self.name if hasattr(self, "name") else "unknown"),
operation=operation,
result=result,
driver_type=self.driver_type,
duration_seconds=duration_seconds,
exemplars=exemplars_from_log_context(),
error_type=error_type,
)
Comment on lines +117 to +139

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider moving the imports to module level or caching the resolved functions. Capture the exporter name and exemplars once before the stream loop.


async def DriverCall(self, request, context):
"""
:meta private:
"""
op = request.method
started = time.perf_counter()
self.logger.info(
"Operation started",
extra={"operation": op, "driver_type": self.driver_type},
Comment on lines 114 to 149

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The operation label passed to _record_operation_metrics comes from request.method, a client-controlled gRPC field. When __lookup_drivercall rejects an unknown method, the resulting AbortError falls through to the catch-all handler, which records metrics with the untrusted value as a Prometheus label. Each unique method name creates a new time series, allowing an authenticated client to exhaust exporter memory by sending random method names. Fixing the AbortError handling (catching it before the generic handler) would largely mitigate this.

AI-generated, human reviewed

Expand All @@ -132,6 +158,11 @@ async def DriverCall(self, request, context):
else:
result = await to_thread.run_sync(method, *args)

self._record_operation_metrics(
operation=op,
result="success",
duration_seconds=time.perf_counter() - started,
)
self.logger.info(
"Operation completed",
extra={"operation": op, "driver_type": self.driver_type, "result": "success"},
Expand All @@ -141,41 +172,77 @@ async def DriverCall(self, request, context):
result=encode_value(result),
)
except NotImplementedError as e:
self._record_operation_metrics(
operation=op,
result="failure",
duration_seconds=time.perf_counter() - started,
error_type="not_implemented",
)
self.logger.warning(
"Operation failed",
extra={"operation": op, "driver_type": self.driver_type,
"result": "failure", "error_type": "not_implemented"},
)
await context.abort(StatusCode.UNIMPLEMENTED, str(e))
except ValueError as e:
self._record_operation_metrics(
operation=op,
result="failure",
duration_seconds=time.perf_counter() - started,
error_type="validation_error",
)
self.logger.warning(
"Operation failed",
extra={"operation": op, "driver_type": self.driver_type,
"result": "failure", "error_type": "validation_error"},
)
await context.abort(StatusCode.INVALID_ARGUMENT, str(e))
except TimeoutError as e:
self._record_operation_metrics(
operation=op,
result="failure",
duration_seconds=time.perf_counter() - started,
error_type="timeout",
)
self.logger.warning(
"Operation failed",
extra={"operation": op, "driver_type": self.driver_type,
"result": "failure", "error_type": "timeout"},
)
await context.abort(StatusCode.DEADLINE_EXCEEDED, str(e))
except ConnectionError as e:
self._record_operation_metrics(
operation=op,
result="failure",
duration_seconds=time.perf_counter() - started,
error_type="connection_error",
)
self.logger.warning(
"Operation failed",
extra={"operation": op, "driver_type": self.driver_type,
"result": "failure", "error_type": "connection_error"},
)
await context.abort(StatusCode.UNAVAILABLE, str(e))
except OSError as e:
self._record_operation_metrics(
operation=op,
result="failure",
duration_seconds=time.perf_counter() - started,
error_type="device_error",
)
self.logger.warning(
"Operation failed",
extra={"operation": op, "driver_type": self.driver_type,
"result": "failure", "error_type": "device_error"},
)
await context.abort(StatusCode.INTERNAL, str(e))
except Exception as e:
self._record_operation_metrics(
operation=op,
result="failure",
duration_seconds=time.perf_counter() - started,
error_type="internal_error",
)
self.logger.warning(
"Operation failed",
extra={"operation": op, "driver_type": self.driver_type,
Comment on lines 239 to 248

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider adding except grpc.aio.AbortError: raise immediately before except Exception in both methods to let abort errors propagate without distorting metrics or status codes.

Comment on lines 172 to 248

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider adding test drivers whose @export methods raise each exception type (e.g., TimeoutError), then assert the correct error_type label appears in the metrics output.

Expand All @@ -188,6 +255,7 @@ async def StreamingDriverCall(self, request, context):
:meta private:
"""
op = request.method
started = time.perf_counter()
self.logger.info(
"Operation started",
extra={"operation": op, "driver_type": self.driver_type},
Expand All @@ -209,46 +277,87 @@ async def StreamingDriverCall(self, request, context):
uuid=str(uuid4()),
result=encode_value(result),
)
self._record_operation_metrics(
operation=op,
result="success",
duration_seconds=time.perf_counter() - started,
)
self.logger.info(
"Operation completed",
extra={"operation": op, "driver_type": self.driver_type, "result": "success"},
)
except NotImplementedError as e:
self._record_operation_metrics(
operation=op,
result="failure",
duration_seconds=time.perf_counter() - started,
error_type="not_implemented",
)
self.logger.warning(
"Operation failed",
extra={"operation": op, "driver_type": self.driver_type,
"result": "failure", "error_type": "not_implemented"},
)
await context.abort(StatusCode.UNIMPLEMENTED, str(e))
except ValueError as e:
self._record_operation_metrics(
operation=op,
result="failure",
duration_seconds=time.perf_counter() - started,
error_type="validation_error",
)
self.logger.warning(
"Operation failed",
extra={"operation": op, "driver_type": self.driver_type,
"result": "failure", "error_type": "validation_error"},
)
await context.abort(StatusCode.INVALID_ARGUMENT, str(e))
except TimeoutError as e:
self._record_operation_metrics(
operation=op,
result="failure",
duration_seconds=time.perf_counter() - started,
error_type="timeout",
)
self.logger.warning(
"Operation failed",
extra={"operation": op, "driver_type": self.driver_type,
"result": "failure", "error_type": "timeout"},
)
await context.abort(StatusCode.DEADLINE_EXCEEDED, str(e))
except ConnectionError as e:
self._record_operation_metrics(
operation=op,
result="failure",
duration_seconds=time.perf_counter() - started,
error_type="connection_error",
)
self.logger.warning(
"Operation failed",
extra={"operation": op, "driver_type": self.driver_type,
"result": "failure", "error_type": "connection_error"},
)
await context.abort(StatusCode.UNAVAILABLE, str(e))
except OSError as e:
self._record_operation_metrics(
operation=op,
result="failure",
duration_seconds=time.perf_counter() - started,
error_type="device_error",
)
self.logger.warning(
"Operation failed",
extra={"operation": op, "driver_type": self.driver_type,
"result": "failure", "error_type": "device_error"},
)
await context.abort(StatusCode.INTERNAL, str(e))
except Exception as e:
self._record_operation_metrics(
operation=op,
result="failure",
duration_seconds=time.perf_counter() - started,
error_type="internal_error",
)
self.logger.warning(
"Operation failed",
extra={"operation": op, "driver_type": self.driver_type,
Expand Down
14 changes: 12 additions & 2 deletions python/packages/jumpstarter/jumpstarter/exporter/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,11 +57,20 @@ class Session(

@contextmanager
def __contextmanager__(self) -> Generator[Self]:
from jumpstarter.logging import set_log_context
from jumpstarter.metrics import get_registry

logging.getLogger().addHandler(self._logging_handler)
self.root_device.reset()
set_log_context(exporter=self.name)
get_registry().inc_active_sessions(exporter=self.name, delta=1.0)
try:
yield self
finally:
try:
get_registry().inc_active_sessions(exporter=self.name, delta=-1.0)
except Exception:
pass
Comment on lines +70 to +73

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Log exception at WARNING or DEBUG level before suppressing it.

try:
self.root_device.close()
except Exception as e:
Expand Down Expand Up @@ -319,14 +328,15 @@ async def StreamingDriverCall(self, request, context):
async def Stream(self, _request_iterator, context):
request = StreamRequestMetadata(**dict(list(context.invocation_metadata()))).request
logger.debug("Streaming(%s)", request)
async with self[request.uuid].Stream(request, context) as stream:
driver = self[request.uuid]
async with driver.Stream(request, context) as stream:
metadata = []
with suppress(TypedAttributeLookupError):
metadata.extend(stream.extra(MetadataStreamAttributes.metadata).items())
await context.send_initial_metadata(metadata)

async with RouterStream(context=context) as remote:
async with forward_stream(remote, stream):
async with forward_stream(remote, stream, metrics_driver_type=driver.driver_type):
event = Event()
context.add_done_callback(lambda _: event.set())
await event.wait()
Expand Down
17 changes: 17 additions & 0 deletions python/packages/jumpstarter/jumpstarter/metrics/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
"""Exporter-local Prometheus metrics (JEP-0013 Phase 2)."""

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see we do use top level comments in some places, but i'd drop the JEP part at least, from the others as well


from .registry import (
DEFAULT_EXEMPLAR_KEYS,
MetricsRegistry,
get_registry,
reset_registry_for_tests,
)
from .server import start_metrics_server

__all__ = [
"DEFAULT_EXEMPLAR_KEYS",
"MetricsRegistry",
"get_registry",
"reset_registry_for_tests",
"start_metrics_server",
]
Loading
Loading