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.

None of that applies to an *adopted* venv. A data scientist makes a venv,
installs what they work with, and then installs Fluksio into it too — at which
point building a second environment beside it is exactly wrong: the packages
the nodes need are already here. So when the engine is running from a venv of
somebody else's, node code runs on it. That venv is theirs: `uv pip sync` is
never pointed at it, because sync means "hold exactly this" and would uninstall
their work along with the engine. It is read-only here, and `pip` is how they
change it.
"""

from __future__ import annotations

import hashlib
import importlib.metadata
import logging
import shutil
import subprocess
import sys
import tempfile
import threading
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 adopted() -> Path | None:
    """The venv this engine was installed into, when node code should use it.

    ``None`` means the managed venv — one the engine builds and owns. The three
    ways that is the answer, in order:

    * ``NODE_VENV=managed`` says so. A container sets this: its venv holds the
      app and nothing of anybody's, so there is nothing to adopt.
    * There is already a managed venv. It may have packages in it that
      somebody installed on purpose, and an upgrade must not take them away.
    * The engine is not running from a venv at all.

    Anything else — ``pip install fluksio`` into the environment you already
    work in — is the case this exists for.
    """
    setting = settings.NODE_VENV.strip()
    if setting == "managed":
        return None
    if setting and setting != "auto":
        return Path(setting)
    if (VENV_DIR / "bin" / "python").exists():
        return None
    prefix = Path(sys.prefix)
    if prefix == Path(sys.base_prefix) or prefix == VENV_DIR:
        return None
    return prefix


def venv_dir() -> Path:
    """Whichever venv node code runs from, adopted or managed."""
    return adopted() or VENV_DIR


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.
    """
    root = adopted()
    if root is not None:
        # NODE_VENV may name a venv or the interpreter inside one; both are
        # useful things to be handed, and they are told apart by shape.
        inner = root / "bin" / "python"
        return str(inner if inner.exists() or root.is_dir() else root)
    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 adopted() is not None:
        return
    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,
    )


#: Held for the length of a `uv pip sync`. Two applies at once — two browser
#: tabs, or an apply landing while the boot-time reconcile is still running —
#: mutate the same venv simultaneously.
_sync_lock = threading.Lock()


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.
    """
    if adopted() is not None:
        # The one thing this must never do. `uv pip sync` makes a venv hold
        # exactly the manifest, so pointed at somebody's own environment it
        # uninstalls their packages — and the engine with them.
        return False, (
            "These packages are not Fluksio's to install: node code runs on "
            f"{venv_python()}, the environment Fluksio itself was installed "
            "into. Install into it with pip or uv, or set NODE_VENV=managed "
            "for a venv the engine owns."
        )
    if not _sync_lock.acquire(timeout=0.0):
        return False, (
            "Another apply is already installing packages. Wait for it to "
            "finish, then try again."
        )
    try:
        return _sync(requirements)
    finally:
        _sync_lock.release()


def _sync(requirements: str) -> tuple[bool, str]:
    """The install itself, under `_sync_lock`."""
    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:
        root = adopted()
        if root is not None:
            # Nothing to reconcile: the environment is somebody else's, and
            # bringing it "in line with the manifest" would empty it.
            logger.info("Node code runs on the adopted venv at %s", root)
            return
        # 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()
    root = venv_dir()
    is_adopted = adopted() is not None

    version = ""
    config = root / "pyvenv.cfg"
    if config.exists():
        for line in config.read_text().splitlines():
            if line.startswith("version"):
                version = line.split("=", 1)[1].strip()
    if not version:
        version = ".".join(str(part) for part in sys.version_info[:3])

    packages: list[ModulePackage] = []
    site = sorted(root.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(root),
        requirements=requirements,
        packages=packages,
        # An adopted venv is never "out of step": the manifest does not
        # describe it, so there is nothing for it to disagree with.
        applied=(
            True
            if is_adopted
            else (
                _marker().read_text() == _digest(requirements)
                if _marker().exists()
                # Nothing asked for and nothing installed is already in step.
                else not requirements.strip()
            )
        ),
        adopted=is_adopted,
    )