def defaultbackend():
with tempfile.NamedTemporaryFile() as f:
connection = create_engine(
"sqlite:///{path}".format(path=f.name),
connect_args={"check_same_thread": False},
poolclass=NullPool,
)
b = Storage(connection)
yield b
b.clear()
def backend():
fd, filepath = tempfile.mkstemp()
connection = create_engine(
"sqlite:///{path}".format(path=filepath),
connect_args={"check_same_thread": False},
poolclass=NullPool,
)
b = Storage(connection)
yield b
b.clear()
os.close(fd)
os.remove(filepath)
def __init__(self, queues, connection=None, num_workers=3):
# Internally, we use concurrent.future.Future to run and track
# job executions. We need to keep track of which future maps to which
# job they were made from, and we use the job_future_mapping dict to do
# so.
if connection is None:
raise ValueError("Connection must be defined")
# Queues that this worker executes tasks for
if type(queues) is not list:
queues = [queues]
self.queues = queues
# Key: future object, Value: job object
self.job_future_mapping = {}
# Key: job_id, Value: future object
self.future_job_mapping = {}
self.storage = Storage(connection)
self.num_workers = num_workers
self.workers = self.start_workers(num_workers=self.num_workers)
self.job_checker = self.start_job_checker()
def execute_job(job_id, db_type, db_url):
"""
Call the function stored in the job.func.
:return: Any
"""
from kolibri.core.tasks.main import make_connection
from kolibri.core.tasks.storage import Storage
connection = make_connection(db_type, db_url)
storage = Storage(connection)
job = storage.get_job(job_id)
setattr(current_state_tracker, "job", job)
func = import_stringified_func(job.func)
args, kwargs = copy.copy(job.args), copy.copy(job.kwargs)
try:
result = func(*args, **kwargs)
except Exception as e:
# If any error occurs, clear the job tracker and reraise
setattr(current_state_tracker, "job", None)
traceback_str = traceback.format_exc()
e.traceback = traceback_str
connection.dispose()
# Close any django connections opened here
django_connection.close()
raise
setattr(current_state_tracker, "job", None)
connection.dispose()
# Close any django connections opened here
django_connection.close()
return result
class Worker(object):
def __init__(self, queues, connection=None, num_workers=3):
# Internally, we use concurrent.future.Future to run and track
# job executions. We need to keep track of which future maps to which
# job they were made from, and we use the job_future_mapping dict to do
# so.
if connection is None:
raise ValueError("Connection must be defined")
# Queues that this worker executes tasks for
if type(queues) is not list:
queues = [queues]
self.queues = queues
# Key: future object, Value: job object
self.job_future_mapping = {}
# Key: job_id, Value: future object
self.future_job_mapping = {}
self.storage = Storage(connection)
self.num_workers = num_workers
self.workers = self.start_workers(num_workers=self.num_workers)
self.job_checker = self.start_job_checker()
def shutdown_workers(self, wait=True):
# First cancel all running jobs
for job_id in self.future_job_mapping:
self.cancel(job_id)
# Now shutdown the workers
self.workers.shutdown(wait=wait)
def start_workers(self, num_workers):
if MULTIPROCESS:
from concurrent.futures import ProcessPoolExecutor
worker_executor = ProcessPoolExecutor
else:
from concurrent.futures import ThreadPoolExecutor
worker_executor = ThreadPoolExecutor
pool = worker_executor(max_workers=num_workers)
return pool
def handle_finished_future(self, future):
# get back the job assigned to the future
job = self.job_future_mapping[future]
# Clean up tracking of this job and its future
del self.job_future_mapping[future]
del self.future_job_mapping[job.job_id]
try:
result = future.result()
except CancelledError as e:
self.report_cancelled(job.job_id)
return
except Exception as e:
self.report_error(job.job_id, e, e.traceback)
return
self.report_success(job.job_id, result)
def shutdown(self, wait=False):
self.job_checker.stop()
self.shutdown_workers(wait=wait)
if wait:
self.job_checker.join()
def start_job_checker(self):
"""
Starts up the job checker thread, that starts scheduled jobs when there are workers free,
and checks for cancellation requests for jobs currently assigned to a worker.
Returns: the Thread object.
"""
t = InfiniteLoopThread(self.check_jobs,
thread_name="JOBCHECKER",
wait_between_runs=0.5)
t.start()
return t
def check_jobs(self):
"""
Check how many workers are currently running.
If fewer workers are running than there are available workers, start a new job!
Returns: None
"""
try:
while len(self.future_job_mapping) < self.num_workers:
self.start_next_job()
except Empty:
logger.debug("No jobs to start.")
for job in self.storage.get_canceling_jobs(self.queues):
job_id = job.job_id
if job_id in self.future_job_mapping:
self.cancel(job_id)
else:
self.report_cancelled(job_id)
def report_cancelled(self, job_id):
self.storage.mark_job_as_canceled(job_id)
def report_success(self, job_id, result):
self.storage.complete_job(job_id)
def report_error(self, job_id, exc, trace):
trace = traceback.format_exc()
logger.error("Job {} raised an exception: {}".format(job_id, trace))
self.storage.mark_job_as_failed(job_id, exc, trace)
def update_progress(self, job_id, progress, total_progress, stage=""):
self.storage.update_job_progress(job_id, progress, total_progress)
def start_next_job(self):
"""
start the next scheduled job to the type of workers spawned by self.start_workers.
:return future:
"""
job = self.storage.get_next_queued_job(self.queues)
if not job:
raise Empty
self.storage.mark_job_as_running(job.job_id)
lambda_to_execute = _reraise_with_traceback(
job.get_lambda_to_execute())
future = self.workers.submit(
lambda_to_execute,
update_progress_func=self.update_progress,
cancel_job_func=self._check_for_cancel,
save_job_meta_func=self.storage.save_job_meta,
)
# assign the futures to a dict, mapping them to a job
self.job_future_mapping[future] = job
self.future_job_mapping[job.job_id] = future
# callback for when the future is now!
future.add_done_callback(self.handle_finished_future)
return future
def cancel(self, job_id):
"""
Request a cancellation from the futures executor pool.
If that didn't work (because it's already running), then mark
a special variable inside the future that we can check
inside a special check_for_cancel function passed to the
job.
:param job_id:
:return:
"""
future = self.future_job_mapping[job_id]
is_future_cancelled = future.cancel()
if is_future_cancelled: # success!
return True
else:
if future.running():
# Already running, but let's mark the future as cancelled
# anyway, to make sure that calling future.result() will raise an error.
# Our cancelling callback will then check this variable to see its state,
# and exit if it's cancelled.
from concurrent.futures._base import CANCELLED
future._state = CANCELLED
return False
else: # probably finished already, too late to cancel!
return False
def _check_for_cancel(self, job_id):
"""
Check if a job has been requested to be cancelled. When called, the calling function can
optionally give the stage it is currently in, so the user has information on where the job
was before it was cancelled.
:param job_id: The job_id to check
:param current_stage: Where the job currently is
:return: raises a UserCancelledError if we find out that we were cancelled.
"""
future = self.future_job_mapping[job_id]
is_cancelled = future._state in [CANCELLED, CANCELLED_AND_NOTIFIED]
if is_cancelled:
raise UserCancelledError()
def __init__(self, queue=DEFAULT_QUEUE, connection=None):
if connection is None:
raise ValueError("Connection must be defined")
self.name = queue
self.storage = Storage(connection)
class Queue(object):
def __init__(self, queue=DEFAULT_QUEUE, connection=None):
if connection is None:
raise ValueError("Connection must be defined")
self.name = queue
self.storage = Storage(connection)
def __len__(self):
return self.storage.count_all_jobs(self.name)
@property
def job_ids(self):
return [job.job_id for job in self.storage.get_all_jobs(self.name)]
@property
def jobs(self):
"""
Return all the jobs scheduled, queued, running, failed or completed.
Returns: A list of all jobs.
"""
return self.storage.get_all_jobs(self.name)
def enqueue(self, func, *args, **kwargs):
"""
Enqueues a function func for execution.
One special parameter is track_progress. If passed in and not None, the func will be passed in a
keyword parameter called update_progress:
def update_progress(progress, total_progress, stage=""):
The running function can call the update_progress function to notify interested parties of the function's
current progress.
Another special parameter is the "cancellable" keyword parameter. When passed in and not None, a special
"check_for_cancel" parameter is passed in. When called, it raises an error when the user has requested a job
to be cancelled.
The caller can also pass in any pickleable object into the "extra_metadata" parameter. This data is stored
within the job and can be retrieved when the job status is queried.
All other parameters are directly passed to the function when it starts running.
:type func: callable or str
:param func: A callable object that will be scheduled for running.
:return: a string representing the job_id.
"""
# if the func is already a job object, just schedule that directly.
if isinstance(func, Job):
job = func
# else, turn it into a job first.
else:
job = Job(func, *args, **kwargs)
job_id = self.storage.enqueue_job(job, self.name)
return job_id
def restart_job(self, job_id):
"""
Given a job_id, restart the job for that id. A job will only be restarted if
in CANCELED or FAILED state.
This first clears the job then creates a new job with the same job_id as
the cleared one.
"""
old_job = self.fetch_job(job_id)
if old_job.state in [State.CANCELED, State.FAILED]:
self.clear_job(job_id)
job = Job(old_job)
job.job_id = job_id
return self.enqueue(job)
else:
raise JobNotRestartable("Cannot restart job with state={}".format(
old_job.state))
def cancel(self, job_id):
"""
Mark a job as canceling, and let the worker pick this up to initiate
the cancel of the job.
:param job_id: the job_id of the Job to cancel.
"""
self.storage.mark_job_as_canceling(job_id)
def fetch_job(self, job_id):
"""
Returns a Job object corresponding to the job_id. From there, you can query for the following attributes:
- function string to run
- its current state (see Job.State for the list of states)
- progress (returning an int), total_progress (returning an int), and percentage_progress
(derived from running job.progress/total_progress)
- the job.exception and job.traceback, if the job's function returned an error
:param job_id: the job_id to get the Job object for
:return: the Job object corresponding to the job_id
"""
return self.storage.get_job(job_id)
def empty(self):
"""
Clear all jobs.
"""
self.storage.clear(force=True, queue=self.name)
def clear(self):
"""
Clear all succeeded, failed, or cancelled jobs.
"""
self.storage.clear(force=False, queue=self.name)
def clear_job(self, job_id):
"""
Clear a job if it has succeeded, failed, or been cancelled.
:type job_id: str
:param job_id: id of job to clear.
"""
self.storage.clear(job_id=job_id, force=False)
def backend():
with connection() as c:
b = Storage(c)
b.clear(force=True)
yield b
b.clear(force=True)
class Worker(object):
def __init__(self, queues, connection=None, num_workers=3):
# Internally, we use concurrent.future.Future to run and track
# job executions. We need to keep track of which future maps to which
# job they were made from, and we use the job_future_mapping dict to do
# so.
if connection is None:
raise ValueError("Connection must be defined")
# Queues that this worker executes tasks for
if type(queues) is not list:
queues = [queues]
self.queues = queues
# Key: future object, Value: job object
self.job_future_mapping = {}
# Key: job_id, Value: future object
self.future_job_mapping = {}
self.storage = Storage(connection)
self.num_workers = num_workers
self.workers = self.start_workers(num_workers=self.num_workers)
self.job_checker = self.start_job_checker()
def shutdown_workers(self, wait=True):
# First cancel all running jobs
# Coerce to a list, as otherwise the iterable can change size
# during iteration, as jobs are cancelled and removed from the mapping
job_ids = list(self.future_job_mapping.keys())
for job_id in job_ids:
logger.info("Canceling job id {}.".format(job_id))
self.cancel(job_id)
# Now shutdown the workers
self.workers.shutdown(wait=wait)
def start_workers(self, num_workers):
pool = PoolExecutor(max_workers=num_workers)
return pool
def handle_finished_future(self, future):
# get back the job assigned to the future
job = self.job_future_mapping[future]
# Clean up tracking of this job and its future
del self.job_future_mapping[future]
del self.future_job_mapping[job.job_id]
try:
result = future.result()
except CancelledError:
self.report_cancelled(job.job_id)
return
except Exception as e:
if hasattr(e, "traceback"):
traceback = e.traceback
else:
traceback = ""
self.report_error(job.job_id, e, traceback)
return
self.report_success(job.job_id, result)
def shutdown(self, wait=True):
logger.info("Asking job schedulers to shut down.")
self.job_checker.stop()
self.shutdown_workers(wait=wait)
if wait:
self.job_checker.join()
def start_job_checker(self):
"""
Starts up the job checker thread, that starts scheduled jobs when there are workers free,
and checks for cancellation requests for jobs currently assigned to a worker.
Returns: the Thread object.
"""
t = InfiniteLoopThread(
self.check_jobs, thread_name="JOBCHECKER", wait_between_runs=0.5
)
t.start()
return t
def check_jobs(self):
"""
Check how many workers are currently running.
If fewer workers are running than there are available workers, start a new job!
Returns: None
"""
try:
while len(self.future_job_mapping) < self.num_workers:
self.start_next_job()
except Empty:
logger.debug("No jobs to start.")
for job in self.storage.get_canceling_jobs(self.queues):
job_id = job.job_id
if job_id in self.future_job_mapping:
self.cancel(job_id)
else:
self.report_cancelled(job_id)
def report_cancelled(self, job_id):
self.storage.mark_job_as_canceled(job_id)
def report_success(self, job_id, result):
self.storage.complete_job(job_id, result=result)
def report_error(self, job_id, exc, trace):
trace = traceback.format_exc()
logger.error("Job {} raised an exception: {}".format(job_id, trace))
self.storage.mark_job_as_failed(job_id, exc, trace)
def update_progress(self, job_id, progress, total_progress, stage=""):
self.storage.update_job_progress(job_id, progress, total_progress)
def start_next_job(self):
"""
start the next scheduled job to the type of workers spawned by self.start_workers.
:return future:
"""
job = self.storage.get_next_queued_job(self.queues)
if not job:
raise Empty
self.storage.mark_job_as_running(job.job_id)
db_type_lookup = {
"sqlite": "sqlite",
"postgresql": "postgres",
}
db_type = db_type_lookup[self.storage.engine.dialect.name]
future = self.workers.submit(
execute_job,
job_id=job.job_id,
db_type=db_type,
db_url=self.storage.engine.url,
)
# assign the futures to a dict, mapping them to a job
self.job_future_mapping[future] = job
self.future_job_mapping[job.job_id] = future
# callback for when the future is now!
future.add_done_callback(self.handle_finished_future)
return future
def cancel(self, job_id):
"""
Request a cancellation from the futures executor pool.
If that didn't work (because it's already running), then mark
a special variable inside the future that we can check
inside a special check_for_cancel function passed to the
job.
:param job_id:
:return:
"""
try:
future = self.future_job_mapping[job_id]
is_future_cancelled = future.cancel()
except KeyError:
# In the case that the future does not even exist, say it has been cancelled.
is_future_cancelled = True
if is_future_cancelled: # success!
return True
else:
if future.running():
# Already running, so we manually mark the future as cancelled
setattr(future, "_is_cancelled", True)
return False
else: # probably finished already, too late to cancel!
return False