Buildx builder context manager

Author: david-tingdahl-nvidia

ci/ci_utils.py

@@ -19,7 +19,7 @@
 import subprocess
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Any
 import re
 
 
@@ -49,6 +49,42 @@ class CudaSmArchitectures(Enum):
 MAX_CONSECUTIVE_IDENTICAL_LOG_LINES = 100
 
 
+class DockerBuildxBuilder:
+    """Context manager for starting and stopping a docker buildx instance.
+       buildx instance is needed to limit the resources used by docker build."""
+
+    def __init__(self, name: str) -> None:
+        self.builder_name = name
+        self.builder_cpus = 6    #TODO: Make this adaptive to the number of cores available
+        self.start_builder()
+
+    def __enter__(self) -> Any:
+        """Start the builder and return it."""
+        self.start_builder()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        """Stop the builder."""
+        self.stop_builder()
+
+    def start_builder(self) -> None:
+        """Start the builder."""
+        # Make sure it is stopped
+        self.stop_builder(check=False)
+        subprocess.run([
+            'docker', 'buildx', 'create', f'--name={self.builder_name}',
+            '--driver=docker-container', f'--driver-opt=cpu-shares={self.builder_cpus}', '--use'
+        ],
+                       check=True)
+
+        print(f'Started buildx builder: {self.builder_name}')
+
+    def stop_builder(self, check: bool = True) -> None:
+        """Stop the builder."""
+        subprocess.run(['docker', 'buildx', 'rm', self.builder_name], check=check)
+        print(f'Stopped buildx builder: {self.builder_name}')
+
+
 def _try_parse_gcc_output_line(line: str) -> Tuple[Optional[str], Optional[int]]:
     """Try to extract file path and line number from gcc/clang/cmake-style output"""
     # Example: /path/to/file.cpp:LINE:
@@ -227,34 +263,37 @@ def build(self) -> None:
         print('=' * 80)
         print('', flush=True)
 
-        cmd = [
-            'docker',
-            'build',
-            '-f',
-            self.dockerfile_path(),
-            '-t',
-            image_name,
-            '--network=host',
-            '--progress=plain',
-        ]
-
-        if parent is not None:
-            parent_name = parent.image_name()
-            cmd += ['--build-arg', f'BASE_IMAGE={parent_name}']
-
-        if self.build_args() is not None:
-            for arg in self.build_args():
-                cmd += ['--build-arg', arg]
-
-        # Add extra docker args from args if provided
-        if self.args.user_build_args is not None:
-            cmd += self.args.user_build_args
-
-        cmd += ['.']
-
-        print(' '.join(cmd))
-
-        _run_and_parse_log(cmd)
+        with DockerBuildxBuilder(name=f'{self.image_name()}_builder') as builder:
+            cmd = [
+                'docker',
+                'buildx',
+                'build',
+                f'--builder={builder.builder_name}',
+                '-f',
+                self.dockerfile_path(),
+                '-t',
+                image_name,
+                '--network=host',
+                '--progress=plain',
+            ]
+
+            if parent is not None:
+                parent_name = parent.image_name()
+                cmd += ['--build-arg', f'BASE_IMAGE={parent_name}']
+
+            if self.build_args() is not None:
+                for arg in self.build_args():
+                    cmd += ['--build-arg', arg]
+
+            # Add extra docker args from args if provided
+            if self.args.user_build_args is not None:
+                cmd += self.args.user_build_args
+
+            cmd += ['.']
+
+            print(' '.join(cmd))
+
+            _run_and_parse_log(cmd)
 
         if self.do_validate_image():
             self._validate()