Function add_library_code

src/fused_mm_sampling/modal_lib/utils.py:136–175

Add the fused_mm_sampling source and benchmarking scripts to the image, then pip-install the package so subprocess-based tools (e.g. ncu) can import it. Layer order matters for caching: pyproject.toml (rarely changes) and dep install go first so that source-only changes don't re-run pip

(image: modal.Image)

def set_volume_caches():
    """Point cache env vars to the Modal volume and enable Triton autotune logging.

    XDG_CACHE_HOME: used by flashinfer, torch.compile, etc.
    TRITON_CACHE_DIR: used by Triton for compiled kernels and autotune results.
    Triton ignores XDG_CACHE_HOME and reads TRITON_CACHE_DIR (or TRITON_HOME) instead.
    TRITON_PRINT_AUTOTUNING: surfaces autotune progress so silent waits
    (cold cache, hangs) are debuggable from the run log.
    """
    os.environ["XDG_CACHE_HOME"] = f"{volume_path}/cache"
    os.environ["TRITON_CACHE_DIR"] = f"{volume_path}/cache/triton"
    os.environ["TRITON_PRINT_AUTOTUNING"] = "1"


def add_library_code(image: modal.Image) -> modal.Image:
    """Add the fused_mm_sampling source and benchmarking scripts to the image,
    then pip-install the package so subprocess-based tools (e.g. ncu) can import it.

    Layer order matters for caching: pyproject.toml (rarely changes) and dep
    install go first so that source-only changes don't re-run pip.
    """
    return (
        image
        # 1. Install deps (cached as long as pyproject.toml is unchanged).
        .add_local_file(
            str(_repo_root / "pyproject.toml"),
            remote_path="/opt/fmms/pyproject.toml",
            copy=True,
        )
        .run_commands(
            "mkdir -p /opt/fmms/src/fused_mm_sampling"
            " && touch /opt/fmms/src/fused_mm_sampling/__init__.py"
            " && cd /opt/fmms && pip install --break-system-packages -e ."
        )
        # 2. Copy source files (changes frequently, but deps layer is cached).
        #    The editable install points to /opt/fmms/src, so the real files
        #    are picked up at runtime.
        .add_local_dir(
            str(_repo_root / "src"),
            remote_path="/opt/fmms/src",
            copy=True,
            ignore=["__pycache__", "*.pyc"],
        )