perf(vanir): optimize signature job concurrency and node resource allocation

deployment/clouddeploy/gke-workers/base/extra/vanir-signatures.yaml

@@ -18,11 +18,11 @@ spec:
             imagePullPolicy: Always
             resources:
               requests:
-                cpu: "4"
-                memory: "63G"
+                cpu: "28"
+                memory: "180Gi"
               limits:
-                cpu: "4"
-                memory: "64G"
+                cpu: "28"
+                memory: "220Gi"
           nodeSelector:
             cloud.google.com/gke-nodepool: highend
           tolerations:

gcp/workers/vanir_signatures/vanir_signatures.py

@@ -232,14 +232,14 @@ def main():
   parser.add_argument(
       '--batch-size',
       type=int,
-      default=500,
+      default=100,
       help='Number of vulnerabilities to process in each batch.')
   parser.add_argument(
       '--max-workers',
       type=int,
-      default=4,
+      default=10,
       help=('Maximum number of parallel workers. Note that total threads '
-            'spawned will be max_workers * max_workers (default 16).'))
+            'spawned will be max_workers * max_workers (default 100).'))
   parser.add_argument(
       '--dry-run', action='store_true', help='Perform a dry run.')
   parser.add_argument(
@@ -284,55 +284,55 @@ def main():
 
   # Note that total threads spawned will be max_workers * max_workers (one pool
   # for batches, one pool within each batch for GCS fetches).
-  with tempfile.TemporaryDirectory() as shared_temp_dir:
-    with futures.ThreadPoolExecutor(max_workers=args.max_workers) as executor:
-
-      def process_with_context(batch):
-        with ndb.Client().context():
-          return process_batch(
-              batch,
-              shared_temp_dir,
-              dry_run=args.dry_run,
-              max_workers=args.max_workers)
-
-      future_to_batch = {}
-      current_batch = []
-
-      logging.info('Streaming vulnerabilities for processing.')
-      for key in query.iter(keys_only=True):
-        current_batch.append(key.id())
-        if len(current_batch) >= args.batch_size:
-          f = executor.submit(process_with_context, current_batch)
-          future_to_batch[f] = current_batch
-          current_batch = []
-
-      # Also add IDs from the retry list
-      if retry_list_data and retry_list_data.value:
-        retry_ids = list(set(retry_list_data.value))
-        logging.info('Adding %d IDs from retry list.', len(retry_ids))
-        for i in range(0, len(retry_ids), args.batch_size):
-          batch = retry_ids[i:i + args.batch_size]
-          f = executor.submit(process_with_context, batch)
-          future_to_batch[f] = batch
-
-      if current_batch:
+  with futures.ThreadPoolExecutor(max_workers=args.max_workers) as executor:
+
+    def process_with_context(batch):
+      with ndb.Client().context(), tempfile.TemporaryDirectory(
+      ) as batch_temp_dir:
+        return process_batch(
+            batch,
+            batch_temp_dir,
+            dry_run=args.dry_run,
+            max_workers=args.max_workers)
+
+    future_to_batch = {}
+    current_batch = []
+
+    logging.info('Streaming vulnerabilities for processing.')
+    for key in query.iter(keys_only=True):
+      current_batch.append(key.id())
+      if len(current_batch) >= args.batch_size:
         f = executor.submit(process_with_context, current_batch)
         future_to_batch[f] = current_batch
-
-      if not future_to_batch:
-        logging.info('No modified vulnerabilities found.')
-      else:
-        logging.info('Processing %d batches of vulnerabilities.',
-                     len(future_to_batch))
-        for future in futures.as_completed(future_to_batch):
-          try:
-            generated, failed_ids = future.result()
-            total_generated_count += generated
-            all_failed_ids.extend(failed_ids)
-            total_processed_count += len(future_to_batch[future])
-          except Exception as e:
-            logging.exception(
-                'Failed to process a batch of vulnerabilities: %s', e)
+        current_batch = []
+
+    # Also add IDs from the retry list
+    if retry_list_data and retry_list_data.value:
+      retry_ids = list(set(retry_list_data.value))
+      logging.info('Adding %d IDs from retry list.', len(retry_ids))
+      for i in range(0, len(retry_ids), args.batch_size):
+        batch = retry_ids[i:i + args.batch_size]
+        f = executor.submit(process_with_context, batch)
+        future_to_batch[f] = batch
+
+    if current_batch:
+      f = executor.submit(process_with_context, current_batch)
+      future_to_batch[f] = current_batch
+
+    if not future_to_batch:
+      logging.info('No modified vulnerabilities found.')
+    else:
+      logging.info('Processing %d batches of vulnerabilities.',
+                   len(future_to_batch))
+      for future in futures.as_completed(future_to_batch):
+        try:
+          generated, failed_ids = future.result()
+          total_generated_count += generated
+          all_failed_ids.extend(failed_ids)
+          total_processed_count += len(future_to_batch[future])
+        except Exception as e:
+          logging.exception('Failed to process a batch of vulnerabilities: %s',
+                            e)
 
   logging.info('Processed %d vulnerabilities, generated %d new signatures.',
                total_processed_count, total_generated_count)