Add AutoDL public base image catalog

.gitea/workflows/manage.yaml

@@ -11,6 +11,18 @@ on:
         description: Required only for provision
         required: false
         default: ""
+      framework:
+        description: Public image framework, for example pytorch
+        required: false
+        default: ""
+      framework_version:
+        description: Public image framework version
+        required: false
+        default: ""
+      cuda_version:
+        description: Public image CUDA version
+        required: false
+        default: ""
       confirm_spend:
         description: Type CREATE to provision a billed instance
         required: false
@@ -25,6 +37,9 @@ jobs:
       TRAINING_RUNNER_REGISTRATION_TOKEN: ${{ secrets.TRAINING_RUNNER_REGISTRATION_TOKEN }}
       OPERATION: ${{ inputs.operation }}
       IMAGE_UUID: ${{ inputs.image_uuid }}
+      FRAMEWORK: ${{ inputs.framework }}
+      FRAMEWORK_VERSION: ${{ inputs.framework_version }}
+      CUDA_VERSION: ${{ inputs.cuda_version }}
       CONFIRM_SPEND: ${{ inputs.confirm_spend }}
     steps:
       - name: Select deployed controller
@@ -44,6 +59,9 @@ jobs:
         run: |
           /usr/local/bin/ces-controller "$OPERATION" \
             --image-uuid "$IMAGE_UUID" \
+            --framework "$FRAMEWORK" \
+            --framework-version "$FRAMEWORK_VERSION" \
+            --cuda-version "$CUDA_VERSION" \
             --confirm-spend "$CONFIRM_SPEND"
 
       - name: Controller host details

README.md

@@ -58,7 +58,7 @@ ces-controller --help
 
 从 Actions 手动运行 `Manage AutoDL GPU runner`：
 
-- `discover`：列出可供 API 使用的私有镜像，不产生 GPU 费用；
+- `discover`：列出账号私有镜像和内置公共基础镜像，不产生 GPU 费用；
 - `provision`：创建 PRO6000、传输模型并注册 GPU Runner；必须填写镜像 UUID，并把确认值
   填为 `CREATE`；
 - `status`：查询实例和 Runner 进程状态；
@@ -70,6 +70,10 @@ ces-controller --help
 控制器固定使用官方 `gitea-runner 3.5.0` Linux AMD64 二进制及其 SHA-256，下载和校验均在
 创建付费实例之前完成。
 
+`discover` 同时输出账号私有镜像和控制器内置的 AutoDL 公共基础镜像目录。`provision` 可以
+直接接收 `image_uuid`，也可以通过 `framework`、`framework_version` 和 `cuda_version` 三项
+精确选择公共镜像；没有精确匹配时会停止，不会自动降级到旧框架或旧 CUDA。
+
 AutoDL GPU Runner 注册为 host runner，并提供：
 
 ```text

controller.py

@@ -38,6 +38,99 @@ MODEL_FILES = [
     "tokenizer.model",
     "tokenizer_config.json",
 ]
+PUBLIC_BASE_IMAGES = [
+    {
+        "image_uuid": "base-image-12be412037",
+        "frameworks": {"pytorch": "1.9.0"},
+        "cuda_version": "11.1",
+        "python_version": "3.8",
+        "os": "ubuntu18.04",
+    },
+    {
+        "image_uuid": "base-image-u9r24vthlk",
+        "frameworks": {"pytorch": "1.10.0"},
+        "cuda_version": "11.3",
+        "python_version": "3.8",
+        "os": "ubuntu20.04",
+    },
+    {
+        "image_uuid": "base-image-l374uiucui",
+        "frameworks": {"pytorch": "1.11.0"},
+        "cuda_version": "11.3",
+        "python_version": "3.8",
+        "os": "ubuntu20.04",
+    },
+    {
+        "image_uuid": "base-image-l2t43iu6uk",
+        "frameworks": {"pytorch": "2.0.0", "tensorrt": "8.5.1"},
+        "cuda_version": "11.8",
+        "python_version": "3.8",
+        "os": "ubuntu20.04",
+    },
+    {
+        "image_uuid": "base-image-0gxqmciyth",
+        "frameworks": {"tensorflow": "2.5.0"},
+        "cuda_version": "11.2",
+        "python_version": "3.8",
+        "os": "ubuntu18.04",
+    },
+    {
+        "image_uuid": "base-image-uxeklgirir",
+        "frameworks": {"tensorflow": "2.9.0"},
+        "cuda_version": "11.2",
+        "python_version": "3.8",
+        "os": "ubuntu20.04",
+    },
+    {
+        "image_uuid": "base-image-4bpg0tt88l",
+        "frameworks": {"tensorflow": "1.15.5"},
+        "cuda_version": "11.4",
+        "python_version": "3.8",
+        "os": None,
+    },
+    {
+        "image_uuid": "base-image-mbr2n4urrc",
+        "frameworks": {"miniconda": "conda3"},
+        "cuda_version": "11.6",
+        "python_version": "3.8",
+        "os": "ubuntu20.04",
+    },
+    {
+        "image_uuid": "base-image-qkkhitpik5",
+        "frameworks": {"miniconda": "conda3"},
+        "cuda_version": "10.2",
+        "python_version": "3.8",
+        "os": "ubuntu18.04",
+    },
+    {
+        "image_uuid": "base-image-h041hn36yt",
+        "frameworks": {"miniconda": "conda3"},
+        "cuda_version": "11.1",
+        "python_version": "3.8",
+        "os": "ubuntu18.04",
+    },
+    {
+        "image_uuid": "base-image-7bn8iqhkb5",
+        "frameworks": {"miniconda": "conda3", "cudagl": "11.3"},
+        "cuda_version": "11.3",
+        "python_version": "3.8",
+        "os": "ubuntu20.04",
+    },
+    {
+        "image_uuid": "base-image-k0vep6kyq8",
+        "frameworks": {"miniconda": "conda3"},
+        "cuda_version": "9.0",
+        "python_version": "3.6",
+        "os": "ubuntu16.04",
+    },
+    {
+        "image_uuid": "base-image-l2843iu23k",
+        "frameworks": {"tensorrt": "8.5.1"},
+        "cuda_version": "11.8",
+        "python_version": "3.8",
+        "os": "ubuntu20.04",
+    },
+]
 
 
 class ControllerError(RuntimeError):
@@ -228,6 +321,36 @@ def instance_spec(image_uuid: str, disk_gb: int, data_centers: list[str]) -> dic
     return spec
 
 
+def select_public_image(framework: str, framework_version: str, cuda_version: str) -> str:
+    normalized_framework = framework.strip().casefold()
+    matches = [
+        image
+        for image in PUBLIC_BASE_IMAGES
+        if image["frameworks"].get(normalized_framework) == framework_version.strip()
+        and image["cuda_version"] == cuda_version.strip()
+    ]
+    if len(matches) != 1:
+        raise ControllerError(
+            "no unique public base image matches "
+            f"framework={framework!r}, version={framework_version!r}, "
+            f"cuda={cuda_version!r}; provide --image-uuid or update the catalog"
+        )
+    return str(matches[0]["image_uuid"])
+
+
+def requested_image_uuid(args: argparse.Namespace) -> str:
+    if args.image_uuid:
+        return args.image_uuid
+    profile = (args.framework, args.framework_version, args.cuda_version)
+    if any(profile):
+        if not all(profile):
+            raise ControllerError(
+                "--framework, --framework-version, and --cuda-version must be used together"
+            )
+        return select_public_image(*profile)
+    raise ControllerError("an image UUID or a complete public image profile is required")
+
+
 def validate_local(paths: Paths, runner_token: str) -> None:
     if not runner_token:
         raise ControllerError("TRAINING_RUNNER_REGISTRATION_TOKEN is missing")
@@ -418,7 +541,13 @@ def discover(client: AutoDLClient) -> None:
         }
         for item in images
     ]
-    print(json.dumps({"private_images": safe}, indent=2, ensure_ascii=False))
+    print(
+        json.dumps(
+            {"private_images": safe, "public_base_images": PUBLIC_BASE_IMAGES},
+            indent=2,
+            ensure_ascii=False,
+        )
+    )
 
 
 def provision(args: argparse.Namespace, client: AutoDLClient, paths: Paths) -> None:
@@ -437,7 +566,7 @@ def provision(args: argparse.Namespace, client: AutoDLClient, paths: Paths) -> N
                     f"refusing to create a second billed instance; "
                     f"{previous_uuid} is {previous_status}"
                 )
-    spec = instance_spec(args.image_uuid, args.disk_gb, args.data_center)
+    spec = instance_spec(requested_image_uuid(args), args.disk_gb, args.data_center)
     instance_uuid: str | None = None
     try:
         instance_uuid = client.create(spec)
@@ -520,6 +649,9 @@ def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument("operation", choices=("discover", "provision", "status", "shutdown"))
     parser.add_argument("--image-uuid", default="")
+    parser.add_argument("--framework", default="")
+    parser.add_argument("--framework-version", default="")
+    parser.add_argument("--cuda-version", default="")
     parser.add_argument("--disk-gb", type=int, default=120)
     parser.add_argument("--data-center", action="append", default=[])
     parser.add_argument("--confirm-spend", default="")

tests/test_controller.py

@@ -59,6 +59,20 @@ class ControllerTests(unittest.TestCase):
         with self.assertRaises(controller.ControllerError):
             controller.instance_spec("", 120, [])
 
+    def test_public_image_catalog_has_unique_uuids(self):
+        uuids = [image["image_uuid"] for image in controller.PUBLIC_BASE_IMAGES]
+        self.assertEqual(len(uuids), len(set(uuids)))
+
+    def test_selects_exact_public_image_profile(self):
+        self.assertEqual(
+            controller.select_public_image("PyTorch", "2.0.0", "11.8"),
+            "base-image-l2t43iu6uk",
+        )
+
+    def test_refuses_to_downgrade_an_unknown_public_image_profile(self):
+        with self.assertRaises(controller.ControllerError):
+            controller.select_public_image("pytorch", "2.8.0", "12.8")
+
     def test_runner_has_required_host_labels_and_no_docker(self):
         rendered = controller.runner_yaml()
         self.assertIn('"ubuntu-latest:host"', rendered)