From 66b19983b34f82cd5b79c4fad73b5b6e08d8e0fd Mon Sep 17 00:00:00 2001
From: Theodore Vasiloudis <thvasilo@amazon.com>
Date: Thu, 12 Dec 2024 09:37:42 -0800
Subject: [PATCH] Fix for error introduced by DGL's dependencies for Docker
 images (#1112)

*Issue #, if available:*

*Description of changes:*

* With the release of torchdata 0.10.0, installs of DGL==2.3.0 are
broken.
* We pin the version of torchdata in our images to avoid errors.

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.
---
 docker/local/Dockerfile.local       | 2 ++
 docker/parmetis/Dockerfile.parmetis | 1 +
 docker/sagemaker/Dockerfile.sm      | 2 ++
 3 files changed, 5 insertions(+)

diff --git a/docker/local/Dockerfile.local b/docker/local/Dockerfile.local
index 3886ac1c6c..4b48bd2234 100644
--- a/docker/local/Dockerfile.local
+++ b/docker/local/Dockerfile.local
@@ -60,6 +60,7 @@ RUN TORCH_MAJOR_MINOR=$(echo $TORCH_VERSION | cut -c1-3) && \
     pip install \
     dgl==${DGL_VERSION} -f https://data.dgl.ai/wheels/torch-${TORCH_MAJOR_MINOR}/repo.html \
     ogb==${OGB_VERSION} \
+    torchdata==0.9.0 \
     transformers==${TRANSFORMERS_VERSION} \
     && rm -rf /root/.cache
 
@@ -72,6 +73,7 @@ RUN TORCH_MAJOR_MINOR=$(echo $TORCH_VERSION | cut -c1-3) && \
     -f https://data.dgl.ai/wheels/torch-${TORCH_MAJOR_MINOR}/cu${DGL_CUDA_VERSION}/repo.html \
     ogb==${OGB_VERSION} \
     torch==${TORCH_VERSION} \
+    torchdata==0.9.0 \
     transformers==${TRANSFORMERS_VERSION} \
     && rm -rf /root/.cache
 
diff --git a/docker/parmetis/Dockerfile.parmetis b/docker/parmetis/Dockerfile.parmetis
index 9b16c49217..7ee971c829 100644
--- a/docker/parmetis/Dockerfile.parmetis
+++ b/docker/parmetis/Dockerfile.parmetis
@@ -40,6 +40,7 @@ RUN pip install \
     pydantic==2.7.0 \
     scikit-learn==1.4.2 \
     scipy==1.13.0 \
+    torchdata==0.9.0 \
     pyyaml \
     && rm -rf /root/.cache
 
diff --git a/docker/sagemaker/Dockerfile.sm b/docker/sagemaker/Dockerfile.sm
index 5deb65d4b1..ba96f4deef 100644
--- a/docker/sagemaker/Dockerfile.sm
+++ b/docker/sagemaker/Dockerfile.sm
@@ -28,8 +28,10 @@ RUN pip3 install \
         numpy==1.26.4 \
         ogb==1.3.6 \
         pyarrow==16.1.0 \
+        pydantic==2.7.1 \
         scikit-learn==1.5.0 \
         scipy==1.13.1 \
+        torchdata==0.9.0 \
         transformers==4.28.1 \
     && rm -rf /root/.cache