Skip to content

Commit c11861b

Browse files
author
Yadan Wei
committed
add label to build
Signed-off-by: Yadan Wei <yadanwei@amazon.com>
1 parent 70a2130 commit c11861b

File tree

3 files changed

+129
-7
lines changed

3 files changed

+129
-7
lines changed

.github/actions/build-image/action.yml

Lines changed: 117 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,42 @@ inputs:
2929
dockerfile-path:
3030
description: 'Path to the Dockerfile'
3131
required: true
32+
arch-type:
33+
description: 'Architecture type (e.g., x86, arm64)'
34+
required: false
35+
default: 'x86'
36+
device-type:
37+
description: 'Device type (e.g., gpu, cpu)'
38+
required: false
39+
default: 'gpu'
40+
cuda-version:
41+
description: 'CUDA version (e.g., cu129)'
42+
required: false
43+
default: ''
44+
python-version:
45+
description: 'Python version (e.g., py312)'
46+
required: false
47+
default: ''
48+
os-version:
49+
description: 'OS version (e.g., ubuntu24.04)'
50+
required: false
51+
default: ''
52+
contributor:
53+
description: 'Contributor name'
54+
required: false
55+
default: 'None'
56+
inference-toolkit-version:
57+
description: 'Inference toolkit version (e.g., 2.0.25)'
58+
required: false
59+
default: ''
60+
torchserve-version:
61+
description: 'TorchServe version (e.g., 0.12.0)'
62+
required: false
63+
default: ''
64+
transformers-version:
65+
description: 'Transformers library version (e.g., 4.28.1)'
66+
required: false
67+
default: ''
3268

3369
outputs:
3470
image-uri:
@@ -60,16 +96,90 @@ runs:
6096
- name: Build image
6197
shell: bash
6298
run: |
63-
docker buildx build --progress plain \
64-
--build-arg CACHE_REFRESH="$(date +"%Y-%m-%d")" \
65-
--build-arg BASE_IMAGE="${{ inputs.base-image }}" \
66-
--build-arg CONTAINER_TYPE="${{ inputs.container-type }}" \
67-
--build-arg FRAMEWORK="${{ inputs.framework }}" \
68-
--build-arg FRAMEWORK_VERSION="${{ inputs.framework-version }}" \
99+
# Process label values
100+
LABEL_FRAMEWORK=$(echo "${{ inputs.framework }}" | tr '_' '-')
101+
LABEL_FRAMEWORK_VERSION=$(echo "${{ inputs.framework-version }}" | tr '.' '-')
102+
LABEL_ARCH="${{ inputs.arch-type }}"
103+
LABEL_PYTHON_VERSION=$(echo "${{ inputs.python-version }}" | tr '.' '-')
104+
LABEL_OS_VERSION=$(echo "${{ inputs.os-version }}" | tr '.' '-')
105+
LABEL_CONTRIBUTOR="${{ inputs.contributor }}"
106+
107+
# Build transformers label if version provided
108+
LABEL_TRANSFORMERS=""
109+
if [[ -n "${{ inputs.transformers-version }}" ]]; then
110+
LABEL_TRANSFORMERS=$(echo "${{ inputs.transformers-version }}" | tr '.' '-')
111+
fi
112+
113+
# Build inference toolkit label if both versions provided
114+
LABEL_INFERENCE_TOOLKIT=""
115+
if [[ -n "${{ inputs.inference-toolkit-version }}" && -n "${{ inputs.torchserve-version }}" ]]; then
116+
TOOLKIT_VER=$(echo "${{ inputs.inference-toolkit-version }}" | tr '.' '-')
117+
TORCHSERVE_VER=$(echo "${{ inputs.torchserve-version }}" | tr '.' '-')
118+
LABEL_INFERENCE_TOOLKIT="${TOOLKIT_VER}.torchserve.${TORCHSERVE_VER}"
119+
fi
120+
121+
# Construct device type label
122+
LABEL_DEVICE_TYPE="${{ inputs.device-type }}"
123+
if [[ "${{ inputs.device-type }}" == "gpu" && -n "${{ inputs.cuda-version }}" ]]; then
124+
LABEL_DEVICE_TYPE="${{ inputs.device-type }}.${{ inputs.cuda-version }}"
125+
fi
126+
127+
# Build base command
128+
BUILD_CMD="docker buildx build --progress plain \
129+
--build-arg CACHE_REFRESH=\"\$(date +\"%Y-%m-%d\")\" \
130+
--build-arg BASE_IMAGE=\"${{ inputs.base-image }}\" \
131+
--build-arg CONTAINER_TYPE=\"${{ inputs.container-type }}\" \
132+
--build-arg FRAMEWORK=\"${{ inputs.framework }}\" \
133+
--build-arg FRAMEWORK_VERSION=\"${{ inputs.framework-version }}\""
134+
135+
# Add SageMaker labels if target contains 'sagemaker'
136+
if [[ "${{ inputs.target }}" == *"sagemaker"* ]]; then
137+
BUILD_CMD="${BUILD_CMD} \
138+
--label \"com.amazonaws.ml.engines.sagemaker.dlc.arch.${LABEL_ARCH}=true\" \
139+
--label \"com.amazonaws.ml.engines.sagemaker.dlc.device.${LABEL_DEVICE_TYPE}=true\" \
140+
--label \"com.amazonaws.ml.engines.sagemaker.dlc.framework.${LABEL_FRAMEWORK}.${LABEL_FRAMEWORK_VERSION}=true\" \
141+
--label \"com.amazonaws.ml.engines.sagemaker.dlc.job.${{ inputs.container-type }}=true\""
142+
143+
# Add OS version label if provided
144+
if [[ -n "${{ inputs.os-version }}" ]]; then
145+
BUILD_CMD="${BUILD_CMD} \
146+
--label \"com.amazonaws.ml.engines.sagemaker.dlc.os.${LABEL_OS_VERSION}=true\""
147+
fi
148+
149+
# Add Python version label if provided
150+
if [[ -n "${{ inputs.python-version }}" ]]; then
151+
BUILD_CMD="${BUILD_CMD} \
152+
--label \"com.amazonaws.ml.engines.sagemaker.dlc.python.${LABEL_PYTHON_VERSION}=true\""
153+
fi
154+
155+
# Add contributor label if provided
156+
if [[ -n "${LABEL_CONTRIBUTOR}" ]]; then
157+
BUILD_CMD="${BUILD_CMD} \
158+
--label \"com.amazonaws.ml.engines.sagemaker.dlc.contributor.${LABEL_CONTRIBUTOR}=true\""
159+
fi
160+
161+
# Add transformers library label if provided
162+
if [[ -n "${LABEL_TRANSFORMERS}" ]]; then
163+
BUILD_CMD="${BUILD_CMD} \
164+
--label \"com.amazonaws.ml.engines.sagemaker.dlc.lib.transformers.${LABEL_TRANSFORMERS}=true\""
165+
fi
166+
167+
# Add inference toolkit label if provided
168+
if [[ -n "${LABEL_INFERENCE_TOOLKIT}" ]]; then
169+
BUILD_CMD="${BUILD_CMD} \
170+
--label \"com.amazonaws.ml.engines.sagemaker.dlc.inference-toolkit.${LABEL_INFERENCE_TOOLKIT}=true\""
171+
fi
172+
fi
173+
174+
# Complete the build command
175+
BUILD_CMD="${BUILD_CMD} \
69176
--cache-to=type=inline \
70177
--cache-from=type=registry,ref=${CI_IMAGE_URI} \
71178
--tag ${CI_IMAGE_URI} \
72179
--push \
73180
--target ${{ inputs.target }} \
74-
-f ${{ inputs.dockerfile-path }} .
181+
-f ${{ inputs.dockerfile-path }} ."
182+
183+
# Execute build
184+
eval ${BUILD_CMD}
75185
docker rmi ${CI_IMAGE_URI}

.github/workflows/pr-sglang.yml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,12 @@ jobs:
296296
aws-region: ${{ vars.AWS_REGION }}
297297
tag-pr: ${{ env.FRAMEWORK }}-${{ env.SGLANG_VERSION }}-gpu-${{ env.PYTHON_VERSION }}-${{ env.CUDA_VERSION }}-${{ env.OS_VERSION }}-sagemaker-pr-${{ github.event.pull_request.number }}
298298
dockerfile-path: docker/${{ env.FRAMEWORK }}/Dockerfile
299+
arch-type: x86
300+
device-type: gpu
301+
cuda-version: ${{ env.CUDA_VERSION }}
302+
python-version: ${{ env.PYTHON_VERSION }}
303+
os-version: ${{ env.OS_VERSION }}
304+
contributor: None
299305

300306
set-sagemaker-test-environment:
301307
needs: [check-changes, build-sglang-sagemaker-image]

.github/workflows/pr-vllm.yml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -522,6 +522,12 @@ jobs:
522522
aws-region: ${{ vars.AWS_REGION }}
523523
tag-pr: ${{ env.FRAMEWORK }}-${{ env.VLLM_VERSION }}-gpu-${{ env.PYTHON_VERSION }}-${{ env.CUDA_VERSION }}-${{ env.OS_VERSION }}-sagemaker-pr-${{ github.event.pull_request.number }}
524524
dockerfile-path: docker/${{ env.FRAMEWORK }}/Dockerfile
525+
arch-type: x86
526+
device-type: gpu
527+
cuda-version: ${{ env.CUDA_VERSION }}
528+
python-version: ${{ env.PYTHON_VERSION }}
529+
os-version: ${{ env.OS_VERSION }}
530+
contributor: None
525531

526532
set-sagemaker-test-environment:
527533
needs: [check-changes, build-vllm-sagemaker-image]

0 commit comments

Comments
 (0)