feat: update tiled node insertion structure (#1669)

<sub>Stack created with <a
href="https://github.com/github/gh-stack">GitHub Stacks CLI</a> • <a
href="https://gh.io/stacks-feedback">Give Feedback 💬</a></sub>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Author: ZohebShaikh

pyproject.toml

@@ -76,7 +76,6 @@ dev = [
     "mock",
     "jwcrypto",
     "deepdiff",
-    "tiled[minimal-server]>=0.2.4", # For system-test of dls.py
     "respx",
     "pytest-xdist[psutil]",
     "httpx2>=2.5.0",

src/blueapi/service/authorization.py

@@ -69,7 +69,7 @@ async def require_submit_task(self, instrument_session: str, token: str):
             self._config.submit_task_check,
             {
                 "token": token,
-                "proposal": int(match["proposal"]),
+                "proposal": int(match["proposal_number"]),
                 "visit": int(match["visit"]),
             },
         ):

src/blueapi/service/interface.py

@@ -1,3 +1,4 @@
+import json
 import logging
 from collections.abc import Mapping
 from dataclasses import dataclass
@@ -8,7 +9,9 @@
 from bluesky.callbacks.tiled_writer import TiledWriter
 from bluesky_stomp.messaging import StompClient
 from bluesky_stomp.models import Broker, DestinationBase, MessageTopic
+from fastapi import status
 from tiled.client import from_uri
+from tiled.client.utils import ClientError
 
 from blueapi.cli.scratch import get_python_environment
 from blueapi.config import ApplicationConfig, OIDCConfig, ServiceAccount, StompConfig
@@ -103,13 +106,60 @@ def setup(config: ApplicationConfig) -> None:
     stomp_client()
 
 
+@cache
+def service_account_tiled_client():
+    """One client for every task: its auth refreshes its own token, so unlike a
+    client carrying a caller's pass-through headers it can be reused, which saves
+    the server-info and token requests from_uri makes on every task."""
+    tiled_config = context().tiled_conf
+    assert tiled_config is not None and isinstance(
+        tiled_config.authentication, ServiceAccount
+    )
+    return from_uri(
+        str(tiled_config.url), auth=TiledAuth(tiled_auth=tiled_config.authentication)
+    )
+
+
+def _tiled_session_container(tiled_client, instrument: str, instrument_session: str):
+    """The beamline/proposal/visit container runs are written under.
+
+    It almost always exists already, so it is looked up in one request; only the
+    first task for a visit walks the levels and creates what is missing.
+    """
+    session_blob = access_blob(instrument_session, instrument)
+    proposal = session_blob["proposal"]
+    try:
+        return tiled_client[f"{instrument}/{proposal}/{instrument_session}"]
+    except KeyError:
+        pass
+    # Each level's access blob is the prefix of the full (beamline, proposal,
+    # visit) one, matching the beamline/proposal/session tiers the tiled access
+    # policy expects a container to be tagged with.
+    level_access_blobs = {
+        instrument: {"beamline": instrument},
+        proposal: {"beamline": instrument, "proposal": proposal},
+        instrument_session: session_blob,
+    }
+    for key, blob in level_access_blobs.items():
+        try:
+            tiled_client = tiled_client.create_container(
+                key=key, access_tags=[json.dumps(blob)]
+            )
+        except ClientError as e:
+            if e.response.status_code != status.HTTP_409_CONFLICT:
+                raise
+            tiled_client = tiled_client[key]  # created already, e.g. by another task
+    return tiled_client
+
+
 def teardown() -> None:
     worker().stop()
     if (stomp_client_ref := stomp_client()) is not None:
         stomp_client_ref.disconnect()
     context.cache_clear()
     worker.cache_clear()
     stomp_client.cache_clear()
+    service_account_tiled_client.cache_clear()
 
 
 def _publish_event_streams(
@@ -166,7 +216,7 @@ def submit_task(
         # We raise an InvalidConfigError on setting tiled_conf if this isn't set
         assert md
         metadata["tiled_access_tags"] = [
-            access_blob(task_request.instrument_session, md.instrument)
+            json.dumps(access_blob(task_request.instrument_session, md.instrument))
         ]
     task = Task(
         name=task_request.name,
@@ -191,21 +241,26 @@ def begin_task(
     if nt := active_context.numtracker:
         nt.set_headers(pass_through_headers or {})
 
+    if task.task_id is None:
+        return task
+
     subscribers = []
     if tiled_config := active_context.tiled_conf:
         # Tiled queries the root node, so must create an authorized client
         if isinstance(tiled_config.authentication, ServiceAccount):
-            tiled_client = from_uri(
-                str(tiled_config.url),
-                auth=TiledAuth(tiled_auth=tiled_config.authentication),
-            )
+            tiled_client = service_account_tiled_client()
         else:
             tiled_client = from_uri(
                 str(tiled_config.url),
                 api_key=tiled_config.authentication,
                 headers=pass_through_headers,
             )
-
+        if (task_ := get_task_by_id(task_id=task.task_id)) is not None:
+            tiled_client = _tiled_session_container(
+                tiled_client,
+                active_context.run_engine.md["instrument"],
+                task_.task.metadata["instrument_session"],
+            )
         tiled_writer_token = active_context.run_engine.subscribe(
             TiledWriter(tiled_client, batch_size=1)
         )
@@ -227,13 +282,12 @@ def remove_callback_when_task_finished(
         )
         subscribers.append((active_worker.worker_events, remove_callback))
 
-    if task.task_id is not None:
-        try:
-            active_worker.begin_task(task.task_id)
-        except:
-            for channel, token in subscribers:
-                channel.unsubscribe(token)
-            raise
+    try:
+        active_worker.begin_task(task.task_id)
+    except Exception:
+        for channel, token in subscribers:
+            channel.unsubscribe(token)
+        raise
     return task
 
 

src/blueapi/utils/__init__.py

@@ -31,7 +31,10 @@
 Args = ParamSpec("Args")
 Return = TypeVar("Return")
 
-INSTRUMENT_SESSION_RE = re.compile(r"^[a-z]{2}(?P<proposal>\d+)-(?P<visit>\d+)$")
+# e.g. "cm12345-1" -> proposal="cm12345", code="cm", proposal_number="12345", visit="1"
+INSTRUMENT_SESSION_RE = re.compile(
+    r"^(?P<proposal>(?P<code>[a-z]{2})(?P<proposal_number>\d+))-(?P<visit>\d+)$"
+)
 
 
 def report_successful_devices(

src/blueapi/utils/serialization.py

@@ -1,4 +1,3 @@
-import json
 from typing import Any
 
 from pydantic import BaseModel
@@ -29,16 +28,16 @@ def serialize(obj: Any) -> Any:
         return obj
 
 
-def access_blob(instrument_session: str, beamline: str) -> str:
-    m = utils.INSTRUMENT_SESSION_RE.match(instrument_session)
-    if m is None:
+def access_blob(instrument_session: str, beamline: str) -> dict[str, Any]:
+    if (match := utils.INSTRUMENT_SESSION_RE.match(instrument_session)) is None:
         raise ValueError(
             "Unable to extract proposal and visit from "
             f"instrument session {instrument_session}"
         )
-    blob = {
-        "proposal": int(m["proposal"]),
-        "visit": int(m["visit"]),
+    return {
+        # The full proposal code (e.g. "cm12345"), not just its number - the
+        # tiled access policy strips the letters itself where it needs them.
+        "proposal": match["proposal"],
+        "visit": int(match["visit"]),
         "beamline": beamline,
     }
-    return json.dumps(blob)

tests/system_tests/compose.yaml

@@ -65,10 +65,9 @@ services:
       - label=disable
 
   tiled:
-    image: ghcr.io/bluesky/tiled:0.2.18
+    image: ghcr.io/diamondlightsource/dls-tiled:0.5.0
     network_mode: host # Port 8407
     environment:
-      - PYTHONPATH=/deploy/
       - VITE_LOGIN_REDIRECT=false
     volumes:
       - ./services/tiled_config:/deploy/config
@@ -80,7 +79,7 @@ services:
       - label=disable
 
   opa:
-    image: openpolicyagent/opa:edge-static-debug
+    image: openpolicyagent/opa:1.21.2-static-debug  # debug: entrypoint.sh needs a shell
     network_mode: host # Port 8181
     volumes:
       - "./services/opa_config:/mnt"

tests/system_tests/services/opa_config/config.yaml

@@ -5,7 +5,7 @@ services:
 bundles:
   diamond-policies:
     service: ghcr
-    resource: ghcr.io/diamondlightsource/authz-policy:0.0.24
+    resource: ghcr.io/diamondlightsource/authz-policy:0.0.25
     polling:
       min_delay_seconds: 30
       max_delay_seconds: 120

tests/system_tests/services/tiled_config/__init__.py

@@ -20,7 +20,7 @@ trees:
       init_if_not_exists: true
 
 access_control:
-  access_policy: "config.dls:DiamondOpenPolicyAgentAuthorizationPolicy"
+  access_policy: "dls_tiled.policy:DiamondOpenPolicyAgentAuthorizationPolicy"
   args:
     authorization_provider: "http://localhost:8181/v1/data/diamond/policy/"
     token_audience: account