[CLI] Adding opportunity to see basic cluster logs (#14334)

* pinning starsessions

* pinning starsessions

* adding strict back to requirements.txt

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Duplicated

* Basic implementation

* Basic implementation

* Basic implementation

* Basic implementation

* Common things moved to log helpers file

* Decomposing logs reader classes for reusing

* Setting colors for log levels

* Manifest trimming

* Changes added to CHANGELOG

* Prettifications

* Prettifications

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Logs function name change

* Logs function name change

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* attempt to fix the pydanitc import

* Tests + command name fixes

* Extending tests

* Adding limit argument

* Unmerging CI fix

* Unmerging CI fix

* Adding fields for errors

* Adding log level fixed field width

* Adding absent typing + exeptions raising

* Adding socket error logging

* Addressing comments on cluster list function return value

* Addressing comments on adding e2e tests

* Adding version range for arrow package in reqs

* New unit tests

* arrow time parsing callback modified + unit tests

* helpers updated

* helpers updated

* helpers updated

* One more test

* CMD test fix

* CMD test fix

* CMD test fix

* CMD test fix

* CMD test fix

* LightningClient mocking

* Flaky test removed

Co-authored-by: hhsecond <sherin@grid.ai>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>

Author: 4 people

requirements/app/base.txt

@@ -6,3 +6,4 @@ fsspec>=2022.01.0, <=2022.7.1
 s3fs>=2022.1.0, <=2022.7.1
 croniter # for now until we found something more robust.
 traitlets<5.2.0 # Traitlets 5.2.X fails: https://github.com/ipython/traitlets/issues/741
+arrow>=1.2.0, <=1.2.2

src/lightning_app/CHANGELOG.md

@@ -14,6 +14,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Add support for Lightning AI BYOC cluster management ([#13835](https://github.com/Lightning-AI/lightning/pull/13835))
 
 
+- Add support to see Lightning AI BYOC cluster logs ([#14334](https://github.com/Lightning-AI/lightning/pull/14334))
+
+
 - Add support to run Lightning apps on Lightning AI BYOC clusters ([#13894](https://github.com/Lightning-AI/lightning/pull/13894))
 
 

src/lightning_app/cli/cmd_clusters.py

@@ -93,10 +93,14 @@ def create(
 
         click.echo(f"${resp.id} cluster is ${resp.status.phase}")
 
-    def list(self):
+    def get_clusters(self):
         resp = self.api_client.cluster_service_list_clusters(phase_not_in=[V1ClusterState.DELETED])
+        return ClusterList(resp.clusters)
+
+    def list(self):
+        clusters = self.get_clusters()
         console = Console()
-        console.print(ClusterList(resp.clusters).as_table())
+        console.print(clusters.as_table())
 
     def delete(self, cluster_id: str = None, force: bool = False, wait: bool = False):
         if force:

src/lightning_app/cli/lightning_cli.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 from typing import List, Tuple, Union
 
+import arrow
 import click
 import requests
 import rich
@@ -13,6 +14,7 @@
 
 from lightning_app import __version__ as ver
 from lightning_app.cli import cmd_init, cmd_install, cmd_pl_init, cmd_react_ui_init
+from lightning_app.cli.cmd_clusters import AWSClusterManager
 from lightning_app.cli.lightning_cli_create import create
 from lightning_app.cli.lightning_cli_delete import delete
 from lightning_app.cli.lightning_cli_list import get_list
@@ -21,10 +23,12 @@
 from lightning_app.runners.runtime_type import RuntimeType
 from lightning_app.utilities.app_logs import _app_logs_reader
 from lightning_app.utilities.cli_helpers import (
+    _arrow_time_callback,
     _format_input_env_variables,
     _retrieve_application_url_and_available_commands,
 )
 from lightning_app.utilities.cloud import _get_project
+from lightning_app.utilities.cluster_logs import _cluster_logs_reader
 from lightning_app.utilities.enum import OpenAPITags
 from lightning_app.utilities.install_components import register_all_external_components
 from lightning_app.utilities.login import Auth
@@ -141,6 +145,94 @@ def logs(app_name: str, components: List[str], follow: bool) -> None:
         rich.print(f"[{color}]{log_event.component_name}[/{color}] {date} {log_event.message}")
 
 
+@show.group()
+def cluster():
+    """Groups cluster commands inside show."""
+    pass
+
+
+@cluster.command(name="logs")
+@click.argument("cluster_name", required=True)
+@click.option(
+    "--from",
+    "from_time",
+    default="24 hours ago",
+    help="The starting timestamp to query cluster logs from. Human-readable (e.g. '48 hours ago') or ISO 8601 "
+    "(e.g. '2022-08-23 12:34') formats.",
+    callback=_arrow_time_callback,
+)
+@click.option(
+    "--to",
+    "to_time",
+    default="0 seconds ago",
+    callback=_arrow_time_callback,
+    help="The end timestamp / relative time increment to query logs for. This is ignored when following logs (with "
+    "-f/--follow). The same format as --from option has.",
+)
+@click.option("--limit", default=1000, help="The max number of log lines returned.")
+@click.option("-f", "--follow", required=False, is_flag=True, help="Wait for new logs, to exit use CTRL+C.")
+def cluster_logs(cluster_name: str, to_time: arrow.Arrow, from_time: arrow.Arrow, limit: int, follow: bool) -> None:
+    """Show cluster logs.
+
+    Example uses:
+
+        Print cluster logs:
+
+            $ lightning show cluster logs my-cluster
+
+
+        Print cluster logs and wait for new logs:
+
+            $ lightning show cluster logs my-cluster --follow
+
+
+        Print cluster logs, from 48 hours ago to now:
+
+            $ lightning show cluster logs my-cluster --from "48 hours ago"
+
+
+        Print cluster logs, 10 most recent lines:
+
+            $ lightning show cluster logs my-cluster --limit 10
+    """
+
+    client = LightningClient()
+    cluster_manager = AWSClusterManager()
+    existing_cluster_list = cluster_manager.get_clusters()
+
+    clusters = {cluster.name: cluster.id for cluster in existing_cluster_list.clusters}
+
+    if not clusters:
+        raise click.ClickException("You don't have any clusters.")
+
+    if not cluster_name:
+        raise click.ClickException(
+            f"You have not specified any clusters. Please select one of available: [{', '.join(clusters.keys())}]"
+        )
+
+    if cluster_name not in clusters:
+        raise click.ClickException(
+            f"The cluster '{cluster_name}' does not exist."
+            f" Please select one of the following: [{', '.join(clusters.keys())}]"
+        )
+
+    log_reader = _cluster_logs_reader(
+        client=client,
+        cluster_id=clusters[cluster_name],
+        start=from_time.int_timestamp,
+        end=to_time.int_timestamp,
+        limit=limit,
+        follow=follow,
+    )
+
+    colors = {"error": "red", "warn": "yellow", "info": "green"}
+
+    for log_event in log_reader:
+        date = log_event.timestamp.strftime("%m/%d/%Y %H:%M:%S")
+        color = colors.get(log_event.labels.level, "green")
+        rich.print(f"[{color}]{log_event.labels.level:5}[/{color}] {date} {log_event.message.rstrip()}")
+
+
 @_main.command()
 def login():
     """Log in to your lightning.ai account."""

src/lightning_app/utilities/app_logs.py

@@ -1,15 +1,14 @@
 import json
 import queue
-import sys
 from dataclasses import dataclass
-from datetime import datetime, timedelta
-from json import JSONDecodeError
+from datetime import timedelta
 from threading import Thread
 from typing import Callable, Iterator, List, Optional
 
 import dateutil.parser
 from websocket import WebSocketApp
 
+from lightning_app.utilities.log_helpers import _error_callback, _OrderedLogEntry
 from lightning_app.utilities.logs_socket_api import _LightningLogsSocketAPI
 from lightning_app.utilities.network import LightningClient
 
@@ -27,18 +26,10 @@ class _LogEventLabels:
 
 
 @dataclass
-class _LogEvent:
-    message: str
-    timestamp: datetime
+class _LogEvent(_OrderedLogEntry):
     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.
@@ -65,17 +56,6 @@ def callback(ws_app: WebSocketApp, msg: str):
     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,
@@ -127,7 +107,7 @@ def _app_logs_reader(
         pass
 
     except KeyboardInterrupt:
-        # User pressed CTRL+C to exit, we sould respect that
+        # User pressed CTRL+C to exit, we should respect that
         pass
 
     finally:

src/lightning_app/utilities/cli_helpers.py

@@ -1,6 +1,8 @@
 import re
 from typing import Dict, Optional
 
+import arrow
+import click
 import requests
 
 from lightning_app.core.constants import APP_SERVER_PORT
@@ -119,3 +121,15 @@ def _retrieve_application_url_and_available_commands(app_id_or_name_or_url: Opti
                     raise Exception(f"The server didn't process the request properly. Found {resp.json()}")
                 return lightningapp.status.url, _extract_command_from_openapi(resp.json())
     return None, None
+
+
+def _arrow_time_callback(
+    _ctx: "click.core.Context", _param: "click.core.Option", value: str, arw_now=arrow.utcnow()
+) -> arrow.Arrow:
+    try:
+        return arw_now.dehumanize(value)
+    except ValueError:
+        try:
+            return arrow.get(value)
+        except (ValueError, TypeError):
+            raise click.ClickException(f"cannot parse time {value}")

src/lightning_app/utilities/cluster_logs.py

@@ -0,0 +1,108 @@
+import json
+import queue
+from dataclasses import dataclass
+from threading import Thread
+from typing import Callable, Iterator, Optional
+
+import arrow
+import dateutil.parser
+from websocket import WebSocketApp
+
+from lightning_app.utilities.log_helpers import _error_callback, _OrderedLogEntry
+from lightning_app.utilities.logs_socket_api import _ClusterLogsSocketAPI
+from lightning_app.utilities.network import LightningClient
+
+
+@dataclass
+class _ClusterLogEventLabels:
+    cluster_id: str
+    grid_url: str
+    hostname: str
+    level: str
+    logger: str
+    path: Optional[str] = None
+    workspace: Optional[str] = None
+    identifier: Optional[str] = None
+    issuer: Optional[str] = None
+    error: Optional[str] = None
+    errorVerbose: Optional[str] = None
+
+
+@dataclass
+class _ClusterLogEvent(_OrderedLogEntry):
+    labels: _ClusterLogEventLabels
+
+
+def _push_log_events_to_read_queue_callback(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 = _ClusterLogEventLabels(**event_dict["labels"])
+
+        if "message" in event_dict:
+            message = event_dict["message"]
+            timestamp = dateutil.parser.isoparse(event_dict["timestamp"])
+            event = _ClusterLogEvent(
+                message=message,
+                timestamp=timestamp,
+                labels=labels,
+            )
+            read_queue.put(event)
+
+    return callback
+
+
+def _cluster_logs_reader(
+    client: LightningClient,
+    cluster_id: str,
+    start: arrow.Arrow,
+    end: arrow.Arrow,
+    limit: int,
+    follow: bool,
+    on_error_callback: Optional[Callable] = None,
+) -> Iterator[_ClusterLogEvent]:
+
+    logs_api_client = _ClusterLogsSocketAPI(client.api_client)
+    read_queue = queue.PriorityQueue()
+
+    # We will use a socket inside a thread to read logs,
+    # to follow our typical reading pattern
+    log_socket = logs_api_client.create_cluster_logs_socket(
+        cluster_id=cluster_id,
+        start=start,
+        end=end,
+        limit=limit,
+        on_message_callback=_push_log_events_to_read_queue_callback(read_queue),
+        on_error_callback=on_error_callback or _error_callback,
+    )
+
+    log_thread = Thread(target=log_socket.run_forever)
+
+    # Establish connection and begin pushing logs to the print queue
+    log_thread.start()
+
+    # Print logs from queue when log event is available
+    try:
+        while True:
+            log_event = read_queue.get(timeout=None if follow else 1.0)
+            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 should respect that
+        pass
+
+    finally:
+        # Close connection - it will cause run_forever() to finish -> thread as finishes as well
+        log_socket.close()
+
+        # The socket was closed, we can just wait for thread to finish.
+        log_thread.join()

src/lightning_app/utilities/log_helpers.py

@@ -0,0 +1,34 @@
+import logging
+from dataclasses import dataclass
+from datetime import datetime
+from json import JSONDecodeError
+
+from websocket import WebSocketApp
+
+logger = logging.getLogger(__name__)
+
+
+# This is a superclass to inherit log entry classes from it:
+# it implements magic methods to sort logs by timestamps.
+@dataclass
+class _OrderedLogEntry:
+    message: str
+    timestamp: datetime
+
+    def __ge__(self, other: "_OrderedLogEntry") -> bool:
+        return self.timestamp >= other.timestamp
+
+    def __gt__(self, other: "_OrderedLogEntry") -> bool:
+        return self.timestamp > other.timestamp
+
+
+# A general error callback for log reading, prints most common types of possible errors.
+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",
+    }
+    logger.error(f"⚡ Error while reading logs ({errors.get(type(error), 'Unknown')}), {error}")
+    ws_app.close()