[BioNeMo] Use decomposer register widget (#3784)

update other bionemo examples

Fixes # .

### Description

Update bionemo examples and tutorial for 2.7 with decomposer register
widget

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [ ] New tests added to cover the changes.
- [ ] Quick tests passed locally by running `./runtest.sh`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated.

Signed-off-by: Holger Roth <[email protected]>
Co-authored-by: Chester Chen <[email protected]>
Co-authored-by: Yuan-Ting Hsieh (謝沅廷) <[email protected]>

Author: 3 people

examples/advanced/bionemo/downstream/downstream_nvflare.ipynb

@@ -49,7 +49,7 @@
    "source": [
     "%%capture --no-display --no-stderr cell_output\n",
     "! pip install fuzzywuzzy PyTDC --no-dependencies  # install tdc without dependencies to avoid version conflicts in the BioNeMo container\n",
-    "! pip install nvflare~=2.6rc\n",
+    "! pip install \"nvflare~=2.7.0rc\"\n",
     "\n",
     "import os\n",
     "import warnings\n",

examples/advanced/bionemo/downstream/sabdab/run_sim_sabdab.py

@@ -20,6 +20,7 @@
 
 from nvflare import FilterType
 from nvflare.app_common.launchers.subprocess_launcher import SubprocessLauncher
+from nvflare.app_common.widgets.decomposer_reg import DecomposerRegister
 from nvflare.app_common.workflows.fedavg import FedAvg
 from nvflare.app_opt.pt.job_config.base_fed_job import BaseFedJob
 from nvflare.job_config.script_runner import BaseScriptRunner
@@ -38,6 +39,7 @@ def main(args):
         num_rounds=args.num_rounds,
     )
     job.to_server(controller)
+    job.to_server(DecomposerRegister(["nvflare.app_opt.pt.decomposers.TensorDecomposer"]))
 
     checkpoint_path = load(f"esm2/{args.model}:2.0")
     print(f"Downloaded {args.model} to {checkpoint_path}")
@@ -87,6 +89,7 @@ def main(args):
             BioNeMoParamsFilter(precision), client_name, tasks=["train", "validate"], filter_type=FilterType.TASK_DATA
         )
         job.to(BioNeMoStateDictFilter(), client_name, tasks=["train", "validate"], filter_type=FilterType.TASK_RESULT)
+        job.to(DecomposerRegister(["nvflare.app_opt.pt.decomposers.TensorDecomposer"]), client_name)
 
     job.export_job("./exported_jobs")
     job.simulator_run(f"/tmp/nvflare/bionemo/sabdab/{job.name}", gpu=args.sim_gpus)

examples/advanced/bionemo/downstream/scl/run_sim_scl.py

@@ -13,20 +13,17 @@
 # limitations under the License.
 
 import argparse
-import os
-import sys
 
 from bionemo.core.data.load import load
+from bionemo_filters import BioNeMoParamsFilter, BioNeMoStateDictFilter
 
 from nvflare import FilterType
 from nvflare.app_common.launchers.subprocess_launcher import SubprocessLauncher
+from nvflare.app_common.widgets.decomposer_reg import DecomposerRegister
 from nvflare.app_common.workflows.fedavg import FedAvg
 from nvflare.app_opt.pt.job_config.base_fed_job import BaseFedJob
 from nvflare.job_config.script_runner import BaseScriptRunner
 
-sys.path.append(os.path.join(os.getcwd(), ".."))  # include parent folder in path
-from bionemo_filters import BioNeMoParamsFilter, BioNeMoStateDictFilter
-
 
 def main(args):
     # Create BaseFedJob with initial model
@@ -38,6 +35,7 @@ def main(args):
         num_rounds=args.num_rounds,
     )
     job.to_server(controller)
+    job.to_server(DecomposerRegister(["nvflare.app_opt.pt.decomposers.TensorDecomposer"]))
 
     checkpoint_path = load(f"esm2/{args.model}:2.0")
     print(f"Downloaded {args.model} to {checkpoint_path}")
@@ -81,6 +79,7 @@ def main(args):
             BioNeMoParamsFilter(precision), client_name, tasks=["train", "validate"], filter_type=FilterType.TASK_DATA
         )
         job.to(BioNeMoStateDictFilter(), client_name, tasks=["train", "validate"], filter_type=FilterType.TASK_RESULT)
+        job.to(DecomposerRegister(["nvflare.app_opt.pt.decomposers.TensorDecomposer"]), client_name)
 
     job.export_job("./exported_jobs")
     job.simulator_run(f"/tmp/nvflare/bionemo/scl/{job.name}", gpu=args.sim_gpus)
@@ -91,9 +90,7 @@ def main(args):
     parser.add_argument("--num_clients", type=int, help="Number of clients", required=False, default=1)
     parser.add_argument("--num_rounds", type=int, help="Number of rounds", required=False, default=30)
     parser.add_argument("--local_steps", type=int, help="Number of rounds", required=False, default=10)
-    parser.add_argument(
-        "--train_script", type=str, help="Training script", required=False, default="../finetune_esm2.py"
-    )
+    parser.add_argument("--train_script", type=str, help="Training script", required=False, default="finetune_esm2.py")
     parser.add_argument("--exp_name", type=str, help="Job name prefix", required=False, default="fedavg")
     parser.add_argument("--model", choices=["8m", "650m", "3b"], help="ESM2 model", required=False, default="8m")
     parser.add_argument(

examples/advanced/bionemo/downstream/tap/run_sim_tap.py

@@ -20,6 +20,7 @@
 
 from nvflare import FilterType
 from nvflare.app_common.launchers.subprocess_launcher import SubprocessLauncher
+from nvflare.app_common.widgets.decomposer_reg import DecomposerRegister
 from nvflare.app_common.workflows.fedavg import FedAvg
 from nvflare.app_opt.pt.job_config.base_fed_job import BaseFedJob
 from nvflare.job_config.script_runner import BaseScriptRunner
@@ -38,6 +39,7 @@ def main(args):
         num_rounds=args.num_rounds,
     )
     job.to_server(controller)
+    job.to_server(DecomposerRegister(["nvflare.app_opt.pt.decomposers.TensorDecomposer"]))
 
     checkpoint_path = load(f"esm2/{args.model}:2.0")
     print(f"Downloaded {args.model} to {checkpoint_path}")
@@ -85,6 +87,7 @@ def main(args):
             BioNeMoParamsFilter(precision), client_name, tasks=["train", "validate"], filter_type=FilterType.TASK_DATA
         )
         job.to(BioNeMoStateDictFilter(), client_name, tasks=["train", "validate"], filter_type=FilterType.TASK_RESULT)
+        job.to(DecomposerRegister(["nvflare.app_opt.pt.decomposers.TensorDecomposer"]), client_name)
         job.to(
             BioNeMoExcludeParamsFilter(exclude_vars="regression_head"),
             client_name,

examples/tutorials/self-paced-training/part-5_federated_learning_applications_in_industries/chapter-11_federated_learning_in_healthcare_lifescience/11.2_drug_discovery/11.2.1_drug_discovery_bionemo/protein_property_prediction_with_bionemo.ipynb

@@ -36,11 +36,13 @@
    "cell_type": "code",
    "execution_count": null,
    "id": "b4ab20c8",
-   "metadata": {},
+   "metadata": {
+    "scrolled": true
+   },
    "outputs": [],
    "source": [
     "%%capture --no-display --no-stderr cell_output\n",
-    "! pip install nvflare>=2.6\n",
+    "! pip install \"nvflare~=2.7.0rc\"\n",
     "! pip install biopython --no-dependencies\n",
     "\n",
     "import io\n",
@@ -242,7 +244,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# for this to work run the task_fitting notebook first in ../nvflare_with_bionemo/task_fitting/task_fitting.ipynb in order to download the SCL dataset, each client will run on the same GPU.\n",
+    "# Each client will run on the same GPU.\n",
     "!python run_sim_scl.py --num_clients=3 --num_rounds=1 --local_steps=5000 --exp_name \"local\" --model \"8m\" --sim_gpus=\"0\""
    ]
   },

examples/tutorials/self-paced-training/part-5_federated_learning_applications_in_industries/chapter-11_federated_learning_in_healthcare_lifescience/11.2_drug_discovery/11.2.1_drug_discovery_bionemo/run_sim_scl.py

@@ -19,6 +19,7 @@
 
 from nvflare import FilterType
 from nvflare.app_common.launchers.subprocess_launcher import SubprocessLauncher
+from nvflare.app_common.widgets.decomposer_reg import DecomposerRegister
 from nvflare.app_common.workflows.fedavg import FedAvg
 from nvflare.app_opt.pt.job_config.base_fed_job import BaseFedJob
 from nvflare.job_config.script_runner import BaseScriptRunner
@@ -34,6 +35,7 @@ def main(args):
         num_rounds=args.num_rounds,
     )
     job.to_server(controller)
+    job.to_server(DecomposerRegister(["nvflare.app_opt.pt.decomposers.TensorDecomposer"]))
 
     checkpoint_path = load(f"esm2/{args.model}:2.0")
     print(f"Downloaded {args.model} to {checkpoint_path}")
@@ -77,6 +79,7 @@ def main(args):
             BioNeMoParamsFilter(precision), client_name, tasks=["train", "validate"], filter_type=FilterType.TASK_DATA
         )
         job.to(BioNeMoStateDictFilter(), client_name, tasks=["train", "validate"], filter_type=FilterType.TASK_RESULT)
+        job.to(DecomposerRegister(["nvflare.app_opt.pt.decomposers.TensorDecomposer"]), client_name)
 
     job.export_job("./exported_jobs")
     job.simulator_run(f"/tmp/nvflare/bionemo/scl/{job.name}", gpu=args.sim_gpus)