Skip to content

Commit 34f0ac8

Browse files
hooray03hooray03
committed
ci(metax): add image contracts and benchmark gate
1 parent d9a22d3 commit 34f0ac8

15 files changed

Lines changed: 465 additions & 1 deletion

File tree

‎.github/configs/metax.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,4 +65,5 @@ env_names:
6565
train: "base"
6666
hetero_train: ""
6767
inference: ""
68+
serve: ""
6869
rl: ""

‎.github/workflows/build_image_metax.yml‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
# Copyright 2026 FlagOS Contributors
2+
# Licensed under the Apache License, Version 2.0.
3+
4+
name: Build Docker Images - MetaX
5+
6+
on:
7+
workflow_dispatch:
8+
inputs:
9+
task:
10+
description: MetaX image to build
11+
required: true
12+
type: choice
13+
options: [train, inference, all]
14+
default: train
15+
push:
16+
description: Push candidate to shared Harbor
17+
required: true
18+
type: boolean
19+
default: true
20+
run_tests:
21+
description: Run the current MetaX matrix after build
22+
required: true
23+
type: boolean
24+
default: true
25+
no_cache:
26+
description: Disable Docker cache
27+
required: true
28+
type: boolean
29+
default: false
30+
tag:
31+
description: Optional tag override
32+
required: false
33+
type: string
34+
base_image:
35+
description: Validated MetaX base image (prefer an immutable digest)
36+
required: true
37+
type: string
38+
default: harbor.baai.ac.cn/flagscale/megatron-lm-with-te:202603231839
39+
40+
permissions:
41+
contents: read
42+
43+
concurrency:
44+
group: build-image-metax-${{ github.ref }}-${{ inputs.task }}
45+
cancel-in-progress: false
46+
47+
env:
48+
REGISTRY: harbor.baai.ac.cn
49+
REGISTRY_NAMESPACE: flagos-dev
50+
CURRENT_TRAIN_IMAGE: harbor.baai.ac.cn/flagscale/megatron-lm-with-te:202603231839
51+
CURRENT_INFERENCE_IMAGE: harbor.baai.ac.cn/flagscale/megatron-lm-with-te:202603231839
52+
53+
jobs:
54+
prepare:
55+
name: Prepare MetaX candidate
56+
runs-on: ubuntu-latest
57+
outputs:
58+
matrix: ${{ steps.params.outputs.matrix }}
59+
candidate: ${{ steps.params.outputs.candidate }}
60+
train_image: ${{ steps.params.outputs.train_image }}
61+
inference_image: ${{ steps.params.outputs.inference_image }}
62+
steps:
63+
- name: Resolve build parameters
64+
id: params
65+
shell: bash
66+
run: |
67+
set -euo pipefail
68+
short_sha="${GITHUB_SHA::7}"
69+
tag='${{ inputs.tag }}'
70+
if [ -z "$tag" ]; then
71+
tag="${short_sha}-metax-dev"
72+
fi
73+
candidate="${REGISTRY}/${REGISTRY_NAMESPACE}/flagscale-${{ inputs.task }}:${tag}"
74+
train_image="$CURRENT_TRAIN_IMAGE"
75+
inference_image="$CURRENT_INFERENCE_IMAGE"
76+
case '${{ inputs.task }}' in
77+
train) train_image="$candidate" ;;
78+
inference) inference_image="$candidate" ;;
79+
all) train_image="$candidate"; inference_image="$candidate" ;;
80+
esac
81+
echo 'matrix={"task":["${{ inputs.task }}"]}' >> "$GITHUB_OUTPUT"
82+
echo "candidate=$candidate" >> "$GITHUB_OUTPUT"
83+
echo "train_image=$train_image" >> "$GITHUB_OUTPUT"
84+
echo "inference_image=$inference_image" >> "$GITHUB_OUTPUT"
85+
86+
runtime_contract:
87+
name: Validate MetaX base runtime
88+
needs: prepare
89+
runs-on: flagscale-metax-c550-gpu2-8c-256g
90+
steps:
91+
- name: Validate devices and framework imports
92+
shell: bash
93+
run: |
94+
set -euo pipefail
95+
base_image='${{ inputs.base_image }}'
96+
docker pull "$base_image"
97+
docker run --rm \
98+
--ipc=host --group-add video \
99+
--device=/dev/dri --device=/dev/mxcd --device=/dev/infiniband \
100+
--entrypoint bash "$base_image" -lc '
101+
set -euo pipefail
102+
python - <<"PY"
103+
import torch
104+
import transformer_engine
105+
from megatron.core.models.gpt import GPTModel
106+
107+
count = torch.cuda.device_count()
108+
print("torch:", torch.__version__)
109+
print("devices:", count)
110+
print("transformer_engine:", transformer_engine.__file__)
111+
print("megatron GPTModel:", GPTModel)
112+
assert count == 8, f"expected 8 MetaX devices, found {count}"
113+
PY
114+
cat >/tmp/metax_collective.py <<"PY"
115+
import os
116+
import torch
117+
import torch.distributed as dist
118+
119+
rank = int(os.environ["LOCAL_RANK"])
120+
torch.cuda.set_device(rank)
121+
dist.init_process_group("nccl")
122+
value = torch.tensor([rank + 1.0], device=f"cuda:{rank}")
123+
dist.all_reduce(value)
124+
assert value.item() == 3.0, value
125+
print(f"rank={rank} all_reduce={value.item()}")
126+
dist.destroy_process_group()
127+
PY
128+
torchrun --standalone --nnodes=1 --nproc-per-node=2 /tmp/metax_collective.py
129+
'
130+
131+
- name: Validate MetaX vLLM contract
132+
if: inputs.task == 'inference' || inputs.task == 'all'
133+
shell: bash
134+
run: |
135+
set -euo pipefail
136+
base_image='${{ inputs.base_image }}'
137+
docker run --rm \
138+
--ipc=host --group-add video \
139+
--device=/dev/dri --device=/dev/mxcd \
140+
--entrypoint python "$base_image" -c '
141+
import importlib.metadata as metadata
142+
import vllm
143+
print("vllm:", metadata.version("vllm"))
144+
print("module:", vllm.__file__)
145+
'
146+
147+
build:
148+
name: Build MetaX candidate
149+
needs: [prepare, runtime_contract]
150+
uses: ./.github/workflows/build_image_common.yml
151+
secrets: inherit
152+
with:
153+
platform: metax
154+
runs_on: '["flagscale-metax-c550-gpu2-8c-256g"]'
155+
matrix: ${{ needs.prepare.outputs.matrix }}
156+
target: dev
157+
build_images: true
158+
no_cache: ${{ inputs.no_cache }}
159+
push: ${{ inputs.push }}
160+
registry: harbor.baai.ac.cn
161+
train_image: ${{ needs.prepare.outputs.candidate }}
162+
inference_image: ${{ needs.prepare.outputs.candidate }}
163+
all_image: ${{ needs.prepare.outputs.candidate }}
164+
base_image: ${{ inputs.base_image }}
165+
build_args: BASE_IMAGE=${{ inputs.base_image }}
166+
167+
tests:
168+
name: Test MetaX candidate
169+
needs: [prepare, build]
170+
if: inputs.push && inputs.run_tests
171+
uses: ./.github/workflows/all_tests_common.yml
172+
secrets: inherit
173+
with:
174+
platform: metax
175+
ci_train_image: ${{ needs.prepare.outputs.train_image }}
176+
ci_inference_image: ${{ needs.prepare.outputs.inference_image }}

‎docker/metax/Dockerfile.all‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
# FlagScale combined train/inference image for MetaX C550. This image is valid
2+
# only when the vendor training runtime and MetaX vLLM coexist in one base.
3+
ARG BASE_IMAGE=harbor.baai.ac.cn/flagscale/megatron-lm-with-te:202603231839
4+
ARG FLAGSCALE_HOME=/opt/flagscale
5+
ARG FLAGSCALE_CONDA=/opt/conda
6+
7+
FROM ${BASE_IMAGE} AS base
8+
9+
ARG FLAGSCALE_HOME
10+
ARG FLAGSCALE_CONDA
11+
ENV FLAGSCALE_HOME=${FLAGSCALE_HOME} \
12+
FLAGSCALE_CONDA=${FLAGSCALE_CONDA} \
13+
PATH="${FLAGSCALE_CONDA}/bin:${PATH}"
14+
15+
COPY tools/install/metax/env.sh /etc/profile.d/flagscale-env.sh
16+
WORKDIR /workspace
17+
18+
FROM base AS dev
19+
COPY tools/install /workspace/tools/install
20+
COPY requirements /workspace/requirements
21+
RUN chmod +x /workspace/tools/install/*.sh \
22+
/workspace/tools/install/utils/*.sh \
23+
/workspace/tools/install/metax/*.sh && \
24+
/workspace/tools/install/install.sh --platform metax --task all \
25+
--pkg-mgr pip --no-system --no-dev && \
26+
rm -rf /workspace/tools /workspace/requirements
27+
CMD ["/bin/bash"]
28+
29+
FROM dev AS release
30+
CMD ["/bin/bash"]

‎docker/metax/Dockerfile.inference‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
# FlagScale inference image for MetaX C550.
2+
ARG BASE_IMAGE=harbor.baai.ac.cn/flagscale/megatron-lm-with-te:202603231839
3+
ARG FLAGSCALE_HOME=/opt/flagscale
4+
ARG FLAGSCALE_CONDA=/opt/conda
5+
6+
FROM ${BASE_IMAGE} AS base
7+
8+
ARG FLAGSCALE_HOME
9+
ARG FLAGSCALE_CONDA
10+
ENV FLAGSCALE_HOME=${FLAGSCALE_HOME} \
11+
FLAGSCALE_CONDA=${FLAGSCALE_CONDA} \
12+
PATH="${FLAGSCALE_CONDA}/bin:${PATH}"
13+
14+
COPY tools/install/metax/env.sh /etc/profile.d/flagscale-env.sh
15+
WORKDIR /workspace
16+
17+
FROM base AS dev
18+
COPY tools/install /workspace/tools/install
19+
COPY requirements /workspace/requirements
20+
RUN chmod +x /workspace/tools/install/*.sh \
21+
/workspace/tools/install/utils/*.sh \
22+
/workspace/tools/install/metax/*.sh && \
23+
/workspace/tools/install/install.sh --platform metax --task inference \
24+
--pkg-mgr pip --no-system --no-dev && \
25+
rm -rf /workspace/tools /workspace/requirements
26+
CMD ["/bin/bash"]
27+
28+
FROM dev AS release
29+
CMD ["/bin/bash"]

‎docker/metax/Dockerfile.train‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
# FlagScale training image for MetaX C550.
2+
ARG BASE_IMAGE=harbor.baai.ac.cn/flagscale/megatron-lm-with-te:202603231839
3+
ARG FLAGSCALE_HOME=/opt/flagscale
4+
ARG FLAGSCALE_CONDA=/opt/conda
5+
6+
FROM ${BASE_IMAGE} AS base
7+
8+
ARG FLAGSCALE_HOME
9+
ARG FLAGSCALE_CONDA
10+
ENV FLAGSCALE_HOME=${FLAGSCALE_HOME} \
11+
FLAGSCALE_CONDA=${FLAGSCALE_CONDA} \
12+
PATH="${FLAGSCALE_CONDA}/bin:${PATH}"
13+
14+
COPY tools/install/metax/env.sh /etc/profile.d/flagscale-env.sh
15+
WORKDIR /workspace
16+
17+
FROM base AS dev
18+
COPY tools/install /workspace/tools/install
19+
COPY requirements /workspace/requirements
20+
RUN chmod +x /workspace/tools/install/*.sh \
21+
/workspace/tools/install/utils/*.sh \
22+
/workspace/tools/install/metax/*.sh && \
23+
/workspace/tools/install/install.sh --platform metax --task train \
24+
--pkg-mgr pip --no-system --no-dev && \
25+
rm -rf /workspace/tools /workspace/requirements
26+
CMD ["/bin/bash"]
27+
28+
FROM dev AS release
29+
CMD ["/bin/bash"]

‎requirements/metax/base.txt‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Copyright 2026 FlagOS Contributors
2+
# Licensed under the Apache License, Version 2.0.
3+
4+
# MetaX platform dependencies. The base image owns the vendor Torch runtime.
5+
-r ../common.txt
6+
7+
datasets==4.5.0
8+
diffusers==0.36.0
9+
draccus==0.11.5
10+
einops==0.8.2
11+
epath==0.7
12+
flask_cors==6.0.2
13+
multi-storage-client==0.42.0
14+
qwen_vl_utils==0.0.14
15+
webdataset==1.0.2

‎requirements/metax/inference.txt‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
# Copyright 2026 FlagOS Contributors
2+
# Licensed under the Apache License, Version 2.0.
3+
4+
# vLLM must be supplied by the validated MetaX base image. Do not install a
5+
# CUDA wheel or replace the vendor Torch runtime from this contract.
6+
-r ./base.txt
7+
8+
sentencepiece==0.2.1
9+
transformers==4.57.6
10+
tiktoken==0.12.0

‎requirements/metax/serve.txt‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
# Copyright 2026 FlagOS Contributors
2+
# Licensed under the Apache License, Version 2.0.
3+
4+
-r ./inference.txt
5+
6+
websocket-client==1.8.0
7+
websockets==15.0.1
8+
msgpack==1.1.0

‎requirements/metax/train.txt‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
# Copyright 2026 FlagOS Contributors
2+
# Licensed under the Apache License, Version 2.0.
3+
4+
-r ./base.txt
5+
6+
sentencepiece==0.2.1
7+
transformers==4.57.6
8+
tiktoken==0.12.0
9+
peft==0.18.1

‎tests/test_utils/config/platforms/metax.yaml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,8 @@ c550:
3838
train:
3939
qwen3: ["0_6b_metax"]
4040
hetero_train: {}
41-
benchmark: {} # qwen3: ["0_6b_metax"]
41+
benchmark:
42+
qwen3: ["0_6b_metax"]
4243
inference: {}
4344
serve: {}
4445
unit:

0 commit comments

Comments
 (0)