run_utils.py

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)