Add shutdown method to Flare Client API #3152

YuanTingHsieh · 2025-01-17T23:43:22Z

Issue

Flare API does not close the Flare agent and CellPipe resources
LauncherExecutor does not shutdown its threadpoolexecutor
SubprocessLauncher have a racing issue when multiple threads try to stop/start the process and monitor the process

Description

Adds the shutdown functionality to the Flare API.
In addition, introduces the context concept so prepare for future if multiple connections need to be made.
Existing API usage is still compatible.
Makes corresponding changes to tracking and lightning API
Adds shutdown of threadpoolexecutor in LauncherExecutor
Adds a lock in SubprocessLauncher

New API Usage:

import nvflare.client as flare
from nvflare.client import FLModel

with flare.init() as ctx:
    input_model = flare.receive(ctx=ctx)
    # do some training
    flare.send(FLModel(xxx), ctx=ctx)

Types of changes

Non-breaking change (fix or new feature that would not break existing functionality).
Breaking change (fix or new feature that would cause existing functionality to change).
New tests added to cover the changes.
Quick tests passed locally by running ./runtest.sh.
In-line docstrings updated.
Documentation updated.

yanchengnv

See my comments and questions.

nvflare/client/__init__.py

nvflare/client/api.py

nvflare/client/ex_process/api.py

examples/hello-world/hello-pt/src/hello-pt_cifar10_fl.py

yanchengnv

We need to be able to shutdown the API.

nvflare/client/api.py

yanchengnv

See my comments for corrections and improvement.

nvflare/client/api.py

yanchengnv · 2025-01-30T15:29:44Z

nvflare/client/api.py

-    return client_api.receive(timeout)
+    global default_context
+    local_ctx = ctx if ctx else default_context
+    return local_ctx.api.receive(timeout)


Should check local_ctx is not None. If none, raise exception.

yanchengnv · 2025-01-30T15:30:01Z

nvflare/client/api.py

-    return client_api.send(model, clear_cache)
+    global default_context
+    local_ctx = ctx if ctx else default_context
+    return local_ctx.api.send(model, clear_cache)


Should check local_ctx is not None. If none, raise exception.

yanchengnv · 2025-01-30T15:30:17Z

nvflare/client/api.py

-    return client_api.system_info()
+    global default_context
+    local_ctx = ctx if ctx else default_context
+    return local_ctx.api.system_info()


Should check local_ctx is not None. If none, raise exception.

yanchengnv · 2025-01-30T15:30:25Z

nvflare/client/api.py

-    return client_api.get_config()
+    global default_context
+    local_ctx = ctx if ctx else default_context
+    return local_ctx.api.get_config()


Should check local_ctx is not None. If none, raise exception.

yanchengnv · 2025-01-30T15:33:29Z

nvflare/client/api_context.py

+    def _create_client_api(self, api_type: ClientAPIType) -> APISpec:
+        """Creates a new client_api based on the provided API type."""
+        if api_type == ClientAPIType.IN_PROCESS_API:
+            return data_bus.get_data(CLIENT_API_KEY)


What if it doesn't exist in data_bus? Should raise exception.

yanchengnv · 2025-01-30T15:39:47Z

nvflare/client/ex_process/api.py

@@ -206,3 +205,7 @@ def clear(self):
        model_registry = self.get_model_registry()
        model_registry.clear()
        self.receive_called = False
+
+    def shutdown(self):
+        model_registry = self.get_model_registry()


What if the app calls shutdown multiple times?
How do you ensure that cell agent is closed? You need to keep flare_agent in "self" and then call self.flare_agent.stop().

yanchengnv · 2025-01-30T15:40:36Z

nvflare/client/in_process/api.py

+
+    def shutdown(self):
+        self.stop = True
+        self.stop_reason = "API shutdown called."


Any resources to clean up? Anything to stop?

yanchengnv · 2025-01-30T15:43:35Z

nvflare/client/task_registry.py

@@ -114,6 +114,6 @@ def clear(self) -> None:
    def __str__(self):
        return f"{self.__class__.__name__}(config: {self.config.get_config()})"

-    def __del__(self):
+    def shutdown(self):


Why does task registry have flare_agent? This design seems very odd. It is created when init the API, hence it should belong to the API, and shutdown there too.

In general, the creator of a resource is responsible for its lifecycle (create and shutdown/close/stop/destroy).

yanchengnv · 2025-01-30T15:46:26Z

nvflare/client/tracking.py

-class SummaryWriter:
+class _BaseWriter:
+    def __init__(self, ctx: Optional[APIContext] = None):
+        global default_context


Should minimize use of gloabl.
You can do: ctx = api.get_default_ctx().

YuanTingHsieh requested review from yanchengnv and nvidianz January 17, 2025 23:55

yanchengnv reviewed Jan 20, 2025

View reviewed changes

nvidianz reviewed Jan 21, 2025

View reviewed changes

examples/hello-world/hello-pt/src/hello-pt_cifar10_fl.py Outdated Show resolved Hide resolved

yanchengnv requested changes Jan 24, 2025

View reviewed changes

nvflare/client/api.py Outdated Show resolved Hide resolved

YuanTingHsieh force-pushed the add_flare_client_context branch from 0ece3c2 to ca7b69c Compare January 28, 2025 05:09

YuanTingHsieh changed the title ~~Enhance Flare Client API to Support Multiple Contexts and Server Connections~~ Add shutdown method to Flare Client API Jan 28, 2025

YuanTingHsieh force-pushed the add_flare_client_context branch 2 times, most recently from 83f8165 to 94544fe Compare January 30, 2025 05:47

yanchengnv requested changes Jan 30, 2025

View reviewed changes

holgerroth self-requested a review February 5, 2025 02:54

YuanTingHsieh added 8 commits February 6, 2025 21:05

Add FlareClientContext

1b5ab95

Address comments

15c3464

Add shutdown_timeout and check for stopped/aborted in PipeHandler

43d4bd8

Register TensorDecomposer automatically

089a12c

Fix SubprocessLauncher racing

30c9e5d

remove unused file

6ef8464

undo

2ca2ace

Address review comments

df38c1c

YuanTingHsieh force-pushed the add_flare_client_context branch from 71ffe16 to df38c1c Compare February 7, 2025 05:05

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add shutdown method to Flare Client API #3152

Add shutdown method to Flare Client API #3152

YuanTingHsieh commented Jan 17, 2025 •

edited

Loading

yanchengnv left a comment

yanchengnv left a comment

yanchengnv left a comment

yanchengnv Jan 30, 2025

yanchengnv Jan 30, 2025

yanchengnv Jan 30, 2025

yanchengnv Jan 30, 2025

yanchengnv Jan 30, 2025

yanchengnv Jan 30, 2025

yanchengnv Jan 30, 2025

yanchengnv Jan 30, 2025 •

edited

Loading

yanchengnv Jan 30, 2025

Add shutdown method to Flare Client API #3152

Are you sure you want to change the base?

Add shutdown method to Flare Client API #3152

Conversation

YuanTingHsieh commented Jan 17, 2025 • edited Loading

Issue

Description

Types of changes

yanchengnv left a comment

Choose a reason for hiding this comment

yanchengnv left a comment

Choose a reason for hiding this comment

yanchengnv left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yanchengnv Jan 30, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

YuanTingHsieh commented Jan 17, 2025 •

edited

Loading

yanchengnv Jan 30, 2025 •

edited

Loading