konduktor-nightly 0.1.0.dev20250209104336__py3-none-any.whl → 0.1.0.dev20250313070642__py3-none-any.whl

konduktor/utils/log_utils.py

@@ -0,0 +1,251 @@
+# Proprietary Changes made for Trainy under the Trainy Software License
+# Original source: skypilot: https://github.com/skypilot-org/skypilot
+# which is Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import io
+import multiprocessing
+import os
+import subprocess
+import sys
+import types
+from typing import List, Optional, Tuple, Type, Union
+
+import prettytable
+
+from konduktor.utils import subprocess_utils
+
+
+class LineProcessor(object):
+    """A processor for log lines."""
+
+    def __enter__(self) -> None:
+        pass
+
+    def process_line(self, log_line: str) -> None:
+        pass
+
+    def __exit__(
+        self,
+        except_type: Optional[Type[BaseException]],
+        except_value: Optional[BaseException],
+        traceback: Optional[types.TracebackType],
+    ) -> None:
+        del except_type, except_value, traceback  # unused
+        pass
+
+
+class _ProcessingArgs:
+    """Arguments for processing logs."""
+
+    def __init__(
+        self,
+        log_path: str,
+        stream_logs: bool,
+        start_streaming_at: str = '',
+        end_streaming_at: Optional[str] = None,
+        skip_lines: Optional[List[str]] = None,
+        replace_crlf: bool = False,
+        line_processor: Optional[LineProcessor] = None,
+        streaming_prefix: Optional[str] = None,
+    ) -> None:
+        self.log_path = log_path
+        self.stream_logs = stream_logs
+        self.start_streaming_at = start_streaming_at
+        self.end_streaming_at = end_streaming_at
+        self.skip_lines = skip_lines
+        self.replace_crlf = replace_crlf
+        self.line_processor = line_processor
+        self.streaming_prefix = streaming_prefix
+
+
+def _handle_io_stream(io_stream, out_stream, args: _ProcessingArgs):
+    """Process the stream of a process."""
+    out_io = io.TextIOWrapper(
+        io_stream, encoding='utf-8', newline='', errors='replace', write_through=True
+    )
+
+    start_streaming_flag = False
+    end_streaming_flag = False
+    streaming_prefix = args.streaming_prefix if args.streaming_prefix else ''
+    line_processor = (
+        LineProcessor() if args.line_processor is None else args.line_processor
+    )
+
+    out = []
+    with open(args.log_path, 'a', encoding='utf-8') as fout:
+        with line_processor:
+            while True:
+                line = out_io.readline()
+                if not line:
+                    break
+                # start_streaming_at logic in processor.process_line(line)
+                if args.replace_crlf and line.endswith('\r\n'):
+                    # Replace CRLF with LF to avoid ray logging to the same
+                    # line due to separating lines with '\n'.
+                    line = line[:-2] + '\n'
+                if args.skip_lines is not None and any(
+                    skip in line for skip in args.skip_lines
+                ):
+                    continue
+                if args.start_streaming_at in line:
+                    start_streaming_flag = True
+                if args.end_streaming_at is not None and args.end_streaming_at in line:
+                    # Keep executing the loop, only stop streaming.
+                    # E.g., this is used for `sky bench` to hide the
+                    # redundant messages of `sky launch` while
+                    # saving them in log files.
+                    end_streaming_flag = True
+                if args.stream_logs and start_streaming_flag and not end_streaming_flag:
+                    print(streaming_prefix + line, end='', file=out_stream, flush=True)
+                if args.log_path != '/dev/null':
+                    fout.write(line)
+                    fout.flush()
+                line_processor.process_line(line)
+                out.append(line)
+    return ''.join(out)
+
+
+def process_subprocess_stream(proc, args: _ProcessingArgs) -> Tuple[str, str]:
+    """Redirect the process's filtered stdout/stderr to both stream and file"""
+    if proc.stderr is not None:
+        # Asyncio does not work as the output processing can be executed in a
+        # different thread.
+        # selectors is possible to handle the multiplexing of stdout/stderr,
+        # but it introduces buffering making the output not streaming.
+        with multiprocessing.pool.ThreadPool(processes=1) as pool:
+            err_args = copy.copy(args)
+            err_args.line_processor = None
+            stderr_fut = pool.apply_async(
+                _handle_io_stream, args=(proc.stderr, sys.stderr, err_args)
+            )
+            # Do not launch a thread for stdout as the rich.status does not
+            # work in a thread, which is used in
+            # log_utils.RayUpLineProcessor.
+            stdout = _handle_io_stream(proc.stdout, sys.stdout, args)
+            stderr = stderr_fut.get()
+    else:
+        stdout = _handle_io_stream(proc.stdout, sys.stdout, args)
+        stderr = ''
+    return stdout, stderr
+
+
+def create_table(field_names: List[str], **kwargs) -> prettytable.PrettyTable:
+    """Creates table with default style."""
+    border = kwargs.pop('border', False)
+    align = kwargs.pop('align', 'l')
+    table = prettytable.PrettyTable(
+        align=align, border=border, field_names=field_names, **kwargs
+    )
+    table.left_padding_width = 0
+    table.right_padding_width = 2
+    return table
+
+
+def run_with_log(
+    cmd: Union[List[str], str],
+    log_path: str,
+    *,
+    require_outputs: bool = False,
+    stream_logs: bool = False,
+    start_streaming_at: str = '',
+    end_streaming_at: Optional[str] = None,
+    skip_lines: Optional[List[str]] = None,
+    shell: bool = False,
+    with_ray: bool = False,
+    process_stream: bool = True,
+    line_processor: Optional[LineProcessor] = None,
+    streaming_prefix: Optional[str] = None,
+    **kwargs,
+) -> Union[int, Tuple[int, str, str]]:
+    """Runs a command and logs its output to a file.
+
+    Args:
+        cmd: The command to run.
+        log_path: The path to the log file.
+        stream_logs: Whether to stream the logs to stdout/stderr.
+        require_outputs: Whether to return the stdout/stderr of the command.
+        process_stream: Whether to post-process the stdout/stderr of the
+            command, such as replacing or skipping lines on the fly. If
+            enabled, lines are printed only when '\r' or '\n' is found.
+
+    Returns the returncode or returncode, stdout and stderr of the command.
+      Note that the stdout and stderr is already decoded.
+    """
+    assert process_stream or not require_outputs, (
+        process_stream,
+        require_outputs,
+        'require_outputs should be False when process_stream is False',
+    )
+
+    log_path = os.path.expanduser(log_path)
+    dirname = os.path.dirname(log_path)
+    os.makedirs(dirname, exist_ok=True)
+    # Redirect stderr to stdout when using ray, to preserve the order of
+    # stdout and stderr.
+    stdout_arg = stderr_arg = None
+    if process_stream:
+        stdout_arg = subprocess.PIPE
+        stderr_arg = subprocess.PIPE if not with_ray else subprocess.STDOUT
+    # Use stdin=subprocess.DEVNULL by default, as allowing inputs will mess up
+    # the terminal output when typing in the terminal that starts the API
+    # server.
+    stdin = kwargs.pop('stdin', subprocess.DEVNULL)
+    with subprocess.Popen(
+        cmd,
+        stdout=stdout_arg,
+        stderr=stderr_arg,
+        start_new_session=True,
+        shell=shell,
+        stdin=stdin,
+        **kwargs,
+    ) as proc:
+        try:
+            subprocess_utils.kill_process_daemon(proc.pid)
+            stdout = ''
+            stderr = ''
+
+            if process_stream:
+                if skip_lines is None:
+                    skip_lines = []
+                # Skip these lines caused by `-i` option of bash. Failed to
+                # find other way to turn off these two warning.
+                # https://stackoverflow.com/questions/13300764/how-to-tell-bash-not-to-issue-warnings-cannot-set-terminal-process-group-and # noqa: E501
+                # `ssh -T -i -tt` still cause the problem.
+                skip_lines += [
+                    'bash: cannot set terminal process group',
+                    'bash: no job control in this shell',
+                ]
+                # We need this even if the log_path is '/dev/null' to ensure the
+                # progress bar is shown.
+                # NOTE: Lines are printed only when '\r' or '\n' is found.
+                args = _ProcessingArgs(
+                    log_path=log_path,
+                    stream_logs=stream_logs,
+                    start_streaming_at=start_streaming_at,
+                    end_streaming_at=end_streaming_at,
+                    skip_lines=skip_lines,
+                    line_processor=line_processor,
+                    # Replace CRLF when the output is logged to driver by ray.
+                    replace_crlf=with_ray,
+                    streaming_prefix=streaming_prefix,
+                )
+                stdout, stderr = process_subprocess_stream(proc, args)
+            proc.wait()
+            if require_outputs:
+                return proc.returncode, stdout, stderr
+            return proc.returncode
+        except KeyboardInterrupt:
+            # Kill the subprocess directly, otherwise, the underlying
+            # process will only be killed after the python program exits,
+            # causing the stream handling stuck at `readline`.
+            subprocess_utils.kill_children_processes()
+            raise

konduktor/utils/loki_utils.py

@@ -0,0 +1,85 @@
+"""Loki utils: query/tail logs from Loki"""
+# TODO(asaiacai): eventually support querying
+# centralized loki that lives outside the cluster
+
+import asyncio
+import json
+import urllib.parse
+
+import colorama
+import kr8s
+import websockets
+
+from konduktor import logging
+
+logger = logging.get_logger(__name__)
+
+LOKI_PORT = 3100
+WEBSOCKET_TIMEOUT = 10
+INFINITY = 999999
+
+
+async def _read_loki_logs(loki_url: str, timeout: int, job_name: str, worker_id: int):
+    ws = await asyncio.wait_for(websockets.connect(loki_url), timeout=WEBSOCKET_TIMEOUT)
+    logger.info(
+        f'{colorama.Fore.YELLOW}Tailing logs from Loki. '
+        f'Forwarding to port {LOKI_PORT}. Press Ctrl+C to stop. '
+        f'{colorama.Style.RESET_ALL}'
+    )
+    try:
+        while True:
+            message = await asyncio.wait_for(ws.recv(), timeout=timeout)
+            try:
+                payload = json.loads(message)
+                for stream in payload['streams']:
+                    if stream['values'][0][1] is not None:
+                        print(
+                            f"{colorama.Fore.CYAN}{colorama.Style.BRIGHT} "
+                            f"(job_name={job_name} worker_id={worker_id})"
+                            f"{colorama.Style.RESET_ALL} {stream['values'][0][1]}",
+                            flush=True,
+                        )
+            except json.JSONDecodeError:
+                logger.warning(f'Failed to decode log skipping: {message}')
+                logger.debug(f'Dropped log: {message}')
+                continue
+    except asyncio.exceptions.TimeoutError:
+        logger.debug('Websocket timed-out, closing the connection!')
+
+
+def tail_loki_logs_ws(
+    job_name: str, worker_id: int = 0, num_logs: int = 1000, follow: bool = True
+):
+    if num_logs > 5000:
+        # TODO(asaiacai): we should not have a limit on the number of logs, but rather
+        # let the user specify any number of lines, and we can print the last N lines.
+        # this can be done in chunks. Potentially, we can query range
+        # until we reach the end of the log and then invoke tail again.
+        # Also include checks that the job is running/ever ran.
+        raise ValueError('num_logs must be less than or equal to 5000')
+    loki_url = f'ws://localhost:{LOKI_PORT}/loki/api/v1/tail'
+    params = {
+        'query': urllib.parse.quote(
+            f'{{k8s_job_name="{job_name}-workers-0"}} '
+            f' | batch_kubernetes_io_job_completion_index = `{worker_id}`'
+        ),
+        'limit': num_logs,
+        'delay': 5,
+        # TODO(asaiacai): need to auto-generate the start and end times.
+    }
+
+    query_string = '&'.join(f'{key}={value}' for key, value in params.items())
+    loki_url += f'?{query_string}'
+
+    loki_svc = kr8s.objects.Service.get('loki', namespace='loki')
+    timeout = INFINITY if follow else WEBSOCKET_TIMEOUT
+    with kr8s.portforward.PortForward(loki_svc, LOKI_PORT):
+        asyncio.run(_read_loki_logs(loki_url, timeout, job_name, worker_id))
+
+
+# TODO(asaiacai): write a query_range function to get all the
+# logs for a job for not tailing option
+
+# Run the WebSocket log tailing function
+if __name__ == '__main__':
+    tail_loki_logs_ws('tune-c3c8', worker_id=0, follow=False)

konduktor/utils/rich_utils.py

@@ -0,0 +1,123 @@
+# Proprietary Changes made for Trainy under the Trainy Software License
+# Original source: skypilot: https://github.com/skypilot-org/skypilot
+# which is Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Rich status spinner utils."""
+
+import contextlib
+import logging
+import threading
+from typing import Union
+
+import rich.console as rich_console
+
+console = rich_console.Console(soft_wrap=True)
+_status = None
+_status_nesting_level = 0
+
+_logging_lock = threading.RLock()
+
+
+class _NoOpConsoleStatus:
+    """An empty class for multi-threaded console.status."""
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    def update(self, text):
+        pass
+
+    def stop(self):
+        pass
+
+    def start(self):
+        pass
+
+
+class _RevertibleStatus:
+    """A wrapper for status that can revert to previous message after exit."""
+
+    def __init__(self, message: str):
+        if _status is not None:
+            self.previous_message = _status.status
+        else:
+            self.previous_message = None
+        self.message = message
+
+    def __enter__(self):
+        global _status_nesting_level
+        _status.update(self.message)
+        _status_nesting_level += 1
+        _status.__enter__()
+        return _status
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        global _status_nesting_level, _status
+        _status_nesting_level -= 1
+        if _status_nesting_level <= 0:
+            _status_nesting_level = 0
+            if _status is not None:
+                _status.__exit__(exc_type, exc_val, exc_tb)
+                _status = None
+        else:
+            _status.update(self.previous_message)
+
+    def update(self, *args, **kwargs):
+        _status.update(*args, **kwargs)
+
+    def stop(self):
+        _status.stop()
+
+    def start(self):
+        _status.start()
+
+
+@contextlib.contextmanager
+def safe_logger():
+    logged = False
+    with _logging_lock:
+        if _status is not None and _status._live.is_started:  # pylint: disable=protected-access
+            _status.stop()
+            yield
+            logged = True
+            _status.start()
+    if not logged:
+        yield
+
+
+class RichSafeStreamHandler(logging.StreamHandler):
+    def emit(self, record: logging.LogRecord) -> None:
+        with safe_logger():
+            return super().emit(record)
+
+
+def force_update_status(msg: str):
+    """Update the status message even if sky_logging.is_silent() is true."""
+    if threading.current_thread() is threading.main_thread() and _status is not None:
+        _status.update(msg)
+
+
+def safe_status(msg: str) -> Union['rich_console.Status', _NoOpConsoleStatus]:
+    """A wrapper for multi-threaded console.status."""
+    from konduktor import logging  # pylint: disable=import-outside-toplevel
+
+    global _status
+    if (
+        threading.current_thread() is threading.main_thread()
+        and not logging.is_silent()
+    ):
+        if _status is None:
+            _status = console.status(msg, refresh_per_second=8)
+        return _RevertibleStatus(msg)
+    return _NoOpConsoleStatus()