Skip to content
omicsedgebioPublic

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Repository files navigation

HelixFlow

A compact, production-shaped reference implementation for reproducible RNA-seq model training and inference across local, containerized, Kubernetes, and GCP environments.

The goal of this repository is not to build the most sophisticated cancer classifier. The model is intentionally small. The interesting part is the platform around the model: reproducible data ingestion, versioned artifacts, containerized execution, Kubernetes batch scheduling, optional GPU placement, inference deployment, infrastructure-as-code, CI, and workflow portability.

Why this project exists

Scientific ML workloads often start as notebooks or one-off scripts and then become shared infrastructure. That transition creates recurring engineering problems:

  • researchers need reproducible environments without learning cluster internals;
  • training is finite batch work, while inference is a long-running service;
  • CPU, memory, GPU, and storage requirements differ across pipeline stages;
  • model artifacts need traceability back to code, configuration, and data;
  • expensive GPU capacity should be requested deliberately, not used as a default;
  • infrastructure changes need the same review discipline as application code;
  • the same scientific workflow should be portable between local and cluster execution.

This repo demonstrates one way to address those problems without hiding the trade-offs.

Dataset

The real-data path uses the public UCI Gene Expression Cancer RNA-Seq dataset, a subset of TCGA PANCAN HiSeq data:

  • 801 tumor samples
  • 20,531 expression features
  • five classes: BRCA, KIRC, COAD, LUAD, PRAD
  • no missing values in the published dataset

Dataset: https://archive.ics.uci.edu/dataset/401/gene+expression+cancer+rna+seq

For CI and quick local testing, the repository also provides a deterministic synthetic expression dataset so tests do not depend on network access.

Architecture

flowchart LR
    A[UCI RNA-seq data] --> B[Data preparation]
    B --> C[Feature selection + scaling]
    C --> D[PyTorch training]
    D --> E[Versioned model artifacts]
    E --> F[FastAPI inference]

    G[Nextflow] --> C
    H[Kubernetes Job] --> D
    I[GPU node pool] -. optional .-> H
    J[Kubernetes Deployment] --> F
    K[Service + HPA] --> J
    L[Terraform / GCP] --> I
    L --> J
    M[GitHub Actions] --> B
    M --> D
    M --> L
Loading

Repository layout

.
├── src/helixflow/     Python package: data, model, training, inference, API
├── tests/                   unit + training smoke tests
├── k8s/                     Job, GPU Job, Deployment, Service, HPA, PVCs
├── terraform/               reusable GKE module + dev environment
├── nextflow/                portable training workflow
├── scripts/                 manifest validation
├── docs/                    architecture and operations notes
├── .github/workflows/       CI for Python, Docker, Kubernetes YAML, Terraform
├── Dockerfile             CPU/inference image
├── Dockerfile.gpu         CUDA-enabled training image
├── docker-compose.yml
└── Makefile

Quick start

1. Create an environment

python3.11 -m venv .venv
source .venv/bin/activate
pip install -e '.[dev]'

2. Run a network-free smoke test

make smoke

This generates synthetic expression data, trains a small model on CPU, and writes a complete artifact bundle under artifacts/model/.

3. Run against the public RNA-seq dataset

python -m helixflow.data --out-dir data
python -m helixflow.train \
  --data data/processed/data.csv \
  --labels data/processed/labels.csv \
  --output-dir artifacts/model \
  --device auto

The real UCI archive is about 69.5 MB compressed, so it is intentionally downloaded at runtime rather than committed to Git.

What is a model artifact here?

Training writes an immutable bundle containing:

artifacts/model/
├── model.pt
├── selector.joblib
├── scaler.joblib
├── label_encoder.joblib
├── selected_features.json
├── metadata.json
└── classification_report.json

metadata.json captures the training configuration, Git SHA, device, Torch version, and evaluation metrics. The inference service loads the complete bundle rather than assuming preprocessing is recreated separately.

That design prevents a common ML failure mode: deploying model weights while accidentally changing preprocessing between training and inference.

Why a Kubernetes Job for training?

Training is run-to-completion work. A Kubernetes Job gives the workload a clear success/failure lifecycle and retry semantics.

apiVersion: batch/v1
kind: Job
...

The CPU example requests only CPU and memory. The GPU variant uses a CUDA-enabled PyTorch image (Dockerfile.gpu) and explicitly requests:

resources:
  limits:
    nvidia.com/gpu: "1"

It also targets a GPU-labelled node pool and tolerates the GPU taint. That keeps ordinary CPU workloads from landing on expensive accelerator nodes. The GPU manifest references the :gpu image tag so a CPU-only inference image cannot accidentally be scheduled as a CUDA training workload.

Why a Deployment for inference?

Inference is a continuously available service, so it uses a Deployment rather than a Job:

  • replica management;
  • rolling replacement;
  • readiness and liveness probes;
  • resource requests/limits;
  • HorizontalPodAutoscaler support.

A Service provides stable networking in front of ephemeral Pods.

GPU scheduling philosophy

The repository makes GPU use explicit and optional.

The Terraform environment defaults to:

enable_gpu_pool = false

That is deliberate. A training workload should first demonstrate that GPU acceleration is useful before infrastructure starts paying for idle accelerator capacity.

When enabled, the GPU node pool is labelled workload=gpu and tainted so that only workloads with the matching toleration are eligible for placement.

For a larger platform, the next step would be queueing/gang scheduling with tooling such as Kueue or another batch scheduler rather than allowing every researcher to submit unconstrained GPU jobs directly.

Terraform design

Terraform is separated into:

  • a reusable modules/gke module;
  • an environment-specific envs/dev configuration;
  • optional GCS backend configuration for remote state.

The module creates:

  • custom VPC/subnet;
  • GKE cluster without the default node pool;
  • autoscaled CPU node pool;
  • optional autoscaled GPU node pool;
  • versioned GCS artifact bucket;
  • Workload Identity configuration.

The Google provider is pinned to the current 7.x line rather than leaving provider resolution unconstrained.

Example:

cd terraform/envs/dev
cp terraform.tfvars.example terraform.tfvars
terraform init
terraform fmt -check -recursive ../..
terraform validate
terraform plan

The repo intentionally does not automatically apply Terraform from CI. CI validates and plans are intended for review; production infrastructure changes should have an explicit approval boundary.

Nextflow portability

nextflow/main.nf contains the scientific workflow definition while nextflow/nextflow.config controls execution.

Local container execution:

nextflow run nextflow/main.nf -profile local

Kubernetes execution:

nextflow run nextflow/main.nf -profile kubernetes \
  --container ghcr.io/omicsedgebio/helixflow:latest

The scientific command does not change when the execution backend changes. That is the main reason workflow engines are useful for researcher-facing platforms: the researcher thinks about inputs and parameters; the platform owns execution details.

MLflow integration

MLflow is optional rather than hard-coupled to the trainer.

pip install -e '.[mlflow]'
export MLFLOW_TRACKING_URI=http://mlflow.example.internal:5000
export MLFLOW_EXPERIMENT_NAME=rna-seq-classifier

When configured, training logs parameters, Git SHA, evaluation metrics, and the complete artifact directory. Without MLFLOW_TRACKING_URI, training remains fully functional and writes artifacts locally.

That keeps experiment tracking useful without making it a single point of failure for development.

Inference service

Run locally:

uvicorn helixflow.api:app --host 0.0.0.0 --port 8080

Health endpoint:

curl http://localhost:8080/healthz

Containerized:

docker compose up --build

The API is intentionally small. The main purpose is to demonstrate the operational boundary between a model artifact and a stable inference service.

Throughput benchmarking

A simple batch benchmark is included:

python -m helixflow.benchmark \
  --artifact-dir artifacts/model \
  --input data/processed/data.csv \
  --repeats 25

In a real inference platform, this would be extended to measure:

  • p50/p95/p99 latency;
  • throughput versus batch size;
  • CPU/GPU utilization;
  • memory pressure;
  • model loading time;
  • concurrency;
  • throughput per dollar.

The key principle is to profile the bottleneck before adding compute. Low GPU utilization may be caused by CPU preprocessing, storage I/O, batching, or synchronization rather than insufficient GPU capacity.

CI/CD

GitHub Actions has three independent jobs:

Python

  • install package;
  • lint with Ruff;
  • unit tests;
  • Kubernetes YAML parsing;
  • synthetic end-to-end train smoke test.

Container

  • build the production Docker image.

Terraform

  • formatting check;
  • provider initialization without a backend;
  • terraform validate.

This deliberately separates application failures from container or infrastructure failures so the source of a CI failure is obvious.

Researcher-facing workflow

A researcher should not need to know about taints, node pools, PVCs, or Terraform to train a model.

The desired interface is closer to:

nextflow run workflow \
  --data cohort.csv \
  --labels labels.csv

The platform is responsible for translating the workload into the appropriate container, resource request, scheduler, storage, and experiment-tracking configuration.

That separation is important when migrating existing research workloads to Kubernetes: preserve the workflow researchers already understand while changing the infrastructure underneath it incrementally.

Production extensions

This repository intentionally stops before pretending a single-node demo is a complete production ML platform. Natural next steps would include:

  • GCS Fuse / object-store-aware data access rather than generic PVCs;
  • dedicated MLflow tracking and artifact services;
  • Workload Identity-bound Kubernetes service accounts;
  • Kueue/Volcano-style queued GPU scheduling and quotas;
  • distributed PyTorch DDP jobs;
  • Ray/Anyscale for high-throughput distributed inference;
  • model registry promotion policies;
  • canary or shadow inference deployment;
  • Prometheus/Grafana and GPU/DCGM telemetry;
  • OPA/Gatekeeper or Kyverno policy controls;
  • separate dev/staging/prod Terraform state;
  • image signing and SBOM generation;
  • cost attribution by namespace/team/workload.

Engineering decisions

See:

License

Project code is MIT licensed. The UCI dataset is not redistributed by this repository; users download it from UCI and should follow its CC BY 4.0 terms and citation guidance.

Reference validation

HelixFlow was exercised end-to-end using the UCI TCGA PANCAN RNA-seq dataset:

  • 801 tumor samples
  • 20,531 gene-expression features
  • 5 classes: BRCA, COAD, KIRC, LUAD, and PRAD

For the reference run:

  • Stratified 80/20 train/test split
  • Zero-variance filtering fit on the training split only
  • Univariate feature selection to 512 features
  • Standard scaling
  • Lightweight PyTorch MLP
  • 20 training epochs
  • CPU execution on Apple Silicon

Held-out performance:

  • Accuracy: 99.38%
  • Macro-F1: 0.9947

Local batch inference benchmark:

  • 200,250 predictions
  • 4.64 seconds
  • Approximately 43,191 predictions/second on CPU

The model is intentionally lightweight. HelixFlow is primarily an infrastructure and reproducibility project: the workload exists to exercise training orchestration, artifact versioning, containerization, Kubernetes scheduling, inference deployment, CI/CD, and cloud infrastructure rather than to establish a state-of-the-art biological classifier.

Reproducible model artifact

Each training run persists the complete preprocessing and inference bundle:

  • model.pt
  • variance_filter.joblib
  • selector.joblib
  • scaler.joblib
  • label_encoder.joblib
  • selected_features.json
  • metadata.json
  • classification_report.json

metadata.json records the Git SHA, PyTorch version, training configuration, execution device, and evaluation metrics.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages