fix(BA-2754): configure mock accelerator (#6324)

Author: devyubin

src/ai/backend/install/cli.py

@@ -40,7 +40,15 @@
 from . import __version__
 from .common import detect_os
 from .context import DevContext, PackageContext, current_log
-from .types import CliArgs, DistInfo, InstallInfo, InstallModes, InstallVariable, PrerequisiteError
+from .types import (
+    Accelerator,
+    CliArgs,
+    DistInfo,
+    InstallInfo,
+    InstallModes,
+    InstallVariable,
+    PrerequisiteError,
+)
 
 top_tasks: WeakSet[asyncio.Task] = WeakSet()
 
@@ -402,7 +410,10 @@ def __init__(
         self._enabled_menus.add(InstallModes.CONFIGURE)
         assert mode is not None
         self._mode = mode
-        self.install_variable = InstallVariable(public_facing_address=args.public_facing_address)
+        self.install_variable = InstallVariable(
+            public_facing_address=args.public_facing_address,
+            accelerator=Accelerator(args.accelerator) if args.accelerator is not None else None,
+        )
 
     def compose(self) -> ComposeResult:
         yield Label(id="heading")
@@ -526,6 +537,7 @@ def __init__(self, args: CliArgs | None = None) -> None:
                 show_guide=False,
                 non_interactive=False,
                 public_facing_address="127.0.0.1",
+                accelerator=None,
             )
         self._args = args
 
@@ -609,6 +621,13 @@ async def action_shutdown(self, message: str | None = None, exit_code: int = 0)
     default=False,
     help="Show the post-install guide using INSTALL-INFO if present.",
 )
+@click.option(
+    "--accelerator",
+    type=click.Choice([a.value for a in Accelerator], case_sensitive=False),
+    default=None,
+    show_default=True,
+    help="Select accelerator plugin (cuda, cuda_mock, cuda_mig_mock, rocm_mock, none)",
+)
 @click.option(
     "--headless",
     is_flag=True,
@@ -631,6 +650,7 @@ def main(
     non_interactive: bool,
     headless: bool,
     public_facing_address: str,
+    accelerator: str,
 ) -> None:
     """The installer"""
     # check sudo permission
@@ -648,6 +668,7 @@ def main(
         show_guide=show_guide,
         non_interactive=non_interactive,
         public_facing_address=public_facing_address,
+        accelerator=accelerator,
     )
     app = InstallerApp(args)
     app.run(headless=headless)

src/ai/backend/install/context.py

@@ -48,6 +48,7 @@
 from .http import wget
 from .python import check_python
 from .types import (
+    Accelerator,
     DistInfo,
     HalfstackConfig,
     HostPortPair,
@@ -105,6 +106,10 @@ def __init__(
     def hydrate_install_info(self) -> InstallInfo:
         raise NotImplementedError
 
+    @abstractmethod
+    async def _configure_mock_accelerator(self, accelerator: Accelerator) -> None:
+        raise NotImplementedError
+
     def add_post_guide(self, guide: PostGuide) -> None:
         self._post_guides.append(guide)
 
@@ -461,6 +466,7 @@ async def configure_manager(self) -> None:
     async def configure_agent(self) -> None:
         halfstack = self.install_info.halfstack_config
         service = self.install_info.service_config
+        accelerator = self.install_info.accelerator
         toml_path = self.copy_config("agent.toml")
         self.sed_in_place_multi(
             toml_path,
@@ -485,25 +491,27 @@ async def configure_agent(self) -> None:
         Path(self.install_info.service_config.agent_var_base_path).mkdir(
             parents=True, exist_ok=True
         )
-        # enable the CUDA plugin (open-source version)
-        # The agent will show an error log if the CUDA is not available in the system and report
-        # "cuda.devices = 0" as the agent capacity, but it will still run.
+        if accelerator is not None:
+            if accelerator == Accelerator.CUDA:
+                plugin_list = ['"ai.backend.accelerator.cuda_open"']
+            elif accelerator in (
+                Accelerator.CUDA_MOCK,
+                Accelerator.CUDA_MIG_MOCK,
+                Accelerator.ROCM_MOCK,
+            ):
+                plugin_list = ['"ai.backend.accelerator.mock"']
+            else:
+                plugin_list = []
+
+            await self._configure_mock_accelerator(accelerator)
+        else:
+            plugin_list = []
+
         self.sed_in_place(
             toml_path,
-            re.compile("^(# )?allow-compute-plugins = .*", flags=re.M),
-            'allow-compute-plugins = ["ai.backend.accelerator.cuda_open"]',
+            re.compile(r"^(# )?allow-compute-plugins = .*", flags=re.M),
+            f"allow-compute-plugins = [{', '.join(plugin_list)}]",
         )
-        # TODO: let the installer enable the CUDA plugin only when it verifies CUDA availability or
-        #       via an explicit installer option/config.
-        r"""
-        if [ $ENABLE_CUDA -eq 1 ]; then
-          sed_inplace "s/# allow-compute-plugins =.*/allow-compute-plugins = [\"ai.backend.accelerator.cuda_open\"]/" ./agent.toml
-        elif [ $ENABLE_CUDA_MOCK -eq 1 ]; then
-          sed_inplace "s/# allow-compute-plugins =.*/allow-compute-plugins = [\"ai.backend.accelerator.mock\"]/" ./agent.toml
-        else
-          sed_inplace "s/# allow-compute-plugins =.*/allow-compute-plugins = []/" ./agent.toml
-        fi
-        """
 
     async def configure_storage_proxy(self) -> None:
         halfstack = self.install_info.halfstack_config
@@ -870,6 +878,7 @@ def hydrate_install_info(self) -> InstallInfo:
             last_updated=datetime.now(tzutc()),
             halfstack_config=halfstack_config,
             service_config=service_config,
+            accelerator=self.install_variable.accelerator,
         )
 
     def copy_config(self, template_name: str) -> Path:
@@ -895,10 +904,23 @@ async def install(self) -> None:
         await install_editable_webui(self)
         await self.install_halfstack()
 
-    async def _configure_mock_accelerator(self) -> None:
+    async def _configure_mock_accelerator(self, accelerator: Accelerator) -> None:
         """
         cp "configs/accelerator/mock-accelerator.toml" mock-accelerator.toml
         """
+        mapping = {
+            Accelerator.CUDA_MOCK: "configs/accelerator/mock-accelerator.toml",
+            Accelerator.CUDA_MIG_MOCK: "configs/accelerator/cuda-mock-mig.toml",
+            Accelerator.ROCM_MOCK: "configs/accelerator/rocm-mock.toml",
+        }
+
+        src = mapping.get(accelerator)
+        if not src:
+            return
+
+        dst = Path("mock-accelerator.toml")
+        print(f"[Installer] Copying accelerator config: {src} -> {dst}")
+        shutil.copy(src, dst)
 
     async def configure(self) -> None:
         self.log_header("Configuring manager...")
@@ -986,6 +1008,7 @@ def hydrate_install_info(self) -> InstallInfo:
             last_updated=datetime.now(tzutc()),
             halfstack_config=halfstack_config,
             service_config=service_config,
+            accelerator=self.install_variable.accelerator,
         )
 
     def copy_config(self, template_name: str) -> Path:
@@ -1148,6 +1171,24 @@ async def install(self) -> None:
         self.log_header("Installing databases (halfstack)...")
         await self.install_halfstack()
 
+    async def _configure_mock_accelerator(self, accelerator: Accelerator) -> None:
+        """
+        cp "configs/accelerator/mock-accelerator.toml" mock-accelerator.toml
+        """
+        mapping = {
+            Accelerator.CUDA_MOCK: "configs/accelerator/mock-accelerator.toml",
+            Accelerator.CUDA_MIG_MOCK: "configs/accelerator/cuda-mock-mig.toml",
+            Accelerator.ROCM_MOCK: "configs/accelerator/rocm-mock.toml",
+        }
+
+        src = mapping.get(accelerator)
+        if not src:
+            return
+
+        dst = Path("mock-accelerator.toml")
+        print(f"[Installer] Copying accelerator config: {src} -> {dst}")
+        shutil.copy(src, dst)
+
     async def configure(self) -> None:
         self.log_header("Configuring manager...")
         await self.configure_manager()

src/ai/backend/install/types.py

@@ -4,7 +4,7 @@
 import enum
 from datetime import datetime
 from pathlib import Path
-from typing import cast
+from typing import Optional, cast
 
 from pydantic import BaseModel, Field
 from rich.console import ConsoleRenderable, RichCast
@@ -53,6 +53,7 @@ class CliArgs:
     show_guide: bool
     non_interactive: bool
     public_facing_address: str
+    accelerator: Optional[str] = None
 
 
 class PrerequisiteError(RichCast, Exception):
@@ -84,13 +85,21 @@ class DistInfo(BaseModel):
     image_refs: list[str] = Field(default_factory=list)
 
 
+class Accelerator(enum.StrEnum):
+    CUDA = "cuda"
+    CUDA_MOCK = "cuda_mock"
+    CUDA_MIG_MOCK = "cuda_mig_mock"
+    ROCM_MOCK = "rocm_mock"
+
+
 class InstallInfo(BaseModel):
     version: str
     type: InstallType
     last_updated: datetime
     base_path: Path
     halfstack_config: HalfstackConfig
     service_config: ServiceConfig
+    accelerator: Optional[Accelerator] = None
 
 
 @dataclasses.dataclass()
@@ -169,3 +178,4 @@ class ServiceConfig:
 @dataclasses.dataclass
 class InstallVariable:
     public_facing_address: str = "127.0.0.1"
+    accelerator: Optional[Accelerator] = None