1-submit_colabfold_run.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Copyright 2022 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#      http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Quick Start: Colabfold inference pipeline with Cloud Batch and Workflows\n",
        "\n",
        "This notebook demonstrates how to submit inference pipeline runs.\n",
        "\n",
        "You use the utility functions in the `workflow_executor` module to configure and submit the runs. The `workflow_executor` module contains two functions:\n",
        "- `prepare_args_for_experiment` - This function formats the runtime parameters for the Google Workflows workflows that implements the pipeline. It also sets default values for a number of runtime parameters\n",
        "- `execute_workflow` - This function executes the Google Workflows workflow.\n",
        "\n",
        "This is a complete list of required and optional parameters accepted by the functions:\n",
        "\n",
        "```\n",
        "    project_id: str\n",
        "    region: str\n",
        "    input_dir: str\n",
        "    image_uri: str\n",
        "    job_gcs_path: str\n",
        "    labels: dict\n",
        "    machine_type: str = 'n1-standard-4'\n",
        "    cpu_milli: int = 8000\n",
        "    memory_mib: int = 30000\n",
        "    boot_disk_mib: int = 200000\n",
        "    gpu_type: str = \"nvidia-tesla-t4\"\n",
        "    gpu_count: int = 1\n",
        "    job_gcsfuse_local_dir: str = '/mnt/disks/gcs/colabfold'\n",
        "    parallelism: int = 8\n",
        "    template_mode: str = \"none\"\n",
        "    use_cpu: bool = False\n",
        "    use_gpu_relax: bool = False\n",
        "    use_amber: bool = False\n",
        "    msa_mode: str = 'mmseqs2_uniref_env'\n",
        "    model_type: str = 'auto'\n",
        "    num_models: int = 5\n",
        "    num_recycle: int = 3\n",
        "    custom_template_path: str = None\n",
        "    overwrite_existing_results: bool = False\n",
        "    rank_by: str = 'auto'\n",
        "    pair_mode: str = 'unpaired_paired'\n",
        "    stop_at_score: int = 100\n",
        "    zip_results: bool = False\n",
        "```"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Install python libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Install packages\n",
        "! pip install -U google-cloud-firestore google-cloud-workflows google-cloud-storage"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Reload the kernel before proceeding\n",
        "%load_ext autoreload\n",
        "%autoreload 2"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Execute Workflow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from src import workflow_executor"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Please set the following variables according to the setup of your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "project_id = '<YOUR PROJECT ID>'    # Project ID. Example: \"my_project_id\"\n",
        "region = '<YOUR REGION>'    # Region where resources will be created. Example: \"us-central1\"\n",
        "\n",
        "input_dir = '<YOUR BUCKET NAME>/<YOUR FOLDER>'   # GCS path where you will upload FASTA files.\n",
        "                                                 # Example: 'my_bucket/input_folder'\n",
        "image_uri = 'gcr.io/<YOUR PROJECT NAME>/colabfold-batch'    # Image built to execute Colabfold\n",
        "job_gcs_path = '<YOUR BUCKET NAME>'     # Bucket name where the resulting artifacts will be created.\n",
        "                                        # Example: 'my_bucket'\n",
        "\n",
        "labels = {}     # Labels to identify your job"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Copy local FASTA files to the GCS path."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "local_input_dir = '/path/to/my/files'   # Local directory where your FASTA files are located"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Copy local files to GCS\n",
        "! gsutil -m cp {local_input_dir}/*.fasta gs://{input_dir}"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Execute the following cell to start the Colabfold execution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Prepare the environment for execution\n",
        "args = workflow_executor.prepare_args_for_experiment(\n",
        "    project_id = project_id,\n",
        "    region = region,\n",
        "    input_dir = input_dir,\n",
        "    image_uri = image_uri,\n",
        "    job_gcs_path = job_gcs_path,\n",
        "    labels = labels\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Execute the workflow\n",
        "workflow_executor.execute_workflow(\n",
        "    workflow_name='colabfold-workflow',\n",
        "    args=args\n",
        ")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "include_colab_link": true,
      "name": "AlphaFold2_batch.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "base",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.12"
    },
    "vscode": {
      "interpreter": {
        "hash": "d4d1e4263499bec80672ea0156c357c1ee493ec2b1c70f0acce89fc37c4a6abe"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}