Daemon transformed Biographica's graph machine learning and transcriptomic foundation model machine learning (ML) workflow with Amazon SageMaker AI to orchestrate model training, hyperparameter tuning and scalable deployments serving end-users.

The Client

Biographica is a cutting-edge biotechnology firm revolutionising agriculture through advanced machine learning (ML) and genomics. They specialise in using AI technology to rapidly identify and prioritise high-value genetic targets for crop gene editing, tackling bottlenecks in developing more productive, sustainable, nutritious, and climate-resilient food sources.

The Challenge

Biographica's scientists and ML engineers needed a process to leverage high-performance computing in the cloud and simplify their workflows. The process should enable them tocollaborate and run their AI experiments for each ML model in a scalable and asynchronous manner, with minimal changes to their current processes.

For their advanced research, Biographica employs three state-of-the-art AI models, each with vastly different complexities and architectures. These include specialised models for protein function analysis, genetic regulation & literature synthesis , each of which has a different underlying architecture and requires specific computational resources and data formats.

Biographica's AI workloads fall into two categories depending on the use case.

1. ML engineer use case - easy, configurable training at scale: In the first use case, ML engineers need to build better and more complex crop gene editing models. This process was constrained by the need to do manual setup, configuration, remote login, data provision, and experimentation in both local and remote environments. ML engineers aimed for better model performance through parallel training and hyperparameter tuning, but this was difficult for a small team managing their own infrastructure.

2. Biological scientist use case - scalable and transparent inference: In the second use case, Biographica's scientists (computational biologists and other biological scientists) use trained models to identify targets for gene editing. The current method requires scientists to log in to virtual machines, run inference, and download results. This is error-prone and costly, as machines must be active even when not in use. Scientists need a way to run their process without manual VM setup, data transfer, and script execution.

Our Approach

Design principles

Containerisation and SageMaker
The gold standard for both scalable inference and training across platforms is containerisation. Daemon opted to use containerised training and inference jobs, along with containerised data preparation and post-processing pipelines. SageMaker AI is AWS's flagship ML environment with out-of-the-box support for on-demand containerised ML workloads, and access to high-performance Nvidia GPU-enabled cloud compute.

Minimal impact on user experience
SageMaker AI's flexibility allows Biographica's staff to run their experiments without the added complexity of managing the execution environment. Biographica staff could choose with little friction between local machines, do-it-yourself (cloud VM) and managed services, maintaining their original workflows while the system transparently handles data and compute.

Training

Transparent training adapters and workflow interfaces
We developed an Amazon SageMaker AI-compatible adapter for Biographica's existing training interfaces, enabling graph- and sequence-based models to leverage cloud-based training across multiple high-compute cloud servers, with results synchronised into an MLFlow model repository, with data flows managed by Dagster as per Biographica's current practice, and job monitoring integrated into Biographica's command-line tooling.

Scaling hyperparameter tuning
We implemented custom hyperparameter tuning orchestration, using the open-source library Optuna, together with SageMaker AI Training for high-performance parallel training. Optuna is a feature-rich open-source optimisation library that also allows us to support Sagemaker AI Plans. These were introduced alongside on-demand instances, allowing Biographica to pre-purchase and guarantee access to GPU capacity. This capability is also being used in production model training by Biographica.

Inference

On-demand inference for scientists
We streamlined Biographica's workflow to allow scientists to run their graph ML inferences without the additional overhead of logging in to remote machines. As inference jobs were generally conducted ad-hoc, we utilised Amazon SageMaker AI Batch Transform to scale to handle large input batches. This removed the overhead of putting up and shutting down costly on-demand infrastructure to run inference.

A familiar interface
Built to run from a scientist's local machine, the tool features a familiar interactive command-line interface. It handles the packaging of models from MLFlow, the preparation of input data for SageMaker (using SageMaker AI Processing), and job execution, all while providing complete visibility into the process and organising the results in a familiar way.

The Outcome

Daemon's strategic implementation of Amazon SageMaker AI for Biographica's advanced AI use-cases has driven efficiency, enhancing research capabilities, and optimising resource utilisation across their critical models.

Right-sized spend via on-demand cloud processing: Inference is now fully on-demand, removing the costs associated with always-on cloud virtual machines, and training is architected around a containerised on-demand model. Cloud spend is kept low, by exactly matching what is needed.

Model performance by training at scale: The combined training solution provides Biographica with a strategic advantage by allowing the company to run scaled-up training and tuning jobs to achieve optimal performance for its foundation models, strengthening Biographica's scientific contributions and potentially leading to more impactful discoveries.

Accelerated Model Development: By streamlining the model training process Biographica can now develop, test, and compare new AI models more rapidly, shortening AI solution development cycles. The solution's adherence to Biographica’s internal workflows ensures a smooth, consistent, and user-friendly experience across the team.

Accelerated biological experimentation: Removing the overhead with infrastructure management during large-scale inference has enabled Biographica's scientists to run experiments faster and pursue new scientific discoveries and insights via self-serve infrastructure.

Testimonials

Daemon delivered highly performant out-the-box functionality for Biographica to scale our ML experimentation while keeping cloud resource spend controlled. Working with them was a low friction, high ROI experience and we’re still building on the tooling they delivered to this day.

Dominic Hall, Co-founder & CTO, Biographica

Labels
AWS
SageMaker Batch Transform
SageMaker Training
SageMaker Training Plans
Hyperparameter tuning
Optuna
ML (machine learning)
Graph machine learning
Transcriptomics
Foundation model training
AI pipeline automation
MLFlow
HPC (high performance computing)