diff --git a/tutorials-and-examples/nvidia-nim/blueprints/README.md b/tutorials-and-examples/nvidia-nim/blueprints/README.md new file mode 100644 index 000000000..1c7f51c47 --- /dev/null +++ b/tutorials-and-examples/nvidia-nim/blueprints/README.md @@ -0,0 +1,8 @@ +### NVIDIA NIM Blueprints on GKE + +Here you will find the NVIDIA NIM Blueprints that can be provisioned to run on GKE. These are good for proof of concepts only. + +- [ ] [Generative Virtual Screening for Drug Discovery](./drugdiscovery/README.md) uses 3 NIMs + - AlphaFold2 + - MolMIM + - DiffDock diff --git a/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery/README.md b/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery/README.md new file mode 100644 index 000000000..64442ed11 --- /dev/null +++ b/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery/README.md @@ -0,0 +1,177 @@ +# Generative Virtual Screening for Drug Discovery on GKE + +This guide outlines the steps to deploy NVIDIA's NIM blueprint for [Generative Virtual screening for Drug Discovery](https://build.nvidia.com/nvidia/generative-virtual-screening-for-drug-discovery) on a Google Kubernetes Engine (GKE) cluster. Three NIMs - AlphaFold2, MolMIM & DiffDock are used to demonstrate Protein folding, molecular generation and protein docking. + +## Prerequisites + +* **GCloud SDK:** Ensure you have the Google Cloud SDK installed and configured. +* **Project:** A Google Cloud project with billing enabled. +* **NGC API Key:** An API key from NVIDIA NGC. +* **kubectl:** kubectl command-line tool installed and configured. + +Clone the repo before proceeding further: +`git clone https://github.com/GoogleCloudPlatform/ai-on-gke` +`cd ai-on-gke/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery` + +## Deployment Steps + +1. **Set Project and Variables:** + + ```bash + + gcloud config set project "" + export PROJECT_ID=$(gcloud config get project) + export CLUSTER_NAME="gke-nimbp-genscr" + export NODE_POOL_NAME="gke-nimbp-genscr-np" + export ZONE="" #us-east5-b + export MACHINE_TYPE= "" #"a2-ultragpu-1g" + export ACCELERATOR_TYPE="" #"nvidia-a100-80gb" + export ACCELERATOR_COUNT="1" + export NODE_POOL_NODES=3 + export NGC_API_KEY="" + + ``` + +2. **Create GKE Cluster:** This creates the initial cluster with a default node pool for management tasks. GPU nodes will be added in the next step. + + ```bash + + gcloud container clusters create "${CLUSTER_NAME}" \ + --project="${PROJECT_ID}" \ + --num-nodes=1 --location="${ZONE}" \ + --machine-type=e2-standard-16 \ + --addons=GcpFilestoreCsiDriver + + ``` + +3. **Create GPU Node Pool:** This creates a node pool with GPU machines optimized for BioNeMo workloads. + + ```bash + + gcloud container node-pools create "${NODE_POOL_NAME}" \ + --cluster="${CLUSTER_NAME}" \ + --location="${ZONE}" \ + --node-locations="${ZONE}" \ + --num-nodes="${NODE_POOL_NODES}" \ + --machine-type="${MACHINE_TYPE}" \ + --accelerator="type=${ACCELERATOR_TYPE},count=${ACCELERATOR_COUNT},gpu-driver-version=LATEST" \ + --placement-type="COMPACT" \ + --disk-type="pd-ssd" \ + --disk-size="500GB" + + ``` + +4. **Get Cluster Credentials:** + + ```bash + + gcloud container clusters get-credentials "${CLUSTER_NAME}" --location="${ZONE}" + + ``` + +5. **Set kubectl Alias (Optional):** + + ```bash + + alias k=kubectl + + ``` + +6. **Create NGC Secret:** Creates the secret for pulling images from NVIDIA NGC. **Important:** Replace the placeholder API key with your actual NGC API Key. + + ```bash + + k create secret docker-registry secret-nvcr \ + --docker-username=\$oauthtoken \ + --docker-password="${NGC_API_KEY}" \ + --docker-server="nvcr.io" + + ``` + +7. **Deploy BioNeMo Services:** Deploy AlphaFold2, MolMIM, and DiffDock. + + ```bash + + k create -f nim-storage-filestore.yaml + k create -f nim-bionemo-generative-virtual-screening.yaml + + ``` + +8. **Port Forwarding (for local testing):** These commands forward the service ports to your local machine for testing. Replace the pod names with the actual names of your deployed pods. + + ```bash + + POD_BIONEMO_ALPHAFOLD=$(k get pods -o go-template --template '{{range .items}}{{.metadata.name}}{{"\n"}}{{end}}' | grep '^alphafold2') + k port-forward pod/$POD_BIONEMO_ALPHAFOLD :8000 + + POD_BIONEMO_MOLMIM=$(k get pods -o go-template --template '{{range .items}}{{.metadata.name}}{{"\n"}}{{end}}' | grep '^molmim') + k port-forward pod/$POD_BIONEMO_MOLMIM :8000 + + POD_BIONEMO_DIFFDOCK=$(k get pods -o go-template --template '{{range .items}}{{.metadata.name}}{{"\n"}}{{end}}' | grep '^diffdock') + k port-forward pod/$POD_BIONEMO_DIFFDOCK :8000 + + ``` + +9. **Test Deployments:** Use `curl` or other tools to test the deployed services by sending requests to the forwarded ports. Examples are provided in the `dev.sh` file. + + ```bash + + # AlphaFold2 + curl \ + --max-time 900 \ + -X POST \ + -i \ + "http://localhost:62921/protein-structure/alphafold2/predict-structure-from-sequence" \ + -H 'accept: application/json' \ + -H 'Content-Type: application/json' \ + -d '{ + "sequence": "MVPSAGQLALFALGIVLAACQALENSTSPLSADPPVAAAVVSHFNDCPDSHTQFCFHGTCRFLVQEDKPACVCHSGYVGARCEHADLLAVVAASQKKQAITALVVVSIVALAVLIITCVLIHCCQVRKHCEWCRALICRHEKPSALLKGRTACCHSETVV", + "databases": [ + "small_bfd" + ] + }' + + ``` + + ```bash + + # MolMIM + curl -X POST \ + -H 'Content-Type: application/json' \ + -d '{ + "smi": "CC1(C2C1C(N(C2)C(=O)C(C(C)(C)C)NC(=O)C(F)(F)F)C(=O)NC(CC3CCNC3=O)C#N)C", + "num_molecules": 5, + "algorithm": "CMA-ES", + "property_name": "QED", + "min_similarity": 0.7, + "iterations": 10 + }' \ + "http://localhost:49530/generate" + + ``` + +10. **Test end to end:** + Update the AF2_HOST, MOLMIM_HOST and DIFFDOCK_HOST variables with port numbers in `test-generative-virtual-screening.py` file. + + ```bash + + python3 -m venv venv + source venv/bin/activate + pip3 install requests + python3 test-generative-virtual-screening.py + deactivate + + ``` + +## Cleanup + + To delete the cluster and all associated resources: + + ```bash + + k delete secret secret-nvcr + k delete -f nim-bionemo-generative-virtual-screening.yaml + k delete -f nim-storage-filestore.yaml + gcloud container clusters delete "${CLUSTER_NAME}" --location="${ZONE}" + + ``` diff --git a/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery/nim-bionemo-generative-virtual-screening.yaml b/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery/nim-bionemo-generative-virtual-screening.yaml new file mode 100644 index 000000000..77264257e --- /dev/null +++ b/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery/nim-bionemo-generative-virtual-screening.yaml @@ -0,0 +1,189 @@ +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: alphafold2 +spec: + replicas: 1 + selector: + matchLabels: + app: alphafold2 + template: + metadata: + labels: + app: alphafold2 + spec: + containers: + - name: alphafold2 + image: nvcr.io/nim/deepmind/alphafold2:2.1.0 + securityContext: + runAsUser: 0 + ports: + - containerPort: 8000 + volumeMounts: + - name: alphafold2-cache + mountPath: /data + env: + - name: NGC_API_KEY + valueFrom: + secretKeyRef: + name: ngc-api-key # Replace with your actual secret name + key: NGC_API_KEY + - name: NIM_CACHE_PATH + value: /data + - name: NIM_MODEL_ORG + value: nim + - name: NIM_MODEL_TEAM + value: deepmind + resources: + requests: + nvidia.com/gpu: 1 + limits: + nvidia.com/gpu: 1 + volumes: + - name: alphafold2-cache + persistentVolumeClaim: + claimName: nim-storage-filestore + imagePullSecrets: + - name: secret-nvcr + +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: molmim +spec: + replicas: 1 + selector: + matchLabels: + app: molmim + template: + metadata: + labels: + app: molmim + spec: + containers: + - name: molmim + image: nvcr.io/nim/nvidia/molmim:1.0.0 + securityContext: + runAsUser: 0 + ports: + - containerPort: 8000 + volumeMounts: + - name: molmim-cache + mountPath: /data + env: + - name: NGC_API_KEY + valueFrom: + secretKeyRef: + name: ngc-api-key # Replace with your actual secret name + key: NGC_API_KEY + - name: NIM_CACHE_PATH + value: /data + resources: + requests: + nvidia.com/gpu: 1 + limits: + nvidia.com/gpu: 1 + volumes: + - name: molmim-cache + persistentVolumeClaim: + claimName: nim-storage-filestore + imagePullSecrets: + - name: secret-nvcr + +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: diffdock +spec: + replicas: 1 + selector: + matchLabels: + app: diffdock + template: + metadata: + labels: + app: diffdock + spec: + containers: + - name: diffdock + image: nvcr.io/nim/mit/diffdock:2.0 + securityContext: + runAsUser: 0 + ports: + - containerPort: 8000 + volumeMounts: + - name: diffdock-cache + mountPath: /data + env: + - name: NGC_API_KEY + valueFrom: + secretKeyRef: + name: ngc-api-key # Replace with your actual secret name + key: NGC_API_KEY + - name: NVIDIA_VISIBLE_DEVICES + value: "0" # Set the default value here + - name: NIM_CACHE_PATH + value: /data + resources: + requests: + nvidia.com/gpu: 1 + limits: + nvidia.com/gpu: 1 + imagePullSecrets: + - name: secret-nvcr + volumes: + - name: diffdock-cache + persistentVolumeClaim: + claimName: nim-storage-filestore + +--- +apiVersion: v1 +kind: Service +metadata: + name: alphafold2-service +spec: + selector: + app: alphafold2 + ports: + - protocol: TCP + port: 8081 + targetPort: 8000 + +--- +apiVersion: v1 +kind: Service +metadata: + name: diffdock-service +spec: + selector: + app: diffdock + ports: + - protocol: TCP + port: 8082 + targetPort: 8000 + +--- +apiVersion: v1 +kind: Service +metadata: + name: molmim-service +spec: + selector: + app: molmim + ports: + - protocol: TCP + port: 8083 + targetPort: 8000 + +--- +# Create a Secret to hold the NGC API key +apiVersion: v1 +kind: Secret +metadata: + name: ngc-api-key +type: Opaque +stringData: + NGC_API_KEY: \ No newline at end of file diff --git a/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery/nim-storage-filestore.yaml b/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery/nim-storage-filestore.yaml new file mode 100644 index 000000000..6e0434b1d --- /dev/null +++ b/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery/nim-storage-filestore.yaml @@ -0,0 +1,24 @@ +apiVersion: storage.k8s.io/v1 +kind: StorageClass +metadata: + name: nim-storage-filestore +provisioner: filestore.csi.storage.gke.io +volumeBindingMode: Immediate +allowVolumeExpansion: true +parameters: + tier: BASIC_SSD + network: default +--- +kind: PersistentVolumeClaim +apiVersion: v1 +metadata: + name: nim-storage-filestore + namespace: default +spec: + accessModes: + - ReadWriteMany + storageClassName: nim-storage-filestore + resources: + requests: + storage: 2.5Ti + diff --git a/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery/test-generative-virtual-screening.py b/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery/test-generative-virtual-screening.py new file mode 100644 index 000000000..ecd076d9a --- /dev/null +++ b/tutorials-and-examples/nvidia-nim/blueprints/drugdiscovery/test-generative-virtual-screening.py @@ -0,0 +1,110 @@ +import requests + +AF2_HOST = "http://localhost:" +MOLMIM_HOST = "http://localhost:" +DIFFDOCK_HOST = "http://localhost:" + + +def check_readiness(service_name, url): + try: + response = requests.get(url) + response.raise_for_status() # Raise an exception for bad status codes + print(f"{service_name} readiness check successful at {url}.") + return True + except requests.exceptions.RequestException as e: + print(f"{service_name} readiness check failed: {e}") + return False + + +check_readiness("AlphaFold2", f"{AF2_HOST}/v1/health/ready") +check_readiness("MolMIM", f"{MOLMIM_HOST}/v1/health/ready") +check_readiness("DiffDock", f"{DIFFDOCK_HOST}/v1/health/ready") + + +## AlphaFold2 NIM +protein_sequence = "MVPSAGQLALFALGIVLAACQALENSTSPLSADPPVAAAVVSHFNDCPDSHTQFCFHGTCRFLVQEDKPACVCHSGYVGARCEHADLLAVVAASQKKQAITALVVVSIVALAVLIITCVLIHCCQVRKHCEWCRALICRHEKPSALLKGRTACCHSETVV" + +print("Protein folding with AlphaFold2...") # Progress update +try: + af2_response = requests.post( + f"{AF2_HOST}/protein-structure/alphafold2/predict-structure-from-sequence", + json={ + "sequence": protein_sequence, + "databases": ["small_bfd"], + }, + ) + af2_response.raise_for_status() # Check for HTTP errors + folded_protein = af2_response.json()[0] + print("Protein folding successfully.") # Progress update + +except requests.exceptions.RequestException as e: + print(f"AlphaFold2 request failed: {e}") + exit(1) # Exit if AlphaFold2 fails +except IndexError: + print("Invalid response from AlphaFold2 (missing structure).") + exit(1) + +## MolMIM NIM +molecule = "CC1(C2C1C(N(C2)C(=O)C(C(C)(C)C)NC(=O)C(F)(F)F)C(=O)NC(CC3CCNC3=O)C#N)C" + +print("Generating molecules with MolMIM...") # Progress update +try: + molmim_response = requests.post( + f"{MOLMIM_HOST}/generate", + json={ + "smi": molecule, + "num_molecules": 5, + "algorithm": "CMA-ES", + "property_name": "QED", + "min_similarity": 0.7, # Ignored if algorithm is not "CMA-ES". + "iterations": 10, + }, + ) + molmim_response.raise_for_status() # Check for HTTP errors + molmim_response = molmim_response.json() + + # Validate the response (check if 'generated' key exists and is a list) + if "generated" not in molmim_response or not isinstance( + molmim_response["generated"], list + ): + raise ValueError("Invalid response format from MolMIM") + + generated_ligands = "\n".join([v["smiles"] for v in molmim_response["generated"]]) + print(generated_ligands) + print("Molecular generation successfully.") # Progress update + +except requests.exceptions.RequestException as e: + print(f"MolMIM request failed: {e}") + exit(1) # Exit if MolMIM fails +except ValueError as e: + print(f"Error processing MolMIM response: {e}") + exit(1) + +## DiffDock NIM +diffdock_response = requests.post( + f"{DIFFDOCK_HOST}/molecular-docking/diffdock/generate", + json={ + "protein": folded_protein, + "ligand": generated_ligands, + "ligand_file_type": "txt", + "num_poses": 10, + "time_divisions": 20, + "num_steps": 18, + }, +).json() + +print("Protein-ligand docking with DiffDock...") # Progress update + +# Example assuming diffdock_response["ligand_positions"] is a list of lists. +try: + for i in range(len(diffdock_response["ligand_positions"])): + print("*" * 80) + print(f"Docking result {i+1}:") # More descriptive output + print( + diffdock_response["ligand_positions"][i][0] + ) # Print the first pose for now. Handle multiple poses if needed. + print("Protein-ligand docking successful") # Progress update + +except requests.exceptions.RequestException as e: + print(f"Error processing DiffDock results: Invalid response format: {e}") + exit(1) # Exit if MolMIM fails