<commit_before># -*- coding: utf-8 -*-
"""Download and extract the IXI Hammersmith Hospital 3T dataset

url: http://brain-development.org/ixi-dataset/
ref: IXI – Information eXtraction from Images (EPSRC GR/S21533/02)

"""

import urllib
import os.path
import tarfile
import pandas as pd
import glob
import SimpleITK as sitk
import numpy as np

DOWNLOAD_IMAGES = True
EXTRACT_IMAGES = True
PROCESS_OTHER = True
RESAMPLE_IMAGES = True
CLEAN_UP = True

def resample_image(itk_image, out_spacing=[1.0, 1.0, 1.0], is_label=False):

    original_spacing = itk_image.GetSpacing()
    original_size = itk_image.GetSize()

    out_size = [int(np.round(original_size[0]*(original_spacing[0]/out_spacing[0]))),
                int(np.round(original_size[1]*(original_spacing[1]/out_spacing[1]))),
                int(np.round(original_size[2]*(original_spacing[2]/out_spacing[2])))]

    resample = sitk.ResampleImageFilter()
    resample.SetOutputSpacing(out_spacing)
    resample.SetSize(out_size)
    resample.SetOutputDirection(itk_image.GetDirection())
    resample.SetOutputOrigin(itk_image.GetOrigin())
    resample.SetTransform(sitk.Transform())
    resample.SetDefaultPixelValue(itk_image.GetPixelIDValue())

    if is_label:
        resample.SetInterpolator(sitk.sitkNearestNeighbor)
    else:
        resample.SetInterpolator(sitk.sitkBSpline)

    return resample.Execute(itk_image)

def reslice_image(itk_image, itk_ref, is_label=False):

    resample = sitk.ResampleImageFilter()
    resample.SetReferenceImage(itk_ref)

    if is_label:
        resample.SetInterpolator(sitk.sitkNearestNeighbor)
    else:
        resample.SetInterpolator(sitk.sitkBSpline)

    return resample.Execute(itk_image)

urls = {}
urls['t1'] = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-T1.tar'
urls['t2'] = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-T2.tar'
urls['pd'] = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-PD.tar'
urls['mra'] = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-MRA.tar'
urls['demographic'] = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI.xls'

fnames = {}
fnames['t1'] = 't1.tar'
fnames['t2'] = 't2.tar'
fnames['pd'] = 'pd.tar'
fnames['mra'] = 'mra.tar'
fnames['demographic'] = 'demographic.xls'

if DOWNLOAD_IMAGES:
    # Download all IXI data
    for key, url in urls.items():

        if not os.path.isfile(fnames[key]):
            print('Downloading {} from {}'.format(fnames[key], url))
            curr_file = urllib.FancyURLopener()
            curr_file.retrieve(url, fnames[key])
        else:
            print('File {} already exists. Skipping download.'.format(fnames[key]))

if EXTRACT_IMAGES:
    # Extract the HH subset of IXI
    for key, fname in fnames.items():

        if (fname.endswith('.tar')):
            print('Extracting IXI HH data from {}.'.format(fnames[key]))
            output_dir = os.path.join('./orig/', key)

            if not os.path.exists(output_dir):
                os.makedirs(output_dir)

            t = tarfile.open(fname, 'r')
            for member in t.getmembers():
                if '-HH-' in member.name:
                    t.extract(member, output_dir)

if PROCESS_OTHER:
    # Process the demographic xls data and save to csv
    xls = pd.ExcelFile('demographic.xls')
    print(xls.sheet_names)

    df = xls.parse('Table')
    for index, row in df.iterrows():
        IXI_id = 'IXI{:03d}'.format(row['IXI_ID'])
        df.loc[index, 'IXI_ID'] = IXI_id

        t1_exists = len(glob.glob('./orig/t1/{}*.nii.gz'.format(IXI_id)))
        t2_exists = len(glob.glob('./orig/t2/{}*.nii.gz'.format(IXI_id)))
        pd_exists = len(glob.glob('./orig/pd/{}*.nii.gz'.format(IXI_id)))
        mra_exists = len(glob.glob('./orig/mra/{}*.nii.gz'.format(IXI_id)))

        # Check if each entry is complete and drop if not
        #if not t1_exists and not t2_exists and not pd_exists and not mra_exists:
        if not (t1_exists and t2_exists and pd_exists and mra_exists):
            df.drop(index, inplace=True)

    # Write to csv file
    df.to_csv('demographic_HH.csv', index=False)

if RESAMPLE_IMAGES:
    # Resample the IXI HH T2 images to 1mm isotropic and reslice all others to it
    df = pd.read_csv('demographic_HH.csv', dtype=object, keep_default_na=False, na_values=[]).as_matrix()

    for i in df:
        IXI_id = i[0]
        print ('Resampling {}'.format(IXI_id))

        t1_fn = glob.glob('./orig/t1/{}*.nii.gz'.format(IXI_id))[0]
        t2_fn = glob.glob('./orig/t2/{}*.nii.gz'.format(IXI_id))[0]
        pd_fn = glob.glob('./orig/pd/{}*.nii.gz'.format(IXI_id))[0]
        mra_fn = glob.glob('./orig/mra/{}*.nii.gz'.format(IXI_id))[0]

        t1 = sitk.ReadImage(t1_fn)
        t2 = sitk.ReadImage(t2_fn)
        pd = sitk.ReadImage(pd_fn)
        mra = sitk.ReadImage(mra_fn)

        ## Resample to 1mm isotropic resolution
        t2_1mm = resample_image(t2)
        t1_1mm = reslice_image(t1, t2_1mm)
        pd_1mm = reslice_image(pd, t2_1mm)
        mra_1mm = reslice_image(mra, t2_1mm)

        output_dir = os.path.join('./1mm/', IXI_id)
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        print ('T1: {} {}'.format(t1_1mm.GetSize(), t1_1mm.GetSpacing()))
        print ('T2: {} {}'.format(t2_1mm.GetSize(), t2_1mm.GetSpacing()))
        print ('PD: {} {}'.format(pd_1mm.GetSize(), pd_1mm.GetSpacing()))
        print ('MRA: {} {}'.format(mra_1mm.GetSize(), mra_1mm.GetSpacing()))

        sitk.WriteImage(t1_1mm, os.path.join(output_dir, 'T1_1mm.nii.gz'))
        sitk.WriteImage(t2_1mm, os.path.join(output_dir, 'T2_1mm.nii.gz'))
        sitk.WriteImage(pd_1mm, os.path.join(output_dir, 'PD_1mm.nii.gz'))
        sitk.WriteImage(mra_1mm, os.path.join(output_dir, 'MRA_1mm.nii.gz'))

        # Resample to 2mm isotropic resolution
        t2_2mm = resample_image(t2, out_spacing=[2.0, 2.0, 2.0])
        t1_2mm = reslice_image(t1, t2_2mm)
        pd_2mm = reslice_image(pd, t2_2mm)
        mra_2mm = reslice_image(mra, t2_2mm)

        output_dir = os.path.join('./2mm/', IXI_id)
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        print ('T1: {} {}'.format(t2_2mm.GetSize(), t1_2mm.GetSpacing()))
        print ('T2: {} {}'.format(t2_2mm.GetSize(), t2_2mm.GetSpacing()))
        print ('PD: {} {}'.format(pd_2mm.GetSize(), pd_2mm.GetSpacing()))
        print ('MRA: {} {}'.format(mra_2mm.GetSize(), mra_2mm.GetSpacing()))

        sitk.WriteImage(t1_2mm, os.path.join(output_dir, 'T1_2mm.nii.gz'))
        sitk.WriteImage(t2_2mm, os.path.join(output_dir, 'T2_2mm.nii.gz'))
        sitk.WriteImage(pd_2mm, os.path.join(output_dir, 'PD_2mm.nii.gz'))
        sitk.WriteImage(mra_2mm, os.path.join(output_dir, 'MRA_2mm.nii.gz'))

if CLEAN_UP:
    # Remove the .tar files
    for key, fname in fnames.items():
        if (fname.endswith('.tar')):
            os.remove(fname)

    # Remove all data in original resolution
    os.system('rm -rf orig')

        <commit_msg>py2/3 compatability<commit_after># -*- coding: utf-8 -*-
"""Download and extract the IXI Hammersmith Hospital 3T dataset

url: http://brain-development.org/ixi-dataset/
ref: IXI – Information eXtraction from Images (EPSRC GR/S21533/02)

"""
from future.standard_library import install_aliases  # py 2/3 compatability
install_aliases()

from urllib.request import FancyURLopener

import os.path
import tarfile
import pandas as pd
import glob
import SimpleITK as sitk
import numpy as np

DOWNLOAD_IMAGES = True
EXTRACT_IMAGES = True
PROCESS_OTHER = True
RESAMPLE_IMAGES = True
CLEAN_UP = True

def resample_image(itk_image, out_spacing=[1.0, 1.0, 1.0], is_label=False):

    original_spacing = itk_image.GetSpacing()
    original_size = itk_image.GetSize()

    out_size = [int(np.round(original_size[0]*(original_spacing[0]/out_spacing[0]))),
                int(np.round(original_size[1]*(original_spacing[1]/out_spacing[1]))),
                int(np.round(original_size[2]*(original_spacing[2]/out_spacing[2])))]

    resample = sitk.ResampleImageFilter()
    resample.SetOutputSpacing(out_spacing)
    resample.SetSize(out_size)
    resample.SetOutputDirection(itk_image.GetDirection())
    resample.SetOutputOrigin(itk_image.GetOrigin())
    resample.SetTransform(sitk.Transform())
    resample.SetDefaultPixelValue(itk_image.GetPixelIDValue())

    if is_label:
        resample.SetInterpolator(sitk.sitkNearestNeighbor)
    else:
        resample.SetInterpolator(sitk.sitkBSpline)

    return resample.Execute(itk_image)

def reslice_image(itk_image, itk_ref, is_label=False):

    resample = sitk.ResampleImageFilter()
    resample.SetReferenceImage(itk_ref)

    if is_label:
        resample.SetInterpolator(sitk.sitkNearestNeighbor)
    else:
        resample.SetInterpolator(sitk.sitkBSpline)

    return resample.Execute(itk_image)

urls = {}
urls['t1'] = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-T1.tar'
urls['t2'] = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-T2.tar'
urls['pd'] = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-PD.tar'
urls['mra'] = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-MRA.tar'
urls['demographic'] = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI.xls'

fnames = {}
fnames['t1'] = 't1.tar'
fnames['t2'] = 't2.tar'
fnames['pd'] = 'pd.tar'
fnames['mra'] = 'mra.tar'
fnames['demographic'] = 'demographic.xls'

if DOWNLOAD_IMAGES:
    # Download all IXI data
    for key, url in urls.items():

        if not os.path.isfile(fnames[key]):
            print('Downloading {} from {}'.format(fnames[key], url))
            curr_file = FancyURLopener()
            curr_file.retrieve(url, fnames[key])
        else:
            print('File {} already exists. Skipping download.'.format(fnames[key]))

if EXTRACT_IMAGES:
    # Extract the HH subset of IXI
    for key, fname in fnames.items():

        if (fname.endswith('.tar')):
            print('Extracting IXI HH data from {}.'.format(fnames[key]))
            output_dir = os.path.join('./orig/', key)

            if not os.path.exists(output_dir):
                os.makedirs(output_dir)

            t = tarfile.open(fname, 'r')
            for member in t.getmembers():
                if '-HH-' in member.name:
                    t.extract(member, output_dir)

if PROCESS_OTHER:
    # Process the demographic xls data and save to csv
    xls = pd.ExcelFile('demographic.xls')
    print(xls.sheet_names)

    df = xls.parse('Table')
    for index, row in df.iterrows():
        IXI_id = 'IXI{:03d}'.format(row['IXI_ID'])
        df.loc[index, 'IXI_ID'] = IXI_id

        t1_exists = len(glob.glob('./orig/t1/{}*.nii.gz'.format(IXI_id)))
        t2_exists = len(glob.glob('./orig/t2/{}*.nii.gz'.format(IXI_id)))
        pd_exists = len(glob.glob('./orig/pd/{}*.nii.gz'.format(IXI_id)))
        mra_exists = len(glob.glob('./orig/mra/{}*.nii.gz'.format(IXI_id)))

        # Check if each entry is complete and drop if not
        #if not t1_exists and not t2_exists and not pd_exists and not mra_exists:
        if not (t1_exists and t2_exists and pd_exists and mra_exists):
            df.drop(index, inplace=True)

    # Write to csv file
    df.to_csv('demographic_HH.csv', index=False)

if RESAMPLE_IMAGES:
    # Resample the IXI HH T2 images to 1mm isotropic and reslice all others to it
    df = pd.read_csv('demographic_HH.csv', dtype=object, keep_default_na=False, na_values=[]).as_matrix()

    for i in df:
        IXI_id = i[0]
        print ('Resampling {}'.format(IXI_id))

        t1_fn = glob.glob('./orig/t1/{}*.nii.gz'.format(IXI_id))[0]
        t2_fn = glob.glob('./orig/t2/{}*.nii.gz'.format(IXI_id))[0]
        pd_fn = glob.glob('./orig/pd/{}*.nii.gz'.format(IXI_id))[0]
        mra_fn = glob.glob('./orig/mra/{}*.nii.gz'.format(IXI_id))[0]

        t1 = sitk.ReadImage(t1_fn)
        t2 = sitk.ReadImage(t2_fn)
        pd = sitk.ReadImage(pd_fn)
        mra = sitk.ReadImage(mra_fn)

        ## Resample to 1mm isotropic resolution
        t2_1mm = resample_image(t2)
        t1_1mm = reslice_image(t1, t2_1mm)
        pd_1mm = reslice_image(pd, t2_1mm)
        mra_1mm = reslice_image(mra, t2_1mm)

        output_dir = os.path.join('./1mm/', IXI_id)
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        print ('T1: {} {}'.format(t1_1mm.GetSize(), t1_1mm.GetSpacing()))
        print ('T2: {} {}'.format(t2_1mm.GetSize(), t2_1mm.GetSpacing()))
        print ('PD: {} {}'.format(pd_1mm.GetSize(), pd_1mm.GetSpacing()))
        print ('MRA: {} {}'.format(mra_1mm.GetSize(), mra_1mm.GetSpacing()))

        sitk.WriteImage(t1_1mm, os.path.join(output_dir, 'T1_1mm.nii.gz'))
        sitk.WriteImage(t2_1mm, os.path.join(output_dir, 'T2_1mm.nii.gz'))
        sitk.WriteImage(pd_1mm, os.path.join(output_dir, 'PD_1mm.nii.gz'))
        sitk.WriteImage(mra_1mm, os.path.join(output_dir, 'MRA_1mm.nii.gz'))

        # Resample to 2mm isotropic resolution
        t2_2mm = resample_image(t2, out_spacing=[2.0, 2.0, 2.0])
        t1_2mm = reslice_image(t1, t2_2mm)
        pd_2mm = reslice_image(pd, t2_2mm)
        mra_2mm = reslice_image(mra, t2_2mm)

        output_dir = os.path.join('./2mm/', IXI_id)
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        print ('T1: {} {}'.format(t2_2mm.GetSize(), t1_2mm.GetSpacing()))
        print ('T2: {} {}'.format(t2_2mm.GetSize(), t2_2mm.GetSpacing()))
        print ('PD: {} {}'.format(pd_2mm.GetSize(), pd_2mm.GetSpacing()))
        print ('MRA: {} {}'.format(mra_2mm.GetSize(), mra_2mm.GetSpacing()))

        sitk.WriteImage(t1_2mm, os.path.join(output_dir, 'T1_2mm.nii.gz'))
        sitk.WriteImage(t2_2mm, os.path.join(output_dir, 'T2_2mm.nii.gz'))
        sitk.WriteImage(pd_2mm, os.path.join(output_dir, 'PD_2mm.nii.gz'))
        sitk.WriteImage(mra_2mm, os.path.join(output_dir, 'MRA_2mm.nii.gz'))

if CLEAN_UP:
    # Remove the .tar files
    for key, fname in fnames.items():
        if (fname.endswith('.tar')):
            os.remove(fname)

    # Remove all data in original resolution
    os.system('rm -rf orig')

        <|endoftext|>


<commit_before>import re
from parsimonious.grammar import Grammar
import parsimonious.exceptions

class ParseError(Exception):
    pass

def is_bdl_scenario(contents):
    return contents.startswith("#!benchDL\n")

def convert(contents, env):
    ast = transform(lex(contents))
    return substitute(ast, env)

def get_includes(ir):
    if isinstance(ir, list):
        return reduce(lambda a, x: a + get_includes(x), ir, [])
    if isinstance(ir, dict):
        if ir["function"] == "include_resource" and isinstance(ir["args"][1], str):
            return [ir["args"][0:2]]
    return []

def get_num_of_workers(ir):
    if isinstance(ir, list):
        return reduce(lambda a, x: a + get_num_of_workers(x), ir, 0)
    if isinstance(ir, dict):
        if ir["function"] == "pool" and isinstance(ir["args"]["size"], int):
            return ir["args"]["size"]
    return 0

def add_indents(text):
    brackets = []
    indents = [0]
    result = []
    comma = None
    for line in text.splitlines():
        old_len = len(brackets)

        tline = line
        while len(tline) > 0:
            if tline.startswith("\\\\"):
                tline = tline[2:]
                continue
            if tline.startswith("\\\""):
                tline = tline[2:]
                continue

            char = tline[0]
            tline = tline[1:]

            if comma != None and char == comma:
                comma = None
            elif char in ['"', "'"]:
                comma = char

            if comma == None:
                if (char == "]" and brackets[-1] == "[") or (char == ")" and brackets[-1] == "("):
                    brackets.pop()
                if char == "(" or char == "[":
                    brackets.append(char)
                if char == "#":
                    break

        if old_len == 0:
            i = get_indent(line)
            if (i == -1) or (i == indents[-1]):
                result.append(line)
                continue
            if i > indents[-1]:
                result.append("_INDENT_" + line)
                indents.append(i)
                continue
            while i < indents[-1]:
                indents.pop()
                line = "_DEDENT_ " + line
            if i != indents[-1]:
                raise ParseError("Incorrect indentation")

        result.append(line)

    line = ""
    while 0 < indents[-1]:
        indents.pop()
        line = "_DEDENT_ " + line

    if line != "":
        result.append(line)

    return "\n".join(result)

def get_indent(line):
    spaces = len(re.match(r"\s*", line).group())
    return -1 if (len(line) <= spaces) or (line[spaces] == "#") else spaces

def substitute(ir, env):
    if isinstance(ir, list):
        return map(lambda a: substitute(a, env), ir)
    if isinstance(ir, dict):
        ir = {k: substitute(v, env) for k, v in ir.items()}

        if "function" in ir:
            if ir["function"] in ["var", "numvar"] and ir["args"][0] in env:
                return env[ir["args"][0]]
            elif ir["function"] in ["var", "numvar"] and len(ir["args"]) > 1:
                return ir["args"][1]

    return ir

def transform(ast):
    if ast.expr_name == "number":
        if ast.text[-1] in "KMG":
            num = ast.text[0:-1]
            if ast.text[-1] == "K":
                mult = 1000
            elif ast.text[-1] == "M":
                mult = 1000000
            else:
                mult = 1000000000
        else:
            num = ast.text
            mult = 1
        return [float(num)*mult] if "." in num else [int(num)*mult]
    elif ast.expr_name == "boolean":
        return [ast.text=="true"]
    elif ast.expr_name == "string":
        return [ast.text[1:-1]] # "something"
    elif ast.expr_name == "atom":
        return [ast.text]
    else:
        lis = reduce(lambda a, x: a + transform(x), ast.children, [])
        if ast.expr_name == "list":
            return [lis]
        if ast.expr_name == "multiline":
            return [{"function": lis[0], "args":lis[1], "children": lis[2:]}]
        if ast.expr_name == "single":
            return [{"function": lis[0], "args":lis[1:]}]
        elif ast.expr_name == "map":
            return [dict(lis)]
        elif ast.expr_name in ["tuple", "kv"]:
            if len(lis) == 2:
                return [tuple(lis)]
            else:
                return [(lis[0], lis[1:])]
        return lis

def lex(text):
    grammar = Grammar("""\
    entry = _ (statement _)* _
    statement = multiline / single
    multiline = atom _ args _ ":" _ "_INDENT_" _ (statement _)+ "_DEDENT_"
    single = atom _ args
    atom = ~"[a-z][0-9a-zA-Z_]*" / ("'" ~"[^']*" "'")
    _ = ~"\s*" (~"#[^\\r\\n]*\s*")*
    args = ( _ map ) / ( _ "(" _ term (_ "," _ term)* _ ")" ) / (_ "(" _ ")")
    map = "(" _ kv (_ "," _ kv)* _ ")"
    list = ( _ "[" _ term (_ "," _ term)* _ "]" ) / ( _ "[" _ "]")
    kv = term _ "=" _ term _
    term = unumber / logic_op / single / list / string / atom / number
    logic_op = (string / number) _ ("<=" / ">=" / "<" / ">" / "==") _ (string / number)
    string = '"' ~'[^"]*' '"'
    number = ~"[0-9]+(\.[0-9]+)?(e\-?[0-9]+)?[GKM]?"
    unumber = (number / single) _ atom
    """)
    try:
        return grammar.parse(add_indents(text))
    except parsimonious.exceptions.ParseError as e:
        raise ParseError(e)
<commit_msg>Fixing quotes<commit_after>import re
from parsimonious.grammar import Grammar
import parsimonious.exceptions

class ParseError(Exception):
    pass

def is_bdl_scenario(contents):
    return contents.startswith("#!benchDL\n")

def convert(contents, env):
    ast = transform(lex(contents))
    return substitute(ast, env)

def get_includes(ir):
    if isinstance(ir, list):
        return reduce(lambda a, x: a + get_includes(x), ir, [])
    if isinstance(ir, dict):
        if ir["function"] == "include_resource" and isinstance(ir["args"][1], str):
            return [ir["args"][0:2]]
    return []

def get_num_of_workers(ir):
    if isinstance(ir, list):
        return reduce(lambda a, x: a + get_num_of_workers(x), ir, 0)
    if isinstance(ir, dict):
        if ir["function"] == "pool" and isinstance(ir["args"]["size"], int):
            return ir["args"]["size"]
    return 0

def add_indents(text):
    brackets = []
    indents = [0]
    result = []
    comma = None
    for line in text.splitlines():
        old_len = len(brackets)

        tline = line
        while len(tline) > 0:
            if tline.startswith("\\\\"):
                tline = tline[2:]
                continue
            if tline.startswith("\\\""):
                tline = tline[2:]
                continue

            char = tline[0]
            tline = tline[1:]

            if comma != None and char == comma:
                comma = None
            elif char in ['"', "'"]:
                comma = char

            if comma == None:
                if (char == "]" and brackets[-1] == "[") or (char == ")" and brackets[-1] == "("):
                    brackets.pop()
                if char == "(" or char == "[":
                    brackets.append(char)
                if char == "#":
                    break

        if old_len == 0:
            i = get_indent(line)
            if (i == -1) or (i == indents[-1]):
                result.append(line)
                continue
            if i > indents[-1]:
                result.append("_INDENT_" + line)
                indents.append(i)
                continue
            while i < indents[-1]:
                indents.pop()
                line = "_DEDENT_ " + line
            if i != indents[-1]:
                raise ParseError("Incorrect indentation")

        result.append(line)

    line = ""
    while 0 < indents[-1]:
        indents.pop()
        line = "_DEDENT_ " + line

    if line != "":
        result.append(line)

    return "\n".join(result)

def get_indent(line):
    spaces = len(re.match(r"\s*", line).group())
    return -1 if (len(line) <= spaces) or (line[spaces] == "#") else spaces

def substitute(ir, env):
    if isinstance(ir, list):
        return map(lambda a: substitute(a, env), ir)
    if isinstance(ir, dict):
        ir = {k: substitute(v, env) for k, v in ir.items()}

        if "function" in ir:
            if ir["function"] in ["var", "numvar"] and ir["args"][0] in env:
                return env[ir["args"][0]]
            elif ir["function"] in ["var", "numvar"] and len(ir["args"]) > 1:
                return ir["args"][1]

    return ir

def transform(ast):
    if ast.expr_name == "number":
        if ast.text[-1] in "KMG":
            num = ast.text[0:-1]
            if ast.text[-1] == "K":
                mult = 1000
            elif ast.text[-1] == "M":
                mult = 1000000
            else:
                mult = 1000000000
        else:
            num = ast.text
            mult = 1
        return [float(num)*mult] if "." in num else [int(num)*mult]
    elif ast.expr_name == "boolean":
        return [ast.text=="true"]
    elif ast.expr_name == "string":
        return [ast.text[1:-1]] # "something"
    elif ast.expr_name == "atom":
        return [ast.text]
    else:
        lis = reduce(lambda a, x: a + transform(x), ast.children, [])
        if ast.expr_name == "list":
            return [lis]
        if ast.expr_name == "multiline":
            return [{"function": lis[0], "args":lis[1], "children": lis[2:]}]
        if ast.expr_name == "single":
            return [{"function": lis[0], "args":lis[1:]}]
        elif ast.expr_name == "map":
            return [dict(lis)]
        elif ast.expr_name in ["tuple", "kv"]:
            if len(lis) == 2:
                return [tuple(lis)]
            else:
                return [(lis[0], lis[1:])]
        return lis

def lex(text):
    grammar = Grammar("""\
    entry = _ (statement _)* _
    statement = multiline / single
    multiline = atom _ args _ ":" _ "_INDENT_" _ (statement _)+ "_DEDENT_"
    single = atom _ args
    atom = ~"[a-z][0-9a-zA-Z_]*" / ("'" ~"[^']*" "'")
    _ = ~"\s*" (~"#[^\\r\\n]*\s*")*
    args = ( _ map ) / ( _ "(" _ term (_ "," _ term)* _ ")" ) / (_ "(" _ ")")
    map = "(" _ kv (_ "," _ kv)* _ ")"
    list = ( _ "[" _ term (_ "," _ term)* _ "]" ) / ( _ "[" _ "]")
    kv = term _ "=" _ term _
    term = unumber / logic_op / single / list / string / atom / number
    logic_op = (string / number) _ ("<=" / ">=" / "<" / ">" / "==") _ (string / number)
    string = '"' ~r'(\\\\"|[^\\\\"])*' '"'
    number = ~"[0-9]+(\.[0-9]+)?(e\-?[0-9]+)?[GKM]?"
    unumber = (number / single) _ atom
    """)
    try:
        return grammar.parse(add_indents(text))
    except parsimonious.exceptions.ParseError as e:
        raise ParseError(e)
<|endoftext|>


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

"""Simple DNS server comparison benchmarking tool.

Designed to assist system administrators in selection and prioritization.
"""

__author__ = 'tstromberg@google.com (Thomas Stromberg)'

import csv
import math
import operator
import random
import sys
import charts
import dns.rcode

def CalculateListAverage(values):
  """Computes the arithmetic mean of a list of numbers."""
  return sum(values) / float(len(values))

def DrawTextBar(value, max_value, max_width=53):
  """Return a simple ASCII bar graph, making sure it fits within max_width.

  Args:
    value: integer or float representing the value of this bar.
    max_value: integer or float representing the largest bar.
    max_width: How many characters this graph can use (int)

  Returns:
    string
  """

  hash_width = max_value / max_width
  return int(math.ceil(value/hash_width)) * '#'

def WeightedDistribution(elements):
  """Simple exponential distribution with a mean of 10%."""
  maximum = len(elements) - 1
  desired_mean = maximum * 0.10
  lambd = 1.0 / desired_mean
  picks = []

  while len(picks) < maximum:
    index = int(random.expovariate(lambd))
    if index < maximum:
      picks.append(elements[index])
  return picks

def ChunkSelect(elements, count):
  start = random.randint(0, len(elements) - count)
  return elements[start:start + count]

class NameBench(object):
  """The main benchmarking class."""

  def __init__(self, nameservers, run_count=2, test_count=30):
    """Constructor.

    Args:
      nameservers: a list of NameServerData objects
      test_domains_path: Path to the list of domains to use (str)
      run_count: How many test-runs to perform on each nameserver (int)
      test_count: How many DNS lookups to test in each test-run (int)
    """
    self.test_count = test_count
    self.run_count = run_count
    self.nameservers = nameservers
    self.results = {}

  def LoadTestData(self, filename, select_mode='weighted'):
    input_data = open(filename).readlines()

    if select_mode == 'weighted':
      selected = WeightedDistribution(input_data)[0:self.test_count]
    elif select_mode == 'chunk':
      selected = ChunkSelect(input_data, self.test_count)
    elif select_mode == 'random':
      selected = random.sample(input_data)[0:self.test_count]

    self.test_data = []
    for line in selected:
      selection = line.rstrip()
      if ' ' in selection:
        self.test_data.append(selection.split(' ')[0:2])
      else:
        self.test_data.append(('A', self.GenerateFqdn(selection)))

    print '- Generated %s tests from %s using %s mode' % (len(self.test_data), filename, select_mode)
    return self.test_data

  def GenerateFqdn(self, domain):
    oracle = random.randint(0, 100)
    if oracle < 70:
      return 'www.%s.' % domain
    elif oracle < 97:
      return '%s.' % domain
    else:
      return 'cache-%s.%s.' % (random.randint(0,10000), domain)

  def BenchmarkNameServer(self, nameserver, tests):
    """Record results for a single run on a nameserver.

    Args:
      nameserver: NameServerData object
      tests: a list of tuples in the form of [(record_type, record_name)]

    Returns:
      results: list of tuples, including data for each request made.
    """
    results = []
    for (req_type, record) in tests:
      # test records can include a (RANDOM) in the string for cache busting.
      if '(RANDOM)' in record:
        record = record.replace('(RANDOM)', str(random.random() * 10))
      (response, duration) = nameserver.TimedRequest(req_type, record)[0:2]
      results.append((record, req_type, duration, response))
    return results

  def Run(self):
    """Manage all attempts."""
    for attempt in range(self.run_count):
      sys.stdout.write('* Benchmarking %s nameservers with %s records each (%s of %s).' %
                       (len(self.nameservers), self.test_count, attempt+1, self.run_count))
      for ns in self.nameservers:
        if ns not in self.results:
          self.results[ns] = []
        sys.stdout.write('.')
        sys.stdout.flush()
        self.results[ns].append(self.BenchmarkNameServer(ns, self.test_data))
      sys.stdout.write('\n')
    sys.stdout.flush()

  def ComputeAverages(self):
    """Process all runs for all hosts, yielding an average for each host."""
    for ns in self.results:
      record_count = 0
      failure_count = 0
      run_averages = []

      for test_run in self.results[ns]:
        record_count += len(test_run)
        failure_count += len([x[3] for x in test_run if x[3] == -1])
        duration = sum([x[2] for x in test_run])
        run_averages.append(duration / len(test_run))

      # This appears to be a safe use of averaging averages
      overall_average = CalculateListAverage(run_averages)
      yield (ns, overall_average, run_averages, failure_count)

  def FastestNameServerResult(self):
    """Process all runs for all hosts, yielding an average for each host."""
    return [(x[0], min(x[1])) for x in self.DigestedResults()]

  def BestOverallNameServer(self):
    sorted_averages = sorted(self.ComputeAverages(), key=operator.itemgetter(1))
    return sorted_averages[0][0]

  def NearestNameServers(self, count=2):
    min_responses = sorted(self.FastestNameServerResult(),
                           key=operator.itemgetter(1))
    return [x[0] for x in min_responses][0:count]

  def DisplayResults(self):
    """Display all of the results in an ASCII-graph format."""

    print ''
    print 'Lowest latency for an individual query (in milliseconds):'
    print '-'* 78
    min_responses = sorted(self.FastestNameServerResult(),
                           key=operator.itemgetter(1))
    slowest_result = min_responses[-1][1]
    for result in min_responses:
      (ns, duration) = result
      textbar = DrawTextBar(duration, slowest_result)
      print '%-16.16s %s %2.2f' % (ns.name, textbar, duration)

    print ''
    print 'Overall Mean Request Duration (in milliseconds):'
    print '-'* 78
    sorted_averages = sorted(self.ComputeAverages(), key=operator.itemgetter(1))
    max_result = sorted_averages[-1][1]
    timeout_seen = False
    for result in sorted_averages:
      (ns, overall_mean, unused_run_means, failure_count) = result
      if failure_count:
        timeout_seen = True
        note = ' (%sT)' % failure_count
      else:
        note = ''
      textbar = DrawTextBar(overall_mean, max_result)
      print '%-16.16s %s %2.0f%s' % (ns.name, textbar, overall_mean, note)
    if timeout_seen:
      print '* (#T) represents the number of timeouts experienced during testing.'
    print ''

    print 'Detailed Mean Request Duration Chart URL'
    print '-' * 78
    runs_data = [(x[0].name, x[2]) for x in sorted_averages]
    print charts.PerRunDurationBarGraph(runs_data)
    print ''
    print 'Detailed Request Duration Distribution Chart URL'
    print '-' * 78
    print charts.DistributionLineGraph(self.DigestedResults())

  def DigestedResults(self):
    """Return a tuple of nameserver and all associated durations."""
    duration_data = []
    for ns in self.results:
      durations = []
      for test_run_results in self.results[ns]:
        durations += [x[2] for x in test_run_results]
      duration_data.append((ns, durations))
    return duration_data

  def SaveResultsToCsv(self, filename):
    """Write out a CSV file with detailed results on each request.

    Args:
      filename: full path on where to save results (string)

    Sample output:
    nameserver, test_number, test, type, duration, answer_count, ttl
    """
    csv_file = open(filename, 'w')
    output = csv.writer(csv_file)
    output.writerow(['IP', 'Name', 'Check Duration', 'Test #', 'Record', 'Record Type', 'Duration', 'TTL', 'Answer Count', 'Response'])
    for ns in self.results:
      for (test_run, test_results) in enumerate(self.results[ns]):
        for (record, req_type, duration, response) in test_results:
          answer_text = ''
          answer_count = -1
          ttl = -1
          if response:
            if response.answer:
              answer_count = len(response.answer)
              answers = [' + '.join(map(str, x.items)) for x in response.answer]
              answer_text = ' -> '.join(answers)
              ttl = response.answer[0].ttl
            else:
              answer_text = dns.rcode.to_text(response.rcode())
          output.writerow([ns.ip, ns.name, ns.check_duration, test_run, record, req_type, duration, ttl, answer_count, answer_text])

    csv_file.close()

<commit_msg>raise an error if a bad select_mode is given<commit_after>#!/usr/bin/env python
# Copyright 2009 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Simple DNS server comparison benchmarking tool.

Designed to assist system administrators in selection and prioritization.
"""

__author__ = 'tstromberg@google.com (Thomas Stromberg)'

import csv
import math
import operator
import random
import sys
import charts
import dns.rcode

def CalculateListAverage(values):
  """Computes the arithmetic mean of a list of numbers."""
  return sum(values) / float(len(values))

def DrawTextBar(value, max_value, max_width=53):
  """Return a simple ASCII bar graph, making sure it fits within max_width.

  Args:
    value: integer or float representing the value of this bar.
    max_value: integer or float representing the largest bar.
    max_width: How many characters this graph can use (int)

  Returns:
    string
  """

  hash_width = max_value / max_width
  return int(math.ceil(value/hash_width)) * '#'

def WeightedDistribution(elements):
  """Simple exponential distribution with a mean of 10%."""
  maximum = len(elements) - 1
  desired_mean = maximum * 0.10
  lambd = 1.0 / desired_mean
  picks = []

  while len(picks) < maximum:
    index = int(random.expovariate(lambd))
    if index < maximum:
      picks.append(elements[index])
  return picks

def ChunkSelect(elements, count):
  start = random.randint(0, len(elements) - count)
  return elements[start:start + count]

class NameBench(object):
  """The main benchmarking class."""

  def __init__(self, nameservers, run_count=2, test_count=30):
    """Constructor.

    Args:
      nameservers: a list of NameServerData objects
      test_domains_path: Path to the list of domains to use (str)
      run_count: How many test-runs to perform on each nameserver (int)
      test_count: How many DNS lookups to test in each test-run (int)
    """
    self.test_count = test_count
    self.run_count = run_count
    self.nameservers = nameservers
    self.results = {}

  def LoadTestData(self, filename, select_mode='weighted'):
    input_data = open(filename).readlines()

    if select_mode == 'weighted':
      selected = WeightedDistribution(input_data)[0:self.test_count]
    elif select_mode == 'chunk':
      selected = ChunkSelect(input_data, self.test_count)
    elif select_mode == 'random':
      selected = random.sample(input_data)[0:self.test_count]
    else:
      raise ValueError('Invalid select_mode: %s' % select_mode)

    self.test_data = []
    for line in selected:
      selection = line.rstrip()
      if ' ' in selection:
        self.test_data.append(selection.split(' ')[0:2])
      else:
        self.test_data.append(('A', self.GenerateFqdn(selection)))

    print '- Generated %s tests from %s using %s mode' % (len(self.test_data), filename, select_mode)
    return self.test_data

  def GenerateFqdn(self, domain):
    oracle = random.randint(0, 100)
    if oracle < 70:
      return 'www.%s.' % domain
    elif oracle < 97:
      return '%s.' % domain
    else:
      return 'cache-%s.%s.' % (random.randint(0,10000), domain)

  def BenchmarkNameServer(self, nameserver, tests):
    """Record results for a single run on a nameserver.

    Args:
      nameserver: NameServerData object
      tests: a list of tuples in the form of [(record_type, record_name)]

    Returns:
      results: list of tuples, including data for each request made.
    """
    results = []
    for (req_type, record) in tests:
      # test records can include a (RANDOM) in the string for cache busting.
      if '(RANDOM)' in record:
        record = record.replace('(RANDOM)', str(random.random() * 10))
      (response, duration) = nameserver.TimedRequest(req_type, record)[0:2]
      results.append((record, req_type, duration, response))
    return results

  def Run(self):
    """Manage all attempts."""
    for attempt in range(self.run_count):
      sys.stdout.write('* Benchmarking %s nameservers with %s records each (%s of %s).' %
                       (len(self.nameservers), self.test_count, attempt+1, self.run_count))
      for ns in self.nameservers:
        if ns not in self.results:
          self.results[ns] = []
        sys.stdout.write('.')
        sys.stdout.flush()
        self.results[ns].append(self.BenchmarkNameServer(ns, self.test_data))
      sys.stdout.write('\n')
    sys.stdout.flush()

  def ComputeAverages(self):
    """Process all runs for all hosts, yielding an average for each host."""
    for ns in self.results:
      record_count = 0
      failure_count = 0
      run_averages = []

      for test_run in self.results[ns]:
        record_count += len(test_run)
        failure_count += len([x[3] for x in test_run if x[3] == -1])
        duration = sum([x[2] for x in test_run])
        run_averages.append(duration / len(test_run))

      # This appears to be a safe use of averaging averages
      overall_average = CalculateListAverage(run_averages)
      yield (ns, overall_average, run_averages, failure_count)

  def FastestNameServerResult(self):
    """Process all runs for all hosts, yielding an average for each host."""
    return [(x[0], min(x[1])) for x in self.DigestedResults()]

  def BestOverallNameServer(self):
    sorted_averages = sorted(self.ComputeAverages(), key=operator.itemgetter(1))
    return sorted_averages[0][0]

  def NearestNameServers(self, count=2):
    min_responses = sorted(self.FastestNameServerResult(),
                           key=operator.itemgetter(1))
    return [x[0] for x in min_responses][0:count]

  def DisplayResults(self):
    """Display all of the results in an ASCII-graph format."""

    print ''
    print 'Lowest latency for an individual query (in milliseconds):'
    print '-'* 78
    min_responses = sorted(self.FastestNameServerResult(),
                           key=operator.itemgetter(1))
    slowest_result = min_responses[-1][1]
    for result in min_responses:
      (ns, duration) = result
      textbar = DrawTextBar(duration, slowest_result)
      print '%-16.16s %s %2.2f' % (ns.name, textbar, duration)

    print ''
    print 'Overall Mean Request Duration (in milliseconds):'
    print '-'* 78
    sorted_averages = sorted(self.ComputeAverages(), key=operator.itemgetter(1))
    max_result = sorted_averages[-1][1]
    timeout_seen = False
    for result in sorted_averages:
      (ns, overall_mean, unused_run_means, failure_count) = result
      if failure_count:
        timeout_seen = True
        note = ' (%sT)' % failure_count
      else:
        note = ''
      textbar = DrawTextBar(overall_mean, max_result)
      print '%-16.16s %s %2.0f%s' % (ns.name, textbar, overall_mean, note)
    if timeout_seen:
      print '* (#T) represents the number of timeouts experienced during testing.'
    print ''

    print 'Detailed Mean Request Duration Chart URL'
    print '-' * 78
    runs_data = [(x[0].name, x[2]) for x in sorted_averages]
    print charts.PerRunDurationBarGraph(runs_data)
    print ''
    print 'Detailed Request Duration Distribution Chart URL'
    print '-' * 78
    print charts.DistributionLineGraph(self.DigestedResults())

  def DigestedResults(self):
    """Return a tuple of nameserver and all associated durations."""
    duration_data = []
    for ns in self.results:
      durations = []
      for test_run_results in self.results[ns]:
        durations += [x[2] for x in test_run_results]
      duration_data.append((ns, durations))
    return duration_data

  def SaveResultsToCsv(self, filename):
    """Write out a CSV file with detailed results on each request.

    Args:
      filename: full path on where to save results (string)

    Sample output:
    nameserver, test_number, test, type, duration, answer_count, ttl
    """
    csv_file = open(filename, 'w')
    output = csv.writer(csv_file)
    output.writerow(['IP', 'Name', 'Check Duration', 'Test #', 'Record', 'Record Type', 'Duration', 'TTL', 'Answer Count', 'Response'])
    for ns in self.results:
      for (test_run, test_results) in enumerate(self.results[ns]):
        for (record, req_type, duration, response) in test_results:
          answer_text = ''
          answer_count = -1
          ttl = -1
          if response:
            if response.answer:
              answer_count = len(response.answer)
              answers = [' + '.join(map(str, x.items)) for x in response.answer]
              answer_text = ' -> '.join(answers)
              ttl = response.answer[0].ttl
            else:
              answer_text = dns.rcode.to_text(response.rcode())
          output.writerow([ns.ip, ns.name, ns.check_duration, test_run, record, req_type, duration, ttl, answer_count, answer_text])

    csv_file.close()

<|endoftext|>


<commit_before>"""Task and actions classes."""
import subprocess, sys
import StringIO
import traceback

from doit import logger

# Exceptions
class InvalidTask(Exception):
    """Invalid task instance. User error on specifying the task."""
    pass

# TODO rename this?
class TaskFailed(Exception):
    """Task execution was not successful."""
    pass

class TaskError(Exception):
    """Error while trying to execute task."""
    pass

# Actions
class BaseAction(object):
    """Base class for all actions"""

    # must implement:
    # def execute(self, capture_stdout=False, capture_stderr=False)

    pass

class CmdAction(BaseAction):
    """
    Command line action. Spawns a new process.

    @ivar action(str): Command to be passed to the shell subprocess.
         It may contain python mapping strings with the keys: dependencies,
         changed and targets. ie. "zip %(targets)s %(changed)s"
    @ivar task(Task): reference to task that contains this action
    """

    def __init__(self, action):
        assert isinstance(action,str), "CmdAction must be a string."
        self.action = action
        self.task = None

    def execute(self, capture_stdout=False, capture_stderr=False):
        """
        Execute command action

        @param capture_stdout(bool): Capture standard output
        @param capture_err(bool): Capture standard error

        @raise TaskError: If subprocess return code is greater than 125
        @raise TaskFailed: If subprocess return code isn't zero (and
        not greater than 125)
        """
        # set Popen stream parameters
        if not capture_stdout:
            stdout = None
        else:
            stdout = subprocess.PIPE
        if not capture_stderr:
            stderr = None
        else:
            stderr = subprocess.PIPE

        action = self.expand_action()

        # spawn task process
        process = subprocess.Popen(action,stdout=stdout,
                                   stderr=stderr, shell=True)

        # log captured stream
        out,err = process.communicate()
        if out:
            logger.log('stdout',out)
        if err:
            logger.log('stderr',err)

        # task error - based on:
        # http://www.gnu.org/software/bash/manual/bashref.html#Exit-Status
        # it doesnt make so much difference to return as Error or Failed anyway
        if process.returncode > 125:
            raise TaskError("Command error: '%s' returned %s" %
                            (action,process.returncode))

        # task failure
        if process.returncode != 0:
            raise TaskFailed("Command failed: '%s' returned %s" %
                             (action,process.returncode))

    def expand_action(self):
        """expand action string using task meta informations
        @returns (string) - expanded string after substitution
        """
        if not self.task:
            return self.action

        subs_dict = {'targets' : " ".join(self.task.targets),
                     'dependencies': " ".join(self.task.dependencies)}
        # just included changed if it is set
        if self.task.dep_changed is not None:
            subs_dict['changed'] = " ".join(self.task.dep_changed)
        return self.action % subs_dict

    def __str__(self):
        return "Cmd: %s" % self.expand_action()

    def __repr__(self):
        return "<CmdAction: '%s'>"  % self.expand_action()

class PythonAction(BaseAction):
    """Python action. Execute a python callable.

    @ivar py_callable: (callable) Python callable that returns a boolean
    result depending on the success of the action
    @ivar args: (sequence)  Extra arguments to be passed to py_callable
    @ivar kwargs: (dict) Extra keyword arguments to be passed to py_callable
    """
    def __init__(self, py_callable, args=None, kwargs=None):

        self.py_callable = py_callable

        if args is None:
            self.args = []
        else:
            self.args = args

        if kwargs is None:
            self.kwargs = {}
        else:
            self.kwargs = kwargs

        # check valid parameters
        if not callable(self.py_callable):
            raise InvalidTask("PythonAction must be a 'callable'.")
        if type(self.args) is not tuple and type(self.args) is not list:
            msg = "args must be a 'tuple' or a 'list'. got '%s'."
            raise InvalidTask(msg % self.args)
        if type(self.kwargs) is not dict:
            msg = "kwargs must be a 'dict'. got '%s'"
            raise InvalidTask(msg % self.kwargs)

    def execute(self, capture_stdout=False, capture_stderr=False):
        """
        Execute command action

        @param capture_stdout (bool): Capture standard output
        @param capture_err (bool): Capture standard error

        @raise TaskFailed: If py_callable returns False
        """
        # set std stream
        if capture_stdout:
            old_stdout = sys.stdout
            sys.stdout = StringIO.StringIO()

        if capture_stderr:
            old_stderr = sys.stderr
            sys.stderr = StringIO.StringIO()

        # execute action / callable
        try:
            # Python2.4
            try:
                result = self.py_callable(*self.args,**self.kwargs)
            # in python 2.4 SystemExit and KeyboardInterrupt subclass
            # from Exception.
            except (SystemExit, KeyboardInterrupt), exp:
                raise

            except Exception, exception:
                error = TaskError(exception)
                error.originalException = traceback.format_exception(\
                    exception.__class__, exception,sys.exc_info()[2])
                raise error
        finally:
            # restore std streams /log captured streams
            if capture_stdout:
                logger.log('stdout',sys.stdout.getvalue())
                sys.stdout.close()
                sys.stdout = old_stdout

            if capture_stderr:
                logger.log('stderr',sys.stderr.getvalue())
                sys.stderr.close()
                sys.stderr = old_stderr

        # if callable returns false. Task failed
        if not result:
            raise TaskFailed("Python Task failed: '%s' returned %s" %
                             (self.py_callable, result))

    def __str__(self):
        # get object description excluding runtime memory address
        return "Python: %s"% str(self.py_callable)[1:].split(' at ')[0]

    def __repr__(self):
        return "<PythonAction: '%s'>"% (repr(self.py_callable))

def create_action(action):
    """
    Create action using proper constructor based on the parameter type

    @param action: Action to be created
    @type action: L{BaseAction} subclass object, str, tuple or callable
    @raise InvalidTask: If action parameter type isn't valid
    """
    if isinstance(action, BaseAction):
        return action

    if type(action) is str:
        return CmdAction(action)

    if type(action) is tuple:
        return PythonAction(*action)

    if callable(action):
        return PythonAction(action)

    msg = "Invalid task action type. got %s"
    raise InvalidTask(msg % (action.__class__))

class Task(object):
    """Task

    @ivar name string
    @ivar actions: (list) of L{BaseAction}
    @ivar dependencies list of all dependencies
    @ivar targets list of targets
    @ivar folder_dep: (list - string)
    @ivar task_dep: (list - string)
    @ivar file_dep: (list - string)
    @ivar dep_changed (list - string): list of file-dependencies that changed
          (are not up_to_date). this must be set before
    @ivar run_once: (bool) task without dependencies should run
    @ivar is_subtask: (bool) indicate this task is a subtask.
    @ivar doc: (string) task documentation
    """

    def __init__(self,name, actions, dependencies=(),targets=(),setup=None,
                 is_subtask=False, doc=None):
        """sanity checks and initialization"""
        # dependencies parameter must be a list
        if not ((isinstance(dependencies,list)) or
                (isinstance(dependencies,tuple))):
            msg = ("%s. paramater 'dependencies' must be a list or " +
                   "tuple got:'%s'%s")
            raise InvalidTask(msg%(name, str(dependencies),type(dependencies)))

        # targets parameter must be a list
        if not(isinstance(targets,list) or isinstance(targets,tuple)):
            msg = ("%s. paramater 'targets' must be a list or tuple " +
                   "got:'%s'%s")
            raise InvalidTask(msg % (name, str(targets),type(targets)))

        self.name = name
        self.dependencies = dependencies
        self.targets = targets
        self.setup = setup
        self.run_once = False
        self.is_subtask = is_subtask
        self.dep_changed = None

        if actions is None:
            self.actions = []
        elif type(actions) is list:
            self.actions = [create_action(a) for a in actions]
        else:
        ## DEPREACTED ON 0.4 - REMOVE ON 0.5
            #raise Exception("DEPRECATED must be list")
            msg = "DEPRECATION WARNING: task %s actions must be a list."
            print msg % self.name
            self.actions = [create_action(actions)]
        ## DEPRECATION END

        # Store just first non-empty line as documentation string
        if doc is None:
            self.doc = ''
        else:
            for line in doc.splitlines():
                striped = line.strip()
                if striped:
                    self.doc = striped
                    break
            else:
                self.doc = ''

        # set self as task for all actions
        for action in self.actions:
            action.task = self

        # there are 3 kinds of dependencies: file, task, and folder
        self.folder_dep = []
        self.task_dep = []
        self.file_dep = []
        for dep in self.dependencies:
            # True on the list. set run_once
            if isinstance(dep,bool):
                if not dep:
                    msg = ("%s. bool paramater in 'dependencies' "+
                           "must be True got:'%s'")
                    raise InvalidTask(msg%(name, str(dep)))
                self.run_once = True
            # folder dep ends with a '/'
            elif dep.endswith('/'):
                self.folder_dep.append(dep)
            # task dep starts with a ':'
            elif dep.startswith(':'):
                self.task_dep.append(dep[1:])
            # file dep
            elif isinstance(dep,str):
                self.file_dep.append(dep)

        # run_once can't be used together with file dependencies
        if self.run_once and self.file_dep:
            msg = ("%s. task cant have file and dependencies and True " +
                   "at the same time. (just remove True)")
            raise InvalidTask(msg % name)

    def execute(self, capture_stdout=False, capture_stderr=False):
        """Executes the task.

        @raise TaskFailed: If raised when executing an action
        @raise TaskError: If raised when executing an action
        """
        for action in self.actions:
            action.execute(capture_stdout, capture_stderr)

    def title(self):
        """String representation on output.

        @rtype: str
        @return: Task name and actions
        """
        return "%s => %s"% (self.name, str(self))

    def __str__(self):
        if self.actions:
            return "\n\t".join([str(action) for action in self.actions])

        # A task that contains no actions at all
        # is used as group task
        return "Group: %s" % ", ".join(self.dependencies)

    def __repr__(self):
        return "<Task: %s>"% self.name

def dict_to_task(task_dict):
    """Create a task instance from dictionary.

    The dictionary has the same format as returned by task-generators
    from dodo files.

    @param task_dict (dict): task representation as a dict.
    @raise InvalidTask: If unexpected fields were passed in task_dict
    """
    # TASK_ATTRS: sequence of know attributes(keys) of a task dict.
    TASK_ATTRS = ('name','actions','dependencies','targets','setup', 'doc')
    # FIXME check field 'name'

    # === DEPRECATED on 0.4 (to be removed on 0.5): START
    # check required fields
    if 'action' in task_dict and 'actions' not in task_dict:
        #raise Exception("DEPRECATED action")
        task_dict['actions'] = task_dict['action']
        del task_dict['action']
        print ("DEPRECATION WARNING: Task %s contains 'action' key. "
               "This will be deprecated in future versions, "
               "please use 'actions' instead" % task_dict['name'])
    # === DEPRECATION: END

    if 'actions' not in task_dict:
        raise InvalidTask("Task %s must contain 'actions' field. %s" %
                          (task_dict['name'],task_dict))

    # user friendly. dont go ahead with invalid input.
    for key in task_dict.keys():
        if key not in TASK_ATTRS:
            raise InvalidTask("Task %s contain invalid field: %s"%
                              (task_dict['name'],key))

    return Task(**task_dict)
<commit_msg>Python actions now also get extra parameters for targets, dependencies and changed files. This parameters are only passed if the callable expects them (it is defined in the argspec or accepts keyword arguments) and if no values were passed or default were defined explicitly.<commit_after>"""Task and actions classes."""
import subprocess, sys
import StringIO
import traceback
import inspect

from doit import logger

# Exceptions
class InvalidTask(Exception):
    """Invalid task instance. User error on specifying the task."""
    pass

# TODO rename this?
class TaskFailed(Exception):
    """Task execution was not successful."""
    pass

class TaskError(Exception):
    """Error while trying to execute task."""
    pass

# Actions
class BaseAction(object):
    """Base class for all actions"""

    # must implement:
    # def execute(self, capture_stdout=False, capture_stderr=False)

    pass

class CmdAction(BaseAction):
    """
    Command line action. Spawns a new process.

    @ivar action(str): Command to be passed to the shell subprocess.
         It may contain python mapping strings with the keys: dependencies,
         changed and targets. ie. "zip %(targets)s %(changed)s"
    @ivar task(Task): reference to task that contains this action
    """

    def __init__(self, action):
        assert isinstance(action,str), "CmdAction must be a string."
        self.action = action
        self.task = None

    def execute(self, capture_stdout=False, capture_stderr=False):
        """
        Execute command action

        @param capture_stdout(bool): Capture standard output
        @param capture_err(bool): Capture standard error

        @raise TaskError: If subprocess return code is greater than 125
        @raise TaskFailed: If subprocess return code isn't zero (and
        not greater than 125)
        """
        # set Popen stream parameters
        if not capture_stdout:
            stdout = None
        else:
            stdout = subprocess.PIPE
        if not capture_stderr:
            stderr = None
        else:
            stderr = subprocess.PIPE

        action = self.expand_action()

        # spawn task process
        process = subprocess.Popen(action,stdout=stdout,
                                   stderr=stderr, shell=True)

        # log captured stream
        out,err = process.communicate()
        if out:
            logger.log('stdout',out)
        if err:
            logger.log('stderr',err)

        # task error - based on:
        # http://www.gnu.org/software/bash/manual/bashref.html#Exit-Status
        # it doesnt make so much difference to return as Error or Failed anyway
        if process.returncode > 125:
            raise TaskError("Command error: '%s' returned %s" %
                            (action,process.returncode))

        # task failure
        if process.returncode != 0:
            raise TaskFailed("Command failed: '%s' returned %s" %
                             (action,process.returncode))

    def expand_action(self):
        """expand action string using task meta informations
        @returns (string) - expanded string after substitution
        """
        if not self.task:
            return self.action

        subs_dict = {'targets' : " ".join(self.task.targets),
                     'dependencies': " ".join(self.task.dependencies)}
        # just included changed if it is set
        if self.task.dep_changed is not None:
            subs_dict['changed'] = " ".join(self.task.dep_changed)
        return self.action % subs_dict

    def __str__(self):
        return "Cmd: %s" % self.expand_action()

    def __repr__(self):
        return "<CmdAction: '%s'>"  % self.expand_action()

class PythonAction(BaseAction):
    """Python action. Execute a python callable.

    @ivar py_callable: (callable) Python callable that returns a boolean
    result depending on the success of the action
    @ivar args: (sequence)  Extra arguments to be passed to py_callable
    @ivar kwargs: (dict) Extra keyword arguments to be passed to py_callable
    @ivar task(Task): reference to task that contains this action
    """
    def __init__(self, py_callable, args=None, kwargs=None):

        self.py_callable = py_callable
        self.task = None

        if args is None:
            self.args = []
        else:
            self.args = args

        if kwargs is None:
            self.kwargs = {}
        else:
            self.kwargs = kwargs

        # check valid parameters
        if not callable(self.py_callable):
            raise InvalidTask("PythonAction must be a 'callable'.")
        if type(self.args) is not tuple and type(self.args) is not list:
            msg = "args must be a 'tuple' or a 'list'. got '%s'."
            raise InvalidTask(msg % self.args)
        if type(self.kwargs) is not dict:
            msg = "kwargs must be a 'dict'. got '%s'"
            raise InvalidTask(msg % self.kwargs)

    def execute(self, capture_stdout=False, capture_stderr=False):
        """
        Execute command action

        @param capture_stdout (bool): Capture standard output
        @param capture_err (bool): Capture standard error

        @raise TaskFailed: If py_callable returns False
        """
        # set std stream
        if capture_stdout:
            old_stdout = sys.stdout
            sys.stdout = StringIO.StringIO()

        if capture_stderr:
            old_stderr = sys.stderr
            sys.stderr = StringIO.StringIO()

        # prepare action arguments
        argspec = inspect.getargspec(self.py_callable)

        if argspec.keywords:
            kwargs = {'targets': self.task.targets,
                      'dependencies': self.task.dependencies,
                      'changed': self.task.dep_changed}

            # use passed arguments (including their default value) by default
            for index, key in enumerate(kwargs.keys()):
                if (argspec.args and key in argspec.args
                    and len(self.args) > index
                    or (argspec.defaults and len(argspec.defaults) > index)):
                    del kwargs[key]
            kwargs.update(self.kwargs)
        else:
            kwargs = self.kwargs

        # execute action / callable
        try:
            # Python2.4
            try:
                result = self.py_callable(*self.args,**kwargs)
            # in python 2.4 SystemExit and KeyboardInterrupt subclass
            # from Exception.
            except (SystemExit, KeyboardInterrupt), exp:
                raise

            except Exception, exception:
                error = TaskError(exception)
                error.originalException = traceback.format_exception(\
                    exception.__class__, exception,sys.exc_info()[2])
                raise error
        finally:
            # restore std streams /log captured streams
            if capture_stdout:
                logger.log('stdout',sys.stdout.getvalue())
                sys.stdout.close()
                sys.stdout = old_stdout

            if capture_stderr:
                logger.log('stderr',sys.stderr.getvalue())
                sys.stderr.close()
                sys.stderr = old_stderr

        # if callable returns false. Task failed
        if not result:
            raise TaskFailed("Python Task failed: '%s' returned %s" %
                             (self.py_callable, result))

    def __str__(self):
        # get object description excluding runtime memory address
        return "Python: %s"% str(self.py_callable)[1:].split(' at ')[0]

    def __repr__(self):
        return "<PythonAction: '%s'>"% (repr(self.py_callable))

def create_action(action):
    """
    Create action using proper constructor based on the parameter type

    @param action: Action to be created
    @type action: L{BaseAction} subclass object, str, tuple or callable
    @raise InvalidTask: If action parameter type isn't valid
    """
    if isinstance(action, BaseAction):
        return action

    if type(action) is str:
        return CmdAction(action)

    if type(action) is tuple:
        return PythonAction(*action)

    if callable(action):
        return PythonAction(action)

    msg = "Invalid task action type. got %s"
    raise InvalidTask(msg % (action.__class__))

class Task(object):
    """Task

    @ivar name string
    @ivar actions: (list) of L{BaseAction}
    @ivar dependencies list of all dependencies
    @ivar targets list of targets
    @ivar folder_dep: (list - string)
    @ivar task_dep: (list - string)
    @ivar file_dep: (list - string)
    @ivar dep_changed (list - string): list of file-dependencies that changed
          (are not up_to_date). this must be set before
    @ivar run_once: (bool) task without dependencies should run
    @ivar is_subtask: (bool) indicate this task is a subtask.
    @ivar doc: (string) task documentation
    """

    def __init__(self,name, actions, dependencies=(),targets=(),setup=None,
                 is_subtask=False, doc=None):
        """sanity checks and initialization"""
        # dependencies parameter must be a list
        if not ((isinstance(dependencies,list)) or
                (isinstance(dependencies,tuple))):
            msg = ("%s. paramater 'dependencies' must be a list or " +
                   "tuple got:'%s'%s")
            raise InvalidTask(msg%(name, str(dependencies),type(dependencies)))

        # targets parameter must be a list
        if not(isinstance(targets,list) or isinstance(targets,tuple)):
            msg = ("%s. paramater 'targets' must be a list or tuple " +
                   "got:'%s'%s")
            raise InvalidTask(msg % (name, str(targets),type(targets)))

        self.name = name
        self.dependencies = dependencies
        self.targets = targets
        self.setup = setup
        self.run_once = False
        self.is_subtask = is_subtask
        self.dep_changed = None

        if actions is None:
            self.actions = []
        elif type(actions) is list:
            self.actions = [create_action(a) for a in actions]
        else:
        ## DEPREACTED ON 0.4 - REMOVE ON 0.5
            #raise Exception("DEPRECATED must be list")
            msg = "DEPRECATION WARNING: task %s actions must be a list."
            print msg % self.name
            self.actions = [create_action(actions)]
        ## DEPRECATION END

        # Store just first non-empty line as documentation string
        if doc is None:
            self.doc = ''
        else:
            for line in doc.splitlines():
                striped = line.strip()
                if striped:
                    self.doc = striped
                    break
            else:
                self.doc = ''

        # set self as task for all actions
        for action in self.actions:
            action.task = self

        # there are 3 kinds of dependencies: file, task, and folder
        self.folder_dep = []
        self.task_dep = []
        self.file_dep = []
        for dep in self.dependencies:
            # True on the list. set run_once
            if isinstance(dep,bool):
                if not dep:
                    msg = ("%s. bool paramater in 'dependencies' "+
                           "must be True got:'%s'")
                    raise InvalidTask(msg%(name, str(dep)))
                self.run_once = True
            # folder dep ends with a '/'
            elif dep.endswith('/'):
                self.folder_dep.append(dep)
            # task dep starts with a ':'
            elif dep.startswith(':'):
                self.task_dep.append(dep[1:])
            # file dep
            elif isinstance(dep,str):
                self.file_dep.append(dep)

        # run_once can't be used together with file dependencies
        if self.run_once and self.file_dep:
            msg = ("%s. task cant have file and dependencies and True " +
                   "at the same time. (just remove True)")
            raise InvalidTask(msg % name)

    def execute(self, capture_stdout=False, capture_stderr=False):
        """Executes the task.

        @raise TaskFailed: If raised when executing an action
        @raise TaskError: If raised when executing an action
        """
        for action in self.actions:
            action.execute(capture_stdout, capture_stderr)

    def title(self):
        """String representation on output.

        @rtype: str
        @return: Task name and actions
        """
        return "%s => %s"% (self.name, str(self))

    def __str__(self):
        if self.actions:
            return "\n\t".join([str(action) for action in self.actions])

        # A task that contains no actions at all
        # is used as group task
        return "Group: %s" % ", ".join(self.dependencies)

    def __repr__(self):
        return "<Task: %s>"% self.name

def dict_to_task(task_dict):
    """Create a task instance from dictionary.

    The dictionary has the same format as returned by task-generators
    from dodo files.

    @param task_dict (dict): task representation as a dict.
    @raise InvalidTask: If unexpected fields were passed in task_dict
    """
    # TASK_ATTRS: sequence of know attributes(keys) of a task dict.
    TASK_ATTRS = ('name','actions','dependencies','targets','setup', 'doc')
    # FIXME check field 'name'

    # === DEPRECATED on 0.4 (to be removed on 0.5): START
    # check required fields
    if 'action' in task_dict and 'actions' not in task_dict:
        #raise Exception("DEPRECATED action")
        task_dict['actions'] = task_dict['action']
        del task_dict['action']
        print ("DEPRECATION WARNING: Task %s contains 'action' key. "
               "This will be deprecated in future versions, "
               "please use 'actions' instead" % task_dict['name'])
    # === DEPRECATION: END

    if 'actions' not in task_dict:
        raise InvalidTask("Task %s must contain 'actions' field. %s" %
                          (task_dict['name'],task_dict))

    # user friendly. dont go ahead with invalid input.
    for key in task_dict.keys():
        if key not in TASK_ATTRS:
            raise InvalidTask("Task %s contain invalid field: %s"%
                              (task_dict['name'],key))

    return Task(**task_dict)
<|endoftext|>


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

"""Helper functions for working with Caffe2 networks (i.e., operator graphs)."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from collections import OrderedDict
import cPickle as pickle
import logging
import numpy as np
import os
import pprint
import yaml

from caffe2.python import core
from caffe2.python import workspace
from core.config import cfg
from utils.io import save_object
import utils.c2 as c2_utils

logger = logging.getLogger(__name__)

def initialize_from_weights_file(model, weights_file, broadcast=True):
    """Initialize a model from weights stored in a pickled dictionary. If
    multiple GPUs are used, the loaded weights are synchronized on all GPUs,
    unless 'broadcast' is False.
    """
    initialize_gpu_from_weights_file(model, weights_file, gpu_id=0)
    if broadcast:
        broadcast_parameters(model)

def initialize_gpu_from_weights_file(model, weights_file, gpu_id=0):
    """Initialize a network with ops on a specific GPU.

    If you use CUDA_VISIBLE_DEVICES to target specific GPUs, Caffe2 will
    automatically map logical GPU ids (starting from 0) to the physical GPUs
    specified in CUDA_VISIBLE_DEVICES.
    """
    logger.info('Loading weights from: {}'.format(weights_file))
    ws_blobs = workspace.Blobs()
    with open(weights_file, 'r') as f:
        src_blobs = pickle.load(f)
    if 'cfg' in src_blobs:
        saved_cfg = yaml.load(src_blobs['cfg'])
        configure_bbox_reg_weights(model, saved_cfg)
    if 'blobs' in src_blobs:
        # Backwards compat--dictionary used to be only blobs, now they are
        # stored under the 'blobs' key
        src_blobs = src_blobs['blobs']
    # Initialize weights on GPU gpu_id only
    unscoped_param_names = OrderedDict()  # Print these out in model order
    for blob in model.params:
        unscoped_param_names[c2_utils.UnscopeName(str(blob))] = True
    with c2_utils.NamedCudaScope(gpu_id):
        for unscoped_param_name in unscoped_param_names.keys():
            if (unscoped_param_name.find(']_') >= 0 and
                    unscoped_param_name not in src_blobs):
                # Special case for sharing initialization from a pretrained
                # model:
                # If a blob named '_[xyz]_foo' is in model.params and not in
                # the initialization blob dictionary, then load source blob
                # 'foo' into destination blob '_[xyz]_foo'
                src_name = unscoped_param_name[
                    unscoped_param_name.find(']_') + 2:]
            else:
                src_name = unscoped_param_name
            if src_name not in src_blobs:
                logger.info('{:s} not found'.format(src_name))
                continue
            dst_name = core.ScopedName(unscoped_param_name)
            has_momentum = src_name + '_momentum' in src_blobs
            has_momentum_str = ' [+ momentum]' if has_momentum else ''
            logger.debug(
                '{:s}{:} loaded from weights file into {:s}: {}'.format(
                    src_name, has_momentum_str, dst_name, src_blobs[src_name]
                    .shape
                )
            )
            if dst_name in ws_blobs:
                # If the blob is already in the workspace, make sure that it
                # matches the shape of the loaded blob
                ws_blob = workspace.FetchBlob(dst_name)
                assert ws_blob.shape == src_blobs[src_name].shape, \
                    ('Workspace blob {} with shape {} does not match '
                     'weights file shape {}').format(
                        src_name,
                        ws_blob.shape,
                        src_blobs[src_name].shape)
            workspace.FeedBlob(
                dst_name,
                src_blobs[src_name].astype(np.float32, copy=False))
            if has_momentum:
                workspace.FeedBlob(
                    dst_name + '_momentum',
                    src_blobs[src_name + '_momentum'].astype(
                        np.float32, copy=False))

    # We preserve blobs that are in the weights file but not used by the current
    # model. We load these into CPU memory under the '__preserve__/' namescope.
    # These blobs will be stored when saving a model to a weights file. This
    # feature allows for alternating optimization of Faster R-CNN in which blobs
    # unused by one step can still be preserved forward and used to initialize
    # another step.
    for src_name in src_blobs.keys():
        if (src_name not in unscoped_param_names and
                not src_name.endswith('_momentum') and
                src_blobs[src_name] is not None):
            with c2_utils.CpuScope():
                workspace.FeedBlob(
                    '__preserve__/{:s}'.format(src_name), src_blobs[src_name])
                logger.debug(
                    '{:s} preserved in workspace (unused)'.format(src_name))

def save_model_to_weights_file(weights_file, model):
    """Stash model weights in a dictionary and pickle them to a file. We map
    GPU device scoped names to unscoped names (e.g., 'gpu_0/conv1_w' ->
    'conv1_w').
    """
    logger.info(
        'Saving parameters and momentum to {}'.format(
            os.path.abspath(weights_file)))
    blobs = {}
    # Save all parameters
    for param in model.params:
        scoped_name = str(param)
        unscoped_name = c2_utils.UnscopeName(scoped_name)
        if unscoped_name not in blobs:
            logger.debug(' {:s} -> {:s}'.format(scoped_name, unscoped_name))
            blobs[unscoped_name] = workspace.FetchBlob(scoped_name)
    # Save momentum
    for param in model.TrainableParams():
        scoped_name = str(param) + '_momentum'
        unscoped_name = c2_utils.UnscopeName(scoped_name)
        if unscoped_name not in blobs:
            logger.debug(' {:s} -> {:s}'.format(scoped_name, unscoped_name))
            blobs[unscoped_name] = workspace.FetchBlob(scoped_name)
    # Save preserved blobs
    for scoped_name in workspace.Blobs():
        if scoped_name.startswith('__preserve__/'):
            unscoped_name = c2_utils.UnscopeName(scoped_name)
            if unscoped_name not in blobs:
                logger.debug(
                    ' {:s} -> {:s} (preserved)'.format(
                        scoped_name, unscoped_name))
                blobs[unscoped_name] = workspace.FetchBlob(scoped_name)
    cfg_yaml = yaml.dump(cfg)
    save_object(dict(blobs=blobs, cfg=cfg_yaml), weights_file)

def broadcast_parameters(model):
    """Copy parameter blobs from GPU 0 over the corresponding parameter blobs
    on GPUs 1 through cfg.NUM_GPUS - 1.
    """
    if cfg.NUM_GPUS == 1:
        # no-op if only running on a single GPU
        return

    def _do_broadcast(all_blobs):
        assert len(all_blobs) % cfg.NUM_GPUS == 0, \
            ('Unexpected value for NUM_GPUS. Make sure you are not '
             'running single-GPU inference with NUM_GPUS > 1.')
        blobs_per_gpu = int(len(all_blobs) / cfg.NUM_GPUS)
        for i in range(blobs_per_gpu):
            blobs = [p for p in all_blobs[i::blobs_per_gpu]]
            data = workspace.FetchBlob(blobs[0])
            logger.debug('Broadcasting {} to'.format(str(blobs[0])))
            for i, p in enumerate(blobs[1:]):
                logger.debug(' |-> {}'.format(str(p)))
                with c2_utils.CudaScope(i + 1):
                    workspace.FeedBlob(p, data)

    _do_broadcast(model.params)
    _do_broadcast([b + '_momentum' for b in model.TrainableParams()])

def sum_multi_gpu_blob(blob_name):
    """Return the sum of a scalar blob held on multiple GPUs."""
    val = 0
    for i in range(cfg.NUM_GPUS):
        val += float(workspace.FetchBlob('gpu_{}/{}'.format(i, blob_name)))
    return val

def average_multi_gpu_blob(blob_name):
    """Return the average of a scalar blob held on multiple GPUs."""
    return sum_multi_gpu_blob(blob_name) / cfg.NUM_GPUS

def print_net(model, namescope='gpu_0'):
    """Print the model network."""
    logger.info('Printing model: {}'.format(model.net.Name()))
    op_list = model.net.Proto().op
    for op in op_list:
        input_name = op.input
        # For simplicity: only print the first output
        # Not recommended if there are split layers
        output_name = str(op.output[0])
        op_type = op.type
        op_name = op.name

        if namescope is None or output_name.startswith(namescope):
            # Only print the forward pass network
            if output_name.find('grad') >= 0 or output_name.find('__m') >= 0:
                continue

            output_shape = workspace.FetchBlob(output_name).shape
            first_blob = True
            op_label = op_type + (op_name if op_name == '' else ':' + op_name)
            suffix = ' ------- (op: {})'.format(op_label)
            for j in range(len(input_name)):
                if input_name[j] in model.params:
                    continue
                input_blob = workspace.FetchBlob(input_name[j])
                if isinstance(input_blob, np.ndarray):
                    input_shape = input_blob.shape
                    logger.info('{:28s}: {:20s} => {:28s}: {:20s}{}'.format(
                        c2_utils.UnscopeName(str(input_name[j])),
                        '{}'.format(input_shape),
                        c2_utils.UnscopeName(str(output_name)),
                        '{}'.format(output_shape),
                        suffix))
                    if first_blob:
                        first_blob = False
                        suffix = ' ------|'
    logger.info('End of model: {}'.format(model.net.Name()))

def configure_bbox_reg_weights(model, saved_cfg):
    """Compatibility for old models trained with bounding box regression
    mean/std normalization (instead of fixed weights).
    """
    if 'MODEL' not in saved_cfg or 'BBOX_REG_WEIGHTS' not in saved_cfg.MODEL:
        logger.warning('Model from weights file was trained before config key '
                       'MODEL.BBOX_REG_WEIGHTS was added. Forcing '
                       'MODEL.BBOX_REG_WEIGHTS = (1., 1., 1., 1.) to ensure '
                       'correct **inference** behavior.')
        # Generally we don't allow modifying the config, but this is a one-off
        # hack to support some very old models
        is_immutable = cfg.is_immutable()
        cfg.immutable(False)
        cfg.MODEL.BBOX_REG_WEIGHTS = (1., 1., 1., 1.)
        cfg.immutable(is_immutable)
        logger.info('New config:')
        logger.info(pprint.pformat(cfg))
        assert not model.train, (
            'This model was trained with an older version of the code that '
            'used bounding box regression mean/std normalization. It can no '
            'longer be used for training. To upgrade it to a trainable model '
            'please use fb/compat/convert_bbox_reg_normalized_model.py.'
        )
<commit_msg>Allow printing net even when output blobs are absent<commit_after># Copyright (c) 2017-present, Facebook, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##############################################################################

"""Helper functions for working with Caffe2 networks (i.e., operator graphs)."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from collections import OrderedDict
import cPickle as pickle
import logging
import numpy as np
import os
import pprint
import yaml

from caffe2.python import core
from caffe2.python import workspace
from core.config import cfg
from utils.io import save_object
import utils.c2 as c2_utils

logger = logging.getLogger(__name__)

def initialize_from_weights_file(model, weights_file, broadcast=True):
    """Initialize a model from weights stored in a pickled dictionary. If
    multiple GPUs are used, the loaded weights are synchronized on all GPUs,
    unless 'broadcast' is False.
    """
    initialize_gpu_from_weights_file(model, weights_file, gpu_id=0)
    if broadcast:
        broadcast_parameters(model)

def initialize_gpu_from_weights_file(model, weights_file, gpu_id=0):
    """Initialize a network with ops on a specific GPU.

    If you use CUDA_VISIBLE_DEVICES to target specific GPUs, Caffe2 will
    automatically map logical GPU ids (starting from 0) to the physical GPUs
    specified in CUDA_VISIBLE_DEVICES.
    """
    logger.info('Loading weights from: {}'.format(weights_file))
    ws_blobs = workspace.Blobs()
    with open(weights_file, 'r') as f:
        src_blobs = pickle.load(f)
    if 'cfg' in src_blobs:
        saved_cfg = yaml.load(src_blobs['cfg'])
        configure_bbox_reg_weights(model, saved_cfg)
    if 'blobs' in src_blobs:
        # Backwards compat--dictionary used to be only blobs, now they are
        # stored under the 'blobs' key
        src_blobs = src_blobs['blobs']
    # Initialize weights on GPU gpu_id only
    unscoped_param_names = OrderedDict()  # Print these out in model order
    for blob in model.params:
        unscoped_param_names[c2_utils.UnscopeName(str(blob))] = True
    with c2_utils.NamedCudaScope(gpu_id):
        for unscoped_param_name in unscoped_param_names.keys():
            if (unscoped_param_name.find(']_') >= 0 and
                    unscoped_param_name not in src_blobs):
                # Special case for sharing initialization from a pretrained
                # model:
                # If a blob named '_[xyz]_foo' is in model.params and not in
                # the initialization blob dictionary, then load source blob
                # 'foo' into destination blob '_[xyz]_foo'
                src_name = unscoped_param_name[
                    unscoped_param_name.find(']_') + 2:]
            else:
                src_name = unscoped_param_name
            if src_name not in src_blobs:
                logger.info('{:s} not found'.format(src_name))
                continue
            dst_name = core.ScopedName(unscoped_param_name)
            has_momentum = src_name + '_momentum' in src_blobs
            has_momentum_str = ' [+ momentum]' if has_momentum else ''
            logger.debug(
                '{:s}{:} loaded from weights file into {:s}: {}'.format(
                    src_name, has_momentum_str, dst_name, src_blobs[src_name]
                    .shape
                )
            )
            if dst_name in ws_blobs:
                # If the blob is already in the workspace, make sure that it
                # matches the shape of the loaded blob
                ws_blob = workspace.FetchBlob(dst_name)
                assert ws_blob.shape == src_blobs[src_name].shape, \
                    ('Workspace blob {} with shape {} does not match '
                     'weights file shape {}').format(
                        src_name,
                        ws_blob.shape,
                        src_blobs[src_name].shape)
            workspace.FeedBlob(
                dst_name,
                src_blobs[src_name].astype(np.float32, copy=False))
            if has_momentum:
                workspace.FeedBlob(
                    dst_name + '_momentum',
                    src_blobs[src_name + '_momentum'].astype(
                        np.float32, copy=False))

    # We preserve blobs that are in the weights file but not used by the current
    # model. We load these into CPU memory under the '__preserve__/' namescope.
    # These blobs will be stored when saving a model to a weights file. This
    # feature allows for alternating optimization of Faster R-CNN in which blobs
    # unused by one step can still be preserved forward and used to initialize
    # another step.
    for src_name in src_blobs.keys():
        if (src_name not in unscoped_param_names and
                not src_name.endswith('_momentum') and
                src_blobs[src_name] is not None):
            with c2_utils.CpuScope():
                workspace.FeedBlob(
                    '__preserve__/{:s}'.format(src_name), src_blobs[src_name])
                logger.debug(
                    '{:s} preserved in workspace (unused)'.format(src_name))

def save_model_to_weights_file(weights_file, model):
    """Stash model weights in a dictionary and pickle them to a file. We map
    GPU device scoped names to unscoped names (e.g., 'gpu_0/conv1_w' ->
    'conv1_w').
    """
    logger.info(
        'Saving parameters and momentum to {}'.format(
            os.path.abspath(weights_file)))
    blobs = {}
    # Save all parameters
    for param in model.params:
        scoped_name = str(param)
        unscoped_name = c2_utils.UnscopeName(scoped_name)
        if unscoped_name not in blobs:
            logger.debug(' {:s} -> {:s}'.format(scoped_name, unscoped_name))
            blobs[unscoped_name] = workspace.FetchBlob(scoped_name)
    # Save momentum
    for param in model.TrainableParams():
        scoped_name = str(param) + '_momentum'
        unscoped_name = c2_utils.UnscopeName(scoped_name)
        if unscoped_name not in blobs:
            logger.debug(' {:s} -> {:s}'.format(scoped_name, unscoped_name))
            blobs[unscoped_name] = workspace.FetchBlob(scoped_name)
    # Save preserved blobs
    for scoped_name in workspace.Blobs():
        if scoped_name.startswith('__preserve__/'):
            unscoped_name = c2_utils.UnscopeName(scoped_name)
            if unscoped_name not in blobs:
                logger.debug(
                    ' {:s} -> {:s} (preserved)'.format(
                        scoped_name, unscoped_name))
                blobs[unscoped_name] = workspace.FetchBlob(scoped_name)
    cfg_yaml = yaml.dump(cfg)
    save_object(dict(blobs=blobs, cfg=cfg_yaml), weights_file)

def broadcast_parameters(model):
    """Copy parameter blobs from GPU 0 over the corresponding parameter blobs
    on GPUs 1 through cfg.NUM_GPUS - 1.
    """
    if cfg.NUM_GPUS == 1:
        # no-op if only running on a single GPU
        return

    def _do_broadcast(all_blobs):
        assert len(all_blobs) % cfg.NUM_GPUS == 0, \
            ('Unexpected value for NUM_GPUS. Make sure you are not '
             'running single-GPU inference with NUM_GPUS > 1.')
        blobs_per_gpu = int(len(all_blobs) / cfg.NUM_GPUS)
        for i in range(blobs_per_gpu):
            blobs = [p for p in all_blobs[i::blobs_per_gpu]]
            data = workspace.FetchBlob(blobs[0])
            logger.debug('Broadcasting {} to'.format(str(blobs[0])))
            for i, p in enumerate(blobs[1:]):
                logger.debug(' |-> {}'.format(str(p)))
                with c2_utils.CudaScope(i + 1):
                    workspace.FeedBlob(p, data)

    _do_broadcast(model.params)
    _do_broadcast([b + '_momentum' for b in model.TrainableParams()])

def sum_multi_gpu_blob(blob_name):
    """Return the sum of a scalar blob held on multiple GPUs."""
    val = 0
    for i in range(cfg.NUM_GPUS):
        val += float(workspace.FetchBlob('gpu_{}/{}'.format(i, blob_name)))
    return val

def average_multi_gpu_blob(blob_name):
    """Return the average of a scalar blob held on multiple GPUs."""
    return sum_multi_gpu_blob(blob_name) / cfg.NUM_GPUS

def print_net(model, namescope='gpu_0'):
    """Print the model network."""
    logger.info('Printing model: {}'.format(model.net.Name()))
    op_list = model.net.Proto().op
    for op in op_list:
        input_name = op.input
        # For simplicity: only print the first output
        # Not recommended if there are split layers
        output_name = str(op.output[0])
        op_type = op.type
        op_name = op.name

        if namescope is None or output_name.startswith(namescope):
            # Only print the forward pass network
            if output_name.find('grad') >= 0 or output_name.find('__m') >= 0:
                continue

            try:
                # Under some conditions (e.g., dynamic memory optimization)
                # it is possible that the network frees some blobs when they are
                # no longer needed. Handle this case...
                output_shape = workspace.FetchBlob(output_name).shape
            except BaseException:
                output_shape = '<unknown>'

            first_blob = True
            op_label = op_type + (op_name if op_name == '' else ':' + op_name)
            suffix = ' ------- (op: {})'.format(op_label)
            for j in range(len(input_name)):
                if input_name[j] in model.params:
                    continue
                input_blob = workspace.FetchBlob(input_name[j])
                if isinstance(input_blob, np.ndarray):
                    input_shape = input_blob.shape
                    logger.info('{:28s}: {:20s} => {:28s}: {:20s}{}'.format(
                        c2_utils.UnscopeName(str(input_name[j])),
                        '{}'.format(input_shape),
                        c2_utils.UnscopeName(str(output_name)),
                        '{}'.format(output_shape),
                        suffix))
                    if first_blob:
                        first_blob = False
                        suffix = ' ------|'
    logger.info('End of model: {}'.format(model.net.Name()))

def configure_bbox_reg_weights(model, saved_cfg):
    """Compatibility for old models trained with bounding box regression
    mean/std normalization (instead of fixed weights).
    """
    if 'MODEL' not in saved_cfg or 'BBOX_REG_WEIGHTS' not in saved_cfg.MODEL:
        logger.warning('Model from weights file was trained before config key '
                       'MODEL.BBOX_REG_WEIGHTS was added. Forcing '
                       'MODEL.BBOX_REG_WEIGHTS = (1., 1., 1., 1.) to ensure '
                       'correct **inference** behavior.')
        # Generally we don't allow modifying the config, but this is a one-off
        # hack to support some very old models
        is_immutable = cfg.is_immutable()
        cfg.immutable(False)
        cfg.MODEL.BBOX_REG_WEIGHTS = (1., 1., 1., 1.)
        cfg.immutable(is_immutable)
        logger.info('New config:')
        logger.info(pprint.pformat(cfg))
        assert not model.train, (
            'This model was trained with an older version of the code that '
            'used bounding box regression mean/std normalization. It can no '
            'longer be used for training. To upgrade it to a trainable model '
            'please use fb/compat/convert_bbox_reg_normalized_model.py.'
        )
<|endoftext|>


<commit_before>import os
from flask import Flask, render_template, Response
import requests
import json

ROOT_DIR = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '..'))
DATA_DIR = os.path.join(ROOT_DIR, 'data')

app = Flask(__name__)

'''
=== Views ==
'''
@app.route('/')
def home():
    issues = get_issue_list()
    return render_template('index.html', issues=issues)

#@app.route('/<lib>')
#def detail(lib=None):
#    return render_template('lib_details.html', lib=lib)

@app.route('/get_data_1/')
def get_data_1():
  jsonfile = os.path.join(DATA_DIR, 'libdata_2013-08-28.json')

  with open(jsonfile, 'rb') as infile:
    ddata = json.load(infile)

  jsondata = json.dumps(ddata)

  #Construct response with JSONdata
  resp = Response (jsondata, status=200,
      mimetype='application/json')

  return resp

'''
=== Helper functions ==
Can move these to seperate file
'''

def get_issue_list():
    root_url = 'https://api.github.com'
    url = root_url + '/repos/amanahuja/lib-stats/issues'
    r = requests.get(url)
    results = json.loads(r.content)

    #If list is small, get closed issues as well.
    if len(results) < 5:
        url = url + '?state=closed'
        r = requests.get(url)
        results = results + json.loads(r.content)
    issue_list = []

    #limit dictionary
    for ii in results:
        issue = {}
        issue['title'] = ii['title']
        issue['body'] = ii['body']
        issue['url'] = ii['html_url']
        issue['state'] = ii['state']
        issue_list.append(issue)

    return issue_list

def get_lib_list():
  pass

if __name__ == '__main__':
  """
  todo: Not running as module, so use Test Configuration
  """
  port = int(os.environ.get('PORT', 5000))
  app.debug = True
  #app.run(host='amanahuja.me', port=port)
  app.run(port=port)

<commit_msg>new day, new data<commit_after>import os
from flask import Flask, render_template, Response
import requests
import json

ROOT_DIR = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '..'))
DATA_DIR = os.path.join(ROOT_DIR, 'data')

app = Flask(__name__)

'''
=== Views ==
'''
@app.route('/')
def home():
    issues = get_issue_list()
    return render_template('index.html', issues=issues)

#@app.route('/<lib>')
#def detail(lib=None):
#    return render_template('lib_details.html', lib=lib)

@app.route('/get_data_1/')
def get_data_1():
  jsonfile = os.path.join(DATA_DIR, 'libdata_2013-08-29.json')

  with open(jsonfile, 'rb') as infile:
    ddata = json.load(infile)

  jsondata = json.dumps(ddata)

  #Construct response with JSONdata
  resp = Response (jsondata, status=200,
      mimetype='application/json')

  return resp

'''
=== Helper functions ==
Can move these to seperate file
'''

def get_issue_list():
    root_url = 'https://api.github.com'
    url = root_url + '/repos/amanahuja/lib-stats/issues'
    r = requests.get(url)
    results = json.loads(r.content)

    #If list is small, get closed issues as well.
    if len(results) < 5:
        url = url + '?state=closed'
        r = requests.get(url)
        results = results + json.loads(r.content)
    issue_list = []

    #limit dictionary
    for ii in results:
        issue = {}
        issue['title'] = ii['title']
        issue['body'] = ii['body']
        issue['url'] = ii['html_url']
        issue['state'] = ii['state']
        issue_list.append(issue)

    return issue_list

def get_lib_list():
  pass

if __name__ == '__main__':
  """
  todo: Not running as module, so use Test Configuration
  """
  port = int(os.environ.get('PORT', 5000))
  app.debug = True
  #app.run(host='amanahuja.me', port=port)
  app.run(port=port)

<|endoftext|>