Merge pull request #14 from NSLS2/tw-validation

Add Data Validation in Tiled

Author: JunAishima

.env

@@ -0,0 +1 @@
+PREFECT_LOGGING_EXTRA_LOGGERS=validator

Dockerfile

@@ -7,6 +7,7 @@ RUN apt-get -y update && \
 
 COPY pixi.toml .
 COPY pixi.lock .
+COPY .env .
 # use `--locked` to ensure the lockfile is up to date with pixi.toml
 RUN pixi install --locked
 # create the shell-hook bash script to activate the environment

data_validation.py

@@ -1,24 +1,71 @@
-import time as ttime
+import time
 
 from prefect import flow, get_run_logger, task
 from prefect.blocks.system import Secret
-from tiled.client import from_profile
+
+from bluesky_tiled_plugins.writing.validator import validate
+from tiled.client import from_uri
+
+
+@task(retries=2, retry_delay_seconds=10)
+def get_run(uid, api_key=None):
+    tiled_client = from_uri("https://tiled.nsls2.bnl.gov", api_key=api_key)
+    run = tiled_client["cms/raw"][uid]
+    return run
+
+
+@task(retries=2, retry_delay_seconds=10)
+def get_run_migration(uid, api_key=None):
+    tiled_client = from_uri("https://tiled.nsls2.bnl.gov", api_key=api_key)
+    run = tiled_client["cms/migration"][uid]
+    return run
+
+
+@task
+def read_stream(run, stream):
+    stream_data = run[stream].read()
+    return stream_data
 
 
 @task(retries=2, retry_delay_seconds=10)
-def read_all_streams(beamline_acronym, uid):
+def read_all_streams(uid, api_key=None):
     logger = get_run_logger()
-    api_key = Secret.load("tiled-cms-api-key", _sync=True).get()
-    tiled_client = from_profile("nsls2", api_key=api_key)
-    run = tiled_client[beamline_acronym]["raw"][uid]
+    run = get_run(uid, api_key=api_key)
     logger.info(f"Validating uid {run.start['uid']}")
-    start_time = ttime.monotonic()
+    start_time = time.monotonic()
     for stream in run:
         logger.info(f"{stream}:")
-        stream_start_time = ttime.monotonic()
-        stream_data = run[stream].read()
-        stream_elapsed_time = ttime.monotonic() - stream_start_time
+        stream_start_time = time.monotonic()
+        stream_data = read_stream(run, stream)
+        stream_elapsed_time = time.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
+    elapsed_time = time.monotonic() - start_time
     logger.info(f"{elapsed_time = }")
+
+
+@task(retries=3, retry_delay_seconds=20)
+def data_validation_task(uid, api_key=None):
+    """Task to validate the data structure and accessibility in Tiled
+
+    Parameters
+    ----------
+        uid : str
+            The UID of the run to validate
+        beamline_acronym : str, optional
+            The acronym of the beamline (default is "cms")
+    """
+
+    logger = get_run_logger()
+    logger.info("Connecting to Tiled client for beamline cms")
+    run = get_run_migration(uid, api_key=api_key)
+    logger.info(f"Validating uid {uid}")
+    start_time = time.monotonic()
+    validate(run, fix_errors=True, try_reading=True, raise_on_error=True)
+    elapsed_time = time.monotonic() - start_time
+    logger.info(f"Finished validating data; {elapsed_time = }")
+
+
+@flow(log_prints=True)
+def data_validation_flow(uid, api_key=None):
+    data_validation_task(uid, api_key=api_key)

end_of_run_workflow.py

@@ -1,9 +1,18 @@
+import os
 from prefect import flow, get_run_logger, task
 from prefect.task_runners import ConcurrentTaskRunner
 
 #from analysis import run_analysis
-from data_validation import read_all_streams
+from data_validation import read_all_streams, data_validation_task
 from linker import create_symlinks
+from dotenv import load_dotenv
+
+
+def get_api_key_from_env():
+    with open("/srv/container.secret", "r") as secrets:
+        load_dotenv(stream=secrets)
+    api_key = os.environ["TILED_API_KEY"]
+    return api_key
 
 
 @task
@@ -13,23 +22,27 @@ def log_completion():
 
 
 @flow(task_runner=ConcurrentTaskRunner())
-def end_of_run_workflow(stop_doc):
+def end_of_run_workflow(stop_doc, api_key=None):
     logger = get_run_logger()
     uid = stop_doc["run_start"]
+    if not api_key:
+        api_key = get_api_key_from_env()
 
     # Launch validation, analysis, and linker tasks concurrently
-    linker_task = create_symlinks.submit(uid)
+    linker_task = create_symlinks.submit(uid, api_key=api_key)
     logger.info("Launched linker task")
 
-    validation_task = read_all_streams.submit("cms", uid)
-    logger.info("Launched validation task")
+    read_streams_task = read_all_streams.submit(uid, api_key=api_key)
+    validation_task = data_validation_task.submit(uid, api_key=api_key)
+    logger.info("Launched validation tasks")
 
     # analysis_task = run_analysis(raw_ref=uid)
     # logger.info("Launched analysis task")
 
     # Wait for all tasks to comple
     logger.info("Waiting for tasks to complete")
     linker_task.result()
+    read_streams_task.result()
     validation_task.result()
     # analysis_task.result()
     log_completion()

linker.py

@@ -1,16 +1,9 @@
 from prefect import task, get_run_logger
-from prefect.blocks.system import Secret
 from pathlib import Path
-from tiled.client import from_profile
 import os
 import glob
+from data_validation import get_run
 
-#tiled_client = from_uri('https://tiled.nsls2.bnl.gov')
-api_key = Secret.load("tiled-cms-api-key", _sync=True).get()
-tiled_client = from_profile("nsls2", api_key=api_key)['cms']
-tiled_client_raw = tiled_client["raw"]
-
-#logger = logging.getLogger()
 
 
 def detector_mapping(detector):
@@ -33,7 +26,7 @@ def chmod_and_chown(path, *, uid=None, gid=None, mode=0o775):
     #    os.chown(path, uid, gid)
 
 @task
-def create_symlinks(ref):
+def create_symlinks(ref, api_key=None):
     """
     Parameters
     ----------
@@ -43,7 +36,7 @@ def create_symlinks(ref):
     """
     logger = get_run_logger()
 
-    hrf = tiled_client_raw[ref]
+    hrf = get_run(ref, api_key=api_key)
     for name, doc in hrf.documents():
         if name == "start":
             if doc.get('experiment_project'):

pixi.toml

@@ -5,7 +5,7 @@ platforms = ["linux-64"]
 
 [dependencies]
 prefect = "3.*"
-tiled-client = ">=0.2.3"
+tiled = ">=0.2.3"
 python = "<3.12"
 bluesky-tiled-plugins = ">=2"
 prefect-docker = "*"