Repository: instadeepai/Mava

Issue:
[TEST] Jax MAPPO System Integration
### What do you want to test?
Jax MAPPO full integration test

### Outline of test structure
* Unit tests (if possible)
* Test component interactions
* Ensure not crashing during standard executor and trainer steps

### Definition of done
Passing checks, cover all basic component interactions, 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/launcher.py b/mava/systems/jax/launcher.py
--- a/mava/systems/jax/launcher.py
+++ b/mava/systems/jax/launcher.py
@@ -119,6 +119,7 @@ def add(
 
         if self._multi_process:
             with self._program.group(name):
+                node_fn = copy_node_fn(node_fn)
                 node = self._program.add_node(node_type(node_fn, *arguments))
             return node
         else:


Test patch:
diff --git a/tests/jax/components/building/distributor_test.py b/tests/jax/components/building/distributor_test.py
--- a/tests/jax/components/building/distributor_test.py
+++ b/tests/jax/components/building/distributor_test.py
@@ -47,7 +47,7 @@ def __init__(self) -> None:
         self.program_launched = False
 
     def data_server(self) -> List[Any]:
-        """Data server to test no multi_process in on_building_program_nodes method
+        """Data server to test on_building_program_nodes method
 
         Returns:
             tables: fake table composed of reverb_server tables
@@ -63,13 +63,13 @@ def data_server(self) -> List[Any]:
         ]
 
     def parameter_server(self) -> str:
-        """parameter_server to test no multi_process in on_building_program_nodes"""
+        """parameter_server to test on_building_program_nodes"""
         return "Parameter Server Test"
 
     def executor(
         self, executor_id: str, data_server_client: Any, parameter_server_client: Any
     ) -> str:
-        """Executor to test no multi_process in on_building_program_nodes method"""
+        """Executor to test on_building_program_nodes method"""
         if executor_id == "evaluator":
             return "Evaluator Test"
 
@@ -79,7 +79,7 @@ def executor(
     def trainer(
         self, trainer_id: str, data_server_client: Any, parameter_server_client: Any
     ) -> str:
-        """Trainer to test no multi_process in on_building_program_nodes method"""
+        """Trainer to test on_building_program_nodes method"""
         return "Trainer Test"
 
     def launch(self) -> None:
@@ -119,27 +119,34 @@ def test_on_building_program_nodes_multi_process(
         "evaluator",
         "trainer",
     ]
+
+    data_server_fn = mock_builder.store.program._program._groups["data_server"][
+        -1
+    ]._priority_tables_fn
+    assert str(repr(data_server_fn).split(" ")[2].split(".")[-1]) == "data_server"
     assert (
-        mock_builder.store.program._program._groups["data_server"][
+        mock_builder.store.program._program._groups["parameter_server"][
             -1
-        ]._priority_tables_fn
-        == mock_builder.data_server
-    )
-    assert (
-        mock_builder.store.program._program._groups["parameter_server"][-1]._constructor
-        == mock_builder.parameter_server
+        ]._constructor()
+        == "Parameter Server Test"
     )
     assert (
-        mock_builder.store.program._program._groups["executor"][-1]._constructor
-        == mock_builder.executor
+        mock_builder.store.program._program._groups["executor"][-1]._constructor(
+            "executor", "fake_data_server", "fake_parameter_server"
+        )
+        == "Executor Test"
     )
     assert (
-        mock_builder.store.program._program._groups["evaluator"][-1]._constructor
-        == mock_builder.executor
+        mock_builder.store.program._program._groups["evaluator"][-1]._constructor(
+            "evaluator", "fake_data_server", "fake_parameter_server"
+        )
+        == "Evaluator Test"
     )
     assert (
-        mock_builder.store.program._program._groups["trainer"][-1]._constructor
-        == mock_builder.trainer
+        mock_builder.store.program._program._groups["trainer"][-1]._constructor(
+            "trainer", "fake_data_server", "fake_parameter_server"
+        )
+        == "Trainer Test"
     )
 
     with pytest.raises(Exception):
@@ -162,23 +169,28 @@ def test_on_building_program_nodes_multi_process_no_evaluator(
         "executor",
         "trainer",
     ]
+    data_server_fn = mock_builder.store.program._program._groups["data_server"][
+        -1
+    ]._priority_tables_fn
+    assert str(repr(data_server_fn).split(" ")[2].split(".")[-1]) == "data_server"
+
     assert (
-        mock_builder.store.program._program._groups["data_server"][
+        mock_builder.store.program._program._groups["parameter_server"][
             -1
-        ]._priority_tables_fn
-        == mock_builder.data_server
-    )
-    assert (
-        mock_builder.store.program._program._groups["parameter_server"][-1]._constructor
-        == mock_builder.parameter_server
+        ]._constructor()
+        == "Parameter Server Test"
     )
     assert (
-        mock_builder.store.program._program._groups["executor"][-1]._constructor
-        == mock_builder.executor
+        mock_builder.store.program._program._groups["executor"][-1]._constructor(
+            "executor", "fake_data_server", "fake_parameter_server"
+        )
+        == "Executor Test"
     )
     assert (
-        mock_builder.store.program._program._groups["trainer"][-1]._constructor
-        == mock_builder.trainer
+        mock_builder.store.program._program._groups["trainer"][-1]._constructor(
+            "trainer", "fake_data_server", "fake_parameter_server"
+        )
+        == "Trainer Test"
     )
 
     with pytest.raises(Exception):
diff --git a/tests/jax/core_system_components/launcher_test.py b/tests/jax/core_system_components/launcher_test.py
--- a/tests/jax/core_system_components/launcher_test.py
+++ b/tests/jax/core_system_components/launcher_test.py
@@ -15,62 +15,14 @@
 
 """Tests for launcher class for Jax-based Mava systems"""
 
-from typing import Any, Callable, List
-
 import launchpad as lp
 import pytest
-from reverb import Client, item_selectors, pybind, rate_limiters
-from reverb import server as reverb_server
+from reverb import Client, pybind
 
 from mava.systems.jax.launcher import Launcher, NodeType
 from tests.jax.components.building.distributor_test import MockBuilder
 
 
-@pytest.fixture
-def mock_data_server_fn() -> Callable:
-    """call data_server function"""
-
-    def data_server() -> List[Any]:
-        """data_server
-
-        Returns:
-                tables: fake table composed of reverb_server tables
-        """
-        return [
-            reverb_server.Table(
-                name="table_0",
-                sampler=item_selectors.Prioritized(priority_exponent=1),
-                remover=item_selectors.Fifo(),
-                max_size=1000,
-                rate_limiter=rate_limiters.MinSize(1),
-            )
-        ]
-
-    return data_server
-
-
-@pytest.fixture
-def mock_parameter_server_fn() -> Callable:
-    """call parameter_server function"""
-
-    def parameter_server() -> str:
-        """Fake parameter server function"""
-        return "test_parameter_server"
-
-    return parameter_server
-
-
-@pytest.fixture
-def mock_parameter_server_second_fn() -> Callable:
-    """call the second parameter_server function"""
-
-    def parameter_server_second() -> str:
-        """Another fake parameter server function"""
-        return "test_parameter_server_second_mock"
-
-    return parameter_server_second
-
-
 @pytest.fixture
 def mock_builder() -> MockBuilder:
     """Mock builder"""
@@ -120,7 +72,7 @@ def test_initiator_non_multi_process() -> None:
     assert not hasattr(launcher, "_program")
 
 
-def test_add_multi_process(mock_data_server_fn: Callable) -> None:
+def test_add_multi_process(mock_builder: MockBuilder) -> None:
     """Test add method in the Launcher for the case of multi process
 
     Args:
@@ -128,16 +80,19 @@ def test_add_multi_process(mock_data_server_fn: Callable) -> None:
     """
     launcher = Launcher(multi_process=True)
     data_server = launcher.add(
-        mock_data_server_fn,
+        mock_builder.data_server,
         node_type=NodeType.reverb,
         name="data_server_test",
     )
 
     assert list(launcher._program._groups.keys()) == ["data_server_test"]
-    assert (
-        launcher._program._groups["data_server_test"][-1]._priority_tables_fn
-        == mock_data_server_fn
-    )
+
+    # Make sure the node have data_server method
+    data_server_fn = launcher._program._groups["data_server_test"][
+        -1
+    ]._priority_tables_fn
+    assert str(repr(data_server_fn).split(" ")[2].split(".")[-1]) == "data_server"
+
     assert [data_server] == launcher._program._groups["data_server_test"][
         -1
     ]._created_handles
@@ -146,9 +101,7 @@ def test_add_multi_process(mock_data_server_fn: Callable) -> None:
     assert not hasattr(launcher, "_node_dict")
 
 
-def test_add_multi_process_two_add_calls(
-    mock_data_server_fn: Callable, mock_parameter_server_fn: Callable
-) -> None:
+def test_add_multi_process_two_add_calls(mock_builder: MockBuilder) -> None:
     """Test calling add more than one time method in the Launcher for the case of multi process
 
     Args:
@@ -157,12 +110,12 @@ def test_add_multi_process_two_add_calls(
     """
     launcher = Launcher(multi_process=True)
     data_server = launcher.add(
-        mock_data_server_fn,
+        mock_builder.data_server,
         node_type=NodeType.reverb,
         name="data_server_test",
     )
     parameter_server = launcher.add(
-        mock_parameter_server_fn,
+        mock_builder.parameter_server,
         node_type=NodeType.courier,
         name="parameter_server_test",
     )
@@ -172,17 +125,19 @@ def test_add_multi_process_two_add_calls(
         "parameter_server_test",
     ]
 
-    assert (
-        launcher._program._groups["data_server_test"][-1]._priority_tables_fn
-        == mock_data_server_fn
-    )
+    # Make sure the node have data_server method
+    data_server_fn = launcher._program._groups["data_server_test"][
+        -1
+    ]._priority_tables_fn
+    assert str(repr(data_server_fn).split(" ")[2].split(".")[-1]) == "data_server"
+
     assert [data_server] == launcher._program._groups["data_server_test"][
         -1
     ]._created_handles
 
     assert (
-        launcher._program._groups["parameter_server_test"][-1]._constructor
-        == mock_parameter_server_fn
+        launcher._program._groups["parameter_server_test"][-1]._constructor()
+        == "Parameter Server Test"
     )
     assert [parameter_server] == launcher._program._groups["parameter_server_test"][
         -1
@@ -192,9 +147,7 @@ def test_add_multi_process_two_add_calls(
     assert not hasattr(launcher, "_node_dict")
 
 
-def test_add_multi_process_two_add_same_name(
-    mock_parameter_server_fn: Callable, mock_parameter_server_second_fn: Callable
-) -> None:
+def test_add_multi_process_two_add_same_name(mock_builder: MockBuilder) -> None:
     """Test calling twice add for two nodes with same name for the case of multi process
 
     Args:
@@ -204,24 +157,24 @@ def test_add_multi_process_two_add_same_name(
     launcher = Launcher(multi_process=True)
 
     parameter_server_1 = launcher.add(
-        mock_parameter_server_fn,
+        mock_builder.parameter_server,
         node_type=NodeType.courier,
         name="parameter_server_test",
     )
     parameter_server_2 = launcher.add(
-        mock_parameter_server_second_fn,
+        mock_builder.parameter_server,
         node_type=NodeType.courier,
         name="parameter_server_test",
     )
 
     assert list(launcher._program._groups.keys()) == ["parameter_server_test"]
     assert (
-        launcher._program._groups["parameter_server_test"][0]._constructor
-        == mock_parameter_server_fn
+        launcher._program._groups["parameter_server_test"][0]._constructor()
+        == "Parameter Server Test"
     )
     assert (
-        launcher._program._groups["parameter_server_test"][1]._constructor
-        == mock_parameter_server_second_fn
+        launcher._program._groups["parameter_server_test"][1]._constructor()
+        == "Parameter Server Test"
     )
     assert [parameter_server_1] == launcher._program._groups["parameter_server_test"][
         0
diff --git a/tests/jax/systems/ippo_system_test.py b/tests/jax/systems/ippo_system_test.py
--- a/tests/jax/systems/ippo_system_test.py
+++ b/tests/jax/systems/ippo_system_test.py
@@ -1,8 +1,28 @@
+# python3
+# Copyright 2021 InstaDeep Ltd. 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import functools
+from datetime import datetime
+from typing import Any
 
-import acme
+import launchpad as lp
 import optax
 import pytest
+from launchpad.launch.test_multi_threading import (
+    address_builder as test_address_builder,
+)
 
 from mava.systems.jax import ippo
 from mava.systems.jax.system import System
@@ -15,16 +35,14 @@
 
 @pytest.fixture
 def test_full_system() -> System:
-    """Add description here."""
+    """Full mava system fixture for testing"""
     return ippo.IPPOSystem()
 
 
-# TODO: fix test
-@pytest.mark.skip(reason="test is currently breaking ci pipeline")
-def test_except_trainer(
+def test_ippo(
     test_full_system: System,
 ) -> None:
-    """Test if the parameter server instantiates processes as expected."""
+    """Full integration test of ippo system."""
 
     # Environment.
     environment_factory = functools.partial(
@@ -34,11 +52,17 @@ def test_except_trainer(
     )
 
     # Networks.
-    network_factory = ippo.make_default_networks
+    def network_factory(*args: Any, **kwargs: Any) -> Any:
+        return ippo.make_default_networks(  # type: ignore
+            policy_layer_sizes=(32, 32),
+            critic_layer_sizes=(64, 64),
+            *args,
+            **kwargs,
+        )
 
     # Checkpointer appends "Checkpoints" to checkpoint_dir.
     base_dir = "~/mava"
-    mava_id = "12345"
+    mava_id = str(datetime.now())
     checkpoint_subpath = f"{base_dir}/{mava_id}"
 
     # Log every [log_every] seconds.
@@ -54,7 +78,8 @@ def test_except_trainer(
 
     # Optimizer.
     optimizer = optax.chain(
-        optax.clip_by_global_norm(40.0), optax.scale_by_adam(), optax.scale(-1e-4)
+        optax.clip_by_global_norm(40.0),
+        optax.adam(1e-4),
     )
 
     # Build the system
@@ -62,28 +87,27 @@ def test_except_trainer(
         environment_factory=environment_factory,
         network_factory=network_factory,
         logger_factory=logger_factory,
-        checkpoint_subpath=checkpoint_subpath,
+        experiment_path=checkpoint_subpath,
         optimizer=optimizer,
-        executor_parameter_update_period=20,
-        multi_process=False,
+        executor_parameter_update_period=1,
+        multi_process=True,
         run_evaluator=True,
         num_executors=1,
+        max_queue_size=500,
         use_next_extras=False,
-        sample_batch_size=2,
+        sample_batch_size=5,
+        nodes_on_gpu=[],
+        lp_launch_type=lp.LaunchType.TEST_MULTI_THREADING,
     )
 
-    (
-        data_server,
-        parameter_server,
-        executor,
-        evaluator,
-        trainer,
-    ) = test_full_system._builder.store.system_build
-
-    assert isinstance(executor, acme.core.Worker)
+    (trainer_node,) = test_full_system._builder.store.program._program._groups[
+        "trainer"
+    ]
+    trainer_node.disable_run()
+    test_address_builder.bind_addresses([trainer_node])
 
-    # Step the executor
-    executor.run_episode()
+    test_full_system.launch()
+    trainer_run = trainer_node.create_handle().dereference()
 
-    # Step the trainer
-    trainer.step()
+    for _ in range(5):
+        trainer_run.step()