Bug 1637752: Add support for generating multiple pools from a taskgraph style template; r=aki
authorTom Prince <mozilla@hocat.ca>
Thu, 14 May 2020 00:42:16 +0000
changeset 228 8431e7feea0764fdb843a689fdcb0c7bb0c26477
parent 227 1ceb3a410714366d66f85165bdabd6c440c4d7f8
child 229 1fe502e14c9ec3726c0c263e50a5eb1392be8008
push id164
push usermozilla@hocat.ca
push dateThu, 14 May 2020 00:50:31 +0000
treeherderci-admin@1fe502e14c9e [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersaki
bugs1637752
Bug 1637752: Add support for generating multiple pools from a taskgraph style template; r=aki Differential Revision: https://phabricator.services.mozilla.com/D74981
src/ciadmin/check/check_privileged.py
src/ciadmin/generate/ciconfig/worker_pools.py
src/ciadmin/generate/worker_pools.py
src/fxci/cli.py
src/fxci/worker_secrets.py
src/fxci/workers.py
--- a/src/ciadmin/check/check_privileged.py
+++ b/src/ciadmin/check/check_privileged.py
@@ -1,17 +1,19 @@
 # -*- coding: utf-8 -*-
 
 # This Source Code Form is subject to the terms of the Mozilla Public License,
 # v. 2.0. If a copy of the MPL was not distributed with this file, You can
 # obtain one at http://mozilla.org/MPL/2.0/.
 
 import pytest
 
+from ciadmin.generate.ciconfig.environment import Environment
 from ciadmin.generate.ciconfig.worker_pools import WorkerPool as WorkerPoolConfig
+from ciadmin.generate.worker_pools import generate_pool_variants
 from tcadmin.resources import WorkerPool
 from tcadmin.util.sessions import with_aiohttp_session
 
 
 @pytest.mark.asyncio
 @with_aiohttp_session
 async def check_privileged_is_untrusted(generated):
     """
@@ -49,36 +51,36 @@ def is_level3_worker(pool):
 
 
 @pytest.mark.asyncio
 @with_aiohttp_session
 async def check_trusted_level_3_workers():
     """
     Ensures that any trusted images are only used in level 3 workers.
     """
-
+    environment = await Environment.current()
     worker_pools = await WorkerPoolConfig.fetch_all()
-    for pool in worker_pools:
+    for pool in generate_pool_variants(worker_pools, environment):
         trusted = "trusted" in pool.config.get("image", "")
         pool_group, pool_id = pool.pool_id.split("/", 1)
         assert not all([trusted, not is_level3_worker(pool)]), (
             f"{pool.pool_id} has trusted CoT keys, "
             "but does not appear to be restricted to level 3 tasks"
         )
 
 
 @pytest.mark.asyncio
 @with_aiohttp_session
 async def check_level_3_worker_security():
     """
     Ensures that all level 3 workers have appropriate security groups.
     """
-
+    environment = await Environment.current()
     worker_pools = await WorkerPoolConfig.fetch_all()
-    for pool in worker_pools:
+    for pool in generate_pool_variants(worker_pools, environment):
         trusted_security = (
             ("trusted" == pool.config.get("security", "untrusted"))
             # GCP workers have networker security at the project/provider level
             or ("level3" in pool.provider_id)
         )
         assert not all([not trusted_security, is_level3_worker(pool)]), (
             f"{pool.pool_id} is a level 3 worker but has "
             f"unrestricted network access"
--- a/src/ciadmin/generate/ciconfig/worker_pools.py
+++ b/src/ciadmin/generate/ciconfig/worker_pools.py
@@ -9,19 +9,20 @@ import attr
 from .get import get_ciconfig_file
 
 
 @attr.s(frozen=True)
 class WorkerPool:
     pool_id = attr.ib(type=str)
     description = attr.ib(type=str)
     owner = attr.ib(type=str)
+    email_on_error = attr.ib(type=bool)
     provider_id = attr.ib(type=str)
     config = attr.ib()
-    email_on_error = attr.ib(type=bool)
+    variants = attr.ib(factory=lambda: [{}])
 
     @pool_id.validator
     def _check_pool_id(self, attribute, value):
         if value.count("/") != 1:
             raise ValueError(
                 "Worker pool_id must be of the form `provisionerId/workerPool`, "
                 f"not {value}"
             )
--- a/src/ciadmin/generate/worker_pools.py
+++ b/src/ciadmin/generate/worker_pools.py
@@ -1,16 +1,18 @@
 # -*- coding: utf-8 -*-
 
 # This Source Code Form is subject to the terms of the Mozilla Public License,
 # v. 2.0. If a copy of the MPL was not distributed with this file, You can
 # obtain one at http://mozilla.org/MPL/2.0/.
 
 import copy
 
+import attr
+
 from tcadmin.resources import WorkerPool
 
 from ..util.keyed_by import evaluate_keyed_by
 from ..util.templates import merge
 from .ciconfig.environment import Environment
 from .ciconfig.get import get_ciconfig_file
 from .ciconfig.worker_images import WorkerImage
 from .ciconfig.worker_pools import WorkerPool as ConfigWorkerPool
@@ -180,28 +182,60 @@ async def make_worker_pool(environment, 
         description=wp.description,
         owner=wp.owner,
         providerId=wp.provider_id,
         config=config,
         emailOnError=wp.email_on_error,
     )
 
 
+def generate_pool_variants(worker_pools, environment):
+    """
+    Generate the list of worker pools by evaluting them at all the specified
+    variants.
+    """
+
+    def update_config(config, name, variant, environment):
+        config = copy.deepcopy(config)
+        attributes = {"environment": environment}
+        attributes.update(variant)
+        for key in ["image", "maxCapacity", "minCapacity", "security"]:
+            if key in config:
+                value = evaluate_keyed_by(config[key], name, attributes)
+                if value is not None:
+                    config[key] = value
+                else:
+                    del config[key]
+        return config
+
+    for wp in worker_pools:
+        for variant in wp.variants:
+            name = wp.pool_id.format(**variant)
+            yield attr.evolve(
+                wp,
+                pool_id=name,
+                config=update_config(
+                    wp.config, name, variant, environment=environment.name
+                ),
+                variants=[{}],
+            )
+
+
 async def update_resources(resources):
     """
     Manage the worker-pool configurations
     """
     worker_pools = await ConfigWorkerPool.fetch_all()
     worker_images = await WorkerImage.fetch_all()
 
     resources.manage("WorkerPool=.*")
 
     worker_defaults = (await get_ciconfig_file("worker-pools.yml")).get(
         "worker-defaults"
     )
     environment = await Environment.current()
 
-    for wp in worker_pools:
+    for wp in generate_pool_variants(worker_pools, environment):
         apwt = await make_worker_pool(
             environment, resources, wp, worker_images, worker_defaults
         )
         if apwt:
             resources.add(apwt)
--- a/src/fxci/cli.py
+++ b/src/fxci/cli.py
@@ -135,16 +135,36 @@ async def check_worker_secrets(options):
 @app.argument("--update", help="Update existing worker secrets to match template.")
 @run_async
 async def create_worker_secrets(options):
     from .worker_secrets import create_worker_secrets
 
     await create_worker_secrets(update=options["update"])
 
 
+@app.command(
+    "generate-worker-pools",
+    help="Generate the complete list of worker pools definitions.",
+    description="This command expands the definied worker pools, "
+    "by evaluating all the specifed variants. This does *not* expand the definitions "
+    "to be useable by taskcluster (use `ci-admin generate` for that), but generates "
+    "input suitable for ci-admin without variants.",
+)
+@ciconfig_arguments(app)
+@format_arguments(app)
+@app.argument("--grep", help="Regular expression to limit the worker pools listed.")
+@run_async
+async def generate_worker_pools(options):
+    from .workers import generate_worker_pools
+
+    await generate_worker_pools(
+        grep=options["grep"], formatter=options["formatter"],
+    )
+
+
 @app.command("replay-hg-push", help="Retrigger the on-push tasks of an mercurial push")
 @app.argument("alias", help="The project alias of the push to retrigger.")
 @app.argument("revision", help="The tip revision of the push to retrigger.")
 @ciconfig_arguments(app, use_environment=False)
 @run_async
 async def replay_hg_push(options):
     from .hg_pushes import replay_hg_push
 
--- a/src/fxci/worker_secrets.py
+++ b/src/fxci/worker_secrets.py
@@ -20,33 +20,36 @@ solution for this.
 """
 
 import sys
 
 from taskcluster import optionsFromEnvironment
 from taskcluster.aio import Secrets
 from taskcluster.utils import fromNow
 
+from ciadmin.generate.ciconfig.environment import Environment
 from ciadmin.generate.ciconfig.worker_pools import WorkerPool
+from ciadmin.generate.worker_pools import generate_pool_variants
 from tcadmin.util.sessions import aiohttp_session, with_aiohttp_session
 
 
 async def list_secrets():
     secrets_api = Secrets(optionsFromEnvironment(), session=aiohttp_session())
     secrets = set()
     await secrets_api.list(
         paginationHandler=lambda response: secrets.update(response["secrets"])
     )
     return secrets
 
 
 async def docker_worker_pools():
     docker_pools = set()
+    environment = await Environment.current()
     worker_pools = await WorkerPool.fetch_all()
-    for worker_pool in worker_pools:
+    for worker_pool in generate_pool_variants(worker_pools, environment):
         implementation = worker_pool.config.get("implementation", "docker-worker")
         if implementation == "docker-worker":
             docker_pools.add(worker_pool.pool_id)
 
     return docker_pools
 
 
 @with_aiohttp_session
--- a/src/fxci/workers.py
+++ b/src/fxci/workers.py
@@ -1,17 +1,55 @@
 # -*- coding: utf-8 -*-
 
 # This Source Code Form is subject to the terms of the Mozilla Public License,
 # v. 2.0. If a copy of the MPL was not distributed with this file, You can
 # obtain one at http://mozilla.org/MPL/2.0/.
 
+import re
+
+import attr
 from taskcluster import WorkerManager, optionsFromEnvironment
 
+from ciadmin.generate.ciconfig.environment import Environment
+from ciadmin.generate.ciconfig.worker_pools import WorkerPool
+from ciadmin.generate.worker_pools import generate_pool_variants
+
 
 def list_workers(worker_pool, states, *, formatter):
     workers = []
     wm = WorkerManager(optionsFromEnvironment())
     wm.listWorkersForWorkerPool(
         worker_pool, paginationHandler=lambda r: workers.extend(r["workers"])
     )
 
     formatter([worker for worker in workers if worker["state"] in states])
+
+
+async def generate_worker_pools(*, formatter, grep):
+    worker_pools = await WorkerPool.fetch_all()
+    environment = await Environment.current()
+
+    if grep:
+        grep = re.compile(grep)
+
+    def match(pool):
+        if grep:
+            return grep.match(pool.pool_id)
+        return True
+
+    def to_json(worker_pool):
+        # attr.asdict generates a dictionary that matches the order of
+        # attributes in WorkerPool, so we ask the formatter to not sort keys.
+        # However, `pool.config` does not have that structure, so we explictly
+        # sort it here.
+        result = attr.asdict(worker_pool, filter=lambda a, _: a.name != "variants")
+        result["config"] = dict(sorted(result["config"].items()))
+        return result
+
+    formatter(
+        [
+            to_json(pool)
+            for pool in sorted(generate_pool_variants(worker_pools, environment))
+            if match(pool)
+        ],
+        sort_keys=False,
+    )