Skip to content

Allow targeting a self-hosted Koina server - #239

Open
BioGeek wants to merge 3 commits into
instadeepai:mainfrom
BioGeek:feat/koina-server-url
Open

Allow targeting a self-hosted Koina server#239
BioGeek wants to merge 3 commits into
instadeepai:mainfrom
BioGeek:feat/koina-server-url

Conversation

@BioGeek

@BioGeek BioGeek commented Aug 15, 2026

Copy link
Copy Markdown
Contributor

Problem

The Koina-backed features build their client with no server argument:

model = koinapy.Koina(self.intensity_model_name)          # fragment_match.py
model = koinapy.Koina(self.prosit_intensity_model_name)   # chimeric.py
koina_model = koinapy.Koina(self.irt_model_name)          # retention_time.py

so every intensity and iRT prediction goes to the public endpoint at koina.wilhelmlab.org, and there is no configuration hook to change it. That makes bulk work impractical.

Change

Adds koina.server_url and koina.ssl, defaulting to the public endpoint, so behaviour is unchanged for anyone not self-hosting:

winnow diagnose-calibration \
    koina.server_url=localhost:8500 \
    koina.ssl=false \
    dataset.spectrum_path_or_directory=data/spectra.mgf \
    dataset.predictions_path=data/preds.csv
  • FragmentMatchFeatures, ChimericFeatures and RetentionTimeFeature take koina_server_url / koina_ssl and pass them to koinapy.Koina, falling back to a new DEFAULT_KOINA_SERVER_URL.
  • ProbabilityCalibrator.apply_koina_server_overrides sets them on restored features, mirroring the existing apply_koina_model_input_overrides. With a pretrained calibrator the features come from the checkpoint, not from calibrator.yaml, so config wiring alone would have no effect. It means the pretrained general model can be used against a local server without retraining.
  • apply_koina_intensity_config applies it where it already applies the model-input overrides, so every command reading the koina block picks it up with no CLI changes.
  • calibrator.yaml interpolates the keys for the training path.

ssl=false is required because a self-hosted Triton serves gRPC without TLS while the public endpoint requires it, so the two settings have to travel together.

Tests

  • tests/configs/test_koina_config.py — the shipped config defaults to the public endpoint, every feature naming a Koina model also names a server, and the interpolations are the expected ${koina.*} references.
  • tests/calibration/test_calibrator.pyapply_koina_server_overrides retargets a Koina feature, is a no-op when given neither value, accepts one value without clobbering the other, and skips features that do not call Koina.
  • tests/calibration/features/test_koina_server.py — all three feature classes default to no server, store a self-hosted one, and can be retargeted by assignment after construction.

Docs

docs/configuration.md gains a "Self-hosting Koina" section with the overrides, why ssl must be disabled locally, and a note that the settings apply to checkpoint-restored features too. The Koina config reference lists the new keys.

The Koina-backed features construct their client as `koinapy.Koina(model_name)`,
which always resolves to the public endpoint at koina.wilhelmlab.org. Scoring a
large dataset therefore sends every fragment-match and iRT prediction to a shared
public service -- for a 780k-spectrum benchmark that is impractical and impolite,
and there is currently no way to point Winnow at a local Koina/Triton instance.

Add `koina.server_url` and `koina.ssl`, defaulting to the public endpoint so
behaviour is unchanged for anyone not self-hosting:

* `FragmentMatchFeatures`, `ChimericFeatures` and `RetentionTimeFeature` take
  `koina_server_url` / `koina_ssl` and pass them to `koinapy.Koina`.
* `ProbabilityCalibrator.apply_koina_server_overrides` sets them on restored
  features, mirroring the existing `apply_koina_model_input_overrides`. This is
  what makes the pretrained general model usable against a local server, since
  its features come from the checkpoint rather than from config.
* `apply_koina_intensity_config` applies the override where it already applies
  the model-input one, so every command that reads the `koina` block picks it up
  without CLI changes.
* `calibrator.yaml` interpolates the new keys for the training path.

Verified `koina.server_url=localhost:8500 koina.ssl=false` composes against
`diagnose_calibration`, which previously failed with ConfigCompositionException.
Adds a "Self-hosting Koina" section to the configuration guide with the two
overrides and the reason ssl must be disabled for a local Triton, and lists the
new keys in the Koina config reference.

Tests cover the three pieces the feature depends on: the feature classes store
the server settings, apply_koina_server_overrides retargets restored features
(including the no-op and partial cases, and skipping features that do not call
Koina), and the shipped configs expose the keys and interpolate them into every
feature that names a Koina model.
@BioGeek
BioGeek requested a review from JemmaLDaniel August 15, 2026 14:44
CI caught an AttributeError restoring a RetentionTimeFeature from legacy pickle
state: unpickling only repopulates what the saved state contained, so a feature
from a checkpoint predating these settings has neither attribute and raises as
soon as the Koina client is built.

That is the normal path for a pretrained calibrator, not just a test fixture, and
it fails in a second way that is quieter: apply_koina_server_overrides skips
objects that lack the attributes, so such a feature would be passed over without
comment and keep calling the public endpoint -- exactly what the override exists
to prevent.

Class-level defaults fix both: instances resolve them without the attribute being
in their __dict__, hasattr sees them, and per-instance assignment still wins.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant