add jax testcontainer and whisperjax folder

Author: Josephrp

modules/jax/testcontainers/jax/__init__.py

@@ -0,0 +1,68 @@
+import logging
+import urllib.request
+from urllib.error import URLError
+
+from core.testcontainers.core.container import DockerContainer
+from core.testcontainers.core.waiting_utils import wait_container_is_ready
+
+class JAXContainer(DockerContainer):
+    """
+    JAX container for GPU-accelerated numerical computing and machine learning.
+
+    Example:
+
+        .. doctest::
+
+            >>> import jax
+            >>> from testcontainers.jax import JAXContainer
+
+            >>> with JAXContainer("nvcr.io/nvidia/jax:23.08-py3") as jax_container:
+            ...     # Connect to the container
+            ...     jax_container.connect()
+            ...     
+            ...     # Run a simple JAX computation
+            ...     result = jax.numpy.add(1, 1)
+            ...     assert result == 2
+    """
+
+    def __init__(self, image="nvcr.io/nvidia/jax:23.08-py3", **kwargs):
+        super().__init__(image, **kwargs)
+        self.with_exposed_ports(8888)  # Expose Jupyter notebook port
+        self.with_env("NVIDIA_VISIBLE_DEVICES", "all")
+        self.with_env("CUDA_VISIBLE_DEVICES", "all")
+        self.with_kwargs(runtime="nvidia")  # Use NVIDIA runtime for GPU support
+
+    @wait_container_is_ready(URLError)
+    def _connect(self):
+        url = f"http://{self.get_container_host_ip()}:{self.get_exposed_port(8888)}"
+        res = urllib.request.urlopen(url)
+        if res.status != 200:
+            raise Exception(f"Failed to connect to JAX container. Status: {res.status}")
+
+    def connect(self):
+        """
+        Connect to the JAX container and ensure it's ready.
+        """
+        self._connect()
+        logging.info("Successfully connected to JAX container")
+
+    def get_jupyter_url(self):
+        """
+        Get the URL for accessing the Jupyter notebook server.
+        """
+        return f"http://{self.get_container_host_ip()}:{self.get_exposed_port(8888)}"
+
+    def run_jax_command(self, command):
+        """
+        Run a JAX command inside the container.
+        """
+        exec_result = self.exec(f"python -c '{command}'")
+        return exec_result
+
+    def start(self):
+        """
+        Start the JAX container.
+        """
+        super().start()
+        logging.info(f"JAX container started. Jupyter URL: {self.get_jupyter_url()}")
+        return self