This repository has been archived by the owner on Nov 2, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 1
/
aws_batch.py
155 lines (124 loc) · 5.89 KB
/
aws_batch.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
import logging
import boto3
from .base import PushWorker
class AWSBatchPushWorker(PushWorker):
supports_state_sync = True
push_type = 'aws_batch'
def __init__(self, *args, default_job_queue=None, default_job_definition=None, **kwargs):
self.logger = logging.getLogger('AWSBatchPushWorker')
super(AWSBatchPushWorker, self).__init__(*args, **kwargs)
self.batch_client = boto3.client('batch')
self.default_job_queue = default_job_queue
self.default_job_definition = default_job_definition
def get_log_url(self, task_instance):
try:
push_state = task_instance.push_state
if push_state and 'taskArn' in push_state:
task_id = re.match(r'task/(.+)', push_state['taskArn']).groups()[0]
return '{}/{}/{}'.format(push_state['jobName'][:50], push_state['jobId'], task_id)
except:
self.logger.exception('Exception getting AWS Batch log URL')
return None
def sync_task_instance_states(self, session, dry_run, task_instances, now):
jobs = dict()
for task_instance in task_instances:
jobs[task_instance.push_state['jobId']] = task_instance
response = self.batch_client.describe_jobs(jobs=list(jobs.keys())) ## TODO: batch by 100
## TODO: tasks that go past 24 hour period ?
## TODO: timeout pushed tasks?
## TODO: tasks that are 'pushed' but not in AWS batch?
for job in response['jobs']:
if job['status'] in ['SUBMITTED','PENDING','RUNNABLE']:
status = 'pushed'
elif job['status'] in ['STARTING','RUNNING']:
status = 'running'
elif job['status'] == 'SUCCEEDED':
status = 'success'
elif job['status'] == 'FAILED':
status = 'failed'
task_instance = jobs[job['jobId']]
if task_instance.status != status:
task_instance.push_state = job
if status == 'failed':
task_instance.fail(session, self.taskflow, now=now, dry_run=dry_run)
else:
task_instance.status = status
if not dry_run:
session.commit()
def get_job_name(self, workflow, task, task_instance):
if workflow != None:
return '{}__{}__{}__{}'.format(
workflow.name,
task_instance.workflow_instance_id,
task.name,
task_instance.id)
else:
return '{}__{}'.format(
task.name,
task_instance.id)
def push_task_instances(self, session, dry_run, task_instances):
for task_instance in task_instances:
try:
task = self.taskflow.get_task(task_instance.task_name)
workflow = None
if task.workflow != None:
workflow = self.taskflow.get_workflow(task.workflow.name)
if task == None:
raise Exception('Task `{}` not found'.format(task_instance.task_name))
parameters = {
'task': task.name,
'task_instance': str(task_instance.id)
}
if workflow != None:
parameters['workflow'] = workflow.name
parameters['workflow_instance'] = str(task_instance.workflow_instance_id)
if 'job_queue' in task_instance.params and task_instance.params['job_queue']:
job_queue = task_instance.params['job_queue']
elif 'job_queue' in task.params and task.params['job_queue']:
job_queue = task.params['job_queue']
else:
job_queue = self.default_job_queue
if 'job_definition' in task_instance.params and task_instance.params['job_definition']:
job_definition = task_instance.params['job_definition']
elif 'job_definition' in task.params and task.params['job_definition']:
job_definition = task.params['job_definition']
else:
job_definition = self.default_job_definition
environment = [
{
'name': 'TASKFLOW_TASK',
'value': task.name
},
{
'name': 'TASKFLOW_TASK_INSTANCE_ID',
'value': str(task_instance.id)
}
]
if workflow != None:
environment.append({
'name': 'TASKFLOW_WORKFLOW',
'value': workflow.name
})
environment.append({
'name': 'TASKFLOW_WORKFLOW_INSTANCE_ID',
'value': str(task_instance.workflow_instance_id)
})
job_name = self.get_job_name(workflow, task, task_instance)
job_name = job_name[-128:] ## AWS Batch max job name is 128 characters
self.logger.info('Submitting job: %s %s %s', job_name, job_queue, job_definition)
if not dry_run:
response = self.batch_client.submit_job(
jobName=job_name,
jobQueue=job_queue,
jobDefinition=job_definition,
parameters=parameters,
containerOverrides={
'environment': environment
})
task_instance.state = 'pushed'
task_instance.push_state = response
if not dry_run:
session.commit()
except Exception:
self.logger.exception('Exception submitting %s %s', task_instance.task_name, task_instance.id)
session.rollback()