hydraflow 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

hydraflow/__init__.py

@@ -1,7 +1,7 @@
 """Integrate Hydra and MLflow to manage and track machine learning experiments."""
 
 from .config import select_config, select_overrides
-from .context import chdir_artifact, chdir_hydra_output, log_run, start_run
+from .context import chdir_artifact, log_run, start_run
 from .mlflow import list_runs, search_runs, set_experiment
 from .run_collection import RunCollection
 from .utils import (
@@ -17,7 +17,6 @@ from .utils import (
 __all__ = [
     "RunCollection",
     "chdir_artifact",
-    "chdir_hydra_output",
     "get_artifact_dir",
     "get_artifact_path",
     "get_hydra_output_dir",

hydraflow/context.py

@@ -13,6 +13,7 @@ import mlflow.artifacts
 from hydra.core.hydra_config import HydraConfig
 
 from hydraflow.mlflow import log_params
+from hydraflow.utils import get_artifact_dir
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
@@ -69,24 +70,26 @@ def log_run(
         raise
 
     finally:
-        log_hydra(output_dir)
+        log_text(output_dir)
 
 
-def log_hydra(output_dir: Path) -> None:
-    """Log hydra logs of the current run as artifacts.
+def log_text(directory: Path, pattern: str = "*.log") -> None:
+    """Log text files in the given directory as artifacts.
+
+    Append the text files to the existing text file in the artifact directory.
 
     Args:
-        output_dir (Path): The output directory of the Hydra job.
+        directory (Path): The directory to find the logs in.
+        pattern (str): The pattern to match the logs.
 
     """
-    uri = mlflow.get_artifact_uri()
-    artifact_dir = Path(mlflow.artifacts.download_artifacts(uri))
+    artifact_dir = get_artifact_dir()
 
-    for file_hydra in output_dir.glob("*.log"):
-        if not file_hydra.is_file():
+    for file in directory.glob(pattern):
+        if not file.is_file():
             continue
 
-        file_artifact = artifact_dir / file_hydra.name
+        file_artifact = artifact_dir / file.name
         if file_artifact.exists():
             text = file_artifact.read_text()
             if not text.endswith("\n"):
@@ -94,8 +97,8 @@ def log_hydra(output_dir: Path) -> None:
         else:
             text = ""
 
-        text += file_hydra.read_text()
-        mlflow.log_text(text, file_hydra.name)
+        text += file.read_text()
+        mlflow.log_text(text, file.name)
 
 
 @contextmanager
@@ -174,29 +177,7 @@ def start_run(  # noqa: PLR0913
 
 
 @contextmanager
-def chdir_hydra_output() -> Iterator[Path]:
-    """Change the current working directory to the hydra output directory.
-
-    This context manager changes the current working directory to the hydra output
-    directory. It ensures that the directory is changed back to the original
-    directory after the context is exited.
-    """
-    curdir = Path.cwd()
-    path = HydraConfig.get().runtime.output_dir
-
-    os.chdir(path)
-    try:
-        yield Path(path)
-
-    finally:
-        os.chdir(curdir)
-
-
-@contextmanager
-def chdir_artifact(
-    run: Run,
-    artifact_path: str | None = None,
-) -> Iterator[Path]:
+def chdir_artifact(run: Run | None = None) -> Iterator[Path]:
     """Change the current working directory to the artifact directory of the given run.
 
     This context manager changes the current working directory to the artifact
@@ -204,19 +185,16 @@ def chdir_artifact(
     to the original directory after the context is exited.
 
     Args:
-        run (Run): The run to get the artifact directory from.
-        artifact_path (str | None): The artifact path.
+        run (Run | None): The run to get the artifact directory from.
 
     """
     curdir = Path.cwd()
-    path = mlflow.artifacts.download_artifacts(
-        run_id=run.info.run_id,
-        artifact_path=artifact_path,
-    )
+    artifact_dir = get_artifact_dir(run)
+
+    os.chdir(artifact_dir)
 
-    os.chdir(path)
     try:
-        yield Path(path)
+        yield artifact_dir
 
     finally:
         os.chdir(curdir)

hydraflow/mlflow.py

@@ -16,7 +16,6 @@ Key Features:
 
 from __future__ import annotations
 
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 import joblib
@@ -28,8 +27,11 @@ from mlflow.tracking.fluent import SEARCH_MAX_RESULTS_PANDAS, _get_experiment_id
 
 from hydraflow.config import iter_params
 from hydraflow.run_collection import RunCollection
+from hydraflow.utils import get_artifact_dir
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from mlflow.entities.experiment import Experiment
 
 
@@ -211,16 +213,10 @@ def _list_runs(
 
     for name in experiment_names:
         if experiment := mlflow.get_experiment_by_name(name):
-            loc = experiment.artifact_location
-
-            if isinstance(loc, str):
-                if loc.startswith("file:"):
-                    path = Path(mlflow.artifacts.download_artifacts(loc))
-                elif Path(loc).is_dir():
-                    path = Path(loc)
-                else:
-                    continue  # no cov
+            uri = experiment.artifact_location
 
+            if isinstance(uri, str):
+                path = get_artifact_dir(uri=uri)
                 run_ids.extend(file.stem for file in path.iterdir() if file.is_dir())
 
     it = (joblib.delayed(mlflow.get_run)(run_id) for run_id in run_ids)

hydraflow/run_collection.py

@@ -236,7 +236,7 @@ class RunCollection:
 
     def filter(
         self,
-        config: object | None = None,
+        config: object | Callable[[Run], bool] | None = None,
         *,
         override: bool = False,
         select: list[str] | None = None,
@@ -257,11 +257,13 @@ class RunCollection:
         - Membership checks for lists of values.
         - Range checks for tuples of two values (inclusive of both the lower
           and upper bound).
+        - Callable that takes a `Run` object and returns a boolean value.
 
         Args:
-            config (object | None): The configuration object to filter the runs.
-                This can be any object that provides key-value pairs through
-                the `iter_params` function.
+            config (object | Callable[[Run], bool] | None): The configuration object
+                to filter the runs. This can be any object that provides key-value
+                pairs through the `iter_params` function, or a callable that
+                takes a `Run` object and returns a boolean value.
             override (bool): If True, override the configuration object with the
                 provided key-value pairs.
             select (list[str] | None): The list of parameters to select.
@@ -711,7 +713,7 @@ def _param_matches(run: Run, key: str, value: Any) -> bool:
 
 def filter_runs(
     runs: list[Run],
-    config: object | None = None,
+    config: object | Callable[[Run], bool] | None = None,
     *,
     override: bool = False,
     select: list[str] | None = None,
@@ -735,9 +737,11 @@ def filter_runs(
 
     Args:
         runs (list[Run]): The list of runs to filter.
-        config (object | None, optional): The configuration object to filter the
-            runs. This can be any object that provides key-value pairs through
-            the `iter_params` function. Defaults to None.
+        config (object | Callable[[Run], bool] | None, optional): The
+            configuration object to filter the runs. This can be any object
+            that provides key-value pairs through the `iter_params` function.
+            This can also be a callable that takes a `Run` object and returns
+            a boolean value. Defaults to None.
         override (bool, optional): If True, filter the runs based on
             the overrides. Defaults to False.
         select (list[str] | None, optional): The list of parameters to select.
@@ -750,15 +754,19 @@ def filter_runs(
         A list of runs that match the specified configuration and key-value pairs.
 
     """
-    if override:
-        config = select_overrides(config)
-    elif select:
-        config = select_config(config, select)
-
-    for key, value in chain(iter_params(config), kwargs.items()):
-        runs = [run for run in runs if _param_matches(run, key, value)]
-        if not runs:
-            return []
+    if callable(config):
+        runs = [run for run in runs if config(run)]
+
+    else:
+        if override:
+            config = select_overrides(config)
+        elif select:
+            config = select_config(config, select)
+
+        for key, value in chain(iter_params(config), kwargs.items()):
+            runs = [run for run in runs if _param_matches(run, key, value)]
+            if not runs:
+                return []
 
     if status is None:
         return runs

hydraflow/utils.py

@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import shutil
+import urllib.parse
+import urllib.request
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -16,30 +18,40 @@ if TYPE_CHECKING:
     from collections.abc import Iterable
 
 
-def get_artifact_dir(run: Run | None = None) -> Path:
+def get_artifact_dir(run: Run | None = None, uri: str | None = None) -> Path:
     """Retrieve the artifact directory for the given run.
 
     This function uses MLflow to get the artifact directory for the given run.
 
     Args:
         run (Run | None): The run object. Defaults to None.
+        uri (str | None): The URI of the artifact. Defaults to None.
 
     Returns:
         The local path to the directory where the artifacts are downloaded.
 
     """
-    uri = mlflow.get_artifact_uri() if run is None else run.info.artifact_uri
+    if run is not None and uri is not None:
+        raise ValueError("Cannot provide both run and uri")
+
+    if run is None and uri is None:
+        uri = mlflow.get_artifact_uri()
+    elif run:
+        uri = run.info.artifact_uri
 
     if not isinstance(uri, str):
         raise NotImplementedError
 
     if uri.startswith("file:"):
-        return Path(mlflow.artifacts.download_artifacts(uri))
+        return file_uri_to_path(uri)
+
+    return Path(uri)
 
-    if Path(uri).is_dir():
-        return Path(uri)
 
-    raise NotImplementedError
+def file_uri_to_path(uri: str) -> Path:
+    """Convert a file URI to a local path."""
+    path = urllib.parse.urlparse(uri).path
+    return Path(urllib.request.url2pathname(path))  # for Windows
 
 
 def get_artifact_path(run: Run | None, path: str) -> Path:

{hydraflow-0.6.0.dist-info → hydraflow-0.6.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hydraflow
-Version: 0.6.0
+Version: 0.6.2
 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.6.2.dist-info/RECORD

@@ -0,0 +1,14 @@
+hydraflow/__init__.py,sha256=VPIPNNCyjMAkWBbdvB7Ltwe3QWoc2FwuqkV8uJM5JoM,809
+hydraflow/config.py,sha256=MNX9da5bPVDcjnpji7Cm9ndK6ura92pt361m4PRh6_E,4326
+hydraflow/context.py,sha256=3xfKhMozkKFqtWeOp9Gie0A5o5URMta4US6iVD5TcLU,6002
+hydraflow/mlflow.py,sha256=imD3XL0RTlpnKrkyvO8FNy_Bv6hwSfLiOu1yJuL40ck,8773
+hydraflow/param.py,sha256=yu1aMNXRLegXGDL-68vwIkfeDF9CaU784WZENGLwl7Q,4572
+hydraflow/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+hydraflow/run_collection.py,sha256=w_GZdc_6yviwzRWLndWDSWB4DKyGyA9di9d9UpkkLZo,27926
+hydraflow/run_data.py,sha256=dpyyfnuH9mCtIZeigMo1iFQo9bafMdEL4i4uI2l0UqY,1525
+hydraflow/run_info.py,sha256=Jf5wrIjRLIV1-k-obHDqwKHa6j_ZonrY8od-rXlbtMo,1024
+hydraflow/utils.py,sha256=a9i5PEJn8Ssowv9dqHadAihZXlsqtVjHZ9MZvkPq1bY,4747
+hydraflow-0.6.2.dist-info/METADATA,sha256=9a3blsQ91rNP1Ql4kFDc7tZxDMbdK5PzEAfP9ZyUY6A,4700
+hydraflow-0.6.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hydraflow-0.6.2.dist-info/licenses/LICENSE,sha256=IGdDrBPqz1O0v_UwCW-NJlbX9Hy9b3uJ11t28y2srmY,1062
+hydraflow-0.6.2.dist-info/RECORD,,

hydraflow-0.6.0.dist-info/RECORD

@@ -1,14 +0,0 @@
-hydraflow/__init__.py,sha256=9XO9FD3uiTTPN6X6UAC9FtkJjEqUQZNqpoAmSrjUHfI,855
-hydraflow/config.py,sha256=MNX9da5bPVDcjnpji7Cm9ndK6ura92pt361m4PRh6_E,4326
-hydraflow/context.py,sha256=rc43zvE2ueki0zEzorCMIthD9cho_PkbLLJYF9WgDqY,6562
-hydraflow/mlflow.py,sha256=h2S_A2wElr_1lAq0D1wkoEfdtDZpPuWFNRcO8mV_VrA,8932
-hydraflow/param.py,sha256=yu1aMNXRLegXGDL-68vwIkfeDF9CaU784WZENGLwl7Q,4572
-hydraflow/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-hydraflow/run_collection.py,sha256=2GRVOy87_2SPjHuCzzUvRNugO_grtFUVjtTfhznwBAc,27444
-hydraflow/run_data.py,sha256=dpyyfnuH9mCtIZeigMo1iFQo9bafMdEL4i4uI2l0UqY,1525
-hydraflow/run_info.py,sha256=Jf5wrIjRLIV1-k-obHDqwKHa6j_ZonrY8od-rXlbtMo,1024
-hydraflow/utils.py,sha256=oXjcyfQBbPzJNTh3_CbZfl23zgJS-mbNM9GAWBwsn8c,4349
-hydraflow-0.6.0.dist-info/METADATA,sha256=xUib1EsbG3Es5jFx0cSkF1QItfTuciBHYM1040GqFzA,4700
-hydraflow-0.6.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hydraflow-0.6.0.dist-info/licenses/LICENSE,sha256=IGdDrBPqz1O0v_UwCW-NJlbX9Hy9b3uJ11t28y2srmY,1062
-hydraflow-0.6.0.dist-info/RECORD,,