This is an automated email from the ASF dual-hosted git repository. eamonford pushed a commit to branch rabbitmq-fix in repository https://gitbox.apache.org/repos/asf/incubator-sdap-ingester.git
commit 0eed399f12fde9257ceac83b28fe991af5d9eafe Author: Eamon Ford <[email protected]> AuthorDate: Wed Jul 8 20:16:29 2020 -0500 better error handling better exception handling error handling the healthchecks now raise exceptions if they rail propagate child worker exceptions up to main process exc handling error handling use pysolr solr history bug fixes use asyncio in the collection ingester --- granule_ingester/conda-requirements.txt | 2 +- granule_ingester/docker/entrypoint.sh | 3 +- .../granule_ingester/consumer/Consumer.py | 40 +++++-- granule_ingester/granule_ingester/main.py | 45 +++++--- .../granule_ingester/pipeline/Pipeline.py | 119 ++++++++++++++------- .../reading_processors/TileReadingProcessor.py | 20 ++-- .../granule_ingester/writers/DataStore.py | 1 + .../granule_ingester/writers/SolrStore.py | 79 +++++++++----- granule_ingester/requirements.txt | 5 +- granule_ingester/tests/pipeline/test_Pipeline.py | 9 +- 10 files changed, 213 insertions(+), 110 deletions(-) diff --git a/granule_ingester/conda-requirements.txt b/granule_ingester/conda-requirements.txt index b2af149..fafd6f3 100644 --- a/granule_ingester/conda-requirements.txt +++ b/granule_ingester/conda-requirements.txt @@ -6,5 +6,5 @@ xarray pyyaml==5.3.1 requests==2.23.0 aiohttp==3.6.2 -aio-pika +aio-pika==6.6.1 tenacity diff --git a/granule_ingester/docker/entrypoint.sh b/granule_ingester/docker/entrypoint.sh index b703ee3..2b6174a 100644 --- a/granule_ingester/docker/entrypoint.sh +++ b/granule_ingester/docker/entrypoint.sh @@ -9,4 +9,5 @@ python /sdap/granule_ingester/main.py \ $([[ ! -z "$CASSANDRA_PORT" ]] && echo --cassandra_port=$CASSANDRA_PORT) \ $([[ ! -z "$CASSANDRA_USERNAME" ]] && echo --cassandra_username=$CASSANDRA_USERNAME) \ $([[ ! -z "$CASSANDRA_PASSWORD" ]] && echo --cassandra_password=$CASSANDRA_PASSWORD) \ - $([[ ! -z "$SOLR_HOST_AND_PORT" ]] && echo --solr_host_and_port=$SOLR_HOST_AND_PORT) + $([[ ! -z "$SOLR_HOST_AND_PORT" ]] && echo --solr_host_and_port=$SOLR_HOST_AND_PORT) \ + $([[ ! -z "$ZK_HOST_AND_PORT" ]] && echo --zk_host_and_port=$ZK_HOST_AND_PORT) diff --git a/granule_ingester/granule_ingester/consumer/Consumer.py b/granule_ingester/granule_ingester/consumer/Consumer.py index 75d347a..d5f1d97 100644 --- a/granule_ingester/granule_ingester/consumer/Consumer.py +++ b/granule_ingester/granule_ingester/consumer/Consumer.py @@ -17,6 +17,8 @@ import logging import aio_pika +from granule_ingester.exceptions import PipelineBuildingError, PipelineRunningError, RabbitMQLostConnectionError, \ + RabbitMQFailedHealthCheckError, LostConnectionError from granule_ingester.healthcheck import HealthCheck from granule_ingester.pipeline import Pipeline @@ -39,7 +41,7 @@ class Consumer(HealthCheck): self._connection_string = "amqp://{username}:{password}@{host}/".format(username=rabbitmq_username, password=rabbitmq_password, host=rabbitmq_host) - self._connection = None + self._connection: aio_pika.Connection = None async def health_check(self) -> bool: try: @@ -47,10 +49,10 @@ class Consumer(HealthCheck): await connection.close() return True except: - logger.error("Cannot connect to RabbitMQ! Connection string was {}".format(self._connection_string)) - return False + raise RabbitMQFailedHealthCheckError(f"Cannot connect to RabbitMQ! " + f"Connection string was {self._connection_string}") - async def _get_connection(self): + async def _get_connection(self) -> aio_pika.Connection: return await aio_pika.connect_robust(self._connection_string) async def __aenter__(self): @@ -73,16 +75,34 @@ class Consumer(HealthCheck): data_store_factory=data_store_factory, metadata_store_factory=metadata_store_factory) await pipeline.run() - message.ack() + await message.ack() + except PipelineBuildingError as e: + await message.reject() + logger.exception(f"Failed to build the granule-processing pipeline. This message will be dropped " + f"from RabbitMQ. The exception was:\n{e}") + except PipelineRunningError as e: + await message.reject() + logger.exception(f"Processing the granule failed. It will not be retried. The exception was:\n{e}") + except LostConnectionError: + # Let main() handle this + raise except Exception as e: - message.reject(requeue=True) - logger.error("Processing message failed. Message will be re-queued. The exception was:\n{}".format(e)) + await message.reject(requeue=True) + logger.exception(f"Processing message failed. Message will be re-queued. The exception was:\n{e}") async def start_consuming(self): channel = await self._connection.channel() await channel.set_qos(prefetch_count=1) queue = await channel.declare_queue(self._rabbitmq_queue, durable=True) - - async with queue.iterator() as queue_iter: - async for message in queue_iter: + queue_iter = queue.iterator() + async for message in queue_iter: + try: await self._received_message(message, self._data_store_factory, self._metadata_store_factory) + except aio_pika.exceptions.MessageProcessError: + # Do not try to close() the queue iterator! If we get here, that means the RabbitMQ + # connection has died, and attempting to close the queue will only raise another exception. + raise RabbitMQLostConnectionError("Lost connection to RabbitMQ while processing a granule.") + except Exception as e: + await queue_iter.close() + await channel.close() + raise e diff --git a/granule_ingester/granule_ingester/main.py b/granule_ingester/granule_ingester/main.py index 9010e33..fd5a6b9 100644 --- a/granule_ingester/granule_ingester/main.py +++ b/granule_ingester/granule_ingester/main.py @@ -16,13 +16,14 @@ import argparse import asyncio import logging +import sys from functools import partial from typing import List from granule_ingester.consumer import Consumer +from granule_ingester.exceptions import FailedHealthCheckError, LostConnectionError from granule_ingester.healthcheck import HealthCheck -from granule_ingester.writers import CassandraStore -from granule_ingester.writers import SolrStore +from granule_ingester.writers import CassandraStore, SolrStore def cassandra_factory(contact_points, port, username, password): @@ -31,8 +32,8 @@ def cassandra_factory(contact_points, port, username, password): return store -def solr_factory(solr_host_and_port): - store = SolrStore(solr_host_and_port) +def solr_factory(solr_host_and_port, zk_host_and_port): + store = SolrStore(zk_url=zk_host_and_port) if zk_host_and_port else SolrStore(solr_url=solr_host_and_port) store.connect() return store @@ -44,7 +45,7 @@ async def run_health_checks(dependencies: List[HealthCheck]): return True -async def main(): +async def main(loop): parser = argparse.ArgumentParser(description='Listen to RabbitMQ for granule ingestion instructions, and process ' 'and ingest a granule for each message that comes through.') parser.add_argument('--rabbitmq_host', @@ -84,6 +85,8 @@ async def main(): default='http://localhost:8983', metavar='HOST:PORT', help='Solr host and port. (Default: http://localhost:8983)') + parser.add_argument('--zk_host_and_port', + metavar="HOST:PORT") parser.add_argument('-v', '--verbose', action='store_true', @@ -106,7 +109,10 @@ async def main(): cassandra_password = args.cassandra_password cassandra_contact_points = args.cassandra_contact_points cassandra_port = args.cassandra_port + cassandra_username = args.cassandra_username + cassandra_password = args.cassandra_password solr_host_and_port = args.solr_host_and_port + zk_host_and_port = args.zk_host_and_port consumer = Consumer(rabbitmq_host=args.rabbitmq_host, rabbitmq_username=args.rabbitmq_username, @@ -117,19 +123,28 @@ async def main(): cassandra_port, cassandra_username, cassandra_password), - metadata_store_factory=partial(solr_factory, solr_host_and_port)) - if await run_health_checks([CassandraStore(cassandra_contact_points, - cassandra_port, - cassandra_username, - cassandra_password), - SolrStore(solr_host_and_port), - consumer]): + metadata_store_factory=partial(solr_factory, solr_host_and_port, zk_host_and_port)) + try: + solr_store = SolrStore(zk_url=zk_host_and_port) if zk_host_and_port else SolrStore(solr_url=solr_host_and_port) + await run_health_checks([CassandraStore(cassandra_contact_points, + cassandra_port, + cassandra_username, + cassandra_password), + solr_store, + consumer]) async with consumer: logger.info("All external dependencies have passed the health checks. Now listening to message queue.") await consumer.start_consuming() - else: - logger.error("Quitting because not all dependencies passed the health checks.") + except FailedHealthCheckError as e: + logger.error(f"Quitting because not all dependencies passed the health checks: {e}") + except LostConnectionError as e: + logger.error(f"{e} Any messages that were being processed have been re-queued. Quitting.") + except Exception as e: + logger.exception(f"Shutting down because of an unrecoverable error:\n{e}") + finally: + sys.exit(1) if __name__ == '__main__': - asyncio.run(main()) + loop = asyncio.get_event_loop() + loop.run_until_complete(main(loop)) diff --git a/granule_ingester/granule_ingester/pipeline/Pipeline.py b/granule_ingester/granule_ingester/pipeline/Pipeline.py index 8f2dd6f..2181da2 100644 --- a/granule_ingester/granule_ingester/pipeline/Pipeline.py +++ b/granule_ingester/granule_ingester/pipeline/Pipeline.py @@ -13,16 +13,22 @@ # See the License for the specific language governing permissions and # limitations under the License. - import logging +import os +import pickle import time +from multiprocessing import Manager from typing import List -import aiomultiprocess import xarray as xr import yaml +from aiomultiprocess import Pool +from aiomultiprocess.types import ProxyException from nexusproto import DataTile_pb2 as nexusproto +from tblib import pickling_support +from yaml.scanner import ScannerError +from granule_ingester.exceptions import PipelineBuildingError from granule_ingester.granule_loaders import GranuleLoader from granule_ingester.pipeline.Modules import modules as processor_module_mappings from granule_ingester.processors.TileProcessor import TileProcessor @@ -31,19 +37,23 @@ from granule_ingester.writers import DataStore, MetadataStore logger = logging.getLogger(__name__) -MAX_QUEUE_SIZE = 2 ** 15 - 1 +# The aiomultiprocessing library has a bug where it never closes out the pool if there are more than a certain +# number of items to process. The exact number is unknown, but 2**8-1 is safe. +MAX_CHUNK_SIZE = 2 ** 8 - 1 _worker_data_store: DataStore = None _worker_metadata_store: MetadataStore = None _worker_processor_list: List[TileProcessor] = None _worker_dataset = None +_shared_memory = None -def _init_worker(processor_list, dataset, data_store_factory, metadata_store_factory): +def _init_worker(processor_list, dataset, data_store_factory, metadata_store_factory, shared_memory): global _worker_data_store global _worker_metadata_store global _worker_processor_list global _worker_dataset + global _shared_memory # _worker_data_store and _worker_metadata_store open multiple TCP sockets from each worker process; # however, these sockets will be automatically closed by the OS once the worker processes die so no need to worry. @@ -51,19 +61,21 @@ def _init_worker(processor_list, dataset, data_store_factory, metadata_store_fac _worker_metadata_store = metadata_store_factory() _worker_processor_list = processor_list _worker_dataset = dataset + _shared_memory = shared_memory async def _process_tile_in_worker(serialized_input_tile: str): - global _worker_data_store - global _worker_metadata_store - global _worker_processor_list - global _worker_dataset + try: + input_tile = nexusproto.NexusTile.FromString(serialized_input_tile) + processed_tile = _recurse(_worker_processor_list, _worker_dataset, input_tile) - input_tile = nexusproto.NexusTile.FromString(serialized_input_tile) - processed_tile = _recurse(_worker_processor_list, _worker_dataset, input_tile) - if processed_tile: - await _worker_data_store.save_data(processed_tile) - await _worker_metadata_store.save_metadata(processed_tile) + if processed_tile: + await _worker_data_store.save_data(processed_tile) + await _worker_metadata_store.save_metadata(processed_tile) + except Exception as e: + pickling_support.install(e) + _shared_memory.error = pickle.dumps(e) + raise def _recurse(processor_list: List[TileProcessor], @@ -88,40 +100,56 @@ class Pipeline: self._data_store_factory = data_store_factory self._metadata_store_factory = metadata_store_factory - @classmethod - def from_string(cls, config_str: str, data_store_factory, metadata_store_factory): - config = yaml.load(config_str, yaml.FullLoader) - return cls._build_pipeline(config, - data_store_factory, - metadata_store_factory, - processor_module_mappings) + self._max_concurrency = int(os.getenv('MAX_CONCURRENCY', 16)) + + # Create a SyncManager so that we can to communicate exceptions from the + # worker processes back to the main process. + self._manager = Manager() + + def __del__(self): + self._manager.shutdown() @classmethod - def from_file(cls, config_path: str, data_store_factory, metadata_store_factory): - with open(config_path) as config_file: - config = yaml.load(config_file, yaml.FullLoader) + def from_string(cls, config_str: str, data_store_factory, metadata_store_factory): + try: + config = yaml.load(config_str, yaml.FullLoader) + cls._validate_config(config) return cls._build_pipeline(config, data_store_factory, metadata_store_factory, processor_module_mappings) + except yaml.scanner.ScannerError: + raise PipelineBuildingError("Cannot build pipeline because of a syntax error in the YAML.") + + # TODO: this method should validate the config against an actual schema definition + @staticmethod + def _validate_config(config: dict): + if type(config) is not dict: + raise PipelineBuildingError("Cannot build pipeline because the config is not valid YAML.") + @classmethod def _build_pipeline(cls, config: dict, data_store_factory, metadata_store_factory, module_mappings: dict): - granule_loader = GranuleLoader(**config['granule']) + try: + granule_loader = GranuleLoader(**config['granule']) - slicer_config = config['slicer'] - slicer = cls._parse_module(slicer_config, module_mappings) + slicer_config = config['slicer'] + slicer = cls._parse_module(slicer_config, module_mappings) - tile_processors = [] - for processor_config in config['processors']: - module = cls._parse_module(processor_config, module_mappings) - tile_processors.append(module) + tile_processors = [] + for processor_config in config['processors']: + module = cls._parse_module(processor_config, module_mappings) + tile_processors.append(module) - return cls(granule_loader, slicer, data_store_factory, metadata_store_factory, tile_processors) + return cls(granule_loader, slicer, data_store_factory, metadata_store_factory, tile_processors) + except KeyError as e: + raise PipelineBuildingError(f"Cannot build pipeline because {e} is missing from the YAML.") + except Exception: + raise PipelineBuildingError("Cannot build pipeline.") @classmethod def _parse_module(cls, module_config: dict, module_mappings: dict): @@ -136,23 +164,36 @@ class Pipeline: return processor_module async def run(self): + + logger.info(f"Running pipeline with {self._max_concurrency} threads per process") async with self._granule_loader as (dataset, granule_name): start = time.perf_counter() - async with aiomultiprocess.Pool(initializer=_init_worker, - initargs=(self._tile_processors, - dataset, - self._data_store_factory, - self._metadata_store_factory)) as pool: + + shared_memory = self._manager.Namespace() + async with Pool(initializer=_init_worker, + initargs=(self._tile_processors, + dataset, + self._data_store_factory, + self._metadata_store_factory, + shared_memory), + maxtasksperchild=self._max_concurrency, + childconcurrency=self._max_concurrency) as pool: serialized_tiles = [nexusproto.NexusTile.SerializeToString(tile) for tile in self._slicer.generate_tiles(dataset, granule_name)] # aiomultiprocess is built on top of the stdlib multiprocessing library, which has the limitation that # a queue can't have more than 2**15-1 tasks. So, we have to batch it. - for chunk in type(self)._chunk_list(serialized_tiles, MAX_QUEUE_SIZE): - await pool.map(_process_tile_in_worker, chunk) + for chunk in self._chunk_list(serialized_tiles, MAX_CHUNK_SIZE): + try: + await pool.map(_process_tile_in_worker, chunk) + except ProxyException: + pool.terminate() + # Give the shared memory manager some time to write the exception + # await asyncio.sleep(1) + raise pickle.loads(shared_memory.error) end = time.perf_counter() logger.info("Pipeline finished in {} seconds".format(end - start)) @staticmethod - def _chunk_list(items, chunk_size): + def _chunk_list(items, chunk_size: int): return [items[i:i + chunk_size] for i in range(0, len(items), chunk_size)] diff --git a/granule_ingester/granule_ingester/processors/reading_processors/TileReadingProcessor.py b/granule_ingester/granule_ingester/processors/reading_processors/TileReadingProcessor.py index 14a44f5..8b69ad2 100644 --- a/granule_ingester/granule_ingester/processors/reading_processors/TileReadingProcessor.py +++ b/granule_ingester/granule_ingester/processors/reading_processors/TileReadingProcessor.py @@ -21,6 +21,7 @@ import numpy as np import xarray as xr from nexusproto import DataTile_pb2 as nexusproto +from granule_ingester.exceptions import TileProcessingError from granule_ingester.processors.TileProcessor import TileProcessor @@ -31,20 +32,17 @@ class TileReadingProcessor(TileProcessor, ABC): self.latitude = latitude self.longitude = longitude - # Common optional properties - self.temp_dir = None - self.metadata = None - # self.temp_dir = self.environ['TEMP_DIR'] - # self.metadata = self.environ['META'] - def process(self, tile, dataset: xr.Dataset, *args, **kwargs): - dimensions_to_slices = type(self)._convert_spec_to_slices(tile.summary.section_spec) + try: + dimensions_to_slices = self._convert_spec_to_slices(tile.summary.section_spec) - output_tile = nexusproto.NexusTile() - output_tile.CopyFrom(tile) - output_tile.summary.data_var_name = self.variable_to_read + output_tile = nexusproto.NexusTile() + output_tile.CopyFrom(tile) + output_tile.summary.data_var_name = self.variable_to_read - return self._generate_tile(dataset, dimensions_to_slices, output_tile) + return self._generate_tile(dataset, dimensions_to_slices, output_tile) + except Exception: + raise TileProcessingError("Could not generate tiles from the granule.") @abstractmethod def _generate_tile(self, dataset: xr.Dataset, dimensions_to_slices: Dict[str, slice], tile): diff --git a/granule_ingester/granule_ingester/writers/DataStore.py b/granule_ingester/granule_ingester/writers/DataStore.py index 889d41e..a64399b 100644 --- a/granule_ingester/granule_ingester/writers/DataStore.py +++ b/granule_ingester/granule_ingester/writers/DataStore.py @@ -7,6 +7,7 @@ from granule_ingester.healthcheck import HealthCheck class DataStore(HealthCheck, ABC): + @abstractmethod def save_data(self, nexus_tile: nexusproto.NexusTile) -> None: pass diff --git a/granule_ingester/granule_ingester/writers/SolrStore.py b/granule_ingester/granule_ingester/writers/SolrStore.py index 9d6a7f0..276a988 100644 --- a/granule_ingester/granule_ingester/writers/SolrStore.py +++ b/granule_ingester/granule_ingester/writers/SolrStore.py @@ -13,64 +13,87 @@ # See the License for the specific language governing permissions and # limitations under the License. - -from asyncio import AbstractEventLoop - +import asyncio +import functools +import json import logging +from asyncio import AbstractEventLoop from datetime import datetime from pathlib import Path from typing import Dict -import aiohttp +import pysolr +from kazoo.handlers.threading import KazooTimeoutError +from kazoo.exceptions import NoNodeError from nexusproto.DataTile_pb2 import * -from tenacity import * +from granule_ingester.exceptions import SolrFailedHealthCheckError, SolrLostConnectionError from granule_ingester.writers.MetadataStore import MetadataStore logger = logging.getLogger(__name__) +def run_in_executor(f): + @functools.wraps(f) + def inner(*args, **kwargs): + loop = asyncio.get_running_loop() + return loop.run_in_executor(None, lambda: f(*args, **kwargs)) + + return inner + + class SolrStore(MetadataStore): - def __init__(self, host_and_port='http://localhost:8983'): + def __init__(self, solr_url=None, zk_url=None): super().__init__() self.TABLE_NAME = "sea_surface_temp" self.iso: str = '%Y-%m-%dT%H:%M:%SZ' - self._host_and_port = host_and_port + self._solr_url = solr_url + self._zk_url = zk_url self.geo_precision: int = 3 - self.collection: str = "nexustiles" + self._collection: str = "nexustiles" self.log: logging.Logger = logging.getLogger(__name__) self.log.setLevel(logging.DEBUG) - self._session = None + self._solr = None + + def _get_connection(self) -> pysolr.Solr: + if self._zk_url: + zk = pysolr.ZooKeeper(f"{self._zk_url}") + collections = {} + for c in zk.zk.get_children("collections"): + collections.update(json.loads(zk.zk.get("collections/{}/state.json".format(c))[0].decode("ascii"))) + zk.collections = collections + return pysolr.SolrCloud(zk, self._collection, always_commit=True) + elif self._solr_url: + return pysolr.Solr(f'{self._solr_url}/solr/{self._collection}', always_commit=True) + else: + raise RuntimeError("You must provide either solr_host or zookeeper_host.") def connect(self, loop: AbstractEventLoop = None): - self._session = aiohttp.ClientSession(loop=loop) + self._solr = self._get_connection() async def health_check(self): try: - async with aiohttp.ClientSession() as session: - response = await session.get('{}/solr/{}/admin/ping'.format(self._host_and_port, self.collection)) - if response.status == 200: - return True - else: - logger.error("Solr health check returned status {}.".format(response.status)) - except aiohttp.ClientConnectionError as e: - logger.error("Cannot connect to Solr!") - - return False + connection = self._get_connection() + connection.ping() + except pysolr.SolrError: + raise SolrFailedHealthCheckError("Cannot connect to Solr!") + except NoNodeError: + raise SolrFailedHealthCheckError("Connected to Zookeeper but cannot connect to Solr!") + except KazooTimeoutError: + raise SolrFailedHealthCheckError("Cannot connect to Zookeeper!") async def save_metadata(self, nexus_tile: NexusTile) -> None: solr_doc = self._build_solr_doc(nexus_tile) + await self._save_document(solr_doc) - await self._save_document(self.collection, solr_doc) - - @retry(stop=stop_after_attempt(5)) - async def _save_document(self, collection: str, doc: dict): - url = '{}/solr/{}/update/json/docs?commit=true'.format(self._host_and_port, collection) - response = await self._session.post(url, json=doc) - if response.status < 200 or response.status >= 400: - raise RuntimeError("Saving data to Solr failed with HTTP status code {}".format(response.status)) + @run_in_executor + def _save_document(self, doc: dict): + try: + self._solr.add([doc]) + except pysolr.SolrError: + raise SolrLostConnectionError("Lost connection to Solr, and cannot save tiles.") def _build_solr_doc(self, tile: NexusTile) -> Dict: summary: TileSummary = tile.summary diff --git a/granule_ingester/requirements.txt b/granule_ingester/requirements.txt index 4d9d4cb..9b06860 100644 --- a/granule_ingester/requirements.txt +++ b/granule_ingester/requirements.txt @@ -1,3 +1,6 @@ cassandra-driver==3.23.0 -aiomultiprocess +aiomultiprocess==0.7.0 aioboto3 +tblib==1.6.0 +pysolr==3.9.0 +kazoo==2.8.0 \ No newline at end of file diff --git a/granule_ingester/tests/pipeline/test_Pipeline.py b/granule_ingester/tests/pipeline/test_Pipeline.py index c18bf8b..34e66c6 100644 --- a/granule_ingester/tests/pipeline/test_Pipeline.py +++ b/granule_ingester/tests/pipeline/test_Pipeline.py @@ -29,10 +29,11 @@ class TestPipeline(unittest.TestCase): pass relative_path = "../config_files/ingestion_config_testfile.yaml" - file_path = os.path.join(os.path.dirname(__file__), relative_path) - pipeline = Pipeline.from_file(config_path=str(file_path), - data_store_factory=MockDataStore, - metadata_store_factory=MockMetadataStore) + with open(os.path.join(os.path.dirname(__file__), relative_path)) as file: + yaml_str = file.read() + pipeline = Pipeline.from_string(config_str=yaml_str, + data_store_factory=MockDataStore, + metadata_store_factory=MockMetadataStore) self.assertEqual(pipeline._data_store_factory, MockDataStore) self.assertEqual(pipeline._metadata_store_factory, MockMetadataStore)
