<commit_before>import numpy
import operator

from functools import reduce
from six import moves

from chainer import cuda
from chainer import function
from chainer.functions.connection import convolution_2d
from chainer.utils import conv
from chainer.utils import conv_nd
from chainer.utils import type_check

if cuda.cudnn_enabled:
    cudnn = cuda.cudnn
    libcudnn = cuda.cudnn.cudnn
    _cudnn_version = libcudnn.getVersion()
    _fwd_pref = libcudnn.CUDNN_CONVOLUTION_FWD_SPECIFY_WORKSPACE_LIMIT
    if _cudnn_version >= 4000:
        _bwd_filter_pref = \
            libcudnn.CUDNN_CONVOLUTION_BWD_FILTER_SPECIFY_WORKSPACE_LIMIT
        _bwd_data_pref = \
            libcudnn.CUDNN_CONVOLUTION_BWD_DATA_SPECIFY_WORKSPACE_LIMIT

def _tuple(x, N):
    if hasattr(x, '__getitem__'):
        assert len(x) == N
        return x
    return (x,) * N

class ConvolutionND(function.Function):

    def __init__(self, ndim, stride=1, pad=0, use_cudnn=True, cover_all=False):
        self.ndim = ndim
        self.stride = _tuple(stride, ndim)
        self.pad = _tuple(pad, ndim)
        self.use_cudnn = use_cudnn
        self.cover_all = cover_all

    def check_type_forward(self, in_types):
        n_in = in_types.size()
        type_check.expect(2 <= n_in, n_in <= 3)

        x_type = in_types[0]
        w_type = in_types[1]
        type_check.expect(
            x_type.dtype.kind == 'f',
            w_type.dtype.kind == 'f',
            x_type.ndim == self.ndim + 2,
            w_type.ndim == self.ndim + 2,
            x_type.shape[1] == w_type.shape[1],
        )

        if n_in.eval() == 3:
            b_type = in_types[2]
            type_check.expect(
                b_type.dtype == x_type.dtype,
                b_type.ndim == 1,
                b_type.shape[0] == w_type.shape[0],
            )

    def forward_cpu(self, inputs):
        x, W = inputs[:2]
        b = inputs[2] if len(inputs) == 3 else None
        ndim = self.ndim
        ksize = W.shape[2:]
        stride = self.stride
        pad = self.pad

        # Make patch array.
        self.col = conv_nd.im2col_nd_cpu(
            x, ksize, stride, pad, cover_all=self.cover_all)

        # Compute correlation.
        axes = tuple(moves.range(1, ndim+2))  # (1, 2, ..., N+1)
        y = numpy.tensordot(self.col, W, (axes, axes)).astype(x.dtype)

        # Apply bias if given.
        if b is not None:
            y += b

        # Roll c_O before the second in (n, y_1, y_2, ..., y_N, c_O).
        return numpy.rollaxis(y, ndim+1, 1),

    def forward_gpu(self, inputs):
        x, W = inputs[:2]
        b = inputs[2] if len(inputs) == 3 else None
        out_c = W.shape[0]      # (c_O, _, k_1, k_2, ..., k_N)
        ksize = W.shape[2:]
        n, c = x.shape[:2]      # (n, c_I, d_1, d_2, ..., d_N)
        dims = x.shape[2:]
        stride = self.stride
        pad = self.pad
        ndim = self.ndim
        colon = slice(None)

        # Compute output image's dimensions.
        outs = tuple(
            [conv.get_conv_outsize(d, k, s, p, cover_all=self.cover_all)
             for (d, k, s, p) in zip(dims, ksize, stride, pad)])

        # Make empty array for result.
        y_shape = (n, out_c) + outs  # (n, c_O, out_1, out_2, ..., out_N)
        y = cuda.empty(y_shape, dtype=x.dtype)
        # Implementation using cuDNN.
        if (not self.cover_all and cuda.cudnn_enabled and self.use_cudnn and
            ndim > 1 and convolution_2d._check_cudnn_acceptable_type(
                x.dtype, W.dtype)):
            x = cuda.cupy.ascontiguousarray(x)
            W = cuda.cupy.ascontiguousarray(W)
            if b is not None:
                b = cuda.cupy.ascontiguousarray(b)

            handle = cudnn.get_handle()
            x_desc = cudnn.create_tensor_descriptor(x)
            y_desc = cudnn.create_tensor_descriptor(y)

            self.filter_desc = cudnn.create_filter_descriptor(W)
            self.conv_desc = cudnn.create_convolution_descriptor(
                pad, stride)
            if b is not None:
                b_index = (None, colon) + (None,) * ndim
                self.bias_desc = cudnn.create_tensor_descriptor(b[b_index])

            workspace_size = cuda.get_max_workspace_size()
            workspace = cuda.empty((workspace_size,), dtype='b')
            algo = libcudnn.getConvolutionForwardAlgorithm(
                handle, x_desc.value, self.filter_desc.value,
                self.conv_desc.value, y_desc.value, _fwd_pref,
                workspace_size)

            oz_dtype = 'd' if x.dtype == 'd' else 'f'
            one = numpy.array(1, dtype=oz_dtype).ctypes
            zero = numpy.array(0, dtype=oz_dtype).ctypes
            libcudnn.convolutionForward(
                handle, one.data, x_desc.value, x.data.ptr,
                self.filter_desc.value, W.data.ptr, self.conv_desc.value,
                algo, workspace.data.ptr, workspace_size, zero.data,
                y_desc.value, y.data.ptr)

            # TODO(takagi) Support unshared bias
            if b is not None:
                cudnn.add_tensor(
                    handle, one.data, self.bias_desc.value, b.data.ptr,
                    one.data, y_desc.value, y.data.ptr)
        # Implementation using im2col.
        else:
            # Make patch array.
            self.col = conv_nd.im2col_nd_gpu(
                x, ksize, stride, pad, cover_all=self.cover_all)

            # Compute correlation.
            W_mat = W.reshape(out_c, -1)
            col_mats = self.col.reshape(n, -1, reduce(operator.mul, outs))
            y_mats = y.reshape(n, out_c, -1)
            # TODO(takagi): Use streams or batch gemm
            for i in moves.range(n):
                y_mats[i] = W_mat.dot(col_mats[i])

            # Apply bias if given.
            # TODO(takagi): Support unshared bias
            if b is not None:
                index = (colon,) + (None,) * ndim  # (:, None, ..., None)
                y += b[index]

        return y,

    def backward_cpu(self, inputs, grad_outputs):
        x, W = inputs[:2]
        b = inputs[2] if len(inputs) == 3 else None
        gy = grad_outputs[0]   # (n, c_O, out_1, out_2, ..., out_N)
        dims = x.shape[2:]     # (n, c_I, d_1, d_2, ..., d_N)
        stride = self.stride
        pad = self.pad
        ndim = self.ndim

        # Compute filter weight gradient.
        # (n, _, out_1, out_2, ..., out_N)
        out_axes = (0,) + tuple(moves.range(2, ndim+2))
        # (n, _, _, ..., _, out_1, out_2, ..., out_N)
        col_axes = (0,) + tuple(moves.range(ndim+2, ndim*2+2))
        gW = numpy.tensordot(
            gy, self.col, (out_axes, col_axes)).astype(W.dtype)

        # Compute patch array gradient.
        gcol = numpy.tensordot(W, gy, (0, 1)).astype(x.dtype)
        gcol = numpy.rollaxis(gcol, ndim + 1)

        # Compute input gradient.
        gx = conv_nd.col2im_nd_cpu(gcol, stride, pad, dims)

        # Compute bias gradient if given and return gradients.
        if b is None:
            return gx, gW
        else:
            # (n, _, out_1, out_2, ..., out_N)
            axis = (0,) + tuple(moves.range(2, ndim+2))
            gb = gy.sum(axis=axis)
            return gx, gW, gb

    def backward_gpu(self, inputs, grad_putputs):
        x, W = inputs[:2]
        b = inputs[2] if len(inputs) == 3 else None
        gy = grad_putputs[0]    # (n, c_O, out_1, out_2, ..., out_N)
        out_c = gy.shape[1]
        outs = gy.shape[2:]
        n, c = x.shape[:2]      # (n, c_I, d_1, d_2, ..., d_N)
        dims = x.shape[2:]
        ksize = W.shape[2:]     # (_, _, k_1, k_2, ..., k_N)
        stride = self.stride
        pad = self.pad
        ndim = self.ndim

        # Compute filter weight gradient.
        gW = cuda.empty_like(W)
        # Implementation using cuDNN.
        if (not self.cover_all and cuda.cudnn_enabled and self.use_cudnn and
            ndim > 1 and convolution_2d._check_cudnn_acceptable_type(
                x.dtype, W.dtype)):
            x = cuda.cupy.ascontiguousarray(x)
            W = cuda.cupy.ascontiguousarray(W)
            gy = cuda.cupy.ascontiguousarray(gy)

            handle = cudnn.get_handle()
            x_desc = cudnn.create_tensor_descriptor(x)
            gy_desc = cudnn.create_tensor_descriptor(gy)
            oz_dtype = 'd' if x.dtype == 'd' else 'f'
            one = numpy.array(1, dtype=oz_dtype).ctypes
            zero = numpy.array(0, dtype=oz_dtype).ctypes
            gx = cuda.cupy.empty_like(x)

            if _cudnn_version >= 4000:
                workspace_size = cuda.get_max_workspace_size()
                workspace = cuda.cupy.empty((workspace_size,), dtype='b')

                algo = libcudnn.getConvolutionBackwardFilterAlgorithm(
                    handle, x_desc.value, gy_desc.value,
                    self.conv_desc.value, self.filter_desc.value,
                    _bwd_filter_pref, workspace_size)
                libcudnn.convolutionBackwardFilter_v3(
                    handle, one.data, x_desc.value, x.data.ptr,
                    gy_desc.value, gy.data.ptr, self.conv_desc.value,
                    algo, workspace.data.ptr, workspace_size,
                    zero.data, self.filter_desc.value, gW.data.ptr)

                algo = libcudnn.getConvolutionBackwardDataAlgorithm(
                    handle, self.filter_desc.value, gy_desc.value,
                    self.conv_desc.value, x_desc.value, _bwd_data_pref,
                    workspace_size)
                libcudnn.convolutionBackwardData_v3(
                    handle, one.data, self.filter_desc.value, W.data.ptr,
                    gy_desc.value, gy.data.ptr, self.conv_desc.value,
                    algo, workspace.data.ptr, workspace_size,
                    zero.data, x_desc.value, gx.data.ptr)
            else:
                libcudnn.convolutionBackwardFilter_v2(
                    handle, one.data, x_desc.value, x.data.ptr,
                    gy_desc.value, gy.data.ptr, self.conv_desc.value,
                    zero.data, self.filter_desc.value, gW.data.ptr)
                libcudnn.convolutionBackwardData_v2(
                    handle, one.data, self.filter_desc.value, W.data.ptr,
                    gy_desc.value, gy.data.ptr, self.conv_desc.value,
                    zero.data, x_desc.value, gx.data.ptr)

            if b is not None:
                gb = cuda.cupy.empty_like(b)
                libcudnn.convolutionBackwardBias(
                    handle, one.data, gy_desc.value, gy.data.ptr,
                    zero.data, self.bias_desc.value, gb.data.ptr)
        # Implementation using col2im.
        else:
            gW_mat = gW.reshape(out_c, reduce(operator.mul, ksize, c))
            col_mats = self.col.reshape(
                n, reduce(operator.mul, ksize, c), reduce(operator.mul, outs))
            gy_mats = gy.reshape(n, out_c, reduce(operator.mul, outs))
            # TODO(takagi): Use streams or batch gemm
            gW_mat[...] = 0
            for i in moves.range(n):
                gW_mat += cuda.cupy.dot(gy_mats[i], col_mats[i].T)

            # Compute patch array gradient.
            W_mat = W.reshape(out_c, -1)
            gcol = cuda.empty_like(self.col)
            gcol_mats = gcol.reshape(
                n, reduce(operator.mul, ksize, c), reduce(operator.mul, outs))
            for i in moves.range(n):
                gcol_mats[i] = cuda.cupy.dot(W_mat.T, gy_mats[i])

            # Compute input gradient.
            gx = conv_nd.col2im_nd_gpu(gcol, stride, pad, dims)

            # Compute bias gradient if given.
            if b is not None:
                # (n, _, out_1, out_2, ..., out_N)
                axis = (0,) + tuple(moves.range(2, ndim+2))
                gb = gy.sum(axis=axis)

        if b is None:
            return gx, gW
        else:
            return gx, gW, gb

def convolution_nd(x, W, b=None, stride=1, pad=0, use_cudnn=True,
                   cover_all=False):
    """N-dimensional convolution function.

    This is an implementation of N-dimensional convolution which is generalized
    two-dimensional convolution in ConvNets. It takes three variables: the
    input image ``x``, the filter weight ``W`` and the bias vector ``b``.

    Notation: here is a notation for dimensionalities.

    - :math:`N` is the number of spatial dimensions.
    - :math:`n` is the batch size.
    - :math:`c_I` and :math:`c_O` are the number of the input and output,
      respectively.
    - :math:`d_1, d_2, ..., d_N` are the size of each axis of the input image,
      respectively.
    - :math:`k_1, k_2, ..., k_N` are the size of each axis of the filters,
      respectively.

    Args:
        x (~chainer.Variable): Input variable of shape
            :math:`(n, c_I, d_1, d_2, ..., d_N)`.
        W (~chainer.Variable): Weight variable of shape
            :math:`(c_O, c_I, k_1, k_2, ..., k_N)`.
        b (~chainer.Variable): Bias variable of length :math:`c_O`
            (optional).
        stride (int or tuple of ints): Stride of filter applications
            :math:`(s_1, s_2, ..., s_N)`. ``stride=s`` is equivalent to
            ``(s, s, ..., s)``.
        pad (int or tuple of ints): Spatial padding width for input arrays
            :math:`(p_1, p_2, ..., p_N)`. ``pad=p`` is equivalent to
            ``(p, p, ..., p)``.
        use_cudnn (bool): If ``True``, then this function uses cuDNN if
            available. cuDNN supports more than one-dimensional convolution.
        cover_all (bool): If ``True``, all spatial locations are convoluted
            into some output pixels. It may make the output size larger.

    Returns:
        ~chainer.Variable: Output variable.

    .. seealso:: :class:`ConvolutionND`, :func:`convolution_2d`
    """
    ndim = len(x.data.shape[2:])
    func = ConvolutionND(ndim, stride, pad, use_cudnn, cover_all)
    if b is None:
        return func(x, W)
    else:
        return func(x, W, b)
<commit_msg>Keep spaces around operators.<commit_after>import numpy
import operator

from functools import reduce
from six import moves

from chainer import cuda
from chainer import function
from chainer.functions.connection import convolution_2d
from chainer.utils import conv
from chainer.utils import conv_nd
from chainer.utils import type_check

if cuda.cudnn_enabled:
    cudnn = cuda.cudnn
    libcudnn = cuda.cudnn.cudnn
    _cudnn_version = libcudnn.getVersion()
    _fwd_pref = libcudnn.CUDNN_CONVOLUTION_FWD_SPECIFY_WORKSPACE_LIMIT
    if _cudnn_version >= 4000:
        _bwd_filter_pref = \
            libcudnn.CUDNN_CONVOLUTION_BWD_FILTER_SPECIFY_WORKSPACE_LIMIT
        _bwd_data_pref = \
            libcudnn.CUDNN_CONVOLUTION_BWD_DATA_SPECIFY_WORKSPACE_LIMIT

def _tuple(x, N):
    if hasattr(x, '__getitem__'):
        assert len(x) == N
        return x
    return (x,) * N

class ConvolutionND(function.Function):

    def __init__(self, ndim, stride=1, pad=0, use_cudnn=True, cover_all=False):
        self.ndim = ndim
        self.stride = _tuple(stride, ndim)
        self.pad = _tuple(pad, ndim)
        self.use_cudnn = use_cudnn
        self.cover_all = cover_all

    def check_type_forward(self, in_types):
        n_in = in_types.size()
        type_check.expect(2 <= n_in, n_in <= 3)

        x_type = in_types[0]
        w_type = in_types[1]
        type_check.expect(
            x_type.dtype.kind == 'f',
            w_type.dtype.kind == 'f',
            x_type.ndim == self.ndim + 2,
            w_type.ndim == self.ndim + 2,
            x_type.shape[1] == w_type.shape[1],
        )

        if n_in.eval() == 3:
            b_type = in_types[2]
            type_check.expect(
                b_type.dtype == x_type.dtype,
                b_type.ndim == 1,
                b_type.shape[0] == w_type.shape[0],
            )

    def forward_cpu(self, inputs):
        x, W = inputs[:2]
        b = inputs[2] if len(inputs) == 3 else None
        ndim = self.ndim
        ksize = W.shape[2:]
        stride = self.stride
        pad = self.pad

        # Make patch array.
        self.col = conv_nd.im2col_nd_cpu(
            x, ksize, stride, pad, cover_all=self.cover_all)

        # Compute correlation.
        axes = tuple(moves.range(1, ndim + 2))  # (1, 2, ..., N+1)
        y = numpy.tensordot(self.col, W, (axes, axes)).astype(x.dtype)

        # Apply bias if given.
        if b is not None:
            y += b

        # Roll c_O before the second in (n, y_1, y_2, ..., y_N, c_O).
        return numpy.rollaxis(y, ndim + 1, 1),

    def forward_gpu(self, inputs):
        x, W = inputs[:2]
        b = inputs[2] if len(inputs) == 3 else None
        out_c = W.shape[0]      # (c_O, _, k_1, k_2, ..., k_N)
        ksize = W.shape[2:]
        n, c = x.shape[:2]      # (n, c_I, d_1, d_2, ..., d_N)
        dims = x.shape[2:]
        stride = self.stride
        pad = self.pad
        ndim = self.ndim
        colon = slice(None)

        # Compute output image's dimensions.
        outs = tuple(
            [conv.get_conv_outsize(d, k, s, p, cover_all=self.cover_all)
             for (d, k, s, p) in zip(dims, ksize, stride, pad)])

        # Make empty array for result.
        y_shape = (n, out_c) + outs  # (n, c_O, out_1, out_2, ..., out_N)
        y = cuda.empty(y_shape, dtype=x.dtype)
        # Implementation using cuDNN.
        if (not self.cover_all and cuda.cudnn_enabled and self.use_cudnn and
            ndim > 1 and convolution_2d._check_cudnn_acceptable_type(
                x.dtype, W.dtype)):
            x = cuda.cupy.ascontiguousarray(x)
            W = cuda.cupy.ascontiguousarray(W)
            if b is not None:
                b = cuda.cupy.ascontiguousarray(b)

            handle = cudnn.get_handle()
            x_desc = cudnn.create_tensor_descriptor(x)
            y_desc = cudnn.create_tensor_descriptor(y)

            self.filter_desc = cudnn.create_filter_descriptor(W)
            self.conv_desc = cudnn.create_convolution_descriptor(
                pad, stride)
            if b is not None:
                b_index = (None, colon) + (None,) * ndim
                self.bias_desc = cudnn.create_tensor_descriptor(b[b_index])

            workspace_size = cuda.get_max_workspace_size()
            workspace = cuda.empty((workspace_size,), dtype='b')
            algo = libcudnn.getConvolutionForwardAlgorithm(
                handle, x_desc.value, self.filter_desc.value,
                self.conv_desc.value, y_desc.value, _fwd_pref,
                workspace_size)

            oz_dtype = 'd' if x.dtype == 'd' else 'f'
            one = numpy.array(1, dtype=oz_dtype).ctypes
            zero = numpy.array(0, dtype=oz_dtype).ctypes
            libcudnn.convolutionForward(
                handle, one.data, x_desc.value, x.data.ptr,
                self.filter_desc.value, W.data.ptr, self.conv_desc.value,
                algo, workspace.data.ptr, workspace_size, zero.data,
                y_desc.value, y.data.ptr)

            # TODO(takagi) Support unshared bias
            if b is not None:
                cudnn.add_tensor(
                    handle, one.data, self.bias_desc.value, b.data.ptr,
                    one.data, y_desc.value, y.data.ptr)
        # Implementation using im2col.
        else:
            # Make patch array.
            self.col = conv_nd.im2col_nd_gpu(
                x, ksize, stride, pad, cover_all=self.cover_all)

            # Compute correlation.
            W_mat = W.reshape(out_c, -1)
            col_mats = self.col.reshape(n, -1, reduce(operator.mul, outs))
            y_mats = y.reshape(n, out_c, -1)
            # TODO(takagi): Use streams or batch gemm
            for i in moves.range(n):
                y_mats[i] = W_mat.dot(col_mats[i])

            # Apply bias if given.
            # TODO(takagi): Support unshared bias
            if b is not None:
                index = (colon,) + (None,) * ndim  # (:, None, ..., None)
                y += b[index]

        return y,

    def backward_cpu(self, inputs, grad_outputs):
        x, W = inputs[:2]
        b = inputs[2] if len(inputs) == 3 else None
        gy = grad_outputs[0]   # (n, c_O, out_1, out_2, ..., out_N)
        dims = x.shape[2:]     # (n, c_I, d_1, d_2, ..., d_N)
        stride = self.stride
        pad = self.pad
        ndim = self.ndim

        # Compute filter weight gradient.
        # (n, _, out_1, out_2, ..., out_N)
        out_axes = (0,) + tuple(moves.range(2, ndim + 2))
        # (n, _, _, ..., _, out_1, out_2, ..., out_N)
        col_axes = (0,) + tuple(moves.range(ndim + 2, ndim * 2 + 2))
        gW = numpy.tensordot(
            gy, self.col, (out_axes, col_axes)).astype(W.dtype)

        # Compute patch array gradient.
        gcol = numpy.tensordot(W, gy, (0, 1)).astype(x.dtype)
        gcol = numpy.rollaxis(gcol, ndim + 1)

        # Compute input gradient.
        gx = conv_nd.col2im_nd_cpu(gcol, stride, pad, dims)

        # Compute bias gradient if given and return gradients.
        if b is None:
            return gx, gW
        else:
            # (n, _, out_1, out_2, ..., out_N)
            axis = (0,) + tuple(moves.range(2, ndim + 2))
            gb = gy.sum(axis=axis)
            return gx, gW, gb

    def backward_gpu(self, inputs, grad_putputs):
        x, W = inputs[:2]
        b = inputs[2] if len(inputs) == 3 else None
        gy = grad_putputs[0]    # (n, c_O, out_1, out_2, ..., out_N)
        out_c = gy.shape[1]
        outs = gy.shape[2:]
        n, c = x.shape[:2]      # (n, c_I, d_1, d_2, ..., d_N)
        dims = x.shape[2:]
        ksize = W.shape[2:]     # (_, _, k_1, k_2, ..., k_N)
        stride = self.stride
        pad = self.pad
        ndim = self.ndim

        # Compute filter weight gradient.
        gW = cuda.empty_like(W)
        # Implementation using cuDNN.
        if (not self.cover_all and cuda.cudnn_enabled and self.use_cudnn and
            ndim > 1 and convolution_2d._check_cudnn_acceptable_type(
                x.dtype, W.dtype)):
            x = cuda.cupy.ascontiguousarray(x)
            W = cuda.cupy.ascontiguousarray(W)
            gy = cuda.cupy.ascontiguousarray(gy)

            handle = cudnn.get_handle()
            x_desc = cudnn.create_tensor_descriptor(x)
            gy_desc = cudnn.create_tensor_descriptor(gy)
            oz_dtype = 'd' if x.dtype == 'd' else 'f'
            one = numpy.array(1, dtype=oz_dtype).ctypes
            zero = numpy.array(0, dtype=oz_dtype).ctypes
            gx = cuda.cupy.empty_like(x)

            if _cudnn_version >= 4000:
                workspace_size = cuda.get_max_workspace_size()
                workspace = cuda.cupy.empty((workspace_size,), dtype='b')

                algo = libcudnn.getConvolutionBackwardFilterAlgorithm(
                    handle, x_desc.value, gy_desc.value,
                    self.conv_desc.value, self.filter_desc.value,
                    _bwd_filter_pref, workspace_size)
                libcudnn.convolutionBackwardFilter_v3(
                    handle, one.data, x_desc.value, x.data.ptr,
                    gy_desc.value, gy.data.ptr, self.conv_desc.value,
                    algo, workspace.data.ptr, workspace_size,
                    zero.data, self.filter_desc.value, gW.data.ptr)

                algo = libcudnn.getConvolutionBackwardDataAlgorithm(
                    handle, self.filter_desc.value, gy_desc.value,
                    self.conv_desc.value, x_desc.value, _bwd_data_pref,
                    workspace_size)
                libcudnn.convolutionBackwardData_v3(
                    handle, one.data, self.filter_desc.value, W.data.ptr,
                    gy_desc.value, gy.data.ptr, self.conv_desc.value,
                    algo, workspace.data.ptr, workspace_size,
                    zero.data, x_desc.value, gx.data.ptr)
            else:
                libcudnn.convolutionBackwardFilter_v2(
                    handle, one.data, x_desc.value, x.data.ptr,
                    gy_desc.value, gy.data.ptr, self.conv_desc.value,
                    zero.data, self.filter_desc.value, gW.data.ptr)
                libcudnn.convolutionBackwardData_v2(
                    handle, one.data, self.filter_desc.value, W.data.ptr,
                    gy_desc.value, gy.data.ptr, self.conv_desc.value,
                    zero.data, x_desc.value, gx.data.ptr)

            if b is not None:
                gb = cuda.cupy.empty_like(b)
                libcudnn.convolutionBackwardBias(
                    handle, one.data, gy_desc.value, gy.data.ptr,
                    zero.data, self.bias_desc.value, gb.data.ptr)
        # Implementation using col2im.
        else:
            gW_mat = gW.reshape(out_c, reduce(operator.mul, ksize, c))
            col_mats = self.col.reshape(
                n, reduce(operator.mul, ksize, c), reduce(operator.mul, outs))
            gy_mats = gy.reshape(n, out_c, reduce(operator.mul, outs))
            # TODO(takagi): Use streams or batch gemm
            gW_mat[...] = 0
            for i in moves.range(n):
                gW_mat += cuda.cupy.dot(gy_mats[i], col_mats[i].T)

            # Compute patch array gradient.
            W_mat = W.reshape(out_c, -1)
            gcol = cuda.empty_like(self.col)
            gcol_mats = gcol.reshape(
                n, reduce(operator.mul, ksize, c), reduce(operator.mul, outs))
            for i in moves.range(n):
                gcol_mats[i] = cuda.cupy.dot(W_mat.T, gy_mats[i])

            # Compute input gradient.
            gx = conv_nd.col2im_nd_gpu(gcol, stride, pad, dims)

            # Compute bias gradient if given.
            if b is not None:
                # (n, _, out_1, out_2, ..., out_N)
                axis = (0,) + tuple(moves.range(2, ndim + 2))
                gb = gy.sum(axis=axis)

        if b is None:
            return gx, gW
        else:
            return gx, gW, gb

def convolution_nd(x, W, b=None, stride=1, pad=0, use_cudnn=True,
                   cover_all=False):
    """N-dimensional convolution function.

    This is an implementation of N-dimensional convolution which is generalized
    two-dimensional convolution in ConvNets. It takes three variables: the
    input image ``x``, the filter weight ``W`` and the bias vector ``b``.

    Notation: here is a notation for dimensionalities.

    - :math:`N` is the number of spatial dimensions.
    - :math:`n` is the batch size.
    - :math:`c_I` and :math:`c_O` are the number of the input and output,
      respectively.
    - :math:`d_1, d_2, ..., d_N` are the size of each axis of the input image,
      respectively.
    - :math:`k_1, k_2, ..., k_N` are the size of each axis of the filters,
      respectively.

    Args:
        x (~chainer.Variable): Input variable of shape
            :math:`(n, c_I, d_1, d_2, ..., d_N)`.
        W (~chainer.Variable): Weight variable of shape
            :math:`(c_O, c_I, k_1, k_2, ..., k_N)`.
        b (~chainer.Variable): Bias variable of length :math:`c_O`
            (optional).
        stride (int or tuple of ints): Stride of filter applications
            :math:`(s_1, s_2, ..., s_N)`. ``stride=s`` is equivalent to
            ``(s, s, ..., s)``.
        pad (int or tuple of ints): Spatial padding width for input arrays
            :math:`(p_1, p_2, ..., p_N)`. ``pad=p`` is equivalent to
            ``(p, p, ..., p)``.
        use_cudnn (bool): If ``True``, then this function uses cuDNN if
            available. cuDNN supports more than one-dimensional convolution.
        cover_all (bool): If ``True``, all spatial locations are convoluted
            into some output pixels. It may make the output size larger.

    Returns:
        ~chainer.Variable: Output variable.

    .. seealso:: :class:`ConvolutionND`, :func:`convolution_2d`
    """
    ndim = len(x.data.shape[2:])
    func = ConvolutionND(ndim, stride, pad, use_cudnn, cover_all)
    if b is None:
        return func(x, W)
    else:
        return func(x, W, b)
<|endoftext|>


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

import errno
import gevent.monkey
import gevent.pool
import netaddr
import prettytable
from pscan import exceptions as exc
import resource
import socket
gevent.monkey.patch_all()

soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (hard, hard))

min_port = 1
max_port = 65535
pool_size = soft

class Port(object):
    """Contains information for a specific port on a host."""

    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.status = "Closed"

class Host(object):
    """Contains information of a host and its ports to be scanned."""

    def __init__(self, ip, ports=[]):
        self.ip = netaddr.IPAddress(ip)
        self.ports = [Port(self, port) for port in ports]

    def get_ports(self):
        return self.ports

class Scan(object):
    """Creates the scanner object. Importable as a module."""

    def __init__(self, hosts, ports):
        self.pool = gevent.pool.Pool(pool_size)
        self.protocol = None
        self.inet = None
        self.result = prettytable.PrettyTable(["Port", "Protocol",
                                               "State", "Service"])
        if not ports:
            self.ports = range(min_port, max_port+1)
        else:
            try:
                self.ports = self._parse_ports(ports)
            except Exception:
                raise exc.InvalidPortRange
        try:
            self.hosts = self._parse_hosts(hosts)
        except Exception:
            raise exc.InvalidIPRange

    def _parse_hosts(self, hosts):
        """Parses the list of hosts and creates corresponding objects."""

        ip_list = [str(ip) for ip in list(netaddr.IPNetwork(hosts))]
        if netaddr.valid_ipv4(ip_list[0]):
            self.inet = socket.AF_INET
        elif netaddr.valid_ipv6(ip_list[0]):
            self.inet = socket.AF_INET6
        return [Host(ip, self.ports) for ip in ip_list]

    def _parse_ports(self, ports):
        """Parses the range of port numbers for validity."""

        ports = ports.split('-')
        minport = int(ports[0])
        maxport = minport
        if len(ports) == 2:
            maxport = int(ports[1])
        elif len(ports) > 2:
            raise
        if minport < min_port or maxport > max_port or minport > maxport:
            raise
        return range(minport, maxport+1)

    def _scan(self, f):
        for host in self.hosts:
            for port in host.get_ports():
                self.pool.spawn(f, port)
            self.pool.join()

    def _tcp(self, port):
        try:
            sock = socket.socket(self.inet, socket.SOCK_STREAM)
            sock.connect((str(port.host.ip), port.port))
        except Exception as e:
            if e.errno == errno.ECONNREFUSED:
                port.status = "Closed"
            else:
                print(e)
        else:
            port.status = "Open"
        sock.close()

    def tcp(self):
        """Performs a TCP scan."""

        self.protocol = "TCP"
        self._scan(self._tcp)

    def udp(self):
        """Performs a UDP scan."""

        self.protocol = "UDP"
        print("NOTE: The code for UDP Scanning is not complete."
              " Hence the scan result may not be correct.\n")
        # self._scan(self._udp)

    def _get_service(self, portnum):
        """Returns the service name for a given port."""

        try:
            service = socket.getservbyport(portnum)
        except Exception:
            service = "unknown"
        return service

    def show(self):
        """Displays the scanned results."""

        for host in self.hosts:
            print("Showing results for target: " + str(host.ip))
            single_port = False
            display_table = False
            ports = host.get_ports()
            if len(ports) == 1:
                single_port = True
                display_table = True
            for port in ports:
                if single_port:
                    self.result.add_row([port.port, self.protocol,
                                         port.status,
                                         self._get_service(port.port)])
                elif "Open" in port.status:
                    self.result.add_row([port.port, self.protocol,
                                         port.status,
                                         self._get_service(port.port)])
                    display_table = True
            if not display_table:
                print("All " + str(len(ports)) +
                      " scanned ports are closed on the target.")
            else:
                print(self.result)
            self.result.clear_rows()
<commit_msg>Setting timeout for socket<commit_after>#   Copyright 2016 Rudrajit Tapadar
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import errno
import gevent.monkey
import gevent.pool
import netaddr
import prettytable
from pscan import exceptions as exc
import resource
import socket
gevent.monkey.patch_all()

soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (hard, hard))

min_port = 1
max_port = 65535
pool_size = soft

class Port(object):
    """Contains information for a specific port on a host."""

    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.status = "Closed"

class Host(object):
    """Contains information of a host and its ports to be scanned."""

    def __init__(self, ip, ports=[]):
        self.ip = netaddr.IPAddress(ip)
        self.ports = [Port(self, port) for port in ports]

    def get_ports(self):
        return self.ports

class Scan(object):
    """Creates the scanner object. Importable as a module."""

    def __init__(self, hosts, ports):
        self.pool = gevent.pool.Pool(pool_size)
        self.protocol = None
        self.inet = None
        self.result = prettytable.PrettyTable(["Port", "Protocol",
                                               "State", "Service"])
        if not ports:
            self.ports = range(min_port, max_port+1)
        else:
            try:
                self.ports = self._parse_ports(ports)
            except Exception:
                raise exc.InvalidPortRange
        try:
            self.hosts = self._parse_hosts(hosts)
        except Exception:
            raise exc.InvalidIPRange

    def _parse_hosts(self, hosts):
        """Parses the list of hosts and creates corresponding objects."""

        ip_list = [str(ip) for ip in list(netaddr.IPNetwork(hosts))]
        if netaddr.valid_ipv4(ip_list[0]):
            self.inet = socket.AF_INET
        elif netaddr.valid_ipv6(ip_list[0]):
            self.inet = socket.AF_INET6
        return [Host(ip, self.ports) for ip in ip_list]

    def _parse_ports(self, ports):
        """Parses the range of port numbers for validity."""

        ports = ports.split('-')
        minport = int(ports[0])
        maxport = minport
        if len(ports) == 2:
            maxport = int(ports[1])
        elif len(ports) > 2:
            raise
        if minport < min_port or maxport > max_port or minport > maxport:
            raise
        return range(minport, maxport+1)

    def _scan(self, f):
        for host in self.hosts:
            for port in host.get_ports():
                self.pool.spawn(f, port)
            self.pool.join()

    def _tcp(self, port):
        try:
            sock = socket.socket(self.inet, socket.SOCK_STREAM)
            sock.settimeout(1)
            sock.connect((str(port.host.ip), port.port))
        except Exception as e:
            if e.errno == errno.ECONNREFUSED:
                port.status = "Closed"
            else:
                pass
        else:
            port.status = "Open"
        sock.close()

    def tcp(self):
        """Performs a TCP scan."""

        self.protocol = "TCP"
        self._scan(self._tcp)

    def udp(self):
        """Performs a UDP scan."""

        self.protocol = "UDP"
        print("NOTE: The code for UDP Scanning is not complete."
              " Hence the scan result may not be correct.\n")
        # self._scan(self._udp)

    def _get_service(self, portnum):
        """Returns the service name for a given port."""

        try:
            service = socket.getservbyport(portnum)
        except Exception:
            service = "unknown"
        return service

    def show(self):
        """Displays the scanned results."""

        for host in self.hosts:
            print("Showing results for target: " + str(host.ip))
            single_port = False
            display_table = False
            ports = host.get_ports()
            if len(ports) == 1:
                single_port = True
                display_table = True
            for port in ports:
                if single_port:
                    self.result.add_row([port.port, self.protocol,
                                         port.status,
                                         self._get_service(port.port)])
                elif "Open" in port.status:
                    self.result.add_row([port.port, self.protocol,
                                         port.status,
                                         self._get_service(port.port)])
                    display_table = True
            if not display_table:
                print("All " + str(len(ports)) +
                      " scanned ports are closed on the target.")
            else:
                print(self.result)
            self.result.clear_rows()
<|endoftext|>


<commit_before>import atexit
import errno
import fnmatch
import itertools
import operator
import os
import shutil
import sys
import uuid
import warnings
from os.path import expanduser
from os.path import expandvars
from os.path import isabs
from os.path import sep
from posixpath import sep as posix_sep

from _pytest.warning_types import PytestWarning

if sys.version_info[:2] >= (3, 6):
    from pathlib import Path, PurePath
else:
    from pathlib2 import Path, PurePath

__all__ = ["Path", "PurePath"]

LOCK_TIMEOUT = 60 * 60 * 3

get_lock_path = operator.methodcaller("joinpath", ".lock")

def ensure_reset_dir(path):
    """
    ensures the given path is an empty directory
    """
    if path.exists():
        rm_rf(path)
    path.mkdir()

def rm_rf(path):
    """Remove the path contents recursively, even if some elements
    are read-only.
    """

    def chmod_w(p):
        import stat

        mode = os.stat(str(p)).st_mode
        os.chmod(str(p), mode | stat.S_IWRITE)

    def force_writable_and_retry(function, p, excinfo):
        p = Path(p)

        # for files, we need to recursively go upwards
        # in the directories to ensure they all are also
        # writable
        if p.is_file():
            for parent in p.parents:
                chmod_w(parent)
                # stop when we reach the original path passed to rm_rf
                if parent == path:
                    break

        chmod_w(p)
        try:
            # retry the function that failed
            function(str(p))
        except Exception as e:
            warnings.warn(PytestWarning("(rm_rf) error removing {}: {}".format(p, e)))

    shutil.rmtree(str(path), onerror=force_writable_and_retry)

def find_prefixed(root, prefix):
    """finds all elements in root that begin with the prefix, case insensitive"""
    l_prefix = prefix.lower()
    for x in root.iterdir():
        if x.name.lower().startswith(l_prefix):
            yield x

def extract_suffixes(iter, prefix):
    """
    :param iter: iterator over path names
    :param prefix: expected prefix of the path names
    :returns: the parts of the paths following the prefix
    """
    p_len = len(prefix)
    for p in iter:
        yield p.name[p_len:]

def find_suffixes(root, prefix):
    """combines find_prefixes and extract_suffixes
    """
    return extract_suffixes(find_prefixed(root, prefix), prefix)

def parse_num(maybe_num):
    """parses number path suffixes, returns -1 on error"""
    try:
        return int(maybe_num)
    except ValueError:
        return -1

def _force_symlink(root, target, link_to):
    """helper to create the current symlink

    it's full of race conditions that are reasonably ok to ignore
    for the context of best effort linking to the latest testrun

    the presumption being thatin case of much parallelism
    the inaccuracy is going to be acceptable
    """
    current_symlink = root.joinpath(target)
    try:
        current_symlink.unlink()
    except OSError:
        pass
    try:
        current_symlink.symlink_to(link_to)
    except Exception:
        pass

def make_numbered_dir(root, prefix):
    """create a directory with an increased number as suffix for the given prefix"""
    for i in range(10):
        # try up to 10 times to create the folder
        max_existing = max(map(parse_num, find_suffixes(root, prefix)), default=-1)
        new_number = max_existing + 1
        new_path = root.joinpath("{}{}".format(prefix, new_number))
        try:
            new_path.mkdir()
        except Exception:
            pass
        else:
            _force_symlink(root, prefix + "current", new_path)
            return new_path
    else:
        raise EnvironmentError(
            "could not create numbered dir with prefix "
            "{prefix} in {root} after 10 tries".format(prefix=prefix, root=root)
        )

def create_cleanup_lock(p):
    """crates a lock to prevent premature folder cleanup"""
    lock_path = get_lock_path(p)
    try:
        fd = os.open(str(lock_path), os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o644)
    except OSError as e:
        if e.errno == errno.EEXIST:
            raise EnvironmentError(
                "cannot create lockfile in {path}".format(path=p)
            ) from e

        else:
            raise
    else:
        pid = os.getpid()
        spid = str(pid).encode()
        os.write(fd, spid)
        os.close(fd)
        if not lock_path.is_file():
            raise EnvironmentError("lock path got renamed after successful creation")
        return lock_path

def register_cleanup_lock_removal(lock_path, register=atexit.register):
    """registers a cleanup function for removing a lock, by default on atexit"""
    pid = os.getpid()

    def cleanup_on_exit(lock_path=lock_path, original_pid=pid):
        current_pid = os.getpid()
        if current_pid != original_pid:
            # fork
            return
        try:
            lock_path.unlink()
        except (OSError, IOError):
            pass

    return register(cleanup_on_exit)

def maybe_delete_a_numbered_dir(path):
    """removes a numbered directory if its lock can be obtained and it does not seem to be in use"""
    lock_path = None
    try:
        lock_path = create_cleanup_lock(path)
        parent = path.parent

        garbage = parent.joinpath("garbage-{}".format(uuid.uuid4()))
        path.rename(garbage)
        rm_rf(garbage)
    except (OSError, EnvironmentError):
        #  known races:
        #  * other process did a cleanup at the same time
        #  * deletable folder was found
        #  * process cwd (Windows)
        return
    finally:
        # if we created the lock, ensure we remove it even if we failed
        # to properly remove the numbered dir
        if lock_path is not None:
            try:
                lock_path.unlink()
            except (OSError, IOError):
                pass

def ensure_deletable(path, consider_lock_dead_if_created_before):
    """checks if a lock exists and breaks it if its considered dead"""
    if path.is_symlink():
        return False
    lock = get_lock_path(path)
    if not lock.exists():
        return True
    try:
        lock_time = lock.stat().st_mtime
    except Exception:
        return False
    else:
        if lock_time < consider_lock_dead_if_created_before:
            lock.unlink()
            return True
        else:
            return False

def try_cleanup(path, consider_lock_dead_if_created_before):
    """tries to cleanup a folder if we can ensure it's deletable"""
    if ensure_deletable(path, consider_lock_dead_if_created_before):
        maybe_delete_a_numbered_dir(path)

def cleanup_candidates(root, prefix, keep):
    """lists candidates for numbered directories to be removed - follows py.path"""
    max_existing = max(map(parse_num, find_suffixes(root, prefix)), default=-1)
    max_delete = max_existing - keep
    paths = find_prefixed(root, prefix)
    paths, paths2 = itertools.tee(paths)
    numbers = map(parse_num, extract_suffixes(paths2, prefix))
    for path, number in zip(paths, numbers):
        if number <= max_delete:
            yield path

def cleanup_numbered_dir(root, prefix, keep, consider_lock_dead_if_created_before):
    """cleanup for lock driven numbered directories"""
    for path in cleanup_candidates(root, prefix, keep):
        try_cleanup(path, consider_lock_dead_if_created_before)
    for path in root.glob("garbage-*"):
        try_cleanup(path, consider_lock_dead_if_created_before)

def make_numbered_dir_with_cleanup(root, prefix, keep, lock_timeout):
    """creates a numbered dir with a cleanup lock and removes old ones"""
    e = None
    for i in range(10):
        try:
            p = make_numbered_dir(root, prefix)
            lock_path = create_cleanup_lock(p)
            register_cleanup_lock_removal(lock_path)
        except Exception as exc:
            e = exc
        else:
            consider_lock_dead_if_created_before = p.stat().st_mtime - lock_timeout
            cleanup_numbered_dir(
                root=root,
                prefix=prefix,
                keep=keep,
                consider_lock_dead_if_created_before=consider_lock_dead_if_created_before,
            )
            return p
    assert e is not None
    raise e

def resolve_from_str(input, root):
    assert not isinstance(input, Path), "would break on py2"
    root = Path(root)
    input = expanduser(input)
    input = expandvars(input)
    if isabs(input):
        return Path(input)
    else:
        return root.joinpath(input)

def fnmatch_ex(pattern, path):
    """FNMatcher port from py.path.common which works with PurePath() instances.

    The difference between this algorithm and PurePath.match() is that the latter matches "**" glob expressions
    for each part of the path, while this algorithm uses the whole path instead.

    For example:
        "tests/foo/bar/doc/test_foo.py" matches pattern "tests/**/doc/test*.py" with this algorithm, but not with
        PurePath.match().

    This algorithm was ported to keep backward-compatibility with existing settings which assume paths match according
    this logic.

    References:
    * https://bugs.python.org/issue29249
    * https://bugs.python.org/issue34731
    """
    path = PurePath(path)
    iswin32 = sys.platform.startswith("win")

    if iswin32 and sep not in pattern and posix_sep in pattern:
        # Running on Windows, the pattern has no Windows path separators,
        # and the pattern has one or more Posix path separators. Replace
        # the Posix path separators with the Windows path separator.
        pattern = pattern.replace(posix_sep, sep)

    if sep not in pattern:
        name = path.name
    else:
        name = str(path)
        if path.is_absolute() and not os.path.isabs(pattern):
            pattern = "*{}{}".format(os.sep, pattern)
    return fnmatch.fnmatch(name, pattern)

def parts(s):
    parts = s.split(sep)
    return {sep.join(parts[: i + 1]) or sep for i in range(len(parts))}
<commit_msg>Fix typos<commit_after>import atexit
import errno
import fnmatch
import itertools
import operator
import os
import shutil
import sys
import uuid
import warnings
from os.path import expanduser
from os.path import expandvars
from os.path import isabs
from os.path import sep
from posixpath import sep as posix_sep

from _pytest.warning_types import PytestWarning

if sys.version_info[:2] >= (3, 6):
    from pathlib import Path, PurePath
else:
    from pathlib2 import Path, PurePath

__all__ = ["Path", "PurePath"]

LOCK_TIMEOUT = 60 * 60 * 3

get_lock_path = operator.methodcaller("joinpath", ".lock")

def ensure_reset_dir(path):
    """
    ensures the given path is an empty directory
    """
    if path.exists():
        rm_rf(path)
    path.mkdir()

def rm_rf(path):
    """Remove the path contents recursively, even if some elements
    are read-only.
    """

    def chmod_w(p):
        import stat

        mode = os.stat(str(p)).st_mode
        os.chmod(str(p), mode | stat.S_IWRITE)

    def force_writable_and_retry(function, p, excinfo):
        p = Path(p)

        # for files, we need to recursively go upwards
        # in the directories to ensure they all are also
        # writable
        if p.is_file():
            for parent in p.parents:
                chmod_w(parent)
                # stop when we reach the original path passed to rm_rf
                if parent == path:
                    break

        chmod_w(p)
        try:
            # retry the function that failed
            function(str(p))
        except Exception as e:
            warnings.warn(PytestWarning("(rm_rf) error removing {}: {}".format(p, e)))

    shutil.rmtree(str(path), onerror=force_writable_and_retry)

def find_prefixed(root, prefix):
    """finds all elements in root that begin with the prefix, case insensitive"""
    l_prefix = prefix.lower()
    for x in root.iterdir():
        if x.name.lower().startswith(l_prefix):
            yield x

def extract_suffixes(iter, prefix):
    """
    :param iter: iterator over path names
    :param prefix: expected prefix of the path names
    :returns: the parts of the paths following the prefix
    """
    p_len = len(prefix)
    for p in iter:
        yield p.name[p_len:]

def find_suffixes(root, prefix):
    """combines find_prefixes and extract_suffixes
    """
    return extract_suffixes(find_prefixed(root, prefix), prefix)

def parse_num(maybe_num):
    """parses number path suffixes, returns -1 on error"""
    try:
        return int(maybe_num)
    except ValueError:
        return -1

def _force_symlink(root, target, link_to):
    """helper to create the current symlink

    it's full of race conditions that are reasonably ok to ignore
    for the context of best effort linking to the latest test run

    the presumption being that in case of much parallelism
    the inaccuracy is going to be acceptable
    """
    current_symlink = root.joinpath(target)
    try:
        current_symlink.unlink()
    except OSError:
        pass
    try:
        current_symlink.symlink_to(link_to)
    except Exception:
        pass

def make_numbered_dir(root, prefix):
    """create a directory with an increased number as suffix for the given prefix"""
    for i in range(10):
        # try up to 10 times to create the folder
        max_existing = max(map(parse_num, find_suffixes(root, prefix)), default=-1)
        new_number = max_existing + 1
        new_path = root.joinpath("{}{}".format(prefix, new_number))
        try:
            new_path.mkdir()
        except Exception:
            pass
        else:
            _force_symlink(root, prefix + "current", new_path)
            return new_path
    else:
        raise EnvironmentError(
            "could not create numbered dir with prefix "
            "{prefix} in {root} after 10 tries".format(prefix=prefix, root=root)
        )

def create_cleanup_lock(p):
    """crates a lock to prevent premature folder cleanup"""
    lock_path = get_lock_path(p)
    try:
        fd = os.open(str(lock_path), os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o644)
    except OSError as e:
        if e.errno == errno.EEXIST:
            raise EnvironmentError(
                "cannot create lockfile in {path}".format(path=p)
            ) from e

        else:
            raise
    else:
        pid = os.getpid()
        spid = str(pid).encode()
        os.write(fd, spid)
        os.close(fd)
        if not lock_path.is_file():
            raise EnvironmentError("lock path got renamed after successful creation")
        return lock_path

def register_cleanup_lock_removal(lock_path, register=atexit.register):
    """registers a cleanup function for removing a lock, by default on atexit"""
    pid = os.getpid()

    def cleanup_on_exit(lock_path=lock_path, original_pid=pid):
        current_pid = os.getpid()
        if current_pid != original_pid:
            # fork
            return
        try:
            lock_path.unlink()
        except (OSError, IOError):
            pass

    return register(cleanup_on_exit)

def maybe_delete_a_numbered_dir(path):
    """removes a numbered directory if its lock can be obtained and it does not seem to be in use"""
    lock_path = None
    try:
        lock_path = create_cleanup_lock(path)
        parent = path.parent

        garbage = parent.joinpath("garbage-{}".format(uuid.uuid4()))
        path.rename(garbage)
        rm_rf(garbage)
    except (OSError, EnvironmentError):
        #  known races:
        #  * other process did a cleanup at the same time
        #  * deletable folder was found
        #  * process cwd (Windows)
        return
    finally:
        # if we created the lock, ensure we remove it even if we failed
        # to properly remove the numbered dir
        if lock_path is not None:
            try:
                lock_path.unlink()
            except (OSError, IOError):
                pass

def ensure_deletable(path, consider_lock_dead_if_created_before):
    """checks if a lock exists and breaks it if its considered dead"""
    if path.is_symlink():
        return False
    lock = get_lock_path(path)
    if not lock.exists():
        return True
    try:
        lock_time = lock.stat().st_mtime
    except Exception:
        return False
    else:
        if lock_time < consider_lock_dead_if_created_before:
            lock.unlink()
            return True
        else:
            return False

def try_cleanup(path, consider_lock_dead_if_created_before):
    """tries to cleanup a folder if we can ensure it's deletable"""
    if ensure_deletable(path, consider_lock_dead_if_created_before):
        maybe_delete_a_numbered_dir(path)

def cleanup_candidates(root, prefix, keep):
    """lists candidates for numbered directories to be removed - follows py.path"""
    max_existing = max(map(parse_num, find_suffixes(root, prefix)), default=-1)
    max_delete = max_existing - keep
    paths = find_prefixed(root, prefix)
    paths, paths2 = itertools.tee(paths)
    numbers = map(parse_num, extract_suffixes(paths2, prefix))
    for path, number in zip(paths, numbers):
        if number <= max_delete:
            yield path

def cleanup_numbered_dir(root, prefix, keep, consider_lock_dead_if_created_before):
    """cleanup for lock driven numbered directories"""
    for path in cleanup_candidates(root, prefix, keep):
        try_cleanup(path, consider_lock_dead_if_created_before)
    for path in root.glob("garbage-*"):
        try_cleanup(path, consider_lock_dead_if_created_before)

def make_numbered_dir_with_cleanup(root, prefix, keep, lock_timeout):
    """creates a numbered dir with a cleanup lock and removes old ones"""
    e = None
    for i in range(10):
        try:
            p = make_numbered_dir(root, prefix)
            lock_path = create_cleanup_lock(p)
            register_cleanup_lock_removal(lock_path)
        except Exception as exc:
            e = exc
        else:
            consider_lock_dead_if_created_before = p.stat().st_mtime - lock_timeout
            cleanup_numbered_dir(
                root=root,
                prefix=prefix,
                keep=keep,
                consider_lock_dead_if_created_before=consider_lock_dead_if_created_before,
            )
            return p
    assert e is not None
    raise e

def resolve_from_str(input, root):
    assert not isinstance(input, Path), "would break on py2"
    root = Path(root)
    input = expanduser(input)
    input = expandvars(input)
    if isabs(input):
        return Path(input)
    else:
        return root.joinpath(input)

def fnmatch_ex(pattern, path):
    """FNMatcher port from py.path.common which works with PurePath() instances.

    The difference between this algorithm and PurePath.match() is that the latter matches "**" glob expressions
    for each part of the path, while this algorithm uses the whole path instead.

    For example:
        "tests/foo/bar/doc/test_foo.py" matches pattern "tests/**/doc/test*.py" with this algorithm, but not with
        PurePath.match().

    This algorithm was ported to keep backward-compatibility with existing settings which assume paths match according
    this logic.

    References:
    * https://bugs.python.org/issue29249
    * https://bugs.python.org/issue34731
    """
    path = PurePath(path)
    iswin32 = sys.platform.startswith("win")

    if iswin32 and sep not in pattern and posix_sep in pattern:
        # Running on Windows, the pattern has no Windows path separators,
        # and the pattern has one or more Posix path separators. Replace
        # the Posix path separators with the Windows path separator.
        pattern = pattern.replace(posix_sep, sep)

    if sep not in pattern:
        name = path.name
    else:
        name = str(path)
        if path.is_absolute() and not os.path.isabs(pattern):
            pattern = "*{}{}".format(os.sep, pattern)
    return fnmatch.fnmatch(name, pattern)

def parts(s):
    parts = s.split(sep)
    return {sep.join(parts[: i + 1]) or sep for i in range(len(parts))}
<|endoftext|>


<commit_before>"""
Replaces page names in the links file with their corresponding IDs, eliminates links containing
non-existing pages, and replaces redirects with the pages to which they redirect.

Output is written to stdout.
"""

from __future__ import print_function

import io
import sys
import gzip
from sets import Set

# Validate inputs
if len(sys.argv) < 4:
  print('[ERROR] Not enough arguments provided!')
  print('[INFO] Usage: {0} <pages_file> <redirects_file> <links_file>'.format(sys.argv[0]))
  sys.exit()

PAGES_FILE = sys.argv[1]
REDIRECTS_FILE = sys.argv[2]
LINKS_FILE = sys.argv[3]

if not PAGES_FILE.endswith('.gz'):
  print('[ERROR] Pages file must be gzipped.')
  sys.exit()

if not REDIRECTS_FILE.endswith('.gz'):
  print('[ERROR] Redirects file must be gzipped.')
  sys.exit()

if not LINKS_FILE.endswith('.gz'):
  print('[ERROR] Links file must be gzipped.')
  sys.exit()

# Create a set of all page IDs and a dictionary of page titles to their corresponding IDs.
ALL_PAGE_IDS = Set()
PAGE_TITLES_TO_IDS = {}
for line in io.BufferedReader(gzip.open(PAGES_FILE, 'r')):
  [page_id, page_title, is_redirect] = line.rstrip('\n').split('\t')
  ALL_PAGE_IDS.add(page_id)
  PAGE_TITLES_TO_IDS[page_title] = page_id

# Create a dictionary of page IDs to the target page ID to which they redirect.
REDIRECTS = {}
for line in io.BufferedReader(gzip.open(REDIRECTS_FILE, 'r')):
  [source_page_id, target_page_id] = line.rstrip('\n').split('\t')
  REDIRECTS[source_page_id] = target_page_id

# Loop through each line in the links file, replacing titles with IDs, applying redirects, and
# removing nonexistent pages, writing the result to stdout.
for line in io.BufferedReader(gzip.open(LINKS_FILE, 'r')):
  [source_page_id, target_page_title] = line.rstrip('\n').split('\t')

  source_page_exists = source_page_id in ALL_PAGE_IDS

  if source_page_exists:
    source_page_id = REDIRECTS.get(source_page_id, source_page_id)

    target_page_id = PAGE_TITLES_TO_IDS.get(target_page_title)

    if target_page_id is not None and source_page_id != target_page_id:
      target_page_id = REDIRECTS.get(target_page_id, target_page_id)
      print('\t'.join([source_page_id, target_page_id]))
<commit_msg>Rename unused variable to underscore<commit_after>"""
Replaces page names in the links file with their corresponding IDs, eliminates links containing
non-existing pages, and replaces redirects with the pages to which they redirect.

Output is written to stdout.
"""

from __future__ import print_function

import io
import sys
import gzip
from sets import Set

# Validate inputs
if len(sys.argv) < 4:
  print('[ERROR] Not enough arguments provided!')
  print('[INFO] Usage: {0} <pages_file> <redirects_file> <links_file>'.format(sys.argv[0]))
  sys.exit()

PAGES_FILE = sys.argv[1]
REDIRECTS_FILE = sys.argv[2]
LINKS_FILE = sys.argv[3]

if not PAGES_FILE.endswith('.gz'):
  print('[ERROR] Pages file must be gzipped.')
  sys.exit()

if not REDIRECTS_FILE.endswith('.gz'):
  print('[ERROR] Redirects file must be gzipped.')
  sys.exit()

if not LINKS_FILE.endswith('.gz'):
  print('[ERROR] Links file must be gzipped.')
  sys.exit()

# Create a set of all page IDs and a dictionary of page titles to their corresponding IDs.
ALL_PAGE_IDS = Set()
PAGE_TITLES_TO_IDS = {}
for line in io.BufferedReader(gzip.open(PAGES_FILE, 'r')):
  [page_id, page_title, _] = line.rstrip('\n').split('\t')
  ALL_PAGE_IDS.add(page_id)
  PAGE_TITLES_TO_IDS[page_title] = page_id

# Create a dictionary of page IDs to the target page ID to which they redirect.
REDIRECTS = {}
for line in io.BufferedReader(gzip.open(REDIRECTS_FILE, 'r')):
  [source_page_id, target_page_id] = line.rstrip('\n').split('\t')
  REDIRECTS[source_page_id] = target_page_id

# Loop through each line in the links file, replacing titles with IDs, applying redirects, and
# removing nonexistent pages, writing the result to stdout.
for line in io.BufferedReader(gzip.open(LINKS_FILE, 'r')):
  [source_page_id, target_page_title] = line.rstrip('\n').split('\t')

  source_page_exists = source_page_id in ALL_PAGE_IDS

  if source_page_exists:
    source_page_id = REDIRECTS.get(source_page_id, source_page_id)

    target_page_id = PAGE_TITLES_TO_IDS.get(target_page_title)

    if target_page_id is not None and source_page_id != target_page_id:
      target_page_id = REDIRECTS.get(target_page_id, target_page_id)
      print('\t'.join([source_page_id, target_page_id]))
<|endoftext|>


<commit_before>#
#  Copyright 2019 The FATE Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import os
import shutil
from ruamel import yaml
from datetime import datetime

from fate_flow.db.db_models import DB, MachineLearningModelInfo as MLModel
from fate_flow.pipelined_model import pipelined_model
from fate_arch.common.base_utils import json_loads, json_dumps
from fate_arch.common.file_utils import get_project_base_directory
from fate_flow.utils import model_utils
from federatedml.protobuf.model_migrate.model_migrate import model_migration

def gen_model_file_path(model_id, model_version):
    return os.path.join(get_project_base_directory(), "model_local_cache", model_id, model_version)

def compare_roles(request_conf_roles: dict, run_time_conf_roles: dict):
    if request_conf_roles.keys() == run_time_conf_roles.keys():
        varify_format = True
        varify_equality = True
        for key in request_conf_roles.keys():
            varify_format = varify_format and (len(request_conf_roles[key]) == len(run_time_conf_roles[key])) and (isinstance(request_conf_roles[key], list))
            request_conf_roles_set = set(str(item) for item in request_conf_roles[key])
            run_time_conf_roles_set = set(str(item) for item in run_time_conf_roles[key])
            varify_equality = varify_equality and (request_conf_roles_set == run_time_conf_roles_set)
        if not varify_format:
            raise Exception("The structure of roles data of local configuration is different from "
                            "model runtime configuration's. Migration aborting.")
        else:
            return varify_equality
    raise Exception("The structure of roles data of local configuration is different from "
                    "model runtime configuration's. Migration aborting.")

def import_from_files(config: dict):
    model = pipelined_model.PipelinedModel(model_id=config["model_id"],
                                           model_version=config["model_version"])
    if config['force']:
        model.force = True
    model.unpack_model(config["file"])

def import_from_db(config: dict):
    model_path = gen_model_file_path(config["model_id"], config["model_version"])
    if config['force']:
        os.rename(model_path, model_path + '_backup_{}'.format(datetime.now().strftime('%Y%m%d%H%M')))

def migration(config_data: dict):
    try:
        party_model_id = model_utils.gen_party_model_id(model_id=config_data["model_id"],
                                                        role=config_data["local"]["role"],
                                                        party_id=config_data["local"]["party_id"])
        model = pipelined_model.PipelinedModel(model_id=party_model_id,
                                               model_version=config_data["model_version"])
        if not model.exists():
            raise Exception("Can not found {} {} model local cache".format(config_data["model_id"],
                                                                           config_data["model_version"]))
        with DB.connection_context():
            if MLModel.get_or_none(MLModel.f_model_version == config_data["unify_model_version"]):
                raise Exception("Unify model version {} has been occupied in database. "
                                "Please choose another unify model version and try again.".format(
                    config_data["unify_model_version"]))

        model_data = model.collect_models(in_bytes=True)
        if "pipeline.pipeline:Pipeline" not in model_data:
            raise Exception("Can not found pipeline file in model.")

        migrate_model = pipelined_model.PipelinedModel(model_id=model_utils.gen_party_model_id(model_id=model_utils.gen_model_id(config_data["migrate_role"]),
                                                                                role=config_data["local"]["role"],
                                                                                party_id=config_data["local"]["migrate_party_id"]),
                                                       model_version=config_data["unify_model_version"])

        # migrate_model.create_pipelined_model()
        shutil.copytree(src=model.model_path, dst=migrate_model.model_path)

        pipeline = migrate_model.read_component_model('pipeline', 'pipeline')['Pipeline']

        # Utilize Pipeline_model collect model data. And modify related inner information of model
        train_runtime_conf = json_loads(pipeline.train_runtime_conf)
        train_runtime_conf["role"] = config_data["migrate_role"]
        train_runtime_conf["job_parameters"]["model_id"] = model_utils.gen_model_id(train_runtime_conf["role"])
        train_runtime_conf["job_parameters"]["model_version"] = migrate_model.model_version
        train_runtime_conf["initiator"] = config_data["migrate_initiator"]

        # update pipeline.pb file
        pipeline.train_runtime_conf = json_dumps(train_runtime_conf, byte=True)
        pipeline.model_id = bytes(train_runtime_conf["job_parameters"]["model_id"], "utf-8")
        pipeline.model_version = bytes(train_runtime_conf["job_parameters"]["model_version"], "utf-8")

        # save updated pipeline.pb file
        migrate_model.save_pipeline(pipeline)
        shutil.copyfile(os.path.join(migrate_model.model_path, "pipeline.pb"),
                        os.path.join(migrate_model.model_path, "variables", "data", "pipeline", "pipeline", "Pipeline"))

        # modify proto
        with open(os.path.join(migrate_model.model_path, 'define', 'define-meta.yaml'), 'r') as fin:
            define_yaml = yaml.safe_load(fin)

        for key, value in define_yaml['model_proto'].items():
            if key == 'pipeline':
                continue
            for v in value.keys():
                buffer_obj = migrate_model.read_component_model(key, v)
                module_name = define_yaml['component_define'].get(key, {}).get('module_name')
                modified_buffer = model_migration(model_contents=buffer_obj,
                                                  module_name=module_name,
                                                  old_guest_list=config_data['role']['guest'],
                                                  new_guest_list=config_data['migrate_role']['guest'],
                                                  old_host_list=config_data['role']['host'],
                                                  new_host_list=config_data['migrate_role']['host'],
                                                  old_arbiter_list=config_data.get('role', {}).get('arbiter', None),
                                                  new_arbiter_list=config_data.get('migrate_role', {}).get('arbiter', None))
                migrate_model.save_component_model(component_name=key, component_module_name=module_name,
                                                   model_alias=v, model_buffers=modified_buffer)

        archive_path = migrate_model.packaging_model()

        # with DB.connection_context():
        #     MLModel.create(
        #         f_role=config_data["local"]["role"],
        #         f_party_id=config_data["local"]["migrate_party_id"],
        #         f_roles=config_data["migrate_role"],
        #         f_job_id=train_runtime_conf["job_parameters"]["model_version"],
        #         f_model_id=train_runtime_conf["job_parameters"]["model_id"],
        #         f_model_version=train_runtime_conf["job_parameters"]["model_version"],
        #         f_initiator_role=config_data["migrate_initiator"]["role"],
        #         f_initiator_party_id=config_data["migrate_initiator"]["party_id"],
        #         f_runtime_conf=train_runtime_conf,
        #         f_work_mode=train_runtime_conf["job_parameters"]["work_mode"],
        #         f_dsl=json_loads(pipeline.train_dsl),
        #         f_migrated=1,
        #         f_job_status='complete'
        #     )

        return (0, f"Migrating model successfully. " \
                  "The configuration of model has been modified automatically. " \
                  "New model id is: {}, model version is: {}. " \
                  "Model files can be found at '{}'.".format(train_runtime_conf["job_parameters"]["model_id"],
                                                             migrate_model.model_version,
                                                             os.path.abspath(archive_path)),
                {"model_id": migrate_model.model_id,
                 "model_version": migrate_model.model_version,
                 "path": os.path.abspath(archive_path)})

    except Exception as e:
        return 100, str(e), {}
<commit_msg>to dash<commit_after>#
#  Copyright 2019 The FATE Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import os
import shutil
from ruamel import yaml
from datetime import datetime

from fate_flow.db.db_models import DB, MachineLearningModelInfo as MLModel
from fate_flow.pipelined_model import pipelined_model
from fate_arch.common.base_utils import json_loads, json_dumps
from fate_arch.common.file_utils import get_project_base_directory
from fate_flow.utils import model_utils
from federatedml.protobuf.model_migrate.model_migrate import model_migration

def gen_model_file_path(model_id, model_version):
    return os.path.join(get_project_base_directory(), "model_local_cache", model_id, model_version)

def compare_roles(request_conf_roles: dict, run_time_conf_roles: dict):
    if request_conf_roles.keys() == run_time_conf_roles.keys():
        varify_format = True
        varify_equality = True
        for key in request_conf_roles.keys():
            varify_format = varify_format and (len(request_conf_roles[key]) == len(run_time_conf_roles[key])) and (isinstance(request_conf_roles[key], list))
            request_conf_roles_set = set(str(item) for item in request_conf_roles[key])
            run_time_conf_roles_set = set(str(item) for item in run_time_conf_roles[key])
            varify_equality = varify_equality and (request_conf_roles_set == run_time_conf_roles_set)
        if not varify_format:
            raise Exception("The structure of roles data of local configuration is different from "
                            "model runtime configuration's. Migration aborting.")
        else:
            return varify_equality
    raise Exception("The structure of roles data of local configuration is different from "
                    "model runtime configuration's. Migration aborting.")

def import_from_files(config: dict):
    model = pipelined_model.PipelinedModel(model_id=config["model_id"],
                                           model_version=config["model_version"])
    if config['force']:
        model.force = True
    model.unpack_model(config["file"])

def import_from_db(config: dict):
    model_path = gen_model_file_path(config["model_id"], config["model_version"])
    if config['force']:
        os.rename(model_path, model_path + '_backup_{}'.format(datetime.now().strftime('%Y%m%d%H%M')))

def migration(config_data: dict):
    try:
        party_model_id = model_utils.gen_party_model_id(model_id=config_data["model_id"],
                                                        role=config_data["local"]["role"],
                                                        party_id=config_data["local"]["party_id"])
        model = pipelined_model.PipelinedModel(model_id=party_model_id,
                                               model_version=config_data["model_version"])
        if not model.exists():
            raise Exception("Can not found {} {} model local cache".format(config_data["model_id"],
                                                                           config_data["model_version"]))
        with DB.connection_context():
            if MLModel.get_or_none(MLModel.f_model_version == config_data["unify_model_version"]):
                raise Exception("Unify model version {} has been occupied in database. "
                                "Please choose another unify model version and try again.".format(
                    config_data["unify_model_version"]))

        model_data = model.collect_models(in_bytes=True)
        if "pipeline.pipeline:Pipeline" not in model_data:
            raise Exception("Can not found pipeline file in model.")

        migrate_model = pipelined_model.PipelinedModel(model_id=model_utils.gen_party_model_id(model_id=model_utils.gen_model_id(config_data["migrate_role"]),
                                                                                role=config_data["local"]["role"],
                                                                                party_id=config_data["local"]["migrate_party_id"]),
                                                       model_version=config_data["unify_model_version"])

        # migrate_model.create_pipelined_model()
        shutil.copytree(src=model.model_path, dst=migrate_model.model_path)

        pipeline = migrate_model.read_component_model('pipeline', 'pipeline')['Pipeline']

        # Utilize Pipeline_model collect model data. And modify related inner information of model
        train_runtime_conf = json_loads(pipeline.train_runtime_conf)
        train_runtime_conf["role"] = config_data["migrate_role"]
        train_runtime_conf["job_parameters"]["model_id"] = model_utils.gen_model_id(train_runtime_conf["role"])
        train_runtime_conf["job_parameters"]["model_version"] = migrate_model.model_version
        train_runtime_conf["initiator"] = config_data["migrate_initiator"]

        # update pipeline.pb file
        pipeline.train_runtime_conf = json_dumps(train_runtime_conf, byte=True)
        pipeline.model_id = bytes(train_runtime_conf["job_parameters"]["model_id"], "utf-8")
        pipeline.model_version = bytes(train_runtime_conf["job_parameters"]["model_version"], "utf-8")

        # save updated pipeline.pb file
        migrate_model.save_pipeline(pipeline)
        shutil.copyfile(os.path.join(migrate_model.model_path, "pipeline.pb"),
                        os.path.join(migrate_model.model_path, "variables", "data", "pipeline", "pipeline", "Pipeline"))

        # modify proto
        with open(os.path.join(migrate_model.model_path, 'define', 'define_meta.yaml'), 'r') as fin:
            define_yaml = yaml.safe_load(fin)

        for key, value in define_yaml['model_proto'].items():
            if key == 'pipeline':
                continue
            for v in value.keys():
                buffer_obj = migrate_model.read_component_model(key, v)
                module_name = define_yaml['component_define'].get(key, {}).get('module_name')
                modified_buffer = model_migration(model_contents=buffer_obj,
                                                  module_name=module_name,
                                                  old_guest_list=config_data['role']['guest'],
                                                  new_guest_list=config_data['migrate_role']['guest'],
                                                  old_host_list=config_data['role']['host'],
                                                  new_host_list=config_data['migrate_role']['host'],
                                                  old_arbiter_list=config_data.get('role', {}).get('arbiter', None),
                                                  new_arbiter_list=config_data.get('migrate_role', {}).get('arbiter', None))
                migrate_model.save_component_model(component_name=key, component_module_name=module_name,
                                                   model_alias=v, model_buffers=modified_buffer)

        archive_path = migrate_model.packaging_model()

        # with DB.connection_context():
        #     MLModel.create(
        #         f_role=config_data["local"]["role"],
        #         f_party_id=config_data["local"]["migrate_party_id"],
        #         f_roles=config_data["migrate_role"],
        #         f_job_id=train_runtime_conf["job_parameters"]["model_version"],
        #         f_model_id=train_runtime_conf["job_parameters"]["model_id"],
        #         f_model_version=train_runtime_conf["job_parameters"]["model_version"],
        #         f_initiator_role=config_data["migrate_initiator"]["role"],
        #         f_initiator_party_id=config_data["migrate_initiator"]["party_id"],
        #         f_runtime_conf=train_runtime_conf,
        #         f_work_mode=train_runtime_conf["job_parameters"]["work_mode"],
        #         f_dsl=json_loads(pipeline.train_dsl),
        #         f_migrated=1,
        #         f_job_status='complete'
        #     )

        return (0, f"Migrating model successfully. " \
                  "The configuration of model has been modified automatically. " \
                  "New model id is: {}, model version is: {}. " \
                  "Model files can be found at '{}'.".format(train_runtime_conf["job_parameters"]["model_id"],
                                                             migrate_model.model_version,
                                                             os.path.abspath(archive_path)),
                {"model_id": migrate_model.model_id,
                 "model_version": migrate_model.model_version,
                 "path": os.path.abspath(archive_path)})

    except Exception as e:
        return 100, str(e), {}
<|endoftext|>