hydraflow 0.9.0__tar.gz → 0.9.1__tar.gz

{hydraflow-0.9.0 → hydraflow-0.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hydraflow
-Version: 0.9.0
+Version: 0.9.1
 Summary: Hydraflow integrates Hydra and MLflow to manage and track machine learning experiments.
 Project-URL: Documentation, https://daizutabi.github.io/hydraflow/
 Project-URL: Source, https://github.com/daizutabi/hydraflow

{hydraflow-0.9.0 → hydraflow-0.9.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "hydraflow"
-version = "0.9.0"
+version = "0.9.1"
 description = "Hydraflow integrates Hydra and MLflow to manage and track machine learning experiments."
 readme = "README.md"
 license = { file = "LICENSE" }

{hydraflow-0.9.0 → hydraflow-0.9.1}/src/hydraflow/__init__.py

@@ -5,6 +5,10 @@ from hydraflow.core.io import (
     get_artifact_dir,
     get_artifact_path,
     get_hydra_output_dir,
+    iter_artifact_paths,
+    iter_artifacts_dirs,
+    iter_experiment_dirs,
+    iter_run_dirs,
     load_config,
     remove_run,
 )
@@ -18,6 +22,10 @@ __all__ = [
     "get_artifact_dir",
     "get_artifact_path",
     "get_hydra_output_dir",
+    "iter_artifact_paths",
+    "iter_artifacts_dirs",
+    "iter_experiment_dirs",
+    "iter_run_dirs",
     "list_run_ids",
     "list_run_paths",
     "list_runs",

{hydraflow-0.9.0 → hydraflow-0.9.1}/src/hydraflow/core/io.py

@@ -15,7 +15,7 @@ from mlflow.entities import Run
 from omegaconf import DictConfig, ListConfig, OmegaConf
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Iterable, Iterator
 
 
 def file_uri_to_path(uri: str) -> Path:
@@ -147,3 +147,74 @@ def remove_run(run: Run | Iterable[Run]) -> None:
         return
 
     shutil.rmtree(get_artifact_dir(run).parent)
+
+
+def get_root_dir(uri: str | Path | None = None) -> Path:
+    """Get the root directory for the MLflow tracking server."""
+    if uri is not None:
+        return Path(uri).absolute()
+
+    uri = mlflow.get_tracking_uri()
+
+    if uri.startswith("file:"):
+        return file_uri_to_path(uri)
+
+    return Path(uri).absolute()
+
+
+def get_experiment_name(path: Path) -> str | None:
+    """Get the experiment name from the meta file."""
+    metafile = path / "meta.yaml"
+    if not metafile.exists():
+        return None
+    lines = metafile.read_text().splitlines()
+    for line in lines:
+        if line.startswith("name:"):
+            return line.split(":")[1].strip()
+    return None
+
+
+def iter_experiment_dirs(
+    experiment_names: str | list[str] | None = None,
+    root_dir: str | Path | None = None,
+) -> Iterator[Path]:
+    """Iterate over the experiment directories in the root directory."""
+    if isinstance(experiment_names, str):
+        experiment_names = [experiment_names]
+
+    root_dir = get_root_dir(root_dir)
+    for path in root_dir.iterdir():
+        if path.is_dir() and path.name not in [".trash", "0"]:
+            if name := get_experiment_name(path):
+                if experiment_names is None or name in experiment_names:
+                    yield path
+
+
+def iter_run_dirs(
+    experiment_names: str | list[str] | None = None,
+    root_dir: str | Path | None = None,
+) -> Iterator[Path]:
+    """Iterate over the run directories in the root directory."""
+    for experiment_dir in iter_experiment_dirs(experiment_names, root_dir):
+        for path in experiment_dir.iterdir():
+            if path.is_dir() and (path / "artifacts").exists():
+                yield path
+
+
+def iter_artifacts_dirs(
+    experiment_names: str | list[str] | None = None,
+    root_dir: str | Path | None = None,
+) -> Iterator[Path]:
+    """Iterate over the artifacts directories in the root directory."""
+    for path in iter_run_dirs(experiment_names, root_dir):
+        yield path / "artifacts"
+
+
+def iter_artifact_paths(
+    artifact_path: str | Path,
+    experiment_names: str | list[str] | None = None,
+    root_dir: str | Path | None = None,
+) -> Iterator[Path]:
+    """Iterate over the artifact paths in the root directory."""
+    for path in iter_artifacts_dirs(experiment_names, root_dir):
+        yield path / artifact_path

hydraflow-0.9.1/tests/core/io/test_iter_dirs.py

@@ -0,0 +1,82 @@
+from pathlib import Path
+
+import mlflow
+import pytest
+
+pytestmark = pytest.mark.xdist_group(name="group1")
+
+
+@pytest.fixture(scope="module")
+def root_dir(chdir):
+    return Path("mlruns").absolute()
+
+
+@pytest.fixture(scope="module", autouse=True)
+def setup(chdir):
+    mlflow.set_experiment("e1")
+    with mlflow.start_run():
+        mlflow.log_text("1", "text.txt")
+    with mlflow.start_run():
+        mlflow.log_text("2", "text.txt")
+    mlflow.set_experiment("e2")
+    with mlflow.start_run():
+        mlflow.log_text("3", "text.txt")
+    with mlflow.start_run():
+        mlflow.log_text("4", "text.txt")
+    with mlflow.start_run():
+        mlflow.log_text("5", "text.txt")
+
+
+def test_root_dir(root_dir: Path):
+    from hydraflow.core.io import get_root_dir
+
+    assert get_root_dir(root_dir) == root_dir
+    assert get_root_dir(root_dir.name) == root_dir
+    assert get_root_dir() == root_dir
+
+
+def test_iter_experiment_dirs():
+    from hydraflow.core.io import get_experiment_name, iter_experiment_dirs
+
+    names = [get_experiment_name(p) for p in iter_experiment_dirs()]
+    assert sorted(names) == ["e1", "e2"]  # type: ignore
+
+
+def test_iter_experiment_dirs_filter():
+    from hydraflow.core.io import get_experiment_name, iter_experiment_dirs
+
+    it = iter_experiment_dirs(experiment_names="e1")
+    assert [get_experiment_name(p) for p in it] == ["e1"]
+
+
+def test_get_experiment_name_none(root_dir: Path):
+    from hydraflow.core.io import get_experiment_name
+
+    assert get_experiment_name(root_dir.parent) is None
+
+
+def test_get_experiment_name_metafile_none(root_dir: Path):
+    from hydraflow.core.io import get_experiment_name
+
+    (root_dir / "meta.yaml").touch()
+    assert get_experiment_name(root_dir) is None
+
+
+def test_iter_run_dirs():
+    from hydraflow.core.io import iter_run_dirs
+
+    assert len(list(iter_run_dirs())) == 5
+
+
+def test_iter_artifacts_dirs():
+    from hydraflow.core.io import iter_artifacts_dirs
+
+    assert len(list(iter_artifacts_dirs())) == 5
+
+
+def test_iter_artifact_paths():
+    from hydraflow.core.io import iter_artifact_paths
+
+    it = iter_artifact_paths("text.txt")
+    text = sorted("".join(p.read_text() for p in it))
+    assert text == ["1", "2", "3", "4", "5"]

{hydraflow-0.9.0 → hydraflow-0.9.1}/tests/core/test_mlflow.py

@@ -81,3 +81,10 @@ def test_list_run_paths(experiment: Experiment):
 
     dirs = list_run_paths(experiment.name, "artifacts")
     assert all(d.is_dir() for d in dirs)
+
+
+@pytest.mark.parametrize("uri", [None, "test_mlflow"])
+def test_root_dir(experiment, uri):
+    from hydraflow.core.io import get_root_dir
+
+    assert get_root_dir(uri) == Path.cwd() / "test_mlflow"