mdadash 0.0.1__py3-none-any.whl

mdadash/__init__.py

@@ -0,0 +1,9 @@
+"""
+mdadash
+Dashboard for tracking and analyzing live MD simulations with streaming.
+"""
+
+# Add imports here
+from importlib.metadata import version
+
+__version__ = version("mdadash")

mdadash/backend/example_api.py

@@ -0,0 +1,3 @@
+def example_function(number: float) -> float:
+    """Return `2*number`. This is and example for tests."""
+    return 2 * number

mdadash/backend/kernel/core.py

@@ -0,0 +1,210 @@
+import asyncio
+
+import comm
+import MDAnalysis as mda
+
+
+class CommHandler:
+    """Comm Handler
+
+    This class is responsible for handling all the communication to and from
+    this kernel. This is the class that interfaces with the KernelManager on
+    the server side.
+
+    """
+
+    def __init__(self):
+        self._comm = None
+        self._handlers = {}
+        comm.get_comm_manager().register_target(
+            "kernel_comm_handler", self._handle_comm_open
+        )
+
+    def register_handler(self, msg_type: str, handler_func: callable) -> None:
+        """Register a handler function for a message type"""
+        self._handlers[msg_type] = handler_func
+
+    def send(self, msg: dict) -> None:
+        """Send a message (response) back on the comm
+
+        Parameters
+        ----------
+        msg: dict
+            A message dictionary
+
+        """
+        if self._comm is not None:
+            self._comm.send(msg)
+        else:
+            raise ValueError("comm is not open yet")  # pragma: no cover
+
+    def _handle_comm_open(self, _comm: comm.base_comm.BaseComm, _msg):
+        """Internal: Handler when the comm is opened (comm_open)"""
+        self._comm = _comm
+        # set the handler for comm messages (comm_msg)
+        self._comm.on_msg(self._handle_msg)
+
+    def _handle_msg(self, msg):
+        """Internal: Dispatch the message to the registered handler"""
+        content_data = msg["content"]["data"]
+        msg_type = content_data["msg_type"]
+        if msg_type in self._handlers:
+            self._handlers[msg_type](content_data["data"])
+        else:
+            error_msg = f"{msg_type} does not have a registered handler"
+            self.send({"status": "error", "message": error_msg})
+            raise ValueError(error_msg)
+
+
+class UniverseManager:
+    """Universe Manager
+
+    This class is responsible for managing all MDAnalysis universes. It has
+    handlers to interact with the MD simulation. These handlers are invoked by
+    comm messages sent from the server.
+
+    This also provides an iterable and indexable access to the individual
+    universes.
+
+    """
+
+    def __init__(self):
+        self._universes = []
+        self._iter_loop_task = None
+        self._iter_loop_running = False
+        self._iter_loop_resumed = asyncio.Event()
+        self._iter_loop_resumed.clear()
+
+    def __iter__(self) -> iter:
+        """To support iteration"""
+        return iter(self._universes)
+
+    def __len__(self) -> int:
+        """Number of universes"""
+        return len(self._universes)
+
+    def __getitem__(self, index: int):
+        """Return universe based on index"""
+        # numeric index based array access
+        _max = len(self._universes)
+        if 0 <= index < _max:
+            return self._universes[index]
+        raise ValueError(f"Invalid index {index} of {_max} items")
+
+    def init_n_universes(self, n: int) -> None:
+        """Initialize array for n universes
+
+        Parameters
+        ----------
+        n: int
+            Number of universes to initialize
+
+        """
+        self._universes = [None] * n
+
+    def connect_to_simulations(self, universe_configs: list[dict]) -> None:
+        """Connect to MD simulations
+
+        Parameters
+        ----------
+        universe_configs: list[dict]
+            A list of configurations for universe(s) creation.
+            Each dict has universe related config like topology, trajectory,
+            imdclient params, user-defined kwargs etc
+
+        """
+        try:
+            for uid, config in enumerate(universe_configs):
+                kwargs = {}
+                topology = config.get("topology")
+                trajectory = config.get("trajectory")
+                for key, value in config.items():
+                    if key in ("topology", "trajectory", "kwargs"):
+                        continue
+                    if value is not None:
+                        kwargs[key] = value
+                for name, value in config["kwargs"]:
+                    if name.strip():
+                        kwargs[name] = value
+                # create universe
+                u = mda.Universe(
+                    topology,
+                    trajectory,
+                    **kwargs,
+                )
+                if uid == 0:
+                    self._send_tsdata(u)
+                self._universes[uid] = u
+            # start iter loop for trajectories
+            self._iter_loop_resumed.clear()
+            self._iter_loop_running = True
+            self._iter_loop_task = asyncio.create_task(self._iter_loop())
+            comm_handler.send({"status": "ok"})
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            comm_handler.send({"status": "error", "message": str(e)})
+
+    def _send_tsdata(self, u: mda.Universe):
+        """Internal: Send timestep data out"""
+        comm_handler.send(
+            {
+                "tsinfo": {
+                    "frame": u.trajectory.frame,
+                    "tsdata": u.trajectory.ts.data,
+                }
+            }
+        )
+
+    def disconnect_from_simulations(self, _data: dict) -> None:
+        """Disconnect from MD simulations"""
+        self._iter_loop_running = False
+        self._iter_loop_task.cancel()
+        for u in self._universes:
+            u.trajectory.close()
+        comm_handler.send({"status": "ok"})
+
+    def pause_simulations(self, _data: dict) -> None:
+        """Pause MD simulations"""
+        self._iter_loop_resumed.clear()
+        comm_handler.send({"status": "ok"})
+
+    def resume_simulations(self, _data: dict) -> None:
+        """Resume MD simulations"""
+        self._iter_loop_resumed.set()
+        comm_handler.send({"status": "ok"})
+
+    def _trajectory_next(self, u):
+        """Internal: Iterate trajectory by 1 frame"""
+        return u.trajectory.next()
+
+    async def _iter_loop(self):
+        """Internal: Iteration loop for trajectories"""
+        try:
+            while self._iter_loop_running:
+                await self._iter_loop_resumed.wait()
+                for uid, u in enumerate(self._universes):
+                    try:
+                        # iterate in thread to not block on a network call here
+                        await asyncio.to_thread(self._trajectory_next, u)
+                        if uid == 0:
+                            self._send_tsdata(u)
+                        # await asyncio.sleep(0)
+                    except StopIteration as e:  # pragma: no cover
+                        print(e)
+                    await asyncio.sleep(0)
+        except asyncio.CancelledError:
+            pass
+
+
+def init_n_universes(data: dict) -> None:
+    um.init_n_universes(data.get("n"))
+
+
+um = UniverseManager()
+comm_handler = CommHandler()
+comm_handler.register_handler("init_n_universes", init_n_universes)
+comm_handler.register_handler("connect_to_simulations", um.connect_to_simulations)
+comm_handler.register_handler(
+    "disconnect_from_simulations", um.disconnect_from_simulations
+)
+comm_handler.register_handler("pause_simulations", um.pause_simulations)
+comm_handler.register_handler("resume_simulations", um.resume_simulations)

mdadash/backend/kernel/manager.py

@@ -0,0 +1,305 @@
+import asyncio
+import logging
+import queue
+import sys
+import uuid
+
+import socketio
+from jupyter_client import AsyncKernelManager
+
+from ..state.manager import StateManager
+
+logger = logging.getLogger(__name__)
+
+
+# pylint: disable=too-many-instance-attributes
+class KernelManager:
+    """Kernel Manager
+
+    This class is responsible for managing the AsyncKernelManager (async kernel)
+    that runs all the MDAnalysis code. It takes care of starting the async
+    kernel, stopping it and communicating with it. It interfaces with the
+    CommHandler on the kernel side for messaging.
+
+    """
+
+    def __init__(self, sm: StateManager, sio: socketio.AsyncServer):
+        self.sm = sm
+        self.sio = sio
+        self.km = AsyncKernelManager(kernel_name="python3")
+        self.kc = None
+        self._pending_futures = {}
+        self._is_running = False
+        self.comm_id = uuid.uuid4().hex
+        self.listen_task = None
+
+    async def start(self) -> None:
+        """Start the async kernel"""
+        # start the kernel
+        await self.km.start_kernel()
+        # create a client
+        self.kc = self.km.client()
+        self.kc.start_channels()
+        await self.kc.wait_for_ready()
+        # create task to listen on iopub and shell channels
+        self.listen_task = asyncio.create_task(self._start_listening())
+        # initialize the kernel core
+        self.kc.execute("from mdadash.backend.kernel import core")
+        # open comms with the kernel
+        self._comm_open()
+        self._is_running = True
+        # initialize n universes in kernel universe manager
+        await self.send_message(
+            "init_n_universes", {"n": len(self.sm.universe_configs)}
+        )
+
+    async def stop(self) -> None:
+        """Stop the async kernel"""
+        self._is_running = False
+        # wait for listen task to completely exit
+        await self.listen_task
+        self.kc.stop_channels()
+        self.kc = None
+        # shutdown kernel gracefully
+        await self.km.shutdown_kernel(now=False)
+        await self.km.cleanup_resources()
+
+    async def _start_listening(self):
+        """Internal: Create separate listen tasks for iopub and shell"""
+        await asyncio.gather(
+            self._listen_iopub_channel(),
+            self._listen_shell_channel(),
+        )
+
+    async def _emit_tsdata(self, tsinfo):
+        """Internal: Emit timestep data"""
+        tsdata = tsinfo["tsdata"]
+        step = tsdata.get("step", None)
+        total_steps = self.sm.universe_configs[0].get("total_steps", None)
+        done = (step / total_steps) * 100 if step and total_steps else None
+        timestep_info = {
+            "frame": tsinfo.get("frame", None),
+            "time": tsdata.get("time", None),
+            "step": step,
+            "done": done,
+            "energies": {
+                "temperature": tsdata.get("temperature", None),
+                "total_energy": tsdata.get("total_energy", None),
+                "potential_energy": tsdata.get("potential_energy", None),
+                "van_der_walls_energy": tsdata.get("van_der_walls_energy", None),
+                "coulomb_energy": tsdata.get("coulomb_energy", None),
+                "bonds_energy": tsdata.get("bonds_energy", None),
+                "angles_energy": tsdata.get("angles_energy", None),
+                "dihedrals_energy": tsdata.get("dihedrals_energy", None),
+                "improper_dihedrals_energy": tsdata.get(
+                    "improper_dihedrals_energy", None
+                ),
+            },
+        }
+        await self.sio.emit("timestepInfo", timestep_info)
+
+    # pylint: disable=too-many-branches
+    async def _listen_iopub_channel(self):
+        """Internal: Listen on iopub channel"""
+        while self._is_running:
+            try:
+                msg = await self.kc.iopub_channel.get_msg(timeout=0.1)
+                msg_type = msg["header"]["msg_type"]
+                content = msg["content"]
+                parent_id = msg.get("parent_header", {}).get("msg_id")
+                # check if a pending future can be resolved with msg
+                resolve_future = False
+                if parent_id and parent_id in self._pending_futures:
+                    future = self._pending_futures[parent_id]
+                    if not future.done():
+                        resolve_future = True
+                # handle different msg_type's
+                if msg_type == "comm_msg":
+                    data = msg["content"]["data"]
+                    if "tsinfo" in data:
+                        await self._emit_tsdata(data["tsinfo"])
+                    elif resolve_future:
+                        future.set_result(data)
+                        continue
+                elif msg_type == "stream":
+                    if resolve_future:
+                        future.set_result(msg["content"]["text"])
+                        continue
+                    # redirect kernel stdout and stderr to this server output
+                    if content["name"] == "stdout" or content["name"] == "stderr":
+                        output = content["text"]
+                        file = sys.stdout if content["name"] == "stdout" else sys.stderr
+                        print(
+                            f"KERNEL ({content['name']}): {output}", end="", file=file
+                        )
+                elif msg_type == "error":
+                    # redirect kernel errors to server output
+                    print(f"KERNEL (error): {content['ename']}: {content['evalue']}")
+                    if resolve_future:
+                        future.set_result(content["evalue"])
+                        continue
+                else:
+                    logger.debug("IOPUB: %s", msg)
+                # TODO: handle other message types
+            except (asyncio.TimeoutError, queue.Empty):
+                continue
+
+    async def _listen_shell_channel(self):
+        """Internal: Listen on shell channel"""
+        while self._is_running:
+            try:
+                msg = await self.kc.shell_channel.get_msg(timeout=0.1)
+                # msg_type = msg["header"]["msg_type"]
+                # content = msg["content"]
+                logger.debug("SHELL: %s", msg)
+            except (asyncio.TimeoutError, queue.Empty):
+                continue
+
+    def _comm_open(self):
+        """Internal: Open comms with the kernel"""
+        content = {
+            "comm_id": self.comm_id,
+            "target_name": "kernel_comm_handler",
+            "data": {"msg_type": "handshake"},
+        }
+        open_msg = self.kc.session.msg("comm_open", content=content)
+        self.kc.shell_channel.send(open_msg)
+
+    async def send_message(self, msg_type: str, data: dict) -> None:
+        """Send message to kernel and don't await a response
+
+        Parameters
+        ----------
+        msg_type: str
+            A message type string that the kernel has a handler registered for
+
+        data: dict
+            Dict that gets passed to the handler in the kernel
+
+        """
+        content = {
+            "comm_id": self.comm_id,
+            "target_name": "kernel_comm_handler",
+            "data": {"msg_type": msg_type, "data": data},
+        }
+        data_msg = self.kc.session.msg("comm_msg", content=content)
+        self.kc.shell_channel.send(data_msg)
+
+    async def send_message_await_response(
+        self, msg_type: str, data: dict = None, timeout: int = 5
+    ) -> dict | None:
+        """Send message to kernel and wait for a response (async)
+
+        Parameters
+        ----------
+        msg_type: str
+            A message type string that the kernel has a handler registered for
+
+        data: dict
+            Dict that gets passed to the handler in the kernel
+
+        timeout: int
+            Timeout in seconds
+
+        """
+        content = {
+            "comm_id": self.comm_id,
+            "target_name": "kernel_comm_handler",
+            "data": {"msg_type": msg_type, "data": data},
+        }
+        data_msg = self.kc.session.msg("comm_msg", content=content)
+        msg_id = data_msg["header"]["msg_id"]
+        # add to the _pending_futures to that it gets resolved when the
+        # response arrives on the iopub channel
+        future = asyncio.get_running_loop().create_future()
+        self._pending_futures[msg_id] = future
+        self.kc.shell_channel.send(data_msg)
+        try:
+            return await asyncio.wait_for(future, timeout=timeout)
+        except asyncio.TimeoutError as e:  # pragma: no cover
+            raise TimeoutError("Timed out waiting for kernel response") from e
+        finally:
+            self._pending_futures.pop(msg_id, None)
+
+    async def execute_code(self, code: str, timeout: int = 5) -> str:
+        """Execute code in the kernel
+
+        Parameters
+        ----------
+        code: str
+            Code to execute in the kernel
+
+        timeout: int
+            Timeout in seconds
+
+        """
+        msg_id = self.kc.execute(code)
+        future = asyncio.get_running_loop().create_future()
+        self._pending_futures[msg_id] = future
+        try:
+            return await asyncio.wait_for(future, timeout=timeout)
+        except asyncio.TimeoutError as e:  # pragma: no cover
+            raise TimeoutError("Timed out waiting for kernel execute response") from e
+        finally:
+            self._pending_futures.pop(msg_id, None)
+
+    async def connect_to_simulations(self) -> dict:
+        """Connect to the MD simulation
+
+        Returns
+        -------
+        response: dict
+            Response dict indicating status
+
+        """
+        response = await self.send_message_await_response(
+            "connect_to_simulations", self.sm.universe_configs
+        )
+        if response["status"] == "ok":
+            self.sm.running_state["connected"] = True
+            self.sm.running_state["running"] = False
+        return response
+
+    async def disconnect_from_simulations(self) -> dict:
+        """Disconnect from the MD simulation
+
+        Returns
+        -------
+        response: dict
+            Response dict indicating status
+
+        """
+        response = await self.send_message_await_response(
+            "disconnect_from_simulations", {}
+        )
+        if response["status"] == "ok":
+            self.sm.running_state["connected"] = False
+        return response
+
+    async def pause_simulations(self) -> dict:
+        """Pause MD simulations
+
+        Returns
+        -------
+        response: dict
+            Response dict indicating status
+
+        """
+        response = await self.send_message_await_response("pause_simulations", {})
+        if response["status"] == "ok":
+            self.sm.running_state["running"] = False
+        return response
+
+    async def resume_simulations(self) -> dict:
+        """Resume MD simulations
+
+        Returns
+        -------
+        response: dict
+            Response dict indicating status
+
+        """
+        response = await self.send_message_await_response("resume_simulations", {})
+        if response["status"] == "ok":
+            self.sm.running_state["running"] = True
+        return response