plato-learn 1.1__py3-none-any.whl

plato/utils/reinforcement_learning/rl_server.py

@@ -0,0 +1,113 @@
+"""
+A federated learning server with RL Agent.
+"""
+
+import asyncio
+import logging
+from abc import abstractmethod
+
+from plato.servers import fedavg
+
+
+class RLServer(fedavg.Server):
+    """A federated learning server with an RL Agent."""
+
+    def __init__(
+        self,
+        agent,
+        model=None,
+        datasource=None,
+        algorithm=None,
+        trainer=None,
+        callbacks=None,
+    ):
+        super().__init__(
+            model=model,
+            datasource=datasource,
+            algorithm=algorithm,
+            trainer=trainer,
+            callbacks=callbacks,
+        )
+        self.agent = agent
+
+    def reset(self):
+        """Resetting the model, trainer, and algorithm on the server."""
+        logging.info(
+            "Reconfiguring the server for episode %d", self.agent.current_episode
+        )
+
+        self.model = None
+        self.trainer = None
+        self.algorithm = None
+        self.init_trainer()
+
+        self.current_round = 0
+
+    async def aggregate_deltas(self, updates, deltas_received):
+        """Aggregate weight updates from the clients using smart weighting."""
+        self.update_state()
+
+        # Extract the total number of samples
+        num_samples = [update.report.num_samples for update in updates]
+        self.total_samples = sum(num_samples)
+
+        # Perform weighted averaging
+        avg_update = {
+            name: self.trainer.zeros(weights.shape)
+            for name, weights in deltas_received[0].items()
+        }
+
+        # e.g., wait for the new action from RL agent
+        # if the action affects the global aggregation
+        self.agent.num_samples = num_samples
+        await self.agent.prep_agent_update()
+        await self.update_action()
+
+        # Use adaptive weighted average
+        for i, update in enumerate(deltas_received):
+            for name, delta in update.items():
+                if delta.type() == "torch.LongTensor":
+                    avg_update[name] += delta * self.smart_weighting[i][0]
+                else:
+                    avg_update[name] += delta * self.smart_weighting[i]
+
+            # Yield to other tasks in the server
+            await asyncio.sleep(0)
+
+        return avg_update
+
+    async def update_action(self):
+        """Updating the RL agent's actions."""
+        if self.agent.current_step == 0:
+            logging.info("[RL Agent] Preparing initial action...")
+            self.agent.prep_action()
+        else:
+            await self.agent.action_updated.wait()
+            self.agent.action_updated.clear()
+
+        self.apply_action()
+
+    def update_state(self):
+        """Wrap up the state update to RL Agent."""
+        # Pass new state to RL Agent
+        self.agent.new_state = self.prep_state()
+        self.agent.process_env_update()
+
+    async def wrap_up(self) -> None:
+        """Wrapping up when each round of training is done."""
+        self.save_to_checkpoint()
+
+        if self.agent.reset_env:
+            self.agent.reset_env = False
+            self.reset()
+        if self.agent.finished:
+            await self._close()
+
+    @abstractmethod
+    def prep_state(self):
+        """Wrap up the state update to RL Agent."""
+        return
+
+    @abstractmethod
+    def apply_action(self):
+        """Apply action update from RL Agent to FL Env."""

plato/utils/rl_env.py

@@ -0,0 +1,154 @@
+"""
+An environment of the reinforcement learning agent for tuning parameters
+during the training of federated learning.
+This environment follows the gym interface, in order to use stable-baselines3:
+https://github.com/DLR-RM/stable-baselines3.
+
+To create and use other custom environments, check out:
+https://stable-baselines3.readthedocs.io/en/master/guide/custom_env.html.
+"""
+
+import asyncio
+import logging
+
+import gym
+import numpy as np
+from plato.config import Config
+from gym import spaces
+
+
+class RLEnv(gym.Env):
+    """The environment of federated learning."""
+
+    metadata = {"render.modes": ["fl"]}
+
+    def __init__(self, rl_agent):
+        super().__init__()
+
+        self.rl_agent = rl_agent
+        self.time_step = 0
+        self.state = None
+        self.is_episode_done = False
+
+        # An RL env waits for the event that it gets the current state from RL agent
+        self.state_got = asyncio.Event()
+
+        # An RL agent waits for the event that the RL env finishes step()
+        # so that it can start a new FL round
+        self.step_done = asyncio.Event()
+
+        # Normalize action space and make it symmetric when continuous.
+        # The reasons behind:
+        # https://stable-baselines3.readthedocs.io/en/master/guide/rl_tips.html#tips-and-tricks-when-creating-a-custom-environment
+        n_actions = 1
+        self.action_space = spaces.Box(
+            low=-1, high=1, shape=(n_actions,), dtype="float32"
+        )
+
+        # Use only global model accurarcy as state for now
+        self.n_states = 1
+        # Also normalize observation space for better RL training
+        self.observation_space = spaces.Box(
+            low=-1, high=1, shape=(self.n_states,), dtype="float32"
+        )
+
+        self.state = [0 for i in range(self.n_states)]
+
+    def reset(self):
+        if self.rl_agent.rl_episode >= Config().algorithm.rl_episodes:
+            while True:
+                # Give RL agent some time to close connections and exit
+                current_loop = asyncio.get_event_loop()
+                task = current_loop.create_task(asyncio.sleep(1))
+                current_loop.run_until_complete(task)
+
+        logging.info("Reseting RL environment.")
+
+        self.time_step = 0
+        # Let the RL agent restart FL training
+        self.rl_agent.reset_rl_env()
+
+        self.rl_agent.new_episode_begin.set()
+
+        self.state = [0 for i in range(self.n_states)]
+        return np.array(self.state)
+
+    def step(self, action):
+        """One step of reinforcement learning."""
+        assert self.action_space.contains(action), "%r (%s) invalid" % (
+            action,
+            type(action),
+        )
+        self.time_step += 1
+        reward = float(0)
+        self.is_episode_done = False
+
+        # For testing code
+        current_edge_agg_num = self.time_step
+
+        # Rescale the action from [-1, 1] to [1, 2, ... , 9]
+        # The action is the number of aggregations on edge servers
+        # current_edge_agg_num = int((action + 2) * (action + 2))
+
+        logging.info("RL Agent: Start time step #%s...", self.time_step)
+        logging.info(
+            "Each edge server will run %s rounds of local aggregation.",
+            current_edge_agg_num,
+        )
+
+        # Pass the tuned parameter to RL agent
+        self.rl_agent.get_tuned_para(current_edge_agg_num, self.time_step)
+
+        # Wait for state
+        current_loop = asyncio.get_event_loop()
+        get_state_task = current_loop.create_task(self.wait_for_state())
+        current_loop.run_until_complete(get_state_task)
+        # print('State:', self.state)
+
+        self.normalize_state()
+        # print('Normalized state:', self.state)
+
+        reward = self.get_reward()
+        info = {}
+
+        self.rl_agent.cumulative_reward += reward
+
+        # Signal the RL agent to start next time step (next round of FL)
+        self.step_done.set()
+
+        return np.array([self.state]), reward, self.is_episode_done, info
+
+    async def wait_for_state(self):
+        """Wait for getting the current state."""
+        await self.state_got.wait()
+        assert self.time_step == self.rl_agent.current_round
+        self.state_got.clear()
+
+    def get_state(self, state, is_episode_done):
+        """
+        Get transitted state from RL agent.
+        This function is called by RL agent.
+        """
+        self.state = state
+        self.is_episode_done = is_episode_done
+        # Signal the RL env that it gets the current state
+        self.state_got.set()
+        print("RL env: Get state", state)
+        self.rl_agent.is_rl_tuned_para_got = False
+
+    def normalize_state(self):
+        """Normalize each element of state to [-1,1]."""
+        self.state = 2 * (self.state - 0.5)
+
+    def get_reward(self):
+        """Get reward based on the state."""
+        accuracy = self.state
+        # Use accuracy as reward for now.
+        reward = accuracy
+        return reward
+
+    def render(self, mode="rl"):
+        pass
+
+    def close(self):
+        pass

plato/utils/s3.py

@@ -0,0 +1,141 @@
+"""
+Utilities to transmit Python objects to and from an S3-compatible object storage service.
+"""
+
+import pickle
+from typing import Any
+
+import boto3
+import botocore
+import requests
+
+from plato.config import Config
+
+
+class S3:
+    """Manages the utilities to transmit Python objects to and from an S3-compatibile
+    object storage service.
+    """
+
+    def __init__(self, endpoint=None, access_key=None, secret_key=None, bucket=None):
+        """All S3-related credentials, such as the access key and the secret key,
+        are either to be stored in ~/.aws/credentials by using the 'aws configure'
+        command, passed into the constructor as parameters, or specified in the
+        `server` section of the configuration file.
+        """
+        self.endpoint = endpoint
+        self.bucket = bucket
+        self.key_prefix = ""
+        self.access_key = access_key
+        self.secret_key = secret_key
+
+        if hasattr(Config().server, "s3_endpoint_url"):
+            self.endpoint = Config().server.s3_endpoint_url
+
+        if hasattr(Config().server, "s3_bucket"):
+            self.bucket = Config().server.s3_bucket
+
+        if hasattr(Config().server, "access_key"):
+            self.access_key = Config().server.access_key
+
+        if hasattr(Config().server, "secret_key"):
+            self.secret_key = Config().server.secret_key
+
+        if self.bucket is None:
+            raise ValueError("The S3 storage service has not been properly configured.")
+
+        if "s3://" in self.bucket:
+            bucket_part = self.bucket[5:]
+            str_list = bucket_part.split("/")
+            self.bucket = str_list[0]
+            if len(str_list) > 1:
+                self.key_prefix = bucket_part[len(self.bucket) :]
+
+        if self.access_key is not None and self.secret_key is not None:
+            self.s3_client = boto3.client(
+                "s3",
+                endpoint_url=self.endpoint,
+                aws_access_key_id=self.access_key,
+                aws_secret_access_key=self.secret_key,
+            )
+        else:
+            # the access key and secret key are stored locally in ~/.aws/credentials
+            self.s3_client = boto3.client("s3", endpoint_url=self.endpoint)
+
+        # Does the bucket exist?
+        try:
+            self.s3_client.head_bucket(Bucket=self.bucket)
+        except botocore.exceptions.ClientError:
+            try:
+                self.s3_client.create_bucket(Bucket=self.bucket)
+            except botocore.exceptions.ClientError as s3_exception:
+                raise ValueError("Fail to create a bucket.") from s3_exception
+
+    def send_to_s3(self, object_key, object_to_send) -> str:
+        """Sends an object to an S3-compatible object storage service.
+
+        Returns: A presigned URL for use later to retrieve the data.
+        """
+        object_key = self.key_prefix + "/" + object_key
+        try:
+            # Does the object key exist already in S3?
+            self.s3_client.head_object(Bucket=self.bucket, Key=object_key)
+        except botocore.exceptions.ClientError:
+            try:
+                # Only send the object if the key does not exist yet
+                data = pickle.dumps(object_to_send)
+                put_url = self.s3_client.generate_presigned_url(
+                    ClientMethod="put_object",
+                    Params={"Bucket": self.bucket, "Key": object_key},
+                    ExpiresIn=300,
+                )
+                response = requests.put(put_url, data=data)
+
+                if response.status_code != 200:
+                    raise ValueError(
+                        f"Error occurred sending data: status code = {response.status_code}"
+                    ) from None
+
+            except botocore.exceptions.ClientError as error:
+                raise ValueError(
+                    f"Error occurred sending data to S3: {error}"
+                ) from error
+
+            except botocore.exceptions.ParamValidationError as error:
+                raise ValueError(f"Incorrect parameters: {error}") from error
+
+    def receive_from_s3(self, object_key) -> Any:
+        """Retrieves an object from an S3-compatible object storage service.
+
+        All S3-related credentials, such as the access key and the secret key,
+        are assumed to be stored in ~/.aws/credentials by using the 'aws configure'
+        command.
+
+        Returns: The object to be retrieved.
+        """
+        object_key = self.key_prefix + "/" + object_key
+        get_url = self.s3_client.generate_presigned_url(
+            ClientMethod="get_object",
+            Params={"Bucket": self.bucket, "Key": object_key},
+            ExpiresIn=300,
+        )
+        response = requests.get(get_url)
+
+        if response.status_code == 200:
+            return pickle.loads(response.content)
+
+        raise ValueError(
+            f"Error occurred sending data: request status code = {response.status_code}"
+        )
+
+    def delete_from_s3(self, object_key):
+        """Deletes an object using its key from S3."""
+        __ = self.s3_client.delete_object(Bucket=self.bucket, Key=object_key)
+
+    def lists(self):
+        """Retrieves keys to all the objects in the S3 bucket."""
+        response = self.s3_client.list_objects_v2(Bucket=self.bucket)
+        keys = []
+        for obj in response["Contents"]:
+            keys.append(obj["Key"])
+        return keys

plato/utils/trainer_utils.py

@@ -0,0 +1,21 @@
+"""
+The necessary tools used by trainers.
+"""
+
+
+def freeze_model(model, layer_names=None):
+    """Freeze a part of the model."""
+    if layer_names is not None:
+        frozen_params = []
+        for name, param in model.named_parameters():
+            if any(param_name in name for param_name in layer_names):
+                param.requires_grad = False
+                frozen_params.append(name)
+
+
+def activate_model(model, layer_names=None):
+    """Activate a part of the model."""
+    if layer_names is not None:
+        for name, param in model.named_parameters():
+            if any(param_name in name for param_name in layer_names):
+                param.requires_grad = True

plato/utils/unary_encoding.py

@@ -0,0 +1,47 @@
+"""Implements unary encoding, used by Google's RAPPOR, as the local differential privacy mechanism.
+
+References:
+
+Wang, et al. "Optimizing Locally Differentially Private Protocols," ATC USENIX 2017.
+
+Erlingsson, et al. "RAPPOR: Randomized Aggregatable Privacy-Preserving Ordinal Response,"
+ACM CCS 2014.
+
+"""
+
+import numpy as np
+
+
+def encode(x: np.ndarray):
+    x[x > 0] = 1
+    x[x <= 0] = 0
+    return x
+
+
+def randomize(bit_array: np.ndarray, epsilon):
+    """
+    The default unary encoding method is symmetric.
+    """
+    assert isinstance(bit_array, np.ndarray)
+    return symmetric_unary_encoding(bit_array, epsilon)
+
+
+def symmetric_unary_encoding(bit_array: np.ndarray, epsilon):
+    p = np.e ** (epsilon / 2) / (np.e ** (epsilon / 2) + 1)
+    q = 1 / (np.e ** (epsilon / 2) + 1)
+    return produce_randomized_response(bit_array, p, q)
+
+
+def optimized_unary_encoding(bit_array: np.ndarray, epsilon):
+    p = 1 / 2
+    q = 1 / (np.e**epsilon + 1)
+    return produce_randomized_response(bit_array, p, q)
+
+
+def produce_randomized_response(bit_array: np.ndarray, p, q=None):
+    """Implements randomized response as the perturbation method."""
+    q = 1 - p if q is None else q
+
+    p_binomial = np.random.binomial(1, p, bit_array.shape)
+    q_binomial = np.random.binomial(1, q, bit_array.shape)
+    return np.where(bit_array == 1, p_binomial, q_binomial)

plato_learn-1.1.dist-info/METADATA

@@ -0,0 +1,35 @@
+Metadata-Version: 2.4
+Name: plato-learn
+Version: 1.1
+Summary: Plato: a research framework for federated learning
+Project-URL: Homepage, https://github.com/TL-System/plato
+Project-URL: Repository, https://github.com/TL-System/plato
+Project-URL: Documentation, https://platodocs.netlify.app/
+License-Expression: Apache-2.0
+License-File: LICENSE
+Requires-Python: >=3.13
+Requires-Dist: aiohttp
+Requires-Dist: boto3
+Requires-Dist: datasets
+Requires-Dist: evaluate
+Requires-Dist: gym
+Requires-Dist: lightly
+Requires-Dist: numpy
+Requires-Dist: opacus
+Requires-Dist: python-socketio
+Requires-Dist: pyyaml
+Requires-Dist: requests
+Requires-Dist: tenseal
+Requires-Dist: timm
+Requires-Dist: torch
+Requires-Dist: torch-optimizer
+Requires-Dist: torchvision
+Requires-Dist: transformers
+Requires-Dist: zstd
+Provides-Extra: dev
+Requires-Dist: pytest; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# Plato: A New Framework for Scalable Federated Learning Research
+
+Welcome to *Plato*, a software framework to facilitate scalable, reproducible, and extensible federated learning research. Please refer to the documentation website, available in `/documentation`, for more details on installing, running and deploying Plato.