Update towards template repo

Dockerfile

@@ -1,4 +1,4 @@
-FROM ghcr.io/prefix-dev/pixi:0.57.0
+FROM ghcr.io/prefix-dev/pixi:0.71.3
 
 ENV TZ="America/New_York"
 

data_validation.py

@@ -1,28 +1,42 @@
-from prefect import task, flow, get_run_logger
-from prefect.blocks.system import Secret
 import time as ttime
-from tiled.client import from_profile
+
+from dotenv import dotenv_values
+from prefect import flow, get_run_logger, task
+from tiled.client import from_uri
+
+BEAMLINE_OR_ENDSTATION = "tst"
+
+
+def get_api_key_from_env():
+    config = dotenv_values("/srv/container.secret")
+    return config.get("TILED_API_KEY")
+
+
+@task(retries=2, retry_delay_seconds=10)
+def get_run(uid, api_key=None):
+    if not api_key:
+        api_key = get_api_key_from_env()
+    cl = from_uri("https://tiled.nsls2.bnl.gov", api_key=api_key)
+    return cl[f"{BEAMLINE_OR_ENDSTATION}/raw"][uid]
 
 
 @task(retries=2, retry_delay_seconds=10)
-def read_all_streams(uid, beamline_acronym):
+def read_stream(run, stream):
+    return run[stream].read()
+
+
+@flow
+def data_validation(uid, api_key=None):
     logger = get_run_logger()
-    api_key = Secret.load("tiled-tst-api-key").get()
-    cl = from_profile("nsls2", api_key=api_key)
-    run = cl["tst"]["raw"][uid]
-    logger.info(f"Validating uid {run.start['uid']}")
+    logger.info(f"Validating uid {uid}")
     start_time = ttime.monotonic()
-    for stream in run:
+    run_client = get_run(uid, api_key=api_key)
+    for stream in run_client:
         logger.info(f"{stream}:")
         stream_start_time = ttime.monotonic()
-        stream_data = run[stream].read()
+        stream_data = read_stream(run_client, stream)  # noqa: F841
         stream_elapsed_time = ttime.monotonic() - stream_start_time
         logger.info(f"{stream} elapsed_time = {stream_elapsed_time}")
         logger.info(f"{stream} nbytes = {stream_data.nbytes:_}")
     elapsed_time = ttime.monotonic() - start_time
     logger.info(f"{elapsed_time = }")
-
-
-@flow
-def data_validation(uid):
-    read_all_streams(uid, beamline_acronym="tst")

end_of_run_workflow.py

@@ -1,20 +1,101 @@
-from prefect import task, flow, get_run_logger
-from data_validation import data_validation
+import traceback
+
+from data_validation import data_validation, get_run
+from prefect import flow, get_run_logger, task
+from prefect.blocks.notifications import SlackWebhook
+from prefect.context import FlowRunContext
 from test_extra_client import get_other_docs
+
+CATEGORY_NAME = "tst"
+SLACK_GENERAL = "mon-prefect"
+SLACK_BLUESKY = "mon-bluesky"
+SLACK_STATUS = "mon-prefect-tst"
+
+
+def slack(func):
+    def end_of_run_workflow(stop_doc, api_key=None, dry_run=False):
+        logger = get_run_logger()
+        flow_run_name = FlowRunContext.get().flow_run.dict().get("name")
+
+        def _load_block(name):
+            try:
+                return SlackWebhook.load(name)
+            except Exception:
+                logger.exception(f"Failed to load Slack webhook block '{name}'")
+                return None
+
+        def _notify(webhook, message, description):
+            if webhook is None:
+                logger.warning(
+                    f"Skipping {description} notification: webhook not available"
+                )
+                return
+            try:
+                webhook.notify(message)
+            except Exception:
+                logger.exception(f"Failed to send {description} notification")
+
+        mon_prefect = _load_block(SLACK_GENERAL)
+        mon_prefect_tst = _load_block(SLACK_STATUS)
+        mon_bluesky = _load_block(SLACK_BLUESKY)
+
+        uid = stop_doc.get("run_start", "unknown")
+        scan_id = "unknown"
+        try:
+            run = get_run(uid, api_key=api_key)
+            scan_id = run.start["scan_id"]
+            if stop_doc.get("exit_status") == "fail":
+                _notify(
+                    mon_bluesky,
+                    f":bangbang: {CATEGORY_NAME} bluesky-run failed. "
+                    f"(*{flow_run_name}*)\n ```run_start: {uid}\n"
+                    f"scan_id: {scan_id}``` ```reason: "
+                    f"{stop_doc.get('reason', 'none')}```",
+                    SLACK_BLUESKY,
+                )
+        except Exception:
+            logger.exception(
+                f"Exception while checking {uid}/{scan_id} for scan exit status"
+            )
+
+        try:
+            result = func(stop_doc, api_key=api_key, dry_run=dry_run)
+            _notify(
+                mon_prefect_tst,
+                f":white_check_mark: {CATEGORY_NAME} flow-run successful. "
+                f"(*{flow_run_name}*)\n ```run_start: {uid}\n"
+                f"scan_id: {scan_id}```",
+                SLACK_STATUS,
+            )
+            return result
+        except Exception as error:
+            tb = traceback.format_exception_only(type(error), error)
+            message = (
+                f":bangbang: {CATEGORY_NAME} flow-run failed. "
+                f"(*{flow_run_name}*)\n ```run_start: {uid}\n"
+                f"scan_id: {scan_id}``` ```{tb[-1]}```"
+            )
+            _notify(mon_prefect_tst, message, SLACK_STATUS)
+            _notify(mon_prefect, message, SLACK_GENERAL)
+            raise
+
+    return end_of_run_workflow
+
+
 # from long_flow import long_flow
 
 
 @task
-def log_completion():
+def log_completion(dry_run=False):
     logger = get_run_logger()
-    logger.info("Complete")
+    logger.info(f"Complete! dry_run: {dry_run}")
 
 
 @flow
-def end_of_run_workflow(stop_doc):
+@slack
+def end_of_run_workflow(stop_doc, api_key=None, dry_run=False):
     uid = stop_doc["run_start"]
-    # hello_world()
-    data_validation(uid, return_state=True)
-    get_other_docs(uid)
-    # long_flow(iterations=100, sleep_length=10)
-    log_completion()
+    data_validation(uid, api_key=api_key)
+    get_other_docs(uid, api_key=api_key)
+    # long_flow(iterations=100, sleep_length=10)  # keep in to potentially run as a test in the future
+    log_completion(dry_run=dry_run)