Spaces:
Runtime error
Runtime error
File size: 3,731 Bytes
6daa4cc 28fb149 6daa4cc 28fb149 6daa4cc 28fb149 6daa4cc 28fb149 6daa4cc 28fb149 6daa4cc 28fb149 6daa4cc 28fb149 6daa4cc 28fb149 6daa4cc 28fb149 6daa4cc 28fb149 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
import datetime
import os
import shutil
import subprocess
import tempfile
import uuid
import logging
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"
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
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 __name__ == "__main__":
# Testing code
set_env_variables()
configure_logging()
work_dir = os.path.abspath("../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,
)
|