graphviper 0.0.1__py3-none-any.whl

graphviper/dask/__init__.py

@@ -0,0 +1 @@
+from .client import local_client, slurm_cluster_client

graphviper/dask/_scheduler.py

@@ -0,0 +1,208 @@
+"""
+MIT License
+
+Copyright (c) 2022 Jonathan Simon Kenyon
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+from collections import defaultdict
+from distributed import SchedulerPlugin
+from dask.core import reverse_dict
+from dask.base import tokenize
+from dask.order import graph_metrics, ndependencies
+import click
+from distributed.diagnostics.plugin import SchedulerPlugin
+import numpy as np
+
+
+def unravel_deps(hlg_deps, name, unravelled_deps=None):
+    """Recursively construct a set of all dependencies for a specific task."""
+
+    if unravelled_deps is None:
+        unravelled_deps = set()
+
+    for dep in hlg_deps[name]:
+        unravelled_deps |= {dep}
+        unravel_deps(hlg_deps, dep, unravelled_deps)
+
+    return unravelled_deps
+
+
+def get_node_depths(dependencies, root_nodes, metrics):
+    node_depths = {}
+
+    for k in dependencies.keys():
+        # Get dependencies per node.
+        deps = unravel_deps(dependencies, k)
+        # Associate nodes with root nodes.
+        roots = root_nodes & deps
+        offset = metrics[k][-1]
+        node_depths[k] = max(metrics[r][-1] - offset for r in roots) if roots else 0
+
+    return node_depths
+
+
+class schedular(SchedulerPlugin):
+    def __init__(self, autorestrictor, local_cache):
+        self.autorestrictor = autorestrictor
+        self.local_cache = local_cache
+        super().__init__()
+
+    def add_worker(self, scheduler, worker):
+        if self.local_cache:
+            # Set the resource label to the ip of the node that the worker is on, so that tasks that require a specific node can be assigned to the correct worker.
+            ip = worker[worker.rfind("/") + 1 : worker.rfind(":")]
+            scheduler.add_resources(worker=worker, resources={ip: 1})
+
+    def update_graph(self, scheduler, dsk=None, keys=None, restrictions=None, **kw):
+        if self.autorestrictor:
+            print("Using autorestrictor")
+            """Processes dependencies to assign tasks to specific workers."""
+            workers = list(scheduler.workers.keys())
+            n_worker = len(workers)
+
+            tasks = scheduler.tasks
+            dependencies = kw["dependencies"]
+
+            # print('In update_graph :', scheduler, ',*,', dsk, ',*,', keys , ',*,', restrictions , ',*,', kw)
+            if dependencies:
+                dependents = reverse_dict(dependencies)
+
+                # print('reversed dict:', dependents)
+
+                _, total_dependencies = ndependencies(dependencies, dependents)
+                # TODO: Avoid calling graph metrics.
+                metrics = graph_metrics(dependencies, dependents, total_dependencies)
+
+                # Terminal nodes have no dependents, root nodes have no dependencies.
+                # Horizontal partition nodes are initialized as the terminal nodes.
+                part_nodes = {k for (k, v) in dependents.items() if not v}
+                root_nodes = {k for (k, v) in dependencies.items() if not v}
+
+                # Figure out the depth of every task. Depth is defined as maximum
+                # distance from a root node. TODO: Optimize get_node_depths.
+
+                node_depths = get_node_depths(dependencies, root_nodes, metrics)
+                # try:
+                max_depth = max(node_depths.values())
+                # except:
+                #        print('&&&&& dependencies, root_nodes, metrics',node_depths,',*,',dependencies, root_nodes, metrics)
+
+                # If we have fewer partition nodes than workers, we cannot utilise all
+                # the workers and are likely dealing with a reduction. We work our way
+                # back through the graph, starting at the deepest terminal nodes, and
+                # try to find a depth at which there was enough work to utilise all
+                # workers.
+                while (len(part_nodes) < n_worker) & (max_depth > 0):
+                    _part_nodes = part_nodes.copy()
+                    for pn in _part_nodes:
+                        if node_depths[pn] == max_depth:
+                            part_nodes ^= set((pn,))
+                            part_nodes |= dependencies[pn]
+                    max_depth -= 1
+                    if max_depth <= 0:
+                        return  # In this case, there in nothing we can do - fall back.
+
+                part_roots = {}
+                part_dependencies = {}
+                part_dependents = {}
+
+                for pn in part_nodes:
+                    # Get dependencies per partition node.
+                    part_dependencies[pn] = unravel_deps(dependencies, pn)
+                    # Get dependents per partition node.
+                    part_dependents[pn] = unravel_deps(dependents, pn)
+                    # Associate partition nodes with root nodes.
+                    part_roots[pn] = root_nodes & part_dependencies[pn]
+
+                # Create a unique token for each set of partition roots. TODO: This is
+                # very strict. What about nodes with very similar roots? Tokenization
+                # may be overkill too.
+                root_tokens = {tokenize(*sorted(v)): v for v in part_roots.values()}
+
+                hash_map = defaultdict(set)
+                group_offset = 0
+
+                # Associate partition roots with a specific group if they are not a
+                # subset of another, larger root set.
+                for k, v in root_tokens.items():
+                    if any(v < vv for vv in root_tokens.values()):  # Strict subset.
+                        continue
+                    else:
+                        hash_map[k] |= set([group_offset])
+                        group_offset += 1
+
+                # If roots were a subset, they should share the group of their
+                # superset/s.
+                for k, v in root_tokens.items():
+                    if not v:  # Special case - no dependencies. Handled below.
+                        continue
+                    shared_roots = {
+                        kk: None for kk, vv in root_tokens.items() if v < vv
+                    }
+                    if shared_roots:
+                        hash_map[k] = set().union(
+                            *[hash_map[kk] for kk in shared_roots.keys()]
+                        )
+
+                task_groups = defaultdict(set)
+
+                for pn in part_nodes:
+                    pdp = part_dependencies[pn]
+                    pdn = part_dependents[pn]
+
+                    if pdp:
+                        groups = hash_map[tokenize(*sorted(part_roots[pn]))]
+                    else:  # Special case - no dependencies.
+                        groups = {group_offset}
+                        group_offset += 1
+
+                    for g in groups:
+                        task_groups[g] |= pdp | pdn | {pn}
+
+                worker_loads = {wkr: 0 for wkr in workers}
+
+                for task_group in task_groups.values():
+                    assignee = min(worker_loads, key=worker_loads.get)
+                    worker_loads[assignee] += len(task_group)
+
+                    for task_name in task_group:
+                        try:
+                            task = tasks[task_name]
+                        except KeyError:  # Keys may not have an assosciated task.
+                            continue
+
+                        # print('^^^^^^',dir(task))
+                        # if task._worker_restrictions is None:
+                        #    task._worker_restrictions = set()
+                        # task._worker_restrictions |= {assignee}
+                        # task._loose_restrictions = False
+
+                        if task.worker_restrictions is None:
+                            task.worker_restrictions = set()
+                        task.worker_restrictions |= {assignee}
+                        task.loose_restrictions = False
+
+
+@click.command()
+@click.option("--autorestrictor", default=False)
+@click.option("--local_cache", default=False)
+def dask_setup(scheduler, autorestrictor, local_cache):
+    plugin = schedular(autorestrictor, local_cache)
+    scheduler.add_plugin(plugin)

graphviper/dask/_worker.py

@@ -0,0 +1,71 @@
+import click
+
+from graphviper._logger import _setup_worker_logger
+
+
+class _worker:
+    def __init__(self, local_cache, log_parms):
+        # print("init local cache")
+        self.local_cache = local_cache
+
+        print("log_parms", log_parms)
+        # /.lustre/aoc/projects/ngvla/viper/ngvla_sim/viper_
+        self.log_to_term = log_parms["log_to_term"]
+        self.log_to_file = log_parms["log_to_file"]
+        self.log_file = log_parms["log_file"]
+        self.log_level = log_parms["log_level"]
+
+    def get_logger(self):
+        return self.logger
+
+    def setup(self, worker):
+        """
+        Run when the plugin is attached to a worker. This happens when the plugin is registered
+        and attached to existing workers, or when a worker is created after the plugin has been
+        registered.
+        """
+
+        self.logger = _setup_worker_logger(
+            self.log_to_term,
+            self.log_to_file,
+            self.log_file,
+            self.log_level,
+            str(worker.id),
+        )
+        self.logger.debug(
+            "Logger created on worker " + str(worker.id) + ",*," + str(worker.address)
+        )
+        # Documentation https://distributed.dask.org/en/stable/worker.html#distributed.worker.Worker
+        self.worker = worker
+
+        if self.local_cache:
+            ip = worker.address[
+                worker.address.rfind("/") + 1 : worker.address.rfind(":")
+            ]
+            self.logger.debug(str(worker.id) + ",*," + ip)
+            worker.state.available_resources = {
+                **worker.state.available_resources,
+                **{ip: 1},
+            }
+            # print(worker.state.available_resources)
+
+
+# https://github.com/dask/distributed/issues/4169
+@click.command()
+@click.option("--local_cache", default=False)
+# @click.option("--log_parms", default={'log_to_term':True,'log_to_file':False,'log_file':'viper_', 'log_level':'DEBUG'})
+@click.option("--log_to_term", default=True)
+@click.option("--log_to_file", default=False)
+@click.option("--log_file", default="viper_")
+@click.option("--log_level", default="INFO")
+async def dask_setup(
+    worker, local_cache, log_to_term, log_to_file, log_file, log_level
+):
+    log_parms = {
+        "log_to_term": log_to_term,
+        "log_to_file": log_to_file,
+        "log_file": log_file,
+        "log_level": log_level,
+    }
+    plugin = _worker(local_cache, log_parms)
+    await worker.client.register_worker_plugin(plugin, name="viper_worker")

graphviper/dask/client.py

@@ -0,0 +1,248 @@
+import warnings, time, os, psutil, multiprocessing, re
+import dask
+import copy
+import os
+import logging
+import graphviper
+import distributed
+from graphviper.parameter_checking.check_logger_parms import (
+    check_logger_parms,
+    check_worker_logger_parms,
+)
+from graphviper.logger import setup_logger, get_logger
+from graphviper.dask._worker import (
+    _worker,
+)  # _worker_logger_plugin
+
+
+def local_client(
+    cores=None,
+    memory_limit=None,
+    autorestrictor=False,
+    dask_local_dir=None,
+    local_dir=None,
+    wait_for_workers=True,
+    log_parms={},
+    worker_log_parms={},
+):
+    """
+    local_dir setting is only useful for testing since this function creates a local cluster. slurm_cluster_client should be used for a multinode cluster.
+
+    https://github.com/dask/dask/issues/5577
+    log_parms['log_to_term'] = True/False
+    log_parms['log_file'] = True/False
+    log_parms['log_level'] =
+    """
+
+    _log_parms = copy.deepcopy(log_parms)
+    _worker_log_parms = copy.deepcopy(worker_log_parms)
+
+    assert check_logger_parms(
+        _log_parms
+    ), "######### ERROR: initialize_processing log_parms checking failed."
+
+    if _worker_log_parms is not None:
+        assert check_worker_logger_parms(
+            _worker_log_parms
+        ), "######### ERROR: initialize_processing log_parms checking failed."
+
+    if local_dir:
+        os.environ["VIPER_LOCAL_DIR"] = local_dir
+        local_cache = True
+    else:
+        local_cache = False
+
+    # print(_log_parms)
+    setup_logger(**_log_parms)
+    logger = get_logger()
+
+    _set_up_dask(dask_local_dir)
+
+    viper_path = graphviper.__path__[0]
+    if local_cache or autorestrictor:
+        dask.config.set(
+            {
+                "distributed.scheduler.preload": os.path.join(
+                    viper_path, "_concurrency/_dask/_scheduler.py"
+                )
+            }
+        )
+        dask.config.set(
+            {
+                "distributed.scheduler.preload-argv": [
+                    "--local_cache",
+                    local_cache,
+                    "--autorestrictor",
+                    autorestrictor,
+                ]
+            }
+        )
+
+    """ This method of assigning a worker plugin does not seem to work when using dask_jobqueue. Consequently using client.register_worker_plugin so that the method of assigning a worker plugin is the same for local_client and slurm_cluster_client.
+    if local_cache or _worker_log_parms:
+        dask.config.set({"distributed.worker.preload": os.path.join(viper_path,'_utils/_worker.py')})
+        dask.config.set({"distributed.worker.preload-argv": ["--local_cache",local_cache,"--log_to_term",_worker_log_parms['log_to_term'],"--log_to_file",_worker_log_parms['log_to_file'],"--log_file",_worker_log_parms['log_file'],"--log_level",_worker_log_parms['log_level']]})
+    """
+    # setup distributed based multiprocessing environment
+    if cores is None:
+        cores = multiprocessing.cpu_count()
+    if memory_limit is None:
+        memory_limit = (
+            str(round(((psutil.virtual_memory().available / (1024**2))) / cores))
+            + "MB"
+        )
+    cluster = distributed.LocalCluster(
+        n_workers=cores, threads_per_worker=1, processes=True, memory_limit=memory_limit
+    )  # , silence_logs=logging.ERROR #,resources={'GPU': 2}
+    client = distributed.Client(cluster)
+    client.get_versions(check=True)
+
+    """
+    When constructing a graph that has local cache enabled all workers need to be up and running.
+    """
+    if local_cache or wait_for_workers:
+        client.wait_for_workers(n_workers=cores)
+
+    if local_cache or _worker_log_parms:
+        plugin = _worker(local_cache, _worker_log_parms)
+        client.register_worker_plugin(plugin, name="viper_worker")
+
+    logger.info("Created client " + str(client))
+
+    return client
+
+
+def slurm_cluster_client(
+    workers_per_node,
+    cores_per_node,
+    memory_per_node,
+    number_of_nodes,
+    queue,
+    interface,
+    python_env_dir,
+    dask_local_dir,
+    dask_log_dir,
+    exclude_nodes="nmpost090",
+    dashboard_port=9000,
+    local_dir=None,
+    autorestrictor=False,
+    wait_for_workers=True,
+    log_parms={},
+    worker_log_parms={},
+):
+    """
+    local_cache setting is only useful for testing since this function creates a local cluster. slurm_cluster_client should be used for a multinode cluster.
+
+    https://github.com/dask/dask/issues/5577
+    log_parms['log_to_term'] = True/False
+    log_parms['log_file'] = True/False
+    log_parms['log_level'] =
+
+    interface eth0, ib0
+    python "/mnt/condor/jsteeb/viper_py/bin/python"
+    dask_local_dir "/mnt/condor/jsteeb"
+    dask_log_dir "/.lustre/aoc/projects/ngvla/viper/ngvla_sim",
+    """
+
+    from dask_jobqueue import SLURMCluster
+    from distributed import Client, config, performance_report
+
+    _log_parms = copy.deepcopy(log_parms)
+    _worker_log_parms = copy.deepcopy(worker_log_parms)
+
+    assert _check_logger_parms(
+        _log_parms
+    ), "######### ERROR: initialize_processing log_parms checking failed."
+    assert _check_worker_logger_parms(
+        _worker_log_parms
+    ), "######### ERROR: initialize_processing log_parms checking failed."
+
+    if local_dir:
+        os.environ["VIPER_LOCAL_DIR"] = local_dir
+        local_cache = True
+    else:
+        local_cache = False
+
+    # Viper logger for code that is not part of the Dask graph. The worker logger is setup in the _viper_worker plugin.
+    from viper._utils._logger import setup_logger
+
+    setup_logger(**_log_parms)
+    logger = get_logger()
+
+    _set_up_dask(dask_local_dir)
+
+    viper_path = graphviper.__path__.__dict__["_path"][0]
+    if local_cache or autorestrictor:
+        dask.config.set(
+            {
+                "distributed.scheduler.preload": os.path.join(
+                    viper_path, "_concurrency/_dask/_scheduler.py"
+                )
+            }
+        )
+        dask.config.set(
+            {
+                "distributed.scheduler.preload-argv": [
+                    "--local_cache",
+                    local_cache,
+                    "--autorestrictor",
+                    autorestrictor,
+                ]
+            }
+        )
+
+    """ This method of assigning a worker plugin does not seem to work when using dask_jobqueue. Consequently using client.register_worker_plugin so that the method of assigning a worker plugin is the same for local_client and slurm_cluster_client.
+    if local_cache or _worker_log_parms:
+        dask.config.set({"distributed.worker.preload": os.path.join(viper_path,'_utils/_worker.py')})
+        dask.config.set({"distributed.worker.preload-argv": ["--local_cache",local_cache,"--log_to_term",_worker_log_parms['log_to_term'],"--log_to_file",_worker_log_parms['log_to_file'],"--log_file",_worker_log_parms['log_file'],"--log_level",_worker_log_parms['log_level']]})
+    """
+
+    cluster = SLURMCluster(
+        processes=workers_per_node,
+        cores=cores_per_node,
+        interface=interface,
+        memory=memory_per_node,
+        walltime="24:00:00",
+        queue=queue,
+        name="viper",
+        python=python_env_dir,  # "/mnt/condor/jsteeb/viper_py/bin/python", #"/.lustre/aoc/projects/ngvla/viper/viper_py_env/bin/python",
+        local_directory=dask_local_dir,  # "/mnt/condor/jsteeb",
+        log_directory=dask_log_dir,
+        job_extra_directives=["--exclude=" + exclude_nodes],
+        # job_extra_directives=["--exclude=nmpost087,nmpost089,nmpost088"],
+        scheduler_options={"dashboard_address": ":" + str(dashboard_port)},
+    )  # interface='ib0'
+
+    client = Client(cluster)
+
+    cluster.scale(workers_per_node * number_of_nodes)
+
+    """
+    When constructing a graph that has local cache enabled all workers need to be up and running.
+    """
+    if local_cache or wait_for_workers:
+        client.wait_for_workers(n_workers=workers_per_node * number_of_nodes)
+
+    if local_cache or _worker_log_parms:
+        plugin = _worker(local_cache, _worker_log_parms)
+        client.register_worker_plugin(plugin, name="viper_worker")
+
+    logger.info("Created client " + str(client))
+
+    return client
+
+
+def _set_up_dask(local_directory):
+    if local_directory:
+        dask.config.set({"temporary_directory": local_directory})
+    dask.config.set({"distributed.scheduler.allowed-failures": 10})
+    dask.config.set({"distributed.scheduler.work-stealing": True})
+    dask.config.set({"distributed.scheduler.unknown-task-duration": "99m"})
+    dask.config.set({"distributed.worker.memory.pause": False})
+    dask.config.set({"distributed.worker.memory.terminate": False})
+    # dask.config.set({"distributed.worker.memory.recent-to-old-time": '999s'})
+    dask.config.set({"distributed.comm.timeouts.connect": "3600s"})
+    dask.config.set({"distributed.comm.timeouts.tcp": "3600s"})
+    dask.config.set({"distributed.nanny.environ.OMP_NUM_THREADS": 1})
+    dask.config.set({"distributed.nanny.environ.MKL_NUM_THREADS": 1})
+    # https://docs.dask.org/en/stable/how-to/customize-initialization.html

graphviper/graph_tools/__init__.py

@@ -0,0 +1,8 @@
+from .coordinate_utils import (
+    make_time_coord,
+    make_frequency_coord,
+    make_parallel_coord,
+    interpolate_data_coords_onto_parallel_coords,
+)
+from .map import map
+from .reduce import reduce