From fc043e78bd0e7cd9787d2c312ca48ee3a75ef4ea Mon Sep 17 00:00:00 2001 From: rohith Date: Fri, 7 Aug 2026 17:04:36 -0700 Subject: [PATCH 1/4] feat: add mmseqs server usage --- pyproject.toml | 2 + .../ml/preprocessing/msa/generating.py | 64 +- src/atomworks/ml/preprocessing/msa/server.py | 605 ++++++++++++++++++ src/atomworks_cli/generate.py | 142 +++- tests/ml/preprocessing/msa/__init__.py | 0 tests/ml/preprocessing/msa/test_server.py | 454 +++++++++++++ 6 files changed, 1240 insertions(+), 27 deletions(-) create mode 100644 src/atomworks/ml/preprocessing/msa/server.py create mode 100644 tests/ml/preprocessing/msa/__init__.py create mode 100644 tests/ml/preprocessing/msa/test_server.py diff --git a/pyproject.toml b/pyproject.toml index fa0057c5..53ae9672 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,8 @@ dependencies = [ "tqdm>=4.65.0", # Fast, extensible progress bar for loops and more # ... CLI & config management "typer>=0.12.5", # Modern CLI framework + # ... networking + "requests>=2.32", # HTTP client for the remote MSA server backend # ... linear algebra, maths & ml "numpy>=1.25.0", "scipy>=1.13.1", diff --git a/src/atomworks/ml/preprocessing/msa/generating.py b/src/atomworks/ml/preprocessing/msa/generating.py index de001f43..0317c551 100644 --- a/src/atomworks/ml/preprocessing/msa/generating.py +++ b/src/atomworks/ml/preprocessing/msa/generating.py @@ -38,6 +38,7 @@ import time from os import PathLike from pathlib import Path +from typing import Literal import pandas as pd @@ -47,6 +48,7 @@ from atomworks.ml.preprocessing.msa.filtering import HHFilterConfig, MSAFilterConfig, filter_msas from atomworks.ml.preprocessing.msa.finding import find_msas from atomworks.ml.preprocessing.msa.organizing import MSAOrganizationConfig, organize_msas +from atomworks.ml.preprocessing.msa.server import MSAServerConfig, make_msas_mmseqs_server from atomworks.ml.utils.misc import hash_sequence LOCAL_DB_PATH_GPU = _load_env_var("COLABFOLD_LOCAL_DB_PATH_GPU") @@ -150,24 +152,30 @@ class MSAGenerationConfig: working ColabFold script parameters. Args: + backend: Where the search runs. "local" uses the local MMseqs2 installation and ColabFold + databases; "server" submits the sequences to a remote ColabFold-style MMseqs2 server. sharding_pattern: Directory sharding pattern for file organization. output_extension: File extension and compression for output files. - use_env: Whether to include environmental (metagenomic) database. - gpu: Whether to use GPU acceleration. - gpu_server: Whether to use GPU server (requires gpu=True). - num_iterations: Number of MMseqs2 search iterations. - max_seqs: Maximum number of cluster centers (NOT total sequences) in the MSA. - threads: Number of CPU threads to use. - use_local_temp_dir: Whether to use local temporary directory for intermediate files. - max_final_sequences: Maximum number of sequences in the final MSA after HHFilter. + use_env: Whether to include environmental (metagenomic) database. Applies to both backends; + with backend="server" it overrides `server.use_env`. + gpu: Whether to use GPU acceleration. Local backend only. + gpu_server: Whether to use GPU server (requires gpu=True). Local backend only. + num_iterations: Number of MMseqs2 search iterations. Local backend only. + max_seqs: Maximum number of cluster centers (NOT total sequences) in the MSA. Local backend only. + threads: Number of CPU threads to use. Local backend only. + use_local_temp_dir: Whether to use local temporary directory for intermediate files. Local backend only. + max_final_sequences: Maximum number of sequences in the final MSA after HHFilter, or None to skip + filtering. Note that HHfilter is a local binary, which a "server" backend user may not have. check_existing: Whether to check for existing MSAs before generation. existing_msa_dirs: Directories to check for existing MSAs. If None, uses LOCAL_MSA_DIRS env var. - search_config: Advanced MMseqs2 search configuration. + search_config: Advanced MMseqs2 search configuration. Local backend only. + server: Remote MSA server configuration. Server backend only. References: * Mirdita, M. et al. (2022). ColabFold: making protein folding accessible to all. *Nature Methods*, 19, 679-682. """ + backend: Literal["local", "server"] = "local" sharding_pattern: str = "/0:2/" output_extension: str = MSAFileExtension.A3M_GZ.value use_env: bool = True @@ -177,17 +185,25 @@ class MSAGenerationConfig: max_seqs: int = 10000 threads: int = 32 use_local_temp_dir: bool = True - max_final_sequences: int = 10000 + max_final_sequences: int | None = 10000 check_existing: bool = False existing_msa_dirs: list[PathLike] | None = None search_config: MMseqs2SearchConfig = dataclasses.field(default_factory=lambda: MMseqs2SearchConfig()) + server: MSAServerConfig = dataclasses.field(default_factory=lambda: MSAServerConfig()) def __post_init__(self): + if self.backend not in ("local", "server"): + raise ValueError(f"Unknown MSA backend: {self.backend!r}. Must be one of 'local', 'server'.") + # If we're using GPU, also use the GPU server by default if self.gpu and not self.gpu_server: logger.info("GPU is enabled, setting gpu_server to True") self.gpu_server = True + if self.backend == "server": + # Keep the two backends' notion of which databases to search in sync + self.server.use_env = self.use_env + def _get_database_path(gpu: bool = False) -> Path: """ @@ -712,7 +728,7 @@ def make_msas_mmseqs( num_iterations: int = 3, max_seqs: int = 10_000, use_local_temp_dir: bool = True, - max_final_sequences: int = 10_000, + max_final_sequences: int | None = 10_000, sharding_pattern: str = "/0:2/", output_extension: str = MSAFileExtension.A3M_GZ.value, search_config: MMseqs2SearchConfig | None = None, @@ -727,7 +743,7 @@ def make_msas_mmseqs( num_iterations: Number of search iterations. max_seqs: Maximum number of cluster centers. use_local_temp_dir: Whether to use local temporary directory for intermediate files. - max_final_sequences: Maximum number of sequences in final MSAs after filtering. + max_final_sequences: Maximum number of sequences in final MSAs after filtering, or None to skip filtering. sharding_pattern: Directory sharding pattern (e.g., "/0:2/"). output_extension: Output file extension (.a3m, .a3m.gz, .a3m.zst, .afa, .afa.gz, .afa.zst). search_config: Advanced MMseqs2 search configuration. @@ -845,6 +861,9 @@ def make_msas_from_csv( ) -> None: """Generate MSAs from sequences in a CSV file. + Dispatches to the local MMseqs2 pipeline or to a remote MMseqs2 server, depending on + ``config.backend``. + Args: csv_file: Path to CSV file containing protein sequences. output_dir: Directory where organized MSA files will be saved. @@ -863,6 +882,12 @@ def make_msas_from_csv( .. code-block:: python make_msas_from_csv("data.csv", "output_msas/", sequence_column="sequence") + + Generate MSAs with a remote MMseqs2 server instead of local databases: + + .. code-block:: python + + make_msas_from_csv("sequences.csv", "output_msas/", config=MSAGenerationConfig(backend="server")) """ df = pd.read_csv(csv_file) @@ -882,6 +907,21 @@ def make_msas_from_csv( if config is None: config = MSAGenerationConfig() + if config.backend == "server": + # The server backend does its own existence check: unlike the local backend it also looks in + # `output_dir`, so it works without LOCAL_MSA_DIRS being set + make_msas_mmseqs_server( + sequences=sequences, + output_dir=output_dir, + config=config.server, + sharding_pattern=config.sharding_pattern, + output_extension=config.output_extension, + max_final_sequences=config.max_final_sequences, + check_existing=config.check_existing, + existing_msa_dirs=config.existing_msa_dirs, + ) + return + # Filter existing sequences if requested if config.check_existing: logger.info(f"Finding existing MSAs among {len(sequences)} sequences...") diff --git a/src/atomworks/ml/preprocessing/msa/server.py b/src/atomworks/ml/preprocessing/msa/server.py new file mode 100644 index 00000000..d73c4c4f --- /dev/null +++ b/src/atomworks/ml/preprocessing/msa/server.py @@ -0,0 +1,605 @@ +"""Generate ColabFold-style multiple sequence alignments (MSAs) with a remote MMseqs2 server. + +This is the remote counterpart to :py:mod:`atomworks.ml.preprocessing.msa.generating`: instead of +running MMseqs2 locally against the ~1 TB ColabFold database set, sequences are submitted to a +ColabFold-compatible MMseqs2 server (the public ``https://api.colabfold.com`` by default, or any +self-hosted deployment) and the resulting a3m files are downloaded. No local databases, no mmseqs +binary and no GPU are required. + +The output is byte-compatible with the local backend: one ``.a3m.gz`` per input sequence in a +hash-sharded directory, where ```` is :py:func:`~atomworks.ml.utils.misc.hash_sequence` of the +query. Everything downstream (finding, filtering, loading) is therefore unchanged. + +Examples: + Generate MSAs against the public ColabFold server: + + .. code-block:: python + + from atomworks.ml.preprocessing.msa.server import make_msas_mmseqs_server + + sequences = ["MSYIWRQLGSPTVAITLSVSTVIYVTVICPIVFIHLFGDHL...", "MKKKEVEKDDLIENASRVASCISIFLIIASTTMYIFIGLKI..."] + make_msas_mmseqs_server(sequences, "output_msas/") + + Generate MSAs against a self-hosted server that requires basic authentication: + + .. code-block:: python + + from atomworks.ml.preprocessing.msa.server import MSAServerConfig, make_msas_mmseqs_server + + config = MSAServerConfig(host_url="https://msa.internal", username="me", password="secret") + make_msas_mmseqs_server(sequences, "output_msas/", config=config) + +Note: + MSAs produced by a ColabFold server carry UniRef accessions and alignment statistics in their + headers, but no ``TaxID=`` field. Multimer MSA pairing in AtomWorks keys on ``TaxID=``, so + multimers built from server MSAs are effectively unpaired. + +References: + * Mirdita, M. et al. (2022). ColabFold: making protein folding accessible to all. *Nature Methods*, 19, 679-682. + * `ColabFold MMseqs2 API client`_ - Reference implementation of the wire protocol. + + .. _ColabFold MMseqs2 API client: https://github.com/sokrypton/ColabFold/blob/main/colabfold/colabfold.py +""" + +import dataclasses +import functools +import logging +import math +import os +import random +import shutil +import tarfile +import tempfile +import time +from collections.abc import Callable +from os import PathLike +from pathlib import Path +from typing import Any, TypeVar + +import requests +from tqdm import tqdm + +from atomworks.enums import MSAFileExtension +from atomworks.ml.preprocessing.msa.filtering import HHFilterConfig, MSAFilterConfig, filter_msas +from atomworks.ml.preprocessing.msa.finding import find_msas, get_msa_dirs_from_env +from atomworks.ml.preprocessing.msa.organizing import MSAOrganizationConfig, organize_msas +from atomworks.ml.utils.misc import hash_sequence + +logger = logging.getLogger(__name__) + +DEFAULT_MSA_SERVER_URL = "https://api.colabfold.com" +"""Public ColabFold MMseqs2 API endpoint.""" + +UNIREF_A3M_FILENAME = "uniref.a3m" +"""Name of the UniRef alignment inside the server's result tarball.""" + +ENV_A3M_FILENAME = "bfd.mgnify30.metaeuk30.smag30.a3m" +"""Name of the metagenomic (environmental) alignment inside the server's result tarball.""" + +_FIRST_QUERY_ID = 101 +"""ColabFold servers expect numeric FASTA headers; queries are numbered from this value.""" + +_PENDING_STATUSES = ("UNKNOWN", "PENDING", "RUNNING") +_RESUBMIT_STATUSES = ("UNKNOWN", "RATELIMIT") + +_T = TypeVar("_T") + + +@dataclasses.dataclass +class MSAServerConfig: + """Configuration for a remote ColabFold-style MMseqs2 server. + + Args: + host_url: Base URL of the MMseqs2 server. + use_env: Whether to include the metagenomic (environmental) database. + use_filter: Whether to let the server filter the alignment. + username: Username for HTTP basic auth. Falls back to ``MSA_SERVER_USERNAME``. + password: Password for HTTP basic auth. Falls back to ``MSA_SERVER_PASSWORD``. + api_key_header: Name of the header carrying an API key (mutually exclusive with basic auth). + api_key_value: Value of the API key header. Falls back to ``MSA_SERVER_API_KEY``. + user_agent: ``User-Agent`` sent with every request. The public server asks clients to identify themselves. + request_timeout: Per-request timeout, in seconds. + poll_interval: Lower and upper bound (in seconds) of the jittered delay between status polls. + max_retries: Maximum number of consecutive failed network calls (or job resubmissions) before giving up. + retry_delay: Delay between retries of a failed network call, in seconds. + batch_size: Maximum number of sequences submitted in a single ticket. + + Raises: + ValueError: If both basic auth and an API key are configured, or if a numeric field is out of range. + """ + + host_url: str = DEFAULT_MSA_SERVER_URL + use_env: bool = True + use_filter: bool = True + username: str | None = None + password: str | None = None + api_key_header: str | None = None + api_key_value: str | None = None + user_agent: str = "atomworks" + request_timeout: float = 6.02 + poll_interval: tuple[float, float] = (5.0, 10.0) + max_retries: int = 5 + retry_delay: float = 5.0 + batch_size: int = 50 + + def __post_init__(self) -> None: + """Resolve credentials from the environment and validate the configuration.""" + # NOTE: we read the environment directly (rather than via `atomworks.constants._load_env_var`) + # because unset credentials are the common case and should not emit a warning. + if self.username is None: + self.username = os.environ.get("MSA_SERVER_USERNAME") + if self.password is None: + self.password = os.environ.get("MSA_SERVER_PASSWORD") + + has_basic_auth = self.username is not None or self.password is not None + if self.api_key_value is None and not has_basic_auth: + self.api_key_value = os.environ.get("MSA_SERVER_API_KEY") + + if has_basic_auth and self.api_key_value is not None: + raise ValueError( + "Cannot use HTTP basic auth (username/password) and an API key header at the same time. " + "Provide one or the other." + ) + if self.api_key_value is not None and not self.api_key_header: + raise ValueError("An API key value was given without `api_key_header`; the header name is required.") + + self.host_url = self.host_url.rstrip("/") + + low, high = self.poll_interval + if low < 0 or high < low: + raise ValueError(f"`poll_interval` must be a non-negative (low, high) pair, got {self.poll_interval}") + if self.batch_size < 1: + raise ValueError(f"`batch_size` must be at least 1, got {self.batch_size}") + if self.max_retries < 1: + raise ValueError(f"`max_retries` must be at least 1, got {self.max_retries}") + + +def _select_mode(config: MSAServerConfig) -> str: + """Map the configured databases and filtering onto a ColabFold server search mode.""" + if config.use_filter: + return "env" if config.use_env else "all" + return "env-nofilter" if config.use_env else "nofilter" + + +def _build_query_fasta(sequences: list[str]) -> str: + """Build the FASTA payload for a batch, using the numeric headers the server expects.""" + return "".join(f">{_FIRST_QUERY_ID + i}\n{sequence}\n" for i, sequence in enumerate(sequences)) + + +def _request_kwargs(config: MSAServerConfig) -> dict[str, Any]: + """Build the shared `requests` keyword arguments (headers, auth, timeout) for a call.""" + headers = {"User-Agent": config.user_agent} + if config.api_key_header and config.api_key_value: + headers[config.api_key_header] = config.api_key_value + + kwargs: dict[str, Any] = {"headers": headers, "timeout": config.request_timeout} + if config.username is not None or config.password is not None: + kwargs["auth"] = (config.username or "", config.password or "") + return kwargs + + +def _parse_json_response(response: requests.Response) -> dict[str, Any]: + """Parse a ticket/status response, degrading to an ``ERROR`` status if the server didn't reply with JSON.""" + try: + payload = response.json() + except ValueError: + logger.error(f"MSA server did not reply with JSON (HTTP {response.status_code}): {response.text[:500]}") + return {"status": "ERROR", "message": f"HTTP {response.status_code}: {response.text[:500]}"} + + if not isinstance(payload, dict): + return {"status": "ERROR", "message": f"Unexpected JSON payload: {payload!r}"} + return payload + + +def _submit(sequences: list[str], mode: str, config: MSAServerConfig) -> dict[str, Any]: + """Submit a batch of sequences to the server. + + Args: + sequences: Sequences to align (a single batch). + mode: Server search mode, see :py:func:`_select_mode`. + config: Server configuration. + + Returns: + The decoded ticket payload, containing at least a ``status`` and (on success) an ``id``. + """ + response = requests.post( + f"{config.host_url}/ticket/msa", + data={"q": _build_query_fasta(sequences), "mode": mode}, + **_request_kwargs(config), + ) + return _parse_json_response(response) + + +def _status(ticket_id: str, config: MSAServerConfig) -> dict[str, Any]: + """Query the status of a submitted ticket.""" + response = requests.get(f"{config.host_url}/ticket/{ticket_id}", **_request_kwargs(config)) + return _parse_json_response(response) + + +def _download(ticket_id: str, dest: PathLike, config: MSAServerConfig) -> None: + """Download the result tarball of a completed ticket to `dest`.""" + with requests.get( + f"{config.host_url}/result/download/{ticket_id}", stream=True, **_request_kwargs(config) + ) as response: + response.raise_for_status() + with open(dest, "wb") as f: + for chunk in response.iter_content(chunk_size=1024 * 1024): + f.write(chunk) + + +def _with_retries(call: Callable[[], _T], config: MSAServerConfig, description: str) -> _T: + """Call `call`, retrying transient network failures up to `config.max_retries` times. + + Args: + call: Zero-argument callable performing a single network request. + config: Server configuration (supplies the retry count and delay). + description: Human-readable description of the call, used in log and error messages. + + Returns: + The return value of `call`. + + Raises: + RuntimeError: If every attempt failed. + """ + last_error: Exception | None = None + for attempt in range(1, config.max_retries + 1): + try: + return call() + except Exception as e: # any network failure is worth retrying + last_error = e + logger.warning(f"Error while {description} (attempt {attempt}/{config.max_retries}): {e}") + if attempt < config.max_retries: + time.sleep(config.retry_delay) + + raise RuntimeError(f"Failed while {description} after {config.max_retries} attempts") from last_error + + +def _sleep_between_polls(config: MSAServerConfig) -> float: + """Sleep for a jittered interval between status polls, returning the number of seconds slept.""" + delay = random.uniform(*config.poll_interval) + time.sleep(delay) + return delay + + +def _raise_on_error_status(result: dict[str, Any]) -> None: + """Raise if the server reported a terminal status. + + Raises: + RuntimeError: If the ticket is in the ``ERROR`` or ``MAINTENANCE`` state. + """ + status = result.get("status") + if status == "ERROR": + message = result.get("message") or "no message given" + raise RuntimeError( + f"MSA server returned an error: {message}. " + "This usually means the query was malformed or too long for the server." + ) + if status == "MAINTENANCE": + raise RuntimeError("MSA server is undergoing maintenance; please retry later.") + + +def _run_batch(sequences: list[str], tar_path: Path, config: MSAServerConfig) -> None: + """Submit one batch of sequences, poll until the job completes, and download the result tarball. + + Args: + sequences: Sequences to align (a single batch, already deduplicated). + tar_path: Destination for the downloaded ``out.tar.gz``. + config: Server configuration. + + Raises: + RuntimeError: If the server reports an error, is under maintenance, or the job never completes. + """ + mode = _select_mode(config) + submit = functools.partial(_submit, sequences, mode, config) + + with tqdm( + desc=f"MMseqs2 server ({len(sequences)} sequences)", unit="s", bar_format="{l_bar}{bar}| {n:.0f}s" + ) as bar: + for resubmission in range(1, config.max_retries + 1): + result = _with_retries(submit, config, "submitting sequences to the MSA server") + while result.get("status") in _RESUBMIT_STATUSES: + _sleep_between_polls(config) + result = _with_retries(submit, config, "submitting sequences to the MSA server") + _raise_on_error_status(result) + + ticket_id = result.get("id") + if not ticket_id: + raise RuntimeError(f"MSA server accepted the job but returned no ticket id: {result!r}") + poll = functools.partial(_status, ticket_id, config) + + while result.get("status") in _PENDING_STATUSES: + slept = _sleep_between_polls(config) + bar.update(slept) + result = _with_retries(poll, config, f"polling ticket {ticket_id}") + + _raise_on_error_status(result) + if result.get("status") == "COMPLETE": + _with_retries( + functools.partial(_download, ticket_id, tar_path, config), + config, + f"downloading results for ticket {ticket_id}", + ) + return + + logger.warning( + f"Unexpected status {result.get('status')!r} for ticket {ticket_id}; " + f"resubmitting ({resubmission}/{config.max_retries})" + ) + + raise RuntimeError(f"MSA server job did not complete after {config.max_retries} submissions") + + +def _extract_tarball(tar_path: Path, dest_dir: Path) -> None: + """Extract the regular files of a result tarball into `dest_dir`, flattening any directory structure.""" + with tarfile.open(tar_path, "r:gz") as tar: + for member in tar.getmembers(): + if not member.isfile(): + continue + source = tar.extractfile(member) + if source is None: + continue + with source, open(dest_dir / Path(member.name).name, "wb") as f: + shutil.copyfileobj(source, f) + + +def _split_multi_query_a3m(a3m_path: Path) -> dict[int, list[str]]: + """Split a multi-query a3m returned by the server into per-query blocks of lines. + + The server concatenates the alignment of every query in a batch into a single file, separated by + null bytes and re-introduced by the numeric query header (``>101``, ``>102``, ...). + + Args: + a3m_path: Path to the multi-query a3m file. + + Returns: + Mapping of numeric query id to the lines of that query's alignment (headers included). + + Raises: + ValueError: If alignment lines appear before any query header. + """ + blocks: dict[int, list[str]] = {} + query_id: int | None = None + expect_header = True + + with open(a3m_path) as f: + for raw_line in f: + line = raw_line + if "\x00" in line: + # A null byte marks the boundary between two queries' alignments + line = line.replace("\x00", "") + expect_header = True + if not line.strip(): + continue + if expect_header and line.startswith(">"): + query_id = int(line[1:].rstrip()) + expect_header = False + blocks.setdefault(query_id, []) + if query_id is None: + raise ValueError( + f"Malformed a3m from the MSA server: alignment lines before a query header in {a3m_path}" + ) + blocks[query_id].append(line if line.endswith("\n") else f"{line}\n") + + return blocks + + +def _drop_leading_query_record(block: list[str]) -> list[str]: + """Drop the repeated query header and sequence from the head of a per-query alignment block.""" + return block[2:] if len(block) >= 2 and block[0].startswith(">") else block + + +def _write_a3m_files( + blocks_per_file: list[dict[int, list[str]]], sequences: list[str], out_dir: Path +) -> dict[str, Path]: + """Concatenate per-database alignments and write one flat ``.a3m`` per query sequence. + + Args: + blocks_per_file: Per-query line blocks, one entry per downloaded a3m file (UniRef first). + sequences: The batch of query sequences, in submission order. + out_dir: Directory the flat a3m files are written to. + + Returns: + Mapping of query sequence to the path of its a3m file. + + Raises: + RuntimeError: If the server returned no alignment at all for one of the queries. + """ + sequence_to_path: dict[str, Path] = {} + + for i, sequence in enumerate(sequences): + query_id = _FIRST_QUERY_ID + i + lines: list[str] = [] + for file_index, blocks in enumerate(blocks_per_file): + block = blocks.get(query_id) + if block is None: + logger.warning(f"MSA server returned no alignment for query {query_id} in result file {file_index}") + continue + # The query itself is repeated at the top of every per-database alignment; keep it only once + lines.extend(block if not lines else _drop_leading_query_record(block)) + + if not lines: + raise RuntimeError(f"MSA server returned no alignment for query {query_id}") + + # Rewrite the numeric query header to the sequence hash, matching the local backend's output + sequence_hash = hash_sequence(sequence) + lines[0] = f">{sequence_hash}\n" + + path = out_dir / f"{sequence_hash}{MSAFileExtension.A3M.value}" + path.write_text("".join(lines)) + sequence_to_path[sequence] = path + + return sequence_to_path + + +def run_mmseqs2_server( + sequences: str | list[str], + output_dir: PathLike, + config: MSAServerConfig | None = None, +) -> dict[str, Path]: + """Align sequences with a remote MMseqs2 server and write one flat ``.a3m`` per sequence. + + This is the low-level entrypoint; it does no sharding, compression or filtering. Most callers + want :py:func:`make_msas_mmseqs_server` instead. + + Args: + sequences: A single protein sequence string or list of protein sequences. + output_dir: Directory the flat a3m files are written to. Created if it doesn't exist. + config: Server configuration. If None, uses defaults (the public ColabFold server). + + Returns: + Mapping of each unique input sequence to the path of its a3m file. + + Raises: + RuntimeError: If the server reports an error or the results are incomplete. + """ + if isinstance(sequences, str): + sequences = [sequences] + if config is None: + config = MSAServerConfig() + + # Duplicate sequences would waste a server slot each; they map to the same output file anyway + unique_sequences = list(dict.fromkeys(sequences)) + + out_path = Path(output_dir) + out_path.mkdir(parents=True, exist_ok=True) + + sequence_to_path: dict[str, Path] = {} + n_batches = math.ceil(len(unique_sequences) / config.batch_size) + logger.info( + f"Requesting MSAs for {len(unique_sequences)} unique sequences from {config.host_url} " + f"in {n_batches} batch(es) (mode: {_select_mode(config)})" + ) + + for batch_index, start in enumerate(range(0, len(unique_sequences), config.batch_size), start=1): + batch = unique_sequences[start : start + config.batch_size] + logger.info(f"Submitting batch {batch_index}/{n_batches} ({len(batch)} sequences)") + + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_path = Path(tmp_dir) + tar_path = tmp_path / "out.tar.gz" + + start_time = time.time() + _run_batch(batch, tar_path, config) + logger.info(f"Batch {batch_index}/{n_batches} completed in {time.time() - start_time:.1f} seconds") + + _extract_tarball(tar_path, tmp_path) + + uniref_a3m = tmp_path / UNIREF_A3M_FILENAME + if not uniref_a3m.exists(): + raise RuntimeError(f"MSA server result did not contain the expected {UNIREF_A3M_FILENAME}") + + a3m_files = [uniref_a3m] + if config.use_env: + env_a3m = tmp_path / ENV_A3M_FILENAME + if env_a3m.exists(): + a3m_files.append(env_a3m) + else: + logger.warning(f"MSA server result did not contain {ENV_A3M_FILENAME}; using UniRef hits only") + + blocks_per_file = [_split_multi_query_a3m(a3m_file) for a3m_file in a3m_files] + sequence_to_path.update(_write_a3m_files(blocks_per_file, batch, out_path)) + + return sequence_to_path + + +def _msa_dirs_to_check(output_dir: Path, existing_msa_dirs: list[PathLike] | None) -> list[PathLike]: + """Resolve the directories to search for already-generated MSAs. + + The output directory is always checked (it is the natural cache); explicitly requested + directories are checked in addition, falling back to ``LOCAL_MSA_DIRS`` when none are given. + """ + if existing_msa_dirs is None: + existing_msa_dirs = get_msa_dirs_from_env(raise_if_not_set=False) or [] + return [output_dir, *existing_msa_dirs] + + +def make_msas_mmseqs_server( + sequences: str | list[str], + output_dir: PathLike, + *, + config: MSAServerConfig | None = None, + sharding_pattern: str = "/0:2/", + output_extension: str = MSAFileExtension.A3M_GZ.value, + max_final_sequences: int | None = None, + check_existing: bool = True, + existing_msa_dirs: list[PathLike] | None = None, +) -> None: + """Generate MSAs from protein sequences using a remote MMseqs2 server. + + Signature-compatible with :py:func:`~atomworks.ml.preprocessing.msa.generating.make_msas_mmseqs` + (the local backend), so the two are interchangeable. Output is written to the same hash-sharded, + compressed layout. + + Args: + sequences: A single protein sequence string or list of protein sequences. + output_dir: Path to the output directory where MSA files will be saved. + config: Server configuration. If None, uses defaults (the public ColabFold server). + sharding_pattern: Directory sharding pattern (e.g., "/0:2/"). + output_extension: Output file extension (.a3m, .a3m.gz, .a3m.zst, .afa, .afa.gz, .afa.zst). + max_final_sequences: If set, MSAs are filtered down to this many sequences with HHfilter. + Defaults to None: the server already filters, and HHfilter is a local binary that a + remote-backend user may not have installed. + check_existing: Whether to skip sequences that already have an MSA. + existing_msa_dirs: Additional directories to check for existing MSAs. The output directory is + always checked. If None, falls back to the LOCAL_MSA_DIRS env var (when set). + + Examples: + .. code-block:: python + + make_msas_mmseqs_server( + ["MSYIWRQLGSPTVAITLSVSTVIYVTVICPIVFIHLFGDHL...", "MKKKEVEKDDLIENASRVASCISIFLIIASTTMYIFIGLKI..."], + "output_msas/", + ) + """ + if isinstance(sequences, str): + sequences = [sequences] + if config is None: + config = MSAServerConfig() + + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + sequences = list(dict.fromkeys(sequences)) + + if check_existing: + logger.info(f"Finding existing MSAs among {len(sequences)} sequences...") + sequences, _ = find_msas( + sequences, + msa_dirs=_msa_dirs_to_check(output_path, existing_msa_dirs), + shard_depths=[0, 1, 2, 3, 4], + extensions=[MSAFileExtension.A3M, MSAFileExtension.A3M_GZ, MSAFileExtension.A3M_ZST], + ) + if not sequences: + logger.info("All sequences already have MSAs, skipping generation") + return + logger.info(f"Found {len(sequences)} sequences needing MSA generation") + + with tempfile.TemporaryDirectory() as tmp_dir: + run_mmseqs2_server(sequences, tmp_dir, config) + + # Organize MSAs (hash-based sharding and compression) using existing organization functionality + logger.info("Organizing MSA files...") + organize_msas( + tmp_dir, + output_path, + MSAOrganizationConfig( + input_extension=MSAFileExtension.A3M, + output_extension=output_extension, + sharding_pattern=sharding_pattern, + copy_files=False, # Move files instead of copying + ), + ) + + if max_final_sequences is not None: + logger.info(f"Filtering MSA files to max {max_final_sequences} sequences (requires a local hhfilter)...") + filter_msas( + output_path, + output_path, + MSAFilterConfig( + input_extension=output_extension, + output_extension=output_extension, + hhfilter=HHFilterConfig(max_sequences=max_final_sequences), + ), + ) + + logger.info(f"MSA files saved to: {output_path.absolute()}") diff --git a/src/atomworks_cli/generate.py b/src/atomworks_cli/generate.py index 43aa4f8d..b41084c4 100644 --- a/src/atomworks_cli/generate.py +++ b/src/atomworks_cli/generate.py @@ -1,8 +1,9 @@ -"""MSA generation command using MMseqs2.""" +"""MSA generation command using MMseqs2, either locally or against a remote MSA server.""" from __future__ import annotations import logging +from enum import Enum from pathlib import Path import torch @@ -14,15 +15,53 @@ MSAGenerationConfig, make_msas_from_csv, ) +from atomworks.ml.preprocessing.msa.server import DEFAULT_MSA_SERVER_URL, MSAServerConfig from .common import enable_logging app = typer.Typer() logger = logging.getLogger(__name__) +# Options that only make sense for the local MMseqs2 pipeline +_LOCAL_ONLY_OPTIONS = { + "gpu": "--gpu/--no-gpu", + "num_workers": "--num-workers/-j", + "sensitivity": "--sensitivity", + "num_iterations": "--num-iterations/-n", +} + + +class MSABackend(str, Enum): + """Where the MMseqs2 search runs.""" + + LOCAL = "local" + SERVER = "server" + + +def _was_passed_on_the_command_line(ctx: typer.Context, parameter_name: str) -> bool: + """Whether the user explicitly passed a parameter (as opposed to it taking its default value). + + Note: + We compare the parameter source by name rather than importing `click.core.ParameterSource`, + since recent Typer releases vendor Click rather than depending on it. + """ + source = ctx.get_parameter_source(parameter_name) + return source is not None and source.name == "COMMANDLINE" + + +def _reject_local_only_options(ctx: typer.Context) -> None: + """Error out if the user passed local-backend options together with `--backend server`.""" + passed = [flag for name, flag in _LOCAL_ONLY_OPTIONS.items() if _was_passed_on_the_command_line(ctx, name)] + if passed: + raise typer.BadParameter( + f"{', '.join(passed)} {'is' if len(passed) == 1 else 'are'} only supported with '--backend local'; " + "the remote MSA server controls its own search parameters." + ) + @app.command() def generate( + ctx: typer.Context, csv_file: Path = typer.Argument( ..., exists=True, @@ -47,6 +86,11 @@ def generate( "-c", help="Name of column containing sequences (required if CSV has multiple columns)", ), + backend: MSABackend = typer.Option( + MSABackend.LOCAL.value, + "--backend", + help="Run MMseqs2 against local ColabFold databases, or submit to a remote MSA server", + ), # MSAGenerationConfig parameters sharding_pattern: str = typer.Option( "/0:2/", @@ -63,18 +107,21 @@ def generate( gpu: bool | None = typer.Option( None, "--gpu/--no-gpu", - help="Use GPU acceleration (auto-detects if not specified)", + help="Use GPU acceleration (auto-detects if not specified). Local backend only", ), num_iterations: int = typer.Option( 3, "--num-iterations", "-n", - help="Number of MMseqs2 search iterations", + help="Number of MMseqs2 search iterations. Local backend only", ), - max_final_sequences: int = typer.Option( - 10_000, + max_final_sequences: int | None = typer.Option( + None, "--max-final-sequences", - help="Maximum number of sequences in final MSAs", + help=( + "Maximum number of sequences in final MSAs " + "(default: 10000 for the local backend, no HHfilter pass for the server backend)" + ), ), use_env: bool = typer.Option( True, @@ -85,12 +132,38 @@ def generate( 32, "--num-workers", "-j", - help="Number of CPU threads", + help="Number of CPU threads. Local backend only", ), sensitivity: float | None = typer.Option( 8.0, "--sensitivity", - help="MMseqs2 sensitivity (lower = faster, sparser MSAs)", + help="MMseqs2 sensitivity (lower = faster, sparser MSAs). Local backend only", + ), + # MSAServerConfig parameters + server_url: str = typer.Option( + DEFAULT_MSA_SERVER_URL, + "--server-url", + help="Base URL of the MMseqs2 server. Server backend only", + ), + server_username: str | None = typer.Option( + None, + "--server-username", + help="Username for HTTP basic auth (or set MSA_SERVER_USERNAME). Server backend only", + ), + server_password: str | None = typer.Option( + None, + "--server-password", + help="Password for HTTP basic auth (or set MSA_SERVER_PASSWORD). Server backend only", + ), + api_key_header: str | None = typer.Option( + None, + "--api-key-header", + help="Header name to carry an API key, e.g. 'X-API-Key'. Server backend only", + ), + api_key_value: str | None = typer.Option( + None, + "--api-key-value", + help="API key value (or set MSA_SERVER_API_KEY). Server backend only", ), verbose: bool = typer.Option( False, @@ -119,13 +192,24 @@ def generate( atomworks msa generate data.csv output_msas/ --sequence-column seq # With custom parameters - atomworks msa generate sequences.csv output_msas/ --gpu --max-final-sequences 5000 --threads 16 + atomworks msa generate sequences.csv output_msas/ --gpu --max-final-sequences 5000 --num-workers 16 + + # Without local databases, using the public ColabFold MSA server + atomworks msa generate sequences.csv output_msas/ --backend server """ enable_logging(verbose) + is_server = backend is MSABackend.SERVER + if is_server: + _reject_local_only_options(ctx) + # Auto-detect GPU if not specified if gpu is None: - gpu = torch.cuda.is_available() + gpu = False if is_server else torch.cuda.is_available() + + # HHfilter is a local binary; don't require it by default when the search itself ran remotely + if max_final_sequences is None and not is_server: + max_final_sequences = 10_000 # Parse MSA directories if provided msa_dirs = None @@ -137,8 +221,21 @@ def generate( s=sensitivity, ) + try: + server_config = MSAServerConfig( + host_url=server_url, + use_env=use_env, + username=server_username, + password=server_password, + api_key_header=api_key_header, + api_key_value=api_key_value, + ) + except ValueError as e: + raise typer.BadParameter(str(e)) from e + # Create generation config config = MSAGenerationConfig( + backend=backend.value, sharding_pattern=sharding_pattern, output_extension=output_extension, gpu=gpu, @@ -149,6 +246,7 @@ def generate( check_existing=check_existing, existing_msa_dirs=msa_dirs, search_config=search_config, + server=server_config, ) # Display configuration @@ -156,14 +254,19 @@ def generate( typer.echo(f" CSV File: {csv_file}") typer.echo(f" Sequence Column: {sequence_column or 'auto-detect'}") typer.echo(f" Output Directory: {output_dir}") - typer.echo(f" GPU Enabled: {config.gpu}") - typer.echo(f" Max Final Sequences: {config.max_final_sequences}") - typer.echo(f" Iterations: {config.num_iterations}") - typer.echo(f" Threads: {config.threads}") + typer.echo(f" Backend: {config.backend}") + if is_server: + typer.echo(f" Server URL: {config.server.host_url}") + typer.echo(f" Server Auth: {_describe_server_auth(config.server)}") + else: + typer.echo(f" GPU Enabled: {config.gpu}") + typer.echo(f" Iterations: {config.num_iterations}") + typer.echo(f" Threads: {config.threads}") + typer.echo(f" Sensitivity: {config.search_config.s}") + typer.echo(f" Max Final Sequences: {config.max_final_sequences if config.max_final_sequences else 'no filtering'}") typer.echo(f" Use Environmental DB: {config.use_env}") typer.echo(f" Output Extension: {config.output_extension}") typer.echo(f" Sharding Pattern: {config.sharding_pattern}") - typer.echo(f" Sensitivity: {config.search_config.s}") typer.echo(f" Check Existing: {config.check_existing}") if config.check_existing: dirs_display = config.existing_msa_dirs if config.existing_msa_dirs else "LOCAL_MSA_DIRS env var" @@ -176,3 +279,12 @@ def generate( except Exception as e: typer.secho(f"Error during MSA generation: {e!s}", fg=typer.colors.RED) raise typer.Exit(code=1) from e + + +def _describe_server_auth(config: MSAServerConfig) -> str: + """Summarize how requests to the MSA server will be authenticated (without echoing secrets).""" + if config.username is not None or config.password is not None: + return f"basic auth (user: {config.username})" + if config.api_key_value is not None: + return f"API key in header '{config.api_key_header}'" + return "none" diff --git a/tests/ml/preprocessing/msa/__init__.py b/tests/ml/preprocessing/msa/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/ml/preprocessing/msa/test_server.py b/tests/ml/preprocessing/msa/test_server.py new file mode 100644 index 00000000..7ef142da --- /dev/null +++ b/tests/ml/preprocessing/msa/test_server.py @@ -0,0 +1,454 @@ +"""Tests for the remote (ColabFold server) MSA generation backend. + +All HTTP primitives (`_submit`, `_status`, `_download`) are monkeypatched; no sockets are opened. +""" + +import base64 +import gzip +import io +import tarfile +import threading +import urllib.parse +from collections.abc import Iterator +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import pytest + +from atomworks.enums import MSAFileExtension +from atomworks.ml.preprocessing.msa import server as server_module +from atomworks.ml.preprocessing.msa.finding import find_msas +from atomworks.ml.preprocessing.msa.server import ( + ENV_A3M_FILENAME, + UNIREF_A3M_FILENAME, + MSAServerConfig, + _select_mode, + _split_multi_query_a3m, + make_msas_mmseqs_server, + run_mmseqs2_server, +) +from atomworks.ml.utils.misc import hash_sequence + +SEQ_A = "MSYIWRQLGSPTVAITLSVSTVIYVTVICPIVFIHLFGDHL" +SEQ_B = "MKKKEVEKDDLIENASRVASCISIFLIIASTTMYIFIGLKI" +SEQ_C = "MGSSHHHHHHSSGLVPRGSHMASMTGGQQMGRGSEFELRRQ" + + +@pytest.fixture(autouse=True) +def _clean_credentials_env(monkeypatch: pytest.MonkeyPatch) -> None: + """Make sure a developer's own server credentials don't leak into the tests.""" + for var in ("MSA_SERVER_USERNAME", "MSA_SERVER_PASSWORD", "MSA_SERVER_API_KEY"): + monkeypatch.delenv(var, raising=False) + + +@pytest.fixture +def config() -> MSAServerConfig: + """A server config that never actually sleeps.""" + return MSAServerConfig(poll_interval=(0.0, 0.0), retry_delay=0.0) + + +def _a3m_block(query_id: int, sequence: str, hit_prefix: str) -> str: + """Build one query's alignment block as the server would return it.""" + return ( + f">{query_id}\n{sequence}\n" + f">UniRef100_{hit_prefix}{query_id}\t91\t0.814\t1.694E-18\t0\t40\t41\t1\t41\t41\n" + f"{sequence[:-1]}-\n" + ) + + +def _multi_query_a3m(sequences: list[str], hit_prefix: str) -> str: + """Concatenate per-query blocks the way the server does: separated by null bytes.""" + blocks = [_a3m_block(101 + i, sequence, hit_prefix) for i, sequence in enumerate(sequences)] + return "\x00".join(blocks) + + +class FakeServer: + """Stand-in for a ColabFold MMseqs2 server, recording the calls made against it.""" + + def __init__(self, statuses: list[str] | None = None, submit_statuses: list[str] | None = None) -> None: + # Status sequence returned by consecutive `_status` polls; the last one repeats + self.statuses = statuses if statuses is not None else ["COMPLETE"] + self.submit_statuses = submit_statuses if submit_statuses is not None else ["PENDING"] + self.submitted_batches: list[list[str]] = [] + self.submitted_modes: list[str] = [] + self.n_status_calls = 0 + self.n_downloads = 0 + + def submit(self, sequences: list[str], mode: str, config: MSAServerConfig) -> dict[str, str]: + self.submitted_batches.append(list(sequences)) + self.submitted_modes.append(mode) + status = self.submit_statuses[min(len(self.submitted_batches) - 1, len(self.submit_statuses) - 1)] + return {"id": f"ticket-{len(self.submitted_batches)}", "status": status} + + def status(self, ticket_id: str, config: MSAServerConfig) -> dict[str, str]: + status = self.statuses[min(self.n_status_calls, len(self.statuses) - 1)] + self.n_status_calls += 1 + return {"id": ticket_id, "status": status} + + def download(self, ticket_id: str, dest: Path, config: MSAServerConfig) -> None: + self.n_downloads += 1 + sequences = self.submitted_batches[-1] + dest = Path(dest) + payload_dir = dest.parent / f"payload-{self.n_downloads}" + payload_dir.mkdir(exist_ok=True) + + members = {UNIREF_A3M_FILENAME: _multi_query_a3m(sequences, "UNI")} + if config.use_env: + members[ENV_A3M_FILENAME] = _multi_query_a3m(sequences, "ENV") + + with tarfile.open(dest, "w:gz") as tar: + for name, content in members.items(): + member_path = payload_dir / name + member_path.write_text(content) + tar.add(member_path, arcname=name) + + def install(self, monkeypatch: pytest.MonkeyPatch) -> "FakeServer": + monkeypatch.setattr(server_module, "_submit", self.submit) + monkeypatch.setattr(server_module, "_status", self.status) + monkeypatch.setattr(server_module, "_download", self.download) + return self + + +@pytest.fixture +def fake_server(monkeypatch: pytest.MonkeyPatch) -> FakeServer: + return FakeServer().install(monkeypatch) + + +# --- Configuration -------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("use_env", "use_filter", "expected_mode"), + [ + (True, True, "env"), + (False, True, "all"), + (True, False, "env-nofilter"), + (False, False, "nofilter"), + ], +) +def test_select_mode(use_env: bool, use_filter: bool, expected_mode: str) -> None: + assert _select_mode(MSAServerConfig(use_env=use_env, use_filter=use_filter)) == expected_mode + + +def test_config_rejects_basic_auth_and_api_key_together() -> None: + with pytest.raises(ValueError, match="basic auth"): + MSAServerConfig(username="me", password="secret", api_key_header="X-API-Key", api_key_value="abc") + + +def test_config_rejects_api_key_without_header() -> None: + with pytest.raises(ValueError, match="header name is required"): + MSAServerConfig(api_key_value="abc") + + +def test_config_reads_credentials_from_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MSA_SERVER_USERNAME", "env-user") + monkeypatch.setenv("MSA_SERVER_PASSWORD", "env-password") + config = MSAServerConfig() + assert (config.username, config.password) == ("env-user", "env-password") + + +def test_config_strips_trailing_slash_from_host_url() -> None: + assert MSAServerConfig(host_url="https://msa.internal/").host_url == "https://msa.internal" + + +# --- Parsing -------------------------------------------------------------------------- + + +def test_split_multi_query_a3m_splits_on_numeric_headers(tmp_path: Path) -> None: + a3m_file = tmp_path / "uniref.a3m" + a3m_file.write_text(_multi_query_a3m([SEQ_A, SEQ_B, SEQ_C], "UNI")) + + blocks = _split_multi_query_a3m(a3m_file) + + assert sorted(blocks) == [101, 102, 103] + assert blocks[102][0] == ">102\n" + assert blocks[102][1] == f"{SEQ_B}\n" + # Null bytes separating the queries are stripped, not carried into the alignment + assert not any("\x00" in line for lines in blocks.values() for line in lines) + + +def test_run_mmseqs2_server_writes_one_a3m_per_sequence(tmp_path: Path, fake_server: FakeServer, config) -> None: + paths = run_mmseqs2_server([SEQ_A, SEQ_B], tmp_path, config) + + assert set(paths) == {SEQ_A, SEQ_B} + for sequence, path in paths.items(): + lines = path.read_text().splitlines() + assert path.name == f"{hash_sequence(sequence)}.a3m" + # The numeric query header is rewritten to the sequence hash, as the local backend does + assert lines[0] == f">{hash_sequence(sequence)}" + assert lines[1] == sequence + # UniRef and environmental hits are concatenated, with the query kept exactly once + assert [line for line in lines if line.startswith(">")][1:] == [ + f">UniRef100_UNI{101 + list(paths).index(sequence)}\t91\t0.814\t1.694E-18\t0\t40\t41\t1\t41\t41", + f">UniRef100_ENV{101 + list(paths).index(sequence)}\t91\t0.814\t1.694E-18\t0\t40\t41\t1\t41\t41", + ] + + +def test_use_env_false_requests_uniref_only(tmp_path: Path, fake_server: FakeServer) -> None: + config = MSAServerConfig(use_env=False, poll_interval=(0.0, 0.0), retry_delay=0.0) + + paths = run_mmseqs2_server([SEQ_A], tmp_path, config) + + assert fake_server.submitted_modes == ["all"] + assert paths[SEQ_A].read_text().count(">") == 2 # query + one UniRef hit + + +def test_duplicate_sequences_are_submitted_once(tmp_path: Path, fake_server: FakeServer, config) -> None: + paths = run_mmseqs2_server([SEQ_A, SEQ_B, SEQ_A], tmp_path, config) + + assert fake_server.submitted_batches == [[SEQ_A, SEQ_B]] + assert len(paths) == 2 + assert paths[SEQ_A] == tmp_path / f"{hash_sequence(SEQ_A)}.a3m" + + +def test_sequences_are_split_into_batches(tmp_path: Path, fake_server: FakeServer) -> None: + config = MSAServerConfig(batch_size=2, poll_interval=(0.0, 0.0), retry_delay=0.0) + + paths = run_mmseqs2_server([SEQ_A, SEQ_B, SEQ_C], tmp_path, config) + + assert fake_server.submitted_batches == [[SEQ_A, SEQ_B], [SEQ_C]] + assert len(paths) == 3 + + +# --- Polling state machine ------------------------------------------------------------ + + +def test_ratelimited_submission_is_resubmitted(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config) -> None: + fake_server = FakeServer(submit_statuses=["RATELIMIT", "RATELIMIT", "PENDING"]).install(monkeypatch) + + run_mmseqs2_server([SEQ_A], tmp_path, config) + + assert len(fake_server.submitted_batches) == 3 + assert fake_server.n_downloads == 1 + + +def test_job_is_polled_until_complete(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config) -> None: + fake_server = FakeServer(statuses=["PENDING", "RUNNING", "RUNNING", "COMPLETE"]).install(monkeypatch) + + run_mmseqs2_server([SEQ_A], tmp_path, config) + + assert fake_server.n_status_calls == 4 + assert fake_server.n_downloads == 1 + + +def test_error_status_raises(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config) -> None: + fake_server = FakeServer(statuses=["ERROR"]).install(monkeypatch) + + with pytest.raises(RuntimeError, match="MSA server returned an error"): + run_mmseqs2_server([SEQ_A], tmp_path, config) + + assert fake_server.n_downloads == 0 + + +def test_maintenance_status_raises(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config) -> None: + FakeServer(submit_statuses=["MAINTENANCE"]).install(monkeypatch) + + with pytest.raises(RuntimeError, match="maintenance"): + run_mmseqs2_server([SEQ_A], tmp_path, config) + + +def test_network_failures_are_retried_then_raise(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config) -> None: + n_calls = 0 + + def always_fails(*args, **kwargs) -> dict[str, str]: + nonlocal n_calls + n_calls += 1 + raise ConnectionError("connection reset by peer") + + monkeypatch.setattr(server_module, "_submit", always_fails) + + with pytest.raises(RuntimeError, match="after 5 attempts"): + run_mmseqs2_server([SEQ_A], tmp_path, config) + + assert n_calls == config.max_retries + + +def test_transient_network_failure_is_retried(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config) -> None: + fake_server = FakeServer().install(monkeypatch) + real_submit = fake_server.submit + n_calls = 0 + + def flaky_submit(*args, **kwargs) -> dict[str, str]: + nonlocal n_calls + n_calls += 1 + if n_calls == 1: + raise TimeoutError("timed out") + return real_submit(*args, **kwargs) + + monkeypatch.setattr(server_module, "_submit", flaky_submit) + + run_mmseqs2_server([SEQ_A], tmp_path, config) + + assert fake_server.n_downloads == 1 + + +# --- End-to-end (mocked server) ------------------------------------------------------- + + +def test_make_msas_mmseqs_server_writes_sharded_store(tmp_path: Path, fake_server: FakeServer, config) -> None: + output_dir = tmp_path / "msas" + + make_msas_mmseqs_server([SEQ_A, SEQ_B], output_dir, config=config) + + for sequence in (SEQ_A, SEQ_B): + sequence_hash = hash_sequence(sequence) + expected = output_dir / sequence_hash[:2] / f"{sequence_hash}.a3m.gz" + assert expected.exists(), f"missing sharded MSA for {sequence_hash}" + with gzip.open(expected, "rt") as f: + assert f.readline().strip() == f">{sequence_hash}" + + # ... and the store round-trips through the standard MSA lookup + missing, found = find_msas( + [SEQ_A, SEQ_B], msa_dirs=[output_dir], shard_depths=[1], extensions=[MSAFileExtension.A3M_GZ] + ) + assert missing == [] + assert set(found) == {SEQ_A, SEQ_B} + + +def test_make_msas_mmseqs_server_honors_the_sharding_pattern(tmp_path: Path, fake_server: FakeServer, config) -> None: + make_msas_mmseqs_server([SEQ_A], tmp_path, config=config, sharding_pattern="/0:2/2:4/") + + sequence_hash = hash_sequence(SEQ_A) + assert (tmp_path / sequence_hash[:2] / sequence_hash[2:4] / f"{sequence_hash}.a3m.gz").exists() + + +def test_make_msas_mmseqs_server_skips_cached_sequences(tmp_path: Path, fake_server: FakeServer, config) -> None: + output_dir = tmp_path / "msas" + + make_msas_mmseqs_server([SEQ_A], output_dir, config=config) + assert len(fake_server.submitted_batches) == 1 + + # A warm cache makes no further requests... + make_msas_mmseqs_server([SEQ_A], output_dir, config=config) + assert len(fake_server.submitted_batches) == 1 + + # ... and only the uncached sequence is submitted + make_msas_mmseqs_server([SEQ_A, SEQ_B], output_dir, config=config) + assert fake_server.submitted_batches[-1] == [SEQ_B] + + +def test_make_msas_mmseqs_server_ignores_cache_when_check_existing_is_false( + tmp_path: Path, fake_server: FakeServer, config +) -> None: + output_dir = tmp_path / "msas" + + make_msas_mmseqs_server([SEQ_A], output_dir, config=config) + make_msas_mmseqs_server([SEQ_A], output_dir, config=config, check_existing=False) + + assert fake_server.submitted_batches == [[SEQ_A], [SEQ_A]] + + +def test_make_msas_mmseqs_server_accepts_a_single_sequence(tmp_path: Path, fake_server: FakeServer, config) -> None: + make_msas_mmseqs_server(SEQ_A, tmp_path, config=config) + + sequence_hash = hash_sequence(SEQ_A) + assert (tmp_path / sequence_hash[:2] / f"{sequence_hash}.a3m.gz").exists() + + +# --- Wire protocol (real HTTP against a loopback server) -------------------------------- + + +def _tarball_bytes(sequences: list[str]) -> bytes: + """Build an in-memory ``out.tar.gz`` holding a UniRef alignment for `sequences`.""" + buffer = io.BytesIO() + with tarfile.open(fileobj=buffer, mode="w:gz") as tar: + payload = _multi_query_a3m(sequences, "UNI").encode() + info = tarfile.TarInfo(UNIREF_A3M_FILENAME) + info.size = len(payload) + tar.addfile(info, io.BytesIO(payload)) + return buffer.getvalue() + + +class _LoopbackHandler(BaseHTTPRequestHandler): + """Minimal stand-in for the ColabFold API, recording the requests it receives.""" + + received: list[dict[str, object]] = [] + + def log_message(self, fmt: str, *args: object) -> None: # silence the default stderr logging + pass + + def _record(self, body: dict[str, str] | None = None) -> None: + self.received.append({"path": self.path, "headers": dict(self.headers), "body": body}) + + def _reply_json(self, payload: str) -> None: + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(payload.encode()) + + def do_POST(self) -> None: # noqa: N802 - name mandated by BaseHTTPRequestHandler + raw = self.rfile.read(int(self.headers["Content-Length"])) + self._record({k: v[0] for k, v in urllib.parse.parse_qs(raw.decode()).items()}) + self._reply_json('{"id": "ticket-1", "status": "PENDING"}') + + def do_GET(self) -> None: # noqa: N802 - name mandated by BaseHTTPRequestHandler + self._record() + if self.path.startswith("/ticket/"): + self._reply_json('{"id": "ticket-1", "status": "COMPLETE"}') + return + payload = _tarball_bytes([SEQ_A]) + self.send_response(200) + self.send_header("Content-Type", "application/gzip") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + +@pytest.fixture +def loopback_server() -> Iterator[tuple[str, list[dict[str, object]]]]: + """Serve the ColabFold endpoints on localhost, yielding the base URL and the received requests.""" + _LoopbackHandler.received = [] + httpd = ThreadingHTTPServer(("127.0.0.1", 0), _LoopbackHandler) + thread = threading.Thread(target=httpd.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{httpd.server_address[1]}", _LoopbackHandler.received + finally: + httpd.shutdown() + httpd.server_close() + thread.join(timeout=5) + + +def test_wire_protocol_against_a_loopback_server(tmp_path: Path, loopback_server) -> None: + host_url, received = loopback_server + config = MSAServerConfig( + host_url=host_url, + use_env=False, + username="me", + password="secret", + user_agent="atomworks-test", + poll_interval=(0.0, 0.0), + retry_delay=0.0, + ) + + paths = run_mmseqs2_server([SEQ_A], tmp_path, config) + + assert paths[SEQ_A].read_text().startswith(f">{hash_sequence(SEQ_A)}\n{SEQ_A}\n") + + submission, status_poll, download = received + assert submission["path"] == "/ticket/msa" + assert submission["body"] == {"q": f">101\n{SEQ_A}\n", "mode": "all"} + assert status_poll["path"] == "/ticket/ticket-1" + assert download["path"] == "/result/download/ticket-1" + + for request in received: + headers = request["headers"] + assert headers["User-Agent"] == "atomworks-test" + assert headers["Authorization"] == f"Basic {base64.b64encode(b'me:secret').decode()}" + + +def test_api_key_header_is_sent(tmp_path: Path, loopback_server) -> None: + host_url, received = loopback_server + config = MSAServerConfig( + host_url=host_url, + use_env=False, + api_key_header="X-API-Key", + api_key_value="s3cr3t", + poll_interval=(0.0, 0.0), + retry_delay=0.0, + ) + + run_mmseqs2_server([SEQ_A], tmp_path, config) + + assert all(request["headers"]["X-API-Key"] == "s3cr3t" for request in received) + assert all("Authorization" not in request["headers"] for request in received) From f50d6291a38d31f70dc1b7f292a2caf22eba2d90 Mon Sep 17 00:00:00 2001 From: rohith Date: Fri, 7 Aug 2026 17:19:07 -0700 Subject: [PATCH 2/4] remove server tests --- tests/ml/preprocessing/msa/test_server.py | 454 ---------------------- 1 file changed, 454 deletions(-) delete mode 100644 tests/ml/preprocessing/msa/test_server.py diff --git a/tests/ml/preprocessing/msa/test_server.py b/tests/ml/preprocessing/msa/test_server.py deleted file mode 100644 index 7ef142da..00000000 --- a/tests/ml/preprocessing/msa/test_server.py +++ /dev/null @@ -1,454 +0,0 @@ -"""Tests for the remote (ColabFold server) MSA generation backend. - -All HTTP primitives (`_submit`, `_status`, `_download`) are monkeypatched; no sockets are opened. -""" - -import base64 -import gzip -import io -import tarfile -import threading -import urllib.parse -from collections.abc import Iterator -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path - -import pytest - -from atomworks.enums import MSAFileExtension -from atomworks.ml.preprocessing.msa import server as server_module -from atomworks.ml.preprocessing.msa.finding import find_msas -from atomworks.ml.preprocessing.msa.server import ( - ENV_A3M_FILENAME, - UNIREF_A3M_FILENAME, - MSAServerConfig, - _select_mode, - _split_multi_query_a3m, - make_msas_mmseqs_server, - run_mmseqs2_server, -) -from atomworks.ml.utils.misc import hash_sequence - -SEQ_A = "MSYIWRQLGSPTVAITLSVSTVIYVTVICPIVFIHLFGDHL" -SEQ_B = "MKKKEVEKDDLIENASRVASCISIFLIIASTTMYIFIGLKI" -SEQ_C = "MGSSHHHHHHSSGLVPRGSHMASMTGGQQMGRGSEFELRRQ" - - -@pytest.fixture(autouse=True) -def _clean_credentials_env(monkeypatch: pytest.MonkeyPatch) -> None: - """Make sure a developer's own server credentials don't leak into the tests.""" - for var in ("MSA_SERVER_USERNAME", "MSA_SERVER_PASSWORD", "MSA_SERVER_API_KEY"): - monkeypatch.delenv(var, raising=False) - - -@pytest.fixture -def config() -> MSAServerConfig: - """A server config that never actually sleeps.""" - return MSAServerConfig(poll_interval=(0.0, 0.0), retry_delay=0.0) - - -def _a3m_block(query_id: int, sequence: str, hit_prefix: str) -> str: - """Build one query's alignment block as the server would return it.""" - return ( - f">{query_id}\n{sequence}\n" - f">UniRef100_{hit_prefix}{query_id}\t91\t0.814\t1.694E-18\t0\t40\t41\t1\t41\t41\n" - f"{sequence[:-1]}-\n" - ) - - -def _multi_query_a3m(sequences: list[str], hit_prefix: str) -> str: - """Concatenate per-query blocks the way the server does: separated by null bytes.""" - blocks = [_a3m_block(101 + i, sequence, hit_prefix) for i, sequence in enumerate(sequences)] - return "\x00".join(blocks) - - -class FakeServer: - """Stand-in for a ColabFold MMseqs2 server, recording the calls made against it.""" - - def __init__(self, statuses: list[str] | None = None, submit_statuses: list[str] | None = None) -> None: - # Status sequence returned by consecutive `_status` polls; the last one repeats - self.statuses = statuses if statuses is not None else ["COMPLETE"] - self.submit_statuses = submit_statuses if submit_statuses is not None else ["PENDING"] - self.submitted_batches: list[list[str]] = [] - self.submitted_modes: list[str] = [] - self.n_status_calls = 0 - self.n_downloads = 0 - - def submit(self, sequences: list[str], mode: str, config: MSAServerConfig) -> dict[str, str]: - self.submitted_batches.append(list(sequences)) - self.submitted_modes.append(mode) - status = self.submit_statuses[min(len(self.submitted_batches) - 1, len(self.submit_statuses) - 1)] - return {"id": f"ticket-{len(self.submitted_batches)}", "status": status} - - def status(self, ticket_id: str, config: MSAServerConfig) -> dict[str, str]: - status = self.statuses[min(self.n_status_calls, len(self.statuses) - 1)] - self.n_status_calls += 1 - return {"id": ticket_id, "status": status} - - def download(self, ticket_id: str, dest: Path, config: MSAServerConfig) -> None: - self.n_downloads += 1 - sequences = self.submitted_batches[-1] - dest = Path(dest) - payload_dir = dest.parent / f"payload-{self.n_downloads}" - payload_dir.mkdir(exist_ok=True) - - members = {UNIREF_A3M_FILENAME: _multi_query_a3m(sequences, "UNI")} - if config.use_env: - members[ENV_A3M_FILENAME] = _multi_query_a3m(sequences, "ENV") - - with tarfile.open(dest, "w:gz") as tar: - for name, content in members.items(): - member_path = payload_dir / name - member_path.write_text(content) - tar.add(member_path, arcname=name) - - def install(self, monkeypatch: pytest.MonkeyPatch) -> "FakeServer": - monkeypatch.setattr(server_module, "_submit", self.submit) - monkeypatch.setattr(server_module, "_status", self.status) - monkeypatch.setattr(server_module, "_download", self.download) - return self - - -@pytest.fixture -def fake_server(monkeypatch: pytest.MonkeyPatch) -> FakeServer: - return FakeServer().install(monkeypatch) - - -# --- Configuration -------------------------------------------------------------------- - - -@pytest.mark.parametrize( - ("use_env", "use_filter", "expected_mode"), - [ - (True, True, "env"), - (False, True, "all"), - (True, False, "env-nofilter"), - (False, False, "nofilter"), - ], -) -def test_select_mode(use_env: bool, use_filter: bool, expected_mode: str) -> None: - assert _select_mode(MSAServerConfig(use_env=use_env, use_filter=use_filter)) == expected_mode - - -def test_config_rejects_basic_auth_and_api_key_together() -> None: - with pytest.raises(ValueError, match="basic auth"): - MSAServerConfig(username="me", password="secret", api_key_header="X-API-Key", api_key_value="abc") - - -def test_config_rejects_api_key_without_header() -> None: - with pytest.raises(ValueError, match="header name is required"): - MSAServerConfig(api_key_value="abc") - - -def test_config_reads_credentials_from_env(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("MSA_SERVER_USERNAME", "env-user") - monkeypatch.setenv("MSA_SERVER_PASSWORD", "env-password") - config = MSAServerConfig() - assert (config.username, config.password) == ("env-user", "env-password") - - -def test_config_strips_trailing_slash_from_host_url() -> None: - assert MSAServerConfig(host_url="https://msa.internal/").host_url == "https://msa.internal" - - -# --- Parsing -------------------------------------------------------------------------- - - -def test_split_multi_query_a3m_splits_on_numeric_headers(tmp_path: Path) -> None: - a3m_file = tmp_path / "uniref.a3m" - a3m_file.write_text(_multi_query_a3m([SEQ_A, SEQ_B, SEQ_C], "UNI")) - - blocks = _split_multi_query_a3m(a3m_file) - - assert sorted(blocks) == [101, 102, 103] - assert blocks[102][0] == ">102\n" - assert blocks[102][1] == f"{SEQ_B}\n" - # Null bytes separating the queries are stripped, not carried into the alignment - assert not any("\x00" in line for lines in blocks.values() for line in lines) - - -def test_run_mmseqs2_server_writes_one_a3m_per_sequence(tmp_path: Path, fake_server: FakeServer, config) -> None: - paths = run_mmseqs2_server([SEQ_A, SEQ_B], tmp_path, config) - - assert set(paths) == {SEQ_A, SEQ_B} - for sequence, path in paths.items(): - lines = path.read_text().splitlines() - assert path.name == f"{hash_sequence(sequence)}.a3m" - # The numeric query header is rewritten to the sequence hash, as the local backend does - assert lines[0] == f">{hash_sequence(sequence)}" - assert lines[1] == sequence - # UniRef and environmental hits are concatenated, with the query kept exactly once - assert [line for line in lines if line.startswith(">")][1:] == [ - f">UniRef100_UNI{101 + list(paths).index(sequence)}\t91\t0.814\t1.694E-18\t0\t40\t41\t1\t41\t41", - f">UniRef100_ENV{101 + list(paths).index(sequence)}\t91\t0.814\t1.694E-18\t0\t40\t41\t1\t41\t41", - ] - - -def test_use_env_false_requests_uniref_only(tmp_path: Path, fake_server: FakeServer) -> None: - config = MSAServerConfig(use_env=False, poll_interval=(0.0, 0.0), retry_delay=0.0) - - paths = run_mmseqs2_server([SEQ_A], tmp_path, config) - - assert fake_server.submitted_modes == ["all"] - assert paths[SEQ_A].read_text().count(">") == 2 # query + one UniRef hit - - -def test_duplicate_sequences_are_submitted_once(tmp_path: Path, fake_server: FakeServer, config) -> None: - paths = run_mmseqs2_server([SEQ_A, SEQ_B, SEQ_A], tmp_path, config) - - assert fake_server.submitted_batches == [[SEQ_A, SEQ_B]] - assert len(paths) == 2 - assert paths[SEQ_A] == tmp_path / f"{hash_sequence(SEQ_A)}.a3m" - - -def test_sequences_are_split_into_batches(tmp_path: Path, fake_server: FakeServer) -> None: - config = MSAServerConfig(batch_size=2, poll_interval=(0.0, 0.0), retry_delay=0.0) - - paths = run_mmseqs2_server([SEQ_A, SEQ_B, SEQ_C], tmp_path, config) - - assert fake_server.submitted_batches == [[SEQ_A, SEQ_B], [SEQ_C]] - assert len(paths) == 3 - - -# --- Polling state machine ------------------------------------------------------------ - - -def test_ratelimited_submission_is_resubmitted(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config) -> None: - fake_server = FakeServer(submit_statuses=["RATELIMIT", "RATELIMIT", "PENDING"]).install(monkeypatch) - - run_mmseqs2_server([SEQ_A], tmp_path, config) - - assert len(fake_server.submitted_batches) == 3 - assert fake_server.n_downloads == 1 - - -def test_job_is_polled_until_complete(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config) -> None: - fake_server = FakeServer(statuses=["PENDING", "RUNNING", "RUNNING", "COMPLETE"]).install(monkeypatch) - - run_mmseqs2_server([SEQ_A], tmp_path, config) - - assert fake_server.n_status_calls == 4 - assert fake_server.n_downloads == 1 - - -def test_error_status_raises(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config) -> None: - fake_server = FakeServer(statuses=["ERROR"]).install(monkeypatch) - - with pytest.raises(RuntimeError, match="MSA server returned an error"): - run_mmseqs2_server([SEQ_A], tmp_path, config) - - assert fake_server.n_downloads == 0 - - -def test_maintenance_status_raises(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config) -> None: - FakeServer(submit_statuses=["MAINTENANCE"]).install(monkeypatch) - - with pytest.raises(RuntimeError, match="maintenance"): - run_mmseqs2_server([SEQ_A], tmp_path, config) - - -def test_network_failures_are_retried_then_raise(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config) -> None: - n_calls = 0 - - def always_fails(*args, **kwargs) -> dict[str, str]: - nonlocal n_calls - n_calls += 1 - raise ConnectionError("connection reset by peer") - - monkeypatch.setattr(server_module, "_submit", always_fails) - - with pytest.raises(RuntimeError, match="after 5 attempts"): - run_mmseqs2_server([SEQ_A], tmp_path, config) - - assert n_calls == config.max_retries - - -def test_transient_network_failure_is_retried(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config) -> None: - fake_server = FakeServer().install(monkeypatch) - real_submit = fake_server.submit - n_calls = 0 - - def flaky_submit(*args, **kwargs) -> dict[str, str]: - nonlocal n_calls - n_calls += 1 - if n_calls == 1: - raise TimeoutError("timed out") - return real_submit(*args, **kwargs) - - monkeypatch.setattr(server_module, "_submit", flaky_submit) - - run_mmseqs2_server([SEQ_A], tmp_path, config) - - assert fake_server.n_downloads == 1 - - -# --- End-to-end (mocked server) ------------------------------------------------------- - - -def test_make_msas_mmseqs_server_writes_sharded_store(tmp_path: Path, fake_server: FakeServer, config) -> None: - output_dir = tmp_path / "msas" - - make_msas_mmseqs_server([SEQ_A, SEQ_B], output_dir, config=config) - - for sequence in (SEQ_A, SEQ_B): - sequence_hash = hash_sequence(sequence) - expected = output_dir / sequence_hash[:2] / f"{sequence_hash}.a3m.gz" - assert expected.exists(), f"missing sharded MSA for {sequence_hash}" - with gzip.open(expected, "rt") as f: - assert f.readline().strip() == f">{sequence_hash}" - - # ... and the store round-trips through the standard MSA lookup - missing, found = find_msas( - [SEQ_A, SEQ_B], msa_dirs=[output_dir], shard_depths=[1], extensions=[MSAFileExtension.A3M_GZ] - ) - assert missing == [] - assert set(found) == {SEQ_A, SEQ_B} - - -def test_make_msas_mmseqs_server_honors_the_sharding_pattern(tmp_path: Path, fake_server: FakeServer, config) -> None: - make_msas_mmseqs_server([SEQ_A], tmp_path, config=config, sharding_pattern="/0:2/2:4/") - - sequence_hash = hash_sequence(SEQ_A) - assert (tmp_path / sequence_hash[:2] / sequence_hash[2:4] / f"{sequence_hash}.a3m.gz").exists() - - -def test_make_msas_mmseqs_server_skips_cached_sequences(tmp_path: Path, fake_server: FakeServer, config) -> None: - output_dir = tmp_path / "msas" - - make_msas_mmseqs_server([SEQ_A], output_dir, config=config) - assert len(fake_server.submitted_batches) == 1 - - # A warm cache makes no further requests... - make_msas_mmseqs_server([SEQ_A], output_dir, config=config) - assert len(fake_server.submitted_batches) == 1 - - # ... and only the uncached sequence is submitted - make_msas_mmseqs_server([SEQ_A, SEQ_B], output_dir, config=config) - assert fake_server.submitted_batches[-1] == [SEQ_B] - - -def test_make_msas_mmseqs_server_ignores_cache_when_check_existing_is_false( - tmp_path: Path, fake_server: FakeServer, config -) -> None: - output_dir = tmp_path / "msas" - - make_msas_mmseqs_server([SEQ_A], output_dir, config=config) - make_msas_mmseqs_server([SEQ_A], output_dir, config=config, check_existing=False) - - assert fake_server.submitted_batches == [[SEQ_A], [SEQ_A]] - - -def test_make_msas_mmseqs_server_accepts_a_single_sequence(tmp_path: Path, fake_server: FakeServer, config) -> None: - make_msas_mmseqs_server(SEQ_A, tmp_path, config=config) - - sequence_hash = hash_sequence(SEQ_A) - assert (tmp_path / sequence_hash[:2] / f"{sequence_hash}.a3m.gz").exists() - - -# --- Wire protocol (real HTTP against a loopback server) -------------------------------- - - -def _tarball_bytes(sequences: list[str]) -> bytes: - """Build an in-memory ``out.tar.gz`` holding a UniRef alignment for `sequences`.""" - buffer = io.BytesIO() - with tarfile.open(fileobj=buffer, mode="w:gz") as tar: - payload = _multi_query_a3m(sequences, "UNI").encode() - info = tarfile.TarInfo(UNIREF_A3M_FILENAME) - info.size = len(payload) - tar.addfile(info, io.BytesIO(payload)) - return buffer.getvalue() - - -class _LoopbackHandler(BaseHTTPRequestHandler): - """Minimal stand-in for the ColabFold API, recording the requests it receives.""" - - received: list[dict[str, object]] = [] - - def log_message(self, fmt: str, *args: object) -> None: # silence the default stderr logging - pass - - def _record(self, body: dict[str, str] | None = None) -> None: - self.received.append({"path": self.path, "headers": dict(self.headers), "body": body}) - - def _reply_json(self, payload: str) -> None: - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(payload.encode()) - - def do_POST(self) -> None: # noqa: N802 - name mandated by BaseHTTPRequestHandler - raw = self.rfile.read(int(self.headers["Content-Length"])) - self._record({k: v[0] for k, v in urllib.parse.parse_qs(raw.decode()).items()}) - self._reply_json('{"id": "ticket-1", "status": "PENDING"}') - - def do_GET(self) -> None: # noqa: N802 - name mandated by BaseHTTPRequestHandler - self._record() - if self.path.startswith("/ticket/"): - self._reply_json('{"id": "ticket-1", "status": "COMPLETE"}') - return - payload = _tarball_bytes([SEQ_A]) - self.send_response(200) - self.send_header("Content-Type", "application/gzip") - self.send_header("Content-Length", str(len(payload))) - self.end_headers() - self.wfile.write(payload) - - -@pytest.fixture -def loopback_server() -> Iterator[tuple[str, list[dict[str, object]]]]: - """Serve the ColabFold endpoints on localhost, yielding the base URL and the received requests.""" - _LoopbackHandler.received = [] - httpd = ThreadingHTTPServer(("127.0.0.1", 0), _LoopbackHandler) - thread = threading.Thread(target=httpd.serve_forever, daemon=True) - thread.start() - try: - yield f"http://127.0.0.1:{httpd.server_address[1]}", _LoopbackHandler.received - finally: - httpd.shutdown() - httpd.server_close() - thread.join(timeout=5) - - -def test_wire_protocol_against_a_loopback_server(tmp_path: Path, loopback_server) -> None: - host_url, received = loopback_server - config = MSAServerConfig( - host_url=host_url, - use_env=False, - username="me", - password="secret", - user_agent="atomworks-test", - poll_interval=(0.0, 0.0), - retry_delay=0.0, - ) - - paths = run_mmseqs2_server([SEQ_A], tmp_path, config) - - assert paths[SEQ_A].read_text().startswith(f">{hash_sequence(SEQ_A)}\n{SEQ_A}\n") - - submission, status_poll, download = received - assert submission["path"] == "/ticket/msa" - assert submission["body"] == {"q": f">101\n{SEQ_A}\n", "mode": "all"} - assert status_poll["path"] == "/ticket/ticket-1" - assert download["path"] == "/result/download/ticket-1" - - for request in received: - headers = request["headers"] - assert headers["User-Agent"] == "atomworks-test" - assert headers["Authorization"] == f"Basic {base64.b64encode(b'me:secret').decode()}" - - -def test_api_key_header_is_sent(tmp_path: Path, loopback_server) -> None: - host_url, received = loopback_server - config = MSAServerConfig( - host_url=host_url, - use_env=False, - api_key_header="X-API-Key", - api_key_value="s3cr3t", - poll_interval=(0.0, 0.0), - retry_delay=0.0, - ) - - run_mmseqs2_server([SEQ_A], tmp_path, config) - - assert all(request["headers"]["X-API-Key"] == "s3cr3t" for request in received) - assert all("Authorization" not in request["headers"] for request in received) From d6c52cb3f7dd995a964e4ef9f9b96646dba37d12 Mon Sep 17 00:00:00 2001 From: rohith Date: Fri, 7 Aug 2026 17:27:33 -0700 Subject: [PATCH 3/4] remove init --- tests/ml/preprocessing/msa/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 tests/ml/preprocessing/msa/__init__.py diff --git a/tests/ml/preprocessing/msa/__init__.py b/tests/ml/preprocessing/msa/__init__.py deleted file mode 100644 index e69de29b..00000000 From bdc64fdbcc15a7b43fc1a3d88b46786b3f32f533 Mon Sep 17 00:00:00 2001 From: rohith Date: Fri, 7 Aug 2026 17:55:13 -0700 Subject: [PATCH 4/4] chore ruff --- docs/conf.py | 4 ++-- uv.lock | 64 +++++++++++++++++++++++++++------------------------- 2 files changed, 35 insertions(+), 33 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 358c8db6..855b0b14 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -97,5 +97,5 @@ } html_js_files = [ - ('https://scripts.simpleanalyticscdn.com/latest.js', {'async': 'async', 'defer': 'defer'}), -] \ No newline at end of file + ("https://scripts.simpleanalyticscdn.com/latest.js", {"async": "async", "defer": "defer"}), +] diff --git a/uv.lock b/uv.lock index b285e636..758320d0 100644 --- a/uv.lock +++ b/uv.lock @@ -78,7 +78,7 @@ wheels = [ [[package]] name = "atomworks" -version = "2.1.2" +version = "2.2.1" source = { editable = "." } dependencies = [ { name = "biotite" }, @@ -91,6 +91,7 @@ dependencies = [ { name = "pyarrow" }, { name = "pymol-remote" }, { name = "rdkit" }, + { name = "requests" }, { name = "scipy" }, { name = "tqdm" }, { name = "typer" }, @@ -145,40 +146,41 @@ openbabel = [ [package.metadata] requires-dist = [ - { name = "ase", marker = "extra == 'ase'", specifier = ">=3.22.0,<4" }, - { name = "ase-db-backends", marker = "extra == 'ase'", specifier = ">=0.10.0,<1" }, - { name = "beartype", marker = "extra == 'ml'", specifier = ">=0.18.0,<1" }, + { name = "ase", marker = "extra == 'ase'", specifier = ">=3.22.0" }, + { name = "ase-db-backends", marker = "extra == 'ase'", specifier = ">=0.10.0" }, + { name = "beartype", marker = "extra == 'ml'", specifier = ">=0.18.0" }, { name = "biotite", specifier = "==1.4.0" }, - { name = "cython", specifier = ">=3.0.0,<4" }, - { name = "cytoolz", specifier = ">=0.12.3,<1" }, - { name = "einops", marker = "extra == 'ml'", specifier = ">=0.7.0,<1" }, - { name = "hydride", specifier = ">=1.2.3,<2" }, + { name = "cython", specifier = ">=3.0.0" }, + { name = "cytoolz", specifier = ">=0.12.3" }, + { name = "einops", marker = "extra == 'ml'", specifier = ">=0.7.0" }, + { name = "hydride", specifier = ">=1.2.3" }, { name = "ipykernel", marker = "extra == 'dev'", specifier = ">=6.28.0" }, - { name = "jaxtyping", marker = "extra == 'ml'", specifier = ">=0.2.17,<1" }, - { name = "lmdb", marker = "extra == 'ase'", specifier = ">=1.5.0,<2" }, - { name = "matplotlib", marker = "extra == 'docs'", specifier = ">=3.10.0,<4" }, - { name = "numpy", specifier = ">=1.25.0,<3" }, - { name = "openbabel-wheel", marker = "extra == 'openbabel'", specifier = "==3.1.1.22" }, - { name = "pandas", specifier = ">=2.2,<2.4" }, - { name = "py3dmol", specifier = ">=2.2.1,<3" }, - { name = "pyarrow", specifier = "==17.0.0" }, - { name = "pydata-sphinx-theme", marker = "extra == 'docs'", specifier = ">=0.16.1,<1" }, + { name = "jaxtyping", marker = "extra == 'ml'", specifier = ">=0.2.17" }, + { name = "lmdb", marker = "extra == 'ase'", specifier = ">=1.5.0" }, + { name = "matplotlib", marker = "extra == 'docs'", specifier = ">=3.10.0" }, + { name = "numpy", specifier = ">=1.25.0" }, + { name = "openbabel-wheel", marker = "extra == 'openbabel'", specifier = ">=3.1.1.22" }, + { name = "pandas", specifier = ">=2.2" }, + { name = "py3dmol", specifier = ">=2.2.1" }, + { name = "pyarrow", specifier = ">=17.0.0" }, + { name = "pydata-sphinx-theme", marker = "extra == 'docs'", specifier = ">=0.16.1" }, { name = "pymol-remote", specifier = ">=0.0.5" }, - { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.2.0,<9" }, - { name = "pytest-benchmark", marker = "extra == 'dev'", specifier = ">=5.0.0,<6" }, - { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.1.0,<5" }, - { name = "pytest-dotenv", marker = "extra == 'dev'", specifier = ">=0.5.2,<1" }, - { name = "pytest-testmon", marker = "extra == 'dev'", specifier = ">=2.1.1,<3" }, - { name = "pytest-xdist", marker = "extra == 'dev'", specifier = ">=3.6.1,<4" }, - { name = "rdkit", specifier = ">=2024.3.5,<2025.9" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.2.0" }, + { name = "pytest-benchmark", marker = "extra == 'dev'", specifier = ">=5.0.0" }, + { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.1.0" }, + { name = "pytest-dotenv", marker = "extra == 'dev'", specifier = ">=0.5.2" }, + { name = "pytest-testmon", marker = "extra == 'dev'", specifier = ">=2.1.1" }, + { name = "pytest-xdist", marker = "extra == 'dev'", specifier = ">=3.6.1" }, + { name = "rdkit", specifier = ">=2024.3.5" }, + { name = "requests", specifier = ">=2.32" }, { name = "ruff", marker = "extra == 'dev'", specifier = "==0.8.3" }, - { name = "scipy", specifier = ">=1.13.1,<2" }, - { name = "sphinx", marker = "extra == 'docs'", specifier = ">=8.0.0,<9" }, - { name = "sphinx-gallery", marker = "extra == 'docs'", specifier = ">=0.19.0,<1" }, - { name = "torch", marker = "extra == 'ml'", specifier = ">=2.2.0,<2.8" }, - { name = "tqdm", specifier = ">=4.65.0,<5" }, - { name = "typer", specifier = ">=0.12.5,<1" }, - { name = "zstandard", specifier = ">=0.21.0,<1" }, + { name = "scipy", specifier = ">=1.13.1" }, + { name = "sphinx", marker = "extra == 'docs'", specifier = ">=8.0.0" }, + { name = "sphinx-gallery", marker = "extra == 'docs'", specifier = ">=0.19.0" }, + { name = "torch", marker = "extra == 'ml'", specifier = ">=2.2.0" }, + { name = "tqdm", specifier = ">=4.65.0" }, + { name = "typer", specifier = ">=0.12.5" }, + { name = "zstandard", specifier = ">=0.21.0" }, ] provides-extras = ["ml", "ase", "openbabel", "dev", "docs"]