From d5f13caee03cebad7706f948851173ab4e88f367 Mon Sep 17 00:00:00 2001 From: Codex Bot Date: Fri, 2 Oct 2026 14:17:42 +0800 Subject: [PATCH] Add AutoDL public base image catalog --- .gitea/workflows/manage.yaml | 18 +++++ README.md | 6 +- controller.py | 136 ++++++++++++++++++++++++++++++++++- tests/test_controller.py | 14 ++++ 4 files changed, 171 insertions(+), 3 deletions(-) diff --git a/.gitea/workflows/manage.yaml b/.gitea/workflows/manage.yaml index b9dfb12..4068d95 100644 --- a/.gitea/workflows/manage.yaml +++ b/.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 diff --git a/README.md b/README.md index 29200c3..0f8714c 100644 --- a/README.md +++ b/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 diff --git a/controller.py b/controller.py index 371aa97..04013f7 100644 --- a/controller.py +++ b/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="") diff --git a/tests/test_controller.py b/tests/test_controller.py index dd2826c..b37b743 100644 --- a/tests/test_controller.py +++ b/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)