refactor: configure with hydra

Author: chungongyu

install_env.sh

@@ -53,6 +53,7 @@ conda install -y -c nvidia/label/cuda-${cuda_version} \
 conda install -y -c conda-forge \
     biopython \
     einops \
+    hydra-core \
     tensorboard \
     tqdm
 

profold2/command/evaluator.py

@@ -5,6 +5,7 @@
         for further help.
 """
 import os
+from dataclasses import dataclass, make_dataclass
 import logging
 import pickle
 
@@ -17,18 +18,23 @@
 from profold2.model import functional, profiler, snapshot, FeatureBuilder, ReturnValues
 from profold2.utils import exists, timing
 
-from profold2.command.worker import main, autocast_ctx, WorkerModel, WorkerXPU
+from profold2.command import worker
 
 
-def evaluate(rank, args):  # pylint: disable=redefined-outer-name
-  worker = WorkerModel(rank, args)
-  feats, model = worker.load(args.model)
-  features = FeatureBuilder(feats).to(worker.device())
+@dataclass
+class Args(worker.Args):
+  pass
+
+
+def run(rank, args):  # pylint: disable=redefined-outer-name
+  runner = worker.WorkerModel(rank, args)
+  feats, model = runner.load(args.model)
+  features = FeatureBuilder(feats).to(runner.device())
   logging.info('feats: %s', feats)
 
   kwargs = {}
-  if rank.is_available() and WorkerXPU.world_size(args.nnodes) > 1:
-    kwargs['num_replicas'] = WorkerXPU.world_size(args.nnodes)
+  if rank.is_available() and worker.world_size(args.nnodes) > 1:
+    kwargs['num_replicas'] = worker.world_size(args.nnodes)
     kwargs['rank'] = rank.rank
   test_loader = dataset.load(
       data_dir=args.eval_data,
@@ -68,7 +74,7 @@ def data_eval(idx, batch):
     # predict - out is (batch, L * 3, 3)
     with timing(f'Running model on {fasta_name} {fasta_len}', logging.debug):
       with torch.no_grad():
-        with autocast_ctx(args.amp_enabled):
+        with worker.autocast_ctx(args.amp_enabled):
           r = ReturnValues(
               **model(
                   batch=batch,  # pylint: disable=not-callable
@@ -320,34 +326,31 @@ def add_arguments(parser):  # pylint: disable=redefined-outer-name
 
 if __name__ == '__main__':
   import argparse
+  import hydra
 
-  parser = argparse.ArgumentParser()
-
-  # init distributed env
-  parser.add_argument('--nnodes', type=int, default=None, help='number of nodes.')
-  parser.add_argument('--node_rank', type=int, default=0, help='rank of the node.')
-  parser.add_argument(
-      '--local_rank', type=int, default=None, help='local rank of xpu, default=None'
-  )
-  parser.add_argument(
-      '--init_method',
-      type=str,
-      default='file:///tmp/profold2.dist',
-      help='method to initialize the process group, '
-      'default=\'file:///tmp/profold2.dist\''
+  parser = argparse.ArgumentParser(
+      formatter_class=argparse.ArgumentDefaultsHelpFormatter
   )
 
-  # output dir
+  parser.add_argument('-c', '--config', type=str, default=None, help='config file.')
   parser.add_argument(
-      '-o',
-      '--prefix',
-      type=str,
-      default='.',
-      help='prefix of out directory, default=\'.\''
+      'overrides',
+      nargs='*',
+      metavar='KEY=VAL',
+      help='override configs, see: https://hydra.cc'
   )
-  add_arguments(parser)
-  parser.add_argument('-v', '--verbose', action='store_true', help='verbose')
 
   args = parser.parse_args()
-
-  main(args, evaluate)
+  config_dir, config_name = os.path.split(
+      os.path.abspath(args.config)
+  ) if exists(args.config) else (os.getcwd(), None)
+
+  with hydra.initialize_config_dir(
+      version_base=None, config_dir=config_dir, job_name=__file__
+  ):
+    worker.main(
+        make_dataclass('t', [], namespace={
+            'Args': Args,
+            'run': run
+        }), hydra.compose(config_name, args.overrides)
+    )

profold2/command/main.py

@@ -4,75 +4,56 @@
      ```
      for further help.
 """
+import os
 import argparse
 
+import hydra
+
 from profold2.command import (evaluator, predictor, trainer, worker)
-from profold2.utils import env
+from profold2.utils import env, exists
 
-_COMMANDS = [
-    ('train', trainer.train, trainer.add_arguments),
-    ('evaluate', evaluator.evaluate, evaluator.add_arguments),
-    ('predict', predictor.predict, predictor.add_arguments),
-]
+_COMMANDS = [('train', trainer), ('evaluate', evaluator), ('predict', predictor)]
 
 
 def create_args():
   formatter_class = argparse.ArgumentDefaultsHelpFormatter
   parser = argparse.ArgumentParser(formatter_class=formatter_class)
 
-  # distributed args
-  parser.add_argument(
-      '--nnodes',
-      type=int,
-      default=env('SLURM_NNODES', defval=None, dtype=int),
-      help='number of nodes.'
-  )
-  parser.add_argument(
-      '--node_rank',
-      type=int,
-      default=env('SLURM_NODEID', defval=0, dtype=int),
-      help='rank of the node.'
-  )
-  parser.add_argument(
-      '--local_rank',
-      type=int,
-      default=int(env('LOCAL_RANK', defval=0, dtype=int)),
-      help='local rank of xpu.'
-  )
-  parser.add_argument(
-      '--init_method',
-      type=str,
-      default=None,
-      help='method to initialize the process group.'
-  )
-
   # command args
   subparsers = parser.add_subparsers(dest='command', required=True)
-  for cmd, _, add_arguments in _COMMANDS:
+  for cmd, _ in _COMMANDS:
     cmd_parser = subparsers.add_parser(cmd, formatter_class=formatter_class)
-
-    # output dir
     cmd_parser.add_argument(
-        '-o', '--prefix', type=str, default='.', help='prefix of out directory.'
+        '-c', '--config', type=str, default=None, help='config file.'
+    )
+    cmd_parser.add_argument(
+        'overrides',
+        nargs='*',
+        metavar='KEY=VAL',
+        help='override configs, see: https://hydra.cc'
     )
-    add_arguments(cmd_parser)
-    # verbose
-    cmd_parser.add_argument('-v', '--verbose', action='store_true', help='verbose')
 
   return parser.parse_args()
 
 
-def create_fn(args):  # pylint: disable=redefined-outer-name
-  for cmd, fn, _ in _COMMANDS:
+def create_task(args):  # pylint: disable=redefined-outer-name
+  for cmd, task in _COMMANDS:
     if cmd == args.command:
-      return fn
+      return task
   return None
 
 
 def main():
   args = create_args()
-  work_fn = create_fn(args)
-  worker.main(args, work_fn)
+  config_dir, config_name = os.path.split(
+      os.path.abspath(args.config)
+  ) if exists(args.config) else ('.', None)
+
+  with hydra.initialize_config_dir(
+      version_base=None, config_dir=config_dir, job_name=args.command
+  ):
+    task = create_task(args)
+    worker.main(task, hydra.compose(config_name, args.overrides))
 
 
 if __name__ == '__main__':

profold2/command/predictor.py

@@ -5,6 +5,7 @@
         for further help.
 """
 import os
+from dataclasses import dataclass, make_dataclass
 import functools
 import glob
 import json
@@ -23,6 +24,7 @@
 from profold2.model import profiler, snapshot, FeatureBuilder, ReturnValues
 from profold2.utils import exists, timing
 
+from profold2.command import worker
 from profold2.command.worker import main, autocast_ctx, WorkerModel, WorkerXPU
 
 
@@ -128,7 +130,21 @@ def _location_split(model_location):
     yield model_name, (features, model)
 
 
-def predict(rank, args):  # pylint: disable=redefined-outer-name
+@dataclass
+class Args(worker.Args):
+  models: list[str]  # models to be loaded using[model_name=model_location] format
+  model_recycles: int = 0  # number of recycles
+  model_shard_size: Optional[int] = 0  # shard size in the evoformer model
+  map_location: str = 'cpu'  # remapped to an alternative set of devices
+
+  data_dir: Optional[str] = None  # dataset dir
+  data_idx: Optional[str] = None  # dataset idx
+  add_pseudo_linker: bool = False  # enable loading complex data
+
+  max_msa_size: int = 1024
+
+
+def run(rank, args):  # pylint: disable=redefined-outer-name
   model_runners = dict(_load_models(rank, args))
   logging.info('Have %d models: %s', len(model_runners), list(model_runners.keys()))
 
@@ -360,35 +376,31 @@ def add_arguments(parser):  # pylint: disable=redefined-outer-name
 
 if __name__ == '__main__':
   import argparse
+  import hydra
 
-  parser = argparse.ArgumentParser()
-
-  # init distributed env
-  parser.add_argument('--nnodes', type=int, default=None, help='number of nodes.')
-  parser.add_argument('--node_rank', type=int, default=0, help='rank of the node.')
-  parser.add_argument(
-      '--local_rank', type=int, default=None, help='local rank of xpu, default=None'
-  )
-  parser.add_argument(
-      '--init_method',
-      type=str,
-      default='file:///tmp/profold2.dist',
-      help='method to initialize the process group, '
-      'default=\'file:///tmp/profold2.dist\''
+  parser = argparse.ArgumentParser(
+      formatter_class=argparse.ArgumentDefaultsHelpFormatter
   )
 
-  # output dir
+  parser.add_argument('-c', '--config', type=str, default=None, help='config file.')
   parser.add_argument(
-      '-o',
-      '--prefix',
-      type=str,
-      default='.',
-      help='prefix of out directory, default=\'.\''
+      'overrides',
+      nargs='*',
+      metavar='KEY=VAL',
+      help='override configs, see: https://hydra.cc'
   )
-  add_arguments(parser)
-  # verbose
-  parser.add_argument('-v', '--verbose', action='store_true', help='verbose')
 
   args = parser.parse_args()
-
-  main(args, predict)
+  config_dir, config_name = os.path.split(
+      os.path.abspath(args.config)
+  ) if exists(args.config) else (os.getcwd(), None)
+
+  with hydra.initialize_config_dir(
+      version_base=None, config_dir=config_dir, job_name=__file__
+  ):
+    worker.main(
+        make_dataclass('t', [], namespace={
+            'Args': Args,
+            'run': run
+        }), hydra.compose(config_name, args.overrides)
+    )