From a0e1bf61aaa0d69950d6e23bcb338e07df848cdc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Honza=20Kr=C3=A1l?= Date: Thu, 1 Oct 2015 01:25:29 +0200 Subject: [PATCH 1/3] Experimental helper for doing bulk requests in parallel --- elasticsearch/helpers/parallel.py | 66 +++++++++++++++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 elasticsearch/helpers/parallel.py diff --git a/elasticsearch/helpers/parallel.py b/elasticsearch/helpers/parallel.py new file mode 100644 index 00000000..7132e568 --- /dev/null +++ b/elasticsearch/helpers/parallel.py @@ -0,0 +1,66 @@ +from multiprocessing.dummy import Pool +from queue import Empty, Queue + +from threading import Event + +from . import streaming_bulk + +def consume(queue, done): + """ + Create an iterator on top of a Queue. + """ + while True: + try: + yield queue.get(True, .01) + except Empty: + if done.is_set(): + break + +def wrapped_bulk(client, input, output, done, **kwargs): + """ + Wrap a call to streaming_bulk by feeding it data frm a queue and writing + the outputs to another queue. + """ + try: + for result in streaming_bulk(client, consume(input, done), **kwargs): + output.put(result) + except: + done.set() + raise + +def feed_data(actions, input, done): + """ + Feed data from an iterator into a queue. + """ + for a in actions: + input.put(a, True) + + # error short-circuit + if done.is_set(): + break + done.set() + + +def parallel_bulk(client, actions, thread_count=5, **kwargs): + """ + Paralel version of the bulk helper. It runs a thread pool with a thread for + a producer and ``thread_count`` threads for. + """ + done = Event() + input, output = Queue(), Queue() + pool = Pool(thread_count + 1) + + results = [ + pool.apply_async(wrapped_bulk, (client, input, output, done), kwargs) + for _ in range(thread_count)] + pool.apply_async(feed_data, (actions, input, done)) + + while True: + try: + yield output.get(True, .01) + except Empty: + if done.is_set() and all(r.ready() for r in results): + break + + pool.close() + pool.join() From 2208387e0baef57ffa5b41baa98159a1b7b87466 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Honza=20Kr=C3=A1l?= Date: Tue, 6 Oct 2015 02:10:42 +0200 Subject: [PATCH 2/3] Extract the processing of individual bulk request to a separate method --- elasticsearch/helpers/__init__.py | 120 ++++++++++++++++-------------- 1 file changed, 65 insertions(+), 55 deletions(-) diff --git a/elasticsearch/helpers/__init__.py b/elasticsearch/helpers/__init__.py index c1fc81bd..c7c9f439 100644 --- a/elasticsearch/helpers/__init__.py +++ b/elasticsearch/helpers/__init__.py @@ -40,6 +40,10 @@ def expand_action(data): return action, data.get('_source', data) def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer): + """ + Split actions into chunks by number or size, serialize them into strings in + the process. + """ bulk_actions = [] size, action_count = 0, 0 for action, data in actions: @@ -65,6 +69,64 @@ def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer): if bulk_actions: yield bulk_actions +def _process_bulk_chunk(client, bulk_actions, raise_on_exception=True, raise_on_error=True, **kwargs): + """ + Send a bulk request to elasticsearch and process the output. + """ + # if raise on error is set, we need to collect errors per chunk before raising them + errors = [] + + try: + # send the actual request + resp = client.bulk('\n'.join(bulk_actions) + '\n', **kwargs) + except TransportError as e: + # default behavior - just propagate exception + if raise_on_exception: + raise e + + # if we are not propagating, mark all actions in current chunk as failed + err_message = str(e) + exc_errors = [] + + # deserialize the data back, thisis expensive but only run on + # errors if raise_on_exception is false, so shouldn't be a real + # issue + bulk_data = iter(map(client.transport.serializer.loads, bulk_actions)) + while True: + try: + # collect all the information about failed actions + action = next(bulk_data) + op_type, action = action.popitem() + info = {"error": err_message, "status": e.status_code, "exception": e} + if op_type != 'delete': + info['data'] = next(bulk_data) + info.update(action) + exc_errors.append({op_type: info}) + except StopIteration: + break + + # emulate standard behavior for failed actions + if raise_on_error: + raise BulkIndexError('%i document(s) failed to index.' % len(exc_errors), exc_errors) + else: + for err in exc_errors: + yield False, err + return + + # go through request-reponse pairs and detect failures + for op_type, item in map(methodcaller('popitem'), resp['items']): + ok = 200 <= item.get('status', 500) < 300 + if not ok and raise_on_error: + errors.append({op_type: item}) + + if ok or not errors: + # if we are not just recording all errors to be able to raise + # them all at once, yield items individually + yield ok, {op_type: item} + + if errors: + raise BulkIndexError('%i document(s) failed to index.' % len(errors), errors) + def streaming_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1014 * 1024, raise_on_error=True, expand_action_callback=expand_action, raise_on_exception=True, **kwargs): @@ -122,63 +184,11 @@ def streaming_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1014 * should return a tuple containing the action line and the data line (`None` if data line should be omitted). """ - serializer = client.transport.serializer actions = map(expand_action_callback, actions) - # if raise on error is set, we need to collect errors per chunk before raising them - errors = [] - - for bulk_actions in _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer): - try: - # send the actual request - resp = client.bulk('\n'.join(bulk_actions) + '\n', **kwargs) - except TransportError as e: - # default behavior - just propagate exception - if raise_on_exception: - raise e - - # if we are not propagating, mark all actions in current chunk as failed - err_message = str(e) - exc_errors = [] - - # deserialize the data back, thisis expensive but only run on - # errors if raise_on_exception is false, so shouldn't be a real - # issue - bulk_data = iter(map(serializer.loads, bulk_actions)) - while True: - try: - # collect all the information about failed actions - action = next(bulk_data) - op_type, action = action.popitem() - info = {"error": err_message, "status": e.status_code, "exception": e} - if op_type != 'delete': - info['data'] = next(bulk_data) - info.update(action) - exc_errors.append({op_type: info}) - except StopIteration: - break - - # emulate standard behavior for failed actions - if raise_on_error: - raise BulkIndexError('%i document(s) failed to index.' % len(exc_errors), exc_errors) - else: - for err in exc_errors: - yield False, err - continue - - # go through request-reponse pairs and detect failures - for op_type, item in map(methodcaller('popitem'), resp['items']): - ok = 200 <= item.get('status', 500) < 300 - if not ok and raise_on_error: - errors.append({op_type: item}) - - if not errors: - # if we are not just recording all errors to be able to raise - # them all at once, yield items individually - yield ok, {op_type: item} - - if errors: - raise BulkIndexError('%i document(s) failed to index.' % len(errors), errors) + for bulk_actions in _chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer): + for result in _process_bulk_chunk(client, bulk_actions, raise_on_exception, raise_on_error, **kwargs): + yield result def bulk(client, actions, stats_only=False, **kwargs): """ From 2ba03d261ab4cce0689e094b8f4c5730c4d37418 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Honza=20Kr=C3=A1l?= Date: Tue, 6 Oct 2015 02:11:22 +0200 Subject: [PATCH 3/3] greatly simplify parallel_bulk by using Pool.imap --- elasticsearch/helpers/parallel.py | 68 ++++++------------------------- 1 file changed, 13 insertions(+), 55 deletions(-) diff --git a/elasticsearch/helpers/parallel.py b/elasticsearch/helpers/parallel.py index 7132e568..036f8da1 100644 --- a/elasticsearch/helpers/parallel.py +++ b/elasticsearch/helpers/parallel.py @@ -1,66 +1,24 @@ from multiprocessing.dummy import Pool -from queue import Empty, Queue -from threading import Event +from . import _process_bulk_chunk, _chunk_actions, expand_action -from . import streaming_bulk -def consume(queue, done): +def parallel_bulk(client, actions, thread_count=4, chunk_size=500, + max_chunk_bytes=100 * 1014 * 1024, + expand_action_callback=expand_action, **kwargs): """ - Create an iterator on top of a Queue. + Parallel version of the bulk helper. """ - while True: - try: - yield queue.get(True, .01) - except Empty: - if done.is_set(): - break + actions = map(expand_action_callback, actions) -def wrapped_bulk(client, input, output, done, **kwargs): - """ - Wrap a call to streaming_bulk by feeding it data frm a queue and writing - the outputs to another queue. - """ - try: - for result in streaming_bulk(client, consume(input, done), **kwargs): - output.put(result) - except: - done.set() - raise + pool = Pool(thread_count) -def feed_data(actions, input, done): - """ - Feed data from an iterator into a queue. - """ - for a in actions: - input.put(a, True) - - # error short-circuit - if done.is_set(): - break - done.set() - - -def parallel_bulk(client, actions, thread_count=5, **kwargs): - """ - Paralel version of the bulk helper. It runs a thread pool with a thread for - a producer and ``thread_count`` threads for. - """ - done = Event() - input, output = Queue(), Queue() - pool = Pool(thread_count + 1) - - results = [ - pool.apply_async(wrapped_bulk, (client, input, output, done), kwargs) - for _ in range(thread_count)] - pool.apply_async(feed_data, (actions, input, done)) - - while True: - try: - yield output.get(True, .01) - except Empty: - if done.is_set() and all(r.ready() for r in results): - break + for result in pool.imap( + lambda chunk: list(_process_bulk_chunk(client, chunk, **kwargs)), + _chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer) + ): + for item in result: + yield item pool.close() pool.join()