fix: paralle el templates searching

PaddlePaddle · Aug 28, 2024 · a832241 · a832241
1 parent db27a25
commit a832241
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 7 deletions.
diff --git a/apps/protein_folding/helixfold3/helixfold/config/helixfold.yaml b/apps/protein_folding/helixfold3/helixfold/config/helixfold.yaml
@@ -27,7 +27,7 @@ mem:
 nproc_msa: 
   hhblits: 16 # Number of processors used by hhblits
   jackhmmer: 6 # Number of processors used by jackhmmer
-  templates: 16 # Number of processors used by templates searching
+  template_searcher: 16 # Number of processors used by templates searching
 
 # File paths
 

diff --git a/apps/protein_folding/helixfold3/helixfold/data/pipeline_parallel.py b/apps/protein_folding/helixfold3/helixfold/data/pipeline_parallel.py
@@ -152,7 +152,7 @@ def __init__(self,
                nprocs: Mapping[str, int] = {
                   'hhblits': 16,
                   'jackhmmer': 8,
-                  'templates': 16,
+                  'template_searcher': 16,
                },
                mem: Mapping[str, int] = {
                   'hhblits': 8,
@@ -284,7 +284,9 @@ def process(self, input_fasta_path: str, msa_output_dir: str,other_args: Optiona
         query_sequence=input_sequence,
         hits=pdb_template_hits,
         query_pdb_code=None,
-        query_release_date=None)
+        query_release_date=None,
+        nproc=self.nprocs.get('template_searcher', 8)
+    )
 
     sequence_features = make_sequence_features(
         sequence=input_sequence,

diff --git a/apps/protein_folding/helixfold3/helixfold/data/templates.py b/apps/protein_folding/helixfold3/helixfold/data/templates.py
@@ -909,7 +909,8 @@ def get_templates(
       query_sequence: str,
       query_pdb_code: Optional[str],
       query_release_date: Optional[datetime.datetime],
-      hits: Sequence[parsers.TemplateHit]) -> TemplateSearchResult:
+      hits: Sequence[parsers.TemplateHit],
+      nproc: int) -> TemplateSearchResult:
     """Computes the templates for given query sequence."""
 
 
@@ -962,10 +963,11 @@ def get_templates(
       batch_results: Tuple[SingleHitResult]=Parallel(n_jobs=nproc,verbose=False)(
         delayed(
           _process_single_hit,
-        )(hit=hit,
-          query_sequence=query_sequence,
+        )(query_sequence=query_sequence,
+          query_pdb_code=query_pdb_code,
+          hit=hit,
           mmcif_dir=self._mmcif_dir,
-          max_template_date=self._max_template_date,
+          max_template_date=template_cutoff_date,
           release_dates=self._release_dates,
           obsolete_pdbs=self._obsolete_pdbs,
           strict_error_check=self._strict_error_check,