✨ progress bar using ipyvuetify for the jupyter notebook

packages/vaex-core/vaex/misc/progressbar.py

@@ -1,7 +1,12 @@
 from __future__ import print_function
 import os
 import sys
+import threading
 import time
+import multiprocessing
+import vaex.utils
+
+cpu_count = multiprocessing.cpu_count()
 
 
 class ProgressBarBase(object):
@@ -55,7 +60,7 @@ def info(self):
                 timeinfo = "estimated time: % 8.2fs = % 4.1fm = % 2.1fh" % (seconds, minutes, hours)
         else:
             timeinfo = "estimated time: unknown                "
-        return {"percentage":percentage, "timeinfo":timeinfo, "cpu_usage": cpu_usage, "title": self.title}
+        return {"percentage":percentage, "timeinfo":timeinfo, "cpu_usage": cpu_usage, "title": self.title, "cpu_count": cpu_count}
 
     def __repr__(self):
         output = ''
@@ -135,6 +140,68 @@ def finish(self):
         self(self.max_value)
 
 
+
+class Psutil(threading.Thread):
+    disk = ""
+    memory = ""
+    net = ""
+    wants_to_stop = False
+
+    def stop(self):
+        self.wants_to_stop = True
+
+    def run(self):
+        import psutil
+        t0 = time.time()
+        disk_prev = psutil.disk_io_counters()
+        net_prev = psutil.net_io_counters()
+        while not self.wants_to_stop:
+            available = psutil.virtual_memory().available
+            total = psutil.virtual_memory().total
+            used = total - available
+            usedp = int(used / total * 100)
+            availablep = int(available / total * 100)
+
+            total = vaex.utils.filesize_format(total)
+            available = vaex.utils.filesize_format(available)
+            used = vaex.utils.filesize_format(used)
+            self.memory = f'Memory:  total={total} avail={available} ({availablep}%) used={used} ({usedp}%)'
+
+            time.sleep(0.5)
+            disk_curr = psutil.disk_io_counters()
+            net_curr = psutil.net_io_counters()
+            t1 = time.time()
+
+            dt = t1 - t0
+            read = disk_curr.read_bytes - disk_prev.read_bytes
+            readps = vaex.utils.filesize_format(read / dt)
+            write = disk_curr.write_bytes - disk_prev.write_bytes
+            writeps = vaex.utils.filesize_format(write / dt)
+            self.disk = f'Disk:    read={readps}/s write={writeps}/s'
+
+            dt = t1 - t0
+            read = net_curr.bytes_recv -  net_prev.bytes_recv
+            readps = vaex.utils.filesize_format(read / dt)
+            write = net_curr.bytes_sent - net_prev.bytes_sent
+            writeps = vaex.utils.filesize_format(write / dt)
+            self.net = f'Network: read={readps}/s write={writeps}/s'
+
+            disk_prev = disk_curr
+            net_prev = net_curr
+            t0 = t1
+
+        t0 = time.time()
+
+    def __rich_console__(self, console, options):
+        yield self.memory
+        yield self.disk
+        yield self.net
+
+
+
+# TODO: make this configutable with new settings system
+show_extra = vaex.utils.get_env_type(bool, 'VAEX_PROGRESS_EXTRA', False)
+
 class ProgressBarRich(ProgressBarBase):
     def __init__(self, min_value, max_value, title=None, progress=None, indent=0, parent=None):
         super(ProgressBarRich, self).__init__(min_value, max_value)
@@ -159,20 +226,31 @@ def __init__(self, min_value, max_value, title=None, progress=None, indent=0, pa
             self.progress = progress
         if parent is None:
             self.node = rich.tree.Tree(self.progress)
+            # TODO: make this configutable
+            if show_extra:
+                self.info = Psutil(daemon=True)
+                self.info.start()
             from rich.live import Live
-            self.live = Live(self.node, refresh_per_second=5, console=self.console)
+            self.live = Live(self, refresh_per_second=3, console=self.console)
             self.live.start()
         else:
             self.node = parent.add(self.progress)
         # we do 1000 discrete steps
         self.steps = 0
         self.indent = indent
 
+        if len(title) > 60:
+            title = title[:60-3] + "..."
         padding = max(0, 50 - (self.indent * 4) - len(title))
         self.task = self.progress.add_task(f"[red]{title}" + (" " * padding), total=1000, start=False)
         self.started = False
         self.subtasks = []
 
+    def __rich_console__(self, console, options):
+        if show_extra:
+            yield self.info
+        yield self.node
+
     def add_child(self, parent, task, title):
         return ProgressBarRich(self.min_value, self.max_value, title, indent=self.indent+1, parent=self.node)
 
@@ -193,4 +271,108 @@ def update(self, value):
     def finish(self):
         self(self.max_value)
         if self.parent is None:
+            if show_extra:
+                self.info.stop()
             self.live.stop()
+
+
+import ipyvuetify as v
+import traitlets
+class ProgressDashboard(v.VuetifyTemplate):
+    items = traitlets.Any({}).tag(sync=True)
+    extra = traitlets.Any({}).tag(sync=True)
+    selected = traitlets.Unicode(default_value=None, allow_none=True).tag(sync=True)
+    @traitlets.default('template')
+    def _template(self):
+        return '''
+            <template>
+                <div>
+                    <pre>{{extra.memory}}
+{{extra.disk}}
+{{extra.net}}</pre>
+                    <v-treeview open-all dense :items="items">
+                    <template v-slot:prepend="{ item, open }">
+                        <v-progress-circular v-if='!item.finished' rotate="-90" :indeterminate="!item.started" :key="item.percentage" :size="30" :width="5" :value="item.percentage" :color="item.percentage < 100 ? 'primary' : 'green'">{{ item.started ? item.percentage : '' }}</v-progress-circular>
+                        <v-icon color="green" v-if='item.finished'>mdi-checkbox-marked-circle</v-icon>
+                    </template>
+                    <template v-slot:label="{ item, open }">
+                        {{item.title}} | {{item.timeinfo}} |
+                        <span :class="item.cpu_usage / item.cpu_count < 33 ? 'vaex-cpu-low' : (item.cpu_usage / item.cpu_count < 66 ? 'vaex-cpu-medium' : 'vaex-cpu-high')">{{Math.ceil(item.cpu_usage)}}% cpu</span>
+                        <!--
+                        <v-progress-linear v-model="item.cpu_usage / item.cpu_count" color="blue-grey" height="25" style="width: 100px">
+                            <template v-slot:default="{ value }">
+                                <strong>cpu: {{ Math.ceil(value) }}%</strong>
+                            </template>
+                        </v-progress-linear>                        
+                        -->
+                    </template>
+                    </v-treeview>
+                </div>
+            </template>
+            <style id="vaex-progress-dashboard-style">
+                .vaex-cpu-low {
+                    color: red;
+                }
+                .vaex-cpu-medium {
+                    color: orange;
+                }
+                .vaex-cpu-high {
+                    color: green;
+                }
+
+            </style>
+
+        '''
+
+class ProgressBarVuetify(ProgressBarBase):
+    def __init__(self, min_value, max_value, title=None, progress=None, indent=0, parent=None):
+        if len(title) > 60:
+            title = title[:60-3] + "..."
+        super(ProgressBarVuetify, self).__init__(min_value, max_value, title=title)
+        self.parent = parent
+        self.data = {'children': [], 'name': self.title, 'started': False, 'finished': False}
+        self.data.update(self.info())
+        self.parent = parent
+        if parent is None:
+            self.psutil = Psutil(daemon=True)
+            self.psutil.start()
+            self.data
+            self.dashboard = ProgressDashboard(items=[self.data])
+            from IPython.display import display
+            display(self.dashboard)
+        else:
+            self.dashboard = None
+        self.steps = 0
+        self.indent = indent
+
+
+    def add_child(self, parent, task, title):
+        child = ProgressBarVuetify(self.min_value, self.max_value, title, indent=self.indent+1, parent=self)
+        self.data['children'].append(child.data)
+        return child
+
+    def __call__(self, value):
+        if self.value == 0:
+            self.info()  # initialize
+        self.value = value
+        self.update_fraction()
+        info = self.info()
+        info['percentage'] = int(info['percentage'])
+        info['started'] = True
+        info['finished'] = self.value == self.max_value
+        self.data.update(info)
+        if self.parent is None:
+            self.dashboard.extra = {
+                'memory': self.psutil.memory,
+                'disk': self.psutil.disk,
+                'net': self.psutil.net,
+            }
+            self.dashboard.send_state('items')
+
+    def update(self, value):
+        self(value)
+
+    def finish(self):
+        self(self.max_value)
+        if self.parent is None:
+            self.psutil.stop()

packages/vaex-core/vaex/utils.py

@@ -288,12 +288,17 @@ def _progressbar_rich(type=None, title="processing", max_value=1):
     import vaex.misc.progressbar as pb
     return pb.ProgressBarRich(0, 1, title=title)
 
+def _progressbar_vuetify(type=None, title="processing", max_value=1):
+    import vaex.misc.progressbar as pb
+    return pb.ProgressBarVuetify(0, 1, title=title)
+
 
 _progressbar_typemap = {}
 _progressbar_typemap['progressbar2'] = _progressbar_progressbar2
 _progressbar_typemap['vaex'] = _progressbar_vaex
 _progressbar_typemap['widget'] = _progressbar_widget
 _progressbar_typemap['rich'] = _progressbar_rich
+_progressbar_typemap['vuetify'] = _progressbar_vuetify
 
 
 def progressbar(type_name=None, title="processing", max_value=1):