hydraflow 0.14.0__tar.gz → 0.14.2__tar.gz

{hydraflow-0.14.0 → hydraflow-0.14.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hydraflow
-Version: 0.14.0
+Version: 0.14.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
@@ -40,6 +40,7 @@ Requires-Dist: hydra-core>=1.3
 Requires-Dist: mlflow>=2.15
 Requires-Dist: omegaconf>=2.3
 Requires-Dist: python-ulid>=3.0.0
+Requires-Dist: rich>=13.9
 Requires-Dist: typer>=0.15
 Description-Content-Type: text/markdown
 

{hydraflow-0.14.0 → hydraflow-0.14.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "hydraflow"
-version = "0.14.0"
+version = "0.14.2"
 description = "Hydraflow integrates Hydra and MLflow to manage and track machine learning experiments."
 readme = "README.md"
 license = { file = "LICENSE" }
@@ -24,6 +24,7 @@ dependencies = [
   "mlflow>=2.15",
   "omegaconf>=2.3",
   "python-ulid>=3.0.0",
+  "rich>=13.9",
   "typer>=0.15",
 ]
 

{hydraflow-0.14.0 → hydraflow-0.14.2}/src/hydraflow/cli.py

@@ -14,8 +14,8 @@ if TYPE_CHECKING:
 app = typer.Typer(add_completion=False)
 
 
-@app.command(context_settings={"ignore_unknown_options": True})
-def run(
+@app.command("run", context_settings={"ignore_unknown_options": True})
+def _run(
     name: Annotated[str, Argument(help="Job name.", show_default=False)],
     *,
     args: Annotated[
@@ -29,55 +29,73 @@ def run(
 ) -> None:
     """Run a job."""
     from hydraflow.executor.io import get_job
-    from hydraflow.executor.job import iter_batches, iter_calls, iter_runs
 
     args = args or []
     job = get_job(name)
 
+    if not dry_run:
+        import mlflow
+
+        mlflow.set_experiment(job.name)
+
     if job.submit:
         submit(job, args, dry_run=dry_run)
-        raise Exit
 
-    if job.run:
-        args = [*shlex.split(job.run), *args]
-        it = iter_runs(args, iter_batches(job), dry_run=dry_run)
+    elif job.run:
+        run(job, args, dry_run=dry_run)
+
     elif job.call:
-        args = [*shlex.split(job.call), *args]
-        it = iter_calls(args, iter_batches(job), dry_run=dry_run)
+        call(job, args, dry_run=dry_run)
+
     else:
         typer.echo(f"No command found in job: {job.name}.")
         raise Exit(1)
 
+
+def run(job: Job, args: list[str], *, dry_run: bool) -> None:
+    """Run a job."""
+    from hydraflow.executor import aio
+    from hydraflow.executor.job import iter_batches, iter_tasks
+
+    args = [*shlex.split(job.run), *args]
+    it = iter_tasks(args, iter_batches(job))
+
     if not dry_run:
-        import mlflow
+        aio.run(it)
+        raise Exit
 
-        mlflow.set_experiment(job.name)
+    for task in it:
+        typer.echo(shlex.join(task.args))
 
-    for task in it:  # jobs will be executed here
-        if job.run and dry_run:
-            typer.echo(shlex.join(task.args))
-        elif job.call and dry_run:
-            funcname, *args = task.args
-            arg = ", ".join(f"{arg!r}" for arg in args)
-            typer.echo(f"{funcname}([{arg}])")
 
+def call(job: Job, args: list[str], *, dry_run: bool) -> None:
+    """Call a job."""
+    from hydraflow.executor.job import iter_batches, iter_calls
 
-def submit(job: Job, args: list[str], *, dry_run: bool) -> None:
-    """Submit a job."""
-    from hydraflow.executor.job import iter_batches, submit
+    args = [*shlex.split(job.call), *args]
+    it = iter_calls(args, iter_batches(job))
 
     if not dry_run:
-        import mlflow
+        for call in it:
+            call.func()
+        raise Exit
 
-        mlflow.set_experiment(job.name)
+    for task in it:
+        funcname, *args = task.args
+        arg = ", ".join(f"{arg!r}" for arg in args)
+        typer.echo(f"{funcname}([{arg}])")
+
+
+def submit(job: Job, args: list[str], *, dry_run: bool) -> None:
+    """Submit a job."""
+    from hydraflow.executor.job import iter_batches, submit
 
     args = [*shlex.split(job.submit), *args]
     result = submit(args, iter_batches(job), dry_run=dry_run)
 
     if dry_run and isinstance(result, tuple):
-        for line in result[1].splitlines():
-            args = shlex.split(line)
-            typer.echo(shlex.join([*result[0][:-1], *args]))
+        typer.echo(shlex.join(result[0]))
+        typer.echo(result[1])
 
 
 @app.command()

hydraflow-0.14.2/src/hydraflow/executor/aio.py

@@ -0,0 +1,85 @@
+"""Asynchronous execution."""
+
+from __future__ import annotations
+
+import asyncio
+from asyncio.subprocess import PIPE
+from typing import TYPE_CHECKING
+
+from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TaskProgressColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+
+if TYPE_CHECKING:
+    from asyncio.streams import StreamReader
+    from collections.abc import Callable, Iterable
+
+    from hydraflow.executor.job import Task
+
+
+console = Console(log_time=False, log_path=False)
+
+
+def run(iterable: Iterable[Task]) -> int | None:
+    """Run multiple tasks."""
+    with Progress(
+        SpinnerColumn(),
+        TimeElapsedColumn(),
+        BarColumn(),
+        TimeRemainingColumn(),
+        MofNCompleteColumn(),
+        TaskProgressColumn(),
+        console=console,
+    ) as progress:
+
+        def stdout(output: str) -> None:
+            progress.log(output.rstrip())
+
+        def stderr(output: str) -> None:
+            progress.log(f"[red]{output}".rstrip())
+
+        task_id = progress.add_task("")
+
+        for task in iterable:
+            progress.update(task_id, total=task.total)
+
+            coro = arun(task.args, stdout, stderr)
+            returncode = asyncio.run(coro)
+
+            if returncode:
+                return returncode
+
+            progress.update(task_id, completed=task.index + 1)
+
+        return 0
+
+
+async def arun(
+    args: list[str],
+    stdout: Callable[[str], None],
+    stderr: Callable[[str], None],
+) -> int | None:
+    """Run a command asynchronously."""
+    process = await asyncio.create_subprocess_exec(*args, stdout=PIPE, stderr=PIPE)
+    coros = alog(process.stdout, stdout), alog(process.stderr, stderr)  # type:ignore
+    await asyncio.gather(*coros)
+    await process.communicate()
+
+    return process.returncode
+
+
+async def alog(reader: StreamReader, write: Callable[[str], None]) -> None:
+    """Log a stream of output asynchronously."""
+    while True:
+        if reader.at_eof():
+            break
+
+        if out := await reader.readline():
+            write(out.decode())

{hydraflow-0.14.0 → hydraflow-0.14.2}/src/hydraflow/executor/job.py

@@ -25,7 +25,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from subprocess import CompletedProcess
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, overload
+from typing import TYPE_CHECKING
 
 import ulid
 
@@ -86,47 +86,22 @@ def iter_batches(job: Job) -> Iterator[list[str]]:
 
 @dataclass
 class Task:
-    """An executed task."""
+    """A task to be executed."""
 
     args: list[str]
     total: int
-    completed: int
-
-
-@dataclass
-class Run(Task):
-    """An executed run."""
-
-    result: CompletedProcess
+    index: int
 
 
 @dataclass
 class Call(Task):
-    """An executed call."""
+    """A call to be executed."""
 
-    result: Any
+    func: Callable[[], Any]
 
 
-@overload
-def iter_runs(args: list[str], iterable: Iterable[list[str]]) -> Iterator[Run]: ...
-
-
-@overload
-def iter_runs(
-    args: list[str],
-    iterable: Iterable[list[str]],
-    *,
-    dry_run: bool = False,
-) -> Iterator[Task | Run]: ...
-
-
-def iter_runs(
-    args: list[str],
-    iterable: Iterable[list[str]],
-    *,
-    dry_run: bool = False,
-) -> Iterator[Task | Run]:
-    """Execute multiple runs of a job using shell commands."""
+def iter_tasks(args: list[str], iterable: Iterable[list[str]]) -> Iterator[Task]:
+    """Yield tasks of a job to be executed using a shell command."""
     executable, *args = args
     if executable == "python" and sys.platform == "win32":
         executable = sys.executable
@@ -134,48 +109,21 @@ def iter_runs(
     iterable = list(iterable)
     total = len(iterable)
 
-    for completed, args_ in enumerate(iterable, 1):
-        cmd = [executable, *args, *args_]
-        if dry_run:
-            yield Task(cmd, total, completed)
-        else:
-            result = subprocess.run(cmd, check=False)
-            yield Run(cmd, total, completed, result)
-
-
-@overload
-def iter_calls(args: list[str], iterable: Iterable[list[str]]) -> Iterator[Call]: ...
+    for index, args_ in enumerate(iterable):
+        yield Task([executable, *args, *args_], total, index)
 
 
-@overload
-def iter_calls(
-    args: list[str],
-    iterable: Iterable[list[str]],
-    *,
-    dry_run: bool = False,
-) -> Iterator[Task | Call]: ...
-
-
-def iter_calls(
-    args: list[str],
-    iterable: Iterable[list[str]],
-    *,
-    dry_run: bool = False,
-) -> Iterator[Task | Call]:
-    """Execute multiple calls of a job using Python functions."""
+def iter_calls(args: list[str], iterable: Iterable[list[str]]) -> Iterator[Call]:
+    """Yield calls of a job to be executed using a Python function."""
     funcname, *args = args
     func = get_callable(funcname)
 
     iterable = list(iterable)
     total = len(iterable)
 
-    for completed, args_ in enumerate(iterable, 1):
+    for index, args_ in enumerate(iterable):
         cmd = [funcname, *args, *args_]
-        if dry_run:
-            yield Task(cmd, total, completed)
-        else:
-            result = func([*args, *args_])
-            yield Call(cmd, total, completed, result)
+        yield Call(cmd, total, index, lambda x=cmd[1:]: func(x))
 
 
 def submit(

{hydraflow-0.14.0 → hydraflow-0.14.2}/tests/cli/test_run.py

@@ -57,7 +57,14 @@ def test_submit_dry_run():
     args = ["run", "submit", "--dry-run", "a", "--b", "--", "--dry-run"]
     result = runner.invoke(app, args)
     assert result.exit_code == 0
-    assert result.stdout.count("submit.py a --b --dry-run --multirun") == 4
+    assert result.stdout.count("submit.py a --b --dry-run") == 1
+    assert result.stdout.count("--multirun") == 4
+    lines = result.stdout.splitlines()
+    assert len(lines) == 5
+    assert "name=a count=1" in lines[1]
+    assert "name=b count=1" in lines[2]
+    assert "name=c count=5" in lines[3]
+    assert "name=d count=6" in lines[4]
 
 
 @pytest.mark.xdist_group(name="group1")
@@ -83,6 +90,11 @@ def test_run_parallel():
     run_ids = hydraflow.list_run_ids("parallel")
     assert len(run_ids) == 8
 
+    result = runner.invoke(app, ["run", "parallel"])  # skip if already run
+    assert result.exit_code == 0
+    run_ids = hydraflow.list_run_ids("parallel")
+    assert len(run_ids) == 8
+
 
 @pytest.mark.xdist_group(name="group4")
 def test_run_echo():
@@ -107,6 +119,11 @@ def test_submit():
     run_ids = hydraflow.list_run_ids("submit")
     assert len(run_ids) == 4
 
+    result = runner.invoke(app, ["run", "submit"])  # skip if already run
+    assert result.exit_code == 0
+    run_ids = hydraflow.list_run_ids("submit")
+    assert len(run_ids) == 4
+
 
 @pytest.mark.xdist_group(name="group5")
 def test_run_error():

hydraflow-0.14.2/tests/executor/test_aio.py

@@ -0,0 +1,22 @@
+import sys
+
+import pytest
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on Windows")
+def test_run_returncode():
+    from hydraflow.executor.aio import run
+    from hydraflow.executor.job import Task
+
+    task = Task(args=["false"], index=0, total=1)
+
+    assert run([task]) == 1
+
+
+def test_run_stderr():
+    from hydraflow.executor.aio import run
+    from hydraflow.executor.job import Task
+
+    task = Task(args=["python", "-c", "1/0"], index=0, total=1)
+
+    assert run([task]) == 1

{hydraflow-0.14.0 → hydraflow-0.14.2}/tests/executor/test_job.py

@@ -59,40 +59,36 @@ def test_sweep_args(batches, i, x):
     assert batches[i][-3] == x
 
 
-def test_iter_runs(job: Job, tmp_path: Path):
-    from hydraflow.executor.job import iter_batches, iter_runs
+def test_iter_tasks(job: Job, tmp_path: Path):
+    import subprocess
+
+    from hydraflow.executor.job import iter_batches, iter_tasks
 
     path = tmp_path / "output.txt"
     file = Path(__file__).parent / "echo.py"
 
     args = ["python", file.as_posix(), path.as_posix()]
-    x = list(iter_runs(args, iter_batches(job)))
+    for task in iter_tasks(args, iter_batches(job)):
+        subprocess.run(task.args, check=True)
     assert path.read_text() == "b=5 a=1,2 b=6 a=1,2 c=7 a=3,4 c=8 a=3,4"
-    assert x[0].completed == 1
-    assert x[0].result.returncode == 0
-    assert x[1].completed == 2
-    assert x[1].result.returncode == 0
-    assert x[2].completed == 3
-    assert x[2].result.returncode == 0
 
 
 def test_iter_calls(job: Job, capsys: pytest.CaptureFixture):
     from hydraflow.executor.job import iter_batches, iter_calls
 
-    x = list(iter_calls(["typer.echo"], iter_batches(job)))
+    for call in iter_calls(["typer.echo"], iter_batches(job)):
+        call.func()
     out, _ = capsys.readouterr()
     assert "'b=5', 'a=1,2'" in out
     assert "'c=8', 'a=3,4'" in out
-    assert x[0].completed == 1
-    assert x[1].completed == 2
-    assert x[2].completed == 3
 
 
 def test_iter_calls_args(job: Job, capsys: pytest.CaptureFixture):
     from hydraflow.executor.job import iter_batches, iter_calls
 
     job.call = "typer.echo a 'b c'"
-    list(iter_calls(["typer.echo", "a", "b c"], iter_batches(job)))
+    for call in iter_calls(["typer.echo", "a", "b c"], iter_batches(job)):
+        call.func()
     out, _ = capsys.readouterr()
     assert "['a', 'b c', '--multirun'," in out