patch_nixl.sh

#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Patch a NIXL source tree so it can build the BLUE_CACHE backend plugin.
#
# Usage:
#   ./scripts/patch_nixl.sh /path/to/nixl/source
#
# The script is idempotent: running it multiple times on the same tree is safe.

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
NIXL_SRC="${1:-}"

if [ -z "$NIXL_SRC" ]; then
    echo "Usage: $0 <NIXL_SOURCE_DIR>" >&2
    exit 1
fi

if [ ! -d "$NIXL_SRC/src/plugins" ] || [ ! -f "$NIXL_SRC/meson.build" ]; then
    echo "Error: $NIXL_SRC does not look like a NIXL source tree." >&2
    exit 1
fi

NIXL_SRC="$(cd "$NIXL_SRC" && pwd)"
PLUGIN_DIR="$NIXL_SRC/src/plugins/blue_cache"

echo "============================================================"
echo "Patching NIXL source tree: $NIXL_SRC"
echo "Project source: $PROJECT_DIR"
echo "============================================================"

# 1. Install plugin source files (resolve symlinks so the NIXL tree is self-contained).
mkdir -p "$PLUGIN_DIR"
cp -L "$PROJECT_DIR/nixl-plugin/src/"*.cpp "$PLUGIN_DIR/"
cp -L "$PROJECT_DIR/nixl-plugin/src/"*.h "$PLUGIN_DIR/"
cp -L "$PROJECT_DIR/nixl-plugin/src/meson.build" "$PLUGIN_DIR/"
cp -L "$PROJECT_DIR/common/include/dma_transfer.h" "$PLUGIN_DIR/"
echo "[OK] Copied plugin source files to $PLUGIN_DIR"

# 2. Patch root meson.build to include BLUE_CACHE in all_plugins.
python3 - "$NIXL_SRC/meson.build" <<'PY'
import sys
path = sys.argv[1]
with open(path, 'r') as f:
    text = f.read()

marker = "'BLUE_CACHE'"
if marker not in text:
    # Insert into the all_plugins list before the closing bracket.
    old = "all_plugins = ['UCX', 'LIBFABRIC', 'POSIX', 'OBJ', 'GDS', 'GDS_MT', 'MOONCAKE', 'HF3FS', 'GUSLI', 'GPUNETIO', 'UCCL', 'AZURE_BLOB'"
    new = "all_plugins = ['UCX', 'LIBFABRIC', 'POSIX', 'OBJ', 'GDS', 'GDS_MT', 'MOONCAKE', 'HF3FS', 'GUSLI', 'GPUNETIO', 'UCCL', 'AZURE_BLOB', 'BLUE_CACHE'"
    if old in text:
        text = text.replace(old, new, 1)
        with open(path, 'w') as f:
            f.write(text)
        print("[OK] Patched root meson.build: added BLUE_CACHE to all_plugins")
    else:
        print("[WARN] Could not locate all_plugins list in root meson.build; manual check needed")
else:
    print("[OK] Root meson.build already contains BLUE_CACHE")
PY

# 3. Patch src/plugins/meson.build to probe DOCA deps and conditionally build the plugin.
python3 - "$NIXL_SRC/src/plugins/meson.build" <<'PY'
import sys
path = sys.argv[1]
with open(path, 'r') as f:
    text = f.read()

# Ensure DOCA dependency probes exist.
probe_block = """
doca_common_dep = dependency('doca-common', required: false)
doca_comch_dep = dependency('doca-comch', required: false)
doca_dma_dep = dependency('doca-dma', required: false)
"""
if "doca_common_dep = dependency('doca-common'" not in text:
    # Insert before the conditional plugin block we add below, or at end if already present.
    text = text.rstrip() + "\n" + probe_block
    print("[OK] Added DOCA dependency probes to src/plugins/meson.build")
else:
    print("[OK] DOCA dependency probes already present in src/plugins/meson.build")

plugin_block = """
# BlueCache (BLUE_CACHE) plugin for GPU <-> DPU transfers via DOCA Comch + DMA
if enabled_plugins.get('BLUE_CACHE')
    if (not cuda_dep.found() or not doca_common_dep.found()
        or not doca_comch_dep.found() or not doca_dma_dep.found()) and is_explicit_enable
        if not cuda_dep.found()
            error('BLUE_CACHE plugin requested but CUDA dependency not found')
        else
            error('BLUE_CACHE plugin requested but DOCA dependency not found')
        endif
    elif cuda_dep.found() and doca_common_dep.found() and doca_comch_dep.found() and doca_dma_dep.found()
        subdir('blue_cache')
    endif
endif
"""
if "BlueCache (BLUE_CACHE) plugin for GPU <-> DPU transfers" not in text:
    text = text.rstrip() + "\n" + plugin_block
    with open(path, 'w') as f:
        f.write(text)
    print("[OK] Added BLUE_CACHE conditional build to src/plugins/meson.build")
else:
    print("[OK] BLUE_CACHE conditional build already present in src/plugins/meson.build")
PY

# 4. Patch src/core/nixl_plugin_manager.cpp for static plugin registration.
python3 - "$NIXL_SRC/src/core/nixl_plugin_manager.cpp" <<'PY'
import sys
path = sys.argv[1]
with open(path, 'r') as f:
    text = f.read()

marker = "STATIC_PLUGIN_BLUE_CACHE"
if marker in text:
    print("[OK] Static plugin registration already present in nixl_plugin_manager.cpp")
else:
    # Insert before the telemetry buffer registration line.
    old = "    NIXL_REGISTER_STATIC_PLUGIN(Telemetry, BUFFER)"
    new = """#ifdef STATIC_PLUGIN_BLUE_CACHE
    NIXL_REGISTER_STATIC_PLUGIN(Backend, BLUE_CACHE)
#endif

    NIXL_REGISTER_STATIC_PLUGIN(Telemetry, BUFFER)"""
    if old in text:
        text = text.replace(old, new, 1)
        with open(path, 'w') as f:
            f.write(text)
        print("[OK] Patched nixl_plugin_manager.cpp: added BLUE_CACHE static registration")
    else:
        print("[WARN] Could not locate static plugin registration block; manual check needed")
PY

echo ""
echo "============================================================"
echo "Patch complete."
echo "Next steps:"
echo "  cd $NIXL_SRC"
echo "  meson setup build -Denable_plugins=BLUE_CACHE"
echo "  ninja -C build"
echo "============================================================"