Have gradio app call inference.py

Dockerfile

@@ -3,7 +3,7 @@ FROM silterra/diffdock-pocket-dev
 USER $APPUSER
 WORKDIR $HOME/app
 
-COPY --chown=$APPUSER . $HOME/app
+COPY --chown=$APPUSER: . $HOME/app
 
 # Expose port for web service
 ENV PORT=7860

main.py

@@ -1,45 +1,52 @@
-import gradio as gr
-import torch
-
-if False:
-    import requests
-    from torchvision import transforms
-    model = torch.hub.load("pytorch/vision:v0.6.0", "resnet18", pretrained=True).eval()
-    response = requests.get("https://git.io/JJkYN")
-    labels = response.text.split("\n")
+import logging
+import os.path
 
+import gradio as gr
 
-def predict(inp, *args, **kwargs):
-    inp = transforms.ToTensor()(inp).unsqueeze(0)
-    with torch.no_grad():
-        prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
-        confidences = {labels[i]: float(prediction[i]) for i in range(1000)}
-    return confidences
+import run_utils
 
 
-def calculate(*args, **kwargs) -> str:
-    output_file_path = "main_output.txt"
-    with open(output_file_path, "w") as fi:
-        fi.write(f"args: {args}\n")
-        fi.write(f"kwargs: {kwargs}\n")
-    return output_file_path
+def run_wrapper(protein_file, ligand_file, *args, **kwargs) -> str:
+    return run_utils.run_cli_command(protein_file.name, ligand_file.name,
+                                     *args, **kwargs)
 
 
 def run():
     iface = gr.Interface(
-        fn=calculate,
+        fn=run_wrapper,
         inputs=[
             gr.File(label="Protein PDB", file_types=[".pdb"]),
             gr.File(label="Ligand SDF", file_types=[".sdf"]),
-            gr.Number(label="Samples Per Complex", value=4, minimum=1, maximum=100, precision=0),
+            gr.Number(
+                label="Samples Per Complex",
+                value=1,
+                minimum=1,
+                maximum=100,
+                precision=0,
+            ),
             gr.Checkbox(label="Keep Local Structures", value=True),
-            gr.Checkbox(label="Save Visualization", value=True)
+            gr.Checkbox(label="Save Visualisation", value=True),
         ],
-        outputs=gr.File(label="Result")
+        outputs=gr.File(label="Result"),
     )
 
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    iface.launch(server_name="0.0.0.0", server_port=7860, share=False)
+
+
+def set_env_variables():
+    if "DiffDock-Pocket-Dir" not in os.environ:
+        work_dir = os.path.abspath(os.path.join("../DiffDock-Pocket"))
+        if os.path.exists(work_dir):
+            os.environ["DiffDock-Pocket-Dir"] = work_dir
+        else:
+            raise ValueError(f"DiffDock-Pocket-Dir {work_dir} not found")
+
+    if "LOG_LEVEL" not in os.environ:
+        os.environ["LOG_LEVEL"] = "INFO"
 
 
 if __name__ == "__main__":
+    set_env_variables()
+    run_utils.configure_logging()
+
     run()

run_utils.py

@@ -0,0 +1,85 @@
+import datetime
+import os
+import shutil
+import subprocess
+import tempfile
+import uuid
+
+import logging
+
+
+def configure_logging(level=None):
+    if level is None:
+        level = getattr(logging, os.environ.get("LOG_LEVEL", "INFO"))
+
+    # Note that this sets the universal logger,
+    # which includes other libraries.
+    logging.basicConfig(
+        level=level,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        handlers=[
+            logging.StreamHandler(),  # Outputs logs to stderr by default
+            # If you also want to log to a file, uncomment the following line:
+            # logging.FileHandler('my_app.log', mode='a', encoding='utf-8')
+        ]
+    )
+
+
+def run_cli_command(protein_path: str, ligand: str, samples_per_complex: int,
+                    keep_local_structures: bool, save_visualisation: bool, work_dir=None):
+
+    if work_dir is None:
+        work_dir = os.environ.get("DiffDock-Pocket-Dir",
+                                  os.path.join(os.environ["HOME"], "DiffDock-Pocket"))
+
+    command = ["python3", "inference.py", f"--protein_path={protein_path}", f"--ligand={ligand}",
+               f"--samples_per_complex={samples_per_complex}"]
+
+    # Adding boolean arguments only if they are True
+    if keep_local_structures:
+        command.append("--keep_local_structures")
+    if save_visualisation:
+        command.append("--save_visualisation")
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_dir_path = temp_dir
+        logging.debug(f"temp dir: {temp_dir}")
+        command.append(f"--out_dir={temp_dir_path}")
+
+        # Convert command list to string for printing
+        command_str = " ".join(command)
+        logging.info(f"Executing command: {command_str}")
+
+        # Running the command
+        try:
+            result = subprocess.run(
+                command, cwd=work_dir, check=False, text=True, capture_output=True, env=os.environ
+            )
+            logging.debug(f"Command output:\n{result.stdout}")
+            if result.stderr:
+                logging.error(f"Command error:\n{result.stderr}")
+        except subprocess.CalledProcessError as e:
+            logging.error(f"An error occurred while executing the command: {e}")
+
+        # Zip the output directory
+        # Generate a unique filename using a timestamp and a UUID
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        uuid_tag = str(uuid.uuid4())[0:8]
+        unique_filename = f"output_{timestamp}_{uuid_tag}"
+        zip_base_name = os.path.join("tmp", unique_filename)
+        full_zip_path = shutil.make_archive(zip_base_name, 'zip', temp_dir)
+
+        logging.debug(f"Directory '{temp_dir}' zipped to '{full_zip_path}'")
+
+    return full_zip_path
+
+
+if False and __name__ == "__main__":
+    # Testing code
+    work_dir = os.path.expanduser("~/Projects/DiffDock-Pocket")
+    os.environ["DiffDock-Pocket-Dir"] = work_dir
+    protein_path = os.path.join(work_dir, "example_data", "3dpf_protein.pdb")
+    ligand = os.path.join(work_dir, "example_data", "3dpf_ligand.sdf")
+
+    run_cli_command(protein_path, ligand, samples_per_complex=1,
+                    keep_local_structures=True, save_visualisation=True)