feat: subdirectory for numtracker

docs/how-to/1.4 integrate-with-numtracker.md

@@ -24,3 +24,24 @@ Numtracker will work for any ophyd-async [`StandardDetector`](https://blueskypro
 Blueapi uses an internal [`PathProvider`](https://blueskyproject.io/ophyd-async/main/_api/ophyd_async/ophyd_async.core.html#ophyd_async.core.PathProvider) when integrated with numtracker, it does not support custom providers in from dodal modules, if your beamline definition in dodal is prefixed with a call to `set_path_provider`, it should be removed before attempting to use numtracker:
 
 Numtracker should work with out-of-box plans that take data via runs. Opening a new run will make the `RunEngine` call numtracker and request a fresh data area. Some plans from before numtracker have a decorator that coordinates where detectors write their data: [`@attach_data_session_metadata_decorator`](https://github.com/DiamondLightSource/dodal/blob/10a9a124931901d7666659c6dbe77215d22a8bfd/src/dodal/plan_stubs/data_session.py#L60). This decorator becomes a no-op with numtracker and can be safely removed from plans.
+
+## Writing to a Subdirectory
+
+A task can ask numtracker to place its scans in a subdirectory of the instrument session's data directory by setting the optional `subdirectory` field when it is submitted to `POST /tasks`. It applies to every run in that task; tasks without it are not affected.
+
+```json
+{
+  "name": "count",
+  "params": {"detectors": ["det"]},
+  "instrument_session": "cm12345-1",
+  "subdirectory": "sample1/calibration"
+}
+```
+
+```{note}
+The subdirectory only has an effect if your instrument's numtracker configuration uses the `{subdirectory}` placeholder, e.g. a scan template of `{subdirectory}/{instrument}-{scan_number}`. If it does not, the subdirectory is ignored. See the [numtracker documentation](https://github.com/DiamondLightSource/numtracker) for configuring templates.
+
+It changes where the scan file goes, not where detectors write: detector files are still written to the root of the data session directory using blueapi's `detector_file_template`.
+
+Scan numbers are allocated per instrument, not per subdirectory, so they keep incrementing across subdirectories.
+```

docs/reference/openapi.yaml

@@ -301,6 +301,11 @@ components:
           description: Values for parameters to plan, if any
           title: Params
           type: object
+        subdirectory:
+          description: Subdirectory of the instrument session's data directory to
+            write this task's scans to, if numtracker is configured to use one
+          title: Subdirectory
+          type: string
       required:
       - name
       - instrument_session
@@ -459,7 +464,7 @@ info:
     name: Apache 2.0
     url: https://www.apache.org/licenses/LICENSE-2.0.html
   title: BlueAPI Control
-  version: 1.6.0
+  version: 1.7.0
 openapi: 3.1.0
 paths:
   /api/v1/devices:

src/blueapi/cli/cli.py

@@ -405,6 +405,16 @@ def on_event(
         the session must be valid and active and you must be a member of it."""),
     required=True,
 )
+@click.option(
+    "-s",
+    "--subdirectory",
+    type=str,
+    help=textwrap.dedent("""
+        Subdirectory of the instrument session's data directory to write
+        data to, e.g. a per-sample directory. Passed to numtracker along with
+        the instrument session."""),
+    required=False,
+)
 @click.pass_obj
 @check_connection
 def run_plan(
@@ -414,6 +424,7 @@ def run_plan(
     foreground: bool,
     ws: bool | None,
     instrument_session: str,
+    subdirectory: str | None,
     args: list[Any],
     parameters: TaskParameters,
 ) -> None:
@@ -426,7 +437,11 @@ def run_plan(
 
     client = cast(BlueapiClient, obj["client"])
     task = TaskRequest(
-        name=name, args=args, params=parameters, instrument_session=instrument_session
+        name=name,
+        args=args,
+        params=parameters,
+        instrument_session=instrument_session,
+        subdirectory=subdirectory,
     )
 
     try:

src/blueapi/client/client.py

@@ -67,6 +67,7 @@ class BlueapiClient(ClientProtocol):
     _rest: BlueapiRestClient
     _events: EventBusClient | None
     _instrument_session: str | None = None
+    _subdirectory: str | None = None
     _callbacks: dict[int, OnAnyEvent]
     _callback_id: itertools.count
 
@@ -137,11 +138,31 @@ def instrument_session(self) -> str:
     def instrument_session(self, session: str):
         log.debug("Setting instrument_session to %s", session)
         self._instrument_session = session
+        # A subdirectory belongs to a session, don't carry it over to a new one
+        if self._subdirectory is not None:
+            log.info(
+                "Clearing subdirectory %s as instrument_session changed",
+                self._subdirectory,
+            )
+            self._subdirectory = None
 
     def with_instrument_session(self, session: str) -> Self:
         self.instrument_session = session
         return self
 
+    @property
+    def subdirectory(self) -> str | None:
+        return self._subdirectory
+
+    @subdirectory.setter
+    def subdirectory(self, subdirectory: str | None):
+        log.debug("Setting subdirectory to %s", subdirectory)
+        self._subdirectory = subdirectory
+
+    def with_subdirectory(self, subdirectory: str | None) -> Self:
+        self.subdirectory = subdirectory
+        return self
+
     @start_as_current_span(TRACER)
     @deprecated("plans property")
     def get_plans(self) -> PlanResponse:

src/blueapi/client/plan_cache.py

@@ -80,6 +80,7 @@ def __call__(self, *args, **kwargs) -> Any:
             args=list(args),
             params=kwargs,
             instrument_session=self._client.instrument_session,
+            subdirectory=self._client.subdirectory,
         )
         try:
             match self._client.run_task(req):

src/blueapi/client/protocols.py

@@ -20,6 +20,9 @@ def run_task(
     @property
     def instrument_session(self) -> str: ...
 
+    @property
+    def subdirectory(self) -> str | None: ...
+
 
 class RestClientProtocol(Protocol):
     def get_devices(self) -> DeviceResponse: ...

src/blueapi/config.py

@@ -335,7 +335,7 @@ class ApplicationConfig(BlueapiBaseModel):
     """
 
     #: API version to publish in OpenAPI schema
-    REST_API_VERSION: ClassVar[str] = "1.6.0"
+    REST_API_VERSION: ClassVar[str] = "1.7.0"
 
     LICENSE_INFO: ClassVar[dict[str, str]] = {
         "name": "Apache 2.0",

src/blueapi/service/interface.py

@@ -64,7 +64,11 @@ def context() -> BlueskyContext:
     if (nt_conf := config().numtracker) is not None and (nt := numtracker()):
 
         async def _update_scan_num(md: dict[str, Any]) -> int:
-            scan = await nt.create_scan(md["instrument_session"], md["instrument"])
+            scan = await nt.create_scan(
+                md["instrument_session"],
+                md["instrument"],
+                md.get("data_session_subdirectory"),
+            )
             md["data_session_directory"] = str(scan.scan.directory.path)
             md["detector_file_template"] = nt_conf.detector_file_template
             md["scan_file"] = scan.scan.scan_file
@@ -232,6 +236,8 @@ def submit_task(
         metadata = {}
 
     metadata["instrument_session"] = task_request.instrument_session
+    if task_request.subdirectory is not None:
+        metadata["data_session_subdirectory"] = task_request.subdirectory
     if context().tiled_conf is not None:
         md = config().env.metadata
         # We raise an InvalidConfigError on setting tiled_conf if this isn't set

src/blueapi/service/model.py

@@ -75,6 +75,12 @@ class TaskRequest(BlueapiBaseModel):
     instrument_session: str = Field(
         description="Instrument session associated with this task",
     )
+    subdirectory: str | SkipJsonSchema[None] = Field(
+        description="Subdirectory of the instrument session's data directory to "
+        "write this task's scans to, if numtracker is configured to use one",
+        default=None,
+        exclude_if=lambda subdirectory: subdirectory is None,
+    )
 
 
 class DeviceRequest(BlueapiBaseModel):

src/blueapi/service/numtracker.py

@@ -1,7 +1,6 @@
 import logging
 from collections.abc import Mapping
 from pathlib import Path
-from textwrap import dedent
 
 import httpx
 from pydantic import Field, HttpUrl
@@ -11,6 +10,31 @@
 LOGGER = logging.getLogger(__name__)
 
 
+CREATE_SCAN_MUTATION = """
+mutation(
+    $instrument: String!,
+    $instrumentSession: String!,
+    $sub: Subdirectory
+){
+    scan(
+        instrument: $instrument,
+        instrumentSession: $instrumentSession,
+        sub: $sub
+        ) {
+        directory{
+            instrumentSession
+            instrument
+            path
+        }
+        scanFile
+        scanNumber
+    }
+}
+"""
+# Collapse whitespace so the query is sent compactly; safe as it has no string literals
+CREATE_SCAN_MUTATION = " ".join(CREATE_SCAN_MUTATION.split())
+
+
 class DirectoryPath(BlueapiBaseModel):
     """
     Directory location and associated metadata
@@ -63,7 +87,7 @@ def set_headers(self, headers: Mapping[str, str]) -> None:
         self._headers = headers
 
     async def create_scan(
-        self, instrument_session: str, instrument: str
+        self, instrument_session: str, instrument: str, sub: str | None = None
     ) -> NumtrackerScanMutationResponse:
         """
         Create a new scan with numtracker.
@@ -73,25 +97,18 @@ async def create_scan(
             e.g. cm12345-1
 
             instrument: The instrument to write data on e.g. i22
+
+            sub: Optional subdirectory of the session directory to write the
+            scan to e.g. sample1/calibration
         """
 
         query = {
-            "query": dedent(f"""
-            mutation{{
-                scan(
-                    instrument: "{instrument}",
-                    instrumentSession: "{instrument_session}"
-                    ) {{
-                    directory{{
-                        instrumentSession
-                        instrument
-                        path
-                    }}
-                    scanFile
-                    scanNumber
-                }}
-            }}
-            """)
+            "query": CREATE_SCAN_MUTATION,
+            "variables": {
+                "instrument": instrument,
+                "instrumentSession": instrument_session,
+                "sub": sub,
+            },
         }
 
         async with httpx.AsyncClient() as client:

src/blueapi/utils/path_provider.py

@@ -38,22 +38,32 @@ def __call__(self, device_name: str | None = None) -> PathInfo:
         example: "detector_file_template": "custom-{device_name}-{scan_id}".
 
         If you do not provide a data_session_directory it will default to "/tmp".
+
+        If the start document has a scan_file in a subdirectory of the data session
+        directory (e.g. "sample1/ixx-1234"), detector files are written alongside it
+        in that subdirectory, which is created if it does not exist.
         """
         if not self._docs:
             raise BlueskyRunStructureError(
                 "Start document not found. This call must be made inside a run."
             )
         else:
-            template = self._docs[-1].get("detector_file_template")
+            doc = self._docs[-1]
+            template = doc.get("detector_file_template")
             if not template:
                 raise ValueError("detector_file_template must be set in metadata")
-            sub_path = template.format_map(
-                self._docs[-1] | {"device_name": device_name}
-            )
-            data_session_directory = Path(
-                self._docs[-1].get("data_session_directory", "/tmp")
+            sub_path = template.format_map(doc | {"device_name": device_name})
+            data_session_directory = Path(doc.get("data_session_directory", "/tmp"))
+            # numtracker's scan_file is relative to data_session_directory, write
+            # detector files in the same directory as it ("." if there is no
+            # scan_file or it has no subdirectory)
+            scan_file_directory = Path(doc.get("scan_file", "")).parent
+            return PathInfo(
+                directory_path=data_session_directory / scan_file_directory,
+                filename=sub_path,
+                # Negative depth creates only the directories scan_file adds
+                create_dir_depth=-len(scan_file_directory.parts),
             )
-            return PathInfo(directory_path=data_session_directory, filename=sub_path)
 
 
 class BlueskyRunStructureError(Exception):

src/blueapi/worker/task.py

@@ -117,6 +117,8 @@ def do_task(self, ctx: BlueskyContext) -> None:
         )
         plan = ctx.plans[self.name].plan_function
         prepared_args, prepared_kwargs = self.prepare_params(ctx)
+        # RunEngine.md persists between tasks, subdirectory must not carry over
+        ctx.run_engine.md.pop("data_session_subdirectory", None)
         ctx.run_engine.md.update(self.metadata)
         result = ctx.run_engine(plan(*prepared_args, **prepared_kwargs))
         if isinstance(result, tuple):  # pragma: no cover

tests/conftest.py

@@ -21,6 +21,7 @@
 
 from blueapi.config import ApplicationConfig, OIDCConfig
 from blueapi.service.model import Cache
+from blueapi.service.numtracker import CREATE_SCAN_MUTATION
 
 
 @pytest.fixture(scope="function")
@@ -341,24 +342,10 @@ def mock_jwks_fetch(json_web_keyset: JWK):
 
 
 @pytest.fixture
-def nt_query() -> dict[str, str]:
+def nt_query() -> dict[str, Any]:
     return {
-        "query": dedent("""
-            mutation{
-                scan(
-                    instrument: "p46",
-                    instrumentSession: "ab123"
-                    ) {
-                    directory{
-                        instrumentSession
-                        instrument
-                        path
-                    }
-                    scanFile
-                    scanNumber
-                }
-            }
-            """)
+        "query": dedent(CREATE_SCAN_MUTATION),
+        "variables": {"instrument": "p46", "instrumentSession": "ab123", "sub": None},
     }
 
 

tests/system_tests/services/startup.py

@@ -277,15 +277,20 @@ def configure_numtracker():
     response = requests.post(
         NT_URL,
         json={
-            "query": """mutation {
-              configure(instrument: "adsim",
-                        config: {directory: "/tmp/",
-                        scan: "{instrument}-{scan_number}",
-                        detector: "{instrument}-{scan_number}-{detector}",
-                        scanNumber: 43}) {
+            "query": """mutation($instrument: String!, $config: ConfigurationUpdates!) {
+              configure(instrument: $instrument, config: $config) {
                 scanTemplate
               }
-            }"""
+            }""",
+            "variables": {
+                "instrument": "adsim",
+                "config": {
+                    "directory": "/tmp/",
+                    "scan": "{instrument}-{scan_number}",
+                    "detector": "{instrument}-{scan_number}-{detector}",
+                    "scanNumber": 43,
+                },
+            },
         },
         headers={"authorization": "Bearer " + access_token},
     )