1# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2# SPDX-License-Identifier: Apache-2.0
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import atexit
17import enum
18import inspect
19import json
20import logging
21import multiprocessing
22import multiprocessing.managers
23import os
24import threading
25from typing import Any, ClassVar, Literal
26
27import opentelemetry.metrics as api_metrics
28import opentelemetry.trace as api_trace
29import psutil
30
31from .. import utils
32
33# MSC telemetry prefers publishing raw samples when possible to support arbitrary post-hoc aggregations.
34#
35# Some setups, however, may need resampling to reduce sample volume. The resampling methods we use
36# sacrifice temporal resolution to preserve other information. Which method is used depends on if
37# the expected post-hoc aggregate function is decomposable:
38#
39# * Decomposable aggregate functions (e.g. count, sum, min, max).
40# * Use client-side aggregation.
41# * E.g. preserve request + response counts.
42# * Non-decomposable aggregate functions (e.g. average, percentile).
43# * Use decimation by an integer factor or last value.
44# * E.g. preserve the shape of the latency distribution (unlike tail sampling).
45
46_METRICS_EXPORTER_MAPPING = {
47 "console": "opentelemetry.sdk.metrics.export.ConsoleMetricExporter",
48 "otlp": "opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter",
49 # "Private" until it's decided whether this will be official.
50 "_otlp_msal": "multistorageclient.telemetry.metrics.exporters.otlp_msal._OTLPMSALMetricExporter",
51 "_otlp_mtls_vault": "multistorageclient.telemetry.metrics.exporters.otlp_mtls_vault._OTLPmTLSVaultMetricExporter",
52}
53
54_TRACE_EXPORTER_MAPPING = {
55 "console": "opentelemetry.sdk.trace.export.ConsoleSpanExporter",
56 "otlp": "opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter",
57 # "Private" until it's decided whether this will be official.
58 "_otlp_msal": "multistorageclient.telemetry.traces.exporters.otlp_msal._OTLPMSALSpanExporter",
59}
60
61logger = logging.getLogger(__name__)
62
63
[docs]
64class Telemetry:
65 """
66 Provides telemetry resources.
67
68 Instances shouldn't be copied between processes. Not fork-safe or pickleable.
69
70 Instances can be shared between processes by registering with a :py:class:`multiprocessing.managers.BaseManager` and using proxy objects.
71 """
72
73 # Metrics are named `multistorageclient.{property}(.{aggregation})?`.
74 #
75 # For example:
76 #
77 # - multistorageclient.data_size
78 # - Gauge for data size per individual operation.
79 # - For distributions (e.g. post-hoc histograms + heatmaps).
80 # - multistorageclient.data_size.sum
81 # - Counter (sum) for data size across all operations.
82 # - For aggregates (e.g. post-hoc data rate calculations).
83
84 # https://opentelemetry.io/docs/specs/semconv/general/naming#metrics
85 class GaugeName(enum.Enum):
86 LATENCY = "multistorageclient.latency"
87 DATA_SIZE = "multistorageclient.data_size"
88 DATA_RATE = "multistorageclient.data_rate"
89 FILE_DESCRIPTOR_DURATION = "multistorageclient.file_descriptor.duration"
90
91 # https://opentelemetry.io/docs/specs/semconv/general/metrics#units
92 _GAUGE_UNIT_MAPPING: ClassVar[dict[GaugeName, str]] = {
93 # Seconds.
94 GaugeName.LATENCY: "s",
95 # Bytes.
96 GaugeName.DATA_SIZE: "By",
97 # Bytes/second.
98 GaugeName.DATA_RATE: "By/s",
99 # Seconds.
100 GaugeName.FILE_DESCRIPTOR_DURATION: "s",
101 }
102
103 # https://opentelemetry.io/docs/specs/semconv/general/naming#metrics
104 class CounterName(enum.Enum):
105 REQUEST_SUM = "multistorageclient.request.sum"
106 RESPONSE_SUM = "multistorageclient.response.sum"
107 DATA_SIZE_SUM = "multistorageclient.data_size.sum"
108
109 # https://opentelemetry.io/docs/specs/semconv/general/metrics#units
110 _COUNTER_UNIT_MAPPING: ClassVar[dict[CounterName, str]] = {
111 # Unitless.
112 CounterName.REQUEST_SUM: "{request}",
113 # Unitless.
114 CounterName.RESPONSE_SUM: "{response}",
115 # Bytes.
116 CounterName.DATA_SIZE_SUM: "By",
117 }
118
119 # https://opentelemetry.io/docs/specs/semconv/general/naming#metrics
120 class UpDownCounterName(enum.Enum):
121 FILE_DESCRIPTOR_OPEN = "multistorageclient.file_descriptor.open"
122
123 # https://opentelemetry.io/docs/specs/semconv/general/metrics#units
124 _UP_DOWN_COUNTER_UNIT_MAPPING: ClassVar[dict[UpDownCounterName, str]] = {
125 UpDownCounterName.FILE_DESCRIPTOR_OPEN: "{file_descriptor}",
126 }
127
128 # Map of config as a sorted JSON string (since dictionaries can't be hashed) to meter provider.
129 _meter_provider_cache: dict[str, api_metrics.MeterProvider]
130 _meter_provider_cache_lock: threading.Lock
131 # Map of config as a sorted JSON string (since dictionaries can't be hashed) to meter.
132 _meter_cache: dict[str, api_metrics.Meter]
133 _meter_cache_lock: threading.Lock
134 # Map of config as a sorted JSON string (since dictionaries can't be hashed) to gauge name to gauge.
135 _gauge_cache: dict[str, dict[GaugeName, api_metrics._Gauge]]
136 _gauge_cache_lock: threading.Lock
137 # Map of config as a sorted JSON string (since dictionaries can't be hashed) to counter name to counter.
138 _counter_cache: dict[str, dict[CounterName, api_metrics.Counter]]
139 _counter_cache_lock: threading.Lock
140 # Map of config as a sorted JSON string (since dictionaries can't be hashed) to up-down counter name to counter.
141 _up_down_counter_cache: dict[str, dict[UpDownCounterName, api_metrics.UpDownCounter]]
142 _up_down_counter_cache_lock: threading.Lock
143 # Map of config as a sorted JSON string (since dictionaries can't be hashed) to tracer provider.
144 _tracer_provider_cache: dict[str, api_trace.TracerProvider]
145 _tracer_provider_cache_lock: threading.Lock
146 # Map of config as a sorted JSON string (since dictionaries can't be hashed) to tracer.
147 _tracer_cache: dict[str, api_trace.Tracer]
148 _tracer_cache_lock: threading.Lock
149
150 def __init__(self):
151 self._meter_provider_cache = {}
152 self._meter_provider_cache_lock = threading.Lock()
153 self._meter_cache = {}
154 self._meter_cache_lock = threading.Lock()
155 self._gauge_cache = {}
156 self._gauge_cache_lock = threading.Lock()
157 self._counter_cache = {}
158 self._counter_cache_lock = threading.Lock()
159 self._up_down_counter_cache = {}
160 self._up_down_counter_cache_lock = threading.Lock()
161 self._tracer_provider_cache = {}
162 self._tracer_provider_cache_lock = threading.Lock()
163 self._tracer_cache = {}
164 self._tracer_cache_lock = threading.Lock()
165
166 def _reinitialize_instance_locks_after_fork(self) -> None:
167 """
168 Reinitialize internal locks after fork to prevent deadlocks.
169
170 This method reinitializes all internal locks. Caches are kept as they
171 may still be valid, but locks must be fresh to avoid deadlocks.
172 """
173 self._meter_provider_cache_lock = threading.Lock()
174 self._meter_cache_lock = threading.Lock()
175 self._gauge_cache_lock = threading.Lock()
176 self._counter_cache_lock = threading.Lock()
177 self._up_down_counter_cache_lock = threading.Lock()
178 self._tracer_provider_cache_lock = threading.Lock()
179 self._tracer_cache_lock = threading.Lock()
180
181 def meter_provider(self, config: dict[str, Any]) -> api_metrics.MeterProvider | None:
182 """
183 Create or return an existing :py:class:`api_metrics.MeterProvider` for a config.
184
185 :param config: ``.opentelemetry.metrics`` config dict.
186 :return: A :py:class:`api_metrics.MeterProvider` or ``None`` if no valid exporter is configured.
187 """
188 config_json = json.dumps(config, sort_keys=True)
189 with self._meter_provider_cache_lock:
190 if config_json in self._meter_provider_cache:
191 return self._meter_provider_cache[config_json]
192 else:
193 if "exporter" in config:
194 try:
195 import opentelemetry.sdk.metrics as sdk_metrics
196 import opentelemetry.sdk.metrics.export as sdk_metrics_export
197
198 from .metrics.readers.diperiodic_exporting import DiperiodicExportingMetricReader
199
200 exporter_type: str = config["exporter"]["type"]
201 exporter_fully_qualified_name = _METRICS_EXPORTER_MAPPING.get(exporter_type, exporter_type)
202 exporter_module_name, exporter_class_name = exporter_fully_qualified_name.rsplit(".", 1)
203 cls = utils.import_class(exporter_class_name, exporter_module_name)
204 exporter_options = config["exporter"].get("options", {}).copy()
205 exporter: sdk_metrics_export.MetricExporter = cls(**exporter_options)
206
207 reader_options = config.get("reader", {}).get("options", {})
208 reader: sdk_metrics_export.MetricReader = DiperiodicExportingMetricReader(
209 **reader_options, exporter=exporter
210 )
211
212 return self._meter_provider_cache.setdefault(
213 config_json, sdk_metrics.MeterProvider(metric_readers=[reader])
214 )
215 except (AttributeError, ImportError):
216 logger.exception("Failed to import OpenTelemetry Python SDK or exporter! Disabling metrics.")
217 return None
218 else:
219 # Don't return a no-op meter provider to avoid unnecessary overhead.
220 logger.error("No exporter configured! Disabling metrics.")
221 return None
222
223 def meter(self, config: dict[str, Any]) -> api_metrics.Meter | None:
224 """
225 Create or return an existing :py:class:`api_metrics.Meter` for a config.
226
227 :param config: ``.opentelemetry.metrics`` config dict.
228 :return: A :py:class:`api_metrics.Meter` or ``None`` if no valid exporter is configured.
229 """
230 config_json = json.dumps(config, sort_keys=True)
231 with self._meter_cache_lock:
232 if config_json in self._meter_cache:
233 return self._meter_cache[config_json]
234 else:
235 meter_provider = self.meter_provider(config=config)
236 if meter_provider is None:
237 return None
238 else:
239 return self._meter_cache.setdefault(
240 config_json, meter_provider.get_meter(name="multistorageclient")
241 )
242
243 def gauge(self, config: dict[str, Any], name: GaugeName) -> api_metrics._Gauge | None:
244 """
245 Create or return an existing :py:class:`api_metrics.Gauge` for a config and gauge name.
246
247 :param config: ``.opentelemetry.metrics`` config dict.
248 :return: A :py:class:`api_metrics.Gauge` or ``None`` if no valid exporter is configured.
249 """
250 config_json = json.dumps(config, sort_keys=True)
251 with self._gauge_cache_lock:
252 if config_json in self._gauge_cache and name in self._gauge_cache[config_json]:
253 return self._gauge_cache[config_json][name]
254 else:
255 meter = self.meter(config=config)
256 if meter is None:
257 return None
258 else:
259 return self._gauge_cache.setdefault(config_json, {}).setdefault(
260 name,
261 meter.create_gauge(name=name.value, unit=Telemetry._GAUGE_UNIT_MAPPING.get(name, "")),
262 )
263
264 def counter(self, config: dict[str, Any], name: CounterName) -> api_metrics.Counter | None:
265 """
266 Create or return an existing :py:class:`api_metrics.Counter` for a config and counter name.
267
268 :param config: ``.opentelemetry.metrics`` config dict.
269 :return: A :py:class:`api_metrics.Counter` or ``None`` if no valid exporter is configured.
270 """
271 config_json = json.dumps(config, sort_keys=True)
272 with self._counter_cache_lock:
273 if config_json in self._counter_cache and name in self._counter_cache[config_json]:
274 return self._counter_cache[config_json][name]
275 else:
276 meter = self.meter(config=config)
277 if meter is None:
278 return None
279 else:
280 return self._counter_cache.setdefault(config_json, {}).setdefault(
281 name,
282 meter.create_counter(name=name.value, unit=Telemetry._COUNTER_UNIT_MAPPING.get(name, "")),
283 )
284
285 def up_down_counter(self, config: dict[str, Any], name: UpDownCounterName) -> api_metrics.UpDownCounter | None:
286 """
287 Create or return an existing :py:class:`api_metrics.UpDownCounter` for a config and counter name.
288
289 :param config: ``.opentelemetry.metrics`` config dict.
290 :param name: Up-down counter name.
291 :return: A :py:class:`api_metrics.UpDownCounter` or ``None`` if no valid exporter is configured.
292 """
293 config_json = json.dumps(config, sort_keys=True)
294 with self._up_down_counter_cache_lock:
295 if config_json in self._up_down_counter_cache and name in self._up_down_counter_cache[config_json]:
296 return self._up_down_counter_cache[config_json][name]
297
298 meter = self.meter(config=config)
299 if meter is None:
300 return None
301
302 return self._up_down_counter_cache.setdefault(config_json, {}).setdefault(
303 name,
304 meter.create_up_down_counter(
305 name=name.value,
306 unit=Telemetry._UP_DOWN_COUNTER_UNIT_MAPPING.get(name, ""),
307 ),
308 )
309
310 def tracer_provider(self, config: dict[str, Any]) -> api_trace.TracerProvider | None:
311 """
312 Create or return an existing :py:class:`api_trace.TracerProvider` for a config.
313
314 :param config: ``.opentelemetry.traces`` config dict.
315 :return: A :py:class:`api_trace.TracerProvider` or ``None`` if no valid exporter is configured.
316 """
317 config_json = json.dumps(config, sort_keys=True)
318 with self._tracer_provider_cache_lock:
319 if config_json in self._tracer_provider_cache:
320 return self._tracer_provider_cache[config_json]
321 else:
322 if "exporter" in config:
323 try:
324 import opentelemetry.sdk.trace as sdk_trace
325 import opentelemetry.sdk.trace.export as sdk_trace_export
326 import opentelemetry.sdk.trace.sampling as sdk_trace_sampling
327
328 exporter_type: str = config["exporter"]["type"]
329 exporter_fully_qualified_name = _TRACE_EXPORTER_MAPPING.get(exporter_type, exporter_type)
330 exporter_module_name, exporter_class_name = exporter_fully_qualified_name.rsplit(".", 1)
331 cls = utils.import_class(exporter_class_name, exporter_module_name)
332 exporter_options = config["exporter"].get("options", {})
333 exporter: sdk_trace_export.SpanExporter = cls(**exporter_options)
334
335 processor: sdk_trace.SpanProcessor = sdk_trace.SynchronousMultiSpanProcessor()
336 processor.add_span_processor(sdk_trace_export.BatchSpanProcessor(span_exporter=exporter))
337
338 # TODO: Add sampler to configuration schema.
339 sampler: sdk_trace_sampling.Sampler = sdk_trace_sampling.ALWAYS_ON
340
341 return self._tracer_provider_cache.setdefault(
342 config_json,
343 sdk_trace.TracerProvider(active_span_processor=processor, sampler=sampler),
344 )
345 except (AttributeError, ImportError):
346 logger.exception("Failed to import OpenTelemetry Python SDK or exporter! Disabling traces.")
347 return None
348 else:
349 logger.error("No exporter configured! Disabling traces.")
350 return None
351
352 def tracer(self, config: dict[str, Any]) -> api_trace.Tracer | None:
353 """
354 Create or return an existing :py:class:`api_trace.Tracer` for a config.
355
356 :param config: ``.opentelemetry.traces`` config dict.
357 :return: A :py:class:`api_trace.Tracer` or ``None`` if no valid exporter is configured.
358 """
359 config_json = json.dumps(config, sort_keys=True)
360 with self._tracer_cache_lock:
361 if config_json in self._tracer_cache:
362 return self._tracer_cache[config_json]
363 else:
364 tracer_provider = self.tracer_provider(config=config)
365 if tracer_provider is None:
366 return None
367 else:
368 return self._tracer_cache.setdefault(
369 config_json, tracer_provider.get_tracer(instrumenting_module_name="multistorageclient")
370 )
371
372
373# To share a single :py:class:`Telemetry` within a process (e.g. local, manager).
374#
375# A manager's server processes shouldn't be forked, so this should be safe.
376_TELEMETRY: Telemetry | None = None
377_TELEMETRY_LOCK = threading.Lock()
378
379
380def _init() -> Telemetry:
381 """
382 Create or return an existing :py:class:`Telemetry`.
383
384 :return: A telemetry instance.
385 """
386 global _TELEMETRY
387 global _TELEMETRY_LOCK # noqa: PLW0602
388
389 with _TELEMETRY_LOCK:
390 if _TELEMETRY is None:
391 _TELEMETRY = Telemetry()
392 return _TELEMETRY
393
394
[docs]
395class TelemetryManager(multiprocessing.managers.BaseManager):
396 """
397 A :py:class:`multiprocessing.managers.BaseManager` for telemetry resources.
398
399 The OpenTelemetry Python SDK isn't fork-safe since telemetry sample buffers can be duplicated.
400
401 In addition, Python ≤3.12 doesn't call exit handlers for forked processes.
402 This causes the OpenTelemetry Python SDK to not flush telemetry before exiting.
403
404 * https://github.com/open-telemetry/opentelemetry-python/issues/4215
405 * https://github.com/open-telemetry/opentelemetry-python/issues/3307
406
407 Forking is multiprocessing's default start method for non-macOS POSIX systems until Python 3.14.
408
409 * https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods
410
411 To fully support multiprocessing, resampling + publishing is handled by
412 a single process that's (ideally) a child of (i.e. directly under) the main process. This:
413
414 * Relieves other processes of this work.
415
416 * Avoids issues with duplicate samples when forking and unpublished samples when exiting forks.
417
418 * Allows cross-process resampling.
419 * Reuses a single connection pool to telemetry backends.
420
421 The downside is it essentially re-introduces global interpreter lock (GIL) with
422 additional IPC overhead. Telemetry operations, however, should be lightweight so
423 this isn't expected to be a problem. Remote data store latency should still be
424 the primary throughput limiter for storage clients.
425
426 :py:class:`multiprocessing.managers.BaseManager` is used for this since it creates
427 a separate server process for shared objects.
428
429 Telemetry resources are provided as
430 `proxy objects <https://docs.python.org/3/library/multiprocessing.html#proxy-objects>`_
431 for location transparency.
432
433 The documentation isn't particularly detailed, but others have written comprehensively on this:
434
435 * https://zpz.github.io/blog/python-mp-manager-1
436 * https://zpz.github.io/blog/python-mp-manager-2
437 * https://zpz.github.io/blog/python-mp-manager-3
438
439 By specification, metric and tracer providers must call shutdown on any
440 underlying metric readers + span processors + exporters.
441
442 * https://opentelemetry.io/docs/specs/otel/metrics/sdk#shutdown
443 * https://opentelemetry.io/docs/specs/otel/trace/sdk#shutdown
444
445 In the OpenTelemetry Python SDK, provider shutdown is called automatically
446 by exit handlers (when they work at least). Consequently, clients should:
447
448 * Only receive proxy objects.
449
450 * Enables metric reader + span processor + exporter re-use across processes.
451
452 * Never call shutdown on the proxy objects.
453
454 * The shutdown exit handler is registered on the manager's server process.
455 * ⚠️ We expect a finite number of providers (i.e. no dynamic configs) so we don't leak them.
456 """
457
458
459def _fully_qualified_name(c: type[Any]) -> str:
460 """
461 Return the fully qualified name for a class (e.g. ``module.Class``).
462
463 For :py:class:`multiprocessing.Manager` type IDs.
464 """
465 return f"{c.__module__}.{c.__qualname__}"
466
467
468# Metrics proxy object setup.
469#
470# ``exposed`` is derived from the instrument class's public functions rather than relying on
471# ``multiprocessing``'s default discovery (``public_methods`` via ``dir()``): an OpenTelemetry build
472# that delegates an instrument's method isn't surfaced by ``dir()``, so the ``AutoProxy`` would drop
473# it and raise ``AttributeError: 'AutoProxy[...]' object has no attribute 'add'`` on record.
474#
475# Unlike ``Span`` below, the metric instrument classes define their own ``__init__``; it must be
476# excluded (public-only), otherwise it would override ``BaseProxy.__init__`` and break proxy setup.
477TelemetryManager.register(
478 typeid=_fully_qualified_name(api_metrics._Gauge),
479 exposed=[
480 name
481 for name, _ in inspect.getmembers(api_metrics._Gauge, predicate=inspect.isfunction)
482 if not name.startswith("_")
483 ],
484)
485TelemetryManager.register(
486 typeid=_fully_qualified_name(api_metrics.Counter),
487 exposed=[
488 name
489 for name, _ in inspect.getmembers(api_metrics.Counter, predicate=inspect.isfunction)
490 if not name.startswith("_")
491 ],
492)
493TelemetryManager.register(
494 typeid=_fully_qualified_name(api_metrics.UpDownCounter),
495 exposed=[
496 name
497 for name, _ in inspect.getmembers(api_metrics.UpDownCounter, predicate=inspect.isfunction)
498 if not name.startswith("_")
499 ],
500)
501TelemetryManager.register(
502 typeid=_fully_qualified_name(api_metrics.Meter),
503 method_to_typeid={
504 api_metrics.Meter.create_gauge.__name__: _fully_qualified_name(api_metrics._Gauge),
505 api_metrics.Meter.create_counter.__name__: _fully_qualified_name(api_metrics.Counter),
506 api_metrics.Meter.create_up_down_counter.__name__: _fully_qualified_name(api_metrics.UpDownCounter),
507 },
508)
509TelemetryManager.register(
510 typeid=_fully_qualified_name(api_metrics.MeterProvider),
511 method_to_typeid={api_metrics.MeterProvider.get_meter.__name__: _fully_qualified_name(api_metrics.Meter)},
512)
513
514# Traces proxy object setup.
515TelemetryManager.register(
516 typeid=_fully_qualified_name(api_trace.Span),
517 # Non-public methods (i.e. ones starting with a ``_``) are omitted by default.
518 #
519 # We need ``__enter__`` and ``__exit__`` so the ``Span`` can be used as a ``ContextManager``.
520 exposed=[name for name, _ in inspect.getmembers(api_trace.Span, predicate=inspect.isfunction)],
521 method_to_typeid={api_trace.Span.__enter__.__name__: _fully_qualified_name(api_trace.Span)},
522)
523TelemetryManager.register(
524 typeid=_fully_qualified_name(api_trace.Tracer),
525 # Can't proxy ``Tracer.start_as_current_span`` since it returns a generator (not pickleable)
526 # and tries to use the process-local global context (in this case, the manager's server process).
527 #
528 # Instead, spans should be constructed by:
529 #
530 # 1. Calling ``opentelemetry.context.get_current()`` to get the process-local global context (pickleable).
531 # 2. Creating a new span with the process-local global context.
532 # 3. Calling ``opentelemetry.trace.use_span()`` with the span and ``end_on_exit=True``.
533 method_to_typeid={api_trace.Tracer.start_span.__name__: _fully_qualified_name(api_trace.Span)},
534)
535TelemetryManager.register(
536 typeid=_fully_qualified_name(api_trace.TracerProvider),
537 method_to_typeid={api_trace.TracerProvider.get_tracer.__name__: _fully_qualified_name(api_trace.Tracer)},
538)
539
540# Telemetry proxy object setup.
541#
542# This should be the only registered type with a ``callable``.
543# It's the only type we create directly with a ``TelemetryManager``.
544TelemetryManager.register(
545 typeid=Telemetry.__name__,
546 callable=_init,
547 method_to_typeid={
548 Telemetry.meter_provider.__name__: _fully_qualified_name(api_metrics.MeterProvider),
549 Telemetry.meter.__name__: _fully_qualified_name(api_metrics.Meter),
550 Telemetry.gauge.__name__: _fully_qualified_name(api_metrics._Gauge),
551 Telemetry.counter.__name__: _fully_qualified_name(api_metrics.Counter),
552 Telemetry.up_down_counter.__name__: _fully_qualified_name(api_metrics.UpDownCounter),
553 Telemetry.tracer_provider.__name__: _fully_qualified_name(api_trace.TracerProvider),
554 Telemetry.tracer.__name__: _fully_qualified_name(api_trace.Tracer),
555 },
556)
557
558
559# Map of init options as a sorted JSON string (since dictionaries can't be hashed) to telemetry proxy.
560_TELEMETRY_PROXIES: dict[str, Telemetry] = {}
561# To share :py:class:`Telemetry` proxy objects within a process (e.g. client, server).
562#
563# Forking isn't expected to happen while this is held (may lead to a deadlock).
564_TELEMETRY_PROXIES_LOCK = threading.Lock()
565
566
567def _reinitialize_locks_after_fork() -> None:
568 """
569 Reinitialize telemetry locks after fork to prevent deadlocks.
570
571 This function is called automatically after a fork to reinitialize all locks.
572 Caches and instances are kept as they may still be valid, but locks must be
573 fresh to avoid deadlocks.
574 """
575 global _TELEMETRY, _TELEMETRY_LOCK, _TELEMETRY_PROXIES_LOCK # noqa: PLW0602
576
577 _TELEMETRY_LOCK = threading.Lock()
578 _TELEMETRY_PROXIES_LOCK = threading.Lock()
579
580 # Reinitialize LOCAL mode telemetry instance locks if it exists
581 if _TELEMETRY is not None:
582 _TELEMETRY._reinitialize_instance_locks_after_fork()
583
584
585if hasattr(os, "register_at_fork"):
586 os.register_at_fork(after_in_child=_reinitialize_locks_after_fork)
587
588
[docs]
589class TelemetryMode(enum.Enum):
590 """
591 How to create a :py:class:`Telemetry` object.
592 """
593
594 #: Keep everything local to the process (not fork-safe).
595 LOCAL = "local"
596 #: Start + connect to a telemetry IPC server.
597 SERVER = "server"
598 #: Connect to a telemetry IPC server.
599 CLIENT = "client"
600
601
602def _telemetry_proxies_key(
603 mode: Literal[TelemetryMode.SERVER, TelemetryMode.CLIENT], address: str | tuple[str, int]
604) -> str:
605 """
606 Get the key for the _TELEMETRY_PROXIES dictionary.
607 """
608 init_options = {"mode": mode.value, "address": address}
609 return json.dumps(init_options, sort_keys=True)
610
611
612def _telemetry_manager_server_port(process_id: int) -> int:
613 """
614 Get the default telemetry manager server port.
615
616 This is PID-based to:
617
618 * Avoid collisions between multiple independent Python interpreters running on the same machine.
619 * Let child processes deterministically find their parent's telemetry manager server.
620
621 :param process_id: Process ID used to calculate the port.
622 """
623
624 # Use the dynamic/private/ephemeral port range.
625 #
626 # https://www.rfc-editor.org/rfc/rfc6335.html#section-6
627 # https://en.wikipedia.org/wiki/List_of_TCP_and_UDP_port_numbers#Dynamic,_private_or_ephemeral_ports
628 #
629 # Modulo the parent/child process PID by the port range length, then add the initial offset.
630 return (2**15 + 2**14) + (process_id % ((2**16) - (2**15 + 2**14)))
631
632
633def _init_server(address: str | tuple[str, int] | None = None) -> Telemetry:
634 """
635 Start + connect to a telemetry IPC server.
636 """
637 global _TELEMETRY_PROXIES # noqa: PLW0602
638 global _TELEMETRY_PROXIES_LOCK # noqa: PLW0602
639
640 address = address or ("127.0.0.1", _telemetry_manager_server_port(process_id=psutil.Process().pid))
641 telemetry_proxies_key = _telemetry_proxies_key(mode=TelemetryMode.SERVER, address=address)
642
643 with _TELEMETRY_PROXIES_LOCK:
644 if telemetry_proxies_key in _TELEMETRY_PROXIES:
645 return _TELEMETRY_PROXIES[telemetry_proxies_key]
646 else:
647 telemetry_manager = TelemetryManager(
648 address=address,
649 # Use spawn instead of the platform-specific default (may be fork) to avoid aforementioned issues with fork.
650 ctx=multiprocessing.get_context(method="spawn"),
651 )
652
653 logger.debug(f"Creating telemetry manager server at {telemetry_manager.address}.")
654 try:
655 telemetry_manager.start()
656 atexit.register(telemetry_manager.shutdown)
657 logger.debug(f"Started telemetry manager server at {telemetry_manager.address}.")
658 except Exception:
659 logger.debug(
660 f"Failed to create telemetry manager server at {telemetry_manager.address}!", exc_info=True
661 )
662 raise
663
664 logger.debug(f"Connecting to telemetry manager server at {telemetry_manager.address}.")
665 try:
666 telemetry_manager.connect()
667 logger.debug(f"Connected to telemetry manager server at {telemetry_manager.address}.")
668 return _TELEMETRY_PROXIES.setdefault(telemetry_proxies_key, telemetry_manager.Telemetry()) # pyright: ignore [reportAttributeAccessIssue]
669 except Exception:
670 logger.debug(
671 f"Failed to connect to telemetry manager server at {telemetry_manager.address}!", exc_info=True
672 )
673 raise
674
675
676def _init_client(address: str | tuple[str, int] | None = None) -> Telemetry:
677 """
678 Connect to a telemetry IPC server.
679 """
680 global _TELEMETRY_PROXIES # noqa: PLW0602
681 global _TELEMETRY_PROXIES_LOCK # noqa: PLW0602
682
683 candidate_addresses: list[str | tuple[str, int]] = []
684
685 if address is not None:
686 candidate_addresses = [address]
687 else:
688 current_process = psutil.Process()
689 # Python processes from leaf to root.
690 python_process_ancestry: list[psutil.Process] = [
691 current_process,
692 # Try the default telemetry manager server port for all ancestor process IDs.
693 #
694 # psutil is used since multiprocessing only exposes the parent process.
695 #
696 # We can't use `itertools.takewhile(lambda ancestor_process: ancestor_process.name() == current_process.name(), ...)`
697 # to limit ourselves to ancestor Python processes by process name since some may not be named
698 # `python{optional version}` in some cases (e.g. may be named `pytest`).
699 *current_process.parents(),
700 ]
701 # Try to connect from leaf to root.
702 candidate_addresses = [
703 ("127.0.0.1", _telemetry_manager_server_port(process_id=process.pid)) for process in python_process_ancestry
704 ]
705
706 for candidate_address in candidate_addresses:
707 telemetry_proxies_key = _telemetry_proxies_key(mode=TelemetryMode.CLIENT, address=candidate_address)
708
709 with _TELEMETRY_PROXIES_LOCK:
710 if telemetry_proxies_key in _TELEMETRY_PROXIES:
711 return _TELEMETRY_PROXIES[telemetry_proxies_key]
712 else:
713 telemetry_manager = TelemetryManager(address=candidate_address)
714
715 logger.debug(f"Connecting to telemetry manager server at {telemetry_manager.address}.")
716 try:
717 telemetry_manager.connect()
718 logger.debug(f"Connected to telemetry manager server at {telemetry_manager.address}.")
719 return _TELEMETRY_PROXIES.setdefault(telemetry_proxies_key, telemetry_manager.Telemetry()) # pyright: ignore [reportAttributeAccessIssue]
720 except Exception:
721 logger.debug(
722 f"Failed to connect to telemetry manager server at {telemetry_manager.address}!",
723 exc_info=True,
724 )
725
726 raise ConnectionError(f"Failed to connect to telemetry manager server at any of {candidate_addresses}!")
727
728
[docs]
729def init(
730 mode: TelemetryMode | None = None,
731 address: str | tuple[str, int] | None = None,
732) -> Telemetry:
733 """
734 Create or return an existing :py:class:`Telemetry` instance or :py:class:`Telemetry` proxy object.
735
736 :param mode: How to create a :py:class:`Telemetry` object. If ``None``, the default heuristic chooses a mode based on the presence of telemetry IPC servers in the process tree.
737 :param address: Telemetry IPC server address. Passed directly to a :py:class:`multiprocessing.managers.BaseManager`. Ignored if the mode is :py:const:`TelemetryMode.LOCAL`.
738 :return: A telemetry instance.
739 """
740
741 if mode is None:
742 # Main process.
743 if multiprocessing.parent_process() is None:
744 # Daemons can't have child processes.
745 #
746 # Try to create a telemetry instance in local mode.
747 #
748 # ⚠️ This may cause CPU contention if the current process is compute-intensive
749 # and a high collect and/or export frequency is used due to global interpreter lock (GIL).
750 if multiprocessing.current_process().daemon:
751 return _init()
752 # Start + connect to a telemetry IPC server.
753 else:
754 return _init_server(address=address)
755 # Child process.
756 else:
757 # Connect to a telemetry IPC server.
758 try:
759 return _init_client(address=address)
760 # No existing telemetry IPC server.
761 except ConnectionError:
762 # Daemons can't have child processes.
763 #
764 # Try to create a telemetry instance in local mode.
765 #
766 # ⚠️ This may cause CPU contention if the current process is compute-intensive
767 # and a high collect and/or export frequency is used due to global interpreter lock (GIL).
768 if multiprocessing.current_process().daemon:
769 return _init()
770 # Start + connect to a telemetry IPC server.
771 else:
772 return _init_server(address=address)
773 elif mode == TelemetryMode.LOCAL:
774 return _init()
775 elif mode == TelemetryMode.SERVER:
776 return _init_server(address=address)
777 elif mode == TelemetryMode.CLIENT:
778 return _init_client(address=address)
779 else:
780 raise ValueError(f"Unsupported telemetry mode: {mode}")