Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

GPUContainerImage schema with os, arch and cache info - move GPU container images to config #5153

Open
wants to merge 13 commits into
base: master
Choose a base branch
from
+149 −19
Open

GPUContainerImage schema with os, arch and cache info - move GPU container images to config #5153

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
4ddd6d9
Add new schema for GPU images
ganeshkumarashok Oct 22, 2024
cb2439d
update schema and version
ganeshkumarashok Oct 23, 2024
0729dd8
Merge github.com:Azure/AgentBaker into aganeshkumar/move_gpu_img_config
ganeshkumarashok Oct 23, 2024
6d11f6a
update schema
ganeshkumarashok Oct 23, 2024
efaf51e
remove repeated version
ganeshkumarashok Oct 23, 2024
2ced72e
Consume GPU version from components.json
ganeshkumarashok Oct 23, 2024
b2cc631
Merge branch 'master' of github.com:Azure/AgentBaker into aganeshkuma…
ganeshkumarashok Oct 24, 2024
a28eab2
New GPUContainerImage schema with os, arch and cache info
ganeshkumarashok Oct 24, 2024
d526a55
Merge branch 'master' of github.com:Azure/AgentBaker into aganeshkuma…
ganeshkumarashok Oct 24, 2024
4076236
Add Readme
ganeshkumarashok Oct 24, 2024
034864b
Use updateMultiArchVersions
ganeshkumarashok Oct 24, 2024
b3c5c41
indent
ganeshkumarashok Oct 24, 2024
60d3638
Constrain to one gpuImageVersion per GPUContainerImage
ganeshkumarashok Oct 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions parts/linux/cloud-init/artifacts/components.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -558,6 +558,36 @@
]
}
],

"GPUContainerImages": [
{
"downloadURL": "mcr.microsoft.com/aks/aks-gpu-cuda:*",
"multiArchVersionsV2": [
{
"renovateTag": "registry=https://mcr.microsoft.com, name=aks/aks-gpu-cuda",
"latestVersion": "550.90.12-20241021233454"
}
],
"cached": true,
"osSelectors":[
{
"os": "ubuntu",
"arch": "amd64"
}
]
},
{
"downloadURL": "mcr.microsoft.com/aks/aks-gpu-grid:*",
"multiArchVersionsV2": [
{
"renovateTag": "registry=https://mcr.microsoft.com, name=aks/aks-gpu-grid",
"latestVersion": "535.161.08-20241021235607"
}
],
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

indent

"cached": false
}
],

"Packages": [
{
"name": "oras",
Expand Down
17 changes: 16 additions & 1 deletion schemas/components.cue
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,27 @@ package components
previousLatestVersion?: #ContainerImagePrefetchOptimization
}

#OSSelector: {
os: string
arch: string
}


#ContainerImage: {
downloadURL: string
amd64OnlyVersions: [...string]
multiArchVersionsV2: [...#VersionV2]
}

#GPUContainerImage: {
downloadURL: string
multiArchVersionsV2: [...#VersionV2]
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any recommendations for alternate names for multiArchVersionsV2? Thought of Version but VersionV2 already exists

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

gpuImageVersion?
I 'd love to have GPU in capital letters but it's not following the naming convention.

cached: bool
osSelectors?: [...#OSSelector]
}

#Images: [...#ContainerImage]
#GPUImages: [...#GPUContainerImage]
#Packages: [...#Package]
#VersionV2: {
k8sVersion?: string
Expand Down Expand Up @@ -67,7 +81,8 @@ package components

#Components: {
ContainerImages: #Images
Packages: #Packages
Packages: #Packages
GPUContainerImages?: #GPUImages
}

#Components
112 changes: 94 additions & 18 deletions vhdbuilder/packer/install-dependencies.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -351,28 +351,104 @@ INSTALLED_RUNC_VERSION=$(runc --version | head -n1 | sed 's/runc version //')
echo " - runc version ${INSTALLED_RUNC_VERSION}" >> ${VHD_LOGS_FILEPATH}
capture_benchmark "${SCRIPT_NAME}_artifact_streaming_download"

if [[ $OS == $UBUNTU_OS_NAME && $(isARM64) != 1 ]]; then # no ARM64 SKU with GPU now
gpu_action="copy"
NVIDIA_DRIVER_IMAGE_SHA="20241008175307"
export NVIDIA_DRIVER_IMAGE_TAG="550.90.12-${NVIDIA_DRIVER_IMAGE_SHA}"
NVIDIA_DRIVER_IMAGE="mcr.microsoft.com/aks/aks-gpu-cuda"

mkdir -p /opt/{actions,gpu}
ctr -n k8s.io image pull $NVIDIA_DRIVER_IMAGE:$NVIDIA_DRIVER_IMAGE_TAG
if grep -q "fullgpu" <<< "$FEATURE_FLAGS"; then
bash -c "$CTR_GPU_INSTALL_CMD $NVIDIA_DRIVER_IMAGE:$NVIDIA_DRIVER_IMAGE_TAG gpuinstall /entrypoint.sh install"
ret=$?
if [[ "$ret" != "0" ]]; then
echo "Failed to install GPU driver, exiting..."
exit $ret
fi
gpu_action=""
declare -A pulled_gpu_images

# Loop over each GPUContainerImage
while IFS= read -r gpuImageToBePulled; do
# Extract 'cached' field and convert it to lowercase
cached=$(echo "${gpuImageToBePulled}" | jq -r '.cached' | tr '[:upper:]' '[:lower:]')

if [[ "$cached" != "true" ]]; then
# Skip images that are not meant to be cached
continue
fi

cat << EOF >> ${VHD_LOGS_FILEPATH}
- nvidia-driver=${NVIDIA_DRIVER_IMAGE_TAG}
EOF
# Extract 'osSelectors' if present
osSelectors=$(echo "${gpuImageToBePulled}" | jq -r '.osSelectors // empty')

shouldPull=0 # Default to not pull

if [[ -n "$osSelectors" ]]; then
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested to put this logic into a function and add unit tests to cover most of the if conditions, so that we don't need to rely on abe2e or RP-e2e to capture issues for us.
One way to do that is put the function into cse_helpers.sh. There is a shellspec unit test file cse_helpers.sh which have some examples to author tests.
The root level readme has some instructions too.

# osSelectors is provided; check if current OS and arch match any entry
while IFS= read -r selector; do
os=$(echo "$selector" | jq -r '.os')
arch=$(echo "$selector" | jq -r '.arch')

if [[ "$os" == "$CURRENT_OS" ]]; then
if [[ "$arch" == "$CPU_ARCH" ]]; then
ganeshkumarashok marked this conversation as resolved.
Show resolved Hide resolved
shouldPull=1
break # Found a matching selector
fi
fi
done <<< "$(echo "$osSelectors" | jq -c '.[]')"
else
# No osSelectors provided; decide whether to pull
# Assuming we pull the image if no osSelectors are specified
shouldPull=1
fi

if [[ "$shouldPull" == "1" ]]; then
# Extract image details
downloadURL=$(echo "${gpuImageToBePulled}" | jq -r '.downloadURL')
imageName=$(echo "$downloadURL" | sed 's/:.*$//')

# Get the latestVersion
latestVersion=$(echo "${gpuImageToBePulled}" | jq -r '.multiArchVersionsV2[0].latestVersion')
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you only put [0], later on when you add another multiArchVersionsV2[1] block in components.json, it won't capture it.
Suggested to loop the whole multiArchVersionsV2 array. And if you only have 1 multiArchVersionsV2 entry in components.json, it will just work like using [0].

Copy link
Collaborator

@Devinwong Devinwong Oct 24, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 example for your reference:
At around ln414.

  MULTI_ARCH_VERSIONS=()
  updateMultiArchVersions "${imageToBePulled}"

Actually you can reuse the function updateMultiArchVersions

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, updated it to reuse that function

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually, under no circumstances would we need to pull two CUDA images for the same VHD

Copy link
Collaborator

@Devinwong Devinwong Oct 24, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If that's the case, maybe create a gpuImageVersion and it's just simple struct instead of an array of struct? Then it's much cleaner and explicit enough. Otherwise one day someone will ask, why we were only picking the first element of an array, but we didn't just create a non-array struct? Does that sound making sense to you? In terms of implementation, your way is absolutely right. I was just thinking 1 more step beyond for future readability and maintainability. When I was migrating to components.json, I had a lot of such questions. 😂

Copy link
Contributor Author

@ganeshkumarashok ganeshkumarashok Oct 24, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense, agreed. Changed it to so we constrain it to one gpuImageVersion per GPUContainerImage:

#GPUContainerImage: {
downloadURL: string
gpuIVersion: #VersionV2
cached: bool
osSelectors?: [...#OSSelector]
}


if [[ -z "$latestVersion" || "$latestVersion" == "null" ]]; then
echo "Error: latestVersion not found for $imageName"
exit 1
fi

fullImage="$imageName:$latestVersion"

# Pull the image
echo "Pulling image: $fullImage"
ctr -n k8s.io image pull "$fullImage"
if [[ $? -ne 0 ]]; then
echo "Failed to pull image: $fullImage"
exit 1
fi

# Record the pulled image
pulled_gpu_images["$imageName"]="$latestVersion"

# Set gpu_action if pulling the aks-gpu-cuda image
if [[ "$imageName" == "mcr.microsoft.com/aks/aks-gpu-cuda" ]]; then
gpu_action="copy"

# Create necessary directories
mkdir -p /opt/{actions,gpu}

# Check for the "fullgpu" feature flag
if grep -q "fullgpu" <<< "$FEATURE_FLAGS"; then
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Again, avoid more than 2 level nested if. It's hard to keep track which level it is for debugging and readability.

Copy link
Contributor Author

@ganeshkumarashok ganeshkumarashok Oct 24, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree - thinking about the alternate way.

But I think this approach is making it a lot more complex than the alternate PR, which is much smaller: #5138

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right. General vs flexible is always a trade-off. If it can fit your mid-future GPU images, I am fine with it too as this will be used by GPU container images.

echo "Installing GPU driver from image: $fullImage"
bash -c "$CTR_GPU_INSTALL_CMD $fullImage gpuinstall /entrypoint.sh install"
ret=$?
if [[ "$ret" != "0" ]]; then
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Again, avoid more than 2 level nested if. It's hard to keep track which level it is for debugging and readability.

echo "Failed to install GPU driver, exiting..."
exit $ret
fi
fi
fi
else
echo "Skipping image $imageName due to osSelector constraints or cached=false."
fi
done <<< "$GPUContainerImages"

# Log the pulled images
if [[ "${#pulled_gpu_images[@]}" -gt 0 ]]; then
echo "Logging pulled GPU images to $VHD_LOGS_FILEPATH"
for imageName in "${!pulled_gpu_images[@]}"; do
imageVersion=${pulled_gpu_images[$imageName]}
echo " - $imageName=$imageVersion" >> "$VHD_LOGS_FILEPATH"
done
else
echo "No GPU images were pulled."
fi


ls -ltr /opt/gpu/* >> ${VHD_LOGS_FILEPATH}

installBpftrace
Expand Down
Loading