Update so prune and reindex work per-stack

Author: kwryankrattiger

spackbot/handlers/mirrors.py

@@ -3,55 +3,87 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
+from sh.contrib import git
+import os
 
 import spackbot.helpers as helpers
-from spackbot.helpers import pr_expected_base, pr_mirror_base_url, pr_shared_mirror
-from spackbot.workers import copy_pr_binaries, update_mirror_index, work_queue
+import spackbot.workers as workers
 
 # If we don't provide a timeout, the default in RQ is 180 seconds
 WORKER_JOB_TIMEOUT = 6 * 60 * 60
 
 logger = helpers.getLogger(__name__)
 
+def list_ci_stacks(branch=helpers.pr_expected_base):
+    with helpers.temp_dir() as cwd:
+        # Shallow clone of spack to use for spack getting the current stacks availabe
+        # in the base branch.
+        git.clone("--branch", branch, "--depth", 1, helpers.spack_upstream, "spack-develop")
+
+        stacks = []
+        pipeline_root = f"{cwd}/spack-develop/share/spack/gitlab/cloud_pipelines/stacks/"
+        for stack in os.listdir(pipeline_root):
+            if os.path.isfile(f"{pipeline_root}/{stack}/spack.yaml"):
+                stacks.append(stack)
+
+        return stacks
 
 async def graduate_pr_binaries(event, gh):
     payload = event.data
 
     base_branch = payload["pull_request"]["base"]["ref"]
     is_merged = payload["pull_request"]["merged"]
 
-    if is_merged and base_branch == pr_expected_base:
+    if is_merged and base_branch == helpers.pr_expected_base:
         pr_number = payload["number"]
         pr_branch = payload["pull_request"]["head"]["ref"]
 
-        shared_mirror_url = f"{pr_mirror_base_url}/{pr_shared_mirror}"
+        shared_mirror_url = f"{helpers.pr_mirror_base_url}/{helpers.pr_shared_mirror}"
 
         logger.info(
             f"PR {pr_number}/{pr_branch} merged to develop, graduating binaries"
         )
 
-        ltask_q = work_queue.get_lqueue()
+        ltask_q = workers.work_queue.get_lqueue()
         copy_job = ltask_q.enqueue(
-            copy_pr_binaries,
+            workers.copy_pr_binaries,
             pr_number,
             pr_branch,
-            shared_mirror_url,
+            f"{shared_mirror_url}",
             job_timeout=WORKER_JOB_TIMEOUT,
         )
+        copy_job.meta["info"] = {
+            "type": "copy",
+            "stack": "all",
+        }
         logger.info(f"Copy job queued: {copy_job.id}")
 
-        pr_number_i = int(pr_number)
-        if pr_number_i % 100 == 0:
+        for stack in list_ci_stacks():
+            pr_stack_mirror = f"{helpers.pr_mirror_bucket}/{stack}"
+            publish_stack_mirror = f"{helpers.publish_mirror_bucket}/{helpers.pr_excepted_base}/{stack}"
             prune_job = ltask_q.enqueue(
+                workers.prune_mirror_duplicates,
+                pr_stack_mirror,
+                publish_stack_mirror,
                 job_timeout=WORKER_JOB_TIMEOUT,
+                depends_on=copy_job
             )
+            prune_job.meta["info"] = {
+                "type": "prune",
+                "stack": stack,
+            }
             logger.info(f"Pruning job queued: {prune_job.id}")
 
-        # If the index job queue has a job queued already, there is no need to
-        # schedule another one
-        update_job = ltask_q.enqueue(
-            update_mirror_index,
-            shared_mirror_url,
-            job_timeout=WORKER_JOB_TIMEOUT,
-        )
-        logger.info(f"Reindex job queued: {update_job.id}")
+            # If the index job queue has a job queued already, there is no need to
+            # schedule another one
+            update_job = ltask_q.enqueue(
+                workers.update_mirror_index,
+                f"{shared_mirror_url}/{stack}",
+                job_timeout=WORKER_JOB_TIMEOUT,
+                depends_on=prune_job
+            )
+            update_job.meta["info"] = {
+                "type": "reindex",
+                "stack": stack,
+            }
+            logger.info(f"Reindex job queued: {update_job.id}")

spackbot/workers.py

@@ -331,7 +331,6 @@ async def fix_style_task(event):
                 event.data["issue"]["comments_url"], {}, data={"body": message}
             )
 
-
 async def find_latest_pipeline(url, headers, session):
     async with session.get(url, headers=headers) as response:
         pipeline_objects = await response.json()
@@ -394,7 +393,7 @@ async def download_spack_lock_files(url, headers, download_dir, session):
 
     return folder_list
 
-async def copy_pr_binaries(pr_number, pr_branch, shared_pr_mirror_url):
+async def copy_pr_binaries(stack, pr_number, pr_branch, shared_pr_mirror_url):
     """Find the latest gitlab pipeline for the PR, get the spack.lock
     for each child pipeline, and for each one, activate the environment
     and issue the spack buildcache sync command to copy between the
@@ -458,19 +457,26 @@ async def copy_pr_binaries(pr_number, pr_branch, shared_pr_mirror_url):
             pr_mirror_url
         ])
 
-def rq_has_reindex():
-    ltask_q = work_queue.get_lqueue()
-    for job in ltask_q.jobs:
-        if "update_mirror_index" in job.func_name:
-            return True
-    return False
-
+# Upate index per stack mirror
 async def update_mirror_index(mirror_url):
     """Use spack buildcache command to update index on remote mirror"""
 
+    # Current job stack
+    job = get_current_job()
+    stack = job.meta["info"]["stack"]
+
+    # Check if another reindex for this stack is queued
+    do_reindex = True
+    ltask_q = work_queue.get_lqueue()
+    for job in ltask_q.jobs:
+        info = job.meta["info"]
+        if info["type"] == "reindex" and info["stack"] == stack:
+            do_reindex = False
+            break
+
     # Check the queue for more reindex jobs, if there are none,
     # run reindex on the graduated PR mirror.
-    if not rq_has_reindex():
+    if do_reindex:
         print(f"Updating binary index at {mirror_url}")
         await helpers.run_in_subprocess([
             "spack",
@@ -484,20 +490,22 @@ async def update_mirror_index(mirror_url):
 
 def hash_from_key(key):
     h = None
+    # 32 chars long
     if key.lower().endswith(("spec.json")):
         h = re.search(".*-([a-zA-Z0-9]+)\.spec\.json", key.lower())
     elif key.lower().endswith(("spack")):
         h = re.search(".*-([a-zA-Z0-9]+)\.spack", key.lower())
     return h
 
-async def prune_mirror_duplicates():
+# Prune per stack mirror
+async def prune_mirror_duplicates(pr_mirror_bucket, publish_mirror_bucket):
     s3 = boto3.resource("s3")
-    pr_bucket = s3.Bucket(helpers.pr_mirror_bucket)
-    publish_bucket = s3.Bucket(helpers.publish_mirror_bucket)
+    pr_bucket = s3.Bucket(pr_mirror_bucket)
+    publish_bucket = s3.Bucket(publish_mirror_bucket)
 
     # Get the hashes in the shared PR bucket
     pr_specs = {}
-    for obj in pr_bucket.filter(Prefix=helpers.pr_shared_bucket).objects():
+    for obj in pr_bucket.objects():
         if not obj.key.endswith(("spec.json")):
             continue
 
@@ -506,12 +514,12 @@ async def prune_mirror_duplicates():
     # Clean up anything that didn't work and log an error. This should never happen
     # but also shouldn't stop the pruning.
     if None in pr_specs:
-        logger.error("Encountered hashless spec.json in shared PR mirror")
+        logger.error("Encountered hashless spec.json in PR mirror")
         pr_specs.remove(None)
 
     # Check in the published base branch bucket for duplicates to delete
     delete_specs = {}
-    for obj in publish_bucket.filter(Prefix=helpers.pr_expected_base).objects():
+    for obj in publish_bucket.objects():
         if not obj.key.endswith(("spec.json")):
             continue
 
@@ -520,13 +528,12 @@ async def prune_mirror_duplicates():
             delete_specs.add(spec_hash)
 
     # Delete all of the objects with marked hashes
-    if delete_specs:
-        logger.info("Pruning PR mirror")
-    for obj in pr_bucket.filter(Prefix=helpers.pr_shared_bucket).objects():
+    for obj in pr_bucket.objects():
         if not obj.key.endswith((".spec.json", ".spack")):
             continue
 
         if hash_from_key(obj.key) in delete_specs:
-            logger.debug(f"""pr mirror pruning is deleteing {obj.key}
-from s3://{helpers.pr_mirror_bucket}/{helpers.pr_shared_bucket}""")
+            logger.debug(
+                f"pr mirror pruning is deleteing {obj.key} from s3://{pr_mirror_bucket}"
+            )
             obj.delete()