2121import backoff
2222import pds .ingress .util .log_util as log_util
2323import requests
24+ import requests_mock
2425from joblib import delayed
2526from joblib import Parallel
2627from more_itertools import chunked as batched
2728from pds .ingress import __version__
2829from pds .ingress .util .auth_util import AuthUtil
30+ from pds .ingress .util .backoff_util import simulate_batch_request_failure
31+ from pds .ingress .util .backoff_util import simulate_ingress_failure
2932from pds .ingress .util .config_util import ConfigUtil
3033from pds .ingress .util .hash_util import md5_for_path
3134from pds .ingress .util .log_util import get_log_level
@@ -193,7 +196,8 @@ def _process_batch(batch_index, request_batch, node_id, force_overwrite, api_gat
193196 except Exception as err :
194197 # If here, the HTTP request error was unrecoverable by a backoff/retry
195198 trimmed_path = ingress_response .get ("trimmed_path" )
196- update_summary_table (SUMMARY_TABLE , "failed" , trimmed_path )
199+ ingress_path = ingress_response .get ("ingress_path" )
200+ update_summary_table (SUMMARY_TABLE , "failed" , ingress_path )
197201
198202 logger .error ("Batch %d : Ingress failed for %s, Reason:\n %s" , batch_index , trimmed_path , str (err ))
199203
@@ -352,10 +356,11 @@ def _prepare_batch_for_ingress(ingress_path_batch, prefix, batch_index, batch_pb
352356@backoff .on_exception (
353357 backoff .expo ,
354358 Exception ,
355- max_time = 120 ,
359+ max_time = 10 , # TODO: for testing only, revert to 120
356360 logger = "request_batch_for_ingress" ,
357361)
358- def request_batch_for_ingress (request_batch , batch_index , node_id , force_overwrite , api_gateway_config ):
362+ @requests_mock .Mocker (kw = "mock_requests" , real_http = True )
363+ def request_batch_for_ingress (request_batch , batch_index , node_id , force_overwrite , api_gateway_config , ** kwargs ):
359364 """
360365 Submits a batch of ingress requests to the PDS Ingress App API.
361366
@@ -409,9 +414,12 @@ def request_batch_for_ingress(request_batch, batch_index, node_id, force_overwri
409414 "x-amz-docs-region" : api_gateway_region ,
410415 }
411416
412- response = requests .post (
413- api_gateway_url , params = params , data = json .dumps (request_batch ), headers = headers , timeout = 600
414- )
417+ # Simulate a random failure for the batch request if configured to do so
418+ with simulate_batch_request_failure (kwargs ["mock_requests" ], api_gateway_url .split ("?" )[0 ]):
419+ response = requests .post (
420+ api_gateway_url , params = params , data = json .dumps (request_batch ), headers = headers , timeout = 600
421+ )
422+
415423 elapsed_time = time .time () - start_time
416424
417425 # Ingress request successful
@@ -428,10 +436,11 @@ def request_batch_for_ingress(request_batch, batch_index, node_id, force_overwri
428436@backoff .on_exception (
429437 backoff .expo ,
430438 Exception ,
431- max_time = 120 ,
439+ max_time = 10 , # TODO: for testing only, revert to 120
432440 logger = "ingress_file_to_s3" ,
433441)
434- def ingress_file_to_s3 (ingress_response , batch_index , batch_pbar ):
442+ #@requests_mock.Mocker(kw="mock_requests", real_http=True)
443+ def ingress_file_to_s3 (ingress_response , batch_index , batch_pbar , ** kwargs ):
435444 """
436445 Copies the local file path using the pre-signed S3 URL returned from the
437446 Ingress Lambda App.
@@ -487,13 +496,16 @@ def ingress_file_to_s3(ingress_response, batch_index, batch_pbar):
487496 batch_pbar , total = os .stat (ingress_path ).st_size , filename = os .path .basename (ingress_path )
488497 )
489498
490- with open (ingress_path , "rb" ) as infile :
491- # Wrap file I/O with our upload bar to automatically track file upload progress
492- wrapped_file = CallbackIOWrapper (upload_pbar .update , infile , "read" )
499+ # Simulate a random failure for the S3 ingress request if configured to do so
500+ with requests_mock .Mocker (real_http = True ) as mock_requests :
501+ with simulate_ingress_failure (mock_requests , s3_ingress_url .split ("?" )[0 ]):
502+ with open (ingress_path , "rb" ) as infile :
503+ # Wrap file I/O with our upload bar to automatically track file upload progress
504+ wrapped_file = CallbackIOWrapper (upload_pbar .update , infile , "read" )
493505
494- # Only send the file data if the file is non-empty
495- response = requests .put (s3_ingress_url , data = wrapped_file if file_length > 0 else b"" , headers = headers )
496- response .raise_for_status ()
506+ # Only send the file data if the file is non-empty
507+ response = requests .put (s3_ingress_url , data = wrapped_file if file_length > 0 else b"" , headers = headers )
508+ response .raise_for_status ()
497509
498510 logger .info ("Batch %d : %s Ingest complete" , batch_index , trimmed_path )
499511 update_summary_table (SUMMARY_TABLE , "uploaded" , ingress_path )
@@ -517,10 +529,11 @@ def ingress_file_to_s3(ingress_response, batch_index, batch_pbar):
517529@backoff .on_exception (
518530 backoff .expo ,
519531 Exception ,
520- max_time = 120 ,
532+ max_time = 10 , # TODO: for testing only, revert to 120
521533 logger = "ingress_multipart_file_to_s3" ,
522534)
523- def ingress_multipart_file_to_s3 (ingress_response , batch_index , batch_pbar ):
535+ @requests_mock .Mocker (kw = "mock_requests" , real_http = True )
536+ def ingress_multipart_file_to_s3 (ingress_response , batch_index , batch_pbar , ** kwargs ):
524537 """
525538 Performs an ingress request for a file that is too large to be uploaded
526539 in a single request. The file is instead uploaded in multiple parts using
@@ -582,9 +595,11 @@ def ingress_multipart_file_to_s3(ingress_response, batch_index, batch_pbar):
582595 upload_pbar , f"{ os .path .basename (ingress_path )} (Part { part_number } /{ len (s3_ingress_urls )} )"
583596 )
584597
585- # Submit a single chunk to AWS
586- response = requests .put (s3_ingress_url , data = next (chunk_iterator ))
587- response .raise_for_status ()
598+ # Simulate a random failure for the S3 ingress request if configured to do so
599+ with simulate_ingress_failure (kwargs ["mock_requests" ], s3_ingress_url .split ("?" )[0 ]):
600+ # Submit a single chunk to AWS
601+ response = requests .put (s3_ingress_url , data = next (chunk_iterator ))
602+ response .raise_for_status ()
588603
589604 completed_parts .append ({"ETag" : response .headers ["ETag" ], "PartNumber" : part_number })
590605 except Exception as err :
@@ -848,16 +863,22 @@ def main(args):
848863 try :
849864 init_batch_progress_bars (args .num_threads )
850865 perform_ingress (request_batchs , node_id , args .force_overwrite , config ["API_GATEWAY" ])
866+ finally :
867+ close_batch_progress_bars ()
851868
852- logger .info ("All batches processed" )
869+ logger .info ("All batches processed" )
853870
871+ try :
854872 if len (SUMMARY_TABLE ["failed" ]) > 0 :
873+ logger .info ("----------------------------------------" )
855874 logger .info ("Reattempting ingress for failed files..." )
875+
856876 failed_ingresses = SUMMARY_TABLE ["failed" ]
857877 batched_failed_ingresses = list (batched (failed_ingresses , batch_size ))
858878 failed_request_batchs = prepare_batches (batched_failed_ingresses , args .prefix )
879+
880+ init_batch_progress_bars (args .num_threads )
859881 perform_ingress (failed_request_batchs , node_id , args .force_overwrite , config ["API_GATEWAY" ])
860- logger .info ("Reattempted ingress complete" )
861882 finally :
862883 close_batch_progress_bars ()
863884
0 commit comments