Add AutoDL public base image catalog

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="")