qcg/service.py

   1 import logging
   2 import time
   3
   4 from django.db import transaction
   5 from django.utils.timezone import now
   6 from pyqcg import QCG
   7 from pyqcg.description import JobDescription
   8 from pyqcg.service import Registry
   9 from pyqcg.utils import Credential, TimePeriod, JobStatus, TaskStatus
  10
  11
  12 logger = logging.getLogger(__name__)
  13
  14
  15 @transaction.atomic
  16 def update_user_data(user, proxy):
  17     ts = time.time()
  18     QCG.start()
  19     from qcg.models import User, Job, Task, Allocation, NodeInfo
  20
  21     credential = Credential(proxy)
  22     registry = Registry(credential)
  23
  24     # put lock on user record (hopefully..?)
  25     user = User.objects.select_for_update().get(pk=user.pk)
  26
  27     changed_filter = {'changed': TimePeriod(after=user.last_update)}
  28
  29     ###################################
  30     # Jobs
  31     ###################################
  32     jts = time.time()
  33     jobs = registry.jobs(stats=[v for k, v in vars(JobStatus).iteritems() if not k.startswith('__')],
  34                          **changed_filter)
  35     jte = time.time()
  36
  37     jobs_count = 0
  38     for qcg_job in jobs:
  39         params = Job.qcg_map(qcg_job, user)
  40         job_id = params.pop('job_id')
  41
  42         Job.objects.update_or_create(job_id=job_id, defaults=params)
  43         jobs_count += 1
  44
  45     ###################################
  46     # Tasks
  47     ###################################
  48     tts = time.time()
  49     tasks = registry.tasks(stats=[v for k, v in vars(TaskStatus).iteritems() if not k.startswith('__')],
  50                            **changed_filter)
  51     tte = time.time()
  52
  53     jobs_cache = {j.job_id: j for j in Job.objects.filter(owner=user)}
  54     task_count = 0
  55     for qcg_task in tasks:
  56         params = Task.qcg_map(qcg_task, jobs_cache)
  57         task_id = params.pop('task_id')
  58
  59         task, created = Task.objects.update_or_create(job__job_id=qcg_task.job_id, task_id=task_id, defaults=params)
  60
  61         if not created:
  62             task.allocations.all().delete()
  63
  64         for qcg_alloc in qcg_task.allocations:
  65             alloc = task.allocations.create(**Allocation.qcg_map(qcg_alloc))
  66
  67             for qcg_node in qcg_alloc.nodes:
  68                 alloc.nodes.create(**NodeInfo.qcg_map(qcg_node))
  69
  70         task_count += 1
  71
  72     # release user lock
  73     user.last_update = now()
  74     user.save()
  75
  76     elapsed = time.time() - ts
  77     elapsed_jobs = jte - jts
  78     elapsed_tasks = tte - tts
  79     elapsed_py = elapsed - elapsed_jobs - elapsed_tasks
  80     logger.info('(%.3f) USER = %s, JOBS = %s (%.3f), TASKS = %s (%.3f), TIME = %.3f',
  81                 elapsed, user, jobs_count, elapsed_jobs, task_count, elapsed_tasks, elapsed_py)
  82
  83
  84 def submit_job(params, proxy):
  85     # print params
  86
  87     QCG.start()
  88     desc = JobDescription(Credential(proxy))
  89
  90     direct_map = ('env_variables', 'executable', 'arguments', 'note', 'grant', 'hosts', 'properties', 'queue', 'procs',
  91                   'wall_time', 'memory', 'memory_per_slot', 'modules', 'input', 'stage_in', 'native', 'notify',
  92                   'preprocess', 'postprocess', 'persistent')
  93
  94     for name in direct_map:
  95         if params[name]:
  96             setattr(desc, name, params[name])
  97
  98     if params['application']:
  99         desc.set_application(*params['application'])
 100     if params['nodes']:
 101         desc.set_nodes(*params['nodes'])
 102     if params['reservation']:
 103         desc.set_reservation(params['reservation'])
 104     if params['watch_output']:
 105         desc.set_watch_output(params['watch_output'], params['watch_output_pattern'])
 106     # TODO script
 107     # TODO stage_out
 108     # TODO monitoring
 109
 110     # for prop in direct_map + ('application', 'nodes', 'env_variables', 'reservation', 'watch_output'):
 111     #     print prop, type(getattr(desc, prop)), repr(getattr(desc, prop))
 112
 113     # print desc.xml_description
 114
 115     # job = desc.submit()
 116
 117     # return job.job_id