<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""

"""

from __future__ import print_function
import os
import sys
import logging
from functools import reduce
import numpy as np
import pyopencl as cl
from .utils.timing import decallmethods, timings
from .utils.dtype import *

__all__ = ["Extensions"]

logger = logging.getLogger(__name__)
logging.basicConfig(filename="spam.log", filemode='w',
                    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
                    level=logging.DEBUG)

@decallmethods(timings)
class CLEnv(object):

    def __init__(self, fast_math=True):
        self.fast_math = fast_math
        self.ctx = cl.create_some_context()
        self.queue = cl.CommandQueue(self.ctx)

@decallmethods(timings)
class CLModule(object):

    def __init__(self, env):
        self.env = env

        # read the kernel source code
        dirname = os.path.dirname(__file__)
        abspath = os.path.abspath(dirname)
        path = os.path.join(abspath, "src")
        src_file = "libcl_gravity.cl"
        fname = os.path.join(path, src_file)
        with open(fname, 'r') as fobj:
            src = fobj.read()
        self.src = src
        self.path = path

    def build(self, junroll=8):
        prec = "double" if REAL is np.float64 else "single"
        logger.debug("Building %s precision CL extension module.", prec)

        # setting options
        options = " -I {path}".format(path=self.path)
        options += " -D JUNROLL={junroll}".format(junroll=junroll)
        if prec is "double":
            options += " -D DOUBLE"
        if self.env.fast_math:
            options += " -cl-fast-relaxed-math"

        # building program
        self.program = cl.Program(self.env.ctx, self.src).build(options=options)

        logger.debug("done.")
        return self

    def __getattr__(self, name):
        return CLKernel(self.env, getattr(self.program, name))

@decallmethods(timings)
class CLKernel(object):

    def __init__(self, env, kernel):
        self.env = env
        self.kernel = kernel
        self._lsize = None
        self._lsize_max = None
        self._gsize = None
        self.dev_buff = {}
        self.input_buffer = {}
        self.output_buffer = {}
        self.output_shape = {}

    @property
    def local_size(self):
        lsize = self._lsize
        return lsize if isinstance(lsize, tuple) else (lsize,)

    @local_size.setter
    def local_size(self, lsize):
        self._lsize = lsize
        self._lsize_max = lsize

    @property
    def global_size(self):
        gsize = self._gsize
        return gsize if isinstance(gsize, tuple) else (gsize,)

    @global_size.setter
    def global_size(self, ni):
#        gsize = 0
#        ngroups = 1
#        lsize_max = self._lsize_max
#        while gsize < ni:
#            lsize = min(lsize_max, ((ni-1)//ngroups + 1))
#            gsize = lsize * ngroups
#            ngroups *= 2

        lsize_max = self._lsize_max
#        ngroups = ((ni-1)//lsize_max + 1)
#        ngroups = ((ni-1) + lsize_max)//lsize_max
#        ngroups = int((ni-1)**0.5 + 1)
        ngroups = int(ni**0.5)
        lsize = ((ni-1)//ngroups + 1)
        gsize = lsize * ngroups

        self._lsize = lsize
        self._gsize = gsize

    def set_local_memory(self, i, arg):
        def foo(x, y): return x * y
#        size = np.dtype(REAL).itemsize * reduce(foo, self.local_size)
        size = np.dtype(REAL).itemsize * reduce(foo, (self._lsize_max,))
        arg = cl.LocalMemory(size * arg)
        self.kernel.set_arg(i, arg)

    def set_int(self, i, arg):
        arg = UINT(arg)
        self.kernel.set_arg(i, arg)

    def set_float(self, i, arg):
        arg = REAL(arg)
        self.kernel.set_arg(i, arg)

    def set_input_buffer(self, i, arr):
#        def allocate_buffer(arr):
#            memf = cl.mem_flags
#            mapf = cl.map_flags
#            dev_buf = cl.Buffer(self.env.ctx,
#                                memf.READ_ONLY,
#                                size=arr.nbytes)
#            pin_buf = cl.Buffer(self.env.ctx,
#                                memf.READ_ONLY | memf.ALLOC_HOST_PTR,
#                                size=arr.nbytes)
#            (in_buf, ev) = cl.enqueue_map_buffer(self.env.queue, pin_buf,
#                                                 mapf.WRITE, 0, arr.shape,
#                                                 REAL, "C")
##            return (dev_buf, in_buf)
#            return (pin_buf, in_buf)
#
#        if not i in self.dev_buff:
#            (self.dev_buff[i], self.input_buffer[i]) = allocate_buffer(arr)
#            logger.debug("%s: allocating buffer for input arg #%d - %s.",
#                         self.kernel.function_name, i, self.dev_buff[i])
#        if len(arr) > len(self.input_buffer[i]):
#            (self.dev_buff[i], self.input_buffer[i]) = allocate_buffer(arr)
#            logger.debug("%s: reallocating buffer for input arg #%d - %s.",
#                         self.kernel.function_name, i, self.dev_buff[i])

#        inbuff = self.input_buffer[i][:len(arr)]
#        inbuff[:] = arr
#        cl.enqueue_copy(self.env.queue, self.dev_buff[i], inbuff)

#        self.input_buffer[i][:len(arr)] = arr
#        cl.enqueue_copy(self.env.queue, self.dev_buff[i], self.input_buffer[i][:len(arr)])

#        cl.enqueue_copy(self.env.queue, self.dev_buff[i], np.ascontiguousarray(arr, dtype=REAL))

        memf = cl.mem_flags
        self.dev_buff[i] = cl.Buffer(self.env.ctx,
                                     memf.READ_ONLY | memf.COPY_HOST_PTR,
                                     hostbuf=np.array(arr, order='C'))
#        self.dev_buff[i] = cl.Buffer(self.env.ctx,
#                                     memf.READ_ONLY | memf.USE_HOST_PTR,
#                                     hostbuf=np.array(arr, order='C'))

        arg = self.dev_buff[i]
        self.kernel.set_arg(i, arg)

    def set_output_buffer(self, i, arr):
        def allocate_buffer(arr):
            memf = cl.mem_flags
            mapf = cl.map_flags
            dev_buf = cl.Buffer(self.env.ctx,
                                memf.WRITE_ONLY,
                                size=arr.nbytes)
            pin_buf = cl.Buffer(self.env.ctx,
                                memf.WRITE_ONLY | memf.ALLOC_HOST_PTR,
                                size=arr.nbytes)
            (out_buf, ev) = cl.enqueue_map_buffer(self.env.queue, pin_buf,
                                                  mapf.READ, 0, arr.shape,
                                                  REAL, "C")
#            return (dev_buf, out_buf)
            return (pin_buf, out_buf)

        if not i in self.dev_buff:
            (self.dev_buff[i], self.output_buffer[i]) = allocate_buffer(arr)
            logger.debug("%s: allocating buffer for output arg #%d - %s.",
                         self.kernel.function_name, i, self.dev_buff[i])
        if arr.shape > self.output_buffer[i].shape:
            (self.dev_buff[i], self.output_buffer[i]) = allocate_buffer(arr)
            logger.debug("%s: reallocating buffer for output arg #%d - %s.",
                         self.kernel.function_name, i, self.dev_buff[i])

        self.output_shape[i] = arr.shape
        arg = self.dev_buff[i]
        self.kernel.set_arg(i, arg)

    def run(self):
        cl.enqueue_nd_range_kernel(self.env.queue,
                                   self.kernel,
                                   self.global_size,
                                   self.local_size,
                                  ).wait()

    def get_result(self):
#        def getter(item):
#            (i, shape) = item
#            return self.dev_buff[i].get_host_array(shape, REAL)

        def getter(item):
            (i, shape) = item
            cl.enqueue_copy(self.env.queue, self.output_buffer[i], self.dev_buff[i])
            return self.output_buffer[i][:shape[0]]

        return [getter(item) for item in self.output_shape.items()]

@decallmethods(timings)
class CEnv(object):

    def __init__(self, fast_math=True):
        self.fast_math = fast_math

@decallmethods(timings)
class CModule(object):

    def __init__(self, env):
        self.env = env

    def build(self):
        prec = "double" if REAL is np.float64 else "single"
        logger.debug("Building %s precision C extension module.", prec)

        if prec is "double":
            from tupan.lib import libc64_gravity as program
        else:
            from tupan.lib import libc32_gravity as program
        self.program = program

        logger.debug("done.")
        return self

    def __getattr__(self, name):
        return CKernel(self.env, getattr(self.program, name))

@decallmethods(timings)
class CKernel(object):

    def __init__(self, env, kernel):
        self.env = env
        self.kernel = kernel
        self.dev_args = {}
        self.res_ids = {}

    def set_int(self, i, arg):
        self.dev_args[i] = UINT(arg)

    def set_float(self, i, arg):
        self.dev_args[i] = REAL(arg)

    def set_input_buffer(self, i, arg):
        self.dev_args[i] = arg

    def set_output_buffer(self, i, arg):
        self.dev_args[i] = arg
        self.res_ids[i] = i

    def set_local_memory(self, i, arg):
        pass

    def run(self):
        args = self.dev_args.values()
        self.dev_result = self.kernel(*args)    # self.dev_result == None

    def get_result(self):       # XXX: FIXME
        return [self.dev_args[i] for i in self.res_ids.values()]

libkernels = {}
libkernels["c"] = CModule(CEnv(fast_math=True)).build()
libkernels["cl"] = CLModule(CLEnv(fast_math=True)).build(junroll=8)

def get_extension(use_sp=False, use_cl=False):
    libname = "cl" if use_cl else "c"
    prec = "single" if use_sp else "double"
    logger.debug("Using %s precision %s extension module.", prec, libname.upper())
    return libkernels[libname]

use_cl = True if "--use_cl" in sys.argv else False
use_sp = True if "--use_sp" in sys.argv else False

kernels = get_extension(use_sp=use_sp, use_cl=use_cl)

########## end of file ##########
<commit_msg>a better way to calculate the global size for launch a CL kernel (III).<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""

"""

from __future__ import print_function
import os
import sys
import logging
from functools import reduce
import numpy as np
import pyopencl as cl
from .utils.timing import decallmethods, timings
from .utils.dtype import *

__all__ = ["Extensions"]

logger = logging.getLogger(__name__)
logging.basicConfig(filename="spam.log", filemode='w',
                    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
                    level=logging.DEBUG)

@decallmethods(timings)
class CLEnv(object):

    def __init__(self, fast_math=True):
        self.fast_math = fast_math
        self.ctx = cl.create_some_context()
        self.queue = cl.CommandQueue(self.ctx)

@decallmethods(timings)
class CLModule(object):

    def __init__(self, env):
        self.env = env

        # read the kernel source code
        dirname = os.path.dirname(__file__)
        abspath = os.path.abspath(dirname)
        path = os.path.join(abspath, "src")
        src_file = "libcl_gravity.cl"
        fname = os.path.join(path, src_file)
        with open(fname, 'r') as fobj:
            src = fobj.read()
        self.src = src
        self.path = path

    def build(self, junroll=8):
        prec = "double" if REAL is np.float64 else "single"
        logger.debug("Building %s precision CL extension module.", prec)

        # setting options
        options = " -I {path}".format(path=self.path)
        options += " -D JUNROLL={junroll}".format(junroll=junroll)
        if prec is "double":
            options += " -D DOUBLE"
        if self.env.fast_math:
            options += " -cl-fast-relaxed-math"

        # building program
        self.program = cl.Program(self.env.ctx, self.src).build(options=options)

        logger.debug("done.")
        return self

    def __getattr__(self, name):
        return CLKernel(self.env, getattr(self.program, name))

@decallmethods(timings)
class CLKernel(object):

    def __init__(self, env, kernel):
        self.env = env
        self.kernel = kernel
        self._lsize = None
        self._lsize_max = None
        self._gsize = None
        self.dev_buff = {}
        self.input_buffer = {}
        self.output_buffer = {}
        self.output_shape = {}

    @property
    def local_size(self):
        lsize = self._lsize
        return lsize if isinstance(lsize, tuple) else (lsize,)

    @local_size.setter
    def local_size(self, lsize):
        self._lsize = lsize
        self._lsize_max = lsize

    @property
    def global_size(self):
        gsize = self._gsize
        return gsize if isinstance(gsize, tuple) else (gsize,)

    @global_size.setter
    def global_size(self, ni):
#        gsize = 0
#        lsize = 0
#        ngroups = 1
#        lsize_max = self._lsize_max
#        while gsize < ni and lsize < lsize_max:
#            ngroups += 1
##            lsize = ((ni-1)//ngroups + 1)
#            lsize = ((ni-1)//ngroups + 1)
#            gsize = lsize * ngroups

        gsize = ((ni-1)//8 + 1) * 8

#        lsize_max = self._lsize_max
##        ngroups = ((ni-1)//lsize_max + 1)
##        ngroups = ((ni-1) + lsize_max)//lsize_max
##        ngroups = int((ni-1)**0.5 + 1)
##        ngroups = int(ni**0.5)
##        lsize = ((ni-1)//ngroups + 1)
#        lsize = min(lsize_max, (ni+1)//2)
#        ngroups = (ni+1)//lsize
#        gsize = lsize * ngroups

#        self._lsize = lsize
        self._gsize = gsize

    def set_local_memory(self, i, arg):
        def foo(x, y): return x * y
#        size = np.dtype(REAL).itemsize * reduce(foo, self.local_size)
#        size = np.dtype(REAL).itemsize * reduce(foo, (self._lsize_max,))
        size = 8 * reduce(foo, (self._lsize_max,))
        arg = cl.LocalMemory(size * arg)
        self.kernel.set_arg(i, arg)

    def set_int(self, i, arg):
        arg = UINT(arg)
        self.kernel.set_arg(i, arg)

    def set_float(self, i, arg):
        arg = REAL(arg)
        self.kernel.set_arg(i, arg)

    def set_input_buffer(self, i, arr):
#        def allocate_buffer(arr):
#            memf = cl.mem_flags
#            mapf = cl.map_flags
#            dev_buf = cl.Buffer(self.env.ctx,
#                                memf.READ_ONLY,
#                                size=arr.nbytes)
#            pin_buf = cl.Buffer(self.env.ctx,
#                                memf.READ_ONLY | memf.ALLOC_HOST_PTR,
#                                size=arr.nbytes)
#            (in_buf, ev) = cl.enqueue_map_buffer(self.env.queue, pin_buf,
#                                                 mapf.WRITE, 0, arr.shape,
#                                                 REAL, "C")
##            return (dev_buf, in_buf)
#            return (pin_buf, in_buf)
#
#        if not i in self.dev_buff:
#            (self.dev_buff[i], self.input_buffer[i]) = allocate_buffer(arr)
#            logger.debug("%s: allocating buffer for input arg #%d - %s.",
#                         self.kernel.function_name, i, self.dev_buff[i])
#        if len(arr) > len(self.input_buffer[i]):
#            (self.dev_buff[i], self.input_buffer[i]) = allocate_buffer(arr)
#            logger.debug("%s: reallocating buffer for input arg #%d - %s.",
#                         self.kernel.function_name, i, self.dev_buff[i])

#        inbuff = self.input_buffer[i][:len(arr)]
#        inbuff[:] = arr
#        cl.enqueue_copy(self.env.queue, self.dev_buff[i], inbuff)

#        self.input_buffer[i][:len(arr)] = arr
#        cl.enqueue_copy(self.env.queue, self.dev_buff[i], self.input_buffer[i][:len(arr)])

#        cl.enqueue_copy(self.env.queue, self.dev_buff[i], np.ascontiguousarray(arr, dtype=REAL))

        memf = cl.mem_flags
        self.dev_buff[i] = cl.Buffer(self.env.ctx,
                                     memf.READ_ONLY | memf.COPY_HOST_PTR,
                                     hostbuf=np.array(arr, order='C'))
#        self.dev_buff[i] = cl.Buffer(self.env.ctx,
#                                     memf.READ_ONLY | memf.USE_HOST_PTR,
#                                     hostbuf=np.array(arr, order='C'))

        arg = self.dev_buff[i]
        self.kernel.set_arg(i, arg)

    def set_output_buffer(self, i, arr):
        def allocate_buffer(arr):
            memf = cl.mem_flags
            mapf = cl.map_flags
            dev_buf = cl.Buffer(self.env.ctx,
                                memf.WRITE_ONLY,
                                size=arr.nbytes)
            pin_buf = cl.Buffer(self.env.ctx,
                                memf.WRITE_ONLY | memf.ALLOC_HOST_PTR,
                                size=arr.nbytes)
            (out_buf, ev) = cl.enqueue_map_buffer(self.env.queue, pin_buf,
                                                  mapf.READ, 0, arr.shape,
                                                  REAL, "C")
#            return (dev_buf, out_buf)
            return (pin_buf, out_buf)

        if not i in self.dev_buff:
            (self.dev_buff[i], self.output_buffer[i]) = allocate_buffer(arr)
            logger.debug("%s: allocating buffer for output arg #%d - %s.",
                         self.kernel.function_name, i, self.dev_buff[i])
        if arr.shape > self.output_buffer[i].shape:
            (self.dev_buff[i], self.output_buffer[i]) = allocate_buffer(arr)
            logger.debug("%s: reallocating buffer for output arg #%d - %s.",
                         self.kernel.function_name, i, self.dev_buff[i])

        self.output_shape[i] = arr.shape
        arg = self.dev_buff[i]
        self.kernel.set_arg(i, arg)

    def run(self):
        cl.enqueue_nd_range_kernel(self.env.queue,
                                   self.kernel,
                                   self.global_size,
                                   None,#self.local_size,
                                  ).wait()

    def get_result(self):
#        def getter(item):
#            (i, shape) = item
#            return self.dev_buff[i].get_host_array(shape, REAL)

        def getter(item):
            (i, shape) = item
            cl.enqueue_copy(self.env.queue, self.output_buffer[i], self.dev_buff[i])
            return self.output_buffer[i][:shape[0]]

        return [getter(item) for item in self.output_shape.items()]

@decallmethods(timings)
class CEnv(object):

    def __init__(self, fast_math=True):
        self.fast_math = fast_math

@decallmethods(timings)
class CModule(object):

    def __init__(self, env):
        self.env = env

    def build(self):
        prec = "double" if REAL is np.float64 else "single"
        logger.debug("Building %s precision C extension module.", prec)

        if prec is "double":
            from tupan.lib import libc64_gravity as program
        else:
            from tupan.lib import libc32_gravity as program
        self.program = program

        logger.debug("done.")
        return self

    def __getattr__(self, name):
        return CKernel(self.env, getattr(self.program, name))

@decallmethods(timings)
class CKernel(object):

    def __init__(self, env, kernel):
        self.env = env
        self.kernel = kernel
        self.dev_args = {}
        self.res_ids = {}
        self._lsize = None
        self._lsize_max = None
        self._gsize = None

    def set_int(self, i, arg):
        self.dev_args[i] = UINT(arg)

    def set_float(self, i, arg):
        self.dev_args[i] = REAL(arg)

    def set_input_buffer(self, i, arg):
        self.dev_args[i] = arg

    def set_output_buffer(self, i, arg):
        self.dev_args[i] = arg
        self.res_ids[i] = i

    def set_local_memory(self, i, arg):
        pass

    def run(self):
        args = self.dev_args.values()
        self.dev_result = self.kernel(*args)    # self.dev_result == None

    def get_result(self):       # XXX: FIXME
        return [self.dev_args[i] for i in self.res_ids.values()]

libkernels = {}
libkernels["c"] = CModule(CEnv(fast_math=True)).build()
libkernels["cl"] = CLModule(CLEnv(fast_math=True)).build(junroll=8)

def get_extension(use_sp=False, use_cl=False):
    libname = "cl" if use_cl else "c"
    prec = "single" if use_sp else "double"
    logger.debug("Using %s precision %s extension module.", prec, libname.upper())
    return libkernels[libname]

use_cl = True if "--use_cl" in sys.argv else False
use_sp = True if "--use_sp" in sys.argv else False

kernels = get_extension(use_sp=use_sp, use_cl=use_cl)

########## end of file ##########
<|endoftext|>


<commit_before>__version__ = '2.29.2'
<commit_msg>Bump version<commit_after>__version__ = '2.29.3'
<|endoftext|>


<commit_before>"""
Distributed locks with Redis
Redis doc: http://redis.io/topics/distlock
"""
from __future__ import division
from datetime import datetime
import random
import time
import uuid

import redis

DEFAULT_RETRY_TIMES = 3
DEFAULT_RETRY_DELAY = 200
DEFAULT_TTL = 100000
CLOCK_DRIFT_FACTOR = 0.01

# Reference:  http://redis.io/topics/distlock
# Section Correct implementation with a single instance
RELEASE_LUA_SCRIPT = """
    if redis.call("get",KEYS[1]) == ARGV[1] then
        return redis.call("del",KEYS[1])
    else
        return 0
    end
"""

class RedLockError(Exception):
    pass

class RedLockFactory(object):
    """
    A Factory class that helps reuse multiple Redis connections.
    """

    def __init__(self, connection_details):
        """

        """
        self.redis_nodes = []

        for conn in connection_details:
            if isinstance(conn, redis.StrictRedis):
                node = conn
            elif 'url' in conn:
                url = conn.pop('url')
                node = redis.StrictRedis.from_url(url, **conn)
            else:
                node = redis.StrictRedis(**conn)
            node._release_script = node.register_script(RELEASE_LUA_SCRIPT)
            self.redis_nodes.append(node)
            self.quorum = len(self.redis_nodes) // 2 + 1

    def create_lock(self, resource, **kwargs):
        """
        Create a new RedLock object and reuse stored Redis clients.
        All the kwargs it received would be passed to the RedLock's __init__
        function.
        """
        lock = RedLock(resource=resource, created_by_factory=True, **kwargs)
        lock.redis_nodes = self.redis_nodes
        lock.quorum = self.quorum
        lock.factory = self
        return lock

class RedLock(object):

    """
    A distributed lock implementation based on Redis.
    It shares a similar API with the `threading.Lock` class in the
    Python Standard Library.
    """

    def __init__(self, resource, connection_details=None,
                 retry_times=DEFAULT_RETRY_TIMES,
                 retry_delay=DEFAULT_RETRY_DELAY,
                 ttl=DEFAULT_TTL,
                 created_by_factory=False):

        self.resource = resource
        self.retry_times = retry_times
        self.retry_delay = retry_delay
        self.ttl = ttl

        if created_by_factory:
            self.factory = None
            return

        self.redis_nodes = []
        # If the connection_details parameter is not provided,
        # use redis://127.0.0.1:6379/0
        if connection_details is None:
            connection_details = [{
                'host': 'localhost',
                'port': 6379,
                'db': 0,
            }]

        for conn in connection_details:
            if isinstance(conn, redis.StrictRedis):
                node = conn
            elif 'url' in conn:
                url = conn.pop('url')
                node = redis.StrictRedis.from_url(url, **conn)
            else:
                node = redis.StrictRedis(**conn)
            node._release_script = node.register_script(RELEASE_LUA_SCRIPT)
            self.redis_nodes.append(node)
        self.quorum = len(self.redis_nodes) // 2 + 1

    def __enter__(self):
        acquired, validity = self.acquire_with_validity()
        if not acquired:
            raise RedLockError('failed to acquire lock')
        return validity

    def __exit__(self, exc_type, exc_value, traceback):
        self.release()

    def _total_ms(self, delta):
        """
        Get the total number of milliseconds in a timedelta object with
        microsecond precision.
        """
        delta_seconds = delta.seconds + delta.days * 24 * 3600
        return (delta.microseconds + delta_seconds * 10**6) / 10**3

    def acquire_node(self, node):
        """
        acquire a single redis node
        """
        try:
            return node.set(self.resource, self.lock_key, nx=True, px=self.ttl)
        except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError):
            return False

    def release_node(self, node):
        """
        release a single redis node
        """
        # use the lua script to release the lock in a safe way
        try:
            node._release_script(keys=[self.resource], args=[self.lock_key])
        except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError):
            pass

    def acquire(self):
        acquired, validity = self._acquire()
        return acquired

    def acquire_with_validity(self):
        return self._acquire()

    def _acquire(self):

        # lock_key should be random and unique
        self.lock_key = uuid.uuid4().hex

        for retry in range(self.retry_times):
            acquired_node_count = 0
            start_time = datetime.utcnow()

            # acquire the lock in all the redis instances sequentially
            for node in self.redis_nodes:
                if self.acquire_node(node):
                    acquired_node_count += 1

            end_time = datetime.utcnow()
            elapsed_milliseconds = self._total_ms(end_time - start_time)

            # Add 2 milliseconds to the drift to account for Redis expires
            # precision, which is 1 milliescond, plus 1 millisecond min drift
            # for small TTLs.
            drift = (self.ttl * CLOCK_DRIFT_FACTOR) + 2

            validity = self.ttl - (elapsed_milliseconds + drift)
            if acquired_node_count >= self.quorum and validity > 0:
                return True, validity
            else:
                for node in self.redis_nodes:
                    self.release_node(node)
                time.sleep(random.randint(0, self.retry_delay) / 1000)
        return False, 0

    def release(self):
        for node in self.redis_nodes:
            self.release_node(node)

class ReentrantRedLock(RedLock):
    def __init__(self, *args, **kwargs):
        super(ReentrantRedLock, self).__init__(*args, **kwargs)
        self._acquired = 0

    def acquire(self):
        if self._acquired == 0:
            result = super(ReentrantRedLock, self).acquire()
            if result:
                self._acquired += 1
            return result
        else:
            self._acquired += 1
            return True

    def release(self):
        if self._acquired > 0:
            self._acquired -= 1
            if self._acquired == 0:
                return super(ReentrantRedLock, self).release()
            return True
        return False
<commit_msg>Fix retry counting<commit_after>"""
Distributed locks with Redis
Redis doc: http://redis.io/topics/distlock
"""
from __future__ import division
from datetime import datetime
import random
import time
import uuid

import redis

DEFAULT_RETRY_TIMES = 3
DEFAULT_RETRY_DELAY = 200
DEFAULT_TTL = 100000
CLOCK_DRIFT_FACTOR = 0.01

# Reference:  http://redis.io/topics/distlock
# Section Correct implementation with a single instance
RELEASE_LUA_SCRIPT = """
    if redis.call("get",KEYS[1]) == ARGV[1] then
        return redis.call("del",KEYS[1])
    else
        return 0
    end
"""

class RedLockError(Exception):
    pass

class RedLockFactory(object):
    """
    A Factory class that helps reuse multiple Redis connections.
    """

    def __init__(self, connection_details):
        """

        """
        self.redis_nodes = []

        for conn in connection_details:
            if isinstance(conn, redis.StrictRedis):
                node = conn
            elif 'url' in conn:
                url = conn.pop('url')
                node = redis.StrictRedis.from_url(url, **conn)
            else:
                node = redis.StrictRedis(**conn)
            node._release_script = node.register_script(RELEASE_LUA_SCRIPT)
            self.redis_nodes.append(node)
            self.quorum = len(self.redis_nodes) // 2 + 1

    def create_lock(self, resource, **kwargs):
        """
        Create a new RedLock object and reuse stored Redis clients.
        All the kwargs it received would be passed to the RedLock's __init__
        function.
        """
        lock = RedLock(resource=resource, created_by_factory=True, **kwargs)
        lock.redis_nodes = self.redis_nodes
        lock.quorum = self.quorum
        lock.factory = self
        return lock

class RedLock(object):

    """
    A distributed lock implementation based on Redis.
    It shares a similar API with the `threading.Lock` class in the
    Python Standard Library.
    """

    def __init__(self, resource, connection_details=None,
                 retry_times=DEFAULT_RETRY_TIMES,
                 retry_delay=DEFAULT_RETRY_DELAY,
                 ttl=DEFAULT_TTL,
                 created_by_factory=False):

        self.resource = resource
        self.retry_times = retry_times
        self.retry_delay = retry_delay
        self.ttl = ttl

        if created_by_factory:
            self.factory = None
            return

        self.redis_nodes = []
        # If the connection_details parameter is not provided,
        # use redis://127.0.0.1:6379/0
        if connection_details is None:
            connection_details = [{
                'host': 'localhost',
                'port': 6379,
                'db': 0,
            }]

        for conn in connection_details:
            if isinstance(conn, redis.StrictRedis):
                node = conn
            elif 'url' in conn:
                url = conn.pop('url')
                node = redis.StrictRedis.from_url(url, **conn)
            else:
                node = redis.StrictRedis(**conn)
            node._release_script = node.register_script(RELEASE_LUA_SCRIPT)
            self.redis_nodes.append(node)
        self.quorum = len(self.redis_nodes) // 2 + 1

    def __enter__(self):
        acquired, validity = self.acquire_with_validity()
        if not acquired:
            raise RedLockError('failed to acquire lock')
        return validity

    def __exit__(self, exc_type, exc_value, traceback):
        self.release()

    def _total_ms(self, delta):
        """
        Get the total number of milliseconds in a timedelta object with
        microsecond precision.
        """
        delta_seconds = delta.seconds + delta.days * 24 * 3600
        return (delta.microseconds + delta_seconds * 10**6) / 10**3

    def acquire_node(self, node):
        """
        acquire a single redis node
        """
        try:
            return node.set(self.resource, self.lock_key, nx=True, px=self.ttl)
        except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError):
            return False

    def release_node(self, node):
        """
        release a single redis node
        """
        # use the lua script to release the lock in a safe way
        try:
            node._release_script(keys=[self.resource], args=[self.lock_key])
        except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError):
            pass

    def acquire(self):
        acquired, validity = self._acquire()
        return acquired

    def acquire_with_validity(self):
        return self._acquire()

    def _acquire(self):

        # lock_key should be random and unique
        self.lock_key = uuid.uuid4().hex

        for retry in range(self.retry_times + 1):
            acquired_node_count = 0
            start_time = datetime.utcnow()

            # acquire the lock in all the redis instances sequentially
            for node in self.redis_nodes:
                if self.acquire_node(node):
                    acquired_node_count += 1

            end_time = datetime.utcnow()
            elapsed_milliseconds = self._total_ms(end_time - start_time)

            # Add 2 milliseconds to the drift to account for Redis expires
            # precision, which is 1 milliescond, plus 1 millisecond min drift
            # for small TTLs.
            drift = (self.ttl * CLOCK_DRIFT_FACTOR) + 2

            validity = self.ttl - (elapsed_milliseconds + drift)
            if acquired_node_count >= self.quorum and validity > 0:
                return True, validity
            else:
                for node in self.redis_nodes:
                    self.release_node(node)
                time.sleep(random.randint(0, self.retry_delay) / 1000)
        return False, 0

    def release(self):
        for node in self.redis_nodes:
            self.release_node(node)

class ReentrantRedLock(RedLock):
    def __init__(self, *args, **kwargs):
        super(ReentrantRedLock, self).__init__(*args, **kwargs)
        self._acquired = 0

    def acquire(self):
        if self._acquired == 0:
            result = super(ReentrantRedLock, self).acquire()
            if result:
                self._acquired += 1
            return result
        else:
            self._acquired += 1
            return True

    def release(self):
        if self._acquired > 0:
            self._acquired -= 1
            if self._acquired == 0:
                return super(ReentrantRedLock, self).release()
            return True
        return False
<|endoftext|>


<commit_before>from pywb.cdx.cdxsource import CDXFile, RedisCDXSource
from pywb.cdx.cdxserver import CDXServer

from pywb.framework.archivalrouter import Route

from pywb.webapp.pywb_init import DirectoryCollsLoader
from pywb.webapp.handlers import WBHandler
from pywb.webapp.live_rewrite_handler import RewriteHandler, LiveResourceException
from pywb.webapp.views import J2TemplateView

from pywb.framework.wbrequestresponse import WbResponse
from pywb.utils.wbexception import NotFoundException
from pywb.rewrite.url_rewriter import UrlRewriter

import functools
import os
import re
import json
import base64
from os.path import expandvars

#=================================================================
jinja_env = J2TemplateView.init_shared_env()

def jinja2_view(template_name):
    def decorator(view_func):
        @functools.wraps(view_func)
        def wrapper(*args, **kwargs):
            response = view_func(*args, **kwargs)

            if isinstance(response, dict):
                template = jinja_env.get_or_select_template(template_name)
                return template.render(**response)
            else:
                return response

        return wrapper

    return decorator

#=================================================================
class DynamicRoute(Route):
    def apply_filters(self, wbrequest, matcher):
        wbrequest.custom_params['output_dir'] = wbrequest.env.get('w_output_dir', '')

        sesh_id = wbrequest.env.get('w_sesh_id', '')
        wbrequest.custom_params['sesh_id'] = sesh_id
        wbrequest.custom_params['user_id'] = wbrequest.env.get('w_user_id', '')
        wbrequest.coll = sesh_id

#=================================================================
class DynCDXServer(CDXServer):
    def _create_cdx_sources(self, paths, config):
        paths = expandvars(paths)
        if paths.startswith('redis://'):
            src = DynCDXRedis(paths)
        else:
            src = DynCDXFile(paths)

        self.sources = [src]

#=================================================================
class DynCDXRedis(RedisCDXSource):
    def load_cdx(self, query):
        path = query.params['output_dir']
        if not path:
            print 'No Path'
            return iter([])

        sesh_id = query.params['sesh_id']
        cdx_key = sesh_id.replace('/', ':') + ':cdxj'

        return self.load_sorted_range(query, cdx_key)

#=================================================================
class DynCDXFile(CDXFile):
    """
    Represents a parametric cdx file
    """
    def __init__(self, filename):
        self.root_path = ''
        self.param_file = filename

    def load_cdx(self, query):
        path = query.params['output_dir']

        if not path:
            print 'No Path'
            return iter([])

        filename = os.path.join(self.root_path,
                                path,
                                self.param_file)

        return self._do_load_file(filename, query)

    def __str__(self):
        return 'Dyn CDX File - ' + self.root_path

#=================================================================
class DynRedisResolver(object):
    def __init__(self, redis, s3_target=None, proxy_target=None, key_prefix='w:'):
        self.redis = redis
        self.s3_target = s3_target
        self.proxy_target = proxy_target
        self.key_prefix = key_prefix

    def __call__(self, filename):
        sesh_id, warc_key = self._split_sesh_warc(filename)
        orig_path = self.redis.hget(sesh_id, warc_key)

        if not orig_path:
            return []

        # if proxy_path set, try proxy path first
        if self.s3_target and self.proxy_target:
            cached_path = orig_path.replace(self.s3_target, self.proxy_target)
            return [cached_path, orig_path]
        else:
            return [orig_path]

    def add_filename(self, filename, remote_url):
        sesh_id, warc_key = self._split_sesh_warc(filename)
        redis_val = self.redis.hset(sesh_id, warc_key, remote_url)

    def _split_sesh_warc(self, filename):
        #TODO: pass sesh_id here...
        parts = filename.rsplit('/')
        #return 'warc:' + parts[-5] + ':' + parts[-3], parts[-1]
        return parts[-5] + ':' + parts[-3] + ':warc', parts[-1]

    def __repr__(self):
        return "DynRedisResolver('{0}')".format(self.redis)

#=================================================================
class DynWBHandler(WBHandler):
    def _init_replay_view(self, config):
        replay = super(DynWBHandler, self)._init_replay_view(config)

        redis_warc_resolver = config.get('redis_warc_resolver')
        if redis_warc_resolver:
            replay.content_loader.path_resolvers.append(redis_warc_resolver)

        return replay

    def get_top_frame_params(self, wbrequest, mod):
        params = (super(DynWBHandler, self).
                  get_top_frame_params(wbrequest, mod))

        manager = wbrequest.env.get('w_manager')

        if manager:
            user, coll = wbrequest.coll.split('/', 1)
            info = manager.get_info(user, coll)
            params['info'] = json.dumps(info)

        return params

    def handle_replay(self, wbrequest, cdx_lines):
        path = wbrequest.custom_params['output_dir']
        #path = os.path.join(path, 'archive')

        cdx_callback = self.index_reader.cdx_load_callback(wbrequest)

        def wrapped_cdx_callback(*args, **kwargs):
            return self._wrap_session_path(path, cdx_callback(*args, **kwargs))

        wrapped_cdx_lines = self._wrap_session_path(path, cdx_lines)

        return self.replay.render_content(wbrequest,
                                          wrapped_cdx_lines,
                                          wrapped_cdx_callback)

    @staticmethod
    def _wrap_session_path(path, cdx_lines):
        for cdx in cdx_lines:
            cdx['filename'] = os.path.join(path, cdx['filename'])
            print(cdx['filename'])
            if cdx.get('orig.filename', '-') != '-':
                cdx['orig.filename'] = os.path.join(path,
                                                    cdx['orig.filename'])
            yield cdx

#=================================================================
class DynRecord(RewriteHandler):
    def __init__(self, config):
        super(DynRecord, self).__init__(config)
        cookie_name = config.get('cookie_name', 'beaker.session.id')
        self.record_path = config.get('record_dir', './')

    def get_top_frame_params(self, wbrequest, mod):
        params = (super(DynRecord, self).
                  get_top_frame_params(wbrequest, mod))

        manager = wbrequest.env.get('w_manager')

        if manager:
            user, coll = wbrequest.coll.split('/', 1)
            info = manager.get_info(user, coll)
            params['info'] = json.dumps(info)

        return params

    def _ignore_proxies(self, wbrequest):
        return wbrequest.env.get('webrec.no_space', False)

    def _live_request_headers(self, wbrequest):
        path = wbrequest.custom_params.get('output_dir')
        if not path:
            path = self.record_path

        if not self._ignore_proxies(wbrequest):
            sesh_id = wbrequest.custom_params.get('sesh_id', wbrequest.coll)
            sesh_id = sesh_id.replace('/', ':')
            user_id = wbrequest.custom_params.get('user_id')

            target = dict(output_dir=path,
                          sesh_id=sesh_id,
                          user_id=user_id)

            req_headers = {'warcprox-meta': json.dumps(target)}
        else:
            req_headers = {}

        # reset HTTP_COOKIE to guarded request_cookie for LiveRewriter
        if 'webrec.request_cookie' in wbrequest.env:
            wbrequest.env['HTTP_COOKIE'] = wbrequest.env['webrec.request_cookie']

        try:
            del wbrequest.env['HTTP_X_PUSH_STATE_REQUEST']
        except:
            pass

        return req_headers

    def _make_response(self, wbrequest, status_headers, gen, is_rewritten):
        if (status_headers.get_statuscode() == '500' and
            status_headers.get_header('X-Archive-Orig-x-warcprox-error') == '500'):
            msg = 'This url could not be recorded: '
            raise LiveResourceException(msg, url=wbrequest.wb_url.url)

        return WbResponse(status_headers, gen)

#=================================================================
class DynUrlRewriter(UrlRewriter):
    def rewrite(self, url, mod=None):
        new_url = super(DynUrlRewriter, self).rewrite(url, mod)
        if new_url == url:
            return new_url

        parts = new_url.split('://', 1)
        if len(parts) < 2:
            return new_url

        rand = base64.b32encode(os.urandom(5))
        parts[1] = rand + '.' + parts[1]

        #print('DYN RW 2')
        #host = parts[2].split('.')[0]
        #parts[1] = host + '.' + parts[1]
        new_url = '://'.join(parts)
        return new_url
<commit_msg>patch: use fallback for any capture exception (eg. self-redirect)<commit_after>from pywb.cdx.cdxsource import CDXFile, RedisCDXSource
from pywb.cdx.cdxserver import CDXServer

from pywb.framework.archivalrouter import Route

from pywb.webapp.pywb_init import DirectoryCollsLoader
from pywb.webapp.handlers import WBHandler
from pywb.webapp.live_rewrite_handler import RewriteHandler, LiveResourceException
from pywb.webapp.views import J2TemplateView
from pywb.webapp.replay_views import CaptureException

from pywb.framework.wbrequestresponse import WbResponse
from pywb.utils.wbexception import NotFoundException
from pywb.rewrite.url_rewriter import UrlRewriter

import functools
import os
import re
import json
import base64
from os.path import expandvars

#=================================================================
jinja_env = J2TemplateView.init_shared_env()

def jinja2_view(template_name):
    def decorator(view_func):
        @functools.wraps(view_func)
        def wrapper(*args, **kwargs):
            response = view_func(*args, **kwargs)

            if isinstance(response, dict):
                template = jinja_env.get_or_select_template(template_name)
                return template.render(**response)
            else:
                return response

        return wrapper

    return decorator

#=================================================================
class DynamicRoute(Route):
    def apply_filters(self, wbrequest, matcher):
        wbrequest.custom_params['output_dir'] = wbrequest.env.get('w_output_dir', '')

        sesh_id = wbrequest.env.get('w_sesh_id', '')
        wbrequest.custom_params['sesh_id'] = sesh_id
        wbrequest.custom_params['user_id'] = wbrequest.env.get('w_user_id', '')
        wbrequest.coll = sesh_id

#=================================================================
class DynCDXServer(CDXServer):
    def _create_cdx_sources(self, paths, config):
        paths = expandvars(paths)
        if paths.startswith('redis://'):
            src = DynCDXRedis(paths)
        else:
            src = DynCDXFile(paths)

        self.sources = [src]

#=================================================================
class DynCDXRedis(RedisCDXSource):
    def load_cdx(self, query):
        path = query.params['output_dir']
        if not path:
            print 'No Path'
            return iter([])

        sesh_id = query.params['sesh_id']
        cdx_key = sesh_id.replace('/', ':') + ':cdxj'

        return self.load_sorted_range(query, cdx_key)

#=================================================================
class DynCDXFile(CDXFile):
    """
    Represents a parametric cdx file
    """
    def __init__(self, filename):
        self.root_path = ''
        self.param_file = filename

    def load_cdx(self, query):
        path = query.params['output_dir']

        if not path:
            print 'No Path'
            return iter([])

        filename = os.path.join(self.root_path,
                                path,
                                self.param_file)

        return self._do_load_file(filename, query)

    def __str__(self):
        return 'Dyn CDX File - ' + self.root_path

#=================================================================
class DynRedisResolver(object):
    def __init__(self, redis, s3_target=None, proxy_target=None, key_prefix='w:'):
        self.redis = redis
        self.s3_target = s3_target
        self.proxy_target = proxy_target
        self.key_prefix = key_prefix

    def __call__(self, filename):
        sesh_id, warc_key = self._split_sesh_warc(filename)
        orig_path = self.redis.hget(sesh_id, warc_key)

        if not orig_path:
            return []

        # if proxy_path set, try proxy path first
        if self.s3_target and self.proxy_target:
            cached_path = orig_path.replace(self.s3_target, self.proxy_target)
            return [cached_path, orig_path]
        else:
            return [orig_path]

    def add_filename(self, filename, remote_url):
        sesh_id, warc_key = self._split_sesh_warc(filename)
        redis_val = self.redis.hset(sesh_id, warc_key, remote_url)

    def _split_sesh_warc(self, filename):
        #TODO: pass sesh_id here...
        parts = filename.rsplit('/')
        #return 'warc:' + parts[-5] + ':' + parts[-3], parts[-1]
        return parts[-5] + ':' + parts[-3] + ':warc', parts[-1]

    def __repr__(self):
        return "DynRedisResolver('{0}')".format(self.redis)

#=================================================================
class DynWBHandler(WBHandler):
    def _init_replay_view(self, config):
        replay = super(DynWBHandler, self)._init_replay_view(config)

        redis_warc_resolver = config.get('redis_warc_resolver')
        if redis_warc_resolver:
            replay.content_loader.path_resolvers.append(redis_warc_resolver)

        return replay

    def get_top_frame_params(self, wbrequest, mod):
        params = (super(DynWBHandler, self).
                  get_top_frame_params(wbrequest, mod))

        manager = wbrequest.env.get('w_manager')

        if manager:
            user, coll = wbrequest.coll.split('/', 1)
            info = manager.get_info(user, coll)
            params['info'] = json.dumps(info)

        return params

    def handle_replay(self, wbrequest, cdx_lines):
        path = wbrequest.custom_params['output_dir']
        #path = os.path.join(path, 'archive')

        try:
            cdx_callback = self.index_reader.cdx_load_callback(wbrequest)

            def wrapped_cdx_callback(*args, **kwargs):
                return self._wrap_session_path(path, cdx_callback(*args, **kwargs))

            wrapped_cdx_lines = self._wrap_session_path(path, cdx_lines)

            return self.replay.render_content(wbrequest,
                                              wrapped_cdx_lines,
                                              wrapped_cdx_callback)
        except CaptureException:
            if (self.fallback_handler and
                not wbrequest.wb_url.is_query() and
                not wbrequest.wb_url.is_identity):
                return self.fallback_handler(wbrequest)
            else:
                raise

    @staticmethod
    def _wrap_session_path(path, cdx_lines):
        for cdx in cdx_lines:
            cdx['filename'] = os.path.join(path, cdx['filename'])
            print(cdx['filename'])
            if cdx.get('orig.filename', '-') != '-':
                cdx['orig.filename'] = os.path.join(path,
                                                    cdx['orig.filename'])
            yield cdx

#=================================================================
class DynRecord(RewriteHandler):
    def __init__(self, config):
        super(DynRecord, self).__init__(config)
        cookie_name = config.get('cookie_name', 'beaker.session.id')
        self.record_path = config.get('record_dir', './')

    def get_top_frame_params(self, wbrequest, mod):
        params = (super(DynRecord, self).
                  get_top_frame_params(wbrequest, mod))

        manager = wbrequest.env.get('w_manager')

        if manager:
            user, coll = wbrequest.coll.split('/', 1)
            info = manager.get_info(user, coll)
            params['info'] = json.dumps(info)

        return params

    def _ignore_proxies(self, wbrequest):
        return wbrequest.env.get('webrec.no_space', False)

    def _live_request_headers(self, wbrequest):
        path = wbrequest.custom_params.get('output_dir')
        if not path:
            path = self.record_path

        if not self._ignore_proxies(wbrequest):
            sesh_id = wbrequest.custom_params.get('sesh_id', wbrequest.coll)
            sesh_id = sesh_id.replace('/', ':')
            user_id = wbrequest.custom_params.get('user_id')

            target = dict(output_dir=path,
                          sesh_id=sesh_id,
                          user_id=user_id)

            req_headers = {'warcprox-meta': json.dumps(target)}
        else:
            req_headers = {}

        # reset HTTP_COOKIE to guarded request_cookie for LiveRewriter
        if 'webrec.request_cookie' in wbrequest.env:
            wbrequest.env['HTTP_COOKIE'] = wbrequest.env['webrec.request_cookie']

        try:
            del wbrequest.env['HTTP_X_PUSH_STATE_REQUEST']
        except:
            pass

        return req_headers

    def _make_response(self, wbrequest, status_headers, gen, is_rewritten):
        if (status_headers.get_statuscode() == '500' and
            status_headers.get_header('X-Archive-Orig-x-warcprox-error') == '500'):
            msg = 'This url could not be recorded: '
            raise LiveResourceException(msg, url=wbrequest.wb_url.url)

        return WbResponse(status_headers, gen)

#=================================================================
class DynUrlRewriter(UrlRewriter):
    def rewrite(self, url, mod=None):
        new_url = super(DynUrlRewriter, self).rewrite(url, mod)
        if new_url == url:
            return new_url

        parts = new_url.split('://', 1)
        if len(parts) < 2:
            return new_url

        rand = base64.b32encode(os.urandom(5))
        parts[1] = rand + '.' + parts[1]

        #print('DYN RW 2')
        #host = parts[2].split('.')[0]
        #parts[1] = host + '.' + parts[1]
        new_url = '://'.join(parts)
        return new_url
<|endoftext|>


<commit_before>
import os.path as op
import mne
from mne.datasets import sample
from mne import read_forward_solution
from mne.minimum_norm import (read_inverse_operator,
                              point_spread_function, cross_talk_function)

from nose.tools import assert_true

data_path = op.join(sample.data_path(download=False), 'MEG', 'sample')
fname_inv = op.join(data_path, 'sample_audvis-meg-oct-6-meg-inv.fif')
fname_fwd = op.join(data_path, 'sample_audvis-meg-oct-6-fwd.fif')

fname_label = [data_path + '/MEG/sample/labels/Aud-rh.label',
               data_path + '/MEG/sample/labels/Aud-lh.label']

snr = 3.0
lambda2 = 1.0 / snr ** 2

@sample.requires_sample_data
def test_psf_ctf():
    """Test computation of PSFs and CTFs for linear estimators
    """

    inverse_operator = read_inverse_operator(fname_inv)
    forward = read_forward_solution(fname_fwd, force_fixed=False,
                                    surf_ori=True)
    labels = [mne.read_label(ss) for ss in fname_label]

    method = 'MNE'
    n_svd_comp = 2

    # Test PSFs (then CTFs)
    for mode in ('sum', 'svd'):
        stc_psf, psf_ev = point_spread_function(inverse_operator,
                                                forward,
                                                method=method,
                                                labels=labels,
                                                lambda2=lambda2,
                                                pick_ori='normal',
                                                mode=mode,
                                                n_svd_comp=n_svd_comp)

        n_vert, n_samples = stc_psf.shape
        should_n_vert = (inverse_operator['src'][1]['vertno'].shape[0] +
                         inverse_operator['src'][0]['vertno'].shape[0])
        if mode == 'svd':
            should_n_samples = len(labels) * n_svd_comp + 1
        else:
            should_n_samples = len(labels) + 1

        assert_true(n_vert == should_n_vert)
        assert_true(n_samples == should_n_samples)

        n_chan, n_samples = psf_ev.data.shape
        assert_true(n_chan == forward['nchan'])

    forward = read_forward_solution(fname_fwd, force_fixed=True, surf_ori=True)

    # Test CTFs
    for mode in ('sum', 'svd'):
        stc_ctf = cross_talk_function(inverse_operator, forward,
                                      labels, method=method,
                                      lambda2=lambda2,
                                      signed=False, mode=mode,
                                      n_svd_comp=n_svd_comp)

        n_vert, n_samples = stc_ctf.shape
        should_n_vert = (inverse_operator['src'][1]['vertno'].shape[0] +
                         inverse_operator['src'][0]['vertno'].shape[0])
        if mode == 'svd':
            should_n_samples = len(labels) * n_svd_comp + 1
        else:
            should_n_samples = len(labels) + 1

        assert_true(n_vert == should_n_vert)
        assert_true(n_samples == should_n_samples)
<commit_msg>FIX : fix paths in tests<commit_after>
import os.path as op
import mne
from mne.datasets import sample
from mne import read_forward_solution
from mne.minimum_norm import (read_inverse_operator,
                              point_spread_function, cross_talk_function)

from nose.tools import assert_true

data_path = op.join(sample.data_path(download=False), 'MEG', 'sample')
fname_inv = op.join(data_path, 'sample_audvis-meg-oct-6-meg-inv.fif')
fname_fwd = op.join(data_path, 'sample_audvis-meg-oct-6-fwd.fif')

fname_label = [op.join(data_path, 'labels', 'Aud-rh.label'),
               op.join(data_path, 'labels', 'Aud-lh.label')]

snr = 3.0
lambda2 = 1.0 / snr ** 2

@sample.requires_sample_data
def test_psf_ctf():
    """Test computation of PSFs and CTFs for linear estimators
    """

    inverse_operator = read_inverse_operator(fname_inv)
    forward = read_forward_solution(fname_fwd, force_fixed=False,
                                    surf_ori=True)
    labels = [mne.read_label(ss) for ss in fname_label]

    method = 'MNE'
    n_svd_comp = 2

    # Test PSFs (then CTFs)
    for mode in ('sum', 'svd'):
        stc_psf, psf_ev = point_spread_function(inverse_operator,
                                                forward,
                                                method=method,
                                                labels=labels,
                                                lambda2=lambda2,
                                                pick_ori='normal',
                                                mode=mode,
                                                n_svd_comp=n_svd_comp)

        n_vert, n_samples = stc_psf.shape
        should_n_vert = (inverse_operator['src'][1]['vertno'].shape[0] +
                         inverse_operator['src'][0]['vertno'].shape[0])
        if mode == 'svd':
            should_n_samples = len(labels) * n_svd_comp + 1
        else:
            should_n_samples = len(labels) + 1

        assert_true(n_vert == should_n_vert)
        assert_true(n_samples == should_n_samples)

        n_chan, n_samples = psf_ev.data.shape
        assert_true(n_chan == forward['nchan'])

    forward = read_forward_solution(fname_fwd, force_fixed=True, surf_ori=True)

    # Test CTFs
    for mode in ('sum', 'svd'):
        stc_ctf = cross_talk_function(inverse_operator, forward,
                                      labels, method=method,
                                      lambda2=lambda2,
                                      signed=False, mode=mode,
                                      n_svd_comp=n_svd_comp)

        n_vert, n_samples = stc_ctf.shape
        should_n_vert = (inverse_operator['src'][1]['vertno'].shape[0] +
                         inverse_operator['src'][0]['vertno'].shape[0])
        if mode == 'svd':
            should_n_samples = len(labels) * n_svd_comp + 1
        else:
            should_n_samples = len(labels) + 1

        assert_true(n_vert == should_n_vert)
        assert_true(n_samples == should_n_samples)
<|endoftext|>


<commit_before>import sqlite3
import json
import logging
import time
import os
import sys
import hashlib
import gzip
import psycopg2
import datetime

def convert(val,typ):
    if typ=="date" and val is not None:
        if type(val) in (str,long) and val.strip() != '':
            val = [int(x) for x in val.split('/')]
            val.reverse()
            val = datetime.date(*val)
        elif type(val) in (int,long):
            val = datetime.datetime.fromtimestamp(val).date()
        else:
            val = None
    elif typ=="jsonb":
        val = json.dumps(val)
    return val

class dump_to_psql(object):

    def process(self,input,output,table,field_definitions):

        good_field_definitions = [(x[0].replace('/','_'),x[1]) for x in field_definitions]

        conn = psycopg2.connect('dbname=obudget')
        c = conn.cursor()
        c.execute("""DROP TABLE if exists %s;""" % table)
        c.execute("""CREATE TABLE %s
                     (%s);""" % (table,",".join("%s %s" % (x[0],x[1]) for x in good_field_definitions)))
        c.execute("""grant select ON %s to redash_reader;""" % table)

        if input.endswith('.gz'):
            infile = gzip.GzipFile(input)
        else:
            infile = file(input)
        to_insert = []
        for line in infile:
            line = line.strip()
            data = json.loads(line)
            values = [convert(data.get(field),typ) for field,typ in field_definitions]
            to_insert.append(values)
        c.executemany("""INSERT INTO %s VALUES(%s)""" % (table,",".join([r"%s"]*len(field_definitions))), to_insert)

        logging.debug("TABLE %s got %s records" % (table,len(to_insert)))

        conn.commit()

        for f in good_field_definitions:
            fieldname = f[0]
            if f[1] == 'jsonb':
                continue
            c.execute("""create index {0}_{1}_idx_asc on {0}({1} asc);""".format(table,fieldname))
            c.execute("""create index {0}_{1}_idx_desc on {0}({1} desc);""".format(table,fieldname))

        conn.commit()
        conn.close()

        file(output,"w").write("OK")

if __name__ == "__main__":
    input = sys.argv[1]
    output = sys.argv[2]
    key_fields = sys.argv[3:]
    processor = dump_to_db().process(input,output,key_fields)
<commit_msg>fix dump to psql<commit_after>import sqlite3
import json
import logging
import time
import os
import sys
import hashlib
import gzip
import psycopg2
import datetime

def convert(val,typ):
    if typ=="date" and val is not None:
        if type(val) in (str,unicode) and val.strip() != '':
            val = [int(x) for x in val.split('/')]
            val.reverse()
            val = datetime.date(*val)
        elif type(val) in (int,long):
            val = datetime.datetime.fromtimestamp(val).date()
        else:
            val = None
    elif typ=="jsonb":
        val = json.dumps(val)
    return val

class dump_to_psql(object):

    def process(self,input,output,table,field_definitions):

        good_field_definitions = [(x[0].replace('/','_'),x[1]) for x in field_definitions]

        conn = psycopg2.connect('dbname=obudget')
        c = conn.cursor()
        c.execute("""DROP TABLE if exists %s;""" % table)
        c.execute("""CREATE TABLE %s
                     (%s);""" % (table,",".join("%s %s" % (x[0],x[1]) for x in good_field_definitions)))
        c.execute("""grant select ON %s to redash_reader;""" % table)

        if input.endswith('.gz'):
            infile = gzip.GzipFile(input)
        else:
            infile = file(input)
        to_insert = []
        for line in infile:
            line = line.strip()
            data = json.loads(line)
            values = [convert(data.get(field),typ) for field,typ in field_definitions]
            to_insert.append(values)
        c.executemany("""INSERT INTO %s VALUES(%s)""" % (table,",".join([r"%s"]*len(field_definitions))), to_insert)

        logging.debug("TABLE %s got %s records" % (table,len(to_insert)))

        conn.commit()

        for f in good_field_definitions:
            fieldname = f[0]
            if f[1] == 'jsonb':
                continue
            c.execute("""create index {0}_{1}_idx_asc on {0}({1} asc);""".format(table,fieldname))
            c.execute("""create index {0}_{1}_idx_desc on {0}({1} desc);""".format(table,fieldname))

        conn.commit()
        conn.close()

        file(output,"w").write("OK")

if __name__ == "__main__":
    input = sys.argv[1]
    output = sys.argv[2]
    key_fields = sys.argv[3:]
    processor = dump_to_db().process(input,output,key_fields)
<|endoftext|>


<commit_before># SPDX-License-Identifier: Apache-2.0
#
# Copyright (C) 2018, Arm Limited and contributors.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import logging
import os
from shlex import quote

from shlex import quote
from datetime import datetime
from devlib.utils.misc import list_to_mask

from lisa.utils import Loggable, ArtifactPath

class Workload(Loggable):
    """
    This is pretty much a wrapper around a command to execute on a target.

    :param target: The Target on which to execute this workload
    :type target: Target

    :param name: Name of the workload. Useful for naming related artefacts.
    :type name: str

    :param res_dir: Directory into which artefacts will be stored
    :type res_dir: str

    :ivar command: Will be called in :meth:`run`. Daughter classes should
      specify its value as needed.

    **Design notes**

    ``__init__`` is there to initialize a given workload, and :meth:`run`
    can be called on it several times, with varying arguments.
    As much work as possible should be delegated to :meth:`run`, so that
    different flavours of the same workload can be run without the hassle of
    creating a superfluous amount of new instances. However, when persistent
    data is involved (e.g. the workload depends on a file), then this data
    should be exposed as an ``__init__`` parameter.

    **Implementation example**::

        class Printer(Workload):
            def __init__(self, target, name, res_dir=None):
                super().__init__(target, name, res_dir)
                self.command = "echo"

            def run(self, cpus=None, cgroup=None, background=False, as_root=False, value=42):
                self.command = "{} {}".format(self.command, shlex.quote(value))
                super().run(cpus, cgroup, background, as_root)

    **Usage example**::

        >>> printer = Printer(target, "test")
        >>> printer.run()
        INFO    : Printer      : Execution start: echo 42
        INFO    : Printer      : Execution complete
        >>> print printer.output
        42\r\n
    """

    required_tools = ['taskset']
    """
    The tools required to execute the workload. See
    :meth:`lisa.target.Target.install_tools`.
    """

    def __init__(self, target, name, res_dir=None):
        self.target = target
        self.name = name
        self.command = None
        self.output = ""

        wlgen_dir = self.target.path.join(target.working_directory,
                                          "lisa", "wlgen")
        self.target.execute("mkdir -p {}/lisa/wlgen".format(quote(wlgen_dir)))

        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        temp_fmt = "{}_{}_XXXXXX".format(self.name, timestamp)
        self.run_dir = target.execute("mktemp -d -p {} {}".format(
                                      quote(wlgen_dir), quote(temp_fmt))).strip()

        logger = self.get_logger()
        logger.info("Creating target's run_dir: %s", self.run_dir)

        res_dir = res_dir if res_dir else target.get_res_dir(
            name='{}-{}'.format(self.__class__.__qualname__, name)
        )
        self.res_dir = res_dir

        self.target.install_tools(self.required_tools)

    def wipe_run_dir(self):
        """
        Wipe all content from the ``run_dir`` target directory.

        .. note :: This function should only be called directly in interactive
            sessions. For other purposes, use :class:`Workload` instances as a
            context manager.
        """
        logger = self.get_logger()
        logger.info("Wiping run_dir [%s]", self.run_dir)
        self.target.execute("rm -rf {}".format(quote(self.run_dir)))

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Wipe the run directory on the target.
        """
        self.wipe_run_dir()

    def run(self, cpus=None, cgroup=None, background=False, as_root=False, timeout=None):
        """
        Execute the workload on the configured target.

        :param cpus: CPUs on which to restrict the workload execution (taskset)
        :type cpus: list(int)

        :param cgroup: cgroup in which to run the workload
        :type cgroup: str

        :param background: Whether to run the workload in background or not
        :type background: bool

        :param as_root: Whether to run the workload as root or not
        :type as_root: bool

        :param timeout: Timeout in seconds for the workload execution.
        :type timeout: int

        :raise devlib.exception.TimeoutError: When the specified ``timeout`` is hit.

        The standard output will be saved into a file in ``self.res_dir``
        """
        logger = self.get_logger()
        if not self.command:
            raise RuntimeError("Workload does not specify any command to execute")

        _command = self.command
        target = self.target

        if cpus:
            taskset_bin = target.which('taskset')
            if not taskset_bin:
                raise RuntimeError("Could not find 'taskset' executable on the target")

            cpumask = list_to_mask(cpus)
            taskset_cmd = '{} {}'.format(quote(taskset_bin), quote('0x{:x}'.format(cpumask)))
            _command = '{} {}'.format(taskset_cmd, _command)

        if cgroup:
            _command = target.cgroups.run_into_cmd(cgroup, _command)

        _command = 'cd {} && {}'.format(quote(self.run_dir), _command)

        logger.info("Execution start: %s", _command)

        if background:
            target.background(_command, as_root=as_root)
        else:
            self.output = target.execute(_command, as_root=as_root, timeout=timeout)
            logger.info("Execution complete")

            logfile = ArtifactPath.join(self.res_dir, 'output.log')
            logger.debug('Saving stdout to %s...', logfile)

            with open(logfile, 'w') as ofile:
                ofile.write(self.output)

# vim :set tabstop=4 shiftwidth=4 textwidth=80 expandtab
<commit_msg>lisa.wlgen.workload: Align run dir wipe log<commit_after># SPDX-License-Identifier: Apache-2.0
#
# Copyright (C) 2018, Arm Limited and contributors.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import logging
import os
from shlex import quote

from shlex import quote
from datetime import datetime
from devlib.utils.misc import list_to_mask

from lisa.utils import Loggable, ArtifactPath

class Workload(Loggable):
    """
    This is pretty much a wrapper around a command to execute on a target.

    :param target: The Target on which to execute this workload
    :type target: Target

    :param name: Name of the workload. Useful for naming related artefacts.
    :type name: str

    :param res_dir: Directory into which artefacts will be stored
    :type res_dir: str

    :ivar command: Will be called in :meth:`run`. Daughter classes should
      specify its value as needed.

    **Design notes**

    ``__init__`` is there to initialize a given workload, and :meth:`run`
    can be called on it several times, with varying arguments.
    As much work as possible should be delegated to :meth:`run`, so that
    different flavours of the same workload can be run without the hassle of
    creating a superfluous amount of new instances. However, when persistent
    data is involved (e.g. the workload depends on a file), then this data
    should be exposed as an ``__init__`` parameter.

    **Implementation example**::

        class Printer(Workload):
            def __init__(self, target, name, res_dir=None):
                super().__init__(target, name, res_dir)
                self.command = "echo"

            def run(self, cpus=None, cgroup=None, background=False, as_root=False, value=42):
                self.command = "{} {}".format(self.command, shlex.quote(value))
                super().run(cpus, cgroup, background, as_root)

    **Usage example**::

        >>> printer = Printer(target, "test")
        >>> printer.run()
        INFO    : Printer      : Execution start: echo 42
        INFO    : Printer      : Execution complete
        >>> print printer.output
        42\r\n
    """

    required_tools = ['taskset']
    """
    The tools required to execute the workload. See
    :meth:`lisa.target.Target.install_tools`.
    """

    def __init__(self, target, name, res_dir=None):
        self.target = target
        self.name = name
        self.command = None
        self.output = ""

        wlgen_dir = self.target.path.join(target.working_directory,
                                          "lisa", "wlgen")
        self.target.execute("mkdir -p {}/lisa/wlgen".format(quote(wlgen_dir)))

        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        temp_fmt = "{}_{}_XXXXXX".format(self.name, timestamp)
        self.run_dir = target.execute("mktemp -d -p {} {}".format(
                                      quote(wlgen_dir), quote(temp_fmt))).strip()

        logger = self.get_logger()
        logger.info("Creating target's run_dir: %s", self.run_dir)

        res_dir = res_dir if res_dir else target.get_res_dir(
            name='{}-{}'.format(self.__class__.__qualname__, name)
        )
        self.res_dir = res_dir

        self.target.install_tools(self.required_tools)

    def wipe_run_dir(self):
        """
        Wipe all content from the ``run_dir`` target directory.

        .. note :: This function should only be called directly in interactive
            sessions. For other purposes, use :class:`Workload` instances as a
            context manager.
        """
        logger = self.get_logger()
        logger.info("Wiping target run directory: {}".format(self.run_dir))
        self.target.execute("rm -rf {}".format(quote(self.run_dir)))

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Wipe the run directory on the target.
        """
        self.wipe_run_dir()

    def run(self, cpus=None, cgroup=None, background=False, as_root=False, timeout=None):
        """
        Execute the workload on the configured target.

        :param cpus: CPUs on which to restrict the workload execution (taskset)
        :type cpus: list(int)

        :param cgroup: cgroup in which to run the workload
        :type cgroup: str

        :param background: Whether to run the workload in background or not
        :type background: bool

        :param as_root: Whether to run the workload as root or not
        :type as_root: bool

        :param timeout: Timeout in seconds for the workload execution.
        :type timeout: int

        :raise devlib.exception.TimeoutError: When the specified ``timeout`` is hit.

        The standard output will be saved into a file in ``self.res_dir``
        """
        logger = self.get_logger()
        if not self.command:
            raise RuntimeError("Workload does not specify any command to execute")

        _command = self.command
        target = self.target

        if cpus:
            taskset_bin = target.which('taskset')
            if not taskset_bin:
                raise RuntimeError("Could not find 'taskset' executable on the target")

            cpumask = list_to_mask(cpus)
            taskset_cmd = '{} {}'.format(quote(taskset_bin), quote('0x{:x}'.format(cpumask)))
            _command = '{} {}'.format(taskset_cmd, _command)

        if cgroup:
            _command = target.cgroups.run_into_cmd(cgroup, _command)

        _command = 'cd {} && {}'.format(quote(self.run_dir), _command)

        logger.info("Execution start: %s", _command)

        if background:
            target.background(_command, as_root=as_root)
        else:
            self.output = target.execute(_command, as_root=as_root, timeout=timeout)
            logger.info("Execution complete")

            logfile = ArtifactPath.join(self.res_dir, 'output.log')
            logger.debug('Saving stdout to %s...', logfile)

            with open(logfile, 'w') as ofile:
                ofile.write(self.output)

# vim :set tabstop=4 shiftwidth=4 textwidth=80 expandtab
<|endoftext|>