Support splitting build jobs into shards (#2)

* Filter targets by product name/shard

Alternate take of #1

* Make dry output more useful.

Some formatting, and linting errors.

* Update readme, changelog and pre-commit hook.

* Sanity checks for bake's shard args.

* Make 0.0.0-dev default for stackable image versions.

* Update src/image_tools/args.py

Co-authored-by: Natalie <nat@nullable.se>

* Add help texts.

---------

Co-authored-by: Razvan-Daniel Mihai <84674+razvan@users.noreply.github.com>

.pre-commit-config.yaml

@@ -14,8 +14,8 @@ repos:
     hooks:
       - id: detect-secrets
         exclude: test/
-  - repo: https://github.com/pre-commit/mirrors-autopep8
-    rev: v2.0.1
+  - repo: https://github.com/hhatto/autopep8
+    rev: v2.0.4
     hooks:
       - id: autopep8
         entry: autopep8

CHANGELOG.md

@@ -3,3 +3,14 @@
 All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
+
+### Added
+
+- Allow `bake` product names to accept a version as suffix separated by "=" ([#2])
+- New `bake` command line options `--shard-index` and `--shard-count` ([#2])
+
+### Changed
+
+- Make stackable image version (`-i`) for `bake` optional and default to `0.0.0-dev` ([#2])
+
+[#2]: https://github.com/stackabletech/image-tools/pull/2

README.md

@@ -27,9 +27,16 @@ Following tools are installed:
 ## Examples
 
     # Build images of the hello-world containers
-    bake -p hello-world -i 0.0.0-dev --organization sandbox
-    # Run preflight checks on the hello-world container images
-    check-container -p hello-world -i 0.0.0-dev
+    bake -p hello-world -i 0.0.0-dev
+
+    # Build only one version [0.37.2] of OPA
+    bake -p opa=0.37.2 -i 0.0.0-dev
+
+    # Build half of all versions defined for OPA
+    bake -p opa -i 0.0.0-dev  --shard-count 2 --shard-index 0
+
+    # Build the other half of all versions defined for OPA
+    bake -p opa -i 0.0.0-dev  --shard-count 2 --shard-index 1
 
 ## Release a new version
 

src/image_tools/args.py

@@ -26,10 +26,16 @@ def bake_args() -> Namespace:
         "-i",
         "--image-version",
         help="Image version",
-        required=True,
+        default='0.0.0-dev',
         type=check_image_version_format,
     )
-    parser.add_argument("-p", "--product", help="Product to build images for")
+    parser.add_argument("-p", "--product",
+                        help="Product to build images for", action='append')
+    parser.add_argument("--shard-count", type=positive_int, default=1,
+                        help="Split the build into N shards, which can be built separately. \
+                        All shards must be built separately, by specifying the --shard-index argument.",)
+    parser.add_argument("--shard-index", type=positive_int, default=0,
+                        help="Build shard number M out of --shard-count. Shards are zero-indexed.")
     parser.add_argument("-u", "--push", help="Push images",
                         action="store_true")
     parser.add_argument("-d", "--dry", help="Dry run.", action="store_true")
@@ -53,7 +59,23 @@ def bake_args() -> Namespace:
         help="Image registry to publish to. Default: docker.stackable.tech",
         default="docker.stackable.tech",
     )
-    return parser.parse_args()
+    result = parser.parse_args()
+
+    if result.shard_index >= result.shard_count:
+        raise ValueError("shard index [{}] cannot be greater or equal than shard count [{}]".format(
+            result.shard_index, result.shard_count))
+    return result
+
+
+def positive_int(value) -> int:
+    try:
+        ivalue = int(value)
+        if ivalue < 0:
+            raise ValueError
+        return ivalue
+    except ValueError:
+        raise ValueError(
+            f"Invalid value [{value}]. Must be an integer greater than or equal to zero.")
 
 
 def check_image_version_format(image_version) -> str:

src/image_tools/bake.py

@@ -130,7 +130,25 @@ def bakefile_product_version_targets(
     }
 
 
-def bake_command(args: Namespace, product_name: str, bakefile) -> Command:
+def targets_for_selector(conf, selected_products: List[str]) -> List[str]:
+    targets = []
+    for selected_product in selected_products or (product['name'] for product in conf.products):
+        product_name, *versions = selected_product.split("=")
+        product = next(
+            (product for product in conf.products if product['name'] == product_name), None)
+        if product is None:
+            raise ValueError(f"Requested unknown product [{product_name}]")
+        for version in versions or (version['product'] for version in product['versions']):
+            targets.append(bakefile_target_name_for_product_version(
+                product_name, version))
+    return targets
+
+
+def filter_targets_for_shard(targets: List[str], shard_count: int, shard_index: int) -> List[str]:
+    return [target for i, target in enumerate(targets) if i % shard_count == shard_index]
+
+
+def bake_command(args: Namespace, targets: List[str], bakefile) -> Command:
     """
     Returns a list of commands that need to be run in order to build and
     publish product images.
@@ -145,14 +163,17 @@ def bake_command(args: Namespace, product_name: str, bakefile) -> Command:
     else:
         target_mode = []
 
+    if args.dry:
+        target_mode = ["--print"]
+
     return Command(
         args=[
             "docker",
             "buildx",
             "bake",
             "--file",
             "-",
-            *([] if product_name is None else [product_name]),
+            *targets,
             *target_mode,
         ],
         stdin=json.dumps(bakefile),
@@ -167,12 +188,17 @@ def main():
 
     bakefile = generate_bakefile(args, conf)
 
-    cmd = bake_command(args, args.product, bakefile)
+    targets = filter_targets_for_shard(targets_for_selector(
+        conf, args.product), args.shard_count, args.shard_index)
 
+    if not targets:
+        print("No targets match this filter")
+        return
+
+    cmd = bake_command(args, targets, bakefile)
     if args.dry:
-        print(cmd)
-    else:
-        run(cmd.args, input=cmd.input, check=True)
+        print(" ".join(cmd.args))
+    run(cmd.args, input=cmd.input, check=True)
 
 
 if __name__ == "__main__":