backend.lib.manager

  1"""
  2The heart of the app - manages jobs and workers
  3"""
  4import threading
  5import signal
  6import time
  7
  8from collections.abc import Generator
  9
 10from backend.lib.proxied_requests import DelegatedRequestHandler
 11from backend.lib.worker import BasicWorker
 12from common.lib.exceptions import JobClaimedException
 13
 14# for now, this is hardcoded - could be dynamic or depending on the queue ID in
 15# the future
 16MAX_JOBS_PER_QUEUE = 1
 17
 18class WorkerManager:
 19	"""
 20	Manages the job queue and worker pool
 21	"""
 22	queue = None
 23	db = None
 24	log = None
 25	modules = None
 26	proxy_delegator = None
 27
 28	worker_pool = {}
 29	job_mapping = {}
 30	pool = []
 31	looping = True
 32	unknown_jobs = set()
 33
 34	def __init__(self, queue, database, logger, modules, as_daemon=True):
 35		"""
 36		Initialize manager
 37
 38		:param queue:  Job queue
 39		:param database:  Database handler
 40		:param logger:  Logger object
 41		:param modules:  Modules cache via ModuleLoader()
 42		:param bool as_daemon:  Whether the manager is being run as a daemon
 43		"""
 44		self.queue = queue
 45		self.db = database
 46		self.log = logger
 47		self.modules = modules
 48		self.proxy_delegator = DelegatedRequestHandler(self.log, self.modules.config)
 49
 50		if as_daemon:
 51			signal.signal(signal.SIGTERM, self.abort)
 52
 53		# datasources are initialized here; the init_datasource functions found in their respective __init__.py files
 54		# are called which, in the case of scrapers, also adds the scrape jobs to the queue.
 55		self.validate_datasources()
 56
 57		# queue jobs for workers that always need one
 58		for worker_name, worker in self.modules.workers.items():
 59			if hasattr(worker, "ensure_job"):
 60				# ensure_job is a class method that returns a dict with job parameters if job should be added
 61				# pass config for some workers (e.g., web studies extensions)
 62				try:
 63					self.log.debug(f"Ensuring job exists for worker {worker_name}")
 64					job_params = worker.ensure_job(config=self.modules.config)
 65				except Exception as e:
 66					self.log.error(f"Error while ensuring job for worker {worker_name}: {e}")
 67					job_params = None
 68
 69				if job_params:
 70					self.queue.add_job(worker_or_type=worker, **job_params)
 71
 72		self.ident = threading.get_ident()
 73		self.log.info("4CAT Started")
 74
 75		# flush module collector log buffer
 76		# the logger is not available when this initialises
 77		# but it is now!
 78		if self.modules.log_buffer:
 79			self.log.warning(self.modules.log_buffer)
 80			self.modules.log_buffer = ""
 81
 82		# it's time
 83		self.loop()
 84
 85	def delegate(self):
 86		"""
 87		Delegate work
 88
 89		Checks for open jobs, and then passes those to dedicated workers, if
 90		slots are available for those workers.
 91		"""
 92		jobs = self.queue.get_all_jobs()
 93
 94		num_active = len(list(self.iterate_active_workers()))
 95		self.log.debug2(f"Running {num_active} active workers")
 96
 97		# clean up workers that have finished processing
 98		# not using iterate_active_workers() here because we're going to change
 99		# the dictionary while iterating through it
100		for queue_id in self.worker_pool:
101			all_workers = self.worker_pool[queue_id]
102			for worker in all_workers:
103				if not worker.is_alive():
104					self.log.debug(f"Terminating worker {worker.job.data['jobtype']}/{worker.job.data['remote_id']}")
105					worker.join()
106					self.worker_pool[queue_id].remove(worker)
107
108			del all_workers
109
110		# check if workers are available for unclaimed jobs
111		for job in jobs:
112			queue_id = job.data["queue_id"]
113			jobtype = job.data["jobtype"]
114
115			if jobtype in self.modules.workers:
116				worker_class = self.modules.workers[jobtype]
117				if queue_id not in self.worker_pool:
118					self.worker_pool[queue_id] = []
119
120				# if a job is of a known type, and that job type has open
121				# worker slots, start a new worker to run it
122				if len(self.worker_pool[queue_id]) < MAX_JOBS_PER_QUEUE:
123					try:
124						job.claim()
125						worker = worker_class(logger=self.log, manager=self, job=job, modules=self.modules)
126						worker.start()
127						log_level = self.log.levels["DEBUG"] if job.data["interval"] else self.log.levels["INFO"]
128						self.log.log(f"Starting new worker for job {job.data['jobtype']}/{job.data['remote_id']}", log_level)
129						self.worker_pool[queue_id].append(worker)
130					except JobClaimedException:
131						# it's fine
132						pass
133			else:
134				if jobtype not in self.unknown_jobs:
135					self.log.error(f"Unknown job type: {jobtype}")
136					self.unknown_jobs.add(jobtype)
137
138		time.sleep(1)
139
140	def loop(self):
141		"""
142		Main loop
143
144		Constantly delegates work, until no longer looping, after which all
145		workers are asked to stop their work. Once that has happened, the loop
146		properly ends.
147		"""
148		while self.looping:
149			try:
150				self.delegate()
151			except KeyboardInterrupt:
152				self.looping = False
153
154		self.log.info("Telling all workers to stop doing whatever they're doing...")
155
156		# request shutdown from all workers except the API
157		# this allows us to use the API to figure out if a certain worker is
158		# hanging during shutdown, for example
159		for queue_id, worker in self.iterate_active_workers():
160			if worker.type == "api":
161				continue
162
163			if hasattr(worker, "request_interrupt"):
164				worker.request_interrupt()
165			else:
166				worker.abort()
167
168		# wait for all workers that we just asked to quit to finish
169		self.log.info("Waiting for all workers to finish...")
170		for queue_id, worker in self.iterate_active_workers():
171			if worker.type == "api":
172				continue
173
174			self.log.info(f"Waiting for worker of type {worker.type}...")
175			worker.join()
176
177		# shut down any remaining workers (i.e. the API)
178		for queue_id, worker in self.iterate_active_workers():
179			worker.request_interrupt()
180			worker.join()
181
182		# abort
183		time.sleep(1)
184		self.log.info("Bye!")
185
186	def validate_datasources(self):
187		"""
188		Validate data sources
189
190		Logs warnings if not all information is present for the configured data
191		sources.
192		"""
193		for datasource in self.modules.datasources:
194			if datasource + "-search" not in self.modules.workers and datasource + "-import" not in self.modules.workers:
195				self.log.error(f"No search worker defined for data source {datasource} or its modules are missing. "
196				               f"Datasets cannot be created for it.")
197
198			self.modules.datasources[datasource]["init"](self.db, self.log, self.queue, datasource, self.modules.config)
199
200	def abort(self, signal=None, stack=None):
201		"""
202		Stop looping the delegator, clean up, and prepare for shutdown
203		"""
204		self.log.info("Received SIGTERM")
205
206		# cancel all interruptible postgres queries
207		# this needs to be done before we stop looping since after that no new
208		# jobs will be claimed, and needs to be done here because the worker's
209		# own database connection is busy executing the query that it should
210		# cancel! so we can't use it to update the job and make it get claimed
211		for job in self.queue.get_all_jobs("cancel-pg-query", restrict_claimable=False):
212			# this will make all these jobs immediately claimable, i.e. queries
213			# will get cancelled asap
214			self.log.debug("Cancelling interruptable Postgres queries for connection %s..." % job.data["remote_id"])
215			job.claim()
216			job.release(delay=0, claim_after=0)
217
218		# wait until all cancel jobs are completed
219		while self.queue.get_all_jobs("cancel-pg-query", restrict_claimable=False):
220			time.sleep(0.25)
221
222		# now stop looping (i.e. accepting new jobs)
223		self.looping = False
224
225	def request_interrupt(self, interrupt_level, job):
226		"""
227		Interrupt a specific job
228
229		This method can be called via e.g. the API, to interrupt a specific
230		job's worker. The worker for the given Job object is searched for and
231		if it exists, its `request_interrupt()` method is called.
232
233		:param int interrupt_level:  Retry later or cancel?
234		:param Job job:  Job object to cancel worker for
235		"""
236		for queue_id, worker in self.iterate_active_workers():
237			if worker.job.data["id"] == job.data["id"]:
238				# first cancel any interruptable postgres queries for this job's worker
239				while True:
240					active_queries = self.queue.get_all_jobs("cancel-pg-query", remote_id=worker.db.appname, restrict_claimable=False)
241					if not active_queries:
242						# all cancellation jobs have been run
243						break
244
245					for cancel_job in active_queries:
246						if cancel_job.is_claimed:
247							continue
248
249						# this will cause the job be run asap
250						cancel_job.claim()
251						cancel_job.release(delay=0, claim_after=0)
252
253					# give the cancel job a moment to run
254					time.sleep(0.25)
255
256				# now all queries are interrupted, formally request an abort
257				self.log.info(f"Requesting interrupt of job {worker.job.data['id']} ({worker.job.data['jobtype']}/{worker.job.data['remote_id']})")
258				worker.request_interrupt(interrupt_level)
259				return
260
261	def iterate_active_workers(self) -> Generator[tuple[str, BasicWorker]]:
262		"""
263		Return all active workers
264
265		Convenience function to avoid having to always use two nested for
266		loops.
267
268		:return:  Generator that yields tuples of (queue_id, worker)
269		"""
270		for queue_id in self.worker_pool:
271			for worker in self.worker_pool[queue_id]:
272				yield queue_id, worker