From 2dbe8fc541f621053241fb3a6224e087a1ba3978 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 5 Feb 2026 13:32:39 -0500 Subject: [PATCH 01/28] update pixi.toml with explicit addition of python-dotenv * already included in lock file, so no update necessary --- pixi.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pixi.toml b/pixi.toml index 50e3351..6438d0d 100644 --- a/pixi.toml +++ b/pixi.toml @@ -8,6 +8,7 @@ prefect = "3.*" python = "<3.14" tiled-client = ">=0.2.3" bluesky-tiled-plugins = ">=2" +python-dotenv = ">=1.2.1,<2" [pypi-dependencies] lixtools = "==2023.1.23.0" From 52c46fcdf4775ab17b206181ec029bf76d335722 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 5 Feb 2026 13:33:25 -0500 Subject: [PATCH 02/28] update code to use dotenv * read TILED_API_KEY in from env.secrets --- data_validation.py | 9 +++++++-- utils.py | 12 ++++++++---- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/data_validation.py b/data_validation.py index 60ad578..7da4597 100644 --- a/data_validation.py +++ b/data_validation.py @@ -1,5 +1,7 @@ +import os + +from dotenv import load_dotenv from prefect import task, flow, get_run_logger -from prefect.blocks.system import Secret import time as ttime from tiled.client import from_profile @@ -7,7 +9,10 @@ @task(retries=2, retry_delay_seconds=10) def read_all_streams(uid, beamline_acronym): logger = get_run_logger() - api_key = Secret.load("tiled-tst-api-key").get() + with open("/srv/env.secrets", "r") as secrets: + load_dotenv(stream=secrets) + api_key = os.environ["TILED_API_KEY"] + logger.info(f"first 4 characters of key: {api_key[:4]}") cl = from_profile("nsls2", api_key=api_key) run = cl["tst"]["raw"][uid] logger.info(f"Validating uid {run.start['uid']}") diff --git a/utils.py b/utils.py index 1cd0072..cbd4935 100644 --- a/utils.py +++ b/utils.py @@ -1,5 +1,6 @@ +from dotenv import load_dotenv +from prefect import get_run_logger from tiled.client import from_profile -from prefect.blocks.system import Secret import os @@ -7,7 +8,10 @@ def get_tiled_client(): - os.environ["TILED_API_KEY"] = Secret.load(f"tiled-{LOCATION}-api-key").get() - tiled_client = from_profile("nsls2")[LOCATION] - os.environ.pop("TILED_API_KEY") + logger = get_run_logger() + with open("/srv/env.secrets", "r") as secrets: + load_dotenv(stream=secrets) + api_key = os.environ["TILED_API_KEY"] + logger.info(f"first 4 characters of key: {api_key:4}") + tiled_client = from_profile("nsls2", api_key=api_key)[LOCATION] return tiled_client From 1f6cb2094241016a46713b8f29fd79453e9bb74a Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 5 Feb 2026 13:34:01 -0500 Subject: [PATCH 03/28] add env.secrets file into container --- prefect.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/prefect.yaml b/prefect.yaml index 065a601..79412b5 100644 --- a/prefect.yaml +++ b/prefect.yaml @@ -33,5 +33,6 @@ deployments: volumes: - /nsls2/data/tst/proposals:/nsls2/data/tst/proposals - /nsls2/software/etc/tiled:/nsls2/software/etc/tiled + - /srv/prefect3-docker-worker-tst/env.secret:/srv/env.secret auto_remove: true name: tst-work-pool-docker From 95ad0462907dc27277601d920c34cd7918029732 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 5 Feb 2026 13:41:46 -0500 Subject: [PATCH 04/28] build container, use update-dotenv code for deployment --- .github/workflows/publish-ghcr.yml | 2 +- prefect.yaml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/publish-ghcr.yml b/.github/workflows/publish-ghcr.yml index 45e4f05..110ad75 100644 --- a/.github/workflows/publish-ghcr.yml +++ b/.github/workflows/publish-ghcr.yml @@ -4,7 +4,7 @@ name: Create and publish a Docker image # Configures this workflow to run every time a change is pushed to the branch called `release`. on: push: - branches: ["main"] + branches: ["update-dotenv"] # Defines two custom environment variables for the workflow. These are used for the Container registry domain, and a name for the Docker image that this workflow builds. env: diff --git a/prefect.yaml b/prefect.yaml index 79412b5..a733b21 100644 --- a/prefect.yaml +++ b/prefect.yaml @@ -10,7 +10,7 @@ pull: directory: /repo - prefect.deployments.steps.git_clone: repository: https://github.com/nsls2/tst-workflows.git - branch: main + branch: update-dotenv deployments: - name: tst-end-of-run-workflow-docker @@ -26,7 +26,7 @@ deployments: job_variables: env: TILED_SITE_PROFILES: /nsls2/software/etc/tiled/profiles - image: ghcr.io/nsls2/tst-workflows:main + image: ghcr.io/nsls2/tst-workflows:update-dotenv image_pull_policy: Always network_mode: slirp4netns userns: "keep-id:uid=402974,gid=402974" # workflow-tst:workflow-tst From ff41adb50098223d53b9036e1de63084fa2e311a Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 5 Feb 2026 13:59:24 -0500 Subject: [PATCH 05/28] update deployment version --- prefect.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/prefect.yaml b/prefect.yaml index a733b21..9f2b159 100644 --- a/prefect.yaml +++ b/prefect.yaml @@ -14,7 +14,7 @@ pull: deployments: - name: tst-end-of-run-workflow-docker - version: 0.1.1 + version: 0.1.2 tags: - tst - main From 1e0fd69223007a6a4936e99d9c34eceec2d840f2 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 5 Feb 2026 14:11:49 -0500 Subject: [PATCH 06/28] fix filename --- data_validation.py | 2 +- prefect.yaml | 2 +- utils.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/data_validation.py b/data_validation.py index 7da4597..c8fb636 100644 --- a/data_validation.py +++ b/data_validation.py @@ -9,7 +9,7 @@ @task(retries=2, retry_delay_seconds=10) def read_all_streams(uid, beamline_acronym): logger = get_run_logger() - with open("/srv/env.secrets", "r") as secrets: + with open("/srv/env.secret", "r") as secrets: load_dotenv(stream=secrets) api_key = os.environ["TILED_API_KEY"] logger.info(f"first 4 characters of key: {api_key[:4]}") diff --git a/prefect.yaml b/prefect.yaml index 9f2b159..70a420d 100644 --- a/prefect.yaml +++ b/prefect.yaml @@ -33,6 +33,6 @@ deployments: volumes: - /nsls2/data/tst/proposals:/nsls2/data/tst/proposals - /nsls2/software/etc/tiled:/nsls2/software/etc/tiled - - /srv/prefect3-docker-worker-tst/env.secret:/srv/env.secret + - /srv/prefect3-docker-worker-tst:/srv auto_remove: true name: tst-work-pool-docker diff --git a/utils.py b/utils.py index cbd4935..71209bf 100644 --- a/utils.py +++ b/utils.py @@ -9,7 +9,7 @@ def get_tiled_client(): logger = get_run_logger() - with open("/srv/env.secrets", "r") as secrets: + with open("/srv/env.secret", "r") as secrets: load_dotenv(stream=secrets) api_key = os.environ["TILED_API_KEY"] logger.info(f"first 4 characters of key: {api_key:4}") From 706b6eb416944b193233487b028375f5140b30dd Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 5 Feb 2026 14:42:35 -0500 Subject: [PATCH 07/28] update location of Tiled secret * corresponding changes in prefect3_worker role --- data_validation.py | 2 +- prefect.yaml | 2 +- utils.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/data_validation.py b/data_validation.py index c8fb636..2a62240 100644 --- a/data_validation.py +++ b/data_validation.py @@ -9,7 +9,7 @@ @task(retries=2, retry_delay_seconds=10) def read_all_streams(uid, beamline_acronym): logger = get_run_logger() - with open("/srv/env.secret", "r") as secrets: + with open("/srv/tiled.secret", "r") as secrets: load_dotenv(stream=secrets) api_key = os.environ["TILED_API_KEY"] logger.info(f"first 4 characters of key: {api_key[:4]}") diff --git a/prefect.yaml b/prefect.yaml index 70a420d..25a6220 100644 --- a/prefect.yaml +++ b/prefect.yaml @@ -33,6 +33,6 @@ deployments: volumes: - /nsls2/data/tst/proposals:/nsls2/data/tst/proposals - /nsls2/software/etc/tiled:/nsls2/software/etc/tiled - - /srv/prefect3-docker-worker-tst:/srv + - /srv/prefect3-docker-worker-tst/tiled:/srv auto_remove: true name: tst-work-pool-docker diff --git a/utils.py b/utils.py index 71209bf..2b477ea 100644 --- a/utils.py +++ b/utils.py @@ -9,7 +9,7 @@ def get_tiled_client(): logger = get_run_logger() - with open("/srv/env.secret", "r") as secrets: + with open("/srv/tiled.secret", "r") as secrets: load_dotenv(stream=secrets) api_key = os.environ["TILED_API_KEY"] logger.info(f"first 4 characters of key: {api_key:4}") From ed405f32158fcade504ed4ab4fa19e7c41e0d49a Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 5 Feb 2026 15:19:37 -0500 Subject: [PATCH 08/28] remove logging (testing dotenv mechanism) --- data_validation.py | 1 - utils.py | 1 - 2 files changed, 2 deletions(-) diff --git a/data_validation.py b/data_validation.py index 2a62240..186d369 100644 --- a/data_validation.py +++ b/data_validation.py @@ -12,7 +12,6 @@ def read_all_streams(uid, beamline_acronym): with open("/srv/tiled.secret", "r") as secrets: load_dotenv(stream=secrets) api_key = os.environ["TILED_API_KEY"] - logger.info(f"first 4 characters of key: {api_key[:4]}") cl = from_profile("nsls2", api_key=api_key) run = cl["tst"]["raw"][uid] logger.info(f"Validating uid {run.start['uid']}") diff --git a/utils.py b/utils.py index 2b477ea..732c67c 100644 --- a/utils.py +++ b/utils.py @@ -12,6 +12,5 @@ def get_tiled_client(): with open("/srv/tiled.secret", "r") as secrets: load_dotenv(stream=secrets) api_key = os.environ["TILED_API_KEY"] - logger.info(f"first 4 characters of key: {api_key:4}") tiled_client = from_profile("nsls2", api_key=api_key)[LOCATION] return tiled_client From eea5a2feebee0f3f779b796caebbaf687ef9caa9 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 5 Feb 2026 15:20:50 -0500 Subject: [PATCH 09/28] LINT: remove unused logging, and thus its import --- utils.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/utils.py b/utils.py index 732c67c..ecefd2d 100644 --- a/utils.py +++ b/utils.py @@ -1,5 +1,4 @@ from dotenv import load_dotenv -from prefect import get_run_logger from tiled.client import from_profile import os @@ -8,7 +7,6 @@ def get_tiled_client(): - logger = get_run_logger() with open("/srv/tiled.secret", "r") as secrets: load_dotenv(stream=secrets) api_key = os.environ["TILED_API_KEY"] From 9b6ba0e97fb16f7d02526ca81de56b10fb34a11a Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Tue, 3 Mar 2026 16:45:56 -0500 Subject: [PATCH 10/28] restructure api key handling * include possibility of calling with API key value to better enable local testing of workflows --- data_validation.py | 12 +++--------- end_of_run_workflow.py | 19 +++++++++++++++++-- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/data_validation.py b/data_validation.py index 186d369..ceeb353 100644 --- a/data_validation.py +++ b/data_validation.py @@ -1,17 +1,11 @@ -import os - -from dotenv import load_dotenv from prefect import task, flow, get_run_logger import time as ttime from tiled.client import from_profile @task(retries=2, retry_delay_seconds=10) -def read_all_streams(uid, beamline_acronym): +def read_all_streams(uid, beamline_acronym, api_key=None): logger = get_run_logger() - with open("/srv/tiled.secret", "r") as secrets: - load_dotenv(stream=secrets) - api_key = os.environ["TILED_API_KEY"] cl = from_profile("nsls2", api_key=api_key) run = cl["tst"]["raw"][uid] logger.info(f"Validating uid {run.start['uid']}") @@ -28,5 +22,5 @@ def read_all_streams(uid, beamline_acronym): @flow -def data_validation(uid): - read_all_streams(uid, beamline_acronym="tst") +def data_validation(uid, api_key=None): + read_all_streams(uid, beamline_acronym="tst", api_key=api_key) diff --git a/end_of_run_workflow.py b/end_of_run_workflow.py index f56291b..d13f6f4 100644 --- a/end_of_run_workflow.py +++ b/end_of_run_workflow.py @@ -1,9 +1,23 @@ +import os + from prefect import task, flow, get_run_logger from data_validation import data_validation from test_extra_client import get_other_docs +from dotenv import load_dotenv # from long_flow import long_flow +def get_api_key_from_env(api_key=None): + logger = get_run_logger() + try: + with open("/srv/tiled.secret", "r") as secrets: + load_dotenv(stream=secrets) + api_key = os.environ["TILED_API_KEY"] + except Exception: + logger.exception("Exception while getting Tiled API key") + return api_key + + @task def log_completion(): logger = get_run_logger() @@ -14,7 +28,8 @@ def log_completion(): def end_of_run_workflow(stop_doc): uid = stop_doc["run_start"] # hello_world() - data_validation(uid, return_state=True) - get_other_docs(uid) + api_key = get_api_key_from_env(api_key=None) + data_validation(uid, return_state=True, api_key=api_key) + get_other_docs(uid, api_key=api_key) # long_flow(iterations=100, sleep_length=10) log_completion() From 2feb28bd5bbdef0046e099e960943f47838384f3 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Tue, 3 Mar 2026 16:51:40 -0500 Subject: [PATCH 11/28] fix Tiled API key handling for utils/test_extra_client --- test_extra_client.py | 2 +- utils.py | 6 +----- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/test_extra_client.py b/test_extra_client.py index 30d0b88..3b0acc8 100644 --- a/test_extra_client.py +++ b/test_extra_client.py @@ -3,7 +3,7 @@ @task -def get_other_docs(uid): +def get_other_docs(uid, api_key=api_key): logger = get_run_logger() result = get_tiled_client()["raw"][uid] for name, doc in result.documents(): diff --git a/utils.py b/utils.py index ecefd2d..82b6538 100644 --- a/utils.py +++ b/utils.py @@ -1,4 +1,3 @@ -from dotenv import load_dotenv from tiled.client import from_profile import os @@ -6,9 +5,6 @@ LOCATION = "tst" -def get_tiled_client(): - with open("/srv/tiled.secret", "r") as secrets: - load_dotenv(stream=secrets) - api_key = os.environ["TILED_API_KEY"] +def get_tiled_client(api_key=None): tiled_client = from_profile("nsls2", api_key=api_key)[LOCATION] return tiled_client From f59fc6f2b4896df5d1a800e9bf446a86158c5b17 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Tue, 3 Mar 2026 17:35:55 -0500 Subject: [PATCH 12/28] pre-commit fixes --- test_extra_client.py | 2 +- utils.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/test_extra_client.py b/test_extra_client.py index 3b0acc8..3885982 100644 --- a/test_extra_client.py +++ b/test_extra_client.py @@ -3,7 +3,7 @@ @task -def get_other_docs(uid, api_key=api_key): +def get_other_docs(uid, api_key=None): logger = get_run_logger() result = get_tiled_client()["raw"][uid] for name, doc in result.documents(): diff --git a/utils.py b/utils.py index 82b6538..076b0fc 100644 --- a/utils.py +++ b/utils.py @@ -1,6 +1,5 @@ from tiled.client import from_profile -import os LOCATION = "tst" From 84cef9e85e76d665298ec0abb72677409ec5e55e Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Mon, 9 Mar 2026 10:15:50 -0400 Subject: [PATCH 13/28] update Tiled secret location --- end_of_run_workflow.py | 2 +- prefect.yaml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/end_of_run_workflow.py b/end_of_run_workflow.py index d13f6f4..e45bdf2 100644 --- a/end_of_run_workflow.py +++ b/end_of_run_workflow.py @@ -10,7 +10,7 @@ def get_api_key_from_env(api_key=None): logger = get_run_logger() try: - with open("/srv/tiled.secret", "r") as secrets: + with open("/srv/container.secret", "r") as secrets: load_dotenv(stream=secrets) api_key = os.environ["TILED_API_KEY"] except Exception: diff --git a/prefect.yaml b/prefect.yaml index 25a6220..e719e2e 100644 --- a/prefect.yaml +++ b/prefect.yaml @@ -33,6 +33,6 @@ deployments: volumes: - /nsls2/data/tst/proposals:/nsls2/data/tst/proposals - /nsls2/software/etc/tiled:/nsls2/software/etc/tiled - - /srv/prefect3-docker-worker-tst/tiled:/srv + - /srv/prefect3-docker-worker-tst/app:/srv auto_remove: true name: tst-work-pool-docker From 29dc6a9e17873811bd810dbbea2b8f4ed0d93e61 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Mon, 9 Mar 2026 10:18:47 -0400 Subject: [PATCH 14/28] change from branch to main * publish-ghcr and the deploy file had the local branch name to use for testing. change those to main for prod --- .github/workflows/publish-ghcr.yml | 2 +- prefect.yaml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/publish-ghcr.yml b/.github/workflows/publish-ghcr.yml index 110ad75..45e4f05 100644 --- a/.github/workflows/publish-ghcr.yml +++ b/.github/workflows/publish-ghcr.yml @@ -4,7 +4,7 @@ name: Create and publish a Docker image # Configures this workflow to run every time a change is pushed to the branch called `release`. on: push: - branches: ["update-dotenv"] + branches: ["main"] # Defines two custom environment variables for the workflow. These are used for the Container registry domain, and a name for the Docker image that this workflow builds. env: diff --git a/prefect.yaml b/prefect.yaml index e719e2e..85c0b07 100644 --- a/prefect.yaml +++ b/prefect.yaml @@ -10,7 +10,7 @@ pull: directory: /repo - prefect.deployments.steps.git_clone: repository: https://github.com/nsls2/tst-workflows.git - branch: update-dotenv + branch: main deployments: - name: tst-end-of-run-workflow-docker @@ -26,7 +26,7 @@ deployments: job_variables: env: TILED_SITE_PROFILES: /nsls2/software/etc/tiled/profiles - image: ghcr.io/nsls2/tst-workflows:update-dotenv + image: ghcr.io/nsls2/tst-workflows:main image_pull_policy: Always network_mode: slirp4netns userns: "keep-id:uid=402974,gid=402974" # workflow-tst:workflow-tst From 447db342f111b8cd9ce2ec20c4e84dc5470bac1e Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Mon, 9 Mar 2026 10:31:10 -0400 Subject: [PATCH 15/28] replace from_profile with from_uri * enable running on hosts without access to our filesystems, such as Github action-based CI --- data_validation.py | 4 ++-- prefect.yaml | 3 --- utils.py | 4 ++-- 3 files changed, 4 insertions(+), 7 deletions(-) diff --git a/data_validation.py b/data_validation.py index ceeb353..c9571f5 100644 --- a/data_validation.py +++ b/data_validation.py @@ -1,12 +1,12 @@ from prefect import task, flow, get_run_logger import time as ttime -from tiled.client import from_profile +from tiled.client import from_uri @task(retries=2, retry_delay_seconds=10) def read_all_streams(uid, beamline_acronym, api_key=None): logger = get_run_logger() - cl = from_profile("nsls2", api_key=api_key) + cl = from_uri("https://tiled.nsls2.bnl.gov", api_key=api_key) run = cl["tst"]["raw"][uid] logger.info(f"Validating uid {run.start['uid']}") start_time = ttime.monotonic() diff --git a/prefect.yaml b/prefect.yaml index 85c0b07..6dba8c9 100644 --- a/prefect.yaml +++ b/prefect.yaml @@ -24,15 +24,12 @@ deployments: schedule: {} work_pool: job_variables: - env: - TILED_SITE_PROFILES: /nsls2/software/etc/tiled/profiles image: ghcr.io/nsls2/tst-workflows:main image_pull_policy: Always network_mode: slirp4netns userns: "keep-id:uid=402974,gid=402974" # workflow-tst:workflow-tst volumes: - /nsls2/data/tst/proposals:/nsls2/data/tst/proposals - - /nsls2/software/etc/tiled:/nsls2/software/etc/tiled - /srv/prefect3-docker-worker-tst/app:/srv auto_remove: true name: tst-work-pool-docker diff --git a/utils.py b/utils.py index 076b0fc..a39a448 100644 --- a/utils.py +++ b/utils.py @@ -1,9 +1,9 @@ -from tiled.client import from_profile +from tiled.client import from_uri LOCATION = "tst" def get_tiled_client(api_key=None): - tiled_client = from_profile("nsls2", api_key=api_key)[LOCATION] + tiled_client = from_uri("https://tiled.nsls2.bnl.gov", api_key=api_key)[LOCATION] return tiled_client From d0be97e3c9a3edbb68267a18ad0acc82de3e4b86 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Mon, 9 Mar 2026 10:39:10 -0400 Subject: [PATCH 16/28] add dry_run option --- data_validation.py | 9 ++++++--- end_of_run_workflow.py | 14 +++++++------- long_flow.py | 9 ++++++--- test_extra_client.py | 5 ++++- 4 files changed, 23 insertions(+), 14 deletions(-) diff --git a/data_validation.py b/data_validation.py index c9571f5..58e00f5 100644 --- a/data_validation.py +++ b/data_validation.py @@ -4,8 +4,11 @@ @task(retries=2, retry_delay_seconds=10) -def read_all_streams(uid, beamline_acronym, api_key=None): +def read_all_streams(uid, beamline_acronym, api_key=None, dry_run=False): logger = get_run_logger() + if dry_run: + logger.info("Dry run: not creating tiled client or checking streams") + return cl = from_uri("https://tiled.nsls2.bnl.gov", api_key=api_key) run = cl["tst"]["raw"][uid] logger.info(f"Validating uid {run.start['uid']}") @@ -22,5 +25,5 @@ def read_all_streams(uid, beamline_acronym, api_key=None): @flow -def data_validation(uid, api_key=None): - read_all_streams(uid, beamline_acronym="tst", api_key=api_key) +def data_validation(uid, api_key=None, dry_run=False): + read_all_streams(uid, beamline_acronym="tst", api_key=api_key, dry_run=dry_run) diff --git a/end_of_run_workflow.py b/end_of_run_workflow.py index e45bdf2..af0aba9 100644 --- a/end_of_run_workflow.py +++ b/end_of_run_workflow.py @@ -19,17 +19,17 @@ def get_api_key_from_env(api_key=None): @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): +def end_of_run_workflow(stop_doc, dry_run=False): uid = stop_doc["run_start"] # hello_world() api_key = get_api_key_from_env(api_key=None) - data_validation(uid, return_state=True, api_key=api_key) - get_other_docs(uid, api_key=api_key) - # long_flow(iterations=100, sleep_length=10) - log_completion() + data_validation(uid, return_state=True, api_key=api_key, dry_run=dry_run) + get_other_docs(uid, api_key=api_key, dry_run=dry_run) + # long_flow(iterations=100, sleep_length=10, dry_run=dry_run) + log_completion(dry_run=dry_run) diff --git a/long_flow.py b/long_flow.py index ca48572..c6d6c6e 100644 --- a/long_flow.py +++ b/long_flow.py @@ -3,16 +3,19 @@ @task -def print_and_sleep(iterations, sleep_length): +def print_and_sleep(iterations, sleep_length, dry_run=False): # Long running task print("Long task...") + if dry_run: + logger.info("Dry run: skipping long task") + return for i in range(int(iterations)): print(f"Iteration number {i}") ttime.sleep(int(sleep_length)) @flow(log_prints=True) -def long_flow(iterations, sleep_length): +def long_flow(iterations, sleep_length, dry_run=False): print("Starting long flow...") - print_and_sleep(iterations, sleep_length) + print_and_sleep(iterations, sleep_length, dry_run=dry_run) print("Done!") diff --git a/test_extra_client.py b/test_extra_client.py index 3885982..742177c 100644 --- a/test_extra_client.py +++ b/test_extra_client.py @@ -3,8 +3,11 @@ @task -def get_other_docs(uid, api_key=None): +def get_other_docs(uid, api_key=None, dry_run=False): logger = get_run_logger() + if dry_run: + logger.info(f"Dry run: not getting other docs") + return result = get_tiled_client()["raw"][uid] for name, doc in result.documents(): logger.info(f"name: {name}, doc: {doc}") From 84b916137384b463cd9f1561c378b64fce89f8e5 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Mon, 9 Mar 2026 10:39:52 -0400 Subject: [PATCH 17/28] update deployment version number --- prefect.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/prefect.yaml b/prefect.yaml index 6dba8c9..ca5b5c6 100644 --- a/prefect.yaml +++ b/prefect.yaml @@ -14,7 +14,7 @@ pull: deployments: - name: tst-end-of-run-workflow-docker - version: 0.1.2 + version: 0.1.3 tags: - tst - main From e32cd6ebfbdf8b26889db39108f5255d30fd120c Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Mon, 9 Mar 2026 10:41:02 -0400 Subject: [PATCH 18/28] add option to call end_of_run_workflow with api_key --- end_of_run_workflow.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/end_of_run_workflow.py b/end_of_run_workflow.py index af0aba9..00d0711 100644 --- a/end_of_run_workflow.py +++ b/end_of_run_workflow.py @@ -25,10 +25,11 @@ def log_completion(dry_run=False): @flow -def end_of_run_workflow(stop_doc, dry_run=False): +def end_of_run_workflow(stop_doc, api_key=None, dry_run=False): uid = stop_doc["run_start"] # hello_world() - api_key = get_api_key_from_env(api_key=None) + if not api_key: + api_key = get_api_key_from_env(api_key=None) data_validation(uid, return_state=True, api_key=api_key, dry_run=dry_run) get_other_docs(uid, api_key=api_key, dry_run=dry_run) # long_flow(iterations=100, sleep_length=10, dry_run=dry_run) From f328b26eb1d237fe5ed4aa47c1819678c392f90d Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Mon, 9 Mar 2026 11:11:19 -0400 Subject: [PATCH 19/28] fix pre-commit issue, add import for logger --- long_flow.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/long_flow.py b/long_flow.py index c6d6c6e..4b9fd5b 100644 --- a/long_flow.py +++ b/long_flow.py @@ -1,9 +1,10 @@ -from prefect import task, flow +from prefect import task, flow, get_run_logger import time as ttime @task def print_and_sleep(iterations, sleep_length, dry_run=False): + logger = get_run_logger() # Long running task print("Long task...") if dry_run: From 4ffc6595d11ee6903d90b325d76ba1f26d542020 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Mon, 9 Mar 2026 11:14:29 -0400 Subject: [PATCH 20/28] fix pre-commit issue --- test_extra_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test_extra_client.py b/test_extra_client.py index 742177c..8a97606 100644 --- a/test_extra_client.py +++ b/test_extra_client.py @@ -6,7 +6,7 @@ def get_other_docs(uid, api_key=None, dry_run=False): logger = get_run_logger() if dry_run: - logger.info(f"Dry run: not getting other docs") + logger.info("Dry run: not getting other docs") return result = get_tiled_client()["raw"][uid] for name, doc in result.documents(): From 4ee1bb510ac3d488731d14f954e7fdeab878580f Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Tue, 10 Mar 2026 16:35:20 -0400 Subject: [PATCH 21/28] make more similar to smi * remove dry_run from reading functions * make get_run() function and use it where Tiled clients are required * remove redundant utils module --- data_validation.py | 16 +++++++++------- end_of_run_workflow.py | 14 +++++--------- test_extra_client.py | 9 +++------ utils.py | 9 --------- 4 files changed, 17 insertions(+), 31 deletions(-) delete mode 100644 utils.py diff --git a/data_validation.py b/data_validation.py index 58e00f5..7b0cf14 100644 --- a/data_validation.py +++ b/data_validation.py @@ -2,15 +2,17 @@ import time as ttime from tiled.client import from_uri +@task +def get_run(uid, api_key=None): + cl = from_uri("https://tiled.nsls2.bnl.gov", api_key=api_key) + run = cl["tst/raw"][uid] + return run + @task(retries=2, retry_delay_seconds=10) -def read_all_streams(uid, beamline_acronym, api_key=None, dry_run=False): +def read_all_streams(uid, beamline_acronym, api_key=None): logger = get_run_logger() - if dry_run: - logger.info("Dry run: not creating tiled client or checking streams") - return - cl = from_uri("https://tiled.nsls2.bnl.gov", api_key=api_key) - run = cl["tst"]["raw"][uid] + run = get_run(uid, api_key=api_key) logger.info(f"Validating uid {run.start['uid']}") start_time = ttime.monotonic() for stream in run: @@ -26,4 +28,4 @@ def read_all_streams(uid, beamline_acronym, api_key=None, dry_run=False): @flow def data_validation(uid, api_key=None, dry_run=False): - read_all_streams(uid, beamline_acronym="tst", api_key=api_key, dry_run=dry_run) + read_all_streams(uid, beamline_acronym="tst", api_key=api_key) diff --git a/end_of_run_workflow.py b/end_of_run_workflow.py index 00d0711..880166e 100644 --- a/end_of_run_workflow.py +++ b/end_of_run_workflow.py @@ -8,13 +8,9 @@ def get_api_key_from_env(api_key=None): - logger = get_run_logger() - try: - with open("/srv/container.secret", "r") as secrets: - load_dotenv(stream=secrets) - api_key = os.environ["TILED_API_KEY"] - except Exception: - logger.exception("Exception while getting Tiled API key") + with open("/srv/container.secret", "r") as secrets: + load_dotenv(stream=secrets) + api_key = os.environ["TILED_API_KEY"] return api_key @@ -30,7 +26,7 @@ def end_of_run_workflow(stop_doc, api_key=None, dry_run=False): # hello_world() if not api_key: api_key = get_api_key_from_env(api_key=None) - data_validation(uid, return_state=True, api_key=api_key, dry_run=dry_run) - get_other_docs(uid, api_key=api_key, dry_run=dry_run) + data_validation(uid, return_state=True, api_key=api_key) + get_other_docs(uid, api_key=api_key) # long_flow(iterations=100, sleep_length=10, dry_run=dry_run) log_completion(dry_run=dry_run) diff --git a/test_extra_client.py b/test_extra_client.py index 8a97606..8910303 100644 --- a/test_extra_client.py +++ b/test_extra_client.py @@ -1,13 +1,10 @@ from prefect import task, get_run_logger -from utils import get_tiled_client +from data_validation import get_run @task -def get_other_docs(uid, api_key=None, dry_run=False): +def get_other_docs(uid, api_key=None): logger = get_run_logger() - if dry_run: - logger.info("Dry run: not getting other docs") - return - result = get_tiled_client()["raw"][uid] + result = get_run(uid, api_key=api_key) for name, doc in result.documents(): logger.info(f"name: {name}, doc: {doc}") diff --git a/utils.py b/utils.py deleted file mode 100644 index a39a448..0000000 --- a/utils.py +++ /dev/null @@ -1,9 +0,0 @@ -from tiled.client import from_uri - - -LOCATION = "tst" - - -def get_tiled_client(api_key=None): - tiled_client = from_uri("https://tiled.nsls2.bnl.gov", api_key=api_key)[LOCATION] - return tiled_client From b55f00daea2ae0b65e7684c633b8c12c66ca50a4 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Tue, 10 Mar 2026 16:41:29 -0400 Subject: [PATCH 22/28] fix pre-commit issues --- data_validation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/data_validation.py b/data_validation.py index 7b0cf14..b534597 100644 --- a/data_validation.py +++ b/data_validation.py @@ -2,6 +2,7 @@ import time as ttime from tiled.client import from_uri + @task def get_run(uid, api_key=None): cl = from_uri("https://tiled.nsls2.bnl.gov", api_key=api_key) From 2919d3753c65611f0e34df9d3c5e56dd895ba2e7 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 12 Mar 2026 10:35:53 -0400 Subject: [PATCH 23/28] extract read_stream(), clean up params for read_all_streams() --- data_validation.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/data_validation.py b/data_validation.py index b534597..6ea39ea 100644 --- a/data_validation.py +++ b/data_validation.py @@ -10,8 +10,13 @@ def get_run(uid, api_key=None): return run +@task +def read_stream(run, stream): + return run[stream].read() + + @task(retries=2, retry_delay_seconds=10) -def read_all_streams(uid, beamline_acronym, api_key=None): +def read_all_streams(uid, api_key=None): logger = get_run_logger() run = get_run(uid, api_key=api_key) logger.info(f"Validating uid {run.start['uid']}") @@ -19,7 +24,7 @@ def read_all_streams(uid, beamline_acronym, api_key=None): for stream in run: logger.info(f"{stream}:") stream_start_time = ttime.monotonic() - stream_data = run[stream].read() + stream_data = read_stream(run, stream) 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:_}") @@ -29,4 +34,4 @@ def read_all_streams(uid, beamline_acronym, api_key=None): @flow def data_validation(uid, api_key=None, dry_run=False): - read_all_streams(uid, beamline_acronym="tst", api_key=api_key) + read_all_streams(uid, api_key=api_key) From 08b04b23585156715ebf688cd904361c691d12be Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 2 Apr 2026 11:52:55 -0400 Subject: [PATCH 24/28] use my fork/ghcr, publish/deploy branch for testing --- .github/workflows/publish-ghcr.yml | 2 +- prefect.yaml | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/publish-ghcr.yml b/.github/workflows/publish-ghcr.yml index 45e4f05..110ad75 100644 --- a/.github/workflows/publish-ghcr.yml +++ b/.github/workflows/publish-ghcr.yml @@ -4,7 +4,7 @@ name: Create and publish a Docker image # Configures this workflow to run every time a change is pushed to the branch called `release`. on: push: - branches: ["main"] + branches: ["update-dotenv"] # Defines two custom environment variables for the workflow. These are used for the Container registry domain, and a name for the Docker image that this workflow builds. env: diff --git a/prefect.yaml b/prefect.yaml index ca5b5c6..9336ff9 100644 --- a/prefect.yaml +++ b/prefect.yaml @@ -9,8 +9,8 @@ pull: - prefect.deployments.steps.set_working_directory: directory: /repo - prefect.deployments.steps.git_clone: - repository: https://github.com/nsls2/tst-workflows.git - branch: main + repository: https://github.com/junaishima/tst-workflows.git + branch: update-dotenv deployments: - name: tst-end-of-run-workflow-docker @@ -24,7 +24,7 @@ deployments: schedule: {} work_pool: job_variables: - image: ghcr.io/nsls2/tst-workflows:main + image: ghcr.io/junaishima/tst-workflows:update-dotenv image_pull_policy: Always network_mode: slirp4netns userns: "keep-id:uid=402974,gid=402974" # workflow-tst:workflow-tst From a78103165f8cf569d31b13688d2ee18407d4a9ff Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 2 Apr 2026 12:08:46 -0400 Subject: [PATCH 25/28] only get dotenv API key when getting the run from Tiled --- data_validation.py | 3 +++ end_of_run_workflow.py | 2 -- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/data_validation.py b/data_validation.py index 6ea39ea..4328570 100644 --- a/data_validation.py +++ b/data_validation.py @@ -1,10 +1,13 @@ from prefect import task, flow, get_run_logger import time as ttime from tiled.client import from_uri +from end_of_run_workflow import get_api_key_from_env @task 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) run = cl["tst/raw"][uid] return run diff --git a/end_of_run_workflow.py b/end_of_run_workflow.py index 880166e..b7c2fc6 100644 --- a/end_of_run_workflow.py +++ b/end_of_run_workflow.py @@ -24,8 +24,6 @@ def log_completion(dry_run=False): def end_of_run_workflow(stop_doc, api_key=None, dry_run=False): uid = stop_doc["run_start"] # hello_world() - if not api_key: - api_key = get_api_key_from_env(api_key=None) data_validation(uid, return_state=True, api_key=api_key) get_other_docs(uid, api_key=api_key) # long_flow(iterations=100, sleep_length=10, dry_run=dry_run) From b5799967a8ee60185da7dde79eb8201f70a5d6b9 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 2 Apr 2026 12:10:06 -0400 Subject: [PATCH 26/28] test fix-pass-api-key fix --- .github/workflows/publish-ghcr.yml | 2 +- prefect.yaml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/publish-ghcr.yml b/.github/workflows/publish-ghcr.yml index 110ad75..1491564 100644 --- a/.github/workflows/publish-ghcr.yml +++ b/.github/workflows/publish-ghcr.yml @@ -4,7 +4,7 @@ name: Create and publish a Docker image # Configures this workflow to run every time a change is pushed to the branch called `release`. on: push: - branches: ["update-dotenv"] + branches: ["fix-pass-api-key"] # Defines two custom environment variables for the workflow. These are used for the Container registry domain, and a name for the Docker image that this workflow builds. env: diff --git a/prefect.yaml b/prefect.yaml index 9336ff9..1b1d528 100644 --- a/prefect.yaml +++ b/prefect.yaml @@ -10,7 +10,7 @@ pull: directory: /repo - prefect.deployments.steps.git_clone: repository: https://github.com/junaishima/tst-workflows.git - branch: update-dotenv + branch: fix-pass-api-key deployments: - name: tst-end-of-run-workflow-docker @@ -24,7 +24,7 @@ deployments: schedule: {} work_pool: job_variables: - image: ghcr.io/junaishima/tst-workflows:update-dotenv + image: ghcr.io/junaishima/tst-workflows:fix-pass-api-key image_pull_policy: Always network_mode: slirp4netns userns: "keep-id:uid=402974,gid=402974" # workflow-tst:workflow-tst From e7fe1c3e883848e5a119accb525d5d117356ba88 Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Thu, 2 Apr 2026 12:22:31 -0400 Subject: [PATCH 27/28] move function to get API key to data_validation * no need to import it from end_of_run_workflow() --- data_validation.py | 10 +++++++++- end_of_run_workflow.py | 8 -------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/data_validation.py b/data_validation.py index 4328570..e7403d4 100644 --- a/data_validation.py +++ b/data_validation.py @@ -1,7 +1,15 @@ +import os from prefect import task, flow, get_run_logger import time as ttime from tiled.client import from_uri -from end_of_run_workflow import get_api_key_from_env +from dotenv import load_dotenv + + +def get_api_key_from_env(api_key=None): + with open("/srv/container.secret", "r") as secrets: + load_dotenv(stream=secrets) + api_key = os.environ["TILED_API_KEY"] + return api_key @task diff --git a/end_of_run_workflow.py b/end_of_run_workflow.py index b7c2fc6..a622463 100644 --- a/end_of_run_workflow.py +++ b/end_of_run_workflow.py @@ -3,17 +3,9 @@ from prefect import task, flow, get_run_logger from data_validation import data_validation from test_extra_client import get_other_docs -from dotenv import load_dotenv # from long_flow import long_flow -def get_api_key_from_env(api_key=None): - with open("/srv/container.secret", "r") as secrets: - load_dotenv(stream=secrets) - api_key = os.environ["TILED_API_KEY"] - return api_key - - @task def log_completion(dry_run=False): logger = get_run_logger() From 2ef3774e9f9644826bd156ebca9b7d8bffee4e0f Mon Sep 17 00:00:00 2001 From: Jun Aishima Date: Wed, 8 Apr 2026 13:37:45 -0400 Subject: [PATCH 28/28] pre-commit lint removal --- end_of_run_workflow.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/end_of_run_workflow.py b/end_of_run_workflow.py index a622463..fcb7ea2 100644 --- a/end_of_run_workflow.py +++ b/end_of_run_workflow.py @@ -1,5 +1,3 @@ -import os - from prefect import task, flow, get_run_logger from data_validation import data_validation from test_extra_client import get_other_docs