feat: convert telemetry service to docket perpetual function (#19754)

zzstoatzz · claude · web-flow · commit 500a3d732e65 · 2025-12-12T14:39:21.000-06:00
Co-authored-by: Claude &lt;noreply@anthropic.com&gt;
diff --git a/src/prefect/server/services/base.py b/src/prefect/server/services/base.py
@@ -40,14 +40,12 @@ def _known_service_modules() -> list[ModuleType]:
     from prefect.server.services import (
         scheduler,
         task_run_recorder,
-        telemetry,
     )
 
     return [
         # Orchestration services
         scheduler,
         task_run_recorder,
-        telemetry,
         # Events services
         event_logger,
         event_persister,
diff --git a/src/prefect/server/services/telemetry.py b/src/prefect/server/services/telemetry.py
@@ -1,143 +1,114 @@
 """
-The Telemetry service.
+The Telemetry service. Sends anonymous data to Prefect to help us improve.
 """
 
-import asyncio
+import logging
 import os
 import platform
-from typing import Any, Optional
+from datetime import timedelta
 from uuid import uuid4
 
 import httpx
+from docket import Perpetual
 
 import prefect
-from prefect.server.database import PrefectDBInterface
-from prefect.server.database.dependencies import db_injector
+from prefect.logging import get_logger
+from prefect.server.database import PrefectDBInterface, provide_database_interface
 from prefect.server.models import configuration
 from prefect.server.schemas.core import Configuration
-from prefect.server.services.base import (
-    LoopService,
-    RunInEphemeralServers,
-    RunInWebservers,
-)
+from prefect.server.services.perpetual_services import perpetual_service
 from prefect.settings import PREFECT_DEBUG_MODE
 from prefect.settings.context import get_current_settings
-from prefect.settings.models.server.services import ServicesBaseSetting
 from prefect.types._datetime import now
 
+logger: logging.Logger = get_logger(__name__)
+
 
-class Telemetry(RunInEphemeralServers, RunInWebservers, LoopService):
+async def _fetch_or_set_telemetry_session(
+    db: PrefectDBInterface,
+) -> tuple[str, str]:
     """
-    Sends anonymous data to Prefect to help us improve
+    Fetch or create a telemetry session in the configuration table.
 
-    It can be toggled off with the PREFECT_SERVER_ANALYTICS_ENABLED setting.
+    Returns:
+        tuple of (session_start_timestamp, session_id)
     """
+    async with db.session_context(begin_transaction=True) as session:
+        telemetry_session = await configuration.read_configuration(
+            session, "TELEMETRY_SESSION"
+        )
 
-    loop_seconds: float = 600
+        if telemetry_session is None:
+            logger.debug("No telemetry session found, setting")
+            session_id = str(uuid4())
+            session_start_timestamp = now("UTC").isoformat()
+
+            telemetry_session = Configuration(
+                key="TELEMETRY_SESSION",
+                value={
+                    "session_id": session_id,
+                    "session_start_timestamp": session_start_timestamp,
+                },
+            )
 
-    @classmethod
-    def service_settings(cls) -> ServicesBaseSetting:
-        raise NotImplementedError("Telemetry service does not have settings")
+            await configuration.write_configuration(session, telemetry_session)
+        else:
+            logger.debug("Session information retrieved from database")
+            session_id = telemetry_session.value["session_id"]
+            session_start_timestamp = telemetry_session.value["session_start_timestamp"]
 
-    @classmethod
-    def environment_variable_name(cls) -> str:
-        return "PREFECT_SERVER_ANALYTICS_ENABLED"
+    logger.debug(f"Telemetry Session: {session_id}, {session_start_timestamp}")
+    return (session_start_timestamp, session_id)
 
-    @classmethod
-    def enabled(cls) -> bool:
-        return get_current_settings().server.analytics_enabled
 
-    def __init__(self, loop_seconds: Optional[int] = None, **kwargs: Any):
-        super().__init__(loop_seconds=loop_seconds, **kwargs)
-        self.telemetry_environment: str = os.environ.get(
-            "PREFECT_API_TELEMETRY_ENVIRONMENT", "production"
-        )
+@perpetual_service(
+    enabled_getter=lambda: get_current_settings().server.analytics_enabled,
+    run_in_ephemeral=True,
+    run_in_webserver=True,
+)
+async def send_telemetry_heartbeat(
+    perpetual: Perpetual = Perpetual(automatic=True, every=timedelta(seconds=600)),
+) -> None:
+    """
+    Sends anonymous telemetry data to Prefect to help us improve.
 
-    @db_injector
-    async def _fetch_or_set_telemetry_session(self, db: PrefectDBInterface):
-        """
-        This method looks for a telemetry session in the configuration table. If there
-        isn't one, it sets one. It then sets `self.session_id` and
-        `self.session_start_timestamp`.
-
-        Telemetry sessions last until the database is reset.
-        """
-        async with db.session_context(begin_transaction=True) as session:
-            telemetry_session = await configuration.read_configuration(
-                session, "TELEMETRY_SESSION"
+    It can be toggled off with the PREFECT_SERVER_ANALYTICS_ENABLED setting.
+    """
+    from prefect.client.constants import SERVER_API_VERSION
+
+    db = provide_database_interface()
+    session_start_timestamp, session_id = await _fetch_or_set_telemetry_session(db=db)
+    telemetry_environment = os.environ.get(
+        "PREFECT_API_TELEMETRY_ENVIRONMENT", "production"
+    )
+
+    heartbeat = {
+        "source": "prefect_server",
+        "type": "heartbeat",
+        "payload": {
+            "platform": platform.system(),
+            "architecture": platform.machine(),
+            "python_version": platform.python_version(),
+            "python_implementation": platform.python_implementation(),
+            "environment": telemetry_environment,
+            "ephemeral_server": bool(os.getenv("PREFECT__SERVER_EPHEMERAL", False)),
+            "api_version": SERVER_API_VERSION,
+            "prefect_version": prefect.__version__,
+            "session_id": session_id,
+            "session_start_timestamp": session_start_timestamp,
+        },
+    }
+
+    try:
+        async with httpx.AsyncClient() as client:
+            result = await client.post(
+                "https://sens-o-matic.prefect.io/",
+                json=heartbeat,
+                headers={"x-prefect-event": "prefect_server"},
             )
-
-            if telemetry_session is None:
-                self.logger.debug("No telemetry session found, setting")
-                session_id = str(uuid4())
-                session_start_timestamp = now("UTC").isoformat()
-
-                telemetry_session = Configuration(
-                    key="TELEMETRY_SESSION",
-                    value={
-                        "session_id": session_id,
-                        "session_start_timestamp": session_start_timestamp,
-                    },
-                )
-
-                await configuration.write_configuration(session, telemetry_session)
-
-                self.session_id = session_id
-                self.session_start_timestamp = session_start_timestamp
-            else:
-                self.logger.debug("Session information retrieved from database")
-                self.session_id: str = telemetry_session.value["session_id"]
-                self.session_start_timestamp: str = telemetry_session.value[
-                    "session_start_timestamp"
-                ]
-        self.logger.debug(
-            f"Telemetry Session: {self.session_id}, {self.session_start_timestamp}"
+        result.raise_for_status()
+    except Exception as exc:
+        logger.error(
+            f"Failed to send telemetry: {exc}",
+            exc_info=PREFECT_DEBUG_MODE.value(),
         )
-        return (self.session_start_timestamp, self.session_id)
-
-    async def run_once(self) -> None:
-        """
-        Sends a heartbeat to the sens-o-matic
-        """
-        from prefect.client.constants import SERVER_API_VERSION
-
-        if not hasattr(self, "session_id"):
-            await self._fetch_or_set_telemetry_session()
-
-        heartbeat = {
-            "source": "prefect_server",
-            "type": "heartbeat",
-            "payload": {
-                "platform": platform.system(),
-                "architecture": platform.machine(),
-                "python_version": platform.python_version(),
-                "python_implementation": platform.python_implementation(),
-                "environment": self.telemetry_environment,
-                "ephemeral_server": bool(os.getenv("PREFECT__SERVER_EPHEMERAL", False)),
-                "api_version": SERVER_API_VERSION,
-                "prefect_version": prefect.__version__,
-                "session_id": self.session_id,
-                "session_start_timestamp": self.session_start_timestamp,
-            },
-        }
-
-        try:
-            async with httpx.AsyncClient() as client:
-                result = await client.post(
-                    "https://sens-o-matic.prefect.io/",
-                    json=heartbeat,
-                    headers={"x-prefect-event": "prefect_server"},
-                )
-            result.raise_for_status()
-        except Exception as exc:
-            self.logger.error(
-                f"Failed to send telemetry: {exc}\nShutting down telemetry service...",
-                # The traceback is only needed if doing deeper debugging, otherwise
-                # this looks like an impactful server error
-                exc_info=PREFECT_DEBUG_MODE.value(),
-            )
-            await self.stop(block=False)
-
-
-if __name__ == "__main__":
-    asyncio.run(Telemetry(handle_signals=True).start())
diff --git a/tests/cli/test_server_services.py b/tests/cli/test_server_services.py
@@ -126,8 +126,9 @@ def test_list_services(self):
                 "Available Services",
                 "Scheduler",
                 "PREFECT_SERVER_SERVICES_SCHEDULER_ENABLED",
-                "Telemetry",
-                "PREFECT_SERVER_ANALYTICS_ENABLED",
+                "TaskRunRecorder",
+                # May be truncated in table display
+                "PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER",
             ],
             expected_code=0,
         )
diff --git a/tests/server/services/test_perpetual_services.py b/tests/server/services/test_perpetual_services.py
@@ -38,6 +38,23 @@ def test_foreman_service_registered():
     assert "monitor_worker_health" in service_names
 
 
+def test_telemetry_service_registered():
+    """Test that telemetry perpetual service is registered."""
+    service_names = [config.function.__name__ for config in _PERPETUAL_SERVICES]
+    assert "send_telemetry_heartbeat" in service_names
+
+
+def test_telemetry_runs_in_all_modes():
+    """Test that telemetry is configured to run in ephemeral and webserver modes."""
+    config = next(
+        c
+        for c in _PERPETUAL_SERVICES
+        if c.function.__name__ == "send_telemetry_heartbeat"
+    )
+    assert config.run_in_ephemeral is True
+    assert config.run_in_webserver is True
+
+
 def test_get_perpetual_services_returns_all_in_default_mode():
     """Test that get_perpetual_services returns all services in default mode."""
     services = get_perpetual_services(ephemeral=False, webserver_only=False)
diff --git a/tests/server/services/test_service_subsets.py b/tests/server/services/test_service_subsets.py
@@ -7,14 +7,12 @@
 from prefect.server.services.base import RunInEphemeralServers, RunInWebservers, Service
 from prefect.server.services.scheduler import RecentDeploymentsScheduler, Scheduler
 from prefect.server.services.task_run_recorder import TaskRunRecorder
-from prefect.server.services.telemetry import Telemetry
 
 
 def test_the_all_service_subset():
     """The following services should be enabled on background servers or full-featured
     API servers"""
     assert set(Service.all_services()) == {
-        Telemetry,
         # Orchestration services
         RecentDeploymentsScheduler,
         Scheduler,
@@ -34,7 +32,6 @@ def test_the_all_service_subset():
 def test_run_in_ephemeral_servers():
     """The following services should be enabled on ephemeral servers"""
     assert set(RunInEphemeralServers.all_services()) == {
-        Telemetry,
         # Orchestration services
         TaskRunRecorder,
         # Events services
@@ -52,7 +49,6 @@ def test_run_in_ephemeral_servers():
 def test_run_in_webservers():
     """The following services should be enabled on webservers"""
     assert set(RunInWebservers.all_services()) == {
-        Telemetry,
         # Events services
         Distributor,
         # Logs services
diff --git a/tests/server/services/test_telemetry.py b/tests/server/services/test_telemetry.py
@@ -6,7 +6,10 @@
 from httpx import Response
 
 import prefect
-from prefect.server.services.telemetry import Telemetry
+from prefect.server.services.telemetry import (
+    _fetch_or_set_telemetry_session,
+    send_telemetry_heartbeat,
+)
 
 
 @pytest.fixture
@@ -32,8 +35,7 @@ def error_sens_o_matic_mock():
 async def test_sens_o_matic_called_correctly(sens_o_matic_mock):
     from prefect.client.constants import SERVER_API_VERSION
 
-    telemetry = Telemetry()
-    await telemetry.start(loops=1)
+    await send_telemetry_heartbeat()
 
     assert sens_o_matic_mock.called
     assert sens_o_matic_mock.call_count == 1
@@ -54,34 +56,35 @@ async def test_sens_o_matic_called_correctly(sens_o_matic_mock):
     assert payload["api_version"] == SERVER_API_VERSION
     assert payload["prefect_version"] == prefect.__version__
 
-    assert payload["session_id"] == telemetry.session_id
-    assert payload["session_start_timestamp"] == telemetry.session_start_timestamp
+    # Session info should be present
+    assert payload["session_id"]
+    assert payload["session_start_timestamp"]
 
 
-async def test_sets_and_fetches_session_information(sens_o_matic_mock):
-    telemetry = Telemetry()
-    await telemetry.start(loops=1)
+async def test_sets_and_fetches_session_information():
+    from prefect.server.database import provide_database_interface
 
-    # set it on the first call
-    sid = telemetry.session_id
-    sts = telemetry.session_start_timestamp
-    assert sid
-    assert sts
+    db = provide_database_interface()
 
-    # retrieve from the db if process restarted
-    telemetry_2 = Telemetry()
-    await telemetry_2.start(loops=1)
-    assert telemetry_2.session_id == sid
-    assert telemetry_2.session_start_timestamp == sts
+    # Get session info from first call
+    session_start_timestamp_1, session_id_1 = await _fetch_or_set_telemetry_session(
+        db=db
+    )
+    assert session_id_1
+    assert session_start_timestamp_1
 
+    # Retrieve from the db if process restarted (same session info)
+    session_start_timestamp_2, session_id_2 = await _fetch_or_set_telemetry_session(
+        db=db
+    )
+    assert session_id_2 == session_id_1
+    assert session_start_timestamp_2 == session_start_timestamp_1
 
-async def test_errors_shutdown_service(error_sens_o_matic_mock, caplog):
-    # When telemetry encounters an error on any loop the service is stopped
-    telemetry = Telemetry()
 
-    await telemetry.start(loops=5)
+async def test_errors_do_not_crash_service(error_sens_o_matic_mock, caplog):
+    # When telemetry encounters an error, it logs but doesn't crash
+    await send_telemetry_heartbeat()
 
-    # The service should only be hit once
     assert error_sens_o_matic_mock.called
     assert error_sens_o_matic_mock.call_count == 1
 
