blob: d08b8fcc07857e7ccfa5693201a896a72e58d032 [file] [log] [blame]
#!/usr/bin/env python3
# Copyright (C) 2019 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
''' Worker main loop. Pulls jobs from the DB and runs them in the sandbox
It also handles timeouts and graceful container termination.
'''
import logging
import os
import random
import signal
import socket
import subprocess
import threading
import time
import traceback
from config import DB, JOB_TIMEOUT_SEC
from common_utils import req, utc_now_iso, init_logging
from common_utils import ConcurrentModificationError, SCOPES
CUR_DIR = os.path.dirname(__file__)
SCOPES.append('https://www.googleapis.com/auth/firebase.database')
SCOPES.append('https://www.googleapis.com/auth/userinfo.email')
WORKER_NAME = '%s-%s' % (os.getenv('WORKER_HOST', 'local').split('-')[-1],
socket.gethostname())
sigterm = threading.Event()
def try_acquire_job(job_id):
''' Transactionally acquire the given job.
Returns the job JSON object if it managed to acquire and put it into the
STARTED state, None if another worker got there first.
'''
logging.debug('Trying to acquire job %s', job_id)
uri = '%s/jobs/%s.json' % (DB, job_id)
job, etag = req('GET', uri, req_etag=True)
if job['status'] != 'QUEUED':
return None # Somebody else took it
try:
job['status'] = 'STARTED'
job['time_started'] = utc_now_iso()
job['worker'] = WORKER_NAME
req('PUT', uri, body=job, etag=etag)
return job
except ConcurrentModificationError:
return None
def make_worker_obj(status, job_id=None):
return {
'job_id': job_id,
'status': status,
'last_update': utc_now_iso(),
'host': os.getenv('WORKER_HOST', '')
}
def worker_loop():
''' Pulls a job from the queue and runs it invoking run_job.py '''
uri = '%s/jobs_queued.json?orderBy="$key"&limitToLast=10' % DB
jobs = req('GET', uri)
if not jobs:
return
# Work out the worker number from the hostname. We try to distribute the load
# (via the time.sleep below) so that we fill first all the worker-1 of each
# vm, then worker-2 and so on. This is designed so that if there is only one
# CL (hence N jobs) in the queue, each VM gets only one job, maximizing the
# cpu efficiency of each VM.
try:
worker_num = int(socket.gethostname().split('-')[-1])
except ValueError:
worker_num = 1
# Transactionally acquire a job. Deal with races (two workers trying to
# acquire the same job).
job = None
job_id = None
for job_id in sorted(jobs.keys(), reverse=True):
job = try_acquire_job(job_id)
if job is not None:
break
logging.info('Raced while trying to acquire job %s, retrying', job_id)
time.sleep(worker_num * 2 + random.random())
if job is None:
logging.error('Failed to acquire a job')
return
logging.info('Starting job %s', job_id)
# Update the db, move the job to the running queue.
patch_obj = {
'jobs_queued/' + job_id: {}, # = DELETE
'jobs_running/' + job_id: {
'worker': WORKER_NAME
},
'workers/' + WORKER_NAME: make_worker_obj('RUNNING', job_id=job_id)
}
req('PATCH', '%s.json' % DB, body=patch_obj)
cmd = [os.path.join(CUR_DIR, 'run_job.py'), job_id]
# Propagate the worker's PERFETTO_ vars and merge with the job-specific vars.
env = dict(os.environ, **{k: str(v) for (k, v) in job['env'].items()})
job_runner = subprocess.Popen(cmd, env=env)
# Run the job in a python subprocess, to isolate the main loop from logs
# uploader failures.
res = None
cancelled = False
timed_out = False
time_started = time.time()
time_last_db_poll = time_started
polled_status = 'STARTED'
while res is None:
time.sleep(0.25)
res = job_runner.poll()
now = time.time()
if now - time_last_db_poll > 10: # Throttle DB polling.
polled_status = req('GET', '%s/jobs/%s/status.json' % (DB, job_id))
time_last_db_poll = now
if now - time_started > JOB_TIMEOUT_SEC:
logging.info('Job %s timed out, terminating', job_id)
timed_out = True
job_runner.terminate()
if (sigterm.is_set() or polled_status != 'STARTED') and not cancelled:
logging.info('Job %s cancelled, terminating', job_id)
cancelled = True
job_runner.terminate()
status = ('INTERRUPTED' if sigterm.is_set() else 'CANCELLED' if cancelled else
'TIMED_OUT' if timed_out else 'COMPLETED' if res == 0 else 'FAILED')
logging.info('Job %s %s with code %s', job_id, status, res)
# Update the DB, unless the job has been cancelled. The "is not None"
# condition deals with a very niche case, that is, avoid creating a partial
# job entry after doing a full clear of the DB (which is super rare, happens
# only when re-deploying the CI).
if polled_status is not None:
patch = {
'jobs/%s/status' % job_id: status,
'jobs/%s/exit_code' % job_id: {} if res is None else res,
'jobs/%s/time_ended' % job_id: utc_now_iso(),
'jobs_running/%s' % job_id: {}, # = DELETE
}
req('PATCH', '%s.json' % (DB), body=patch)
def sig_handler(_, __):
logging.warning('Interrupted by signal, exiting worker')
sigterm.set()
def main():
init_logging()
logging.info('Worker started')
signal.signal(signal.SIGTERM, sig_handler)
signal.signal(signal.SIGINT, sig_handler)
while not sigterm.is_set():
logging.debug('Starting poll cycle')
try:
worker_loop()
req('PUT',
'%s/workers/%s.json' % (DB, WORKER_NAME),
body=make_worker_obj('IDLE'))
except:
logging.error('Exception in worker loop:\n%s', traceback.format_exc())
if sigterm.is_set():
break
# Synchronize sleeping with the wall clock. This is so all VMs wake up at
# the same time. See comment on distributing load above in this file.
poll_time_sec = 5
time.sleep(poll_time_sec - (time.time() % poll_time_sec))
# The use case here is the VM being terminated by the GCE infrastructure.
# We mark the worker as terminated and the job as cancelled so we don't wait
# forever for it.
logging.warning('Exiting the worker loop, got signal: %s', sigterm.is_set())
req('PUT',
'%s/workers/%s.json' % (DB, WORKER_NAME),
body=make_worker_obj('TERMINATED'))
if __name__ == '__main__':
main()