Source code for pycif.plugins.obsoperators.standard.transforms.utils.dask.graph_status

import json
import logging
import os
import time
from datetime import datetime

from dask.callbacks import Callback




class GraphStatus(Callback):
    """Record the progress of the transforms and dump it to ``dask_status.js``.

    The file is loaded by ``dask_graph.html`` when the page is opened or
    reloaded. It is written when the computation starts, then at most every
    ``every_seconds`` seconds and/or every ``every_ntasks`` completed
    transforms (0 disables each criterion), and at the end of the
    computation, including when it fails.

    Only local schedulers (synchronous, threads, processes) trigger the
    recording; a distributed client does not.
    """

    def __init__(self, names, rundir, every_seconds=30, every_ntasks=0):
        super().__init__()

        # names maps Dask keys to graph node ids
        self.names = names
        self.out_file = os.path.join(rundir, "dask_status.js")
        self.every_seconds = every_seconds
        self.every_ntasks = every_ntasks

        self.tasks = {}
        self.t0 = time.monotonic()
        self.last_write = self.t0
        self.ndone_since_write = 0

    def _start(self, dsk):
        self.t0 = time.monotonic()
        self._write()

    def _pretask(self, key, dsk, state):
        if key not in self.names:
            return
        self.tasks[self.names[key]] = ["running", self._now(), None]
        self._maybe_write()

    def _posttask(self, key, result, dsk, state, id):
        if key not in self.names:
            return
        status = self.tasks.setdefault(self.names[key], [None, None, None])
        status[0] = "done"
        status[2] = self._now()
        self.ndone_since_write += 1
        self._maybe_write()

    def _finish(self, dsk, state, errored):
        # Tasks still running when the computation fails either raised
        # the error or were abandoned by Dask
        if errored:
            for status in self.tasks.values():
                if status[0] == "running":
                    status[0] = "failed"
        self._write(final=True, errored=errored)

    def _now(self):
        return round(time.monotonic() - self.t0, 3)

    def _maybe_write(self):
        if (self.every_ntasks
                and self.ndone_since_write >= self.every_ntasks) \
                or (self.every_seconds
                    and time.monotonic() - self.last_write
                    >= self.every_seconds):
            self._write()

    def _write(self, final=False, errored=False):
        status = {
            "updated": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "final": final,
            "errored": errored,
            "tasks": self.tasks,
        }

        # Write then rename, so that the page never reads a partial file.
        # Never let a failed write break the computation
        tmp_file = f"{self.out_file}.tmp"
        try:
            with open(tmp_file, "w") as f:
                f.write(f"window.cifStatus = {json.dumps(status)};\n")
            os.replace(tmp_file, self.out_file)
        except OSError as e:
            logging.warning(f"Could not write the Dask graph status: {e}")

        self.last_write = time.monotonic()
        self.ndone_since_write = 0