fix: ramdisk permission

PaddlePaddle · Aug 28, 2024 · 484e3bc · 484e3bc
1 parent 7dc26a7
commit 484e3bc
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 2 deletions.
diff --git a/apps/protein_folding/helixfold3/helixfold/config/helixfold.yaml b/apps/protein_folding/helixfold3/helixfold/config/helixfold.yaml
@@ -18,7 +18,8 @@ ramdisk:  # Force to load database to ram disk
   uniref90: false # 67 GB
   mgnify: false # 64 GB
 
-
+mem:
+  hhblits: 64 # in GB
 
 nproc_msa: 
   hhblits: 16 # Number of processors used by hhblits

diff --git a/apps/protein_folding/helixfold3/helixfold/data/pipeline_parallel.py b/apps/protein_folding/helixfold3/helixfold/data/pipeline_parallel.py
@@ -152,10 +152,14 @@ def __init__(self,
                nprocs: Mapping[str, int] = {
                   'hhblits': 16,
                   'jackhmmer': 8,
+               },
+               mem: Mapping[str, int] = {
+                  'hhblits': 8,
                }):
     """Initializes the data pipeline. Constructs a feature dict for a given FASTA file."""
     self._use_small_bfd = use_small_bfd
     self.nprocs=nprocs
+    self.mem = mem
     self.jackhmmer_uniref90_runner = jackhmmer.Jackhmmer(
         binary_path=jackhmmer_binary_path,
         database_path=uniref90_database_path, n_cpu=self.nprocs.get('jackhmmer', 8))
@@ -166,7 +170,7 @@ def __init__(self,
     else:
       self.bfd_runner = hhblits.HHBlits(
           binary_path=hhblits_binary_path,
-          databases=[bfd_database_path, uniclust30_database_path], n_cpu=self.nprocs.get('hhblits', 8))
+          databases=[bfd_database_path, uniclust30_database_path], n_cpu=self.nprocs.get('hhblits', 8), maxmem=self.mem.get('hhblits', 8))
     self.jackhmmer_mgnify_runner = jackhmmer.Jackhmmer(
         binary_path=jackhmmer_binary_path,
         database_path=mgnify_database_path, n_cpu=self.nprocs.get('jackhmmer', 8))

diff --git a/apps/protein_folding/helixfold3/helixfold/data/tools/hhblits.py b/apps/protein_folding/helixfold3/helixfold/data/tools/hhblits.py
@@ -36,6 +36,7 @@ def __init__(self,
                binary_path: str,
                databases: Sequence[str],
                n_cpu: int = 4,
+               maxmem: int =8,
                n_iter: int = 3,
                e_value: float = 0.001,
                maxseq: int = 1_000_000,
@@ -83,6 +84,7 @@ def __init__(self,
         raise ValueError(f'Could not find HHBlits database {database_path}')
 
     self.n_cpu = n_cpu
+    self.maxmem = maxmem
     self.n_iter = n_iter
     self.e_value = e_value
     self.maxseq = maxseq
@@ -107,6 +109,7 @@ def query(self, input_fasta_path: str) -> List[Mapping[str, Any]]:
           self.binary_path,
           '-i', input_fasta_path,
           '-cpu', str(self.n_cpu),
+          '-maxmem', str(self.maxmem),
           '-oa3m', a3m_path,
           '-o', '/dev/null',
           '-n', str(self.n_iter),

diff --git a/apps/protein_folding/helixfold3/helixfold/inference.py b/apps/protein_folding/helixfold3/helixfold/inference.py
@@ -145,6 +145,7 @@ def load_to_dev_shm(file_path: str, ramdisk_path: str = "/dev/shm") -> str:
     target_path = os.path.join(ramdisk_path, pathlib.Path(file_path).name)
     with timing(f'loading {file_path} -> {target_path}'):
         shutil.copy(file_path, target_path)
+        os.chmod(target_path,777)
 
     return target_path
 
@@ -189,6 +190,7 @@ def get_msa_templates_pipeline(cfg: DictConfig) -> Dict:
         use_small_bfd=cfg.use_small_bfd,
         use_precomputed_msas=use_precomputed_msas,
         nprocs=cfg.nproc_msa,
+        mem=cfg.mem
         )