pts: variant index in polars

orchestration/src/orchestration/dags/config/clusters.yaml

@@ -28,33 +28,6 @@ clusters:
     init_actions_uris:
       - '{{pts_init_action_uri}}'
 
-  # The variant index (PySpark, ported from Gentropy) joins the VEP output with gnomAD and
-  # FoldX. It is CPU-bound (parsing ~110 GB of VEP JSON): 16 n1-highmem-16 workers finished in
-  # ~6 min at 17 GB peak executor heap, ~25 GB disk per node and no spill, so 4 workers take
-  # proportionally longer for the same compute. A fixed primary pool, because the shared
-  # autoscaling policy pins primaries at 2 and only adds spot secondaries, which lose shuffle
-  # data on preemption. pd-ssd throughput scales with size, so worker disks stay at 256 GB.
-  # The driver collects nothing, so the master stays small.
-  pts_variant:
-    image_version: '{{dataproc_image_version_pts}}'
-    num_masters: 1
-    num_workers: 4
-    worker_machine_type: n1-highmem-16
-    master_machine_type: n1-standard-8
-    master_disk_size: 128
-    worker_disk_size: 256
-    worker_disk_type: pd-ssd
-    idle_delete_ttl: 3600
-    internal_ip_only: false
-    metadata:
-      PTS_REF: 'pts%40v{{pts_version}}'
-    properties:
-      'spark:spark.sql.adaptive.enabled': 'true'
-      'spark:spark.sql.parquet.compression.codec': 'zstd'
-      'spark:spark.sql.ansi.enabled': 'false'
-    init_actions_uris:
-      - '{{pts_init_action_uri}}'
-
   gentropy:
     image_version: '{{dataproc_image_version_gentropy}}'
     num_masters: 1

orchestration/src/orchestration/dags/config/ppp/pts.override.yaml

@@ -2,8 +2,8 @@
 steps:
   # PPP adds the curated LoF annotations to the variant index.
   variant:
-    - name: pyspark variant index
-      pyspark: variant
+    - name: transform variant index
+      transformer: variant
       source:
         vep: intermediate/variant_annotation
         variant_annotations:
@@ -15,7 +15,6 @@ steps:
         output: output/variant
       settings:
         hash_threshold: 300
-        output_partitions: 25
 
   # PPP validates its own study and credible-set inputs instead of the platform molQTL sources.
   study:

orchestration/tests/test_lof_curation_config.py

@@ -20,7 +20,7 @@ def test_lof_ingestion_and_consumer_are_enabled_together(monkeypatch, ppp):
     assert 'pts_variant_ppp' not in variant['depends_on']
     assert 'pts_variant_ppp' in variant['depends_on_ppp']
     [task] = config.pts.config['steps']['variant']
-    assert task['pyspark'] == 'variant'
+    assert task['transformer'] == 'variant'
     assert task['source']['amino_acid_annotations'] == ['intermediate/variant/foldx']
     annotations = task['source']['variant_annotations']
     assert 'input/variant/gnomad' in annotations

orchestration/tests/test_variant_step.py

@@ -1,19 +1,17 @@
-"""The PySpark variant index runs on its own Dataproc cluster, after FoldX on a VM."""
+"""The variant index runs on a VM, after FoldX."""
 
 from orchestration.dags.config.unified_pipeline import UnifiedPipelineConfig
-from orchestration.models.pts_step import PTSDataprocStep, pts_step_from_config
+from orchestration.models.pts_step import pts_step_from_config
 
 
-def test_variant_runs_on_its_dataproc_cluster() -> None:
+def test_variant_and_foldx_run_on_vms() -> None:
     config = UnifiedPipelineConfig()
-    step = pts_step_from_config('pts_variant', config)
-    assert isinstance(step, PTSDataprocStep)
-    assert step.cluster_definition.cluster_type == 'pts_variant'
+    assert pts_step_from_config('pts_variant', config).is_gce
+    assert pts_step_from_config('pts_foldx', config).is_gce
 
 
-def test_foldx_runs_on_a_vm_before_the_variant_index() -> None:
+def test_foldx_runs_before_the_variant_index() -> None:
     config = UnifiedPipelineConfig()
-    assert pts_step_from_config('pts_foldx', config).is_gce
     assert config.step_definition('pts_foldx')['depends_on'] == ['pis_variant']
     # pis_variant is reached through pts_foldx, so the index does not list it again.
     assert config.step_definition('pts_variant')['depends_on'] == ['pts_foldx', 'gentropy_variant_annotation']

pts/config.yaml

@@ -27,8 +27,8 @@ steps:
 
   #: VARIANT STEP :################################################################################
   variant:
-    - name: pyspark variant index
-      pyspark: variant
+    - name: transform variant index
+      transformer: variant
       source:
         vep: intermediate/variant_annotation
         variant_annotations:
@@ -39,7 +39,6 @@ steps:
         output: output/variant
       settings:
         hash_threshold: 300
-        output_partitions: 25
   ##################################################################################################
 
   #: DISEASE STEP :#################################################################################