flyteplugins-nextflow 0.0.0a0__py3-none-any.whl

flyteplugins/nextflow/__init__.py

@@ -0,0 +1,28 @@
+"""
+Run Nextflow pipelines on Flyte.
+
+The Nextflow head runs inside a Flyte task, and the `nf-flyte` Nextflow plugin turns every
+Nextflow task into a child action of that task, so a whole pipeline is a single Flyte run.
+
+```python
+import flyte
+from flyteplugins.nextflow import nextflow_image, run_nextflow
+
+env = flyte.TaskEnvironment(name="nextflow", image=nextflow_image())
+
+@env.task
+async def demo() -> flyte.io.Dir:
+    return await run_nextflow("nf-core/demo", revision="1.2.0", profile="test", outdir="results")
+```
+"""
+
+__all__ = [
+    "DEFAULT_NEXTFLOW_VERSION",
+    "DEFAULT_NF_FLYTE_VERSION",
+    "NextflowError",
+    "nextflow_image",
+    "run_nextflow",
+]
+
+from flyteplugins.nextflow._image import DEFAULT_NEXTFLOW_VERSION, DEFAULT_NF_FLYTE_VERSION, nextflow_image
+from flyteplugins.nextflow._run import NextflowError, run_nextflow

flyteplugins/nextflow/_image.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import flyte
+
+# Nextflow 25.10.6 bundles nf-amazon 3.4.5, the version nf-flyte is built against
+DEFAULT_NEXTFLOW_VERSION = "25.10.6"
+DEFAULT_NF_FLYTE_VERSION = "0.1.0"
+
+NXF_HOME = "/opt/nxf"
+_PLUGIN_ZIP_DIR = "/opt/nf-flyte"
+
+
+def nextflow_image(
+    *,
+    nextflow_version: str = DEFAULT_NEXTFLOW_VERSION,
+    nf_flyte: str | Path = DEFAULT_NF_FLYTE_VERSION,
+    base: flyte.Image | None = None,
+    name: str = "nextflow",
+) -> flyte.Image:
+    """
+    Image for the task that runs the Nextflow head: a JRE, Nextflow and the nf-flyte plugin.
+
+    Pipeline processes don't run in this image; each one runs in its own `container`.
+
+    Args:
+        nextflow_version: Nextflow release to install.
+        nf_flyte: nf-flyte plugin to install, either a version from the Nextflow plugin
+            registry or the path to a locally built plugin zip (`nf-flyte-<version>.zip`).
+        base: Image to add Nextflow to (default: the Flyte debian base image).
+        name: Image name.
+    """
+    image = base or flyte.Image.from_debian_base(name=name)
+    image = image.with_apt_packages("openjdk-17-jre-headless", "curl", "unzip").with_env_vars(
+        {"NXF_HOME": NXF_HOME, "NXF_VER": nextflow_version, "NXF_ANSI_LOG": "false"}
+    )
+
+    if isinstance(nf_flyte, Path):
+        version = _plugin_version_from_zip(nf_flyte)
+        image = image.with_source_file(nf_flyte, f"{_PLUGIN_ZIP_DIR}/")
+        install_plugin = [
+            f"mkdir -p {NXF_HOME}/plugins/nf-flyte-{version}",
+            f"unzip -q {_PLUGIN_ZIP_DIR}/{nf_flyte.name} -d {NXF_HOME}/plugins/nf-flyte-{version}",
+        ]
+    else:
+        version = nf_flyte
+        install_plugin = [f"nextflow plugin install nf-flyte@{version}"]
+
+    return image.with_env_vars({"NF_FLYTE_VERSION": version}).with_commands(
+        [
+            "curl -fsSL https://get.nextflow.io | bash && mv nextflow /usr/local/bin/nextflow",
+            # Download the Nextflow runtime and plugins at build time, not on every run. Without a
+            # version, `plugin install` picks the release bundled with this Nextflow version (e.g.
+            # nf-amazon 3.4.5 on 25.10.6, 3.9.2 on 26.04.x), not the newest in the registry.
+            # nf-cloudcache keeps resume state in the work dir (see run_nextflow).
+            "nextflow info && nextflow plugin install nf-amazon,nf-cloudcache",
+            *install_plugin,
+            f"chmod -R a+rwX {NXF_HOME}",
+        ]
+    )
+
+
+def _plugin_version_from_zip(path: Path) -> str:
+    prefix = "nf-flyte-"
+    if not (path.name.startswith(prefix) and path.suffix == ".zip"):
+        raise ValueError(f"Expected an nf-flyte plugin zip named 'nf-flyte-<version>.zip', got '{path.name}'")
+    # no existence check: task modules are imported again inside the task pod, where the zip
+    # doesn't exist. A missing zip fails the image build instead.
+    return path.name[len(prefix) : -len(".zip")]

flyteplugins/nextflow/_run.py

@@ -0,0 +1,358 @@
+from __future__ import annotations
+
+import asyncio
+import html
+import json
+import os
+import re
+import signal
+import tempfile
+import uuid
+from collections import deque
+from pathlib import Path
+from typing import Any, Mapping, Sequence
+
+from flyteplugins.nextflow._image import DEFAULT_NF_FLYTE_VERSION
+
+import flyte
+from flyte.io import Dir
+from flyte.storage import join
+
+_LOG_TAIL_LINES = 80
+_READ_CHUNK = 64 * 1024
+# Kubernetes gives a terminating pod 30s by default before SIGKILL
+_SHUTDOWN_GRACE_SECONDS = 20
+
+# Nextflow report files rendered into the task's Flyte report, by tab name
+_REPORTS = {"Nextflow report": "report.html", "Timeline": "timeline.html", "DAG": "dag.html"}
+
+
+class NextflowError(RuntimeError):
+    """Raised when `nextflow run` exits with a non-zero status."""
+
+    def __init__(self, exit_code: int, output_tail: Sequence[str], log_tail: Sequence[str]):
+        self.exit_code = exit_code
+        message = [f"nextflow exited with status {exit_code}", *output_tail]
+        if log_tail:
+            message += ["--- .nextflow.log (tail) ---", *log_tail]
+        super().__init__("\n".join(message))
+
+
+async def run_nextflow(
+    pipeline: str,
+    *,
+    revision: str | None = None,
+    profile: str | None = None,
+    params: Mapping[str, Any] | None = None,
+    outdir: str | None = None,
+    work_dir: str | None = None,
+    config: str | Path | None = None,
+    resume: bool = False,
+    report: bool = False,
+    extra_args: Sequence[str] = (),
+) -> Dir | None:
+    """
+    Run a Nextflow pipeline from inside a Flyte task. On a Flyte cluster, every Nextflow task
+    runs as a child action of the calling task, through the nf-flyte executor; the calling task
+    must use `nextflow_image()`. In a local run, Nextflow uses its own executors instead
+    (`local` unless `config` or `profile` selects another, e.g. `docker`).
+
+    Args:
+        pipeline: Pipeline to run: a repository (e.g. `nf-core/rnaseq`), URL or local path.
+        revision: Pipeline revision (git tag, branch or commit).
+        profile: Comma separated Nextflow config profiles, e.g. `test`.
+        params: Pipeline parameters, passed through a params file.
+        outdir: Output directory, passed as `--outdir`. A relative path is placed under this
+            task's raw data prefix. The directory is returned as a `Dir`.
+        work_dir: Work directory; on a cluster, an object store location Flyte can read.
+            Defaults to a `nextflow-work` prefix under this task's raw data prefix, which is shared
+            by all attempts of the task, so a retried task resumes where the failed attempt stopped.
+            Set a fixed location to resume across runs.
+        config: Extra Nextflow config: a path to a config file, or config text.
+        resume: Resume a previous run that used the same `work_dir`. Always on when the task
+            is being retried.
+        report: Render Nextflow's execution report, timeline and DAG into the task's Flyte
+            report. The task must be declared with `report=True`.
+        extra_args: Extra arguments for `nextflow run`.
+
+    Returns:
+        The `outdir` as a `Dir`, or None when no `outdir` is given.
+    """
+    tctx = flyte.ctx()
+    if tctx is None:
+        raise RuntimeError("run_nextflow() must be called from inside a Flyte task")
+    remote = tctx.is_in_cluster()
+
+    storage_root = stable_raw_data_prefix(
+        tctx.raw_data_path.path, tctx.task_action.run_name, tctx.task_action.name, tctx.attempt_number
+    )
+    work_dir = work_dir or join(storage_root, "nextflow-work")
+    if outdir is not None and "://" not in outdir and not os.path.isabs(outdir):
+        outdir = join(storage_root, outdir)
+    if remote and "://" not in work_dir:
+        raise ValueError(f"The nf-flyte executor needs an object store work directory Flyte can read, got '{work_dir}'")
+
+    session_id = session_id_for(work_dir)
+    resume = resume or tctx.attempt_number > 0
+
+    with tempfile.TemporaryDirectory(prefix="nextflow-") as launch_dir:
+        launch = Path(launch_dir)
+        configs = []
+        if remote:
+            (launch / "flyte.config").write_text(
+                flyte_config(os.environ.get("NF_FLYTE_VERSION", DEFAULT_NF_FLYTE_VERSION))
+            )
+            configs.append(launch / "flyte.config")
+        if config is not None:
+            configs.append(_user_config(config, launch))
+
+        params_file = None
+        if params:
+            params_file = launch / "params.json"
+            params_file.write_text(json.dumps(dict(params)))
+
+        cmd = build_command(
+            pipeline,
+            work_dir=work_dir,
+            configs=configs,
+            revision=revision,
+            profile=profile,
+            params_file=params_file,
+            outdir=outdir,
+            resume_session=session_id if resume else None,
+            run_name=None
+            if "-name" in extra_args
+            else nextflow_run_name(tctx.task_action.run_name, tctx.task_action.name, tctx.attempt_number),
+            reports_dir=launch if report else None,
+            extra_args=extra_args,
+        )
+        env = nextflow_env(
+            os.environ,
+            session_id=session_id,
+            cache_path=join(work_dir, ".nextflow-cache"),
+            run_name=tctx.task_action.run_name if remote else None,
+            action_name=tctx.task_action.name if remote else None,
+        )
+        try:
+            exit_code, output_tail = await _stream(cmd, cwd=launch, env=env)
+        finally:
+            if report:
+                await _publish_reports(launch)
+        if exit_code != 0:
+            raise NextflowError(exit_code, output_tail, _tail(launch / ".nextflow.log"))
+
+    return Dir.from_existing_remote(outdir) if outdir is not None else None
+
+
+def stable_raw_data_prefix(raw_data_path: str, run_name: str, action_name: str, attempt: int) -> str:
+    """
+    The part of the task's raw data path shared by all of its attempts. Flyte gives each attempt
+    its own raw data prefix ending in `<run>-<action>-<attempt>`; resuming after a retry needs the
+    work dir to stay put, so drop that last segment.
+    """
+    path = raw_data_path.rstrip("/")
+    parent, _, last = path.rpartition("/")
+    if parent and last == f"{run_name}-{action_name}-{attempt}":
+        return parent
+    return path
+
+
+def session_id_for(work_dir: str) -> str:
+    """
+    Nextflow session id for a work dir. Resuming needs the id of the previous session, which
+    Nextflow normally looks up in `.nextflow/history` under the launch dir; that doesn't survive
+    the task pod, so derive the id from the work dir instead.
+    """
+    return str(uuid.uuid5(uuid.NAMESPACE_URL, work_dir.rstrip("/")))
+
+
+def nextflow_run_name(run_name: str, action_name: str, attempt: int) -> str:
+    """
+    Nextflow normally picks an unused run name from its history file; with history disabled
+    (`NXF_IGNORE_RESUME_HISTORY`), the run name has to be given. It must start with a letter and
+    contain only letters, digits, `-` and `_`, at most 80 characters.
+    """
+    name = re.sub(r"[^a-z0-9]+", "-", f"flyte-{run_name}-{action_name}-{attempt}".lower()).strip("-")
+    return name[:80].rstrip("-")
+
+
+def flyte_config(nf_flyte_version: str) -> str:
+    """Nextflow config that routes every process to the nf-flyte executor."""
+    return f"""\
+plugins {{
+    id 'nf-flyte@{nf_flyte_version}'
+}}
+
+process.executor = 'flyte'
+"""
+
+
+def build_command(
+    pipeline: str,
+    *,
+    work_dir: str,
+    configs: Sequence[Path],
+    revision: str | None = None,
+    profile: str | None = None,
+    params_file: Path | None = None,
+    outdir: str | None = None,
+    resume_session: str | None = None,
+    run_name: str | None = None,
+    reports_dir: Path | None = None,
+    extra_args: Sequence[str] = (),
+) -> list[str]:
+    cmd = ["nextflow"]
+    for c in configs:
+        cmd += ["-c", str(c)]
+    cmd += ["run", pipeline, "-w", work_dir]
+    if revision:
+        cmd += ["-r", revision]
+    if profile:
+        cmd += ["-profile", profile]
+    if params_file is not None:
+        cmd += ["-params-file", str(params_file)]
+    if run_name:
+        cmd += ["-name", run_name]
+    if resume_session:
+        cmd += ["-resume", resume_session]
+    if reports_dir is not None:
+        cmd += [
+            "-with-report", str(reports_dir / _REPORTS["Nextflow report"]),
+            "-with-timeline", str(reports_dir / _REPORTS["Timeline"]),
+            "-with-dag", str(reports_dir / _REPORTS["DAG"]),
+        ]  # fmt: skip
+    cmd += list(extra_args)
+    if outdir is not None:
+        cmd += ["--outdir", outdir]
+    return cmd
+
+
+def nextflow_env(
+    base: Mapping[str, str],
+    *,
+    session_id: str,
+    cache_path: str,
+    run_name: str | None = None,
+    action_name: str | None = None,
+) -> dict[str, str]:
+    """
+    Environment for the Nextflow head. The session id and cloud cache keep resume state outside
+    the launch dir; nf-flyte reads the run and parent action it adds task actions to.
+    """
+    env = dict(base)
+    env["NXF_UUID"] = session_id
+    if "://" in cache_path:
+        env["NXF_CLOUDCACHE_PATH"] = cache_path
+    else:
+        # the cloud cache only takes object store paths; locally, move `.nextflow` instead
+        env["NXF_CACHE_DIR"] = cache_path
+    # resume by the explicit session id: don't require it in the (per-pod) history file
+    env["NXF_IGNORE_RESUME_HISTORY"] = "true"
+    if run_name and action_name:
+        env["NF_FLYTE_RUN_NAME"] = run_name
+        env["NF_FLYTE_PARENT_ACTION"] = action_name
+    env.setdefault("NXF_ANSI_LOG", "false")
+    return env
+
+
+def _user_config(config: str | Path, launch: Path) -> Path:
+    if isinstance(config, Path):
+        return config
+    if "\n" not in config and config.endswith(".config") and Path(config).is_file():
+        return Path(config)
+    path = launch / "user.config"
+    path.write_text(config)
+    return path
+
+
+async def _stream(cmd: Sequence[str], *, cwd: Path, env: Mapping[str, str]) -> tuple[int, list[str]]:
+    print("+", " ".join(cmd), flush=True)
+    proc = await asyncio.create_subprocess_exec(
+        *cmd,
+        cwd=str(cwd),
+        env=dict(env),
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.STDOUT,
+    )
+    tail: deque[str] = deque(maxlen=_LOG_TAIL_LINES)
+    uninstall = _forward_sigterm(proc)
+    try:
+        await _pump_lines(proc.stdout, tail)
+        return await proc.wait(), list(tail)
+    except asyncio.CancelledError:
+        # keep relaying output while Nextflow shuts down: its abort messages belong in the logs
+        await _terminate(proc, _pump_lines(proc.stdout, tail))
+        raise
+    finally:
+        uninstall()
+
+
+def _forward_sigterm(proc: asyncio.subprocess.Process):
+    """
+    Pass SIGTERM on to Nextflow so it can shut down cleanly (abort its task actions, save its
+    cache) when the task pod is deleted, e.g. on abort. Without a handler, the Flyte runtime
+    process dies on SIGTERM and never reaches the cancellation path.
+    """
+
+    def forward():
+        if proc.returncode is None:
+            proc.send_signal(signal.SIGTERM)
+
+    loop = asyncio.get_running_loop()
+    try:
+        loop.add_signal_handler(signal.SIGTERM, forward)
+    except (NotImplementedError, RuntimeError, ValueError):
+        # not the main thread's loop, or not supported on this platform
+        return lambda: None
+    return lambda: loop.remove_signal_handler(signal.SIGTERM)
+
+
+async def _terminate(proc: asyncio.subprocess.Process, drain):
+    if proc.returncode is None:
+        proc.terminate()
+    try:
+        await asyncio.wait_for(asyncio.gather(drain, proc.wait()), timeout=_SHUTDOWN_GRACE_SECONDS)
+    except asyncio.TimeoutError:
+        proc.kill()
+
+
+async def _pump_lines(stream: asyncio.StreamReader | None, tail: deque[str]):
+    """Echo output line by line. Reads in chunks: `readline()` fails on lines over 64 KiB."""
+    assert stream is not None
+    buf = b""
+    while chunk := await stream.read(_READ_CHUNK):
+        buf += chunk
+        *lines, buf = buf.split(b"\n")
+        for raw in lines:
+            _emit(raw, tail)
+    if buf:
+        _emit(buf, tail)
+
+
+def _emit(raw: bytes, tail: deque[str]):
+    line = raw.decode(errors="replace").rstrip("\r")
+    print(line, flush=True)
+    tail.append(line)
+
+
+async def _publish_reports(launch: Path):
+    import flyte.report
+
+    published = False
+    for tab, name in _REPORTS.items():
+        path = launch / name
+        if path.exists():
+            # the reports are full HTML documents; a tab's content goes inside a <div>
+            doc = html.escape(path.read_text(errors="replace"), quote=True)
+            flyte.report.get_tab(tab).replace(
+                f'<iframe srcdoc="{doc}" style="width:100%;height:85vh;border:0"></iframe>'
+            )
+            published = True
+    if published:
+        await flyte.report.flush.aio()
+
+
+def _tail(path: Path) -> list[str]:
+    if not path.exists():
+        return []
+    return path.read_text(errors="replace").splitlines()[-_LOG_TAIL_LINES:]

flyteplugins_nextflow-0.0.0a0.dist-info/METADATA

@@ -0,0 +1,99 @@
+Metadata-Version: 2.4
+Name: flyteplugins-nextflow
+Version: 0.0.0a0
+Summary: Run Nextflow pipelines on Flyte, with every Nextflow task as a Flyte action
+Author-email: Kyle Hazen <kyle@union.ai>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: flyte
+
+# flyteplugins-nextflow
+
+Run [Nextflow](https://www.nextflow.io/) pipelines on Flyte.
+
+The Nextflow head runs inside a Flyte task. The [nf-flyte](https://github.com/unionai/nf-flyte)
+Nextflow plugin runs every Nextflow task as a child action of that task. A pipeline
+therefore shows up in Flyte as one run with one action per Nextflow task:
+
+- Each process runs in its own `container`, with its `cpus`, `memory`, `disk`,
+  `accelerator` and `time` directives. The image only needs `bash`: Flyte stages the task's files.
+- Nextflow keeps its own retries (`errorStrategy`) and `-resume`; Flyte caches task results across runs.
+- Aborting the Flyte run aborts the pipeline and its running tasks.
+
+## Installation
+
+```bash
+pip install flyteplugins-nextflow
+```
+
+## Quick start
+
+```python
+import flyte
+from flyte.io import Dir
+from flyteplugins.nextflow import nextflow_image, run_nextflow
+
+env = flyte.TaskEnvironment(name="nextflow", image=nextflow_image())
+
+
+@env.task
+async def demo() -> Dir:
+    return await run_nextflow("nf-core/demo", revision="1.2.0", profile="test", outdir="results")
+```
+
+`run_nextflow` accepts:
+
+| Argument     | Description |
+|--------------|-------------|
+| `pipeline`   | Repository (`nf-core/rnaseq`), URL or path to run. |
+| `revision`   | Git tag, branch or commit (`-r`). |
+| `profile`    | Config profiles (`-profile`). |
+| `params`     | Pipeline parameters, passed as a params file. |
+| `outdir`     | `--outdir`. A relative path goes under the task's raw data prefix. The directory is returned as a `Dir`. |
+| `work_dir`   | Work directory. Default: `nextflow-work` under the task's raw data prefix. |
+| `config`     | Extra Nextflow config, as a file path or config text. |
+| `resume`     | Resume the previous run that used the same `work_dir`. |
+| `report`     | Render Nextflow's execution report, timeline and DAG into the task's Flyte report. The task must be declared with `report=True`. |
+| `extra_args` | Any other `nextflow run` arguments. |
+
+If Nextflow fails, `run_nextflow` raises a `NextflowError` containing the end of the
+Nextflow output and the end of `.nextflow.log`.
+
+### Resume
+
+Nextflow's resume state lives in the work directory, not in the task pod:
+- the task cache (the cloud cache on object stores, `.nextflow` locally)
+- a session ID derived from the work directory
+
+This has two effects:
+- **Retries resume automatically.** The default work directory is shared by all attempts of
+  a task, so a retried task passes `-resume` and picks up where the failed attempt stopped.
+- **Resuming across runs needs a fixed `work_dir`.** Pass the same `work_dir` to each run,
+  and `resume=True` from the second run on:
+
+  ```python
+  await run_nextflow("nf-core/rnaseq", work_dir="s3://my-bucket/rnaseq-work", resume=True)
+  ```
+
+### Local runs
+
+When the task runs locally (`flyte run --local`), Nextflow uses its own executors instead
+of nf-flyte: `local` by default, or whatever `profile` or `config` selects, such as `docker`.
+This needs `nextflow` on the `PATH`.
+
+## Image
+
+`nextflow_image()` builds the image for the task that runs the Nextflow head. It contains
+a JRE, Nextflow (`nextflow_version`, default `25.10.6`) and the nf-flyte plugin. To use a
+locally built plugin, pass its zip:
+
+```python
+from pathlib import Path
+
+image = nextflow_image(nf_flyte=Path("~/nf-flyte/build/distributions/nf-flyte-0.1.0.zip").expanduser())
+```
+
+## Requirements
+
+- Flyte v2. On a cluster, the work directory must be on an object store Flyte can read (S3, GCS or Azure); the default, under the task's raw data prefix, always is.
+- Every process container needs `bash`.

flyteplugins_nextflow-0.0.0a0.dist-info/RECORD

@@ -0,0 +1,7 @@
+flyteplugins/nextflow/__init__.py,sha256=4kbhKMYm909CnCw1aCdAQhkRg0FnAA3a49YAwWc6GAA,846
+flyteplugins/nextflow/_image.py,sha256=x96z35fqG11WjGHJxBilucbs6r-fqnUc0LsMbolN4Go,2958
+flyteplugins/nextflow/_run.py,sha256=8l0sFqFwTaYbDjk-oHp-6S4O7RhzXvuhyfmN-PAuugw,13024
+flyteplugins_nextflow-0.0.0a0.dist-info/METADATA,sha256=o8mJbpxrIL10Wl-si36LMy9Eovvhttys9kjqxg4hI8o,3892
+flyteplugins_nextflow-0.0.0a0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+flyteplugins_nextflow-0.0.0a0.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_nextflow-0.0.0a0.dist-info/RECORD,,

flyteplugins_nextflow-0.0.0a0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flyteplugins_nextflow-0.0.0a0.dist-info/top_level.txt

@@ -0,0 +1 @@
+flyteplugins