liblaf-cherries 0.0.14__py3-none-any.whl → 0.1.0__py3-none-any.whl

liblaf/cherries/__init__.pyi

@@ -1,89 +1,116 @@
-from . import info, pathutils, plugin, presets, utils
-from ._config import BaseConfig
+from . import config, integration, meta, pathutils, presets, utils
 from ._run import end, run, start
+from .config import (
+    AssetKind,
+    BaseConfig,
+    MetaAsset,
+    PathProvider,
+    get_assets,
+    get_inputs,
+    get_outputs,
+    input,  # noqa: A004
+    output,
+)
+from .integration import (
+    AddTag,
+    AddTags,
+    End,
+    Experiment,
+    LogAsset,
+    LogCode,
+    LogMetric,
+    LogMetrics,
+    LogOther,
+    LogOthers,
+    LogParam,
+    LogParams,
+    Plugin,
+    Start,
+    add_tag,
+    add_tags,
+    log_asset,
+    log_code,
+    log_metric,
+    log_metrics,
+    log_other,
+    log_others,
+    log_param,
+    log_params,
+)
 from .pathutils import (
     as_os_path,
     as_path,
     as_posix,
-    config,
     data,
     entrypoint,
+    exp_dir,
     git_root,
     git_root_safe,
     inputs,
     outputs,
     params,
     path,
-    run_dir,
     src,
 )
-from .plugin import (
-    End,
-    LogArtifact,
-    LogArtifacts,
-    LoggingEnd,
-    LoggingStart,
-    LogMetric,
-    LogParam,
-    Plugin,
-    Run,
-    RunStatus,
-    Start,
-    log_artifact,
-    log_artifacts,
-    log_input,
-    log_inputs,
-    log_metric,
-    log_output,
-    log_outputs,
-    log_param,
-    log_src,
-    set_tag,
-)
 
 __all__ = [
+    "AddTag",
+    "AddTags",
+    "AssetKind",
     "BaseConfig",
     "End",
-    "LogArtifact",
-    "LogArtifacts",
+    "Experiment",
+    "LogAsset",
+    "LogCode",
     "LogMetric",
+    "LogMetrics",
+    "LogOther",
+    "LogOthers",
     "LogParam",
-    "LoggingEnd",
-    "LoggingStart",
+    "LogParams",
+    "MetaAsset",
+    "PathProvider",
     "Plugin",
-    "Run",
-    "RunStatus",
     "Start",
+    "add_tag",
+    "add_tags",
     "as_os_path",
     "as_path",
     "as_posix",
     "config",
     "data",
     "end",
+    "end",
+    "end",
     "entrypoint",
+    "exp_dir",
+    "get_assets",
+    "get_inputs",
+    "get_outputs",
     "git_root",
     "git_root_safe",
-    "info",
+    "input",
     "inputs",
-    "log_artifact",
-    "log_artifacts",
-    "log_input",
-    "log_inputs",
+    "integration",
+    "log_asset",
+    "log_code",
     "log_metric",
-    "log_output",
-    "log_outputs",
+    "log_metrics",
+    "log_other",
+    "log_others",
     "log_param",
-    "log_src",
+    "log_params",
+    "meta",
+    "output",
     "outputs",
     "params",
     "path",
     "pathutils",
-    "plugin",
     "presets",
     "run",
-    "run_dir",
-    "set_tag",
+    "run",
     "src",
     "start",
+    "start",
     "utils",
 ]

liblaf/cherries/_run.py

@@ -3,37 +3,49 @@ from typing import get_type_hints
 
 import pydantic
 
+from liblaf.cherries import config, integration, presets
 from liblaf.cherries import pathutils as _path
-from liblaf.cherries import plugin, presets
 
 
-def run[C: pydantic.BaseModel, T](main: Callable[[C], T]) -> T:
-    run: plugin.Run = start()
+def run[C: pydantic.BaseModel, T](
+    main: Callable[[], T] | Callable[[C], T],
+    *,
+    play: bool = False,
+    preset: presets.Preset = presets.default,
+) -> T:
+    exp: integration.Experiment = start(preset=preset, play=play)
     type_hints: dict[str, type[C]] = get_type_hints(main)
-    cls: type[C] = next(iter(type_hints.values()))
-    cfg: C = cls()
-    run.log_param("cherries.config", cfg.model_dump(mode="json"))
+    del type_hints["return"]
+    args: list[C] = []
+    if len(type_hints) == 1:
+        cls: type[C] = next(iter(type_hints.values()))
+        cfg: C = cls()
+        args.append(cfg)
+        for path in config.get_inputs(cfg):
+            exp.log_input(path)
     try:
-        ret: T = main(cfg)
-    except BaseException as e:
-        if isinstance(e, KeyboardInterrupt):
-            run.end(plugin.RunStatus.KILLED)
-            raise
-        run.end(plugin.RunStatus.FAILED)
-        raise
-    else:
-        run.end()
-        return ret
-
-
-def start() -> plugin.Run:
-    run: plugin.Run = presets.default()
-    run.start()
-    run.log_src(_path.entrypoint(absolute=True))
-    run.set_tag("cherries.entrypoint", _path.entrypoint(absolute=False))
-    run.set_tag("cherries.run-dir", _path.run_dir(absolute=False))
-    return plugin.run
+        result: T = main(*args)
+    finally:
+        if len(type_hints) == 1:
+            cfg: C = args[0]
+            for path in config.get_outputs(cfg):
+                exp.log_output(path)
+        exp.end()
+    return result
+
+
+def start(
+    preset: presets.Preset = presets.default, *, play: bool = False
+) -> integration.Experiment:
+    if play:
+        preset = presets.playground
+    exp: integration.Experiment = preset(integration.exp)
+    exp.start()
+    exp.log_other("cherries.entrypoint", _path.entrypoint(absolute=False))
+    exp.log_other("cherries.exp-dir", _path.exp_dir(absolute=False))
+    exp.log_other("cherries.start-time", exp.start_time)
+    return exp
 
 
 def end() -> None:
-    plugin.run.end()
+    integration.exp.end()

liblaf/cherries/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.0.14'
-__version_tuple__ = version_tuple = (0, 0, 14)
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)

liblaf/cherries/config/__init__.pyi

@@ -0,0 +1,23 @@
+from ._asset import (
+    AssetKind,
+    MetaAsset,
+    PathProvider,
+    get_assets,
+    get_inputs,
+    get_outputs,
+    input,  # noqa: A004
+    output,
+)
+from ._config import BaseConfig
+
+__all__ = [
+    "AssetKind",
+    "BaseConfig",
+    "MetaAsset",
+    "PathProvider",
+    "get_assets",
+    "get_inputs",
+    "get_outputs",
+    "input",
+    "output",
+]

liblaf/cherries/config/_asset.py

@@ -0,0 +1,72 @@
+import enum
+from collections.abc import Callable, Iterable
+from pathlib import Path
+from typing import Any
+
+import pydantic
+
+from liblaf import grapes
+from liblaf.cherries import pathutils as _path
+from liblaf.cherries.typed import PathLike
+
+
+class AssetKind(enum.StrEnum):
+    INPUT = enum.auto()
+    OUTPUT = enum.auto()
+
+
+type PathProvider = (
+    PathLike | Iterable[PathLike] | Callable[[PathLike], PathLike | Iterable[PathLike]]
+)
+
+
+class MetaAsset:
+    kind: AssetKind
+    _extra: PathProvider | None = None
+
+    def __init__(self, kind: AssetKind, extra: PathProvider | None = None) -> None:
+        self.kind = kind
+        self._extra = extra
+
+    def get_extra(self, value: Path) -> list[Path]:
+        if self._extra is None:
+            return []
+        if callable(self._extra):
+            extra: Iterable[PathLike] = grapes.as_iterable(self._extra(value))
+            return [Path(p) for p in extra]
+        extra: Iterable[PathLike] = grapes.as_iterable(self._extra)
+        return [Path(p) for p in extra]
+
+
+def get_assets(cfg: pydantic.BaseModel, kind: AssetKind) -> list[Path]:
+    assets: list[Path] = []
+    for name, info in type(cfg).model_fields.items():
+        value: Any = getattr(cfg, name)
+        if isinstance(value, pydantic.BaseModel):
+            assets.extend(get_assets(value, kind))
+        for meta in info.metadata:
+            if isinstance(meta, MetaAsset) and meta.kind == kind:
+                value: Path = Path(value)
+                assets.append(value)
+                assets.extend(meta.get_extra(value))
+    return assets
+
+
+def get_inputs(cfg: pydantic.BaseModel) -> list[Path]:
+    return get_assets(cfg, AssetKind.INPUT)
+
+
+def get_outputs(cfg: pydantic.BaseModel) -> list[Path]:
+    return get_assets(cfg, AssetKind.OUTPUT)
+
+
+def input(path: PathLike, extra: PathProvider | None = None, **kwargs) -> Path:  # noqa: A001
+    field_info: pydantic.fields.FieldInfo = pydantic.Field(_path.data(path), **kwargs)  # pyright: ignore[reportAssignmentType]
+    field_info.metadata.append(MetaAsset(kind=AssetKind.INPUT, extra=extra))
+    return field_info  # pyright: ignore[reportReturnType]
+
+
+def output(path: PathLike, extra: PathProvider | None = None, **kwargs) -> Path:
+    field_info: pydantic.fields.FieldInfo = pydantic.Field(_path.data(path), **kwargs)  # pyright: ignore[reportAssignmentType]
+    field_info.metadata.append(MetaAsset(kind=AssetKind.OUTPUT, extra=extra))
+    return field_info  # pyright: ignore[reportReturnType]

liblaf/cherries/config/_config.py

@@ -0,0 +1,6 @@
+import pydantic_settings as ps
+
+
+class BaseConfig(ps.BaseSettings):
+    model_config = ps.SettingsConfigDict(cli_parse_args=True)
+    # TODO: add support for config files

liblaf/cherries/integration/__init__.pyi

@@ -0,0 +1,66 @@
+from . import comet, dvc, git, logging
+from ._abc import (
+    AddTag,
+    AddTags,
+    End,
+    LogAsset,
+    LogCode,
+    LogMetric,
+    LogMetrics,
+    LogOther,
+    LogOthers,
+    LogParam,
+    LogParams,
+    Plugin,
+    Start,
+)
+from ._exp import (
+    Experiment,
+    add_tag,
+    add_tags,
+    end,
+    exp,
+    log_asset,
+    log_code,
+    log_metric,
+    log_metrics,
+    log_other,
+    log_others,
+    log_param,
+    log_params,
+    start,
+)
+
+__all__ = [
+    "AddTag",
+    "AddTags",
+    "End",
+    "Experiment",
+    "LogAsset",
+    "LogCode",
+    "LogMetric",
+    "LogMetrics",
+    "LogOther",
+    "LogOthers",
+    "LogParam",
+    "LogParams",
+    "Plugin",
+    "Start",
+    "add_tag",
+    "add_tags",
+    "comet",
+    "dvc",
+    "end",
+    "exp",
+    "git",
+    "log_asset",
+    "log_code",
+    "log_metric",
+    "log_metrics",
+    "log_other",
+    "log_others",
+    "log_param",
+    "log_params",
+    "logging",
+    "start",
+]

liblaf/cherries/integration/_abc.py

@@ -0,0 +1,144 @@
+import bisect
+import functools
+import operator
+from collections.abc import Iterable, Mapping
+from typing import Any
+
+import attrs
+from loguru import logger
+
+from liblaf.cherries.typed import PathLike
+
+
+@attrs.define
+@functools.total_ordering
+class Plugin[**P, T]:
+    priority: int = attrs.field(default=0, kw_only=True)
+    _children: list["Plugin"] = attrs.field(
+        factory=list, eq=False, order=False, alias="children"
+    )
+
+    def __attrs_post_init__(self) -> None:
+        self._children.sort(key=operator.attrgetter("priority"))
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
+        ret: T | None = None
+        for child in self._children:
+            try:
+                ret = child(*args, **kwargs)
+            except BaseException as err:
+                if isinstance(err, (KeyboardInterrupt, SystemExit)):
+                    raise
+                logger.exception(child)
+        return ret  # pyright: ignore[reportReturnType]
+
+    def __lt__(self, other: "Plugin") -> bool:
+        if not isinstance(other, Plugin):
+            return NotImplemented
+        return self.priority < other.priority
+
+    def __eq__(self, other: "Plugin") -> bool:  # pyright: ignore[reportIncompatibleMethodOverride]
+        if not isinstance(other, Plugin):
+            return NotImplemented
+        return self.priority == other.priority
+
+    @property
+    def children(self) -> list["Plugin"]:
+        return self._children
+
+    def add(self, *children: "Plugin") -> None:
+        for c in children:
+            bisect.insort(self._children, c, key=operator.attrgetter("priority"))
+
+    def extend(self, children: Iterable["Plugin"]) -> None:
+        self.add(*children)
+
+    def remove(self, child: "Plugin") -> None:
+        self._children.remove(child)
+
+
+@attrs.define
+class AddTag(Plugin):
+    def __call__(self, tag: str, /, **kwargs) -> None:
+        super().__call__(tag, **kwargs)
+
+
+@attrs.define
+class AddTags(Plugin):
+    def __call__(self, tags: Iterable[str], /, **kwargs) -> None:
+        super().__call__(tags, **kwargs)
+
+
+@attrs.define
+class End(Plugin):
+    def __call__(self, **kwargs) -> None:
+        super().__call__(**kwargs)
+
+
+@attrs.define
+class LogAsset(Plugin):
+    def __call__(self, path: PathLike, /, prefix: str | None = None, **kwargs) -> None:
+        super().__call__(path, prefix=prefix, **kwargs)
+
+
+@attrs.define
+class LogCode(Plugin):
+    def __call__(self, path: PathLike, /, **kwargs) -> None:
+        super().__call__(path, **kwargs)
+
+
+@attrs.define
+class LogMetric(Plugin):
+    def __call__(
+        self,
+        key: str,
+        value: float,
+        /,
+        step: int | None = None,
+        epoch: int | None = None,
+        **kwargs,
+    ) -> None:
+        super().__call__(key, value, step=step, epoch=epoch, **kwargs)
+
+
+@attrs.define
+class LogMetrics(Plugin):
+    def __call__(
+        self,
+        metrics: Mapping[str, Any],
+        /,
+        step: int | None = None,
+        epoch: int | None = None,
+        **kwargs,
+    ) -> None:
+        super().__call__(metrics, step=step, epoch=epoch, **kwargs)
+
+
+@attrs.define
+class LogOther(Plugin):
+    def __call__(self, key: str, value: Any, /, **kwargs) -> None:
+        super().__call__(key, value, **kwargs)
+
+
+@attrs.define
+class LogOthers(Plugin):
+    def __call__(self, others: Mapping[str, Any], /, **kwargs) -> None:
+        super().__call__(others, **kwargs)
+
+
+@attrs.define
+class LogParam(Plugin):
+    def __call__(self, name: str, value: Any, /, **kwargs) -> None:
+        super().__call__(name, value, **kwargs)
+
+
+@attrs.define
+class LogParams(Plugin):
+    def __call__(self, params: Mapping[str, Any], /, **kwargs) -> None:
+        super().__call__(params, **kwargs)
+
+
+@attrs.define
+class Start(Plugin):
+    def __call__(self, **kwargs) -> None:
+        super().__call__(**kwargs)

liblaf/cherries/integration/_exp.py

@@ -0,0 +1,97 @@
+import datetime
+import functools
+from os import PathLike
+from pathlib import Path
+
+import attrs
+import comet_ml as comet
+
+from liblaf.cherries import pathutils as _path
+
+from ._abc import (
+    AddTag,
+    AddTags,
+    End,
+    LogAsset,
+    LogCode,
+    LogMetric,
+    LogMetrics,
+    LogOther,
+    LogOthers,
+    LogParam,
+    LogParams,
+    Start,
+)
+
+
+@attrs.define
+class Experiment:
+    add_tag: AddTag = attrs.field(factory=AddTag, init=False)
+    add_tags: AddTags = attrs.field(factory=AddTags, init=False)
+    end: End = attrs.field(factory=End, init=False)
+    log_asset: LogAsset = attrs.field(factory=LogAsset, init=False)
+    log_code: LogCode = attrs.field(factory=LogCode, init=False)
+    log_metric: LogMetric = attrs.field(factory=LogMetric, init=False)
+    log_metrics: LogMetrics = attrs.field(factory=LogMetrics, init=False)
+    log_other: LogOther = attrs.field(factory=LogOther, init=False)
+    log_others: LogOthers = attrs.field(factory=LogOthers, init=False)
+    log_param: LogParam = attrs.field(factory=LogParam, init=False)
+    log_params: LogParams = attrs.field(factory=LogParams, init=False)
+    start: Start = attrs.field(factory=Start, init=False)
+
+    @property
+    def comet(self) -> comet.CometExperiment:
+        return comet.get_running_experiment()  # pyright: ignore[reportReturnType]
+
+    @functools.cached_property
+    def exp_dir(self) -> Path:
+        return _path.exp_dir(absolute=True)
+
+    @property
+    def exp_key(self) -> str:
+        return self.comet.get_key()
+
+    @property
+    def exp_name(self) -> str:
+        return self.comet.get_name()
+
+    @exp_name.setter
+    def exp_name(self, value: str) -> None:
+        self.comet.set_name(value)
+
+    @property
+    def exp_url(self) -> str:
+        return self.comet.url  # pyright: ignore[reportReturnType]
+
+    @property
+    def project_id(self) -> str:
+        return self.comet.project_id  # pyright: ignore[reportReturnType]
+
+    @property
+    def project_name(self) -> str:
+        return self.comet.project_name
+
+    @functools.cached_property
+    def start_time(self) -> datetime.datetime:
+        return datetime.datetime.now().astimezone()
+
+    def log_input(self, path: PathLike, /, **kwargs) -> None:
+        self.log_asset(path, prefix="inputs/", **kwargs)
+
+    def log_output(self, path: PathLike, /, **kwargs) -> None:
+        self.log_asset(path, prefix="outputs/", **kwargs)
+
+
+exp = Experiment()
+add_tag: AddTag = exp.add_tag
+add_tags: AddTags = exp.add_tags
+end: End = exp.end
+log_asset: LogAsset = exp.log_asset
+log_code: LogCode = exp.log_code
+log_metric: LogMetric = exp.log_metric
+log_metrics: LogMetrics = exp.log_metrics
+log_other: LogOther = exp.log_other
+log_others: LogOthers = exp.log_others
+log_param: LogParam = exp.log_param
+log_params: LogParams = exp.log_params
+start: Start = exp.start