Skills › Scientific Agent Skills › diffdock › setup_check.py
setup_check.py
276 lines · skills/diffdock/scripts/setup_check.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3DiffDock Environment Setup Checker
5This script verifies that the DiffDock environment is properly configured
6and all dependencies are available.
8Usage:
9 python /path/to/skill/scripts/setup_check.py
10 python /path/to/skill/scripts/setup_check.py --verbose
11"""
13import argparse
14import sys
15import os
16from pathlib import Path
19def check_python_version():
20 """Check Python version."""
21 import sys
22 version = sys.version_info
24 print("Checking Python version...")
25 if version.major == 3 and version.minor == 9:
26 print(f" [OK] Python {version.major}.{version.minor}.{version.micro}")
27 return True
28 else:
29 print(f" [FAIL] Python {version.major}.{version.minor}.{version.micro} "
30 f"(upstream environment.yml targets Python 3.9.18; other minors are not verified)")
31 return False
34def check_package(package_name, import_name=None, version_attr='__version__'):
35 """Check if a Python package is installed."""
36 if import_name is None:
37 import_name = package_name
39 try:
40 module = __import__(import_name)
41 version = module
42 for attr in version_attr.split('.'):
43 version = getattr(version, attr, None)
44 if version is None:
45 version = 'unknown'
46 break
47 print(f" [OK] {package_name:20s} (version: {version})")
48 return True
49 except ImportError:
50 print(f" [FAIL] {package_name:20s} (not installed)")
51 return False
54def check_pytorch():
55 """Check PyTorch installation and CUDA availability."""
56 print("\nChecking PyTorch...")
57 try:
58 import torch
59 print(f" [OK] PyTorch version: {torch.__version__}")
61 # Check CUDA
62 if torch.cuda.is_available():
63 print(f" [OK] CUDA available: {torch.cuda.get_device_name(0)}")
64 print(f" - CUDA version: {torch.version.cuda}")
65 print(f" - Number of GPUs: {torch.cuda.device_count()}")
66 return True, True
67 else:
68 print(f" [WARN] CUDA not available (will run on CPU)")
69 return True, False
70 except ImportError:
71 print(f" [FAIL] PyTorch not installed")
72 return False, False
75def check_pytorch_geometric():
76 """Check PyTorch Geometric installation."""
77 print("\nChecking PyTorch Geometric...")
78 packages = [
79 ('torch-geometric', 'torch_geometric'),
80 ('torch-scatter', 'torch_scatter'),
81 ('torch-sparse', 'torch_sparse'),
82 ('torch-cluster', 'torch_cluster'),
83 ]
85 all_ok = True
86 for pkg_name, import_name in packages:
87 if not check_package(pkg_name, import_name):
88 all_ok = False
90 return all_ok
93def check_core_dependencies():
94 """Check core DiffDock dependencies."""
95 print("\nChecking core dependencies...")
97 dependencies = [
98 ('numpy', 'numpy'),
99 ('scipy', 'scipy'),
100 ('pandas', 'pandas'),
101 ('rdkit', 'rdkit', '__version__'),
102 ('biopython', 'Bio', '__version__'),
103 ('pytorch-lightning', 'pytorch_lightning'),
104 ('PyYAML', 'yaml'),
105 ('ProDy', 'prody'),
106 ('e3nn', 'e3nn'),
107 ]
109 all_ok = True
110 for dep in dependencies:
111 pkg_name = dep[0]
112 import_name = dep[1]
113 version_attr = dep[2] if len(dep) > 2 else '__version__'
115 if not check_package(pkg_name, import_name, version_attr):
116 all_ok = False
118 return all_ok
121def check_esm():
122 """Check ESM (protein language model) installation."""
123 print("\nChecking fair-esm (ESM2 embeddings for all inputs; ESMFold for sequences)...")
124 try:
125 import esm
126 if not hasattr(esm, 'FastaBatchedDataset') or not hasattr(getattr(esm, 'pretrained', None), 'load_model_and_alphabet'):
127 print(" [FAIL] Wrong/incomplete esm module; DiffDock requires fair-esm==2.0.0")
128 return False
129 print(f" [OK] ESM installed (version: {esm.__version__ if hasattr(esm, '__version__') else 'unknown'})")
130 return True
131 except ImportError:
132 print(f" [FAIL] fair-esm not installed (required for PDB embeddings and protein sequence folding)")
133 print(f" Use the upstream environment.yml fair-esm[esmfold]==2.0.0 pin")
134 return False
137def check_diffdock_installation():
138 """Check if DiffDock is properly installed/cloned."""
139 print("\nChecking DiffDock installation...")
141 # Look for key files
142 key_files = [
143 'inference.py',
144 'default_inference_args.yaml',
145 'environment.yml',
146 ]
148 found_files = []
149 missing_files = []
151 for filename in key_files:
152 if Path(filename).is_file():
153 found_files.append(filename)
154 else:
155 missing_files.append(filename)
157 if found_files:
158 print(f" [OK] Found DiffDock files in current directory:")
159 for f in found_files:
160 print(f" - {f}")
161 else:
162 print(f" [WARN] DiffDock files not found in current directory")
163 print(f" Current directory: {os.getcwd()}")
164 print(f" Make sure you're in the DiffDock repository root")
166 # Check for model checkpoints
167 model_dir = Path('./workdir/v1.1/score_model')
168 confidence_dir = Path('./workdir/v1.1/confidence_model')
170 checkpoint_files = [model_dir / 'model_parameters.yml', model_dir / 'best_ema_inference_epoch_model.pt', confidence_dir / 'model_parameters.yml', confidence_dir / 'best_model_epoch75.pt']
171 if all(path.is_file() and path.stat().st_size > 0 for path in checkpoint_files):
172 print(f" [OK] Model checkpoints found")
173 else:
174 print(f" [WARN] Model checkpoints not found in ./workdir/v1.1/")
175 print(f" Models are downloaded on first run only if the score-model directory is absent; repair partial downloads")
177 if missing_files:
178 print(f" [FAIL] Missing checkout files: {', '.join(missing_files)}")
179 return not missing_files
182def print_installation_instructions():
183 """Print installation instructions if setup is incomplete."""
184 print("\n" + "="*80)
185 print("Installation Instructions")
186 print("="*80)
188 print("""
189If DiffDock is not installed, follow these steps:
1911. Clone the repository:
192 git clone --branch v1.1.3 --depth 1 https://github.com/gcorso/DiffDock.git
193 cd DiffDock
1952. Create conda environment:
196 conda env create --file environment.yml
197 conda activate diffdock
1993. Verify installation:
200 python /path/to/skill/scripts/setup_check.py
202For Docker installation:
203 docker pull rbgcsail/diffdock
204 docker run -it --gpus all --entrypoint /bin/bash rbgcsail/diffdock
205 micromamba activate diffdock
207For more information, visit: https://github.com/gcorso/DiffDock
208 """)
211def print_performance_notes(has_cuda):
212 """Print performance notes based on available hardware."""
213 print("\n" + "="*80)
214 print("Performance Notes")
215 print("="*80)
217 if has_cuda:
218 print("[OK] GPU detected; benchmark runtime and memory on your inputs.")
219 else:
220 print("[WARN] No GPU detected; PDB inference may be SIGNIFICANTLY slower on CPU.")
221 print("Sequence folding in v1.1.3 calls CUDA unconditionally; provide a PDB on CPU.")
222 print("First use can build lookup tables and download docking, ESM2 and ESMFold weights.")
223 print("Import checks do not establish checkpoint compatibility or successful docking.")
226def main():
227 parser = argparse.ArgumentParser(
228 description='Check DiffDock environment setup',
229 formatter_class=argparse.RawDescriptionHelpFormatter
230 )
232 parser.add_argument('--verbose', '-v', action='store_true',
233 help='Show detailed version information')
235 args = parser.parse_args()
237 print("="*80)
238 print("DiffDock Environment Setup Checker")
239 print("="*80)
241 checks = []
243 # Run all checks
244 checks.append(("Python version", check_python_version()))
246 pytorch_ok, has_cuda = check_pytorch()
247 checks.append(("PyTorch", pytorch_ok))
249 checks.append(("PyTorch Geometric", check_pytorch_geometric()))
250 checks.append(("Core dependencies", check_core_dependencies()))
251 checks.append(("ESM", check_esm()))
252 checks.append(("DiffDock files", check_diffdock_installation()))
254 # Summary
255 print("\n" + "="*80)
256 print("Summary")
257 print("="*80)
259 all_passed = all(result for _, result in checks)
261 for check_name, result in checks:
262 status = "[OK] PASS" if result else "[FAIL] FAIL"
263 print(f" {status:8s} - {check_name}")
265 if all_passed:
266 print("\n[OK] Dependency imports and checkout checks passed; verify pinned versions and run a small docking smoke test.")
267 print_performance_notes(has_cuda)
268 return 0
269 else:
270 print("\n[FAIL] Some checks failed. Please install missing dependencies.")
271 print_installation_instructions()
272 return 1
275if __name__ == '__main__':
276 sys.exit(main())