예제 #1
0
def clear_all():
    c = IPythonClient()
    if c is None:
        return
    try:
        v = c[:]
        mods = clear(v)
    finally:
        c.close()
    return mods
예제 #2
0
def clear_all():
    c = IPythonClient()
    if c is None:
        return
    try:
        v = c[:]
        mods = clear(v)
    finally:
        c.close()
    return mods
예제 #3
0
def cleanup_all(module_name, prefix):
    """ Connects to all engines and runs ``cleanup()`` on them. """
    c = IPythonClient()
    if c is None:
        return
    try:
        v = c[:]
        cleanup(v, module_name, prefix)
    finally:
        c.close()
예제 #4
0
def cleanup_all(module_name, prefix):
    """ Connects to all engines and runs ``cleanup()`` on them. """
    c = IPythonClient()
    if c is None:
        return
    try:
        v = c[:]
        cleanup(v, module_name, prefix)
    finally:
        c.close()
예제 #5
0
    def __init__(self, client=None, targets=None):

        if client is None:
            self.client = IPythonClient()
            self.owns_client = True
        else:
            self.client = client
            self.owns_client = False

        self.view = self.client[:]
        self.nengines = len(self.view)

        self.all_targets = sorted(self.view.targets)
        if targets is None:
            self.targets = self.all_targets
        else:
            self.targets = []
            for target in targets:
                if target not in self.all_targets:
                    raise ValueError("Target %r not registered" % target)
                else:
                    self.targets.append(target)
        self.targets = sorted(self.targets)

        # local imports
        self.view.execute("from functools import reduce; "
                          "from importlib import import_module; "
                          "import distarray.localapi; "
                          "import distarray.localapi.mpiutils; "
                          "import distarray.utils; "
                          "import distarray.localapi.proxyize as proxyize; "
                          "from distarray.localapi.proxyize import Proxy; "
                          "import numpy")

        self.context_key = self._setup_context_key()

        # setup proxyize which is used by context.apply in the rest of the
        # setup.
        cmd = "proxyize = proxyize.Proxyize()"
        self.view.execute(cmd)

        self._base_comm = self._make_base_comm()
        self._comm_from_targets = {
            tuple(sorted(self.view.targets)): self._base_comm
        }
        self.comm = self.make_subcomm(self.targets)

        if not BaseContext._CLEANUP:
            BaseContext._CLEANUP = (atexit.register(ipython_cleanup.clear_all),
                                    atexit.register(
                                        ipython_cleanup.cleanup_all,
                                        '__main__', DISTARRAY_BASE_NAME))
예제 #6
0
    def __init__(self, client=None, targets=None):

        if client is None:
            self.client = IPythonClient()
            self.owns_client = True
        else:
            self.client = client
            self.owns_client = False

        self.view = self.client[:]
        self.nengines = len(self.view)

        self.all_targets = sorted(self.view.targets)
        if targets is None:
            self.targets = self.all_targets
        else:
            self.targets = []
            for target in targets:
                if target not in self.all_targets:
                    raise ValueError("Target %r not registered" % target)
                else:
                    self.targets.append(target)
        self.targets = sorted(self.targets)

        # local imports
        self.view.execute("from functools import reduce; "
                          "from importlib import import_module; "
                          "import distarray.localapi; "
                          "import distarray.localapi.mpiutils; "
                          "import distarray.utils; "
                          "import distarray.localapi.proxyize as proxyize; "
                          "from distarray.localapi.proxyize import Proxy; "
                          "import numpy")

        self.context_key = self._setup_context_key()

        # setup proxyize which is used by context.apply in the rest of the
        # setup.
        cmd = "proxyize = proxyize.Proxyize()"
        self.view.execute(cmd)

        self._base_comm = self._make_base_comm()
        self._comm_from_targets = {tuple(sorted(self.view.targets)): self._base_comm}
        self.comm = self.make_subcomm(self.targets)

        if not BaseContext._CLEANUP:
            BaseContext._CLEANUP = (atexit.register(ipython_cleanup.clear_all),
                                    atexit.register(ipython_cleanup.cleanup_all,
                                                    '__main__',
                                                    DISTARRAY_BASE_NAME))
예제 #7
0
class IPythonContext(BaseContext):
    """
    Context class that uses IPython.parallel.

    See the docstring for  `BaseContext` for more information about Contexts.

    See also
    --------
    BaseContext
    """
    def __init__(self, client=None, targets=None):

        if client is None:
            self.client = IPythonClient()
            self.owns_client = True
        else:
            self.client = client
            self.owns_client = False

        self.view = self.client[:]
        self.nengines = len(self.view)

        self.all_targets = sorted(self.view.targets)
        if targets is None:
            self.targets = self.all_targets
        else:
            self.targets = []
            for target in targets:
                if target not in self.all_targets:
                    raise ValueError("Target %r not registered" % target)
                else:
                    self.targets.append(target)
        self.targets = sorted(self.targets)

        # local imports
        self.view.execute("from functools import reduce; "
                          "from importlib import import_module; "
                          "import distarray.localapi; "
                          "import distarray.localapi.mpiutils; "
                          "import distarray.utils; "
                          "import distarray.localapi.proxyize as proxyize; "
                          "from distarray.localapi.proxyize import Proxy; "
                          "import numpy")

        self.context_key = self._setup_context_key()

        # setup proxyize which is used by context.apply in the rest of the
        # setup.
        cmd = "proxyize = proxyize.Proxyize()"
        self.view.execute(cmd)

        self._base_comm = self._make_base_comm()
        self._comm_from_targets = {
            tuple(sorted(self.view.targets)): self._base_comm
        }
        self.comm = self.make_subcomm(self.targets)

        if not BaseContext._CLEANUP:
            BaseContext._CLEANUP = (atexit.register(ipython_cleanup.clear_all),
                                    atexit.register(
                                        ipython_cleanup.cleanup_all,
                                        '__main__', DISTARRAY_BASE_NAME))

    def make_subcomm(self, new_targets):

        if new_targets != sorted(new_targets):
            raise ValueError("targets must be in sorted order.")

        try:
            return self._comm_from_targets[tuple(new_targets)]
        except KeyError:
            pass

        def _make_new_comm(rank_list, base_comm):
            import distarray.localapi.mpiutils as mpiutils
            res = mpiutils.create_comm_with_list(rank_list, base_comm)
            return proxyize(res)  # noqa

        new_comm = self.apply(_make_new_comm, (new_targets, self._base_comm),
                              targets=self.view.targets)[0]

        self._comm_from_targets[tuple(new_targets)] = new_comm
        return new_comm

    def _make_base_comm(self):
        """
        Returns a proxy for an MPI communicator that encompasses all targets in
        self.view.targets (not self.targets, which can be a subset).
        """
        def get_rank():
            from distarray.localapi.mpiutils import get_comm_private
            return get_comm_private().Get_rank()

        # self.view's engines must encompass all ranks in the MPI communicator,
        # i.e., everything in rank_map.values().
        def get_size():
            from distarray.localapi.mpiutils import get_comm_private
            return get_comm_private().Get_size()

        # get a mapping of IPython engine ID to MPI rank
        rank_from_target = self.view.apply_async(get_rank).get_dict()
        ranks = [rank_from_target[target] for target in self.view.targets]

        comm_size = self.view.apply_async(get_size).get()[0]
        if set(rank_from_target.values()) != set(range(comm_size)):
            raise ValueError('Engines in view must encompass all MPI ranks.')

        # create a new communicator with the subset of ranks. Note that
        # create_comm_with_list() must be called on all engines, not just those
        # involved in the new communicator.  This is because
        # create_comm_with_list() issues a collective MPI operation.
        def _make_new_comm(rank_list):
            import distarray.localapi.mpiutils as mpiutils
            new_comm = mpiutils.create_comm_with_list(rank_list)
            if not mpiutils.get_base_comm():
                mpiutils.set_base_comm(new_comm)
            return proxyize(new_comm)  # noqa

        return self.apply(_make_new_comm,
                          args=(ranks, ),
                          targets=self.view.targets)[0]

    # Key management routines:
    def cleanup(self):
        """ Delete keys that this context created from all the engines. """
        # TODO: FIXME: cleanup needs updating to work with proxy objects.
        ipython_cleanup.cleanup(view=self.view,
                                module_name='__main__',
                                prefix=self.context_key)

    def close(self):
        self.cleanup()

        def free_subcomm(subcomm):
            subcomm.Free()

        for targets, subcomm in self._comm_from_targets.items():
            self.apply(free_subcomm, (subcomm, ), targets=targets)
        if self.owns_client:
            self.client.close()
        self.comm = None

    # End of key management routines.

    def _execute(self, lines, targets):
        return self.view.execute(lines, targets=targets, block=True)

    def _push(self, d, targets):
        return self.view.push(d, targets=targets, block=True)

    def apply(self,
              func,
              args=None,
              kwargs=None,
              targets=None,
              autoproxyize=False):
        """
        Analogous to IPython.parallel.view.apply_sync

        Parameters
        ----------
        func : function
        args : tuple
            positional arguments to func
        kwargs : dict
            key word arguments to func
        targets : sequence of integers
            engines func is to be run on.
        autoproxyize: bool, default False
            If True, implicitly return a Proxy object from the function.

        Returns
        -------
            return a list of the results on the each engine.
        """
        def func_wrapper(func, apply_nonce, context_key, args, kwargs,
                         autoproxyize):
            """
            Function which calls the applied function after grabbing all the
            arguments on the engines that are passed in as names of the form
            `__distarray__<some uuid>`.
            """
            from importlib import import_module
            import types
            from distarray.metadata_utils import arg_kwarg_proxy_converter
            from distarray.localapi import LocalArray

            main = import_module('__main__')
            main.proxyize.set_state(apply_nonce)

            # Modify func to change the namespace it executes in.
            # but builtins don't have __code__, __globals__, etc.
            if not isinstance(func, types.BuiltinFunctionType):
                # get func's building  blocks first
                func_code = func.__code__
                func_name = func.__name__
                func_defaults = func.__defaults__
                func_closure = func.__closure__

                # build the func's new execution environment
                main.__dict__.update({'context_key': context_key})
                new_func_globals = main.__dict__
                # create the new func
                func = types.FunctionType(func_code, new_func_globals,
                                          func_name, func_defaults,
                                          func_closure)

            args, kwargs = arg_kwarg_proxy_converter(args, kwargs)
            result = func(*args, **kwargs)

            if autoproxyize and isinstance(result, LocalArray):
                return main.proxyize(result)
            else:
                return result

        # default arguments
        args = () if args is None else args
        kwargs = {} if kwargs is None else kwargs

        args = tuple(a.key if isinstance(a, DistArray) else a for a in args)
        kwargs = {
            k: (v.key if isinstance(v, DistArray) else v)
            for k, v in kwargs.items()
        }

        apply_nonce = nonce()
        wrapped_args = (func, apply_nonce, self.context_key, args, kwargs,
                        autoproxyize)

        targets = self.targets if targets is None else targets

        with self.view.temp_flags(targets=targets):
            return self.view.apply_sync(func_wrapper, *wrapped_args)

    def push_function(self, key, func, targets=None):
        targets = targets or self.targets
        self._push({key: func}, targets=targets)
예제 #8
0
class IPythonContext(BaseContext):

    """
    Context class that uses IPython.parallel.

    See the docstring for  `BaseContext` for more information about Contexts.

    See also
    --------
    BaseContext
    """

    def __init__(self, client=None, targets=None):

        if client is None:
            self.client = IPythonClient()
            self.owns_client = True
        else:
            self.client = client
            self.owns_client = False

        self.view = self.client[:]
        self.nengines = len(self.view)

        self.all_targets = sorted(self.view.targets)
        if targets is None:
            self.targets = self.all_targets
        else:
            self.targets = []
            for target in targets:
                if target not in self.all_targets:
                    raise ValueError("Target %r not registered" % target)
                else:
                    self.targets.append(target)
        self.targets = sorted(self.targets)

        # local imports
        self.view.execute("from functools import reduce; "
                          "from importlib import import_module; "
                          "import distarray.localapi; "
                          "import distarray.localapi.mpiutils; "
                          "import distarray.utils; "
                          "import distarray.localapi.proxyize as proxyize; "
                          "from distarray.localapi.proxyize import Proxy; "
                          "import numpy")

        self.context_key = self._setup_context_key()

        # setup proxyize which is used by context.apply in the rest of the
        # setup.
        cmd = "proxyize = proxyize.Proxyize()"
        self.view.execute(cmd)

        self._base_comm = self._make_base_comm()
        self._comm_from_targets = {tuple(sorted(self.view.targets)): self._base_comm}
        self.comm = self.make_subcomm(self.targets)

        if not BaseContext._CLEANUP:
            BaseContext._CLEANUP = (atexit.register(ipython_cleanup.clear_all),
                                    atexit.register(ipython_cleanup.cleanup_all,
                                                    '__main__',
                                                    DISTARRAY_BASE_NAME))

    def make_subcomm(self, new_targets):

        if new_targets != sorted(new_targets):
            raise ValueError("targets must be in sorted order.")

        try:
            return self._comm_from_targets[tuple(new_targets)]
        except KeyError:
            pass

        def _make_new_comm(rank_list, base_comm):
            import distarray.localapi.mpiutils as mpiutils
            res = mpiutils.create_comm_with_list(rank_list, base_comm)
            return proxyize(res)  # noqa

        new_comm = self.apply(_make_new_comm, (new_targets, self._base_comm),
                              targets=self.view.targets)[0]

        self._comm_from_targets[tuple(new_targets)] = new_comm
        return new_comm

    def _make_base_comm(self):
        """
        Returns a proxy for an MPI communicator that encompasses all targets in
        self.view.targets (not self.targets, which can be a subset).
        """

        def get_rank():
            from distarray.localapi.mpiutils import get_comm_private
            return get_comm_private().Get_rank()

        # self.view's engines must encompass all ranks in the MPI communicator,
        # i.e., everything in rank_map.values().
        def get_size():
            from distarray.localapi.mpiutils import get_comm_private
            return get_comm_private().Get_size()

        # get a mapping of IPython engine ID to MPI rank
        rank_from_target = self.view.apply_async(get_rank).get_dict()
        ranks = [ rank_from_target[target] for target in self.view.targets ]

        comm_size = self.view.apply_async(get_size).get()[0]
        if set(rank_from_target.values()) != set(range(comm_size)):
            raise ValueError('Engines in view must encompass all MPI ranks.')

        # create a new communicator with the subset of ranks. Note that
        # create_comm_with_list() must be called on all engines, not just those
        # involved in the new communicator.  This is because
        # create_comm_with_list() issues a collective MPI operation.
        def _make_new_comm(rank_list):
            import distarray.localapi.mpiutils as mpiutils
            new_comm = mpiutils.create_comm_with_list(rank_list)
            if not mpiutils.get_base_comm():
                mpiutils.set_base_comm(new_comm)
            return proxyize(new_comm)  # noqa

        return self.apply(_make_new_comm, args=(ranks,),
                          targets=self.view.targets)[0]

    # Key management routines:
    def cleanup(self):
        """ Delete keys that this context created from all the engines. """
        # TODO: FIXME: cleanup needs updating to work with proxy objects.
        ipython_cleanup.cleanup(view=self.view, module_name='__main__',
                                prefix=self.context_key)

    def close(self):
        self.cleanup()
        def free_subcomm(subcomm):
            subcomm.Free()
        for targets, subcomm in self._comm_from_targets.items():
            self.apply(free_subcomm, (subcomm,), targets=targets)
        if self.owns_client:
            self.client.close()
        self.comm = None

    # End of key management routines.

    def _execute(self, lines, targets):
        return self.view.execute(lines, targets=targets, block=True)

    def _push(self, d, targets):
        return self.view.push(d, targets=targets, block=True)

    def apply(self, func, args=None, kwargs=None, targets=None, autoproxyize=False):
        """
        Analogous to IPython.parallel.view.apply_sync

        Parameters
        ----------
        func : function
        args : tuple
            positional arguments to func
        kwargs : dict
            key word arguments to func
        targets : sequence of integers
            engines func is to be run on.
        autoproxyize: bool, default False
            If True, implicitly return a Proxy object from the function.

        Returns
        -------
            return a list of the results on the each engine.
        """

        def func_wrapper(func, apply_nonce, context_key, args, kwargs, autoproxyize):
            """
            Function which calls the applied function after grabbing all the
            arguments on the engines that are passed in as names of the form
            `__distarray__<some uuid>`.
            """
            from importlib import import_module
            import types
            from distarray.metadata_utils import arg_kwarg_proxy_converter
            from distarray.localapi import LocalArray

            main = import_module('__main__')
            main.proxyize.set_state(apply_nonce)

            # Modify func to change the namespace it executes in.
            # but builtins don't have __code__, __globals__, etc.
            if not isinstance(func, types.BuiltinFunctionType):
                # get func's building  blocks first
                func_code = func.__code__
                func_name = func.__name__
                func_defaults = func.__defaults__
                func_closure = func.__closure__

                # build the func's new execution environment
                main.__dict__.update({'context_key': context_key})
                new_func_globals = main.__dict__
                # create the new func
                func = types.FunctionType(func_code, new_func_globals,
                                          func_name, func_defaults,
                                          func_closure)

            args, kwargs = arg_kwarg_proxy_converter(args, kwargs)
            result = func(*args, **kwargs)

            if autoproxyize and isinstance(result, LocalArray):
                return main.proxyize(result)
            else:
                return result

        # default arguments
        args = () if args is None else args
        kwargs = {} if kwargs is None else kwargs

        args = tuple(a.key if isinstance(a, DistArray) else a for a in args)
        kwargs = {k: (v.key if isinstance(v, DistArray) else v) for k, v in kwargs.items()}

        apply_nonce = nonce()
        wrapped_args = (func, apply_nonce, self.context_key, args, kwargs, autoproxyize)

        targets = self.targets if targets is None else targets

        with self.view.temp_flags(targets=targets):
            return self.view.apply_sync(func_wrapper, *wrapped_args)

    def push_function(self, key, func, targets=None):
        targets = targets or self.targets
        self._push({key: func}, targets=targets)