<commit_before>#!/usr/bin/env python
"""
A set of basic Audio and DSP functions using audioread and numpy
Mostly the core dsp tools collected from librosa, but with fewer dependencies
"""
import os
import warnings
import math

import audioread
import numpy as np

def load(path, sr=22050, mono=True, offset=0.0, duration=None,
         dtype=np.float32):
    """
    Load an audio file as a numpy floating point time series.

    Args:
        path: string
            path to the input file.
            Any format supported by `audioread` will work.
        sr: number > 0 [scalar]
            target sampling rate
            'None' uses the native sampling rate
        mono: bool
            convert signal to mono
        offset: float
            start reading after this time (in seconds)
        duration: float
            only load up to this much audio (in seconds)
        dtype: numeric type
            data type of `y`
    Returns:
        y: np.ndarray [shape=(n,) or (2, n)]
            audio time series
        sr: number > 0 [scalar]
            sampling rate of `y`
    """
    y = []
    with audioread.audio_open(os.path.realpath(path)) as input_file:
        sr_native = input_file.samplerate
        n_channels = input_file.channels

        start = int(np.round(sr_native * offset)) * n_channels

        if duration is None:
            end = np.inf
        else:
            end = start + (int(np.round(sr_native * duration)) * n_channels)

        n = 0

        for frame in input_file:
            frame = buf_to_float(frame, dtype=dtype)
            n_prev = n
            n = n + len(frame)

            if n < start:
                # offset is after the current frame
                # keep reading
                continue

            if end < n_prev:
                # we're off the end.  stop reading
                break

            if end < n:
                # the end is in this frame.  crop.
                frame = frame[:end - n_prev]

            if n_prev <= start <= n:
                # beginning is in this frame
                frame = frame[(start - n_prev):]

            # tack on the current frame
            y.append(frame)

    if y:
        y = np.concatenate(y)

        if n_channels > 1:
            y = y.reshape((-1, 2)).T
            if mono:
                y = to_mono(y)

        if sr is not None:
            y = resample(y, sr_native, sr)

        else:
            sr = sr_native

    # Final cleanup for dtype and contiguity
    y = np.ascontiguousarray(y, dtype=dtype)

    return (y, sr)

def write_wav(path, y, sr, norm=True):
    """
    Output a time series as a .wav file
    Args:
        path : str
            path to save the output wav file
        y : np.ndarray [shape=(n,) or (2,n)]
            audio time series (mono or stereo)
        sr : int > 0 [scalar]
            sampling rate of `y`
        norm : boolean [scalar]
            enable amplitude normalization
    """
    # Validate the buffer.  Stereo is okay here.
    valid_audio(y, mono=False)
    # Normalize
    if norm:
        y = y / np.max(np.abs(y))
    else:
        y = y
    # Convert to 16bit int
    wav = buf_to_int(y)
    # Check for stereo
    if wav.ndim > 1 and wav.shape[0] == 2:
        wav = wav.T
    # Save
    from scipy.io import wavfile
    wavfile.write(path, sr, wav)

def to_mono(y):
    """
    Force an audio signal down to mono.
    Args:
        y : np.ndarray [shape=(2,n) or shape=(n,)]
            audio time series, either stereo or mono
    Returns:
        y_mono : np.ndarray [shape=(n,)]
            `y` as a monophonic time-series
    """
    valid_audio(y, mono=False)
    if y.ndim > 1:
        y = np.mean(y, axis=0)
    return y

def resample(y, orig_sr, target_sr, fix=True, scale=False):
    """Resample a time series from orig_sr to target_sr
    Args:
        y : np.ndarray [shape=(n,) or shape=(2, n)]
            audio time series.  Can be mono or stereo.
        orig_sr : number > 0 [scalar]
            original sampling rate of `y`
        target_sr : number > 0 [scalar]
            target sampling rate
        fix : bool
            adjust the length of the resampled signal to be of size exactly
            `ceil(target_sr * len(y) / orig_sr)`
        scale : bool
            Scale the resampled signal so that `y` and `y_hat` have
            approximately equal total energy.
    Returns
        y_hat : np.ndarray [shape=(n * target_sr / orig_sr,)]
            `y` resampled from `orig_sr` to `target_sr`
    """
    if y.ndim > 1:
        return np.vstack([resample(yi, orig_sr, target_sr, fix=fix)
                          for yi in y])
    valid_audio(y, mono=True)
    if orig_sr == target_sr:
        return y
    ratio = float(target_sr) / orig_sr
    n_samples = int(np.ceil(y.shape[-1] * ratio))
    try:
        # Try to use scikits.samplerate if available
        import scikits.samplerate as samplerate
        y_hat = samplerate.resample(y.T, ratio, 'sinc_best').T
    except ImportError:
        warnings.warn('Could not import scikits.samplerate. '
                      'Falling back to scipy.signal')
        import scipy.signal
        y_hat = scipy.signal.resample(y, n_samples, axis=-1)
    if fix:
        y_hat = fix_length(y_hat, n_samples)
    if scale:
        y_hat /= np.sqrt(ratio)
    return np.ascontiguousarray(y_hat, dtype=y.dtype)

def get_duration(y=None, sr=22050):
    """
    Compute the duration (in seconds) of an audio time series.
    Args:
        y: np.ndarray [shape=(n,), (2, n)] or None
            audio time series
        sr: number > 0 [scalar]
            audio sampling rate of `y`
    Returns:
        d: float >= 0
            Duration (in seconds) of the input time series.
    """
    valid_audio(y, mono=False)
    if y.ndim == 1:
        n_samples = len(y)
    else:
        n_samples = y.shape[-1]
    return float(n_samples) / sr

def fix_length(data, size, axis=-1, **kwargs):
    """
    Fix the length of array `data` to exactly `size`.
    If `data.shape[axis] < n`, pad according to kwargs.mode

    Args:
        data: np.ndarray
            array to be length-adjusted
        size: int >= 0 [scalar]
            desired length of the array
        axis: int, <= data.ndim
            axis along which to fix length
        kwargs:
            Parameters to np.pad()
    Returns:
        np.ndarray [shape=data.shape]
            `data` either trimmed or padded to length `size`
            along the specified axis.
    """
    kwargs.setdefault('mode', 'constant')
    n = data.shape[axis]
    if n > size:
        slices = [slice(None)] * data.ndim
        slices[axis] = slice(0, size)
        return data[slices]
    elif n < size:
        lengths = [(0, 0)] * data.ndim
        lengths[axis] = (0, size - n)
        return np.pad(data, lengths, **kwargs)
    return data

def buf_to_float(x, n_bytes=2, dtype=np.float32):
    """
    Convert an integer buffer to floating point values.
    Used for loading integer-valued wav data into numpy arrays.
    Args:
        x : np.ndarray [dtype=int]
            The integer-valued data buffer
        n_bytes : int [1, 2, 4]
            The number of bytes per sample in `x`
        dtype : numeric type
            The target output type (default: 32-bit float)
    Returns:
        x_float : np.ndarray [dtype=float]
            The input data buffer cast to floating point
    """
    # Invert the scale of the data
    scale = 1./float(1 << ((8 * n_bytes) - 1))
    # Construct the format string
    fmt = '<i{:d}'.format(n_bytes)
    # Rescale and format the data buffer
    return scale * np.frombuffer(x, fmt).astype(dtype)

def buf_to_int(x, n_bytes=2):
    """
    Convert a floating point buffer into integer values.
    This is primarily useful as an intermediate step in wav output.
    Args:
        x: np.ndarray [dtype=float]
            Floating point data buffer
        n_bytes: int [1, 2, 4]
            Number of bytes per output sample
    Returns:
        x_int : np.ndarray [dtype=int]
            The original buffer cast to integer type.
    """
    if n_bytes not in [1, 2, 4]:
        raise ValueError('n_bytes must be one of {1, 2, 4}')
    # What is the scale of the input data?
    scale = float(1 << ((8 * n_bytes) - 1))
    # Construct a format string
    fmt = '<i{:d}'.format(n_bytes)
    # Rescale and cast the data
    return (x * scale).astype(fmt)

def valid_audio(y, mono=True):
    """
    Validate whether a variable contains valid audio data.
        Args:
        y : np.ndarray
          The input data to validate
        mono : bool
      Wheth:
        valid : bool
            True if all tests pass
    Raises:
        ParameterError
            If `y` fails to meet the following criteria:
                - `type(y)` is `np.ndarray`
                - `mono == True` and `y.ndim` is not 1
                - `mono == False` and `y.ndim` is not 1 or 2
                - `np.isfinite(y).all()` is not True
    """
    if not isinstance(y, np.ndarray):
        raise ValueError('Data must be of type numpy.ndarray')
    if mono and y.ndim != 1:
        raise ValueError('Invalid shape for monophonic audio: '
                         'ndim={:d}, shape={}'.format(y.ndim, y.shape))
    elif y.ndim > 2:
        raise ValueError('Invalid shape for audio: '
                         'ndim={:d}, shape={}'.format(y.ndim, y.shape))
    if not np.isfinite(y).all():
        raise ValueError('Audio buffer is not finite everywhere')
    return True

def normalize(S, norm=np.inf, axis=0):
    """
    Normalize the columns or rows of a matrix
    Args:
        S : np.ndarray [shape=(d, n)]
            The matrix to normalize
        norm : {np.inf, -np.inf, 0, float > 0, None}
            - `np.inf`  : maximum absolute value
            - `-np.inf` : mininum absolute value
            - `0`    : number of non-zeros
            - float  : corresponding l_p norm.
                See `scipy.linalg.norm` for details.
            - None : no normalization is performed
        axis : int [scalar]
            Axis along which to compute the norm.
            `axis=0` will normalize columns, `axis=1` will normalize rows.
            `axis=None` will normalize according to the entire matrix.
    Returns:
        S_norm : np.ndarray [shape=S.shape]
            Normalized matrix
    """
    # All norms only depend on magnitude, let's do that first
    mag = np.abs(S)

    if norm == np.inf:
        length = np.max(mag, axis=axis, keepdims=True)

    elif norm == -np.inf:
        length = np.min(mag, axis=axis, keepdims=True)

    elif norm == 0:
        length = np.sum(mag > 0, axis=axis, keepdims=True)

    elif np.issubdtype(type(norm), np.number) and norm > 0:
        length = np.sum(mag ** norm, axis=axis, keepdims=True)**(1./norm)

    elif norm is None:
        return S

    else:
        raise ValueError('Unsupported norm: {}'.format(repr(norm)))

    # Avoid div-by-zero
    SMALL_FLOAT = 1e-20
    length[length < SMALL_FLOAT] = 1.0

    return S / length

def fractional_delay(y, time, sr):
    """
    Prepend zeros to delay signal by time seconds
    Use linear interpolation through convolution for fractional delay
    """
    samples = time * sr
    if isinstance(samples, int) or samples.is_integer():
        return delay(y, time, sr)

    f, i = math.modf(samples)
    y = delay(y, i, sr, mode='samples')
    import scipy.signal
    y = scipy.signal.convolve(y, [f, i-f], "same")
    return y

def delay(y, time, sr, mode='time'):
    """
    Prepend zeros to delay signal by time seconds
    """
    if mode == 'time':
        samples = time * sr
    else:
        samples = time
    # fractional delay by upsampling has been commented out for now
    #decimals = sum(c != '0' for c in str(round(samples % 1, 4))[2:])
    #new_sr = None
    #if decimals > 0:
    #    new_sr = sr * (decimals + 1)
    #    samples = time * new_sr
    #    y = resample(y, sr, new_sr)
    y = np.pad(
        y,
        pad_width=[round(samples), 0],
        mode='constant',
        constant_values=0
    )
    #if new_sr:
    #    y = resample(y, new_sr, sr)
    return y

def sum_signals(signals):
    """
    Sum together a list of mono signals
    append zeros to match the longest array
    """
    if not signals:
        return np.array([])
    max_length = max(len(sig) for sig in signals)
    y = np.zeros(max_length)
    for sig in signals:
        padded = np.zeros(max_length)
        padded[0:len(sig)] = sig
        y += padded
    return y

def channel_merge(channels):
    """
    vstack a list of mono signals
    append zeros to match the longest signal
    """
    longest = max(len(channel) for channel in channels)
    padded_channels = []
    for i, channel in enumerate(channels):
        if len(channel) < longest:
            padded = np.zeros(longest)
            padded[0:len(channel)] = channel
            channel = padded
        padded_channels.append(channel)
    y = np.vstack(padded_channels)
    return y
<commit_msg>add headroom to wav_write normalize<commit_after>#!/usr/bin/env python
"""
A set of basic Audio and DSP functions using audioread and numpy
Mostly the core dsp tools collected from librosa, but with fewer dependencies
"""
import os
import warnings
import math

import audioread
import numpy as np

def load(path, sr=22050, mono=True, offset=0.0, duration=None,
         dtype=np.float32):
    """
    Load an audio file as a numpy floating point time series.

    Args:
        path: string
            path to the input file.
            Any format supported by `audioread` will work.
        sr: number > 0 [scalar]
            target sampling rate
            'None' uses the native sampling rate
        mono: bool
            convert signal to mono
        offset: float
            start reading after this time (in seconds)
        duration: float
            only load up to this much audio (in seconds)
        dtype: numeric type
            data type of `y`
    Returns:
        y: np.ndarray [shape=(n,) or (2, n)]
            audio time series
        sr: number > 0 [scalar]
            sampling rate of `y`
    """
    y = []
    with audioread.audio_open(os.path.realpath(path)) as input_file:
        sr_native = input_file.samplerate
        n_channels = input_file.channels

        start = int(np.round(sr_native * offset)) * n_channels

        if duration is None:
            end = np.inf
        else:
            end = start + (int(np.round(sr_native * duration)) * n_channels)

        n = 0

        for frame in input_file:
            frame = buf_to_float(frame, dtype=dtype)
            n_prev = n
            n = n + len(frame)

            if n < start:
                # offset is after the current frame
                # keep reading
                continue

            if end < n_prev:
                # we're off the end.  stop reading
                break

            if end < n:
                # the end is in this frame.  crop.
                frame = frame[:end - n_prev]

            if n_prev <= start <= n:
                # beginning is in this frame
                frame = frame[(start - n_prev):]

            # tack on the current frame
            y.append(frame)

    if y:
        y = np.concatenate(y)

        if n_channels > 1:
            y = y.reshape((-1, 2)).T
            if mono:
                y = to_mono(y)

        if sr is not None:
            y = resample(y, sr_native, sr)

        else:
            sr = sr_native

    # Final cleanup for dtype and contiguity
    y = np.ascontiguousarray(y, dtype=dtype)

    return (y, sr)

def write_wav(path, y, sr, norm=True):
    """
    Output a time series as a .wav file
    Args:
        path : str
            path to save the output wav file
        y : np.ndarray [shape=(n,) or (2,n)]
            audio time series (mono or stereo)
        sr : int > 0 [scalar]
            sampling rate of `y`
        norm : boolean [scalar]
            enable amplitude normalization
    """
    # Validate the buffer.  Stereo is okay here.
    valid_audio(y, mono=False)
    # Normalize
    if norm:
        headroom = 0.00003
        y = y / (np.max(np.abs(y)) + headroom)
    else:
        y = y
    # Convert to 16bit int
    wav = buf_to_int(y)
    # Check for stereo
    if wav.ndim > 1 and wav.shape[0] == 2:
        wav = wav.T
    # Save
    from scipy.io import wavfile
    wavfile.write(path, sr, wav)

def to_mono(y):
    """
    Force an audio signal down to mono.
    Args:
        y : np.ndarray [shape=(2,n) or shape=(n,)]
            audio time series, either stereo or mono
    Returns:
        y_mono : np.ndarray [shape=(n,)]
            `y` as a monophonic time-series
    """
    valid_audio(y, mono=False)
    if y.ndim > 1:
        y = np.mean(y, axis=0)
    return y

def resample(y, orig_sr, target_sr, fix=True, scale=False):
    """Resample a time series from orig_sr to target_sr
    Args:
        y : np.ndarray [shape=(n,) or shape=(2, n)]
            audio time series.  Can be mono or stereo.
        orig_sr : number > 0 [scalar]
            original sampling rate of `y`
        target_sr : number > 0 [scalar]
            target sampling rate
        fix : bool
            adjust the length of the resampled signal to be of size exactly
            `ceil(target_sr * len(y) / orig_sr)`
        scale : bool
            Scale the resampled signal so that `y` and `y_hat` have
            approximately equal total energy.
    Returns
        y_hat : np.ndarray [shape=(n * target_sr / orig_sr,)]
            `y` resampled from `orig_sr` to `target_sr`
    """
    if y.ndim > 1:
        return np.vstack([resample(yi, orig_sr, target_sr, fix=fix)
                          for yi in y])
    valid_audio(y, mono=True)
    if orig_sr == target_sr:
        return y
    ratio = float(target_sr) / orig_sr
    n_samples = int(np.ceil(y.shape[-1] * ratio))
    try:
        # Try to use scikits.samplerate if available
        import scikits.samplerate as samplerate
        y_hat = samplerate.resample(y.T, ratio, 'sinc_best').T
    except ImportError:
        warnings.warn('Could not import scikits.samplerate. '
                      'Falling back to scipy.signal')
        import scipy.signal
        y_hat = scipy.signal.resample(y, n_samples, axis=-1)
    if fix:
        y_hat = fix_length(y_hat, n_samples)
    if scale:
        y_hat /= np.sqrt(ratio)
    return np.ascontiguousarray(y_hat, dtype=y.dtype)

def get_duration(y=None, sr=22050):
    """
    Compute the duration (in seconds) of an audio time series.
    Args:
        y: np.ndarray [shape=(n,), (2, n)] or None
            audio time series
        sr: number > 0 [scalar]
            audio sampling rate of `y`
    Returns:
        d: float >= 0
            Duration (in seconds) of the input time series.
    """
    valid_audio(y, mono=False)
    if y.ndim == 1:
        n_samples = len(y)
    else:
        n_samples = y.shape[-1]
    return float(n_samples) / sr

def fix_length(data, size, axis=-1, **kwargs):
    """
    Fix the length of array `data` to exactly `size`.
    If `data.shape[axis] < n`, pad according to kwargs.mode

    Args:
        data: np.ndarray
            array to be length-adjusted
        size: int >= 0 [scalar]
            desired length of the array
        axis: int, <= data.ndim
            axis along which to fix length
        kwargs:
            Parameters to np.pad()
    Returns:
        np.ndarray [shape=data.shape]
            `data` either trimmed or padded to length `size`
            along the specified axis.
    """
    kwargs.setdefault('mode', 'constant')
    n = data.shape[axis]
    if n > size:
        slices = [slice(None)] * data.ndim
        slices[axis] = slice(0, size)
        return data[slices]
    elif n < size:
        lengths = [(0, 0)] * data.ndim
        lengths[axis] = (0, size - n)
        return np.pad(data, lengths, **kwargs)
    return data

def buf_to_float(x, n_bytes=2, dtype=np.float32):
    """
    Convert an integer buffer to floating point values.
    Used for loading integer-valued wav data into numpy arrays.
    Args:
        x : np.ndarray [dtype=int]
            The integer-valued data buffer
        n_bytes : int [1, 2, 4]
            The number of bytes per sample in `x`
        dtype : numeric type
            The target output type (default: 32-bit float)
    Returns:
        x_float : np.ndarray [dtype=float]
            The input data buffer cast to floating point
    """
    # Invert the scale of the data
    scale = 1./float(1 << ((8 * n_bytes) - 1))
    # Construct the format string
    fmt = '<i{:d}'.format(n_bytes)
    # Rescale and format the data buffer
    return scale * np.frombuffer(x, fmt).astype(dtype)

def buf_to_int(x, n_bytes=2):
    """
    Convert a floating point buffer into integer values.
    This is primarily useful as an intermediate step in wav output.
    Args:
        x: np.ndarray [dtype=float]
            Floating point data buffer
        n_bytes: int [1, 2, 4]
            Number of bytes per output sample
    Returns:
        x_int : np.ndarray [dtype=int]
            The original buffer cast to integer type.
    """
    if n_bytes not in [1, 2, 4]:
        raise ValueError('n_bytes must be one of {1, 2, 4}')
    # What is the scale of the input data?
    scale = float(1 << ((8 * n_bytes) - 1))
    # Construct a format string
    fmt = '<i{:d}'.format(n_bytes)
    # Rescale and cast the data
    return (x * scale).astype(fmt)

def valid_audio(y, mono=True):
    """
    Validate whether a variable contains valid audio data.
        Args:
        y : np.ndarray
          The input data to validate
        mono : bool
      Wheth:
        valid : bool
            True if all tests pass
    Raises:
        ParameterError
            If `y` fails to meet the following criteria:
                - `type(y)` is `np.ndarray`
                - `mono == True` and `y.ndim` is not 1
                - `mono == False` and `y.ndim` is not 1 or 2
                - `np.isfinite(y).all()` is not True
    """
    if not isinstance(y, np.ndarray):
        raise ValueError('Data must be of type numpy.ndarray')
    if mono and y.ndim != 1:
        raise ValueError('Invalid shape for monophonic audio: '
                         'ndim={:d}, shape={}'.format(y.ndim, y.shape))
    elif y.ndim > 2:
        raise ValueError('Invalid shape for audio: '
                         'ndim={:d}, shape={}'.format(y.ndim, y.shape))
    if not np.isfinite(y).all():
        raise ValueError('Audio buffer is not finite everywhere')
    return True

def normalize(S, norm=np.inf, axis=0):
    """
    Normalize the columns or rows of a matrix
    Args:
        S : np.ndarray [shape=(d, n)]
            The matrix to normalize
        norm : {np.inf, -np.inf, 0, float > 0, None}
            - `np.inf`  : maximum absolute value
            - `-np.inf` : mininum absolute value
            - `0`    : number of non-zeros
            - float  : corresponding l_p norm.
                See `scipy.linalg.norm` for details.
            - None : no normalization is performed
        axis : int [scalar]
            Axis along which to compute the norm.
            `axis=0` will normalize columns, `axis=1` will normalize rows.
            `axis=None` will normalize according to the entire matrix.
    Returns:
        S_norm : np.ndarray [shape=S.shape]
            Normalized matrix
    """
    # All norms only depend on magnitude, let's do that first
    mag = np.abs(S)

    if norm == np.inf:
        length = np.max(mag, axis=axis, keepdims=True)

    elif norm == -np.inf:
        length = np.min(mag, axis=axis, keepdims=True)

    elif norm == 0:
        length = np.sum(mag > 0, axis=axis, keepdims=True)

    elif np.issubdtype(type(norm), np.number) and norm > 0:
        length = np.sum(mag ** norm, axis=axis, keepdims=True)**(1./norm)

    elif norm is None:
        return S

    else:
        raise ValueError('Unsupported norm: {}'.format(repr(norm)))

    # Avoid div-by-zero
    SMALL_FLOAT = 1e-20
    length[length < SMALL_FLOAT] = 1.0

    return S / length

def fractional_delay(y, time, sr):
    """
    Prepend zeros to delay signal by time seconds
    Use linear interpolation through convolution for fractional delay
    """
    samples = time * sr
    if isinstance(samples, int) or samples.is_integer():
        return delay(y, time, sr)

    f, i = math.modf(samples)
    y = delay(y, i, sr, mode='samples')
    import scipy.signal
    y = scipy.signal.convolve(y, [f, i-f], "same")
    return y

def delay(y, time, sr, mode='time'):
    """
    Prepend zeros to delay signal by time seconds
    """
    if mode == 'time':
        samples = time * sr
    else:
        samples = time
    # fractional delay by upsampling has been commented out for now
    #decimals = sum(c != '0' for c in str(round(samples % 1, 4))[2:])
    #new_sr = None
    #if decimals > 0:
    #    new_sr = sr * (decimals + 1)
    #    samples = time * new_sr
    #    y = resample(y, sr, new_sr)
    y = np.pad(
        y,
        pad_width=[round(samples), 0],
        mode='constant',
        constant_values=0
    )
    #if new_sr:
    #    y = resample(y, new_sr, sr)
    return y

def sum_signals(signals):
    """
    Sum together a list of mono signals
    append zeros to match the longest array
    """
    if not signals:
        return np.array([])
    max_length = max(len(sig) for sig in signals)
    y = np.zeros(max_length)
    for sig in signals:
        padded = np.zeros(max_length)
        padded[0:len(sig)] = sig
        y += padded
    return y

def channel_merge(channels):
    """
    vstack a list of mono signals
    append zeros to match the longest signal
    """
    longest = max(len(channel) for channel in channels)
    padded_channels = []
    for i, channel in enumerate(channels):
        if len(channel) < longest:
            padded = np.zeros(longest)
            padded[0:len(channel)] = channel
            channel = padded
        padded_channels.append(channel)
    y = np.vstack(padded_channels)
    return y
<|endoftext|>


<commit_before>#
#  DIALS viewer_frame
#
#  Copyright (C) 2014 Diamond Light Source
#
#  Author: Luis Fuentes-Montero (Luiso)
#
#  This code is distributed under the BSD license, a copy of which is
#  included in the root directory of this package."

import wx
from dials.viewer.viewer_utilities \
     import np_to_bmp

from dials.viewer.reflection_data_navigator import table_s_navigator

class ReflectionFrame(wx.Frame):
  def __init__(self, *args, **kwargs):
    wx.Frame.__init__(self, *args, **kwargs)

    self.MaxImageSizeX = 280
    self.MaxImageSizeY = 150

    nxt_ref = "Reflection -->"
    btn_nxt_refl = wx.Button(self, -1, nxt_ref)
    prv_ref = "<-- Reflection"
    btn_prv_refl = wx.Button(self, -1, prv_ref)
    btn_read = wx.Button(self, -1, "Reflection #")

    radio1 = wx.RadioButton(self, -1, "data, background, mask"
                           , style = wx.RB_GROUP)
    radio2 = wx.RadioButton(self, -1, "3 layers of data")
    radio3 = wx.RadioButton(self, -1, "3 layers of background")
    radio4 = wx.RadioButton(self, -1, "3 layers of mask")

    btn_nxt_slice = wx.Button(self, -1, "slice  -->")
    btn_prv_slice = wx.Button(self, -1, "<--  slice")

    self.data_txt_01 = wx.StaticText(self, -1, "(data_txt)", size = (800, 16))
    self.data_txt_02 = wx.StaticText(self, -1, "(data_txt)", size = (800, 16))
    self.data_txt_03 = wx.StaticText(self, -1, "(data_txt)", size = (800, 16))

    self.Image = [None, None, None]

    for indx in range(len(self.Image)):
      self.Image[indx] = wx.StaticBitmap(self, bitmap = wx.EmptyBitmap(
                                         self.MaxImageSizeX, self.MaxImageSizeY))

    self.text01 = wx.TextCtrl(self, -1, "enter number", size = (100, -1))

    u_box = wx.BoxSizer(wx.VERTICAL)
    div_h_box = wx.BoxSizer(wx.HORIZONTAL)
    div_h_box.Add(btn_prv_refl, 0, wx.CENTER | wx.ALL,3)
    div_h_box.Add(btn_nxt_refl, 0, wx.CENTER | wx.ALL,3)
    u_box.Add(div_h_box)

    div_h_box = wx.BoxSizer(wx.HORIZONTAL)
    div_h_box.Add(btn_read, 0, wx.CENTER | wx.ALL,3)
    div_h_box.Add(self.text01, 0, wx.CENTER | wx.ALL,3)
    u_box.Add(div_h_box)

    u_box.Add(radio1, 0, wx.LEFT | wx.ALL,3)
    u_box.Add(radio2, 0, wx.LEFT | wx.ALL,3)
    u_box.Add(radio3, 0, wx.LEFT | wx.ALL,3)
    u_box.Add(radio4, 0, wx.LEFT | wx.ALL,3)

    div_h_box = wx.BoxSizer(wx.HORIZONTAL)
    div_h_box.Add(btn_prv_slice, 0, wx.CENTER | wx.ALL,3)
    div_h_box.Add(btn_nxt_slice, 0, wx.CENTER | wx.ALL,3)
    u_box.Add(div_h_box)

    h_box = wx.BoxSizer(wx.HORIZONTAL)

    h_box.Add(u_box)

    div_h_box = wx.BoxSizer(wx.HORIZONTAL)
    div_h_box.Add(self.Image[0], 0
            , wx.ALIGN_CENTER_HORIZONTAL | wx.ALL | wx.ADJUST_MINSIZE, 5)

    div_h_box.Add(self.Image[1], 0
            , wx.ALIGN_CENTER_HORIZONTAL | wx.ALL | wx.ADJUST_MINSIZE, 5)

    div_h_box.Add(self.Image[2], 0
            , wx.ALIGN_CENTER_HORIZONTAL | wx.ALL | wx.ADJUST_MINSIZE, 5)
    div_u_box = wx.BoxSizer(wx.VERTICAL)
    div_u_box.Add(div_h_box)
    div_u_box.Add(self.data_txt_01, 0, wx.CENTER | wx.ALL,3)
    div_u_box.Add(self.data_txt_02, 0, wx.CENTER | wx.ALL,3)
    div_u_box.Add(self.data_txt_03, 0, wx.CENTER | wx.ALL,3)
    h_box.Add(div_u_box)

    self.frame_scale = 0.4
    self.opt = 0

    self.sizing_counter = 0
    self.SetSizerAndFit(h_box)

    btn_nxt_refl.Bind(wx.EVT_BUTTON, self._DisplayNext_refl)
    btn_prv_refl.Bind(wx.EVT_BUTTON, self._DisplayPrev_refl)
    btn_nxt_slice.Bind(wx.EVT_BUTTON, self._DisplayNext_slice)
    btn_prv_slice.Bind(wx.EVT_BUTTON, self._DisplayPrev_slice)

    self.Bind(wx.EVT_RADIOBUTTON, self._OnRadio1, radio1)
    self.Bind(wx.EVT_RADIOBUTTON, self._OnRadio2, radio2)
    self.Bind(wx.EVT_RADIOBUTTON, self._OnRadio3, radio3)
    self.Bind(wx.EVT_RADIOBUTTON, self._OnRadio4, radio4)

    self.Bind(wx.EVT_BUTTON, self._read_num_from_txt, btn_read)

    self.Bind(wx.EVT_SIZE, self._OnSize)

    self.bmp = np_to_bmp()
    self.arr_img = [None, None, None]
    self.arr_ttl = [None, None, None]
    self.My_Img = [None, None, None]
    self.wx_Img = [None, None, None]
    wx.EVT_CLOSE(self, self._On_Close_Window)

  def tabl_to_frame(self, loc_tabl):
    self.tabl = table_s_navigator(loc_tabl)
    self._DisplayPrev_refl()
  def _DisplayNext_refl(self, event = None):
    self.tabl.next_Reflection()
    self._My_Update()
  def _DisplayPrev_refl(self, event = None):
    self.tabl.Previous_Reflection()
    self._My_Update()
  def _DisplayNext_slice(self, event = None):
    self.tabl.next_slice()
    self._My_Update()
  def _DisplayPrev_slice(self, event = None):
    self.tabl.Previous_slice()
    self._My_Update()

  def _read_num_from_txt(self, event = None):
    # TODO check if the user entered a number
    a = int(self.text01.GetValue())
    self.tabl.Jump_to_Reflection(a)
    self._My_Update()

  def _OnRadio1(self, event = None):
    self.opt = 0
    self._My_Update(request_new_data = True)

  def _OnRadio2(self, event = None):
    self.opt = 1
    self._My_Update(request_new_data = True)

  def _OnRadio3(self, event = None):
    self.opt = 2
    self._My_Update(request_new_data = True)

  def _OnRadio4(self, event = None):
    self.opt = 3
    self._My_Update(request_new_data = True)

  def _OnSize(self, event = None):
    if( self.sizing_counter > 5 ):
      siz_data = self.GetSize()
      #print "New size of window =", siz_data
      optm_aspec_ratio = 4.21
      #print "siz_data = ", siz_data[0], siz_data[1]
      #print "aspect ratio = ", float(siz_data[0])/ float(siz_data[1])
      aspec_ratio = float(siz_data[0])/ float(siz_data[1])

      if(aspec_ratio > optm_aspec_ratio):
        #print "use float(siz_data[1] (Height) to calculate new size"
        self.frame_scale = float(siz_data[1]) * 0.5 / 295.0
        #(1227, 295)
      else:
        #print "use float(siz_data[0] (with) to calculate new size"
        self.frame_scale = float(siz_data[0]) * 0.5 / 1227.0
        #(1227, 295)
      self._My_Update(request_new_data = False)
      #print "resizing"
      #print "aspec_ratio =", aspec_ratio
    else:
      self.sizing_counter += 1

  def _My_Update(self, request_new_data = True):

    if( request_new_data == True ):
      self.arr_img, self.arr_ttl = self.tabl(opt = self.opt)
      ref_max = self.tabl.Get_Max(self.opt)
      #box_lmt = self.tabl.Get_bbox()
      self.xyz_px = self.tabl.Get_xyz()
      self.bbox_px = self.tabl.Get_bbox()
      self.hkl_data = self.tabl.Get_hkl()
      for indx in range(len(self.arr_img)):
        if( self.opt == 0 and indx == len(self.arr_img) - 1 ):
          Imax = 10
        else:
          Imax = ref_max
        self.wx_Img[indx] = self.bmp(np_img_2d = self.arr_img[indx]
                          , Intst_max = Imax, ofst = self.bbox_px
                          , xyz = self.xyz_px, title = self.arr_ttl[indx])
      self.text01.SetValue(str(self.tabl.Get_ref_num()))

    for indx in range(len(self.arr_img)):
      if( self.arr_img[indx] == None ):
        empty_flag = True
      else:
        empty_flag = False
      self.My_Img[indx] = self.bmp.from_wx_image_to_wx_bitmap(self.wx_Img[indx]
                        , self.frame_scale, empty = empty_flag)
      self.Image[indx].SetBitmap(self.My_Img[indx])
      if( self.xyz_px != None ):
        my_str = ( " Centroid Position = ( "
                 + str(self.xyz_px[0]) + ", "
                 + str(self.xyz_px[1]) + ", "
                 + str(self.xyz_px[2]) + ") ")
        self.data_txt_01.SetLabel(my_str)
      else:
        self.data_txt_01.SetLabel(" No (x, y, z) Data")

      if( self.bbox_px != None ):
        my_str_01 = ( " Bbox = ( "
                 + str(self.bbox_px[0]) + ", "
                 + str(self.bbox_px[1]) + ", "
                 + str(self.bbox_px[2]) + ", "
                 + str(self.bbox_px[3]) + ", "
                 + str(self.bbox_px[4]) + ", "
                 + str(self.bbox_px[5]) + ") ")
        self.data_txt_02.SetLabel(my_str_01)
      else:
        self.data_txt_02.SetLabel(" No Bbox Data")

      try:

        hkl_string = 'HKL:    ' + str(self.hkl_data[0]) + \
                            '    ' + str(self.hkl_data[1]) + \
                            '    ' + str(self.hkl_data[2]) + ' '
      except:
        hkl_string = ' No HKL data'
      self.data_txt_03.SetLabel(hkl_string)

    self.Layout()
    self.Refresh()

  def _On_Close_Window(self, event):
    self.Destroy()
<commit_msg>more stable (maybe)<commit_after>#
#  DIALS viewer_frame
#
#  Copyright (C) 2014 Diamond Light Source
#
#  Author: Luis Fuentes-Montero (Luiso)
#
#  This code is distributed under the BSD license, a copy of which is
#  included in the root directory of this package."

import wx
from dials.viewer.viewer_utilities \
     import np_to_bmp

from dials.viewer.reflection_data_navigator import table_s_navigator

class ReflectionFrame(wx.Frame):
  def __init__(self, *args, **kwargs):
    wx.Frame.__init__(self, *args, **kwargs)

    #self.MaxImageSizeX = 280
    #self.MaxImageSizeY = 150
    self.MaxImageSizeX = 14
    self.MaxImageSizeY = 7

    nxt_ref = "Reflection -->"
    btn_nxt_refl = wx.Button(self, -1, nxt_ref)
    prv_ref = "<-- Reflection"
    btn_prv_refl = wx.Button(self, -1, prv_ref)
    btn_read = wx.Button(self, -1, "Reflection #")

    radio1 = wx.RadioButton(self, -1, "data, background, mask"
                           , style = wx.RB_GROUP)
    radio2 = wx.RadioButton(self, -1, "3 layers of data")
    radio3 = wx.RadioButton(self, -1, "3 layers of background")
    radio4 = wx.RadioButton(self, -1, "3 layers of mask")

    btn_nxt_slice = wx.Button(self, -1, "slice  -->")
    btn_prv_slice = wx.Button(self, -1, "<--  slice")

    self.data_txt_01 = wx.StaticText(self, -1, "(data_txt)", size = (800, 16))
    self.data_txt_02 = wx.StaticText(self, -1, "(data_txt)", size = (800, 16))
    self.data_txt_03 = wx.StaticText(self, -1, "(data_txt)", size = (800, 16))

    self.Image = [None, None, None]

    for indx in range(len(self.Image)):
      self.Image[indx] = wx.StaticBitmap(self, bitmap = wx.EmptyBitmap(
                                         self.MaxImageSizeX, self.MaxImageSizeY))

    self.text01 = wx.TextCtrl(self, -1, "enter number", size = (100, -1))

    u_box = wx.BoxSizer(wx.VERTICAL)
    div_h_box = wx.BoxSizer(wx.HORIZONTAL)
    div_h_box.Add(btn_prv_refl, 0, wx.CENTER | wx.ALL,3)
    div_h_box.Add(btn_nxt_refl, 0, wx.CENTER | wx.ALL,3)
    u_box.Add(div_h_box)

    div_h_box = wx.BoxSizer(wx.HORIZONTAL)
    div_h_box.Add(btn_read, 0, wx.CENTER | wx.ALL,3)
    div_h_box.Add(self.text01, 0, wx.CENTER | wx.ALL,3)
    u_box.Add(div_h_box)

    u_box.Add(radio1, 0, wx.LEFT | wx.ALL,3)
    u_box.Add(radio2, 0, wx.LEFT | wx.ALL,3)
    u_box.Add(radio3, 0, wx.LEFT | wx.ALL,3)
    u_box.Add(radio4, 0, wx.LEFT | wx.ALL,3)

    div_h_box = wx.BoxSizer(wx.HORIZONTAL)
    div_h_box.Add(btn_prv_slice, 0, wx.CENTER | wx.ALL,3)
    div_h_box.Add(btn_nxt_slice, 0, wx.CENTER | wx.ALL,3)
    u_box.Add(div_h_box)

    h_box = wx.BoxSizer(wx.HORIZONTAL)

    h_box.Add(u_box)

    div_h_box = wx.BoxSizer(wx.HORIZONTAL)
    div_h_box.Add(self.Image[0], 0
            , wx.ALIGN_CENTER_HORIZONTAL | wx.ALL | wx.ADJUST_MINSIZE, 5)

    div_h_box.Add(self.Image[1], 0
            , wx.ALIGN_CENTER_HORIZONTAL | wx.ALL | wx.ADJUST_MINSIZE, 5)

    div_h_box.Add(self.Image[2], 0
            , wx.ALIGN_CENTER_HORIZONTAL | wx.ALL | wx.ADJUST_MINSIZE, 5)
    div_u_box = wx.BoxSizer(wx.VERTICAL)
    div_u_box.Add(div_h_box)
    div_u_box.Add(self.data_txt_01, 0, wx.CENTER | wx.ALL,3)
    div_u_box.Add(self.data_txt_02, 0, wx.CENTER | wx.ALL,3)
    div_u_box.Add(self.data_txt_03, 0, wx.CENTER | wx.ALL,3)
    h_box.Add(div_u_box)

    self.frame_scale = 0.4
    self.opt = 0

    self.sizing_counter = 0
    self.SetSizerAndFit(h_box)

    btn_nxt_refl.Bind(wx.EVT_BUTTON, self._DisplayNext_refl)
    btn_prv_refl.Bind(wx.EVT_BUTTON, self._DisplayPrev_refl)
    btn_nxt_slice.Bind(wx.EVT_BUTTON, self._DisplayNext_slice)
    btn_prv_slice.Bind(wx.EVT_BUTTON, self._DisplayPrev_slice)

    self.Bind(wx.EVT_RADIOBUTTON, self._OnRadio1, radio1)
    self.Bind(wx.EVT_RADIOBUTTON, self._OnRadio2, radio2)
    self.Bind(wx.EVT_RADIOBUTTON, self._OnRadio3, radio3)
    self.Bind(wx.EVT_RADIOBUTTON, self._OnRadio4, radio4)

    self.Bind(wx.EVT_BUTTON, self._read_num_from_txt, btn_read)

    self.Bind(wx.EVT_SIZE, self._OnSize)

    self.bmp = np_to_bmp()
    self.arr_img = [None, None, None]
    self.arr_ttl = [None, None, None]
    self.My_Img = [None, None, None]
    self.wx_Img = [None, None, None]
    wx.EVT_CLOSE(self, self._On_Close_Window)

  def tabl_to_frame(self, loc_tabl):
    self.tabl = table_s_navigator(loc_tabl)
    self._DisplayPrev_refl()
  def _DisplayNext_refl(self, event = None):
    self.tabl.next_Reflection()
    self._My_Update()
  def _DisplayPrev_refl(self, event = None):
    self.tabl.Previous_Reflection()
    self._My_Update()
  def _DisplayNext_slice(self, event = None):
    self.tabl.next_slice()
    self._My_Update()
  def _DisplayPrev_slice(self, event = None):
    self.tabl.Previous_slice()
    self._My_Update()

  def _read_num_from_txt(self, event = None):
    # TODO check if the user entered a number
    a = int(self.text01.GetValue())
    self.tabl.Jump_to_Reflection(a)
    self._My_Update()

  def _OnRadio1(self, event = None):
    self.opt = 0
    self._My_Update(request_new_data = True)

  def _OnRadio2(self, event = None):
    self.opt = 1
    self._My_Update(request_new_data = True)

  def _OnRadio3(self, event = None):
    self.opt = 2
    self._My_Update(request_new_data = True)

  def _OnRadio4(self, event = None):
    self.opt = 3
    self._My_Update(request_new_data = True)

  def _OnSize(self, event = None):
    if( self.sizing_counter > 5 ):
      siz_data = self.GetSize()
      #print "New size of window =", siz_data
      optm_aspec_ratio = 4.21
      #print "siz_data = ", siz_data[0], siz_data[1]
      #print "aspect ratio = ", float(siz_data[0])/ float(siz_data[1])
      aspec_ratio = float(siz_data[0])/ float(siz_data[1])

      if(aspec_ratio > optm_aspec_ratio):
        #print "use float(siz_data[1] (Height) to calculate new size"
        self.frame_scale = float(siz_data[1]) * 0.5 / 295.0
        #(1227, 295)
      else:
        #print "use float(siz_data[0] (with) to calculate new size"
        self.frame_scale = float(siz_data[0]) * 0.5 / 1227.0
        #(1227, 295)
      self._My_Update(request_new_data = False)
      #print "resizing"
      #print "aspec_ratio =", aspec_ratio
    else:
      self.sizing_counter += 1

  def _My_Update(self, request_new_data = True):

    if( request_new_data == True ):
      self.arr_img, self.arr_ttl = self.tabl(opt = self.opt)
      ref_max = self.tabl.Get_Max(self.opt)
      #box_lmt = self.tabl.Get_bbox()
      self.xyz_px = self.tabl.Get_xyz()
      self.bbox_px = self.tabl.Get_bbox()
      self.hkl_data = self.tabl.Get_hkl()
      for indx in range(len(self.arr_img)):
        if( self.opt == 0 and indx == len(self.arr_img) - 1 ):
          Imax = 10
        else:
          Imax = ref_max
        self.wx_Img[indx] = self.bmp(np_img_2d = self.arr_img[indx]
                          , Intst_max = Imax, ofst = self.bbox_px
                          , xyz = self.xyz_px, title = self.arr_ttl[indx])
      self.text01.SetValue(str(self.tabl.Get_ref_num()))

    for indx in range(len(self.arr_img)):
      if( self.arr_img[indx] == None ):
        empty_flag = True
      else:
        empty_flag = False
      self.My_Img[indx] = self.bmp.from_wx_image_to_wx_bitmap(self.wx_Img[indx]
                        , self.frame_scale, empty = empty_flag)
      self.Image[indx].SetBitmap(self.My_Img[indx])
      if( self.xyz_px != None ):
        my_str = ( " Centroid Position = ( "
                 + str(self.xyz_px[0]) + ", "
                 + str(self.xyz_px[1]) + ", "
                 + str(self.xyz_px[2]) + ") ")
        self.data_txt_01.SetLabel(my_str)
      else:
        self.data_txt_01.SetLabel(" No (x, y, z) Data")

      if( self.bbox_px != None ):
        my_str_01 = ( " Bbox = ( "
                 + str(self.bbox_px[0]) + ", "
                 + str(self.bbox_px[1]) + ", "
                 + str(self.bbox_px[2]) + ", "
                 + str(self.bbox_px[3]) + ", "
                 + str(self.bbox_px[4]) + ", "
                 + str(self.bbox_px[5]) + ") ")
        self.data_txt_02.SetLabel(my_str_01)
      else:
        self.data_txt_02.SetLabel(" No Bbox Data")

      try:

        hkl_string = 'HKL:    ' + str(self.hkl_data[0]) + \
                            '    ' + str(self.hkl_data[1]) + \
                            '    ' + str(self.hkl_data[2]) + ' '
      except:
        hkl_string = ' No HKL data'
      self.data_txt_03.SetLabel(hkl_string)

    self.Layout()
    self.Refresh()

  def _On_Close_Window(self, event):
    self.Destroy()
<|endoftext|>


<commit_before>import os

from charmhelpers.core.hookenv import (
    config,
    open_port,
    status_set,
)

from charmhelpers.core.host import (
    chownr,
    service_restart,
)

import charmhelpers.core.host as ch_host

from charmhelpers.core.templating import render

from charmhelpers.fetch import (
    apt_install,
    install_remote,
)

from charms.layer import nginx

from charms.reactive import (
    remove_state,
    set_state,
    when,
    when_not,
)

def base_tracker_releases():
    if config('base-tracker-releases'):
        return config('base-tracker-releases').split()
    return None

def upstream_tracker_releases():
    if config('upstream-tracker-releases'):
        return config('upstream-tracker-releases').split()
    return None

template_map = {
    'generate-uca-reports': {
        'target': '/usr/local/bin/generate-uca-reports',
        'owner': 'root',
        'group': 'root',
        'perms': 0o755,
        'context': {
            'base_releases': base_tracker_releases(),
            'upstream_releases': upstream_tracker_releases(),
        },
    },
    'index.html': {
        'target': '/usr/share/nginx/www/index.html',
        'owner': 'www-data',
        'group': 'www-data',
        'perms': 0o644,
        'context': {
            'base_releases': base_tracker_releases(),
            'upstream_releases': upstream_tracker_releases(),
        },
    },
    'uca-tracker.service': {
        'target': '/etc/systemd/system/uca-tracker.service',
        'owner': 'root',
        'group': 'root',
        'perms': 0o644,
        'context': {},
    },
}

@when_not('uca-tracker.installed')
def install_uca_tracker():
    apt_install(['python3-git', 'python3-yaml', 'python3-httplib2',
                 'python3-launchpadlib', 'python3-swift'], fatal=True)

    install_remote('lp:ubuntu-reports', dest='/opt')

    for template, target_info in template_map.items():
        render(source=template,
               target=target_info['target'],
               owner=target_info['owner'],
               group=target_info['group'],
               perms=target_info['perms'],
               context=target_info['context'])

    if not os.path.exists('/var/log/uca-tracker'):
        os.mkdir('/var/log/uca-tracker')
    chownr('/usr/share/nginx/www', 'www-data', 'www-data')
    set_state('uca-tracker.installed')

@when('config.changed')
def update_config():
    remove_state('uca-tracker.installed')
    remove_state('uca-tracker.ready')

@when('nginx.available')
@when('uca-tracker.installed')
@when_not('uca-tracker.nginx.ready')
def setup_vhost():
    status_set('maintenance', 'setting up vhost')

    nginx.configure_site(
        'uca-tracker',
        'vhost.conf',
        app_path='/usr/share/nginx/www',
    )

    ch_host.service_restart('nginx')
    set_state('uca-tracker.nginx.ready')

@when('uca-tracker.nginx.ready')
@when_not('uca-tracker.ready')
def tracker_ready():
    open_port(80)

    service_restart('uca-tracker')
    status_set('active', 'ready')
    set_state('uca-tracker.ready')
<commit_msg>Switch to python3-swiftclient dependency<commit_after>import os

from charmhelpers.core.hookenv import (
    config,
    open_port,
    status_set,
)

from charmhelpers.core.host import (
    chownr,
    service_restart,
)

import charmhelpers.core.host as ch_host

from charmhelpers.core.templating import render

from charmhelpers.fetch import (
    apt_install,
    install_remote,
)

from charms.layer import nginx

from charms.reactive import (
    remove_state,
    set_state,
    when,
    when_not,
)

def base_tracker_releases():
    if config('base-tracker-releases'):
        return config('base-tracker-releases').split()
    return None

def upstream_tracker_releases():
    if config('upstream-tracker-releases'):
        return config('upstream-tracker-releases').split()
    return None

template_map = {
    'generate-uca-reports': {
        'target': '/usr/local/bin/generate-uca-reports',
        'owner': 'root',
        'group': 'root',
        'perms': 0o755,
        'context': {
            'base_releases': base_tracker_releases(),
            'upstream_releases': upstream_tracker_releases(),
        },
    },
    'index.html': {
        'target': '/usr/share/nginx/www/index.html',
        'owner': 'www-data',
        'group': 'www-data',
        'perms': 0o644,
        'context': {
            'base_releases': base_tracker_releases(),
            'upstream_releases': upstream_tracker_releases(),
        },
    },
    'uca-tracker.service': {
        'target': '/etc/systemd/system/uca-tracker.service',
        'owner': 'root',
        'group': 'root',
        'perms': 0o644,
        'context': {},
    },
}

@when_not('uca-tracker.installed')
def install_uca_tracker():
    apt_install(['python3-git', 'python3-yaml', 'python3-httplib2',
                 'python3-launchpadlib', 'python3-swiftclient'], fatal=True)

    install_remote('lp:ubuntu-reports', dest='/opt')

    for template, target_info in template_map.items():
        render(source=template,
               target=target_info['target'],
               owner=target_info['owner'],
               group=target_info['group'],
               perms=target_info['perms'],
               context=target_info['context'])

    if not os.path.exists('/var/log/uca-tracker'):
        os.mkdir('/var/log/uca-tracker')
    chownr('/usr/share/nginx/www', 'www-data', 'www-data')
    set_state('uca-tracker.installed')

@when('config.changed')
def update_config():
    remove_state('uca-tracker.installed')
    remove_state('uca-tracker.ready')

@when('nginx.available')
@when('uca-tracker.installed')
@when_not('uca-tracker.nginx.ready')
def setup_vhost():
    status_set('maintenance', 'setting up vhost')

    nginx.configure_site(
        'uca-tracker',
        'vhost.conf',
        app_path='/usr/share/nginx/www',
    )

    ch_host.service_restart('nginx')
    set_state('uca-tracker.nginx.ready')

@when('uca-tracker.nginx.ready')
@when_not('uca-tracker.ready')
def tracker_ready():
    open_port(80)

    service_restart('uca-tracker')
    status_set('active', 'ready')
    set_state('uca-tracker.ready')
<|endoftext|>


<commit_before>from __future__ import print_function, division, absolute_import #, unicode_literals # not casa compatible
from builtins import bytes, dict, object, range, map, input#, str # not casa compatible
from future.utils import itervalues, viewitems, iteritems, listvalues, listitems
from io import open

import pickle
from evla_mcast.controller import Controller
import rfpipe

import logging
ch = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s %(levelname)8s %(name)s | %(message)s')
ch.setFormatter(formatter)
logger = logging.getLogger('realfast_controllers')
logger.setLevel(logging.INFO)

vys_cfile = '/home/cbe-master/realfast/soft/vysmaw_apps/vys.conf'
default_preffile = '/lustre/evla/test/realfast/realfast.yml'
default_vys_timeout = 10  # seconds more than segment length
distributed_host = 'cbe-node-01'

class realfast_controller(Controller):

    def __init__(self, preffile=default_preffile, inprefs={},
                 vys_timeout=default_vys_timeout):
        """ Creates controller object that can act on a scan configuration.
        Inherits a "run" method that starts asynchronous operation.
        """

        super(realfast_controller, self).__init__()
        self.preffile = preffile
        self.inprefs = inprefs
        self.vys_timeout = vys_timeout
        self.jobs = []

    def handle_config(self, config):
        """ Triggered when obs comes in.
        Downstream logic starts here.
        """

        self.summarize(config)

        if self.runsearch(config):
            logger.info('Config looks good. Generating rfpipe state...')
            st = rfpipe.state.State(config=config, preffile=self.preffile,
                                    inprefs=self.inprefs)
            logger.info('Starting pipeline...')
            jobs = rfpipe.pipeline.pipeline_scan_distributed(st, segments=[0],
                                                             host=distributed_host,
                                                             cfile=vys_cfile,
                                                             vys_timeout=self.vys_timeout)
            self.jobs += jobs
        else:
            logger.info("Config not suitable for realfast. Skipping.")

        removed = 0
        for job in self.jobs:
            if job.status == 'finished':
                _ = self.jobs.remove(job)
                removed += 1
        if removed:
            logger.info('Removed {0} finished jobs from job queue.'.format(removed))

    def handle_finish(self, dataset):
        """ Triggered when obs doc defines end of a script.
        """

        logger.info('End of scheduling block message received.')

    def runsearch(self, config):
        """ Test whether configuration specifies a config that realfast should search
        """

        # find config properties of interest
        intent = config.scan_intent
        antennas = config.get_antennas()
        antnames = [str(ant.name) for ant in antennas]
        subbands = config.get_subbands()
        inttimes = [subband.hw_time_res for subband in subbands]
        pols = [subband.pp for subband in subbands]
        nchans = [subband.spectralChannels for subband in subbands]
        chansizes = [subband.bw/subband.spectralChannels for subband in subbands]
        reffreqs = [subband.sky_center_freq*1e6 for subband in subbands]

        # Do not process if...
        # 1) chansize changes between subbands
        if not all([chansizes[0] == chansize for chansize in chansizes]):
            logger.warn("Channel size changes between subbands: {0}"
                        .format(chansizes))
            return False

        return True

    def summarize(self, config):
        """ Print summary info for config
        """

        try:
            logger.info(':: ConfigID {0} ::'.format(config.configId))
            logger.info('\tScan {0}, source {1}, intent {2}'
                        .format(config.scanNo, config.source,
                                config.scan_intent))
            logger.info('\t(RA, Dec) = ({0}, {1})'
                        .format(config.ra_deg, config.dec_deg))
            sbs = config.get_subbands()
            freqs = [sb.sky_center_freq for sb in sbs]
            logger.info('\tFreq: {0} - {1}'.format(min(freqs), max(freqs)))
            nchan = sum([sb.spectralChannels for sb in sbs])
            sb0 = sbs[0]
            logger.info('\t(nspw, chan/spw, nchan) = ({0}, {1}, {2})'
                        .format(len(sbs), sb0.spectralChannels, nchan))
            if not all([sb0.spectralChannels == sb.spectralChannels
                        for sb in sbs]):
                logger.info('\tNot all spw have same number of channels.')
                logger.info('\t(nant, npol) = ({0}, {1})'
                            .format(config.numAntenna, config.npol))
                logger.info('\tStartMJD {0}, duration {1} s. '
                            'HW/Final time resolution = {2}/{3} s'
                            .format(config.startTime,
                                    24*3600*(config.stopTime-config.startTime),
                                    sb0.hw_time_res, sb0.final_time_res))
        except:
            logger.warn("Failed to fully parse config to print summary."
                        "Proceeding.")

class config_controller(Controller):

    def __init__(self, pklfile):
        """ Creates controller object that saves scan configs.
        Inherits a "run" method that starts asynchronous operation.
        """

        super(config_controller, self).__init__()
        self.pklfile = pklfile

    def handle_config(self, config):
        """ Triggered when obs comes in.
        Downstream logic starts here.
        """

        logger.info('Received complete configuration for {0}, '
                    'scan {1}, source {2}, intent {3}'
                    .format(config.scanId, config.scanNo, config.source,
                            config.scan_intent))

        with open(self.pklfile, 'ab') as pkl:
            pickle.dump(config, pkl)
<commit_msg>playing with logging, but most loggers seem to be disabled<commit_after>from __future__ import print_function, division, absolute_import #, unicode_literals # not casa compatible
from builtins import bytes, dict, object, range, map, input#, str # not casa compatible
from future.utils import itervalues, viewitems, iteritems, listvalues, listitems
from io import open

import pickle
from evla_mcast.controller import Controller
import rfpipe

import logging
ch = logging.StreamHandler()
ch.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s %(levelname)8s %(name)s | %(message)s')
ch.setFormatter(formatter)
logging.getLogger().addHandler(ch)
logging.getLogger().setLevel(logging.INFO)
logger = logging.getLogger('realfast_controller')

vys_cfile = '/home/cbe-master/realfast/soft/vysmaw_apps/vys.conf'
default_preffile = '/lustre/evla/test/realfast/realfast.yml'
default_vys_timeout = 10  # seconds more than segment length
distributed_host = 'cbe-node-01'

class realfast_controller(Controller):

    def __init__(self, preffile=default_preffile, inprefs={},
                 vys_timeout=default_vys_timeout):
        """ Creates controller object that can act on a scan configuration.
        Inherits a "run" method that starts asynchronous operation.
        """

        super(realfast_controller, self).__init__()
        self.preffile = preffile
        self.inprefs = inprefs
        self.vys_timeout = vys_timeout
        self.jobs = []

    def handle_config(self, config):
        """ Triggered when obs comes in.
        Downstream logic starts here.
        """

        self.summarize(config)

        if self.runsearch(config):
            logger.info('Config looks good. Generating rfpipe state...')
            st = rfpipe.state.State(config=config, preffile=self.preffile,
                                    inprefs=self.inprefs)
            logger.info('Starting pipeline...')
            jobs = rfpipe.pipeline.pipeline_scan_distributed(st, segments=[0],
                                                             host=distributed_host,
                                                             cfile=vys_cfile,
                                                             vys_timeout=self.vys_timeout)
            self.jobs += jobs
        else:
            logger.info("Config not suitable for realfast. Skipping.")

        removed = 0
        for job in self.jobs:
            if job.status == 'finished':
                _ = self.jobs.remove(job)
                removed += 1
        if removed:
            logger.info('Removed {0} finished jobs from job queue.'.format(removed))

    def handle_finish(self, dataset):
        """ Triggered when obs doc defines end of a script.
        """

        logger.info('End of scheduling block message received.')

    def runsearch(self, config):
        """ Test whether configuration specifies a config that realfast should search
        """

        # find config properties of interest
        intent = config.scan_intent
        antennas = config.get_antennas()
        antnames = [str(ant.name) for ant in antennas]
        subbands = config.get_subbands()
        inttimes = [subband.hw_time_res for subband in subbands]
        pols = [subband.pp for subband in subbands]
        nchans = [subband.spectralChannels for subband in subbands]
        chansizes = [subband.bw/subband.spectralChannels for subband in subbands]
        reffreqs = [subband.sky_center_freq*1e6 for subband in subbands]

        # Do not process if...
        # 1) chansize changes between subbands
        if not all([chansizes[0] == chansize for chansize in chansizes]):
            logger.warn("Channel size changes between subbands: {0}"
                        .format(chansizes))
            return False

        return True

    def summarize(self, config):
        """ Print summary info for config
        """

        try:
            logger.info(':: ConfigID {0} ::'.format(config.configId))
            logger.info('\tScan {0}, source {1}, intent {2}'
                        .format(config.scanNo, config.source,
                                config.scan_intent))
            logger.info('\t(RA, Dec) = ({0}, {1})'
                        .format(config.ra_deg, config.dec_deg))
            sbs = config.get_subbands()
            freqs = [sb.sky_center_freq for sb in sbs]
            logger.info('\tFreq: {0} - {1}'.format(min(freqs), max(freqs)))
            nchan = sum([sb.spectralChannels for sb in sbs])
            sb0 = sbs[0]
            logger.info('\t(nspw, chan/spw, nchan) = ({0}, {1}, {2})'
                        .format(len(sbs), sb0.spectralChannels, nchan))
            if not all([sb0.spectralChannels == sb.spectralChannels
                        for sb in sbs]):
                logger.info('\tNot all spw have same number of channels.')
                logger.info('\t(nant, npol) = ({0}, {1})'
                            .format(config.numAntenna, config.npol))
                logger.info('\tStartMJD {0}, duration {1} s. '
                            'HW/Final time resolution = {2}/{3} s'
                            .format(config.startTime,
                                    24*3600*(config.stopTime-config.startTime),
                                    sb0.hw_time_res, sb0.final_time_res))
        except:
            logger.warn("Failed to fully parse config to print summary."
                        "Proceeding.")

class config_controller(Controller):

    def __init__(self, pklfile):
        """ Creates controller object that saves scan configs.
        Inherits a "run" method that starts asynchronous operation.
        """

        super(config_controller, self).__init__()
        self.pklfile = pklfile

    def handle_config(self, config):
        """ Triggered when obs comes in.
        Downstream logic starts here.
        """

        logger.info('Received complete configuration for {0}, '
                    'scan {1}, source {2}, intent {3}'
                    .format(config.scanId, config.scanNo, config.source,
                            config.scan_intent))

        with open(self.pklfile, 'ab') as pkl:
            pickle.dump(config, pkl)
<|endoftext|>


<commit_before>"""
API described here: http://apfmon.lancs.ac.uk/help
"""

import requests
import json
import time
import traceback

from pandaharvester.harvesterconfig import harvester_config
from pandaharvester.harvestercore import core_utils
from pandaharvester import panda_pkg_info
from pandaharvester.harvestermisc import generic_utils
from pandaharvester.harvestercore.work_spec import WorkSpec
from pandaharvester.harvestermisc.info_utils import PandaQueuesDict

_base_logger = core_utils.setup_logger('apfmon')
NO_CE = 'noCE'

def apfmon_active(method, *args, **kwargs):
    if cls.__active:
        method(*args, **kwargs)
    else:
        return

def clean_ce(ce):
    return ce.split('.')[0].split('://')[-1]

class Apfmon(object):

    def __init__(self, queue_config_mapper):

        try:
            self.__active = harvester_config.apfmon.active
        except:
            self.__active = False

        try:
            self.__worker_timeout = harvester_config.apfmon.worker_timeout
        except:
            self.__worker_timeout = 0.2

        try:
            self.__label_timeout = harvester_config.apfmon.worker_timeout
        except:
            self.__label_timeout = 1

        # TODO: make proper exception handling and defaults
        try:
            self.harvester_id = harvester_config.master.harvester_id
        except:
            self.harvester_id = 'DUMMY'

        try:
            self.base_url = harvester_config.apfmon.base_url
        except:
            self.base_url = 'http://apfmon.lancs.ac.uk/api'

        self.queue_config_mapper = queue_config_mapper

    def create_factory(self):
        """
        Creates or updates a harvester instance to APF Mon. Should be done at startup of the instance.
        """

        start_time = time.time()
        tmp_log = core_utils.make_logger(_base_logger, 'harvester_id={0}'.format(self.harvester_id),
                                         method_name='create_factory')

        if not self.__active:
            tmp_log.debug('APFMon reporting not enabled')
            return

        try:
            tmp_log.debug('start')

            url = '{0}/factories/{1}'.format(self.base_url, self.harvester_id)

            f = {'url': 'url_to_logs',
                 'email': 'atlas-adc-harvester-central-support@cern.ch',
                 'version': panda_pkg_info.release_version}
            payload = json.dumps(f)

            r = requests.put(url, data=payload, timeout=self.__label_timeout)
            tmp_log.debug('registration ended with {0} {1}'.format(r.status_code, r.text))
            end_time = time.time()
            tmp_log.debug('done (took {0})'.format(end_time - start_time))
        except:
            tmp_log.error('Excepted with: {0}'.format(traceback.format_exc()))

    def create_labels(self):
        """
        Creates or updates a collection of labels (=panda queue+CE)
        """
        start_time = time.time()
        tmp_log = core_utils.make_logger(_base_logger, 'harvester_id={0}'.format(self.harvester_id),
                                         method_name='create_labels')

        if not self.__active:
            tmp_log.debug('APFMon reporting not enabled')
            return

        try:
            tmp_log.debug('start')

            url = '{0}/labels'.format(self.base_url)

            # get the active queues from the config mapper
            all_sites = self.queue_config_mapper.get_active_queues().keys()
            panda_queues_dict = PandaQueuesDict()

            # publish the active queues to APF mon in shards
            for sites in generic_utils.create_shards(all_sites, 20):
                labels = []
                for site in sites:
                    try:
                        site_info = panda_queues_dict.get(site, dict())
                        if not site_info:
                            tmp_log.warning('No site info for {0}'.format(site))
                            continue

                        if not site_info['queues']:
                            # when no CEs associated to a queue, e.g. P1, HPCs, etc. Try to see if there is something
                            # in local configuration, otherwise set it to a dummy value
                            try:
                                ce = self.queue_config_mapper.queueConfig[site].submitter['ceEndpoint']
                            except:
                                ce = NO_CE
                            queues = [{'ce_endpoint': ce}]
                        else:
                            queues = site_info['queues']

                        for queue in queues:
                            try:
                                ce = clean_ce(queue['ce_endpoint'])
                            except:
                                ce = ''

                            try:
                                ce_queue_id = queue['ce_queue_id']
                            except KeyError:
                                ce_queue_id = 0

                            labels.append({'name': '{0}-{1}'.format(site, ce),
                                           'wmsqueue': site,
                                           'ce_queue_id': ce_queue_id,
                                           'factory': self.harvester_id})
                    except:
                        tmp_log.error('Excepted for site {0} with: {1}'.format(site, traceback.format_exc()))
                        continue

                payload = json.dumps(labels)

                r = requests.put(url, data=payload, timeout=self.__label_timeout)
                tmp_log.debug('label creation for {0} ended with {1} {2}'.format(sites, r.status_code, r.text))

            end_time = time.time()
            tmp_log.debug('done (took {0})'.format(end_time - start_time))
        except:
            tmp_log.error('Excepted with: {0}'.format(traceback.format_exc()))

    def massage_label_data(self, data):

        tmp_log = core_utils.make_logger(_base_logger, 'harvester_id={0}'.format(self.harvester_id),
                                         method_name='massage_label_data')
        if not data:
            return data

        try:
            any = data['ANY']
            agg = {}
            for rtype in data:
                if rtype == 'ANY':
                    continue
                else:
                    for value in data[rtype]:
                        agg.setdefault(value, 0)
                        agg[value] += data[rtype][value]

            if agg:
                data['ANY'] = agg
            else:
                data['ANY'] = any

            tmp_log.debug('Massaged to data: {0}'.format(data))

        except Exception:
            tmp_log.debug('Exception in data: {0}'.format(data))

        return data

    def update_label(self, site, msg, data):
        """
        Updates a label (=panda queue+CE)
        """
        start_time = time.time()
        tmp_log = core_utils.make_logger(_base_logger, 'harvester_id={0}'.format(self.harvester_id),
                                         method_name='update_label')

        if not self.__active:
            tmp_log.debug('APFMon reporting not enabled')
            return

        try:
            tmp_log.debug('start')
            data = self.massage_label_data(data)

            # get the active queues from the config mapper
            all_sites = self.queue_config_mapper.get_active_queues().keys()
            panda_queues_dict = PandaQueuesDict()

            site_info = panda_queues_dict.get(site, dict())
            if not site_info:
                tmp_log.warning('No site info for {0}'.format(site))
                return

            if not site_info['queues']:
                # when no CEs associated to a queue, e.g. P1, HPCs, etc. Try to see if there is something
                # in local configuration, otherwise set it to a dummy value
                try:
                    ce = self.queue_config_mapper.queueConfig[site].submitter['ceEndpoint']
                except:
                    ce = NO_CE
                queues = [{'ce_endpoint': ce}]
            else:
                queues = site_info['queues']

            for queue in queues:
                try:
                    try:
                        ce = clean_ce(queue['ce_endpoint'])
                    except:
                        ce = ''

                    label_data = {'status': msg, 'data': data}
                    label = '{0}-{1}'.format(site, ce)
                    label_id = '{0}:{1}'.format(self.harvester_id, label)
                    url = '{0}/labels/{1}'.format(self.base_url, label_id)

                    r = requests.post(url, data=json.dumps(label_data), timeout=self.__label_timeout)
                    tmp_log.debug('label update for {0} ended with {1} {2}'.format(label, r.status_code, r.text))
                except:
                    tmp_log.error('Excepted for site {0} with: {1}'.format(label, traceback.format_exc()))

            end_time = time.time()
            tmp_log.debug('done (took {0})'.format(end_time - start_time))
        except:
            tmp_log.error('Excepted with: {0}'.format(traceback.format_exc()))

    def create_workers(self, worker_spec_list):
        """
        Creates a worker
        """
        start_time = time.time()
        tmp_log = core_utils.make_logger(_base_logger, 'harvester_id={0}'.format(self.harvester_id),
                                         method_name='create_workers')

        if not self.__active:
            tmp_log.debug('APFMon reporting not enabled')
            return

        try:
            tmp_log.debug('start')

            url = '{0}/jobs'.format(self.base_url)

            for worker_spec_shard in generic_utils.create_shards(worker_spec_list, 20):
                apfmon_workers = []
                for worker_spec in worker_spec_shard:
                    batch_id = worker_spec.batchID
                    factory = self.harvester_id
                    computingsite = worker_spec.computingSite
                    try:
                        ce = clean_ce(worker_spec.computingElement)
                    except AttributeError:
                        ce = ''

                    # extract the log URLs
                    stdout_url = ''
                    stderr_url = ''
                    log_url = ''
                    jdl_url = ''

                    work_attribs = worker_spec.workAttributes
                    if work_attribs:
                        if 'stdOut' in work_attribs:
                            stdout_url = work_attribs['stdOut']
                            jdl_url = '{0}.jdl'.format(stdout_url[:-4])
                        if 'stdErr' in work_attribs:
                            stderr_url = work_attribs['stdErr']
                        if 'batchLog' in work_attribs:
                            log_url = work_attribs['batchLog']

                    apfmon_worker = {'cid': batch_id,
                                     'factory': factory,
                                     'label': '{0}-{1}'.format(computingsite, ce),
                                     'jdlurl': jdl_url,
                                     'stdouturl': stdout_url,
                                     'stderrurl': stderr_url,
                                     'logurl': log_url
                                     }
                    tmp_log.debug('packed worker: {0}'.format(apfmon_worker))
                    apfmon_workers.append(apfmon_worker)

                payload = json.dumps(apfmon_workers)
                r = requests.put(url, data=payload, timeout=self.__worker_timeout)
                tmp_log.debug('worker creation for {0} ended with {1} {2}'.format(apfmon_workers, r.status_code, r.text))

            end_time = time.time()
            tmp_log.debug('done (took {0})'.format(end_time - start_time))
        except:
            tmp_log.error('Excepted with: {0}'.format(traceback.format_exc()))

    def convert_status(self, harvester_status):
        """
        convert harvester status to APFMon status
        :param harvester_status
        :return: list with apfmon_status. Usually it's just one status, except for exiting&done
        """
        if harvester_status == 'submitted':
            return ['created']
        if harvester_status in ['running', 'idle']:
            return ['running']
        if harvester_status in ['missed', 'failed', 'cancelled']:
            return ['fault']
        if harvester_status == 'finished':
            return ['exiting', 'done']

    def update_worker(self, worker_spec, worker_status):
        """
        Updates the state of a worker. This can also be done directly from the wrapper, assuming there is outbound
        connectivity on the worker node
        """
        start_time = time.time()
        tmp_log = core_utils.make_logger(_base_logger, 'harvester_id={0}'.format(self.harvester_id),
                                         method_name='update_workers')

        if not self.__active:
            tmp_log.debug('APFMon reporting not enabled')
            return

        try:
            tmp_log.debug('start')

            batch_id = worker_spec.batchID
            factory = self.harvester_id

            url = '{0}/jobs/{1}:{2}'.format(self.base_url, factory, batch_id)

            apfmon_status = self.convert_status(worker_status)
            apfmon_worker = {}

            for status in apfmon_status:
                apfmon_worker['state'] = status

                if status == 'exiting':
                    # return code
                    apfmon_worker['rc'] = 0
                    if hasattr(worker_spec, 'pandaid_list') and worker_spec.pandaid_list:
                        apfmon_worker['ids'] = ','.join(str(x) for x in worker_spec.pandaid_list)

                tmp_log.debug('updating worker {0}: {1}'.format(batch_id, apfmon_worker))

                r = requests.post(url, data=apfmon_worker, timeout=self.__worker_timeout)
                tmp_log.debug('worker update for {0} ended with {1} {2}'.format(batch_id, r.status_code, r.text))

            end_time = time.time()
            tmp_log.debug('done (took {0})'.format(end_time - start_time))
        except:
            tmp_log.error('Excepted with: {0}'.format(traceback.format_exc()))

if __name__== "__main__":

    """
    Quick tests
    """
    from pandaharvester.harvestercore.queue_config_mapper import QueueConfigMapper
    queue_config_mapper = QueueConfigMapper()

    apfmon = Apfmon(queue_config_mapper)
    apfmon.create_factory()
    apfmon.create_labels()

    worker_a = WorkSpec()
    worker_a.batchID = 1
    worker_a.computingSite = 'CERN-PROD-DEV_UCORE'
    worker_a.computingElement = 'bla1'
    worker_a.workAttributes = {"batchLog": "https://aipanda024.cern.ch/condor_logs/18-07-19_09/grid.9659.0.log", "stdErr": "https://aipanda024.cern.ch/condor_logs/18-07-19_09/grid.9659.0.err", "stdOut": "https://aipanda024.cern.ch/condor_logs/18-07-19_09/grid.9659.0.out"}
    worker_a.pandaid_list = [1234, 5678]

    worker_b = WorkSpec()
    worker_b.batchID = 2
    worker_b.computingSite = 'CERN-PROD-DEV_UCORE'
    worker_b.computingElement = 'bla2'
    worker_b.workAttributes = {"batchLog": "https://aipanda024.cern.ch/condor_logs/18-07-19_09/grid.9659.0.log", "stdErr": "https://aipanda024.cern.ch/condor_logs/18-07-19_09/grid.9659.0.err", "stdOut": "https://aipanda024.cern.ch/condor_logs/18-07-19_09/grid.9659.0.out"}

    workers = [worker_a, worker_b]

    apfmon.create_workers(workers)
    worker_a.status = 'running'
    worker_b.status = 'running'
    apfmon.update_workers(workers)
    worker_a.status = 'finished'
    worker_b.status = 'failed'
    apfmon.update_workers(workers)
<commit_msg>apfmon: skipping workers without batchID, i.e. not submitted properly<commit_after>"""
API described here: http://apfmon.lancs.ac.uk/help
"""

import requests
import json
import time
import traceback

from pandaharvester.harvesterconfig import harvester_config
from pandaharvester.harvestercore import core_utils
from pandaharvester import panda_pkg_info
from pandaharvester.harvestermisc import generic_utils
from pandaharvester.harvestercore.work_spec import WorkSpec
from pandaharvester.harvestermisc.info_utils import PandaQueuesDict

_base_logger = core_utils.setup_logger('apfmon')
NO_CE = 'noCE'

def apfmon_active(method, *args, **kwargs):
    if cls.__active:
        method(*args, **kwargs)
    else:
        return

def clean_ce(ce):
    return ce.split('.')[0].split('://')[-1]

class Apfmon(object):

    def __init__(self, queue_config_mapper):

        try:
            self.__active = harvester_config.apfmon.active
        except:
            self.__active = False

        try:
            self.__worker_timeout = harvester_config.apfmon.worker_timeout
        except:
            self.__worker_timeout = 0.2

        try:
            self.__label_timeout = harvester_config.apfmon.worker_timeout
        except:
            self.__label_timeout = 1

        # TODO: make proper exception handling and defaults
        try:
            self.harvester_id = harvester_config.master.harvester_id
        except:
            self.harvester_id = 'DUMMY'

        try:
            self.base_url = harvester_config.apfmon.base_url
        except:
            self.base_url = 'http://apfmon.lancs.ac.uk/api'

        self.queue_config_mapper = queue_config_mapper

    def create_factory(self):
        """
        Creates or updates a harvester instance to APF Mon. Should be done at startup of the instance.
        """

        start_time = time.time()
        tmp_log = core_utils.make_logger(_base_logger, 'harvester_id={0}'.format(self.harvester_id),
                                         method_name='create_factory')

        if not self.__active:
            tmp_log.debug('APFMon reporting not enabled')
            return

        try:
            tmp_log.debug('start')

            url = '{0}/factories/{1}'.format(self.base_url, self.harvester_id)

            f = {'url': 'url_to_logs',
                 'email': 'atlas-adc-harvester-central-support@cern.ch',
                 'version': panda_pkg_info.release_version}
            payload = json.dumps(f)

            r = requests.put(url, data=payload, timeout=self.__label_timeout)
            tmp_log.debug('registration ended with {0} {1}'.format(r.status_code, r.text))
            end_time = time.time()
            tmp_log.debug('done (took {0})'.format(end_time - start_time))
        except:
            tmp_log.error('Excepted with: {0}'.format(traceback.format_exc()))

    def create_labels(self):
        """
        Creates or updates a collection of labels (=panda queue+CE)
        """
        start_time = time.time()
        tmp_log = core_utils.make_logger(_base_logger, 'harvester_id={0}'.format(self.harvester_id),
                                         method_name='create_labels')

        if not self.__active:
            tmp_log.debug('APFMon reporting not enabled')
            return

        try:
            tmp_log.debug('start')

            url = '{0}/labels'.format(self.base_url)

            # get the active queues from the config mapper
            all_sites = self.queue_config_mapper.get_active_queues().keys()
            panda_queues_dict = PandaQueuesDict()

            # publish the active queues to APF mon in shards
            for sites in generic_utils.create_shards(all_sites, 20):
                labels = []
                for site in sites:
                    try:
                        site_info = panda_queues_dict.get(site, dict())
                        if not site_info:
                            tmp_log.warning('No site info for {0}'.format(site))
                            continue

                        if not site_info['queues']:
                            # when no CEs associated to a queue, e.g. P1, HPCs, etc. Try to see if there is something
                            # in local configuration, otherwise set it to a dummy value
                            try:
                                ce = self.queue_config_mapper.queueConfig[site].submitter['ceEndpoint']
                            except:
                                ce = NO_CE
                            queues = [{'ce_endpoint': ce}]
                        else:
                            queues = site_info['queues']

                        for queue in queues:
                            try:
                                ce = clean_ce(queue['ce_endpoint'])
                            except:
                                ce = ''

                            try:
                                ce_queue_id = queue['ce_queue_id']
                            except KeyError:
                                ce_queue_id = 0

                            labels.append({'name': '{0}-{1}'.format(site, ce),
                                           'wmsqueue': site,
                                           'ce_queue_id': ce_queue_id,
                                           'factory': self.harvester_id})
                    except:
                        tmp_log.error('Excepted for site {0} with: {1}'.format(site, traceback.format_exc()))
                        continue

                payload = json.dumps(labels)

                r = requests.put(url, data=payload, timeout=self.__label_timeout)
                tmp_log.debug('label creation for {0} ended with {1} {2}'.format(sites, r.status_code, r.text))

            end_time = time.time()
            tmp_log.debug('done (took {0})'.format(end_time - start_time))
        except:
            tmp_log.error('Excepted with: {0}'.format(traceback.format_exc()))

    def massage_label_data(self, data):

        tmp_log = core_utils.make_logger(_base_logger, 'harvester_id={0}'.format(self.harvester_id),
                                         method_name='massage_label_data')
        if not data:
            return data

        try:
            any = data['ANY']
            agg = {}
            for rtype in data:
                if rtype == 'ANY':
                    continue
                else:
                    for value in data[rtype]:
                        agg.setdefault(value, 0)
                        agg[value] += data[rtype][value]

            if agg:
                data['ANY'] = agg
            else:
                data['ANY'] = any

            tmp_log.debug('Massaged to data: {0}'.format(data))

        except Exception:
            tmp_log.debug('Exception in data: {0}'.format(data))

        return data

    def update_label(self, site, msg, data):
        """
        Updates a label (=panda queue+CE)
        """
        start_time = time.time()
        tmp_log = core_utils.make_logger(_base_logger, 'harvester_id={0}'.format(self.harvester_id),
                                         method_name='update_label')

        if not self.__active:
            tmp_log.debug('APFMon reporting not enabled')
            return

        try:
            tmp_log.debug('start')
            data = self.massage_label_data(data)

            # get the active queues from the config mapper
            all_sites = self.queue_config_mapper.get_active_queues().keys()
            panda_queues_dict = PandaQueuesDict()

            site_info = panda_queues_dict.get(site, dict())
            if not site_info:
                tmp_log.warning('No site info for {0}'.format(site))
                return

            if not site_info['queues']:
                # when no CEs associated to a queue, e.g. P1, HPCs, etc. Try to see if there is something
                # in local configuration, otherwise set it to a dummy value
                try:
                    ce = self.queue_config_mapper.queueConfig[site].submitter['ceEndpoint']
                except:
                    ce = NO_CE
                queues = [{'ce_endpoint': ce}]
            else:
                queues = site_info['queues']

            for queue in queues:
                try:
                    try:
                        ce = clean_ce(queue['ce_endpoint'])
                    except:
                        ce = ''

                    label_data = {'status': msg, 'data': data}
                    label = '{0}-{1}'.format(site, ce)
                    label_id = '{0}:{1}'.format(self.harvester_id, label)
                    url = '{0}/labels/{1}'.format(self.base_url, label_id)

                    r = requests.post(url, data=json.dumps(label_data), timeout=self.__label_timeout)
                    tmp_log.debug('label update for {0} ended with {1} {2}'.format(label, r.status_code, r.text))
                except:
                    tmp_log.error('Excepted for site {0} with: {1}'.format(label, traceback.format_exc()))

            end_time = time.time()
            tmp_log.debug('done (took {0})'.format(end_time - start_time))
        except:
            tmp_log.error('Excepted with: {0}'.format(traceback.format_exc()))

    def create_workers(self, worker_spec_list):
        """
        Creates a worker
        """
        start_time = time.time()
        tmp_log = core_utils.make_logger(_base_logger, 'harvester_id={0}'.format(self.harvester_id),
                                         method_name='create_workers')

        if not self.__active:
            tmp_log.debug('APFMon reporting not enabled')
            return

        try:
            tmp_log.debug('start')

            url = '{0}/jobs'.format(self.base_url)

            for worker_spec_shard in generic_utils.create_shards(worker_spec_list, 20):
                apfmon_workers = []
                for worker_spec in worker_spec_shard:
                    batch_id = worker_spec.batchID
                    worker_id = worker_spec.workerID
                    if not batch_id:
                        tmp_log.debug('no batchID found for workerID {0}... skipping'.format(worker_id))
                        continue
                    factory = self.harvester_id
                    computingsite = worker_spec.computingSite
                    try:
                        ce = clean_ce(worker_spec.computingElement)
                    except AttributeError:
                        tmp_log.debug('no CE found for workerID {0} batchID {1}'.format(worker_id, batch_id))
                        ce = ''

                    # extract the log URLs
                    stdout_url = ''
                    stderr_url = ''
                    log_url = ''
                    jdl_url = ''

                    work_attribs = worker_spec.workAttributes
                    if work_attribs:
                        if 'stdOut' in work_attribs:
                            stdout_url = work_attribs['stdOut']
                            jdl_url = '{0}.jdl'.format(stdout_url[:-4])
                        if 'stdErr' in work_attribs:
                            stderr_url = work_attribs['stdErr']
                        if 'batchLog' in work_attribs:
                            log_url = work_attribs['batchLog']

                    apfmon_worker = {'cid': batch_id,
                                     'factory': factory,
                                     'label': '{0}-{1}'.format(computingsite, ce),
                                     'jdlurl': jdl_url,
                                     'stdouturl': stdout_url,
                                     'stderrurl': stderr_url,
                                     'logurl': log_url
                                     }
                    tmp_log.debug('packed worker: {0}'.format(apfmon_worker))
                    apfmon_workers.append(apfmon_worker)

                payload = json.dumps(apfmon_workers)
                r = requests.put(url, data=payload, timeout=self.__worker_timeout)
                tmp_log.debug('worker creation for {0} ended with {1} {2}'.format(apfmon_workers, r.status_code, r.text))

            end_time = time.time()
            tmp_log.debug('done (took {0})'.format(end_time - start_time))
        except:
            tmp_log.error('Excepted with: {0}'.format(traceback.format_exc()))

    def convert_status(self, harvester_status):
        """
        convert harvester status to APFMon status
        :param harvester_status
        :return: list with apfmon_status. Usually it's just one status, except for exiting&done
        """
        if harvester_status == 'submitted':
            return ['created']
        if harvester_status in ['running', 'idle']:
            return ['running']
        if harvester_status in ['missed', 'failed', 'cancelled']:
            return ['fault']
        if harvester_status == 'finished':
            return ['exiting', 'done']

    def update_worker(self, worker_spec, worker_status):
        """
        Updates the state of a worker. This can also be done directly from the wrapper, assuming there is outbound
        connectivity on the worker node
        """
        start_time = time.time()
        tmp_log = core_utils.make_logger(_base_logger, 'harvester_id={0}'.format(self.harvester_id),
                                         method_name='update_workers')

        if not self.__active:
            tmp_log.debug('APFMon reporting not enabled')
            return

        try:
            tmp_log.debug('start')

            batch_id = worker_spec.batchID
            factory = self.harvester_id

            url = '{0}/jobs/{1}:{2}'.format(self.base_url, factory, batch_id)

            apfmon_status = self.convert_status(worker_status)
            apfmon_worker = {}

            for status in apfmon_status:
                apfmon_worker['state'] = status

                if status == 'exiting':
                    # return code
                    apfmon_worker['rc'] = 0
                    if hasattr(worker_spec, 'pandaid_list') and worker_spec.pandaid_list:
                        apfmon_worker['ids'] = ','.join(str(x) for x in worker_spec.pandaid_list)

                tmp_log.debug('updating worker {0}: {1}'.format(batch_id, apfmon_worker))

                r = requests.post(url, data=apfmon_worker, timeout=self.__worker_timeout)
                tmp_log.debug('worker update for {0} ended with {1} {2}'.format(batch_id, r.status_code, r.text))

            end_time = time.time()
            tmp_log.debug('done (took {0})'.format(end_time - start_time))
        except:
            tmp_log.error('Excepted with: {0}'.format(traceback.format_exc()))

if __name__== "__main__":

    """
    Quick tests
    """
    from pandaharvester.harvestercore.queue_config_mapper import QueueConfigMapper
    queue_config_mapper = QueueConfigMapper()

    apfmon = Apfmon(queue_config_mapper)
    apfmon.create_factory()
    apfmon.create_labels()

    worker_a = WorkSpec()
    worker_a.batchID = 1
    worker_a.computingSite = 'CERN-PROD-DEV_UCORE'
    worker_a.computingElement = 'bla1'
    worker_a.workAttributes = {"batchLog": "https://aipanda024.cern.ch/condor_logs/18-07-19_09/grid.9659.0.log", "stdErr": "https://aipanda024.cern.ch/condor_logs/18-07-19_09/grid.9659.0.err", "stdOut": "https://aipanda024.cern.ch/condor_logs/18-07-19_09/grid.9659.0.out"}
    worker_a.pandaid_list = [1234, 5678]

    worker_b = WorkSpec()
    worker_b.batchID = 2
    worker_b.computingSite = 'CERN-PROD-DEV_UCORE'
    worker_b.computingElement = 'bla2'
    worker_b.workAttributes = {"batchLog": "https://aipanda024.cern.ch/condor_logs/18-07-19_09/grid.9659.0.log", "stdErr": "https://aipanda024.cern.ch/condor_logs/18-07-19_09/grid.9659.0.err", "stdOut": "https://aipanda024.cern.ch/condor_logs/18-07-19_09/grid.9659.0.out"}

    workers = [worker_a, worker_b]

    apfmon.create_workers(workers)
    worker_a.status = 'running'
    worker_b.status = 'running'
    apfmon.update_workers(workers)
    worker_a.status = 'finished'
    worker_b.status = 'failed'
    apfmon.update_workers(workers)
<|endoftext|>