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,13 @@ def request_batch_for_ingress(request_batch, batch_index, node_id, force_overwri
409414 "x-amz-docs-region" : api_gateway_region ,
410415 }
411416
417+ # Simulate a random failure for the batch request if configured to do so
418+ simulate_batch_request_failure (kwargs ["mock_requests" ], api_gateway_url .split ("?" )[0 ])
419+
412420 response = requests .post (
413421 api_gateway_url , params = params , data = json .dumps (request_batch ), headers = headers , timeout = 600
414422 )
423+
415424 elapsed_time = time .time () - start_time
416425
417426 # Ingress request successful
@@ -428,10 +437,11 @@ def request_batch_for_ingress(request_batch, batch_index, node_id, force_overwri
428437@backoff .on_exception (
429438 backoff .expo ,
430439 Exception ,
431- max_time = 120 ,
440+ max_time = 10 , # TODO: for testing only, revert to 120
432441 logger = "ingress_file_to_s3" ,
433442)
434- def ingress_file_to_s3 (ingress_response , batch_index , batch_pbar ):
443+ @requests_mock .Mocker (kw = "mock_requests" , real_http = True )
444+ def ingress_file_to_s3 (ingress_response , batch_index , batch_pbar , ** kwargs ):
435445 """
436446 Copies the local file path using the pre-signed S3 URL returned from the
437447 Ingress Lambda App.
@@ -487,6 +497,9 @@ def ingress_file_to_s3(ingress_response, batch_index, batch_pbar):
487497 batch_pbar , total = os .stat (ingress_path ).st_size , filename = os .path .basename (ingress_path )
488498 )
489499
500+ # Simulate a random failure for the S3 ingress request if configured to do so
501+ simulate_ingress_failure (kwargs ["mock_requests" ], s3_ingress_url .split ("?" )[0 ])
502+
490503 with open (ingress_path , "rb" ) as infile :
491504 # Wrap file I/O with our upload bar to automatically track file upload progress
492505 wrapped_file = CallbackIOWrapper (upload_pbar .update , infile , "read" )
@@ -517,10 +530,11 @@ def ingress_file_to_s3(ingress_response, batch_index, batch_pbar):
517530@backoff .on_exception (
518531 backoff .expo ,
519532 Exception ,
520- max_time = 120 ,
533+ max_time = 10 , # TODO: for testing only, revert to 120
521534 logger = "ingress_multipart_file_to_s3" ,
522535)
523- def ingress_multipart_file_to_s3 (ingress_response , batch_index , batch_pbar ):
536+ @requests_mock .Mocker (kw = "mock_requests" , real_http = True )
537+ def ingress_multipart_file_to_s3 (ingress_response , batch_index , batch_pbar , ** kwargs ):
524538 """
525539 Performs an ingress request for a file that is too large to be uploaded
526540 in a single request. The file is instead uploaded in multiple parts using
@@ -582,6 +596,9 @@ def ingress_multipart_file_to_s3(ingress_response, batch_index, batch_pbar):
582596 upload_pbar , f"{ os .path .basename (ingress_path )} (Part { part_number } /{ len (s3_ingress_urls )} )"
583597 )
584598
599+ # Simulate a random failure for the S3 ingress request if configured to do so
600+ simulate_ingress_failure (kwargs ["mock_requests" ], s3_ingress_url .split ("?" )[0 ])
601+
585602 # Submit a single chunk to AWS
586603 response = requests .put (s3_ingress_url , data = next (chunk_iterator ))
587604 response .raise_for_status ()
@@ -848,16 +865,22 @@ def main(args):
848865 try :
849866 init_batch_progress_bars (args .num_threads )
850867 perform_ingress (request_batchs , node_id , args .force_overwrite , config ["API_GATEWAY" ])
868+ finally :
869+ close_batch_progress_bars ()
851870
852- logger .info ("All batches processed" )
871+ logger .info ("All batches processed" )
853872
873+ try :
854874 if len (SUMMARY_TABLE ["failed" ]) > 0 :
875+ logger .info ("----------------------------------------" )
855876 logger .info ("Reattempting ingress for failed files..." )
877+
856878 failed_ingresses = SUMMARY_TABLE ["failed" ]
857879 batched_failed_ingresses = list (batched (failed_ingresses , batch_size ))
858880 failed_request_batchs = prepare_batches (batched_failed_ingresses , args .prefix )
881+
882+ init_batch_progress_bars (args .num_threads )
859883 perform_ingress (failed_request_batchs , node_id , args .force_overwrite , config ["API_GATEWAY" ])
860- logger .info ("Reattempted ingress complete" )
861884 finally :
862885 close_batch_progress_bars ()
863886
0 commit comments