Fix missing worker Prometheus metrics after early startup scrapes

distributed/http/worker/prometheus/core.py

@@ -171,7 +171,7 @@ def collect_crick(self) -> Iterator[Metric]:
         # The following metrics will export NaN, if the corresponding digests are None
         if not self.crick_available:
             return
-        assert self.server.digests
+        assert self.server.digests is not None
 
         yield GaugeMetricFamily(
             self.build_name("tick_duration_median"),
@@ -262,9 +262,11 @@ def __init__(self, *args, dask_server=None, **kwargs):
             PrometheusHandler._collector.server = self.server
             return
 
-        PrometheusHandler._collector = WorkerMetricCollector(self.server)
-        # Register collector
-        prometheus_client.REGISTRY.register(PrometheusHandler._collector)
+        collector = WorkerMetricCollector(self.server)
+        # Registration collects metrics and may fail. Only cache a registered
+        # collector so a later request can retry after a failure.
+        prometheus_client.REGISTRY.register(collector)
+        PrometheusHandler._collector = collector
 
     def get(self):
         self.write(prometheus_client.generate_latest())

distributed/http/worker/tests/test_worker_http.py

@@ -15,6 +15,71 @@
 )
 
 
+@gen_cluster(nthreads=[])
+async def test_prometheus_before_worker_registered(s, monkeypatch):
+    prometheus_client = pytest.importorskip("prometheus_client")
+    pytest.importorskip("crick")
+    from distributed.http.worker.prometheus.core import PrometheusHandler
+
+    registry = prometheus_client.CollectorRegistry(auto_describe=True)
+    generate_latest = prometheus_client.generate_latest
+    monkeypatch.setattr(prometheus_client, "REGISTRY", registry)
+    monkeypatch.setattr(
+        prometheus_client, "generate_latest", lambda: generate_latest(registry)
+    )
+    monkeypatch.setattr(PrometheusHandler, "_collector", None)
+
+    class StartupWorker(Worker):
+        async def _register_with_scheduler(self):
+            # HTTP is already listening, but periodic callbacks have not populated
+            # the worker's digests yet. Prometheus may discover this endpoint now.
+            assert self.digests is not None
+            assert not self.digests
+            metrics = await fetch_metrics_sample_names(
+                self.http_server.port, prefix="dask_worker_"
+            )
+            assert "dask_worker_threads" in metrics
+            assert "dask_worker_tick_duration_median_seconds" in metrics
+            await super()._register_with_scheduler()
+
+    async with StartupWorker(s.address, dashboard_address=":0") as worker:
+        for _ in range(2):
+            metrics = await fetch_metrics_sample_names(
+                worker.http_server.port, prefix="dask_worker_"
+            )
+            assert "dask_worker_threads" in metrics
+
+
+@gen_cluster(nthreads=[("127.0.0.1", 1)])
+async def test_prometheus_registration_retry(s, a, monkeypatch):
+    prometheus_client = pytest.importorskip("prometheus_client")
+    from tornado.simple_httpclient import HTTPStreamClosedError
+
+    from distributed.http.worker.prometheus.core import PrometheusHandler
+
+    registry = prometheus_client.CollectorRegistry(auto_describe=True)
+    generate_latest = prometheus_client.generate_latest
+    monkeypatch.setattr(prometheus_client, "REGISTRY", registry)
+    monkeypatch.setattr(
+        prometheus_client, "generate_latest", lambda: generate_latest(registry)
+    )
+    monkeypatch.setattr(PrometheusHandler, "_collector", None)
+
+    with (
+        mock.patch.object(
+            registry, "register", side_effect=RuntimeError("test failure")
+        ),
+        pytest.raises(HTTPStreamClosedError),
+    ):
+        await fetch_metrics_sample_names(a.http_server.port, prefix="dask_worker_")
+
+    for _ in range(2):
+        metrics = await fetch_metrics_sample_names(
+            a.http_server.port, prefix="dask_worker_"
+        )
+        assert "dask_worker_threads" in metrics
+
+
 @gen_cluster(
     client=True,
     nthreads=[("127.0.0.1", 1)],