Merge branch 'setup-improvements' of https://github.com/aqlaboratory/openfold...

Merge branch 'setup-improvements' of https://github.com/aqlaboratory/openfold into setup-improvements

environment.yml

@@ -23,6 +23,7 @@ dependencies:
   - modelcif==0.7
   - awscli
   - ml-collections
+  - mkl==2024.0.0
   - aria2
   - mkl==2024.0
   - git

run_pretrained_openfold.py

@@ -190,6 +190,11 @@ def main(args):
             custom_config_dict = json.load(f)
         config.update_from_flattened_dict(custom_config_dict)
 
+    if args.experiment_config_json: 
+        with open(args.experiment_config_json, 'r') as f:
+            custom_config_dict = json.load(f)
+        config.update_from_flattened_dict(custom_config_dict)
+
     if args.trace_model:
         if not config.data.predict.fixed_size:
             raise ValueError(

train_openfold.py

@@ -382,7 +382,7 @@ def main(args):
         callbacks.append(lr_monitor)
 
     loggers = []
-    is_rank_zero = int(os.environ.get("PMI_RANK")) == 0
+    is_rank_zero = args.mpi_plugin and (int(os.environ.get("PMI_RANK")) == 0)
     if(args.wandb):
         if args.mpi_plugin and is_rank_zero:
             wandb_init_dict = dict(