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