From 89ed5948eed862ddcabf7755e7d87ad778e27b79 Mon Sep 17 00:00:00 2001 From: Troy Sankey Date: Mon, 15 Nov 2021 09:22:30 -0500 Subject: [PATCH 1/2] feat: Add logic to enable loading pepped events into Amplitude DENG-???? --- edx_prefectutils/amplitude.py | 347 ++++++++++++++++++++++++++++++++++ 1 file changed, 347 insertions(+) create mode 100644 edx_prefectutils/amplitude.py diff --git a/edx_prefectutils/amplitude.py b/edx_prefectutils/amplitude.py new file mode 100644 index 0000000..007b018 --- /dev/null +++ b/edx_prefectutils/amplitude.py @@ -0,0 +1,347 @@ +import backoff +#from prefect import task +#from prefect.engine import signals +from prefect.utilities.logging import get_logger + +AMPLITUDE_REQUEST_HEADERS = { + 'Content-Type': 'application/json', + 'Accept': '*/*' +} + +def batch_generator(iterable, batch_size=1000): + """ + This is just a simple batching algorithm, implemented as a generator. + + Args: + iterable (iterable): all the individual elements to group into batches. + batch_size (int): the number of elements in each batch. + + Returns: + A generator of batches, where each batch is effectively an iterable a + slice of the input. + """ + total_length = len(iterable) + for batch_start in range(0, total_length, batch_size): + batch_end = min(batch_start + batch_size, total_length) + yield iterable[batch_start:batch_end] + + +class LazyResultFetcher(list): + """ + This creates a lazy list-like object which represents a complete snowflake result set by deferring fetching batches + of results until they are indexed. + + Usage: Invoke cursor.execute(query), create a new LazyResultFetcher providing the cursor to the constructor, and treat + the resulting LazyResultFetcher object as a list containing every row in the complete result set, with the caveat that + len() will return the count of the actual contents. Keep indexing higher and higher until you get an IndexError. + + Slicing this object always "downgrades" the result to a standard list, so we can safely pass slices around to other + worker threads without worrying that indexing those slices will inadvertently cause snowflake batch fetches. + """ + + def __init__(self, cursor, fetch_size=10000): + self._cursor = cursor + self._fetch_size = fetch_size + self._results_exhausted = False + self._total_rowcount = cursor.rowcount + self._fetched_rowcount = 0 + super(LazyResultFetcher, self).__init__() + + def actual_length(): + """ + """ + return len(self) + self._total_rowcount - self._fetched_rowcount + + def __getitem__(self, arg): + """ + This behaves the exact same way as standard list indexing, except this allows indexing past the end of the list! + This function will make a best effort to return data for the given index by fetching batches of snowflake + records and accumulating them up to the index. + + Raises: + TypeError: if the given arg isn't of type slice or int. + IndexError: if the given arg represents a list index which isn't backed by snowflake data. + """ + logger = get_logger() + + # First, determine the maximum index requested, which is what drives how many more batches to fetch from + # snowflake. + max_idx = None + if isinstance(arg, slice): + if arg.stop is None: + max_idx = len(self) - 1 + else: + max_idx = arg.stop + elif isinstance(arg, int): + max_idx = arg + else: + raise TypeError("list index must be of type slice or int.") + + # Next, fetch just the right amount of result batches to be able to index max_idx. If there are no more results + # to tap, then we'll raise an IndexError at this point. + while max_idx >= len(self) and not self._results_exhausted: + fetched_events = self._cursor.fetchmany(self._fetch_size) + if fetched_events: + self.extend(fetched_events) + self._fetched_rowcount += len(fetched_events) + if self._fetched_rowcount == self._total_rowcount: + self._results_exhausted = True + else: + # We should never get to this point, but if we have then something unexpected caused the number of + # fetched rows to become exhausted before we fetch the expected "total" rowcount. One thing that can + # cause this to happen is if the caller manually fetched any results without letting this LazyResultFetcher + # object do it. + LOGGER.warning("Fetched fewer than expected results. Check for skipped events!") + self._results_exhausted = True + if max_idx >= len(self): + raise IndexError("list index out of range, there are no more Snowflake results to fetch for this index.") + + # Finally, perform the actual indexing. + return super(LazyResultFetcher, self).__getitem__(arg) + + +def _fetch_and_consolidate_and_batch(cursor, batch_size=1000, fetch_size=10000): + """ + Given a snowflake cursor referencing the results of a query returning an entire processing segment of amplitude + events, allocate events to workers such that the loading speed into amplitude is optimized. + + Argument `cursor` must not have had any rows fetched from it yet, and should represent the results of a query which: + + - Makes consecutive the events with the same user_id/device_id segment, i.e. coalesce(user_id, device_id). + - Sorts user_id/device_id segments in descending order by rowcount of segments. + - Sorts events by timestamp within each user_id/device_id segment. + + This function will consolidate segments smaller than batch_size events such that the resulting consolidated segments + don't exceed batch_size events. Then, it will batch any segments still larger than batch_size. The final product + should look like an iterable of iterable of iterable of events. I.e.: + + - First level loops over user/device segment (each of these will be allocated to a different worker thread). + - Second level loops over batch_size batches of events (this is a requirement imposed by the Amplitude API). + - Third level loops over individual events. + + Memory Efficiency: In practice, this function won't fetch all events for the current processing segment at once; + instead it will fetch small amounts until there are enough events to go around for all the worker threads. The + best-case-scenario memory consumption would be about N_worker_threads * batch_size * event_size, which is tiny. + However, the worse-case-scenario is that somehow a large proportion of the events in this processing segment + corresponds to a single user, in which case the memory consumption would equal that of loading the entire processing + segment into memory (multiple gigabytes). We could address memory consumption by adding even more complexity (and + possibly bugs) to the code, but this is an exceedingly rare scenario that isn't worth optimizing for. + + Args: + cursor (Snowflake cursor): Cursor of an executed query from which to fetch results. + batch_size (int): Maximum number of events to include in a single Amplitude API call. This cannot be greater + than 1000 since that is currently the maximum imposed by Amplitude. + fetch_size (int): The "small" number of events to fetch from snowflake each time we need to ask for more. + + Returns: + iterable of iterable of iterable of events. + """ + accumulated_events = LazyResultFetcher(cursor, fetch_size=fetch_size) + last_idx_for_worker = None + while True: + if accumulated_events.actual_length() <= batch_size: + # All we have left is one or more segments that completely fit into one batch, so send it all! + last_idx_for_worker = len(accumulated_events) + else: + if accumulated_events[0].user_device == accumulated_events[batch_size - 1].user_device: + # We have a segment which is larger than a single bucket. We should find the last event for that + # segment and send the entire segment to a worker. + target_user_device = accumulated_events[0].user_device + idx = batch_size - 1 + while True: + # Is there a more efficient way than indexing one event at a time? Yes, but simple iteration is not + # our performance bottleneck, and anything more complicated can introduce bugs. + idx += 1 + try: + if accumulated_events[idx].user_device != target_user_device: + last_idx_for_worker = idx + break + except IndexError: + last_idx_for_worker = idx + break + else: + # One or more segments completely fit into one batch, so lets consolidate them and send them. Exclude + # the last segment so that we send no more than one-batch-worth of events by finding the first event for + # the last segment and using that as a cutoff: + cut_user_device = accumulated_events[batch_size - 1].user_device + idx = batch_size - 1 + while True: + idx -= 1 + if accumulated_events[idx].user_device != cut_user_device: + last_idx_for_worker = idx + 1 + break + + # Now we have `last_idx_for_worker` set to the point where we want to slice off events to send to a worker + # thread. Also, it is important that we delete that slice from the accumulated_events so that the python + # garbage collector will free all that memory as soon as the corresponding worker thread returns. + events_for_next_worker = accumulated_events[:last_idx_for_worker] + del accumulated_events[:last_idx_for_worker] + + # Finally, batch the events such that each batch is small enough to fit inside a single Amplitude API call. + yield batch_generator(events_for_next_worker, batch_size) + + +def _get_old_high_watermark(sf_credentials: SFCredentials, sf_role: str): + """ + Get the amplitude event ID of the event immediately following the very last event sent to Amplitude's API. + + Raises: + snowflake.connector.ProgrammingError: if the loader log table is inaccessible for some reason. + """ + sf_connection = create_snowflake_connection(sf_credentials, sf_role) + query = """ + select max(amplitude_event_id) + from prod.amplitude_events.amplitude_events_loader_log + """ + cursor.execute(query) + old_high_watermark = cursor.fetchone()[0] + if not old_high_watermark: + old_high_watermark = 0 + + +def _get_new_high_watermark(sf_credentials: SFCredentials, sf_role: str): + """ + TODO: make this more robust (e.g. what if the table does not exist? what if the table is empty?) + """ + sf_connection = create_snowflake_connection(sf_credentials, sf_role) + cursor = sf_connection.cursor() + # Get the current maximum amplitude_event_id. + query = """ + select max(amplitude_event_id) + from prod.amplitude_events.amplitude_events + """ + cursor.execute(query) + return cursor.fetchone()[0] + + +def get_processing_segments(): + # Get the old (i.e. current) high watermark. + old_high_watermark = _get_old_high_watermark() + + # Get the new high watermark, i.e. the current maximum amplitude_event_id. + new_high_watermark = _get_new_high_watermark(sf_credentials, sf_role) + + # Get a list of processing segments which cover all events between the old and + # new high watermark. The result is a list of start/end tuples describing + # amplitude_event_id ranges, where each range represents a single processing + # segment with the configured batch size. The start ID is inclusive, whereas + # the end ID is exclusive. + # + # e.g. if old=1000000, new=3600000, bucketsize=1000000, then: + # processing_segments = [ + # (1000000, 2000000), + # (2000000, 3000000), + # (3000000, 3600001), + # ] + # + # In the above example, the last batch is smaller than 1m events because there + # are no more events to include. + return [ + (processing_segment_start, min(processing_segment_start + processing_bucket_size, new_high_watermark + 1)) + for processing_segment_start in range(old_high_watermark, new_high_watermark, processing_bucket_size) + ] + + +def load_events_into_amplitude(amplitude_api_key: str, processing_segments: list) + # Iterate over processing segments in chronological order, handling each one + # one at a time, and completely freeing the memory of the previous one before + # fetching the next. + for processing_segment in processing_segments: + # Prep a cursor to fetch the events for the current processing segment. + # + # For events with amplitude_event_id within the current processing_segment, further segment the events by the + # following dimension: + # + # coalesce(user_id, device_id) + # + # Then, sort events by segment size and sort by timestamp within each segment. + segmented_by_device_user_query = ( + f""" + select coalesce(user_id, device_id) as user_device, + count(*) over (partition by user_device) as user_device_event_count, + * + from events_to_load + where amplitude_event_id >= {processing_segment_start} + and amplitude_event_id < {processing_segment_end} + order by user_device_event_count asc, + user_device asc, + timestamp asc + """ + ) + segmented_by_device_user = cursor.execute(segmented_by_device_user_query) + + # Make a generator which can fetch events and slice them for worker threads. + segmented_by_device_user_and_batched = _fetch_and_consolidate_and_batch(cursor) + + # Only allow a fixed number (64, for example) of simultaneous network calls. + available_senders = threading.Semaphore(64) + + @backoff.on_exception( + backoff.expo, + ( + requests.exceptions.HTTPError, # i.e. the status is 4xx or 5xx. + requests.exceptions.ConnectionError, # This includes connection timeout. + requests.exceptions.Timeout, # This includes read timeout. + # In other data pipelines, retrying on read timeout could be dangerous because the events from the + # failed request my have been successfully ingested, so retrying can result in duplicate events. + # However, we do not mind retrying on read timeout with Amplitude loading because Amplitude will de-dupe + # the events anyway. + ), + ) + def send_batch(amplitude_events_batch): + """ + Actually make a Batch API request to send a batch of events. + + This is meant to be executed by a worker thread only. Since we only want to allow a finite number of worker + threads to be running API calls simultaneously, the API call is surrounded by logic to acquire/release a + semaphore lock. It's important that this semaphore logic exists *inside* the retry loop because we don't want + to block any other workers while this one is in the middle of backing off. + + Args: + amplitude_events_batch (iterable of events): 1000 or fewer events formatted for consumption by Amplitude. + """ + available_senders.acquire() + r = requests.post( + "https://api2.amplitude.com/2/httpapi", + json={ + "api_key": amplitude_api_key, + "events": amplitude_events_batch, + }, + headers=AMPLITUDE_REQUEST_HEADERS, + ) + r.raise_for_status() + available_senders.release() + + def load_device_segment(device_segment): + """ + Load all the events for a given device within the current processing + segment. + + Args: + segment(iterable of iterable of events): + All events for single user or device, batched by 1000. + """ + for api_batch in device_segment: + send_batch(api_batch) + + # Launch a thread pool of 128 (twice that of the semaphore limit) so that + # we have a few extra threads to pick up the slack whenever a thread gets + # stuck backing off due to EPDS rate limiting. In theory, any time we have + # fewer than the maximum number of blocking network calls, we are running + # slower than capacity. + with ThreadPoolExecutor(max_workers=128) as executor: + executor.map(load_device_segment, segmented_by_device_user_and_batched) + + # Update the high watermark. + # We do not need to wrap this in retry logic, since that is already + # implemented in the prefect source: + # https://github.com/PrefectHQ/prefect/blob/da07d255/src/prefect/client/client.py#L755-L756 + prefect.backend.set_key_value( + key="amplitude_high_watermark", + value=new_high_watermark + ) + + # This shouldn't be necessary, but make sure that we aren't deferring any garbage collection since the next + # iteration may start to allocate a lot of memory. + del segmented_by_device_user_and_batched + gc.collect() From 2f57f68ba310fc8686b6e93e67d6a0bb4ef78784 Mon Sep 17 00:00:00 2001 From: Troy Sankey Date: Mon, 15 Nov 2021 16:27:20 -0500 Subject: [PATCH 2/2] feat: FIXUP --- edx_prefectutils/amplitude.py | 428 ++++++++++++++++------------------ edx_prefectutils/snowflake.py | 111 ++++++++- tests/test_snowflake.py | 62 ++++- 3 files changed, 378 insertions(+), 223 deletions(-) diff --git a/edx_prefectutils/amplitude.py b/edx_prefectutils/amplitude.py index 007b018..f925d3d 100644 --- a/edx_prefectutils/amplitude.py +++ b/edx_prefectutils/amplitude.py @@ -1,115 +1,55 @@ +import gc +import itertools +import threading +from concurrent.futures import ThreadPoolExecutor +from typing import Generator, Iterable, List, Sequence + import backoff -#from prefect import task -#from prefect.engine import signals -from prefect.utilities.logging import get_logger +import requests +import snowflake.connector +from prefect import task + +from .snowflake import (LazySnowflakeResultList, SFCredentials, + create_snowflake_connection) AMPLITUDE_REQUEST_HEADERS = { 'Content-Type': 'application/json', 'Accept': '*/*' } +PROCESSING_BUCKET_SIZE = 1000000 +MAX_CONCURRENT_AMPLITUDE_CONNECTIONS = 64 + -def batch_generator(iterable, batch_size=1000): +def _batch_generator(sequence: Sequence, batch_size: int = 1000) -> Generator[Sequence, None, None]: """ This is just a simple batching algorithm, implemented as a generator. Args: - iterable (iterable): all the individual elements to group into batches. + sequence (Sequence): all the individual elements to group into batches. Must support len() and random indexing. batch_size (int): the number of elements in each batch. Returns: - A generator of batches, where each batch is effectively an iterable a + A generator of batches, where each batch is effectively an iterable slice of the input. """ - total_length = len(iterable) + total_length = len(sequence) for batch_start in range(0, total_length, batch_size): batch_end = min(batch_start + batch_size, total_length) - yield iterable[batch_start:batch_end] - - -class LazyResultFetcher(list): - """ - This creates a lazy list-like object which represents a complete snowflake result set by deferring fetching batches - of results until they are indexed. - - Usage: Invoke cursor.execute(query), create a new LazyResultFetcher providing the cursor to the constructor, and treat - the resulting LazyResultFetcher object as a list containing every row in the complete result set, with the caveat that - len() will return the count of the actual contents. Keep indexing higher and higher until you get an IndexError. - - Slicing this object always "downgrades" the result to a standard list, so we can safely pass slices around to other - worker threads without worrying that indexing those slices will inadvertently cause snowflake batch fetches. - """ + yield sequence[batch_start:batch_end] - def __init__(self, cursor, fetch_size=10000): - self._cursor = cursor - self._fetch_size = fetch_size - self._results_exhausted = False - self._total_rowcount = cursor.rowcount - self._fetched_rowcount = 0 - super(LazyResultFetcher, self).__init__() - def actual_length(): - """ - """ - return len(self) + self._total_rowcount - self._fetched_rowcount - - def __getitem__(self, arg): - """ - This behaves the exact same way as standard list indexing, except this allows indexing past the end of the list! - This function will make a best effort to return data for the given index by fetching batches of snowflake - records and accumulating them up to the index. - - Raises: - TypeError: if the given arg isn't of type slice or int. - IndexError: if the given arg represents a list index which isn't backed by snowflake data. - """ - logger = get_logger() - - # First, determine the maximum index requested, which is what drives how many more batches to fetch from - # snowflake. - max_idx = None - if isinstance(arg, slice): - if arg.stop is None: - max_idx = len(self) - 1 - else: - max_idx = arg.stop - elif isinstance(arg, int): - max_idx = arg - else: - raise TypeError("list index must be of type slice or int.") - - # Next, fetch just the right amount of result batches to be able to index max_idx. If there are no more results - # to tap, then we'll raise an IndexError at this point. - while max_idx >= len(self) and not self._results_exhausted: - fetched_events = self._cursor.fetchmany(self._fetch_size) - if fetched_events: - self.extend(fetched_events) - self._fetched_rowcount += len(fetched_events) - if self._fetched_rowcount == self._total_rowcount: - self._results_exhausted = True - else: - # We should never get to this point, but if we have then something unexpected caused the number of - # fetched rows to become exhausted before we fetch the expected "total" rowcount. One thing that can - # cause this to happen is if the caller manually fetched any results without letting this LazyResultFetcher - # object do it. - LOGGER.warning("Fetched fewer than expected results. Check for skipped events!") - self._results_exhausted = True - if max_idx >= len(self): - raise IndexError("list index out of range, there are no more Snowflake results to fetch for this index.") - - # Finally, perform the actual indexing. - return super(LazyResultFetcher, self).__getitem__(arg) - - -def _fetch_and_consolidate_and_batch(cursor, batch_size=1000, fetch_size=10000): +def _fetch_and_consolidate_and_batch( + connection: snowflake.connector.SnowflakeConnection, query: str, batch_size: int = 1000, fetch_size: int = 10000 +) -> Iterable[Iterable[Sequence[dict]]]: """ - Given a snowflake cursor referencing the results of a query returning an entire processing segment of amplitude - events, allocate events to workers such that the loading speed into amplitude is optimized. + Given a snowflake query returning an entire processing segment of amplitude events, fetch and allocate events to + workers such that the loading speed into amplitude is optimized. - Argument `cursor` must not have had any rows fetched from it yet, and should represent the results of a query which: + Argument `query` must: - - Makes consecutive the events with the same user_id/device_id segment, i.e. coalesce(user_id, device_id). - - Sorts user_id/device_id segments in descending order by rowcount of segments. - - Sorts events by timestamp within each user_id/device_id segment. + - Make consecutive the events with the same user_id/device_id segment, i.e. coalesce(user_id, device_id). + - Sort user_id/device_id segments in descending order by rowcount of segments. + - Sort events by increasing timestamp within each user_id/device_id segment. This function will consolidate segments smaller than batch_size events such that the resulting consolidated segments don't exceed batch_size events. Then, it will batch any segments still larger than batch_size. The final product @@ -128,7 +68,8 @@ def _fetch_and_consolidate_and_batch(cursor, batch_size=1000, fetch_size=10000): possibly bugs) to the code, but this is an exceedingly rare scenario that isn't worth optimizing for. Args: - cursor (Snowflake cursor): Cursor of an executed query from which to fetch results. + connection (snowflake.connector.SnowflakeConnection): A snowflake connection object. + query (str): A snowflake query. batch_size (int): Maximum number of events to include in a single Amplitude API call. This cannot be greater than 1000 since that is currently the maximum imposed by Amplitude. fetch_size (int): The "small" number of events to fetch from snowflake each time we need to ask for more. @@ -136,59 +77,58 @@ def _fetch_and_consolidate_and_batch(cursor, batch_size=1000, fetch_size=10000): Returns: iterable of iterable of iterable of events. """ - accumulated_events = LazyResultFetcher(cursor, fetch_size=fetch_size) - last_idx_for_worker = None + accumulated_events = LazySnowflakeResultList(connection, query, fetch_size=fetch_size) + end_idx_for_worker = None while True: - if accumulated_events.actual_length() <= batch_size: + if accumulated_events.logical_length() <= batch_size: # All we have left is one or more segments that completely fit into one batch, so send it all! - last_idx_for_worker = len(accumulated_events) - else: - if accumulated_events[0].user_device == accumulated_events[batch_size - 1].user_device: - # We have a segment which is larger than a single bucket. We should find the last event for that - # segment and send the entire segment to a worker. - target_user_device = accumulated_events[0].user_device - idx = batch_size - 1 - while True: - # Is there a more efficient way than indexing one event at a time? Yes, but simple iteration is not - # our performance bottleneck, and anything more complicated can introduce bugs. - idx += 1 - try: - if accumulated_events[idx].user_device != target_user_device: - last_idx_for_worker = idx - break - except IndexError: - last_idx_for_worker = idx - break - else: - # One or more segments completely fit into one batch, so lets consolidate them and send them. Exclude - # the last segment so that we send no more than one-batch-worth of events by finding the first event for - # the last segment and using that as a cutoff: - cut_user_device = accumulated_events[batch_size - 1].user_device - idx = batch_size - 1 - while True: - idx -= 1 - if accumulated_events[idx].user_device != cut_user_device: - last_idx_for_worker = idx + 1 + end_idx_for_worker = len(accumulated_events) + elif accumulated_events[0].user_device == accumulated_events[batch_size].user_device: + # We have a segment which is larger than a single bucket. We should find the last event for that + # segment and send the entire segment to a worker. + target_user_device = accumulated_events[0].user_device + idx = batch_size + while True: + # Is there a more efficient way than indexing one event at a time? Yes, but simple iteration is not + # our performance bottleneck, and anything more complicated can introduce bugs. + idx += 1 + try: + if accumulated_events[idx].user_device != target_user_device: + end_idx_for_worker = idx break - - # Now we have `last_idx_for_worker` set to the point where we want to slice off events to send to a worker + except IndexError: + end_idx_for_worker = idx + break + else: + # One or more segments completely fit into one batch, so lets consolidate them and send them. Exclude + # the last segment so that we send no more than one-batch-worth of events by finding the first event for + # the last segment and using that as a cutoff: + cut_user_device = accumulated_events[batch_size].user_device + idx = batch_size + while True: + idx -= 1 + if accumulated_events[idx].user_device != cut_user_device: + end_idx_for_worker = idx + 1 + break + + # Now we have `end_idx_for_worker` set to the point where we want to slice off events to send to a worker # thread. Also, it is important that we delete that slice from the accumulated_events so that the python # garbage collector will free all that memory as soon as the corresponding worker thread returns. - events_for_next_worker = accumulated_events[:last_idx_for_worker] - del accumulated_events[:last_idx_for_worker] + events_for_next_worker = accumulated_events[:end_idx_for_worker] + del accumulated_events[:end_idx_for_worker] # Finally, batch the events such that each batch is small enough to fit inside a single Amplitude API call. - yield batch_generator(events_for_next_worker, batch_size) + yield _batch_generator(events_for_next_worker, batch_size) -def _get_old_high_watermark(sf_credentials: SFCredentials, sf_role: str): +def _get_old_high_watermark(connection: snowflake.connector.SnowflakeConnection) -> int: """ Get the amplitude event ID of the event immediately following the very last event sent to Amplitude's API. Raises: - snowflake.connector.ProgrammingError: if the loader log table is inaccessible for some reason. + snowflake.connector.ProgrammingError: if the loader log table is inaccessible. """ - sf_connection = create_snowflake_connection(sf_credentials, sf_role) + cursor = connection.cursor() query = """ select max(amplitude_event_id) from prod.amplitude_events.amplitude_events_loader_log @@ -197,15 +137,20 @@ def _get_old_high_watermark(sf_credentials: SFCredentials, sf_role: str): old_high_watermark = cursor.fetchone()[0] if not old_high_watermark: old_high_watermark = 0 + return old_high_watermark -def _get_new_high_watermark(sf_credentials: SFCredentials, sf_role: str): +def _get_new_high_watermark(connection: snowflake.connector.SnowflakeConnection) -> int: """ - TODO: make this more robust (e.g. what if the table does not exist? what if the table is empty?) + Get the last amplitude_event_id (inclusive) which should be processed as part of this flow. + + Args: + TODO + + Raises: + snowflake.connector.ProgrammingError: if the amplitude_events table is inaccessible. """ - sf_connection = create_snowflake_connection(sf_credentials, sf_role) - cursor = sf_connection.cursor() - # Get the current maximum amplitude_event_id. + cursor = connection.cursor() query = """ select max(amplitude_event_id) from prod.amplitude_events.amplitude_events @@ -214,39 +159,130 @@ def _get_new_high_watermark(sf_credentials: SFCredentials, sf_role: str): return cursor.fetchone()[0] -def get_processing_segments(): - # Get the old (i.e. current) high watermark. - old_high_watermark = _get_old_high_watermark() +def _set_high_watermark(connection: snowflake.connector.SnowflakeConnection, new_high_watermark: int): + """ + Add a new watermark record to the loader log. + + Args: + TODO + + Raises: + snowflake.connector.ProgrammingError: if the loader log table is inaccessible. + """ + cursor = connection.cursor() + query = f""" + insert into prod.amplitude_events.amplitude_events_loader_log (timestamp, amplitude_event_id) + values (current_timestamp()::timestamp_tz, {new_high_watermark}) + """ + cursor.execute(query) + connection.commit() + + +@backoff.on_exception( + backoff.expo, + ( + requests.exceptions.HTTPError, # i.e. the status is 4xx or 5xx. + requests.exceptions.ConnectionError, # This includes connection timeout. + requests.exceptions.Timeout, # This includes read timeout. + # In other data pipelines, retrying on read timeout could be dangerous because the events from the + # failed request my have been successfully ingested, so retrying can result in duplicate events. + # However, we do not mind retrying on read timeout with Amplitude loading because Amplitude will de-dupe + # the events anyway. + ), +) +def _send_batch_to_amplitude( + amplitude_events_batch: Sequence[dict], finite_senders_lock: threading.Semaphore, amplitude_api_key: str +): + """ + Actually make a Batch API request to send a batch of events. - # Get the new high watermark, i.e. the current maximum amplitude_event_id. - new_high_watermark = _get_new_high_watermark(sf_credentials, sf_role) + This is meant to be executed by a worker thread only. Since we only want to allow a finite number of worker + threads to be running API calls simultaneously, the API call is surrounded by logic to acquire/release a + semaphore lock. It's important that this semaphore logic exists *inside* the retry loop because we don't + want to block any other workers while this one is in the middle of backing off. - # Get a list of processing segments which cover all events between the old and - # new high watermark. The result is a list of start/end tuples describing - # amplitude_event_id ranges, where each range represents a single processing - # segment with the configured batch size. The start ID is inclusive, whereas - # the end ID is exclusive. + Args: + amplitude_events_batch (iterable of events): 1000 or fewer events formatted for consumption by + Amplitude. + + Raises: + TODO + """ + with finite_senders_lock: + r = requests.post( + "https://api2.amplitude.com/2/httpapi", + json={ + "api_key": amplitude_api_key, + "events": amplitude_events_batch, + "options": { + "min_id_length": 1, # valid edX user IDs start in the single-digis. + }, + }, + headers=AMPLITUDE_REQUEST_HEADERS, + ) + r.raise_for_status() + + +def _send_user_device_segment( + user_device_segment: Iterable[Sequence[dict]], finite_senders_lock: threading.Semaphore, amplitude_api_key: str +): + """ + Load all the events for a given device within the current processing segment into Amplitude. + + Args: + segment(iterable of iterable of events): All events for single user or device, batched by 1000. + + Raises: + TODO + """ + for api_batch in user_device_segment: + _send_batch_to_amplitude(api_batch, finite_senders_lock, amplitude_api_key) + + +@task +def get_processing_segments(sf_credentials: SFCredentials, sf_role: str) -> List[tuple[int, int]]: + """ + """ + sf_connection = create_snowflake_connection(sf_credentials, sf_role) + + # Get the old (i.e. current) high watermark. This is the ID of the last event that was loaded. + old_high_watermark = _get_old_high_watermark(sf_connection) + + # Get the new high watermark, i.e. the current maximum amplitude_event_id. Thi sis the last event that we plan to + # load. + new_high_watermark = _get_new_high_watermark(sf_connection) + + # Get a list of processing segments which cover all events between the old and new high watermark. The result is a + # list of start/end tuples describing amplitude_event_id ranges, where each range represents a single processing + # segment with the configured batch size. The start ID of each segment is inclusive, whereas the end ID is + # exclusive; this is pythonic. # - # e.g. if old=1000000, new=3600000, bucketsize=1000000, then: + # e.g. if old=999999, new=3599999, bucketsize=1000000, then: # processing_segments = [ # (1000000, 2000000), # (2000000, 3000000), - # (3000000, 3600001), + # (3000000, 3600000), # ] # - # In the above example, the last batch is smaller than 1m events because there - # are no more events to include. + # In the above example, the last batch is smaller than 1m events because there are no more events to include. return [ - (processing_segment_start, min(processing_segment_start + processing_bucket_size, new_high_watermark + 1)) - for processing_segment_start in range(old_high_watermark, new_high_watermark, processing_bucket_size) + (processing_segment_start, min(processing_segment_start + PROCESSING_BUCKET_SIZE, new_high_watermark + 1)) + for processing_segment_start in range(old_high_watermark + 1, new_high_watermark + 1, PROCESSING_BUCKET_SIZE) ] -def load_events_into_amplitude(amplitude_api_key: str, processing_segments: list) +@task +def load_events_into_amplitude( + sf_credentials: SFCredentials, sf_role: str, amplitude_api_key: str, processing_segments: List[tuple[int, int]] +): + """ + """ # Iterate over processing segments in chronological order, handling each one # one at a time, and completely freeing the memory of the previous one before # fetching the next. for processing_segment in processing_segments: + sf_connection = create_snowflake_connection(sf_credentials, sf_role) + # Prep a cursor to fetch the events for the current processing segment. # # For events with amplitude_event_id within the current processing_segment, further segment the events by the @@ -255,8 +291,8 @@ def load_events_into_amplitude(amplitude_api_key: str, processing_segments: list # coalesce(user_id, device_id) # # Then, sort events by segment size and sort by timestamp within each segment. - segmented_by_device_user_query = ( - f""" + processing_segment_start, processing_segment_end = processing_segment + segmented_by_user_device_query = f""" select coalesce(user_id, device_id) as user_device, count(*) over (partition by user_device) as user_device_event_count, * @@ -267,81 +303,31 @@ def load_events_into_amplitude(amplitude_api_key: str, processing_segments: list user_device asc, timestamp asc """ - ) - segmented_by_device_user = cursor.execute(segmented_by_device_user_query) # Make a generator which can fetch events and slice them for worker threads. - segmented_by_device_user_and_batched = _fetch_and_consolidate_and_batch(cursor) - - # Only allow a fixed number (64, for example) of simultaneous network calls. - available_senders = threading.Semaphore(64) - - @backoff.on_exception( - backoff.expo, - ( - requests.exceptions.HTTPError, # i.e. the status is 4xx or 5xx. - requests.exceptions.ConnectionError, # This includes connection timeout. - requests.exceptions.Timeout, # This includes read timeout. - # In other data pipelines, retrying on read timeout could be dangerous because the events from the - # failed request my have been successfully ingested, so retrying can result in duplicate events. - # However, we do not mind retrying on read timeout with Amplitude loading because Amplitude will de-dupe - # the events anyway. - ), + segmented_by_user_device_and_batched = _fetch_and_consolidate_and_batch( + sf_connection, segmented_by_user_device_query ) - def send_batch(amplitude_events_batch): - """ - Actually make a Batch API request to send a batch of events. - - This is meant to be executed by a worker thread only. Since we only want to allow a finite number of worker - threads to be running API calls simultaneously, the API call is surrounded by logic to acquire/release a - semaphore lock. It's important that this semaphore logic exists *inside* the retry loop because we don't want - to block any other workers while this one is in the middle of backing off. - Args: - amplitude_events_batch (iterable of events): 1000 or fewer events formatted for consumption by Amplitude. - """ - available_senders.acquire() - r = requests.post( - "https://api2.amplitude.com/2/httpapi", - json={ - "api_key": amplitude_api_key, - "events": amplitude_events_batch, - }, - headers=AMPLITUDE_REQUEST_HEADERS, + # Only allow a fixed number of simultaneous network calls. + finite_senders_lock = threading.Semaphore(MAX_CONCURRENT_AMPLITUDE_CONNECTIONS) + + # Launch a thread pool with twice as many threads as the the connection limit so that we have a few extra + # threads to pick up the slack whenever a thread gets stuck backing off due to EPDS rate limiting. In theory, + # any time we have fewer than the maximum number of blocking network calls, we are running slower than capacity. + with ThreadPoolExecutor(max_workers=MAX_CONCURRENT_AMPLITUDE_CONNECTIONS * 2) as executor: + executor.map( + _send_user_device_segment, + segmented_by_user_device_and_batched, + itertools.repeat(finite_senders_lock), + itertools.repeat(amplitude_api_key), ) - r.raise_for_status() - available_senders.release() - def load_device_segment(device_segment): - """ - Load all the events for a given device within the current processing - segment. - - Args: - segment(iterable of iterable of events): - All events for single user or device, batched by 1000. - """ - for api_batch in device_segment: - send_batch(api_batch) - - # Launch a thread pool of 128 (twice that of the semaphore limit) so that - # we have a few extra threads to pick up the slack whenever a thread gets - # stuck backing off due to EPDS rate limiting. In theory, any time we have - # fewer than the maximum number of blocking network calls, we are running - # slower than capacity. - with ThreadPoolExecutor(max_workers=128) as executor: - executor.map(load_device_segment, segmented_by_device_user_and_batched) - - # Update the high watermark. - # We do not need to wrap this in retry logic, since that is already - # implemented in the prefect source: - # https://github.com/PrefectHQ/prefect/blob/da07d255/src/prefect/client/client.py#L755-L756 - prefect.backend.set_key_value( - key="amplitude_high_watermark", - value=new_high_watermark - ) + # Update the high watermark. processing_segment_end represents the next index after the last index loaded, so + # subtract 1 to get a valid watermark. + _set_high_watermark(sf_connection, processing_segment_end - 1) # This shouldn't be necessary, but make sure that we aren't deferring any garbage collection since the next # iteration may start to allocate a lot of memory. - del segmented_by_device_user_and_batched + del segmented_by_user_device_and_batched gc.collect() diff --git a/edx_prefectutils/snowflake.py b/edx_prefectutils/snowflake.py index 724501f..f57d767 100644 --- a/edx_prefectutils/snowflake.py +++ b/edx_prefectutils/snowflake.py @@ -3,6 +3,7 @@ """ import os from collections import namedtuple +from collections.abc import MutableSequence from typing import List, TypedDict import backoff @@ -512,7 +513,7 @@ def get_batched_rows_from_snowflake( where: str = None, ): """ - Query batches of rows from snowflake as a generator. + Export a table from snowflake and get batches of rows as a generator. Args: sf_credentials (SFCredentials): Snowflake public key credentials in the format @@ -558,3 +559,111 @@ def get_batched_rows_from_snowflake( while len(batch) > 0: yield [SnowFlakeRow(*row) for row in batch] batch = cursor.fetchmany(batch_size) + + +class LazySnowflakeResultList(MutableSequence): + """ + Given a Snowflake connection and an arbitrary query, this creates a lazy list-like object which represents a + complete result set by deferring fetching batches of results until they are indexed. + + Usage: + 1. Create a new LazySnowflakeResultList providing the snowflake connection and a query to the constructor. + 2. Call self.execute() on the resulting object. + 3. Treat the resulting object as a list containing every row in the complete result set. + + The higher you index, the more data gets fetched; eventually, if you index high enough, you will get an IndexError. + As you work through the list and process the data, consider deleting elements that aren't needed anymore: + + del lazy_snowflake_result_list[:1000] # delete first 1k records, encouraging the GC to deallocate that memory. + + Caveats: + - Avoid indexing the list any order other than low to high. Any other indexing pattern may result in a massive + number of batch fetches and memory allocation, defeating the purpose of this class. + - Slicing or otherwise indexing this object always "downgrades" the result to a standard list or element, so we can + safely pass output slices around without worrying that subsequent indexing of those output slices will + inadvertently mess up the internal state of the LazySnowflakeResultList. + """ + def __init__(self, connection: snowflake.connector.SnowflakeConnection, query: str, fetch_size: int = 10000): + self._buffer = list() + self._connection = connection + self._query = query + self._fetch_size = fetch_size + + self._fetched_rowcount = 0 + self._cursor = self._connection.cursor() # TODO: should we use a dict cursor? + self._cursor.execute(query) + + def __len__(self) -> int: + """ + Get the logical rowcount of this object. In other words, what would len(self._buffer) return if all rows were + prefetched? + + The reason this is different than self._cursor.rowcount is because by the time this function is called we may + already have deleted some rows from the internal list (e.g. to deallocate memory which isn't needed anymore). + """ + return len(self._buffer) + self._cursor.rowcount - self._fetched_rowcount + + def __getitem__(self, arg): + """ + This behaves the exact same way as standard list indexing, except this allows indexing past the end of the list! + This function will make a best effort to return data for the given index by fetching batches of snowflake + records and accumulating them up to the index. + + Raises: + TypeError: if the given arg isn't of type slice or int. + IndexError: if the given arg represents a list index which isn't backed by snowflake data. + """ + logger = get_logger() + + # First, determine the maximum index requested, which is what drives how many more batches to fetch from + # snowflake. + max_idx = None + if isinstance(arg, slice): + if arg.stop is None: + max_idx = len(self) - 1 + else: + max_idx = arg.stop + elif isinstance(arg, int): + max_idx = arg + else: + raise TypeError("list index must be of type slice or int.") + + # Next, make a best effort to fetch just the right amount of result batches to be able to index max_idx. + # + # Note: The requested index may be greater than the actual number of results, but here we make no attempt to + # remedy that. Later in the code an IndexError will be raised. + while max_idx >= len(self._buffer) and not self._cursor.is_closed(): + fetched_events = self._cursor.fetchmany(self._fetch_size) + if fetched_events: + logger.debug("Fetched %s more results.", len(fetched_events)) + self._buffer.extend(fetched_events) + self._fetched_rowcount += len(fetched_events) + if self._fetched_rowcount == self._cursor.rowcount: + logger.debug("Closing cursor due to all expected rows being fetched.") + self._cursor.close() + else: + # We should never get to this point, but if we have then something unexpected caused the number of + # fetched rows to become exhausted before we fetch the expected "total" rowcount. One thing that can + # cause this to happen is if the caller manually fetched any results without letting this + # LazySnowflakeResultList object do it. + logger.warning("Fetched fewer than expected results. Check for skipped events!") + self._cursor.close() + + # Finally, perform the actual indexing. + return self._buffer[arg] + + ###### + # All remaining instance methods are passthrough methods to expose the underlying buffer without any customization. + ###### + + def __delitem__(self, i): + del self._buffer[i] + + def __setitem__(self, i, v): + self._buffer[i] = v + + def insert(self, i, v): + self._buffer.insert(i, v) + + def __str__(self): + return str(self._buffer) diff --git a/tests/test_snowflake.py b/tests/test_snowflake.py index 7864419..f87772d 100755 --- a/tests/test_snowflake.py +++ b/tests/test_snowflake.py @@ -1,11 +1,13 @@ #!/usr/bin/env python - """ Tests for Snowflake utils in the `edx_prefectutils` package. """ +from unittest import TestCase + import mock import pytest +from ddt import data, ddt, unpack from prefect.core import Flow from prefect.engine import signals from prefect.utilities.debug import raise_on_exception @@ -398,3 +400,61 @@ def test_load_s3_data_to_snowflake_data_disable_check(mock_sf_connection): mock_call ] ) + + +@ddt +class TestLazySnowflakeResultList(TestCase): + @data( + ([[1, 2, 3], [4, 5, 6], [7, 8], [], []], 3), + ([[1, 2, 3], [4, 5, 6], [7, 8, 9], [], []], 3), + ) + @unpack + def test_success(self, batches, fetch_size): + """ + Test multiple aspects of successful scenarios. + """ + # Flatten the batches to make a list which resembles the final results. + expected_results = [val for sublist in batches for val in sublist] + + mock_cursor = mock.Mock() + mock_cursor.rowcount = len(expected_results) + mock_cursor.is_closed.side_effect = lambda: True if mock_cursor.close.mock_calls else False + mock_cursor.fetchmany.side_effect = batches + mock_connection = mock.Mock() + mock_connection.cursor.return_value = mock_cursor + test_sql = "select * from does_not_matter" + test_fetcher = snowflake.LazySnowflakeResultList( + mock_connection, test_sql, fetch_size=fetch_size + ) + + # After object construction, we should already have called cursor.execute(). + mock_cursor.execute.assert_called_once_with(test_sql) + + # The length should be reported as the logical length, rather than the actual length (which would still be 0 at + # this point, before any rows have been fetched. + assert len(test_fetcher) == len(expected_results) + + # Actually make sure no fetches were attempted yet, despite already knowing the final length. + assert not mock_cursor.fetchmany.mock_calls + + # Coerce a fetch. + assert test_fetcher[0] == expected_results[0] + assert test_fetcher[2] == expected_results[2] + + # Make sure only a single fetch was made, despite multiple indexings. All rows indexed so far are within the + # first fetch batch. + assert len(mock_cursor.fetchmany.mock_calls) == 1 + + # While we are still in the middle of fetching rows, attempt deleting some rows from the beginning and make sure + # it continues to behave the exact same way as a standard list. + assert len(test_fetcher) == len(expected_results) + del test_fetcher[:2] + del expected_results[:2] + assert len(test_fetcher) == len(expected_results) + assert test_fetcher[0] == expected_results[0] + + # Coerce two more fetches and make sure the cursor was closed. + assert test_fetcher[4] == expected_results[4] + assert len(mock_cursor.fetchmany.mock_calls) == 3 + assert len(test_fetcher) == len(expected_results) + assert mock_cursor.close.mock_calls