Keep one build's telemetry login from deciding other builds' runs

README.md

@@ -49,8 +49,11 @@ until you restart the kernel, load it by path with
 Supported US and UK build commands always start the local telemetry emitter
 service. It reports live progress to the hosted collector when the operator's
 existing Hugging Face login is accepted; otherwise it retains the events
-locally and the build continues. This live event path is independent of the
-staging run files described below.
+locally and the build continues. Every build on a machine shares one local
+queue. A build's login decides only that build's own run. A run whose service
+has exited is delivered later by another build's service whose login the
+collector accepts for it. This live event path is independent of the staging
+run files described below.
 
 US fiscal refresh builds also write pre-release staging run files **by
 default**. Progress JSON is uploaded to `policyengine/populace-us-staging`

changelog.d/telemetry-foreign-run-local-only.fixed.md

@@ -0,0 +1 @@
+A build's telemetry service no longer disables hosted telemetry for other builds on the same machine. Every service shares one local spool. Before, a service with no Hugging Face login, or with a login the collector refused, permanently marked every other build's queued run local-only. Now a service touches another build's run only after that build's service has exited, and only one service does so at a time. Each service holds an advisory lock file next to the spool. A service that cannot take its lock does not start, and that build continues without telemetry. A run with no lock file, from an older checkout, waits until it has been idle for 15 minutes. A missing, refused or non-owning login decides only the service's own run; another build's run stays pending for a service whose login the collector accepts for it, and a service that was refused for it offers it again once its own login changes. A refused login is sent to the collector once, not on every pass, and the shutdown drain waits only for the build's own run.

packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/constants.py

@@ -42,6 +42,11 @@
     "events (HTTP {status}), so they will not be retried. The dataset build will "
     "continue."
 )
+OWN_LEASE_UNAVAILABLE_MESSAGE: Final = (
+    "Microcosm telemetry is off for this run: its service could not take the "
+    "lock that marks the run as live, and without it another build's service "
+    "could upload the run under its own login. The dataset build will continue."
+)
 COLLECTOR_URL_HTTPS_ERROR: Final = "collector URL must be an HTTPS origin"
 COLLECTOR_URL_ORIGIN_ERROR: Final = (
     "collector URL must be an origin without credentials or path data"
@@ -60,6 +65,17 @@
 DATABASE_TIMEOUT_SECONDS: Final = 5
 PRUNE_INTERVAL_SECONDS: Final = 60.0
 
+LEASE_DIRECTORY_SUFFIX: Final = ".leases"
+LEASE_FILE_SUFFIX: Final = ".lock"
+LEASE_RETRY_SECONDS: Final = 0.01
+MAX_LEASE_OPEN_ATTEMPTS: Final = 8
+OWN_LEASE_TIMEOUT_SECONDS: Final = 1.0
+LEASE_SWEEP_INTERVAL_SECONDS: Final = 600.0
+#: A run with no lease file (its producer predates leases, or its free lease was
+#: swept) counts as orphaned once it has gone this long without an update. A
+#: live service appends a heartbeat every ``DEFAULT_HEARTBEAT_SECONDS``.
+ORPHAN_IDLE_SECONDS: Final = 900.0
+
 DEFAULT_HEARTBEAT_SECONDS: Final = 60.0
 DEFAULT_DRAIN_SECONDS: Final = 15.0
 MINIMUM_HEARTBEAT_SECONDS: Final = 1.0

packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/leases.py

@@ -0,0 +1,226 @@
+"""Advisory leases that tell a live producer's run from an orphaned one.
+
+Every emitter service on a host shares one spool, so the spool also holds other
+builds' runs. A run belongs to the service that produced it while that service
+lives. The collector gives a run to the Hugging Face user who registers it
+first, so another service delivering a live run would register it under its own
+credential, and the collector's answer to that credential says nothing about
+the run's own.
+
+Each service therefore holds an exclusive ``flock`` on its producer's lease file
+for its whole life. The kernel drops the lock when the process exits by any
+route, SIGKILL included, so a lease that can be taken means the producing
+service is gone and its run may be adopted. The lease files live beside the
+spool, outside the database, so no spool migration is needed and checkouts that
+predate leases can still share the spool.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import time
+from collections.abc import Callable, Iterable
+from pathlib import Path
+
+try:
+    import fcntl
+except ImportError:  # pragma: no cover - the service needs Unix sockets anyway
+    fcntl = None
+
+from microcosm.build.telemetry_emitter_service.constants import (
+    LEASE_DIRECTORY_SUFFIX,
+    LEASE_FILE_SUFFIX,
+    LEASE_RETRY_SECONDS,
+    MAX_LEASE_OPEN_ATTEMPTS,
+)
+
+
+class OwnLeaseUnavailableError(RuntimeError):
+    """A service could not take its own producer's lease."""
+
+
+class ProducerLease:
+    """An exclusive lock on one producer's lease file, held until released."""
+
+    def __init__(self, path: Path, descriptor: int) -> None:
+        self.path = path
+        self._descriptor: int | None = descriptor
+
+    @property
+    def held(self) -> bool:
+        return self._descriptor is not None
+
+    def release(self, *, remove: bool = False) -> None:
+        """Drop the lock, first unlinking the file when ``remove`` is set.
+
+        The file is unlinked while the lock is still held, so no other process
+        can be holding it then. Anyone who opened the old file meanwhile finds,
+        once it gets the lock, that the path no longer names that file, and
+        tries again.
+        """
+
+        if self._descriptor is None:
+            return
+        try:
+            if remove:
+                self.path.unlink(missing_ok=True)
+        finally:
+            os.close(self._descriptor)
+            self._descriptor = None
+
+    def __enter__(self) -> ProducerLease:
+        return self
+
+    def __exit__(self, *exc_info: object) -> None:
+        self.release()
+
+
+class ProducerLeases:
+    """The lease files of every producer that shares one spool."""
+
+    def __init__(self, directory: Path | str) -> None:
+        self.directory = Path(directory)
+
+    @classmethod
+    def beside(cls, spool_path: Path | str) -> ProducerLeases:
+        """Return the leases for the spool at ``spool_path``."""
+
+        spool_path = Path(spool_path)
+        return cls(spool_path.with_name(spool_path.name + LEASE_DIRECTORY_SUFFIX))
+
+    @property
+    def supported(self) -> bool:
+        """Whether this platform has advisory file locks."""
+
+        return fcntl is not None
+
+    def path(self, run_id: str, producer_id: str) -> Path:
+        """Return the lease file of one producer.
+
+        Identifiers come from the stored registration, so the file name is a
+        digest of them rather than the identifiers themselves.
+        """
+
+        digest = hashlib.sha256(f"{run_id}\0{producer_id}".encode()).hexdigest()
+        return self.directory / f"{digest}{LEASE_FILE_SUFFIX}"
+
+    def exists(self, run_id: str, producer_id: str) -> bool:
+        return self.path(run_id, producer_id).exists()
+
+    def try_acquire(
+        self,
+        run_id: str,
+        producer_id: str,
+        *,
+        create: bool,
+    ) -> ProducerLease | None:
+        """Take a producer's lease without waiting.
+
+        Returns ``None`` while another process holds it. With ``create`` unset,
+        a missing lease file raises ``FileNotFoundError`` rather than creating
+        one, so the caller can tell a producer that never had a lease (an older
+        checkout) from one that has gone. Any other failure to test the lease
+        raises ``OSError``.
+        """
+
+        if fcntl is None:
+            raise OSError("advisory file locks are not available")
+        if create:
+            self.directory.mkdir(mode=0o700, parents=True, exist_ok=True)
+        return _lock(self.path(run_id, producer_id), create=create)
+
+    def hold(
+        self,
+        run_id: str,
+        producer_id: str,
+        *,
+        timeout_seconds: float,
+        clock: Callable[[], float] = time.monotonic,
+        sleep: Callable[[float], None] = time.sleep,
+    ) -> ProducerLease | None:
+        """Take this service's own lease, waiting out brief holders and errors.
+
+        Only a sweeper or an adopter holds another producer's lease, and only
+        for a moment, and a lock the system refuses once may be granted on the
+        next attempt, so a short wait suffices. Returns ``None`` if the lease is
+        still not held at the deadline, or on a platform without advisory
+        locks. A failed attempt can leave a lease file nobody holds, which
+        reads as an exited producer's, so the caller must not serve its run
+        without the lease.
+        """
+
+        if fcntl is None:
+            return None
+        deadline = clock() + timeout_seconds
+        while True:
+            try:
+                lease = self.try_acquire(run_id, producer_id, create=True)
+            except OSError:
+                lease = None
+            if lease is not None or clock() >= deadline:
+                return lease
+            sleep(LEASE_RETRY_SECONDS)
+
+    def sweep(self, keep: Iterable[tuple[str, str]]) -> int:
+        """Remove free lease files of producers that have nothing pending.
+
+        A held lease belongs to a live service and is never removed, nor is the
+        lease of a run in ``keep`` (the pending runs), so an orphan stays
+        adoptable at once. Returns the number of files removed.
+        """
+
+        if fcntl is None or not self.directory.is_dir():
+            return 0
+        kept = {self.path(run_id, producer_id).name for run_id, producer_id in keep}
+        removed = 0
+        for path in self.directory.glob(f"*{LEASE_FILE_SUFFIX}"):
+            if path.name in kept:
+                continue
+            try:
+                lease = _lock(path, create=False)
+            except OSError:
+                continue
+            if lease is not None:
+                lease.release(remove=True)
+                removed += 1
+        return removed
+
+
+def _lock(path: Path, *, create: bool) -> ProducerLease | None:
+    flags = os.O_RDWR | getattr(os, "O_CLOEXEC", 0)
+    if create:
+        flags |= os.O_CREAT
+    for _ in range(MAX_LEASE_OPEN_ATTEMPTS):
+        descriptor = os.open(path, flags, 0o600)
+        try:
+            fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except BlockingIOError:
+            os.close(descriptor)
+            return None
+        except BaseException:
+            os.close(descriptor)
+            raise
+        try:
+            current = _names_open_file(path, descriptor)
+        except BaseException:
+            # Never leave a lock held on a descriptor nobody will close.
+            os.close(descriptor)
+            raise
+        if current:
+            return ProducerLease(path, descriptor)
+        # Another process removed this file between our open and our lock.
+        os.close(descriptor)
+        if not create and not path.exists():
+            raise FileNotFoundError(path)
+    # The path keeps being replaced under us: someone is taking it.
+    return None
+
+
+def _names_open_file(path: Path, descriptor: int) -> bool:
+    try:
+        on_disk = path.stat()
+    except FileNotFoundError:
+        return False
+    opened = os.fstat(descriptor)
+    return (on_disk.st_dev, on_disk.st_ino) == (opened.st_dev, opened.st_ino)

packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/main.py

@@ -4,12 +4,15 @@
 
 import argparse
 import json
+import sys
 from pathlib import Path
 
 from microcosm.build.telemetry_emitter_service.collector import CollectorDelivery
 from microcosm.build.telemetry_emitter_service.constants import (
     DEFAULT_HEARTBEAT_SECONDS,
+    OWN_LEASE_UNAVAILABLE_MESSAGE,
 )
+from microcosm.build.telemetry_emitter_service.leases import OwnLeaseUnavailableError
 from microcosm.build.telemetry_emitter_service.resources import ProcessTreeSampler
 from microcosm.build.telemetry_emitter_service.runtime import EmitterService
 from microcosm.build.telemetry_emitter_service.spool import EventSpool
@@ -38,16 +41,28 @@ def main(argv: list[str] | None = None) -> int:
     args = build_parser().parse_args(argv)
     registration = json.loads(args.registration_json)
     spool = EventSpool(args.spool)
+    # The delivery takes this producer's lease, which must happen before the
+    # service is ready: from then on the run can have events. Without the
+    # lease the service does not start, and the build goes on without it.
+    try:
+        delivery = CollectorDelivery(
+            spool,
+            registration,
+            development_collector_url=args.development_collector_url,
+        )
+    except OwnLeaseUnavailableError:
+        print(OWN_LEASE_UNAVAILABLE_MESSAGE, file=sys.stderr, flush=True)
+        return 1
     service = EmitterService(
         socket_path=args.socket,
         registration=registration,
         spool=spool,
-        delivery=CollectorDelivery(
-            spool,
-            development_collector_url=args.development_collector_url,
-        ),
+        delivery=delivery,
         sampler=ProcessTreeSampler(args.parent_pid),
         heartbeat_seconds=args.heartbeat_seconds,
     )
-    service.run()
+    try:
+        service.run()
+    finally:
+        delivery.close()
     return 0

packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/runtime.py

@@ -211,8 +211,15 @@ def _worker(self) -> None:
         self._drain()
 
     def _drain(self) -> None:
+        # Only this build's run holds the service up: other builds' runs are
+        # their own services' to deliver, or are adopted while this one runs.
+        run_id = str(self.registration["run_id"])
+        producer_id = str(self.registration["producer_id"])
         deadline = time.monotonic() + self.drain_seconds
-        while self.spool.has_deliverable() and time.monotonic() < deadline:
+        while (
+            self.spool.run_has_deliverable(run_id, producer_id)
+            and time.monotonic() < deadline
+        ):
             if not self.delivery.flush_once():
                 remaining = deadline - time.monotonic()
                 if remaining > 0:

packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/spool.py

@@ -170,6 +170,39 @@ def has_deliverable(self) -> bool:
         with self._lock, self._session_factory() as session:
             return session.scalar(statement) is not None
 
+    def run_has_deliverable(self, run_id: str, producer_id: str) -> bool:
+        """Return whether one producer still has events eligible for delivery."""
+
+        statement = (
+            select(TelemetryEventRecord.event_id)
+            .join(TelemetryEventRecord.run)
+            .where(
+                TelemetryRunRecord.run_id == run_id,
+                TelemetryRunRecord.producer_id == producer_id,
+                TelemetryRunRecord.upload_state == UPLOAD_STATE_PENDING,
+            )
+            .limit(1)
+        )
+        with self._lock, self._session_factory() as session:
+            return session.scalar(statement) is not None
+
+    def last_updated(self, run_id: str, producer_id: str) -> datetime | None:
+        """Return when one producer's run last changed, or ``None`` if unknown."""
+
+        statement = select(TelemetryRunRecord.updated_at).where(
+            TelemetryRunRecord.run_id == run_id,
+            TelemetryRunRecord.producer_id == producer_id,
+        )
+        with self._lock, self._session_factory() as session:
+            updated_at = session.scalar(statement)
+        if updated_at is None:
+            return None
+        try:
+            parsed = datetime.fromisoformat(updated_at)
+        except ValueError:
+            return None
+        return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=UTC)
+
     def batch(
         self,
         run_id: str,