Execution of Dashboard widgets

[PardhavMaradani]

Hi @orbeckst, @jeremyleung521, @amruthesht, @HeydenLabASU,

This is a discussion regarding the core execution of the dashboard widgets on the server side. #1 showed wireframes of how the widgets are added, updated, displayed and deleted from the dashboard UI.

The details below focus purely on the execution of widgets on the server side and doesn't include a lot of details about the actual interface with UI.

WidgetManager

The lifecycle of all the widgets is managed by a WidgetManager (singleton) on the server side. This is responsible for:

• Providing the list of widgets available (name, description) to the UI
• Creating a widget instance when a widget is added (selected from the list shown under Add Widget or when an existing widget is duplicated)
• Managing the lifecycle of each widget instance - running it, passing the output to the UI, handling changes to inputs made through UI, delete, etc

WidgetBase

There is a WidgetBase, which is the base class for all widgets. This class implements __init_subclass__ method to register all widget classes (all of which derive from WidgetBase) automatically with the WidgetManager. It is this class list that is used by the WidgetManager to provide the list of available widgets to the UI (from the name and description class atrributes of the corresponding widget classes), create instances using these classes and maintain the widget instances list (each widget instance is identified by a uuid), use this list to run and manage these instances.

Trajectory Iteration Loop

There is a core loop that gets created along with the universe creation (i.e., when connected to the live simulation) as a separate thread or an asyncio task. It uses events to handle Pause / Resume / Connect / Disconnect from the UI. trajectory.next() is used to iterate the trajectory in this loop. The buffer of Timestep's, which is used for batching is updated in this loop (see details in next section). The run_widgets() method of the WidgetManager is also run in this loop (see details in later section).

Timestep Buffer for Batch Analysis

The Timesteps buffer (global) is maintained as a collections.deque where the maxlen is what is configured in the Settings page. A copy of the the current timestep is made before it is added to this buffer. We need to make this buffer available to the various analysis widgets in a seamless manner so that it can be accessed in a natural way - for example: u.trajectory[-1] should change the timestep to the last but one frame with respect to the current frame, etc. Currently, streamed trajectories cannot be directly indexed this way and throw a Streamed trajectories must be an indexed using a slice exception. One way I was able to achieve this while exploring was by creating a simple wrapper (within the dashboard code) for the trajectory attribute and implementing the __getitem__ to set the timestep from this buffer. (A __getattr__ to pass through other attributes and a __dir__ to keep code completions working were also needed) Without this, passing the buffer to individual widget code and having them access it in a usable way would be very tedious.

Show example code

Here is an example of how the BufferedTrajectory wrapper could look like. This internally handles the circular buffer, provides direct negative indexing, updates buffer automatically with .next() iteration as shown below:

from collections import deque

class BufferedTrajectory:
    def __init__(self, trajectory):
        self._trajectory = trajectory
        self._maxlen = 100
        self._create_buffer(self._maxlen)
        BufferedTrajectory.next.__doc__ = type(trajectory).next.__doc__

    def _create_buffer(self, maxlen):
        self._buffer = deque(maxlen=self._maxlen)

    def __getitem__(self, index):
        if index > 0:
            raise ValueError("index should be <= 0")
        if index <= -1 * self._maxlen:
            raise ValueError("Out of range of buffer")
        ts = self._buffer[index - 1]
        self._trajectory.ts = ts
        return ts

    def __getattr__(self, name):
        wrapped = object.__getattribute__(self, "_trajectory")
        return getattr(wrapped, name)

    def __dir__(self):
        return dir(self._trajectory)

    def next(self):
        ts = self._trajectory.next()
        self._buffer.append(ts.copy())
        return ts

Here is how it is setup:

if not isinstance(u.trajectory, BufferedTrajectory):
    u.trajectory = BufferedTrajectory(u.trajectory)

WidgetManager.run_widgets()

The run_widgets method in the WidgetManager is the one responsible for executing all the widgets and passing over the output to the UI. It runs through all the widget instances, checks which of them are runnable at a given moment (based on whether the widget execution mode is per-frame or batch), checks whether the widget needs to be run in serial / parallel (through a is_parallelizable attribute) and in the later case, get a delayed task (joblib's delayed) that it adds to the list of parallel tasks. A separate list of instances which need to be run in serial is also created. All the parallel tasks are run using joblib.Parallel in a separate thread (which is joined later) before the serial instances are run. The n_jobs will be based on a configurable value in the Settings page. Note that these attributes for the execution mode and serial / parallel can also be made as inputs in the widget and changed at runtime as long as the underlying code supports them.

Paralellization with Streamed Trajectories

More details regarding the parallelization case: Streamed trajectories cannot be pickled today. We get a IMDReader does not support pickling exception when we try to do this. Removing this exception here, reveals that the core problem arises becuase of the network socket in imdclient with Python's cannot pickle 'socket' object exception. Pickling support is required if we need to use the universe as is within the parallelized task. Without this, the analysis code will have to be refactored quite a bit to remove all dependencies on the universe and make them independently executable. The best way I was able to get streamed trajectories picklable was by adding a custom __getstate__ method in IMDReader to remove the _imdclient attribute from the state (on the unpickling side, in __setstate__, set the _imdclient to None and added a warning). I was then able to create parallelized tasks easily with much more natural and simpler code based on our existing examples (as the analysis code don't / shouldn't need to iterate the streamed trajectory).

As stated at the beginning, the discussion above is about the core execution of the widgets on the server side and doesn't include a lot of details about communication with the UI etc, which are for a later time.

Could you please let me know if you have any thoughts, suggestions or feedback on the above. Thanks

[HeydenLabASU]

First of all, fantastic work on the widgets. The WidgetManager and WidgetBase organization looks very clean and thought out. This far exceeds my expectations. At this stage "a quick and dirty hack to get things running and refine later" would have been sufficient, but it looks you are building things right away with the big picture in mind.

Regarding pickling and parallelization, I am a little confused.
With "by adding a custom __getstate__ method in IMDReader to remove the _imdclient attribute from the state", did you mean removing the socket from a copy of the state? If yes, that sounds fine.
There is probably overhead involved with generating the copy, but enabling parallel processing should be worth it.

[PardhavMaradani]

did you mean removing the socket from a copy of the state?

Yes, I meant removing the unpicklable _imdclient from the copy of the state returned by __getstate__. Here is the exact code:

Show Code

    def __getstate__(self):
        state = self.__dict__.copy()
        del state["_imdclient"]
        # add warning about removing _imdclient from pickled state
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self._imdclient = None
        # add warning that _imdclient cannot be unpickled and being reset

So as to not add a dependency on MDAnalysis code / release, I verified that this can be patched from within the dashboard code if needed and also that it ran fine with joblib. (On a slightly related note, even if we had full support for AnalysisBase-based classes for streamed trajectories later, we'd still need serialization support to use the multiprocessing / dask backends)

[amruthesht]

This might not necessarily be the 'safest' way to do this, but is definitely a great hack! We can discuss more alternatives maybe during our weekly meetings.

[amruthesht]

Great work and start on the Dashboard backend. It seems you have given it enough thought and the org looks good to me.

The Timesteps buffer (global) is maintained as a collections.deque where the maxlen is what is configured in the Settings page. A copy of the the current timestep is made before it is added to this buffer. We need to make this buffer available to the various analysis widgets in a seamless manner so that it can be accessed in a natural way - for example: u.trajectory[-1] should change the timestep to the last but one frame with respect to the current frame, etc. Currently, streamed trajectories cannot be directly indexed this way and throw a Streamed trajectories must be an indexed using a slice exception. One way I was able to achieve this while exploring was by creating a simple wrapper (within the dashboard code) for the trajectory attribute and implementing the getitem to set the timestep from this buffer. (A getattr to pass through other attributes and a dir to keep code completions working were also needed) Without this, passing the buffer to individual widget code and having them access it in a usable way would be very tedious.

I like this deque buffer for analysis that require access to previous timestep data. I would however suggest skipping this in the first pass and focusing on simple single-frame analysis. We can discuss more but I would suggest inheriting more functions accessible to a regular trajectory attribute from the base MDA class, when implanting this in the future.

Trajectory Iteration Loop
There is a core loop that gets created along with the universe creation (i.e., when connected to the live simulation) as a separate thread or an asyncio task. It uses events to handle Pause / Resume / Connect / Disconnect from the UI. trajectory.next() is used to iterate the trajectory in this loop. The buffer of Timestep's, which is used for batching is updated in this loop (see details in next section). The run_widgets() method of the WidgetManager is also run in this loop (see details in later section).

For this master loop, is there an event manager or similar class that logs pause/resume etc. from the user and passes it to the loop?
Along similar lines -- did you give any thought to if any universe creation/destruction will happen after a disconnect is called or if you intend to preserve the universe so everything can be handled seamlessly for continued analysis/visualization? For the latter you might have to use imdclient functions that might not always be readily available through IMDReader

Looking forward to how this setup communicates with the UI of the dashboard ; and how you tackle related challenges, good luck!

[PardhavMaradani]

I would however suggest skipping this in the first pass and focusing on simple single-frame analysis.

As discussed during the call, this is a simple wrapper and can be added once the core functionality is done to support batching.

For this master loop, is there an event manager or similar class that logs pause/resume etc. from the user and passes it to the loop?

Yes, I was thinking of using the standard threading / asyncio events in the loop to control this behavior. The user actions from the UI set these events accordingly on the server side to achieve this.

Along similar lines -- did you give any thought to if any universe creation/destruction will happen after a disconnect is called or if you intend to preserve the universe so everything can be handled seamlessly for continued analysis/visualization?

So far, I was thinking that we will create / re-create the universe every time we connect / re-connect. The WidgetManager and all the widget instances will continue to exist as is and the new universe gets passed along during the run. It was for these use cases, we had the post_disconnect and pre_connect events in the proposal that widget classes can register to keep their analysis seamless (like clearing plot values, etc if needed). I will make a note and think through this again. Thanks

[PardhavMaradani]

All the relevant details from this discussion are now part of the documentation:

• Architecture: https://mdadash.readthedocs.io/en/latest/architecture.html
• Batching: https://mdadash.readthedocs.io/en/latest/batching.html
• Parallelization: https://mdadash.readthedocs.io/en/latest/parallelization.html