nnlogging 0.1.0__py3-none-any.whl

nnlogging/__init__.py

@@ -0,0 +1,7 @@
+"""nnlog"""
+
+from . import shell
+from . import utils
+
+
+__all__ = ["shell", "utils"]

nnlogging/shell.py

@@ -0,0 +1,307 @@
+"""nnlog.shell"""
+
+import logging
+from collections.abc import Mapping
+from logging import Handler, Logger
+
+from aim import Run
+from aim.sdk.types import AimObject
+from rich.console import Console as RichConsole
+from rich.console import ConsoleRenderable
+from rich.progress import Progress as RichProgress
+from rich.progress import TaskID
+
+from .utils import (
+    get_aim_run,
+    get_rich_handler,
+    get_rich_progress,
+)
+
+__all__ = ["Shell"]
+
+
+class Shell:
+    name: str = "nnlog"
+
+    def __init__(
+        self,
+        name: str | None = None,
+        level: int = logging.INFO,
+        propagate: bool = False,
+    ):
+        self.name = name or self.name
+        self.level: int = level
+        self.propagate: bool = propagate
+
+        self.logger: Logger = logging.getLogger(self.name)
+        self.consoles: dict[str, RichConsole] = dict()
+        self.handlers: dict[str, Handler] = dict()
+        self.progresses: dict[str, RichProgress] = dict()
+        self.tasks: dict[str, dict[str, TaskID]] = dict()
+        self.aim_run: Run | None = None
+
+        self.configure_logger()
+
+    def configure_logger(
+        self,
+        level: int | None = None,
+        propagate: bool | None = None,
+    ):
+        if level is not None:
+            self.level = level
+        if propagate is not None:
+            self.propagate = propagate
+        self.logger.setLevel(self.level)
+        self.logger.propagate = self.propagate
+
+    def add_console(
+        self,
+        consoles: Mapping[str, RichConsole],
+    ):
+        self.consoles.update(consoles)
+
+    def add_handler(
+        self,
+        handlers: Mapping[str, Handler],
+    ):
+        self.handlers.update(handlers)
+
+    def add_progress(
+        self,
+        progresses: Mapping[str, RichProgress],
+    ):
+        self.progresses.update(progresses)
+
+    def build_handler_from_console(
+        self,
+        *names: str,
+        **kwargs,
+    ):
+        if len(names) == 0:
+            names = tuple(n for n in self.consoles.keys())
+
+        handlers = {
+            name: get_rich_handler(
+                console=self.consoles[name],
+                name=name,
+                **kwargs,
+            )
+            for name in names
+        }
+        self.add_handler(handlers)
+        self.sync_handlers()
+
+    def build_progress_from_console(
+        self,
+        *names: str,
+        **kwargs,
+    ):
+        if len(names) == 0:
+            names = tuple(n for n in self.consoles.keys())
+
+        progresses = {
+            name: get_rich_progress(
+                console=self.consoles[name],
+                **kwargs,
+            )
+            for name in names
+        }
+        self.add_progress(progresses)
+
+    def sync_handlers(self):
+        self.logger.handlers.clear()
+        for hdlr in self.handlers.values():
+            self.logger.addHandler(hdlr)
+
+    def add_task(
+        self,
+        task_name: str,
+        description: str,
+        start: bool = True,
+        total: float | None = 100,
+        completed: float = 0,
+    ):
+        self.tasks[task_name] = dict()
+        for progress_name in self.progresses.keys():
+            task_id = self.progresses[progress_name].add_task(
+                description=description,
+                start=start,
+                total=total,
+                completed=completed,  # pyright: ignore[reportArgumentType]
+            )
+            self.tasks[task_name].update({progress_name: task_id})
+
+    def start_progress(self, *names: str):
+        if len(names) == 0:
+            names = tuple(n for n in self.progresses.keys())
+
+        for name in names:
+            self.progresses[name].start()
+
+    def stop_progress(self, *names: str):
+        if len(names) == 0:
+            names = tuple(n for n in self.progresses.keys())
+
+        for progress_name in names:
+            self.progresses[progress_name].stop()
+            del self.progresses[progress_name]
+            for task_name in self.tasks.keys():
+                del self.tasks[task_name][progress_name]
+
+    def update_task(
+        self,
+        task_name: str,
+        advance: float | None = None,
+        total: float | None = None,
+        completed: float | None = None,
+    ):
+        for progress_name, progress in self.progresses.items():
+            task_id = self.tasks[task_name][progress_name]
+            progress.update(
+                task_id=task_id,
+                advance=advance,
+                total=total,
+                completed=completed,
+            )
+
+    def log(
+        self,
+        level: int,
+        msg: object,
+        *args: object,
+        stacklevel: int = 2,
+        **kwargs,
+    ):
+        if isinstance(msg, ConsoleRenderable):
+            self.debug(f"Parsing {msg} to consoles ...")
+            for console in self.consoles.values():
+                console.print(msg, *args, **kwargs)
+            return
+
+        self.logger.log(
+            level=level,
+            msg=msg,
+            *args,
+            stacklevel=stacklevel,
+            **kwargs,
+        )
+
+    def debug(
+        self,
+        msg: object,
+        *args: object,
+        stacklevel: int = 3,
+        **kwargs,
+    ):
+        self.log(
+            logging.DEBUG,
+            msg,
+            *args,
+            stacklevel=stacklevel,
+            **kwargs,
+        )
+
+    def info(
+        self,
+        msg: object,
+        *args: object,
+        stacklevel: int = 3,
+        **kwargs,
+    ):
+        self.log(
+            logging.INFO,
+            msg,
+            *args,
+            stacklevel=stacklevel,
+            **kwargs,
+        )
+
+    def warn(
+        self,
+        msg: object,
+        *args: object,
+        stacklevel: int = 3,
+        **kwargs,
+    ):
+        self.log(
+            logging.WARN,
+            msg,
+            *args,
+            stacklevel=stacklevel,
+            **kwargs,
+        )
+
+    def error(
+        self,
+        msg: object,
+        *args: object,
+        stacklevel: int = 3,
+        **kwargs,
+    ):
+        self.log(
+            logging.ERROR,
+            msg,
+            *args,
+            stacklevel=stacklevel,
+            **kwargs,
+        )
+
+    def exception(
+        self,
+        msg: object,
+        *args,
+        exc_info: bool = True,
+        stacklevel: int = 3,
+        **kwargs,
+    ):
+        self.log(
+            logging.ERROR,
+            msg,
+            *args,
+            exc_info=exc_info,
+            stacklevel=stacklevel,
+            **kwargs,
+        )
+
+    def add_aimrun(
+        self,
+        **kwargs,
+    ):
+        self.aim_run = get_aim_run(**kwargs)
+
+    def update_aimrun_metadata(
+        self,
+        label: str,
+        metadata: dict[str, object],
+    ):
+        if self.aim_run is None:
+            self.error(
+                "`aim_run` is not initialized",
+                stacklevel=4,
+                exc_info=ValueError(f"`aim_run` is {self.aim_run}"),
+            )
+        else:
+            self.aim_run[label] = metadata
+
+    def track(
+        self,
+        value: object,
+        name: str | None = None,
+        step: int | None = None,
+        epoch: int | None = None,
+        context: AimObject | None = None,
+    ):
+        if self.aim_run is None:
+            self.error(
+                "`aim_run` is not initialized",
+                stacklevel=4,
+                exc_info=ValueError(f"`aim_run` is {self.aim_run}"),
+            )
+        else:
+            self.aim_run.track(
+                value,
+                name=name,
+                step=step,
+                epoch=epoch,
+                context=context,
+            )

nnlogging/utils.py

@@ -0,0 +1,212 @@
+"""nnlog.utils"""
+
+import logging
+from collections.abc import Collection
+from datetime import datetime
+from os import PathLike
+from types import NoneType
+from typing import Callable, Literal, TextIO
+
+from rich.console import Console as RichConsole
+from rich.highlighter import Highlighter
+from rich.progress import ProgressColumn
+from rich.text import Text
+from aim import Repo, Run
+
+type FormatTimeCallable = Callable[[datetime], Text]
+type StrPath = str | PathLike[str]
+
+
+__all__ = [
+    "get_rich_console",
+    "get_rich_handler",
+    "get_rich_progress",
+    "get_aim_run",
+]
+
+
+def get_rich_console(
+    sink: TextIO | None = None,
+    width: int | None = None,
+    height: int | None = None,
+    tab_size: int = 4,
+    soft_wrap: bool = True,
+    color_system: Literal["auto", "standard", "truecolor"] | None = "auto",
+    markup: bool = True,
+    highlighter: Highlighter | None = None,
+    record: bool = False,
+):
+    from sys import stderr
+
+    if sink is None:
+        sink = stderr
+
+    if highlighter is not None:
+        highlight = True
+    else:
+        from rich.highlighter import NullHighlighter
+
+        highlight = False
+        highlighter = NullHighlighter()
+
+    console = RichConsole(
+        file=sink,
+        width=width,
+        height=height,
+        tab_size=tab_size,
+        soft_wrap=soft_wrap,
+        color_system=color_system,
+        markup=markup,
+        highlight=highlight,
+        highlighter=highlighter,
+        record=record,
+    )
+
+    return console
+
+
+def get_rich_handler(
+    console: RichConsole | None = None,
+    name: str | None = None,
+    fmt: str | logging.Formatter = "%(message)s",
+    level: str | int = logging.INFO,
+    show_time: bool = True,
+    log_time_format: str | FormatTimeCallable = "[%x %X]",
+    omit_repeated_times: bool = True,
+    show_level: bool = True,
+    show_path: bool = True,
+    enable_link_path: bool = True,
+    markup: bool = True,
+    highlighter: Highlighter | None = None,
+    rich_tracebacks: bool = True,
+    tracebacks_width: int | None = None,
+    tracebacks_code_width: int | None = 88,
+    tracebacks_extra_lines: int = 3,
+    tracebacks_word_wrap: bool = True,
+    tracebacks_max_frames: int = 100,
+    tracebacks_show_locals: bool = False,
+    locals_max_length: int = 10,
+    locals_max_string: int = 80,
+):
+    from rich.logging import RichHandler
+
+    if console is None:
+        console = get_rich_console()
+
+    if isinstance(fmt, str):
+        fmt = logging.Formatter(fmt)
+
+    handler = RichHandler(
+        console=console,
+        level=level,
+        show_time=show_time,
+        log_time_format=log_time_format,
+        omit_repeated_times=omit_repeated_times,
+        show_level=show_level,
+        show_path=show_path,
+        enable_link_path=enable_link_path,
+        markup=markup,
+        highlighter=highlighter,
+        rich_tracebacks=rich_tracebacks,
+        tracebacks_width=tracebacks_width,
+        tracebacks_code_width=tracebacks_code_width,
+        tracebacks_extra_lines=tracebacks_extra_lines,
+        tracebacks_word_wrap=tracebacks_word_wrap,
+        tracebacks_max_frames=tracebacks_max_frames,
+        tracebacks_show_locals=tracebacks_show_locals,
+        locals_max_length=locals_max_length,
+        locals_max_string=locals_max_string,
+    )
+    handler.name = name
+    handler.setFormatter(fmt=fmt)
+    return handler
+
+
+def get_rich_progress(
+    console: RichConsole | None = None,
+    columns: Collection[str | ProgressColumn] | None = None,
+    expand: bool = False,
+    auto_refresh: bool = True,
+    refresh_per_second: float = 10,
+    speed_estimate_period: float = 3600,
+    transient: bool = False,
+    redirect_stderr: bool = True,
+    redirect_stdout: bool = True,
+):
+    from rich.progress import Progress
+
+    def get_default_rich_progress_columns():
+        from rich.progress import (
+            BarColumn,
+            SpinnerColumn,
+            TaskProgressColumn,
+            TextColumn,
+            TimeRemainingColumn,
+        )
+
+        return (
+            SpinnerColumn(
+                spinner_name="dots",
+                style="progress.spinner",
+                finished_text=" ",
+            ),
+            TextColumn(
+                text_format="{task.description}",
+                style="progress.description",
+                justify="left",
+                markup=True,
+            ),
+            BarColumn(bar_width=40),
+            TaskProgressColumn(
+                text_format="{task.percentage:>3.0f}%",
+                text_format_no_percentage="",
+                style="progress.percentage",
+                justify="right",
+                markup=True,
+                show_speed=True,
+            ),
+            TimeRemainingColumn(
+                compact=False,
+                elapsed_when_finished=False,
+            ),
+        )
+
+    if console is None:
+        console = get_rich_console()
+
+    if columns is None:
+        columns = get_default_rich_progress_columns()
+
+    return Progress(
+        *columns,
+        console=console,
+        auto_refresh=auto_refresh,
+        refresh_per_second=refresh_per_second,
+        speed_estimate_period=speed_estimate_period,
+        transient=transient,
+        expand=expand,
+        redirect_stderr=redirect_stderr,
+        redirect_stdout=redirect_stdout,
+    )
+
+
+def get_aim_run(
+    run_hash: str | None = None,
+    repo: str | Repo | None = None,
+    read_only: bool = False,
+    experiment: str | None = None,
+    force_resume: bool = False,
+    system_tracking_interval: float | None = 10,
+    log_system_params: bool = False,
+    capture_terminal_logs: bool = False,
+):
+    return Run(
+        run_hash=run_hash,
+        repo=repo,
+        read_only=read_only,
+        experiment=experiment,
+        force_resume=force_resume,
+        system_tracking_interval=system_tracking_interval,
+        log_system_params=log_system_params,
+        capture_terminal_logs=capture_terminal_logs,
+    )

nnlogging-0.1.0.dist-info/METADATA

@@ -0,0 +1,24 @@
+Metadata-Version: 2.4
+Name: nnlogging
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: <3.13,>=3.12
+Requires-Dist: aim>=3.29.1
+Requires-Dist: rich>=14.1.0
+Provides-Extra: dev
+Requires-Dist: twine>=6.2.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# nnlog
+
+---
+
+## *Workflow Overview*
+
+![workflow](assets/workflow-overview.png)
+
+---
+
+## *Functions Overview*
+
+![functions](assets/functions-overview.png)

nnlogging-0.1.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+nnlogging/__init__.py,sha256=rK6DT6Jtc63odOR8D5xvxRmBtg3EoCkOOC9-2xIaxKg,84
+nnlogging/shell.py,sha256=Y7yv-paD0IS-31_HwnvRA70SY6tE6lERVM2ryc2BgtU,7559
+nnlogging/utils.py,sha256=dgRi4Q9moF6RtyYz_W762o9dqPC8UvhfCr0JwJifppg,5928
+nnlogging-0.1.0.dist-info/METADATA,sha256=NWzHYl_h1frGimfhzW70YwaGDkmCm3Qp1DAZziy5rjE,433
+nnlogging-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+nnlogging-0.1.0.dist-info/RECORD,,

nnlogging-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any