hydraflow 0.1.0__py3-none-any.whl

hydraflow/__init__.py

@@ -0,0 +1,28 @@
+from .context import Info, chdir_artifact, log_run, watch
+from .mlflow import set_experiment
+from .run import (
+    filter_by_config,
+    get_artifact_dir,
+    get_artifact_path,
+    get_artifact_uri,
+    get_by_config,
+    get_param_dict,
+    get_param_names,
+    get_run_id,
+)
+
+__all__ = [
+    "Info",
+    "chdir_artifact",
+    "filter_by_config",
+    "get_artifact_dir",
+    "get_artifact_path",
+    "get_artifact_uri",
+    "get_by_config",
+    "get_param_dict",
+    "get_param_names",
+    "get_run_id",
+    "log_run",
+    "set_experiment",
+    "watch",
+]

hydraflow/config.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from typing import Any
+
+
+def iter_params(config: object, prefix: str = "") -> Iterator[tuple[str, Any]]:
+    if not isinstance(config, DictConfig | ListConfig):
+        config = OmegaConf.create(config)  # type: ignore
+
+    if isinstance(config, DictConfig):
+        for key, value in config.items():
+            if isinstance(value, (DictConfig, ListConfig)):
+                yield from iter_params(value, f"{prefix}{key}.")
+
+            else:
+                yield f"{prefix}{key}", value
+
+    elif isinstance(config, ListConfig):
+        for index, value in enumerate(config):
+            if isinstance(value, (DictConfig, ListConfig)):
+                yield from iter_params(value, f"{prefix}{index}.")
+
+            else:
+                yield f"{prefix}{index}", value

hydraflow/context.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import os
+import time
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import mlflow
+from hydra.core.hydra_config import HydraConfig
+from watchdog.events import FileModifiedEvent, FileSystemEventHandler
+from watchdog.observers import Observer
+
+from hydraflow.mlflow import log_params
+from hydraflow.run import get_artifact_path
+from hydraflow.util import uri_to_path
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterator
+
+    from mlflow.entities.run import Run
+    from pandas import Series
+
+
+@dataclass
+class Info:
+    output_dir: Path
+    artifact_dir: Path
+
+
+@contextmanager
+def log_run(
+    config: object,
+    *,
+    synchronous: bool | None = None,
+) -> Iterator[Info]:
+    log_params(config, synchronous=synchronous)
+
+    hc = HydraConfig.get()
+    output_dir = Path(hc.runtime.output_dir)
+    uri = mlflow.get_artifact_uri()
+    location = Info(output_dir, uri_to_path(uri))
+
+    # Save '.hydra' config directory first.
+    output_subdir = output_dir / (hc.output_subdir or "")
+    mlflow.log_artifacts(output_subdir.as_posix(), hc.output_subdir)
+
+    try:
+        yield location
+
+    finally:
+        # Save output_dir including '.hydra' config directory.
+        mlflow.log_artifacts(output_dir.as_posix())
+
+
+@contextmanager
+def watch(
+    func: Callable[[Path], None],
+    dir: Path | str = "",
+    timeout: int = 600,
+) -> Iterator[None]:
+    if not dir:
+        uri = mlflow.get_artifact_uri()
+        dir = uri_to_path(uri)
+
+    handler = Handler(func)
+    observer = Observer()
+    observer.schedule(handler, dir, recursive=True)
+    observer.start()
+
+    try:
+        yield
+
+    finally:
+        elapsed = 0
+        while not observer.event_queue.empty():
+            time.sleep(0.2)
+            elapsed += 0.2
+            if elapsed > timeout:
+                break
+
+        observer.stop()
+        observer.join()
+
+
+class Handler(FileSystemEventHandler):
+    def __init__(self, func: Callable[[Path], None]) -> None:
+        self.func = func
+
+    def on_modified(self, event: FileModifiedEvent) -> None:
+        file = Path(event.src_path)
+        if file.is_file():
+            self.func(file)
+
+
+@contextmanager
+def chdir_artifact(
+    run: Run | Series | str,
+    artifact_path: str | None = None,
+) -> Iterator[Path]:
+    curdir = Path.cwd()
+
+    artifact_dir = get_artifact_path(run, artifact_path)
+
+    os.chdir(artifact_dir)
+    try:
+        yield artifact_dir
+    finally:
+        os.chdir(curdir)

hydraflow/mlflow.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import mlflow
+from hydra.core.hydra_config import HydraConfig
+
+from hydraflow.config import iter_params
+
+
+def set_experiment() -> None:
+    hc = HydraConfig.get()
+    mlflow.set_tracking_uri("")
+    mlflow.set_experiment(hc.job.name)
+
+
+def log_params(config: object, *, synchronous: bool | None = None) -> None:
+    for key, value in iter_params(config):
+        mlflow.log_param(key, value, synchronous=synchronous)

hydraflow/run.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, overload
+
+import mlflow
+import numpy as np
+from mlflow.entities.run import Run
+from mlflow.tracking import artifact_utils
+from omegaconf import DictConfig, OmegaConf
+
+from hydraflow.config import iter_params
+from hydraflow.util import uri_to_path
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from pandas import DataFrame, Series
+
+
+@overload
+def filter_by_config(runs: list[Run], config: object) -> list[Run]: ...
+
+
+@overload
+def filter_by_config(runs: DataFrame, config: object) -> DataFrame: ...
+
+
+def filter_by_config(runs: list[Run] | DataFrame, config: object):
+    if isinstance(runs, list):
+        return filter_by_config_list(runs, config)
+
+    return filter_by_config_dataframe(runs, config)
+
+
+def _is_equal(run: Run, key: str, value: Any) -> bool:
+    param = run.data.params.get(key, value)
+    if param is None:
+        return False
+
+    return type(value)(param) == value
+
+
+def filter_by_config_list(runs: list[Run], config: object) -> list[Run]:
+    for key, value in iter_params(config):
+        runs = [run for run in runs if _is_equal(run, key, value)]
+
+    return runs
+
+
+def filter_by_config_dataframe(runs: DataFrame, config: object) -> DataFrame:
+    index = np.ones(len(runs), dtype=bool)
+
+    for key, value in iter_params(config):
+        name = f"params.{key}"
+        if name in runs:
+            series = runs[name]
+            is_value = -series.isna()
+            param = series.fillna(value).astype(type(value))
+            index &= is_value & (param == value)
+
+    return runs[index]
+
+
+@overload
+def get_by_config(runs: list[Run], config: object) -> Run: ...
+
+
+@overload
+def get_by_config(runs: DataFrame, config: object) -> Series: ...
+
+
+def get_by_config(runs: list[Run] | DataFrame, config: object):
+    runs = filter_by_config(runs, config)
+
+    if len(runs) == 1:
+        return runs[0] if isinstance(runs, list) else runs.iloc[0]
+
+    msg = f"filtered runs has not length of 1.: {len(runs)}"
+    raise ValueError(msg)
+
+
+def drop_unique_params(runs: DataFrame) -> DataFrame:
+    def select(column: str) -> bool:
+        return not column.startswith("params.") or len(runs[column].unique()) > 1
+
+    columns = [select(column) for column in runs.columns]
+    return runs.iloc[:, columns]
+
+
+def get_param_names(runs: DataFrame) -> list[str]:
+    def get_name(column: str) -> str:
+        if column.startswith("params."):
+            return column.split(".", maxsplit=1)[-1]
+
+        return ""
+
+    columns = [get_name(column) for column in runs.columns]
+    return [column for column in columns if column]
+
+
+def get_param_dict(runs: DataFrame) -> dict[str, list[str]]:
+    params = {}
+    for name in get_param_names(runs):
+        params[name] = list(runs[f"params.{name}"].unique())
+
+    return params
+
+
+def get_run_id(run: Run | Series | str) -> str:
+    if isinstance(run, Run):
+        return run.info.run_id
+    if isinstance(run, str):
+        return run
+    return run.run_id
+
+
+def get_artifact_uri(run: Run | Series | str, artifact_path: str | None = None) -> str:
+    if isinstance(run, Run):
+        uri = run.info.artifact_uri
+    elif isinstance(run, str):
+        uri = artifact_utils.get_artifact_uri(run_id=run)
+    else:
+        uri = run.artifact_uri
+
+    if artifact_path:
+        uri = f"{uri}/{artifact_path}"
+
+    return uri  # type: ignore
+
+
+def get_artifact_dir(run: Run | Series | str) -> Path:
+    uri = get_artifact_uri(run)
+    return uri_to_path(uri)
+
+
+def get_artifact_path(
+    run: Run | Series | str,
+    artifact_path: str | None = None,
+) -> Path:
+    artifact_dir = get_artifact_dir(run)
+    return artifact_dir / artifact_path if artifact_path else artifact_dir
+
+
+def load_config(run: Run | Series | str, output_subdir: str = ".hydra") -> DictConfig:
+    run_id = get_run_id(run)
+
+    try:
+        path = mlflow.artifacts.download_artifacts(
+            run_id=run_id,
+            artifact_path=f"{output_subdir}/config.yaml",
+        )
+    except OSError:
+        return DictConfig({})
+
+    return OmegaConf.load(path)  # type: ignore
+
+
+def get_hydra_output_dir(run: Run | Series | str) -> Path:
+    path = get_artifact_dir(run) / ".hydra/hydra.yaml"
+
+    if path.exists():
+        hc = OmegaConf.load(path)
+        return Path(hc.hydra.runtime.output_dir)
+
+    raise FileNotFoundError
+
+
+def log_hydra_output_dir(run: Run | Series | str) -> None:
+    output_dir = get_hydra_output_dir(run)
+    run_id = run if isinstance(run, str) else run.info.run_id
+    mlflow.log_artifacts(output_dir.as_posix(), run_id=run_id)

hydraflow/util.py

@@ -0,0 +1,11 @@
+import platform
+from pathlib import Path
+from urllib.parse import urlparse
+
+
+def uri_to_path(uri: str) -> Path:
+    path = urlparse(uri).path
+    if platform.system() == "Windows" and path.startswith("/"):
+        path = path[1:]
+
+    return Path(path)

hydraflow-0.1.0.dist-info/METADATA

@@ -0,0 +1,29 @@
+Metadata-Version: 2.3
+Name: hydraflow
+Version: 0.1.0
+Summary: Hydra with MLflow
+Project-URL: Source, https://github.com/daizutabi/hydraflow
+Project-URL: Issues, https://github.com/daizutabi/hydraflow/issues
+Author-email: daizutabi <daizutabi@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Development Status :: 4 - Beta
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Documentation
+Classifier: Topic :: Software Development :: Documentation
+Requires-Python: >=3.10
+Requires-Dist: hydra-core
+Requires-Dist: mlflow
+Requires-Dist: watchdog
+Provides-Extra: dev
+Requires-Dist: pytest-clarity; extra == 'dev'
+Requires-Dist: pytest-cov; extra == 'dev'
+Requires-Dist: pytest-randomly; extra == 'dev'
+Requires-Dist: pytest-xdist; extra == 'dev'
+Requires-Dist: setuptools; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# hydraflow

hydraflow-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+hydraflow/__init__.py,sha256=9RaHPTloOOJYPUKKfPuK_wxKDr_J9A3rJ_gr-bLABD0,559
+hydraflow/config.py,sha256=b3Plh_lmq94loZNw9QP2asd6thCLyTzzYSutH0cONXA,964
+hydraflow/context.py,sha256=zBmbZWNLxUF2IDDPregPnR_sh3utmFwFJaneSsBsLDM,2558
+hydraflow/mlflow.py,sha256=yDZ_oB1IZdCNNqHm_0LxdZ1Nld28IkW8Xl7NMhWLApE,453
+hydraflow/run.py,sha256=XTAD_fd-ivvZ4tbjQLHrf6u5eAGRrrhqvExiZQcFnX8,4591
+hydraflow/util.py,sha256=HTymDLqa2UzCw3kNjqHDaAZNdRMnrEAWhCJ7_ZD7ffA,264
+hydraflow-0.1.0.dist-info/METADATA,sha256=WuryvAC_8MrC-UerPqbvcWxgBn9ABrnysQ0aRYimw3A,1021
+hydraflow-0.1.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+hydraflow-0.1.0.dist-info/licenses/LICENSE,sha256=IGdDrBPqz1O0v_UwCW-NJlbX9Hy9b3uJ11t28y2srmY,1062
+hydraflow-0.1.0.dist-info/RECORD,,

hydraflow-0.1.0.dist-info/WHEEL

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

hydraflow-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Daizu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.