qcg/utils.py

   1 import os
   2 import string
   3 import random
   4
   5 from django.core.paginator import Paginator
   6 from django.utils.formats import date_format
   7 from django.utils.timezone import localtime
   8 from pyqcg import QCG
   9 from pyqcg.utils import Credential
  10 from pyqcg.description import JobDescription
  11
  12 from filex.ftp import FTPOperation
  13 from qcg import constants
  14
  15
  16 def get_attributes(obj, attrs):
  17     return {name: getattr(obj, name) for name in attrs if getattr(obj, name) is not None}
  18
  19
  20 def username_from_dn(dn):
  21     _, username = dn.rsplit('=', 1)
  22
  23     return username
  24
  25
  26 def try_parse_int(s, default=None, base=10):
  27     try:
  28         return int(s, base)
  29     except (TypeError, ValueError):
  30         return default
  31
  32
  33 def paginator_context(request, objects, per_page=constants.PER_PAGE):
  34     paginator = Paginator(objects, per_page)
  35
  36     page_num = try_parse_int(request.GET.get('page'), 1)
  37     if not (1 <= page_num <= paginator.num_pages):
  38         page_num = 1
  39
  40     pages_range = range(max(2, min(page_num - 2, paginator.num_pages - 4)),
  41                         min(max(page_num + 2, 5), paginator.num_pages) + 1)
  42
  43     return {'page': paginator.page(page_num), 'num_pages': paginator.num_pages, 'pages_range': pages_range}
  44
  45
  46 def localtime_str(datetime):
  47     return date_format(localtime(datetime), 'DATETIME_FORMAT')
  48
  49
  50 def random_id(size=8, chars=string.ascii_uppercase + string.digits):
  51     return ''.join(random.choice(chars) for _ in range(size))
  52
  53
  54 def chunks(seq, size):
  55     return (seq[pos:pos + size] for pos in xrange(0, len(seq), size))
  56
  57
  58 def to_job_desc(params, proxy):
  59     QCG.start()
  60     desc = JobDescription(Credential(proxy))
  61
  62     direct_map = ('env_variables', 'executable', 'arguments', 'note', 'grant', 'hosts', 'properties', 'queue', 'procs',
  63                   'wall_time', 'memory', 'memory_per_slot', 'modules', 'input', 'stage_in', 'native', 'notify',
  64                   'preprocess', 'postprocess', 'persistent')
  65
  66     for name in direct_map:
  67         if params[name]:
  68             setattr(desc, name, params[name])
  69
  70     if params['application']:
  71         desc.set_application(*params['application'])
  72         desc.stage_in += [params['master_file']]
  73         desc.arguments.insert(0, os.path.basename(params['master_file']))
  74     if params['script']:
  75         ftp = FTPOperation(proxy)
  76
  77         ftp.mkdir(constants.QCG_DATA_URL, parents=True)
  78         url = os.path.join(constants.QCG_DATA_URL, 'script.{}.sh'.format(random_id()))
  79         ftp.put(url)
  80
  81         for chunk in chunks(params['script'], 4096):
  82             ftp.stream.put(chunk)
  83         ftp.stream.put(None)
  84         ftp.wait()
  85
  86         desc.executable = url
  87     if params['nodes']:
  88         desc.set_nodes(*params['nodes'])
  89     if params['reservation']:
  90         desc.set_reservation(params['reservation'])
  91     if params['watch_output']:
  92         desc.set_watch_output(params['watch_output'], params['watch_output_pattern'])
  93     # TODO monitoring
  94
  95     return desc
  96
  97
  98 def to_form_data(xml):
  99     # prevent circular import errors
 100     from forms import JobDescriptionForm
 101
 102     QCG.start()
 103     desc = JobDescription()
 104     desc.xml_description = xml
 105
 106     direct_map = ('env_variables', 'executable', 'arguments', 'note', 'grant', 'hosts', 'properties', 'queue', 'procs',
 107                   'wall_time', 'memory', 'memory_per_slot', 'modules', 'input', 'stage_in', 'native', 'persistent')
 108
 109     params = {}
 110     for name in direct_map:
 111         attr = getattr(desc, name)
 112         if isinstance(attr, bool) or attr:
 113             params[name] = attr
 114
 115     if desc.application is not None:
 116         app_name, app_ver = desc.application
 117         params['application'] = app_name if app_ver is None else app_name + '/' + app_name
 118         stage_in = params['stage_in']
 119         params['stage_in'], params['master_file'] = stage_in[:-1], stage_in[-1]
 120         params['arguments'] = params['arguments'][1:]
 121     if desc.nodes is not None:
 122         params['nodes'] = ':'.join(map(str, desc.nodes))
 123     if desc.reservation is not None:
 124         res_id, res_type = desc.reservation
 125         params['reservation'] = res_id
 126     if desc.notify is not None:
 127         params['notify_type'], params['notify_address'] = desc.notify.split(':')
 128     if desc.watch_output is not None:
 129         watch_output, params['watch_output_pattern'] = desc.watch_output
 130         params['watch_output_type'], params['watch_output_address'] = watch_output.split(':')
 131     if desc.preprocess is not None:
 132         if desc.preprocess.startswith('gsiftp://'):
 133             params['preprocess_type'] = JobDescriptionForm.Process.SCRIPT
 134             params['preprocess_script'] = desc.preprocess
 135         else:
 136             params['preprocess_type'] = JobDescriptionForm.Process.CMD
 137             params['preprocess_cmd'] = desc.preprocess
 138     if desc.postprocess is not None:
 139         if desc.postprocess.startswith('gsiftp://'):
 140             params['postprocess_type'] = JobDescriptionForm.Process.SCRIPT
 141             params['postprocess_script'] = desc.postprocess
 142         else:
 143             params['postprocess_type'] = JobDescriptionForm.Process.CMD
 144             params['postprocess_cmd'] = desc.postprocess
 145
 146     return params