Add parallel exporter

deployment/clouddeploy/gke-workers/base/exporter.yaml

@@ -5,10 +5,19 @@ metadata:
 spec:
   schedule: "0 */1 * * *"
   concurrencyPolicy: Forbid
+  # If the previous job overruns by more than a minute,
+  # Don't try to immediately queue up the next job
+  startingDeadlineSeconds: 60
   jobTemplate:
     spec:
       template:
         spec:
+          tolerations:
+          - key: workloadType
+            operator: Equal
+            value: highend
+          nodeSelector:
+            cloud.google.com/gke-nodepool: highend
           containers:
           - name: exporter
             image: exporter
@@ -18,13 +27,13 @@ spec:
                 name: "ssd"
             resources:
               requests:
-                cpu: 1
-                memory: "6G"
-              limits:
-                cpu: 1
+                cpu: "6"
                 memory: "8G"
+              limits:
+                cpu: "6"
+                memory: "10G"
           restartPolicy: OnFailure
           volumes:
             - name: "ssd"
-              hostPath:
-                path: "/mnt/disks/ssd0"
+              emptyDir:
+                sizeLimit: 60Gi

docker/exporter/Dockerfile

@@ -15,5 +15,7 @@
 FROM gcr.io/oss-vdb/worker
 
 COPY exporter.py /usr/local/bin
+COPY export_runner.py /usr/local/bin
 RUN chmod 755 /usr/local/bin/exporter.py
-ENTRYPOINT ["exporter.py"]
+RUN chmod 755 /usr/local/bin/export_runner.py
+ENTRYPOINT ["export_runner.py"]

docker/exporter/export_runner.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""OSV Exporter."""
+import argparse
+import concurrent.futures
+import logging
+import os
+import subprocess
+
+from google.cloud import ndb
+
+import osv
+import osv.logs
+
+DEFAULT_WORK_DIR = '/work'
+DEFAULT_EXPORT_BUCKET = 'osv-vulnerabilities'
+DEFAULT_EXPORT_PROCESSES = 7
+
+
+def main():
+  parser = argparse.ArgumentParser(description='Exporter')
+  parser.add_argument(
+      '--work_dir', help='Working directory', default=DEFAULT_WORK_DIR)
+  parser.add_argument(
+      '--bucket',
+      help='Bucket name to export to',
+      default=DEFAULT_EXPORT_BUCKET)
+  parser.add_argument(
+      '--processes',
+      help='Maximum number of parallel exports, default to number of cpu cores',
+      # If 0 or None, use the DEFAULT_EXPORT_PROCESSES value
+      default=os.cpu_count() or DEFAULT_EXPORT_PROCESSES)
+  args = parser.parse_args()
+
+  tmp_dir = os.path.join(args.work_dir, 'tmp')
+  os.makedirs(tmp_dir, exist_ok=True)
+  os.environ['TMPDIR'] = tmp_dir
+
+  query = osv.Bug.query(projection=[osv.Bug.ecosystem], distinct=True)
+  ecosystems = [bug.ecosystem[0] for bug in query if bug.ecosystem] + ["list"]
+
+  with concurrent.futures.ThreadPoolExecutor(
+      max_workers=args.processes) as executor:
+    for eco in ecosystems:
+      executor.submit(spawn_ecosystem_exporter, args.work_dir, args.bucket, eco)
+
+
+def spawn_ecosystem_exporter(work_dir: str, bucket: str, eco: str):
+  """
+  Spawns the ecosystem specific exporter.
+  """
+  logging.info("Starting export of ecosystem: %s", eco)
+  proc = subprocess.Popen([
+      "exporter.py", "--work_dir", work_dir, "--bucket", bucket, "--ecosystem",
+      eco
+  ])
+  return_code = proc.wait()
+  if return_code != 0:
+    logging.error("Export of %s failed with Exit Code: %d", eco, return_code)
+
+
+if __name__ == '__main__':
+  _ndb_client = ndb.Client()
+  osv.logs.setup_gcp_logging('exporter-runner')
+  with _ndb_client.context():
+    main()

docker/exporter/exporter.py

@@ -17,7 +17,6 @@
 import concurrent.futures
 import logging
 import os
-import shutil
 import tempfile
 import zipfile
 from typing import List
@@ -39,21 +38,21 @@
 class Exporter:
   """Exporter."""
 
-  def __init__(self, work_dir, export_bucket):
+  def __init__(self, work_dir, export_bucket, ecosystem):
     self._work_dir = work_dir
     self._export_bucket = export_bucket
+    self._ecosystem = ecosystem
 
   def run(self):
     """Run exporter."""
-    query = osv.Bug.query(projection=[osv.Bug.ecosystem], distinct=True)
-    ecosystems = [bug.ecosystem[0] for bug in query if bug.ecosystem]
-
-    for ecosystem in ecosystems:
+    if self._ecosystem == "list":
+      query = osv.Bug.query(projection=[osv.Bug.ecosystem], distinct=True)
+      ecosystems = [bug.ecosystem[0] for bug in query if bug.ecosystem]
       with tempfile.TemporaryDirectory() as tmp_dir:
-        self._export_ecosystem_to_bucket(ecosystem, tmp_dir)
-
-    with tempfile.TemporaryDirectory() as tmp_dir:
-      self._export_ecosystem_list_to_bucket(ecosystems, tmp_dir)
+        self._export_ecosystem_list_to_bucket(ecosystems, tmp_dir)
+    else:
+      with tempfile.TemporaryDirectory() as tmp_dir:
+        self._export_ecosystem_to_bucket(self._ecosystem, tmp_dir)
 
   def upload_single(self, bucket, source_path, target_path):
     """Upload a single file to a GCS bucket."""
@@ -139,19 +138,17 @@ def main():
       '--bucket',
       help='Bucket name to export to',
       default=DEFAULT_EXPORT_BUCKET)
+  parser.add_argument(
+      '--ecosystem',
+      required=True,
+      help='Ecosystem to upload, pass the value "list" ' +
+      'to export the ecosystem.txt file')
   args = parser.parse_args()
 
   tmp_dir = os.path.join(args.work_dir, 'tmp')
-  # Temp files are on the persistent local SSD,
-  # and they do not get removed when GKE sends a SIGTERM to stop the pod.
-  # Manually clear the tmp_dir folder of any leftover files
-  # TODO(michaelkedar): use an ephemeral disk for temp storage.
-  if os.path.exists(tmp_dir):
-    shutil.rmtree(tmp_dir)
-  os.makedirs(tmp_dir, exist_ok=True)
   os.environ['TMPDIR'] = tmp_dir
 
-  exporter = Exporter(args.work_dir, args.bucket)
+  exporter = Exporter(args.work_dir, args.bucket, args.ecosystem)
   exporter.run()