chunk_partners.py (3317B)
1 # This Source Code Form is subject to the terms of the Mozilla Public 2 # License, v. 2.0. If a copy of the MPL was not distributed with this 3 # file, You can obtain one at http://mozilla.org/MPL/2.0/. 4 """ 5 Chunk the partner repack tasks by subpartner and locale 6 """ 7 8 import copy 9 10 from mozbuild.chunkify import chunkify 11 from taskgraph.transforms.base import TransformSequence 12 from taskgraph.util.dependencies import get_primary_dependency 13 14 from gecko_taskgraph.util.partners import ( 15 apply_partner_priority, 16 get_repack_ids_by_platform, 17 ) 18 19 transforms = TransformSequence() 20 transforms.add(apply_partner_priority) 21 22 23 @transforms.add 24 def chunk_partners(config, jobs): 25 for job in jobs: 26 dep_job = get_primary_dependency(config, job) 27 assert dep_job 28 29 build_platform = dep_job.attributes["build_platform"] 30 repack_id = dep_job.task.get("extra", {}).get("repack_id") 31 repack_ids = dep_job.task.get("extra", {}).get("repack_ids") 32 copy_repack_ids = job.pop("copy-repack-ids", False) 33 34 if copy_repack_ids: 35 assert repack_ids, f"dep_job {dep_job.label} doesn't have repack_ids!" 36 job.setdefault("extra", {})["repack_ids"] = repack_ids 37 yield job 38 # first downstream of the repack task, no chunking or fanout has been done yet 39 elif not any([repack_id, repack_ids]): 40 platform_repack_ids = get_repack_ids_by_platform(config, build_platform) 41 # we chunk mac signing 42 if config.kind in ( 43 "release-partner-repack-signing", 44 "release-eme-free-repack-signing", 45 "release-eme-free-repack-mac-signing", 46 "release-partner-repack-mac-signing", 47 ): 48 repacks_per_chunk = job.get("repacks-per-chunk") 49 chunks, remainder = divmod(len(platform_repack_ids), repacks_per_chunk) 50 if remainder: 51 chunks = int(chunks + 1) 52 for this_chunk in range(1, chunks + 1): 53 chunk = chunkify(platform_repack_ids, this_chunk, chunks) 54 partner_job = copy.deepcopy(job) 55 partner_job.setdefault("extra", {}).setdefault("repack_ids", chunk) 56 partner_job["extra"]["repack_suffix"] = str(this_chunk) 57 yield partner_job 58 # linux and windows we fan out immediately to one task per partner-sub_partner-locale 59 else: 60 for repack_id in platform_repack_ids: 61 partner_job = copy.deepcopy(job) # don't overwrite dict values here 62 partner_job.setdefault("extra", {}) 63 partner_job["extra"]["repack_id"] = repack_id 64 yield partner_job 65 # fan out chunked mac signing for repackage 66 elif repack_ids: 67 for repack_id in repack_ids: 68 partner_job = copy.deepcopy(job) 69 partner_job.setdefault("extra", {}).setdefault("repack_id", repack_id) 70 yield partner_job 71 # otherwise we've fully fanned out already, continue by passing repack_id on 72 else: 73 partner_job = copy.deepcopy(job) 74 partner_job.setdefault("extra", {}).setdefault("repack_id", repack_id) 75 yield partner_job