Source code for thesis.core.nipype.interfaces.fsl

"""
Custom FSL Nipype interfaces.

Provides patched or extended versions of FSL Nipype interfaces to work around
upstream bugs or limitations, and adds GPU-accelerated variants where FSL
ships separate GPU binaries.

CUDA / GPU validation is performed **once at CLI startup** by
:func:`thesis.core.gpu.check_gpu`, not here.  By the time a workflow is
built, ``config.hardware.gpu_enabled`` already reflects whether a compatible
GPU setup was found.  This module's only responsibility is to map that flag
to the correct binary name.

Classes:
    ProbTrackX2:    CPU ProbTrackX2 with corrected stderr-based error detection.
    ProbTrackX2GPU: GPU-accelerated ProbTrackX2.

Functions:
    get_probtrackx2_interface: Return the best available ProbTrackX2 class.
"""

from typing import Optional, Type

from nipype.interfaces.base import traits
from nipype.interfaces.fsl import ProbTrackX2 as _UpstreamProbTrackX2

from thesis.core.gpu import present_gpu_binaries
from thesis.core.logging import get_logger

logger = get_logger(__name__)

__all__ = ["ProbTrackX2", "ProbTrackX2GPU", "get_probtrackx2_interface"]

# ---------------------------------------------------------------------------
# Binary name resolution (presence only — CUDA validation is the CLI's job)
# ---------------------------------------------------------------------------

# The candidate list and the presence-search live once in ``thesis.core.gpu``;
# this module reuses :func:`present_gpu_binaries` so the two never drift.


def _find_gpu_binary() -> Optional[str]:
    """
    Return the name of the first GPU probtrackx2 binary found on the system.

    Searches ``$FSLDIR/bin/`` first, then ``$PATH`` (via
    :func:`thesis.core.gpu.present_gpu_binaries`).  This is a **presence check
    only** — CUDA compatibility is validated separately by
    :func:`thesis.core.gpu.check_gpu` at CLI startup.

    Returns:
        Binary name string (first present candidate in preference order), or
        ``None`` if no GPU binary is installed.
    """
    present = present_gpu_binaries()
    return present[0] if present else None


# Resolved once at module-import time.
_GPU_BINARY: Optional[str] = _find_gpu_binary()


# ---------------------------------------------------------------------------
# CPU interface (stderr fix)
# ---------------------------------------------------------------------------




class ProbTrackX2(_UpstreamProbTrackX2):
    """
    ProbTrackX2 with a corrected ``_run_interface`` stderr check.

    The upstream nipype implementation raises ``RuntimeError`` whenever the
    command writes *anything* to stderr — including harmless OS warnings::

        bash: /path/to/libtinfo.so.6: no version information available

    These appear on some HPC clusters due to conda / system library mismatches
    and are entirely benign.  The upstream check therefore causes spurious
    workflow failures even when ``probtrackx2`` exits with return code 0.

    This subclass re-raises only for genuine command failures (non-zero return
    code).  All other behaviour — sample symlinking, ``targets.txt`` writing,
    output listing — is inherited unchanged from the upstream class.
    """

    def _run_interface(self, runtime):
        try:
            return super()._run_interface(runtime)
        except RuntimeError:
            # Upstream raises on any stderr content regardless of return code.
            # Re-raise only when the command actually failed.
            if getattr(runtime, "returncode", None) != 0:
                raise
        return runtime



# ---------------------------------------------------------------------------
# GPU interface
# ---------------------------------------------------------------------------


class ProbTrackX2GPUInputSpec(_UpstreamProbTrackX2.input_spec):
    """Input specification for ProbTrackX2GPU."""

    use_gpu = traits.Bool(
        True,
        usedefault=True,
        desc=(
            "Marks this node as requiring a GPU slot. "
            "GPU serialization is handled at the scheduler level via "
            "``n_gpu_procs=1`` in ``plugin_args``."
        ),
    )




class ProbTrackX2GPU(ProbTrackX2):
    """
    GPU-accelerated ProbTrackX2 (``probtrackx2_gpu11.0`` / ``probtrackx2_gpu``).

    Accepts the same inputs and produces the same outputs as the CPU
    :class:`ProbTrackX2`.  The binary is selected at module-import time by
    :func:`_find_gpu_binary` (presence check only).

    CUDA compatibility is validated at CLI startup via
    :func:`thesis.core.gpu.check_gpu` — if a compatible GPU is not available,
    ``config.hardware.gpu_enabled`` is set to ``False`` before any workflow
    is built, so this class will never be instantiated on an incompatible node.

    GPU serialization is handled at the scheduler level via ``n_gpu_procs=1``
    in ``plugin_args`` — no process-level lock is needed, so this class adds no
    ``_run_interface`` override and runs ``probtrackx2_gpu`` directly via the
    inherited :class:`ProbTrackX2` implementation.

    Use :func:`get_probtrackx2_interface` to select between CPU and GPU
    transparently based on the config flag.
    """

    _cmd: str = _GPU_BINARY or "probtrackx2_gpu"
    input_spec = ProbTrackX2GPUInputSpec



# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------




def get_probtrackx2_interface(use_gpu: bool = True) -> Type[ProbTrackX2]:
    """
    Return the appropriate ``ProbTrackX2`` interface class.

    When ``use_gpu=True`` and a GPU binary was found at module-import time,
    returns :class:`ProbTrackX2GPU`.  Otherwise returns CPU :class:`ProbTrackX2`.

    CUDA availability is **not** re-checked here — it was already validated
    (and ``use_gpu`` corrected if necessary) at CLI startup by
    :func:`thesis.core.gpu.check_gpu`.

    Args:
        use_gpu: Use the GPU variant when ``True`` and the binary is present.

    Returns:
        Uninstantiated class; pass directly to ``nipype.Node``.
    """
    if use_gpu and _GPU_BINARY is not None:
        logger.debug("Using GPU probtrackx2 binary: {}", _GPU_BINARY)
        return ProbTrackX2GPU

    return ProbTrackX2