Backport/release 1.39 to master (#1767)

docker/build_image.py

@@ -3,14 +3,27 @@ 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
+
+import json
 
 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:
@@ -359,7 +372,8 @@ class StableGPUImageBuilder(Builder):
             extra_content = extra_content.replace('{python_version}',
                                                   self.args.python_version)
             extra_content += """
-RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
+RUN export PIP_EXTRA_INDEX_URL=https://pypi.org/simple && \
+    pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
 """
         version_args = (
             f'{self.args.torch_version} {self.args.torchvision_version} {self.args.torchaudio_version} '
@@ -420,7 +434,8 @@ class LatestGPUImageBuilder(StableGPUImageBuilder):
             extra_content = extra_content.replace('{python_version}',
                                                   self.args.python_version)
         extra_content += """
-RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
+RUN export PIP_EXTRA_INDEX_URL=https://pypi.org/simple && \
+    pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
 """
         version_args = (
             f'{self.args.torch_version} {self.args.torchvision_version} {self.args.torchaudio_version} '
@@ -467,6 +482,370 @@ 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, json.JSONDecodeError) 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)."""
+        # Keep this script compact: it runs inside the base image via python -c.
+        script = (
+            'import json,os,pathlib,subprocess,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'
+            'def _dpkg_ver(*names):\n'
+            ' for n in names:\n'
+            '  try:\n'
+            '   r=subprocess.run(["dpkg-query","-W","-f=${Version}",n],'
+            'capture_output=True,text=True)\n'
+            '   if r.returncode==0 and r.stdout.strip():\n'
+            '    return r.stdout.strip()\n'
+            '  except Exception:\n'
+            '   pass\n'
+            ' return None\n'
+            'def _dpkg_scan(prefixes):\n'
+            ' try:\n'
+            '  r=subprocess.run(["dpkg-query","-W","-f=${Package}\\t${Version}\\n"],'
+            'capture_output=True,text=True)\n'
+            ' except Exception:\n'
+            '  return {}\n'
+            ' found={}\n'
+            ' for line in (r.stdout or "").splitlines():\n'
+            '  if "\\t" not in line: continue\n'
+            '  pkg,ver=line.split("\\t",1)\n'
+            '  for pref in prefixes:\n'
+            '   if pkg==pref or pkg.startswith(pref+"-"):\n'
+            '    found.setdefault(pref,ver)\n'
+            ' return found\n'
+            'pkgs=_dpkg_scan(("rccl","miopen"))\n'
+            'info["system.library.rccl"]=_dpkg_ver("rccl") or pkgs.get("rccl")\n'
+            'info["system.library.miopen"]=('
+            '_dpkg_ver("miopen-hip","miopen") or pkgs.get("miopen"))\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')
+        # Keep dpkg package versions as-is (may contain '~', e.g. 2.27.7.70201-81~22.04).
+        rccl_ver = probed.get('system.library.rccl')
+        miopen_ver = probed.get('system.library.miopen')
+
+        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,
+            'system.library.rccl': rccl_ver or None,
+            'system.library.miopen': miopen_ver or 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):
 
     _DEFAULT_TORCH_VERSION = '2.9.0'
@@ -638,7 +1017,8 @@ class AscendImageBuilder(StableGPUImageBuilder):
 
     def generate_dockerfile(self) -> str:
         extra_content = """
-RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
+RUN export PIP_EXTRA_INDEX_URL=https://pypi.org/simple && \
+    pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
 """
         with open('docker/Dockerfile.ascend', 'r') as f:
             content = f.read()
@@ -708,6 +1088,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()
 
@@ -719,6 +1105,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: