foldit-plugin-sdk 0.1.5

Foldit plugin SDK - owns plugin.proto, the protocol types, and the Plugin trait; exposes a cbindgen C-ABI and pyo3 Python bindings
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

"""Checkpoint discovery utilities for foldit-runner plugins.

These are pure *detection* helpers: they locate a weight file / model
directory and raise a clean ``FileNotFoundError`` when it is absent. They
do NOT download anything; fetching is the plugin's own ``download_weights``
op (see :mod:`foldit_plugin_sdk.weights`). The error a plugin surfaces
when a load is attempted before the weights exist points the user at that
op rather than at any out-of-band CLI command.

Usage:
    from foldit_plugin_sdk.checkpoint_utils import (
        find_checkpoint, ensure_checkpoint_exists,
    )

    # For foundry models (glob pattern)
    checkpoint = find_checkpoint(
        cache_dir, pattern="rf3_*.ckpt", model_name="RoseTTAFold3"
    )

    # For direct path models
    ensure_checkpoint_exists(
        path="<plugin_dir>/assets/weights/esm2/esm2_t36_3B_UR50D.pt",
        model_name="ESM2",
    )
"""

import os
from pathlib import Path
from typing import List, Optional

from .logging_config import get_logger

logger = get_logger(__name__)

# Appended to every "not found" message: weights are fetched in-process by
# the plugin's own download op, not by any external command.
_DOWNLOAD_HINT = "Run this plugin's 'download_weights' op to fetch it."


def find_checkpoint(
    cache_dir: str,
    pattern: str,
    model_name: str,
    search_dirs: Optional[List[str]] = None,
) -> str:
    """Find a model checkpoint matching a glob pattern.

    Searches for checkpoint files matching the given pattern in one or more
    subdirectories of the cache directory.

    Args:
        cache_dir: Base cache directory (e.g., <plugin_dir>/assets/weights)
        pattern: Glob pattern to match (e.g., "rf3_*.ckpt")
        model_name: Human-readable model name for error messages
        search_dirs: List of subdirectories to search (default: ["rc_foundry"])

    Returns:
        Path to the first matching checkpoint file

    Raises:
        FileNotFoundError: If no matching checkpoint is found

    Example:
        checkpoint = find_checkpoint(
            cache_dir="<plugin_dir>/assets/weights",
            pattern="rf3_*.ckpt",
            model_name="RoseTTAFold3"
        )
    """
    if search_dirs is None:
        search_dirs = ["rc_foundry"]

    searched_paths = []

    for subdir in search_dirs:
        search_path = Path(cache_dir) / subdir
        searched_paths.append(str(search_path))

        if search_path.exists():
            matches = list(search_path.glob(pattern))
            if matches:
                checkpoint_path = str(matches[0])
                logger.info("Found %s checkpoint: %s", model_name, checkpoint_path)
                return checkpoint_path

    # No checkpoint found - raise helpful error
    raise FileNotFoundError(
        f"{model_name} checkpoint not found matching pattern '{pattern}'\n"
        f"Searched in: {', '.join(searched_paths)}\n"
        f"{_DOWNLOAD_HINT}"
    )


def ensure_checkpoint_exists(path: str, model_name: str) -> str:
    """Ensure a checkpoint file exists at the given path.

    Args:
        path: Expected path to the checkpoint file
        model_name: Human-readable model name for error messages

    Returns:
        The path (unchanged) if file exists

    Raises:
        FileNotFoundError: If the checkpoint file doesn't exist

    Example:
        path = ensure_checkpoint_exists(
            path="<plugin_dir>/assets/weights/esm2/esm2_t36_3B_UR50D.pt",
            model_name="ESM2",
        )
    """
    expanded_path = os.path.expanduser(path)

    if os.path.exists(expanded_path):
        logger.debug("Checkpoint exists: %s", expanded_path)
        return expanded_path

    raise FileNotFoundError(
        f"{model_name} checkpoint not found at {expanded_path}\n{_DOWNLOAD_HINT}"
    )


def ensure_directory_exists(path: str, model_name: str) -> str:
    """Ensure a model directory exists (for HuggingFace-style models).

    Args:
        path: Expected path to the model directory
        model_name: Human-readable model name for error messages

    Returns:
        The path (unchanged) if directory exists

    Raises:
        FileNotFoundError: If the directory doesn't exist

    Example:
        path = ensure_directory_exists(
            path="<plugin_dir>/assets/weights/my-hf-model",
            model_name="MyHFModel",
        )
    """
    expanded_path = os.path.expanduser(path)

    if os.path.isdir(expanded_path):
        logger.debug("Model directory exists: %s", expanded_path)
        return expanded_path

    raise FileNotFoundError(
        f"{model_name} model not found at {expanded_path}\n{_DOWNLOAD_HINT}"
    )