An automated machine learning modeling pipeline (SurvAutoML) designed to train, evaluate, and track survival models on tabular and clinical time-to-event datasets. While originally built to ingest extracted features from our kidney transplantation data, the pipeline architecture is data-agnostic and can be applied to generic clinical survival analysis tasks.
This project was developed as part of the IDEN project, which ended in 2026, and the code is provided here for archival, reproducibility, and general use.
- Data-Agnostic AutoML: Accepts generic tabular survival datasets; easy to run out of the box using the provided synthetic example dataset without requiring access to restricted data.
- Multi-Model Survival Analysis: Automatically trains and evaluates a diverse suite of survival algorithms including Cox Regression, DeepSurv, Random Survival Forests, and Survival Gradient Boosting.
- Integrated Experiment Tracking: Utilizes a local, Docker-orchestrated MLflow instance (backed by PostgreSQL) to track runs, hyperparameter sweeps, evaluation metrics, and model artifacts seamlessly.
- ICD-10-GM Embeddings Support: Features an optional data preparation step to map and embed diagnostic codes using pre-trained transformer models prior to model training.
- Streamlined Export & Deployment: Includes automated scripts to select and export the best-performing models. Exported models are designed to be deployed instantly using our companion project, SurvAutoMLDeploy.
- Reproducible Infrastructure: Fully containerized development environment
utilizing Devcontainers, strictly pinned micromamba/conda environments,
and a
justtask runner for consistent development.
- Experiment Tracking: MLflow
- Task Runner: Just
- Data & Model Versioning: DVC
- Modeling & Data Processing: Python, PyTorch (CUDA-enabled), Pandas
- Environment: Docker (Devcontainer + Docker Compose) & micromamba
TxReg_Modellierung/
├── MLproject # MLflow project definitions, parameters, and entry points
├── envs/ # Conda environment specifications and lock files
├── scratching/ # Exploratory Jupyter notebooks
├── src/
│ └── python/ # Core pipeline scripts (train, eval, export, embeddings)
├── tests/ # Unit and integration tests (pytest)
├── workdir/ # Target directory for input data, embeddings, logs
│ ├── results/ # ⚠️ Local artifact output (DVC-tracked)
│ └── *.dvc # DVC tracking files for input features and embeddings
└── .devcontainer/ # Devcontainer setup
# (spins up MLflow alongside the dev environment)
This project uses a Devcontainer coupled with Docker Compose, which completely automates the setup of the Python environment, GPU support, and the local MLflow tracking server.
- Install Docker and VS Code.
- Install the Dev Containers extension in VS Code.
- Clone this repository and open it in VS Code. Prompt it to "Reopen in Container".
- Note: The Devcontainer automatically spins up the MLflow UI. You can access it at http://localhost:5000 once the container is running.
- Alternatively, if not using the devcontainer, ensure you have
micromambainstalled, build the environment fromenvs/main-lock.yml, and manually configure theconda_envkey in theMLprojectfile.
Runs and artifact generation are managed through MLflow commands.
-
Provide the Data: Place your survival dataset into
workdir/. -
Provide the ICD-10 Catalogue (Optional): If generating diagnostic embeddings, place
ICD-10-GM-Katalog_202022.xlsx(available from BfArM / BIFG) intoworkdir/. -
Generate Embeddings (Optional): (
workdir/example/icd10_gm_embeddings.npz)mkdir -p workdir/example mlflow run . --experiment-name "Embedding" -e prepare_embedding \ -P datafile=example_data/txreg_artificial_dataset.csv.gz \ -P output_dir=./workdir/example \ -P catalogue="./workdir/ICD-10-GM-Katalog_202022.xlsx"
-
Train Models: Execute training runs for specific models. Track your runs live in the MLflow UI.
export ICD10_EMBEDDING_PATH=$(realpath workdir/example/icd10_gm_embeddings.npz) datafile="example_data/txreg_artificial_dataset.csv.gz" # datafile="workdir/metabric_data.csv" experiment_name="ExampleRun" trails=5 # Number of hyperparameter trials, should be around 100 or more mkdir -p workdir/example/logs # Run one model mlflow run . \ -e train_eval \ -P datafile="$datafile" \ -P modelname=cox \ -P trials="$trails" \ --experiment-name "$experiment_name" \ 2>&1 | tee workdir/example/logs/cox.log # Other options for modelname: age_cox, coxbackward, deepsurv, km, # randomsurvforest, ranger, survgradientboosting, survreg, # ipcwlogisticreg
-
Export the Best Model: Locate the parent run ID from the MLflow UI or logs, and export the model artifacts for deployment via SurvAutoMLDeploy.
# download the model with reduced features? - featureselection="True" mlflow run . --experiment-name "Orga" -e export_model \ -P output_dir=workdir/example/cox \ -P featureselection="True" \ -P parent_run="123456789abcdef" # id of the train_eval run
The runs follow the following structure. The parent run is the train_eval
run. It has three children: train and two eval. The train run stores the
the hyperparameter search results (see the artifacts) and the feature
selection resaults. It produces a model artifact with all features and
one with a reduced feature set. For each trial run by the hyperparameter search,
there is a child run of train that stores the results across the CV folds.
For each of the two model artifacts, there is a child run of eval that stores
the evaluation on the test set. It also shows the feature importance
and the SHAP results.
See the scratch_notebook directory and the export_model run to see
how to programmtically use the data in mlflow.
Instead of the synthetic data, you can also use other data like the datasets from the pycox library or your own.
After running just run python:
from pycox.datasets import metabric
df = metabric.read_df()
# columns need to be named time and event
df.rename(columns={"duration": "time", "event": "event"}, inplace=True)
df["time"] = df["time"].round().astype(int) # integer times values work better
df.to_csv("workdir/metabric_data.csv", index=False)
# Ctrl+DIf you hit "ValueError: time must be smaller than largest observed time point" from cumulative_dynamic_auc (or concordance_index_ipcw / brier_score / integrated_brier_score), this is not a bug in the CV setup — it's an inherent constraint of IPCW-based metrics in sksurv.
These metrics estimate a censoring distribution (IPCW) from the training fold and then evaluate it at the test fold's observed times. If any test time exceeds the largest time seen in that training fold, the censoring distribution has no support there and IPCW cannot be computed, so sksurv raises instead of extrapolating. You need to censor your dataset at using administrative censoring.
See this for Instructions on how to use DVC at the h_da.
Note that this repository expects to be next to the TxRegAnalysis repository,
(same parent directory) as it uses DVC to pull the input data from there.
The input and output data can be then be restored with dvc pull and
dvc checkout.
We use just as a command runner to simplify development tasks and tests.
# Check code for style and syntax issues (flake8, pylint)
just lint
# Automatically format code (isort, autoflake, black)
just format
# Run the test suite and generate a coverage report
just test
# Update locked dependencies and render environment files
just lockdepsThis code is licensed under the MIT License. See the LICENSE file for details.
Data Notice: While this framework was initially developed to model sensitive kidney transplant data from the Deutsches Transplantationsregister, the pipeline itself is completely data-agnostic. No confidential patient data is included in this repository.
To explore or verify the pipeline without legal access to the registry data, please use the provided which reproduces key feature structures without exposing real patient relationships.
This research project was funded by the Federal Ministry of Research, Technology and Space (Bundesministerium für Forschung, Technologie und Raumfahrt, BMFTR) under project grant 13FH019KX1 and the German federal state of Hesse.
Further details regarding the project can be found on the FORSCHUNG.HAW project page.