Repository: instadeepai/Mava

Issue:
[FEATURE] Core abstract system class for Jax based systems
### Feature
Abstract call to describe new system design for Jax systems.

### Proposal
The design should include the following functions:
* Configure - to specific systems components
* Update - to update system components
* Add - to add system components
* Launch - to run the system

### Testing
A detailed description of how this feature will be tested.
All of the above class functions should be tested.
#### configure: 
- input is a config file - a simpleNamespace of dataclasses
- return SimpleNamespace of components
#### update:
- input is a component and str name
- return None (just side effect of updating, i.e. "overwriting" component)
#### add:
- input is a component and str name
- return None
#### launch:
- Input is number of executors, boolean for multi-processing, nodes specification on device and system name.
- return None (simple runs system)

### Definition of done
Class is done when all of the above functions have been implemented, typed and decorated.

### Mandatory checklist before making a PR
* [x] The success criteria laid down in “Definition of done” are met.
* [x] Code is documented - docstrings for methods and classes, static types for arguments.
* [x] Code is tested - unit, integration and/or functional tests are added.
* [x] Documentation is updated - README, CONTRIBUTING, or other documentation.
* [x] All functional tests are green.



Discussion:

Resolving patch:
diff --git a/mava/core_jax.py b/mava/core_jax.py
new file mode 100644
--- /dev/null
+++ b/mava/core_jax.py
@@ -0,0 +1,71 @@
+# 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.
+
+
+"""Core Mava interfaces for Jax systems."""
+
+import abc
+from types import SimpleNamespace
+from typing import Any, List
+
+
+class BaseSystem(abc.ABC):
+    """Abstract system object."""
+
+    @abc.abstractmethod
+    def design(self) -> SimpleNamespace:
+        """System design specifying the list of components to use.
+
+        Returns:
+            system callback components
+        """
+
+    @abc.abstractmethod
+    def update(self, component: Any) -> None:
+        """Update a component that has already been added to the system.
+
+        Args:
+            component : system callback component
+        """
+
+    @abc.abstractmethod
+    def add(self, component: Any) -> None:
+        """Add a new component to the system.
+
+        Args:
+            component : system callback component
+        """
+
+    @abc.abstractmethod
+    def configure(self, **kwargs: Any) -> None:
+        """Configure system hyperparameters."""
+
+    @abc.abstractmethod
+    def launch(
+        self,
+        num_executors: int,
+        nodes_on_gpu: List[str],
+        multi_process: bool = True,
+        name: str = "system",
+    ) -> None:
+        """Run the system.
+
+        Args:
+            num_executors : number of executor processes to run in parallel
+            nodes_on_gpu : which processes to run on gpu
+            multi_process : whether to run single or multi process, single process runs
+                are primarily for debugging
+            name : name of the system
+        """


Test patch:
diff --git a/mava/core_jax_test.py b/mava/core_jax_test.py
new file mode 100644
--- /dev/null
+++ b/mava/core_jax_test.py
@@ -0,0 +1,43 @@
+# 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.
+
+
+"""Tests for core Mava interfaces for Jax systems."""
+
+from typing import Any
+
+import pytest
+
+from mava.core_jax import BaseSystem
+
+
+def test_exception_for_incomplete_child_class() -> None:
+    """Test if error is thrown for missing abstract class overwrites."""
+    with pytest.raises(TypeError):
+
+        class TestIncompleteDummySystem(BaseSystem):
+            def update(self, component: Any) -> None:
+                """Dummy update"""
+                pass
+
+            def add(self, component: Any) -> None:
+                """Dummy add"""
+                pass
+
+            def configure(self, **kwargs: Any) -> None:
+                """Dummy configure"""
+                pass
+
+        TestIncompleteDummySystem()  # type: ignore