flwr-nightly 1.8.0.dev20240315__py3-none-any.whl → 1.11.0.dev20240813__py3-none-any.whl

flwr/simulation/ray_transport/ray_client_proxy.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2021 Flower Labs GmbH. 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.
@@ -20,11 +20,16 @@ from logging import ERROR
 from typing import Optional
 
 from flwr import common
-from flwr.client import ClientFn
+from flwr.client import ClientFnExt
 from flwr.client.client_app import ClientApp
 from flwr.client.node_state import NodeState
-from flwr.common import Message, Metadata, RecordSet
-from flwr.common.constant import MessageType, MessageTypeLegacy
+from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
+from flwr.common.constant import (
+    NUM_PARTITIONS_KEY,
+    PARTITION_ID_KEY,
+    MessageType,
+    MessageTypeLegacy,
+)
 from flwr.common.logger import log
 from flwr.common.recordset_compat import (
     evaluateins_to_recordset,
@@ -43,17 +48,30 @@ from flwr.simulation.ray_transport.ray_actor import VirtualClientEngineActorPool
 class RayActorClientProxy(ClientProxy):
     """Flower client proxy which delegates work using Ray."""
 
-    def __init__(
-        self, client_fn: ClientFn, cid: str, actor_pool: VirtualClientEngineActorPool
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        client_fn: ClientFnExt,
+        node_id: int,
+        partition_id: int,
+        num_partitions: int,
+        actor_pool: VirtualClientEngineActorPool,
     ):
-        super().__init__(cid)
+        super().__init__(cid=str(node_id))
+        self.node_id = node_id
+        self.partition_id = partition_id
 
         def _load_app() -> ClientApp:
             return ClientApp(client_fn=client_fn)
 
         self.app_fn = _load_app
         self.actor_pool = actor_pool
-        self.proxy_state = NodeState()
+        self.proxy_state = NodeState(
+            node_id=node_id,
+            node_config={
+                PARTITION_ID_KEY: str(partition_id),
+                NUM_PARTITIONS_KEY: str(num_partitions),
+            },
+        )
 
     def _submit_job(self, message: Message, timeout: Optional[float]) -> Message:
         """Sumbit a message to the ActorPool."""
@@ -62,16 +80,19 @@ class RayActorClientProxy(ClientProxy):
         # Register state
         self.proxy_state.register_context(run_id=run_id)
 
-        # Retrieve state
-        state = self.proxy_state.retrieve_context(run_id=run_id)
+        # Retrieve context
+        context = self.proxy_state.retrieve_context(run_id=run_id)
+        partition_id_str = str(context.node_config[PARTITION_ID_KEY])
 
         try:
             self.actor_pool.submit_client_job(
-                lambda a, a_fn, mssg, cid, state: a.run.remote(a_fn, mssg, cid, state),
-                (self.app_fn, message, self.cid, state),
+                lambda a, a_fn, mssg, partition_id, context: a.run.remote(
+                    a_fn, mssg, partition_id, context
+                ),
+                (self.app_fn, message, partition_id_str, context),
             )
             out_mssg, updated_context = self.actor_pool.get_client_result(
-                self.cid, timeout
+                partition_id_str, timeout
             )
 
             # Update state
@@ -103,11 +124,10 @@ class RayActorClientProxy(ClientProxy):
                 message_id="",
                 group_id=str(group_id) if group_id is not None else "",
                 src_node_id=0,
-                dst_node_id=int(self.cid),
+                dst_node_id=self.node_id,
                 reply_to_message="",
-                ttl=str(timeout) if timeout else "",
+                ttl=timeout if timeout else DEFAULT_TTL,
                 message_type=message_type,
-                partition_id=int(self.cid),
             ),
         )