From ec9e2523b6bd05f346c883c231a118bf97e92a4e Mon Sep 17 00:00:00 2001 From: michaelkedar Date: Mon, 25 May 2026 06:54:44 +0000 Subject: [PATCH] feat: create new work pools for reimports & cves --- .../gke-workers/base/importer-reconciler.yaml | 2 + .../gke-workers/base/importer.yaml | 2 + .../gke-workers/base/kustomization.yaml | 2 + .../gke-workers/base/workers-cves.yaml | 84 +++++++++++++++++++ .../gke-workers/base/workers-reimport.yaml | 84 +++++++++++++++++++ .../oss-vdb-test/kustomization.yaml | 2 + .../oss-vdb-test/workers-cves.yaml | 20 +++++ .../oss-vdb-test/workers-reimport.yaml | 20 +++++ .../environments/oss-vdb/kustomization.yaml | 2 + .../environments/oss-vdb/workers-cves.yaml | 14 ++++ .../oss-vdb/workers-reimport.yaml | 16 ++++ .../environments/oss-vdb-test/main.tf | 5 ++ .../terraform/environments/oss-vdb/main.tf | 5 ++ .../terraform/modules/osv/pubsub_tasks.tf | 24 ++++++ deployment/terraform/modules/osv/variables.tf | 6 ++ gcp/workers/recoverer/recoverer.py | 4 +- 16 files changed, 290 insertions(+), 2 deletions(-) create mode 100644 deployment/clouddeploy/gke-workers/base/workers-cves.yaml create mode 100644 deployment/clouddeploy/gke-workers/base/workers-reimport.yaml create mode 100644 deployment/clouddeploy/gke-workers/environments/oss-vdb-test/workers-cves.yaml create mode 100644 deployment/clouddeploy/gke-workers/environments/oss-vdb-test/workers-reimport.yaml create mode 100644 deployment/clouddeploy/gke-workers/environments/oss-vdb/workers-cves.yaml create mode 100644 deployment/clouddeploy/gke-workers/environments/oss-vdb/workers-reimport.yaml diff --git a/deployment/clouddeploy/gke-workers/base/importer-reconciler.yaml b/deployment/clouddeploy/gke-workers/base/importer-reconciler.yaml index 806fd140a0f..176ef4bf5fd 100644 --- a/deployment/clouddeploy/gke-workers/base/importer-reconciler.yaml +++ b/deployment/clouddeploy/gke-workers/base/importer-reconciler.yaml @@ -27,6 +27,8 @@ spec: value: "1.0" - name: TRACE_SAMPLE_RATE # for the individual vulnerability entries value: "0.05" + - name: REIMPORT_TASK_POOL + value: "reimport" resources: requests: cpu: "1" diff --git a/deployment/clouddeploy/gke-workers/base/importer.yaml b/deployment/clouddeploy/gke-workers/base/importer.yaml index 034d3f20c9f..b1da9d729da 100644 --- a/deployment/clouddeploy/gke-workers/base/importer.yaml +++ b/deployment/clouddeploy/gke-workers/base/importer.yaml @@ -30,6 +30,8 @@ spec: value: "1.0" - name: TRACE_SAMPLE_RATE # for the individual vulnerability entries value: "0.05" + - name: REIMPORT_TASK_POOL + value: "reimport" resources: requests: cpu: "1" diff --git a/deployment/clouddeploy/gke-workers/base/kustomization.yaml b/deployment/clouddeploy/gke-workers/base/kustomization.yaml index 4d21838b682..d58e81525ee 100644 --- a/deployment/clouddeploy/gke-workers/base/kustomization.yaml +++ b/deployment/clouddeploy/gke-workers/base/kustomization.yaml @@ -1,6 +1,8 @@ resources: - workers.yaml - scaler.yaml +- workers-cves.yaml +- workers-reimport.yaml - importer.yaml - importer-deleter.yaml - importer-reconciler.yaml diff --git a/deployment/clouddeploy/gke-workers/base/workers-cves.yaml b/deployment/clouddeploy/gke-workers/base/workers-cves.yaml new file mode 100644 index 00000000000..2cef068fed6 --- /dev/null +++ b/deployment/clouddeploy/gke-workers/base/workers-cves.yaml @@ -0,0 +1,84 @@ +# Copyright 2026 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. + +apiVersion: autoscaling/v2 +kind: HorizontalPodAutoscaler +metadata: + name: pubsub-cves +spec: + minReplicas: 0 + maxReplicas: 100 + metrics: + - external: + metric: + name: pubsub.googleapis.com|subscription|num_undelivered_messages + selector: + matchLabels: + resource.labels.subscription_id: cves + target: + type: AverageValue + averageValue: 10 # each worker can handle 10 tasks + type: External + scaleTargetRef: + apiVersion: apps/v1 + kind: Deployment + name: workers-cves + behavior: + scaleUp: + stabilizationWindowSeconds: 0 + policies: + - type: Percent + value: 10 + periodSeconds: 300 + - type: Pods + value: 20 + periodSeconds: 300 + selectPolicy: Max +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: workers-cves +spec: + replicas: 0 + selector: + matchLabels: + name: workers-cves + template: + metadata: + labels: + name: workers-cves + spec: + containers: + - name: workers-cves + image: worker + imagePullPolicy: Always + env: + - name: GITTER_HOST + value: http://gitter-service:8888 + - name: PUBSUB_SUBSCRIPTION + value: cves + - name: DATASTORE_DATABASE_ID + value: "" # default + - name: FAILED_TASKS_TOPIC + value: failed-tasks + - name: NOTIFY_PYPI + value: "false" + resources: + requests: + cpu: "0.9" + memory: "1.2Gi" + limits: + cpu: "1.5" + memory: "1.3Gi" diff --git a/deployment/clouddeploy/gke-workers/base/workers-reimport.yaml b/deployment/clouddeploy/gke-workers/base/workers-reimport.yaml new file mode 100644 index 00000000000..8701661c485 --- /dev/null +++ b/deployment/clouddeploy/gke-workers/base/workers-reimport.yaml @@ -0,0 +1,84 @@ +# Copyright 2026 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. + +apiVersion: autoscaling/v2 +kind: HorizontalPodAutoscaler +metadata: + name: pubsub-reimport +spec: + minReplicas: 0 + maxReplicas: 100 + metrics: + - external: + metric: + name: pubsub.googleapis.com|subscription|num_undelivered_messages + selector: + matchLabels: + resource.labels.subscription_id: reimport + target: + type: AverageValue + averageValue: 10 # each worker can handle 10 tasks + type: External + scaleTargetRef: + apiVersion: apps/v1 + kind: Deployment + name: workers-reimport + behavior: + scaleUp: + stabilizationWindowSeconds: 0 + policies: + - type: Percent + value: 10 + periodSeconds: 300 + - type: Pods + value: 20 + periodSeconds: 300 + selectPolicy: Max +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: workers-reimport +spec: + replicas: 0 + selector: + matchLabels: + name: workers-reimport + template: + metadata: + labels: + name: workers-reimport + spec: + containers: + - name: workers-reimport + image: worker + imagePullPolicy: Always + env: + - name: GITTER_HOST + value: http://gitter-service:8888 + - name: PUBSUB_SUBSCRIPTION + value: reimport + - name: DATASTORE_DATABASE_ID + value: "" # default + - name: FAILED_TASKS_TOPIC + value: failed-tasks + - name: NOTIFY_PYPI + value: "false" + resources: + requests: + cpu: "0.9" + memory: "1.2Gi" + limits: + cpu: "1.5" + memory: "1.3Gi" diff --git a/deployment/clouddeploy/gke-workers/environments/oss-vdb-test/kustomization.yaml b/deployment/clouddeploy/gke-workers/environments/oss-vdb-test/kustomization.yaml index 91165997cf7..7ec61ca37bb 100644 --- a/deployment/clouddeploy/gke-workers/environments/oss-vdb-test/kustomization.yaml +++ b/deployment/clouddeploy/gke-workers/environments/oss-vdb-test/kustomization.yaml @@ -5,6 +5,8 @@ resources: patches: - path: workers.yaml - path: scaler.yaml +- path: workers-cves.yaml +- path: workers-reimport.yaml - path: importer.yaml - path: importer-deleter.yaml - path: importer-reconciler.yaml diff --git a/deployment/clouddeploy/gke-workers/environments/oss-vdb-test/workers-cves.yaml b/deployment/clouddeploy/gke-workers/environments/oss-vdb-test/workers-cves.yaml new file mode 100644 index 00000000000..4120fbb8ac2 --- /dev/null +++ b/deployment/clouddeploy/gke-workers/environments/oss-vdb-test/workers-cves.yaml @@ -0,0 +1,20 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: workers-cves +spec: + template: + spec: + tolerations: + - key: workloadType + operator: Equal + value: worker-pool + nodeSelector: + workloadType: worker-pool + containers: + - name: workers-cves + env: + - name: GOOGLE_CLOUD_PROJECT + value: oss-vdb-test + - name: OSV_VULNERABILITIES_BUCKET + value: osv-test-vulnerabilities diff --git a/deployment/clouddeploy/gke-workers/environments/oss-vdb-test/workers-reimport.yaml b/deployment/clouddeploy/gke-workers/environments/oss-vdb-test/workers-reimport.yaml new file mode 100644 index 00000000000..6b5668aeb73 --- /dev/null +++ b/deployment/clouddeploy/gke-workers/environments/oss-vdb-test/workers-reimport.yaml @@ -0,0 +1,20 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: workers-reimport +spec: + template: + spec: + tolerations: + - key: workloadType + operator: Equal + value: worker-pool + nodeSelector: + workloadType: worker-pool + containers: + - name: workers-reimport + env: + - name: GOOGLE_CLOUD_PROJECT + value: oss-vdb-test + - name: OSV_VULNERABILITIES_BUCKET + value: osv-test-vulnerabilities diff --git a/deployment/clouddeploy/gke-workers/environments/oss-vdb/kustomization.yaml b/deployment/clouddeploy/gke-workers/environments/oss-vdb/kustomization.yaml index c95f9ffc101..d4a87f2fd2b 100644 --- a/deployment/clouddeploy/gke-workers/environments/oss-vdb/kustomization.yaml +++ b/deployment/clouddeploy/gke-workers/environments/oss-vdb/kustomization.yaml @@ -2,6 +2,8 @@ resources: - ../../base patches: - path: workers.yaml +- path: workers-cves.yaml +- path: workers-reimport.yaml - path: importer.yaml - path: importer-deleter.yaml - path: importer-reconciler.yaml diff --git a/deployment/clouddeploy/gke-workers/environments/oss-vdb/workers-cves.yaml b/deployment/clouddeploy/gke-workers/environments/oss-vdb/workers-cves.yaml new file mode 100644 index 00000000000..6a908a28835 --- /dev/null +++ b/deployment/clouddeploy/gke-workers/environments/oss-vdb/workers-cves.yaml @@ -0,0 +1,14 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: workers-cves +spec: + template: + spec: + containers: + - name: workers-cves + env: + - name: GOOGLE_CLOUD_PROJECT + value: oss-vdb + - name: OSV_VULNERABILITIES_BUCKET + value: osv-vulnerabilities diff --git a/deployment/clouddeploy/gke-workers/environments/oss-vdb/workers-reimport.yaml b/deployment/clouddeploy/gke-workers/environments/oss-vdb/workers-reimport.yaml new file mode 100644 index 00000000000..4962c338db1 --- /dev/null +++ b/deployment/clouddeploy/gke-workers/environments/oss-vdb/workers-reimport.yaml @@ -0,0 +1,16 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: workers-reimport +spec: + template: + spec: + containers: + - name: workers-reimport + env: + - name: GOOGLE_CLOUD_PROJECT + value: oss-vdb + - name: OSV_VULNERABILITIES_BUCKET + value: osv-vulnerabilities + - name: NOTIFY_PYPI + value: "true" diff --git a/deployment/terraform/environments/oss-vdb-test/main.tf b/deployment/terraform/environments/oss-vdb-test/main.tf index ed67fd17c25..348a09f89a7 100644 --- a/deployment/terraform/environments/oss-vdb-test/main.tf +++ b/deployment/terraform/environments/oss-vdb-test/main.tf @@ -46,6 +46,11 @@ module "osv_test" { website_domain = "test.osv.dev" api_url = "api.test.osv.dev" esp_version = "2.55.1" + + extra_work_pools = [ + "reimport", + "cves", + ] } module "k8s_cron_alert" { diff --git a/deployment/terraform/environments/oss-vdb/main.tf b/deployment/terraform/environments/oss-vdb/main.tf index 183ba84a7ae..b5f42286f07 100644 --- a/deployment/terraform/environments/oss-vdb/main.tf +++ b/deployment/terraform/environments/oss-vdb/main.tf @@ -46,6 +46,11 @@ module "osv" { website_domain = "osv.dev" api_url = "api.osv.dev" esp_version = "2.55.1" + + extra_work_pools = [ + "reimport", + "cves", + ] } module "oss_fuzz" { diff --git a/deployment/terraform/modules/osv/pubsub_tasks.tf b/deployment/terraform/modules/osv/pubsub_tasks.tf index 55089f4008a..254adcedf91 100644 --- a/deployment/terraform/modules/osv/pubsub_tasks.tf +++ b/deployment/terraform/modules/osv/pubsub_tasks.tf @@ -41,6 +41,30 @@ resource "google_pubsub_subscription" "default_work" { filter = "attributes.work_pool = \"default\"" } +resource "google_pubsub_subscription" "work_pools" { + for_each = toset(var.extra_work_pools) + project = var.project_id + name = each.value + topic = google_pubsub_topic.tasks.id + message_retention_duration = "604800s" + ack_deadline_seconds = 600 + + dead_letter_policy { + dead_letter_topic = google_pubsub_topic.failed_tasks.id + max_delivery_attempts = 5 + } + + expiration_policy { + ttl = "" # never expires + } + + labels = { + goog-dm = "pubsub" + } + + filter = "attributes.work_pool = \"${each.value}\"" +} + resource "google_pubsub_topic" "pypi_bridge" { project = var.project_id name = "pypi-bridge" diff --git a/deployment/terraform/modules/osv/variables.tf b/deployment/terraform/modules/osv/variables.tf index 4c0e9f753c1..e663656cccc 100644 --- a/deployment/terraform/modules/osv/variables.tf +++ b/deployment/terraform/modules/osv/variables.tf @@ -67,3 +67,9 @@ variable "website_domain" { type = string description = "Domain to serve the OSV website on. Domain ownership and DNS settings must be manually configured." } + +variable "extra_work_pools" { + type = list(string) + description = "Additional Pub/Sub worker pool subscriptions to create" + default = [] +} diff --git a/gcp/workers/recoverer/recoverer.py b/gcp/workers/recoverer/recoverer.py index ff2896bdceb..4e335eb9d14 100644 --- a/gcp/workers/recoverer/recoverer.py +++ b/gcp/workers/recoverer/recoverer.py @@ -137,8 +137,8 @@ def handle_gcs_missing(message: pubsub_v1.types.PubsubMessage) -> bool: deleted='false', skip_hash_check='true', req_timestamp=str(int(time.time())), - work_pool='default', - ) # TODO(michaelkedar): replace with reimport pool when created + work_pool='reimport', + ) return True