From 39da42b466c61833486106d53ef654f932ae94ef Mon Sep 17 00:00:00 2001
From: YuanTingHsieh <yuantingh@nvidia.com>
Date: Thu, 25 Jun 2026 13:41:34 -0700
Subject: [PATCH 1/5] Add torchrun CPU Client API integration test

---
 .../app/config/config_fed_client.conf         | 48 +++++++++++++
 .../app/config/config_fed_server.conf         | 53 ++++++++++++++
 .../app/custom/net.py                         | 24 +++++++
 .../app/custom/torchrun_client.py             | 71 +++++++++++++++++++
 .../jobs/pt_client_api_torchrun_cpu/meta.conf | 11 +++
 .../standalone_job/client_api.yml             | 15 ++++
 6 files changed, 222 insertions(+)
 create mode 100644 tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/config/config_fed_client.conf
 create mode 100644 tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/config/config_fed_server.conf
 create mode 100644 tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/net.py
 create mode 100644 tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py
 create mode 100644 tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/meta.conf

diff --git a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/config/config_fed_client.conf b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/config/config_fed_client.conf
new file mode 100644
index 0000000000..fd72615989
--- /dev/null
+++ b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/config/config_fed_client.conf
@@ -0,0 +1,48 @@
+{
+  format_version = 2
+  app_script = "torchrun_client.py"
+  app_config = ""
+  executors = [
+    {
+      tasks = [
+        "train"
+      ]
+      executor {
+        path = "nvflare.app_opt.pt.client_api_launcher_executor.PTClientAPILauncherExecutor"
+        args {
+          launcher_id = "launcher"
+          pipe_id = "pipe"
+          heartbeat_timeout = 60
+          params_exchange_format = "pytorch"
+          params_transfer_type = "FULL"
+          train_with_evaluation = false
+        }
+      }
+    }
+  ]
+  task_data_filters = []
+  task_result_filters = []
+  components = [
+    {
+      id = "launcher"
+      path = "nvflare.app_common.launchers.subprocess_launcher.SubprocessLauncher"
+      args {
+        script = "python3 -m torch.distributed.run --standalone --nproc_per_node=2 custom/{app_script} {app_config}"
+        launch_once = false
+        shutdown_timeout = 30
+      }
+    }
+    {
+      id = "pipe"
+      path = "nvflare.fuel.utils.pipe.cell_pipe.CellPipe"
+      args {
+        mode = "PASSIVE"
+        site_name = "{SITE_NAME}"
+        token = "{JOB_ID}"
+        root_url = "{ROOT_URL}"
+        secure_mode = "{SECURE_MODE}"
+        workspace_dir = "{WORKSPACE}"
+      }
+    }
+  ]
+}
diff --git a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/config/config_fed_server.conf b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/config/config_fed_server.conf
new file mode 100644
index 0000000000..5764045c5c
--- /dev/null
+++ b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/config/config_fed_server.conf
@@ -0,0 +1,53 @@
+{
+  format_version = 2
+  task_data_filters = []
+  task_result_filters = []
+  model_class_path = "net.Net"
+  workflows = [
+    {
+      id = "scatter_and_gather"
+      path = "nvflare.app_common.workflows.scatter_and_gather.ScatterAndGather"
+      args {
+        min_clients = 2
+        num_rounds = 1
+        start_round = 0
+        wait_time_after_min_received = 0
+        aggregator_id = "aggregator"
+        persistor_id = "persistor"
+        shareable_generator_id = "shareable_generator"
+        train_task_name = "train"
+        train_timeout = 0
+      }
+    }
+  ]
+  components = [
+    {
+      id = "persistor"
+      path = "nvflare.app_opt.pt.file_model_persistor.PTFileModelPersistor"
+      args {
+        model {
+          path = "{model_class_path}"
+        }
+      }
+    }
+    {
+      id = "shareable_generator"
+      path = "nvflare.app_common.shareablegenerators.full_model_shareable_generator.FullModelShareableGenerator"
+      args {}
+    }
+    {
+      id = "aggregator"
+      path = "nvflare.app_common.aggregators.intime_accumulate_model_aggregator.InTimeAccumulateWeightedAggregator"
+      args {
+        expected_data_kind = "WEIGHTS"
+      }
+    }
+    {
+      id = "model_selector"
+      path = "nvflare.app_common.widgets.intime_model_selector.IntimeModelSelector"
+      args {
+        key_metric = "accuracy"
+      }
+    }
+  ]
+}
diff --git a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/net.py b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/net.py
new file mode 100644
index 0000000000..ea077b0c6c
--- /dev/null
+++ b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/net.py
@@ -0,0 +1,24 @@
+# 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.
+
+import torch.nn as nn
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(2, 1)
+
+    def forward(self, x):
+        return self.fc1(x)
diff --git a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py
new file mode 100644
index 0000000000..3407197d00
--- /dev/null
+++ b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py
@@ -0,0 +1,71 @@
+# 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.
+
+import torch
+import torch.distributed as dist
+
+import nvflare.client as flare
+
+
+def _broadcast_object_from_rank0(obj):
+    objects = [obj if dist.get_rank() == 0 else None]
+    dist.broadcast_object_list(objects, src=0)
+    return objects[0]
+
+
+def main():
+    dist.init_process_group(backend="gloo")
+
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    if world_size != 2:
+        raise RuntimeError(f"expected two torchrun ranks, got world_size={world_size}")
+
+    flare.init(rank=rank)
+
+    received_model = flare.receive()
+    if rank == 0:
+        if received_model is None:
+            raise RuntimeError("rank 0 expected an FLModel from NVFlare")
+    elif received_model is not None:
+        raise RuntimeError("nonzero rank should not receive directly from NVFlare")
+
+    input_model = _broadcast_object_from_rank0(received_model)
+    if input_model is None or input_model.params is None:
+        raise RuntimeError("distributed broadcast did not provide the FLModel to all ranks")
+
+    rank_contribution = torch.tensor(float(rank + 1))
+    dist.all_reduce(rank_contribution, op=dist.ReduceOp.SUM)
+    if rank_contribution.item() != 3.0:
+        raise RuntimeError(f"expected rank contribution sum 3.0, got {rank_contribution.item()}")
+
+    params = {name: tensor.detach().clone() + rank_contribution.item() for name, tensor in input_model.params.items()}
+    output_model = flare.FLModel(
+        params=params,
+        metrics={
+            "accuracy": rank_contribution.item(),
+            "torchrun_world_size": world_size,
+        },
+        meta={
+            "NUM_STEPS_CURRENT_ROUND": world_size,
+            "torchrun_rank": rank,
+        },
+    )
+
+    flare.send(output_model)
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/meta.conf b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/meta.conf
new file mode 100644
index 0000000000..ea9077e0fb
--- /dev/null
+++ b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/meta.conf
@@ -0,0 +1,11 @@
+{
+  name = "pt_client_api_torchrun_cpu"
+  resource_spec {}
+  deploy_map {
+    app = [
+      "@ALL"
+    ]
+  }
+  min_clients = 2
+  mandatory_clients = []
+}
diff --git a/tests/integration_test/data/test_configs/standalone_job/client_api.yml b/tests/integration_test/data/test_configs/standalone_job/client_api.yml
index 8aed53082a..7291087043 100644
--- a/tests/integration_test/data/test_configs/standalone_job/client_api.yml
+++ b/tests/integration_test/data/test_configs/standalone_job/client_api.yml
@@ -67,6 +67,21 @@ tests:
           "data": { "run_finished": True }
     setup:
       - python tools/prepare_cifar10.py --root /tmp/nvflare/data
+  - test_name: "run pt-client-api-torchrun-cpu"
+    event_sequence:
+      - "trigger":
+          "type": "server_log"
+          "data": "Server started"
+        "actions": [ "submit_job pt_client_api_torchrun_cpu" ]
+        "result":
+          "type": "job_submit_success"
+      - "trigger":
+          "type": "run_state"
+          "data": { "run_finished": True }
+        "actions": [ "ensure_current_job_done" ]
+        "result":
+          "type": "run_state"
+          "data": { "run_finished": True }
   - test_name: "run pt-client-api-launch-once"
     event_sequence:
       - "trigger":

From d5d26e38192c46f8a3840ef501333de9f9a9c6b7 Mon Sep 17 00:00:00 2001
From: YuanTingHsieh <yuantingh@nvidia.com>
Date: Thu, 25 Jun 2026 14:04:42 -0700
Subject: [PATCH 2/5] Clarify Client API rank handling for torchrun

---
 examples/advanced/multi-gpu/pt/README.md      | 15 ++++-
 examples/advanced/multi-gpu/pt/client.py      | 30 +++++----
 nvflare/client/api.py                         |  5 +-
 tests/unit_test/client/ex_process/api_test.py | 63 ++++++++++++++++++-
 4 files changed, 95 insertions(+), 18 deletions(-)

diff --git a/examples/advanced/multi-gpu/pt/README.md b/examples/advanced/multi-gpu/pt/README.md
index f5ba7575b4..302d80bd40 100644
--- a/examples/advanced/multi-gpu/pt/README.md
+++ b/examples/advanced/multi-gpu/pt/README.md
@@ -84,6 +84,20 @@ Set `--nproc_per_node` to the number of GPUs you want to use:
 python3 -m torch.distributed.run --nproc_per_node=4 client.py
 ```
 
+The client script uses the global distributed rank for NVFlare Client API and
+the local rank for CUDA device placement:
+
+```python
+global_rank = dist.get_rank()
+local_rank = int(os.environ["LOCAL_RANK"])
+
+torch.cuda.set_device(local_rank)
+flare.init(rank=global_rank)
+```
+
+`torchrun` sets both `RANK` and `LOCAL_RANK`. `RANK` is unique across the whole
+distributed job, while `LOCAL_RANK` is only unique on the current node.
+
 ### Multiple Clients on Same Machine
 When running multiple clients on the same machine, use different master ports:
 ```python
@@ -118,4 +132,3 @@ python -c "import torch; print(torch.cuda.device_count())"
 - [PyTorch DDP Tutorial](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html)
 - [torch.distributed.run Documentation](https://pytorch.org/docs/stable/distributed.html#launch-utility)
 - [NVFlare Documentation](https://nvflare.readthedocs.io/)
-
diff --git a/examples/advanced/multi-gpu/pt/client.py b/examples/advanced/multi-gpu/pt/client.py
index 4253daaae3..7e1808bb93 100644
--- a/examples/advanced/multi-gpu/pt/client.py
+++ b/examples/advanced/multi-gpu/pt/client.py
@@ -19,6 +19,8 @@
     python -m torch.distributed.run --nnodes=1 --nproc_per_node=2 --master_port=7777 client.py
 """
 
+import os
+
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -61,10 +63,14 @@ def evaluate(input_weights, device, dataloader):
 def main():
     # Initialize distributed process group
     dist.init_process_group("nccl")
-    rank = int(dist.get_rank())
-    device = f"cuda:{rank}"
-    torch.cuda.set_device(rank)
-    print(f"DDP rank {rank} initialized on {device}")
+    global_rank = int(dist.get_rank())
+    if "LOCAL_RANK" not in os.environ:
+        raise RuntimeError("LOCAL_RANK is required for torchrun/DDP GPU training. Launch with torchrun.")
+    local_rank = int(os.environ["LOCAL_RANK"])
+
+    device = f"cuda:{local_rank}"
+    torch.cuda.set_device(local_rank)
+    print(f"DDP global_rank={global_rank}, local_rank={local_rank} initialized on {device}")
 
     # Data setup
     transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
@@ -83,13 +89,13 @@ def main():
     criterion = nn.CrossEntropyLoss()
 
     # (2) initializes NVFlare client API
-    flare.init(rank=f"{rank}")
+    flare.init(rank=global_rank)
 
-    print(f"flare init DDP rank {rank} initialized on {device}")
+    print(f"flare init DDP global_rank={global_rank}, local_rank={local_rank} initialized on {device}")
     # (3) gets FLModel from NVFlare
     while flare.is_running():
         input_model = flare.receive()
-        if rank == 0:
+        if global_rank == 0:
             print(f"\n[Round={input_model.current_round}, Site={flare.get_site_name()}]")
             # (4) loads model from NVFlare
             net.load_state_dict(input_model.params)
@@ -97,10 +103,10 @@ def main():
 
         # Wrap model with DDP
         net.to(device)
-        ddp_model = DDP(net, device_ids=[rank])
+        ddp_model = DDP(net, device_ids=[local_rank])
 
         # Sync model across ranks
-        if rank == 0:
+        if global_rank == 0:
             torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)
 
         dist.barrier()
@@ -120,14 +126,14 @@ def main():
                 optimizer.step()
 
                 running_loss += loss.item()
-                if rank == 0 and i % 2000 == 1999:
+                if global_rank == 0 and i % 2000 == 1999:
                     print(f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}")
                     running_loss = 0.0
 
-        print(f"Rank {rank}: Finished Training")
+        print(f"Rank {global_rank}: Finished Training")
 
         # Only rank 0 sends model back
-        if rank == 0:
+        if global_rank == 0:
             # All processes should see same parameters as they all start from same
             # random parameters and gradients are synchronized in backward passes.
             # Therefore, saving it in one process is sufficient.
diff --git a/nvflare/client/api.py b/nvflare/client/api.py
index 570433f35a..521d315b45 100644
--- a/nvflare/client/api.py
+++ b/nvflare/client/api.py
@@ -53,8 +53,9 @@ def init(rank: Optional[Union[str, int]] = None, config_file: Optional[str] = No
     """Initializes NVFlare Client API environment.
 
     Args:
-        rank (str): local rank of the process.
-            It is only useful when the training script has multiple worker processes. (for example multi GPU)
+        rank (str): rank of the process for Client API control-path behavior.
+            In distributed training, use the global process rank (for example torchrun's RANK),
+            not the device-local rank used for GPU placement.
         config_file (str): client api configuration.
 
     Returns:
diff --git a/tests/unit_test/client/ex_process/api_test.py b/tests/unit_test/client/ex_process/api_test.py
index 1ecad75f59..0e36ac34f5 100644
--- a/tests/unit_test/client/ex_process/api_test.py
+++ b/tests/unit_test/client/ex_process/api_test.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import copy
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 
@@ -135,8 +135,6 @@ def test_get_streaming_idle_timeout_default_and_explicit():
 
 def test_ex_process_api_passes_submit_result_timeout_to_agent(monkeypatch):
     """ExProcessClientAPI.init() must pass task-exchange timeouts from config to FlareAgentWithFLModel."""
-    from unittest.mock import MagicMock
-
     from nvflare.client.config import ConfigKey
     from nvflare.client.ex_process.api import ExProcessClientAPI
 
@@ -265,6 +263,65 @@ def test_ex_process_send_failure_raises_and_preserves_params():
     assert api.receive_called is True
 
 
+def test_ex_process_api_uses_rank_env_for_default_rank(monkeypatch):
+    """RANK, not LOCAL_RANK, determines the default Client API rank."""
+    from nvflare.client.ex_process.api import ExProcessClientAPI
+
+    captured = {}
+
+    def _capture_model_registry(client_config, rank, flare_agent):
+        captured["rank"] = rank
+        captured["flare_agent"] = flare_agent
+        return MagicMock()
+
+    monkeypatch.setenv("RANK", "1")
+    monkeypatch.setenv("LOCAL_RANK", "0")
+    monkeypatch.setattr("nvflare.client.ex_process.api._create_client_config", lambda config: ClientConfig({}))
+    monkeypatch.setattr("nvflare.client.ex_process.api.ModelRegistry", _capture_model_registry)
+
+    api = ExProcessClientAPI(config_file="fake_config.json")
+    api._configure_subprocess_logging = lambda client_config: None
+    api.init()
+
+    assert captured["rank"] == "1"
+    assert captured["flare_agent"] is None
+    assert api.flare_agent is None
+
+
+def test_ex_process_api_defaults_to_rank0_when_rank_env_missing(monkeypatch):
+    """LOCAL_RANK alone must not make the process a non-control Client API rank."""
+    from nvflare.client.ex_process.api import ExProcessClientAPI
+
+    captured = {}
+
+    class _CapturingAgent:
+        def __init__(self, *args, **kwargs):
+            captured["agent_created"] = True
+
+        def start(self):
+            captured["agent_started"] = True
+
+    def _capture_model_registry(client_config, rank, flare_agent):
+        captured["rank"] = rank
+        captured["flare_agent"] = flare_agent
+        return MagicMock()
+
+    monkeypatch.delenv("RANK", raising=False)
+    monkeypatch.setenv("LOCAL_RANK", "1")
+    monkeypatch.setattr("nvflare.client.ex_process.api._create_client_config", lambda config: ClientConfig({}))
+    monkeypatch.setattr("nvflare.client.ex_process.api.FlareAgentWithFLModel", _CapturingAgent)
+    monkeypatch.setattr("nvflare.client.ex_process.api.ModelRegistry", _capture_model_registry)
+
+    api = ExProcessClientAPI(config_file="fake_config.json")
+    api._configure_subprocess_logging = lambda client_config: None
+    api.init()
+
+    assert captured["rank"] == "0"
+    assert captured["agent_created"] is True
+    assert captured["agent_started"] is True
+    assert captured["flare_agent"] is not None
+
+
 # ── _downgrade_rotating_handlers tests ────────────────────────────────────────
 
 _ROTATING_CONFIG = {

From 7ab2a1020d34249e5676c53cd1722cda5256a2b8 Mon Sep 17 00:00:00 2001
From: YuanTingHsieh <yuantingh@nvidia.com>
Date: Thu, 25 Jun 2026 15:19:04 -0700
Subject: [PATCH 3/5] Align torchrun test with rank-zero send contract

---
 .../pt_client_api_torchrun_cpu/app/custom/torchrun_client.py   | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py
index 3407197d00..dcc33545c5 100644
--- a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py
+++ b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py
@@ -63,7 +63,8 @@ def main():
         },
     )
 
-    flare.send(output_model)
+    if rank == 0:
+        flare.send(output_model)
     dist.destroy_process_group()
 
 

From 235c9e80c1908326c19cdc4f4bb26f85140bb3aa Mon Sep 17 00:00:00 2001
From: YuanTingHsieh <yuantingh@nvidia.com>
Date: Tue, 30 Jun 2026 13:33:44 -0700
Subject: [PATCH 4/5] Address torchrun Client API review comments

---
 examples/advanced/multi-gpu/pt/README.md      |  2 +
 .../app/custom/torchrun_client.py             | 71 ++++++++++---------
 2 files changed, 40 insertions(+), 33 deletions(-)

diff --git a/examples/advanced/multi-gpu/pt/README.md b/examples/advanced/multi-gpu/pt/README.md
index 302d80bd40..221a4c86d4 100644
--- a/examples/advanced/multi-gpu/pt/README.md
+++ b/examples/advanced/multi-gpu/pt/README.md
@@ -88,6 +88,8 @@ The client script uses the global distributed rank for NVFlare Client API and
 the local rank for CUDA device placement:
 
 ```python
+import os
+
 global_rank = dist.get_rank()
 local_rank = int(os.environ["LOCAL_RANK"])
 
diff --git a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py
index dcc33545c5..6bdbb28408 100644
--- a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py
+++ b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py
@@ -27,45 +27,50 @@ def _broadcast_object_from_rank0(obj):
 def main():
     dist.init_process_group(backend="gloo")
 
-    rank = dist.get_rank()
-    world_size = dist.get_world_size()
-    if world_size != 2:
-        raise RuntimeError(f"expected two torchrun ranks, got world_size={world_size}")
+    try:
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+        if world_size != 2:
+            raise RuntimeError(f"expected two torchrun ranks, got world_size={world_size}")
 
-    flare.init(rank=rank)
+        flare.init(rank=rank)
 
-    received_model = flare.receive()
-    if rank == 0:
-        if received_model is None:
-            raise RuntimeError("rank 0 expected an FLModel from NVFlare")
-    elif received_model is not None:
-        raise RuntimeError("nonzero rank should not receive directly from NVFlare")
+        received_model = flare.receive()
+        if rank == 0:
+            if received_model is None:
+                raise RuntimeError("rank 0 expected an FLModel from NVFlare")
+        elif received_model is not None:
+            raise RuntimeError("nonzero rank should not receive directly from NVFlare")
 
-    input_model = _broadcast_object_from_rank0(received_model)
-    if input_model is None or input_model.params is None:
-        raise RuntimeError("distributed broadcast did not provide the FLModel to all ranks")
+        input_model = _broadcast_object_from_rank0(received_model)
+        if input_model is None or input_model.params is None:
+            raise RuntimeError("distributed broadcast did not provide the FLModel to all ranks")
 
-    rank_contribution = torch.tensor(float(rank + 1))
-    dist.all_reduce(rank_contribution, op=dist.ReduceOp.SUM)
-    if rank_contribution.item() != 3.0:
-        raise RuntimeError(f"expected rank contribution sum 3.0, got {rank_contribution.item()}")
+        rank_contribution = torch.tensor(float(rank + 1))
+        dist.all_reduce(rank_contribution, op=dist.ReduceOp.SUM)
+        if rank_contribution.item() != 3.0:
+            raise RuntimeError(f"expected rank contribution sum 3.0, got {rank_contribution.item()}")
 
-    params = {name: tensor.detach().clone() + rank_contribution.item() for name, tensor in input_model.params.items()}
-    output_model = flare.FLModel(
-        params=params,
-        metrics={
-            "accuracy": rank_contribution.item(),
-            "torchrun_world_size": world_size,
-        },
-        meta={
-            "NUM_STEPS_CURRENT_ROUND": world_size,
-            "torchrun_rank": rank,
-        },
-    )
+        params = {
+            name: tensor.detach().clone() + rank_contribution.item() for name, tensor in input_model.params.items()
+        }
+        output_model = flare.FLModel(
+            params=params,
+            metrics={
+                "accuracy": rank_contribution.item(),
+                "torchrun_world_size": world_size,
+            },
+            meta={
+                "NUM_STEPS_CURRENT_ROUND": world_size,
+                "torchrun_rank": rank,
+            },
+        )
 
-    if rank == 0:
-        flare.send(output_model)
-    dist.destroy_process_group()
+        if rank == 0:
+            flare.send(output_model)
+        dist.barrier()
+    finally:
+        dist.destroy_process_group()
 
 
 if __name__ == "__main__":

From 0bd0795aa303d3029ce790fc818fa80767dc10c5 Mon Sep 17 00:00:00 2001
From: YuanTingHsieh <yuantingh@nvidia.com>
Date: Wed, 1 Jul 2026 14:57:27 -0700
Subject: [PATCH 5/5] Address review findings for torchrun Client API test

- Move dist.barrier() before flare.send() in the torchrun test script: with launch_once=false the Client API exits the process inside flare.send() after the result upload, so the trailing barrier was unreachable on rank 0 and rank 1 always crashed (launcher recorded COMPLETE_FAILED every run)
- Build the result FLModel only on rank 0 and drop unused metrics/meta (accuracy, torchrun_world_size, torchrun_rank)
- Remove the dead IntimeModelSelector from the test server config (it skips round 0 and the job only runs round 0)
- Apply the corrected rank docstring to APISpec.init, ExProcessClientAPI.init and InProcessClientAPI.init
- Apply the global/local rank split to both pt-ddp-docker example copies to match the updated multi-gpu example
- Broadcast the running state from rank 0 in the DDP examples so nonzero ranks exit the training loop cleanly at job end

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 examples/advanced/multi-gpu/pt/client.py      | 12 +++++-
 .../pt-ddp-docker/app_site-1/custom/client.py | 40 +++++++++++++------
 .../pt-ddp-docker/app_site-2/custom/client.py | 40 +++++++++++++------
 nvflare/client/api_spec.py                    |  5 ++-
 nvflare/client/ex_process/api.py              |  5 ++-
 nvflare/client/in_process/api.py              |  5 ++-
 .../app/config/config_fed_server.conf         |  7 ----
 .../app/custom/torchrun_client.py             | 26 +++++-------
 8 files changed, 84 insertions(+), 56 deletions(-)

diff --git a/examples/advanced/multi-gpu/pt/client.py b/examples/advanced/multi-gpu/pt/client.py
index 7e1808bb93..708e832c9f 100644
--- a/examples/advanced/multi-gpu/pt/client.py
+++ b/examples/advanced/multi-gpu/pt/client.py
@@ -92,8 +92,16 @@ def main():
     flare.init(rank=global_rank)
 
     print(f"flare init DDP global_rank={global_rank}, local_rank={local_rank} initialized on {device}")
-    # (3) gets FLModel from NVFlare
-    while flare.is_running():
+    # (3) gets FLModel from NVFlare.
+    # Only rank 0 talks to NVFlare, so flare.is_running() never turns False on
+    # nonzero ranks — rank 0 must broadcast the running state so all ranks exit
+    # the loop together at job end.
+    while True:
+        running = [flare.is_running() if global_rank == 0 else None]
+        dist.broadcast_object_list(running, src=0)
+        if not running[0]:
+            break
+
         input_model = flare.receive()
         if global_rank == 0:
             print(f"\n[Round={input_model.current_round}, Site={flare.get_site_name()}]")
diff --git a/examples/docker/jobs/pt-ddp-docker/app_site-1/custom/client.py b/examples/docker/jobs/pt-ddp-docker/app_site-1/custom/client.py
index 2aeb235851..d2e18c3244 100644
--- a/examples/docker/jobs/pt-ddp-docker/app_site-1/custom/client.py
+++ b/examples/docker/jobs/pt-ddp-docker/app_site-1/custom/client.py
@@ -19,6 +19,8 @@
     python -m torch.distributed.run --nnodes=1 --nproc_per_node=2 --master_port=7777 client.py
 """
 
+import os
+
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -61,10 +63,14 @@ def evaluate(input_weights, device, dataloader):
 def main():
     # Initialize distributed process group
     dist.init_process_group("nccl")
-    rank = int(dist.get_rank())
-    device = f"cuda:{rank}"
-    torch.cuda.set_device(rank)
-    print(f"DDP rank {rank} initialized on {device}")
+    global_rank = int(dist.get_rank())
+    if "LOCAL_RANK" not in os.environ:
+        raise RuntimeError("LOCAL_RANK is required for torchrun/DDP GPU training. Launch with torchrun.")
+    local_rank = int(os.environ["LOCAL_RANK"])
+
+    device = f"cuda:{local_rank}"
+    torch.cuda.set_device(local_rank)
+    print(f"DDP global_rank={global_rank}, local_rank={local_rank} initialized on {device}")
 
     # Data setup
     transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
@@ -83,13 +89,21 @@ def main():
     criterion = nn.CrossEntropyLoss()
 
     # (2) initializes NVFlare client API
-    flare.init(rank=f"{rank}")
+    flare.init(rank=global_rank)
+
+    print(f"flare init DDP global_rank={global_rank}, local_rank={local_rank} initialized on {device}")
+    # (3) gets FLModel from NVFlare.
+    # Only rank 0 talks to NVFlare, so flare.is_running() never turns False on
+    # nonzero ranks — rank 0 must broadcast the running state so all ranks exit
+    # the loop together at job end.
+    while True:
+        running = [flare.is_running() if global_rank == 0 else None]
+        dist.broadcast_object_list(running, src=0)
+        if not running[0]:
+            break
 
-    print(f"flare init DDP rank {rank} initialized on {device}")
-    # (3) gets FLModel from NVFlare
-    while flare.is_running():
         input_model = flare.receive()
-        if rank == 0:
+        if global_rank == 0:
             print(f"\n[Round={input_model.current_round}, Site={flare.get_site_name()}]")
             # (4) loads model from NVFlare
             net.load_state_dict(input_model.params)
@@ -105,7 +119,7 @@ def main():
 
         # Wrap model with DDP — all ranks now have identical parameters
         net.to(device)
-        ddp_model = DDP(net, device_ids=[rank])
+        ddp_model = DDP(net, device_ids=[local_rank])
         optimizer = optim.SGD(ddp_model.parameters(), lr=0.001, momentum=0.9)
 
         # Training loop
@@ -122,14 +136,14 @@ def main():
                 optimizer.step()
 
                 running_loss += loss.item()
-                if rank == 0 and i % 2000 == 1999:
+                if global_rank == 0 and i % 2000 == 1999:
                     print(f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}")
                     running_loss = 0.0
 
-        print(f"Rank {rank}: Finished Training")
+        print(f"Rank {global_rank}: Finished Training")
 
         # Only rank 0 sends model back
-        if rank == 0:
+        if global_rank == 0:
             # All processes should see same parameters as they all start from same
             # random parameters and gradients are synchronized in backward passes.
             # Therefore, saving it in one process is sufficient.
diff --git a/examples/docker/jobs/pt-ddp-docker/app_site-2/custom/client.py b/examples/docker/jobs/pt-ddp-docker/app_site-2/custom/client.py
index 2aeb235851..d2e18c3244 100644
--- a/examples/docker/jobs/pt-ddp-docker/app_site-2/custom/client.py
+++ b/examples/docker/jobs/pt-ddp-docker/app_site-2/custom/client.py
@@ -19,6 +19,8 @@
     python -m torch.distributed.run --nnodes=1 --nproc_per_node=2 --master_port=7777 client.py
 """
 
+import os
+
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -61,10 +63,14 @@ def evaluate(input_weights, device, dataloader):
 def main():
     # Initialize distributed process group
     dist.init_process_group("nccl")
-    rank = int(dist.get_rank())
-    device = f"cuda:{rank}"
-    torch.cuda.set_device(rank)
-    print(f"DDP rank {rank} initialized on {device}")
+    global_rank = int(dist.get_rank())
+    if "LOCAL_RANK" not in os.environ:
+        raise RuntimeError("LOCAL_RANK is required for torchrun/DDP GPU training. Launch with torchrun.")
+    local_rank = int(os.environ["LOCAL_RANK"])
+
+    device = f"cuda:{local_rank}"
+    torch.cuda.set_device(local_rank)
+    print(f"DDP global_rank={global_rank}, local_rank={local_rank} initialized on {device}")
 
     # Data setup
     transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
@@ -83,13 +89,21 @@ def main():
     criterion = nn.CrossEntropyLoss()
 
     # (2) initializes NVFlare client API
-    flare.init(rank=f"{rank}")
+    flare.init(rank=global_rank)
+
+    print(f"flare init DDP global_rank={global_rank}, local_rank={local_rank} initialized on {device}")
+    # (3) gets FLModel from NVFlare.
+    # Only rank 0 talks to NVFlare, so flare.is_running() never turns False on
+    # nonzero ranks — rank 0 must broadcast the running state so all ranks exit
+    # the loop together at job end.
+    while True:
+        running = [flare.is_running() if global_rank == 0 else None]
+        dist.broadcast_object_list(running, src=0)
+        if not running[0]:
+            break
 
-    print(f"flare init DDP rank {rank} initialized on {device}")
-    # (3) gets FLModel from NVFlare
-    while flare.is_running():
         input_model = flare.receive()
-        if rank == 0:
+        if global_rank == 0:
             print(f"\n[Round={input_model.current_round}, Site={flare.get_site_name()}]")
             # (4) loads model from NVFlare
             net.load_state_dict(input_model.params)
@@ -105,7 +119,7 @@ def main():
 
         # Wrap model with DDP — all ranks now have identical parameters
         net.to(device)
-        ddp_model = DDP(net, device_ids=[rank])
+        ddp_model = DDP(net, device_ids=[local_rank])
         optimizer = optim.SGD(ddp_model.parameters(), lr=0.001, momentum=0.9)
 
         # Training loop
@@ -122,14 +136,14 @@ def main():
                 optimizer.step()
 
                 running_loss += loss.item()
-                if rank == 0 and i % 2000 == 1999:
+                if global_rank == 0 and i % 2000 == 1999:
                     print(f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}")
                     running_loss = 0.0
 
-        print(f"Rank {rank}: Finished Training")
+        print(f"Rank {global_rank}: Finished Training")
 
         # Only rank 0 sends model back
-        if rank == 0:
+        if global_rank == 0:
             # All processes should see same parameters as they all start from same
             # random parameters and gradients are synchronized in backward passes.
             # Therefore, saving it in one process is sufficient.
diff --git a/nvflare/client/api_spec.py b/nvflare/client/api_spec.py
index c2e85a57cf..9859f3dfc0 100644
--- a/nvflare/client/api_spec.py
+++ b/nvflare/client/api_spec.py
@@ -58,8 +58,9 @@ def init(self, rank: Optional[str] = None):
         """Initializes NVFlare Client API environment.
 
         Args:
-            rank (str): local rank of the process.
-                It is only useful when the training script has multiple worker processes. (for example multi GPU)
+            rank (str): rank of the process for Client API control-path behavior.
+                In distributed training, use the global process rank (for example torchrun's RANK),
+                not the device-local rank used for GPU placement.
 
         Returns:
             None
diff --git a/nvflare/client/ex_process/api.py b/nvflare/client/ex_process/api.py
index 057dd08bb2..87982d4e89 100644
--- a/nvflare/client/ex_process/api.py
+++ b/nvflare/client/ex_process/api.py
@@ -160,8 +160,9 @@ def init(self, rank: Optional[str] = None):
         """Initializes NVFlare Client API environment.
 
         Args:
-            rank (str): local rank of the process.
-                It is only useful when the training script has multiple worker processes. (for example multi GPU)
+            rank (str): rank of the process for Client API control-path behavior.
+                In distributed training, use the global process rank (for example torchrun's RANK),
+                not the device-local rank used for GPU placement.
         """
 
         if rank is None:
diff --git a/nvflare/client/in_process/api.py b/nvflare/client/in_process/api.py
index 56a7a4af85..e1ce18c412 100644
--- a/nvflare/client/in_process/api.py
+++ b/nvflare/client/in_process/api.py
@@ -71,8 +71,9 @@ def init(self, rank: Optional[str] = None, config: Optional[Dict] = None):
 
         Args:
             config (Union[str, Dict]): config dictionary.
-            rank (str): local rank of the process.
-                It is only useful when the training script has multiple worker processes. (for example multi GPU)
+            rank (str): rank of the process for Client API control-path behavior.
+                In distributed training, use the global process rank (for example torchrun's RANK),
+                not the device-local rank used for GPU placement.
         """
 
         self.rank = rank
diff --git a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/config/config_fed_server.conf b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/config/config_fed_server.conf
index 5764045c5c..b29019e4fc 100644
--- a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/config/config_fed_server.conf
+++ b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/config/config_fed_server.conf
@@ -42,12 +42,5 @@
         expected_data_kind = "WEIGHTS"
       }
     }
-    {
-      id = "model_selector"
-      path = "nvflare.app_common.widgets.intime_model_selector.IntimeModelSelector"
-      args {
-        key_metric = "accuracy"
-      }
-    }
   ]
 }
diff --git a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py
index 6bdbb28408..2c3dda6870 100644
--- a/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py
+++ b/tests/integration_test/data/jobs/pt_client_api_torchrun_cpu/app/custom/torchrun_client.py
@@ -51,24 +51,20 @@ def main():
         if rank_contribution.item() != 3.0:
             raise RuntimeError(f"expected rank contribution sum 3.0, got {rank_contribution.item()}")
 
-        params = {
-            name: tensor.detach().clone() + rank_contribution.item() for name, tensor in input_model.params.items()
-        }
-        output_model = flare.FLModel(
-            params=params,
-            metrics={
-                "accuracy": rank_contribution.item(),
-                "torchrun_world_size": world_size,
-            },
-            meta={
-                "NUM_STEPS_CURRENT_ROUND": world_size,
-                "torchrun_rank": rank,
-            },
-        )
+        # With launch_once=false the Client API exits the process inside flare.send()
+        # right after the result upload, so all collectives must complete before the
+        # send, and the send must be the last thing rank 0 does.
+        dist.barrier()
 
         if rank == 0:
+            params = {
+                name: tensor.detach().clone() + rank_contribution.item() for name, tensor in input_model.params.items()
+            }
+            output_model = flare.FLModel(
+                params=params,
+                meta={"NUM_STEPS_CURRENT_ROUND": world_size},
+            )
             flare.send(output_model)
-        dist.barrier()
     finally:
         dist.destroy_process_group()