Abstract
Modern machine learning pipelines increasingly rely on heterogeneous task graphs that mix data preprocessing, feature engineering, model training, evaluation, and serving-oriented transformations. These pipelines exhibit nontrivial interactions between data locality, accelerator topology, and runtime variability, particularly when multiple workflows contend for shared GPU resources. A scheduler that ignores locality may overuse remote storage paths or interconnect links, while a scheduler that ignores GPU affinity may incur avoidable device-to-device transfers and kernel launch serialization. At the same time, long-running GPU-bound stages can reduce cluster responsiveness unless preemption is supported in a manner that respects the limited preemptibility of GPU execution and the high cost of state migration. This paper presents a task-graph scheduling approach for ML pipelines that jointly targets data locality, GPU affinity, and practical preemption. The design models each task’s data footprint, read and write dependencies, and preferred device neighborhoods, while representing accelerator interconnects and storage tiers as a weighted topology graph. Scheduling decisions combine criticality-aware ordering with topology-aware placement and an explicit accounting of preemption overhead. Preemption is realized at safe boundaries with lightweight checkpointing and bounded rollback, enabling latency-sensitive jobs to obtain service without destabilizing throughput-oriented workloads. The proposed framework is described as a set of algorithms and runtime mechanisms that can be implemented over common cluster managers and GPU runtimes. Analytical considerations and experimental methodology are provided to characterize the scheduler’s efficiency, fairness, and sensitivity to pipeline structure and contention.