expr_tracker 0.1.0__tar.gz

expr_tracker-0.1.0/.gitignore

@@ -0,0 +1,12 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+.python-version
+wandb/
+.env
+# Virtual environments
+.venv

expr_tracker-0.1.0/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: expr_tracker
+Version: 0.1.0
+Summary: Add your description here
+Author-email: HSPK <whxway@whu.edu.cn>
+Requires-Python: >=3.10
+Requires-Dist: jsonlines>=4.0.0
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: slark>=0.1.27
+Requires-Dist: wandb>=0.21.0

expr_tracker-0.1.0/pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "expr_tracker"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+authors = [{ name = "HSPK", email = "whxway@whu.edu.cn" }]
+requires-python = ">=3.10"
+dependencies = [
+    "jsonlines>=4.0.0",
+ "loguru>=0.7.3",
+ "slark>=0.1.27",
+ "wandb>=0.21.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[dependency-groups]
+dev = [
+    "ipykernel>=6.30.1",
+]

expr_tracker-0.1.0/src/expr_tracker/__init__.py

@@ -0,0 +1,9 @@
+from . import tracker
+from .alert import alert
+
+init = tracker.init
+finish = tracker.finish
+log = tracker.log
+info = tracker.info
+
+__all__ = ["init", "finish", "log", "info", "alert"]

expr_tracker-0.1.0/src/expr_tracker/alert.py

@@ -0,0 +1,55 @@
+from typing import Literal
+
+
+class LarkBackend:
+    def __init__(self, **kwargs):
+        from slark import Lark
+
+        self.lark = Lark(**kwargs)
+
+    def publish_alert(
+        self,
+        title: str,
+        text: str,
+        subtitle: str | None = None,
+        level: Literal["info", "warning", "error"] = "info",
+        traceback: str | None = None,
+    ):
+        if level == "info" or level == "warning":
+            self.lark.webhook.post_success_card(
+                msg=text, title=title, subtitle=subtitle
+            )
+        elif level == "error":
+            self.lark.webhook.post_error_card(
+                msg=text, traceback=traceback or "", title=title, subtitle=subtitle
+            )
+        else:
+            raise ValueError(
+                f"Unsupported level: {level}. Supported levels are 'info', 'warning', and 'error'."
+            )
+
+
+def alert(
+    title: str,
+    text: str,
+    subtitle: str | None = None,
+    traceback: str | None = None,
+    level: Literal["info", "warning", "error"] = "info",
+    backends: list[Literal["lark", "slack", "email"]] = ["lark"],
+):
+    """
+    Send an alert message to specified backends.
+
+    Args:
+        title (str): The title of the alert.
+        text (str): The main content of the alert.
+        subtitle (str | None): Optional subtitle for the alert.
+        level (Literal["info", "warning", "error"]): The severity level of the alert.
+        backends (list[Literal["lark", "slack", "email"]]): List of backends to send the alert to.
+    """
+    for backend in backends:
+        if backend == "lark":
+            lark_backend = LarkBackend()
+            lark_backend.publish_alert(title, text, subtitle, level, traceback)
+        else:
+            raise NotImplementedError(f"Backend '{backend}' is not implemented.")

expr_tracker-0.1.0/src/expr_tracker/jsonl.py

@@ -0,0 +1,64 @@
+from loguru import logger
+from pathlib import Path
+
+
+def custom_json_encoder(obj):
+    if isinstance(obj, Path):
+        return str(obj)
+    raise TypeError(f"Object of type {obj.__class__.__name__} is not JSON serializable")
+
+
+class JsonlTracker:
+    def __init__(self):
+        pass
+
+    def init(
+        self,
+        project: str,
+        name: str | None = None,
+        config: dict | None = None,
+        dir: str | None = None,
+        print_to_screen: bool = False,
+        print_handle=print,
+        **kwargs,
+    ):
+        self.project = project
+        self.name = name
+        if dir is None:
+            dir = "./tracker/jsonl"
+        self.log_dir = Path(dir) / self.project / self.name
+        self.config_fp = self.log_dir / "config.json"
+        self.log_fp = self.log_dir / "metrics.jsonl"
+        self.log_dir.mkdir(parents=True, exist_ok=True)
+        if self.config_fp.exists():
+            logger.warning(
+                f"Config file {self.config_fp} already exists. It will be overwritten."
+            )
+        if config is not None:
+            with open(self.config_fp, "w") as f:
+                import json
+
+                json.dump(config, f, indent=4, default=custom_json_encoder)
+        self.print_to_screen = print_to_screen
+        self.print_handle = print_handle
+        self.current_step = (
+            0 if not self.log_fp.exists() else len(self.log_fp.read_text().splitlines())
+        )
+
+    def log(self, metrics: dict, step: int | None = None):
+        import jsonlines
+
+        if not self.log_fp.parent.exists():
+            self.log_fp.parent.mkdir(parents=True, exist_ok=True)
+
+        if step is not None:
+            self.current_step = step
+        metrics = {"_step": self.current_step, **metrics}
+        with jsonlines.open(self.log_fp, mode="a") as writer:
+            writer.write(metrics)
+        self.current_step += 1
+        if self.print_to_screen:
+            self.print_handle(f"{metrics}")
+
+    def finish(self):
+        pass

expr_tracker-0.1.0/src/expr_tracker/tracker.py

@@ -0,0 +1,155 @@
+from typing import Literal
+from contextvars import ContextVar
+from loguru import logger
+
+_tracker: ContextVar["Tracker | None"] = ContextVar("tracker", default=None)
+
+
+def get_backend(backend: str | object):
+    if isinstance(backend, str):
+        backend = backend.lower()
+        if backend == "wandb":
+            try:
+                import wandb
+
+                return wandb
+            except ImportError:
+                raise ImportError("WandB backend is not installed.")
+        elif backend == "jsonl":
+            from .jsonl import JsonlTracker
+
+            return JsonlTracker()
+        elif backend == "trackio":
+            try:
+                import trackio
+
+                return trackio
+            except ImportError:
+                raise ImportError("TrackIO backend is not installed.")
+    else:
+        return backend
+
+
+class Tracker:
+    def __init__(
+        self,
+        project: str,
+        name: str | None = None,
+        entity: str | None = None,
+        dir: str | None = None,
+        notes: str | None = None,
+        tags: list[str] | None = None,
+        resume: bool | Literal["allow", "never", "must", "auto"] | None = "allow",
+        config: dict | None = None,
+        backends: list[Literal["wandb", "jsonl", "trackio"]] = [
+            "wandb",
+            "trackio",
+            "jsonl",
+        ],
+        backend_kwargs: dict[str, dict] | None = None,
+        **kwargs,
+    ):
+        if kwargs:
+            logger.info(
+                f"Unrecognized keyword arguments passed to Tracker: {kwargs}. "
+                "These will be ignored."
+            )
+        self.backend = {b: get_backend(b) for b in backends}
+        if config is None:
+            config = {}
+        if backend_kwargs is None:
+            backend_kwargs = {}
+        for b, t in self.backend.items():
+            if b == "trackio":
+                t.init(
+                    project=project,
+                    name=name,
+                    config=config.update(
+                        {
+                            "trackio.notes": notes,
+                            "trackio.tags": tags,
+                            "trackio.entity": entity,
+                            "trackio.resume": resume,
+                        }
+                    ),
+                    **backend_kwargs.get(b, {}),
+                )
+            else:
+                t.init(
+                    project=project,
+                    name=name,
+                    entity=entity,
+                    dir=dir,
+                    notes=notes,
+                    tags=tags,
+                    resume=resume,
+                    id=name,
+                    config=config,
+                    **backend_kwargs.get(b, {}),
+                )
+
+
+def init(
+    project: str,
+    name: str | None = None,
+    entity: str | None = None,
+    dir: str | None = None,
+    notes: str | None = None,
+    tags: list[str] | None = None,
+    resume: bool | Literal["allow", "never", "must", "auto"] | None = "allow",
+    config: dict | None = None,
+    backends: list[Literal["wandb", "jsonl", "trackio"]] = ["wandb", "jsonl"],
+    backend_kwargs: dict[str, dict] | None = None,
+    **kwargs,
+) -> Tracker:
+    """
+    Initialize the tracker with the given parameters.
+    """
+    if _tracker.get() is not None:
+        raise RuntimeError("Tracker is already initialized. Call finish() first.")
+
+    tracker = Tracker(
+        project=project,
+        name=name,
+        entity=entity,
+        dir=dir,
+        notes=notes,
+        tags=tags,
+        resume=resume,
+        config=config,
+        backends=backends,
+        backend_kwargs=backend_kwargs,
+        **kwargs,
+    )
+    _tracker.set(tracker)
+    return tracker
+
+
+def log(metrics: dict, step: int | None = None):
+    tracker = _tracker.get()
+    if tracker is None:
+        raise RuntimeError("Tracker is not initialized. Call init() first.")
+    for b in tracker.backend.values():
+        b.log(metrics, step=step)
+
+
+def finish():
+    tracker = _tracker.get()
+    if tracker is None:
+        raise RuntimeError("Tracker is not initialized. Call init() first.")
+    for b in tracker.backend.values():
+        b.finish()
+    _tracker.set(None)
+
+
+def info():
+    tracker = _tracker.get()
+    if tracker is None:
+        raise RuntimeError("Tracker is not initialized. Call init() first.")
+    info = {}
+    for b, t in tracker.backend.items():
+        if b == "wandb":
+            info["wandb"] = {"url": t.run.url}
+        if b == "jsonl":
+            info["jsonl"] = {"log_dir": t.log_dir.as_posix()}
+    return info

expr_tracker-0.1.0/tests/test_tracker.py

@@ -0,0 +1,14 @@
+def test_expr_tracker():
+    import expr_tracker as et
+
+    et.init(
+        project="test_project",
+        name="test_run",
+        dir="./tests/logs",
+        config={"param": 42},
+        print_to_screen=True,
+        backends=["jsonl", "wandb"],
+    )
+    for i in range(5):
+        et.log({"metric": i * 10}, step=i)
+    et.finish()