From 6b6fcead08c468d4c8966b6eabf5055f1a9d160c Mon Sep 17 00:00:00 2001 From: Kevin Date: Wed, 8 Jul 2026 20:58:33 -0400 Subject: [PATCH 1/8] Add recipe tensor streaming helper --- .../data_scientist_guide/recipe_api.rst | 4 ++ examples/advanced/kaplan-meier-he/job.py | 2 +- examples/advanced/tensor-stream/README.md | 9 +-- examples/advanced/tensor-stream/job.py | 5 +- .../hello-world/hello-log-streaming/README.md | 8 +-- nvflare/recipe/spec.py | 55 +++++++++++++++++++ 6 files changed, 70 insertions(+), 13 deletions(-) diff --git a/docs/user_guide/data_scientist_guide/recipe_api.rst b/docs/user_guide/data_scientist_guide/recipe_api.rst index 69c1f2fc2f..33335d0ee3 100644 --- a/docs/user_guide/data_scientist_guide/recipe_api.rst +++ b/docs/user_guide/data_scientist_guide/recipe_api.rst @@ -133,6 +133,10 @@ Shared helpers: Add the default Recipe log-streaming components. If no file names are given, the recipe streams ``log.json``. +``recipe.enable_tensor_streaming(format="pytorch", tasks=None, tensor_send_timeout=30.0, wait_send_task_data_all_clients_timeout=300.0)`` + Add matching tensor-streaming components to the server and all generated + client apps. The format must match the recipe's ``server_expected_format``. + Utility helpers: ``add_experiment_tracking(recipe, tracking_type, tracking_config=None, client_side=False, server_side=True, clients=None)`` diff --git a/examples/advanced/kaplan-meier-he/job.py b/examples/advanced/kaplan-meier-he/job.py index bac844c93a..59906d81b8 100644 --- a/examples/advanced/kaplan-meier-he/job.py +++ b/examples/advanced/kaplan-meier-he/job.py @@ -169,7 +169,7 @@ def main(): ) # Export job - recipe.job.export_job(args.job_dir) + recipe.export(args.job_dir) # Run recipe if args.startup_kit_location: diff --git a/examples/advanced/tensor-stream/README.md b/examples/advanced/tensor-stream/README.md index 5030e1bf24..4b355afa8a 100644 --- a/examples/advanced/tensor-stream/README.md +++ b/examples/advanced/tensor-stream/README.md @@ -145,14 +145,15 @@ This configuration is shared between standalone training and federated learning, ### Tensor Streaming Setup -The job uses `TensorServerStreamer` and `TensorClientStreamer` components: +Enable the matching server and client tensor-streaming components through the recipe: ```python -recipe.job.to_server(TensorServerStreamer(), "tensor_server_streamer") -recipe.job.to_clients(TensorClientStreamer(), "tensor_client_streamer") +recipe.enable_tensor_streaming() ``` -These streamers handle the efficient transmission of large model tensors between server and clients. +The helper also accepts `format`, `tasks`, `tensor_send_timeout`, and +`wait_send_task_data_all_clients_timeout` when the defaults need to be adjusted. +It applies the shared settings consistently to both sides. ## Tensor Streaming Components diff --git a/examples/advanced/tensor-stream/job.py b/examples/advanced/tensor-stream/job.py index 1abc6010ef..3ce23414ff 100644 --- a/examples/advanced/tensor-stream/job.py +++ b/examples/advanced/tensor-stream/job.py @@ -20,8 +20,6 @@ from model import get_model from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe -from nvflare.app_opt.tensor_stream.client import TensorClientStreamer -from nvflare.app_opt.tensor_stream.server import TensorServerStreamer from nvflare.client.config import ExchangeFormat from nvflare.recipe import SimEnv, add_experiment_tracking @@ -52,8 +50,7 @@ def main(): ) add_experiment_tracking(recipe, tracking_type="tensorboard") - recipe.job.to_server(TensorServerStreamer(), "tensor_server_streamer") - recipe.job.to_clients(TensorClientStreamer(), "tensor_client_streamer") + recipe.enable_tensor_streaming() env = SimEnv(num_clients=n_clients) run = recipe.execute(env) diff --git a/examples/hello-world/hello-log-streaming/README.md b/examples/hello-world/hello-log-streaming/README.md index 4e8db9097e..fa98fe8c05 100644 --- a/examples/hello-world/hello-log-streaming/README.md +++ b/examples/hello-world/hello-log-streaming/README.md @@ -106,13 +106,13 @@ job.to_server(JobLogReceiver(idle_timeout=15.0)) ### Using the Recipe API -When using a recipe (e.g. `NumpyFedAvgRecipe`), access the underlying job via -`recipe.job`: +When using a recipe (e.g. `NumpyFedAvgRecipe`), enable log streaming through +the recipe helper. Pass `"log.txt"` explicitly to stream the standard text log; +without arguments, the helper streams `log.json`. ```python recipe = NumpyFedAvgRecipe(...) -recipe.job.to_clients(JobLogStreamer()) -recipe.job.to_server(JobLogReceiver()) +recipe.enable_log_streaming("log.txt") ``` ## How It Works diff --git a/nvflare/recipe/spec.py b/nvflare/recipe/spec.py index 63e9006e88..a535c0d93e 100644 --- a/nvflare/recipe/spec.py +++ b/nvflare/recipe/spec.py @@ -542,6 +542,61 @@ def enable_log_streaming(self, *file_names: str) -> None: self._add_to_client_apps(JobLogStreamer(log_file_name=name)) self.job.to_server(JobLogReceiver()) + def enable_tensor_streaming( + self, + format: str = "pytorch", + tasks: Optional[List[str]] = None, + tensor_send_timeout: float = 30.0, + wait_send_task_data_all_clients_timeout: float = 300.0, + ) -> None: + """Enable tensor streaming between the server and all client apps. + + The same exchange format, task names, and per-transfer timeout are used + on both sides. The format must match the recipe's + ``server_expected_format`` when the recipe declares one. + + Args: + format: Tensor exchange format. Defaults to ``"pytorch"`` + (``ExchangeFormat.PYTORCH``). + tasks: Task names whose tensors should be streamed. ``None`` uses + the streamers' default ``["train"]``. + tensor_send_timeout: Timeout in seconds for each tensor transfer. + wait_send_task_data_all_clients_timeout: Maximum time in seconds + for the server to wait for all clients to receive task tensors. + + Raises: + ValueError: If ``format`` does not match a declared + ``server_expected_format``. + """ + from nvflare.app_opt.tensor_stream.client import TensorClientStreamer + from nvflare.app_opt.tensor_stream.server import TensorServerStreamer + + server_expected_format = getattr(self, "server_expected_format", None) + if server_expected_format is not None and format != server_expected_format: + raise ValueError( + f"tensor streaming format {format!r} must match server_expected_format {server_expected_format!r}" + ) + + server_tasks = list(tasks) if tasks is not None else None + client_tasks = list(tasks) if tasks is not None else None + self.job.to_server( + TensorServerStreamer( + format=format, + tasks=server_tasks, + tensor_send_timeout=tensor_send_timeout, + wait_send_task_data_all_clients_timeout=wait_send_task_data_all_clients_timeout, + ), + id="tensor_server_streamer", + ) + self._add_to_client_apps( + TensorClientStreamer( + format=format, + tasks=client_tasks, + tensor_send_timeout=tensor_send_timeout, + ), + id="tensor_client_streamer", + ) + def add_decomposers(self, decomposers: List[Union[str, Decomposer]]): """Add decomposers to the job From e3bb9715c7158a491430c5ea386341505535ec79 Mon Sep 17 00:00:00 2001 From: Kevin Date: Wed, 8 Jul 2026 21:00:22 -0400 Subject: [PATCH 2/8] Update recipe experiment tracking guidance --- .../advanced/experiment-tracking/README.md | 28 ++-- .../mlflow/hello-pt-mlflow-client/README.md | 126 +++++++++++++----- .../experiment-tracking/wandb/README.md | 10 +- 3 files changed, 120 insertions(+), 44 deletions(-) diff --git a/examples/advanced/experiment-tracking/README.md b/examples/advanced/experiment-tracking/README.md index 2ffd8f7316..5b12f228ab 100644 --- a/examples/advanced/experiment-tracking/README.md +++ b/examples/advanced/experiment-tracking/README.md @@ -213,17 +213,27 @@ add_experiment_tracking(recipe, "mlflow", tracking_config={ ### Client-Side Tracking (Decentralized) ```python -from nvflare.app_opt.tracking.mlflow.mlflow_receiver import MLflowReceiver - -# Add tracking to specific clients -for site_name in ["site-1", "site-2"]: - receiver = MLflowReceiver( - tracking_uri=f"file:///tmp/{site_name}/mlruns", - kw_args={"experiment_name": f"{site_name}-experiment"} - ) - recipe.job.to(receiver, site_name, id="mlflow_receiver") +# Add the receiver to every client, with no server-side receiver. +# tracking_uri=None creates a local MLflow store in each site's job workspace. +add_experiment_tracking( + recipe, + "mlflow", + tracking_config={ + "tracking_uri": None, + "kw_args": {"experiment_name": "local-client-experiment"}, + }, + client_side=True, + server_side=False, +) ``` +If sites need different tracking URIs, credentials, or experiment names, construct +the recipe with per-site client apps (for example, +`per_site_config={"site-1": {}, "site-2": {}}`) and call +`add_experiment_tracking()` once per configuration with `clients=[site_name]`. +The `clients` argument requires that per-site topology; it intentionally rejects a +recipe whose client app is deployed through the default `@ALL` target. + --- ## Additional Resources diff --git a/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/README.md b/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/README.md index e907823ef3..f1abab7f47 100644 --- a/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/README.md +++ b/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/README.md @@ -10,19 +10,19 @@ This example demonstrates **site-specific (decentralized) MLflow tracking** usin - Server handles MLflow authentication **Client-Side** (this example): -- Each client has its own MLflow instance +- Each client has its own local MLflow store - Site-specific metric tracking -- Each site manages its own MLflow server +- Each site can optionally point at its own MLflow server ## Overview This example shows how to configure per-client MLflow tracking: ```python +from model import SimpleNetwork + from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe -from nvflare.app_opt.tracking.mlflow.mlflow_receiver import MLflowReceiver -from nvflare.apis.analytix import ANALYTIC_EVENT_TYPE -from nvflare.recipe import SimEnv +from nvflare.recipe import SimEnv, add_experiment_tracking # Create recipe WITHOUT default server-side tracking recipe = FedAvgRecipe( @@ -36,27 +36,31 @@ recipe = FedAvgRecipe( train_script="client.py", ) -# Add MLflow receiver to each client -for i in range(2): - site_name = f"site-{i + 1}" - tracking_uri = f"file:///tmp/nvflare/jobs/workdir/{site_name}/mlruns" - - receiver = MLflowReceiver( - tracking_uri=tracking_uri, - events=[ANALYTIC_EVENT_TYPE], # Listen to LOCAL events (not federated) - kw_args={"experiment_name": f"site-{site_name}-experiment"} - ) - - recipe.job.to(receiver, site_name, id="mlflow_receiver") +# Add an MLflow receiver to every client, but not to the server. +# With tracking_uri=None, each receiver creates a store in its local job workspace. +add_experiment_tracking( + recipe, + "mlflow", + tracking_config={ + "tracking_uri": None, + "kw_args": { + "experiment_name": "nvflare-fedavg-experiment", + "run_name": "nvflare-fedavg-client", + }, + }, + client_side=True, + server_side=False, +) env = SimEnv(num_clients=2) run = recipe.execute(env) ``` **Key points**: -- By default, no server-side tracking is enabled (analytics_receiver defaults to None) -- `events=[ANALYTIC_EVENT_TYPE]` - Listen to local events (not `fed.` events) -- Each site gets its own `tracking_uri` +- `client_side=True, server_side=False` keeps metric handling on the clients. +- The helper configures receivers to listen to local events rather than federated `fed.` events. +- `tracking_uri=None` resolves to a separate MLflow store in each client's job workspace. +- Sites can use explicit local paths or remote MLflow servers when their deployment requires them. ## Setup @@ -87,16 +91,22 @@ python job.py Since each site has its own MLflow receiver, metrics are stored separately: +With the example's `tracking_uri=None`, the store is created under each site's +job-result directory as `//mlflow`. Under the default +simulation workspace, use `find` to locate the generated directory: + ### View Site-1 Metrics: ```bash -mlflow ui --backend-store-uri /tmp/nvflare/jobs/workdir/site-1/mlruns +find /tmp/nvflare/jobs/workdir/fedavg_mlflow_client/site-1 -type d -name mlflow +mlflow ui --backend-store-uri ``` Open browser to `http://localhost:5000` ### View Site-2 Metrics: ```bash -mlflow ui --backend-store-uri /tmp/nvflare/jobs/workdir/site-2/mlruns --port 5001 +find /tmp/nvflare/jobs/workdir/fedavg_mlflow_client/site-2 -type d -name mlflow +mlflow ui --backend-store-uri --port 5001 ``` Open browser to `http://localhost:5001` @@ -166,29 +176,81 @@ To avoid confusion: ### Change Tracking Location ```python -tracking_uri = f"file:///my/custom/path/{site_name}/mlruns" +add_experiment_tracking( + recipe, + "mlflow", + tracking_config={"tracking_uri": "file:///my/custom/path/mlruns"}, + client_side=True, + server_side=False, +) ``` +In a distributed deployment, the same local path is resolved independently on +each site. In simulation, all clients share the host filesystem, so use the +default `tracking_uri=None` behavior or configure distinct paths if the stores +must remain physically separate. + +For different configuration at each site, construct the recipe with per-site +client apps and then target each site explicitly: + +```python +sites = ["site-1", "site-2"] +recipe = FedAvgRecipe( + ..., + per_site_config={site: {} for site in sites}, +) + +for site in sites: + add_experiment_tracking( + recipe, + "mlflow", + tracking_config={ + "tracking_uri": f"file:///my/custom/path/{site}/mlruns", + "kw_args": {"experiment_name": f"{site}-experiment"}, + }, + client_side=True, + server_side=False, + clients=[site], + ) +``` + +Targeted `clients=[...]` placement requires `per_site_config` when the recipe is +constructed; it cannot split an existing `@ALL` client app after the fact. + ### Add Experiment Tags ```python -receiver = MLflowReceiver( - tracking_uri=tracking_uri, - events=[ANALYTIC_EVENT_TYPE], - kw_args={ - "experiment_name": f"{site_name}-experiment", - "run_name": f"{site_name}-run-001", - "experiment_tags": {"site": site_name, "privacy": "local"}, - } +add_experiment_tracking( + recipe, + "mlflow", + tracking_config={ + "tracking_uri": None, + "kw_args": { + "experiment_name": "local-client-experiment", + "run_name": "client-run-001", + "experiment_tags": {"privacy": "local"}, + }, + }, + client_side=True, + server_side=False, ) ``` ### Remote MLflow Server ```python -tracking_uri = "http://mlflow-server.site-1.local:5000" +add_experiment_tracking( + recipe, + "mlflow", + tracking_config={"tracking_uri": "http://mlflow-server.local:5000"}, + client_side=True, + server_side=False, +) ``` +If each site uses a different server or credentials, use the per-site pattern +above and pass the appropriate configuration with `clients=[site]`. + --- ## Additional Resources diff --git a/examples/advanced/experiment-tracking/wandb/README.md b/examples/advanced/experiment-tracking/wandb/README.md index e580f5b8aa..ff660f736a 100644 --- a/examples/advanced/experiment-tracking/wandb/README.md +++ b/examples/advanced/experiment-tracking/wandb/README.md @@ -152,9 +152,13 @@ add_experiment_tracking(recipe, "wandb", tracking_config=wandb_config) ```python # Each client logs to its own WandB run -for site_name in ["site-1", "site-2"]: - receiver = WandBReceiver(**client_config) - recipe.job.to(receiver, site_name, id="wandb_receiver") +add_experiment_tracking( + recipe, + "wandb", + tracking_config=client_config, + client_side=True, + server_side=False, +) ``` **Event Flow**: From ca6c15613a8cd608bd625d2ada59c13030bf34f4 Mon Sep 17 00:00:00 2001 From: Kevin Date: Wed, 8 Jul 2026 21:01:52 -0400 Subject: [PATCH 3/8] Add final global evaluation recipe helper --- .../data_scientist_guide/recipe_api.rst | 4 ++ nvflare/recipe/__init__.py | 9 ++- nvflare/recipe/utils.py | 60 +++++++++++++++++++ .../auto-fl-research/tasks/cifar10/job.py | 25 +------- .../auto-fl-research/tasks/vlm_med/job.py | 25 +------- 5 files changed, 74 insertions(+), 49 deletions(-) diff --git a/docs/user_guide/data_scientist_guide/recipe_api.rst b/docs/user_guide/data_scientist_guide/recipe_api.rst index 33335d0ee3..94b59b649b 100644 --- a/docs/user_guide/data_scientist_guide/recipe_api.rst +++ b/docs/user_guide/data_scientist_guide/recipe_api.rst @@ -149,6 +149,10 @@ Utility helpers: Add cross-site evaluation to a training recipe when the recipe/framework supports it. +``add_final_global_evaluation(recipe, participating_clients=None, validation_timeout=6000)`` + Add a final evaluation of a PyTorch recipe's persisted global model without + asking clients to submit their local models. + Per-Site And Metadata Helpers ----------------------------- diff --git a/nvflare/recipe/__init__.py b/nvflare/recipe/__init__.py index 46efdcd900..2484c58b52 100644 --- a/nvflare/recipe/__init__.py +++ b/nvflare/recipe/__init__.py @@ -18,7 +18,13 @@ from .prod_env import ProdEnv from .run import Run from .sim_env import SimEnv -from .utils import add_cross_site_evaluation, add_experiment_tracking, set_per_site_config, set_recipe_meta +from .utils import ( + add_cross_site_evaluation, + add_experiment_tracking, + add_final_global_evaluation, + set_per_site_config, + set_recipe_meta, +) __all__ = [ "SimEnv", @@ -27,6 +33,7 @@ "Run", "add_experiment_tracking", "add_cross_site_evaluation", + "add_final_global_evaluation", "set_per_site_config", "set_recipe_meta", "FedAvgRecipe", diff --git a/nvflare/recipe/utils.py b/nvflare/recipe/utils.py index 763133b919..8de122af26 100644 --- a/nvflare/recipe/utils.py +++ b/nvflare/recipe/utils.py @@ -299,6 +299,66 @@ def add_experiment_tracking( recipe._add_to_client_apps(client_receiver, clients=clients, id="client_receiver") +def add_final_global_evaluation( + recipe: Recipe, + participating_clients: Optional[List[str]] = None, + validation_timeout: int = 6000, +) -> None: + """Evaluate a PyTorch recipe's final global model on selected clients. + + Unlike full cross-site evaluation, this helper does not ask clients to + submit their local models. It locates the recipe's persisted global model + and sends only that model for validation after training. + + Args: + recipe: PyTorch recipe to augment with final global model evaluation. + participating_clients: Optional client names to run validation. If not + provided, all clients connected when the controller starts are used. + validation_timeout: Timeout in seconds for validation tasks. + + Raises: + ValueError: If the recipe is not PyTorch or has no model persistor. + RuntimeError: If a cross-site evaluation workflow is already configured. + """ + from nvflare.app_common.widgets.validation_json_generator import ValidationJsonGenerator + from nvflare.app_common.workflows.cross_site_model_eval import CrossSiteModelEval + from nvflare.app_opt.pt.file_model_locator import PTFileModelLocator + from nvflare.job_config.script_runner import FrameworkType + + if getattr(recipe, "_cse_added", False) or _has_cross_site_eval_workflow(recipe.job): + raise RuntimeError("a cross-site evaluation workflow is already configured for this recipe") + + if getattr(recipe, "framework", None) != FrameworkType.PYTORCH: + raise ValueError("final global evaluation currently supports PyTorch recipes only") + + comp_ids = getattr(recipe.job, "comp_ids", None) + if not isinstance(comp_ids, dict): + raise ValueError("final global evaluation requires a recipe that tracks component IDs") + + model_locator_id = comp_ids.get("locator_id", "") + if not model_locator_id: + persistor_id = comp_ids.get("persistor_id", "") + if not persistor_id: + raise ValueError("final global evaluation requires a PyTorch model persistor") + model_locator_id = recipe.job.to_server( + PTFileModelLocator(pt_persistor_id=persistor_id), id="final_model_locator" + ) + if not isinstance(model_locator_id, str) or not model_locator_id: + raise RuntimeError("failed to register the final global model locator") + comp_ids["locator_id"] = model_locator_id + + recipe.job.to_server(ValidationJsonGenerator()) + recipe.job.to_server( + CrossSiteModelEval( + model_locator_id=model_locator_id, + submit_model_task_name="", + validation_timeout=validation_timeout, + participating_clients=participating_clients, + ) + ) + recipe._cse_added = True + + def add_cross_site_evaluation( recipe: Recipe, submit_model_timeout: int = 600, diff --git a/research/auto-fl-research/tasks/cifar10/job.py b/research/auto-fl-research/tasks/cifar10/job.py index 8c19daefc9..57d9780ae1 100644 --- a/research/auto-fl-research/tasks/cifar10/job.py +++ b/research/auto-fl-research/tasks/cifar10/job.py @@ -50,11 +50,8 @@ ) from nvflare.apis.dxo import DataKind -from nvflare.app_common.widgets.validation_json_generator import ValidationJsonGenerator -from nvflare.app_common.workflows.cross_site_model_eval import CrossSiteModelEval -from nvflare.app_opt.pt.file_model_locator import PTFileModelLocator from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe -from nvflare.recipe import SimEnv, add_experiment_tracking +from nvflare.recipe import SimEnv, add_experiment_tracking, add_final_global_evaluation def define_parser(): @@ -212,26 +209,6 @@ def parse_final_eval_clients(client_spec: str, n_clients: int): return clients -def add_final_global_evaluation(recipe, participating_clients): - comp_ids = getattr(recipe.job, "comp_ids", {}) - model_locator_id = comp_ids.get("locator_id", "") - - if not model_locator_id: - persistor_id = comp_ids.get("persistor_id", "") - if not persistor_id: - raise ValueError("Final evaluation requires a PyTorch model persistor, but no persistor_id was found") - model_locator_id = recipe.job.to_server(PTFileModelLocator(pt_persistor_id=persistor_id)) - - recipe.job.to_server(ValidationJsonGenerator()) - recipe.job.to_server( - CrossSiteModelEval( - model_locator_id=model_locator_id, - submit_model_task_name="", - participating_clients=participating_clients, - ) - ) - - def get_aggregator(args): kind = args.aggregator if kind == "weighted": diff --git a/research/auto-fl-research/tasks/vlm_med/job.py b/research/auto-fl-research/tasks/vlm_med/job.py index 6fff51134a..f9c4efca2e 100644 --- a/research/auto-fl-research/tasks/vlm_med/job.py +++ b/research/auto-fl-research/tasks/vlm_med/job.py @@ -37,11 +37,8 @@ ) from nvflare.apis.dxo import DataKind -from nvflare.app_common.widgets.validation_json_generator import ValidationJsonGenerator -from nvflare.app_common.workflows.cross_site_model_eval import CrossSiteModelEval -from nvflare.app_opt.pt.file_model_locator import PTFileModelLocator from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe -from nvflare.recipe import SimEnv, add_experiment_tracking +from nvflare.recipe import SimEnv, add_experiment_tracking, add_final_global_evaluation def define_parser(): @@ -139,26 +136,6 @@ def parse_final_eval_clients(client_spec: str, n_clients: int): return clients -def add_final_global_evaluation(recipe, participating_clients): - comp_ids = getattr(recipe.job, "comp_ids", {}) - model_locator_id = comp_ids.get("locator_id", "") - - if not model_locator_id: - persistor_id = comp_ids.get("persistor_id", "") - if not persistor_id: - raise ValueError("Final evaluation requires a PyTorch model persistor, but no persistor_id was found") - model_locator_id = recipe.job.to_server(PTFileModelLocator(pt_persistor_id=persistor_id)) - - recipe.job.to_server(ValidationJsonGenerator()) - recipe.job.to_server( - CrossSiteModelEval( - model_locator_id=model_locator_id, - submit_model_task_name="", - participating_clients=participating_clients, - ) - ) - - def write_result_dir_sidecar(result_dir: str): sidecar_path = os.environ.get("AUTOFL_RESULT_DIR_FILE") if not sidecar_path: From 37af2cb3e5d2e02b1414fe4696dd72414e3b9f0e Mon Sep 17 00:00:00 2001 From: Kevin Date: Wed, 8 Jul 2026 21:03:41 -0400 Subject: [PATCH 4/8] Remove public recipe job access guidance --- docs/design/nvflare_cli.md | 8 ++----- .../data_scientist_guide/job_recipe.rst | 4 ++-- examples/advanced/recipe-k8s/job.py | 5 ++--- nvflare/app_opt/tf/tf_validator.py | 21 ++++--------------- nvflare/recipe/utils.py | 20 +++--------------- research/brats18/job.py | 2 +- 6 files changed, 14 insertions(+), 46 deletions(-) diff --git a/docs/design/nvflare_cli.md b/docs/design/nvflare_cli.md index cc34f48cf9..685c57f941 100644 --- a/docs/design/nvflare_cli.md +++ b/docs/design/nvflare_cli.md @@ -251,14 +251,10 @@ Current state: users have been using interactive console shell commands (`cat`, When client-side log streaming to the server is enabled, client logs are treated as server-side stored job logs. `nvflare job logs` does not connect to client sites directly and does not execute shell commands on client machines. It asks the server for the job log content that the server has locally for the requested site target. -Client log streaming is enabled by the job, not by `nvflare job logs`. The portable job-level pattern is: +Client log streaming is enabled by the job, not by `nvflare job logs`. For Recipe jobs, use: ```python -from nvflare.app_common.logging.job_log_receiver import JobLogReceiver -from nvflare.app_common.logging.job_log_streamer import JobLogStreamer - -recipe.job.to_clients(JobLogStreamer()) -recipe.job.to_server(JobLogReceiver()) +recipe.enable_log_streaming("log.txt") ``` `JobLogStreamer` runs in each client job process, tails the configured job log diff --git a/docs/user_guide/data_scientist_guide/job_recipe.rst b/docs/user_guide/data_scientist_guide/job_recipe.rst index 7c2b623666..06c0e9c82b 100644 --- a/docs/user_guide/data_scientist_guide/job_recipe.rst +++ b/docs/user_guide/data_scientist_guide/job_recipe.rst @@ -227,8 +227,8 @@ Recipe Metadata --------------- Use ``set_recipe_meta`` to add generated job metadata from a recipe without -mutating ``recipe.job.job.meta_props`` directly. The helper sets one -``JobMetaKey`` metadata entry at a time: +mutating nested generated-job metadata directly. The helper sets one ``JobMetaKey`` +metadata entry at a time: .. code-block:: python diff --git a/examples/advanced/recipe-k8s/job.py b/examples/advanced/recipe-k8s/job.py index 4337c16a5a..5596e26068 100644 --- a/examples/advanced/recipe-k8s/job.py +++ b/examples/advanced/recipe-k8s/job.py @@ -197,9 +197,8 @@ def create_recipe(args: argparse.Namespace) -> FedAvgRecipe: # The client training script and the server-side PyTorch persistor both # import the model definition, so bundle it into every generated app. model_path = str(Path(__file__).with_name("model.py")) - recipe.job.add_file_to_server(model_path) - for site_name in client_sites: - recipe.job.add_file_to(model_path, site_name) + recipe.add_server_file(model_path) + recipe.add_client_file(model_path, clients=list(client_sites)) # Scheduler-facing resource requirements stay separate from Kubernetes # container settings. K8sJobLauncher maps num_of_gpus to nvidia.com/gpu. diff --git a/nvflare/app_opt/tf/tf_validator.py b/nvflare/app_opt/tf/tf_validator.py index 5e2558ba8f..f9e6f577d7 100644 --- a/nvflare/app_opt/tf/tf_validator.py +++ b/nvflare/app_opt/tf/tf_validator.py @@ -55,31 +55,18 @@ def __init__( **Default recommendation**: Use Client API pattern (like PyTorch) for consistency and simplicity. TFValidator is provided for flexibility and backward compatibility with NumPy-style workflows. - Example (Component-based with TFValidator): + Example (Component-based with the lower-level Job API): ```python - from nvflare.app_opt.tf.recipes import FedAvgRecipe from nvflare.app_opt.tf.tf_validator import TFValidator - from nvflare.recipe.utils import add_cross_site_evaluation + from nvflare.job_config.api import FedJob - # Create recipe - recipe = FedAvgRecipe( - name="my-job", - min_clients=2, - num_rounds=3, - model=my_model, - train_script="client.py" - ) - - # Add CSE - add_cross_site_evaluation(recipe) - - # Manually add TFValidator if you prefer component-based validation + job = FedJob(name="my-job", min_clients=2) validator = TFValidator( model=my_model, data_loader=test_loader, metric_fn=lambda model, loader: {"accuracy": model.evaluate(loader)[1]} ) - recipe.job.to_clients(validator, tasks=["validate"]) + job.to_clients(validator, tasks=["validate"]) ``` Example (Client API pattern - recommended): diff --git a/nvflare/recipe/utils.py b/nvflare/recipe/utils.py index 8de122af26..98642217f6 100644 --- a/nvflare/recipe/utils.py +++ b/nvflare/recipe/utils.py @@ -446,23 +446,9 @@ def add_cross_site_evaluation( add_cross_site_evaluation(recipe) ``` - Example (TensorFlow - Component-based alternative): - ```python - from nvflare.app_opt.tf.recipes import FedAvgRecipe - from nvflare.app_opt.tf.tf_validator import TFValidator - from nvflare.recipe.utils import add_cross_site_evaluation - - recipe = FedAvgRecipe( - name="my-job", min_clients=2, num_rounds=3, - model=MyTFModel(), train_script="client.py" - ) - - add_cross_site_evaluation(recipe) - - # Optional: manually add TFValidator for component-based validation - validator = TFValidator(model=my_model, data_loader=test_loader) - recipe.job.to_clients(validator, tasks=["validate"]) - ``` + TensorFlow component-based validators are executors, not plain components. + Use the lower-level Job API when explicit ``TFValidator`` placement is required; + Recipe-based jobs should use the Client API pattern above. Args: recipe: Recipe instance to augment with cross-site evaluation. diff --git a/research/brats18/job.py b/research/brats18/job.py index 9b41faeaa0..ac56d3226a 100644 --- a/research/brats18/job.py +++ b/research/brats18/job.py @@ -107,7 +107,7 @@ def main(): # Use WEIGHT_DIFF for efficiency and consistency (required for SVTPrivacy filter) params_transfer_type=TransferType.DIFF, ) - recipe.job.to_server({"server": {"heart_beat_timeout": DEFAULT_HEARTBEAT_TIMEOUT}}) + recipe.add_server_config({"server": {"heart_beat_timeout": DEFAULT_HEARTBEAT_TIMEOUT}}) # Enable TensorBoard tracking add_experiment_tracking(recipe, tracking_type="tensorboard") From d81e03fef9925885cc082934488a57435ff264a0 Mon Sep 17 00:00:00 2001 From: Kevin Date: Thu, 9 Jul 2026 07:39:22 -0400 Subject: [PATCH 5/8] Address recipe API review feedback --- nvflare/recipe/spec.py | 16 +++- nvflare/recipe/utils.py | 12 ++- tests/unit_test/recipe/spec_test.py | 83 +++++++++++++++++++ tests/unit_test/recipe/utils_test.py | 117 +++++++++++++++++++++++++++ 4 files changed, 225 insertions(+), 3 deletions(-) diff --git a/nvflare/recipe/spec.py b/nvflare/recipe/spec.py index a535c0d93e..e7473bce23 100644 --- a/nvflare/recipe/spec.py +++ b/nvflare/recipe/spec.py @@ -565,12 +565,23 @@ def enable_tensor_streaming( for the server to wait for all clients to receive task tensors. Raises: - ValueError: If ``format`` does not match a declared - ``server_expected_format``. + TypeError: If ``tasks`` is not a list of strings. + ValueError: If ``tasks`` is empty, if ``format`` does not match a + declared ``server_expected_format``. + RuntimeError: If tensor streaming was already enabled. """ from nvflare.app_opt.tensor_stream.client import TensorClientStreamer from nvflare.app_opt.tensor_stream.server import TensorServerStreamer + if getattr(self, "_tensor_streaming_added", False): + raise RuntimeError("tensor streaming has already been enabled for this recipe") + + if tasks is not None: + if not isinstance(tasks, list) or not all(isinstance(task, str) for task in tasks): + raise TypeError(f"tasks must be a list of str, got {tasks!r}") + if not tasks: + raise ValueError("tasks must not be empty; use None for the default train task") + server_expected_format = getattr(self, "server_expected_format", None) if server_expected_format is not None and format != server_expected_format: raise ValueError( @@ -596,6 +607,7 @@ def enable_tensor_streaming( ), id="tensor_client_streamer", ) + self._tensor_streaming_added = True def add_decomposers(self, decomposers: List[Union[str, Decomposer]]): """Add decomposers to the job diff --git a/nvflare/recipe/utils.py b/nvflare/recipe/utils.py index 98642217f6..9edee49098 100644 --- a/nvflare/recipe/utils.py +++ b/nvflare/recipe/utils.py @@ -317,7 +317,9 @@ def add_final_global_evaluation( validation_timeout: Timeout in seconds for validation tasks. Raises: - ValueError: If the recipe is not PyTorch or has no model persistor. + TypeError: If ``participating_clients`` is not a list of strings. + ValueError: If ``participating_clients`` is empty, the recipe is not + PyTorch, or the recipe has no model persistor. RuntimeError: If a cross-site evaluation workflow is already configured. """ from nvflare.app_common.widgets.validation_json_generator import ValidationJsonGenerator @@ -331,6 +333,14 @@ def add_final_global_evaluation( if getattr(recipe, "framework", None) != FrameworkType.PYTORCH: raise ValueError("final global evaluation currently supports PyTorch recipes only") + if participating_clients is not None: + if not isinstance(participating_clients, list) or not all( + isinstance(client, str) for client in participating_clients + ): + raise TypeError(f"participating_clients must be a list of str, got {participating_clients!r}") + if not participating_clients: + raise ValueError("participating_clients must not be empty; use None to evaluate on all clients") + comp_ids = getattr(recipe.job, "comp_ids", None) if not isinstance(comp_ids, dict): raise ValueError("final global evaluation requires a recipe that tracks component IDs") diff --git a/tests/unit_test/recipe/spec_test.py b/tests/unit_test/recipe/spec_test.py index d8c1730858..c3758fb35f 100644 --- a/tests/unit_test/recipe/spec_test.py +++ b/tests/unit_test/recipe/spec_test.py @@ -406,6 +406,89 @@ def __init__(self, class_names): assert captured["server_obj"] is not captured["client_obj"] +class TestRecipeTensorStreaming: + def _make_recipe(self): + from nvflare.recipe.spec import Recipe + + return Recipe(FedJob(name="test_tensor_streaming", min_clients=1)) + + def test_enable_tensor_streaming_adds_matching_components(self): + from nvflare.apis.job_def import ALL_SITES + from nvflare.app_opt.tensor_stream.client import TensorClientStreamer + from nvflare.app_opt.tensor_stream.server import TensorServerStreamer + from nvflare.client.config import ExchangeFormat + + recipe = self._make_recipe() + recipe.server_expected_format = ExchangeFormat.PYTORCH + tasks = ["train", "validate"] + + recipe.enable_tensor_streaming( + format=ExchangeFormat.PYTORCH, + tasks=tasks, + tensor_send_timeout=45.0, + wait_send_task_data_all_clients_timeout=600.0, + ) + + server_streamer = recipe.job._deploy_map["server"].app_config.components["tensor_server_streamer"] + client_streamer = recipe.job._deploy_map[ALL_SITES].app_config.components["tensor_client_streamer"] + assert isinstance(server_streamer, TensorServerStreamer) + assert isinstance(client_streamer, TensorClientStreamer) + assert server_streamer.format == client_streamer.format == ExchangeFormat.PYTORCH + assert server_streamer.tasks == client_streamer.tasks == tasks + assert server_streamer.tasks is not tasks + assert client_streamer.tasks is not tasks + assert server_streamer.tensor_send_timeout == client_streamer.tensor_send_timeout == 45.0 + assert server_streamer.wait_task_data_sent_to_all_clients_timeout == 600.0 + assert recipe._tensor_streaming_added is True + + def test_enable_tensor_streaming_uses_default_train_task(self): + from nvflare.apis.job_def import ALL_SITES + + recipe = self._make_recipe() + recipe.enable_tensor_streaming() + + server_streamer = recipe.job._deploy_map["server"].app_config.components["tensor_server_streamer"] + client_streamer = recipe.job._deploy_map[ALL_SITES].app_config.components["tensor_client_streamer"] + assert server_streamer.tasks == client_streamer.tasks == ["train"] + + @pytest.mark.parametrize( + "tasks, error_type", + [ + ("train", TypeError), + (("train",), TypeError), + ([1], TypeError), + ([], ValueError), + ], + ) + def test_enable_tensor_streaming_validates_tasks(self, tasks, error_type): + recipe = self._make_recipe() + + with pytest.raises(error_type, match="tasks must"): + recipe.enable_tensor_streaming(tasks=tasks) + + assert recipe.job._deploy_map == {} + + def test_enable_tensor_streaming_rejects_mismatched_format(self): + from nvflare.client.config import ExchangeFormat + + recipe = self._make_recipe() + recipe.server_expected_format = ExchangeFormat.NUMPY + + with pytest.raises(ValueError, match="must match server_expected_format"): + recipe.enable_tensor_streaming(format=ExchangeFormat.PYTORCH) + + assert recipe.job._deploy_map == {} + + def test_enable_tensor_streaming_rejects_duplicate_call(self): + recipe = self._make_recipe() + recipe.enable_tensor_streaming() + + with pytest.raises(RuntimeError, match="already been enabled"): + recipe.enable_tensor_streaming() + + assert len(recipe.job._deploy_map["server"].app_config.components) == 1 + + class _DummyExecEnv: def __init__(self): self.extra = {} diff --git a/tests/unit_test/recipe/utils_test.py b/tests/unit_test/recipe/utils_test.py index 73ec1f8069..d61ffa8a00 100644 --- a/tests/unit_test/recipe/utils_test.py +++ b/tests/unit_test/recipe/utils_test.py @@ -17,6 +17,7 @@ import tempfile from datetime import datetime from decimal import Decimal +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -243,6 +244,11 @@ def test_add_cross_site_evaluation_importable_from_recipe(self): assert callable(add_cross_site_evaluation) + def test_add_final_global_evaluation_importable_from_recipe(self): + from nvflare.recipe import add_final_global_evaluation + + assert callable(add_final_global_evaluation) + def test_set_per_site_config_importable_from_recipe(self): """set_per_site_config must be importable from the top-level nvflare.recipe package.""" from nvflare.recipe import set_per_site_config @@ -525,6 +531,117 @@ def __init__(self, *args, **kwargs): assert captured_kwargs["participating_clients"] == participating_clients +class TestFinalGlobalEvaluation: + def _make_recipe(self, comp_ids=None, framework=None): + from nvflare.job_config.script_runner import FrameworkType + + job = MagicMock() + job._deploy_map = {} + job.comp_ids = comp_ids + return SimpleNamespace(job=job, framework=framework or FrameworkType.PYTORCH) + + def test_adds_locator_generator_and_final_eval_controller(self): + from nvflare.app_common.widgets.validation_json_generator import ValidationJsonGenerator + from nvflare.app_common.workflows.cross_site_model_eval import CrossSiteModelEval + from nvflare.app_opt.pt.file_model_locator import PTFileModelLocator + from nvflare.recipe import add_final_global_evaluation + + recipe = self._make_recipe({"persistor_id": "persistor"}) + recipe.job.to_server.side_effect = ["final_model_locator", None, None] + + add_final_global_evaluation(recipe, participating_clients=["site-1"], validation_timeout=42) + + calls = recipe.job.to_server.call_args_list + assert isinstance(calls[0].args[0], PTFileModelLocator) + assert calls[0].kwargs == {"id": "final_model_locator"} + assert isinstance(calls[1].args[0], ValidationJsonGenerator) + controller = calls[2].args[0] + assert isinstance(controller, CrossSiteModelEval) + assert controller._model_locator_id == "final_model_locator" + assert controller._submit_model_task_name == "" + assert controller._participating_clients == ["site-1"] + assert controller._validation_timeout == 42 + assert recipe.job.comp_ids["locator_id"] == "final_model_locator" + assert recipe._cse_added is True + + def test_reuses_existing_model_locator(self): + from nvflare.app_common.workflows.cross_site_model_eval import CrossSiteModelEval + from nvflare.recipe import add_final_global_evaluation + + recipe = self._make_recipe({"persistor_id": "persistor", "locator_id": "existing_locator"}) + + add_final_global_evaluation(recipe) + + assert recipe.job.to_server.call_count == 2 + controller = recipe.job.to_server.call_args_list[-1].args[0] + assert isinstance(controller, CrossSiteModelEval) + assert controller._model_locator_id == "existing_locator" + assert controller._participating_clients is None + + @pytest.mark.parametrize( + "participating_clients, error_type", + [ + ("site-1", TypeError), + (("site-1",), TypeError), + ([1], TypeError), + ([], ValueError), + ], + ) + def test_validates_participating_clients(self, participating_clients, error_type): + from nvflare.recipe import add_final_global_evaluation + + recipe = self._make_recipe({"persistor_id": "persistor"}) + + with pytest.raises(error_type, match="participating_clients must"): + add_final_global_evaluation(recipe, participating_clients=participating_clients) + + recipe.job.to_server.assert_not_called() + + def test_rejects_duplicate_configuration(self): + from nvflare.recipe import add_final_global_evaluation + + recipe = self._make_recipe({"persistor_id": "persistor"}) + recipe._cse_added = True + + with pytest.raises(RuntimeError, match="already configured"): + add_final_global_evaluation(recipe) + + def test_requires_pytorch_recipe(self): + from nvflare.job_config.script_runner import FrameworkType + from nvflare.recipe import add_final_global_evaluation + + recipe = self._make_recipe({"persistor_id": "persistor"}, framework=FrameworkType.NUMPY) + + with pytest.raises(ValueError, match="supports PyTorch"): + add_final_global_evaluation(recipe) + + @pytest.mark.parametrize("comp_ids", [None, [], "persistor"]) + def test_requires_component_id_mapping(self, comp_ids): + from nvflare.recipe import add_final_global_evaluation + + recipe = self._make_recipe(comp_ids) + + with pytest.raises(ValueError, match="tracks component IDs"): + add_final_global_evaluation(recipe) + + def test_requires_model_persistor(self): + from nvflare.recipe import add_final_global_evaluation + + recipe = self._make_recipe({}) + + with pytest.raises(ValueError, match="requires a PyTorch model persistor"): + add_final_global_evaluation(recipe) + + def test_rejects_failed_model_locator_registration(self): + from nvflare.recipe import add_final_global_evaluation + + recipe = self._make_recipe({"persistor_id": "persistor"}) + recipe.job.to_server.return_value = None + + with pytest.raises(RuntimeError, match="failed to register"): + add_final_global_evaluation(recipe) + + class TestAddExperimentTrackingClients: """Test client targeting and per-site configs in add_experiment_tracking.""" From 1245aefef5a771fdea9828505733ec0c5ffd3941 Mon Sep 17 00:00:00 2001 From: Kevin Date: Thu, 9 Jul 2026 16:47:39 -0400 Subject: [PATCH 6/8] Add recipe-derived MLflow defaults --- .../data_scientist_guide/recipe_api.rst | 5 +- .../mlflow/hello-pt-mlflow-client/README.md | 23 ++++---- .../mlflow/hello-pt-mlflow-client/job.py | 11 +--- nvflare/recipe/utils.py | 23 ++++++-- tests/unit_test/recipe/utils_test.py | 55 +++++++++++++++++++ 5 files changed, 92 insertions(+), 25 deletions(-) diff --git a/docs/user_guide/data_scientist_guide/recipe_api.rst b/docs/user_guide/data_scientist_guide/recipe_api.rst index 94b59b649b..880b017df5 100644 --- a/docs/user_guide/data_scientist_guide/recipe_api.rst +++ b/docs/user_guide/data_scientist_guide/recipe_api.rst @@ -143,7 +143,10 @@ Utility helpers: Add supported experiment tracking receivers such as TensorBoard, MLflow, or Weights & Biases. With ``client_side=True``, ``clients`` limits which sites receive the client-side receiver; call once per site with different - ``tracking_config`` values for per-site tracking destinations. + ``tracking_config`` values for per-site tracking destinations. For MLflow, + omitting ``tracking_config`` uses local storage and defaults the experiment + and run-name suffix to ``-experiment`` and + ``-Client``. ``add_cross_site_evaluation(recipe, submit_model_timeout=600, validation_timeout=6000, participating_clients=None)`` Add cross-site evaluation to a training recipe when the recipe/framework diff --git a/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/README.md b/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/README.md index f1abab7f47..8c3d2f40a3 100644 --- a/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/README.md +++ b/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/README.md @@ -36,18 +36,11 @@ recipe = FedAvgRecipe( train_script="client.py", ) -# Add an MLflow receiver to every client, but not to the server. -# With tracking_uri=None, each receiver creates a store in its local job workspace. +# Add an MLflow receiver to every client, but not to the server. With no +# tracking_config, each receiver creates a local store and uses recipe-derived names. add_experiment_tracking( recipe, "mlflow", - tracking_config={ - "tracking_uri": None, - "kw_args": { - "experiment_name": "nvflare-fedavg-experiment", - "run_name": "nvflare-fedavg-client", - }, - }, client_side=True, server_side=False, ) @@ -56,10 +49,20 @@ env = SimEnv(num_clients=2) run = recipe.execute(env) ``` +Only the placement flags are needed here because this example intentionally +tracks on clients. For default server-side MLflow tracking, the complete call is: + +```python +add_experiment_tracking(recipe, "mlflow") +``` + **Key points**: - `client_side=True, server_side=False` keeps metric handling on the clients. - The helper configures receivers to listen to local events rather than federated `fed.` events. -- `tracking_uri=None` resolves to a separate MLflow store in each client's job workspace. +- Omitting `tracking_config` uses `tracking_uri=None`, which resolves to a separate MLflow store in each client's job workspace. +- The default `experiment_name` and configured run-name suffix are + `fedavg_mlflow_client-experiment` and `fedavg_mlflow_client-Client`; the + receiver adds site and job identifiers to the final run name. - Sites can use explicit local paths or remote MLflow servers when their deployment requires them. ## Setup diff --git a/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/job.py b/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/job.py index cd168dddd9..9ea9165612 100644 --- a/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/job.py +++ b/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/job.py @@ -46,18 +46,11 @@ def define_parser(): train_script="client.py", ) - # Add MLflow tracking to all clients (client-side only, no server aggregation) - # Each client will track its own local metrics independently + # Add MLflow tracking to all clients (client-side only, no server aggregation). + # With no tracking_config, each client uses a local store and recipe-derived names. add_experiment_tracking( recipe, "mlflow", - tracking_config={ - "tracking_uri": None, # Will auto-configure per site - "kw_args": { - "experiment_name": "nvflare-fedavg-experiment", - "run_name": "nvflare-fedavg-client", - }, - }, client_side=True, server_side=False, ) diff --git a/nvflare/recipe/utils.py b/nvflare/recipe/utils.py index d14515f9f3..0b886eb54c 100644 --- a/nvflare/recipe/utils.py +++ b/nvflare/recipe/utils.py @@ -238,7 +238,9 @@ def add_experiment_tracking( Args: recipe: Recipe instance to augment with experiment tracking. tracking_type: Type of tracking to enable ("mlflow", "tensorboard", or "wandb"). - tracking_config: Optional configuration dict for the tracking receiver. + tracking_config: Optional configuration dict for the tracking receiver. For MLflow, + omitting this uses a local file store and derives ``experiment_name`` and + ``run_name`` from the recipe name. client_side: If True, add tracking to clients (each client tracks locally). server_side: If True, add tracking to server (aggregates metrics from all clients). Default: True. clients: Optional list of client names for client-side tracking. If None, the @@ -251,11 +253,11 @@ def add_experiment_tracking( all-clients topology or unknown site names, targeted placement raises ValueError. Examples: - # Server-side tracking (default - federated metrics) - add_experiment_tracking(recipe, "mlflow", {"tracking_uri": "..."}) + # Server-side MLflow tracking with local storage and recipe-derived names + add_experiment_tracking(recipe, "mlflow") # Client-side tracking only (each client tracks independently) - add_experiment_tracking(recipe, "mlflow", {...}, client_side=True, server_side=False) + add_experiment_tracking(recipe, "mlflow", client_side=True, server_side=False) # Both server and client tracking add_experiment_tracking(recipe, "mlflow", {...}, client_side=True, server_side=True) @@ -270,10 +272,21 @@ def add_experiment_tracking( client_side=True, server_side=False, clients=["site-2"], ) """ - tracking_config = tracking_config or {} if tracking_type not in TRACKING_REGISTRY: raise ValueError(f"Invalid tracking type: {tracking_type}") + tracking_config = copy.deepcopy(tracking_config) if tracking_config else {} + if tracking_type == "mlflow": + kw_args = tracking_config.get("kw_args") + if kw_args is None: + kw_args = {} + tracking_config["kw_args"] = kw_args + elif not isinstance(kw_args, dict): + raise TypeError(f"MLflow kw_args must be a dict, got {type(kw_args).__name__}") + recipe_name = getattr(recipe, "name", None) or getattr(recipe.job, "name", None) or "nvflare" + kw_args.setdefault("experiment_name", f"{recipe_name}-experiment") + kw_args.setdefault("run_name", f"{recipe_name}-Client") + if not server_side and not client_side: raise ValueError("At least one of server_side or client_side must be True") diff --git a/tests/unit_test/recipe/utils_test.py b/tests/unit_test/recipe/utils_test.py index d61ffa8a00..eebebdbbe1 100644 --- a/tests/unit_test/recipe/utils_test.py +++ b/tests/unit_test/recipe/utils_test.py @@ -662,6 +662,61 @@ def dummy_tracking(self, monkeypatch): def _make_recipe(self, name="test_tracking_clients"): return Recipe(FedJob(name=name, min_clients=1)) + @pytest.fixture + def dummy_mlflow(self, monkeypatch): + """Replace MLflow with a dependency-free receiver while retaining MLflow defaults.""" + import nvflare.recipe.utils as utils_mod + + monkeypatch.setitem( + utils_mod.TRACKING_REGISTRY, + "mlflow", + {"package": "json", "receiver_module": "argparse", "receiver_class": "Namespace"}, + ) + + def test_mlflow_defaults_derive_from_recipe_name(self, dummy_mlflow): + from nvflare.recipe.utils import add_experiment_tracking + + recipe = self._make_recipe("named_job") + + add_experiment_tracking(recipe, "mlflow") + + receiver = recipe.job._deploy_map["server"].app_config.components["receiver"] + assert receiver.kw_args == { + "experiment_name": "named_job-experiment", + "run_name": "named_job-Client", + } + + def test_mlflow_client_tracking_can_omit_config(self, dummy_mlflow): + from nvflare.apis.analytix import ANALYTIC_EVENT_TYPE + from nvflare.apis.job_def import ALL_SITES + from nvflare.recipe.utils import add_experiment_tracking + + recipe = self._make_recipe("client_job") + + add_experiment_tracking(recipe, "mlflow", client_side=True, server_side=False) + + receiver = recipe.job._deploy_map[ALL_SITES].app_config.components["client_receiver"] + assert receiver.kw_args == { + "experiment_name": "client_job-experiment", + "run_name": "client_job-Client", + } + assert receiver.events == [ANALYTIC_EVENT_TYPE] + + def test_mlflow_defaults_preserve_explicit_values_and_input(self, dummy_mlflow): + from nvflare.recipe.utils import add_experiment_tracking + + recipe = self._make_recipe("named_job") + config = {"kw_args": {"experiment_name": "custom-experiment"}} + + add_experiment_tracking(recipe, "mlflow", config) + + receiver = recipe.job._deploy_map["server"].app_config.components["receiver"] + assert receiver.kw_args == { + "experiment_name": "custom-experiment", + "run_name": "named_job-Client", + } + assert config == {"kw_args": {"experiment_name": "custom-experiment"}} + def test_client_side_tracking_specific_clients(self, dummy_tracking): from nvflare.apis.analytix import ANALYTIC_EVENT_TYPE from nvflare.apis.job_def import ALL_SITES From 823705f12d708923766977de0bd1258cb2f1289e Mon Sep 17 00:00:00 2001 From: Kevin Date: Fri, 10 Jul 2026 08:09:35 -0400 Subject: [PATCH 7/8] Use SQLite for default MLflow tracking --- .../mlflow/hello-pt-mlflow-client/README.md | 14 ++-- .../tracking/mlflow/mlflow_receiver.py | 9 +- .../slow/experiment_tracking_recipes_test.py | 84 +++++++++++++++---- .../app_opt/tracking/mlflow_receiver_test.py | 31 +++++++ 4 files changed, 112 insertions(+), 26 deletions(-) create mode 100644 tests/unit_test/app_opt/tracking/mlflow_receiver_test.py diff --git a/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/README.md b/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/README.md index 8c3d2f40a3..ec511786ee 100644 --- a/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/README.md +++ b/examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/README.md @@ -94,22 +94,22 @@ python job.py Since each site has its own MLflow receiver, metrics are stored separately: -With the example's `tracking_uri=None`, the store is created under each site's -job-result directory as `//mlflow`. Under the default -simulation workspace, use `find` to locate the generated directory: +With the example's `tracking_uri=None`, a SQLite store is created under each +site's job-result directory as `//mlflow.db`. Under the +default simulation workspace, use `find` to locate the generated database: ### View Site-1 Metrics: ```bash -find /tmp/nvflare/jobs/workdir/fedavg_mlflow_client/site-1 -type d -name mlflow -mlflow ui --backend-store-uri +find /tmp/nvflare/jobs/workdir/fedavg_mlflow_client/site-1 -type f -name mlflow.db +mlflow ui --backend-store-uri sqlite:////absolute/path/to/site-1/mlflow.db ``` Open browser to `http://localhost:5000` ### View Site-2 Metrics: ```bash -find /tmp/nvflare/jobs/workdir/fedavg_mlflow_client/site-2 -type d -name mlflow -mlflow ui --backend-store-uri --port 5001 +find /tmp/nvflare/jobs/workdir/fedavg_mlflow_client/site-2 -type f -name mlflow.db +mlflow ui --backend-store-uri sqlite:////absolute/path/to/site-2/mlflow.db --port 5001 ``` Open browser to `http://localhost:5001` diff --git a/nvflare/app_opt/tracking/mlflow/mlflow_receiver.py b/nvflare/app_opt/tracking/mlflow/mlflow_receiver.py index f54c2dc179..aad3de21e6 100644 --- a/nvflare/app_opt/tracking/mlflow/mlflow_receiver.py +++ b/nvflare/app_opt/tracking/mlflow/mlflow_receiver.py @@ -62,8 +62,9 @@ def __init__( """MLflowReceiver receives log events from clients and deliver them to the MLflow tracking server. Args: - tracking_uri (Optional[str], optional): MLflow tracking server URI. When this is not specified, the metrics will be written to the local file system. - If the tracking URI is specified, the MLflow tracking server must started before running the job. Defaults to None. + tracking_uri (Optional[str], optional): MLflow tracking server URI. When this is not specified, metrics + are written to a local SQLite database in the job result directory. If the tracking URI is specified, + the MLflow tracking server must be started before running the job. Defaults to None. kw_args (Optional[dict], optional): keyword arguments: "experiment_name" (str): Specifies the experiment name. If not specified, the default name of "FLARE FL Experiment" will be used. "run_name" (str): Specifies the run name @@ -105,8 +106,8 @@ def _get_tracking_uri(self, fl_ctx: FLContext): workspace = fl_ctx.get_workspace() job_id = fl_ctx.get_job_id() - mlflow_root = os.path.abspath(os.path.join(workspace.get_result_root(job_id), "mlflow")) - return f"file://{mlflow_root}" + mlflow_db = os.path.abspath(os.path.join(workspace.get_result_root(job_id), "mlflow.db")) + return f"sqlite:///{mlflow_db}" def initialize(self, fl_ctx: FLContext): """Initializes MlflowClient for each site. diff --git a/tests/integration_test/slow/experiment_tracking_recipes_test.py b/tests/integration_test/slow/experiment_tracking_recipes_test.py index a08ecd281f..fbcea3ea92 100644 --- a/tests/integration_test/slow/experiment_tracking_recipes_test.py +++ b/tests/integration_test/slow/experiment_tracking_recipes_test.py @@ -74,10 +74,22 @@ def client_script_dir(self): def model_path(self): return os.path.join(self.client_script_dir, "model.py") - def _add_model_to_apps(self, recipe): - recipe.job.add_file_to_server(self.model_path) + @property + def mlflow_client_script_path(self): + return os.path.join( + REPO_ROOT, + "examples/advanced/experiment-tracking/mlflow/hello-pt-mlflow-client/client.py", + ) + + @property + def mlflow_client_model_path(self): + return os.path.join(os.path.dirname(self.mlflow_client_script_path), "model.py") + + def _add_model_to_apps(self, recipe, model_path=None): + model_path = model_path or self.model_path + recipe.job.add_file_to_server(model_path) for site_name in ("site-1", "site-2"): - recipe.job.add_file_to(self.model_path, site_name) + recipe.job.add_file_to(model_path, site_name) def _assert_exported_client_configs_have_executors(self, recipe, export_root: str): recipe.export(export_root) @@ -155,23 +167,65 @@ def test_mlflow_tracking_integration(self, cifar10_data_root): min_clients=2, num_rounds=1, model={"class_path": "model.SimpleNetwork", "args": {}}, - train_script=self.client_script_path, + train_script=self.mlflow_client_script_path, per_site_config=self._per_site_config(cifar10_data_root, tmpdir), ) - self._add_model_to_apps(recipe) + self._add_model_to_apps(recipe, self.mlflow_client_model_path) + + # Exercise the zero-config MLflow path: local storage and names derived from the recipe. + add_experiment_tracking(recipe, "mlflow") + + # Run and verify completion plus the actual MLflow experiment and runs. + run = recipe.execute(env) + assert run.get_result() is not None + assert os.path.exists(run.get_result()) + + mlflow_db = os.path.join(run.get_result(), "server", "simulate_job", "mlflow.db") + assert os.path.isfile(mlflow_db) + + from mlflow.tracking import MlflowClient + + client = MlflowClient(tracking_uri=f"sqlite:///{mlflow_db}") + experiment = client.get_experiment_by_name("test_mlflow-experiment") + assert experiment is not None - # Add MLflow tracking - mlflow_uri = os.path.join(mlflow_dir, "mlruns") - add_experiment_tracking( - recipe, - "mlflow", - tracking_config={ - "tracking_uri": f"file:///{mlflow_uri}", - "kw_args": {"experiment_name": "test", "run_name": "test"}, - }, + runs = client.search_runs([experiment.experiment_id]) + assert len(runs) == 2 + assert all(run.info.run_name.endswith("test_mlflow-Client") for run in runs) + + def test_mlflow_client_tracking_defaults_integration(self, cifar10_data_root): + """Test zero-config client-side MLflow tracking creates one local store per site.""" + import tempfile + + with tempfile.TemporaryDirectory() as tmpdir: + workspace_root = os.path.join(tmpdir, "test_mlflow_client") + env = SimEnv(clients=["site-1", "site-2"], workspace_root=workspace_root) + recipe = FedAvgRecipe( + name="test_mlflow_client", + min_clients=2, + num_rounds=1, + model={"class_path": "model.SimpleNetwork", "args": {}}, + train_script=self.mlflow_client_script_path, + per_site_config=self._per_site_config(cifar10_data_root, tmpdir), ) + self._add_model_to_apps(recipe, self.mlflow_client_model_path) + + add_experiment_tracking(recipe, "mlflow", client_side=True, server_side=False) - # Run and verify completion run = recipe.execute(env) assert run.get_result() is not None assert os.path.exists(run.get_result()) + + from mlflow.tracking import MlflowClient + + for site_name in ("site-1", "site-2"): + mlflow_db = os.path.join(run.get_result(), site_name, "simulate_job", "mlflow.db") + assert os.path.isfile(mlflow_db) + + client = MlflowClient(tracking_uri=f"sqlite:///{mlflow_db}") + experiment = client.get_experiment_by_name("test_mlflow_client-experiment") + assert experiment is not None + + runs = client.search_runs([experiment.experiment_id]) + assert len(runs) == 1 + assert runs[0].info.run_name.endswith("test_mlflow_client-Client") diff --git a/tests/unit_test/app_opt/tracking/mlflow_receiver_test.py b/tests/unit_test/app_opt/tracking/mlflow_receiver_test.py new file mode 100644 index 0000000000..68867b2f50 --- /dev/null +++ b/tests/unit_test/app_opt/tracking/mlflow_receiver_test.py @@ -0,0 +1,31 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock + +from nvflare.app_opt.tracking.mlflow.mlflow_receiver import MLflowReceiver + + +def test_default_tracking_uri_uses_local_sqlite(tmp_path): + result_root = tmp_path / "simulate_job" + workspace = MagicMock() + workspace.get_result_root.return_value = str(result_root) + fl_ctx = MagicMock() + fl_ctx.get_workspace.return_value = workspace + fl_ctx.get_job_id.return_value = "simulate_job" + + receiver = MLflowReceiver() + + assert receiver._get_tracking_uri(fl_ctx) == f"sqlite:///{result_root / 'mlflow.db'}" + workspace.get_result_root.assert_called_once_with("simulate_job") From 5c8fd8266d03335ed04415b24de6b5e2ad21f701 Mon Sep 17 00:00:00 2001 From: Kevin Date: Fri, 10 Jul 2026 11:35:32 -0400 Subject: [PATCH 8/8] Address final recipe API review feedback --- .../data_scientist_guide/recipe_api.rst | 5 +- nvflare/recipe/spec.py | 2 + nvflare/recipe/utils.py | 11 ++- .../slow/experiment_tracking_recipes_test.py | 2 +- tests/unit_test/recipe/utils_test.py | 76 +++++++++++++++---- 5 files changed, 76 insertions(+), 20 deletions(-) diff --git a/docs/user_guide/data_scientist_guide/recipe_api.rst b/docs/user_guide/data_scientist_guide/recipe_api.rst index 880b017df5..1638b0a0c2 100644 --- a/docs/user_guide/data_scientist_guide/recipe_api.rst +++ b/docs/user_guide/data_scientist_guide/recipe_api.rst @@ -145,8 +145,9 @@ Utility helpers: receive the client-side receiver; call once per site with different ``tracking_config`` values for per-site tracking destinations. For MLflow, omitting ``tracking_config`` uses local storage and defaults the experiment - and run-name suffix to ``-experiment`` and - ``-Client``. + name to ``-experiment``. The run-name suffix defaults to + ``-Server`` for server-side tracking and + ``-Client`` for client-side tracking. ``add_cross_site_evaluation(recipe, submit_model_timeout=600, validation_timeout=6000, participating_clients=None)`` Add cross-site evaluation to a training recipe when the recipe/framework diff --git a/nvflare/recipe/spec.py b/nvflare/recipe/spec.py index e7473bce23..4359252d95 100644 --- a/nvflare/recipe/spec.py +++ b/nvflare/recipe/spec.py @@ -192,6 +192,8 @@ def __init__(self, job: FedJob): """ self.job = job self._helper_per_site_config = None + self._tensor_streaming_added = False + self._cse_added = False def process_env(self, env: ExecEnv): """Process environment-specific configuration. diff --git a/nvflare/recipe/utils.py b/nvflare/recipe/utils.py index 82b284ef5b..bf9cb17426 100644 --- a/nvflare/recipe/utils.py +++ b/nvflare/recipe/utils.py @@ -368,7 +368,6 @@ def add_experiment_tracking( raise TypeError(f"MLflow kw_args must be a dict, got {type(kw_args).__name__}") recipe_name = getattr(recipe, "name", None) or getattr(recipe.job, "name", None) or "nvflare" kw_args.setdefault("experiment_name", f"{recipe_name}-experiment") - kw_args.setdefault("run_name", f"{recipe_name}-Client") if not server_side and not client_side: raise ValueError("At least one of server_side or client_side must be True") @@ -392,7 +391,10 @@ def add_experiment_tracking( # Add server-side tracking if server_side: - receiver = receiver_class(**tracking_config) + server_config = copy.deepcopy(tracking_config) + if tracking_type == "mlflow": + server_config["kw_args"].setdefault("run_name", f"{recipe_name}-Server") + receiver = receiver_class(**server_config) recipe.job.to_server(receiver, "receiver") # Add client-side tracking @@ -400,6 +402,8 @@ def add_experiment_tracking( # For client-side tracking, need to configure local events # Deep copy to avoid shared mutable state (tracking_config may contain nested dicts) client_config = copy.deepcopy(tracking_config) + if tracking_type == "mlflow": + client_config["kw_args"].setdefault("run_name", f"{recipe_name}-Client") # Override events to track local analytics (not federated) if "events" not in client_config: client_config["events"] = [ANALYTIC_EVENT_TYPE] @@ -425,7 +429,8 @@ def add_final_global_evaluation( recipe: PyTorch recipe to augment with final global model evaluation. participating_clients: Optional client names to run validation. If not provided, all clients connected when the controller starts are used. - validation_timeout: Timeout in seconds for validation tasks. + validation_timeout: Timeout in seconds for validation tasks. Defaults to + 6000, matching ``CrossSiteModelEval``'s existing default. Raises: TypeError: If ``participating_clients`` is not a list of strings. diff --git a/tests/integration_test/slow/experiment_tracking_recipes_test.py b/tests/integration_test/slow/experiment_tracking_recipes_test.py index fbcea3ea92..172dda3f7d 100644 --- a/tests/integration_test/slow/experiment_tracking_recipes_test.py +++ b/tests/integration_test/slow/experiment_tracking_recipes_test.py @@ -191,7 +191,7 @@ def test_mlflow_tracking_integration(self, cifar10_data_root): runs = client.search_runs([experiment.experiment_id]) assert len(runs) == 2 - assert all(run.info.run_name.endswith("test_mlflow-Client") for run in runs) + assert all(run.info.run_name.endswith("test_mlflow-Server") for run in runs) def test_mlflow_client_tracking_defaults_integration(self, cifar10_data_root): """Test zero-config client-side MLflow tracking creates one local store per site.""" diff --git a/tests/unit_test/recipe/utils_test.py b/tests/unit_test/recipe/utils_test.py index 5638ea9fc4..97bbdbdc2d 100644 --- a/tests/unit_test/recipe/utils_test.py +++ b/tests/unit_test/recipe/utils_test.py @@ -601,12 +601,37 @@ def _make_recipe(self, comp_ids=None, framework=None): job.comp_ids = comp_ids return SimpleNamespace(job=job, framework=framework or FrameworkType.PYTORCH) - def test_adds_locator_generator_and_final_eval_controller(self): - from nvflare.app_common.widgets.validation_json_generator import ValidationJsonGenerator + @pytest.fixture + def recording_cross_site_eval(self, monkeypatch): + from nvflare.app_common.workflows import cross_site_model_eval + from nvflare.app_common.workflows.cross_site_model_eval import CrossSiteModelEval + + captured_kwargs = [] + + class RecordingCrossSiteModelEval(CrossSiteModelEval): + def __init__(self, *args, **kwargs): + captured_kwargs.append(dict(kwargs)) + super().__init__(*args, **kwargs) + + monkeypatch.setattr(cross_site_model_eval, "CrossSiteModelEval", RecordingCrossSiteModelEval) + return captured_kwargs, RecordingCrossSiteModelEval + + def test_default_validation_timeout_matches_controller(self): + import inspect + from nvflare.app_common.workflows.cross_site_model_eval import CrossSiteModelEval + from nvflare.recipe import add_final_global_evaluation + + helper_default = inspect.signature(add_final_global_evaluation).parameters["validation_timeout"].default + controller_default = inspect.signature(CrossSiteModelEval).parameters["validation_timeout"].default + assert helper_default == controller_default == 6000 + + def test_adds_locator_generator_and_final_eval_controller(self, recording_cross_site_eval): + from nvflare.app_common.widgets.validation_json_generator import ValidationJsonGenerator from nvflare.app_opt.pt.file_model_locator import PTFileModelLocator from nvflare.recipe import add_final_global_evaluation + captured_kwargs, controller_type = recording_cross_site_eval recipe = self._make_recipe({"persistor_id": "persistor"}) recipe.job.to_server.side_effect = ["final_model_locator", None, None] @@ -617,27 +642,37 @@ def test_adds_locator_generator_and_final_eval_controller(self): assert calls[0].kwargs == {"id": "final_model_locator"} assert isinstance(calls[1].args[0], ValidationJsonGenerator) controller = calls[2].args[0] - assert isinstance(controller, CrossSiteModelEval) - assert controller._model_locator_id == "final_model_locator" - assert controller._submit_model_task_name == "" - assert controller._participating_clients == ["site-1"] - assert controller._validation_timeout == 42 + assert isinstance(controller, controller_type) + assert captured_kwargs == [ + { + "model_locator_id": "final_model_locator", + "submit_model_task_name": "", + "validation_timeout": 42, + "participating_clients": ["site-1"], + } + ] assert recipe.job.comp_ids["locator_id"] == "final_model_locator" assert recipe._cse_added is True - def test_reuses_existing_model_locator(self): - from nvflare.app_common.workflows.cross_site_model_eval import CrossSiteModelEval + def test_reuses_existing_model_locator(self, recording_cross_site_eval): from nvflare.recipe import add_final_global_evaluation + captured_kwargs, controller_type = recording_cross_site_eval recipe = self._make_recipe({"persistor_id": "persistor", "locator_id": "existing_locator"}) add_final_global_evaluation(recipe) assert recipe.job.to_server.call_count == 2 controller = recipe.job.to_server.call_args_list[-1].args[0] - assert isinstance(controller, CrossSiteModelEval) - assert controller._model_locator_id == "existing_locator" - assert controller._participating_clients is None + assert isinstance(controller, controller_type) + assert captured_kwargs == [ + { + "model_locator_id": "existing_locator", + "submit_model_task_name": "", + "validation_timeout": 6000, + "participating_clients": None, + } + ] @pytest.mark.parametrize( "participating_clients, error_type", @@ -744,7 +779,7 @@ def test_mlflow_defaults_derive_from_recipe_name(self, dummy_mlflow): receiver = recipe.job._deploy_map["server"].app_config.components["receiver"] assert receiver.kw_args == { "experiment_name": "named_job-experiment", - "run_name": "named_job-Client", + "run_name": "named_job-Server", } def test_mlflow_client_tracking_can_omit_config(self, dummy_mlflow): @@ -774,10 +809,23 @@ def test_mlflow_defaults_preserve_explicit_values_and_input(self, dummy_mlflow): receiver = recipe.job._deploy_map["server"].app_config.components["receiver"] assert receiver.kw_args == { "experiment_name": "custom-experiment", - "run_name": "named_job-Client", + "run_name": "named_job-Server", } assert config == {"kw_args": {"experiment_name": "custom-experiment"}} + def test_mlflow_both_sides_use_side_specific_default_run_names(self, dummy_mlflow): + from nvflare.apis.job_def import ALL_SITES + from nvflare.recipe.utils import add_experiment_tracking + + recipe = self._make_recipe("both_sides") + + add_experiment_tracking(recipe, "mlflow", client_side=True) + + server_receiver = recipe.job._deploy_map["server"].app_config.components["receiver"] + client_receiver = recipe.job._deploy_map[ALL_SITES].app_config.components["client_receiver"] + assert server_receiver.kw_args["run_name"] == "both_sides-Server" + assert client_receiver.kw_args["run_name"] == "both_sides-Client" + def test_client_side_tracking_specific_clients(self, dummy_tracking): from nvflare.apis.analytix import ANALYTIC_EVENT_TYPE from nvflare.apis.job_def import ALL_SITES