sdk/python/arvados/__init__.py

   1 import gflags
   2 import httplib
   3 import httplib2
   4 import logging
   5 import os
   6 import pprint
   7 import sys
   8 import types
   9 import subprocess
  10 import json
  11 import UserDict
  12 import re
  13 import hashlib
  14 import string
  15 import bz2
  16 import zlib
  17 import fcntl
  18 import time
  19 import threading
  20
  21 EMPTY_BLOCK_LOCATOR = 'd41d8cd98f00b204e9800998ecf8427e+0'
  22
  23 from api import *
  24 from stream import *
  25 from collection import *
  26 from keep import *
  27 import util
  28
  29 class errors:
  30     class SyntaxError(Exception):
  31         pass
  32     class AssertionError(Exception):
  33         pass
  34     class NotFoundError(Exception):
  35         pass
  36     class CommandFailedError(Exception):
  37         pass
  38     class KeepWriteError(Exception):
  39         pass
  40     class NotImplementedError(Exception):
  41         pass
  42
  43 def task_set_output(self,s):
  44     api('v1').job_tasks().update(uuid=self['uuid'],
  45                                  body={
  46             'output':s,
  47             'success':True,
  48             'progress':1.0
  49             }).execute()
  50
  51 _current_task = None
  52 def current_task():
  53     global _current_task
  54     if _current_task:
  55         return _current_task
  56     t = api('v1').job_tasks().get(uuid=os.environ['TASK_UUID']).execute()
  57     t = UserDict.UserDict(t)
  58     t.set_output = types.MethodType(task_set_output, t)
  59     t.tmpdir = os.environ['TASK_WORK']
  60     _current_task = t
  61     return t
  62
  63 _current_job = None
  64 def current_job():
  65     global _current_job
  66     if _current_job:
  67         return _current_job
  68     t = api('v1').jobs().get(uuid=os.environ['JOB_UUID']).execute()
  69     t = UserDict.UserDict(t)
  70     t.tmpdir = os.environ['JOB_WORK']
  71     _current_job = t
  72     return t
  73
  74 def getjobparam(*args):
  75     return current_job()['script_parameters'].get(*args)
  76
  77 class JobTask(object):
  78     def __init__(self, parameters=dict(), runtime_constraints=dict()):
  79         print "init jobtask %s %s" % (parameters, runtime_constraints)
  80
  81 class job_setup:
  82     @staticmethod
  83     def one_task_per_input_file(if_sequence=0, and_end_task=True):
  84         if if_sequence != current_task()['sequence']:
  85             return
  86         job_input = current_job()['script_parameters']['input']
  87         cr = CollectionReader(job_input)
  88         for s in cr.all_streams():
  89             for f in s.all_files():
  90                 task_input = f.as_manifest()
  91                 new_task_attrs = {
  92                     'job_uuid': current_job()['uuid'],
  93                     'created_by_job_task_uuid': current_task()['uuid'],
  94                     'sequence': if_sequence + 1,
  95                     'parameters': {
  96                         'input':task_input
  97                         }
  98                     }
  99                 api('v1').job_tasks().create(body=new_task_attrs).execute()
 100         if and_end_task:
 101             api('v1').job_tasks().update(uuid=current_task()['uuid'],
 102                                        body={'success':True}
 103                                        ).execute()
 104             exit(0)
 105
 106     @staticmethod
 107     def one_task_per_input_stream(if_sequence=0, and_end_task=True):
 108         if if_sequence != current_task()['sequence']:
 109             return
 110         job_input = current_job()['script_parameters']['input']
 111         cr = CollectionReader(job_input)
 112         for s in cr.all_streams():
 113             task_input = s.tokens()
 114             new_task_attrs = {
 115                 'job_uuid': current_job()['uuid'],
 116                 'created_by_job_task_uuid': current_task()['uuid'],
 117                 'sequence': if_sequence + 1,
 118                 'parameters': {
 119                     'input':task_input
 120                     }
 121                 }
 122             api('v1').job_tasks().create(body=new_task_attrs).execute()
 123         if and_end_task:
 124             api('v1').job_tasks().update(uuid=current_task()['uuid'],
 125                                        body={'success':True}
 126                                        ).execute()
 127             exit(0)
 128
 129