app/backend/fluksio/flow/modules.py

"""The python packages node code may import, in a venv of the user's own.

A pip manifest lives beside the flows in the same git repository, so what a
deployment installed is versioned with what uses it. The packages themselves go
into a venv on the data volume rather than into the engine's environment: a
pin here can never shadow — or be shadowed by — what the app itself runs on,
and the worker processes that import them need nothing from the app.

``uv pip sync`` rather than install, so a line taken out of the manifest is
uninstalled. The manifest is written only after a sync succeeds, which is all
the rollback a failed resolve needs.
"""

from __future__ import annotations

import hashlib
import importlib.metadata
import logging
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING

from fluksio.core.config import settings
from fluksio.flow.schemas import ModulePackage, ModulesInfo

if TYPE_CHECKING:
    from fluksio.flow.store import FlowStore

logger = logging.getLogger(__name__)

#: Beside the flow store rather than in it: this is installed state, not source.
VENV_DIR = settings.FLOWS_DIR.parent / "user-venv"

#: A resolve that takes longer than this is not going to finish.
SYNC_TIMEOUT = 300


def uv_bin() -> str:
    """Where ``uv`` is, without depending on what PATH happens to hold.

    It is a dependency, so it is installed beside the interpreter running this
    — which is what a systemd unit naming an absolute ExecStart, or a container
    entrypoint, would otherwise miss.
    """
    beside = Path(sys.executable).with_name("uv")
    if beside.exists():
        return str(beside)
    return shutil.which("uv") or "uv"


def venv_python() -> str:
    """The interpreter node code runs on.

    Falls back to the engine's own when there is no venv — an installation that
    could not build one still runs python nodes, it just cannot add packages
    to them.
    """
    path = VENV_DIR / "bin" / "python"
    return str(path) if path.exists() else sys.executable


def _marker() -> Path:
    return VENV_DIR / ".applied"


def _digest(requirements: str) -> str:
    version = f"{sys.version_info.major}.{sys.version_info.minor}"
    return hashlib.sha256(f"{version}\n{requirements}".encode()).hexdigest()


def ensure_venv() -> None:
    """Create the venv if it is missing or its interpreter has gone."""
    if (VENV_DIR / "bin" / "python").exists():
        return
    VENV_DIR.parent.mkdir(parents=True, exist_ok=True)
    # From the base interpreter, not from the engine's venv: nesting one venv
    # inside another is how a user pin ends up resolving against app packages.
    subprocess.run(
        [
            uv_bin(),
            "venv",
            "--python",
            str(Path(sys.base_prefix, "bin", "python3")),
            str(VENV_DIR),
        ],
        capture_output=True,
        text=True,
        check=True,
        timeout=SYNC_TIMEOUT,
    )


def sync(requirements: str) -> tuple[bool, str]:
    """Make the venv hold exactly these packages. Returns success and uv's output.

    ``uv`` missing, or refusing to make the venv, is a failed apply like any
    other — the caller answers 400 with what came back, which is the only thing
    a person can act on.
    """
    manifest = ""
    try:
        ensure_venv()
        with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as handle:
            handle.write(requirements)
            manifest = handle.name
        result = subprocess.run(
            # Empty means empty: without the flag uv refuses to clear a venv,
            # so deleting the last line would leave the package installed.
            [
                uv_bin(),
                "pip",
                "sync",
                "--allow-empty-requirements",
                "--python",
                venv_python(),
                manifest,
            ],
            capture_output=True,
            text=True,
            timeout=SYNC_TIMEOUT,
        )
    except (OSError, subprocess.SubprocessError) as exc:
        detail = str(getattr(exc, "stderr", None) or exc).strip()
        return False, f"uv could not run: {detail}"
    finally:
        if manifest:
            Path(manifest).unlink(missing_ok=True)

    output = (result.stdout + result.stderr).strip()
    if result.returncode == 0:
        _marker().write_text(_digest(requirements))
    return result.returncode == 0, output


def reconcile(store: FlowStore) -> None:
    """Bring the venv in line with the stored manifest. Blocking.

    Called at startup, where nothing may be fatal: a deployment whose packages
    cannot be installed still runs, with the nodes needing them reporting an
    import error each time they execute.
    """
    try:
        # First, and unconditionally: the workers are started against this
        # interpreter, and it has to be the venv's one before anything is
        # installed into it, not after.
        ensure_venv()
        requirements = store.read_requirements()
        if not requirements.strip() and not _marker().exists():
            return
        if _marker().exists() and _marker().read_text() == _digest(requirements):
            return
        ok, output = sync(requirements)
        if not ok:
            logger.warning("Could not install the stored modules: %s", output)
    except Exception:
        logger.exception("Could not reconcile the module venv")


def info(store: FlowStore) -> ModulesInfo:
    """What is installed, what was asked for, and whether the two agree."""
    requirements = store.read_requirements()
    config = VENV_DIR / "pyvenv.cfg"
    version = ""
    if config.exists():
        for line in config.read_text().splitlines():
            if line.startswith("version"):
                version = line.split("=", 1)[1].strip()

    packages: list[ModulePackage] = []
    site = sorted(VENV_DIR.glob("lib/python*/site-packages"))
    if site:
        packages = sorted(
            (
                ModulePackage(name=dist.metadata["Name"] or "", version=dist.version)
                for dist in importlib.metadata.distributions(path=[str(site[0])])
            ),
            key=lambda package: package.name.lower(),
        )

    return ModulesInfo(
        python_version=version,
        venv_path=str(VENV_DIR),
        requirements=requirements,
        packages=packages,
        applied=(
            _marker().read_text() == _digest(requirements)
            if _marker().exists()
            # Nothing asked for and nothing installed is already in step.
            else not requirements.strip()
        ),
    )