diff --git a/docs/design/nvflare_cli.md b/docs/design/nvflare_cli.md index c72d58a67e..7c34618bd7 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 50734a5145..c420e8afa9 100644 --- a/docs/user_guide/data_scientist_guide/job_recipe.rst +++ b/docs/user_guide/data_scientist_guide/job_recipe.rst @@ -244,8 +244,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/docs/user_guide/data_scientist_guide/recipe_api.rst b/docs/user_guide/data_scientist_guide/recipe_api.rst index f1d02e0f3e..3726cb6211 100644 --- a/docs/user_guide/data_scientist_guide/recipe_api.rst +++ b/docs/user_guide/data_scientist_guide/recipe_api.rst @@ -138,20 +138,32 @@ 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)`` 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. Keep tracking - credentials in the executing site's environment or mounted secret files; - never put them in ``tracking_config``. + ``tracking_config`` values for per-site tracking destinations. For MLflow, + omitting ``tracking_config`` uses local storage and defaults the experiment + name to ``-experiment``. The run-name suffix defaults to + ``-Server`` for server-side tracking and + ``-Client`` for client-side tracking. Keep tracking credentials + in the executing site's environment or mounted secret files; never put them + in ``tracking_config``. ``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 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/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..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 @@ -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,34 @@ 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 no +# tracking_config, each receiver creates a local store and uses recipe-derived names. +add_experiment_tracking( + recipe, + "mlflow", + client_side=True, + server_side=False, +) 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**: -- 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. +- 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 @@ -87,16 +94,22 @@ python job.py Since each site has its own MLflow receiver, metrics are stored separately: +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 -mlflow ui --backend-store-uri /tmp/nvflare/jobs/workdir/site-1/mlruns +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 -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 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` @@ -166,29 +179,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/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/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**: 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/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/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/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/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/nvflare/recipe/__init__.py b/nvflare/recipe/__init__.py index c2bed56623..58387ba039 100644 --- a/nvflare/recipe/__init__.py +++ b/nvflare/recipe/__init__.py @@ -19,7 +19,13 @@ from .run import Run from .secrets import PotentialSecretWarning, UnsupportedSecretRefWarning, secret_file_ref, secret_ref 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", @@ -28,6 +34,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/spec.py b/nvflare/recipe/spec.py index 43c94d9a43..a493c49a91 100644 --- a/nvflare/recipe/spec.py +++ b/nvflare/recipe/spec.py @@ -253,6 +253,8 @@ def __init__(self, job: FedJob): """ self.job = job self._helper_per_site_config = None + self._tensor_streaming_added = False + self._cse_added = False warn_on_potential_secrets(getattr(job, "name", None), context="recipe parameter 'name'") def _warn_potential_secrets_in_params(self): @@ -662,6 +664,73 @@ 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: + 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( + 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", + ) + 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 24db6cc35f..0bb71e28d4 100644 --- a/nvflare/recipe/utils.py +++ b/nvflare/recipe/utils.py @@ -327,10 +327,12 @@ 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. It becomes - part of the generated job definition and must never contain actual credentials; - configure authentication through the executing site's environment or a mounted - secret instead. + 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. The configuration becomes part of the + generated job definition and must never contain actual credentials; configure + authentication through the executing site's environment or a mounted secret + instead. 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 @@ -343,11 +345,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) @@ -368,6 +370,17 @@ def add_experiment_tracking( 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") + if not server_side and not client_side: raise ValueError("At least one of server_side or client_side must be True") @@ -390,7 +403,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 @@ -398,6 +414,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] @@ -408,6 +426,77 @@ 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. Defaults to + 6000, matching ``CrossSiteModelEval``'s existing default. + + Raises: + 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 + 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") + + 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") + + 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, @@ -495,23 +584,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/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: 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") diff --git a/tests/integration_test/slow/experiment_tracking_recipes_test.py b/tests/integration_test/slow/experiment_tracking_recipes_test.py index a08ecd281f..172dda3f7d 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-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.""" + 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") 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 f1ff246d45..f412446316 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 @@ -306,6 +307,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 @@ -595,6 +601,152 @@ 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) + + @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] + + 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, 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, 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, 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", + [ + ("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.""" @@ -615,6 +767,74 @@ 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-Server", + } + + 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-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_tracking_config_warns_for_secret_and_unsupported_ref(self, dummy_tracking): from nvflare.recipe.utils import add_experiment_tracking