Refactor Job API (#2799)

* refactor fed job api

* improve docstrings

* refactor fed job api

* improve docstrings

* polish changes

* fix bugs; check args

* fix getattr

* refactor

* fixes and cleanup

* added ccwf and flower jobs

* added optional model selector for fed-avg

* update hello-world examples

* formatting and updates

* address feedback

* remove unnecessary stuff

* add license text

* detect duplicate executor

---------

Co-authored-by: Yan Cheng <[email protected]>

Author: SYangster

examples/advanced/job_api/pt/cyclic_cc_script_executor_cifar10.py

@@ -14,34 +14,26 @@
 
 from src.net import Net
 
-from nvflare import FedJob, ScriptExecutor
-from nvflare.app_common.ccwf import CyclicClientController, CyclicServerController
+from nvflare.app_common.ccwf.ccwf_job import CCWFJob, CyclicClientConfig, CyclicServerConfig
 from nvflare.app_common.ccwf.comps.simple_model_shareable_generator import SimpleModelShareableGenerator
+from nvflare.app_common.executors.script_executor import ScriptExecutor
 from nvflare.app_opt.pt.file_model_persistor import PTFileModelPersistor
 
 if __name__ == "__main__":
     n_clients = 2
     num_rounds = 3
     train_script = "src/cifar10_fl.py"
 
-    job = FedJob(name="cifar10_cyclic")
+    job = CCWFJob(name="cifar10_cyclic")
 
-    controller = CyclicServerController(num_rounds=num_rounds, max_status_report_interval=300)
-    job.to(controller, "server")
-
-    for i in range(n_clients):
-        executor = ScriptExecutor(
-            task_script_path=train_script, task_script_args=""  # f"--batch_size 32 --data_path /tmp/data/site-{i}"
-        )
-        job.to(executor, f"site-{i}", tasks=["train"], gpu=0)
-
-        # Add client-side controller for cyclic workflow
-        executor = CyclicClientController()
-        job.to(executor, f"site-{i}", tasks=["cyclic_*"])
-
-        # In swarm learning, each client uses a model persistor and shareable_generator
-        job.to(PTFileModelPersistor(model=Net()), f"site-{i}", id="persistor")
-        job.to(SimpleModelShareableGenerator(), f"site-{i}", id="shareable_generator")
+    job.add_cyclic(
+        server_config=CyclicServerConfig(num_rounds=num_rounds, max_status_report_interval=300),
+        client_config=CyclicClientConfig(
+            executor=ScriptExecutor(task_script_path=train_script),
+            persistor=PTFileModelPersistor(model=Net()),
+            shareable_generator=SimpleModelShareableGenerator(),
+        ),
+    )
 
     # job.export_job("/tmp/nvflare/jobs/job_config")
-    job.simulator_run("/tmp/nvflare/jobs/workdir")
+    job.simulator_run("/tmp/nvflare/jobs/workdir", n_clients=n_clients, gpu="0")

examples/advanced/job_api/pt/fedavg_model_learner_xsite_val_cifar10.py

@@ -15,16 +15,18 @@
 import os
 import sys
 
-sys.path.insert(0, os.path.join(os.getcwd(), "..", "..", "advanced", "cifar10"))
+sys.path.insert(0, os.path.join(os.getcwd(), "..", "..", "cifar10"))
 
 from pt.learners.cifar10_model_learner import CIFAR10ModelLearner
 from pt.networks.cifar10_nets import ModerateCNN
 from pt.utils.cifar10_data_splitter import Cifar10DataSplitter
 from pt.utils.cifar10_data_utils import load_cifar10_data
 
-from nvflare import FedAvg, FedJob
+from nvflare import FedJob
 from nvflare.app_common.executors.model_learner_executor import ModelLearnerExecutor
 from nvflare.app_common.workflows.cross_site_model_eval import CrossSiteModelEval
+from nvflare.app_common.workflows.fedavg import FedAvg
+from nvflare.app_opt.pt.job_config.model import PTModel
 
 if __name__ == "__main__":
     n_clients = 2
@@ -53,12 +55,12 @@
     job.to(data_splitter, "server")
 
     # Define the initial global model and send to server
-    job.to(ModerateCNN(), "server")
+    job.to(PTModel(ModerateCNN()), "server")
 
     for i in range(n_clients):
         learner = CIFAR10ModelLearner(train_idx_root=train_split_root, aggregation_epochs=aggregation_epochs, lr=0.01)
         executor = ModelLearnerExecutor(learner_id=job.as_id(learner))
-        job.to(executor, f"site-{i+1}", gpu=0)  # data splitter assumes client names start from 1
+        job.to(executor, f"site-{i+1}")  # data splitter assumes client names start from 1
 
     # job.export_job("/tmp/nvflare/jobs/job_config")
-    job.simulator_run("/tmp/nvflare/jobs/workdir")
+    job.simulator_run("/tmp/nvflare/jobs/workdir", gpu="0")

examples/advanced/job_api/pt/fedavg_script_executor_cifar10.py

@@ -14,7 +14,12 @@
 
 from src.net import Net
 
-from nvflare import FedAvg, FedJob, ScriptExecutor
+from nvflare.app_common.executors.script_executor import ScriptExecutor
+from nvflare.app_common.workflows.fedavg import FedAvg
+from nvflare.app_opt.pt.job_config.model import PTModel
+
+# from nvflare.app_opt.pt.job_config.fed_avg import FedAvgJob
+from nvflare.job_config.api import FedJob
 
 if __name__ == "__main__":
     n_clients = 2
@@ -32,16 +37,18 @@
     # job.to_server(controller)
 
     # Define the initial global model and send to server
-    job.to(Net(), "server")
-    # job.to_server(Net())
+    job.to(PTModel(Net()), "server")
+
+    # Note: We can optionally replace the above code with the FedAvgJob, which is a pattern to simplify FedAvg job creations
+    # job = FedAvgJob(name="cifar10_fedavg", num_rounds=num_rounds, n_clients=n_clients, initial_model=Net())
 
     # Add clients
     for i in range(n_clients):
         executor = ScriptExecutor(
             task_script_path=train_script, task_script_args=""  # f"--batch_size 32 --data_path /tmp/data/site-{i}"
         )
-        job.to(executor, f"site-{i}", gpu=0)
-        # job.to_clients(executor)
+        job.to(executor, target=f"site-{i}")
+    # job.to_clients(executor)
 
     # job.export_job("/tmp/nvflare/jobs/job_config")
-    job.simulator_run("/tmp/nvflare/jobs/workdir")
+    job.simulator_run("/tmp/nvflare/jobs/workdir", gpu="0")

examples/advanced/job_api/pt/fedavg_script_executor_dp_filter_cifar10.py

@@ -14,33 +14,25 @@
 
 from src.net import Net
 
-from nvflare import FedAvg, FedJob, FilterType, ScriptExecutor
+from nvflare import FilterType
+from nvflare.app_common.executors.script_executor import ScriptExecutor
 from nvflare.app_common.filters.percentile_privacy import PercentilePrivacy
+from nvflare.app_opt.pt.job_config.fed_avg import FedAvgJob
 
 if __name__ == "__main__":
     n_clients = 2
     num_rounds = 2
     train_script = "src/cifar10_fl.py"
 
-    job = FedJob(name="cifar10_fedavg_privacy")
-
-    # Define the controller workflow and send to server
-    controller = FedAvg(
-        num_clients=n_clients,
-        num_rounds=num_rounds,
-    )
-    job.to(controller, "server")
-
-    # Define the initial global model and send to server
-    job.to(Net(), "server")
+    job = FedAvgJob(name="cifar10_fedavg_privacy", num_rounds=num_rounds, n_clients=n_clients, initial_model=Net())
 
     for i in range(n_clients):
         executor = ScriptExecutor(task_script_path=train_script, task_script_args="")
-        job.to(executor, f"site-{i}", tasks=["train"], gpu=0)
+        job.to(executor, f"site-{i}", tasks=["train"])
 
         # add privacy filter.
         pp_filter = PercentilePrivacy(percentile=10, gamma=0.01)
         job.to(pp_filter, f"site-{i}", tasks=["train"], filter_type=FilterType.TASK_RESULT)
 
     # job.export_job("/tmp/nvflare/jobs/job_config")
-    job.simulator_run("/tmp/nvflare/jobs/workdir")
+    job.simulator_run("/tmp/nvflare/jobs/workdir", gpu="0")

examples/advanced/job_api/pt/fedavg_script_executor_lightning_cifar10.py

@@ -14,31 +14,22 @@
 
 from src.lit_net import LitNet
 
-from nvflare import FedAvg, FedJob, ScriptExecutor
+from nvflare.app_common.executors.script_executor import ScriptExecutor
+from nvflare.app_opt.pt.job_config.fed_avg import FedAvgJob
 
 if __name__ == "__main__":
     n_clients = 2
     num_rounds = 2
     train_script = "src/cifar10_lightning_fl.py"
 
-    job = FedJob(name="cifar10_fedavg_lightning")
-
-    # Define the controller workflow and send to server
-    controller = FedAvg(
-        num_clients=n_clients,
-        num_rounds=num_rounds,
-    )
-    job.to(controller, "server")
-
-    # Define the initial global model and send to server
-    job.to(LitNet(), "server")
+    job = FedAvgJob(name="cifar10_fedavg_lightning", num_rounds=num_rounds, n_clients=n_clients, initial_model=LitNet())
 
     # Add clients
     for i in range(n_clients):
         executor = ScriptExecutor(
             task_script_path=train_script, task_script_args=""  # f"--batch_size 32 --data_path /tmp/data/site-{i}"
         )
-        job.to(executor, f"site-{i}", gpu=0)
+        job.to(executor, f"site-{i}")
 
     # job.export_job("/tmp/nvflare/jobs/job_config")
-    job.simulator_run("/tmp/nvflare/jobs/workdir")
+    job.simulator_run("/tmp/nvflare/jobs/workdir", gpu="0")

examples/advanced/job_api/pt/swarm_script_executor_cifar10.py

@@ -14,56 +14,30 @@
 
 from src.net import Net
 
-from nvflare import FedJob, ScriptExecutor
 from nvflare.apis.dxo import DataKind
 from nvflare.app_common.aggregators.intime_accumulate_model_aggregator import InTimeAccumulateWeightedAggregator
-from nvflare.app_common.ccwf import (
-    CrossSiteEvalClientController,
-    CrossSiteEvalServerController,
-    SwarmClientController,
-    SwarmServerController,
-)
+from nvflare.app_common.ccwf.ccwf_job import CCWFJob, CrossSiteEvalConfig, SwarmClientConfig, SwarmServerConfig
 from nvflare.app_common.ccwf.comps.simple_model_shareable_generator import SimpleModelShareableGenerator
+from nvflare.app_common.executors.script_executor import ScriptExecutor
 from nvflare.app_opt.pt.file_model_persistor import PTFileModelPersistor
 
 if __name__ == "__main__":
     n_clients = 2
     num_rounds = 3
     train_script = "src/train_eval_submit.py"
 
-    job = FedJob(name="cifar10_swarm")
-
-    controller = SwarmServerController(
-        num_rounds=num_rounds,
+    job = CCWFJob(name="cifar10_swarm")
+    aggregator = InTimeAccumulateWeightedAggregator(expected_data_kind=DataKind.WEIGHTS)
+    job.add_swarm(
+        server_config=SwarmServerConfig(num_rounds=num_rounds),
+        client_config=SwarmClientConfig(
+            executor=ScriptExecutor(task_script_path=train_script, evaluate_task_name="validate"),
+            aggregator=aggregator,
+            persistor=PTFileModelPersistor(model=Net()),
+            shareable_generator=SimpleModelShareableGenerator(),
+        ),
+        cse_config=CrossSiteEvalConfig(eval_task_timeout=300),
     )
-    job.to(controller, "server")
-    controller = CrossSiteEvalServerController(eval_task_timeout=300)
-    job.to(controller, "server")
-
-    # Define the initial server model
-    job.to(Net(), "server")
-
-    for i in range(n_clients):
-        executor = ScriptExecutor(task_script_path=train_script, evaluate_task_name="validate")
-        job.to(executor, f"site-{i}", gpu=0, tasks=["train", "validate", "submit_model"])
-
-        # In swarm learning, each client acts also as an aggregator
-        aggregator = InTimeAccumulateWeightedAggregator(expected_data_kind=DataKind.WEIGHTS)
-
-        # In swarm learning, each client uses a model persistor and shareable_generator
-        persistor = PTFileModelPersistor(model=Net())
-        shareable_generator = SimpleModelShareableGenerator()
-
-        persistor_id = job.as_id(persistor)
-        client_controller = SwarmClientController(
-            aggregator_id=job.as_id(aggregator),
-            persistor_id=persistor_id,
-            shareable_generator_id=job.as_id(shareable_generator),
-        )
-        job.to(client_controller, f"site-{i}", tasks=["swarm_*"])
-
-        client_controller = CrossSiteEvalClientController(persistor_id=persistor_id)
-        job.to(client_controller, f"site-{i}", tasks=["cse_*"])
 
     # job.export_job("/tmp/nvflare/jobs/job_config")
-    job.simulator_run("/tmp/nvflare/jobs/workdir")
+    job.simulator_run("/tmp/nvflare/jobs/workdir", n_clients=n_clients, gpu="0")

examples/advanced/job_api/sklearn/kmeans_script_executor_higgs.py

@@ -18,8 +18,9 @@
 from src.kmeans_assembler import KMeansAssembler
 from src.split_csv import distribute_header_file, split_csv
 
-from nvflare import FedJob, ScriptExecutor
+from nvflare import FedJob
 from nvflare.app_common.aggregators.collect_and_assemble_aggregator import CollectAndAssembleAggregator
+from nvflare.app_common.executors.script_executor import ScriptExecutor
 from nvflare.app_common.shareablegenerators.full_model_shareable_generator import FullModelShareableGenerator
 from nvflare.app_common.workflows.scatter_and_gather import ScatterAndGather
 from nvflare.app_opt.sklearn.joblib_model_param_persistor import JoblibModelParamPersistor