|
|
"""CLI experiment command and sub-commands.""" |
|
|
|
|
|
import json |
|
|
import os |
|
|
import shutil |
|
|
import zipfile |
|
|
from pathlib import Path |
|
|
from typing import List, Optional |
|
|
|
|
|
import requests |
|
|
import typer |
|
|
from rich import print |
|
|
from rich.console import Console |
|
|
from rich.table import Table |
|
|
from typing_extensions import Annotated |
|
|
|
|
|
from folding_studio.config import API_URL, REQUEST_TIMEOUT |
|
|
from folding_studio.utils.headers import get_auth_headers |
|
|
|
|
|
app = typer.Typer( |
|
|
no_args_is_help=True, |
|
|
help=( |
|
|
"Get experiment information and metadata, like its status, " |
|
|
"results or the generated features (msa, templates, etc.).\n" |
|
|
"Read more at https://int-bio-foldingstudio-gcp.nw.r.appspot.com/how-to-guides/af2_openfold/fetch_folding_job_status/." |
|
|
), |
|
|
) |
|
|
|
|
|
experiment_ID_argument = typer.Argument(help="ID of the experiment.") |
|
|
|
|
|
|
|
|
def _download_file_from_signed_url( |
|
|
exp_id: str, |
|
|
endpoint: str, |
|
|
output: Path, |
|
|
force: bool, |
|
|
unzip: bool = False, |
|
|
api_key: str | None = None, |
|
|
) -> None: |
|
|
"""Download a zip file from an experiment id. |
|
|
|
|
|
Args: |
|
|
exp_id (str): Experiment id. |
|
|
endpoint (str): API endpoint to call. |
|
|
output (Path): Output file path. |
|
|
force (bool): Force file writing if it already exists. |
|
|
unzip (bool): Unzip the zip file after downloading. |
|
|
|
|
|
Raises: |
|
|
typer.Exit: If output file path exists but force set to false. |
|
|
typer.Exit: If unzip set to true but the directory already exists and force set to false. |
|
|
typer.Exit: If an error occurred during the initial API call. |
|
|
""" |
|
|
if output.exists() and not force: |
|
|
print( |
|
|
f"Warning: The file '{output}' already exists. Use the --force flag to overwrite it." |
|
|
) |
|
|
raise typer.Exit(code=1) |
|
|
|
|
|
if unzip: |
|
|
if not output.suffix == ".zip": |
|
|
print( |
|
|
"Error: The downloaded file is not a .zip file. Please ensure the correct file format." |
|
|
) |
|
|
raise typer.Exit(code=1) |
|
|
|
|
|
dir_path = output.with_suffix("") |
|
|
if dir_path.exists() and not force: |
|
|
print( |
|
|
f"Warning: The --unzip flag is raised but the directory '{dir_path}' " |
|
|
"already exists. Use the --force flag to overwrite it." |
|
|
) |
|
|
raise typer.Exit(code=1) |
|
|
|
|
|
headers = get_auth_headers(api_key) |
|
|
url = API_URL + endpoint |
|
|
|
|
|
response = requests.get( |
|
|
url, |
|
|
params={"experiment_id": exp_id}, |
|
|
headers=headers, |
|
|
timeout=REQUEST_TIMEOUT, |
|
|
) |
|
|
if not response.ok: |
|
|
print(f"Failed to download the file: {response.content.decode()}.") |
|
|
raise typer.Exit(code=1) |
|
|
|
|
|
file_response = requests.get( |
|
|
response.json()["signed_url"], |
|
|
stream=True, |
|
|
timeout=REQUEST_TIMEOUT, |
|
|
) |
|
|
with output.open("wb") as f: |
|
|
file_response.raw.decode_content = True |
|
|
shutil.copyfileobj(file_response.raw, f) |
|
|
print(f"File downloaded successfully to {output}.") |
|
|
|
|
|
if unzip: |
|
|
dir_path.mkdir(parents=True, exist_ok=True) |
|
|
with zipfile.ZipFile(output, "r") as zip_ref: |
|
|
zip_ref.extractall(dir_path) |
|
|
print(f"Extracted all files to {dir_path}.") |
|
|
|
|
|
|
|
|
@app.command() |
|
|
def status( |
|
|
exp_id: Annotated[str, experiment_ID_argument], |
|
|
api_key: Annotated[str, typer.Option("--api-key", "-k")], |
|
|
): |
|
|
"""Get an experiment status.""" |
|
|
headers = get_auth_headers(api_key) |
|
|
url = API_URL + "getExperimentStatus" |
|
|
response = requests.get( |
|
|
url, |
|
|
params={"experiment_id": exp_id}, |
|
|
headers=headers, |
|
|
timeout=REQUEST_TIMEOUT, |
|
|
) |
|
|
|
|
|
if not response.ok: |
|
|
print(f"An error occurred : {response.content.decode()}") |
|
|
raise typer.Exit(code=1) |
|
|
|
|
|
message = response.json() |
|
|
print(message["status"]) |
|
|
|
|
|
|
|
|
@app.command() |
|
|
def list( |
|
|
limit: Annotated[ |
|
|
int, |
|
|
typer.Option( |
|
|
help=("Max number of experiment to display in the terminal."), |
|
|
), |
|
|
] = 100, |
|
|
output: Annotated[ |
|
|
Optional[Path], |
|
|
typer.Option( |
|
|
"--output", |
|
|
"-o", |
|
|
help=( |
|
|
"Path to the file where the job metadata returned by the server are written." |
|
|
), |
|
|
), |
|
|
] = None, |
|
|
): |
|
|
"""Get all your done and pending experiment ids. The IDs are provided in the order of submission, starting with the most recent.""" |
|
|
headers = get_auth_headers() |
|
|
url = API_URL + "getDoneAndPendingExperiments" |
|
|
response = requests.get( |
|
|
url, |
|
|
headers=headers, |
|
|
timeout=REQUEST_TIMEOUT, |
|
|
) |
|
|
|
|
|
if not response.ok: |
|
|
print(f"An error occurred : {response.content.decode()}") |
|
|
raise typer.Exit(code=1) |
|
|
|
|
|
response_json = response.json() |
|
|
if output: |
|
|
with open(output, "w") as f: |
|
|
json.dump(response_json, f, indent=4) |
|
|
|
|
|
print(f"Done and pending experiments list written to [bold]{output}[/bold]") |
|
|
|
|
|
table = Table(title="Done and pending experiments") |
|
|
|
|
|
table.add_column("Experiment ID", justify="right", style="cyan", no_wrap=True) |
|
|
table.add_column("Status", style="magenta") |
|
|
|
|
|
total_exp_nb = 0 |
|
|
for status, exp_list in response_json.items(): |
|
|
total_exp_nb += len(exp_list) |
|
|
for exp in exp_list: |
|
|
table.add_row(exp, status) |
|
|
if limit < total_exp_nb: |
|
|
print( |
|
|
f"The table below is truncated to the last [bold]{limit}[/bold] submitted experiments. Increase '--limit' to see more." |
|
|
) |
|
|
if not output: |
|
|
print("Use '--output' to get the full list in file format.") |
|
|
else: |
|
|
print(f"See the full list in file format at [bold]{output}[/bold]") |
|
|
break |
|
|
|
|
|
console = Console() |
|
|
console.print(table) |
|
|
|
|
|
|
|
|
@app.command() |
|
|
def features( |
|
|
exp_id: Annotated[str, experiment_ID_argument], |
|
|
output: Annotated[ |
|
|
Optional[Path], |
|
|
typer.Option( |
|
|
help="Local path to download the zip to. Default to '<exp_id>_features.zip'." |
|
|
), |
|
|
] = None, |
|
|
force: Annotated[ |
|
|
bool, |
|
|
typer.Option( |
|
|
help=( |
|
|
"Forces the download to overwrite any existing file " |
|
|
"with the same name in the specified location." |
|
|
) |
|
|
), |
|
|
] = False, |
|
|
unzip: Annotated[ |
|
|
bool, typer.Option(help="Automatically unzip the file after its download.") |
|
|
] = False, |
|
|
): |
|
|
"""Get an experiment features.""" |
|
|
if output is None: |
|
|
output = Path(f"{exp_id}_features.zip") |
|
|
|
|
|
_download_file_from_signed_url( |
|
|
exp_id=exp_id, |
|
|
endpoint="getZippedExperimentFeatures", |
|
|
output=output, |
|
|
force=force, |
|
|
unzip=unzip, |
|
|
) |
|
|
|
|
|
|
|
|
@app.command() |
|
|
def results( |
|
|
exp_id: Annotated[str, experiment_ID_argument], |
|
|
api_key: Annotated[str, typer.Option("--api-key", "-k")], |
|
|
output: Annotated[ |
|
|
Optional[Path], |
|
|
typer.Option( |
|
|
help="Local path to download the zip to. Default to '<exp_id>_results.zip'." |
|
|
), |
|
|
] = None, |
|
|
force: Annotated[ |
|
|
bool, |
|
|
typer.Option( |
|
|
help=( |
|
|
"Forces the download to overwrite any existing file " |
|
|
"with the same name in the specified location." |
|
|
) |
|
|
), |
|
|
] = False, |
|
|
unzip: Annotated[ |
|
|
bool, typer.Option(help="Automatically unzip the file after its download.") |
|
|
] = False, |
|
|
): |
|
|
"""Get an experiment results.""" |
|
|
|
|
|
if output is None: |
|
|
output = Path(f"{exp_id}_results.zip") |
|
|
|
|
|
_download_file_from_signed_url( |
|
|
exp_id=exp_id, |
|
|
endpoint="getZippedExperimentResults", |
|
|
output=output, |
|
|
force=force, |
|
|
unzip=unzip, |
|
|
api_key=api_key, |
|
|
) |
|
|
|
|
|
|
|
|
@app.command() |
|
|
def cancel( |
|
|
exp_id: Annotated[List[str], experiment_ID_argument], |
|
|
): |
|
|
"""Cancel experiments job executions. |
|
|
|
|
|
You can pass one or more experiment id |
|
|
""" |
|
|
headers = get_auth_headers() |
|
|
|
|
|
url = API_URL + "cancelJob" |
|
|
response = requests.post( |
|
|
url, |
|
|
data={"experiment_ids": exp_id}, |
|
|
headers=headers, |
|
|
timeout=REQUEST_TIMEOUT, |
|
|
) |
|
|
|
|
|
if not response.ok: |
|
|
print(f"An error occurred : {response.content.decode()}") |
|
|
raise typer.Exit(code=1) |
|
|
message = response.json() |
|
|
print(message) |
|
|
|
|
|
|
|
|
@app.command() |
|
|
def logs( |
|
|
exp_id: Annotated[str, experiment_ID_argument], |
|
|
output: Annotated[ |
|
|
Optional[Path], |
|
|
typer.Option( |
|
|
help="Local path to download the logs to. Default to '<exp_id>_logs.txt'." |
|
|
), |
|
|
] = None, |
|
|
force: Annotated[ |
|
|
bool, |
|
|
typer.Option( |
|
|
help=( |
|
|
"Forces the download to overwrite any existing file " |
|
|
"with the same name in the specified location." |
|
|
) |
|
|
), |
|
|
] = False, |
|
|
): |
|
|
"""Get an experiment logs.""" |
|
|
if output is None: |
|
|
output = Path(f"{exp_id}_logs.txt") |
|
|
|
|
|
_download_file_from_signed_url( |
|
|
exp_id=exp_id, |
|
|
endpoint="getExperimentLogs", |
|
|
output=output, |
|
|
force=force, |
|
|
) |
|
|
|
