Merge remote-tracking branch 'origin/andrei/telem' into andrei/telem

src/exo/main.py

@@ -24,6 +24,7 @@ from exo.routing.router import Router, get_node_id_keypair
 from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_LOG, EXO_PID_FILE
 from exo.shared.election import Election, ElectionResult
 from exo.shared.logging import logger_cleanup, logger_setup
+from exo.shared.telemetry import TelemetryService
 from exo.shared.types.common import NodeId, SessionId
 from exo.utils import STDIO_FDS
 from exo.utils.channels import Receiver, channel
@@ -42,6 +43,7 @@ class Node:
     election_result_receiver: Receiver[ElectionResult]
     master: Master | None
     api: API | None
+    telemetry: TelemetryService
 
     node_id: NodeId
     offline: bool
@@ -70,6 +72,7 @@ class Node:
             external_outbound=router.sender(topics.LOCAL_EVENTS),
             external_inbound=router.receiver(topics.GLOBAL_EVENTS),
         )
+        telemetry = TelemetryService.create(telemetry_disabled=not args.telemetry)
 
         logger.info(f"Starting node {node_id}")
 
@@ -108,6 +111,7 @@ class Node:
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 api_port=args.api_port,
+                telemetry_sink=telemetry.sink(),
             )
         else:
             worker = None
@@ -146,6 +150,7 @@ class Node:
             er_recv,
             master,
             api,
+            telemetry,
             node_id,
             args.offline,
             args.api_port,
@@ -157,6 +162,7 @@ class Node:
             signal.signal(signal.SIGTERM, lambda _, __: self.shutdown())
             tg.start_soon(self.router.run)
             tg.start_soon(self.event_router.run)
+            tg.start_soon(self.telemetry.run)
             tg.start_soon(self.election.run)
             if self.download_coordinator:
                 tg.start_soon(self.download_coordinator.run)
@@ -264,6 +270,7 @@ class Node:
                                 topics.DOWNLOAD_COMMANDS
                             ),
                             api_port=self._api_port,
+                            telemetry_sink=self.telemetry.sink(),
                         )
                         self._tg.start_soon(self.worker.run)
                     if self.api:
@@ -384,6 +391,7 @@ class Args(FrozenModel):
     no_downloads: bool = False
     offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
     no_batch: bool = False
+    telemetry: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
     legacy_daemon: bool = False
     bootstrap_peers: list[str] = []
@@ -445,6 +453,11 @@ class Args(FrozenModel):
             action="store_true",
             help="Disable continuous batching, use sequential generation",
         )
+        parser.add_argument(
+            "--telemetry",
+            action="store_true",
+            help="Enable telemetry uploads. Disabled by default; disabled mode keeps telemetry in dry-run.",
+        )
         parser.add_argument(
             "--legacy-daemon",
             action="store_true",

src/exo/shared/constants.py

@@ -26,6 +26,11 @@ EXO_CONFIG_HOME = _get_xdg_dir("XDG_CONFIG_HOME", ".config")
 EXO_DATA_HOME = _get_xdg_dir("XDG_DATA_HOME", ".local/share")
 EXO_CACHE_HOME = _get_xdg_dir("XDG_CACHE_HOME", ".cache")
 
+# Exo website API endpoints
+EXO_TELEMETRY_API_URL = os.environ.get(
+    "EXO_TELEMETRY_API_URL", "https://telemetry.exolabs.net/"
+)
+
 # Default models directory (always included as first entry in writable dirs)
 _EXO_DEFAULT_MODELS_DIR_ENV = os.environ.get("EXO_DEFAULT_MODELS_DIR", None)
 EXO_DEFAULT_MODELS_DIR = (
@@ -69,8 +74,6 @@ DASHBOARD_DIR = (
 EXO_LOG_DIR = EXO_CACHE_HOME / "exo_log"
 EXO_LOG = EXO_LOG_DIR / "exo.log"
 EXO_RUNNER_LOG_DIR = EXO_LOG_DIR / "runner_log"
-EXO_RUNNER_STDOUT_LOG = EXO_RUNNER_LOG_DIR / "stdout.log"
-EXO_RUNNER_STDERR_LOG = EXO_RUNNER_LOG_DIR / "stderr.log"
 
 EXO_TEST_LOG = EXO_CACHE_HOME / "exo_test.log"
 EXO_PID_FILE = EXO_CACHE_HOME / "exo.pid"

src/exo/shared/telemetry.py

@@ -0,0 +1,176 @@
+import contextlib
+import hashlib
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Self
+from urllib.parse import urlparse
+
+import httpx
+from anyio import BrokenResourceError, ClosedResourceError, WouldBlock, to_thread
+from loguru import logger
+
+from exo.shared.constants import EXO_TELEMETRY_API_URL
+from exo.utils.channels import Receiver, Sender, channel
+from exo.utils.pydantic_ext import FrozenModel, TaggedModel
+from exo.utils.task_group import TaskGroup
+
+CHANNEL_BOUND_SIZE = 64
+TELEMETRY_HTTP_TIMEOUT_SECONDS = 10.0
+
+
+class BaseTelemetrySubmission(TaggedModel):
+    pass
+
+
+class TestSubmission(BaseTelemetrySubmission):
+    pass
+
+
+class RunnerStderrSubmission(BaseTelemetrySubmission):
+    path: Path
+
+
+TelemetrySubmission = TestSubmission | RunnerStderrSubmission
+
+
+class TelemetryPresignResponse(FrozenModel):
+    key: str
+    upload_url: str
+    expires_in: int
+    max_size: int
+
+
+@dataclass(eq=False)
+class TelemetrySink:
+    """
+    A non-blocking non-throwing bounded wrapper around sender/receiver channels
+    to ensure telemetry never blocks or has adverse side-effects, since telemetry
+    is an optional diagnostic feature and hence should never break the main app.
+    """
+
+    _send: Sender[TelemetrySubmission]
+
+    @classmethod
+    def pair(cls) -> tuple[Self, Receiver[TelemetrySubmission]]:
+        send, recv = channel[TelemetrySubmission](CHANNEL_BOUND_SIZE)
+        return cls(_send=send), recv
+
+    def submit(self, submission: TelemetrySubmission):
+        try:
+            self._send.send_nowait(submission)
+        except WouldBlock:
+            logger.debug("Telemetry submission would block. why so many submissions??")
+        except (BrokenResourceError, ClosedResourceError):
+            logger.debug("Telemetry submission receivers are broken or closed. why??")
+
+    def clone(self) -> "TelemetrySink":
+        return TelemetrySink(_send=self._send.clone())
+
+    def close(self):
+        with contextlib.suppress(BrokenResourceError, ClosedResourceError):
+            self._send.close()
+
+
+@dataclass(eq=False)
+class TelemetryService:
+    telemetry_disabled: bool
+    api_url: str
+    _send: Sender[TelemetrySubmission]
+    _recv: Receiver[TelemetrySubmission]
+    _http_transport: httpx.AsyncBaseTransport | None
+    _tg: TaskGroup = field(default_factory=TaskGroup, init=False)
+
+    @classmethod
+    def create(
+        cls,
+        telemetry_disabled: bool,
+        api_url: str = EXO_TELEMETRY_API_URL,
+        http_transport: httpx.AsyncBaseTransport | None = None,
+    ) -> Self:
+        api_url = urlparse(api_url).geturl().rstrip("/")
+
+        send, recv = channel[TelemetrySubmission](CHANNEL_BOUND_SIZE)
+
+        return cls(
+            telemetry_disabled=telemetry_disabled,
+            api_url=api_url,
+            _send=send,
+            _recv=recv,
+            _http_transport=http_transport,
+        )
+
+    @classmethod
+    def dummy(cls) -> Self:
+        return cls.create(True)
+
+    async def run(self):
+        try:
+            async with self._tg as tg:
+                tg.start_soon(self._process)
+        finally:
+            self._send.close()
+            self._recv.close()
+
+    async def _process(self):
+        with self._recv as submissions:
+            async for submission in submissions:
+                if not self.telemetry_disabled:
+                    try:
+                        await self._process_submission(submission)
+                    except Exception as e:
+                        logger.opt(exception=e).warning(
+                            "Exception when processing telemetry submission"
+                        )
+
+    async def _process_submission(self, submission: TelemetrySubmission):
+        match submission:
+            case TestSubmission():
+                pass
+            case RunnerStderrSubmission(path=path):
+                await self._submit_runner_stderr(path)
+
+    async def _submit_runner_stderr(self, path: Path):
+        data = await to_thread.run_sync(path.read_bytes)
+        if not data:
+            logger.debug(f"Skipping empty runner stderr telemetry file: {path}")
+            return
+
+        sha256 = hashlib.sha256(data).hexdigest()
+
+        async with httpx.AsyncClient(
+            timeout=TELEMETRY_HTTP_TIMEOUT_SECONDS,
+            transport=self._http_transport,
+        ) as client:
+            presign_response = await client.post(
+                f"{self.api_url}/telemetry/runner-log/presign",
+                json={
+                    "sha256": sha256,
+                    "size": len(data),
+                },
+            )
+            presign_response.raise_for_status()
+            presign = TelemetryPresignResponse.model_validate_json(
+                presign_response.text,
+            )
+
+            upload_response = await client.put(
+                presign.upload_url,
+                content=data,
+            )
+            upload_response.raise_for_status()
+
+    def sink(self) -> TelemetrySink:
+        sink, recv = TelemetrySink.pair()
+        if self._tg.is_running():
+            self._tg.start_soon(self._ingest, recv)
+        else:
+            self._tg.queue(self._ingest, recv)
+        return sink
+
+    async def _ingest(self, recv: Receiver[TelemetrySubmission]):
+        try:
+            with recv as submissions:
+                async for submission in submissions:
+                    await self._send.send(submission)
+        except ClosedResourceError:
+            pass

src/exo/shared/tests/test_telemetry.py

@@ -0,0 +1,95 @@
+import hashlib
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+import httpx
+import pytest
+
+from exo.shared.telemetry import RunnerStderrSubmission, TelemetryService
+
+
+@dataclass(frozen=True)
+class RecordedRequest:
+    method: str
+    url: str
+    content: bytes
+
+
+def _queue_submission(
+    service: TelemetryService,
+    submission: RunnerStderrSubmission,
+) -> None:
+    service._send.send_nowait(submission)  # pyright: ignore[reportPrivateUsage]
+    service._send.close()  # pyright: ignore[reportPrivateUsage]
+
+
+@pytest.mark.anyio
+async def test_runner_stderr_upload_hashes_and_uploads_file_bytes(tmp_path: Path):
+    log_bytes = b"runner stderr\nsecond line\n"
+    log_path = tmp_path / "runner.stderr.log"
+    log_path.write_bytes(log_bytes)
+    requests: list[RecordedRequest] = []
+
+    async def handler(request: httpx.Request) -> httpx.Response:
+        requests.append(
+            RecordedRequest(
+                method=request.method,
+                url=str(request.url),
+                content=await request.aread(),
+            )
+        )
+        if request.method == "POST":
+            return httpx.Response(
+                200,
+                json={
+                    "key": "runner_log/test.stderr.log",
+                    "uploadUrl": "https://uploads.example/runner.stderr.log",
+                    "expiresIn": 300,
+                    "maxSize": 52428800,
+                },
+            )
+        if request.method == "PUT":
+            return httpx.Response(200)
+        return httpx.Response(404)
+
+    service = TelemetryService.create(
+        telemetry_disabled=False,
+        api_url="https://telemetry.example/",
+        http_transport=httpx.MockTransport(handler),
+    )
+
+    await service._process_submission(  # pyright: ignore[reportPrivateUsage]
+        RunnerStderrSubmission(path=log_path)
+    )
+
+    assert [r.method for r in requests] == ["POST", "PUT"]
+    assert requests[0].url == "https://telemetry.example/telemetry/runner-log/presign"
+    assert json.loads(requests[0].content) == {
+        "sha256": hashlib.sha256(log_bytes).hexdigest(),
+        "size": len(log_bytes),
+    }
+    assert requests[1].url == "https://uploads.example/runner.stderr.log"
+    assert requests[1].content == log_bytes
+
+
+@pytest.mark.anyio
+async def test_runner_stderr_upload_failure_is_swallowed(tmp_path: Path):
+    log_path = tmp_path / "runner.stderr.log"
+    log_path.write_text("runner stderr\n")
+    requests: list[httpx.Request] = []
+
+    async def handler(request: httpx.Request) -> httpx.Response:
+        requests.append(request)
+        return httpx.Response(500)
+
+    service = TelemetryService.create(
+        telemetry_disabled=False,
+        api_url="https://telemetry.example",
+        http_transport=httpx.MockTransport(handler),
+    )
+    _queue_submission(service, RunnerStderrSubmission(path=log_path))
+
+    await service._process()  # pyright: ignore[reportPrivateUsage]
+
+    assert len(requests) == 1

src/exo/worker/main.py

@@ -15,6 +15,7 @@ from exo.routing.event_router import (
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
 from exo.shared.models.model_cards import ModelId, card_cache
+from exo.shared.telemetry import TelemetrySink
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.commands import (
     DeleteInstance,
@@ -74,6 +75,7 @@ class Worker:
         command_sender: Sender[ForwarderCommand],
         download_command_sender: Sender[ForwarderDownloadCommand],
         api_port: int,
+        telemetry_sink: TelemetrySink,
     ):
         self.node_id: NodeId = node_id
         self.event_receiver = event_receiver
@@ -81,6 +83,7 @@ class Worker:
         self.command_sender = command_sender
         self.download_command_sender = download_command_sender
         self.api_port = api_port
+        self.telemetry_sink = telemetry_sink
 
         self.state: State = State()
         self.runners: dict[RunnerId, RunnerSupervisor] = {}
@@ -122,6 +125,7 @@ class Worker:
             self.event_sender.close()
             self.command_sender.close()
             self.download_command_sender.close()
+            self.telemetry_sink.close()
             for runner in self.runners.values():
                 runner.shutdown()
             self._stopped.set()
@@ -381,6 +385,7 @@ class Worker:
         runner = await RunnerSupervisor.create(
             bound_instance=task.bound_instance,
             event_sender=self.event_sender.clone(),
+            telemetry_sink=self.telemetry_sink.clone(),
         )
         self.runners[task.bound_instance.bound_runner_id] = runner
         self._tg.start_soon(runner.run)

src/exo/worker/runner/supervisor.py

@@ -2,8 +2,9 @@ import codecs
 import contextlib
 import signal
 from dataclasses import dataclass, field
-from os import PathLike
-from typing import Callable, Self
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Self
 
 import anyio
 from anyio import (
@@ -12,9 +13,13 @@ from anyio import (
     CancelScope,
     ClosedResourceError,
 )
+from anyio.streams.text import TextReceiveStream
 from loguru import logger
 
-from exo.shared.constants import EXO_RUNNER_STDERR_LOG, EXO_RUNNER_STDOUT_LOG
+from exo.shared.constants import (
+    EXO_RUNNER_LOG_DIR,
+)
+from exo.shared.telemetry import RunnerStderrSubmission, TelemetrySink
 from exo.shared.types.chunks import ErrorChunk
 from exo.shared.types.events import (
     ChunkGenerated,
@@ -60,10 +65,12 @@ DECODE_TIMEOUT_SECONDS = 5
 
 @dataclass(eq=False)
 class RunnerStdioHandler:
+    _bound_instance: BoundInstance
     _stdout_rx: Receiver[bytes]
     _stderr_rx: Receiver[bytes]
-    _stdout_log: AsyncFile[str]
+    _stderr_log_path: Path
     _stderr_log: AsyncFile[str]
+    _telemetry: TelemetrySink
     diagnostics: RunnerDiagnosticCollector = field(
         default_factory=RunnerDiagnosticCollector
     )
@@ -74,58 +81,68 @@ class RunnerStdioHandler:
     async def create(
         cls,
         *,
+        bound_instance: BoundInstance,
         stdout_rx: Receiver[bytes],
         stderr_rx: Receiver[bytes],
-        stdout_log_path: PathLike[str] = EXO_RUNNER_STDOUT_LOG,
-        stderr_log_path: PathLike[str] = EXO_RUNNER_STDERR_LOG,
+        telemetry_sink: TelemetrySink,
+        runner_log_dir: Path = EXO_RUNNER_LOG_DIR,
     ) -> Self:
-        # these are append only logs used to gather data for log template mining
-        #
-        # TODO: in the future use [Drain3](https://github.com/logpai/Drain3)
-        #       to mine these logs
-        ensure_parent_directory_exists(stdout_log_path)
+        # create file in <log_dir>/<instanceID>/<runnerID>/<timestamp>.stderr.log
+        now = datetime.now(timezone.utc).strftime("%Y-%m-%d_%H-%M-%S_%fZ")
+        stderr_log_path = (
+            runner_log_dir
+            / bound_instance.instance.instance_id
+            / bound_instance.bound_runner_id
+            / f"{now}.stderr.log"
+        )
         ensure_parent_directory_exists(stderr_log_path)
-        stdout_log = await anyio.open_file(stdout_log_path, "a")
-        stderr_log = await anyio.open_file(stderr_log_path, "a")
+        stderr_log = await anyio.open_file(stderr_log_path, "w")
 
-        # instantiate and return
         self = cls(
+            _bound_instance=bound_instance,
             _stdout_rx=stdout_rx,
             _stderr_rx=stderr_rx,
-            _stdout_log=stdout_log,
+            _stderr_log_path=stderr_log_path,
             _stderr_log=stderr_log,
+            _telemetry=telemetry_sink,
         )
         return self
 
     async def run(self):
         try:
             async with self._tg as tg:
-                tg.start_soon(  # pyright: ignore[reportUnknownArgumentType]
-                    self._handle_runner_output,
-                    self._stdout_rx,
-                    self._stdout_log,
-                    lambda line: logger.info(f"Runner stdout: {line}"),  # pyright: ignore[reportUnknownLambdaType]
-                    lambda _: None,  # pyright: ignore[reportUnknownLambdaType]
-                )
-                tg.start_soon(  # pyright: ignore[reportUnknownArgumentType]
-                    self._handle_runner_output,
-                    self._stderr_rx,
-                    self._stderr_log,
-                    lambda line: logger.warning(f"Runner stderr: {line}"),  # pyright: ignore[reportUnknownLambdaType]
-                    self.diagnostics.record_line,
-                )
+                tg.start_soon(self._handle_stdout)
+                tg.start_soon(self._handle_stderr)
         finally:
             with CancelScope(shield=True):
-                await self._stdout_log.aclose()
                 await self._stderr_log.aclose()
 
-    async def _handle_runner_output(
-        self,
-        rx: Receiver[bytes],
-        logfile: AsyncFile[str],
-        log_line: Callable[[str], None],
-        record_diagnostic_line: Callable[[str], None],
-    ):
+                # send off telemetry submission when runner stdio dies;
+                # it may have been for entirely innocuous reasons or
+                # the log may have nothing in it, but its submitted regardless
+                self._telemetry.submit(
+                    RunnerStderrSubmission(
+                        path=self._stderr_log_path,
+                    )
+                )
+                self._telemetry.close()
+
+    def shutdown(self):
+        self._tg.cancel_tasks()
+
+    async def _handle_stdout(self):
+        # We don't expect anything in stdout so even reading this at all is going
+        # to be quite weird; hence handle it by logging error and the received chunk
+
+        rx = TextReceiveStream(self._stdout_rx, encoding="utf-8", errors="replace")
+        try:
+            async with rx:
+                async for chunk in rx:
+                    logger.warning(f"Unexpected runner stdout chunk: {chunk}")
+        except (ClosedResourceError, BrokenResourceError):
+            logger.warning("Runner stdio stream closed before clean EOF")
+
+    async def _handle_stderr(self):
         # The diagnostic collector is deliberately line-level for now. It records
         # bounded stderr context and known failure anchors; the supervisor
         # correlates those hints with the runner exit status before surfacing an
@@ -142,8 +159,8 @@ class RunnerStdioHandler:
                 return
 
             # Send to logger & error recovery task
-            log_line(line)
-            record_diagnostic_line(line)
+            logger.warning(f"Runner stderr: {line}")
+            self.diagnostics.record_line(line)
 
         async def handle_text(text: str):
             nonlocal pending_line
@@ -151,8 +168,8 @@ class RunnerStdioHandler:
             if not text:
                 return
 
-            await logfile.write(text)
-            await logfile.flush()
+            await self._stderr_log.write(text)
+            await self._stderr_log.flush()
 
             # newline buffering
             pending_line += text
@@ -163,15 +180,15 @@ class RunnerStdioHandler:
                 await handle_line(line)
 
         try:
-            with rx:
-                async for chunk in rx:
+            with self._stderr_rx:
+                async for chunk in self._stderr_rx:
                     await handle_text(decoder.decode(chunk, final=False))
         except (ClosedResourceError, BrokenResourceError):
             logger.warning("Runner stdio stream closed before clean EOF")
         finally:
             with CancelScope(shield=True):
                 await handle_text(decoder.decode(b"", final=True))
-                await logfile.flush()
+                await self._stderr_log.flush()
 
                 if pending_line:
                     await handle_line(pending_line)
@@ -205,6 +222,7 @@ class RunnerSupervisor:
         *,
         bound_instance: BoundInstance,
         event_sender: Sender[Event],
+        telemetry_sink: TelemetrySink,
         initialize_timeout: float = 400,
     ) -> Self:
         ev_send, ev_recv = mp_channel[Event | RunnerTerminationError]()
@@ -223,7 +241,10 @@ class RunnerSupervisor:
             daemon=True,
         )
         runner_stdio_handler = await RunnerStdioHandler.create(
-            stdout_rx=runner_process.stdout, stderr_rx=runner_process.stderr
+            bound_instance=bound_instance,
+            stdout_rx=runner_process.stdout,
+            stderr_rx=runner_process.stderr,
+            telemetry_sink=telemetry_sink,
         )
 
         shard_metadata = bound_instance.bound_shard

src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py

@@ -1,9 +1,11 @@
+from pathlib import Path
 from typing import cast
 
 import anyio
 import pytest
 
 from exo.shared.models.model_cards import ModelId
+from exo.shared.telemetry import TelemetryService
 from exo.shared.types.chunks import ErrorChunk
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.events import ChunkGenerated, Event, RunnerStatusUpdated
@@ -36,7 +38,9 @@ class _DeadProcess:
 
 
 @pytest.mark.anyio
-async def test_check_runner_emits_error_chunk_for_inflight_text_generation() -> None:
+async def test_check_runner_emits_error_chunk_for_inflight_text_generation(
+    tmp_path: Path,
+) -> None:
     event_sender, event_receiver = channel[Event]()
     task_sender, _ = mp_channel[Task]()
     cancel_sender, _ = mp_channel[TaskId]()
@@ -50,8 +54,13 @@ async def test_check_runner_emits_error_chunk_for_inflight_text_generation() ->
     )
 
     proc = cast(AsyncProcess, cast(object, _DeadProcess()))
+    telemetry = TelemetryService.dummy()
     handler = await RunnerStdioHandler.create(
-        stdout_rx=proc.stdout, stderr_rx=proc.stderr
+        bound_instance=bound_instance,
+        stdout_rx=proc.stdout,
+        stderr_rx=proc.stderr,
+        telemetry_sink=telemetry.sink(),
+        runner_log_dir=tmp_path,
     )
     supervisor = RunnerSupervisor(
         shard_metadata=bound_instance.bound_shard,