Add automated CMS data stitching workflow

end_of_run_workflow.py

@@ -8,13 +8,26 @@
 from prefect.settings import PREFECT_UI_URL
 
 #from analysis import run_analysis
+from stitch_tasks import run_auto_stitch_anchor, verify_stitch_outputs
 from data_validation import data_validation_task, get_run
 from linker import create_symlinks
 from dotenv import load_dotenv
+from workflow_settings import load_stitch_settings
 
 CATALOG_NAME = "cms"
 
 
+def _slack_run_message(title: str, flow_run_name: str, uid: str, scan_id: int, details: str | None = None) -> str:
+    lines = [
+        f"{title} (*{flow_run_name}*)",
+        f"run_start: {uid}",
+        f"scan_id: {scan_id}",
+    ]
+    if details:
+        lines.append(details)
+    return "\n".join(lines)
+
+
 def slack(func):
     """
     Send a message to mon-prefect and mon-prefect-cs slack channels if the flow-run failed.
@@ -26,8 +39,15 @@ def slack(func):
     the flow. To keep the naming of workflows consistent, the name of this inner function had to match the expected name.
     """
 
-    def end_of_run_workflow(stop_doc, api_key=None, dry_run=False):
-        flow_run_name = FlowRunContext.get().flow_run.dict().get("name")
+    def end_of_run_workflow(
+        stop_doc,
+        api_key=None,
+        dry_run=False,
+        workflow_options=None,
+    ):
+        flow_run = FlowRunContext.get().flow_run
+        flow_run_name = flow_run.dict().get("name")
+        flow_run_link = f"{PREFECT_UI_URL.value()}/flow-runs/{flow_run.id}"
 
         # Load slack credentials that are saved in Prefect.
         mon_prefect = SlackWebhook.load("mon-prefect")
@@ -45,28 +65,48 @@ def end_of_run_workflow(stop_doc, api_key=None, dry_run=False):
         # Send a message to mon-bluesky if bluesky-run failed.
         if stop_doc.get("exit_status") == "fail":
             mon_bluesky.notify(
-                f":bangbang: {CATALOG_NAME} bluesky-run failed. (*{flow_run_name}*)\n ```run_start: {uid}\nscan_id: {scan_id}``` ```reason: {stop_doc.get('reason', 'none')}```"
+                _slack_run_message(
+                    f":bangbang: {CATALOG_NAME} bluesky-run failed.",
+                    flow_run_name,
+                    uid,
+                    scan_id,
+                    details=f"reason: {stop_doc.get('reason', 'none')}",
+                )
             )
 
         try:
-            result = func(stop_doc, api_key=api_key, dry_run=dry_run)
+            result = func(
+                stop_doc,
+                api_key=api_key,
+                dry_run=dry_run,
+                workflow_options=workflow_options,
+            )
 
             # Send a message to mon-prefect-cms if flow-run is successful.
-            message = f":white_check_mark: {CATALOG_NAME} flow-run successful. (*{flow_run_name}*)\n ```run_start: {uid}\nscan_id: {scan_id}```"
+            message = (
+                f":white_check_mark: {CATALOG_NAME} flow-run successful. (*{flow_run_name}*)\n"
+                f"<{flow_run_link}|View flow run>\n"
+                f"```run_start: {uid}\nscan_id: {scan_id}```"
+            )
             mon_prefect_cms.notify(message)
             return result
         except Exception as e:
             tb = traceback.format_exception_only(e)
 
             # Send a message to mon-prefect-cms, mon-prefect if flow-run failed.
-            message = f":bangbang: {CATALOG_NAME} flow-run failed. (*{flow_run_name}*)\n ```run_start: {uid}\nscan_id: {scan_id}``` ```{tb[-1]}```"
+            message = (
+                f":bangbang: {CATALOG_NAME} flow-run failed. (*{flow_run_name}*)\n"
+                f"```run_start: {uid}\nscan_id: {scan_id}```\n"
+                f"```{tb[-1]}```"
+            )
             mon_prefect.notify(message)
             mon_prefect_cms.notify(message)
-            flow_run = FlowRunContext.get().flow_run
             # Add link to flow-run for the message to mon-prefect-cs.
             program_message = (
-                f":bangbang: {CATALOG_NAME} flow-run failed. <{PREFECT_UI_URL.value()}/flow-runs/"
-                + f"flow-run/{flow_run.id}|the flow run link> (*{flow_run_name}*)\n ```run_start: {uid}\nscan_id: {scan_id}``` ```{tb[-1]}```"
+                f":bangbang: {CATALOG_NAME} flow-run failed.\n"
+                f"<{flow_run_link}|the flow run link> (*{flow_run_name}*)\n"
+                f"```run_start: {uid}\nscan_id: {scan_id}```\n"
+                f"```{tb[-1]}```"
             )
             mon_prefect_cs.notify(program_message)
             raise
@@ -82,12 +122,18 @@ def log_completion():
 
 @flow(task_runner=ConcurrentTaskRunner())
 @slack
-def end_of_run_workflow(stop_doc, api_key=None, dry_run=False):
+def end_of_run_workflow(
+    stop_doc,
+    api_key=None,
+    dry_run=False,
+    workflow_options=None,
+):
     load_dotenv()
     logger = get_run_logger()
     uid = stop_doc["run_start"]
+    stitch = load_stitch_settings(workflow_options=workflow_options)
 
-    # Launch validation, analysis, and linker tasks concurrently
+    # Launch core tasks concurrently
     linker_task = create_symlinks.submit(uid, api_key=api_key, dry_run=dry_run)
     logger.info("Launched linker task")
 
@@ -97,9 +143,23 @@ def end_of_run_workflow(stop_doc, api_key=None, dry_run=False):
     # analysis_task = run_analysis(raw_ref=uid)
     # logger.info("Launched analysis task")
 
-    # Wait for all tasks to comple
+    pending = [linker_task, validation_task]
+    stitch_task = None
+
+    if stitch.enabled:
+        stitch_task = run_auto_stitch_anchor.submit(uid, api_key=api_key, stitch_config=stitch.config)
+        logger.info("Launched anchor auto-stitch task")
+        pending.append(stitch_task)
+    else:
+        logger.info("Anchor auto-stitch is disabled for this deployment")
+
     logger.info("Waiting for tasks to complete")
-    linker_task.result()
-    validation_task.result()
-    # analysis_task.result()
+    for t in pending:
+        t.result()
+
+    if stitch.enabled and stitch.verify_outputs and stitch_task is not None:
+        stitch_result = stitch_task.result()
+        if not stitch_result.get("skipped"):
+            verify_stitch_outputs.submit(stitch_result).result()
+
     log_completion()

linker.py

@@ -5,6 +5,8 @@
 from data_validation import get_run
 
 
+PROPOSAL_ROOT = Path("/nsls2/data/cms/proposals")
+
 
 def detector_mapping(detector):
     if detector in {"pilatus300k-1", "pilatus800k-2"}:
@@ -26,8 +28,27 @@ def chmod_and_chown(path, *, uid=None, gid=None, mode=0o775):
     #    os.chown(path, uid, gid)
 
 def make_relative_path(path_value):
+    """Prevent absolute metadata paths from escaping the proposal directory."""
     path = Path(path_value)
-    return Path(*path.parts[1:]) if path.is_absolute() else path
+    relative_path = Path(*path.parts[1:]) if path.is_absolute() else path
+    if ".." in relative_path.parts:
+        raise ValueError(f"Path must not contain traversal segments: {path_value}")
+    return relative_path
+
+def proposal_directory(doc):
+    """Return the proposal directory for a run start document."""
+    return PROPOSAL_ROOT / doc["cycle"] / doc["data_session"]
+
+def experiment_directory(doc):
+    """Return the experiments directory inside the proposal directory."""
+    return proposal_directory(doc) / make_relative_path(doc.get("experiments_directory", "experiments"))
+
+def experiment_alias_directory(doc):
+    """Return the user-facing experiment alias directory, or ``None`` if unset."""
+    path_expr_alias = doc.get("experiment_alias_directory")
+    if not path_expr_alias:
+        return None
+    return experiment_directory(doc) / make_relative_path(path_expr_alias)
 
 @task(retries=2, retry_delay_seconds=10)
 def create_symlinks(ref, api_key=None, dry_run=False):
@@ -53,17 +74,15 @@ def create_symlinks(ref, api_key=None, dry_run=False):
             else:
                 logger.info("Skipping the creation of the link because 'filename' is not set.")
                 return
-            if path_expr_alias := doc.get("experiment_alias_directory"):
-                path_proposal = Path(f"/nsls2/data/cms/proposals/{doc['cycle']}/{doc['data_session']}")
+            path_expr_alias = experiment_alias_directory(doc)
+            if path_expr_alias:
                 # stats = path_proposal.stat()
-                experiments_dir = make_relative_path(doc.get("experiments_directory", "experiments"))
-                path_expr = path_proposal / experiments_dir
+                path_expr = experiment_directory(doc)
                 if dry_run:
                     logger.info(f"Dry run: mkdir {path_expr}")
                 else:
                     path_expr.mkdir(exist_ok=True, parents=True)
                 #chmod_and_chown(path_expr, uid=stats.st_uid, gid=stats.st_gid)
-                path_expr_alias = path_expr / make_relative_path(path_expr_alias)
                 if dry_run:
                     logger.info(f"Dry run: mkdir {path_expr_alias}")
                 else: