hydraflow 0.1.1__tar.gz → 0.1.4__tar.gz

{hydraflow-0.1.1 → hydraflow-0.1.4}/.devcontainer/devcontainer.json

@@ -1,9 +1,7 @@
 {
   "name": "hydraflow",
-  "image": "mcr.microsoft.com/devcontainers/python:3.10-bookworm",
+  "image": "mcr.microsoft.com/vscode/devcontainers/base:ubuntu-22.04",
   "features": {
-    "ghcr.io/devcontainers-contrib/features/ruff:1": {},
-    "ghcr.io/devcontainers-contrib/features/hatch:2": {},
     "ghcr.io/devcontainers-contrib/features/starship:1": {},
     "ghcr.io/va-h/devcontainers-features/uv:1": {}
   },

{hydraflow-0.1.1 → hydraflow-0.1.4}/.devcontainer/starship.toml

@@ -1,6 +1,6 @@
 "$schema" = 'https://starship.rs/config-schema.json'
 
-add_newline = false
+add_newline = true
 
 [username]
 disabled = true

{hydraflow-0.1.1 → hydraflow-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: hydraflow
-Version: 0.1.1
+Version: 0.1.4
 Summary: Hydra with MLflow
 Project-URL: Documentation, https://github.com/daizutabi/hydraflow
 Project-URL: Source, https://github.com/daizutabi/hydraflow

{hydraflow-0.1.1 → hydraflow-0.1.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "hydraflow"
-version = "0.1.1"
+version = "0.1.4"
 description = "Hydra with MLflow"
 readme = "README.md"
 license = "MIT"
@@ -54,7 +54,7 @@ serve = "mkdocs serve --dev-addr localhost:8000 {args}"
 deploy = "mkdocs gh-deploy --force"
 
 [tool.ruff]
-line-length = 88
+line-length = 100
 target-version = "py312"
 
 [tool.ruff.lint]

{hydraflow-0.1.1 → hydraflow-0.1.4}/src/hydraflow/__init__.py

@@ -1,27 +1,35 @@
 from .context import Info, chdir_artifact, log_run, watch
 from .mlflow import set_experiment
-from .run import (
-    filter_by_config,
+from .runs import (
+    Run,
+    Runs,
+    drop_unique_params,
+    filter_runs,
     get_artifact_dir,
     get_artifact_path,
     get_artifact_uri,
-    get_by_config,
     get_param_dict,
     get_param_names,
+    get_run,
     get_run_id,
+    load_config,
 )
 
 __all__ = [
     "Info",
+    "Run",
+    "Runs",
     "chdir_artifact",
-    "filter_by_config",
+    "drop_unique_params",
+    "filter_runs",
     "get_artifact_dir",
     "get_artifact_path",
     "get_artifact_uri",
-    "get_by_config",
     "get_param_dict",
     "get_param_names",
+    "get_run",
     "get_run_id",
+    "load_config",
     "log_run",
     "set_experiment",
     "watch",

{hydraflow-0.1.1 → hydraflow-0.1.4}/src/hydraflow/context.py

@@ -13,7 +13,7 @@ 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.runs import get_artifact_path
 from hydraflow.util import uri_to_path
 
 if TYPE_CHECKING:
@@ -40,14 +40,19 @@ def log_run(
     hc = HydraConfig.get()
     output_dir = Path(hc.runtime.output_dir)
     uri = mlflow.get_artifact_uri()
-    location = Info(output_dir, uri_to_path(uri))
+    info = 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)
 
+    def log_artifact(path: Path) -> None:
+        local_path = (output_dir / path).as_posix()
+        mlflow.log_artifact(local_path)
+
     try:
-        yield location
+        with watch(log_artifact, output_dir):
+            yield info
 
     finally:
         # Save output_dir including '.hydra' config directory.
@@ -55,11 +60,7 @@ def log_run(
 
 
 @contextmanager
-def watch(
-    func: Callable[[Path], None],
-    dir: Path | str = "",
-    timeout: int = 600,
-) -> Iterator[None]:
+def watch(func: Callable[[Path], None], dir: Path | str = "", timeout: int = 60) -> Iterator[None]:
     if not dir:
         uri = mlflow.get_artifact_uri()
         dir = uri_to_path(uri)

{hydraflow-0.1.1 → hydraflow-0.1.4}/src/hydraflow/mlflow.py

@@ -6,10 +6,13 @@ from hydra.core.hydra_config import HydraConfig
 from hydraflow.config import iter_params
 
 
-def set_experiment() -> None:
+def set_experiment(prefix: str = "", suffix: str = "", uri: str | None = None) -> None:
+    if uri:
+        mlflow.set_tracking_uri(uri)
+
     hc = HydraConfig.get()
-    mlflow.set_tracking_uri("")
-    mlflow.set_experiment(hc.job.name)
+    name = f"{prefix}{hc.job.name}{suffix}"
+    mlflow.set_experiment(name)
 
 
 def log_params(config: object, *, synchronous: bool | None = None) -> None:

hydraflow-0.1.4/src/hydraflow/runs.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import cache
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import mlflow
+import numpy as np
+from mlflow.entities.run import Run as Run_
+from mlflow.tracking import artifact_utils
+from omegaconf import DictConfig, OmegaConf
+from pandas import DataFrame, Series
+
+from hydraflow.config import iter_params
+from hydraflow.util import uri_to_path
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+@dataclass
+class Runs:
+    runs: list[Run_] | DataFrame
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({len(self)})"
+
+    def __len__(self) -> int:
+        return len(self.runs)
+
+    def filter(self, config: object) -> Runs:
+        return Runs(filter_runs(self.runs, config))
+
+    def get(self, config: object) -> Run:
+        return Run(get_run(self.runs, config))
+
+    def drop_unique_params(self) -> Runs:
+        if isinstance(self.runs, DataFrame):
+            return Runs(drop_unique_params(self.runs))
+
+        raise NotImplementedError
+
+    def get_param_names(self) -> list[str]:
+        if isinstance(self.runs, DataFrame):
+            return get_param_names(self.runs)
+
+        raise NotImplementedError
+
+    def get_param_dict(self) -> dict[str, list[str]]:
+        if isinstance(self.runs, DataFrame):
+            return get_param_dict(self.runs)
+
+        raise NotImplementedError
+
+
+def filter_runs(runs: list[Run_] | DataFrame, config: object) -> list[Run_] | DataFrame:
+    if isinstance(runs, list):
+        return filter_runs_list(runs, config)
+
+    return filter_runs_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_runs_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_runs_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]
+
+
+def get_run(runs: list[Run_] | DataFrame, config: object) -> Run_ | Series:
+    runs = filter_runs(runs, config)
+
+    if len(runs) == 1:
+        return runs[0] if isinstance(runs, list) else runs.iloc[0]
+
+    msg = f"number of filtered runs is not 1: got {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
+
+
+@dataclass
+class Run:
+    run: Run_ | Series | str
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.run_id!r})"
+
+    @property
+    def run_id(self) -> str:
+        return get_run_id(self.run)
+
+    def artifact_uri(self, artifact_path: str | None = None) -> str:
+        return get_artifact_uri(self.run, artifact_path)
+
+    @property
+    def artifact_dir(self) -> Path:
+        return get_artifact_dir(self.run)
+
+    def artifact_path(self, artifact_path: str | None = None) -> Path:
+        return get_artifact_path(self.run, artifact_path)
+
+    @property
+    def config(self) -> DictConfig:
+        return load_config(self.run)
+
+    def log_hydra_output_dir(self) -> None:
+        log_hydra_output_dir(self.run)
+
+
+def get_run_id(run: Run_ | Series | str) -> str:
+    if isinstance(run, str):
+        return run
+
+    if isinstance(run, Run_):
+        return run.info.run_id
+
+    return run.run_id
+
+
+def get_artifact_uri(run: Run_ | Series | str, artifact_path: str | None = None) -> str:
+    run_id = get_run_id(run)
+    return artifact_utils.get_artifact_uri(run_id, artifact_path)
+
+
+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) -> DictConfig:
+    run_id = get_run_id(run)
+    return _load_config(run_id)
+
+
+@cache
+def _load_config(run_id: str) -> DictConfig:
+    try:
+        path = mlflow.artifacts.download_artifacts(
+            run_id=run_id,
+            artifact_path=".hydra/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-0.1.1/tests → hydraflow-0.1.4/tests/scripts}/watch.py

@@ -1,3 +1,4 @@
+import os
 import sys
 import time
 from pathlib import Path
@@ -14,6 +15,9 @@ def run():
         Path(f"{k}.txt").write_text(f"{k} {time.time()}")
         time.sleep(1)
 
+    os.makedirs("a/b")
+    Path("a/b/c.txt").write_text(f"4 {time.time()}")
+
 
 if __name__ == "__main__":
     run()

{hydraflow-0.1.1 → hydraflow-0.1.4}/tests/test_log_run.py

@@ -12,12 +12,10 @@ from mlflow.entities.run import Run
 
 @pytest.fixture
 def runs(monkeypatch, tmp_path):
-    file = Path("tests/log_run.py").absolute()
+    file = Path("tests/scripts/log_run.py").absolute()
     monkeypatch.chdir(tmp_path)
 
-    subprocess.check_call(
-        [sys.executable, file.as_posix(), "-m", "host=x,y", "port=1,2"]
-    )
+    subprocess.check_call([sys.executable, file.as_posix(), "-m", "host=x,y", "port=1,2"])
 
     mlflow.set_experiment("log_run")
     runs = mlflow.search_runs(output_format="list")
@@ -48,7 +46,7 @@ def read_log(run_id: str) -> str:
 
 
 def test_load_config(run_id: str):
-    from hydraflow.run import load_config
+    from hydraflow.runs import load_config
 
     log = read_log(run_id)
     host, port = log.splitlines()[0].split("START,")[-1].split(",")
@@ -56,9 +54,3 @@ def test_load_config(run_id: str):
     cfg = load_config(run_id)
     assert cfg.host == host.strip()
     assert cfg.port == int(port)
-
-
-def test_load_config_err(run_id: str):
-    from hydraflow.run import load_config
-
-    assert not load_config(run_id, "a")

hydraflow-0.1.1/tests/test_run.py → hydraflow-0.1.4/tests/test_runs.py

@@ -45,47 +45,47 @@ def runs_df(_runs: tuple[list[Run], DataFrame], request: pytest.FixtureRequest):
     return _runs[1]
 
 
-def test_filter_by_config_one(runs: list[Run] | DataFrame):
-    from hydraflow.run import filter_by_config
+def test_filter_one(runs: list[Run] | DataFrame):
+    from hydraflow.runs import filter_runs
 
     assert len(runs) == 6
-    x = filter_by_config(runs, {"p": 1})
+    x = filter_runs(runs, {"p": 1})
     assert len(x) == 1
 
 
-def test_filter_by_config_all(runs: list[Run] | DataFrame):
-    from hydraflow.run import filter_by_config
+def test_filter_all(runs: list[Run] | DataFrame):
+    from hydraflow.runs import filter_runs
 
     assert len(runs) == 6
-    x = filter_by_config(runs, {"q": 0})
+    x = filter_runs(runs, {"q": 0})
     assert len(x) == 6
 
 
-def test_get_by_config_list(runs_list: list[Run]):
-    from hydraflow.run import get_by_config
+def test_get_list(runs_list: list[Run]):
+    from hydraflow.runs import get_run
 
-    run = get_by_config(runs_list, {"p": 4})
+    run = get_run(runs_list, {"p": 4})
     assert isinstance(run, Run)
     assert run.data.params["p"] == "4"
 
 
-def test_get_by_config_df(runs_df: DataFrame):
-    from hydraflow.run import get_by_config
+def test_get_df(runs_df: DataFrame):
+    from hydraflow.runs import get_run
 
-    run = get_by_config(runs_df, {"p": 2})
+    run = get_run(runs_df, {"p": 2})
     assert isinstance(run, Series)
     assert run["params.p"] == "2"
 
 
-def test_get_by_config_error(runs: list[Run] | DataFrame):
-    from hydraflow.run import get_by_config
+def test_get_error(runs: list[Run] | DataFrame):
+    from hydraflow.runs import get_run
 
     with pytest.raises(ValueError):
-        get_by_config(runs, {"q": 0})
+        get_run(runs, {"q": 0})
 
 
 def test_drop_unique_params(runs_df):
-    from hydraflow.run import drop_unique_params
+    from hydraflow.runs import drop_unique_params
 
     assert "params.p" in runs_df
     assert "params.q" in runs_df
@@ -95,7 +95,7 @@ def test_drop_unique_params(runs_df):
 
 
 def test_get_param_names(runs_df: DataFrame):
-    from hydraflow.run import get_param_names
+    from hydraflow.runs import get_param_names
 
     params = get_param_names(runs_df)
     assert len(params) == 2
@@ -104,7 +104,7 @@ def test_get_param_names(runs_df: DataFrame):
 
 
 def test_get_param_dict(runs_df: DataFrame):
-    from hydraflow.run import get_param_dict
+    from hydraflow.runs import get_param_dict
 
     params = get_param_dict(runs_df)
     assert len(params["p"]) == 6
@@ -113,7 +113,7 @@ def test_get_param_dict(runs_df: DataFrame):
 
 @pytest.mark.parametrize("i", range(6))
 def test_get_run_id(i: int, runs_list: list[Run], runs_df: DataFrame):
-    from hydraflow.run import get_run_id
+    from hydraflow.runs import get_run_id
 
     assert get_run_id(runs_list[i]) == get_run_id(runs_df.iloc[i])
     assert get_run_id(runs_list[i]) == get_run_id(runs_df.iloc[i])
@@ -125,7 +125,7 @@ def test_get_run_id(i: int, runs_list: list[Run], runs_df: DataFrame):
 @pytest.mark.parametrize("i", range(6))
 @pytest.mark.parametrize("path", [None, "a"])
 def test_get_artifact_uri(i: int, path, runs_list: list[Run], runs_df: DataFrame):
-    from hydraflow.run import get_artifact_uri, get_run_id
+    from hydraflow.runs import get_artifact_uri, get_run_id
 
     x = get_run_id(runs_list[i])
     y = get_artifact_uri(runs_list[i], path)
@@ -144,7 +144,7 @@ def test_chdir_artifact_list(i: int, runs_list: list[Run]):
 
 
 def test_hydra_output_dir_error(runs_list: list[Run]):
-    from hydraflow.run import get_hydra_output_dir
+    from hydraflow.runs import get_hydra_output_dir
 
     with pytest.raises(FileNotFoundError):
         get_hydra_output_dir(runs_list[0])
@@ -163,7 +163,7 @@ def df():
 
 
 def test_unique_params(df):
-    from hydraflow.run import drop_unique_params
+    from hydraflow.runs import drop_unique_params
 
     df = drop_unique_params(df)
     assert len(df.columns) == 3
@@ -173,16 +173,88 @@ def test_unique_params(df):
 
 
 def test_param_names(df):
-    from hydraflow.run import get_param_names
+    from hydraflow.runs import get_param_names
 
     names = get_param_names(df)
     assert names == ["x", "y", "z"]
 
 
 def test_param_dict(df):
-    from hydraflow.run import get_param_dict
+    from hydraflow.runs import get_param_dict
 
     x = get_param_dict(df)
     assert x["x"] == [1, 2]
     assert x["y"] == [1, 2]
     assert x["z"] == [1]
+
+
+def test_runs_repr(runs):
+    from hydraflow.runs import Runs
+
+    assert repr(Runs(runs)) == "Runs(6)"
+
+
+def test_runs_filter(runs):
+    from hydraflow.runs import Runs
+
+    runs = Runs(runs)
+
+    assert len(runs.filter({})) == 6
+    assert len(runs.filter({"p": 1})) == 1
+    assert len(runs.filter({"q": 0})) == 6
+    assert len(runs.filter({"q": -1})) == 0
+
+
+def test_runs_get(runs):
+    from hydraflow.runs import Run, Runs
+
+    runs = Runs(runs)
+    run = runs.get({"p": 4})
+    assert isinstance(run, Run)
+
+
+def test_runs_drop_unique_params(runs_df):
+    from hydraflow.runs import Runs
+
+    runs = Runs(runs_df)
+    assert runs.runs.shape == (6, 12)  # type: ignore
+    runs = runs.drop_unique_params()
+    assert runs.runs.shape == (6, 11)  # type: ignore
+
+
+def test_runs_get_params_names(runs_df):
+    from hydraflow.runs import Runs
+
+    runs = Runs(runs_df)
+    names = runs.get_param_names()
+    assert len(names) == 2
+    assert "p" in names
+    assert "q" in names
+
+
+def test_runs_get_params_dict(runs_df):
+    from hydraflow.runs import Runs
+
+    runs = Runs(runs_df)
+    params = runs.get_param_dict()
+    assert params["p"] == ["0", "1", "2", "3", "4", "5"]
+    assert params["q"] == ["0"]
+
+
+@pytest.fixture
+def run(runs):
+    from hydraflow.runs import Runs
+
+    return Runs(runs).get({"p": 5})
+
+
+def test_run_id(run):
+    assert run.run_id in repr(run)
+
+
+def test_run_artifact_uri(run):
+    assert run.artifact_uri().startswith("file:")
+
+
+def test_run_artifact_dir(run):
+    assert run.artifact_dir.exists()

{hydraflow-0.1.1 → hydraflow-0.1.4}/tests/test_watch.py

@@ -11,7 +11,7 @@ import pytest
 def test_watch(dir, monkeypatch, tmp_path):
     from hydraflow.context import watch
 
-    file = Path("tests/watch.py").absolute()
+    file = Path("tests/scripts/watch.py").absolute()
     monkeypatch.chdir(tmp_path)
 
     lines = []

hydraflow-0.1.1/src/hydraflow/run.py

@@ -1,172 +0,0 @@
-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)