Repository: instadeepai/Mava

Issue:
[INVESTIGATION] The multi-thread vs multi-process 
### What do you want to investigate?
In MAVA currently we have three options **SIngle process**, **Multi process** and **Multi thread**.

To run the multi-thread case, we need to have `multi_process` attribute equal True, and to  fill the `lp_launch_type` attribute with `lp.LaunchType.LOCAL_MULTI_THREADING`. However, if in the future we are going to consider the case of multi thread we need to identify an attribute for the `multi_thread` independent of  themulti-procss case.

### Definition of done
- Are we only going to keep the single-process and multi-process cases, or will we also keep the code for the multi-thread case?
- If we want to keep the multi-thread case, we should change the launcher code.


Discussion:

Resolving patch:
diff --git a/examples/jax/debugging/simple_spread/feedforward/decentralised/run_ippo_single_process.py b/examples/jax/debugging/simple_spread/feedforward/decentralised/run_ippo_single_process.py
--- a/examples/jax/debugging/simple_spread/feedforward/decentralised/run_ippo_single_process.py
+++ b/examples/jax/debugging/simple_spread/feedforward/decentralised/run_ippo_single_process.py
@@ -18,8 +18,6 @@
 import os
 from datetime import datetime
 from typing import Any
-
-import launchpad as lp
 import optax
 from absl import app, flags
 
@@ -100,7 +98,6 @@ def network_factory(*args: Any, **kwargs: Any) -> Any:
         num_executors=1,
         max_queue_size=500,
         sample_batch_size=5,
-        lp_launch_type=lp.LaunchType.LOCAL_MULTI_PROCESSING,
     )
     system.launch()
 
diff --git a/mava/components/jax/building/distributor.py b/mava/components/jax/building/distributor.py
--- a/mava/components/jax/building/distributor.py
+++ b/mava/components/jax/building/distributor.py
@@ -34,7 +34,6 @@ class DistributorConfig:
     run_evaluator: bool = True
     distributor_name: str = "System"
     terminal: str = "current_terminal"
-    lp_launch_type: Union[str, lp.LaunchType] = lp.LaunchType.LOCAL_MULTI_PROCESSING
     single_process_max_episodes: Optional[int] = None
 
 
@@ -66,7 +65,6 @@ def on_building_program_nodes(self, builder: SystemBuilder) -> None:
             nodes_on_gpu=self.config.nodes_on_gpu,
             name=self.config.distributor_name,
             terminal=self.config.terminal,
-            lp_launch_type=self.config.lp_launch_type,
             single_process_max_episodes=self.config.single_process_max_episodes,
         )
 
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
@@ -43,9 +43,6 @@ def __init__(
         single_process_max_episodes: Optional[int] = None,
         name: str = "System",
         terminal: str = "current_terminal",
-        lp_launch_type: Union[
-            str, lp.LaunchType
-        ] = lp.LaunchType.LOCAL_MULTI_PROCESSING,
     ) -> None:
         """Initialise the launcher.
 
@@ -63,7 +60,6 @@ def __init__(
                 before termination.
             name : launchpad program name.
             terminal : terminal for launchpad processes to be shown on.
-            lp_launch_type: launchpad launch type.
         """
         self._multi_process = multi_process
         self._name = name
@@ -71,7 +67,6 @@ def __init__(
         self._single_process_evaluator_period = single_process_evaluator_period
         self._single_process_max_episodes = single_process_max_episodes
         self._terminal = terminal
-        self._lp_launch_type = lp_launch_type
         if multi_process:
             self._program = lp.Program(name=name)
             self._nodes_on_gpu = nodes_on_gpu
@@ -173,7 +168,7 @@ def launch(self) -> None:
 
             lp.launch(
                 self._program,
-                launch_type=self._lp_launch_type,
+                launch_type=lp.LaunchType.LOCAL_MULTI_PROCESSING,
                 terminal=self._terminal,
                 local_resources=local_resources,
             )


Test patch:
diff --git a/tests/jax/integration/mock_systems.py b/tests/jax/integration/mock_systems.py
--- a/tests/jax/integration/mock_systems.py
+++ b/tests/jax/integration/mock_systems.py
@@ -93,69 +93,6 @@ def network_factory(*args: Any, **kwargs: Any) -> Any:
     return test_system
 
 
-def mock_system_multi_thread() -> System:
-    """Mock of ippo system in the multi thread case"""
-    # Environment.
-    environment_factory = functools.partial(
-        debugging_utils.make_environment,
-        env_name="simple_spread",
-        action_space="discrete",
-    )
-
-    # 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 = str(datetime.now())
-    checkpoint_subpath = f"{base_dir}/{mava_id}"
-
-    # Log every [log_every] seconds.
-    log_every = 10
-    logger_factory = functools.partial(
-        logger_utils.make_logger,
-        directory=base_dir,
-        to_terminal=True,
-        to_tensorboard=True,
-        time_stamp=mava_id,
-        time_delta=log_every,
-    )
-
-    # Optimizer.
-    optimizer = optax.chain(
-        optax.clip_by_global_norm(40.0),
-        optax.adam(1e-4),
-    )
-
-    # Create ippo system
-    test_system = ippo.IPPOSystem()
-
-    # Build the system
-    test_system.build(
-        environment_factory=environment_factory,
-        network_factory=network_factory,
-        logger_factory=logger_factory,
-        experiment_path=checkpoint_subpath,
-        optimizer=optimizer,
-        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=5,
-        nodes_on_gpu=[],
-        lp_launch_type=lp.LaunchType.TEST_MULTI_THREADING,
-    )
-    return test_system
-
-
 def mock_system_multi_process() -> System:
     """Mock of ippo system in the multi process case"""
     # Environment.
@@ -213,6 +150,5 @@ def network_factory(*args: Any, **kwargs: Any) -> Any:
         use_next_extras=False,
         sample_batch_size=5,
         nodes_on_gpu=[],
-        lp_launch_type=lp.LaunchType.LOCAL_MULTI_PROCESSING,
     )
     return test_system
diff --git a/tests/jax/integration/trainer_test.py b/tests/jax/integration/trainer_test.py
--- a/tests/jax/integration/trainer_test.py
+++ b/tests/jax/integration/trainer_test.py
@@ -15,17 +15,12 @@
 
 """Integration test of the Trainer for Jax-based Mava"""
 
-import os
-import signal
-import time
-
 import jax.numpy as jnp
 import pytest
 
 from mava.systems.jax import System
 from tests.jax.integration.mock_systems import (
     mock_system_multi_process,
-    mock_system_multi_thread,
     mock_system_single_process,
 )
 
@@ -42,14 +37,9 @@ def test_system_mp() -> System:
     return mock_system_multi_process()
 
 
-@pytest.fixture
-def test_system_mt() -> System:
-    """A multi thread built system"""
-    return mock_system_multi_thread()
-
-
 def test_trainer_single_process(test_system_sp: System) -> None:
     """Test if the trainer instantiates processes as expected."""
+    # extract the nodes
     (
         data_server,
         parameter_server,
@@ -58,8 +48,7 @@ def test_trainer_single_process(test_system_sp: System) -> None:
         trainer,
     ) = test_system_sp._builder.store.system_build
 
-    for _ in range(5):
-        executor.run_episode()
+    executor.run_episode()
 
     # Before run step method
     for net_key in trainer.store.networks["networks"].keys():
@@ -78,49 +67,29 @@ def test_trainer_single_process(test_system_sp: System) -> None:
             assert not jnp.all(categorical_value_head["w"] == 0)
 
 
-def test_trainer_multi_thread(test_system_mt: System) -> None:
-    """Test if the trainer instantiates processes as expected."""
-    # Disable the run of the trainer node
-    (trainer_node,) = test_system_mt._builder.store.program._program._groups["trainer"]
-    trainer_node.disable_run()
-
-    # launch the system
-    test_system_mt.launch()
-    time.sleep(10)  # wait till the executor has run
-
-    trainer = trainer_node._construct_instance()
-
-    # Before run step function
-    for net_key in trainer.store.networks["networks"].keys():
-        mu = trainer.store.opt_states[net_key][1][0][-1]  # network
-        for categorical_value_head in mu.values():
-            assert jnp.all(categorical_value_head["b"] == 0)
-            assert jnp.all(categorical_value_head["w"] == 0)
-
-    # Step function
-    trainer.step()
-
-    # Check that the trainer update the network
-    for net_key in trainer.store.networks["networks"].keys():
-        mu = trainer.store.opt_states[net_key][1][0][-1]
-        for categorical_value_head in mu.values():
-            assert not jnp.all(categorical_value_head["b"] == 0)
-            assert not jnp.all(categorical_value_head["w"] == 0)
-
-
 def test_trainer_multi_process(test_system_mp: System) -> None:
     """Test if the trainer instantiates processes as expected."""
-    # Disable the run of the trainer node
+    # Disable the nodes
     (trainer_node,) = test_system_mp._builder.store.program._program._groups["trainer"]
+    (executor_node,) = test_system_mp._builder.store.program._program._groups[
+        "executor"
+    ]
+    (evaluator_node,) = test_system_mp._builder.store.program._program._groups[
+        "evaluator"
+    ]
+    (parameter_server_node,) = test_system_mp._builder.store.program._program._groups[
+        "parameter_server"
+    ]
     trainer_node.disable_run()
+    executor_node.disable_run()
+    evaluator_node.disable_run()
+    parameter_server_node.disable_run()
 
-    # pid to help stop the launcher once the test ends
-    pid = os.getpid()
     # launch the system
     test_system_mp.launch()
 
-    time.sleep(10)  # wait till the executor has run
-
+    # Extract the executor and trainer instance
+    executor = executor_node._construct_instance()
     trainer = trainer_node._construct_instance()
 
     # Before run step function
@@ -130,7 +99,9 @@ def test_trainer_multi_process(test_system_mp: System) -> None:
             assert jnp.all(categorical_value_head["b"] == 0)
             assert jnp.all(categorical_value_head["w"] == 0)
 
-    # Step function
+    # Run the executor and the trainer
+    for _ in range(3):
+        executor.run_episode()
     trainer.step()
 
     # Check that the trainer update the network
@@ -139,6 +110,3 @@ def test_trainer_multi_process(test_system_mp: System) -> None:
         for categorical_value_head in mu.values():
             assert not jnp.all(categorical_value_head["b"] == 0)
             assert not jnp.all(categorical_value_head["w"] == 0)
-
-    # stop the launcher
-    os.kill(pid, signal.SIGTERM)