Repository: instadeepai/Mava

Issue:
[TEST] Jax System Core Parameter Client
### What do you want to test?
Jax system core parameter client (in `mava/systems/jax/`)

### Outline of test structure
* Unit tests
* Test components and hooks

### Definition of done
Passing checks, cover all hooks, edge cases considered

### Mandatory checklist before making a PR
* [ ] The success criteria laid down in “Definition of done” are met.
* [ ] Test code is documented - docstrings for methods and classes, static types for arguments.
* [ ] Documentation is updated - README, CONTRIBUTING, or other documentation.


Discussion:

Resolving patch:
diff --git a/mava/systems/jax/parameter_client.py b/mava/systems/jax/parameter_client.py
--- a/mava/systems/jax/parameter_client.py
+++ b/mava/systems/jax/parameter_client.py
@@ -244,6 +244,6 @@ def _copy(self, new_parameters: Dict[str, Any]) -> None:
                     else:
                         self._parameters[key][i] = new_parameters[key][i]
             else:
-                NotImplementedError(
+                raise NotImplementedError(
                     f"Parameter type of {type(new_parameters[key])} not implemented."
                 )


Test patch:
diff --git a/tests/jax/components/building/parameter_client_test.py b/tests/jax/components/building/parameter_client_test.py
--- a/tests/jax/components/building/parameter_client_test.py
+++ b/tests/jax/components/building/parameter_client_test.py
@@ -103,7 +103,9 @@ def mock_builder_with_parameter_client() -> Builder:
     builder.store.trainer_id = "trainer_0"
 
     builder.store.parameter_server_client = ParameterServer(
-        store=SimpleNamespace(get_parameters={"key_1": 1}),
+        store=SimpleNamespace(
+            get_parameters={"trainer_steps": np.array(0, dtype=np.int32)}
+        ),
         components=[],
     )
 
diff --git a/tests/jax/core_system_components/parameter_client_test.py b/tests/jax/core_system_components/parameter_client_test.py
--- a/tests/jax/core_system_components/parameter_client_test.py
+++ b/tests/jax/core_system_components/parameter_client_test.py
@@ -15,4 +15,410 @@
 
 """Tests for parameter client class for Jax-based Mava systems"""
 
-# TODO (Dries)
+import copy
+from types import SimpleNamespace
+from typing import Any, Dict, List, Sequence, Set, Union
+
+import jax
+import numpy as np
+import pytest
+
+from mava.callbacks.base import Callback
+from mava.systems.jax.parameter_client import ParameterClient
+from mava.systems.jax.parameter_server import ParameterServer
+
+
+class MockParameterServer(ParameterServer):
+    def __init__(
+        self,
+        store: SimpleNamespace,
+        components: List[Callback],
+        set_parameter_keys: Union[str, Sequence[str]],
+    ) -> None:
+        """Initialize mock parameter server."""
+        self.store = store
+        self.callbacks = components
+        self.set_parameter_keys = set_parameter_keys
+
+    def _increment_get_parameters(
+        self, names: Union[str, Sequence[str], Set[str]]
+    ) -> None:
+        """Dummy method to update get parameters before updating client"""
+        for name in names:
+            if name.split("_")[0] == "key" or name.split("_")[0] == "allkey":
+                self.store.parameters[name] += 1
+            elif name.split("-")[0] == "networks":
+                self.store.parameters[name]["layer_0"]["weights"] += 1
+                self.store.parameters[name]["layer_0"]["biases"] += 1
+
+    def get_parameters(self, names: Union[str, Sequence[str]]) -> Any:
+        """Dummy method for returning get parameters"""
+        self.store._param_names = names
+
+        # Manually increment all parameters except the set parameters
+        # and add them to store to simulate parameters that have changed.
+        get_names = set(names) - set(self.set_parameter_keys)
+        self._increment_get_parameters(names=get_names)
+        get_params = {name: self.store.parameters[name] for name in names}
+        self.store.get_parameters = get_params
+
+        return self.store.get_parameters
+
+    def set_parameters(self, set_params: Dict[str, Any]) -> None:
+        """Overwrite set parameters method"""
+
+        self.store._set_params = set_params
+
+        for key in set_params:
+            self.store.parameters[key] = copy.deepcopy(set_params[key])
+
+
+def increment_set_parameters(
+    params: Dict[str, Any], names: Union[str, Sequence[str], Set[str]]
+) -> None:
+    """Utility function for incrementing parameter client parameters"""
+    for key in names:
+        if key.split("_")[0] == "key" or key.split("_")[0] == "allkey":
+            params[key] += 1
+        elif key.split("-")[0] == "networks":
+            params[key]["layer_0"]["weights"] += 1
+            params[key]["layer_0"]["biases"] += 1
+
+
+@pytest.fixture
+def mock_parameter_server() -> ParameterServer:
+    """Create mock parameter server"""
+    param_server = MockParameterServer(
+        store=SimpleNamespace(
+            parameters={
+                "key_0": np.array(0, dtype=np.int32),
+                "key_1": np.array(1, dtype=np.float32),
+                "key_2": np.array(2, dtype=np.int32),
+                "key_3": np.array(3, dtype=np.int32),
+                "key_4": np.array(4, dtype=np.int32),
+                "networks-network_key_0": {"layer_0": {"weights": 0, "biases": 0}},
+                "networks-network_key_1": {"layer_0": {"weights": 1, "biases": 1}},
+                "networks-network_key_2": {"layer_0": {"weights": 2, "biases": 2}},
+                "allkey_0": np.array(0, dtype=np.int32),
+            },
+        ),
+        components=[],
+        set_parameter_keys=["key_0", "key_2"],
+    )
+
+    return param_server
+
+
+@pytest.fixture()
+def parameter_client(mock_parameter_server: ParameterServer) -> ParameterClient:
+    """Creates a mock parameter client for testing
+
+    Args:
+        mock_parameter_server: ParameterServer
+
+    Returns:
+        A parameter client object.
+    """
+
+    param_client = ParameterClient(
+        client=mock_parameter_server,
+        parameters={
+            "key_0": np.array(0, dtype=np.int32),
+            "key_1": np.array(1, dtype=np.float32),
+            "key_2": np.array(2, dtype=np.int32),
+            "key_3": np.array(3, dtype=np.int32),
+            "key_4": np.array(4, dtype=np.int32),
+            "networks-network_key_0": {"layer_0": {"weights": 0, "biases": 0}},
+            "networks-network_key_1": {"layer_0": {"weights": 1, "biases": 1}},
+            "networks-network_key_2": {"layer_0": {"weights": 2, "biases": 2}},
+            "allkey_0": np.array(0, dtype=np.int32),
+        },
+        get_keys=[
+            "key_0",
+            "key_1",
+            "key_2",
+            "key_3",
+            "key_4",
+            "networks-network_key_0",
+            "networks-network_key_1",
+            "networks-network_key_2",
+        ],
+        set_keys=[
+            "key_0",
+            "key_2",
+        ],
+        update_period=10,
+    )
+
+    return param_client
+
+
+def test_add_and_wait(parameter_client: ParameterClient) -> None:
+    """Test add and wait method."""
+    parameter_client.add_and_wait(params={"new_key": "new_value"})
+
+    assert parameter_client._client.store._add_to_params == {"new_key": "new_value"}
+
+
+def test_get_and_wait(parameter_client: ParameterClient) -> None:
+    """Test get and wait method."""
+    parameter_client.get_and_wait()
+    # check that all get parameters have been incremented and updated
+    # except for the set parameters
+    assert parameter_client._parameters == {
+        "key_0": np.array(0, dtype=np.int32),
+        "key_1": np.array(2, dtype=np.float32),
+        "key_2": np.array(2, dtype=np.int32),
+        "key_3": np.array(4, dtype=np.int32),
+        "key_4": np.array(5, dtype=np.int32),
+        "networks-network_key_0": {"layer_0": {"weights": 1, "biases": 1}},
+        "networks-network_key_1": {"layer_0": {"weights": 2, "biases": 2}},
+        "networks-network_key_2": {"layer_0": {"weights": 3, "biases": 3}},
+        "allkey_0": np.array(0, dtype=np.int32),
+    }
+
+
+def test_get_all_and_wait(parameter_client: ParameterClient) -> None:
+    """Test get all and wait method."""
+    parameter_client.get_all_and_wait()
+    # check that all parameters have been incremented and updated
+    # except for the set parameters
+    assert parameter_client._parameters == {
+        "key_0": np.array(0, dtype=np.int32),
+        "key_1": np.array(2, dtype=np.float32),
+        "key_2": np.array(2, dtype=np.int32),
+        "key_3": np.array(4, dtype=np.int32),
+        "key_4": np.array(5, dtype=np.int32),
+        "networks-network_key_0": {"layer_0": {"weights": 1, "biases": 1}},
+        "networks-network_key_1": {"layer_0": {"weights": 2, "biases": 2}},
+        "networks-network_key_2": {"layer_0": {"weights": 3, "biases": 3}},
+        "allkey_0": np.array(1, dtype=np.int32),
+    }
+
+
+def test_set_and_wait(parameter_client: ParameterClient) -> None:
+    """Test set and wait method."""
+    # should set parameters from client to server.
+    increment_set_parameters(
+        params=parameter_client._parameters, names=parameter_client._set_keys
+    )
+
+    parameter_client.set_and_wait()
+
+    assert parameter_client._client.store._set_params == {
+        "key_0": np.array(1, dtype=np.int32),
+        "key_2": np.array(3, dtype=np.int32),
+    }
+
+    assert parameter_client._client.store.parameters == {
+        "key_0": np.array(1, dtype=np.int32),
+        "key_1": np.array(1, dtype=np.float32),
+        "key_2": np.array(3, dtype=np.int32),
+        "key_3": np.array(3, dtype=np.int32),
+        "key_4": np.array(4, dtype=np.int32),
+        "networks-network_key_0": {"layer_0": {"weights": 0, "biases": 0}},
+        "networks-network_key_1": {"layer_0": {"weights": 1, "biases": 1}},
+        "networks-network_key_2": {"layer_0": {"weights": 2, "biases": 2}},
+        "allkey_0": np.array(0, dtype=np.int32),
+    }
+
+
+def test_get_async(parameter_client: ParameterClient) -> None:
+    """Test get async method"""
+
+    # test that get_call_counter is incremented
+    parameter_client.get_async()
+    assert parameter_client._get_call_counter == 1
+
+    # set call counter equal to update period and verify the request was made
+    parameter_client._get_call_counter = 10
+    parameter_client.get_async()
+
+    assert parameter_client._parameters == {
+        "key_0": np.array(0, dtype=np.int32),
+        "key_1": np.array(2, dtype=np.float32),
+        "key_2": np.array(2, dtype=np.int32),
+        "key_3": np.array(4, dtype=np.int32),
+        "key_4": np.array(5, dtype=np.int32),
+        "networks-network_key_0": {"layer_0": {"weights": 1, "biases": 1}},
+        "networks-network_key_1": {"layer_0": {"weights": 2, "biases": 2}},
+        "networks-network_key_2": {"layer_0": {"weights": 3, "biases": 3}},
+        "allkey_0": np.array(0, dtype=np.int32),
+    }
+    assert parameter_client._get_call_counter == 0
+    assert parameter_client._get_future is None
+
+
+def test_set_async(parameter_client: ParameterClient) -> None:
+    """Set set async method"""
+
+    # test that set_call_counter is incremented
+    parameter_client.set_async()
+    assert parameter_client._set_call_counter == 1
+
+    # set call counter equal to update period and verify the request was made
+    parameter_client._set_call_counter = 10
+
+    # increment set parameters, call set_async and verify that
+    # parameters were set asynchronously
+    increment_set_parameters(
+        params=parameter_client._parameters, names=parameter_client._set_keys
+    )
+
+    parameter_client.set_async()
+
+    assert parameter_client._client.store._set_params == {
+        "key_0": np.array(1, dtype=np.int32),
+        "key_2": np.array(3, dtype=np.int32),
+    }
+    assert parameter_client._set_call_counter == 0
+
+    # assert that _set_future is None again after request was made
+    parameter_client.set_async()
+    assert parameter_client._set_future is None
+
+
+def test_set_and_get_async(parameter_client: ParameterClient) -> None:
+    """Set set async method"""
+
+    # test that set_and_get_call_counter is incremented
+    parameter_client.set_and_get_async()
+    assert parameter_client._set_get_call_counter == 1
+
+    # set counter equal to update period and verify the request was made
+    parameter_client._set_get_call_counter = 10
+
+    # increment set parameters
+    increment_set_parameters(
+        params=parameter_client._parameters, names=parameter_client._set_keys
+    )
+
+    parameter_client.set_and_get_async()
+
+    assert parameter_client._parameters == {
+        "key_0": np.array(1, dtype=np.int32),
+        "key_1": np.array(2, dtype=np.float32),
+        "key_2": np.array(3, dtype=np.int32),
+        "key_3": np.array(4, dtype=np.int32),
+        "key_4": np.array(5, dtype=np.int32),
+        "networks-network_key_0": {"layer_0": {"weights": 1, "biases": 1}},
+        "networks-network_key_1": {"layer_0": {"weights": 2, "biases": 2}},
+        "networks-network_key_2": {"layer_0": {"weights": 3, "biases": 3}},
+        "allkey_0": np.array(0, dtype=np.int32),
+    }
+
+    assert parameter_client._set_get_call_counter == 0
+
+    parameter_client.set_and_get_async()
+
+    assert parameter_client._set_get_future is None
+
+
+def test_add_async(parameter_client: ParameterClient) -> None:
+    """Test add async method."""
+
+    assert parameter_client._add_future is None
+
+    parameter_client.add_async(params={"new_key": "new_value"})
+    assert parameter_client._add_future is not None
+    assert parameter_client._client.store._add_to_params == {"new_key": "new_value"}
+    assert parameter_client._async_add_buffer == {}
+    assert parameter_client._add_future.done()
+
+    # set add_future to dummy variable other than none
+    parameter_client._add_future._state = "not_done"  # type: ignore
+    parameter_client._async_add_buffer = {"new_key_1": 1}
+
+    parameter_client.add_async(params={"new_key_1": 1})
+
+    assert parameter_client._async_add_buffer == {"new_key_1": 2}
+
+    # set add_future to dummy variable other than none
+    # assert that new parameter is added to async buffer
+    parameter_client._add_future._state = "not_done"  # type: ignore
+    parameter_client._async_add_buffer = {}
+    parameter_client.add_async(params={"new_key_2": 1})
+
+    assert parameter_client._async_add_buffer == {"new_key_2": 1}
+
+
+def test__copy(parameter_client: ParameterClient) -> None:
+    """Test _copy method with different kinds of new parameters"""
+    parameter_client._copy(
+        new_parameters={
+            "networks-network_key_0": {
+                "layer_0": {"weights": "new_weights", "biases": "new_biases"}
+            },
+            "key_2": np.array(20, dtype=np.int32),
+            "key_4": np.array([40], dtype=np.int32),
+        }
+    )
+
+    assert parameter_client._parameters["networks-network_key_0"] == {
+        "layer_0": {"weights": "new_weights", "biases": "new_biases"}
+    }
+    assert parameter_client._parameters["key_2"] == 20
+    assert parameter_client._parameters["key_4"] == 40
+
+
+def test__copy_not_implemented_error(parameter_client: ParameterClient) -> None:
+    """Test that NotImplementedError is raised when a new parameter of the wrong \
+        type is passed in."""
+
+    with pytest.raises(NotImplementedError):
+        parameter_client._copy(
+            new_parameters={"wrong_type_parameter": lambda: "wrong_type"}
+        )
+
+
+def test__copy_device_not_implemented_error(parameter_client: ParameterClient) -> None:
+    """Test that NotImplementedError is raised when a new parameter of the wrong \
+        type is passed in."""
+
+    parameter_client._devices = {"device_1": "dummy_device"}
+
+    with pytest.raises(NotImplementedError):
+        parameter_client._copy(
+            new_parameters={
+                "networks-network_key_0": {
+                    "layer_0": {"weights": "new_weights", "biases": "new_biases"}
+                }
+            },
+        )
+
+
+def test_copy_array_with_device(parameter_client: ParameterClient) -> None:
+    """Test that new parameters are set on a device when a device \
+        is give."""
+    local_devices = jax.local_devices()
+    parameter_client._devices = {"key_1": local_devices[0]}
+
+    parameter_client._copy(new_parameters={"key_1": np.array([0])})
+
+    assert type(parameter_client._parameters["key_1"]).__name__ == "DeviceArray"
+    assert parameter_client._parameters["key_1"] == jax.numpy.array([0])
+
+
+def test_copy_tuple_with_device(parameter_client: ParameterClient) -> None:
+    """Test that new parameters are set on a device when a device \
+        is given."""
+    parameter_client._parameters.update({"tuple_key_0": [1, 2]})
+    local_devices = jax.local_devices()
+    parameter_client._devices = {
+        "tuple_key_0": (local_devices[0], local_devices[0])  # type: ignore
+    }
+
+    parameter_client._copy(
+        new_parameters={"tuple_key_0": (np.array([11]), np.array([22]))}
+    )
+
+    assert (
+        type(parameter_client._parameters["tuple_key_0"][0]).__name__ == "DeviceArray"
+    )
+    assert (
+        type(parameter_client._parameters["tuple_key_0"][1]).__name__ == "DeviceArray"
+    )
+    assert parameter_client._parameters["tuple_key_0"] == [
+        jax.numpy.array([11]),
+        jax.numpy.array([22]),
+    ]