partial cherry pick of 2f7daac

Author: rlizzo

src/lightning_app/testing/testing.py

@@ -1,25 +1,29 @@
 import asyncio
 import json
+import logging
 import os
 import shutil
 import sys
 import tempfile
 import time
+import traceback
 from contextlib import contextmanager
 from subprocess import Popen
 from time import sleep
-from typing import Any, Callable, Dict, Generator, List, Type
+from typing import Any, Callable, Dict, Generator, List, Optional, Type
 
 import requests
 from lightning_cloud.openapi.rest import ApiException
 from requests import Session
 from rich import print
+from rich.color import ANSI_COLOR_NAMES
 
 from lightning_app import LightningApp, LightningFlow
 from lightning_app.cli.lightning_cli import run_app
 from lightning_app.core.constants import LIGHTNING_CLOUD_PROJECT_ID
 from lightning_app.runners.multiprocess import MultiProcessRuntime
 from lightning_app.testing.config import Config
+from lightning_app.utilities.app_logs import _app_logs_reader
 from lightning_app.utilities.cloud import _get_project
 from lightning_app.utilities.enum import CacheCallsKeys
 from lightning_app.utilities.imports import _is_playwright_available, requires
@@ -31,6 +35,9 @@
     from playwright.sync_api import HttpCredentials, sync_playwright
 
 
+_logger = logging.getLogger(__name__)
+
+
 class LightningTestApp(LightningApp):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -259,20 +266,6 @@ def run_app_in_cloud(app_folder: str, app_name: str = "app.py") -> Generator:
                 var scrollingElement = (document.scrollingElement || document.body);
                 scrollingElement.scrollTop = scrollingElement.scrollHeight;
             }, 200);
-
-            if (!window._logs) {
-                window._logs = [];
-            }
-
-            if (window.logTerminals) {
-                Object.entries(window.logTerminals).forEach(
-                    ([key, value]) => {
-                        window.logTerminals[key]._onLightningWritelnHandler = function (data) {
-                            window._logs = window._logs.concat([data]);
-                        }
-                    }
-                );
-            }
             """
         )
 
@@ -286,8 +279,46 @@ def run_app_in_cloud(app_folder: str, app_name: str = "app.py") -> Generator:
             except (playwright._impl._api_types.Error, playwright._impl._api_types.TimeoutError):
                 pass
 
-        def fetch_logs() -> str:
-            return admin_page.evaluate("window._logs;")
+        client = LightningClient()
+        project = _get_project(client)
+        identifiers = []
+        rich_colors = list(ANSI_COLOR_NAMES)
+
+        def fetch_logs(component_names: Optional[List[str]] = None) -> Generator:
+            """This methods creates websockets connection in threads and returns the logs to the main thread."""
+            app_id = admin_page.url.split("/")[-1]
+
+            if not component_names:
+                works = client.lightningwork_service_list_lightningwork(
+                    project_id=project.project_id,
+                    app_id=app_id,
+                ).lightningworks
+                component_names = ["flow"] + [w.name for w in works]
+
+            def on_error_callback(ws_app, *_):
+                print(traceback.print_exc())
+                ws_app.close()
+
+            colors = {c: rich_colors[i + 1] for i, c in enumerate(component_names)}
+            gen = _app_logs_reader(
+                client=client,
+                project_id=project.project_id,
+                app_id=app_id,
+                component_names=component_names,
+                follow=False,
+                on_error_callback=on_error_callback,
+            )
+            max_length = max(len(c.replace("root.", "")) for c in component_names)
+            for log_event in gen:
+                message = log_event.message
+                identifier = f"{log_event.timestamp}{log_event.message}"
+                if identifier not in identifiers:
+                    date = log_event.timestamp.strftime("%m/%d/%Y %H:%M:%S")
+                    identifiers.append(identifier)
+                    color = colors[log_event.component_name]
+                    padding = (max_length - len(log_event.component_name)) * " "
+                    print(f"[{color}]{log_event.component_name}{padding}[/{color}] {date} {message}")
+                yield message
 
         # 5. Print your application ID
         print(
@@ -300,11 +331,6 @@ def fetch_logs() -> str:
             pass
         finally:
             print("##################################################")
-            printed_logs = []
-            for log in fetch_logs():
-                if log not in printed_logs:
-                    printed_logs.append(log)
-                    print(log.split("[0m")[-1])
             button = admin_page.locator('[data-cy="stop"]')
             try:
                 button.wait_for(timeout=3 * 1000)
@@ -314,8 +340,6 @@ def fetch_logs() -> str:
             context.close()
             browser.close()
 
-            client = LightningClient()
-            project = _get_project(client)
             list_lightningapps = client.lightningapp_instance_service_list_lightningapp_instances(project.project_id)
 
             for lightningapp in list_lightningapps.lightningapps:

src/lightning_app/utilities/app_logs.py

@@ -0,0 +1,140 @@
+import json
+import queue
+import sys
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from json import JSONDecodeError
+from threading import Thread
+from typing import Callable, Iterator, List, Optional
+
+import dateutil.parser
+from websocket import WebSocketApp
+
+from lightning_app.utilities.logs_socket_api import _LightningLogsSocketAPI
+from lightning_app.utilities.network import LightningClient
+
+
+@dataclass
+class _LogEventLabels:
+    app: str
+    container: str
+    filename: str
+    job: str
+    namespace: str
+    node_name: str
+    pod: str
+    stream: Optional[str] = None
+
+
+@dataclass
+class _LogEvent:
+    message: str
+    timestamp: datetime
+    component_name: str
+    labels: _LogEventLabels
+
+    def __ge__(self, other: "_LogEvent") -> bool:
+        return self.timestamp >= other.timestamp
+
+    def __gt__(self, other: "_LogEvent") -> bool:
+        return self.timestamp > other.timestamp
+
+
+def _push_log_events_to_read_queue_callback(component_name: str, read_queue: queue.PriorityQueue):
+    """Pushes _LogEvents from websocket to read_queue.
+
+    Returns callback function used with `on_message_callback` of websocket.WebSocketApp.
+    """
+
+    def callback(ws_app: WebSocketApp, msg: str):
+        # We strongly trust that the contract on API will hold atm :D
+        event_dict = json.loads(msg)
+        labels = _LogEventLabels(**event_dict["labels"])
+
+        if "message" in event_dict:
+            message = event_dict["message"]
+            timestamp = dateutil.parser.isoparse(event_dict["timestamp"])
+            event = _LogEvent(
+                message=message,
+                timestamp=timestamp,
+                component_name=component_name,
+                labels=labels,
+            )
+            read_queue.put(event)
+
+    return callback
+
+
+def _error_callback(ws_app: WebSocketApp, error: Exception):
+    errors = {
+        KeyError: "Malformed log message, missing key",
+        JSONDecodeError: "Malformed log message",
+        TypeError: "Malformed log format",
+        ValueError: "Malformed date format",
+    }
+    print(f"Error while reading logs ({errors.get(type(error), 'Unknown')})", file=sys.stderr)
+    ws_app.close()
+
+
+def _app_logs_reader(
+    client: LightningClient,
+    project_id: str,
+    app_id: str,
+    component_names: List[str],
+    follow: bool,
+    on_error_callback: Optional[Callable] = None,
+) -> Iterator[_LogEvent]:
+
+    read_queue = queue.PriorityQueue()
+    logs_api_client = _LightningLogsSocketAPI(client.api_client)
+
+    # We will use a socket per component
+    log_sockets = [
+        logs_api_client.create_lightning_logs_socket(
+            project_id=project_id,
+            app_id=app_id,
+            component=component_name,
+            on_message_callback=_push_log_events_to_read_queue_callback(component_name, read_queue),
+            on_error_callback=on_error_callback or _error_callback,
+        )
+        for component_name in component_names
+    ]
+
+    # And each socket on separate thread pushing log event to print queue
+    #   run_forever() will run until we close() the connection from outside
+    log_threads = [Thread(target=work.run_forever) for work in log_sockets]
+
+    # Establish connection and begin pushing logs to the print queue
+    for th in log_threads:
+        th.start()
+
+    # Print logs from queue when log event is available
+    user_log_start = "<<< BEGIN USER_RUN_FLOW SECTION >>>"
+    start_timestamp = None
+
+    # Print logs from queue when log event is available
+    try:
+        while True:
+            log_event = read_queue.get(timeout=None if follow else 1.0)
+            if user_log_start in log_event.message:
+                start_timestamp = log_event.timestamp + timedelta(seconds=0.5)
+
+            if start_timestamp and log_event.timestamp > start_timestamp:
+                yield log_event
+
+    except queue.Empty:
+        # Empty is raised by queue.get if timeout is reached. Follow = False case.
+        pass
+
+    except KeyboardInterrupt:
+        # User pressed CTRL+C to exit, we sould respect that
+        pass
+
+    finally:
+        # Close connections - it will cause run_forever() to finish -> thread as finishes aswell
+        for socket in log_sockets:
+            socket.close()
+
+        # Because all socket were closed, we can just wait for threads to finish.
+        for th in log_threads:
+            th.join()

src/lightning_app/utilities/logs_socket_api.py

@@ -0,0 +1,85 @@
+from typing import Callable, Optional
+from urllib.parse import urlparse
+
+from lightning_cloud.openapi import ApiClient, AuthServiceApi, V1LoginRequest
+from websocket import WebSocketApp
+
+from lightning_app.utilities.login import Auth
+
+
+class _LightningLogsSocketAPI:
+    def __init__(self, api_client: ApiClient):
+        self.api_client = api_client
+        self._auth = Auth()
+        self._auth.authenticate()
+        self._auth_service = AuthServiceApi(api_client)
+
+    def _get_api_token(self) -> str:
+        token_resp = self._auth_service.auth_service_login(
+            body=V1LoginRequest(
+                username=self._auth.username,
+                api_key=self._auth.api_key,
+            )
+        )
+        return token_resp.token
+
+    @staticmethod
+    def _socket_url(host: str, project_id: str, app_id: str, token: str, component: str) -> str:
+        return (
+            f"wss://{host}/v1/projects/{project_id}/appinstances/{app_id}/logs?"
+            f"token={token}&component={component}&follow=true"
+        )
+
+    def create_lightning_logs_socket(
+        self,
+        project_id: str,
+        app_id: str,
+        component: str,
+        on_message_callback: Callable[[WebSocketApp, str], None],
+        on_error_callback: Optional[Callable[[Exception, str], None]] = None,
+    ) -> WebSocketApp:
+        """Creates and returns WebSocketApp to listen to lightning app logs.
+
+            .. code-block:: python
+                # Synchronous reading, run_forever() is blocking
+                def print_log_msg(ws_app, msg):
+                    print(msg)
+
+
+                flow_logs_socket = client.create_lightning_logs_socket("project_id", "app_id", "flow", print_log_msg)
+                flow_socket.run_forever()
+            .. code-block:: python
+                # Asynchronous reading (with Threads)
+                def print_log_msg(ws_app, msg):
+                    print(msg)
+
+
+                flow_logs_socket = client.create_lightning_logs_socket("project_id", "app_id", "flow", print_log_msg)
+                work_logs_socket = client.create_lightning_logs_socket("project_id", "app_id", "work_1", print_log_msg)
+                flow_logs_thread = Thread(target=flow_logs_socket.run_forever)
+                work_logs_thread = Thread(target=work_logs_socket.run_forever)
+                flow_logs_thread.start()
+                work_logs_thread.start()
+                # .......
+                flow_logs_socket.close()
+                work_logs_thread.close()
+        Arguments:
+            project_id: Project ID.
+            app_id: Application ID.
+            component: Component name eg flow.
+            on_message_callback: Callback object which is called when received data.
+            on_error_callback: Callback object which is called when we get error.
+        Returns:
+            WebSocketApp of the wanted socket
+        """
+        _token = self._get_api_token()
+        clean_ws_host = urlparse(self.api_client.configuration.host).netloc
+        socket_url = self._socket_url(
+            host=clean_ws_host,
+            project_id=project_id,
+            app_id=app_id,
+            token=_token,
+            component=component,
+        )
+
+        return WebSocketApp(socket_url, on_message=on_message_callback, on_error=on_error_callback)

tests/tests_app/utilities/test_app_logs.py

@@ -0,0 +1,11 @@
+from datetime import datetime
+from unittest.mock import MagicMock
+
+from lightning_app.utilities.app_logs import _LogEvent
+
+
+def test_log_event():
+    event_1 = _LogEvent("", datetime.now(), MagicMock(), MagicMock())
+    event_2 = _LogEvent("", datetime.now(), MagicMock(), MagicMock())
+    assert event_1 < event_2
+    assert event_1 <= event_2