sdk/python/arvados.py

   1 import gflags
   2 import httplib2
   3 import logging
   4 import os
   5 import pprint
   6 import sys
   7 import types
   8
   9 from apiclient import errors
  10 from apiclient.discovery import build
  11
  12 class CredentialsFromEnv:
  13     @staticmethod
  14     def http_request(self, uri, **kwargs):
  15         if 'headers' not in kwargs:
  16             kwargs['headers'] = {}
  17         kwargs['headers']['Authorization'] = 'OAuth2 %s' % os.environ['ARVADOS_API_TOKEN']
  18         return self.orig_http_request(uri, **kwargs)
  19     def authorize(self, http):
  20         http.orig_http_request = http.request
  21         http.request = types.MethodType(self.http_request, http)
  22         return http
  23
  24 url = ('https://%s/discovery/v1/apis/'
  25        '{api}/{apiVersion}/rest' % os.environ['ARVADOS_API_HOST'])
  26 credentials = CredentialsFromEnv()
  27 http = httplib2.Http()
  28 http = credentials.authorize(http)
  29 http.disable_ssl_certificate_validation=True
  30 service = build("arvados", "v1", http=http, discoveryServiceUrl=url)
  31
  32 def current_task():
  33     t = service.job_tasks().get(uuid=os.environ['TASK_UUID']).execute()
  34     return t
  35
  36 def current_job():
  37     t = service.jobs().get(uuid=os.environ['JOB_UUID']).execute()
  38     return t
  39
  40 class JobTask:
  41     def __init__(self, parameters=dict(), resource_limits=dict()):
  42         print "init jobtask %s %s" % (parameters, resource_limits)
  43
  44 class job_setup:
  45     @staticmethod
  46     def one_task_per_input_file(if_sequence=0, and_end_task=True):
  47         if if_sequence != current_job()['sequence']:
  48             return
  49         job_input = current_job()['parameters']['input']
  50         p = subprocess.Popen(["whls", job_input],
  51                              stdout=subprocess.PIPE,
  52                              stdin=None, stderr=None,
  53                              shell=False, close_fds=True)
  54         for f in p.stdout.read().split("\n"):
  55             if f != '':
  56                 task_input = job_input + '/' + f
  57                 new_task_attrs = {
  58                     'job_uuid': current_job()['uuid'],
  59                     'parameters': {
  60                         'input':task_input
  61                         }
  62                     }
  63                 service.jobs_tasks().create(job_task=new_task_attrs)
  64         if and_end_task:
  65             service.job_tasks().update(uuid=current_task()['uuid'],
  66                                        job_task={'success':True})
  67             exit 0