PROTOTYPE: move L2G training and prediction from gentropy into PTS

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

@@ -50,8 +50,6 @@ clusters:
     secret_map:
       # this results in /var/run/secrets/hfhub-key.json with {"HF_TOKEN": "hfhub-key_value"} as content
       HF_TOKEN: hfhub-key
-      # this results in /var/run/secrets/wandb-key.json with {"WANDB_API_KEY": "wandb-key_value"} as content
-      WANDB_API_KEY: wandb-key
 
   gentropy_colocalisation:
     image_version: '2.2'

orchestration/src/orchestration/dags/config/gentropy.yaml

@@ -210,99 +210,6 @@ steps:
       # OUTPUTS
       step.feature_matrix_path: '{{release_uri}}/intermediate/l2g_feature_matrix'
 
-  l2g_train_test_split:
-    params:
-      step: locus_to_gene_train_test_split
-      step.session.write_mode: overwrite
-      +step.session.extended_spark_conf: >-
-        {spark.sql.autoBroadcastJoinThreshold:'-1'}
-      # INPUTS
-      step.credible_set_path: '{{release_uri}}/output/credible_set'
-      step.feature_matrix_path: '{{release_uri}}/intermediate/l2g_feature_matrix'
-      step.gold_standard_curation_path: '{{release_uri}}/input/l2g/gold_standard.json'
-      # The test set is pinned across releases rather than resampled per run, so
-      # models stay comparable. `step.test_size` is ignored when this is set.
-      # pis_l2g_predefined_test_for_split stages it into the release input directory.
-      step.predefined_test_parquet_path: '{{release_uri}}/input/l2g/predefined_test_for_split'
-      # OUTPUTS
-      step.train_parquet_path: '{{release_uri}}/intermediate/l2g_train_split'
-      step.test_parquet_path: '{{release_uri}}/intermediate/l2g_test_split'
-      step.split_stats_path: '{{release_uri}}/intermediate/l2g_train_test_split_stats.json'
-
-  l2g_training:
-    params:
-      step: locus_to_gene
-      step.session.write_mode: overwrite
-      step.run_mode: train
-      step.wandb_run_name: '{{l2g_training_version}}'
-      # Metrics come from the held-out test set; the saved model is then refit on
-      # train + test so it is not deployed having discarded the held-out labels.
-      step.train_on_full_dataset: true
-      step.hf_hub_repo_id: opentargets/locus_to_gene_{{l2g_training_version}}
-      step.hf_model_commit_message: 'chore: update model base model for {{l2g_training_version}} run'
-      +step.session.extended_spark_conf: >-
-        {spark.kryoserializer.buffer.max:500m, spark.sql.autoBroadcastJoinThreshold:'-1'}
-      # INPUTS
-      step.train_parquet_path: '{{release_uri}}/intermediate/l2g_train_split'
-      step.test_parquet_path: '{{release_uri}}/intermediate/l2g_test_split'
-      # OUTPUTS
-      step.model_path: '{{release_uri}}/etc/model/locus_to_gene_model/classifier.skops'
-      step.hf_credentials_path: '/var/run/secrets/hfhub-key'
-      step.wandb_credentials_path: '/var/run/secrets/wandb-key'
-
-  l2g_prediction:
-    cluster: false
-    google_batch_index_specs:
-      pointer: 'gentropy_step'
-      max_task_count: 1000 # Number of tasks running in a single google batch job.
-      generator_specs:
-        generator_options:
-          # Job partitioning depends on the number of credible_set parquet files found in the input glob.
-          input_glob: '{{release_uri}}/output/credible_set/**.parquet'
-          # With output_partitions = 1, each input is mapped to 1 output.
-          output_prefix: '{{release_uri}}/intermediate/l2g_prediction'
-    google_batch:
-      job:
-        task_group:
-          parallelism: 200
-          task_environments:
-            # Populated at runtime by BatchIndexOperator with dictionaries
-            # containing $INPUT_PARTITION and $OUTPUT_PARTITION
-            environments: []
-          task_config:
-            shared_environment:
-              secrets:
-                mapping:
-                  HF_TOKEN:
-                    secret_id: hfhub-key
-            max_retry_count: 2
-            # exit_codes is what makes max_retry_count do anything: Batch retries only
-            # codes a lifecycle policy names, defaulting to the reserved 50001-50005, so
-            # an application failure (exit 1) was never retried. Listing codes replaces
-            # that default rather than adding to it, hence the reserved ones repeated.
-            exit_codes: [1, 50001, 50002, 50003, 50004, 50005]
-            max_run_duration: '1h'
-            instance_resource_spec:
-              cpu_milli: 2000
-              memory_mib: 7000
-              boot_disk_mib: 4000
-            runnable_spec:
-              image_uri: 'ghcr.io/opentargets/gentropy:{{gentropy_version}}'
-              entrypoint: /usr/bin/bash
-              script_file: l2g_predict.sh
-              script_variables:
-                l2g_training_version: '{{l2g_training_version}}'
-                feature_matrix_path: '{{release_uri}}/intermediate/l2g_feature_matrix'
-        allocation:
-          instance:
-            machine_type: n1-standard-2
-            provisioning_model: STANDARD
-        logs: {}
-      collect:
-        source_prefix: '{{release_uri}}/intermediate/l2g_prediction'
-        destination_prefix: '{{release_uri}}/output/l2g_prediction'
-        file_glob: '**.parquet'
-
   l2g_evidence:
     params:
       step: locus_to_gene_evidence

orchestration/src/orchestration/dags/config/unified_pipeline.py

@@ -87,7 +87,6 @@ def __init__(self) -> None:
             template_context={
                 'release_uri': self.release_uri,
                 'gentropy_version': up.get('gentropy_version'),
-                'l2g_training_version': self.run.release_name,
                 'vep_version': up.get('vep_version'),
             },
         )

orchestration/src/orchestration/dags/config/unified_pipeline.yaml

@@ -528,8 +528,36 @@ steps:
 
   pts_vep_view:
     depends_on:
-      - gentropy_l2g_prediction
+      - pts_l2g
 
+  pts_l2g:
+    # Measured on the real 26.09-2 inputs, not estimated. Whole-step wall clock:
+    #   n1-highmem-32   39m40s
+    #   n2d-standard-32 23m12s   (n2d cores are ~1.6x faster at interventional TreeSHAP)
+    #   n2d-standard-64 16m06s   <- this
+    #
+    # 64 cores is only 1.44x faster for 1.39x the core-time, because the step is
+    # memory-bandwidth-bound: utilisation falls from 81% to 67% and total CPU consumed RISES
+    # 15.6% for identical work. But the absolute difference is ~4.8 vCPU-hours per run, so
+    # paying it for seven minutes of wall clock is worth it. Do not read the ratio as a reason
+    # to go back to 32 without looking at the absolute.
+    #
+    # Scaling further is not the lever: past this point cores mostly stall on memory. The only
+    # large win left is algorithmic -- xgboost's native pred_contribs is ~200x faster but emits
+    # log-odds margin contributions rather than probability-space SHAP values.
+    #
+    # Do NOT "right-size" this to a highcpu shape. Peak RSS is 28.2 GB, transient, during the
+    # polars scan and join of the 5.5 GB of parquet inputs -- SHAP itself sits at a flat
+    # ~12.7 GB. A 32 GB shape leaves under 4 GB for the OS, page cache and 32 spawning workers,
+    # and thrashes. The 64-core shape's 256 GB is ~9x the measured peak.
+    #
+    # Memory also scales with shap_workers: the pool spawns rather than forks (see
+    # l2g/explain.py), so every worker loads its own copy of xgboost, shap and sklearn.
+    machine_type: n2d-standard-64
+    depends_on:
+      - gentropy_l2g_feature_matrix
+      - pis_l2g
+      - pis_l2g_predefined_test_for_split
 
   # GENTROPY STEPS
   gentropy_biosample:
@@ -577,17 +605,6 @@ steps:
       - gentropy_variant
       - gentropy_enhancer_to_gene
       - pts_interaction
-  gentropy_l2g_train_test_split:
-    depends_on:
-      - pis_l2g
-      - pis_l2g_predefined_test_for_split
-      - gentropy_l2g_feature_matrix
-  gentropy_l2g_training:
-    depends_on:
-      - gentropy_l2g_train_test_split
-  gentropy_l2g_prediction:
-    depends_on:
-      - gentropy_l2g_training
   gentropy_l2g_evidence:
     depends_on:
-      - gentropy_l2g_prediction
+      - pts_l2g

orchestration/src/orchestration/dags/unified_pipeline.py

@@ -368,16 +368,6 @@ def gentropy_step_differs(step_name: str) -> list[Differ]:
             'gentropy_l2g_feature_matrix': {
                 'step.feature_matrix_path': gsp('gentropy_l2g_feature_matrix', 'step.feature_matrix_path'),
             },
-            'gentropy_l2g_train_test_split': {
-                'step.train_parquet_path': gsp('gentropy_l2g_train_test_split', 'step.train_parquet_path'),
-                'step.test_parquet_path': gsp('gentropy_l2g_train_test_split', 'step.test_parquet_path'),
-            },
-            'gentropy_l2g_training': {
-                'step.model_path': gsp('gentropy_l2g_training', 'step.model_path'),
-            },
-            'gentropy_l2g_prediction': {
-                'step.predictions_path': gsp('gentropy_l2g_prediction', 'step.predictions_path'),
-            },
             'gentropy_l2g_evidence': {
                 'step.evidence_output_path': gsp('gentropy_l2g_evidence', 'step.evidence_output_path'),
             },

pts/config.yaml

@@ -1921,6 +1921,96 @@ steps:
       destination: view/association_otf
   ##################################################################################################
 
+  #: L2G STEP :#####################################################################################
+  l2g:
+    - name: transform l2g_train
+      transformer: l2g_train
+      source:
+        feature_matrix: intermediate/l2g_feature_matrix
+        credible_set: output/credible_set
+        gold_standard: input/l2g/gold_standard.json
+        predefined_test: input/l2g/predefined_test_for_split
+      destination:
+        model: etc/model/locus_to_gene_model/classifier.skops
+        background: etc/model/locus_to_gene_model/shap_background.parquet
+        metrics: etc/model/locus_to_gene_model/metrics.json
+        train_split: intermediate/l2g_train_split
+        test_split: intermediate/l2g_test_split
+        split_stats: intermediate/l2g_train_test_split_stats.json
+      settings:
+        # Anchored, not a scratchpad sentinel: otter's Scratchpad is string.Template
+        # substitution, so `${...}` can only ever produce a string. The order is the
+        # order the model is fitted on and the order of the output `features` array.
+        features_list: &l2g_features
+          - eQtlColocClppMaximum
+          - pQtlColocClppMaximum
+          - sQtlColocClppMaximum
+          - eQtlColocH4Maximum
+          - pQtlColocH4Maximum
+          - sQtlColocH4Maximum
+          - eQtlColocClppMaximumNeighbourhood
+          - pQtlColocClppMaximumNeighbourhood
+          - sQtlColocClppMaximumNeighbourhood
+          - eQtlColocH4MaximumNeighbourhood
+          - pQtlColocH4MaximumNeighbourhood
+          - sQtlColocH4MaximumNeighbourhood
+          - distanceSentinelFootprint
+          - distanceSentinelFootprintNeighbourhood
+          - distanceFootprintMean
+          - distanceFootprintMeanNeighbourhood
+          - distanceTssMean
+          - distanceTssMeanNeighbourhood
+          - distanceSentinelTss
+          - distanceSentinelTssNeighbourhood
+          - vepMaximum
+          - vepMaximumNeighbourhood
+          - vepMean
+          - vepMeanNeighbourhood
+          - e2gMean
+          - e2gMeanNeighbourhood
+          - geneCount500kb
+          - proteinGeneCount500kb
+          - credibleSetConfidence
+          - transPQtlColocH4Maximum
+          - transPQtlColocH4MaximumNeighbourhood
+        # Read off the 26.09-2 classifier.skops. n_estimators, gamma and max_delta_step
+        # are in neither gentropy nor this repo; do not "restore" them from gentropy's config.
+        hyperparameters:
+          objective: binary:logistic
+          eval_metric: aucpr
+          random_state: 777
+          n_estimators: 300
+          max_depth: 5
+          min_child_weight: 10
+          eta: 0.05
+          subsample: 0.8
+          colsample_bytree: 0.8
+          reg_alpha: 1
+          reg_lambda: 1.0
+          scale_pos_weight: 0.8
+          gamma: 0
+          max_delta_step: 1
+        train_on_full_dataset: true
+        shap_background_size: 100
+        shap_background_seed: 42
+    - name: transform l2g_predict
+      requires:
+        - transform l2g_train
+      transformer: l2g_predict
+      source:
+        feature_matrix: intermediate/l2g_feature_matrix
+        credible_set: output/credible_set
+        model: etc/model/locus_to_gene_model/classifier.skops
+        background: etc/model/locus_to_gene_model/shap_background.parquet
+      destination: output/l2g_prediction
+      settings:
+        features_list: *l2g_features
+        l2g_threshold: 0.05
+        explain_predictions: true
+        # No `shap_background_size` here on purpose: the masker is sized from the background
+        # parquet the training task wrote, so the size is set once, above.
+  ##################################################################################################
+
   #: VEP_VIEW STEP :################################################################################
   vep_view:
     - name: transform vep view

pts/pyproject.toml

@@ -28,6 +28,10 @@ dependencies = [
   "clinical-mining",
   "huggingface-hub>=0.32.4",
   "networkx>=3.6.1",
+  "xgboost>=3.2.0",
+  "shap>=0.51.0",
+  "scikit-learn>=1.7.0",
+  "skops>=0.11.0",
 ]
 
 [dependency-groups]
@@ -37,6 +41,9 @@ test = ["pytest>=9.0.3", "pytest-asyncio>=1.4.0", "pytest-mock>=3.15.1", "pytest
 [project.optional-dependencies]
 baseline_expression = ["pandas<3", "scanpy>=1.10.0", "cellex"]
 
+[tool.uv]
+environments = ["sys_platform != 'darwin' or platform_machine != 'x86_64'"]
+
 [[tool.uv.index]]
 name = "pytorch-cpu"
 url = "https://download.pytorch.org/whl/cpu"

pts/src/pts/transformers/l2g/__init__.py

@@ -0,0 +1,10 @@
+"""Locus-to-gene training and prediction, ported from gentropy onto polars."""
+
+# xgboost MUST be imported before skops/scikit-learn anywhere in this package.
+# sklearn ships its own LLVM OpenMP runtime (sklearn/.dylibs/libomp.dylib) and
+# libxgboost.dylib links another; loading skops' first and xgboost's second kills
+# the process with SIGSEGV on macOS, at the first fit or predict rather than at
+# import. Linux is unaffected -- both link the one system libgomp -- so this
+# protects the dev machine, not production. Importing the package here means every
+# submodule inherits the ordering instead of each repeating it.
+import xgboost  # noqa: F401