From 5ee9b27220e55fa7c32acfb69f2867b4c9d1f196 Mon Sep 17 00:00:00 2001 From: suluyan Date: Fri, 17 Jul 2026 11:19:02 +0800 Subject: [PATCH] feat(docker): add AMD ROCm image build from vllm-openai-rocm Add image_type=amd that auto-resolves the newest concrete Docker Hub tag, probes rocm/python/torch without GPU, installs modelscope from the build branch, and tags as ubuntu*-rocm*-py*-torch*-{version}. Co-authored-by: Cursor --- .github/workflows/docker-image.yml | 2 +- docker/Dockerfile.amd | 21 ++ docker/build_image.py | 355 ++++++++++++++++++++++++++++- 3 files changed, 376 insertions(+), 2 deletions(-) create mode 100644 docker/Dockerfile.amd diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml index f5c576da..364d1834 100644 --- a/.github/workflows/docker-image.yml +++ b/.github/workflows/docker-image.yml @@ -11,7 +11,7 @@ on: description: 'ModelScope branch to build from(release/x.xx)' required: true image_type: - description: 'The image type to build(base/old/stable/latest)' + description: 'The image type to build(base/old/stable/latest/amd/ascend)' required: true modelscope_version: description: 'ModelScope version to use(x.xx.x)' diff --git a/docker/Dockerfile.amd b/docker/Dockerfile.amd new file mode 100644 index 00000000..f98517ee --- /dev/null +++ b/docker/Dockerfile.amd @@ -0,0 +1,21 @@ +FROM {base_image} + +ARG BASE_IMAGE_TAG={base_image_tag} +LABEL modelscope.base_image="vllm/vllm-openai-rocm:${BASE_IMAGE_TAG}" + +COPY docker/scripts/modelscope_env_init.sh /usr/local/bin/ms_env_init.sh + +ARG CUR_TIME={cur_time} +RUN echo "CUR_TIME=${CUR_TIME}" && echo "BASE_IMAGE_TAG=${BASE_IMAGE_TAG}" + +RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \ + pip config set install.trusted-host mirrors.aliyun.com && \ + cd /tmp && GIT_LFS_SKIP_SMUDGE=1 git clone -b {modelscope_branch} --single-branch https://github.com/modelscope/modelscope.git && \ + cd modelscope && pip install --no-cache-dir . -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \ + cd / && rm -fr /tmp/modelscope && pip cache purge + +ENV VLLM_USE_MODELSCOPE=True +ENV LMDEPLOY_USE_MODELSCOPE=True +ENV MODELSCOPE_CACHE=/mnt/workspace/.cache/modelscope/hub + +SHELL ["/bin/bash", "-c"] diff --git a/docker/build_image.py b/docker/build_image.py index 80e8ed06..8fab7f96 100644 --- a/docker/build_image.py +++ b/docker/build_image.py @@ -1,16 +1,28 @@ import argparse +import json import os import platform import re import subprocess +import urllib.error +import urllib.request from copy import copy from datetime import datetime -from typing import Any +from typing import Any, List, Optional docker_registry = os.environ['DOCKER_REGISTRY'] assert docker_registry, 'You must pass a valid DOCKER_REGISTRY' timestamp = datetime.now() formatted_time = timestamp.strftime('%Y%m%d%H%M%S') +VLLM_ROCM_REPO = 'vllm/vllm-openai-rocm' +_FLOATING_ROCM_TAGS = frozenset({ + 'latest', + 'latest-base', + 'nightly', + 'base-nightly', +}) +_VERSION_TAG_PATTERN = re.compile(r'^v\d+(?:\.\d+)*$') +_NIGHTLY_HASH_PATTERN = re.compile(r'^(?:base-)?nightly-[0-9a-f]{7,40}$') class Builder: @@ -467,6 +479,339 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy return self.run_cmd('docker', 'push', image_tag2) +class AmdImageBuilder(Builder): + """Build ModelScope image on top of vllm/vllm-openai-rocm.""" + + @staticmethod + def _is_specific_release_tag(tag: str) -> bool: + tag = tag.strip() + if not tag or tag.lower() in _FLOATING_ROCM_TAGS: + return False + if tag.endswith('-base'): + return False + if _NIGHTLY_HASH_PATTERN.fullmatch(tag): + return False + return bool(_VERSION_TAG_PATTERN.fullmatch(tag)) + + @staticmethod + def _image_digest(tag_info: dict) -> Optional[str]: + digest = tag_info.get('digest') + if digest: + return digest + for image in tag_info.get('images') or []: + digest = image.get('digest') + if digest: + return digest + return None + + @classmethod + def _fetch_rocm_tags(cls, page_size: int = 100) -> List[dict]: + tags: List[dict] = [] + url = ( + f'https://hub.docker.com/v2/repositories/{VLLM_ROCM_REPO}/tags' + f'?page_size={page_size}&ordering=-last_updated') + while url: + req = urllib.request.Request( + url, headers={'User-Agent': 'modelscope-docker-builder'}) + try: + with urllib.request.urlopen(req, timeout=60) as resp: + payload = json.load(resp) + except urllib.error.URLError as exc: + raise RuntimeError( + f'Failed to query Docker Hub tags for {VLLM_ROCM_REPO}: ' + f'{exc}') from exc + tags.extend(payload.get('results') or []) + url = payload.get('next') + # Only scan the first few pages; release tags are near the top. + if len(tags) >= 300: + break + if not tags: + raise RuntimeError( + f'No tags returned from Docker Hub for {VLLM_ROCM_REPO}') + return tags + + @classmethod + def resolve_latest_rocm_tag(cls) -> str: + """Resolve the newest concrete release tag for vllm-openai-rocm. + + Preference order: + 1. Semver tag (vX.Y.Z) that shares digest with floating ``latest`` + 2. Newest semver tag by Docker Hub ``last_updated`` + """ + tags = cls._fetch_rocm_tags() + by_name = {item['name']: item for item in tags if item.get('name')} + release_tags = [ + item for item in tags + if cls._is_specific_release_tag(item.get('name', '')) + ] + latest_info = by_name.get('latest') + latest_digest = cls._image_digest(latest_info) if latest_info else None + if latest_digest: + matched = [ + item for item in release_tags + if cls._image_digest(item) == latest_digest + ] + if matched: + # Prefer the first match in last_updated order from API. + chosen = matched[0]['name'] + print( + f'Resolved {VLLM_ROCM_REPO} latest digest to release tag: ' + f'{chosen}') + return chosen + + if not release_tags: + raise RuntimeError( + f'No concrete release tags found for {VLLM_ROCM_REPO}') + chosen = release_tags[0]['name'] + print(f'Resolved newest {VLLM_ROCM_REPO} release tag: {chosen}') + return chosen + + def init_args(self, args: Any) -> Any: + # Auto-discover from Docker Hub unless an explicit override is given. + override = getattr(args, 'base_image_tag', None) + if override and str(override).strip() and str(override).strip().lower( + ) not in {'auto', 'latest'}: + args.base_image_tag = str(override).strip() + if not self._is_specific_release_tag(args.base_image_tag): + raise ValueError( + 'base_image_tag override must be a concrete release tag ' + f'(e.g. v0.25.1), got: {args.base_image_tag}') + print(f'Using override AMD ROCm base image tag: ' + f'{args.base_image_tag}') + else: + args.base_image_tag = self.resolve_latest_rocm_tag() + if not args.base_image: + args.base_image = f'{VLLM_ROCM_REPO}:{args.base_image_tag}' + if not args.cuda_version: + args.cuda_version = '0.0.0' + return args + + @staticmethod + def _sanitize_tag(tag: str) -> str: + return re.sub(r'[^A-Za-z0-9._-]+', '-', tag) + + @staticmethod + def _normalize_version(version: str) -> str: + version = version.strip().lstrip('vV') + version = version.split('+')[0].split(' ')[0] + return re.sub(r'[^0-9A-Za-z._-]+', '', version) + + @staticmethod + def _python_tag_from_version(version: str) -> str: + parts = version.strip().split('.') + if len(parts) >= 2 and parts[0].isdigit() and parts[1].isdigit(): + return f'py{parts[0]}{parts[1]}' + return f'py{re.sub(r"[^0-9]", "", version)}' + + @classmethod + def _run_capture(cls, *cmd: str) -> subprocess.CompletedProcess: + return subprocess.run( + list(cmd), capture_output=True, text=True, check=False) + + @classmethod + def _probe_via_entrypoint(cls, base_image: str) -> dict: + """Read versions with docker run --entrypoint (no GPU required).""" + script = ( + 'import json,os,pathlib,sys\n' + 'info={"python":"%d.%d.%d"%sys.version_info[:3]}\n' + 'try:\n' + ' import torch\n' + ' info["torch"]=torch.__version__\n' + ' hip=getattr(torch.version,"hip",None)\n' + ' if hip: info["torch_hip"]=hip\n' + 'except Exception as e:\n' + ' info["torch_error"]=str(e)\n' + 'for p in ("/opt/rocm/.info/version","/opt/rocm/.info/version-dev"):\n' + ' f=pathlib.Path(p)\n' + ' if f.is_file():\n' + ' info["rocm_file"]=f.read_text().strip().splitlines()[0]\n' + ' break\n' + 'for k in ("ROCM_VERSION","HIP_VERSION","TORCH_VERSION"):\n' + ' if os.environ.get(k): info[k.lower()]=os.environ[k]\n' + 'print(json.dumps(info))\n') + for py in ('python3', 'python'): + result = cls._run_capture( + 'docker', 'run', '--rm', '--network', 'none', '--entrypoint', + py, base_image, '-c', script) + if result.returncode == 0 and result.stdout.strip(): + try: + return json.loads(result.stdout.strip().splitlines()[-1]) + except json.JSONDecodeError: + continue + return {} + + @classmethod + def _probe_via_history(cls, base_image: str) -> dict: + """Parse build ARGs from docker history (no container start).""" + result = cls._run_capture('docker', 'history', '--no-trunc', + '--format', '{{.CreatedBy}}', base_image) + if result.returncode != 0: + return {} + text = result.stdout + info = {} + for key, pattern in ( + ('rocm', r'ROCM_VERSION=([0-9]+(?:\.[0-9]+)*)'), + ('python', r'PYTHON_VERSION=([0-9]+(?:\.[0-9]+)*)'), + ('ubuntu', + r'org\.opencontainers\.image\.version=([0-9]+(?:\.[0-9]+)*)'), + ): + matches = re.findall(pattern, text) + if matches: + # docker history lists newest layers first. + info[key] = matches[0] + return info + + @classmethod + def _probe_via_create_cp(cls, base_image: str) -> dict: + """Copy version files out of a created (not started) container.""" + import tempfile + create = cls._run_capture('docker', 'create', base_image) + if create.returncode != 0: + return {} + cid = create.stdout.strip() + info = {} + try: + with tempfile.TemporaryDirectory() as tmp: + dest = os.path.join(tmp, 'version') + for src in ('/opt/rocm/.info/version', + '/opt/rocm/.info/version-dev'): + result = cls._run_capture('docker', 'cp', f'{cid}:{src}', + dest) + if result.returncode == 0 and os.path.isfile(dest): + with open(dest, 'r', encoding='utf-8') as f: + line = f.read().strip().splitlines() + if line: + info['rocm_file'] = line[0].strip() + break + finally: + cls._run_capture('docker', 'rm', '-f', cid) + return info + + @classmethod + def probe_base_image_versions(cls, base_image: str) -> dict: + """Discover rocm/python/torch without needing AMD GPU. + + Methods (in order): + 1. docker run --entrypoint python -c ... (CPU-only, no --device) + 2. docker history --no-trunc parse ROCM_VERSION/PYTHON_VERSION + 3. docker create + docker cp /opt/rocm/.info/version + """ + probed = {} + entry = cls._probe_via_entrypoint(base_image) + history = cls._probe_via_history(base_image) + copied = cls._probe_via_create_cp(base_image) + probed.update(history) + probed.update(copied) + probed.update(entry) + + rocm = ( + probed.get('rocm_file') or probed.get('rocm_version') + or probed.get('rocm') or probed.get('torch_hip') + or probed.get('hip_version')) + python_ver = probed.get('python') + torch_ver = probed.get('torch') or probed.get('torch_version') + ubuntu_ver = probed.get('ubuntu') + + versions = { + 'rocm': cls._normalize_version(rocm) if rocm else None, + 'python': cls._normalize_version(python_ver) if python_ver else None, + 'torch': cls._normalize_version(torch_ver) if torch_ver else None, + 'ubuntu': cls._normalize_version(ubuntu_ver) if ubuntu_ver else None, + } + print('Probed AMD base image versions:') + for key, value in versions.items(): + print(f' {key}: {value or "unknown"}') + return versions + + def generate_dockerfile(self) -> str: + with open('docker/Dockerfile.amd', 'r') as f: + content = f.read() + content = content.replace('{base_image}', self.args.base_image) + content = content.replace('{base_image_tag}', self.args.base_image_tag) + content = content.replace('{modelscope_branch}', + self.args.modelscope_branch) + content = content.replace('{cur_time}', formatted_time) + return content + + def image(self) -> str: + ubuntu = getattr(self.args, 'amd_ubuntu_version', + None) or self.args.ubuntu_version + rocm = getattr(self.args, 'amd_rocm_version', None) + py_tag = getattr(self.args, 'amd_python_tag', None) or getattr( + self.args, 'python_tag', None) + torch = getattr(self.args, 'amd_torch_version', None) + if not (rocm and py_tag and torch): + raise RuntimeError( + 'AMD image tag requires probed rocm/python/torch versions. ' + f'Got rocm={rocm}, python={py_tag}, torch={torch}') + return ( + f'{docker_registry}:ubuntu{ubuntu}-rocm{rocm}-{py_tag}-' + f'torch{torch}-{self.args.modelscope_version}-test') + + def _log_base_image_info(self) -> int: + base_image = self.args.base_image + print('=' * 60) + print(f'AMD ROCm base image: {base_image}') + print(f'AMD ROCm base image tag: {self.args.base_image_tag}') + print('=' * 60) + ret = self.run_cmd('docker', 'pull', base_image) + if ret != 0: + return ret + result = self._run_capture( + 'docker', 'image', 'inspect', base_image, + '--format={{.Id}} {{if index .RepoDigests 0}}' + '{{index .RepoDigests 0}}{{else}}local-only{{end}}') + if result.returncode == 0: + print(f'AMD base image resolved: {result.stdout.strip()}') + else: + print(f'AMD base image inspect warning: {result.stderr.strip()}') + + versions = self.probe_base_image_versions(base_image) + if not versions.get('rocm') or not versions.get('python') or not versions.get( + 'torch'): + print('ERROR: failed to probe rocm/python/torch from base image') + return 1 + self.args.amd_rocm_version = versions['rocm'] + self.args.amd_torch_version = versions['torch'] + self.args.amd_python_tag = self._python_tag_from_version( + versions['python']) + if versions.get('ubuntu'): + self.args.amd_ubuntu_version = versions['ubuntu'] + else: + self.args.amd_ubuntu_version = self.args.ubuntu_version + print( + f'AMD output image tag will be: {self.image()}' + ) + print('=' * 60) + return 0 + + def build(self) -> int: + ret = self._log_base_image_info() + if ret != 0: + return ret + return self.run_cmd('docker', 'build', '-t', self.image(), '-f', + 'Dockerfile', '.') + + def push(self): + image_name = self.image() + ret = self.run_cmd('docker', 'push', image_name) + if ret != 0: + return ret + ubuntu = self.args.amd_ubuntu_version + rocm = self.args.amd_rocm_version + py_tag = self.args.amd_python_tag + torch = self.args.amd_torch_version + image_tag2 = ( + f'{docker_registry}:ubuntu{ubuntu}-rocm{rocm}-{py_tag}-' + f'torch{torch}-{self.args.modelscope_version}-' + f'{formatted_time}-test') + ret = self.run_cmd('docker', 'tag', image_name, image_tag2) + if ret != 0: + return ret + print(f'AMD image timestamp tag: {image_tag2}') + return self.run_cmd('docker', 'push', image_tag2) + + class AscendImageBuilder(StableGPUImageBuilder): _CANN_VERSION_PATTERN = re.compile(r'^\d+(?:\.[0-9A-Za-z]+)+$') @@ -596,6 +941,12 @@ parser.add_argument('--megatron_branch', type=str, default='v0.15.3') parser.add_argument('--mindspeed_branch', type=str, default='core_r0.15.3') parser.add_argument('--soc_version', type=str, default='ascend910_9391') parser.add_argument('--arch', type=str, choices=['x86', 'arm'], default=None) +parser.add_argument( + '--base_image_tag', + type=str, + default=None, + help='Optional AMD ROCm override tag. Default: auto-resolve newest ' + 'concrete vllm/vllm-openai-rocm release tag from Docker Hub.') parser.add_argument('--dry_run', type=int, default=0) args = parser.parse_args() @@ -607,6 +958,8 @@ elif args.image_type.lower() == 'stable': builder_cls = [StableCPUImageBuilder, StableGPUImageBuilder] elif args.image_type.lower() == 'ascend': builder_cls = [AscendImageBuilder] +elif args.image_type.lower() == 'amd': + builder_cls = [AmdImageBuilder] elif args.image_type.lower() == 'latest': builder_cls = [LatestGPUImageBuilder] else: