<commit_before>#!/usr/bin/env python
# (Python >= 2.6 including >= 3)
docs = """
Usage: ./dev-build.py [cmake args] [certain make args]

This script wraps `cmake` and `make`; you can use CMake directly
or use this, whichever you prefer.  This script takes `cmake` and
some `make` flags directly on the command-line, puts the build in
a directory under build/ that is determined by which `cmake` flags you
specified (so that you can store separate intermediate compilation results
for various flag/compiler combinations), runs Lasercake's self-tests,
and if they're successful, copies Lasercake to the main directory.
"""

import os, os.path, sys, subprocess, re, shutil, hashlib, datetime

try:
	import multiprocessing
	MAKE_PARALLEL_JOBS_DEFAULT = multiprocessing.cpu_count() + 1
except (ImportError,NotImplementedError):
	MAKE_PARALLEL_JOBS_DEFAULT = 2

ansi_grey = '\033[30m'
ansi_red = '\033[31m'
ansi_green = '\033[32m'
ansi_yellow = '\033[33m'
ansi_blue = '\033[34m'
ansi_magenta = '\033[35m'
ansi_cyan = '\033[36m'
ansi_white = '\033[37m'
ansi_end = '\033[0m'

os.environ['PATH'] = '/usr/lib/colorgcc/bin:'+os.environ['PATH']

def say(string):
	sys.stdout.write(string)
	sys.stdout.flush()

def escaped_control_characters(string):
	return re.sub(r"[\x01-\x1F\x7F]", r"\\", string)

def say_we_are_calling(string):
	say(ansi_cyan+'% '+escaped_control_characters(string)+ansi_end+'\n')

def hash_list(l):
	hashes = (hashlib.sha256(arg.encode('utf8')).digest() for arg in l)
	return hashlib.sha256(b''.join(hashes)).hexdigest()

def replace_app(src, dst):
	"""
	Copies src to dst, replacing dst, even if dst is a running executable.
	Complications this deals with:
	On Windows, one can move but not remove or edit a running binary.
	On Linux, one can move and unlink but not edit a running binary.
	On Mac, if using .app bundles, the app is a directory rather than a file.
	"""
	old = None
	if os.path.exists(dst):
		unique_str = datetime.datetime.fromtimestamp(
				os.stat(dst).st_mtime).isoformat()
		path, basefile = os.path.split(dst)
		base, ext = os.path.splitext(basefile)
		old = os.path.join(path, base+'.'+unique_str+ext)
		os.rename(dst, old)
	if os.path.isdir(src):
		shutil.copytree(src, dst, symlinks=True)
	else:
		shutil.copy2(src, dst)
	if old != None:
		if os.path.isdir(old):
			shutil.rmtree(old)
		else:
			try: os.unlink(old)
			except OSError:
				say("Couldn't remove old (running?) executable "+os.path.basename(old))

def main():
	os.chdir(os.path.dirname(os.path.join('.', __file__)))
	try: subprocess.call(['cmake', '--version'])
	except OSError:
		say(ansi_red+"Error: 'cmake' not found; please install it."+ansi_end+'\n')
		exit(1)
	try: subprocess.call(['make', '--version'])
	except OSError:
		say(ansi_red+"Error: 'make' not found; please install it."+ansi_end+'\n')
		exit(1)
	try:
		status = subprocess.call(['time', '-f', '', 'true'])
		time_has_f = (status == 0)
		has_time = True
	except OSError:
		has_time = False
	cmake_args = []
	make_args = []
	making_lasercake = True
	running_tests = True
	make_jobs_flags = ['-j'+str(MAKE_PARALLEL_JOBS_DEFAULT)]
	for arg in sys.argv[1:]:
		if arg in ['-?', '-h', '--help', '-help']:
			say(docs)
			sys.exit()
		elif arg == 'no-test':
			running_tests = False
		elif arg == 'VERBOSE=1':
			make_args.append(arg)
		elif re.search(r'^-j[0-9]*$', arg):
			make_jobs_flags = [arg]
		else:
			cmake_args.append(arg)
		if arg == '-DBUILD_SELF_TESTS=OFF':
			running_tests = False
	# Make it more likely the directories are named non-conflictingly
	# even though we delete some special characters from their paths.
	hash_for_this_config = hash_list(cmake_args)
	name_for_this_config = hash_for_this_config[:10] + '_' + (','.join([re.sub(r'[^-+.a-zA-Z0-9]', r'_', arg) for arg in cmake_args]) or 'default')
	build_dir = 'build/'+name_for_this_config
	try: os.remove('CMakeCache.txt')
	except OSError: pass
	try:
		os.makedirs(build_dir)
	except OSError:
		say(ansi_green+'This build dir already exists!'+ansi_end+'\n')
		cmake_args = []
	to_call_cmake = ['cmake', '../../'] + cmake_args
	say_we_are_calling('cd '+build_dir+'; '+'   '.join(to_call_cmake))
	say(ansi_cyan+'''  (^^ not escaped properly in these info messages - doin' it right in python)'''+ansi_end+'\n')
	os.chdir(build_dir)
	subprocess.check_call(to_call_cmake)
	if not has_time:
		time_args = []
	elif time_has_f:
		time_args = ['time', '-f', (ansi_magenta+'`make` took %E'+ansi_end)]
	else:
		time_args = ['time']
	to_call_make = time_args + ['make'] + make_jobs_flags + make_args
	say_we_are_calling('   '.join(to_call_make))
	make_status = subprocess.call(to_call_make)
	if has_time and not time_has_f:
		say(ansi_magenta+'  (time to run `make`).'+ansi_end+'\n')
	# How to print the compiler & flags here? Can we get it from cmake
	# somehow? Add to the cmakelists to write those?
	# but with the -l's and all... hmm.
	if make_status != 0:
		# skip make_jobs_flags (-jN) to get just one file's error
		# messages not mixed up with the others
		to_call_make_again = ['make'] + make_args
		say_we_are_calling('   '.join(to_call_make_again))
		subprocess.call(to_call_make_again)
		say(ansi_red+'build failed'+ansi_end+'\n')
		exit(1)
	is_windows_exe = is_apple_app = False
	app_name = exe_name = None
	if os.path.exists('Lasercake'):
		app_name = exe_name = 'Lasercake'
	if os.path.exists('Lasercake.exe'):
		app_name = exe_name = 'Lasercake.exe'
		is_windows_exe = True
	if os.path.exists('Lasercake.app'):
		app_name = 'Lasercake.app'
		exe_name = 'Lasercake.app/Contents/MacOS/Lasercake'
		is_apple_app = True
		# The bundle isn't complete without calling `make install`
		subprocess.check_call(['make', 'install'])
	if exe_name == None:
		say(ansi_red+"couldn't find Lasercake binary?!"+ansi_end+'\n')
		exit(1)
	# poor estimate that assumes people only cross-compile
	# if it involves Windows:
	is_cross_compiling = (is_windows_exe != (os.name == 'nt'))
	if making_lasercake:
		say(ansi_green+'and you got:\n./'+build_dir+'/'+exe_name+'\n(etc.)'+ansi_end+'\n')
	if is_cross_compiling:
		say(ansi_yellow+'Cross-compiling, so not running tests or copying the binary.'+ansi_end+'\n')
	else:
		if running_tests:
			say(ansi_cyan+'Testing...\n')
			say_we_are_calling('./'+build_dir+'/'+exe_name+' --run-self-tests')
			test_status = subprocess.call(['./'+exe_name, '--run-self-tests'])
			if test_status != 0:
				say(ansi_red+'Tests failed.\n')
				exit(test_status)
			say(ansi_green+'success')
		else:
			say(ansi_yellow+'NOT RUNNING TESTS')
		if making_lasercake:
			say('; copying '+build_dir+'/'+app_name+' to ./'+app_name)
		say(ansi_end+'\n')
		if making_lasercake:
			replace_app(app_name, '../../'+app_name)
	exit()

if __name__ == '__main__':
	main()

<commit_msg>improve dev-build.py some for use with Emscripten<commit_after>#!/usr/bin/env python
# (Python >= 2.6 including >= 3)
docs = """
Usage: ./dev-build.py [cmake args] [certain make args]

This script wraps `cmake` and `make`; you can use CMake directly
or use this, whichever you prefer.  This script takes `cmake` and
some `make` flags directly on the command-line, puts the build in
a directory under build/ that is determined by which `cmake` flags you
specified (so that you can store separate intermediate compilation results
for various flag/compiler combinations), runs Lasercake's self-tests,
and if they're successful, copies Lasercake to the main directory.
"""

import os, os.path, sys, subprocess, re, shutil, hashlib, datetime

try:
	import multiprocessing
	MAKE_PARALLEL_JOBS_DEFAULT = multiprocessing.cpu_count() + 1
except (ImportError,NotImplementedError):
	MAKE_PARALLEL_JOBS_DEFAULT = 2

ansi_grey = '\033[30m'
ansi_red = '\033[31m'
ansi_green = '\033[32m'
ansi_yellow = '\033[33m'
ansi_blue = '\033[34m'
ansi_magenta = '\033[35m'
ansi_cyan = '\033[36m'
ansi_white = '\033[37m'
ansi_end = '\033[0m'

os.environ['PATH'] = '/usr/lib/colorgcc/bin:'+os.environ['PATH']

def say(string):
	sys.stdout.write(string)
	sys.stdout.flush()

def escaped_control_characters(string):
	return re.sub(r"[\x01-\x1F\x7F]", r"\\", string)

def say_we_are_calling(string):
	say(ansi_cyan+'% '+escaped_control_characters(string)+ansi_end+'\n')

def hash_list(l):
	hashes = (hashlib.sha256(arg.encode('utf8')).digest() for arg in l)
	return hashlib.sha256(b''.join(hashes)).hexdigest()

def replace_app(src, dst):
	"""
	Copies src to dst, replacing dst, even if dst is a running executable.
	Complications this deals with:
	On Windows, one can move but not remove or edit a running binary.
	On Linux, one can move and unlink but not edit a running binary.
	On Mac, if using .app bundles, the app is a directory rather than a file.
	"""
	old = None
	if os.path.exists(dst):
		unique_str = datetime.datetime.fromtimestamp(
				os.stat(dst).st_mtime).isoformat()
		path, basefile = os.path.split(dst)
		base, ext = os.path.splitext(basefile)
		old = os.path.join(path, base+'.'+unique_str+ext)
		os.rename(dst, old)
	if os.path.isdir(src):
		shutil.copytree(src, dst, symlinks=True)
	else:
		shutil.copy2(src, dst)
	if old != None:
		if os.path.isdir(old):
			shutil.rmtree(old)
		else:
			try: os.unlink(old)
			except OSError:
				say("Couldn't remove old (running?) executable "+os.path.basename(old))

def main():
	os.chdir(os.path.dirname(os.path.join('.', __file__)))
	try: subprocess.call(['cmake', '--version'])
	except OSError:
		say(ansi_red+"Error: 'cmake' not found; please install it."+ansi_end+'\n')
		exit(1)
	try: subprocess.call(['make', '--version'])
	except OSError:
		say(ansi_red+"Error: 'make' not found; please install it."+ansi_end+'\n')
		exit(1)
	try:
		status = subprocess.call(['time', '-f', '', 'true'])
		time_has_f = (status == 0)
		has_time = True
	except OSError:
		has_time = False
	cmake_args = []
	make_args = []
	making_lasercake = True
	running_tests = True
	make_jobs_flags = ['-j'+str(MAKE_PARALLEL_JOBS_DEFAULT)]
	for arg in sys.argv[1:]:
		if arg in ['-?', '-h', '--help', '-help']:
			say(docs)
			sys.exit()
		elif arg == 'no-test':
			running_tests = False
		elif arg == 'VERBOSE=1':
			make_args.append(arg)
		elif re.search(r'^-j[0-9]*$', arg):
			make_jobs_flags = [arg]
		else:
			cmake_args.append(arg)
		if arg == '-DBUILD_SELF_TESTS=OFF':
			running_tests = False
	# Make it more likely the directories are named non-conflictingly
	# even though we delete some special characters from their paths.
	hash_for_this_config = hash_list(cmake_args)
	name_for_this_config = hash_for_this_config[:10] + '_' + (','.join([re.sub(r'[^-+.a-zA-Z0-9]', r'_', arg) for arg in cmake_args]) or 'default')
	build_dir = 'build/'+name_for_this_config
	try: os.remove('CMakeCache.txt')
	except OSError: pass
	try:
		os.makedirs(build_dir)
	except OSError:
		say(ansi_green+'This build dir already exists!'+ansi_end+'\n')
		cmake_args = []
	to_call_cmake = ['cmake', '../../'] + cmake_args
	say_we_are_calling('cd '+build_dir+'; '+'   '.join(to_call_cmake))
	say(ansi_cyan+'''  (^^ not escaped properly in these info messages - doin' it right in python)'''+ansi_end+'\n')
	os.chdir(build_dir)
	subprocess.check_call(to_call_cmake)
	if not has_time:
		time_args = []
	elif time_has_f:
		time_args = ['time', '-f', (ansi_magenta+'`make` took %E'+ansi_end)]
	else:
		time_args = ['time']
	to_call_make = time_args + ['make'] + make_jobs_flags + make_args
	say_we_are_calling('   '.join(to_call_make))
	make_status = subprocess.call(to_call_make)
	if has_time and not time_has_f:
		say(ansi_magenta+'  (time to run `make`).'+ansi_end+'\n')
	# How to print the compiler & flags here? Can we get it from cmake
	# somehow? Add to the cmakelists to write those?
	# but with the -l's and all... hmm.
	if make_status != 0:
		# skip make_jobs_flags (-jN) to get just one file's error
		# messages not mixed up with the others
		to_call_make_again = ['make'] + make_args
		say_we_are_calling('   '.join(to_call_make_again))
		subprocess.call(to_call_make_again)
		say(ansi_red+'build failed'+ansi_end+'\n')
		exit(1)
	is_windows_exe = is_apple_app = is_js = False
	app_name = exe_name = None
	if os.path.exists('Lasercake'):
		app_name = exe_name = 'Lasercake'
	if os.path.exists('Lasercake.exe'):
		app_name = exe_name = 'Lasercake.exe'
		is_windows_exe = True
	if os.path.exists('Lasercake.js'):
		exe_name = 'Lasercake.js'
		is_js = True
	if os.path.exists('Lasercake.app'):
		app_name = 'Lasercake.app'
		exe_name = 'Lasercake.app/Contents/MacOS/Lasercake'
		is_apple_app = True
		# The bundle isn't complete without calling `make install`
		subprocess.check_call(['make', 'install'])
	if exe_name == None:
		say(ansi_red+"couldn't find Lasercake binary?!"+ansi_end+'\n')
		exit(1)
	# poor estimate that assumes people only cross-compile
	# if it involves Windows (or Javascript, but we can run
	# tests even for Javascript):
	is_cross_compiling = (is_windows_exe != (os.name == 'nt'))
	if making_lasercake:
		say(ansi_green+'and you got:\n./'+build_dir+'/'+exe_name+'\n(etc.)'+ansi_end+'\n')
	if is_cross_compiling:
		say(ansi_yellow+'Cross-compiling, so not running tests or copying the binary.'+ansi_end+'\n')
	else:
		if running_tests:
			say(ansi_cyan+'Testing...\n')
			say_we_are_calling('(cd '+build_dir+'; '+
				('node ' if is_js else '')+'./'+exe_name+' --run-self-tests)')
			test_status = subprocess.call((['node'] if is_js else []) +
				['./'+exe_name, '--run-self-tests'])
			if test_status != 0:
				say(ansi_red+'Tests failed.\n')
				exit(test_status)
			say(ansi_green+'success')
		else:
			say(ansi_yellow+'NOT RUNNING TESTS')
		if is_js:
			making_lasercake = False
		if making_lasercake:
			say('; copying '+build_dir+'/'+app_name+' to ./'+app_name)
		say(ansi_end+'\n')
		if making_lasercake:
			replace_app(app_name, '../../'+app_name)
	exit()

if __name__ == '__main__':
	main()

<|endoftext|>


<commit_before>import logging
import os
import numpy as np
from copy import deepcopy
try:
    from math import gcd
except:  # Moved to math in python 3.5, this is to be 3.4 compatible
    from fractions import gcd
from ..waveform_control import pulsar
from ..waveform_control import element
from ..waveform_control.element import calculate_time_corr
from ..waveform_control import pulse
from ..waveform_control import sequence
from pycqed.measurement.pulse_sequences.standard_elements import multi_pulse_elt

from importlib import reload
reload(pulse)
from ..waveform_control import pulse_library
reload(pulse_library)

station = None
reload(element)
kernel_dir_path = 'kernels/'
# You need to explicitly set this before running any functions from this module
# I guess there are cleaner solutions :)

def single_pulse_seq(pulse_pars=None,
                     verbose=False,
                     distortion_dict=None,
                     return_seq=False):
    '''

    '''
    if pulse_pars is None:
        pulse_pars = {'pulse_type': 'SquarePulse',
                      'pulse_delay': .1e-6,
                      'channel': 'ch3',
                      'amplitude': 0.5,
                      'length': .1e-6,
                      'dead_time_length': 10e-6}
    minus_pulse_pars = {'pulse_type': 'SquarePulse',
                  'pulse_delay': 3e-6 + pulse_pars['length'] + pulse_pars['pulse_delay'],
                  'channel': pulse_pars['channel'],
                  'amplitude': -pulse_pars['amplitude'],
                  'length': pulse_pars['length'],
                  'dead_time_length': 10e-6}

    dead_time_pulse = {'pulse_type': 'SquarePulse',
                       'pulse_delay': (minus_pulse_pars['length']),
                       'channel': pulse_pars['channel'],
                       'amplitude': 0,
                       'length': pulse_pars['dead_time_length']}
                       # 'length': 5e-6}
    seq_name = 'Square_seq'
    seq = sequence.Sequence(seq_name)
    station.pulsar.update_channel_settings()
    el_list = []
    for i, iter in enumerate([0, 1]):  # seq has to have at least 2 elts

        pulse_list = [pulse_pars, minus_pulse_pars,dead_time_pulse]
        # pulse_list = [pulse_pars, dead_time_pulse]

        el = multi_pulse_elt(i, station, pulse_list)
        el_list.append(el)

    if distortion_dict is not None:
        preloaded_kernels_vec = preload_kernels_func(distortion_dict)
    else:
        preloaded_kernels = []
    for i, el in enumerate(el_list):
        if distortion_dict is not None:
            el = distort_and_compensate(el, distortion_dict, preloaded_kernels_vec)
            el_list[i] = el
        seq.append_element(el, trigger_wait=True)
    station.components['AWG'].stop()
    station.pulsar.program_awg(seq, *el_list, verbose=verbose)
    if return_seq:
        return seq, el_list
    else:
        return seq

def chevron_seq_length(lengths, mw_pulse_pars, RO_pars, flux_pulse_pars=None,
                     verbose=False,
                     distortion_dict=None,
                     upload=True,
                     return_seq=False):
    '''

    '''
    preloaded_kernels_vec = preload_kernels_func(distortion_dict)
    if flux_pulse_pars is None:
        flux_pulse_pars = {'pulse_type': 'SquarePulse',
                      'pulse_delay': .1e-6,
                      'channel': 'ch3',
                      'amplitude': 0.5,
                      'length': .1e-6}
    # flux_pulse_pars['amplitude'] = 0.
    minus_flux_pulse_pars = {'pulse_type': 'SquarePulse',
                  'pulse_delay': 0., # will be overwritten
                  'channel': 'ch3',
                  'amplitude': -flux_pulse_pars['amplitude'],
                  'length': flux_pulse_pars['length']}
    original_delay = deepcopy(RO_pars)['pulse_delay']

    seq_name = 'Chevron_seq'
    seq = sequence.Sequence(seq_name)
    station.pulsar.update_channel_settings()
    el_list = []
    pulses = get_pulse_dict_from_pars(mw_pulse_pars)
    for i, lngt in enumerate(lengths):  # seq has to have at least 2 elts
        flux_pulse_pars['length'] = lngt
        minus_flux_pulse_pars['length'] = lngt
        # correcting timings
        pulse_buffer = 50e-9
        flux_pulse_pars['pulse_delay'] = pulse_buffer + (mw_pulse_pars['sigma'] *
                                                         mw_pulse_pars['nr_sigma'])
        msmt_buffer = 50e-9
        RO_pars['pulse_delay'] = msmt_buffer + lngt

        dead_time = 3e-6
        minus_flux_pulse_pars['pulse_delay'] = dead_time + RO_pars['length']
        pulse_list = [pulses['X180'], flux_pulse_pars, RO_pars, minus_flux_pulse_pars]
        # copy first element and set extra wait
        pulse_list[0] = deepcopy(pulse_list[0])
        pulse_list[0]['pulse_delay'] += 0.01e-6 + ((-int(lngt*1e9)) % 50)*1e-9

        el = multi_pulse_elt(i, station, pulse_list)
        el_list.append(el)

    for i, el in enumerate(el_list):
        if distortion_dict is not None:
            el = distort_and_compensate(el, distortion_dict, preloaded_kernels_vec)
            el_list[i] = el
        seq.append_element(el, trigger_wait=True)
    cal_points = 4
    RO_pars['pulse_delay'] = original_delay
    for i in range(int(cal_points/2)):
        pulse_list = [pulses['I'], RO_pars]
        # copy first element and set extra wait
        pulse_list[0] = deepcopy(pulse_list[0])
        pulse_list[0]['pulse_delay'] += 0.01e-6

        el = multi_pulse_elt(len(lengths)+i, station, pulse_list)
        el_list.append(el)
        seq.append_element(el, trigger_wait=True)
    for i in range(int(cal_points/2)):
        pulse_list = [pulses['X180'], RO_pars]
        # copy first element and set extra wait
        pulse_list[0] = deepcopy(pulse_list[0])
        pulse_list[0]['pulse_delay'] += 0.01e-6

        el = multi_pulse_elt(len(lengths)+int(cal_points/2)+i, station, pulse_list)
        el_list.append(el)
        seq.append_element(el, trigger_wait=True)

    if upload:
        station.components['AWG'].stop()
        station.pulsar.program_awg(seq, *el_list, verbose=verbose)
    if return_seq:
        return seq, el_list
    else:
        return seq

def BusT1(times, mw_pulse_pars, RO_pars, flux_pulse_pars=None,
          verbose=False, distortion_dict=None,
          upload=True, return_seq=False):
    '''

    '''
    preloaded_kernels_vec = preload_kernels_func(distortion_dict)
    if flux_pulse_pars is None:
        raise ValueError('Need flux parameters for the gate.')
    minus_flux_pulse_pars = deepcopy(flux_pulse_pars)
    minus_flux_pulse_pars['amplitude']=-minus_flux_pulse_pars['amplitude']

    original_delay = deepcopy(RO_pars)['pulse_delay']

    seq_name = 'BusT1_seq'
    seq = sequence.Sequence(seq_name)
    station.pulsar.update_channel_settings()
    el_list = []
    pulses = get_pulse_dict_from_pars(mw_pulse_pars)
    for i, tt in enumerate(times):
        # correcting timings
        pulse_buffer = 50e-9
        flux_pulse_pars['pulse_delay'] = pulse_buffer + (mw_pulse_pars['sigma'] *
                                                         mw_pulse_pars['nr_sigma'])

        flux_pulse_pars_2 = deepcopy(flux_pulse_pars)
        # flux_pulse_pars_2['amplitude'] = 0.
        flux_pulse_pars_2['pulse_delay'] = tt + flux_pulse_pars['length']

        msmt_buffer = 50e-9
        RO_pars['pulse_delay'] = msmt_buffer + flux_pulse_pars['length']

        dead_time = 3e-6
        minus_flux_pulse_pars['pulse_delay'] = dead_time + RO_pars['length']

        minus_flux_pulse_pars_2 = deepcopy(flux_pulse_pars_2)
        minus_flux_pulse_pars_2['amplitude']=-minus_flux_pulse_pars_2['amplitude']

        pulse_list = [pulses['X180'], flux_pulse_pars, flux_pulse_pars_2,
                      RO_pars, minus_flux_pulse_pars, minus_flux_pulse_pars_2]

        #This ensures fixed point
        pulse_list[0]['pulse_delay'] += 0.01e-6#+ ((-int(tt*1e9)) % 50)*1e-9

        el = multi_pulse_elt(i, station, pulse_list)
        el_list.append(el)

    for i, el in enumerate(el_list):
        if distortion_dict is not None:
            el = distort_and_compensate(el, distortion_dict, preloaded_kernels_vec)
            el_list[i] = el
        seq.append_element(el, trigger_wait=True)
    cal_points = 4
    RO_pars['pulse_delay'] = original_delay
    for i in range(int(cal_points/2)):
        pulse_list = [pulses['I'], RO_pars]
        # copy first element and set extra wait
        pulse_list[0] = deepcopy(pulse_list[0])
        pulse_list[0]['pulse_delay'] += 0.01e-6

        el = multi_pulse_elt(len(times)+i, station, pulse_list)
        el_list.append(el)
        seq.append_element(el, trigger_wait=True)
    for i in range(int(cal_points/2)):
        pulse_list = [pulses['X180'], RO_pars]
        # copy first element and set extra wait
        pulse_list[0] = deepcopy(pulse_list[0])
        pulse_list[0]['pulse_delay'] += 0.01e-6

        el = multi_pulse_elt(len(times)+int(cal_points/2)+i, station, pulse_list)
        el_list.append(el)
        seq.append_element(el, trigger_wait=True)

    if upload:
        station.components['AWG'].stop()
        station.pulsar.program_awg(seq, *el_list, verbose=verbose)
    if return_seq:
        return seq, el_list
    else:
        return seq

def preload_kernels_func(distortion_dict):
    output_dict = {ch: [] for ch in distortion_dict['ch_list']}
    for ch in distortion_dict['ch_list']:
        for kernel in distortion_dict[ch]:
            if kernel is not '':
                print('Loading {}'.format(kernel_dir_path+kernel))
                # print(os.path.isfile('kernels/'+kernel))
                output_dict[ch].append(np.loadtxt(kernel_dir_path+kernel))
    return output_dict

def distort_and_compensate(element, distortion_dict, preloaded_kernels):
    t_vals, outputs_dict = element.waveforms()
    # print(len(t_vals),t_vals[-1])
    for ch in distortion_dict['ch_list']:
        element._channels[ch]['distorted'] = True
        length = len(outputs_dict[ch])
        for kernelvec in preloaded_kernels[ch]:
            outputs_dict[ch] = np.convolve(outputs_dict[ch], kernelvec)[:length]

        element.distorted_wfs[ch] = outputs_dict[ch][:len(t_vals)]
    return element

def get_pulse_dict_from_pars(pulse_pars):
    '''
    Returns a dictionary containing pulse_pars for all the primitive pulses
    based on a single set of pulse_pars.
    Using this function deepcopies the pulse parameters preventing accidently
    editing the input dictionary.

    input args:
        pulse_pars: dictionary containing pulse_parameters
    return:
        pulses: dictionary of pulse_pars dictionaries
    '''
    pi_amp = pulse_pars['amplitude']
    pi2_amp = pulse_pars['amplitude']*pulse_pars['amp90_scale']

    pulses = {'I': deepcopy(pulse_pars),
              'X180': deepcopy(pulse_pars),
              'mX180': deepcopy(pulse_pars),
              'X90': deepcopy(pulse_pars),
              'mX90': deepcopy(pulse_pars),
              'Y180': deepcopy(pulse_pars),
              'mY180': deepcopy(pulse_pars),
              'Y90': deepcopy(pulse_pars),
              'mY90': deepcopy(pulse_pars)}

    pulses['I']['amplitude'] = 0
    pulses['mX180']['amplitude'] = -pi_amp
    pulses['X90']['amplitude'] = pi2_amp
    pulses['mX90']['amplitude'] = -pi2_amp
    pulses['Y180']['phase'] = 90
    pulses['mY180']['phase'] = 90
    pulses['mY180']['amplitude'] = -pi_amp

    pulses['Y90']['amplitude'] = pi2_amp
    pulses['Y90']['phase'] = 90
    pulses['mY90']['amplitude'] = -pi2_amp
    pulses['mY90']['phase'] = 90

    return pulses
<commit_msg>Added param for compensations in kernel tuning sequence<commit_after>import logging
import os
import numpy as np
from copy import deepcopy
try:
    from math import gcd
except:  # Moved to math in python 3.5, this is to be 3.4 compatible
    from fractions import gcd
from ..waveform_control import pulsar
from ..waveform_control import element
from ..waveform_control.element import calculate_time_corr
from ..waveform_control import pulse
from ..waveform_control import sequence
from pycqed.measurement.pulse_sequences.standard_elements import multi_pulse_elt

from importlib import reload
reload(pulse)
from ..waveform_control import pulse_library
reload(pulse_library)

station = None
reload(element)
kernel_dir_path = 'kernels/'
# You need to explicitly set this before running any functions from this module
# I guess there are cleaner solutions :)

def single_pulse_seq(pulse_pars=None,
                     comp_pulse=True,
                     verbose=False,
                     distortion_dict=None,
                     return_seq=False):
    '''

    '''
    if pulse_pars is None:
        pulse_pars = {'pulse_type': 'SquarePulse',
                      'pulse_delay': .1e-6,
                      'channel': 'ch3',
                      'amplitude': 0.5,
                      'length': .1e-6,
                      'dead_time_length': 10e-6}
    minus_pulse_pars = {'pulse_type': 'SquarePulse',
                  'pulse_delay': 3e-6 + pulse_pars['length'] + pulse_pars['pulse_delay'],
                  'channel': pulse_pars['channel'],
                  'amplitude': -pulse_pars['amplitude'],
                  'length': pulse_pars['length'],
                  'dead_time_length': 10e-6}

    dead_time_pulse = {'pulse_type': 'SquarePulse',
                       'pulse_delay': (minus_pulse_pars['length']),
                       'channel': pulse_pars['channel'],
                       'amplitude': 0,
                       'length': pulse_pars['dead_time_length']}
                       # 'length': 5e-6}
    seq_name = 'Square_seq'
    seq = sequence.Sequence(seq_name)
    station.pulsar.update_channel_settings()
    el_list = []
    for i, iter in enumerate([0, 1]):  # seq has to have at least 2 elts

        if comp_pulse:
            pulse_list = [pulse_pars, minus_pulse_pars,dead_time_pulse]
        else:
            pulse_list = [pulse_pars, dead_time_pulse]
        # pulse_list = [pulse_pars, dead_time_pulse]

        el = multi_pulse_elt(i, station, pulse_list)
        el_list.append(el)

    if distortion_dict is not None:
        preloaded_kernels_vec = preload_kernels_func(distortion_dict)
    else:
        preloaded_kernels = []
    for i, el in enumerate(el_list):
        if distortion_dict is not None:
            el = distort_and_compensate(el, distortion_dict, preloaded_kernels_vec)
            el_list[i] = el
        seq.append_element(el, trigger_wait=True)
    station.components['AWG'].stop()
    station.pulsar.program_awg(seq, *el_list, verbose=verbose)
    if return_seq:
        return seq, el_list
    else:
        return seq

def chevron_seq_length(lengths, mw_pulse_pars, RO_pars, flux_pulse_pars=None,
                     verbose=False,
                     distortion_dict=None,
                     upload=True,
                     return_seq=False):
    '''

    '''
    preloaded_kernels_vec = preload_kernels_func(distortion_dict)
    if flux_pulse_pars is None:
        flux_pulse_pars = {'pulse_type': 'SquarePulse',
                      'pulse_delay': .1e-6,
                      'channel': 'ch3',
                      'amplitude': 0.5,
                      'length': .1e-6}
    # flux_pulse_pars['amplitude'] = 0.
    minus_flux_pulse_pars = {'pulse_type': 'SquarePulse',
                  'pulse_delay': 0., # will be overwritten
                  'channel': 'ch3',
                  'amplitude': -flux_pulse_pars['amplitude'],
                  'length': flux_pulse_pars['length']}
    original_delay = deepcopy(RO_pars)['pulse_delay']

    seq_name = 'Chevron_seq'
    seq = sequence.Sequence(seq_name)
    station.pulsar.update_channel_settings()
    el_list = []
    pulses = get_pulse_dict_from_pars(mw_pulse_pars)
    for i, lngt in enumerate(lengths):  # seq has to have at least 2 elts
        flux_pulse_pars['length'] = lngt
        minus_flux_pulse_pars['length'] = lngt
        # correcting timings
        pulse_buffer = 50e-9
        flux_pulse_pars['pulse_delay'] = pulse_buffer + (mw_pulse_pars['sigma'] *
                                                         mw_pulse_pars['nr_sigma'])
        msmt_buffer = 50e-9
        RO_pars['pulse_delay'] = msmt_buffer + lngt

        dead_time = 3e-6
        minus_flux_pulse_pars['pulse_delay'] = dead_time + RO_pars['length']
        pulse_list = [pulses['X180'], flux_pulse_pars, RO_pars, minus_flux_pulse_pars]
        # copy first element and set extra wait
        pulse_list[0] = deepcopy(pulse_list[0])
        pulse_list[0]['pulse_delay'] += 0.01e-6 + ((-int(lngt*1e9)) % 50)*1e-9

        el = multi_pulse_elt(i, station, pulse_list)
        el_list.append(el)

    for i, el in enumerate(el_list):
        if distortion_dict is not None:
            el = distort_and_compensate(el, distortion_dict, preloaded_kernels_vec)
            el_list[i] = el
        seq.append_element(el, trigger_wait=True)
    cal_points = 4
    RO_pars['pulse_delay'] = original_delay
    for i in range(int(cal_points/2)):
        pulse_list = [pulses['I'], RO_pars]
        # copy first element and set extra wait
        pulse_list[0] = deepcopy(pulse_list[0])
        pulse_list[0]['pulse_delay'] += 0.01e-6

        el = multi_pulse_elt(len(lengths)+i, station, pulse_list)
        el_list.append(el)
        seq.append_element(el, trigger_wait=True)
    for i in range(int(cal_points/2)):
        pulse_list = [pulses['X180'], RO_pars]
        # copy first element and set extra wait
        pulse_list[0] = deepcopy(pulse_list[0])
        pulse_list[0]['pulse_delay'] += 0.01e-6

        el = multi_pulse_elt(len(lengths)+int(cal_points/2)+i, station, pulse_list)
        el_list.append(el)
        seq.append_element(el, trigger_wait=True)

    if upload:
        station.components['AWG'].stop()
        station.pulsar.program_awg(seq, *el_list, verbose=verbose)
    if return_seq:
        return seq, el_list
    else:
        return seq

def BusT1(times, mw_pulse_pars, RO_pars, flux_pulse_pars=None,
          verbose=False, distortion_dict=None,
          upload=True, return_seq=False):
    '''

    '''
    preloaded_kernels_vec = preload_kernels_func(distortion_dict)
    if flux_pulse_pars is None:
        raise ValueError('Need flux parameters for the gate.')
    minus_flux_pulse_pars = deepcopy(flux_pulse_pars)
    minus_flux_pulse_pars['amplitude']=-minus_flux_pulse_pars['amplitude']

    original_delay = deepcopy(RO_pars)['pulse_delay']

    seq_name = 'BusT1_seq'
    seq = sequence.Sequence(seq_name)
    station.pulsar.update_channel_settings()
    el_list = []
    pulses = get_pulse_dict_from_pars(mw_pulse_pars)
    for i, tt in enumerate(times):
        # correcting timings
        pulse_buffer = 50e-9
        flux_pulse_pars['pulse_delay'] = pulse_buffer + (mw_pulse_pars['sigma'] *
                                                         mw_pulse_pars['nr_sigma'])

        flux_pulse_pars_2 = deepcopy(flux_pulse_pars)
        # flux_pulse_pars_2['amplitude'] = 0.
        flux_pulse_pars_2['pulse_delay'] = tt + flux_pulse_pars['length']

        msmt_buffer = 50e-9
        RO_pars['pulse_delay'] = msmt_buffer + flux_pulse_pars['length']

        dead_time = 3e-6
        minus_flux_pulse_pars['pulse_delay'] = dead_time + RO_pars['length']

        minus_flux_pulse_pars_2 = deepcopy(flux_pulse_pars_2)
        minus_flux_pulse_pars_2['amplitude']=-minus_flux_pulse_pars_2['amplitude']

        pulse_list = [pulses['X180'], flux_pulse_pars, flux_pulse_pars_2,
                      RO_pars, minus_flux_pulse_pars, minus_flux_pulse_pars_2]

        #This ensures fixed point
        pulse_list[0]['pulse_delay'] += 0.01e-6#+ ((-int(tt*1e9)) % 50)*1e-9

        el = multi_pulse_elt(i, station, pulse_list)
        el_list.append(el)

    for i, el in enumerate(el_list):
        if distortion_dict is not None:
            el = distort_and_compensate(el, distortion_dict, preloaded_kernels_vec)
            el_list[i] = el
        seq.append_element(el, trigger_wait=True)
    cal_points = 4
    RO_pars['pulse_delay'] = original_delay
    for i in range(int(cal_points/2)):
        pulse_list = [pulses['I'], RO_pars]
        # copy first element and set extra wait
        pulse_list[0] = deepcopy(pulse_list[0])
        pulse_list[0]['pulse_delay'] += 0.01e-6

        el = multi_pulse_elt(len(times)+i, station, pulse_list)
        el_list.append(el)
        seq.append_element(el, trigger_wait=True)
    for i in range(int(cal_points/2)):
        pulse_list = [pulses['X180'], RO_pars]
        # copy first element and set extra wait
        pulse_list[0] = deepcopy(pulse_list[0])
        pulse_list[0]['pulse_delay'] += 0.01e-6

        el = multi_pulse_elt(len(times)+int(cal_points/2)+i, station, pulse_list)
        el_list.append(el)
        seq.append_element(el, trigger_wait=True)

    if upload:
        station.components['AWG'].stop()
        station.pulsar.program_awg(seq, *el_list, verbose=verbose)
    if return_seq:
        return seq, el_list
    else:
        return seq

def preload_kernels_func(distortion_dict):
    output_dict = {ch: [] for ch in distortion_dict['ch_list']}
    for ch in distortion_dict['ch_list']:
        for kernel in distortion_dict[ch]:
            if kernel is not '':
                print('Loading {}'.format(kernel_dir_path+kernel))
                # print(os.path.isfile('kernels/'+kernel))
                output_dict[ch].append(np.loadtxt(kernel_dir_path+kernel))
    return output_dict

def distort_and_compensate(element, distortion_dict, preloaded_kernels):
    t_vals, outputs_dict = element.waveforms()
    # print(len(t_vals),t_vals[-1])
    for ch in distortion_dict['ch_list']:
        element._channels[ch]['distorted'] = True
        length = len(outputs_dict[ch])
        for kernelvec in preloaded_kernels[ch]:
            outputs_dict[ch] = np.convolve(outputs_dict[ch], kernelvec)[:length]

        element.distorted_wfs[ch] = outputs_dict[ch][:len(t_vals)]
    return element

def get_pulse_dict_from_pars(pulse_pars):
    '''
    Returns a dictionary containing pulse_pars for all the primitive pulses
    based on a single set of pulse_pars.
    Using this function deepcopies the pulse parameters preventing accidently
    editing the input dictionary.

    input args:
        pulse_pars: dictionary containing pulse_parameters
    return:
        pulses: dictionary of pulse_pars dictionaries
    '''
    pi_amp = pulse_pars['amplitude']
    pi2_amp = pulse_pars['amplitude']*pulse_pars['amp90_scale']

    pulses = {'I': deepcopy(pulse_pars),
              'X180': deepcopy(pulse_pars),
              'mX180': deepcopy(pulse_pars),
              'X90': deepcopy(pulse_pars),
              'mX90': deepcopy(pulse_pars),
              'Y180': deepcopy(pulse_pars),
              'mY180': deepcopy(pulse_pars),
              'Y90': deepcopy(pulse_pars),
              'mY90': deepcopy(pulse_pars)}

    pulses['I']['amplitude'] = 0
    pulses['mX180']['amplitude'] = -pi_amp
    pulses['X90']['amplitude'] = pi2_amp
    pulses['mX90']['amplitude'] = -pi2_amp
    pulses['Y180']['phase'] = 90
    pulses['mY180']['phase'] = 90
    pulses['mY180']['amplitude'] = -pi_amp

    pulses['Y90']['amplitude'] = pi2_amp
    pulses['Y90']['phase'] = 90
    pulses['mY90']['amplitude'] = -pi2_amp
    pulses['mY90']['phase'] = 90

    return pulses
<|endoftext|>


<commit_before>import daemon
import datetime
import logging
import logging.handlers
import multiprocessing
import os
import resource
import signal
import sys
import threading
import time
import traceback
import yaml

from lobster import cmssw, job, util

from pkg_resources import get_distribution

import work_queue as wq

logger = multiprocessing.get_logger()

class ShortPathFormatter(logging.Formatter):
    def format(self, record):
        if len(record.pathname) > 40:
            record.pathname = '...' + record.pathname[-37:]
        # FIXME at some point, Formatter is a new-style class and we can
        # use super
        # return super(ShortPathFormatter, self).format(record)
        return logging.Formatter.format(self, record)

def kill(args):
    logger.info("setting flag to quit at the next checkpoint")
    with open(args.configfile) as configfile:
        config = yaml.load(configfile)

    workdir = config['workdir']
    util.register_checkpoint(workdir, 'KILLED', 'PENDING')

def run(args):
    with open(args.configfile) as configfile:
        config = yaml.load(configfile)

    workdir = config['workdir']
    if not os.path.exists(workdir):
        os.makedirs(workdir)
        util.register_checkpoint(workdir, "version", get_distribution('Lobster').version)
    else:
        util.verify(workdir)

    cmsjob = False
    if config.get('type', 'cmssw') == 'cmssw':
        cmsjob = True

        from ProdCommon.Credential.CredentialAPI import CredentialAPI
        cred = CredentialAPI({'credential': 'Proxy'})
        if cred.checkCredential(Time=60):
            if not 'X509_USER_PROXY' in os.environ:
                os.environ['X509_USER_PROXY'] = cred.credObj.getUserProxy()
        else:
            if config.get('advanced', {}).get('renew proxy', True):
                try:
                    cred.ManualRenewCredential()
                except Exception as e:
                    print("could not renew proxy")
                    sys.exit(1)
            else:
                print("please renew your proxy")
                sys.exit(1)

    print "Saving log to {0}".format(os.path.join(workdir, 'lobster.log'))

    if not args.foreground:
        ttyfile = open(os.path.join(workdir, 'lobster.err'), 'a')
        print "Saving stderr and stdout to {0}".format(os.path.join(workdir, 'lobster.err'))

    if config.get('advanced', {}).get('dump core', False):
        print "Setting core dump size to unlimited"
        resource.setrlimit(resource.RLIMIT_CORE, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))

    signals = daemon.daemon.make_default_signal_map()
    signals[signal.SIGTERM] = lambda num, frame: kill(args)

    with daemon.DaemonContext(
            detach_process=not args.foreground,
            stdout=sys.stdout if args.foreground else ttyfile,
            stderr=sys.stderr if args.foreground else ttyfile,
            working_directory=workdir,
            pidfile=util.get_lock(workdir, args.force),
            prevent_core=False,
            signal_map=signals):

        level = max(1, config.get('advanced', {}).get('log level', 2)) * 10
        fileh = logging.handlers.RotatingFileHandler(os.path.join(workdir, 'lobster.log'), maxBytes=500e6, backupCount=10)
        fileh.setFormatter(ShortPathFormatter("%(asctime)s [%(levelname)5s] - %(pathname)-40s %(lineno)4d: %(message)s"))
        fileh.setLevel(level)

        logger.addHandler(fileh)
        logger.setLevel(level)

        if args.foreground:
            console = logging.StreamHandler()
            console.setLevel(level)
            console.setFormatter(ShortPathFormatter("%(asctime)s [%(levelname)5s] - %(pathname)-40s %(lineno)4d: %(message)s"))
            logger.addHandler(console)

        config['configdir'] = args.configdir
        config['filename'] = args.configfile
        config['startdir'] = args.startdir

        t = threading.Thread(target=sprint, args=(config, workdir, cmsjob))
        t.start()
        t.join()

        logger.info("lobster terminated")

def sprint(config, workdir, cmsjob):
    if cmsjob:
        job_src = cmssw.JobProvider(config)
        actions = cmssw.Actions(config)
    else:
        job_src = job.SimpleJobProvider(config)
        actions = None

    logger.info("using wq from {0}".format(wq.__file__))

    wq.cctools_debug_flags_set("all")
    wq.cctools_debug_config_file(os.path.join(workdir, "work_queue_debug.log"))
    wq.cctools_debug_config_file_size(1 << 29)

    queue = wq.WorkQueue(-1)
    queue.specify_log(os.path.join(workdir, "work_queue.log"))
    queue.specify_name("lobster_" + config["id"])
    queue.specify_keepalive_timeout(300)
    # queue.tune("short-timeout", 600)
    queue.tune("transfer-outlier-factor", 4)
    queue.specify_algorithm(wq.WORK_QUEUE_SCHEDULE_RAND)

    cores = config.get('cores per job', 1)
    logger.info("starting queue as {0}".format(queue.name))
    logger.info("submit workers with: condor_submit_workers -M {0}{1} <num>".format(
        queue.name, ' --cores {0}'.format(cores) if cores > 1 else ''))

    payload = config.get('advanced', {}).get('payload', 10)
    abort_active = False
    abort_threshold = config.get('advanced', {}).get('abort threshold', 400)
    abort_multiplier = config.get('advanced', {}).get('abort multiplier', 4)

    if util.checkpoint(workdir, 'KILLED') == 'PENDING':
        util.register_checkpoint(workdir, 'KILLED', 'RESTART')

    # time in seconds to wait for WQ to return tasks
    interval = 60

    jobits_left = 0
    successful_jobs = 0

    creation_time = 0
    destruction_time = 0

    with open(os.path.join(workdir, "lobster_stats.log"), "a") as statsfile:
        statsfile.write(
                "#timestamp " +
                "total_workers_connected total_workers_joined total_workers_removed " +
                "workers_busy workers_idle " +
                "tasks_running " +
                "total_send_time total_receive_time " +
                "total_create_time total_return_time " +
                "idle_percentage " +
                "capacity " +
                "efficiency " +
                "total_memory " +
                "total_cores " +
                "jobits_left\n")

    bad_exitcodes = job_src.bad_exitcodes

    while not job_src.done():
        jobits_left = job_src.work_left()
        stats = queue.stats_hierarchy

        with open(os.path.join(workdir, "lobster_stats.log"), "a") as statsfile:
            now = datetime.datetime.now()
            statsfile.write(" ".join(map(str,
                [
                    int(int(now.strftime('%s')) * 1e6 + now.microsecond),
                    stats.total_workers_connected,
                    stats.total_workers_joined,
                    stats.total_workers_removed,
                    stats.workers_busy,
                    stats.workers_idle,
                    stats.tasks_running,
                    stats.total_send_time,
                    stats.total_receive_time,
                    creation_time,
                    destruction_time,
                    stats.idle_percentage,
                    stats.capacity,
                    stats.efficiency,
                    stats.total_memory,
                    stats.total_cores,
                    jobits_left
                ]
                )) + "\n"
            )

        if util.checkpoint(workdir, 'KILLED') == 'PENDING':
            util.register_checkpoint(workdir, 'KILLED', str(datetime.datetime.utcnow()))

            # let the job source shut down gracefully
            logger.info("terminating job source")
            job_src.terminate()
            logger.info("terminating gracefully")
            break

        logger.info("{0} out of {1} workers busy; {3} jobs running, {4} waiting; {2} jobits left".format(
                stats.workers_busy,
                stats.workers_busy + stats.workers_ready,
                jobits_left,
                stats.tasks_running,
                stats.tasks_waiting))

        hunger = max(payload + stats.total_workers_connected / 10 - stats.tasks_waiting, 0)

        t = time.time()
        while hunger > 0:
            jobs = job_src.obtain(hunger)

            if jobs == None or len(jobs) == 0:
                break

            hunger -= len(jobs)
            for cores, cmd, id, inputs, outputs in jobs:
                task = wq.Task(cmd)
                task.specify_tag(id)
                task.specify_cores(cores)
                # temporary work-around?
                # task.specify_memory(1000)
                # task.specify_disk(4000)

                for (local, remote, cache) in inputs:
                    if os.path.isfile(local):
                        cache_opt = wq.WORK_QUEUE_CACHE if cache else wq.WORK_QUEUE_NOCACHE
                        task.specify_input_file(str(local), str(remote), cache_opt)
                    elif os.path.isdir(local):
                        task.specify_directory(str(local), str(remote), wq.WORK_QUEUE_INPUT,
                                wq.WORK_QUEUE_CACHE, recursive=True)
                    else:
                        logger.critical("cannot send file to worker: {0}".format(local))
                        raise NotImplementedError

                for (local, remote) in outputs:
                    task.specify_output_file(str(local), str(remote))

                queue.submit(task)
        creation_time += int((time.time() - t) * 1e6)

        job_src.update(queue)
        starttime = time.time()
        task = queue.wait(interval)
        tasks = []
        while task:
            if task.return_status == 0:
                successful_jobs += 1
            elif task.return_status in bad_exitcodes:
                logger.warning("blacklisting host {0} due to bad exit code from job {1}".format(task.hostname, task.tag))
                queue.blacklist(task.hostname)
            tasks.append(task)

            remaining = int(starttime + interval - time.time())
            if queue.stats.tasks_waiting > 0 and remaining > 0:
                task = queue.wait(remaining)
            else:
                task = None
        if len(tasks) > 0:
            try:
                t = time.time()
                job_src.release(tasks)
                destruction_time += int((time.time() - t) * 1e6)
            except:
                tb = traceback.format_exc()
                logger.critical("cannot recover from the following exception:\n" + tb)
                for task in tasks:
                    logger.critical("tried to return task {0} from {1}".format(task.tag, task.hostname))
                raise
        if abort_threshold > 0 and successful_jobs >= abort_threshold and not abort_active:
            logger.info("activating fast abort with multiplier: {0}".format(abort_multiplier))
            abort_active = True
            queue.activate_fast_abort(abort_multiplier)

        # recurring actions are triggered here
        if actions:
            actions.take()
    if jobits_left == 0:
        logger.info("no more work left to do")
<commit_msg>Wait at least 10 seconds for returning tasks.<commit_after>import daemon
import datetime
import logging
import logging.handlers
import multiprocessing
import os
import resource
import signal
import sys
import threading
import time
import traceback
import yaml

from lobster import cmssw, job, util

from pkg_resources import get_distribution

import work_queue as wq

logger = multiprocessing.get_logger()

class ShortPathFormatter(logging.Formatter):
    def format(self, record):
        if len(record.pathname) > 40:
            record.pathname = '...' + record.pathname[-37:]
        # FIXME at some point, Formatter is a new-style class and we can
        # use super
        # return super(ShortPathFormatter, self).format(record)
        return logging.Formatter.format(self, record)

def kill(args):
    logger.info("setting flag to quit at the next checkpoint")
    with open(args.configfile) as configfile:
        config = yaml.load(configfile)

    workdir = config['workdir']
    util.register_checkpoint(workdir, 'KILLED', 'PENDING')

def run(args):
    with open(args.configfile) as configfile:
        config = yaml.load(configfile)

    workdir = config['workdir']
    if not os.path.exists(workdir):
        os.makedirs(workdir)
        util.register_checkpoint(workdir, "version", get_distribution('Lobster').version)
    else:
        util.verify(workdir)

    cmsjob = False
    if config.get('type', 'cmssw') == 'cmssw':
        cmsjob = True

        from ProdCommon.Credential.CredentialAPI import CredentialAPI
        cred = CredentialAPI({'credential': 'Proxy'})
        if cred.checkCredential(Time=60):
            if not 'X509_USER_PROXY' in os.environ:
                os.environ['X509_USER_PROXY'] = cred.credObj.getUserProxy()
        else:
            if config.get('advanced', {}).get('renew proxy', True):
                try:
                    cred.ManualRenewCredential()
                except Exception as e:
                    print("could not renew proxy")
                    sys.exit(1)
            else:
                print("please renew your proxy")
                sys.exit(1)

    print "Saving log to {0}".format(os.path.join(workdir, 'lobster.log'))

    if not args.foreground:
        ttyfile = open(os.path.join(workdir, 'lobster.err'), 'a')
        print "Saving stderr and stdout to {0}".format(os.path.join(workdir, 'lobster.err'))

    if config.get('advanced', {}).get('dump core', False):
        print "Setting core dump size to unlimited"
        resource.setrlimit(resource.RLIMIT_CORE, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))

    signals = daemon.daemon.make_default_signal_map()
    signals[signal.SIGTERM] = lambda num, frame: kill(args)

    with daemon.DaemonContext(
            detach_process=not args.foreground,
            stdout=sys.stdout if args.foreground else ttyfile,
            stderr=sys.stderr if args.foreground else ttyfile,
            working_directory=workdir,
            pidfile=util.get_lock(workdir, args.force),
            prevent_core=False,
            signal_map=signals):

        level = max(1, config.get('advanced', {}).get('log level', 2)) * 10
        fileh = logging.handlers.RotatingFileHandler(os.path.join(workdir, 'lobster.log'), maxBytes=500e6, backupCount=10)
        fileh.setFormatter(ShortPathFormatter("%(asctime)s [%(levelname)5s] - %(pathname)-40s %(lineno)4d: %(message)s"))
        fileh.setLevel(level)

        logger.addHandler(fileh)
        logger.setLevel(level)

        if args.foreground:
            console = logging.StreamHandler()
            console.setLevel(level)
            console.setFormatter(ShortPathFormatter("%(asctime)s [%(levelname)5s] - %(pathname)-40s %(lineno)4d: %(message)s"))
            logger.addHandler(console)

        config['configdir'] = args.configdir
        config['filename'] = args.configfile
        config['startdir'] = args.startdir

        t = threading.Thread(target=sprint, args=(config, workdir, cmsjob))
        t.start()
        t.join()

        logger.info("lobster terminated")

def sprint(config, workdir, cmsjob):
    if cmsjob:
        job_src = cmssw.JobProvider(config)
        actions = cmssw.Actions(config)
    else:
        job_src = job.SimpleJobProvider(config)
        actions = None

    logger.info("using wq from {0}".format(wq.__file__))

    wq.cctools_debug_flags_set("all")
    wq.cctools_debug_config_file(os.path.join(workdir, "work_queue_debug.log"))
    wq.cctools_debug_config_file_size(1 << 29)

    queue = wq.WorkQueue(-1)
    queue.specify_log(os.path.join(workdir, "work_queue.log"))
    queue.specify_name("lobster_" + config["id"])
    queue.specify_keepalive_timeout(300)
    # queue.tune("short-timeout", 600)
    queue.tune("transfer-outlier-factor", 4)
    queue.specify_algorithm(wq.WORK_QUEUE_SCHEDULE_RAND)

    cores = config.get('cores per job', 1)
    logger.info("starting queue as {0}".format(queue.name))
    logger.info("submit workers with: condor_submit_workers -M {0}{1} <num>".format(
        queue.name, ' --cores {0}'.format(cores) if cores > 1 else ''))

    payload = config.get('advanced', {}).get('payload', 10)
    abort_active = False
    abort_threshold = config.get('advanced', {}).get('abort threshold', 400)
    abort_multiplier = config.get('advanced', {}).get('abort multiplier', 4)

    if util.checkpoint(workdir, 'KILLED') == 'PENDING':
        util.register_checkpoint(workdir, 'KILLED', 'RESTART')

    # time in seconds to wait for WQ to return tasks, with minimum wait
    # time in case no more tasks are waiting
    interval = 60
    interval_minimum = 10

    jobits_left = 0
    successful_jobs = 0

    creation_time = 0
    destruction_time = 0

    with open(os.path.join(workdir, "lobster_stats.log"), "a") as statsfile:
        statsfile.write(
                "#timestamp " +
                "total_workers_connected total_workers_joined total_workers_removed " +
                "workers_busy workers_idle " +
                "tasks_running " +
                "total_send_time total_receive_time " +
                "total_create_time total_return_time " +
                "idle_percentage " +
                "capacity " +
                "efficiency " +
                "total_memory " +
                "total_cores " +
                "jobits_left\n")

    bad_exitcodes = job_src.bad_exitcodes

    while not job_src.done():
        jobits_left = job_src.work_left()
        stats = queue.stats_hierarchy

        with open(os.path.join(workdir, "lobster_stats.log"), "a") as statsfile:
            now = datetime.datetime.now()
            statsfile.write(" ".join(map(str,
                [
                    int(int(now.strftime('%s')) * 1e6 + now.microsecond),
                    stats.total_workers_connected,
                    stats.total_workers_joined,
                    stats.total_workers_removed,
                    stats.workers_busy,
                    stats.workers_idle,
                    stats.tasks_running,
                    stats.total_send_time,
                    stats.total_receive_time,
                    creation_time,
                    destruction_time,
                    stats.idle_percentage,
                    stats.capacity,
                    stats.efficiency,
                    stats.total_memory,
                    stats.total_cores,
                    jobits_left
                ]
                )) + "\n"
            )

        if util.checkpoint(workdir, 'KILLED') == 'PENDING':
            util.register_checkpoint(workdir, 'KILLED', str(datetime.datetime.utcnow()))

            # let the job source shut down gracefully
            logger.info("terminating job source")
            job_src.terminate()
            logger.info("terminating gracefully")
            break

        logger.info("{0} out of {1} workers busy; {3} jobs running, {4} waiting; {2} jobits left".format(
                stats.workers_busy,
                stats.workers_busy + stats.workers_ready,
                jobits_left,
                stats.tasks_running,
                stats.tasks_waiting))

        hunger = max(payload + stats.total_workers_connected / 10 - stats.tasks_waiting, 0)

        t = time.time()
        while hunger > 0:
            jobs = job_src.obtain(hunger)

            if jobs == None or len(jobs) == 0:
                break

            hunger -= len(jobs)
            for cores, cmd, id, inputs, outputs in jobs:
                task = wq.Task(cmd)
                task.specify_tag(id)
                task.specify_cores(cores)
                # temporary work-around?
                # task.specify_memory(1000)
                # task.specify_disk(4000)

                for (local, remote, cache) in inputs:
                    if os.path.isfile(local):
                        cache_opt = wq.WORK_QUEUE_CACHE if cache else wq.WORK_QUEUE_NOCACHE
                        task.specify_input_file(str(local), str(remote), cache_opt)
                    elif os.path.isdir(local):
                        task.specify_directory(str(local), str(remote), wq.WORK_QUEUE_INPUT,
                                wq.WORK_QUEUE_CACHE, recursive=True)
                    else:
                        logger.critical("cannot send file to worker: {0}".format(local))
                        raise NotImplementedError

                for (local, remote) in outputs:
                    task.specify_output_file(str(local), str(remote))

                queue.submit(task)
        creation_time += int((time.time() - t) * 1e6)

        job_src.update(queue)
        starttime = time.time()
        task = queue.wait(interval)
        tasks = []
        while task:
            if task.return_status == 0:
                successful_jobs += 1
            elif task.return_status in bad_exitcodes:
                logger.warning("blacklisting host {0} due to bad exit code from job {1}".format(task.hostname, task.tag))
                queue.blacklist(task.hostname)
            tasks.append(task)

            remaining = int(starttime + interval - time.time())
            if (interval - remaining > interval_minimum or queue.stats.tasks_waiting > 0) and remaining > 0:
                task = queue.wait(remaining)
            else:
                task = None
        if len(tasks) > 0:
            try:
                t = time.time()
                job_src.release(tasks)
                destruction_time += int((time.time() - t) * 1e6)
            except:
                tb = traceback.format_exc()
                logger.critical("cannot recover from the following exception:\n" + tb)
                for task in tasks:
                    logger.critical("tried to return task {0} from {1}".format(task.tag, task.hostname))
                raise
        if abort_threshold > 0 and successful_jobs >= abort_threshold and not abort_active:
            logger.info("activating fast abort with multiplier: {0}".format(abort_multiplier))
            abort_active = True
            queue.activate_fast_abort(abort_multiplier)

        # recurring actions are triggered here
        if actions:
            actions.take()
    if jobits_left == 0:
        logger.info("no more work left to do")
<|endoftext|>


<commit_before>"""
Usage:
    launch.py <init_filename>
    launch.py <init_filename> <targets_path>

Arguments:
    initExp_path (required) : path to a YAML file for arguments to launch experiment
    targets_path (optional) : yaml file that contains the targets information

Options:
    -h, --help

Example:

    > cd ~/Developer/NEXT/local/
    > python launch.py strange_fruit_triplet/init.yaml strange_fruit_triplet/targets.yaml

"""

import os
import sys
from collections import OrderedDict
import base64
import yaml
import requests

sys.path.append('../next/lib')
from docopt import docopt

def get_backend():
    """
    Get the backend host and port.

    Defaults to dict(host='localhost', port=8000).
    """
    return {
        'host': os.environ.get('NEXT_BACKEND_GLOBAL_HOST', 'localhost'),
        'port': os.environ.get('NEXT_BACKEND_GLOBAL_PORT', '8000')
    }

def launch(init_filename, targets_filename=None):
    with open(init_filename, 'r') as f:
        init = yaml.load(f)

    if targets_filename:
        with open(targets_filename, 'r') as f:
            targets = yaml.load(f)
            if 'targets' not in init['args']:
                init['args'] = {'targetset': {}}
            if 'targetset' not in init['args']['targets']:
                init['args']['targets']['targetset'] = dict()
            init['args']['targets']['targetset'] = targets

    header = "data:application/{};base64,"
    exp_args = base64.encodestring(yaml.dump(init))
    d = {'args': header.format('x-yaml') + exp_args}

    d = OrderedDict(d)

    header = ['{}:{}'.format(key, len(item)) for key, item in d.items()]
    header = ';'.join(header) + '\n'

    to_send = ''.join([item for _, item in d.items()])

    data = header + to_send
    # -- send the packed experiment definition file
    host = get_backend()
    host_url = "http://{host}:{port}".format(**host)

    r = requests.post(host_url + '/assistant/init/experiment', data=data)
    response = r.json()
    if not response['success']:
        print('An error occurred launching the experiment:')
        print(response['message'])
        sys.exit()

    dashboard_url = host_url + '/dashboard/experiment_dashboard/{}/{}'.format(response['exp_uid'], init['app_id'])
    print('Dashboard URL: {}'.format(dashboard_url))
    print('NEXT Home URL: {}'.format(host_url + '/home'))

    return response

if __name__ == "__main__":
    args = docopt(__doc__)
    launch(args['<init_filename>'], targets_filename=args['<targets_path>'])
<commit_msg>local/launch.py: Clean up API payload encoder<commit_after>"""
Usage:
    launch.py <init_filename>
    launch.py <init_filename> <targets_path>

Arguments:
    initExp_path (required) : path to a YAML file for arguments to launch experiment
    targets_path (optional) : yaml file that contains the targets information

Options:
    -h, --help

Example:

    > cd ~/Developer/NEXT/local/
    > python launch.py strange_fruit_triplet/init.yaml strange_fruit_triplet/targets.yaml

"""

import os
import sys
from collections import OrderedDict
import base64
import yaml
import requests

sys.path.append('../next/lib')
from docopt import docopt

def get_backend():
    """
    Get the backend host and port.

    Defaults to dict(host='localhost', port=8000).
    """
    return {
        'host': os.environ.get('NEXT_BACKEND_GLOBAL_HOST', 'localhost'),
        'port': os.environ.get('NEXT_BACKEND_GLOBAL_PORT', '8000')
    }

def launch(init_filename, targets_filename=None):
    with open(init_filename, 'r') as f:
        init = yaml.load(f)

    if targets_filename:
        with open(targets_filename, 'r') as f:
            targets = yaml.load(f)
            if 'targets' not in init['args']:
                init['args'] = {'targetset': {}}
            if 'targetset' not in init['args']['targets']:
                init['args']['targets']['targetset'] = dict()
            init['args']['targets']['targetset'] = targets

    # -- encode the experiment definition for transmission to the backend
    data_header = "data:application/x-yaml;base64,"

    encoded_args = base64.encodestring(yaml.dump(init))
    encoded_attrs = OrderedDict(args=data_header+encoded_args)

    # generate metadata describing each attr's length; this is prepended to the request data.
    metadata = ';'.join(['{}:{}'.format(k, len(v)) for k, v in encoded_attrs.items()])

    # concat all the encoded attrs together
    body = ''.join([v for _, v in encoded_attrs.items()])

    # what we're actually going to send
    payload = metadata + '\n' + body

    # -- send the packed experiment definition file
    host = get_backend()
    host_url = "http://{host}:{port}".format(**host)

    r = requests.post(host_url + '/assistant/init/experiment', data=payload)
    response = r.json()
    if not response['success']:
        print('An error occurred launching the experiment:')
        print(response['message'])
        sys.exit()

    dashboard_url = host_url + '/dashboard/experiment_dashboard/{}/{}'.format(response['exp_uid'], init['app_id'])
    print('Dashboard URL: {}'.format(dashboard_url))
    print('NEXT Home URL: {}'.format(host_url + '/home'))

    return response

if __name__ == "__main__":
    args = docopt(__doc__)
    launch(args['<init_filename>'], targets_filename=args['<targets_path>'])
<|endoftext|>


<commit_before>#!/usr/bin/env python
# vim: ai ts=4 sts=4 et sw=4

import Queue, sys, datetime

def _logging_method (level):
    return lambda self, msg, *args: self.log(level, msg, *args)

class Component(object):
    @classmethod
    def config_options (cls):
        if hasattr(cls, "__init__"):
            # deep Python voodoo to discover the keyword
            # arguments of the class's constructor method
            return cls.__init__.func_code.co_varnames
        else:
            return ()

    @property
    def router (self):
        if hasattr(self, "_router"):
            return self._router

    def _get_name(self):
        if hasattr(self, '_name'):
            return self._name
        else:
            return unicode(type(self).__name__)

    def _set_name(self, name):
        self._name = name

    name = property(_get_name, _set_name)

    def log(self, level, msg, *args):
        if self.router:
            self.router.logger.write(self, level, msg, *args)

    debug    = _logging_method('debug')
    info     = _logging_method('info')
    warning  = _logging_method('warning')
    error    = _logging_method('error')
    critical = _logging_method('critical')

class Receiver(Component):
    def __init__(self):
        # do we want to put a limit on the queue size?
        # and what do we do if the queue gets full?
        self._queue = Queue.Queue()

    @property
    def message_waiting (self):
        return self._queue.qsize()

    def next_message (self, timeout=0.0):
        try:
            return self._queue.get(bool(timeout), timeout)
        except Queue.Empty:
            return None

    def send(self, message):
        # block until we can add to the queue.
        # it shouldn't be that long.
        self._queue.put(message, True)
<commit_msg>Remove the foolhardy introspection code from component.<commit_after>#!/usr/bin/env python
# vim: ai ts=4 sts=4 et sw=4

import Queue, sys, datetime

def _logging_method (level):
    return lambda self, msg, *args: self.log(level, msg, *args)

class Component(object):
    @property
    def router (self):
        if hasattr(self, "_router"):
            return self._router

    def _get_name(self):
        if hasattr(self, '_name'):
            return self._name
        else:
            return unicode(type(self).__name__)

    def _set_name(self, name):
        self._name = name

    name = property(_get_name, _set_name)

    def log(self, level, msg, *args):
        if self.router:
            self.router.logger.write(self, level, msg, *args)

    debug    = _logging_method('debug')
    info     = _logging_method('info')
    warning  = _logging_method('warning')
    error    = _logging_method('error')
    critical = _logging_method('critical')

class Receiver(Component):
    def __init__(self):
        # do we want to put a limit on the queue size?
        # and what do we do if the queue gets full?
        self._queue = Queue.Queue()

    @property
    def message_waiting (self):
        return self._queue.qsize()

    def next_message (self, timeout=0.0):
        try:
            return self._queue.get(bool(timeout), timeout)
        except Queue.Empty:
            return None

    def send(self, message):
        # block until we can add to the queue.
        # it shouldn't be that long.
        self._queue.put(message, True)
<|endoftext|>


<commit_before>from lxml import etree
import sys
import pdb
from glob import glob
import csv
import os
from os import listdir
from os.path import isfile, join
import subprocess
import base64
import time
import re
import smtplib
import argparse
import mimetypes
from email.mime.multipart import MIMEMultipart
from email import encoders
from email.message import Message
from email.mime.audio import MIMEAudio
from email.mime.base import MIMEBase
from email.mime.image import MIMEImage
from email.mime.text import MIMEText
import bagit

parser = argparse.ArgumentParser(description='DCP FIXITY checker/bagging tool.'
                                 ' Written by Kieran O\'Leary.')
parser.add_argument('input')
parser.add_argument(
                    '-bag',
                    action='store_true',help='bag the dcp_dir if it passes the hash check')
parser.add_argument(
                    '-m',
                    action='store_true',help='send email report')
args = parser.parse_args()

if args.bag:
    bagging = 'enabled'
else:
    bagging = 'disabled'

if args.m:
    email = 'enabled'
else:
    email = 'disabled'
#bagrm =  os.path.abspath('bag-rm.py')
#bagit =  os.path.abspath('bagit.py')
#print bagrm
dcp_dir = args.input
# Two csv functions. One to create a csv, the other to add info to.
def create_csv(csv_file, *args):
    f = open(csv_file, 'wb')
    try:
        writer = csv.writer(f)
        writer.writerow(*args)
    finally:
        f.close()

def append_csv(csv_file, *args):
    f = open(csv_file, 'ab')
    try:
        writer = csv.writer(f)
        writer.writerow(*args)
    finally:
        f.close()

# Create a new .csv file with headings.
# CSV filename will be DCp directory name + time/date.

# CSV will be saved to your Desktop.

csv_report_filename = os.path.basename(dcp_dir) + '_dcp_level' + time.strftime("_%Y_%m_%dT%H_%M_%S")
# CSV will be saved to your Desktop.
csv_report = os.path.expanduser("~/Desktop/%s.csv") % csv_report_filename
create_csv(csv_report, ('DCP NAME', 'DIRECTORY NAME', 'JUDGEMENT'))
for root,dirnames,filenames in os.walk(dcp_dir):
    if ("ASSETMAP.xml"  in filenames) or ("ASSETMAP"  in filenames) :
        dir = root
        #print os.path.basename(os.path.dirname(root))
        filenoext = os.path.splitext(os.path.dirname(root))[0]
        #print filenoext + 'dfsdfjkljoewuiljkdfs'
        # Change directory to directory with video files

        #print filenoext
        # Generate new directory names in AIP
        fixity_dir   = "%s/fixity" % filenoext
        os.makedirs(fixity_dir)
        csvfile = fixity_dir + '/' + os.path.basename(os.path.dirname(root)) + '_item_level' + time.strftime("_%Y_%m_%dT%H_%M_%S") + '.csv'
        print 'csvfile is ', csvfile
        #csv_file = os.path.expanduser("~/Desktop/%s.csv") % csv_filename
        create_csv(csvfile, ('MXF HASH', 'STORED HASH', 'FILENAME', 'JUDGEMENT'))

        # Changing directory makes globbing easier (from my experience anyhow).
        os.chdir(dir)

        # Scan the main DCP directory for an assetmap.
        dcp_files = [f for f in listdir(dir) if isfile(join(dir, f))]
        if 'ASSETMAP' in dcp_files:
            assetmap = 'ASSETMAP'
        elif 'ASSETMAP.xml' in dcp_files:
            assetmap = 'ASSETMAP.xml'

        # Parse the assetmap in order to find the namespace.
        try:
            assetmap_xml = etree.parse(assetmap)
        except SyntaxError:
            append_csv(csvfile,('NOT A VALID ASSETMAP', 'NOT A VALID ASSETMAP', dir,'NOT A VALID ASSETMAP'))
            append_csv(csv_report,(os.path.basename(dir), dir, 'NOT A VALID ASSETMAP'))
            print 'not an assetmap!!!!'
            continue

        print assetmap_xml
        assetmap_namespace = assetmap_xml.xpath('namespace-uri(.)')

        # Get a list of all XML files in the main DCP directory.
        xmlfiles = glob('*.xml')

        # Generate an empty list as there may be multiple PKLs.
        pkl_list = []

        # Loop through xmlfiles in order to find any PKL files.
        for i in xmlfiles:
            try:
                xmlname = etree.parse(i)
            except SyntaxError:
                append_csv(csvfile,('NOT A VALID PKL', 'NOT A VALID PKL', dir,'NOT A VALID PKL'))
                append_csv(csv_report,(os.path.basename(dir), dir, 'NOT A VALID PKL'))
                print 'not a valid PKL!!!!'
                continue
            except KeyError:
                append_csv(csvfile,('PKL APPEARS TO BE MISSING', 'PKL APPEARS TO BE MISSING', dir,'PKL APPEARS TO BE MISSING'))
                append_csv(csv_report,(os.path.basename(dir), dir, 'PKL APPEARS TO BE MISSING'))
                print 'Missing PKL!!!!'
                continue

            is_pkl = xmlname.xpath('namespace-uri(.)')
            print is_pkl
            if 'PKL' in is_pkl:
                pkl_list.append(i)

        if len(pkl_list) == 0:

            append_csv(csvfile,('PKL APPEARS TO BE MISSING', 'PKL APPEARS TO BE MISSING', dir,'PKL APPEARS TO BE MISSING'))
            append_csv(csv_report,(os.path.basename(dir), dir, 'PKL APPEARS TO BE MISSING'))
            continue

        # Generate an empty dictionary that will link the PKL hashes to each UUID.
        pkl_hashes = {}

        # Loop through the PKLs and link each hash to a UUID.
        for i in pkl_list:
            pkl_parse = etree.parse(i)
            pkl_namespace = pkl_parse.xpath('namespace-uri(.)')
            hashes =  pkl_parse.findall('//ns:Hash',namespaces={'ns': pkl_namespace})
            xmluuid =  pkl_parse.findall('//ns:Asset/ns:Id',namespaces={'ns': pkl_namespace})

            counter = 0

            while counter <= len(hashes) -1 : # The -1 is there because of lxml's zero indexing.
                pkl_hashes[xmluuid[counter].text] = hashes[counter].text # {pkl_uuid:pkl_hash}
                counter +=1

        # Begin analysis of assetmap xml.
        counter = 0
        assetmap_paths =  assetmap_xml.findall('//ns:Path',namespaces={'ns': assetmap_namespace})
        assetmap_uuids =  assetmap_xml.findall('//ns:Asset/ns:Id',namespaces={'ns': assetmap_namespace})
        #while counter <= len(assetmap_paths) -1 :

        counter = 0

        file_paths = {}

        while counter <= len(assetmap_paths) -1 :
            print assetmap_paths[counter].text
            if 'file:///' in assetmap_paths[counter].text:
                remove_this = 'file:///'
                assetmap_paths[counter].text =  assetmap_paths[counter].text.replace(remove_this,"")
            elif 'file://' in assetmap_paths[counter].text:
                remove_this = 'file://'
                assetmap_paths[counter].text =  assetmap_paths[counter].text.replace(remove_this,"")

            elif 'file:/' in assetmap_paths[counter].text:
                remove_this = 'file:/'
                assetmap_paths[counter].text =  assetmap_paths[counter].text.replace(remove_this,"")

            file_paths[assetmap_uuids[counter].text] = [assetmap_paths[counter].text] # {assetmapuuid:assetmapfilename}
            counter +=1

        # Removes PKLs from list of files to hash, as these files are not in manifest.
        keys_to_remove = []

        for i in file_paths:
            if file_paths[i][0] in pkl_list:
                keys_to_remove.append(i)
        # PKL files are deleted from the file_paths dictionary.
        for i in keys_to_remove:
            del file_paths[i]

        # Check if there are any files missing from the DCP.

        missing_files = []
        for i in file_paths:
            if not os.path.isfile(file_paths[i][0]): # This checks if the file exists.
                print time.strftime("%Y-%m-%dT%H:%M:%S") + ' - **********' + file_paths[i][0] + ' is missing **********'
                missing_files.append(i)
                # Add missing file info to the csv.
                append_csv(csvfile,('MISSING FILE', pkl_hashes[i], os.path.abspath(file_paths[i][0]),'MISSING FILE'))

        # This removes the missing files from the hashable list.
        for i in missing_files:
            del file_paths[i]
            del pkl_hashes[i]

        # Generate fresh hashes on the actual files in the DCP.
        for i in file_paths:
            print time.strftime("%Y-%m-%dT%H:%M:%S") + ' - Generating fresh hash for ' + file_paths[i][0]
            # Create SHA-1 binary hashes with OPENSSL.

            openssl_hash = subprocess.check_output(['openssl',
                                                    'sha1',
                                                    '-binary',
                                                     file_paths[i][0]])
            # Encode the hashes as base64.
            b64hash =  base64.b64encode(openssl_hash)
            # Append hashes to the list within the file_paths dictionary.
            file_paths[i].append(b64hash)

        # Compare the hashes in the PKL manifest to the fresh hashes.
        hash_mismatches = []
        for i in file_paths:
            if file_paths[i][1] == pkl_hashes[i]:
                print file_paths[i][0] + ' is ok'
                append_csv(csvfile,(file_paths[i][1], pkl_hashes[i], os.path.abspath(file_paths[i][0]),'HASH MATCH'))
            else:
                print file_paths[i][0] + ' mismatch'
                hash_mismatches.append(file_paths[i][0])
                append_csv(csvfile,(file_paths[i][1], pkl_hashes[i], os.path.abspath(file_paths[i][0]),'HASH MISMATCH'))

        if len(hash_mismatches) > 0:
            report = ' but THERE ARE HASH MISMATCHES. SCROLL UP FOR MORE INFO OR CHECK THE CSV'
            baggable = 'n'
        else:
            report = ' and all hashes match.'
            baggable = 'y'

        if len(missing_files) > 0:
            print time.strftime("%Y-%m-%dT%H:%M:%S") + ' - WARNING - THERE ARE FILES MISSING FROM THIS DCP. SCROLL UP FOR MORE INFO OR CHECK THE CSV'
            append_csv(csv_report,(os.path.dirname(root), dir, 'FILES MISSING - CHECK REPORT'))
        else:
            print time.strftime("%Y-%m-%dT%H:%M:%S") + ' - All files are present in your DCP' + report
            append_csv(csv_report,(os.path.dirname(root), dir,'All files present ' + report))

        if bagging == 'enabled':
            if baggable == 'y':
                #pdb.set_trace()
                if os.path.dirname(root) != args.input:

                    print 1
                    dir = os.path.dirname(root)
                    print dir
                    print args.input
                    bag = bagit.make_bag(dir)
                else:
                     print 'bagging not supported for this folder structure right now'
if email == 'enabled':
    emailfrom = ""
    emailto = ['', '']
    #emailto = ", ".join(emailto)
    fileToSend = csv_report
    username = ""
    password = ""

    msg = MIMEMultipart()
    msg["From"] = emailfrom
    msg["To"] = ", ".join(emailto)
    msg["Subject"] = "Hash check complete"
    msg.preamble = "testtesttest"
    body = MIMEText("example email body")
    msg.attach(body)

    ctype, encoding = mimetypes.guess_type(fileToSend)
    if ctype is None or encoding is not None:
        ctype = "application/octet-stream"

    maintype, subtype = ctype.split("/", 1)

    if maintype == "text":
        fp = open(fileToSend)
        # Note: we should handle calculating the charset
        attachment = MIMEText(fp.read(), _subtype=subtype)
        fp.close()
    elif maintype == "image":
        fp = open(fileToSend, "rb")
        attachment = MIMEImage(fp.read(), _subtype=subtype)
        fp.close()
    elif maintype == "audio":
        fp = open(fileToSend, "rb")
        attachment = MIMEAudio(fp.read(), _subtype=subtype)
        fp.close()
    else:
        fp = open(fileToSend, "rb")
        attachment = MIMEBase(maintype, subtype)
        attachment.set_payload(fp.read())
        fp.close()
        encoders.encode_base64(attachment)
    attachment.add_header("Content-Disposition", "attachment", filename=fileToSend)
    msg.attach(attachment)

    server_ssl = smtplib.SMTP_SSL("smtp.gmail.com", 465)
    server_ssl.ehlo() # optional, called by login()
    server_ssl.login(username, password)
    # ssl server doesn't support or need tls, so don't call server_ssl.starttls()
    server_ssl.sendmail(emailfrom, emailto, msg.as_string())
    print msg.as_string()
    #server_ssl.quit()
    server_ssl.close()
    print 'successfully sent the mail'
<commit_msg>dcpfixity - csv sidecar is now optional<commit_after>from lxml import etree
import sys
import pdb
from glob import glob
import csv
import os
from os import listdir
from os.path import isfile, join
import subprocess
import base64
import time
import re
import smtplib
import argparse
import mimetypes
from email.mime.multipart import MIMEMultipart
from email import encoders
from email.message import Message
from email.mime.audio import MIMEAudio
from email.mime.base import MIMEBase
from email.mime.image import MIMEImage
from email.mime.text import MIMEText
import bagit

parser = argparse.ArgumentParser(description='DCP FIXITY checker/bagging tool.'
                                 ' Written by Kieran O\'Leary.')
parser.add_argument('input')
parser.add_argument(
                    '-bag',
                    action='store_true',help='bag the dcp_dir if it passes the hash check')
parser.add_argument(
                    '-csv',
                    action='store_true',help='File level csv is stored as sidecar to dcp directory')
parser.add_argument(
                    '-m',
                    action='store_true',help='send email report')
args = parser.parse_args()

if args.bag:
    bagging = 'enabled'
else:
    bagging = 'disabled'

if args.m:
    email = 'enabled'
else:
    email = 'disabled'
#bagrm =  os.path.abspath('bag-rm.py')
#bagit =  os.path.abspath('bagit.py')
#print bagrm
dcp_dir = args.input
# Two csv functions. One to create a csv, the other to add info to.
def create_csv(csv_file, *args):
    f = open(csv_file, 'wb')
    try:
        writer = csv.writer(f)
        writer.writerow(*args)
    finally:
        f.close()

def append_csv(csv_file, *args):
    f = open(csv_file, 'ab')
    try:
        writer = csv.writer(f)
        writer.writerow(*args)
    finally:
        f.close()

# Create a new .csv file with headings.
# CSV filename will be DCp directory name + time/date.

# CSV will be saved to your Desktop.

csv_report_filename = os.path.basename(dcp_dir) + '_dcp_level' + time.strftime("_%Y_%m_%dT%H_%M_%S")
# CSV will be saved to your Desktop.

csv_report = os.path.expanduser("~/Desktop/%s.csv") % csv_report_filename
create_csv(csv_report, ('DCP NAME', 'DIRECTORY NAME', 'JUDGEMENT'))

if args.csv:
    csv = 'enabled'

else:
    csv_filename = os.path.basename(dcp_dir) + '_file_level' + time.strftime("_%Y_%m_%dT%H_%M_%S")
    csvfile = os.path.expanduser("~/Desktop/%s.csv") % csv_filename
    print 'csvfile is ', csvfile
    #csv_file = os.path.expanduser("~/Desktop/%s.csv") % csv_filename
    create_csv(csvfile, ('MXF HASH', 'STORED HASH', 'FILENAME', 'JUDGEMENT'))
for root,dirnames,filenames in os.walk(dcp_dir):
    if ("ASSETMAP.xml"  in filenames) or ("ASSETMAP"  in filenames) :
        dir = root
        #print os.path.basename(os.path.dirname(root))
        filenoext = os.path.splitext(os.path.dirname(root))[0]
        #print filenoext + 'dfsdfjkljoewuiljkdfs'
        # Change directory to directory with video files

        if csv == 'enabled':
            print 'enableeeeeeed'
            #print filenoext
            # Generate new directory names in AIP
            fixity_dir   = "%s/fixity" % filenoext
            os.makedirs(fixity_dir)
            csvfile = fixity_dir + '/' + os.path.basename(os.path.dirname(root)) + '_item_level' + time.strftime("_%Y_%m_%dT%H_%M_%S") + '.csv'
            print 'csvfile is ', csvfile
            #csv_file = os.path.expanduser("~/Desktop/%s.csv") % csv_filename
            create_csv(csvfile, ('MXF HASH', 'STORED HASH', 'FILENAME', 'JUDGEMENT'))

        # Changing directory makes globbing easier (from my experience anyhow).
        os.chdir(dir)

        # Scan the main DCP directory for an assetmap.
        dcp_files = [f for f in listdir(dir) if isfile(join(dir, f))]
        if 'ASSETMAP' in dcp_files:
            assetmap = 'ASSETMAP'
        elif 'ASSETMAP.xml' in dcp_files:
            assetmap = 'ASSETMAP.xml'

        # Parse the assetmap in order to find the namespace.
        try:
            assetmap_xml = etree.parse(assetmap)
        except SyntaxError:
            append_csv(csvfile,('NOT A VALID ASSETMAP', 'NOT A VALID ASSETMAP', dir,'NOT A VALID ASSETMAP'))
            append_csv(csv_report,(os.path.basename(dir), dir, 'NOT A VALID ASSETMAP'))
            print 'not an assetmap!!!!'
            continue

        print assetmap_xml
        assetmap_namespace = assetmap_xml.xpath('namespace-uri(.)')

        # Get a list of all XML files in the main DCP directory.
        xmlfiles = glob('*.xml')

        # Generate an empty list as there may be multiple PKLs.
        pkl_list = []

        # Loop through xmlfiles in order to find any PKL files.
        for i in xmlfiles:
            try:
                xmlname = etree.parse(i)
            except SyntaxError:
                append_csv(csvfile,('NOT A VALID PKL', 'NOT A VALID PKL', dir,'NOT A VALID PKL'))
                append_csv(csv_report,(os.path.basename(dir), dir, 'NOT A VALID PKL'))
                print 'not a valid PKL!!!!'
                continue
            except KeyError:
                append_csv(csvfile,('PKL APPEARS TO BE MISSING', 'PKL APPEARS TO BE MISSING', dir,'PKL APPEARS TO BE MISSING'))
                append_csv(csv_report,(os.path.basename(dir), dir, 'PKL APPEARS TO BE MISSING'))
                print 'Missing PKL!!!!'
                continue

            is_pkl = xmlname.xpath('namespace-uri(.)')
            print is_pkl
            if 'PKL' in is_pkl:
                pkl_list.append(i)

        if len(pkl_list) == 0:

            append_csv(csvfile,('PKL APPEARS TO BE MISSING', 'PKL APPEARS TO BE MISSING', dir,'PKL APPEARS TO BE MISSING'))
            append_csv(csv_report,(os.path.basename(dir), dir, 'PKL APPEARS TO BE MISSING'))
            continue

        # Generate an empty dictionary that will link the PKL hashes to each UUID.
        pkl_hashes = {}

        # Loop through the PKLs and link each hash to a UUID.
        for i in pkl_list:
            pkl_parse = etree.parse(i)
            pkl_namespace = pkl_parse.xpath('namespace-uri(.)')
            hashes =  pkl_parse.findall('//ns:Hash',namespaces={'ns': pkl_namespace})
            xmluuid =  pkl_parse.findall('//ns:Asset/ns:Id',namespaces={'ns': pkl_namespace})

            counter = 0

            while counter <= len(hashes) -1 : # The -1 is there because of lxml's zero indexing.
                pkl_hashes[xmluuid[counter].text] = hashes[counter].text # {pkl_uuid:pkl_hash}
                counter +=1

        # Begin analysis of assetmap xml.
        counter = 0
        assetmap_paths =  assetmap_xml.findall('//ns:Path',namespaces={'ns': assetmap_namespace})
        assetmap_uuids =  assetmap_xml.findall('//ns:Asset/ns:Id',namespaces={'ns': assetmap_namespace})
        #while counter <= len(assetmap_paths) -1 :

        counter = 0

        file_paths = {}

        while counter <= len(assetmap_paths) -1 :
            print assetmap_paths[counter].text
            if 'file:///' in assetmap_paths[counter].text:
                remove_this = 'file:///'
                assetmap_paths[counter].text =  assetmap_paths[counter].text.replace(remove_this,"")
            elif 'file://' in assetmap_paths[counter].text:
                remove_this = 'file://'
                assetmap_paths[counter].text =  assetmap_paths[counter].text.replace(remove_this,"")

            elif 'file:/' in assetmap_paths[counter].text:
                remove_this = 'file:/'
                assetmap_paths[counter].text =  assetmap_paths[counter].text.replace(remove_this,"")

            file_paths[assetmap_uuids[counter].text] = [assetmap_paths[counter].text] # {assetmapuuid:assetmapfilename}
            counter +=1

        # Removes PKLs from list of files to hash, as these files are not in manifest.
        keys_to_remove = []

        for i in file_paths:
            if file_paths[i][0] in pkl_list:
                keys_to_remove.append(i)
        # PKL files are deleted from the file_paths dictionary.
        for i in keys_to_remove:
            del file_paths[i]

        # Check if there are any files missing from the DCP.

        missing_files = []
        for i in file_paths:
            if not os.path.isfile(file_paths[i][0]): # This checks if the file exists.
                print time.strftime("%Y-%m-%dT%H:%M:%S") + ' - **********' + file_paths[i][0] + ' is missing **********'
                missing_files.append(i)
                # Add missing file info to the csv.
                append_csv(csvfile,('MISSING FILE', pkl_hashes[i], os.path.abspath(file_paths[i][0]),'MISSING FILE'))

        # This removes the missing files from the hashable list.
        for i in missing_files:
            del file_paths[i]
            del pkl_hashes[i]

        # Generate fresh hashes on the actual files in the DCP.
        for i in file_paths:
            print time.strftime("%Y-%m-%dT%H:%M:%S") + ' - Generating fresh hash for ' + file_paths[i][0]
            # Create SHA-1 binary hashes with OPENSSL.

            openssl_hash = subprocess.check_output(['openssl',
                                                    'sha1',
                                                    '-binary',
                                                     file_paths[i][0]])
            # Encode the hashes as base64.
            b64hash =  base64.b64encode(openssl_hash)
            # Append hashes to the list within the file_paths dictionary.
            file_paths[i].append(b64hash)

        # Compare the hashes in the PKL manifest to the fresh hashes.
        hash_mismatches = []
        for i in file_paths:
            if file_paths[i][1] == pkl_hashes[i]:
                print file_paths[i][0] + ' is ok'
                append_csv(csvfile,(file_paths[i][1], pkl_hashes[i], os.path.abspath(file_paths[i][0]),'HASH MATCH'))
            else:
                print file_paths[i][0] + ' mismatch'
                hash_mismatches.append(file_paths[i][0])
                append_csv(csvfile,(file_paths[i][1], pkl_hashes[i], os.path.abspath(file_paths[i][0]),'HASH MISMATCH'))

        if len(hash_mismatches) > 0:
            report = ' but THERE ARE HASH MISMATCHES. SCROLL UP FOR MORE INFO OR CHECK THE CSV'
            baggable = 'n'
        else:
            report = ' and all hashes match.'
            baggable = 'y'

        if len(missing_files) > 0:
            print time.strftime("%Y-%m-%dT%H:%M:%S") + ' - WARNING - THERE ARE FILES MISSING FROM THIS DCP. SCROLL UP FOR MORE INFO OR CHECK THE CSV'
            append_csv(csv_report,(os.path.dirname(root), dir, 'FILES MISSING - CHECK REPORT'))
        else:
            print time.strftime("%Y-%m-%dT%H:%M:%S") + ' - All files are present in your DCP' + report
            append_csv(csv_report,(os.path.dirname(root), dir,'All files present ' + report))

        if bagging == 'enabled':
            if baggable == 'y':
                #pdb.set_trace()
                if os.path.dirname(root) != args.input:

                    print 1
                    dir = os.path.dirname(root)
                    print dir
                    print args.input
                    bag = bagit.make_bag(dir)
                else:
                     print 'bagging not supported for this folder structure right now'
if email == 'enabled':
    emailfrom = ""
    emailto = ['', '']
    #emailto = ", ".join(emailto)
    fileToSend = csv_report
    username = ""
    password = ""

    msg = MIMEMultipart()
    msg["From"] = emailfrom
    msg["To"] = ", ".join(emailto)
    msg["Subject"] = "Hash check complete"
    msg.preamble = "testtesttest"
    body = MIMEText("example email body")
    msg.attach(body)

    ctype, encoding = mimetypes.guess_type(fileToSend)
    if ctype is None or encoding is not None:
        ctype = "application/octet-stream"

    maintype, subtype = ctype.split("/", 1)

    if maintype == "text":
        fp = open(fileToSend)
        # Note: we should handle calculating the charset
        attachment = MIMEText(fp.read(), _subtype=subtype)
        fp.close()
    elif maintype == "image":
        fp = open(fileToSend, "rb")
        attachment = MIMEImage(fp.read(), _subtype=subtype)
        fp.close()
    elif maintype == "audio":
        fp = open(fileToSend, "rb")
        attachment = MIMEAudio(fp.read(), _subtype=subtype)
        fp.close()
    else:
        fp = open(fileToSend, "rb")
        attachment = MIMEBase(maintype, subtype)
        attachment.set_payload(fp.read())
        fp.close()
        encoders.encode_base64(attachment)
    attachment.add_header("Content-Disposition", "attachment", filename=fileToSend)
    msg.attach(attachment)

    server_ssl = smtplib.SMTP_SSL("smtp.gmail.com", 465)
    server_ssl.ehlo() # optional, called by login()
    server_ssl.login(username, password)
    # ssl server doesn't support or need tls, so don't call server_ssl.starttls()
    server_ssl.sendmail(emailfrom, emailto, msg.as_string())
    print msg.as_string()
    #server_ssl.quit()
    server_ssl.close()
    print 'successfully sent the mail'
<|endoftext|>