def abort(self):
"""ask the abortion of the task
The effective abortion will be done by the responsible monitor
and is inherently asynchronous. To wait synchronously for an
abortion one can do as follows:
.. code-block:: python
t = api.schedule(engine, 'mytask', input=42)
t.join('running')
t.abort()
t.join()
"""
with self.engine.begin() as cn:
# will still be marked as running
# the worker kill must do the actual job
update('rework.task').where(id=self.tid).values(
abort=True
).do(cn)
update(
'rework.worker as worker').table('rework.task as task'
).where(
'worker.id = task.worker',
'task.id = %(id)s', id=self.tid
).values(
kill=True
).do(cn)
def fromqueue(cls, engine, wid, domain='default'):
with engine.begin() as cn:
sql = ("select task.id, task.operation "
"from rework.task as task, rework.operation as op "
"where task.status = 'queued' "
" and op.domain = %(domain)s"
" and task.operation = op.id "
"order by task.id "
"for update of task skip locked "
"limit 1")
tid_operation = cn.execute(sql, domain=domain).fetchone()
if tid_operation is None:
return
tid, opid = tid_operation
update(
'rework.task'
).where(
id=tid
).values(
status='running',
worker=wid
).do(cn)
return cls(engine, tid, opid)
def run(self):
with self.engine.begin() as cn:
update(
'rework.task'
).where(
id=self.tid
).values(
started=utcnow()
).do(cn)
try:
name, path = self.engine.execute("""
select name, path
from rework.operation
where rework.operation.id = %(operation)s
""", {'operation': self.operation}
).fetchone()
mod = imp.load_source('module', path)
func = getattr(mod, name)
func(self)
except:
with self.engine.begin() as cn:
update(
'rework.task'
).where(
id=self.tid
).values(
traceback=tb.format_exc()
).do(cn)
finally:
self.finish()
def kill_worker(wid):
if not has_permission('kill'):
abort(403, 'You cannoy do that.')
with engine.begin() as cn:
update('rework.worker').where(id=wid).values(kill=True).do(cn)
return json.dumps(True)
def kill_worker(dburi, worker_id):
""" ask to preemptively kill a given worker to its monitor
If you want to not risk interrupting any ongoing work,
you should use `rework shutdown` instead.
"""
engine = create_engine(find_dburi(dburi))
with engine.begin() as cn:
update('rework.worker').where(id=worker_id).values(kill=True).do(cn)
def shutdown_worker(dburi, worker_id):
""" ask a worker to shut down as soon as it becomes idle
If you want to immediately and unconditionally terminate
a worker, use `rework kill-worker`
"""
engine = create_engine(find_dburi(dburi))
with engine.begin() as cn:
update('rework.worker').where(id=worker_id).values(
shutdown=True).do(cn)
def finish(self):
with self.engine.begin() as cn:
update(
'rework.task'
).where(
id=self.tid
).values(
finished=utcnow(),
status='done'
).do(cn)
def shrink_workers(self):
with self.engine.begin() as cn:
needed = self.queued_tasks(cn)
busy = self.busy_workers(cn)
idle = self.num_workers - len(busy)
# ask idle workers to shutdown
# let' not even try to do anything if
# the task queue is unempty
if not needed and idle > self.minworkers:
candidate = self.idle_worker(cn, busy)
# we now have a case to at least retire one
update('rework.worker').where(id=candidate).values(
shutdown=True).do(cn)
return candidate
def mark_dead_workers(cn, wids, message, traceback=None):
if not wids:
return
# mark workers as dead
update('rework.worker').where('id in %(ids)s', ids=tuple(wids)).values(
running=False, finished=utcnow(), deathinfo=message).do(cn)
# mark tasks as done
update('rework.task as task').table('rework.worker as worker').where(
"task.status != 'done'",
'worker.id = task.worker',
'worker.id in %(ids)s',
ids=tuple(wids),
).values(finished=utcnow(), status='done', abort=True,
traceback=traceback).do(cn)
def run(self):
if self.debugport:
pystuck.run_server(port=self.debugport)
try:
with self.running_status():
self.main_loop()
except Exception:
with self.engine.begin() as cn:
update('rework.worker').where(id=self.wid).values(
traceback=traceback.format_exc()).do(cn)
raise
except SystemExit as exit:
raise
def track_resources(self):
if not self.workers:
return
for wid, proc in self.workers.items():
q = update('rework.worker').where(id=wid).values(
mem=memory_usage(proc.pid), cpu=cpu_usage(proc.pid))
with self.engine.begin() as cn:
q.do(cn)
def test_worker_kill(engine):
with workers(engine) as mon:
wid = mon.wids[0]
with engine.begin() as cn:
update('rework.worker').where(id=wid).values(kill=True).do(cn)
guard(engine,
'select kill from rework.worker where id = {}'.format(wid),
lambda r: r.scalar() == True)
mon.preemptive_kill()
guard(engine,
'select count(id) from rework.worker where running = true',
lambda r: r.scalar() == 0)
assert engine.execute(
'select deathinfo from rework.worker where id = %(wid)s',
wid=wid).scalar().startswith('preemptive kill')
def save_output(self, data, raw=False):
"""saves the output of a task run for later use by other
components. This is to be used from *within* an operation.
The `data` parameter can be any picklable python object.
The `raw` parameter, combined with a bytestring as input data,
allows to bypass the pickling protocol.
"""
if not raw:
data = dumps(data, protocol=2)
with self.engine.begin() as cn:
update(
'rework.task'
).where(
id=self.tid
).values(
output=data
).do(cn)
def test_worker_shutdown(engine):
with workers(engine) as mon:
wid = mon.wids[0]
worker = Worker(engine.url, wid)
assert not worker.shutdown_asked()
with engine.begin() as cn:
update('rework.worker').where(id=wid).values(shutdown=True).do(cn)
assert worker.shutdown_asked()
guard(engine,
'select shutdown from rework.worker where id = {}'.format(wid),
lambda r: r.scalar() == True)
guard(engine,
'select count(id) from rework.worker where running = true',
lambda r: r.scalar() == 0)
assert u'explicit shutdown' == engine.execute(
'select deathinfo from rework.worker where id = %(wid)s',
wid=wid).scalar()
def dead_man_switch(self):
with self.engine.begin() as cn:
update('rework.monitor').where(id=self.monid).values(
lastseen=utcnow().astimezone(TZ)).do(cn)
def death_sql(self, cause):
return update('rework.worker').where(id=self.wid).values(
deathinfo=cause, running=False, finished=utcnow())
def running_sql(self, running):
value = {'running': running, 'debugport': self.debugport or None}
if running:
value['pid'] = os.getpid()
value['started'] = utcnow()
return update('rework.worker').where(id=self.wid).values(**value)
def ask_shutdown(self):
with self.engine.begin() as cn:
update('rework.worker').where(id=self.wid).values(
shutdown=True).do(cn)