hydraflow 0.14.3__py3-none-any.whl → 0.15.0__py3-none-any.whl

hydraflow/__init__.py

@@ -3,35 +3,25 @@
 from hydraflow.core.context import chdir_artifact, log_run, start_run
 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,
 )
 from hydraflow.core.main import main
-from hydraflow.core.mlflow import list_run_ids, list_run_paths, list_runs
-from hydraflow.entities.run_collection import RunCollection
+from hydraflow.core.run import Run
+from hydraflow.core.run_collection import RunCollection
 
 __all__ = [
+    "Run",
     "RunCollection",
     "chdir_artifact",
     "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",
-    "load_config",
     "log_run",
     "main",
-    "remove_run",
     "start_run",
 ]

hydraflow/core/context.py

@@ -10,9 +10,7 @@ from typing import TYPE_CHECKING
 
 from hydra.core.hydra_config import HydraConfig
 
-from hydraflow.core.io import get_artifact_dir
-
-from .mlflow import log_params, log_text
+from .io import get_artifact_dir, log_text
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
@@ -23,38 +21,22 @@ log = logging.getLogger(__name__)
 
 
 @contextmanager
-def log_run(
-    config: object | None,
-    *,
-    synchronous: bool | None = None,
-) -> Iterator[None]:
-    """Log the parameters from the given configuration object.
+def log_run(run: Run) -> Iterator[None]:
+    """Log the parameters from the given configuration instance.
 
-    This context manager logs the parameters from the provided configuration object
+    This context manager logs the parameters from the provided configuration instance
     using MLflow. It also manages the MLflow run context, ensuring that artifacts
     are logged and the run is properly closed.
 
     Args:
-        config (object): The configuration object to log the parameters from.
-        synchronous (bool | None): Whether to log the parameters synchronously.
-            Defaults to None.
+        run (Run): The run instance.
 
     Yields:
         None
 
-    Example:
-        ```python
-        with log_run(config):
-            # Perform operations within the MLflow run context
-            pass
-        ```
-
     """
     import mlflow
 
-    if config:
-        log_params(config, synchronous=synchronous)
-
     hc = HydraConfig.get()
     hydra_dir = Path(hc.runtime.output_dir)
 
@@ -71,12 +53,11 @@ def log_run(
         raise
 
     finally:
-        log_text(hydra_dir)
+        log_text(run, hydra_dir)
 
 
 @contextmanager
 def start_run(
-    config: object,
     *,
     chdir: bool = False,
     run_id: str | None = None,
@@ -87,15 +68,14 @@ def start_run(
     tags: dict[str, str] | None = None,
     description: str | None = None,
     log_system_metrics: bool | None = None,
-    synchronous: bool | None = None,
 ) -> Iterator[Run]:
-    """Start an MLflow run and log parameters using the provided configuration object.
+    """Start an MLflow run and log parameters using the provided configuration instance.
 
     This context manager starts an MLflow run and logs parameters using the specified
-    configuration object. It ensures that the run is properly closed after completion.
+    configuration instance. It ensures that the run is properly closed after completion.
 
     Args:
-        config (object): The configuration object to log parameters from.
+        config (object): The configuration instance to log parameters from.
         chdir (bool): Whether to change the current working directory to the
             artifact directory of the current run. Defaults to False.
         run_id (str | None): The existing run ID. Defaults to None.
@@ -111,7 +91,7 @@ def start_run(
             Defaults to None.
 
     Yields:
-        Run: An MLflow Run object representing the started run.
+        Run: An MLflow Run instance representing the started run.
 
     """
     import mlflow
@@ -127,7 +107,7 @@ def start_run(
             description=description,
             log_system_metrics=log_system_metrics,
         ) as run,
-        log_run(config if run_id is None else None, synchronous=synchronous),
+        log_run(run),
     ):
         if chdir:
             with chdir_artifact(run):
@@ -137,7 +117,7 @@ def start_run(
 
 
 @contextmanager
-def chdir_artifact(run: Run | None = None) -> Iterator[Path]:
+def chdir_artifact(run: Run) -> 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

hydraflow/core/io.py

@@ -3,17 +3,13 @@
 from __future__ import annotations
 
 import fnmatch
-import shutil
 import urllib.parse
 import urllib.request
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from hydra.core.hydra_config import HydraConfig
-from omegaconf import DictConfig, ListConfig, OmegaConf
-
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterable, Iterator
+    from collections.abc import Callable, Iterator
 
     from mlflow.entities import Run
 
@@ -27,24 +23,19 @@ def file_uri_to_path(uri: str) -> Path:
     return Path(urllib.request.url2pathname(path))  # for Windows
 
 
-def get_artifact_dir(run: Run | None = None) -> Path:
+def get_artifact_dir(run: Run) -> 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.
+        run (Run | None): The run instance. Defaults to None.
 
     Returns:
         The local path to the directory where the artifacts are downloaded.
 
     """
-    import mlflow
-
-    if run is None:
-        uri = mlflow.get_artifact_uri()
-    else:
-        uri = run.info.artifact_uri
+    uri = run.info.artifact_uri
 
     if not isinstance(uri, str):
         raise NotImplementedError
@@ -52,105 +43,35 @@ def get_artifact_dir(run: Run | None = None) -> Path:
     return file_uri_to_path(uri)
 
 
-def get_artifact_path(run: Run | None, path: str) -> Path:
-    """Retrieve the artifact path for the given run and path.
-
-    This function uses MLflow to get the artifact path for the given run and path.
-
-    Args:
-        run (Run | None): The run object. Defaults to None.
-        path (str): The path to the artifact.
-
-    Returns:
-        The local path to the artifact.
-
-    """
-    return get_artifact_dir(run) / path
-
-
-def get_hydra_output_dir(run: Run | None = None) -> Path:
-    """Retrieve the Hydra output directory for the given run.
-
-    This function returns the Hydra output directory. If no run is provided,
-    it retrieves the output directory from the current Hydra configuration.
-    If a run is provided, it retrieves the artifact path for the run, loads
-    the Hydra configuration from the downloaded artifacts, and returns the
-    output directory specified in that configuration.
-
-    Args:
-        run (Run | None): The run object. Defaults to None.
-
-    Returns:
-        Path: The path to the Hydra output directory.
-
-    Raises:
-        FileNotFoundError: If the Hydra configuration file is not found
-            in the artifacts.
-
-    """
-    if run is None:
-        hc = HydraConfig.get()
-        return Path(hc.runtime.output_dir)
-
-    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_text(run: Run, from_dir: Path, pattern: str = "*.log") -> None:
+    """Log text files in the given directory as artifacts.
 
-def load_config(run: Run) -> DictConfig:
-    """Load the configuration for a given run.
-
-    This function loads the configuration for the provided Run instance
-    by downloading the configuration file from the MLflow artifacts and
-    loading it using OmegaConf. It returns an empty config if
-    `.hydra/config.yaml` is not found in the run's artifact directory.
+    Append the text files to the existing text file in the artifact directory.
 
     Args:
-        run (Run): The Run instance for which to load the configuration.
-
-    Returns:
-        The loaded configuration as a DictConfig object. Returns an empty
-        DictConfig if the configuration file is not found.
+        run (Run): The run instance.
+        from_dir (Path): The directory to find the logs in.
+        pattern (str): The pattern to match the logs.
 
     """
-    path = get_artifact_dir(run) / ".hydra/config.yaml"
-    return OmegaConf.load(path)  # type: ignore
-
-
-def load_overrides(run: Run) -> ListConfig:
-    """Load the overrides for a given run.
-
-    This function loads the overrides for the provided Run instance
-    by downloading the overrides file from the MLflow artifacts and
-    loading it using OmegaConf. It returns an empty config if
-    `.hydra/overrides.yaml` is not found in the run's artifact directory.
-
-    Args:
-        run (Run): The Run instance for which to load the configuration.
-
-    Returns:
-        The loaded configuration as a DictConfig object. Returns an empty
-        DictConfig if the configuration file is not found.
-
-    """
-    path = get_artifact_dir(run) / ".hydra/overrides.yaml"
-    return sorted(OmegaConf.load(path))  # type: ignore
+    import mlflow
 
+    artifact_dir = get_artifact_dir(run)
 
-def remove_run(run: Run | Iterable[Run]) -> None:
-    """Remove the given run from the MLflow tracking server."""
-    from mlflow.entities import Run
+    for file in from_dir.glob(pattern):
+        if not file.is_file():
+            continue
 
-    if not isinstance(run, Run):
-        for r in run:
-            remove_run(r)
-        return
+        file_artifact = artifact_dir / file.name
+        if file_artifact.exists():
+            text = file_artifact.read_text()
+            if not text.endswith("\n"):
+                text += "\n"
+        else:
+            text = ""
 
-    shutil.rmtree(get_artifact_dir(run).parent)
+        text += file.read_text()
+        mlflow.log_text(text, file.name)
 
 
 def get_experiment_name(path: Path) -> str | None:
@@ -187,43 +108,43 @@ def predicate_experiment_dir(
 
 
 def iter_experiment_dirs(
-    root_dir: str | Path,
+    tracking_dir: str | Path,
     experiment_names: str | list[str] | Callable[[str], bool] | None = None,
 ) -> Iterator[Path]:
-    """Iterate over the experiment directories in the root directory."""
+    """Iterate over the experiment directories in the tracking directory."""
     if isinstance(experiment_names, str):
         experiment_names = [experiment_names]
 
-    for path in Path(root_dir).iterdir():
+    for path in Path(tracking_dir).iterdir():
         if predicate_experiment_dir(path, experiment_names):
             yield path
 
 
 def iter_run_dirs(
-    root_dir: str | Path,
+    tracking_dir: str | Path,
     experiment_names: str | list[str] | Callable[[str], bool] | None = None,
 ) -> Iterator[Path]:
-    """Iterate over the run directories in the root directory."""
-    for experiment_dir in iter_experiment_dirs(root_dir, experiment_names):
+    """Iterate over the run directories in the tracking directory."""
+    for experiment_dir in iter_experiment_dirs(tracking_dir, experiment_names):
         for path in experiment_dir.iterdir():
             if path.is_dir() and (path / "artifacts").exists():
                 yield path
 
 
 def iter_artifacts_dirs(
-    root_dir: str | Path,
+    tracking_dir: str | Path,
     experiment_names: str | list[str] | Callable[[str], bool] | None = None,
 ) -> Iterator[Path]:
-    """Iterate over the artifacts directories in the root directory."""
-    for path in iter_run_dirs(root_dir, experiment_names):
+    """Iterate over the artifacts directories in the tracking directory."""
+    for path in iter_run_dirs(tracking_dir, experiment_names):
         yield path / "artifacts"
 
 
 def iter_artifact_paths(
-    root_dir: str | Path,
+    tracking_dir: str | Path,
     artifact_path: str | Path,
     experiment_names: str | list[str] | Callable[[str], bool] | None = None,
 ) -> Iterator[Path]:
-    """Iterate over the artifact paths in the root directory."""
-    for path in iter_artifacts_dirs(root_dir, experiment_names):
+    """Iterate over the artifact paths in the tracking directory."""
+    for path in iter_artifacts_dirs(tracking_dir, experiment_names):
         yield path / artifact_path

hydraflow/core/main.py

@@ -111,7 +111,7 @@ def main(
                     if run.info.status == finished:
                         return
 
-            with start_run(config, run_id=run_id, chdir=chdir) as run:
+            with start_run(run_id=run_id, chdir=chdir) as run:
                 app(run, config)
 
         return inner_decorator
@@ -126,7 +126,7 @@ def get_run_id(uri: str, config: Any, overrides: list[str] | None) -> str | None
 
     Args:
         uri (str): The URI of the experiment.
-        config (object): The configuration object.
+        config (object): The configuration instance.
         overrides (list[str] | None): The task overrides.
 
     Returns:
@@ -146,7 +146,7 @@ def equals(run_dir: Path, config: Any, overrides: list[str] | None) -> bool:
 
     Args:
         run_dir (Path): The run directory.
-        config (object): The configuration object.
+        config (object): The configuration instance.
         overrides (list[str] | None): The task overrides.
 
     Returns: