Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make ServerApp shutdown at the end. #3002

Merged
merged 2 commits into from
Feb 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion src/py/flwr/server/compat/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,13 +142,15 @@ def start_driver( # pylint: disable=too-many-arguments, too-many-locals
initialized_config,
)

f_stop = threading.Event()
# Start the thread updating nodes
thread = threading.Thread(
target=update_client_manager,
args=(
grpc_driver,
initialized_server.client_manager(),
lock,
f_stop,
),
)
thread.start()
Expand All @@ -159,6 +161,8 @@ def start_driver( # pylint: disable=too-many-arguments, too-many-locals
config=initialized_config,
)

f_stop.set()

# Stop the Driver API server and the thread
with lock:
if driver:
Expand All @@ -177,6 +181,7 @@ def update_client_manager(
driver: GrpcDriver,
client_manager: ClientManager,
lock: threading.Lock,
f_stop: threading.Event,
) -> None:
"""Update the nodes list in the client manager.

Expand All @@ -195,7 +200,7 @@ def update_client_manager(

# Loop until the driver is disconnected
registered_nodes: Dict[int, DriverClientProxy] = {}
while True:
while not f_stop.is_set():
with lock:
# End the while loop if the driver is disconnected
if driver.stub is None:
Expand Down
9 changes: 3 additions & 6 deletions src/py/flwr/server/compat/app_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,15 +50,11 @@ def test_simple_client_manager_update(self) -> None:
driver.get_nodes.return_value = GetNodesResponse(nodes=expected_nodes)
client_manager = SimpleClientManager()
lock = threading.Lock()

f_stop = threading.Event()
# Execute
thread = threading.Thread(
target=update_client_manager,
args=(
driver,
client_manager,
lock,
),
args=(driver, client_manager, lock, f_stop),
daemon=True,
)
thread.start()
Expand All @@ -83,5 +79,6 @@ def test_simple_client_manager_update(self) -> None:
assert node_ids == {node.node_id for node in expected_nodes}
assert updated_node_ids == {node.node_id for node in expected_updated_nodes}

f_stop.set()
# Exit
thread.join()