Bootstrap: docker
From: nvidia/cuda:12.4.0-runtime-ubuntu22.04

%post
	cp -a /etc/apt/sources.list /etc/apt/sources.list.bak
	sed -i "s@http://.*archive.ubuntu.com@http://mirrors.huaweicloud.com@g" /etc/apt/sources.list
	sed -i "s@http://.*security.ubuntu.com@http://mirrors.huaweicloud.com@g" /etc/apt/sources.list
	apt-get update && apt-get install -y python3 python3-pip python-is-python3 git wget build-essential
	pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
	pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu124
	pip install torch_geometric
	pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.5.1+cu124.html
    pip install lightning litdata mosaicml-streaming torchmetrics polars optuna optuna-dashboard zstd pyarrow scanpy scib-metrics leidenalg pacmap scipy scikit-learn tensorboard