I have run the tone_check_training_dag in staging, and the following tasks succeeded: generate_training_data, split_training_data, and copy_hdfs_to_pvc. The train_tone_check task running on CPU failed with an OOM issue shown in the logs below:

To fix the OOM issue reported in T407212#11271755, I increased container_resources.limits.memory from 8Gi to 16Gi in the train_tone_check task. However, the task running on CPU still failed with the error below:

HTTP response body: {"kind":"Status","apiVersion":"v1","metadata":{},"status":"Failure","message":"pods \"train-tone-check-e7hk664\" is forbidden: [maximum memory usage per Container is 8Gi, but limit is 16Gi, maximum memory usage per Pod is 10Gi, but limit is 17179869184]","reason":"Forbidden","details":{"name":"train-tone-check-e7hk664","kind":"pods"},"code":403}

Full logs can be found here. This indicates that we requested 16Gi for the container, but the cluster only allows up to 8Gi per container and 10Gi per pod.

Following T407212#11275074, I ran the tone-check training job locally with model-ready training data to determine memory usage. Both 8Gi and 16Gi were not enough, as the job required over 18Gi. With a 20Gi memory limit, the job ran on CPU without OOM issues, although it took a very long time to complete.

Since our goal is to run model training on GPU and the fix in T406958 is pending verification, we tested the GPU node selector in T373806#11275873 and confirmed successful GPU access. I then ran the tone-check model training task with the amd.com/gpu.vram=64G node selector, and the job completed in 8hrs on an MI210 GPU with 64G VRAM in staging, as shown in the logs below: