apherisfold-cli 0.1.6__py3-none-any.whl

apherisfold_cli/__init__.py

@@ -0,0 +1,3 @@
+from .version import __version__
+
+__all__ = ["__version__"]

apherisfold_cli/commands/__init__.py

@@ -0,0 +1,3 @@
+from . import benchmark, fine_tune, jobs, login, predict, weights, workflow
+
+__all__ = ["benchmark", "fine_tune", "jobs", "login", "predict", "weights", "workflow"]

apherisfold_cli/commands/benchmark.py

@@ -0,0 +1,310 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Annotated, Any
+
+import typer
+from rich import box
+from rich.console import Group
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+
+from apherisfold_cli.commands.predict import _resolve_prediction_capability, load_model_params_schema
+from apherisfold_cli.core.config import load_config, resolve_command_config
+from apherisfold_cli.core.errors import EXIT_USAGE, CLIError
+from apherisfold_cli.core.hub_models import resolve_weight
+from apherisfold_cli.core.hub_requests import (
+    create_benchmark,
+    start_benchmark,
+    upload_benchmark_structure,
+    upload_benchmark_structure_msa,
+)
+from apherisfold_cli.core.io import ensure_output_dir, write_result, write_status
+from apherisfold_cli.core.model_schema import validate_schema_object
+from apherisfold_cli.core.predict_inputs import load_json_object
+from apherisfold_cli.handlers import get_ctx, run_simple_listing
+
+_STRUCTURE_SUFFIXES = {".cif", ".mmcif"}
+
+
+def benchmark(
+    output_dir: Annotated[
+        Path,
+        typer.Option(
+            "--output",
+            "--output-dir",
+            "-o",
+            dir_okay=True,
+            file_okay=False,
+            help="Directory for workflow artifacts and downloaded outputs.",
+        ),
+    ],
+    name: Annotated[str, typer.Option("--name", "-n", help="Human-friendly name for this benchmark run.")],
+    model: Annotated[str, typer.Option("--model", "-m", help="Installed model identifier, for example 'openfold3'.")],
+    weight: Annotated[
+        str,
+        typer.Option("--weight", "-w", help="Weight version shown in 'apherisfold weights list', for example '3.0.0'."),
+    ],
+    input_paths: Annotated[
+        list[Path] | None,
+        typer.Option("--input", help="Repeat for CIF files or pass a directory containing CIF files."),
+    ] = None,
+    model_version: Annotated[
+        str | None,
+        typer.Option(
+            "--model-version",
+            help="Installed model version/build identifier when selection is ambiguous, for example 'openfold3:1.2.3'.",
+        ),
+    ] = None,
+    model_params: Annotated[
+        str | None, typer.Option("--model-params", help="Hub-style modelParams JSON object, inline or as @file.json.")
+    ] = None,
+    model_param_overrides: dict[str, Any] | None = None,
+) -> None:
+    selected_weight = resolve_weight(model, model_version, weight, task="benchmark")
+    selected_capability = "inference"
+    if selected_capability not in selected_weight["capabilities"]:
+        selected_capability = _resolve_prediction_capability(available_capabilities=selected_weight["capabilities"])
+    model_params_schema = load_model_params_schema(
+        model_version_id=str(selected_weight["model_version"]),
+        weight_version=str(selected_weight["weight"]),
+        capability=selected_capability,
+        task="benchmark",
+    )
+    loaded_config = load_config(get_ctx().config_raw)
+    resolved = resolve_command_config(
+        command_name="benchmark",
+        defaults={
+            "name": name,
+            "model": model,
+            "weight": str(selected_weight["weight"]),
+            "input": [],
+            "model_params": {},
+        },
+        loaded_config=loaded_config,
+        cli_overrides={
+            "name": name,
+            "model": model,
+            "weight": str(selected_weight["weight"]),
+            "input": input_paths,
+            "model_params": load_json_object(model_params, option_name="--model-params", task="benchmark")
+            if model_params is not None
+            else None,
+        },
+        allowed_keys={"name", "model", "weight", "input", "model_params"},
+    )
+    raw_inputs = resolved.get("input")
+    if not isinstance(raw_inputs, list) or not raw_inputs:
+        raise CLIError(
+            code=EXIT_USAGE,
+            category="usage",
+            message="Benchmark requires at least one '--input' path.",
+            task="benchmark",
+        )
+    normalized_inputs = [Path(value) for value in raw_inputs]
+    resolved_inputs = _resolve_benchmark_inputs(normalized_inputs)
+
+    ctx = get_ctx()
+    ctx.output = output_dir
+    ctx.extra["write_errors"] = True
+    ensure_output_dir(output_dir, overwrite=ctx.overwrite)
+
+    resolved_model_params = dict(resolved.get("model_params") or {})
+    resolved_model_params.update(model_param_overrides or {})
+    validate_schema_object(
+        resolved_model_params,
+        model_params_schema,
+        task="benchmark",
+        label="model parameter",
+    )
+
+    result_path = (ctx.output / "result.json").expanduser().resolve()
+
+    if ctx.dry_run:
+        result = {
+            "schema_version": "1.0",
+            "kind": "benchmark",
+            "status": "dry-run",
+            "request": {
+                "name": resolved["name"],
+                "model": resolved["model"],
+                "weight": resolved["weight"],
+                "input": [str(path) for path in normalized_inputs],
+                "structures": [str(path) for path in resolved_inputs["structures"]],
+                "paired_msas": [
+                    {"structure": str(structure), "msa": str(msa)} for structure, msa in resolved_inputs["paired_msas"]
+                ],
+                "model_params": resolved_model_params,
+            },
+            "summary": {
+                "queued_structures": len(resolved_inputs["structures"]),
+                "paired_msa_count": len(resolved_inputs["paired_msas"]),
+            },
+            "result_path": str(result_path),
+        }
+        write_status(ctx, task="benchmark", phase="validation", state="complete")
+        write_result(ctx, result)
+        run_simple_listing(
+            task_name="benchmark",
+            payload=result,
+            renderable_output=_benchmark_renderable("Benchmark Dry Run", result),
+        )
+        return
+
+    write_status(ctx, task="benchmark", phase="creating-benchmark", state="running")
+    created = create_benchmark(
+        name=str(resolved["name"]),
+        model_version_id=str(selected_weight["model_version"]),
+        weight_version=str(selected_weight["weight"]),
+        model_params=resolved_model_params,
+    )
+    benchmark_id = str(created["id"])
+
+    for structure_path in resolved_inputs["structures"]:
+        query_id = structure_path.stem
+        write_status(ctx, task="benchmark", phase=f"uploading-structure:{query_id}", state="running")
+        upload_benchmark_structure(
+            benchmark_id=benchmark_id,
+            query_id=query_id,
+            cif_path=structure_path,
+        )
+
+    for structure_path, msa_path in resolved_inputs["paired_msas"]:
+        query_id = structure_path.stem
+        write_status(ctx, task="benchmark", phase=f"uploading-msa:{query_id}", state="running")
+        upload_benchmark_structure_msa(
+            benchmark_id=benchmark_id,
+            query_id=query_id,
+            msa_path=msa_path,
+        )
+
+    write_status(ctx, task="benchmark", phase="starting", state="running")
+    start_benchmark(benchmark_id)
+
+    write_status(ctx, task="benchmark", phase="submitted", state="complete")
+    result = {
+        "schema_version": "1.0",
+        "kind": "benchmark",
+        "status": "submitted",
+        "job_id": benchmark_id,
+        "request": created,
+        "submission": {
+            "name": resolved["name"],
+            "model": resolved["model"],
+            "weight": resolved["weight"],
+            "structures": [str(path) for path in resolved_inputs["structures"]],
+            "paired_msas": [
+                {"structure": str(structure), "msa": str(msa)} for structure, msa in resolved_inputs["paired_msas"]
+            ],
+            "model_params": resolved_model_params,
+        },
+        "summary": {
+            "queued_structures": len(resolved_inputs["structures"]),
+            "paired_msa_count": len(resolved_inputs["paired_msas"]),
+        },
+        "result_path": str(result_path),
+    }
+    write_result(ctx, result)
+    next_cmd = f"apherisfold workflow benchmark get --id {benchmark_id}"
+    run_simple_listing(
+        task_name="benchmark",
+        payload=result,
+        renderable_output=_benchmark_renderable("Benchmark Submitted", result, next_cmd=next_cmd),
+    )
+
+
+def _resolve_benchmark_inputs(input_paths: list[Path]) -> dict[str, list]:
+    structures: list[Path] = []
+    paired_msas: list[tuple[Path, Path]] = []
+    seen_structures: set[Path] = set()
+
+    for raw_path in input_paths:
+        path = raw_path.expanduser()
+        if not path.exists():
+            raise CLIError(
+                code=EXIT_USAGE,
+                category="usage",
+                message=f"Benchmark input path does not exist: {raw_path}",
+                task="benchmark",
+            )
+        if path.is_dir():
+            directory_structures = sorted(
+                child for child in path.iterdir() if child.is_file() and child.suffix.lower() in _STRUCTURE_SUFFIXES
+            )
+            if not directory_structures:
+                raise CLIError(
+                    code=EXIT_USAGE,
+                    category="usage",
+                    message=f"Benchmark input directory does not contain any CIF files: {raw_path}",
+                    task="benchmark",
+                )
+            for structure_path in directory_structures:
+                _append_structure(structure_path, structures, paired_msas, seen_structures)
+            continue
+        if not path.is_file():
+            raise CLIError(
+                code=EXIT_USAGE,
+                category="usage",
+                message=f"Benchmark input is not a regular file or directory: {raw_path}",
+                task="benchmark",
+            )
+        if path.suffix.lower() not in _STRUCTURE_SUFFIXES:
+            raise CLIError(
+                code=EXIT_USAGE,
+                category="usage",
+                message=f"Benchmark input files must be CIF or mmCIF: {raw_path}",
+                task="benchmark",
+            )
+        _append_structure(path, structures, paired_msas, seen_structures)
+
+    return {"structures": structures, "paired_msas": paired_msas}
+
+
+def _append_structure(
+    structure_path: Path,
+    structures: list[Path],
+    paired_msas: list[tuple[Path, Path]],
+    seen_structures: set[Path],
+) -> None:
+    resolved_path = structure_path.resolve()
+    if resolved_path in seen_structures:
+        return
+    seen_structures.add(resolved_path)
+    structures.append(structure_path)
+
+    candidate_msa = structure_path.with_suffix(".a3m")
+    if candidate_msa.exists():
+        paired_msas.append((structure_path, candidate_msa))
+
+
+def _benchmark_renderable(title: str, payload: dict[str, Any], *, next_cmd: str | None = None) -> Group:
+    submission = payload.get("submission") or payload.get("request") or {}
+    summary = payload.get("summary") or {}
+    table = Table(box=box.SIMPLE_HEAD, show_header=True, padding=(0, 1), expand=False)
+    table.add_column("FIELD", style="bold", no_wrap=True)
+    table.add_column("VALUE", overflow="fold")
+    if payload.get("job_id"):
+        table.add_row("JOB ID", str(payload["job_id"]))
+    table.add_row("STATUS", str(payload.get("status", "")))
+    if isinstance(submission, dict):
+        if submission.get("name"):
+            table.add_row("NAME", str(submission["name"]))
+        if submission.get("model"):
+            table.add_row("MODEL", str(submission["model"]))
+        if submission.get("weight"):
+            table.add_row("WEIGHT", str(submission["weight"]))
+    if isinstance(summary, dict):
+        if summary.get("queued_structures") is not None:
+            table.add_row("STRUCTURES", str(summary["queued_structures"]))
+        if summary.get("paired_msa_count") is not None:
+            table.add_row("PAIRED MSAs", str(summary["paired_msa_count"]))
+    if payload.get("result_path"):
+        table.add_row("RESULT", str(payload["result_path"]))
+    panels: list = [Panel(table, title=title, border_style="dim")]
+    if next_cmd:
+        next_table = Table(box=box.SIMPLE_HEAD, show_header=False, padding=(0, 1), expand=False)
+        next_table.add_column("command", style="cyan", overflow="fold")
+        next_table.add_row(Text(next_cmd, style="cyan"))
+        panels.append(Panel(next_table, title="Next", border_style="dim"))
+    return Group(*panels)