Entwicklung einer ML-Lösung für Massenspektrometrie mit Open MLOps

Die Anwendung: Unbekannte Massen-Spektren identifizieren

Ein wesentlicher Bestandteil der Analyse von Massenspektren besteht darin, die vorliegenden Metaboliten zu bestimmen. Dafür werden Datenbanken wie GNPS nach möglichen Matches mit bestehenden Spektren durchsucht. Die gefundenen Treffer sind in der Regel allerdings nicht sehr präzise, was die Suche nach Matches erheblich erschwert. 

Machine Learning Tools wie Spec2Vec oder ML2DeepScore können potentielle Matches leichter finden und direkt evaluieren. Das ist zwar praktisch, jedoch erfordert die Einrichtung eines Machine Learning Systems eingangs eine Menge Aufwand. Man benötigt einen großen Server, ggf. mit automatischer Skalierung und Parallelverarbeitung, um ein vollständiges Spec2Vec-Modell zu trainieren. Das System muss außerdem regelmäßig die aktuellsten Daten abrufen und die Modelle anpassen. 

Für viele Metabolomik-Forscher - und auch für die meisten Forschungsteams - ist es nicht praktikabel, die erforderliche Cloud-Infrastruktur einzurichten und zu pflegen.

Wir haben es uns daher zum Ziel gesetzt, eine REST-API zu entwickeln, mit der Forscher ganz einfach ein Spektrum übermitteln, und von einem aktuellen Machine Learning Modell in unserer Cloud den besten Treffer erhalten.

Da GNPS täglich aktualisiert wird, müssen auch die Modelle der API laufend neu trainiert werden. In diesem Artikel erklären wir, wie unser Open MLOps-Framework uns ermöglicht hat, dieses Projekt umzusetzen.

System Übersicht: Omigami

Unsere Lösung besteht im Wesentlichen aus zwei Komponenten:

1. Daten Pipeline: Neue Daten herunterladen, verarbeiten und die neuen Modelle trainieren;
2. Deployment pipeline: Die neuen Modelle über eine REST-API registrieren und bereitstellen.

Der erste Schritt dreht sich also vor allem um die Verarbeitung der Daten. Wir laden täglich eine mehrere Gigabyte große JSON-Datei herunter, bereiten sie in einem Tabellenformat auf und trainieren dann das Modell.

Im zweiten Schritt stellen wir unsere Modelle als API zur Verfügung. Dieser Schritt besteht hauptsächlich daraus, Seldon mit YAML zu konfigurieren.

So die stark vereinfachte Erklärung. Die beiden Schritte lassen sich etwas feiner in vier Teilkomponenten aufteilen: Herunterladen, Vorverarbeitung, Training und Deployment.

Nachdem neue Daten in die GNPS-Datenbank aufgenommen wurden, folgen also folgende Schritte:

• Wir laden die neuen Daten automatisch herunter;
• Wir konvertieren die JSON-Daten in tabellarische Daten, die sich für das Training eignen;
• Wir trainieren die Modelle anhand der Spec2Vec und MS2DeepScore Algorithmen;
• Wir registrieren die Modelle in MLFlow und speichern sie in S3, wo Seldon sie abruft und bereitstellt.

Das bedeutet, dass jeder Nutzer der API jederzeit auf die neueste Version des Modells zugreifen kann. Das funktioniert so:

• Ein Forscher hat ein Spektrum vorliegen, das er identifizieren möchte;
• Er lädt seine Daten über unsere API hoch;
• Die API sucht anhand des neuesten Modells nach potenziellen Matches mit dem Spektrum;
• Die API übermittelt die gefundenen Matches zurück an den Forscher.

Jeder Schritt lässt sich hervorragend in unser Open MLOps Framework integrieren. Mit Hilfe von Prefect definieren und überwachen wir alle Tasks, und überblicken so den gesamten Prozess. Mit MLFlow behalten wir den Überblick über die trainierten Modelle und verbinden sie mit Seldon. Seldon wiederum stellt die Modelle über eine REST-API bereit und übermittelt die Informationen an den Endnutzer.

Die einzelnen Schritte im Detail

Anfangs haben wir unsere Lösung stark vereinfacht in zwei Komponenten eingeteilt: die Daten Pipeline und die Deployment Pipeline. Dann haben wir das Ganze aufgeschlüsselt: Herunterladen, vorverarbeiten, trainieren, bereitstellen. An dieser Stelle gehen wir hier etwas mehr ins Detail.

Wie wir mit Prefect Tasks unsere Lösung modular gestalten

Eine praktische Funktion von Prefect ist das automatisch generierte Übersichtsdiagramm des gesamten Systems. Du definierst lediglich einzelne Tasks und Prefect fügt sie zusammen.

Unser TrainModel Prefect Task sieht zum Beispiel so aus:

Da wir objektorientiert programmieren, stammt unser TrainModel Task von einem Prefect Task. Wir definieren die Funktionen __init__ und __run__, um das Word2Vec-Modell auf unseren Eingabedaten zu trainieren.

Sobald wir unser Modell trainiert haben, sieht der RegisterModel Task sehr ähnlich aus. Wir benennen das Modell nach dem Zeitpunkt, an dem es das Training abgeschlossen hat, und speichern es in MLFlow.

In ähnlicher Weise legen wir einen Prefect Task für andere Aufgaben fest, z. B. für das Herunterladen der Daten und das Deployment des Modells. Prefect kombiniert alle Tasks zu einem Flow und erstellt einen DAG (Directed Acyclic Graph). Unser DAG sieht in diesem Fall folgendermaßen aus:

Im Detail gibt es acht verschiedene Elemente des Systems:

• Daten herunterladen: Wir ziehen die aktuellsten JSON-Daten aus GNPS.
• Chunks erstellen: Wir teilen die Daten in Chunks ein. Diese werden in Redis geladen und nach ihrer ID gruppiert, damit sie parallel verarbeitet werden können.
• Rohspektren speichern: Wir speichern eine Kopie der einzelnen Spektren.
• Spektren verarbeiten: Wir wandeln die Spektren in “Wort"-Vektoren um, die wir Gensim's Word2Vec geben können.
• Modell trainieren: Wir trainieren ein Word2Vec Modell.
• Modell registrieren: Wir protokollieren die Details des Modells in MLFlow und speichern die Binärdaten in S3.
• Embeddings speichern: Wir erstellen und speichern Referenz-Embeddings aus den Spektren, um sie mit den von den Nutzern hochgeladenen Daten zu vergleichen.
• Modell bereitstellen: Wir aktualisieren unser Seldon Deployment mit dem neuen Modell.

Die Vorteile eines End-to-End-Frameworks wie Open MLOps

Jeder der oben genannten Schritte könnte zwar auch manuell durchgeführt werden, mit den vorgefertigten Komponenten können wir uns jedoch enger an die MLOps-Richtlinien halten. Indem wir ein Framework einsetzen, erfüllt unsere Lösung automatisch folgende Grundsätze:

• Einfache Kollaboration: Open MLOps ermöglicht ein einheitliches Vorgehen im Team. Da jedes Mitglied die gleiche Plattform bereits für andere Projekte verwendet hat, kann sich jeder leicht einbringen (zum Beispiel bei der Deployment-Konfiguration (in YAML) oder der Optimierung einer bestimmten Python Komponente als Prefect Task);
• Reproduzierbarkeit: Wir speichern und tracken alle Trainingsläufe und Modelle, die wir verwenden, in MLFlow. Das heißt, wir können immer herausfinden, wie wir zu einer bestimmten Vorhersage oder einem bestimmten Ergebnis gekommen sind.
• Kontinuität und Automatisierung: Wenn wir Änderungen und Verbesserungen vornehmen, müssen wir uns keine Sorgen über Fehler beim Deployment machen. Unsere Änderungen werden automatisch getestet und live geschaltet.
• Monitoring: Wir können unsere Modelle ganz einfach auf Model Decay überprüfen. Der Zustand unserer Infrastruktur lässt sich leicht überprüfen und wir wissen im Problemfall genau, wo wir nach Fehlern suchen müssen.

Für Engineers, die an manuelle Abläufe gewöhnt sind (z.B. Python-Code zu schreiben, um das Modell mit FastAPI zu verpacken, um die REST-API zu erstellen), kann das Schreiben von YAML zur Konfiguration von Seldon mit einer steilen Lernkurve einhergehen. Insgesamt ist die Konfiguration jedoch einfacher, weniger fehleranfällig und gewährleistet Konsistenz zwischen verschiedenen Deployments.

Diese Herausforderungen stehen uns noch bevor

Natürlich ist ein Framework kein Allheilmittel, das den Entwicklungsprozess auf einmal zum Kinderspiel macht. An einige spezielle Bedürfnisse mussten wir unsere Lösung noch anpassen: 

• Wir haben Redis hinzugefügt, um den täglichen JSON-Blob effizient zwischenzuspeichern, zu gruppieren und in ein Tabellenformat zu konvertieren.
• Außerdem haben wir Authentifizierungs- und Zugriffsregeln hinzugefügt. Open MLOps geht standardmäßig davon aus, dass die REST-API intern verwendet wird. Das mussten wir manuell anpassen, damit sich auch externe Nutzer anmelden und mit einem API-Schlüssel auf die API zugreifen können.
• Wir ermöglichen eine stetig wachsende Datenmenge. Der GNPS-Datensatz ist hochspannend, da er kontinuierlich wächst. Da wir unsere Lösung auf unbestimmte Zeit einsetzen wollen, muss sie mit den täglich größer werdenden Eingabedaten mitwachsen.

Keine dieser Anpassungen ist Teil der aktuellen Standardkonfiguration von Open MLOps. Da es sich hier aber um relativ häufige Herausforderungen von Machine Learning-Lösungen handelt, werden wir Open MLOps perspektivisch um entsprechende Optionen erweitern.

Probiere Open MLOps selbst aus

Wenn du deine eigene Machine Learning-Lösung entwickeln und bereitstellen möchtest, probiere Open MLOps direkt selbst aus. Solltest du dabei auf Schwierigkeiten stoßen, eröffne einfach ein GitHub Issue und wir helfen dir weiter.