Merge pull request #96 from ImogenBits/size_limit

Add docker image size limit

Author: Benezivas

algobattle/docker_util.py

@@ -52,6 +52,7 @@ class DockerConfig(BaseModel):
     build_timeout: float | None = None
     strict_timeouts: bool = False
     set_cpus: str | list[str] | None = None
+    image_size: int | None = None
     generator: RunParameters = RunParameters()
     solver: RunParameters = RunParameters()
     advanced_run_params: "AdvancedRunArgs | None" = None
@@ -82,6 +83,7 @@ def set_docker_config(config: DockerConfig) -> None:
     if config.advanced_build_params is not None:
         Image.build_kwargs = config.advanced_build_params.to_docker_args()
     Program.docker_config = config
+    Image.docker_config = config
 
 
 class ProgramUiProxy(Protocol):
@@ -107,6 +109,8 @@ class Image:
     id: str
     path: Path
 
+    docker_config: ClassVar[DockerConfig] = DockerConfig()
+
     run_kwargs: ClassVar[dict[str, Any]] = {
         "network_mode": "none",
     }
@@ -183,7 +187,15 @@ async def build(
         except APIError as e:
             raise BuildError("Docker APIError thrown while building.", detail=str(e)) from e
 
-        return cls(cast(str, image.id), path=path)
+        self = cls(cast(str, image.id), path=path)
+        size_limit = cls.docker_config.image_size
+        used_size = cast(dict[str, Any], image.attrs).get("Size", 0)
+        if size_limit is not None and used_size > size_limit:
+            try:
+                self.remove()
+            finally:
+                raise BuildError("Built image is too large.", detail=f"Built size: {used_size}, limit: {size_limit}.")
+        return self
 
     def __enter__(self):
         return self