hydraflow 0.2.11__py3-none-any.whl → 0.2.12__py3-none-any.whl

hydraflow/mlflow.py

@@ -22,6 +22,7 @@ from __future__ import annotations
 from pathlib import Path
 from typing import TYPE_CHECKING
 
+import joblib
 import mlflow
 from hydra.core.hydra_config import HydraConfig
 from mlflow.entities import ViewType
@@ -146,7 +147,9 @@ def search_runs(
     return RunCollection(runs)  # type: ignore
 
 
-def list_runs(experiment_names: list[str] | None = None) -> RunCollection:
+def list_runs(
+    experiment_names: str | list[str] | None = None, *, n_jobs: int = 0
+) -> RunCollection:
     """
     List all runs for the specified experiments.
 
@@ -168,8 +171,30 @@ def list_runs(experiment_names: list[str] | None = None) -> RunCollection:
     Returns:
         A `RunCollection` object containing the runs for the specified experiments.
     """
-    if experiment_names == []:
+    if isinstance(experiment_names, str):
+        experiment_names = [experiment_names]
+
+    elif experiment_names == []:
         experiments = mlflow.search_experiments()
         experiment_names = [e.name for e in experiments if e.name != "Default"]
 
-    return search_runs(experiment_names=experiment_names)
+    if n_jobs == 0:
+        return search_runs(experiment_names=experiment_names)
+
+    if experiment_names is None:
+        raise NotImplementedError
+
+    run_ids = []
+
+    for name in experiment_names:
+        if experiment := mlflow.get_experiment_by_name(name):
+            loc = experiment.artifact_location
+
+            if isinstance(loc, str) and loc.startswith("file://"):
+                path = Path(mlflow.artifacts.download_artifacts(loc))
+                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)
+    runs = joblib.Parallel(n_jobs, prefer="threads")(it)
+    runs = sorted(runs, key=lambda run: run.info.start_time)  # type: ignore
+    return RunCollection(runs)  # type: ignore

{hydraflow-0.2.11.dist-info → hydraflow-0.2.12.dist-info}/METADATA

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

{hydraflow-0.2.11.dist-info → hydraflow-0.2.12.dist-info}/RECORD

@@ -3,10 +3,10 @@ hydraflow/asyncio.py,sha256=jdXuEFC6f7L_Dq6beASFZPQSvCnGimVxU-PRFsNc5U0,6241
 hydraflow/config.py,sha256=6TCKNQZ3sSrIEvl245T2udwFuknejyN1dMcIVmOHdrQ,2102
 hydraflow/context.py,sha256=G7JMrG70sgBH2qILXl5nkGWNUoRggj518JWUq0ZiJ9E,7776
 hydraflow/info.py,sha256=Vj2sT66Ric63mmaq7Yu8nDFhsGQYO3MCHrxFpapDufc,3458
-hydraflow/mlflow.py,sha256=Q8RGijSURTjRkEDxzi_2Tk9KOx3QK__al5aArGQriHA,7249
+hydraflow/mlflow.py,sha256=YWfq_d8qG6d3k7GC3hWJ7hYHGMiSZ4CnJEyr4vU8t6U,8118
 hydraflow/progress.py,sha256=uaxS9UnwsikOPB_Ho-5mXdBYz3Zedj-SL9FzT094wjY,6623
 hydraflow/run_collection.py,sha256=V5lGdGHYgsSpBOYGaVEL1mpKJvdiEshBL0KmmZ8qeZo,29161
-hydraflow-0.2.11.dist-info/METADATA,sha256=skZKh1wvBJWSLXE7iB6q2Gv_NsegcuJQjmmv8KgFg8s,4182
-hydraflow-0.2.11.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-hydraflow-0.2.11.dist-info/licenses/LICENSE,sha256=IGdDrBPqz1O0v_UwCW-NJlbX9Hy9b3uJ11t28y2srmY,1062
-hydraflow-0.2.11.dist-info/RECORD,,
+hydraflow-0.2.12.dist-info/METADATA,sha256=BIdR4v_bxeeQn4lYS7d4j331cz57HeOkkBtg6FK6YrQ,4182
+hydraflow-0.2.12.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+hydraflow-0.2.12.dist-info/licenses/LICENSE,sha256=IGdDrBPqz1O0v_UwCW-NJlbX9Hy9b3uJ11t28y2srmY,1062
+hydraflow-0.2.12.dist-info/RECORD,,