[H] Backport mirroring hotfixes (#7130) (#7122, #7123, PR #7124) (#7109, PR #7116)

Author: hannes-ucsc

deployments/prod/environment.py

@@ -1720,5 +1720,9 @@ def env() -> Mapping[str, Optional[str]]:
 
         'AZUL_ENABLE_VERBATIM_RELATIONS': '0',
 
-        'AZUL_MIRROR_BUCKET': 'humancellatlas'
+        'AZUL_ENABLE_MIRRORING': '1',
+
+        'AZUL_MIRROR_BUCKET': 'humancellatlas',
+
+        'AZUL_MIRRORING_CONCURRENCY': '128'
     }

lambdas/indexer/.chalice/config.json.template.py

@@ -27,7 +27,6 @@
     'environment_variables': config.lambda_env,
     'lambda_timeout': config.api_gateway_lambda_timeout,
     'lambda_memory_size': 128,
-    **chalice.vpc_lambda_config(app_name),
     'stages': {
         config.deployment_stage: {
             **chalice.private_api_stage_config(app_name),
@@ -37,36 +36,45 @@
                     'reserved_concurrency': config.contribution_concurrency(retry=False),
                     'lambda_memory_size': 256,
                     'lambda_timeout': config.contribution_lambda_timeout(retry=False),
+                    **chalice.vpc_lambda_config(app_name)
                 },
                 indexer.contribute_retry.name: {
                     'reserved_concurrency': config.contribution_concurrency(retry=True),
                     'lambda_memory_size': 4096,  # FIXME https://github.com/DataBiosphere/azul/issues/2902
-                    'lambda_timeout': config.contribution_lambda_timeout(retry=True)
+                    'lambda_timeout': config.contribution_lambda_timeout(retry=True),
+                    **chalice.vpc_lambda_config(app_name)
                 },
                 indexer.aggregate.name: {
                     'reserved_concurrency': config.aggregation_concurrency(retry=False),
                     'lambda_memory_size': 256,
-                    'lambda_timeout': config.aggregation_lambda_timeout(retry=False)
+                    'lambda_timeout': config.aggregation_lambda_timeout(retry=False),
+                    **chalice.vpc_lambda_config(app_name)
                 },
                 indexer.aggregate_retry.name: {
                     'reserved_concurrency': config.aggregation_concurrency(retry=True),
                     'lambda_memory_size': 6500,
-                    'lambda_timeout': config.aggregation_lambda_timeout(retry=True)
+                    'lambda_timeout': config.aggregation_lambda_timeout(retry=True),
+                    **chalice.vpc_lambda_config(app_name)
                 },
+                indexer.forward_alb_logs.name: chalice.vpc_lambda_config(app_name),
+                indexer.forward_s3_logs.name: chalice.vpc_lambda_config(app_name),
                 **(
                     {
                         indexer.mirror.name: {
                             'reserved_concurrency': config.mirroring_concurrency,
                             'lambda_memory_size': 512,
                             'lambda_timeout': config.mirror_lambda_timeout
+                            # No VPC for this function so as to avoid paying for
+                            # NAT Gateway traffic
                         },
                     }
                     if config.enable_mirroring else
                     {}
                 ),
                 indexer.update_health_cache.name: {
                     'lambda_memory_size': 128,
-                    'lambda_timeout': config.health_cache_lambda_timeout
+                    'lambda_timeout': config.health_cache_lambda_timeout,
+                    **chalice.vpc_lambda_config(app_name)
                 }
             }
         }

scripts/mirror.py

@@ -28,9 +28,6 @@
 def mirror_catalog(catalog: CatalogName, source_globs: set[str], wait: bool):
     azul = AzulClient()
     plugin = azul.repository_plugin(catalog)
-    assert azul.is_queue_empty(config.mirror_queue.name), R(
-        'A mirroring operation is already in progress. The current operation '
-        'must finish before another can begin.')
     fail_queue = config.mirror_queue.to_fail.name
     assert azul.is_queue_empty(fail_queue), R(
         'Cannot begin mirroring because a previous operation failed: '

src/azul/__init__.py

@@ -1807,7 +1807,7 @@ def enable_mirroring(self) -> bool:
         return self._boolean(self.environ['AZUL_ENABLE_MIRRORING'])
 
     @property
-    def external_mirror_bucket(self) -> str | None:
+    def mirror_bucket(self) -> str | None:
         return self.environ.get('AZUL_MIRROR_BUCKET')
 
 

src/azul/deployment.py

@@ -92,6 +92,7 @@
     )
     from mypy_boto3_s3 import (
         S3Client,
+        S3ServiceResource,
     )
     from mypy_boto3_secretsmanager import (
         SecretsManagerClient,
@@ -200,6 +201,10 @@ def region_name(self):
     def s3(self) -> 'S3Client':
         return self.client('s3', azul_logging=True)
 
+    @property
+    def s3_resource(self) -> 'S3ServiceResource':
+        return self.resource('s3', azul_logging=True)
+
     @property
     def securityhub(self) -> 'SecurityHubClient':
         return self.client('securityhub')
@@ -734,7 +739,7 @@ def monitoring_topic_name(self):
 
     @property
     def sqs_resource(self) -> 'SQSServiceResource':
-        return self.resource('sqs')
+        return self.resource('sqs', azul_logging=config.is_in_lambda)
 
     @_cache
     def sqs_queue(self, queue_name: str) -> 'Queue':

src/azul/health.py

@@ -214,7 +214,7 @@ def queues(self):
         """
         Returns information about the SQS queues used by the indexer.
         """
-        sqs = aws.resource('sqs', azul_logging=True)
+        sqs = aws.sqs_resource
         response: MutableJSON = {'up': True}
         for queue in config.all_queue_names:
             try:

src/azul/indexer/lambda_iam_policy.py

@@ -127,7 +127,7 @@
                     ],
                     'Resource': [
                         f'arn:aws:s3:::{resource}'
-                        for bucket in alist(aws.mirror_bucket, config.external_mirror_bucket)
+                        for bucket in alist(aws.mirror_bucket, config.mirror_bucket)
                         for resource in [bucket, f'{bucket}/*']
                     ]
                 }

src/azul/indexer/mirror_controller.py

@@ -57,7 +57,7 @@
 log = logging.getLogger(__name__)
 
 
-@attrs.frozen(auto_attribs=True, kw_only=True)
+@attrs.frozen(kw_only=True)
 class MirrorController(ActionController[MirrorAction]):
     schema_url_func: SchemaUrlFunc
 
@@ -107,31 +107,39 @@ def mirror_source(self, catalog: CatalogName, source_json: JSON):
         plugin = self.repository_plugin(catalog)
         source = plugin.source_ref_cls.from_json(source_json)
         source = plugin.partition_source_for_mirroring(catalog, source)
+        prefix = source.spec.prefix
+        log.info('Queueing %d partitions of source %r in catalog %r',
+                 prefix.num_partitions, str(source.spec), catalog)
 
-        def message(prefix: str) -> SQSMessage:
-            log.info('Mirroring files in partition %r of source %r from catalog %r',
-                     prefix, str(source.spec), catalog)
-            return self.mirror_partition_message(catalog, source, prefix)
+        def message(partition: str) -> SQSMessage:
+            log.debug('Queueing partition %r', partition)
+            return self.mirror_partition_message(catalog, source, partition)
 
-        messages = map(message, source.spec.prefix.partition_prefixes())
+        messages = map(message, prefix.partition_prefixes())
         self.client.queue_mirror_messages(messages)
 
-    def mirror_partition(self, catalog: CatalogName, source_json: JSON, prefix: str):
+    def mirror_partition(self,
+                         catalog: CatalogName,
+                         source_json: JSON,
+                         prefix: str
+                         ):
         plugin = self.repository_plugin(catalog)
         source = plugin.source_ref_cls.from_json(source_json)
         already_mirrored = self.service.list_info_objects(catalog, prefix)
+        files = plugin.list_files(source, prefix)
 
         def messages() -> Iterable[SQSMessage]:
-            for file in plugin.list_files(source, prefix):
+            for file in files:
                 info_key = self.service.info_object_key(file)
                 if info_key in already_mirrored:
-                    log.info('Not mirroring file %r because info object already exists at %r',
-                             file.uuid, info_key)
+                    log.debug('Not queueing file %r because info object already exists', file)
                 else:
-                    log.info('Mirroring file %r', file.uuid)
+                    log.debug('Queueing file %r', file)
                     yield self.mirror_file_message(catalog, source, file)
 
-        self.client.queue_mirror_messages(messages())
+        message_count = self.client.queue_mirror_messages(messages())
+        log.info('Queued %d/%d files in partition %r of source %r in catalog %r',
+                 message_count, len(files), prefix, str(source), catalog)
 
     def mirror_file(self,
                     catalog: CatalogName,
@@ -145,26 +153,36 @@ def mirror_file(self,
                                 and not config.deployment.is_unit_test
                                 and catalog not in config.integration_test_catalogs)
         if file_is_large and not deployment_is_stable:
-            log.info('Not mirroring file %r (%d bytes) to save cost',
-                     file.uuid, file.size)
+            log.info('Not mirroring file to save cost: %r', file)
         else:
             # Ensure we test with multiple parts on lower deployments
             part_size = FilePart.default_size if deployment_is_stable else FilePart.min_size
             if file.size <= part_size:
-                log.info('Mirroring file %r via standard upload', file.uuid)
+                log.info('Mirroring file via standard upload: %r', file)
                 self.service.mirror_file(catalog, file)
-                log.info('Successfully mirrored file %r via standard upload', file.uuid)
+                log.info('Successfully mirrored file via standard upload: %r', file)
             else:
-                log.info('Mirroring file %r via multi-part upload', file.uuid)
+                log.info('Mirroring file via multi-part upload: %r', file)
                 _, digest_type = file.digest()
                 hasher = get_resumable_hasher(digest_type)
                 upload_id = self.service.begin_mirroring_file(catalog, file)
                 first_part = FilePart.first(file, part_size)
-                etag = self.service.mirror_file_part(catalog, file, first_part, upload_id, hasher)
+                log.info('Uploading part #%d of file %r', first_part.index, file)
+                etag = self.service.mirror_file_part(catalog,
+                                                     file,
+                                                     first_part,
+                                                     upload_id,
+                                                     hasher)
                 next_part = first_part.next(file)
                 assert next_part is not None
-                messages = [self.mirror_part_message(catalog, file, next_part, upload_id, [etag], hasher)]
-                self.client.queue_mirror_messages(messages)
+                log.info('Queueing part #%d of file %r', next_part.index, file)
+                message = self.mirror_part_message(catalog,
+                                                   file,
+                                                   next_part,
+                                                   upload_id,
+                                                   [etag],
+                                                   hasher)
+                self.client.queue_mirror_messages([message])
 
     def mirror_file_part(self,
                          catalog: CatalogName,
@@ -177,17 +195,19 @@ def mirror_file_part(self,
         file = self.load_file(catalog, file_json)
         part = FilePart.from_json(part_json)
         hasher = hasher_from_str(hasher_data)
+        log.info('Uploading part #%d of file %r', part.index, file)
         etag = self.service.mirror_file_part(catalog, file, part, upload_id, hasher)
         etags = [*etags, etag]
         next_part = part.next(file)
         if next_part is None:
-            log.info('File %r fully uploaded in %d parts', file.uuid, len(etags))
+            log.info('File fully uploaded in %d parts: %r', len(etags), file)
             message = self.finalize_file_message(catalog,
                                                  file,
                                                  upload_id,
                                                  etags,
                                                  hasher)
         else:
+            log.info('Queueing part #%d of file %r', next_part.index, file)
             message = self.mirror_part_message(catalog,
                                                file,
                                                next_part,
@@ -211,7 +231,7 @@ def finalize_file(self,
                                            upload_id=upload_id,
                                            etags=etags,
                                            hasher=hasher)
-        log.info('Successfully mirrored file %r via multi-part upload', file.uuid)
+        log.info('Successfully mirrored file via multi-part upload: %r', file)
 
     def load_file(self, catalog: CatalogName, file: JSON) -> File:
         return self.client.metadata_plugin(catalog).file_class.from_json(file)