Skip to content

Commit 87abf1c

Browse files
hooray03hooray03
committed
feat(metax): add inference image contract
1 parent d8b9e89 commit 87abf1c

7 files changed

Lines changed: 218 additions & 1 deletion

File tree

‎.github/configs/metax.yml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,13 @@ image_build:
5151
source_refs:
5252
FLAGSCALE_MEGATRON_REF: megatron_lm_fl
5353
FLAGSCALE_TE_REF: transformer_engine_fl
54+
inference:
55+
dockerfile: docker/metax/Dockerfile.inference
56+
image: flagos-dev/flagscale-inference
57+
test_roles: [inference]
58+
base_image: harbor.baai.ac.cn/flagos-dev/vllm-plugin-fl@sha256:c83a7b26e5b72fcffa920f32929a76984c78d5ee5520828b6e50feb9864240ed
59+
source_refs:
60+
VLLM_PLUGIN_FL_REF: vllm_plugin_fl
5461

5562
# Container options for MetaX C550 (MACA runtime)
5663
container_options: >-

‎docker/metax/Dockerfile.inference‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
# FlagScale inference and serve image for MetaX C550.
2+
# The immutable base owns Torch, MACA, vLLM-MetaX and FlagGems. This image
3+
# overlays only the workflow-resolved vllm-plugin-FL source revision.
4+
5+
ARG BASE_IMAGE=harbor.baai.ac.cn/flagos-dev/vllm-plugin-fl@sha256:c83a7b26e5b72fcffa920f32929a76984c78d5ee5520828b6e50feb9864240ed
6+
7+
FROM ${BASE_IMAGE} AS dev
8+
9+
ARG FLAGSCALE_HOME=/opt/flagscale
10+
ARG VLLM_PLUGIN_FL_REF=main
11+
12+
LABEL io.flagscale.vllm-plugin-fl.revision=${VLLM_PLUGIN_FL_REF}
13+
14+
ENV FLAGSCALE_HOME=${FLAGSCALE_HOME} \
15+
FLAGSCALE_DEPS=${FLAGSCALE_HOME}/deps \
16+
FLAGSCALE_VLLM_PLUGIN_REF=${VLLM_PLUGIN_FL_REF} \
17+
VLLM_PLUGINS=fl \
18+
VLLM_FL_PLATFORM=metax
19+
20+
COPY tools/install /workspace/tools/install
21+
COPY requirements /workspace/requirements
22+
23+
# Preserve the immutable vendor runtime contract: the source overlay may add
24+
# its entry point but must not introduce new dependency conflicts.
25+
RUN python -m pip check 2>&1 | sort > /tmp/pip-check.before || true
26+
27+
RUN chmod +x /workspace/tools/install/*.sh \
28+
/workspace/tools/install/utils/*.sh \
29+
/workspace/tools/install/metax/*.sh && \
30+
/workspace/tools/install/install.sh --platform metax --task inference \
31+
--pkg-mgr pip --no-system --no-dev --no-base && \
32+
/workspace/tools/install/install.sh --platform metax --task serve \
33+
--pkg-mgr pip --no-system --no-dev --no-base
34+
35+
RUN python -m pip check 2>&1 | sort > /tmp/pip-check.after || true
36+
37+
RUN diff -u /tmp/pip-check.before /tmp/pip-check.after && \
38+
rm -f /tmp/pip-check.before /tmp/pip-check.after && \
39+
rm -rf /opt/flagscale/deps /workspace/tools /workspace/requirements
40+
41+
WORKDIR /workspace
42+
CMD ["/bin/bash"]
43+
44+
FROM dev AS release

‎requirements/metax/inference.txt‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
# Copyright 2026 FlagOS Contributors
2+
# Licensed under the Apache License, Version 2.0.
3+
4+
# The pinned MetaX vLLM base owns Torch, MACA, vLLM-MetaX and FlagGems.
5+
# This is the sole runtime dependency declared by vllm-plugin-FL; the plugin
6+
# source itself is installed from the workflow-resolved commit.
7+
pyyaml>=6.0,<7

‎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+
# Client-side dependencies used by the serve health and real-request checks.
7+
httpx>=0.27,<1
8+
requests>=2.32,<3

‎tools/install/metax/install_inference.sh‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
#!/bin/bash
2+
3+
# Copyright 2026 FlagOS Contributors
4+
# Licensed under the Apache License, Version 2.0.
5+
6+
set -euo pipefail
7+
8+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
9+
source "$SCRIPT_DIR/../utils/utils.sh"
10+
source "$SCRIPT_DIR/../utils/pkg_utils.sh"
11+
source "$SCRIPT_DIR/../utils/retry_utils.sh"
12+
13+
PROJECT_ROOT=$(get_project_root)
14+
DEBUG="${FLAGSCALE_DEBUG:-false}"
15+
RETRY_COUNT="${FLAGSCALE_RETRY_COUNT:-3}"
16+
FLAGSCALE_HOME="${FLAGSCALE_HOME:-/opt/flagscale}"
17+
FLAGSCALE_DEPS="${FLAGSCALE_DEPS:-$FLAGSCALE_HOME/deps}"
18+
REQ_FILE="$PROJECT_ROOT/requirements/metax/inference.txt"
19+
PLUGIN_REPO="${FLAGSCALE_VLLM_PLUGIN_REPO:-https://github.com/flagos-ai/vllm-plugin-FL.git}"
20+
PLUGIN_REF="${FLAGSCALE_VLLM_PLUGIN_REF:-}"
21+
22+
while [[ $# -gt 0 ]]; do
23+
case $1 in --debug) DEBUG=true; shift ;; *) shift ;; esac
24+
done
25+
26+
checkout_pinned_ref() {
27+
local repo=$1
28+
local ref=$2
29+
local target=$3
30+
31+
[ -z "$ref" ] && { log_error "A pinned vllm-plugin-FL ref is required"; return 1; }
32+
retry -d "$DEBUG" "$RETRY_COUNT" "rm -rf '$target' && \
33+
git init -q '$target' && \
34+
git -C '$target' remote add origin '$repo' && \
35+
git -c http.version=HTTP/1.1 -C '$target' fetch --depth 1 origin '$ref' && \
36+
git -C '$target' checkout -q --detach FETCH_HEAD"
37+
}
38+
39+
install_requirements() {
40+
set_step "Installing MetaX inference requirements"
41+
retry_pip_install -d "$DEBUG" "$REQ_FILE" "$RETRY_COUNT" || return 1
42+
}
43+
44+
install_plugin() {
45+
set_step "Installing resolved vllm-plugin-FL for MetaX"
46+
mkdir -p "$FLAGSCALE_DEPS"
47+
checkout_pinned_ref "$PLUGIN_REPO" "$PLUGIN_REF" \
48+
"$FLAGSCALE_DEPS/vllm-plugin-FL" || return 1
49+
50+
local pip_cmd
51+
pip_cmd=$(get_pip_cmd)
52+
run_cmd -d "$DEBUG" bash -c \
53+
"cd '$FLAGSCALE_DEPS/vllm-plugin-FL' && $pip_cmd install \
54+
--root-user-action=ignore --no-deps --no-build-isolation ." || return 1
55+
}
56+
57+
validate_runtime() {
58+
[ "$DEBUG" = true ] && return 0
59+
VLLM_PLUGINS=fl VLLM_FL_PLATFORM=metax python - <<'PY'
60+
import importlib.metadata as metadata
61+
62+
assert metadata.version("vllm").startswith("0.20.2")
63+
entrypoints = {
64+
entry.name: entry.value
65+
for entry in metadata.entry_points(group="vllm.platform_plugins")
66+
}
67+
assert entrypoints.get("fl") == "vllm_fl:register", entrypoints
68+
print("vllm:", metadata.version("vllm"))
69+
print("vllm-plugin-fl:", metadata.version("vllm-plugin-fl"))
70+
print("platform entrypoints:", entrypoints)
71+
PY
72+
}
73+
74+
main() {
75+
install_requirements || die "MetaX inference requirements failed"
76+
install_plugin || die "vllm-plugin-FL installation failed"
77+
validate_runtime || die "MetaX inference runtime validation failed"
78+
log_success "MetaX inference runtime ready"
79+
}
80+
81+
main

‎tools/install/metax/install_serve.sh‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
#!/bin/bash
2+
3+
# Copyright 2026 FlagOS Contributors
4+
# Licensed under the Apache License, Version 2.0.
5+
6+
set -euo pipefail
7+
8+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
9+
source "$SCRIPT_DIR/../utils/utils.sh"
10+
source "$SCRIPT_DIR/../utils/pkg_utils.sh"
11+
source "$SCRIPT_DIR/../utils/retry_utils.sh"
12+
13+
PROJECT_ROOT=$(get_project_root)
14+
DEBUG="${FLAGSCALE_DEBUG:-false}"
15+
RETRY_COUNT="${FLAGSCALE_RETRY_COUNT:-3}"
16+
REQ_FILE="$PROJECT_ROOT/requirements/metax/serve.txt"
17+
18+
while [[ $# -gt 0 ]]; do
19+
case $1 in --debug) DEBUG=true; shift ;; *) shift ;; esac
20+
done
21+
22+
main() {
23+
set_step "Installing MetaX serve requirements"
24+
retry_pip_install -d "$DEBUG" "$REQ_FILE" "$RETRY_COUNT" || \
25+
die "MetaX serve requirements failed"
26+
27+
[ "$DEBUG" = true ] || python - <<'PY'
28+
import httpx
29+
import requests
30+
31+
print("httpx:", httpx.__version__)
32+
print("requests:", requests.__version__)
33+
PY
34+
log_success "MetaX serve dependencies ready"
35+
}
36+
37+
main

‎tools/install/metax/validate_image_build.sh‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,10 @@ task="${IMAGE_BUILD_TASK:?IMAGE_BUILD_TASK is required}"
1010
base_image="${IMAGE_BUILD_BASE_IMAGE:?IMAGE_BUILD_BASE_IMAGE is required}"
1111
candidate="${IMAGE_BUILD_CANDIDATE_IMAGE:?IMAGE_BUILD_CANDIDATE_IMAGE is required}"
1212

13-
[ "$task" = train ] || exit 0
13+
case "$task" in
14+
train|inference) ;;
15+
*) exit 0 ;;
16+
esac
1417

1518
if [ "$phase" = pre ]; then
1619
docker pull "$base_image"
@@ -47,6 +50,36 @@ torchrun --standalone --nnodes=1 --nproc-per-node=2 /tmp/metax_collective.py
4750
fi
4851

4952
[ "$phase" = post ] || exit 0
53+
54+
if [ "$task" = inference ]; then
55+
docker run --rm \
56+
--ipc=host --group-add video \
57+
--device=/dev/dri --device=/dev/mxcd --device=/dev/infiniband \
58+
--entrypoint python "$candidate" -c '
59+
import importlib.metadata as metadata
60+
import torch
61+
import vllm_fl
62+
from vllm.platforms import current_platform
63+
64+
print("torch:", torch.__version__)
65+
print("vllm:", metadata.version("vllm"))
66+
print("vllm-plugin-fl:", metadata.version("vllm-plugin-fl"))
67+
print("platform:", type(current_platform).__module__, type(current_platform).__name__)
68+
print("vendor:", current_platform.vendor_name)
69+
print("device_type:", current_platform.device_type)
70+
71+
assert torch.cuda.device_count() == 8
72+
assert type(current_platform).__module__ == "vllm_fl.platform"
73+
assert type(current_platform).__name__ == "PlatformFL"
74+
assert current_platform.vendor_name == "metax"
75+
assert current_platform.device_type == "cuda"
76+
77+
value = torch.ones(16, device="cuda:0")
78+
assert value.sum().item() == 16
79+
'
80+
exit 0
81+
fi
82+
5083
docker run --rm \
5184
--ipc=host --group-add video \
5285
--device=/dev/dri --device=/dev/mxcd \

0 commit comments

Comments
 (0)