From b83af66c5c13f3b7c5357c89b9fe981e01a1797d Mon Sep 17 00:00:00 2001 From: qlevavasseur Date: Tue, 20 Jun 2023 16:32:28 +0200 Subject: [PATCH] test pour VM DEV --- CHANGELOG.md | 4 ++ target_bigquery/storage_write.py | 64 +++++++++++++++++++++--------- target_bigquery/target.py | 15 +++++-- target_bigquery/tests/test_sync.py | 57 ++++++++++++++++++++++---- 4 files changed, 111 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5687246..0bebd3f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html), and is generated by [Changie](https://github.com/miniscruff/changie). +## Unreleased +### Fixed +* Manage new workers connections with common offset +* Waiting for workers to finish their jobs bfore closing streams ## 0.6.6 - 2023-04-20 ### Added diff --git a/target_bigquery/storage_write.py b/target_bigquery/storage_write.py index 887e0a0..22cf939 100644 --- a/target_bigquery/storage_write.py +++ b/target_bigquery/storage_write.py @@ -13,6 +13,8 @@ NOTE: This is naive and will vary drastically based on network speed, for example on a GCP VM. """ import os +from threading import Lock +from time import sleep from multiprocessing import Process from multiprocessing.connection import Connection from multiprocessing.dummy import Process as _Thread @@ -46,13 +48,14 @@ from target_bigquery.core import BaseBigQuerySink, BaseWorker, Denormalized, storage_client_factory from target_bigquery.proto_gen import proto_schema_factory_v2 +logger = logging.getLogger(__name__) + # Stream specific constant MAX_IN_FLIGHT = 15 """Maximum number of concurrent requests per worker be processed by grpc before awaiting.""" - Dispatcher = Callable[[types.AppendRowsRequest], writer.AppendRowsFuture] -StreamComponents = Tuple[str, writer.AppendRowsStream, Dispatcher] +StreamComponents = Tuple[str, writer.AppendRowsStream, Lock, Dispatcher] def get_application_stream(client: BigQueryWriteClient, job: "Job") -> StreamComponents: @@ -62,7 +65,8 @@ def get_application_stream(client: BigQueryWriteClient, job: "Job") -> StreamCom write_stream = client.create_write_stream(parent=job.parent, write_stream=write_stream) job.template.write_stream = write_stream.name append_rows_stream = writer.AppendRowsStream(client, job.template) - rv = (write_stream.name, append_rows_stream) + lock = Lock() + rv = (write_stream.name, append_rows_stream, lock) job.stream_notifier.send(rv) return *rv, retry( append_rows_stream.send, @@ -78,7 +82,8 @@ def get_default_stream(client: BigQueryWriteClient, job: "Job") -> StreamCompone **BigQueryWriteClient.parse_table_path(job.parent), stream="_default" ) append_rows_stream = writer.AppendRowsStream(client, job.template) - rv = (job.template.write_stream, append_rows_stream) + lock = Lock() + rv = (job.template.write_stream, append_rows_stream, lock) job.stream_notifier.send(rv) return *rv, retry( append_rows_stream.send, @@ -122,14 +127,23 @@ def generate_template(message: Type[Message]): return template -class Job(NamedTuple): +class Job(): parent: str template: types.AppendRowsRequest stream_notifier: Connection data: types.ProtoRows - offset: int = 0 attempts: int = 1 - + + def __init__(self, + parent, + template, + stream_notifier, + data): + """Initialize the worker process.""" + self.parent = parent + self.template = template + self.stream_notifier = stream_notifier + self.data = data class StorageWriteBatchWorker(BaseWorker): """Worker process for the storage write API.""" @@ -141,6 +155,8 @@ def __init__(self, *args, **kwargs): self.awaiting: List[writer.AppendRowsFuture] = [] self.cache: Dict[str, StreamComponents] = {} self.max_errors_before_recycle = 5 + self.offsets: Dict[str, int] = {} + self.logger=logger def run(self): """Run the worker process.""" @@ -155,19 +171,25 @@ def run(self): break if job is None: break - if job.parent not in self.cache: + if job.parent not in self.cache or self.cache[job.parent][1]._closed: self.cache[job.parent] = self.get_stream_components(client, job) - write_stream, _, dispatch = cast(StreamComponents, self.cache[job.parent]) + self.offsets[job.parent] = 0 + write_stream, _, _, dispatch = cast(StreamComponents, self.cache[job.parent]) + try: + if self.cache[job.parent][1]._closed: + raise Exception("Connection closed before locking.") kwargs = {} if write_stream.endswith("_default"): kwargs["offset"] = None kwargs["path"] = write_stream else: - kwargs["offset"] = job.offset + kwargs["offset"] = self.offsets[job.parent] self.awaiting.append(dispatch(generate_request(job.data, **kwargs))) + except Exception as exc: job.attempts += 1 + self.logger.info(f"job.attempts : {job.attempts}") self.max_errors_before_recycle -= 1 if job.attempts > 3: # TODO: add a metric for this + a DLQ & wrap exception type @@ -186,20 +208,23 @@ def run(self): else: self.log_notifier.send( f"[{self.ext_id}] Sent {len(job.data.serialized_rows)} rows to {write_stream}" - f" with offset {job.offset}." + f" with offset {self.offsets[job.parent]}." ) + self.offsets[job.parent] += len(job.data.serialized_rows) if len(self.awaiting) > MAX_IN_FLIGHT: self.wait() finally: self.queue.task_done() # Wait for all in-flight requests to complete after poison pill + self.logger.info(f"[{self.ext_id}] : {self.offsets}") self.wait(drain=True) self.close_cached_streams() + self.logger.info("Worker process exiting.") self.log_notifier.send("Worker process exiting.") def close_cached_streams(self) -> None: """Close all cached streams.""" - for _, stream, _ in self.cache.values(): + for _, stream, _, _ in self.cache.values(): try: stream.close() except Exception as exc: @@ -240,6 +265,7 @@ class StorageWriteProcessBatchWorker(StorageWriteBatchWorker, Process): class BigQueryStorageWriteSink(BaseBigQuerySink): MAX_WORKERS = os.cpu_count() * 2 + MAX_JOBS_QUEUED = MAX_WORKERS * 2 WORKER_CAPACITY_FACTOR = 10 WORKER_CREATION_MIN_INTERVAL = 1.0 @@ -276,7 +302,6 @@ def __init__( ) self.stream_notification, self.stream_notifier = target.pipe_cls(False) self.template = generate_template(self.proto_schema) - self.offset = 0 @property def proto_schema(self) -> Type[Message]: @@ -300,17 +325,19 @@ def process_record(self, record: Dict[str, Any], context: Dict[str, Any]) -> Non ) def process_batch(self, context: Dict[str, Any]) -> None: + while self.global_queue.qsize() >= self.MAX_JOBS_QUEUED: + self.logger.warn(f"Max jobs enqueued reached ({self.MAX_JOBS_QUEUED})") + sleep(1) + self.global_queue.put( Job( parent=self.parent, template=self.template, data=self.proto_rows, stream_notifier=self.stream_notifier, - offset=self.offset, ) ) self.increment_jobs_enqueued() - self.offset += len(self.proto_rows.serialized_rows) def commit_streams(self) -> None: while self.stream_notification.poll(): @@ -320,17 +347,18 @@ def commit_streams(self) -> None: if not self.open_streams: return self.open_streams = [ - (name, stream) for name, stream in self.open_streams if not name.endswith("_default") + (name, stream, lock) for name, stream, lock in self.open_streams if not name.endswith("_default") ] if self.open_streams: committer = storage_client_factory(self._credentials) - for name, stream in self.open_streams: + for name, stream, _ in self.open_streams: + #self.logger.info(f"YO !!!!!!!! : AFTER LOCK INSTANT KILL {name}") stream.close() committer.finalize_write_stream(name=name) write = committer.batch_commit_write_streams( types.BatchCommitWriteStreamsRequest( parent=self.parent, - write_streams=[name for name, _ in self.open_streams], + write_streams=[name for name, _, _ in self.open_streams], ) ) self.logger.info(f"Batch commit time: {write.commit_time}") diff --git a/target_bigquery/target.py b/target_bigquery/target.py index aafce5f..ecd7d17 100644 --- a/target_bigquery/target.py +++ b/target_bigquery/target.py @@ -9,6 +9,7 @@ # The above copyright notice and this permission notice shall be included in all copies or # substantial portions of the Software. """BigQuery target class.""" +import os import copy import time import uuid @@ -37,15 +38,18 @@ # Defaults for target worker pool parameters MAX_WORKERS = 15 """Maximum number of workers to spawn.""" +MAX_JOBS_QUEUED = 30 +"""Maximum number of jobs placed in the global queue to avoid memory overload.""" WORKER_CAPACITY_FACTOR = 5 """Jobs enqueued must exceed the number of active workers times this number.""" WORKER_CREATION_MIN_INTERVAL = 5 """Minimum time between worker creation attempts.""" - class TargetBigQuery(Target): """Target for BigQuery.""" + _MAX_RECORD_AGE_IN_MINUTES = 5.0 + name = "target-bigquery" config_jsonschema = th.PropertiesList( th.Property( @@ -484,6 +488,7 @@ def get_sink( def drain_one(self, sink: Sink) -> None: # type: ignore """Drain a sink. Includes a hook to manage the worker pool and notifications.""" + #self.logger.info(f"Jobs queued : {self.queue.qsize()} | Max nb jobs queued : {os.cpu_count() * 4} | Nb workers : {len(self.workers)} | Max nb workers : {os.cpu_count() * 2}") self.resize_worker_pool() while self.job_notification.poll(): ext_id = self.job_notification.recv() @@ -519,13 +524,15 @@ def drain_all(self, is_endofpipe: bool = False) -> None: # type: ignore for worker in self.workers: if worker.is_alive(): self.queue.put(None) - for worker in self.workers: + while len(self.workers): worker.join() + worker = self.workers.pop() for sink in self._sinks_active.values(): sink.clean_up() else: + for worker in self.workers: + worker.join() for sink in self._sinks_active.values(): sink.pre_state_hook() - if state: - self._write_state_message(state) + self._write_state_message(state) self._reset_max_record_age() diff --git a/target_bigquery/tests/test_sync.py b/target_bigquery/tests/test_sync.py index 0485b33..c803606 100644 --- a/target_bigquery/tests/test_sync.py +++ b/target_bigquery/tests/test_sync.py @@ -28,13 +28,38 @@ {"type": "RECORD", "stream": "{stream_name}", "record": {"id": 2, "rep_key": 2}, "time_extracted": "2022-11-11T20:40:56.807796+00:00"} {"type": "RECORD", "stream": "{stream_name}", "record": {"id": 3, "rep_key": 3}, "time_extracted": "2022-11-11T20:40:56.808289+00:00"} {"type": "RECORD", "stream": "{stream_name}", "record": {"id": 4, "rep_key": 4}, "time_extracted": "2022-11-11T20:40:56.808782+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 5, "rep_key": 5}, "time_extracted": "2022-11-11T20:40:56.807284+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 6, "rep_key": 6}, "time_extracted": "2022-11-11T20:40:56.807796+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 7, "rep_key": 7}, "time_extracted": "2022-11-11T20:40:56.808289+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 8, "rep_key": 8}, "time_extracted": "2022-11-11T20:40:56.808782+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 9, "rep_key": 9}, "time_extracted": "2022-11-11T20:40:56.808782+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 10, "rep_key": 10}, "time_extracted": "2022-11-11T20:40:56.808782+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 11, "rep_key": 11}, "time_extracted": "2022-11-11T20:40:56.808782+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 12, "rep_key": 12}, "time_extracted": "2022-11-11T20:40:56.808782+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 13, "rep_key": 13}, "time_extracted": "2022-11-11T20:40:56.808782+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 14, "rep_key": 14}, "time_extracted": "2022-11-11T20:40:56.808782+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 15, "rep_key": 15}, "time_extracted": "2022-11-11T20:40:56.808782+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 16, "rep_key": 16}, "time_extracted": "2022-11-11T20:40:56.808782+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 17, "rep_key": 17}, "time_extracted": "2022-11-11T20:40:56.808782+00:00"} +{"type": "RECORD", "stream": "{stream_name}", "record": {"id": 18, "rep_key": 18}, "time_extracted": "2022-11-11T20:40:56.808782+00:00"} """.strip() @pytest.mark.parametrize( "method", - ["batch_job", "streaming_insert", "storage_write_api", "gcs_stage"], - ids=["batch_job", "streaming_insert", "storage_write_api", "gcs_stage"], + + [ + #"batch_job", + #"streaming_insert", + "storage_write_api", + #"gcs_stage" + ], + ids=[ + #"batch_job", + #"streaming_insert", + "storage_write_api", + #"gcs_stage" + ], ) @pytest.mark.parametrize("batch_mode", [False, True], ids=["no_batch_mode", "batch_mode"]) def test_basic_sync(method, batch_mode): @@ -73,7 +98,11 @@ def test_basic_sync(method, batch_mode): **OPTS, }, ) - + target.get_sink_class().MAX_WORKERS = 1 + target.get_sink_class().MAX_JOBS_QUEUED = 2 + target.get_sink_class().WORKER_CAPACITY_FACTOR = 1 + target.get_sink_class().WORKER_CREATION_MIN_INTERVAL = 1 + client = bigquery_client_factory(BigQueryCredentials(json=target.config["credentials_json"])) try: client.query(f"TRUNCATE TABLE {target.config['dataset']}.{table_name}").result() @@ -102,7 +131,7 @@ def test_basic_sync(method, batch_mode): ] assert len(records) == 5 - assert len(records_2) == 5 + assert len(records_2) == 19 assert records == [ { "data": '{"Col_10_string":"9a164a0b-c6c7-46c5-8224-1ea0f7de485d","Col_11_datetime":"2022-11-11T13:40:56.806233+00:00","Col_12_int":1024423793,"Col_13_float":0.71921051655093,"Col_14_string":"c3630f45-660d-42c1-baf9-c58434d68786","Col_15_datetime":"2022-11-11T13:40:56.806238+00:00","Col_16_int":556025207,"Col_17_float":0.3965419777404805,"Col_18_string":"84ebc035-4b0c-40be-a9a3-61cbf7e0f01e","Col_19_datetime":"2022-11-11T13:40:56.806243+00:00","Col_20_int":661621821,"Col_21_float":0.37192361356880477,"Col_22_string":"41613825-9170-4a55-bfd6-c64222292573","Col_23_datetime":"2022-11-11T13:40:56.806248+00:00","Col_24_int":1716807152,"Col_25_float":0.9639895917683756,"Col_26_string":"87025878-196b-4990-9a65-ec06799fe34e","Col_27_datetime":"2022-11-11T13:40:56.806253+00:00","Col_28_int":542678613,"Col_29_float":0.4722333859761568,"Col_2_string":"55fcfce1-f7ce-4f90-8a54-e92de8dae40d","Col_3_datetime":"2022-11-11T13:40:56.806208+00:00","Col_4_int":1475701382,"Col_5_float":0.7920345506520963,"Col_6_string":"744dddf7-d07b-4b09-baf1-8b1e4914632b","Col_7_datetime":"2022-11-11T13:40:56.806225+00:00","Col_8_int":60612870,"Col_9_float":0.35286203712175723,"id":0,"rep_key":0}' @@ -124,15 +153,25 @@ def test_basic_sync(method, batch_mode): @pytest.mark.parametrize( "method", - ["batch_job", "streaming_insert", "gcs_stage", "storage_write_api"], - ids=["batch_job", "streaming_insert", "gcs_stage", "storage_write_api"], + [ + #"batch_job", + #"streaming_insert", + #"gcs_stage", + "storage_write_api" + ], + ids=[ + #"batch_job", + #"streaming_insert", + #"gcs_stage", + "storage_write_api" + ], ) def test_basic_denorm_sync(method): OPTS = { "method": method, "denormalized": True, "generate_view": False, - "batch_size": 2, # force multiple batches + "batch_size": 3, # force multiple batches } table_name = OPTS["method"] @@ -154,6 +193,10 @@ def test_basic_denorm_sync(method): **OPTS, }, ) + target.get_sink_class().MAX_WORKERS = 1 + target.get_sink_class().MAX_JOBS_QUEUED = 2 + target.get_sink_class().WORKER_CAPACITY_FACTOR = 1 + target.get_sink_class().WORKER_CREATION_MIN_INTERVAL = 1 client = bigquery_client_factory(BigQueryCredentials(json=target.config["credentials_json"])) try: