diff --git a/Dockerfiles/murfey-instrument-server b/Dockerfiles/murfey-instrument-server index dcd32d45f..c9eb47111 100644 --- a/Dockerfiles/murfey-instrument-server +++ b/Dockerfiles/murfey-instrument-server @@ -2,7 +2,7 @@ # podman build --build-arg groupid= --build-arg userid= --build-arg groupname= --no-cache -f path/to/Dockerfiles/murfey-instrument-server -t murfey-instrument-server: path/to/python-murfey # Set up the base image to build with -FROM docker.io/library/python:3.12.8-slim-bullseye AS base +FROM docker.io/library/python:3.12.10-slim-bookworm AS base # Install Vim in base image RUN apt-get update && \ diff --git a/Dockerfiles/murfey-rsync b/Dockerfiles/murfey-rsync index 77ef74e27..a93058ba3 100644 --- a/Dockerfiles/murfey-rsync +++ b/Dockerfiles/murfey-rsync @@ -1,7 +1,7 @@ # Template build command # podman build --build-arg groupid= --build-arg userid= --build-arg groupname= --no-cache -f path/to/Dockerfiles/murfey-rsync -FROM docker.io/library/alpine:3.20 +FROM docker.io/library/alpine:3.21 # FROM alpine:3.14 ARG groupid diff --git a/Dockerfiles/murfey-server b/Dockerfiles/murfey-server index bf2b30ce8..dba6330c0 100644 --- a/Dockerfiles/murfey-server +++ b/Dockerfiles/murfey-server @@ -2,7 +2,7 @@ # podman build --build-arg groupid= --build-arg userid= --build-arg groupname= --no-cache -f path/to/Dockerfiles/murfey-server -t murfey-server: path/to/python-murfey # Set up the base image to build with -FROM docker.io/library/python:3.12.8-slim-bullseye AS base +FROM docker.io/library/python:3.12.10-slim-bookworm AS base # Install Vim and PostgreSQL dependencies in base image RUN apt-get update && \ diff --git a/pyproject.toml b/pyproject.toml index 40a6508ca..41b1a2143 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,9 +87,10 @@ GitHub = "https://github.com/DiamondLightSource/python-murfey" "murfey.decrypt_password" = "murfey.cli.decrypt_db_password:run" "murfey.generate_key" = "murfey.cli.generate_crypto_key:run" "murfey.generate_password" = "murfey.cli.generate_db_password:run" +"murfey.generate_route_manifest" = "murfey.cli.generate_route_manifest:run" "murfey.instrument_server" = "murfey.instrument_server:run" "murfey.repost_failed_calls" = "murfey.cli.repost_failed_calls:run" -"murfey.server" = "murfey.server:run" +"murfey.server" = "murfey.server.run:run" "murfey.sessions" = "murfey.cli.db_sessions:run" "murfey.simulate" = "murfey.cli.dummy:run" "murfey.spa_inject" = "murfey.cli.inject_spa_processing:run" @@ -117,6 +118,7 @@ zip-safe = false [tool.setuptools.package-data] "murfey.client.tui" = ["*.css"] +"murfey.util" = ["route_manifest.yaml"] [tool.setuptools.packages.find] where = ["src", "tests"] diff --git a/src/murfey/bootstrap/__main__.py b/src/murfey/bootstrap/__main__.py index 979d8eaf8..109bd3a9e 100644 --- a/src/murfey/bootstrap/__main__.py +++ b/src/murfey/bootstrap/__main__.py @@ -7,9 +7,11 @@ import pathlib import subprocess import sys -from urllib.parse import urlparse +from urllib.parse import ParseResult, urlparse from urllib.request import urlopen +from murfey.util.api import url_path_for + """ A script to simplify installing Murfey on a network-isolated machine. This could in theory be invoked by @@ -64,10 +66,13 @@ def _download_to_file(url: str, outfile: str): # Construct a minimal base path string # Extract the host name for pip installation purposes try: - murfey_url = urlparse(args.server) + murfey_url: ParseResult = urlparse(args.server) except Exception: exit(f"{args.server} is not a valid URL") + murfey_proxy_path = murfey_url.path.rstrip("/") murfey_base = f"{murfey_url.scheme}://{murfey_url.netloc}" + if murfey_proxy_path: + murfey_base = f"{murfey_base}{murfey_proxy_path}" murfey_hostname = murfey_url.netloc.split(":")[0] # Check that Python version is supported @@ -82,7 +87,10 @@ def _download_to_file(url: str, outfile: str): # Step 1: Download pip wheel print() print(f"1/4 -- Connecting to murfey server on {murfey_base}...") - _download_to_file(f"{murfey_base}/bootstrap/pip.whl", "pip.whl") + _download_to_file( + f"{murfey_base}{url_path_for('bootstrap.bootstrap', 'get_pip_wheel')}", + "pip.whl", + ) # Step 2: Get pip to install itself print() @@ -96,7 +104,7 @@ def _download_to_file(url: str, outfile: str): "--trusted-host", murfey_hostname, "-i", - f"{murfey_base}/pypi", + f"{murfey_base}{url_path_for('bootstrap.pypi', 'get_pypi_index')}", "pip", ] ) @@ -116,7 +124,7 @@ def _download_to_file(url: str, outfile: str): "--trusted-host", murfey_hostname, "-i", - f"{murfey_base}/pypi", + f"{murfey_base}{url_path_for('bootstrap.pypi', 'get_pypi_index')}", "--upgrade", "pip", ] @@ -135,7 +143,7 @@ def _download_to_file(url: str, outfile: str): "--trusted-host", murfey_hostname, "-i", - f"{murfey_base}/pypi", + f"{murfey_base}{url_path_for('bootstrap.pypi', 'get_pypi_index')}", "murfey[client]", ] ) diff --git a/src/murfey/cli/generate_route_manifest.py b/src/murfey/cli/generate_route_manifest.py new file mode 100644 index 000000000..d103abe7d --- /dev/null +++ b/src/murfey/cli/generate_route_manifest.py @@ -0,0 +1,138 @@ +""" +CLI to generate a manifest of the FastAPI router paths present in both the instrument +server and backend server to enable lookup of the URLs based on function name. +""" + +import importlib +import inspect +import pkgutil +from argparse import ArgumentParser +from pathlib import Path +from types import ModuleType +from typing import Any + +import yaml +from fastapi import APIRouter + +import murfey + + +def find_routers(name: str) -> dict[str, APIRouter]: + + def _extract_routers_from_module(module: ModuleType): + routers = {} + for name, obj in inspect.getmembers(module): + if isinstance(obj, APIRouter): + module_path = module.__name__ + key = f"{module_path}.{name}" + routers[key] = obj + return routers + + routers = {} + + # Import the module or package + try: + root = importlib.import_module(name) + except ImportError: + raise ImportError( + f"Cannot import '{name}'. Please ensure that you've installed all the " + "dependencies for the client, instrument server, and backend server " + "before running this command." + ) + + # If it's a package, walk through submodules and extract routers from each + if hasattr(root, "__path__"): + module_list = pkgutil.walk_packages(root.__path__, prefix=name + ".") + for _, module_name, _ in module_list: + try: + module = importlib.import_module(module_name) + except ImportError: + raise ImportError( + f"Cannot import '{module_name}'. Please ensure that you've " + "installed all the dependencies for the client, instrument " + "server, and backend server before running this command." + ) + + routers.update(_extract_routers_from_module(module)) + + # Extract directly from single module + else: + routers.update(_extract_routers_from_module(root)) + + return routers + + +def get_route_manifest(routers: dict[str, APIRouter]): + + manifest = {} + + for router_name, router in routers.items(): + routes = [] + for route in router.routes: + path_params = [] + for param in route.dependant.path_params: + param_type = param.type_ if param.type_ is not None else Any + param_info = { + "name": param.name if hasattr(param, "name") else "", + "type": ( + param_type.__name__ + if hasattr(param_type, "__name__") + else str(param_type) + ), + } + path_params.append(param_info) + route_info = { + "path": route.path if hasattr(route, "path") else "", + "function": route.name if hasattr(route, "name") else "", + "path_params": path_params, + "methods": list(route.methods) if hasattr(route, "methods") else [], + } + routes.append(route_info) + manifest[router_name] = routes + return manifest + + +def run(): + # Set up additional args + parser = ArgumentParser() + parser.add_argument( + "--debug", + action="store_true", + default=False, + help=("Outputs the modules being inspected when creating the route manifest"), + ) + args = parser.parse_args() + + # Find routers + print("Finding routers...") + routers = { + **find_routers("murfey.instrument_server.api"), + **find_routers("murfey.server.api"), + } + # Generate the manifest + print("Extracting route information") + manifest = get_route_manifest(routers) + + # Verify + if args.debug: + for router_name, routes in manifest.items(): + print(f"Routes found in {router_name!r}") + for route in routes: + for key, value in route.items(): + print(f"\t{key}: {value}") + print() + + # Save the manifest + murfey_dir = Path(murfey.__path__[0]) + manifest_file = murfey_dir / "util" / "route_manifest.yaml" + with open(manifest_file, "w") as file: + yaml.dump(manifest, file, default_flow_style=False, sort_keys=False) + print( + "Route manifest for instrument and backend servers saved to " + f"{str(manifest_file)!r}" + ) + exit() + + +if __name__ == "__main__": + run() diff --git a/src/murfey/cli/spa_ispyb_messages.py b/src/murfey/cli/spa_ispyb_messages.py index 87264366d..f78388008 100644 --- a/src/murfey/cli/spa_ispyb_messages.py +++ b/src/murfey/cli/spa_ispyb_messages.py @@ -18,10 +18,11 @@ from sqlmodel import create_engine, select from murfey.client.contexts.spa import _get_xml_list_index -from murfey.server import _murfey_id, _register +from murfey.server.feedback import _murfey_id, _register from murfey.server.ispyb import ISPyBSession, TransportManager, get_session_id from murfey.server.murfey_db import url from murfey.util import db +from murfey.util.api import url_path_for from murfey.util.config import get_machine_config, get_microscope, get_security_config @@ -69,7 +70,7 @@ def run(): help="Path to directory containing image files", ) parser.add_argument( - "--suffic", + "--suffix", dest="suffix", required=True, type=str, @@ -203,7 +204,9 @@ def run(): ] ) binning_factor = 1 - server_config = requests.get(f"{args.url}/machine").json() + server_config = requests.get( + f"{args.url}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=args.microscope)}" + ).json() if server_config.get("superres"): # If camera is capable of superres and collection is in superres binning_factor = 2 diff --git a/src/murfey/cli/transfer.py b/src/murfey/cli/transfer.py index 8fd49578e..d5bd7a503 100644 --- a/src/murfey/cli/transfer.py +++ b/src/murfey/cli/transfer.py @@ -3,13 +3,14 @@ import argparse import subprocess from pathlib import Path -from urllib.parse import urlparse +from urllib.parse import ParseResult, urlparse import requests from rich.console import Console from rich.prompt import Confirm from murfey.client import read_config +from murfey.util.api import url_path_for from murfey.util.config import MachineConfig @@ -35,11 +36,11 @@ def run(): args = parser.parse_args() console = Console() - murfey_url = urlparse(args.server, allow_fragments=False) + murfey_url: ParseResult = urlparse(args.server, allow_fragments=False) machine_data = MachineConfig( requests.get( - f"{murfey_url.geturl()}/instruments/{instrument_name}/machine" + f"{murfey_url.geturl()}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=instrument_name)}" ).json() ) if Path(args.source or ".").resolve() in machine_data.data_directories: diff --git a/src/murfey/client/__init__.py b/src/murfey/client/__init__.py index aafc38697..bbc900cfe 100644 --- a/src/murfey/client/__init__.py +++ b/src/murfey/client/__init__.py @@ -26,13 +26,26 @@ from murfey.client.instance_environment import MurfeyInstanceEnvironment from murfey.client.tui.app import MurfeyTUI from murfey.client.tui.status_bar import StatusBar -from murfey.util.client import _get_visit_list, authorised_requests, read_config +from murfey.util.api import url_path_for +from murfey.util.client import authorised_requests, read_config +from murfey.util.models import Visit log = logging.getLogger("murfey.client") requests.get, requests.post, requests.put, requests.delete = authorised_requests() +def _get_visit_list(api_base: ParseResult, instrument_name: str): + proxy_path = api_base.path.rstrip("/") + get_visits_url = api_base._replace( + path=f"{proxy_path}{url_path_for('session_control.router', 'get_current_visits', instrument_name=instrument_name)}" + ) + server_reply = requests.get(get_visits_url.geturl()) + if server_reply.status_code != 200: + raise ValueError(f"Server unreachable ({server_reply.status_code})") + return [Visit.parse_obj(v) for v in server_reply.json()] + + def write_config(config: configparser.ConfigParser): mcch = os.environ.get("MURFEY_CLIENT_CONFIG_HOME") murfey_client_config_home = Path(mcch) if mcch else Path.home() @@ -262,7 +275,9 @@ def run(): rich_handler.setLevel(logging.DEBUG if args.debug else logging.INFO) # Set up websocket app and handler - client_id = requests.get(f"{murfey_url.geturl()}/new_client_id/").json() + client_id = requests.get( + f"{murfey_url.geturl()}{url_path_for('session_control.router', 'new_client_id')}" + ).json() ws = murfey.client.websocket.WSApp( server=args.server, id=client_id["new_id"], @@ -279,7 +294,7 @@ def run(): # Load machine data for subsequent sections machine_data = requests.get( - f"{murfey_url.geturl()}/instruments/{instrument_name}/machine" + f"{murfey_url.geturl()}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=instrument_name)}" ).json() gain_ref: Path | None = None diff --git a/src/murfey/client/analyser.py b/src/murfey/client/analyser.py index 7043598f1..e90da8aad 100644 --- a/src/murfey/client/analyser.py +++ b/src/murfey/client/analyser.py @@ -82,6 +82,9 @@ def __init__( else {} ) + def __repr__(self) -> str: + return f"" + def _find_extension(self, file_path: Path): """ Identifies the file extension and stores that information in the class. diff --git a/src/murfey/client/contexts/clem.py b/src/murfey/client/contexts/clem.py index 37d03f789..fd193eac0 100644 --- a/src/murfey/client/contexts/clem.py +++ b/src/murfey/client/contexts/clem.py @@ -4,7 +4,6 @@ """ import logging -from datetime import datetime from pathlib import Path from typing import Dict, Generator, List, Optional from urllib.parse import quote @@ -14,6 +13,7 @@ from murfey.client.context import Context from murfey.client.instance_environment import MurfeyInstanceEnvironment +from murfey.util.api import url_path_for from murfey.util.client import capture_post, get_machine_config_client # Create logger object @@ -31,13 +31,16 @@ def _file_transferred_to( instrument_name=environment.instrument_name, demo=environment.demo, ) - # rsync basepath and modules are set in the microscope's configuration YAML file - return ( - Path(machine_config.get("rsync_basepath", "")) - / str(datetime.now().year) - / source.name - / file_path.relative_to(source) + + # Construct destination path + base_destination = Path(machine_config.get("rsync_basepath", "")) / Path( + environment.default_destinations[source] ) + # Add visit number to the path if it's not present in default destination + if environment.visit not in environment.default_destinations[source]: + base_destination = base_destination / environment.visit + destination = base_destination / file_path.relative_to(source) + return destination def _get_source( @@ -291,7 +294,7 @@ def post_transfer( post_result = self.process_tiff_series(tiff_dataset, environment) if post_result is False: return False - + logger.info(f"Started preprocessing of TIFF series {series_name!r}") else: logger.debug(f"TIFF series {series_name!r} is still being processed") @@ -351,7 +354,7 @@ def register_lif_file( """ try: # Construct URL to post to post the request to - url = f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/clem/lif_files?lif_file={quote(str(lif_file), safe='')}" + url = f"{environment.url.geturl()}{url_path_for('clem.router', 'register_lif_file', session_id=environment.murfey_session)}?lif_file={quote(str(lif_file), safe='')}" # Validate if not url: logger.error( @@ -381,7 +384,7 @@ def process_lif_file( try: # Construct the URL to post the request to - url = f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/clem/preprocessing/process_raw_lifs?lif_file={quote(str(lif_file), safe='')}" + url = f"{environment.url.geturl()}{url_path_for('clem.router', 'process_raw_lifs', session_id=environment.murfey_session)}?lif_file={quote(str(lif_file), safe='')}" # Validate if not url: logger.error( @@ -408,7 +411,7 @@ def register_tiff_file( """ try: - url = f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/clem/tiff_files?tiff_file={quote(str(tiff_file), safe='')}" + url = f"{environment.url.geturl()}{url_path_for('clem.router', 'register_tiff_file', session_id=environment.murfey_session)}?tiff_file={quote(str(tiff_file), safe='')}" if not url: logger.error( "URL could not be constructed from the environment and file path" @@ -437,7 +440,7 @@ def process_tiff_series( try: # Construct URL for Murfey server to communicate with - url = f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/clem/preprocessing/process_raw_tiffs" + url = f"{environment.url.geturl()}{url_path_for('clem.router', 'process_raw_tiffs', session_id=environment.murfey_session)}" if not url: logger.error( "URL could not be constructed from the environment and file path" diff --git a/src/murfey/client/contexts/fib.py b/src/murfey/client/contexts/fib.py index 393b849b2..43f1ce7a2 100644 --- a/src/murfey/client/contexts/fib.py +++ b/src/murfey/client/contexts/fib.py @@ -10,6 +10,7 @@ from murfey.client.context import Context from murfey.client.instance_environment import MurfeyInstanceEnvironment +from murfey.util.api import url_path_for from murfey.util.client import authorised_requests logger = logging.getLogger("murfey.client.contexts.fib") @@ -95,7 +96,7 @@ def post_transfer( ).name # post gif list to gif making API call requests.post( - f"{str(environment.url.geturl())}/visits/{datetime.now().year}/{environment.visit}/{environment.murfey_session}/make_milling_gif", + f"{environment.url.geturl()}{url_path_for('workflow.correlative_router', 'make_gif', year=datetime.now().year, visit_name=environment.visit, session_id=environment.murfey_session)}", json={ "lamella_number": lamella_number, "images": gif_list, diff --git a/src/murfey/client/contexts/spa.py b/src/murfey/client/contexts/spa.py index c37036f9f..89017f60a 100644 --- a/src/murfey/client/contexts/spa.py +++ b/src/murfey/client/contexts/spa.py @@ -15,6 +15,7 @@ MurfeyID, MurfeyInstanceEnvironment, ) +from murfey.util.api import url_path_for from murfey.util.client import ( authorised_requests, capture_get, @@ -262,7 +263,7 @@ def gather_metadata( binning_factor = 1 if environment: server_config_response = capture_get( - f"{str(environment.url.geturl())}/instruments/{environment.instrument_name}/machine" + f"{str(environment.url.geturl())}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=environment.instrument_name)}" ) if server_config_response is None: return None @@ -404,7 +405,7 @@ def _position_analysis( ] = (None, None, None, None, None, None, None) data_collection_group = ( requests.get( - f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/data_collection_groups" + f"{environment.url.geturl()}{url_path_for('session_info.router', 'get_dc_groups', session_id=environment.murfey_session)}" ) .json() .get(str(source), {}) @@ -426,7 +427,7 @@ def _position_analysis( local_atlas_path, grid_square=str(grid_square), )[str(grid_square)] - gs_url = f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/grid_square/{grid_square}" + gs_url = f"{str(environment.url.geturl())}{url_path_for('session_control.spa_router', 'register_grid_square', session_id=environment.murfey_session, gsid=grid_square)}" gs = grid_square_data( grid_square_metadata_file, grid_square, @@ -467,7 +468,7 @@ def _position_analysis( ) foil_hole = foil_hole_from_file(transferred_file) if foil_hole not in self._foil_holes[grid_square]: - fh_url = f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/grid_square/{grid_square}/foil_hole" + fh_url = f"{str(environment.url.geturl())}{url_path_for('session_control.spa_router', 'register_foil_hole', session_id=environment.murfey_session, gs_name=grid_square)}" if environment.murfey_session is not None: fh = foil_hole_data( grid_square_metadata_file, @@ -566,7 +567,7 @@ def post_transfer( ) if not environment.movie_counters.get(str(source)): movie_counts_get = capture_get( - f"{str(environment.url.geturl())}/num_movies", + f"{environment.url.geturl()}{url_path_for('session_info.router', 'count_number_of_movies')}", ) if movie_counts_get is not None: environment.movie_counters[str(source)] = count( @@ -580,7 +581,7 @@ def post_transfer( eer_fractionation_file = None if file_transferred_to.suffix == ".eer": response = capture_post( - f"{str(environment.url.geturl())}/visits/{environment.visit}/{environment.murfey_session}/eer_fractionation_file", + f"{str(environment.url.geturl())}{url_path_for('file_manip.router', 'write_eer_fractionation_file', visit_name=environment.visit, session_id=environment.murfey_session)}", json={ "eer_path": str(file_transferred_to), "fractionation": environment.data_collection_parameters[ @@ -609,7 +610,7 @@ def post_transfer( ) foil_hole = None - preproc_url = f"{str(environment.url.geturl())}/visits/{environment.visit}/{environment.murfey_session}/spa_preprocess" + preproc_url = f"{str(environment.url.geturl())}{url_path_for('workflow.spa_router', 'request_spa_preprocessing', visit_name=environment.visit, session_id=environment.murfey_session)}" preproc_data = { "path": str(file_transferred_to), "description": "", diff --git a/src/murfey/client/contexts/spa_metadata.py b/src/murfey/client/contexts/spa_metadata.py index 79546a7ca..4f23448f4 100644 --- a/src/murfey/client/contexts/spa_metadata.py +++ b/src/murfey/client/contexts/spa_metadata.py @@ -8,6 +8,7 @@ from murfey.client.context import Context from murfey.client.contexts.spa import _file_transferred_to, _get_source from murfey.client.instance_environment import MurfeyInstanceEnvironment, SampleInfo +from murfey.util.api import url_path_for from murfey.util.client import ( authorised_requests, capture_post, @@ -166,7 +167,7 @@ def post_transfer( environment.samples[source] = SampleInfo( atlas=Path(partial_path), sample=sample ) - url = f"{str(environment.url.geturl())}/visits/{environment.visit}/{environment.murfey_session}/register_data_collection_group" + url = f"{str(environment.url.geturl())}{url_path_for('workflow.router', 'register_dc_group', visit_name=environment.visit, session_id=environment.murfey_session)}" dcg_search_dir = "/".join( p for p in transferred_file.parent.parts if p != environment.visit ) @@ -202,7 +203,7 @@ def post_transfer( for gs, pos_data in gs_pix_positions.items(): if pos_data: capture_post( - f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/grid_square/{gs}", + f"{str(environment.url.geturl())}{url_path_for('session_control.spa_router', 'register_grid_square', session_id=environment.murfey_session, gsid=int(gs))}", json={ "tag": dcg_tag, "x_location": pos_data[0], @@ -221,7 +222,7 @@ def post_transfer( and environment ): # Make sure we have a data collection group before trying to register grid square - url = f"{str(environment.url.geturl())}/visits/{environment.visit}/{environment.murfey_session}/register_data_collection_group" + url = f"{str(environment.url.geturl())}{url_path_for('workflow.router', 'register_dc_group', visit_name=environment.visit, session_id=environment.murfey_session)}" dcg_search_dir = "/".join( p for p in transferred_file.parent.parent.parts @@ -247,11 +248,11 @@ def post_transfer( } capture_post(url, json=dcg_data) - gs_name = transferred_file.stem.split("_")[1] + gs_name = int(transferred_file.stem.split("_")[1]) logger.info( - f"Collecting foil hole positions for {str(transferred_file)} and grid square {int(gs_name)}" + f"Collecting foil hole positions for {str(transferred_file)} and grid square {gs_name}" ) - fh_positions = _foil_hole_positions(transferred_file, int(gs_name)) + fh_positions = _foil_hole_positions(transferred_file, gs_name) source = _get_source(transferred_file, environment=environment) if source is None: return None @@ -270,10 +271,10 @@ def post_transfer( visitless_source = str(visitless_source_images_dirs[-1]) if fh_positions: - gs_url = f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/grid_square/{gs_name}" + gs_url = f"{str(environment.url.geturl())}{url_path_for('session_control.spa_router', 'register_grid_square', session_id=environment.murfey_session, gsid=gs_name)}" gs_info = grid_square_data( transferred_file, - int(gs_name), + gs_name, ) image_path = ( _file_transferred_to(environment, source, Path(gs_info.image)) @@ -295,7 +296,7 @@ def post_transfer( for fh, fh_data in fh_positions.items(): capture_post( - f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/grid_square/{gs_name}/foil_hole", + f"{str(environment.url.geturl())}{url_path_for('session_control.spa_router', 'register_foil_hole', session_id=environment.murfey_session, gs_name=gs_name)}", json={ "name": fh, "x_location": fh_data.x_location, diff --git a/src/murfey/client/contexts/tomo.py b/src/murfey/client/contexts/tomo.py index b69d76d8b..ab502febf 100644 --- a/src/murfey/client/contexts/tomo.py +++ b/src/murfey/client/contexts/tomo.py @@ -17,6 +17,7 @@ MurfeyID, MurfeyInstanceEnvironment, ) +from murfey.util.api import url_path_for from murfey.util.client import ( authorised_requests, capture_post, @@ -109,7 +110,7 @@ def register_tomography_data_collections( ) return try: - dcg_url = f"{str(environment.url.geturl())}/visits/{str(environment.visit)}/{environment.murfey_session}/register_data_collection_group" + dcg_url = f"{str(environment.url.geturl())}{url_path_for('workflow.router', 'register_dc_group', visit_name=environment.visit, session_id=environment.murfey_session)}" dcg_data = { "experiment_type": "tomo", "experiment_type_id": 36, @@ -121,7 +122,7 @@ def register_tomography_data_collections( for tilt_series in self._tilt_series.keys(): if tilt_series not in self._tilt_series_with_pjids: - dc_url = f"{str(environment.url.geturl())}/visits/{environment.visit}/{environment.murfey_session}/start_data_collection" + dc_url = f"{str(environment.url.geturl())}{url_path_for('workflow.router', 'start_dc', visit_name=environment.visit, session_id=environment.murfey_session)}" dc_data = { "experiment_type": "tomography", "file_extension": file_extension, @@ -157,7 +158,7 @@ def register_tomography_data_collections( ) capture_post(dc_url, json=dc_data) - proc_url = f"{str(environment.url.geturl())}/visits/{environment.visit}/{environment.murfey_session}/register_processing_job" + proc_url = f"{str(environment.url.geturl())}{url_path_for('workflow.router', 'register_proc', visit_name=environment.visit, session_id=environment.murfey_session)}" for recipe in ("em-tomo-preprocess", "em-tomo-align"): capture_post( proc_url, @@ -262,7 +263,7 @@ def _add_tilt( f"Tilt series {tilt_series} was previously thought complete but now {file_path} has been seen" ) self._completed_tilt_series.remove(tilt_series) - rerun_url = f"{str(environment.url.geturl())}/visits/{environment.visit}/rerun_tilt_series" + rerun_url = f"{str(environment.url.geturl())}{url_path_for('workflow.tomo_router', 'register_tilt_series_for_rerun', visit_name=environment.visit)}" rerun_data = { "session_id": environment.murfey_session, "tag": tilt_series, @@ -276,7 +277,7 @@ def _add_tilt( if not self._tilt_series.get(tilt_series): logger.info(f"New tilt series found: {tilt_series}") self._tilt_series[tilt_series] = [file_path] - ts_url = f"{str(environment.url.geturl())}/visits/{environment.visit}/tilt_series" + ts_url = f"{str(environment.url.geturl())}{url_path_for('workflow.tomo_router', 'register_tilt_series', visit_name=environment.visit)}" ts_data = { "session_id": environment.murfey_session, "tag": tilt_series, @@ -305,7 +306,7 @@ def _add_tilt( self._tilt_series[tilt_series].append(file_path) if environment: - tilt_url = f"{str(environment.url.geturl())}/visits/{environment.visit}/{environment.murfey_session}/tilt" + tilt_url = f"{str(environment.url.geturl())}{url_path_for('workflow.tomo_router', 'register_tilt', visit_name=environment.visit, session_id=environment.murfey_session)}" tilt_data = { "movie_path": str(file_transferred_to), "tilt_series_tag": tilt_series, @@ -316,7 +317,7 @@ def _add_tilt( eer_fractionation_file = None if environment.data_collection_parameters.get("num_eer_frames"): response = requests.post( - f"{str(environment.url.geturl())}/visits/{environment.visit}/{environment.murfey_session}/eer_fractionation_file", + f"{str(environment.url.geturl())}{url_path_for('file_manip.router', 'write_eer_fractionation_file', visit_name=environment.visit, session_id=environment.murfey_session)}", json={ "num_frames": environment.data_collection_parameters[ "num_eer_frames" @@ -331,7 +332,7 @@ def _add_tilt( }, ) eer_fractionation_file = response.json()["eer_fractionation_file"] - preproc_url = f"{str(environment.url.geturl())}/visits/{environment.visit}/{environment.murfey_session}/tomography_preprocess" + preproc_url = f"{str(environment.url.geturl())}{url_path_for('workflow.tomo_router', 'request_tomography_preprocessing', visit_name=environment.visit, session_id=environment.murfey_session)}" preproc_data = { "path": str(file_transferred_to), "description": "", @@ -491,7 +492,7 @@ def post_transfer( # Always update the tilt series length in the database after an mdoc if environment.murfey_session is not None: - length_url = f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/tilt_series_length" + length_url = f"{str(environment.url.geturl())}{url_path_for('workflow.tomo_router', 'register_tilt_series_length', session_id=environment.murfey_session)}" capture_post( length_url, json={ @@ -508,7 +509,7 @@ def post_transfer( f"The following tilt series are considered complete: {completed_tilts} " f"after {transferred_file}" ) - complete_url = f"{str(environment.url.geturl())}/visits/{environment.visit}/{environment.murfey_session}/completed_tilt_series" + complete_url = f"{str(environment.url.geturl())}{url_path_for('workflow.tomo_router', 'register_completed_tilt_series', visit_name=environment.visit, session_id=environment.murfey_session)}" capture_post( complete_url, json={ @@ -592,7 +593,7 @@ def gather_metadata( binning_factor = 1 if environment: server_config = requests.get( - f"{str(environment.url.geturl())}/instruments/{environment.instrument_name}/machine" + f"{str(environment.url.geturl())}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=environment.instrument_name)}" ).json() if ( server_config.get("superres") diff --git a/src/murfey/client/multigrid_control.py b/src/murfey/client/multigrid_control.py index 4abcb5c0b..03ab2cfaa 100644 --- a/src/murfey/client/multigrid_control.py +++ b/src/murfey/client/multigrid_control.py @@ -21,6 +21,7 @@ from murfey.client.tui.screens import determine_default_destination from murfey.client.watchdir import DirWatcher from murfey.util import posix_path +from murfey.util.api import url_path_for from murfey.util.client import capture_post, get_machine_config_client log = logging.getLogger("murfey.client.mutligrid_control") @@ -62,7 +63,7 @@ def __post_init__(self): requests.delete, headers={"Authorization": f"Bearer {self.token}"} ) machine_data = requests.get( - f"{self.murfey_url}/instruments/{self.instrument_name}/machine" + f"{self.murfey_url}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=self.instrument_name)}" ).json() self.rsync_url = machine_data.get("rsync_url", "") self.rsync_module = machine_data.get("rsync_module", "data") @@ -102,9 +103,9 @@ def __post_init__(self): if self.visit_end_time: current_time = datetime.now() - server_timestamp = requests.get(f"{self.murfey_url}/time").json()[ - "timestamp" - ] + server_timestamp = requests.get( + f"{self.murfey_url}{url_path_for('session_control.router', 'get_current_timestamp')}" + ).json()["timestamp"] self.visit_end_time += current_time - datetime.fromtimestamp( server_timestamp ) @@ -124,7 +125,7 @@ async def dormancy_check(self): ): async with aiohttp.ClientSession() as clientsession: async with clientsession.delete( - f"{self._environment.url.geturl()}/sessions/{self.session_id}", + f"{self._environment.url.geturl()}{url_path_for('session_control.router', 'remove_session', session_id=self.session_id)}", json={"access_token": self.token, "token_type": "bearer"}, ) as response: success = response.status == 200 @@ -164,7 +165,7 @@ def _start_rsyncer_multigrid( log.debug(f"Analysis of {source} is {('enabled' if analyse else 'disabled')}") destination_overrides = destination_overrides or {} machine_data = requests.get( - f"{self._environment.url.geturl()}/instruments/{self.instrument_name}/machine" + f"{self._environment.url.geturl()}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=self.instrument_name)}" ).json() if destination_overrides.get(source): destination = ( @@ -208,10 +209,10 @@ def _start_rsyncer_multigrid( def _rsyncer_stopped(self, source: Path, explicit_stop: bool = False): if explicit_stop: - remove_url = f"{self.murfey_url}/sessions/{self.session_id}/rsyncer?source={quote(str(source), safe='')}" + remove_url = f"{self.murfey_url}{url_path_for('session_control.router', 'delete_rsyncer', session_id=self.session_id)}?source={quote(str(source), safe='')}" requests.delete(remove_url) else: - stop_url = f"{self.murfey_url}/sessions/{self.session_id}/rsyncer_stopped" + stop_url = f"{self.murfey_url}{url_path_for('session_control.router', 'register_stopped_rsyncer', session_id=self.session_id)}" capture_post(stop_url, json={"source": str(source)}) def _finalise_rsyncer(self, source: Path): @@ -227,7 +228,7 @@ def _finalise_rsyncer(self, source: Path): def _restart_rsyncer(self, source: Path): self.rsync_processes[source].restart() - restarted_url = f"{self.murfey_url}/sessions/{self.session_id}/rsyncer_started" + restarted_url = f"{self.murfey_url}{url_path_for('session_control.router', 'register_restarted_rsyncer', session_id=self.session_id)}" capture_post(restarted_url, json={"source": str(source)}) def _request_watcher_stop(self, source: Path): @@ -250,9 +251,7 @@ def _start_rsyncer( log.info(f"starting rsyncer: {source}") if transfer: # Always make sure the destination directory exists - make_directory_url = ( - f"{self.murfey_url}/sessions/{self.session_id}/make_rsyncer_destination" - ) + make_directory_url = f"{self.murfey_url}{url_path_for('file_manip.router', 'make_rsyncer_destination', session_id=self.session_id)}" capture_post(make_directory_url, json={"destination": destination}) if self._environment: self._environment.default_destinations[source] = destination @@ -322,12 +321,10 @@ def rsync_result(update: RSyncerUpdate): secondary=True, ) if restarted: - restarted_url = ( - f"{self.murfey_url}/sessions/{self.session_id}/rsyncer_started" - ) + restarted_url = f"{self.murfey_url}{url_path_for('session_control.router', 'register_restarted_rsyncer', session_id=self.session_id)}" capture_post(restarted_url, json={"source": str(source)}) else: - url = f"{str(self._environment.url.geturl())}/sessions/{str(self._environment.murfey_session)}/rsyncer" + url = f"{str(self._environment.url.geturl())}{url_path_for('session_control.router', 'register_rsyncer', session_id=self._environment.murfey_session)}" rsyncer_data = { "source": str(source), "destination": destination, @@ -440,7 +437,7 @@ def _start_dc(self, json, from_form: bool = False): log.info("Registering tomography processing parameters") if self._environment.data_collection_parameters.get("num_eer_frames"): eer_response = requests.post( - f"{str(self._environment.url.geturl())}/visits/{self._environment.visit}/{self._environment.murfey_session}/eer_fractionation_file", + f"{str(self._environment.url.geturl())}{url_path_for('file_manip.router', 'write_eer_fractionation_file', visit_name=self._environment.visit, session_id=self._environment.murfey_session)}", json={ "num_frames": self._environment.data_collection_parameters[ "num_eer_frames" @@ -457,17 +454,17 @@ def _start_dc(self, json, from_form: bool = False): eer_fractionation_file = eer_response.json()["eer_fractionation_file"] json.update({"eer_fractionation_file": eer_fractionation_file}) capture_post( - f"{self._environment.url.geturl()}/sessions/{self._environment.murfey_session}/tomography_processing_parameters", + f"{self._environment.url.geturl()}{url_path_for('workflow.tomo_router', 'register_tomo_proc_params', session_id=self._environment.murfey_session)}", json=json, ) capture_post( - f"{self._environment.url.geturl()}/visits/{self._environment.visit}/{self._environment.murfey_session}/flush_tomography_processing", + f"{self._environment.url.geturl()}{url_path_for('workflow.tomo_router', 'flush_tomography_processing', visit_name=self._environment.visit, session_id=self._environment.murfey_session)}", json={"rsync_source": str(source)}, ) log.info("Tomography processing flushed") elif isinstance(context, SPAModularContext): - url = f"{str(self._environment.url.geturl())}/visits/{str(self._environment.visit)}/{self.session_id}/register_data_collection_group" + url = f"{str(self._environment.url.geturl())}{url_path_for('workflow.router', 'register_dc_group', visit_name=self._environment.visit, session_id=self.session_id)}" dcg_data = { "experiment_type": "single particle", "experiment_type_id": 37, @@ -506,7 +503,7 @@ def _start_dc(self, json, from_form: bool = False): "phase_plate": json.get("phase_plate", False), } capture_post( - f"{str(self._environment.url.geturl())}/visits/{str(self._environment.visit)}/{self.session_id}/start_data_collection", + f"{str(self._environment.url.geturl())}{url_path_for('workflow.router', 'start_dc', visit_name=self._environment.visit, session_id=self.session_id)}", json=data, ) for recipe in ( @@ -517,7 +514,7 @@ def _start_dc(self, json, from_form: bool = False): "em-spa-refine", ): capture_post( - f"{str(self._environment.url.geturl())}/visits/{str(self._environment.visit)}/{self.session_id}/register_processing_job", + f"{str(self._environment.url.geturl())}{url_path_for('workflow.router', 'register_proc', visit_name=self._environment.visit, session_id=self.session_id)}", json={ "tag": str(source), "source": str(source), @@ -526,7 +523,7 @@ def _start_dc(self, json, from_form: bool = False): ) log.info(f"Posting SPA processing parameters: {json}") response = capture_post( - f"{self._environment.url.geturl()}/sessions/{self.session_id}/spa_processing_parameters", + f"{self._environment.url.geturl()}{url_path_for('workflow.spa_router', 'register_spa_proc_params', session_id=self.session_id)}", json={ **{k: None if v == "None" else v for k, v in json.items()}, "tag": str(source), @@ -535,14 +532,14 @@ def _start_dc(self, json, from_form: bool = False): if response and not str(response.status_code).startswith("2"): log.warning(f"{response.reason}") capture_post( - f"{self._environment.url.geturl()}/visits/{self._environment.visit}/{self.session_id}/flush_spa_processing", + f"{self._environment.url.geturl()}{url_path_for('workflow.spa_router', 'flush_spa_processing', visit_name=self._environment.visit, session_id=self.session_id)}", json={"tag": str(source)}, ) def _increment_file_count( self, observed_files: List[Path], source: str, destination: str ): - url = f"{str(self._environment.url.geturl())}/visits/{str(self._environment.visit)}/increment_rsync_file_count" + url = f"{str(self._environment.url.geturl())}{url_path_for('prometheus.router', 'increment_rsync_file_count', visit_name=self._environment.visit)}" num_data_files = len( [ f @@ -566,7 +563,7 @@ def _increment_transferred_files_prometheus( self, update: RSyncerUpdate, source: str, destination: str ): if update.outcome is TransferResult.SUCCESS: - url = f"{str(self._environment.url.geturl())}/visits/{str(self._environment.visit)}/increment_rsync_transferred_files_prometheus" + url = f"{str(self._environment.url.geturl())}{url_path_for('prometheus.router', 'increment_rsync_transferred_files_prometheus', visit_name=self._environment.visit)}" data_files = ( [update] if update.file_path.suffix in self._data_suffixes @@ -595,7 +592,7 @@ def _increment_transferred_files( ] if not checked_updates: return - url = f"{str(self._environment.url.geturl())}/visits/{str(self._environment.visit)}/increment_rsync_transferred_files" + url = f"{str(self._environment.url.geturl())}{url_path_for('prometheus.router', 'increment_rsync_transferred_files', visit_name=self._environment.visit)}" data_files = [ u for u in updates diff --git a/src/murfey/client/rsync.py b/src/murfey/client/rsync.py index eac9ceec3..2ff0133d8 100644 --- a/src/murfey/client/rsync.py +++ b/src/murfey/client/rsync.py @@ -109,7 +109,7 @@ def __init__( self._statusbar = status_bar def __repr__(self) -> str: - return f" None: exisiting_sessions = requests.get( - f"{self._environment.url.geturl()}/sessions" + f"{self._environment.url.geturl()}{url_path_for('session_control.router', 'get_sessions')}" ).json() if self.visits: self.install_screen(VisitSelection(self.visits), "visit-select-screen") @@ -647,7 +648,7 @@ async def on_mount(self) -> None: else: session_name = "Client connection" resp = capture_post( - f"{self._environment.url.geturl()}/instruments/{self._environment.instrument_name}/clients/{self._environment.client_id}/session", + f"{self._environment.url.geturl()}{url_path_for('session_control.router', 'link_client_to_session', instrument_name=self._environment.instrument_name, client_id=self._environment.client_id)}", json={"session_id": None, "session_name": session_name}, ) if resp: @@ -666,7 +667,7 @@ async def reset(self): sources = "\n".join(str(k) for k in self.rsync_processes.keys()) prompt = f"Remove files from the following:\n {sources} \n" rsync_instances = requests.get( - f"{self._environment.url.geturl()}/sessions/{self._environment.murfey_session}/rsyncers" + f"{self._environment.url.geturl()}{url_path_for('session_control.router', 'get_rsyncers_for_session', session_id=self._environment.murfey_session)}" ).json() prompt += f"Copied {sum(r['files_counted'] for r in rsync_instances)} / {sum(r['files_transferred'] for r in rsync_instances)}" self.install_screen( @@ -690,7 +691,7 @@ async def action_quit(self) -> None: async def action_remove_session(self) -> None: requests.delete( - f"{self._environment.url.geturl()}/sessions/{self._environment.murfey_session}" + f"{self._environment.url.geturl()}{url_path_for('session_control.router', 'remove_session', session_id=self._environment.murfey_session)}" ) if self.rsync_processes: for rp in self.rsync_processes.values(): @@ -704,7 +705,7 @@ async def action_remove_session(self) -> None: def clean_up_quit(self) -> None: requests.delete( - f"{self._environment.url.geturl()}/sessions/{self._environment.murfey_session}" + f"{self._environment.url.geturl()}{url_path_for('session_control.router', 'remove_session', session_id=self._environment.murfey_session)}" ) self.exit() @@ -747,10 +748,10 @@ def _remove_data(self, listener: Callable[..., Awaitable[None] | None], **kwargs removal_rp.stop() log.info(f"rsyncer {rp} rerun with removal") requests.post( - f"{self._environment.url.geturl()}/sessions/{self._environment.murfey_session}/successful_processing" + f"{self._environment.url.geturl()}{url_path_for('session_control.router', 'register_processing_success_in_ispyb', session_id=self._environment.murfey_session)}" ) requests.delete( - f"{self._environment.url.geturl()}/sessions/{self._environment.murfey_session}" + f"{self._environment.url.geturl()}{url_path_for('session_control.router', 'remove_session', session_id=self._environment.murfey_session)}" ) self.exit() diff --git a/src/murfey/client/tui/screens.py b/src/murfey/client/tui/screens.py index 971603d67..04b52c5be 100644 --- a/src/murfey/client/tui/screens.py +++ b/src/murfey/client/tui/screens.py @@ -57,6 +57,7 @@ from murfey.client.rsync import RSyncer from murfey.client.tui.forms import FormDependency from murfey.util import posix_path +from murfey.util.api import url_path_for from murfey.util.client import capture_post, get_machine_config_client, read_config from murfey.util.models import ProcessingParametersSPA, ProcessingParametersTomo @@ -84,7 +85,7 @@ def determine_default_destination( use_suggested_path: bool = True, ) -> str: machine_data = requests.get( - f"{environment.url.geturl()}/instruments/{environment.instrument_name}/machine" + f"{environment.url.geturl()}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=environment.instrument_name)}" ).json() _default = "" if environment.processing_only_mode and environment.sources: @@ -109,7 +110,7 @@ def determine_default_destination( _default = environment.destination_registry[source_name] else: suggested_path_response = capture_post( - url=f"{str(environment.url.geturl())}/visits/{visit}/{environment.murfey_session}/suggested_path", + url=f"{str(environment.url.geturl())}{url_path_for('file_manip.router', 'suggest_path', visit_name=visit, session_id=environment.murfey_session)}", json={ "base_path": f"{destination}/{visit}/{mid_path.parent if include_mid_path else ''}/raw", "touch": touch, @@ -265,7 +266,7 @@ def __init__( def compose(self): machine_data = requests.get( - f"{self.app._environment.url.geturl()}/instruments/{instrument_name}/machine" + f"{self.app._environment.url.geturl()}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=instrument_name)}" ).json() self._dir_tree = _DirectoryTree( str(self._selected_dir), @@ -476,16 +477,16 @@ def _write_params( self.app._start_dc(params) if model == ProcessingParametersTomo: requests.post( - f"{self.app._environment.url.geturl()}/sessions/{self.app._environment.murfey_session}/tomography_processing_parameters", + f"{self.app._environment.url.geturl()}/{url_path_for('workflow.tomo_router', 'register_tomo_proc_params', session_id=self.app._environment.murfey_session)}", json=params, ) elif model == ProcessingParametersSPA: requests.post( - f"{self.app._environment.url.geturl()}/sessions/{self.app._environment.murfey_session}/spa_processing_parameters", + f"{self.app._environment.url.geturl()}{url_path_for('workflow.spa_router', 'register_spa_proc_params', session_id=self.app._environment.murfey_session)}", json=params, ) requests.post( - f"{self.app._environment.url.geturl()}/visits/{self.app._environment.visit}/{self.app._environment.murfey_session}/flush_spa_processing" + f"{self.app._environment.url.geturl()}{url_path_for('workflow.spa_router', 'flush_spa_processing', visit_name=self.app._environment.visit, session_id=self.app._environment.murfey_session)}", ) def on_switch_changed(self, event): @@ -645,14 +646,16 @@ def on_button_pressed(self, event: Button.Pressed): self.app.pop_screen() session_name = "Client connection" self.app._environment.murfey_session = requests.post( - f"{self.app._environment.url.geturl()}/instruments/{self.app._environment.instrument_name}/clients/{self.app._environment.client_id}/session", + f"{self.app._environment.url.geturl()}{url_path_for('session_control.router', 'link_client_to_session', instrument_name=self.app._environment.instrument_name, client_id=self.app._environment.client_id)}", json={"session_id": session_id, "session_name": session_name}, ).json() def _remove_session(self, session_id: int, **kwargs): - requests.delete(f"{self.app._environment.url.geturl()}/sessions/{session_id}") + requests.delete( + f"{self.app._environment.url.geturl()}{url_path_for('session_control.router', 'remove_session', session_id=session_id)}" + ) exisiting_sessions = requests.get( - f"{self.app._environment.url.geturl()}/sessions" + f"{self.app._environment.url.geturl()}{url_path_for('session_control.router', 'get_sessions')}" ).json() self.app.uninstall_screen("session-select-screen") if exisiting_sessions: @@ -674,7 +677,7 @@ def _remove_session(self, session_id: int, **kwargs): else: session_name = "Client connection" resp = capture_post( - f"{self.app._environment.url.geturl()}/instruments/{self._environment.instrument_name}/clients/{self.app._environment.client_id}/session", + f"{self.app._environment.url.geturl()}{url_path_for('session_control.router', 'link_client_to_session', instrument_name=self.app._environment.instrument_name, client_id=self.app._environment.client_id)}", json={"session_id": None, "session_name": session_name}, ) if resp: @@ -696,12 +699,12 @@ def on_button_pressed(self, event: Button.Pressed): self.app._visit = text self.app._environment.visit = text response = requests.post( - f"{self.app._environment.url.geturl()}/visits/{text}", + f"{self.app._environment.url.geturl()}{url_path_for('session_control.router', 'register_client_to_visit', visit_name=text)}", json={"id": self.app._environment.client_id}, ) log.info(f"Posted visit registration: {response.status_code}") machine_data = requests.get( - f"{self.app._environment.url.geturl()}/instruments/{instrument_name}/machine" + f"{self.app._environment.url.geturl()}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=instrument_name)}" ).json() if self._switch_status: @@ -735,7 +738,7 @@ def on_button_pressed(self, event: Button.Pressed): if machine_data.get("upstream_data_directories"): upstream_downloads = requests.get( - f"{self.app._environment.url.geturl()}/sessions/{self.app._environment.murfey_session}/upstream_visits" + f"{self.app._environment.url.geturl()}{url_path_for('session_control.correlative_router', 'find_upstream_visits', session_id=self.app._environment.murfey_session)}" ).json() self.app.install_screen( UpstreamDownloads(upstream_downloads), "upstream-downloads" @@ -763,12 +766,12 @@ def on_button_pressed(self, event: Button.Pressed): self.app._visit = text self.app._environment.visit = text response = requests.post( - f"{self.app._environment.url.geturl()}/visits/{text}", + f"{self.app._environment.url.geturl()}{url_path_for('session_control.router', 'register_client_to_visit', visit_name=text)}", json={"id": self.app._environment.client_id}, ) log.info(f"Posted visit registration: {response.status_code}") machine_data = requests.get( - f"{self.app._environment.url.geturl()}/instruments/{instrument_name}/machine" + f"{self.app._environment.url.geturl()}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=instrument_name)}" ).json() self.app.install_screen( @@ -797,7 +800,7 @@ def on_button_pressed(self, event: Button.Pressed): if machine_data.get("upstream_data_directories"): upstream_downloads = requests.get( - f"{self.app._environment.url.geturl()}/sessions/{self.app._environment.murfey_session}/upstream_visits" + f"{self.app._environment.url.geturl()}{url_path_for('session_control.correlative_router', 'find_upstream_visits', session_id=self.app._environment.murfey_session)}" ).json() self.app.install_screen( UpstreamDownloads(upstream_downloads), "upstream-downloads" @@ -819,7 +822,7 @@ def compose(self): def on_button_pressed(self, event: Button.Pressed): machine_data = requests.get( - f"{self.app._environment.url.geturl()}/instruments/{instrument_name}/machine" + f"{self.app._environment.url.geturl()}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=instrument_name)}" ).json() if machine_data.get("upstream_data_download_directory"): # Create the directory locally to save files to @@ -830,7 +833,7 @@ def on_button_pressed(self, event: Button.Pressed): # Get the paths to the TIFF files generated previously under the same session ID upstream_tiff_paths_response = requests.get( - f"{self.app._environment.url.geturl()}/visits/{event.button.label}/{self.app._environment.murfey_session}/upstream_tiff_paths" + f"{self.app._environment.url.geturl()}{url_path_for('session_control.correlative_router', 'gather_upstream_tiffs', visit_name=event.button.label, session_id=self.app._environment.murfey_session)}" ) upstream_tiff_paths = upstream_tiff_paths_response.json() or [] @@ -839,7 +842,7 @@ def on_button_pressed(self, event: Button.Pressed): (download_dir / tp).parent.mkdir(exist_ok=True, parents=True) # Write TIFF to the specified file path stream_response = requests.get( - f"{self.app._environment.url.geturl()}/visits/{event.button.label}/{self.app._environment.murfey_session}/upstream_tiff/{tp}", + f"{self.app._environment.url.geturl()}{url_path_for('session_control.correlative_router', 'get_tiff', visit_name=event.button.label, session_id=self.app._environment.murfey_session, tiff_path=tp)}", stream=True, ) # Write the file chunk-by-chunk to avoid hogging memory @@ -903,7 +906,7 @@ def on_button_pressed(self, event): f"Gain reference file {posix_path(self._dir_tree._gain_reference)!r} was not successfully transferred to {visit_path}/processing" ) process_gain_response = requests.post( - url=f"{str(self.app._environment.url.geturl())}/sessions/{self.app._environment.murfey_session}/process_gain", + url=f"{str(self.app._environment.url.geturl())}{url_path_for('file_manip.router', 'process_gain', session_id=self.app._environment.murfey_session)}", json={ "gain_ref": str(self._dir_tree._gain_reference), "eer": bool( @@ -1231,10 +1234,10 @@ def file_copied(self, *args, **kwargs): self.query_one(ProgressBar).advance(1) if self.query_one(ProgressBar).progress == self.query_one(ProgressBar).total: requests.post( - f"{self.app._environment.url.geturl()}/sessions/{self.app._environment.murfey_session}/successful_processing" + f"{self.app._environment.url.geturl()}{url_path_for('session_control.router', 'register_processing_success_in_ispyb', session_id=self.app._environment.murfey_session)}" ) requests.delete( - f"{self.app._environment.url.geturl()}/instruments/{self._environment.instrument_name}/clients/{self.app._environment.client_id}/session" + f"{self.app._environment.url.geturl()}{url_path_for('session_control.router', 'remove_session', session_id=self.app._environment.murfey_session)}" ) self.app.exit() @@ -1266,5 +1269,5 @@ def compose(self): def on_mount(self, event): requests.post( - f"{self.app._environment.url.geturl()}/visits/{self.app._environment.visit}/monitoring/1" + f"{self.app._environment.url.geturl()}{url_path_for('prometheus.router', 'change_monitoring_status', visit_name=self.app._environment.visit, on=1)}" ) diff --git a/src/murfey/client/update.py b/src/murfey/client/update.py index 49570bd1b..8d722afe2 100644 --- a/src/murfey/client/update.py +++ b/src/murfey/client/update.py @@ -7,6 +7,7 @@ import requests import murfey +from murfey.util.api import url_path_for def check(api_base: ParseResult, install: bool = True, force: bool = False): @@ -17,7 +18,8 @@ def check(api_base: ParseResult, install: bool = True, force: bool = False): """ proxy_path = api_base.path.rstrip("/") version_check_url = api_base._replace( - path=f"{proxy_path}/version/", query=f"client_version={murfey.__version__}" + path=f"{proxy_path}{url_path_for('bootstrap.version', 'get_version')}", + query=f"client_version={murfey.__version__}", ) server_reply = requests.get(version_check_url.geturl()) if server_reply.status_code != 200: @@ -69,7 +71,10 @@ def install_murfey(api_base: ParseResult, version: str) -> bool: "--trusted-host", api_base.hostname, "-i", - api_base._replace(path=f"{proxy_path}/pypi", query="").geturl(), + api_base._replace( + path=f"{proxy_path}{url_path_for('bootstrap.pypi', 'get_pypi_index')}", + query="", + ).geturl(), f"murfey[client]=={version}", ] ) diff --git a/src/murfey/client/websocket.py b/src/murfey/client/websocket.py index d5be8bcb0..6c87e0127 100644 --- a/src/murfey/client/websocket.py +++ b/src/murfey/client/websocket.py @@ -12,6 +12,7 @@ import websocket from murfey.client.instance_environment import MurfeyInstanceEnvironment +from murfey.util.api import url_path_for log = logging.getLogger("murfey.client.websocket") @@ -22,7 +23,7 @@ class WSApp: def __init__( self, *, server: str, id: int | str | None = None, register_client: bool = True ): - self.id = uuid.uuid4() if id is None else id + self.id = str(uuid.uuid4()) if id is None else id log.info(f"Opening websocket connection for Client {self.id}") websocket.enableTrace(True) @@ -42,9 +43,13 @@ def __init__( # Prepend the proxy path to the new URL path # It will evaluate to "" if nothing's there, and starts with "/" if present ws_url = ( - url._replace(path=f"{proxy_path}/ws/test/{self.id}").geturl() + url._replace( + path=f"{proxy_path}{url_path_for('websocket.ws', 'websocket_endpoint', client_id=self.id)}" + ).geturl() if register_client - else url._replace(path=f"{proxy_path}/ws/connect/{self.id}").geturl() + else url._replace( + path=f"{proxy_path}{url_path_for('websocket.ws', 'websocket_connection_endpoint', client_id=self.id)}" + ).geturl() ) self._ws = websocket.WebSocketApp( ws_url, diff --git a/src/murfey/instrument_server/api.py b/src/murfey/instrument_server/api.py index e45d46a7c..93f49099e 100644 --- a/src/murfey/instrument_server/api.py +++ b/src/murfey/instrument_server/api.py @@ -21,6 +21,7 @@ from murfey.client.rsync import RSyncer from murfey.client.watchdir_multigrid import MultigridDirWatcher from murfey.util import posix_path, sanitise, sanitise_nonpath, secure_path +from murfey.util.api import url_path_for from murfey.util.client import read_config from murfey.util.instrument_models import MultigridWatcherSpec from murfey.util.models import File, Token @@ -92,7 +93,7 @@ async def murfey_server_handshake(token: str, session_id: int | None = None) -> # test provided token against Murfey server murfey_url = urlparse(_get_murfey_url(), allow_fragments=False) handshake_response = requests.get( - f"{murfey_url.geturl()}/validate_token", + f"{murfey_url.geturl()}{url_path_for('auth.router', 'simple_token_validation')}", headers={"Authorization": f"Bearer {token}"}, ) res = handshake_response.status_code == 200 and handshake_response.json().get( @@ -141,12 +142,22 @@ def check_token(session_id: MurfeySessionID): def setup_multigrid_watcher( session_id: MurfeySessionID, watcher_spec: MultigridWatcherSpec ): + # Return 'True' if controllers are already set up if controllers.get(session_id) is not None: return {"success": True} + label = watcher_spec.label for sid, controller in controllers.items(): if controller.dormant: del controllers[sid] + + # Load machine config as dictionary + machine_config: dict[str, Any] = requests.get( + f"{_get_murfey_url()}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=sanitise_nonpath(watcher_spec.instrument_name))}", + headers={"Authorization": f"Bearer {tokens[session_id]}"}, + ).json() + + # Set up the multigrid controll controller controllers[session_id] = MultigridController( [], watcher_spec.visit, @@ -156,22 +167,21 @@ def setup_multigrid_watcher( demo=True, do_transfer=True, processing_enabled=not watcher_spec.skip_existing_processing, - _machine_config=watcher_spec.configuration.dict(), + _machine_config=machine_config, token=tokens.get(session_id, "token"), data_collection_parameters=data_collection_parameters.get(label, {}), rsync_restarts=watcher_spec.rsync_restarts, visit_end_time=watcher_spec.visit_end_time, ) + # Make child directories, if specified watcher_spec.source.mkdir(exist_ok=True) - machine_config = requests.get( - f"{_get_murfey_url()}/instruments/{sanitise_nonpath(watcher_spec.instrument_name)}/machine", - headers={"Authorization": f"Bearer {tokens[session_id]}"}, - ).json() for d in machine_config.get("create_directories", []): (watcher_spec.source / d).mkdir(exist_ok=True) + + # Set up multigrid directory watcher watchers[session_id] = MultigridDirWatcher( watcher_spec.source, - watcher_spec.configuration.dict(), + machine_config, skip_existing_processing=watcher_spec.skip_existing_processing, ) watchers[session_id].subscribe( @@ -326,7 +336,7 @@ def get_possible_gain_references( instrument_name: str, session_id: MurfeySessionID ) -> list[File]: machine_config = requests.get( - f"{_get_murfey_url()}/instruments/{sanitise_nonpath(instrument_name)}/machine", + f"{_get_murfey_url()}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=sanitise_nonpath(instrument_name))}", headers={"Authorization": f"Bearer {tokens[session_id]}"}, ).json() candidates = [] @@ -365,7 +375,7 @@ def upload_gain_reference( # Load machine config and other needed properties machine_config: dict[str, Any] = requests.get( - f"{_get_murfey_url()}/instruments/{sanitise_nonpath(instrument_name)}/machine", + f"{_get_murfey_url()}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=sanitise_nonpath(instrument_name))}", headers={"Authorization": f"Bearer {tokens[session_id]}"}, ).json() @@ -422,14 +432,14 @@ def gather_upstream_tiffs( upstream_tiff_info.download_dir.mkdir(exist_ok=True) upstream_tiff_paths = ( requests.get( - f"{murfey_url.geturl()}/visits/{sanitised_visit_name}/upstream_tiff_paths", + f"{murfey_url.geturl()}{url_path_for('session_control.correlative_router', 'gather_upstream_tiffs', session_id=session_id, visit_name=sanitised_visit_name)}", headers={"Authorization": f"Bearer {tokens[session_id]}"}, ).json() or [] ) for tiff_path in upstream_tiff_paths: tiff_data = requests.get( - f"{murfey_url.geturl()}/visits/{sanitised_visit_name}/upstream_tiff/{tiff_path}", + f"{murfey_url.geturl()}{url_path_for('session_control.correlative_router', 'get_tiff', session_id=session_id, visit_name=sanitised_visit_name, tiff_path=tiff_path)}", stream=True, headers={"Authorization": f"Bearer {tokens[session_id]}"}, ) diff --git a/src/murfey/server/__init__.py b/src/murfey/server/__init__.py index 334c5dcbe..8474f991e 100644 --- a/src/murfey/server/__init__.py +++ b/src/murfey/server/__init__.py @@ -1,2738 +1,14 @@ from __future__ import annotations -import argparse -import logging -import math -import os -import subprocess -import time -from datetime import datetime -from functools import partial, singledispatch -from importlib.metadata import EntryPoint # For type hinting only from importlib.resources import files -from pathlib import Path -from threading import Thread -from typing import Any, Dict, List, Literal, NamedTuple, Tuple +from typing import TYPE_CHECKING -import graypy -import mrcfile -import numpy as np -import uvicorn -from backports.entry_points_selectable import entry_points -from fastapi import Request -from fastapi.templating import Jinja2Templates -from ispyb.sqlalchemy._auto_db_schema import ( - Atlas, - AutoProcProgram, - Base, - DataCollection, - DataCollectionGroup, - ProcessingJob, - ProcessingJobParameter, -) -from rich.logging import RichHandler -from sqlalchemy import func -from sqlalchemy.exc import ( - InvalidRequestError, - OperationalError, - PendingRollbackError, - SQLAlchemyError, -) -from sqlalchemy.orm.exc import ObjectDeletedError -from sqlmodel import Session, create_engine, select -from werkzeug.utils import secure_filename -from workflows.transport.pika_transport import PikaTransport +# Classes are only imported for type checking purposes +if TYPE_CHECKING: + from uvicorn import Server -import murfey -import murfey.server.prometheus as prom -import murfey.util.db as db -from murfey.server.ispyb import ISPyBSession, get_session_id -from murfey.server.murfey_db import url # murfey_db -from murfey.util import LogFilter -from murfey.util.config import ( - MachineConfig, - get_hostname, - get_machine_config, - get_microscope, - get_security_config, -) -from murfey.util.processing_params import default_spa_parameters -from murfey.util.tomo import midpoint + from murfey.server.ispyb import TransportManager -try: - from murfey.server.ispyb import TransportManager # Session -except AttributeError: - pass - - -logger = logging.getLogger("murfey.server") - -template_files = files("murfey") / "templates" -templates = Jinja2Templates(directory=template_files) - -_running_server: uvicorn.Server | None = None +_running_server: Server | None = None _transport_object: TransportManager | None = None - -try: - _url = url(get_security_config()) - engine = create_engine(_url) - murfey_db = Session(engine, expire_on_commit=False) -except Exception: - murfey_db = None - - -class ExtendedRecord(NamedTuple): - record: Base # type: ignore - record_params: List[Base] # type: ignore - - -class JobIDs(NamedTuple): - dcgid: int - dcid: int - pid: int - appid: int - - -def sanitise(in_string: str) -> str: - return in_string.replace("\r\n", "").replace("\n", "") - - -def sanitise_path(in_path: Path) -> Path: - return Path("/".join(secure_filename(p) for p in in_path.parts)) - - -def get_angle(tilt_file_name: str) -> float: - for p in Path(tilt_file_name).name.split("_"): - if "." in p: - return float(p) - raise ValueError(f"Tilt angle not found for file {tilt_file_name}") - - -def check_tilt_series_mc(tilt_series_id: int) -> bool: - results = murfey_db.exec( - select(db.Tilt, db.TiltSeries) - .where(db.Tilt.tilt_series_id == db.TiltSeries.id) - .where(db.TiltSeries.id == tilt_series_id) - ).all() - return ( - all(r[0].motion_corrected for r in results) - and len(results) >= results[0][1].tilt_series_length - and results[0][1].tilt_series_length > 0 - ) - - -def get_all_tilts(tilt_series_id: int) -> List[str]: - complete_results = murfey_db.exec( - select(db.Tilt, db.TiltSeries, db.Session) - .where(db.Tilt.tilt_series_id == db.TiltSeries.id) - .where(db.TiltSeries.id == tilt_series_id) - .where(db.TiltSeries.session_id == db.Session.id) - ).all() - if not complete_results: - return [] - instrument_name = complete_results[0][2].instrument_name - results = [r[0] for r in complete_results] - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - - def _mc_path(mov_path: Path) -> str: - for p in mov_path.parts: - if "-" in p and p.startswith(("bi", "nr", "nt", "cm", "sw")): - visit_name = p - break - else: - raise ValueError(f"No visit found in {mov_path}") - visit_idx = Path(mov_path).parts.index(visit_name) - core = Path(*Path(mov_path).parts[: visit_idx + 1]) - ppath = Path(mov_path) - sub_dataset = "/".join(ppath.relative_to(core).parts[:-1]) - extra_path = machine_config.processed_extra_directory - mrc_out = ( - core - / machine_config.processed_directory_name - / sub_dataset - / extra_path - / "MotionCorr" - / "job002" - / "Movies" - / str(ppath.stem + "_motion_corrected.mrc") - ) - return str(mrc_out) - - return [_mc_path(Path(r.movie_path)) for r in results] - - -def get_job_ids(tilt_series_id: int, appid: int) -> JobIDs: - results = murfey_db.exec( - select( - db.TiltSeries, - db.AutoProcProgram, - db.ProcessingJob, - db.DataCollection, - db.DataCollectionGroup, - db.Session, - ) - .where(db.TiltSeries.id == tilt_series_id) - .where(db.DataCollection.tag == db.TiltSeries.tag) - .where(db.ProcessingJob.id == db.AutoProcProgram.pj_id) - .where(db.AutoProcProgram.id == appid) - .where(db.ProcessingJob.dc_id == db.DataCollection.id) - .where(db.DataCollectionGroup.id == db.DataCollection.dcg_id) - .where(db.Session.id == db.TiltSeries.session_id) - ).all() - return JobIDs( - dcgid=results[0][4].id, - dcid=results[0][3].id, - pid=results[0][2].id, - appid=results[0][1].id, - ) - - -def get_tomo_proc_params(dcg_id: int, *args) -> db.TomographyProcessingParameters: - results = murfey_db.exec( - select(db.TomographyProcessingParameters).where( - db.TomographyProcessingParameters.dcg_id == dcg_id - ) - ).one() - return results - - -def respond_with_template( - request: Request, filename: str, parameters: dict[str, Any] | None = None -): - template_parameters = { - "hostname": get_hostname(), - "microscope": get_microscope(), - "version": murfey.__version__, - # Extra parameters to reconstruct URLs for forwarded requests - "netloc": request.url.netloc, - "proxy_path": "", - } - if parameters: - template_parameters.update(parameters) - return templates.TemplateResponse( - request=request, name=filename, context=template_parameters - ) - - -def run(): - # Set up argument parser - parser = argparse.ArgumentParser(description="Start the Murfey server") - parser.add_argument( - "--host", - help="Listen for incoming connections on a specific interface (IP address or hostname; default: all)", - default="0.0.0.0", - ) - parser.add_argument( - "--port", - help="Listen for incoming TCP connections on this port (default: 8000)", - type=int, - default=8000, - ) - parser.add_argument( - "--workers", help="Number of workers for Uvicorn server", type=int, default=2 - ) - parser.add_argument( - "--demo", - action="store_true", - ) - parser.add_argument( - "--feedback", - action="store_true", - ) - parser.add_argument( - "--temporary", - action="store_true", - ) - parser.add_argument( - "--root-path", - default="", - type=str, - help="Uvicorn root path for use in conjunction with a proxy", - ) - verbosity = parser.add_mutually_exclusive_group() - verbosity.add_argument( - "-q", - "--quiet", - action="store_true", - default=False, - help="Decrease logging output verbosity", - ) - verbosity.add_argument( - "-v", - "--verbose", - action="count", - help="Increase logging output verbosity", - default=0, - ) - # Parse and separate known and unknown args - args, unknown = parser.parse_known_args() - - # Load the security configuration - security_config = get_security_config() - - # Set up GrayLog handler if provided in the configuration - if security_config.graylog_host: - handler = graypy.GELFUDPHandler( - security_config.graylog_host, security_config.graylog_port, level_names=True - ) - root_logger = logging.getLogger() - root_logger.addHandler(handler) - # Install a log filter to all existing handlers. - LogFilter.install() - - if args.demo: - # Run in demo mode with no connections set up - os.environ["MURFEY_DEMO"] = "1" - else: - # Load RabbitMQ configuration and set up the connection - PikaTransport().load_configuration_file(security_config.rabbitmq_credentials) - _set_up_transport("PikaTransport") - - # Set up logging now that the desired verbosity is known - _set_up_logging(quiet=args.quiet, verbosity=args.verbose) - - if not args.temporary and _transport_object: - _transport_object.feedback_queue = security_config.feedback_queue - rabbit_thread = Thread( - target=feedback_listen, - daemon=True, - ) - logger.info("Starting Murfey RabbitMQ thread") - if args.feedback: - rabbit_thread.start() - - logger.info( - f"Starting Murfey server version {murfey.__version__} for beamline {get_microscope()}, listening on {args.host}:{args.port}" - ) - global _running_server - config = uvicorn.Config( - "murfey.server.main:app", - host=args.host, - port=args.port, - log_config=None, - ws_ping_interval=300, - ws_ping_timeout=300, - workers=args.workers, - root_path=args.root_path, - ) - - _running_server = uvicorn.Server(config=config) - _running_server.run() - logger.info("Server shutting down") - - -def shutdown(): - global _running_server - if _running_server: - _running_server.should_exit = True - _running_server.force_exit = True - - -def _set_up_logging(quiet: bool, verbosity: int): - rich_handler = RichHandler(enable_link_path=False) - if quiet: - rich_handler.setLevel(logging.INFO) - log_levels = { - "murfey": logging.INFO, - "uvicorn": logging.WARNING, - "fastapi": logging.INFO, - "starlette": logging.INFO, - "sqlalchemy": logging.WARNING, - } - elif verbosity <= 0: - rich_handler.setLevel(logging.INFO) - log_levels = { - "murfey": logging.DEBUG, - "uvicorn": logging.INFO, - "uvicorn.access": logging.WARNING, - "fastapi": logging.INFO, - "starlette": logging.INFO, - "sqlalchemy": logging.WARNING, - } - elif verbosity <= 1: - rich_handler.setLevel(logging.DEBUG) - log_levels = { - "": logging.INFO, - "murfey": logging.DEBUG, - "uvicorn": logging.INFO, - "fastapi": logging.INFO, - "starlette": logging.INFO, - "sqlalchemy": logging.WARNING, - } - elif verbosity <= 2: - rich_handler.setLevel(logging.DEBUG) - log_levels = { - "": logging.INFO, - "murfey": logging.DEBUG, - "uvicorn": logging.DEBUG, - "fastapi": logging.DEBUG, - "starlette": logging.DEBUG, - "sqlalchemy": logging.WARNING, - } - else: - rich_handler.setLevel(logging.DEBUG) - log_levels = { - "": logging.DEBUG, - "murfey": logging.DEBUG, - "uvicorn": logging.DEBUG, - "fastapi": logging.DEBUG, - "starlette": logging.DEBUG, - "sqlalchemy": logging.DEBUG, - } - - logging.getLogger().addHandler(rich_handler) - for logger_name, log_level in log_levels.items(): - logging.getLogger(logger_name).setLevel(log_level) - - -def _set_up_transport(transport_type: Literal["PikaTransport"]): - global _transport_object - _transport_object = TransportManager(transport_type) - - -def _murfey_id(app_id: int, _db, number: int = 1, close: bool = True) -> List[int]: - murfey_ledger = [db.MurfeyLedger(app_id=app_id) for _ in range(number)] - for ml in murfey_ledger: - _db.add(ml) - _db.commit() - # There is a race condition between the IDs being read back from the database - # after the insert and the insert being synchronised so allow multiple attempts - attempts = 0 - while attempts < 100: - try: - for m in murfey_ledger: - _db.refresh(m) - res = [m.id for m in murfey_ledger if m.id is not None] - break - except (ObjectDeletedError, InvalidRequestError): - pass - attempts += 1 - time.sleep(0.1) - else: - raise RuntimeError( - "Maximum number of attempts exceeded when producing new Murfey IDs" - ) - if close: - _db.close() - return res - - -def _murfey_class2ds( - murfey_ids: List[int], particles_file: str, app_id: int, _db, close: bool = False -): - pj_id = _pj_id(app_id, _db, recipe="em-spa-class2d") - class2ds = [ - db.Class2D( - class_number=i, - particles_file=particles_file, - pj_id=pj_id, - murfey_id=mid, - ) - for i, mid in enumerate(murfey_ids) - ] - for c in class2ds: - _db.add(c) - _db.commit() - if close: - _db.close() - - -def _murfey_class3ds(murfey_ids: List[int], particles_file: str, app_id: int, _db): - pj_id = _pj_id(app_id, _db, recipe="em-spa-class3d") - class3ds = [ - db.Class3D( - class_number=i, - particles_file=str(Path(particles_file).parent), - pj_id=pj_id, - murfey_id=mid, - ) - for i, mid in enumerate(murfey_ids) - ] - for c in class3ds: - _db.add(c) - _db.commit() - _db.close() - - -def _murfey_refine(murfey_id: int, refine_dir: str, tag: str, app_id: int, _db): - pj_id = _pj_id(app_id, _db, recipe="em-spa-refine") - refine3d = db.Refine3D( - tag=tag, - refine_dir=refine_dir, - pj_id=pj_id, - murfey_id=murfey_id, - ) - _db.add(refine3d) - _db.commit() - _db.close() - - -def _2d_class_murfey_ids(particles_file: str, app_id: int, _db) -> Dict[str, int]: - pj_id = ( - _db.exec(select(db.AutoProcProgram).where(db.AutoProcProgram.id == app_id)) - .one() - .pj_id - ) - classes = _db.exec( - select(db.Class2D).where( - db.Class2D.particles_file == particles_file and db.Class2D.pj_id == pj_id - ) - ).all() - return {str(cl.class_number): cl.murfey_id for cl in classes} - - -def _3d_class_murfey_ids(particles_file: str, app_id: int, _db) -> Dict[str, int]: - pj_id = ( - _db.exec(select(db.AutoProcProgram).where(db.AutoProcProgram.id == app_id)) - .one() - .pj_id - ) - classes = _db.exec( - select(db.Class3D).where( - db.Class3D.particles_file == str(Path(particles_file).parent) - and db.Class3D.pj_id == pj_id - ) - ).all() - return {str(cl.class_number): cl.murfey_id for cl in classes} - - -def _refine_murfey_id(refine_dir: str, tag: str, app_id: int, _db) -> Dict[str, int]: - pj_id = ( - _db.exec(select(db.AutoProcProgram).where(db.AutoProcProgram.id == app_id)) - .one() - .pj_id - ) - refined_class = _db.exec( - select(db.Refine3D) - .where(db.Refine3D.refine_dir == refine_dir) - .where(db.Refine3D.pj_id == pj_id) - .where(db.Refine3D.tag == tag) - ).one() - return refined_class.murfey_id - - -def _app_id(pj_id: int, _db) -> int: - return ( - _db.exec(select(db.AutoProcProgram).where(db.AutoProcProgram.pj_id == pj_id)) - .one() - .id - ) - - -def _pj_id(app_id: int, _db, recipe: str = "") -> int: - if recipe: - dc_id = ( - _db.exec( - select(db.AutoProcProgram, db.ProcessingJob) - .where(db.AutoProcProgram.id == app_id) - .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id) - ) - .one()[1] - .dc_id - ) - pj_id = ( - _db.exec( - select(db.ProcessingJob) - .where(db.ProcessingJob.dc_id == dc_id) - .where(db.ProcessingJob.recipe == recipe) - ) - .one() - .id - ) - else: - pj_id = ( - _db.exec(select(db.AutoProcProgram).where(db.AutoProcProgram.id == app_id)) - .one() - .pj_id - ) - return pj_id - - -def _get_spa_params( - app_id: int, _db -) -> Tuple[db.SPARelionParameters, db.SPAFeedbackParameters]: - pj_id = _pj_id(app_id, _db, recipe="em-spa-preprocess") - relion_params = _db.exec( - select(db.SPARelionParameters).where(db.SPARelionParameters.pj_id == pj_id) - ).one() - feedback_params = _db.exec( - select(db.SPAFeedbackParameters).where(db.SPAFeedbackParameters.pj_id == pj_id) - ).one() - _db.expunge(relion_params) - _db.expunge(feedback_params) - return relion_params, feedback_params - - -def _release_2d_hold(message: dict, _db=murfey_db): - relion_params, feedback_params = _get_spa_params(message["program_id"], _db) - if not feedback_params.star_combination_job: - feedback_params.star_combination_job = feedback_params.next_job + ( - 3 if default_spa_parameters.do_icebreaker_jobs else 2 - ) - pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-class2d") - if feedback_params.rerun_class2d: - first_class2d = _db.exec( - select(db.Class2DParameters).where(db.Class2DParameters.pj_id == pj_id) - ).first() - instrument_name = ( - _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - zocalo_message: dict = { - "parameters": { - "particles_file": first_class2d.particles_file, - "class2d_dir": message["job_dir"], - "batch_is_complete": first_class2d.complete, - "batch_size": first_class2d.batch_size, - "particle_diameter": relion_params.particle_diameter, - "mask_diameter": relion_params.mask_diameter or 0, - "combine_star_job_number": feedback_params.star_combination_job, - "autoselect_min_score": feedback_params.class_selection_score or 0, - "autoproc_program_id": message["program_id"], - "nr_iter": default_spa_parameters.nr_iter_2d, - "nr_classes": default_spa_parameters.nr_classes_2d, - "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, - "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, - "picker_id": feedback_params.picker_ispyb_id, - "class_uuids": _2d_class_murfey_ids( - first_class2d.particles_file, message["program_id"], _db - ), - "class2d_grp_uuid": _db.exec( - select(db.Class2DParameters) - .where( - db.Class2DParameters.particles_file - == first_class2d.particles_file - ) - .where(db.Class2DParameters.pj_id == pj_id) - ) - .one() - .murfey_id, - "session_id": message["session_id"], - "node_creator_queue": machine_config.node_creator_queue, - }, - "recipes": ["em-spa-class2d"], - } - if first_class2d.complete: - feedback_params.next_job += ( - 4 if default_spa_parameters.do_icebreaker_jobs else 3 - ) - feedback_params.rerun_class2d = False - _db.add(feedback_params) - if first_class2d.complete: - _db.delete(first_class2d) - _db.commit() - _db.close() - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - else: - feedback_params.hold_class2d = False - _db.add(feedback_params) - _db.commit() - _db.close() - - -def _release_3d_hold(message: dict, _db=murfey_db): - pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") - pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-class3d") - relion_params = _db.exec( - select(db.SPARelionParameters).where( - db.SPARelionParameters.pj_id == pj_id_params - ) - ).one() - feedback_params = _db.exec( - select(db.SPAFeedbackParameters).where( - db.SPAFeedbackParameters.pj_id == pj_id_params - ) - ).one() - class3d_params = _db.exec( - select(db.Class3DParameters).where(db.Class3DParameters.pj_id == pj_id) - ).one() - if class3d_params.run: - instrument_name = ( - _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - zocalo_message: dict = { - "parameters": { - "particles_file": class3d_params.particles_file, - "class3d_dir": class3d_params.class3d_dir, - "batch_size": class3d_params.batch_size, - "symmetry": relion_params.symmetry, - "particle_diameter": relion_params.particle_diameter, - "mask_diameter": relion_params.mask_diameter or 0, - "do_initial_model": False if feedback_params.initial_model else True, - "initial_model_file": feedback_params.initial_model, - "picker_id": feedback_params.picker_ispyb_id, - "class_uuids": _3d_class_murfey_ids( - class3d_params.particles_file, _app_id(pj_id, _db), _db - ), - "class3d_grp_uuid": _db.exec( - select(db.Class3DParameters) - .where( - db.Class3DParameters.particles_file - == class3d_params.particles_file - ) - .where(db.Class3DParameters.pj_id == pj_id) - ) - .one() - .murfey_id, - "nr_iter": default_spa_parameters.nr_iter_3d, - "initial_model_iterations": default_spa_parameters.nr_iter_ini_model, - "nr_classes": default_spa_parameters.nr_classes_3d, - "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, - "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, - "session_id": message["session_id"], - "autoproc_program_id": _app_id( - _pj_id(message["program_id"], _db, recipe="em-spa-class3d"), _db - ), - "node_creator_queue": machine_config.node_creator_queue, - }, - "recipes": ["em-spa-class3d"], - } - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - class3d_params.run = False - _db.add(class3d_params) - else: - feedback_params.hold_class3d = False - _db.add(feedback_params) - _db.commit() - _db.close() - - -def _release_refine_hold(message: dict, _db=murfey_db): - pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") - pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-refine") - relion_params = _db.exec( - select(db.SPARelionParameters).where( - db.SPARelionParameters.pj_id == pj_id_params - ) - ).one() - feedback_params = _db.exec( - select(db.SPAFeedbackParameters).where( - db.SPAFeedbackParameters.pj_id == pj_id_params - ) - ).one() - refine_params = _db.exec( - select(db.RefineParameters) - .where(db.RefineParameters.pj_id == pj_id) - .where(db.RefineParameters.tag == "first") - ).one() - symmetry_refine_params = _db.exec( - select(db.RefineParameters) - .where(db.RefineParameters.pj_id == pj_id) - .where(db.RefineParameters.tag == "symmetry") - ).one() - if refine_params.run: - instrument_name = ( - _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - zocalo_message: dict = { - "parameters": { - "refine_job_dir": refine_params.refine_dir, - "class3d_dir": refine_params.class3d_dir, - "class_number": refine_params.class_number, - "pixel_size": relion_params.angpix, - "particle_diameter": relion_params.particle_diameter, - "mask_diameter": relion_params.mask_diameter or 0, - "symmetry": relion_params.symmetry, - "node_creator_queue": machine_config.node_creator_queue, - "nr_iter": default_spa_parameters.nr_iter_3d, - "picker_id": feedback_params.picker_ispyb_id, - "refined_class_uuid": _refine_murfey_id( - refine_dir=refine_params.refine_dir, - tag=refine_params.tag, - app_id=_app_id(pj_id, _db), - _db=_db, - ), - "refined_grp_uuid": refine_params.murfey_id, - "symmetry_refined_class_uuid": _refine_murfey_id( - refine_dir=symmetry_refine_params.refine_dir, - tag=symmetry_refine_params.tag, - app_id=_app_id(pj_id, _db), - _db=_db, - ), - "symmetry_refined_grp_uuid": symmetry_refine_params.murfey_id, - "session_id": message["session_id"], - "autoproc_program_id": _app_id( - _pj_id(message["program_id"], _db, recipe="em-spa-refine"), _db - ), - }, - "recipes": ["em-spa-refine"], - } - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - refine_params.run = False - _db.add(refine_params) - else: - feedback_params.hold_refine = False - _db.add(feedback_params) - _db.commit() - _db.close() - - -def _register_incomplete_2d_batch(message: dict, _db=murfey_db, demo: bool = False): - """Received first batch from particle selection service""" - # the general parameters are stored using the preprocessing auto proc program ID - logger.info("Registering incomplete particle batch for 2D classification") - instrument_name = ( - _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") - pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-class2d") - relion_params = _db.exec( - select(db.SPARelionParameters).where( - db.SPARelionParameters.pj_id == pj_id_params - ) - ).one() - feedback_params = _db.exec( - select(db.SPAFeedbackParameters).where( - db.SPAFeedbackParameters.pj_id == pj_id_params - ) - ).one() - if feedback_params.hold_class2d: - feedback_params.rerun_class2d = True - _db.add(feedback_params) - _db.commit() - _db.close() - return - feedback_params.next_job = 10 if default_spa_parameters.do_icebreaker_jobs else 7 - feedback_params.hold_class2d = True - relion_options = dict(relion_params) - other_options = dict(feedback_params) - if other_options["picker_ispyb_id"] is None: - logger.info("No ISPyB particle picker ID yet") - feedback_params.hold_class2d = False - _db.add(feedback_params) - _db.commit() - _db.expunge(feedback_params) - return - _db.add(feedback_params) - _db.commit() - _db.expunge(feedback_params) - class2d_message = message.get("class2d_message") - assert isinstance(class2d_message, dict) - if not _db.exec( - select(func.count(db.Class2DParameters.particles_file)) - .where(db.Class2DParameters.particles_file == class2d_message["particles_file"]) - .where(db.Class2DParameters.pj_id == pj_id) - ).one(): - class2d_params = db.Class2DParameters( - pj_id=pj_id, - murfey_id=_murfey_id(message["program_id"], _db)[0], - particles_file=class2d_message["particles_file"], - class2d_dir=class2d_message["class2d_dir"], - batch_size=class2d_message["batch_size"], - complete=False, - ) - _db.add(class2d_params) - _db.commit() - murfey_ids = _murfey_id(message["program_id"], _db, number=50) - _murfey_class2ds( - murfey_ids, class2d_message["particles_file"], message["program_id"], _db - ) - zocalo_message: dict = { - "parameters": { - "particles_file": class2d_message["particles_file"], - "class2d_dir": f"{class2d_message['class2d_dir']}{other_options['next_job']:03}", - "batch_is_complete": False, - "particle_diameter": relion_options["particle_diameter"], - "combine_star_job_number": -1, - "picker_id": other_options["picker_ispyb_id"], - "nr_iter": default_spa_parameters.nr_iter_2d, - "batch_size": default_spa_parameters.batch_size_2d, - "nr_classes": default_spa_parameters.nr_classes_2d, - "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, - "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, - "mask_diameter": 0, - "class_uuids": _2d_class_murfey_ids( - class2d_message["particles_file"], _app_id(pj_id, _db), _db - ), - "class2d_grp_uuid": _db.exec( - select(db.Class2DParameters).where( - db.Class2DParameters.particles_file - == class2d_message["particles_file"] - and db.Class2DParameters.pj_id == pj_id - ) - ) - .one() - .murfey_id, - "session_id": message["session_id"], - "autoproc_program_id": _app_id( - _pj_id(message["program_id"], _db, recipe="em-spa-class2d"), _db - ), - "node_creator_queue": machine_config.node_creator_queue, - }, - "recipes": ["em-spa-class2d"], - } - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send("processing_recipe", zocalo_message, new_connection=True) - logger.info("2D classification requested") - if demo: - logger.info("Incomplete 2D batch registered in demo mode") - if not _db.exec( - select(func.count(db.Class2DParameters.particles_file)).where( - db.Class2DParameters.particles_file == class2d_message["particles_file"] - and db.Class2DParameters.pj_id == pj_id - and db.Class2DParameters.complete - ) - ).one(): - _register_complete_2d_batch(message, _db=_db, demo=demo) - message["class2d_message"]["particles_file"] = ( - message["class2d_message"]["particles_file"] + "_new" - ) - _register_complete_2d_batch(message, _db=_db, demo=demo) - _db.close() - - -def _register_complete_2d_batch(message: dict, _db=murfey_db, demo: bool = False): - """Received full batch from particle selection service""" - instrument_name = ( - _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - class2d_message = message.get("class2d_message") - assert isinstance(class2d_message, dict) - pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") - pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-class2d") - relion_params = _db.exec( - select(db.SPARelionParameters).where( - db.SPARelionParameters.pj_id == pj_id_params - ) - ).one() - feedback_params = _db.exec( - select(db.SPAFeedbackParameters).where( - db.SPAFeedbackParameters.pj_id == pj_id_params - ) - ).one() - _db.expunge(relion_params) - _db.expunge(feedback_params) - if feedback_params.hold_class2d or feedback_params.picker_ispyb_id is None: - feedback_params.rerun_class2d = True - _db.add(feedback_params) - _db.commit() - # If waiting then save the message - if _db.exec( - select(func.count(db.Class2DParameters.particles_file)) - .where(db.Class2DParameters.pj_id == pj_id) - .where( - db.Class2DParameters.particles_file == class2d_message["particles_file"] - ) - ).one(): - class2d_params = _db.exec( - select(db.Class2DParameters) - .where(db.Class2DParameters.pj_id == pj_id) - .where( - db.Class2DParameters.particles_file - == class2d_message["particles_file"] - ) - ).one() - class2d_params.complete = True - _db.add(class2d_params) - _db.commit() - _db.close() - else: - class2d_params = db.Class2DParameters( - pj_id=pj_id, - murfey_id=_murfey_id(message["program_id"], _db)[0], - particles_file=class2d_message["particles_file"], - class2d_dir=class2d_message["class2d_dir"], - batch_size=class2d_message["batch_size"], - ) - _db.add(class2d_params) - _db.commit() - _db.close() - murfey_ids = _murfey_id(_app_id(pj_id, _db), _db, number=50) - _murfey_class2ds( - murfey_ids, class2d_message["particles_file"], _app_id(pj_id, _db), _db - ) - if demo: - _register_class_selection( - {"session_id": message["session_id"], "class_selection_score": 0.5}, - _db=_db, - demo=demo, - ) - elif not feedback_params.class_selection_score: - # For the first batch, start a container and set the database to wait - job_number_after_first_batch = ( - 10 if default_spa_parameters.do_icebreaker_jobs else 7 - ) - if ( - feedback_params.next_job is not None - and feedback_params.next_job < job_number_after_first_batch - ): - feedback_params.next_job = job_number_after_first_batch - if not feedback_params.star_combination_job: - feedback_params.star_combination_job = feedback_params.next_job + ( - 3 if default_spa_parameters.do_icebreaker_jobs else 2 - ) - if _db.exec( - select(func.count(db.Class2DParameters.particles_file)) - .where(db.Class2DParameters.pj_id == pj_id) - .where( - db.Class2DParameters.particles_file == class2d_message["particles_file"] - ) - ).one(): - class_uuids = _2d_class_murfey_ids( - class2d_message["particles_file"], _app_id(pj_id, _db), _db - ) - class2d_grp_uuid = ( - _db.exec( - select(db.Class2DParameters) - .where(db.Class2DParameters.pj_id == pj_id) - .where( - db.Class2DParameters.particles_file - == class2d_message["particles_file"] - ) - ) - .one() - .murfey_id - ) - else: - class_uuids = { - str(i + 1): m - for i, m in enumerate(_murfey_id(_app_id(pj_id, _db), _db, number=50)) - } - class2d_grp_uuid = _murfey_id(_app_id(pj_id, _db), _db)[0] - zocalo_message: dict = { - "parameters": { - "particles_file": class2d_message["particles_file"], - "class2d_dir": f"{class2d_message['class2d_dir']}{feedback_params.next_job:03}", - "batch_is_complete": True, - "particle_diameter": relion_params.particle_diameter, - "mask_diameter": relion_params.mask_diameter or 0, - "combine_star_job_number": feedback_params.star_combination_job, - "autoselect_min_score": 0, - "picker_id": feedback_params.picker_ispyb_id, - "class_uuids": class_uuids, - "class2d_grp_uuid": class2d_grp_uuid, - "nr_iter": default_spa_parameters.nr_iter_2d, - "batch_size": default_spa_parameters.batch_size_2d, - "nr_classes": default_spa_parameters.nr_classes_2d, - "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, - "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, - "session_id": message["session_id"], - "autoproc_program_id": _app_id( - _pj_id(message["program_id"], _db, recipe="em-spa-class2d"), _db - ), - "node_creator_queue": machine_config.node_creator_queue, - }, - "recipes": ["em-spa-class2d"], - } - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - feedback_params.hold_class2d = True - feedback_params.next_job += ( - 4 if default_spa_parameters.do_icebreaker_jobs else 3 - ) - _db.add(feedback_params) - _db.commit() - _db.close() - else: - # Send all other messages on to a container - if _db.exec( - select(func.count(db.Class2DParameters.particles_file)) - .where(db.Class2DParameters.pj_id == pj_id) - .where( - db.Class2DParameters.particles_file == class2d_message["particles_file"] - ) - ).one(): - class_uuids = _2d_class_murfey_ids( - class2d_message["particles_file"], _app_id(pj_id, _db), _db - ) - class2d_grp_uuid = ( - _db.exec( - select(db.Class2DParameters) - .where(db.Class2DParameters.pj_id == pj_id) - .where( - db.Class2DParameters.particles_file - == class2d_message["particles_file"] - ) - ) - .one() - .murfey_id - ) - else: - class_uuids = { - str(i + 1): m - for i, m in enumerate(_murfey_id(_app_id(pj_id, _db), _db, number=50)) - } - class2d_grp_uuid = _murfey_id(_app_id(pj_id, _db), _db)[0] - zocalo_message = { - "parameters": { - "particles_file": class2d_message["particles_file"], - "class2d_dir": f"{class2d_message['class2d_dir']}{feedback_params.next_job:03}", - "batch_is_complete": True, - "particle_diameter": relion_params.particle_diameter, - "mask_diameter": relion_params.mask_diameter or 0, - "combine_star_job_number": feedback_params.star_combination_job, - "autoselect_min_score": feedback_params.class_selection_score or 0, - "picker_id": feedback_params.picker_ispyb_id, - "class_uuids": class_uuids, - "class2d_grp_uuid": class2d_grp_uuid, - "nr_iter": default_spa_parameters.nr_iter_2d, - "batch_size": default_spa_parameters.batch_size_2d, - "nr_classes": default_spa_parameters.nr_classes_2d, - "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, - "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, - "session_id": message["session_id"], - "autoproc_program_id": _app_id( - _pj_id(message["program_id"], _db, recipe="em-spa-class2d"), _db - ), - "node_creator_queue": machine_config.node_creator_queue, - }, - "recipes": ["em-spa-class2d"], - } - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - feedback_params.next_job += ( - 3 if default_spa_parameters.do_icebreaker_jobs else 2 - ) - _db.add(feedback_params) - _db.commit() - _db.close() - - -def _flush_class2d( - session_id: int, - app_id: int, - _db, - relion_params: db.SPARelionParameters | None = None, - feedback_params: db.SPAFeedbackParameters | None = None, -): - instrument_name = ( - _db.exec(select(db.Session).where(db.Session.id == session_id)) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - if not relion_params or feedback_params: - pj_id_params = _pj_id(app_id, _db, recipe="em-spa-preprocess") - if not relion_params: - relion_params = _db.exec( - select(db.SPARelionParameters).where( - db.SPARelionParameters.pj_id == pj_id_params - ) - ).one() - _db.expunge(relion_params) - if not feedback_params: - feedback_params = _db.exec( - select(db.SPAFeedbackParameters).where( - db.SPAFeedbackParameters.pj_id == pj_id_params - ) - ).one() - _db.expunge(feedback_params) - if not relion_params or not feedback_params: - return - pj_id = _pj_id(app_id, _db, recipe="em-spa-class2d") - class2d_db = _db.exec( - select(db.Class2DParameters) - .where(db.Class2DParameters.pj_id == pj_id) - .where(db.Class2DParameters.complete) - ).all() - if not feedback_params.next_job: - feedback_params.next_job = ( - 10 if default_spa_parameters.do_icebreaker_jobs else 7 - ) - if not feedback_params.star_combination_job: - feedback_params.star_combination_job = feedback_params.next_job + ( - 3 if default_spa_parameters.do_icebreaker_jobs else 2 - ) - for saved_message in class2d_db: - # Send all held Class2D messages on with the selection score added - _db.expunge(saved_message) - zocalo_message: dict = { - "parameters": { - "particles_file": saved_message.particles_file, - "class2d_dir": f"{saved_message.class2d_dir}{feedback_params.next_job:03}", - "batch_is_complete": True, - "batch_size": saved_message.batch_size, - "particle_diameter": relion_params.particle_diameter, - "mask_diameter": relion_params.mask_diameter or 0, - "combine_star_job_number": feedback_params.star_combination_job, - "autoselect_min_score": feedback_params.class_selection_score or 0, - "picker_id": feedback_params.picker_ispyb_id, - "class_uuids": _2d_class_murfey_ids( - saved_message.particles_file, _app_id(pj_id, _db), _db - ), - "class2d_grp_uuid": saved_message.murfey_id, - "nr_iter": default_spa_parameters.nr_iter_2d, - "nr_classes": default_spa_parameters.nr_classes_2d, - "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, - "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, - "session_id": session_id, - "autoproc_program_id": _app_id(pj_id, _db), - "node_creator_queue": machine_config.node_creator_queue, - }, - "recipes": ["em-spa-class2d"], - } - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - feedback_params.next_job += ( - 3 if default_spa_parameters.do_icebreaker_jobs else 2 - ) - _db.delete(saved_message) - _db.add(feedback_params) - _db.commit() - - -def _register_class_selection(message: dict, _db=murfey_db, demo: bool = False): - """Received selection score from class selection service""" - pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") - pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-class2d") - relion_params = _db.exec( - select(db.SPARelionParameters).where( - db.SPARelionParameters.pj_id == pj_id_params - ) - ).one() - class2d_db = _db.exec( - select(db.Class2DParameters).where(db.Class2DParameters.pj_id == pj_id) - ).all() - # Add the class selection score to the database - feedback_params = _db.exec( - select(db.SPAFeedbackParameters).where( - db.SPAFeedbackParameters.pj_id == pj_id_params - ) - ).one() - _db.expunge(feedback_params) - - if feedback_params.picker_ispyb_id is None: - selection_stash = db.SelectionStash( - pj_id=pj_id, - class_selection_score=message["class_selection_score"] or 0, - ) - _db.add(selection_stash) - _db.commit() - _db.close() - return - - feedback_params.class_selection_score = message.get("class_selection_score") or 0 - feedback_params.hold_class2d = False - next_job = feedback_params.next_job - if demo: - for saved_message in class2d_db: - # Send all held Class2D messages on with the selection score added - _db.expunge(saved_message) - particles_file = saved_message.particles_file - logger.info("Complete 2D classification registered in demo mode") - _register_3d_batch( - { - "session_id": message["session_id"], - "class3d_message": { - "particles_file": particles_file, - "class3d_dir": "Class3D", - "batch_size": 50000, - }, - }, - _db=_db, - demo=demo, - ) - logger.info("3D classification registered in demo mode") - _register_3d_batch( - { - "session_id": message["session_id"], - "class3d_message": { - "particles_file": particles_file + "_new", - "class3d_dir": "Class3D", - "batch_size": 50000, - }, - }, - _db=_db, - demo=demo, - ) - _register_initial_model( - { - "session_id": message["session_id"], - "initial_model": "InitialModel/job015/model.mrc", - }, - _db=_db, - demo=demo, - ) - next_job += 3 if default_spa_parameters.do_icebreaker_jobs else 2 - feedback_params.next_job = next_job - _db.close() - else: - _flush_class2d( - message["session_id"], - message["program_id"], - _db, - relion_params=relion_params, - feedback_params=feedback_params, - ) - _db.add(feedback_params) - for sm in class2d_db: - _db.delete(sm) - _db.commit() - _db.close() - - -def _find_initial_model(visit: str, machine_config: MachineConfig) -> Path | None: - if machine_config.initial_model_search_directory: - visit_directory = ( - machine_config.rsync_basepath / str(datetime.now().year) / visit - ) - possible_models = [ - p - for p in ( - visit_directory / machine_config.initial_model_search_directory - ).glob("*.mrc") - if "rescaled" not in p.name - ] - if possible_models: - return sorted(possible_models, key=lambda x: x.stat().st_ctime)[-1] - return None - - -def _downscaled_box_size( - particle_diameter: int, pixel_size: float -) -> Tuple[int, float]: - box_size = int(math.ceil(1.2 * particle_diameter)) - box_size = box_size + box_size % 2 - for small_box_pix in ( - 48, - 64, - 96, - 128, - 160, - 192, - 256, - 288, - 300, - 320, - 360, - 384, - 400, - 420, - 450, - 480, - 512, - 640, - 768, - 896, - 1024, - ): - # Don't go larger than the original box - if small_box_pix > box_size: - return box_size, pixel_size - # If Nyquist freq. is better than 8.5 A, use this downscaled box, else step size - small_box_angpix = pixel_size * box_size / small_box_pix - if small_box_angpix < 4.25: - return small_box_pix, small_box_angpix - raise ValueError(f"Box size is too large: {box_size}") - - -def _resize_intial_model( - downscaled_box_size: int, - downscaled_pixel_size: float, - input_path: Path, - output_path: Path, - executables: Dict[str, str], - env: Dict[str, str], -) -> None: - if executables.get("relion_image_handler"): - comp_proc = subprocess.run( - [ - f"{executables['relion_image_handler']}", - "--i", - str(input_path), - "--new_box", - str(downscaled_box_size), - "--rescale_angpix", - str(downscaled_pixel_size), - "--o", - str(output_path), - ], - capture_output=True, - text=True, - env=env, - ) - with mrcfile.open(output_path) as rescaled_mrc: - rescaled_mrc.header.cella = ( - downscaled_pixel_size, - downscaled_pixel_size, - downscaled_pixel_size, - ) - if comp_proc.returncode: - logger.error( - f"Resizing initial model {input_path} failed \n {comp_proc.stdout}" - ) - return None - - -def _register_3d_batch(message: dict, _db=murfey_db, demo: bool = False): - """Received 3d batch from class selection service""" - class3d_message = message.get("class3d_message") - assert isinstance(class3d_message, dict) - instrument_name = ( - _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") - pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-class3d") - relion_params = _db.exec( - select(db.SPARelionParameters).where( - db.SPARelionParameters.pj_id == pj_id_params - ) - ).one() - relion_options = dict(relion_params) - feedback_params = _db.exec( - select(db.SPAFeedbackParameters).where( - db.SPAFeedbackParameters.pj_id == pj_id_params - ) - ).one() - other_options = dict(feedback_params) - - visit_name = ( - _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) - .one() - .visit - ) - - provided_initial_model = _find_initial_model(visit_name, machine_config) - if provided_initial_model and not feedback_params.initial_model: - rescaled_initial_model_path = ( - provided_initial_model.parent - / f"{provided_initial_model.stem}_rescaled_{pj_id}{provided_initial_model.suffix}" - ) - if not rescaled_initial_model_path.is_file(): - _resize_intial_model( - *_downscaled_box_size( - message["particle_diameter"], - relion_options["angpix"], - ), - provided_initial_model, - rescaled_initial_model_path, - machine_config.external_executables, - machine_config.external_environment, - ) - feedback_params.initial_model = str(rescaled_initial_model_path) - other_options["initial_model"] = str(rescaled_initial_model_path) - next_job = feedback_params.next_job - class3d_dir = ( - f"{class3d_message['class3d_dir']}{(feedback_params.next_job+1):03}" - ) - feedback_params.next_job += 1 - _db.add(feedback_params) - _db.commit() - - class3d_grp_uuid = _murfey_id(message["program_id"], _db)[0] - class_uuids = _murfey_id(message["program_id"], _db, number=4) - class3d_params = db.Class3DParameters( - pj_id=pj_id, - murfey_id=class3d_grp_uuid, - particles_file=class3d_message["particles_file"], - class3d_dir=class3d_dir, - batch_size=class3d_message["batch_size"], - ) - _db.add(class3d_params) - _db.commit() - _murfey_class3ds( - class_uuids, - class3d_message["particles_file"], - message["program_id"], - _db, - ) - - if feedback_params.hold_class3d: - # If waiting then save the message - class3d_params = _db.exec( - select(db.Class3DParameters).where(db.Class3DParameters.pj_id == pj_id) - ).one() - class3d_params.run = True - class3d_params.particles_file = class3d_message["particles_file"] - class3d_params.batch_size = class3d_message["batch_size"] - _db.add(class3d_params) - _db.commit() - _db.close() - elif not feedback_params.initial_model: - # For the first batch, start a container and set the database to wait - next_job = feedback_params.next_job - class3d_dir = ( - f"{class3d_message['class3d_dir']}{(feedback_params.next_job+1):03}" - ) - class3d_grp_uuid = _murfey_id(message["program_id"], _db)[0] - class_uuids = _murfey_id(message["program_id"], _db, number=4) - class3d_params = db.Class3DParameters( - pj_id=pj_id, - murfey_id=class3d_grp_uuid, - particles_file=class3d_message["particles_file"], - class3d_dir=class3d_dir, - batch_size=class3d_message["batch_size"], - ) - _db.add(class3d_params) - _db.commit() - _murfey_class3ds( - class_uuids, class3d_message["particles_file"], message["program_id"], _db - ) - - feedback_params.hold_class3d = True - next_job += 2 - feedback_params.next_job = next_job - zocalo_message: dict = { - "parameters": { - "particles_file": class3d_message["particles_file"], - "class3d_dir": class3d_dir, - "batch_size": class3d_message["batch_size"], - "symmetry": relion_options["symmetry"], - "particle_diameter": relion_options["particle_diameter"], - "mask_diameter": relion_options["mask_diameter"] or 0, - "do_initial_model": True, - "picker_id": other_options["picker_ispyb_id"], - "class_uuids": {i + 1: m for i, m in enumerate(class_uuids)}, - "class3d_grp_uuid": class3d_grp_uuid, - "nr_iter": default_spa_parameters.nr_iter_3d, - "initial_model_iterations": default_spa_parameters.nr_iter_ini_model, - "nr_classes": default_spa_parameters.nr_classes_3d, - "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, - "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, - "session_id": message["session_id"], - "autoproc_program_id": _app_id( - _pj_id(message["program_id"], _db, recipe="em-spa-class3d"), _db - ), - "node_creator_queue": machine_config.node_creator_queue, - }, - "recipes": ["em-spa-class3d"], - } - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - _db.add(feedback_params) - _db.commit() - _db.close() - else: - # Send all other messages on to a container - class3d_params = _db.exec( - select(db.Class3DParameters).where(db.Class3DParameters.pj_id == pj_id) - ).one() - zocalo_message = { - "parameters": { - "particles_file": class3d_message["particles_file"], - "class3d_dir": class3d_params.class3d_dir, - "batch_size": class3d_message["batch_size"], - "symmetry": relion_options["symmetry"], - "particle_diameter": relion_options["particle_diameter"], - "mask_diameter": relion_options["mask_diameter"] or 0, - "do_initial_model": False, - "initial_model_file": other_options["initial_model"], - "picker_id": other_options["picker_ispyb_id"], - "class_uuids": _3d_class_murfey_ids( - class3d_params.particles_file, _app_id(pj_id, _db), _db - ), - "class3d_grp_uuid": class3d_params.murfey_id, - "nr_iter": default_spa_parameters.nr_iter_3d, - "initial_model_iterations": default_spa_parameters.nr_iter_ini_model, - "nr_classes": default_spa_parameters.nr_classes_3d, - "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, - "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, - "session_id": message["session_id"], - "autoproc_program_id": _app_id( - _pj_id(message["program_id"], _db, recipe="em-spa-class3d"), _db - ), - "node_creator_queue": machine_config.node_creator_queue, - }, - "recipes": ["em-spa-class3d"], - } - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - feedback_params.hold_class3d = True - _db.add(feedback_params) - _db.commit() - _db.close() - - -def _register_initial_model(message: dict, _db=murfey_db, demo: bool = False): - """Received initial model from 3d classification service""" - pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") - # Add the initial model file to the database - feedback_params = _db.exec( - select(db.SPAFeedbackParameters).where( - db.SPAFeedbackParameters.pj_id == pj_id_params - ) - ).one() - feedback_params.initial_model = message.get("initial_model") - _db.add(feedback_params) - _db.commit() - _db.close() - - -def _flush_tomography_preprocessing(message: dict): - session_id = message["session_id"] - instrument_name = ( - murfey_db.exec(select(db.Session).where(db.Session.id == session_id)) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - stashed_files = murfey_db.exec( - select(db.PreprocessStash) - .where(db.PreprocessStash.session_id == session_id) - .where(db.PreprocessStash.group_tag == message["data_collection_group_tag"]) - ).all() - if not stashed_files: - return - collected_ids = murfey_db.exec( - select( - db.DataCollectionGroup, - ) - .where(db.DataCollectionGroup.session_id == session_id) - .where(db.DataCollectionGroup.tag == message["data_collection_group_tag"]) - ).first() - proc_params = get_tomo_proc_params(collected_ids.id) - if not proc_params: - visit_name = message["visit_name"].replace("\r\n", "").replace("\n", "") - logger.warning( - f"No tomography processing parameters found for Murfey session {sanitise(str(message['session_id']))} on visit {sanitise(visit_name)}" - ) - return - - recipe_name = machine_config.recipes.get("em-tomo-preprocess", "em-tomo-preprocess") - - for f in stashed_files: - collected_ids = murfey_db.exec( - select( - db.DataCollectionGroup, - db.DataCollection, - db.ProcessingJob, - db.AutoProcProgram, - ) - .where(db.DataCollectionGroup.session_id == session_id) - .where(db.DataCollectionGroup.tag == message["data_collection_group_tag"]) - .where(db.DataCollection.dcg_id == db.DataCollectionGroup.id) - .where(db.DataCollection.tag == f.tag) - .where(db.ProcessingJob.dc_id == db.DataCollection.id) - .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id) - .where(db.ProcessingJob.recipe == recipe_name) - ).one() - detached_ids = [c.id for c in collected_ids] - - murfey_ids = _murfey_id(detached_ids[3], murfey_db, number=1, close=False) - p = Path(f.mrc_out) - if not p.parent.exists(): - p.parent.mkdir(parents=True) - movie = db.Movie( - murfey_id=murfey_ids[0], - path=f.file_path, - image_number=f.image_number, - tag=f.tag, - ) - murfey_db.add(movie) - zocalo_message: dict = { - "recipes": [recipe_name], - "parameters": { - "node_creator_queue": machine_config.node_creator_queue, - "dcid": detached_ids[1], - "autoproc_program_id": detached_ids[3], - "movie": f.file_path, - "mrc_out": f.mrc_out, - "pixel_size": proc_params.pixel_size, - "kv": proc_params.voltage, - "image_number": f.image_number, - "microscope": get_microscope(), - "mc_uuid": murfey_ids[0], - "ft_bin": proc_params.motion_corr_binning, - "fm_dose": proc_params.dose_per_frame, - "frame_count": proc_params.frame_count, - "gain_ref": ( - str(machine_config.rsync_basepath / proc_params.gain_ref) - if proc_params.gain_ref - else proc_params.gain_ref - ), - "fm_int_file": proc_params.eer_fractionation_file or "", - }, - } - logger.info( - f"Launching tomography preprocessing with Zocalo message: {zocalo_message}" - ) - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - else: - feedback_callback( - {}, - { - "register": "motion_corrected", - "movie": f.file_path, - "mrc_out": f.mrc_out, - "movie_id": murfey_ids[0], - "program_id": detached_ids[3], - }, - ) - murfey_db.delete(f) - murfey_db.commit() - murfey_db.close() - - -def _flush_grid_square_records(message: dict, _db=murfey_db, demo: bool = False): - tag = message["tag"] - session_id = message["session_id"] - gs_ids = [] - for gs in _db.exec( - select(db.GridSquare) - .where(db.GridSquare.session_id == session_id) - .where(db.GridSquare.tag == tag) - ).all(): - gs_ids.append(gs.id) - if demo: - logger.info(f"Flushing grid square {gs.name}") - for i in gs_ids: - _flush_foil_hole_records(i, _db=_db, demo=demo) - - -def _flush_foil_hole_records(grid_square_id: int, _db=murfey_db, demo: bool = False): - for fh in _db.exec( - select(db.FoilHole).where(db.FoilHole.grid_square_id == grid_square_id) - ).all(): - if demo: - logger.info(f"Flushing foil hole: {fh.name}") - - -def _register_refinement(message: dict, _db=murfey_db, demo: bool = False): - """Received class to refine from 3D classification""" - instrument_name = ( - _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") - pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-refine") - relion_params = _db.exec( - select(db.SPARelionParameters).where( - db.SPARelionParameters.pj_id == pj_id_params - ) - ).one() - relion_options = dict(relion_params) - feedback_params = _db.exec( - select(db.SPAFeedbackParameters).where( - db.SPAFeedbackParameters.pj_id == pj_id_params - ) - ).one() - other_options = dict(feedback_params) - - if feedback_params.hold_refine: - # If waiting then save the message - refine_params = _db.exec( - select(db.RefineParameters) - .where(db.RefineParameters.pj_id == pj_id) - .where(db.RefineParameters.tag == "first") - ).one() - # refine_params.refine_dir is not set as it will be the same as before - refine_params.run = True - refine_params.class3d_dir = message["class3d_dir"] - refine_params.class_number = message["best_class"] - _db.add(refine_params) - _db.commit() - _db.close() - else: - # Send all other messages on to a container - try: - refine_params = _db.exec( - select(db.RefineParameters) - .where(db.RefineParameters.pj_id == pj_id) - .where(db.RefineParameters.tag == "first") - ).one() - symmetry_refine_params = _db.exec( - select(db.RefineParameters) - .where(db.RefineParameters.pj_id == pj_id) - .where(db.RefineParameters.tag == "symmetry") - ).one() - except SQLAlchemyError: - next_job = feedback_params.next_job - refine_dir = f"{message['refine_dir']}{(feedback_params.next_job + 2):03}" - refined_grp_uuid = _murfey_id(message["program_id"], _db)[0] - refined_class_uuid = _murfey_id(message["program_id"], _db)[0] - symmetry_refined_grp_uuid = _murfey_id(message["program_id"], _db)[0] - symmetry_refined_class_uuid = _murfey_id(message["program_id"], _db)[0] - - refine_params = db.RefineParameters( - tag="first", - pj_id=pj_id, - murfey_id=refined_grp_uuid, - refine_dir=refine_dir, - class3d_dir=message["class3d_dir"], - class_number=message["best_class"], - ) - symmetry_refine_params = db.RefineParameters( - tag="symmetry", - pj_id=pj_id, - murfey_id=symmetry_refined_grp_uuid, - refine_dir=refine_dir, - class3d_dir=message["class3d_dir"], - class_number=message["best_class"], - ) - _db.add(refine_params) - _db.add(symmetry_refine_params) - _db.commit() - _murfey_refine( - murfey_id=refined_class_uuid, - refine_dir=refine_dir, - tag="first", - app_id=message["program_id"], - _db=_db, - ) - _murfey_refine( - murfey_id=symmetry_refined_class_uuid, - refine_dir=refine_dir, - tag="symmetry", - app_id=message["program_id"], - _db=_db, - ) - - if relion_options["symmetry"] == "C1": - # Extra Refine, Mask, PostProcess beyond for determined symmetry - next_job += 8 - else: - # Select and Extract particles, then Refine, Mask, PostProcess - next_job += 5 - feedback_params.next_job = next_job - - zocalo_message: dict = { - "parameters": { - "refine_job_dir": refine_params.refine_dir, - "class3d_dir": message["class3d_dir"], - "class_number": message["best_class"], - "pixel_size": relion_options["angpix"], - "particle_diameter": relion_options["particle_diameter"], - "mask_diameter": relion_options["mask_diameter"] or 0, - "symmetry": relion_options["symmetry"], - "node_creator_queue": machine_config.node_creator_queue, - "nr_iter": default_spa_parameters.nr_iter_3d, - "picker_id": other_options["picker_ispyb_id"], - "refined_class_uuid": _refine_murfey_id( - refine_dir=refine_params.refine_dir, - tag=refine_params.tag, - app_id=_app_id(pj_id, _db), - _db=_db, - ), - "refined_grp_uuid": refine_params.murfey_id, - "symmetry_refined_class_uuid": _refine_murfey_id( - refine_dir=symmetry_refine_params.refine_dir, - tag=symmetry_refine_params.tag, - app_id=_app_id(pj_id, _db), - _db=_db, - ), - "symmetry_refined_grp_uuid": symmetry_refine_params.murfey_id, - "session_id": message["session_id"], - "autoproc_program_id": _app_id( - _pj_id(message["program_id"], _db, recipe="em-spa-refine"), _db - ), - }, - "recipes": ["em-spa-refine"], - } - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - feedback_params.hold_refine = True - _db.add(feedback_params) - _db.commit() - _db.close() - - -def _register_bfactors(message: dict, _db=murfey_db, demo: bool = False): - """Received refined class to calculate b-factor""" - instrument_name = ( - _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") - pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-refine") - relion_params = _db.exec( - select(db.SPARelionParameters).where( - db.SPARelionParameters.pj_id == pj_id_params - ) - ).one() - relion_options = dict(relion_params) - feedback_params = _db.exec( - select(db.SPAFeedbackParameters).where( - db.SPAFeedbackParameters.pj_id == pj_id_params - ) - ).one() - - if message["symmetry"] != relion_params.symmetry: - # Currently don't do anything with a symmetrised re-run of the refinement - logger.info( - f"Received symmetrised structure of {sanitise(message['symmetry'])}" - ) - return True - - if not feedback_params.hold_refine: - logger.warning("B-Factors requested but refine hold is off") - return False - - # Add b-factor for refinement run - bfactor_run = db.BFactors( - pj_id=pj_id, - bfactor_directory=f"{message['project_dir']}/Refine3D/bfactor_{message['number_of_particles']}", - number_of_particles=message["number_of_particles"], - resolution=message["resolution"], - ) - _db.add(bfactor_run) - _db.commit() - - # All messages should create b-factor jobs as the refine hold is on at this point - try: - bfactor_params = _db.exec( - select(db.BFactorParameters).where(db.BFactorParameters.pj_id == pj_id) - ).one() - except SQLAlchemyError: - bfactor_params = db.BFactorParameters( - pj_id=pj_id, - project_dir=message["project_dir"], - batch_size=message["number_of_particles"], - refined_grp_uuid=message["refined_grp_uuid"], - refined_class_uuid=message["refined_class_uuid"], - class_reference=message["class_reference"], - class_number=message["class_number"], - mask_file=message["mask_file"], - ) - _db.add(bfactor_params) - _db.commit() - - bfactor_particle_count = default_spa_parameters.bfactor_min_particles - while bfactor_particle_count < bfactor_params.batch_size: - bfactor_run_name = ( - f"{bfactor_params.project_dir}/BFactors/bfactor_{bfactor_particle_count}" - ) - try: - bfactor_run = _db.exec( - select(db.BFactors) - .where(db.BFactors.pj_id == pj_id) - .where(db.BFactors.bfactor_directory == bfactor_run_name) - ).one() - bfactor_run.resolution = 0 - except SQLAlchemyError: - bfactor_run = db.BFactors( - pj_id=pj_id, - bfactor_directory=bfactor_run_name, - number_of_particles=bfactor_particle_count, - resolution=0, - ) - _db.add(bfactor_run) - _db.commit() - - bfactor_particle_count *= 2 - - zocalo_message: dict = { - "parameters": { - "bfactor_directory": bfactor_run.bfactor_directory, - "class_reference": bfactor_params.class_reference, - "class_number": bfactor_params.class_number, - "number_of_particles": bfactor_run.number_of_particles, - "batch_size": bfactor_params.batch_size, - "pixel_size": message["pixel_size"], - "mask": bfactor_params.mask_file, - "particle_diameter": relion_options["particle_diameter"], - "mask_diameter": relion_options["mask_diameter"] or 0, - "node_creator_queue": machine_config.node_creator_queue, - "picker_id": feedback_params.picker_ispyb_id, - "refined_grp_uuid": bfactor_params.refined_grp_uuid, - "refined_class_uuid": bfactor_params.refined_class_uuid, - "session_id": message["session_id"], - "autoproc_program_id": _app_id( - _pj_id(message["program_id"], _db, recipe="em-spa-refine"), _db - ), - }, - "recipes": ["em-spa-bfactor"], - } - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - _db.close() - return True - - -def _save_bfactor(message: dict, _db=murfey_db, demo: bool = False): - """Received b-factor from refinement run""" - pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-refine") - bfactor_run = _db.exec( - select(db.BFactors) - .where(db.BFactors.pj_id == pj_id) - .where(db.BFactors.number_of_particles == message["number_of_particles"]) - ).one() - bfactor_run.resolution = message["resolution"] - _db.add(bfactor_run) - _db.commit() - - # Find all the resolutions in the b-factors table - all_bfactors = _db.exec(select(db.BFactors).where(db.BFactors.pj_id == pj_id)).all() - particle_counts = [bf.number_of_particles for bf in all_bfactors] - resolutions = [bf.resolution for bf in all_bfactors] - - if all(resolutions): - # Calculate b-factor and add to ispyb class table - bfactor_fitting = np.polyfit( - np.log(particle_counts), 1 / np.array(resolutions) ** 2, 1 - ) - refined_class_uuid = message["refined_class_uuid"] - - # Request an ispyb insert of the b-factor fitting parameters - if _transport_object: - _transport_object.send( - "ispyb_connector", - { - "ispyb_command": "buffer", - "buffer_lookup": { - "particle_classification_id": refined_class_uuid, - }, - "buffer_command": { - "ispyb_command": "insert_particle_classification" - }, - "program_id": message["program_id"], - "bfactor_fit_intercept": str(bfactor_fitting[1]), - "bfactor_fit_linear": str(bfactor_fitting[0]), - }, - new_connection=True, - ) - - # Clean up the b-factors table and release the hold - [_db.delete(bf) for bf in all_bfactors] - _db.commit() - _release_refine_hold(message) - _db.close() - - -def feedback_callback(header: dict, message: dict) -> None: - try: - record = None - if "environment" in message: - params = message["recipe"][str(message["recipe-pointer"])].get( - "parameters", {} - ) - message = message["payload"] - message.update(params) - if message["register"] == "motion_corrected": - collected_ids = murfey_db.exec( - select( - db.DataCollectionGroup, - db.DataCollection, - db.ProcessingJob, - db.AutoProcProgram, - ) - .where(db.DataCollection.dcg_id == db.DataCollectionGroup.id) - .where(db.ProcessingJob.dc_id == db.DataCollection.id) - .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id) - .where(db.AutoProcProgram.id == message["program_id"]) - ).one() - session_id = collected_ids[0].session_id - - # Find the autoprocprogram id for the alignment recipe - alignment_ids = murfey_db.exec( - select( - db.DataCollection, - db.ProcessingJob, - db.AutoProcProgram, - ) - .where(db.ProcessingJob.dc_id == db.DataCollection.id) - .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id) - .where(db.DataCollection.id == collected_ids[1].id) - .where(db.ProcessingJob.recipe == "em-tomo-align") - ).one() - - relevant_tilt_and_series = murfey_db.exec( - select(db.Tilt, db.TiltSeries) - .where(db.Tilt.movie_path == message.get("movie")) - .where(db.Tilt.tilt_series_id == db.TiltSeries.id) - .where(db.TiltSeries.session_id == session_id) - ).one() - relevant_tilt = relevant_tilt_and_series[0] - relevant_tilt_series = relevant_tilt_and_series[1] - relevant_tilt.motion_corrected = True - murfey_db.add(relevant_tilt) - murfey_db.commit() - if ( - check_tilt_series_mc(relevant_tilt_series.id) - and not relevant_tilt_series.processing_requested - and relevant_tilt_series.tilt_series_length > 2 - ): - relevant_tilt_series.processing_requested = True - murfey_db.add(relevant_tilt_series) - - instrument_name = ( - murfey_db.exec( - select(db.Session).where(db.Session.id == session_id) - ) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - tilts = get_all_tilts(relevant_tilt_series.id) - ids = get_job_ids(relevant_tilt_series.id, alignment_ids[2].id) - preproc_params = get_tomo_proc_params(ids.dcgid) - stack_file = ( - Path(message["mrc_out"]).parents[3] - / "Tomograms" - / "job006" - / "tomograms" - / f"{relevant_tilt_series.tag}_stack.mrc" - ) - if not stack_file.parent.exists(): - stack_file.parent.mkdir(parents=True) - tilt_offset = midpoint([float(get_angle(t)) for t in tilts]) - zocalo_message = { - "recipes": ["em-tomo-align"], - "parameters": { - "input_file_list": str([[t, str(get_angle(t))] for t in tilts]), - "path_pattern": "", # blank for now so that it works with the tomo_align service changes - "dcid": ids.dcid, - "appid": ids.appid, - "stack_file": str(stack_file), - "dose_per_frame": preproc_params.dose_per_frame, - "frame_count": preproc_params.frame_count, - "kv": preproc_params.voltage, - "tilt_axis": preproc_params.tilt_axis, - "pixel_size": preproc_params.pixel_size, - "manual_tilt_offset": -tilt_offset, - "node_creator_queue": machine_config.node_creator_queue, - }, - } - if _transport_object: - logger.info( - f"Sending Zocalo message for processing: {zocalo_message}" - ) - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - else: - logger.info( - f"No transport object found. Zocalo message would be {zocalo_message}" - ) - - prom.preprocessed_movies.labels(processing_job=collected_ids[2].id).inc() - murfey_db.commit() - murfey_db.close() - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "data_collection_group": - ispyb_session_id = get_session_id( - microscope=message["microscope"], - proposal_code=message["proposal_code"], - proposal_number=message["proposal_number"], - visit_number=message["visit_number"], - db=ISPyBSession(), - ) - if dcg_murfey := murfey_db.exec( - select(db.DataCollectionGroup) - .where(db.DataCollectionGroup.session_id == message["session_id"]) - .where(db.DataCollectionGroup.tag == message.get("tag")) - ).all(): - dcgid = dcg_murfey[0].id - else: - if ispyb_session_id is None: - murfey_dcg = db.DataCollectionGroup( - session_id=message["session_id"], - tag=message.get("tag"), - ) - else: - record = DataCollectionGroup( - sessionId=ispyb_session_id, - experimentType=message["experiment_type"], - experimentTypeId=message["experiment_type_id"], - ) - dcgid = _register(record, header) - atlas_record = Atlas( - dataCollectionGroupId=dcgid, - atlasImage=message.get("atlas", ""), - pixelSize=message.get("atlas_pixel_size", 0), - cassetteSlot=message.get("sample"), - ) - if _transport_object: - atlas_id = _transport_object.do_insert_atlas(atlas_record)[ - "return_value" - ] - murfey_dcg = db.DataCollectionGroup( - id=dcgid, - atlas_id=atlas_id, - session_id=message["session_id"], - tag=message.get("tag"), - ) - murfey_db.add(murfey_dcg) - murfey_db.commit() - murfey_db.close() - if _transport_object: - if dcgid is None: - time.sleep(2) - _transport_object.transport.nack(header, requeue=True) - return None - _transport_object.transport.ack(header) - if dcg_hooks := entry_points().select( - group="murfey.hooks", name="data_collection_group" - ): - try: - for hook in dcg_hooks: - hook.load()(dcgid, session_id=message["session_id"]) - except Exception: - logger.error( - "Call to data collection group hook failed", exc_info=True - ) - return None - elif message["register"] == "atlas_update": - if _transport_object: - _transport_object.do_update_atlas( - message["atlas_id"], - message["atlas"], - message["atlas_pixel_size"], - message["sample"], - ) - _transport_object.transport.ack(header) - return None - elif message["register"] == "data_collection": - logger.debug( - "Received message named 'data_collection' containing the following items:\n" - f"{', '.join([f'{sanitise(key)}: {sanitise(str(value))}' for key, value in message.items()])}" - ) - murfey_session_id = message["session_id"] - ispyb_session_id = get_session_id( - microscope=message["microscope"], - proposal_code=message["proposal_code"], - proposal_number=message["proposal_number"], - visit_number=message["visit_number"], - db=ISPyBSession(), - ) - dcg = murfey_db.exec( - select(db.DataCollectionGroup) - .where(db.DataCollectionGroup.session_id == murfey_session_id) - .where(db.DataCollectionGroup.tag == message["source"]) - ).all() - if dcg: - dcgid = dcg[0].id - # flush_data_collections(message["source"], murfey_db) - else: - logger.warning( - "No data collection group ID was found for image directory " - f"{sanitise(message['image_directory'])} and source " - f"{sanitise(message['source'])}" - ) - if _transport_object: - _transport_object.transport.nack(header, requeue=True) - return None - if dc_murfey := murfey_db.exec( - select(db.DataCollection) - .where(db.DataCollection.tag == message.get("tag")) - .where(db.DataCollection.dcg_id == dcgid) - ).all(): - dcid = dc_murfey[0].id - else: - if ispyb_session_id is None: - murfey_dc = db.DataCollection( - tag=message.get("tag"), - dcg_id=dcgid, - ) - else: - record = DataCollection( - SESSIONID=ispyb_session_id, - experimenttype=message["experiment_type"], - imageDirectory=message["image_directory"], - imageSuffix=message["image_suffix"], - voltage=message["voltage"], - dataCollectionGroupId=dcgid, - pixelSizeOnImage=message["pixel_size"], - imageSizeX=message["image_size_x"], - imageSizeY=message["image_size_y"], - slitGapHorizontal=message.get("slit_width"), - magnification=message.get("magnification"), - exposureTime=message.get("exposure_time"), - totalExposedDose=message.get("total_exposed_dose"), - c2aperture=message.get("c2aperture"), - phasePlate=int(message.get("phase_plate", 0)), - ) - dcid = _register( - record, - header, - tag=( - message.get("tag") - if message["experiment_type"] == "tomography" - else "" - ), - ) - murfey_dc = db.DataCollection( - id=dcid, - tag=message.get("tag"), - dcg_id=dcgid, - ) - murfey_db.add(murfey_dc) - murfey_db.commit() - dcid = murfey_dc.id - murfey_db.close() - if dcid is None and _transport_object: - _transport_object.transport.nack(header, requeue=True) - return None - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "processing_job": - murfey_session_id = message["session_id"] - logger.info("registering processing job") - dc = murfey_db.exec( - select(db.DataCollection, db.DataCollectionGroup) - .where(db.DataCollection.dcg_id == db.DataCollectionGroup.id) - .where(db.DataCollectionGroup.session_id == murfey_session_id) - .where(db.DataCollectionGroup.tag == message["source"]) - .where(db.DataCollection.tag == message["tag"]) - ).all() - if dc: - _dcid = dc[0][0].id - else: - logger.warning( - f"No data collection ID found for {sanitise(message['tag'])}" - ) - if _transport_object: - _transport_object.transport.nack(header, requeue=True) - return None - if pj_murfey := murfey_db.exec( - select(db.ProcessingJob) - .where(db.ProcessingJob.recipe == message["recipe"]) - .where(db.ProcessingJob.dc_id == _dcid) - ).all(): - pid = pj_murfey[0].id - else: - if ISPyBSession() is None: - murfey_pj = db.ProcessingJob(recipe=message["recipe"], dc_id=_dcid) - else: - record = ProcessingJob( - dataCollectionId=_dcid, recipe=message["recipe"] - ) - run_parameters = message.get("parameters", {}) - assert isinstance(run_parameters, dict) - if message.get("job_parameters"): - job_parameters = [ - ProcessingJobParameter(parameterKey=k, parameterValue=v) - for k, v in message["job_parameters"].items() - ] - pid = _register(ExtendedRecord(record, job_parameters), header) - else: - pid = _register(record, header) - murfey_pj = db.ProcessingJob( - id=pid, recipe=message["recipe"], dc_id=_dcid - ) - murfey_db.add(murfey_pj) - murfey_db.commit() - pid = murfey_pj.id - murfey_db.close() - if pid is None and _transport_object: - _transport_object.transport.nack(header, requeue=True) - return None - prom.preprocessed_movies.labels(processing_job=pid) - if not murfey_db.exec( - select(db.AutoProcProgram).where(db.AutoProcProgram.pj_id == pid) - ).all(): - if ISPyBSession() is None: - murfey_app = db.AutoProcProgram(pj_id=pid) - else: - record = AutoProcProgram( - processingJobId=pid, processingStartTime=datetime.now() - ) - appid = _register(record, header) - if appid is None and _transport_object: - _transport_object.transport.nack(header, requeue=True) - return None - murfey_app = db.AutoProcProgram(id=appid, pj_id=pid) - murfey_db.add(murfey_app) - murfey_db.commit() - murfey_db.close() - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "flush_tomography_preprocess": - _flush_tomography_preprocessing(message) - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "spa_processing_parameters": - session_id = message["session_id"] - collected_ids = murfey_db.exec( - select( - db.DataCollectionGroup, - db.DataCollection, - db.ProcessingJob, - db.AutoProcProgram, - ) - .where(db.DataCollectionGroup.session_id == session_id) - .where(db.DataCollectionGroup.tag == message["tag"]) - .where(db.DataCollection.dcg_id == db.DataCollectionGroup.id) - .where(db.ProcessingJob.dc_id == db.DataCollection.id) - .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id) - .where(db.ProcessingJob.recipe == "em-spa-preprocess") - ).one() - pj_id = collected_ids[2].id - if not murfey_db.exec( - select(db.SPARelionParameters).where( - db.SPARelionParameters.pj_id == pj_id - ) - ).all(): - instrument_name = ( - murfey_db.exec( - select(db.Session).where(db.Session.id == session_id) - ) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - params = db.SPARelionParameters( - pj_id=collected_ids[2].id, - angpix=float(message["pixel_size_on_image"]) * 1e10, - dose_per_frame=message["dose_per_frame"], - gain_ref=( - str(machine_config.rsync_basepath / message["gain_ref"]) - if message["gain_ref"] and machine_config.data_transfer_enabled - else message["gain_ref"] - ), - voltage=message["voltage"], - motion_corr_binning=message["motion_corr_binning"], - eer_fractionation_file=message["eer_fractionation_file"], - symmetry=message["symmetry"], - particle_diameter=message["particle_diameter"], - downscale=message["downscale"], - boxsize=message["boxsize"], - small_boxsize=message["small_boxsize"], - mask_diameter=message["mask_diameter"], - ) - feedback_params = db.SPAFeedbackParameters( - pj_id=collected_ids[2].id, - estimate_particle_diameter=not bool(message["particle_diameter"]), - hold_class2d=False, - hold_class3d=False, - class_selection_score=0, - star_combination_job=0, - initial_model="", - next_job=0, - ) - murfey_db.add(params) - murfey_db.add(feedback_params) - murfey_db.commit() - logger.info( - f"SPA processing parameters registered for processing job {collected_ids[2].id}" - ) - murfey_db.close() - else: - logger.info( - f"SPA processing parameters already exist for processing job ID {pj_id}" - ) - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "tomography_processing_parameters": - session_id = message["session_id"] - collected_ids = murfey_db.exec( - select( - db.DataCollectionGroup, - db.DataCollection, - db.ProcessingJob, - db.AutoProcProgram, - ) - .where(db.DataCollectionGroup.session_id == session_id) - .where(db.DataCollectionGroup.tag == message["tag"]) - .where(db.DataCollection.dcg_id == db.DataCollectionGroup.id) - .where(db.DataCollection.tag == message["tilt_series_tag"]) - .where(db.ProcessingJob.dc_id == db.DataCollection.id) - .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id) - .where(db.ProcessingJob.recipe == "em-tomo-preprocess") - ).one() - if not murfey_db.exec( - select(func.count(db.TomographyProcessingParameters.dcg_id)).where( - db.TomographyProcessingParameters.dcg_id == collected_ids[0].id - ) - ).one(): - params = db.TomographyProcessingParameters( - dcg_id=collected_ids[0].id, - pixel_size=float(message["pixel_size_on_image"]) * 10**10, - voltage=message["voltage"], - dose_per_frame=message["dose_per_frame"], - frame_count=message["frame_count"], - tilt_axis=message["tilt_axis"], - motion_corr_binning=message["motion_corr_binning"], - gain_ref=message["gain_ref"], - eer_fractionation_file=message["eer_fractionation_file"], - ) - murfey_db.add(params) - murfey_db.commit() - murfey_db.close() - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "done_incomplete_2d_batch": - _release_2d_hold(message) - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "incomplete_particles_file": - _register_incomplete_2d_batch(message) - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "complete_particles_file": - _register_complete_2d_batch(message) - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "save_class_selection_score": - _register_class_selection(message) - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "done_3d_batch": - _release_3d_hold(message) - if message.get("do_refinement"): - _register_refinement(message) - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "run_class3d": - _register_3d_batch(message) - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "save_initial_model": - _register_initial_model(message) - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "done_particle_selection": - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "done_class_selection": - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "atlas_registered": - _flush_grid_square_records(message) - if _transport_object: - _transport_object.transport.ack(header) - return None - elif message["register"] == "done_refinement": - bfactors_registered = _register_bfactors(message) - if _transport_object: - if bfactors_registered: - _transport_object.transport.ack(header) - else: - _transport_object.transport.nack(header) - return None - elif message["register"] == "done_bfactor": - _save_bfactor(message) - if _transport_object: - _transport_object.transport.ack(header) - return None - elif ( - message["register"] in entry_points().select(group="murfey.workflows").names - ): - # Search for corresponding workflow - workflows: list[EntryPoint] = list( - entry_points().select( - group="murfey.workflows", name=message["register"] - ) - ) # Returns either 1 item or empty list - if not workflows: - logger.error(f"No workflow found for {sanitise(message['register'])}") - if _transport_object: - _transport_object.transport.nack(header, requeue=False) - return None - # Run the workflow if a match is found - workflow: EntryPoint = workflows[0] - result = workflow.load()( - message=message, - murfey_db=murfey_db, - ) - if _transport_object: - if result: - _transport_object.transport.ack(header) - else: - # Send it directly to DLQ without trying to rerun it - _transport_object.transport.nack(header, requeue=False) - if not result: - logger.error( - f"Workflow {sanitise(message['register'])} returned {result}" - ) - return None - logger.error(f"No workflow found for {sanitise(message['register'])}") - if _transport_object: - _transport_object.transport.nack(header, requeue=False) - return None - except PendingRollbackError: - murfey_db.rollback() - murfey_db.close() - logger.warning("Murfey database required a rollback") - if _transport_object: - _transport_object.transport.nack(header, requeue=True) - except OperationalError: - logger.warning("Murfey database error encountered", exc_info=True) - time.sleep(1) - if _transport_object: - _transport_object.transport.nack(header, requeue=True) - except Exception: - logger.warning( - "Exception encountered in server RabbitMQ callback", exc_info=True - ) - if _transport_object: - _transport_object.transport.nack(header, requeue=False) - return None - - -@singledispatch -def _register(record, header: dict, **kwargs): - raise NotImplementedError(f"Not method to register {record} or type {type(record)}") - - -@_register.register -def _(record: Base, header: dict, **kwargs): # type: ignore - if not _transport_object: - logger.error( - f"No transport object found when processing record {record}. Message header: {header}" - ) - return None - try: - if isinstance(record, DataCollection): - return _transport_object.do_insert_data_collection(record, **kwargs)[ - "return_value" - ] - if isinstance(record, DataCollectionGroup): - return _transport_object.do_insert_data_collection_group(record)[ - "return_value" - ] - if isinstance(record, ProcessingJob): - return _transport_object.do_create_ispyb_job(record)["return_value"] - if isinstance(record, AutoProcProgram): - return _transport_object.do_update_processing_status(record)["return_value"] - # session = Session() - # session.add(record) - # session.commit() - # _transport_object.transport.ack(header, requeue=False) - return getattr(record, record.__table__.primary_key.columns[0].name) - - except SQLAlchemyError as e: - logger.error(f"Murfey failed to insert ISPyB record {record}", e, exc_info=True) - # _transport_object.transport.nack(header) - return None - except AttributeError as e: - logger.error( - f"Murfey could not find primary key when inserting record {record}", - e, - exc_info=True, - ) - return None - - -@_register.register -def _(extended_record: ExtendedRecord, header: dict, **kwargs): - if not _transport_object: - raise ValueError( - "Transport object should not be None if a database record is being updated" - ) - return _transport_object.do_create_ispyb_job( - extended_record.record, params=extended_record.record_params - )["return_value"] - - -def feedback_listen(): - if _transport_object: - if not _transport_object.feedback_queue: - _transport_object.feedback_queue = ( - _transport_object.transport._subscribe_temporary( - channel_hint="", callback=None, sub_id=None - ) - ) - _transport_object._connection_callback = partial( - _transport_object.transport.subscribe, - _transport_object.feedback_queue, - feedback_callback, - acknowledgement=True, - ) - _transport_object.transport.subscribe( - _transport_object.feedback_queue, feedback_callback, acknowledgement=True - ) +template_files = files("murfey") / "templates" diff --git a/src/murfey/server/api/__init__.py b/src/murfey/server/api/__init__.py index 65e64bc61..69896bbac 100644 --- a/src/murfey/server/api/__init__.py +++ b/src/murfey/server/api/__init__.py @@ -1,2053 +1,5 @@ -from __future__ import annotations +from fastapi.templating import Jinja2Templates -import asyncio -import datetime -import logging -import os -from functools import lru_cache -from pathlib import Path -from typing import Dict, List, Optional +from murfey.server import template_files -import sqlalchemy -from fastapi import APIRouter, Depends, Request -from fastapi.responses import FileResponse, HTMLResponse -from ispyb.sqlalchemy import Atlas -from ispyb.sqlalchemy import AutoProcProgram as ISPyBAutoProcProgram -from ispyb.sqlalchemy import ( - BLSample, - BLSampleGroup, - BLSampleImage, - BLSession, - BLSubSample, - Proposal, -) -from PIL import Image -from prometheus_client import Counter, Gauge -from pydantic import BaseModel -from sqlalchemy import func -from sqlalchemy.exc import OperationalError -from sqlmodel import col, select -from werkzeug.utils import secure_filename - -import murfey.server.ispyb -import murfey.server.prometheus as prom -import murfey.server.websocket as ws -import murfey.util.eer -from murfey.server import ( - _murfey_id, - _transport_object, - check_tilt_series_mc, - get_all_tilts, - get_angle, - get_hostname, - get_job_ids, - get_machine_config, - get_microscope, - get_tomo_proc_params, - sanitise, - templates, -) -from murfey.server.api.auth import MurfeySessionID, validate_token -from murfey.server.api.spa import _cryolo_model_path -from murfey.server.gain import Camera, prepare_eer_gain, prepare_gain -from murfey.server.murfey_db import murfey_db -from murfey.util import safe_run, secure_path -from murfey.util.config import MachineConfig, from_file, settings -from murfey.util.db import ( - AutoProcProgram, - ClientEnvironment, - DataCollection, - DataCollectionGroup, - FoilHole, - GridSquare, - MagnificationLookup, - Movie, - PreprocessStash, - ProcessingJob, - RsyncInstance, - Session, - SessionProcessingParameters, - SPAFeedbackParameters, - SPARelionParameters, - Tilt, - TiltSeries, -) -from murfey.util.models import ( - BLSampleImageParameters, - BLSampleParameters, - BLSubSampleParameters, - ClientInfo, - CurrentGainRef, - DCGroupParameters, - DCParameters, - FoilHoleParameters, - FractionationParameters, - GainReference, - GridSquareParameters, - MillingParameters, - PostInfo, - ProcessingJobParameters, - ProcessingParametersSPA, - ProcessingParametersTomo, - RegistrationMessage, - RsyncerInfo, - RsyncerSource, - Sample, - SessionInfo, - SPAProcessFile, - SuggestedPathParameters, - TiltInfo, - TiltSeriesGroupInfo, - TiltSeriesInfo, - TomoProcessFile, - Visit, -) -from murfey.util.processing_params import default_spa_parameters -from murfey.util.tomo import midpoint -from murfey.workflows.spa.flush_spa_preprocess import ( - register_foil_hole, - register_grid_square, -) - -log = logging.getLogger("murfey.server.api") - -router = APIRouter(dependencies=[Depends(validate_token)]) - - -# This will be the homepage for a given microscope. -@router.get("/", response_class=HTMLResponse) -async def root(request: Request): - return templates.TemplateResponse( - request=request, - name="home.html", - context={ - "hostname": get_hostname(), - "microscope": get_microscope(), - "version": murfey.__version__, - }, - ) - - -@router.get("/time") -async def get_current_timestamp(): - return {"timestamp": datetime.datetime.now().timestamp()} - - -@router.get("/health/") -def health_check(db=murfey.server.ispyb.DB): - conn = db.connection() - conn.close() - return { - "ispyb_connection": True, - "rabbitmq_connection": _transport_object.transport.is_connected(), - } - - -@router.get("/connections/") -def connections_check(): - return {"connections": list(ws.manager.active_connections.keys())} - - -@router.get("/machine") -def machine_info() -> Optional[MachineConfig]: - instrument_name = os.getenv("BEAMLINE") - if settings.murfey_machine_configuration and instrument_name: - return from_file(Path(settings.murfey_machine_configuration), instrument_name)[ - instrument_name - ] - return None - - -@lru_cache(maxsize=5) -@router.get("/instruments/{instrument_name}/machine") -def machine_info_by_name(instrument_name: str) -> Optional[MachineConfig]: - if settings.murfey_machine_configuration: - return from_file(Path(settings.murfey_machine_configuration), instrument_name)[ - instrument_name - ] - return None - - -@router.get("/mag_table/") -def get_mag_table(db=murfey_db) -> List[MagnificationLookup]: - return db.exec(select(MagnificationLookup)).all() - - -@router.post("/mag_table/") -def add_to_mag_table(rows: List[MagnificationLookup], db=murfey_db): - for r in rows: - db.add(r) - db.commit() - - -@router.delete("/mag_table/{mag}") -def remove_mag_table_row(mag: int, db=murfey_db): - row = db.exec( - select(MagnificationLookup).where(MagnificationLookup.magnification == mag) - ).one() - db.delete(row) - db.commit() - - -@router.get("/instruments/{instrument_name}/instrument_name") -def get_instrument_display_name(instrument_name: str) -> str: - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - if machine_config: - return machine_config.display_name - return "" - - -@router.get("/instruments/{instrument_name}/visits/") -def all_visit_info(instrument_name: str, request: Request, db=murfey.server.ispyb.DB): - visits = murfey.server.ispyb.get_all_ongoing_visits(instrument_name, db) - - if visits: - return_query = [ - { - "Start date": visit.start, - "End date": visit.end, - "Visit name": visit.name, - "Time remaining": str(visit.end - datetime.datetime.now()), - } - for visit in visits - ] # "Proposal title": visit.proposal_title - log.debug( - f"{len(visits)} visits active for {sanitise(instrument_name)=}: {', '.join(v.name for v in visits)}" - ) - return templates.TemplateResponse( - request=request, - name="activevisits.html", - context={"info": return_query, "microscope": instrument_name}, - ) - else: - log.debug(f"No visits identified for {sanitise(instrument_name)=}") - return templates.TemplateResponse( - request=request, - name="activevisits.html", - context={"info": [], "microscope": instrument_name}, - ) - - -@router.post("/visits/{visit_name}") -def register_client_to_visit(visit_name: str, client_info: ClientInfo, db=murfey_db): - client_env = db.exec( - select(ClientEnvironment).where(ClientEnvironment.client_id == client_info.id) - ).one() - session = db.exec(select(Session).where(Session.id == client_env.session_id)).one() - if client_env: - client_env.visit = visit_name - db.add(client_env) - db.commit() - if session: - session.visit = visit_name - db.add(session) - db.commit() - db.close() - return client_info - - -@router.get("/num_movies") -def count_number_of_movies(db=murfey_db) -> Dict[str, int]: - res = db.exec( - select(Movie.tag, func.count(Movie.murfey_id)).group_by(Movie.tag) - ).all() - return {r[0]: r[1] for r in res} - - -@router.post("/sessions/{session_id}/rsyncer") -def register_rsyncer(session_id: int, rsyncer_info: RsyncerInfo, db=murfey_db): - visit_name = db.exec(select(Session).where(Session.id == session_id)).one().visit - rsync_instance = RsyncInstance( - source=rsyncer_info.source, - session_id=rsyncer_info.session_id, - transferring=rsyncer_info.transferring, - destination=rsyncer_info.destination, - tag=rsyncer_info.tag, - ) - db.add(rsync_instance) - db.commit() - db.close() - prom.seen_files.labels(rsync_source=rsyncer_info.source, visit=visit_name) - prom.seen_data_files.labels(rsync_source=rsyncer_info.source, visit=visit_name) - prom.transferred_files.labels(rsync_source=rsyncer_info.source, visit=visit_name) - prom.transferred_files_bytes.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ) - prom.transferred_data_files.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ) - prom.transferred_data_files_bytes.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ) - prom.seen_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).set(0) - prom.transferred_files.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).set(0) - prom.transferred_files_bytes.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).set(0) - prom.seen_data_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).set( - 0 - ) - prom.transferred_data_files.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).set(0) - prom.transferred_data_files_bytes.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).set(0) - return rsyncer_info - - -@router.delete("/sessions/{session_id}/rsyncer") -def delete_rsyncer(session_id: int, source: Path, db=murfey_db): - try: - rsync_instance = db.exec( - select(RsyncInstance) - .where(RsyncInstance.session_id == session_id) - .where(RsyncInstance.source == str(source)) - ).one() - db.delete(rsync_instance) - db.commit() - except Exception: - log.error( - f"Failed to delete rsyncer for source directory {sanitise(str(source))!r} " - f"in session {session_id}.", - exc_info=True, - ) - - -@router.post("/sessions/{session_id}/rsyncer_stopped") -def register_stopped_rsyncer( - session_id: int, rsyncer_source: RsyncerSource, db=murfey_db -): - rsyncer = db.exec( - select(RsyncInstance) - .where(RsyncInstance.session_id == session_id) - .where(RsyncInstance.source == rsyncer_source.source) - ).one() - rsyncer.transferring = False - db.add(rsyncer) - db.commit() - - -@router.post("/sessions/{session_id}/rsyncer_started") -def register_restarted_rsyncer( - session_id: int, rsyncer_source: RsyncerSource, db=murfey_db -): - rsyncer = db.exec( - select(RsyncInstance) - .where(RsyncInstance.session_id == session_id) - .where(RsyncInstance.source == rsyncer_source.source) - ).one() - rsyncer.transferring = True - db.add(rsyncer) - db.commit() - - -@router.get("/sessions/{session_id}/rsyncers", response_model=List[RsyncInstance]) -def get_rsyncers_for_client(session_id: MurfeySessionID, db=murfey_db): - rsync_instances = db.exec( - select(RsyncInstance).where(RsyncInstance.session_id == session_id) - ) - return rsync_instances.all() - - -class SessionClients(BaseModel): - session: Session - clients: List[ClientEnvironment] - - -@router.get("/session/{session_id}") -async def get_session(session_id: MurfeySessionID, db=murfey_db) -> SessionClients: - session = db.exec(select(Session).where(Session.id == session_id)).one() - clients = db.exec( - select(ClientEnvironment).where(ClientEnvironment.session_id == session_id) - ).all() - return SessionClients(session=session, clients=clients) - - -@router.post("/visits/{visit_name}/increment_rsync_file_count") -def increment_rsync_file_count( - visit_name: str, rsyncer_info: RsyncerInfo, db=murfey_db -): - try: - rsync_instance = db.exec( - select(RsyncInstance).where( - RsyncInstance.source == rsyncer_info.source, - RsyncInstance.destination == rsyncer_info.destination, - RsyncInstance.session_id == rsyncer_info.session_id, - ) - ).one() - except Exception: - log.error( - f"Failed to find rsync instance for visit {sanitise(visit_name)} " - "with the following properties: \n" - f"{rsyncer_info.dict()}", - exc_info=True, - ) - return None - rsync_instance.files_counted += rsyncer_info.increment_count - db.add(rsync_instance) - db.commit() - db.close() - prom.seen_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).inc( - rsyncer_info.increment_count - ) - prom.seen_data_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).inc( - rsyncer_info.increment_data_count - ) - - -@router.post("/visits/{visit_name}/increment_rsync_transferred_files") -def increment_rsync_transferred_files( - visit_name: str, rsyncer_info: RsyncerInfo, db=murfey_db -): - rsync_instance = db.exec( - select(RsyncInstance).where( - RsyncInstance.source == rsyncer_info.source, - RsyncInstance.destination == rsyncer_info.destination, - RsyncInstance.session_id == rsyncer_info.session_id, - ) - ).one() - rsync_instance.files_transferred += rsyncer_info.increment_count - db.add(rsync_instance) - db.commit() - db.close() - - -@router.post("/visits/{visit_name}/increment_rsync_transferred_files_prometheus") -def increment_rsync_transferred_files_prometheus( - visit_name: str, rsyncer_info: RsyncerInfo, db=murfey_db -): - prom.transferred_files.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).inc(rsyncer_info.increment_count) - prom.transferred_files_bytes.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).inc(rsyncer_info.bytes) - prom.transferred_data_files.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).inc(rsyncer_info.increment_data_count) - prom.transferred_data_files_bytes.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).inc(rsyncer_info.data_bytes) - - -class ProcessingDetails(BaseModel): - data_collection_group: DataCollectionGroup - data_collections: List[DataCollection] - processing_jobs: List[ProcessingJob] - relion_params: SPARelionParameters - feedback_params: SPAFeedbackParameters - - -@router.get("/sessions/{session_id}/spa_processing_parameters") -def get_spa_proc_param_details( - session_id: MurfeySessionID, db=murfey_db -) -> Optional[List[ProcessingDetails]]: - params = db.exec( - select( - DataCollectionGroup, - DataCollection, - ProcessingJob, - SPARelionParameters, - SPAFeedbackParameters, - ) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollectionGroup.id == DataCollection.dcg_id) - .where(DataCollection.id == ProcessingJob.dc_id) - .where(SPARelionParameters.pj_id == ProcessingJob.id) - .where(SPAFeedbackParameters.pj_id == ProcessingJob.id) - ).all() - if not params: - return None - unique_dcg_indices = [] - dcg_ids = [] - for i, p in enumerate(params): - if p[0].id not in dcg_ids: - dcg_ids.append(p[0].id) - unique_dcg_indices.append(i) - - def _parse(ps, i, dcg_id): - res = [] - for p in ps: - if p[0].id == dcg_id: - if p[i] not in res: - res.append(p[i]) - return res - - return [ - ProcessingDetails( - data_collection_group=params[i][0], - data_collections=_parse(params, 1, d), - processing_jobs=_parse(params, 2, d), - relion_params=_parse(params, 3, d)[0], - feedback_params=_parse(params, 4, d)[0], - ) - for i, d in zip(unique_dcg_indices, dcg_ids) - ] - - -@router.post("/sessions/{session_id}/spa_processing_parameters") -def register_spa_proc_params( - session_id: MurfeySessionID, proc_params: ProcessingParametersSPA, db=murfey_db -): - session_processing_parameters = db.exec( - select(SessionProcessingParameters).where( - SessionProcessingParameters.session_id == session_id - ) - ).all() - if session_processing_parameters: - proc_params.gain_ref = session_processing_parameters[0].gain_ref - proc_params.dose_per_frame = session_processing_parameters[0].dose_per_frame - proc_params.eer_fractionation_file = session_processing_parameters[ - 0 - ].eer_fractionation_file - proc_params.symmetry = session_processing_parameters[0].symmetry - - zocalo_message = { - "register": "spa_processing_parameters", - **dict(proc_params), - "session_id": session_id, - } - if _transport_object: - _transport_object.send(_transport_object.feedback_queue, zocalo_message) - - -@router.get("/sessions/{session_id}/grid_squares") -def get_grid_squares(session_id: MurfeySessionID, db=murfey_db): - grid_squares = db.exec( - select(GridSquare).where(GridSquare.session_id == session_id) - ).all() - tags = {gs.tag for gs in grid_squares} - res = {} - for t in tags: - res[t] = [gs for gs in grid_squares if gs.tag == t] - return res - - -@router.get("/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares") -def get_grid_squares_from_dcg( - session_id: int, dcgid: int, db=murfey_db -) -> List[GridSquare]: - grid_squares = db.exec( - select(GridSquare, DataCollectionGroup) - .where(GridSquare.session_id == session_id) - .where(GridSquare.tag == DataCollectionGroup.tag) - .where(DataCollectionGroup.id == dcgid) - ).all() - return [gs[0] for gs in grid_squares] - - -@router.get( - "/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares/{gsid}/num_movies" -) -def get_number_of_movies_from_grid_square( - session_id: int, dcgid: int, gsid: int, db=murfey_db -) -> int: - movies = db.exec( - select(Movie, FoilHole, GridSquare, DataCollectionGroup) - .where(Movie.foil_hole_id == FoilHole.id) - .where(FoilHole.grid_square_id == GridSquare.id) - .where(GridSquare.name == gsid) - .where(GridSquare.session_id == session_id) - .where(GridSquare.tag == DataCollectionGroup.tag) - .where(DataCollectionGroup.id == dcgid) - ).all() - return len(movies) - - -@router.get( - "/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares/{gsid}/foil_holes" -) -def get_foil_holes_from_grid_square( - session_id: int, dcgid: int, gsid: int, db=murfey_db -) -> List[FoilHole]: - foil_holes = db.exec( - select(FoilHole, GridSquare, DataCollectionGroup) - .where(FoilHole.grid_square_id == GridSquare.id) - .where(GridSquare.name == gsid) - .where(GridSquare.session_id == session_id) - .where(GridSquare.tag == DataCollectionGroup.tag) - .where(DataCollectionGroup.id == dcgid) - ).all() - return [fh[0] for fh in foil_holes] - - -@router.get( - "/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares/{gsid}/foil_holes/{fhid}/num_movies" -) -def get_number_of_movies_from_foil_hole( - session_id: int, dcgid: int, gsid: int, fhid: int, db=murfey_db -) -> int: - movies = db.exec( - select(Movie, FoilHole, GridSquare, DataCollectionGroup) - .where(Movie.foil_hole_id == FoilHole.id) - .where(FoilHole.name == fhid) - .where(FoilHole.grid_square_id == GridSquare.id) - .where(GridSquare.name == gsid) - .where(GridSquare.session_id == session_id) - .where(GridSquare.tag == DataCollectionGroup.tag) - .where(DataCollectionGroup.id == dcgid) - ).all() - return len(movies) - - -@router.post("/sessions/{session_id}/grid_square/{gsid}") -def posted_grid_square( - session_id: MurfeySessionID, - gsid: int, - grid_square_params: GridSquareParameters, - db=murfey_db, -): - return register_grid_square(session_id, gsid, grid_square_params, db) - - -@router.get("/sessions/{session_id}/foil_hole/{fh_name}") -def get_foil_hole( - session_id: MurfeySessionID, fh_name: int, db=murfey_db -) -> Dict[str, int]: - foil_holes = db.exec( - select(FoilHole, GridSquare) - .where(FoilHole.name == fh_name) - .where(FoilHole.session_id == session_id) - .where(GridSquare.id == FoilHole.grid_square_id) - ).all() - return {f[1].tag: f[0].id for f in foil_holes} - - -@router.post("/sessions/{session_id}/grid_square/{gs_name}/foil_hole") -def post_foil_hole( - session_id: MurfeySessionID, - gs_name: int, - foil_hole_params: FoilHoleParameters, - db=murfey_db, -): - log.info( - f"Registering foil hole {foil_hole_params.name} with position {(foil_hole_params.x_location, foil_hole_params.y_location)}" - ) - return register_foil_hole(session_id, gs_name, foil_hole_params, db) - - -@router.post("/sessions/{session_id}/tomography_processing_parameters") -def register_tomo_proc_params( - session_id: MurfeySessionID, proc_params: ProcessingParametersTomo, db=murfey_db -): - session_processing_parameters = db.exec( - select(SessionProcessingParameters).where( - SessionProcessingParameters.session_id == session_id - ) - ).all() - if session_processing_parameters: - proc_params.gain_ref = session_processing_parameters[0].gain_ref - proc_params.dose_per_frame = session_processing_parameters[0].dose_per_frame - proc_params.eer_fractionation_file = session_processing_parameters[ - 0 - ].eer_fractionation_file - - zocalo_message = { - "register": "tomography_processing_parameters", - **dict(proc_params), - "session_id": session_id, - } - if _transport_object: - _transport_object.send(_transport_object.feedback_queue, zocalo_message) - - -class Tag(BaseModel): - tag: str - - -@router.post("/visits/{visit_name}/{session_id}/flush_spa_processing") -def flush_spa_processing( - visit_name: str, session_id: MurfeySessionID, tag: Tag, db=murfey_db -): - zocalo_message = { - "register": "spa.flush_spa_preprocess", - "session_id": session_id, - "tag": tag.tag, - } - if _transport_object: - _transport_object.send(_transport_object.feedback_queue, zocalo_message) - return - - -class Source(BaseModel): - rsync_source: str - - -@router.post("/visits/{visit_name}/{session_id}/flush_tomography_processing") -def flush_tomography_processing( - visit_name: str, session_id: MurfeySessionID, rsync_source: Source, db=murfey_db -): - zocalo_message = { - "register": "flush_tomography_preprocess", - "session_id": session_id, - "visit_name": visit_name, - "data_collection_group_tag": rsync_source.rsync_source, - } - if _transport_object: - _transport_object.send(_transport_object.feedback_queue, zocalo_message) - return - - -@router.post("/visits/{visit_name}/tilt_series") -def register_tilt_series( - visit_name: str, tilt_series_info: TiltSeriesInfo, db=murfey_db -): - session_id = tilt_series_info.session_id - if db.exec( - select(TiltSeries) - .where(TiltSeries.session_id == session_id) - .where(TiltSeries.tag == tilt_series_info.tag) - .where(TiltSeries.rsync_source == tilt_series_info.source) - ).all(): - return - tilt_series = TiltSeries( - session_id=session_id, - tag=tilt_series_info.tag, - rsync_source=tilt_series_info.source, - ) - db.add(tilt_series) - db.commit() - - -@router.post("/sessions/{session_id}/tilt_series_length") -def register_tilt_series_length( - session_id: int, - tilt_series_group: TiltSeriesGroupInfo, - db=murfey_db, -): - tilt_series_db = db.exec( - select(TiltSeries) - .where(col(TiltSeries.tag).in_(tilt_series_group.tags)) - .where(TiltSeries.session_id == session_id) - .where(TiltSeries.rsync_source == tilt_series_group.source) - ).all() - for ts in tilt_series_db: - ts_index = tilt_series_group.tags.index(ts.tag) - ts.tilt_series_length = tilt_series_group.tilt_series_lengths[ts_index] - db.add(ts) - db.commit() - - -@router.post("/visits/{visit_name}/{session_id}/completed_tilt_series") -def register_completed_tilt_series( - visit_name: str, - session_id: MurfeySessionID, - tilt_series_group: TiltSeriesGroupInfo, - db=murfey_db, -): - tilt_series_db = db.exec( - select(TiltSeries) - .where(col(TiltSeries.tag).in_(tilt_series_group.tags)) - .where(TiltSeries.session_id == session_id) - .where(TiltSeries.rsync_source == tilt_series_group.source) - ).all() - for ts in tilt_series_db: - ts_index = tilt_series_group.tags.index(ts.tag) - ts.tilt_series_length = tilt_series_group.tilt_series_lengths[ts_index] - db.add(ts) - db.commit() - for ts in tilt_series_db: - if ( - check_tilt_series_mc(ts.id) - and not ts.processing_requested - and ts.tilt_series_length > 2 - ): - ts.processing_requested = True - db.add(ts) - - collected_ids = db.exec( - select( - DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram - ) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollectionGroup.tag == tilt_series_group.source) - .where(DataCollection.tag == ts.tag) - .where(DataCollection.dcg_id == DataCollectionGroup.id) - .where(ProcessingJob.dc_id == DataCollection.id) - .where(AutoProcProgram.pj_id == ProcessingJob.id) - .where(ProcessingJob.recipe == "em-tomo-align") - ).one() - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)) - .one() - .instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - tilts = get_all_tilts(ts.id) - ids = get_job_ids(ts.id, collected_ids[3].id) - preproc_params = get_tomo_proc_params(ids.dcgid) - - first_tilt = db.exec( - select(Tilt).where(Tilt.tilt_series_id == ts.id) - ).first() - parts = [secure_filename(p) for p in Path(first_tilt.movie_path).parts] - visit_idx = parts.index(visit_name) - core = Path(*Path(first_tilt.movie_path).parts[: visit_idx + 1]) - ppath = Path( - "/".join(secure_filename(p) for p in Path(first_tilt.movie_path).parts) - ) - sub_dataset = "/".join(ppath.relative_to(core).parts[:-1]) - extra_path = machine_config.processed_extra_directory - stack_file = ( - core - / machine_config.processed_directory_name - / sub_dataset - / extra_path - / "Tomograms" - / "job006" - / "tomograms" - / f"{ts.tag}_stack.mrc" - ) - if not stack_file.parent.exists(): - stack_file.parent.mkdir(parents=True) - tilt_offset = midpoint([float(get_angle(t)) for t in tilts]) - zocalo_message = { - "recipes": ["em-tomo-align"], - "parameters": { - "input_file_list": str([[t, str(get_angle(t))] for t in tilts]), - "path_pattern": "", # blank for now so that it works with the tomo_align service changes - "dcid": ids.dcid, - "appid": ids.appid, - "stack_file": str(stack_file), - "dose_per_frame": preproc_params.dose_per_frame, - "frame_count": preproc_params.frame_count, - "kv": preproc_params.voltage, - "tilt_axis": preproc_params.tilt_axis, - "pixel_size": preproc_params.pixel_size, - "manual_tilt_offset": -tilt_offset, - "node_creator_queue": machine_config.node_creator_queue, - }, - } - if _transport_object: - log.info(f"Sending Zocalo message for processing: {zocalo_message}") - _transport_object.send( - "processing_recipe", zocalo_message, new_connection=True - ) - else: - log.info( - f"No transport object found. Zocalo message would be {zocalo_message}" - ) - db.commit() - - -@router.post("/visits/{visit_name}/rerun_tilt_series") -def register_tilt_series_for_rerun( - visit_name: str, tilt_series_info: TiltSeriesInfo, db=murfey_db -): - """Set processing to false for cases where an extra tilt is found for a series""" - session_id = tilt_series_info.session_id - tilt_series_db = db.exec( - select(TiltSeries) - .where(TiltSeries.session_id == session_id) - .where(TiltSeries.tag == tilt_series_info.tag) - .where(TiltSeries.rsync_source == tilt_series_info.source) - ).all() - for ts in tilt_series_db: - ts.processing_requested = False - db.add(ts) - db.commit() - - -@router.get("/sessions/{session_id}/tilt_series/{tilt_series_tag}/tilts") -def get_tilts(session_id: MurfeySessionID, tilt_series_tag: str, db=murfey_db): - res = db.exec( - select(TiltSeries, Tilt) - .where(TiltSeries.tag == tilt_series_tag) - .where(TiltSeries.session_id == session_id) - .where(Tilt.tilt_series_id == TiltSeries.id) - ).all() - tilts: Dict[str, List[str]] = {} - for el in res: - if tilts.get(el[1].rsync_source): - tilts[el[1].rsync_source].append(el[2].movie_path) - else: - tilts[el[1].rsync_source] = [el[2].movie_path] - return tilts - - -@router.post("/visits/{visit_name}/{session_id}/tilt") -async def register_tilt( - visit_name: str, session_id: MurfeySessionID, tilt_info: TiltInfo, db=murfey_db -): - def _add_tilt(): - tilt_series_id = ( - db.exec( - select(TiltSeries) - .where(TiltSeries.tag == tilt_info.tilt_series_tag) - .where(TiltSeries.session_id == session_id) - .where(TiltSeries.rsync_source == tilt_info.source) - ) - .one() - .id - ) - if db.exec( - select(Tilt) - .where(Tilt.movie_path == tilt_info.movie_path) - .where(Tilt.tilt_series_id == tilt_series_id) - ).all(): - return - tilt = Tilt(movie_path=tilt_info.movie_path, tilt_series_id=tilt_series_id) - db.add(tilt) - db.commit() - - try: - _add_tilt() - except OperationalError: - await asyncio.sleep(30) - _add_tilt() - - -@router.get("/instruments/{instrument_name}/visits_raw", response_model=List[Visit]) -def get_current_visits(instrument_name: str, db=murfey.server.ispyb.DB): - log.debug( - f"Received request to look up ongoing visits for {sanitise(instrument_name)}" - ) - return murfey.server.ispyb.get_all_ongoing_visits(instrument_name, db) - - -@router.get("/visit/{visit_name}/samples") -def get_samples(visit_name: str, db=murfey.server.ispyb.DB) -> List[Sample]: - return murfey.server.ispyb.get_sub_samples_from_visit(visit_name, db=db) - - -@router.post("/visit/{visit_name}/sample_group") -def register_sample_group(visit_name: str, db=murfey.server.ispyb.DB) -> dict: - proposal_id = murfey.server.ispyb.get_proposal_id( - visit_name[:2], visit_name.split("-")[0][2:], db=db - ) - record = BLSampleGroup(proposalId=proposal_id) - if _transport_object: - return _transport_object.do_insert_sample_group(record) - return {"success": False} - - -@router.post("/visit/{visit_name}/sample") -def register_sample(visit_name: str, sample_params: BLSampleParameters) -> dict: - record = BLSample() - if _transport_object: - return _transport_object.do_insert_sample(record, sample_params.sample_group_id) - return {"success": False} - - -@router.post("/visit/{visit_name}/subsample") -def register_subsample( - visit_name: str, subsample_params: BLSubSampleParameters -) -> dict: - record = BLSubSample( - blSampleId=subsample_params.sample_id, imgFilePath=subsample_params.image_path - ) - if _transport_object: - return _transport_object.do_insert_subsample(record) - return {"success": False} - - -@router.post("/visit/{visit_name}/sample_image") -def register_sample_image( - visit_name: str, sample_image_params: BLSampleImageParameters -) -> dict: - record = BLSampleImage( - blSampleId=sample_image_params.sample_id, - imageFullPath=sample_image_params.image_path, - ) - if _transport_object: - return _transport_object.do_insert_sample_image(record) - return {"success": False} - - -@router.get("/instruments/{instrument_name}/visits/{visit_name}") -def visit_info( - request: Request, instrument_name: str, visit_name: str, db=murfey.server.ispyb.DB -): - query = ( - db.query(BLSession) - .join(Proposal) - .filter( - BLSession.proposalId == Proposal.proposalId, - BLSession.beamLineName == instrument_name, - BLSession.endDate > datetime.datetime.now(), - BLSession.startDate < datetime.datetime.now(), - ) - .add_columns( - BLSession.startDate, - BLSession.endDate, - BLSession.beamLineName, - Proposal.proposalCode, - Proposal.proposalNumber, - BLSession.visit_number, - Proposal.title, - ) - .all() - ) - if query: - return_query = [ - { - "Start date": id.startDate, - "End date": id.endDate, - "Beamline name": id.beamLineName, - "Visit name": visit_name, - "Time remaining": str(id.endDate - datetime.datetime.now()), - } - for id in query - if id.proposalCode + str(id.proposalNumber) + "-" + str(id.visit_number) - == visit_name - ] # "Proposal title": id.title - return templates.TemplateResponse( - request=request, - name="visit.html", - context={"visit": return_query}, - ) - else: - return None - - -@router.post("/instruments/{instrument_name}/feedback") -async def send_murfey_message(instrument_name: str, msg: RegistrationMessage): - if _transport_object: - _transport_object.send( - _transport_object.feedback_queue, {"register": msg.registration} - ) - - -@router.post("/visits/{visit_name}/{session_id}/spa_preprocess") -async def request_spa_preprocessing( - visit_name: str, - session_id: MurfeySessionID, - proc_file: SPAProcessFile, - db=murfey_db, -): - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)).one().instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - parts = [secure_filename(p) for p in Path(proc_file.path).parts] - visit_idx = parts.index(visit_name) - core = Path("/") / Path(*parts[: visit_idx + 1]) - ppath = Path("/") / Path(*parts) - sub_dataset = ppath.relative_to(core).parts[0] - extra_path = machine_config.processed_extra_directory - for i, p in enumerate(ppath.parts): - if p.startswith("raw"): - movies_path_index = i - break - else: - raise ValueError(f"{proc_file.path} does not contain a raw directory") - mrc_out = ( - core - / machine_config.processed_directory_name - / sub_dataset - / extra_path - / "MotionCorr" - / "job002" - / "Movies" - / "/".join(ppath.parts[movies_path_index + 1 : -1]) - / str(ppath.stem + "_motion_corrected.mrc") - ) - try: - collected_ids = db.exec( - select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollectionGroup.tag == proc_file.tag) - .where(DataCollection.dcg_id == DataCollectionGroup.id) - .where(ProcessingJob.dc_id == DataCollection.id) - .where(AutoProcProgram.pj_id == ProcessingJob.id) - .where(ProcessingJob.recipe == "em-spa-preprocess") - ).one() - params = db.exec( - select(SPARelionParameters, SPAFeedbackParameters) - .where(SPARelionParameters.pj_id == collected_ids[2].id) - .where(SPAFeedbackParameters.pj_id == SPARelionParameters.pj_id) - ).one() - proc_params: dict | None = dict(params[0]) - feedback_params = params[1] - except sqlalchemy.exc.NoResultFound: - proc_params = None - try: - foil_hole_id = ( - db.exec( - select(FoilHole, GridSquare) - .where(FoilHole.name == proc_file.foil_hole_id) - .where(FoilHole.session_id == session_id) - .where(GridSquare.id == FoilHole.grid_square_id) - .where(GridSquare.tag == proc_file.tag) - ) - .one()[0] - .id - ) - except Exception as e: - log.warning( - f"Foil hole ID not found for foil hole {sanitise(str(proc_file.foil_hole_id))}: {e}", - exc_info=True, - ) - foil_hole_id = None - if proc_params: - - detached_ids = [c.id for c in collected_ids] - - murfey_ids = _murfey_id(detached_ids[3], db, number=2, close=False) - - if feedback_params.picker_murfey_id is None: - feedback_params.picker_murfey_id = murfey_ids[1] - db.add(feedback_params) - movie = Movie( - murfey_id=murfey_ids[0], - path=proc_file.path, - image_number=proc_file.image_number, - tag=proc_file.tag, - foil_hole_id=foil_hole_id, - ) - db.add(movie) - db.commit() - db.close() - - if not mrc_out.parent.exists(): - Path(secure_filename(str(mrc_out))).parent.mkdir( - parents=True, exist_ok=True - ) - recipe_name = machine_config.recipes.get( - "em-spa-preprocess", "em-spa-preprocess" - ) - zocalo_message: dict = { - "recipes": [recipe_name], - "parameters": { - "node_creator_queue": machine_config.node_creator_queue, - "dcid": detached_ids[1], - "kv": proc_params["voltage"], - "autoproc_program_id": detached_ids[3], - "movie": proc_file.path, - "mrc_out": str(mrc_out), - "pixel_size": proc_params["angpix"], - "image_number": proc_file.image_number, - "microscope": instrument_name, - "mc_uuid": murfey_ids[0], - "foil_hole_id": foil_hole_id, - "ft_bin": proc_params["motion_corr_binning"], - "fm_dose": proc_params["dose_per_frame"], - "gain_ref": proc_params["gain_ref"], - "picker_uuid": murfey_ids[1], - "session_id": session_id, - "particle_diameter": proc_params["particle_diameter"] or 0, - "fm_int_file": ( - proc_params["eer_fractionation_file"] - if proc_params["eer_fractionation_file"] - else proc_file.eer_fractionation_file - ), - "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, - "cryolo_model_weights": str( - _cryolo_model_path(visit_name, instrument_name) - ), - }, - } - # log.info(f"Sending Zocalo message {zocalo_message}") - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send("processing_recipe", zocalo_message) - else: - log.error( - f"Pe-processing was requested for {sanitise(ppath.name)} but no Zocalo transport object was found" - ) - return proc_file - - else: - for_stash = PreprocessStash( - file_path=str(proc_file.path), - tag=proc_file.tag, - session_id=session_id, - image_number=proc_file.image_number, - mrc_out=str(mrc_out), - eer_fractionation_file=str(proc_file.eer_fractionation_file), - foil_hole_id=foil_hole_id, - ) - db.add(for_stash) - db.commit() - db.close() - - return proc_file - - -@router.post("/visits/{visit_name}/{session_id}/tomography_preprocess") -async def request_tomography_preprocessing( - visit_name: str, - session_id: MurfeySessionID, - proc_file: TomoProcessFile, - db=murfey_db, -): - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)).one().instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - visit_idx = Path(proc_file.path).parts.index(visit_name) - core = Path(*Path(proc_file.path).parts[: visit_idx + 1]) - ppath = Path("/".join(secure_filename(p) for p in Path(proc_file.path).parts)) - sub_dataset = "/".join(ppath.relative_to(core).parts[:-1]) - extra_path = machine_config.processed_extra_directory - mrc_out = ( - core - / machine_config.processed_directory_name - / sub_dataset - / extra_path - / "MotionCorr" - / "job002" - / "Movies" - / str(ppath.stem + "_motion_corrected.mrc") - ) - mrc_out = Path("/".join(secure_filename(p) for p in mrc_out.parts)) - - recipe_name = machine_config.recipes.get("em-tomo-preprocess", "em-tomo-preprocess") - - data_collection = db.exec( - select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollectionGroup.tag == proc_file.group_tag) - .where(DataCollection.tag == proc_file.tag) - .where(DataCollection.dcg_id == DataCollectionGroup.id) - .where(ProcessingJob.dc_id == DataCollection.id) - .where(AutoProcProgram.pj_id == ProcessingJob.id) - .where(ProcessingJob.recipe == recipe_name) - ).all() - if data_collection: - if registered_tilts := db.exec( - select(Tilt).where(Tilt.movie_path == proc_file.path) - ).all(): - if len(registered_tilts) == 1: - if registered_tilts[0].motion_corrected: - return proc_file - dcid = data_collection[0][1].id - appid = data_collection[0][3].id - murfey_ids = _murfey_id(appid, db, number=1, close=False) - if not mrc_out.parent.exists(): - mrc_out.parent.mkdir(parents=True, exist_ok=True) - - session_processing_parameters = db.exec( - select(SessionProcessingParameters).where( - SessionProcessingParameters.session_id == session_id - ) - ).all() - if session_processing_parameters: - proc_file.gain_ref = session_processing_parameters[0].gain_ref - proc_file.dose_per_frame = session_processing_parameters[0].dose_per_frame - proc_file.eer_fractionation_file = session_processing_parameters[ - 0 - ].eer_fractionation_file - - zocalo_message: dict = { - "recipes": [recipe_name], - "parameters": { - "node_creator_queue": machine_config.node_creator_queue, - "dcid": dcid, - # "timestamp": datetime.datetime.now(), - "autoproc_program_id": appid, - "movie": proc_file.path, - "mrc_out": str(mrc_out), - "pixel_size": (proc_file.pixel_size) * 10**10, - "image_number": proc_file.image_number, - "kv": int(proc_file.voltage), - "microscope": instrument_name, - "mc_uuid": murfey_ids[0], - "ft_bin": proc_file.mc_binning, - "fm_dose": proc_file.dose_per_frame, - "frame_count": proc_file.frame_count, - "gain_ref": ( - str(machine_config.rsync_basepath / proc_file.gain_ref) - if proc_file.gain_ref and machine_config.data_transfer_enabled - else proc_file.gain_ref - ), - "fm_int_file": proc_file.eer_fractionation_file, - }, - } - if _transport_object: - zocalo_message["parameters"][ - "feedback_queue" - ] = _transport_object.feedback_queue - _transport_object.send("processing_recipe", zocalo_message) - else: - log.error( - f"Pe-processing was requested for {sanitise(ppath.name)} but no Zocalo transport object was found" - ) - return proc_file - else: - for_stash = PreprocessStash( - file_path=str(proc_file.path), - session_id=session_id, - image_number=proc_file.image_number, - mrc_out=str(mrc_out), - tag=proc_file.tag, - group_tag=proc_file.group_tag, - ) - db.add(for_stash) - db.commit() - db.close() - return proc_file - - -@router.post("/visits/{visit_name}/{session_id}/suggested_path") -def suggest_path( - visit_name: str, session_id: int, params: SuggestedPathParameters, db=murfey_db -): - count: int | None = None - secure_path_parts = [secure_filename(p) for p in params.base_path.parts] - base_path = "/".join(secure_path_parts) - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)).one().instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - if not machine_config: - raise ValueError( - "No machine configuration set when suggesting destination path" - ) - - # Construct the full path to where the dataset is to be saved - check_path = machine_config.rsync_basepath / base_path - - # Check previous year to account for the year rolling over during data collection - if not check_path.parent.exists(): - base_path_parts = base_path.split("/") - for part in base_path_parts: - # Find the path part corresponding to the year - if len(part) == 4 and part.isdigit(): - year_idx = base_path_parts.index(part) - base_path_parts[year_idx] = str(int(part) - 1) - base_path = "/".join(base_path_parts) - check_path_prev = check_path - check_path = machine_config.rsync_basepath / base_path - - # If it's not in the previous year either, it's a genuine error - if not check_path.parent.exists(): - log_message = ( - "Unable to find current visit folder under " - f"{str(check_path_prev.parent)!r} or {str(check_path.parent)!r}" - ) - log.error(log_message) - raise FileNotFoundError(log_message) - - check_path_name = check_path.name - while check_path.exists(): - count = count + 1 if count else 2 - check_path = check_path.parent / f"{check_path_name}{count}" - if params.touch: - check_path.mkdir(mode=0o750) - if params.extra_directory: - (check_path / secure_filename(params.extra_directory)).mkdir(mode=0o750) - return {"suggested_path": check_path.relative_to(machine_config.rsync_basepath)} - - -class Dest(BaseModel): - destination: Path - - -@router.post("/sessions/{session_id}/make_rsyncer_destination") -def make_rsyncer_destination(session_id: int, destination: Dest, db=murfey_db): - secure_path_parts = [secure_filename(p) for p in destination.destination.parts] - destination_path = "/".join(secure_path_parts) - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)).one().instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - if not machine_config: - raise ValueError("No machine configuration set when making rsyncer destination") - full_destination_path = machine_config.rsync_basepath / destination_path - for parent_path in full_destination_path.parents: - parent_path.mkdir(mode=0o750, exist_ok=True) - return destination - - -@router.get("/sessions/{session_id}/data_collection_groups") -def get_dc_groups( - session_id: MurfeySessionID, db=murfey_db -) -> Dict[str, DataCollectionGroup]: - data_collection_groups = db.exec( - select(DataCollectionGroup).where(DataCollectionGroup.session_id == session_id) - ).all() - return {dcg.tag: dcg for dcg in data_collection_groups} - - -@router.get("/sessions/{session_id}/data_collection_groups/{dcgid}/data_collections") -def get_data_collections( - session_id: MurfeySessionID, dcgid: int, db=murfey_db -) -> List[DataCollection]: - data_collections = db.exec( - select(DataCollection).where(DataCollection.dcg_id == dcgid) - ).all() - return data_collections - - -@router.post("/visits/{visit_name}/{session_id}/register_data_collection_group") -def register_dc_group( - visit_name, session_id: MurfeySessionID, dcg_params: DCGroupParameters, db=murfey_db -): - ispyb_proposal_code = visit_name[:2] - ispyb_proposal_number = visit_name.split("-")[0][2:] - ispyb_visit_number = visit_name.split("-")[-1] - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)).one().instrument_name - ) - log.info(f"Registering data collection group on microscope {instrument_name}") - if dcg_murfey := db.exec( - select(DataCollectionGroup) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollectionGroup.tag == dcg_params.tag) - ).all(): - dcg_murfey[0].atlas = dcg_params.atlas - dcg_murfey[0].sample = dcg_params.sample - dcg_murfey[0].atlas_pixel_size = dcg_params.atlas_pixel_size - - if _transport_object: - if dcg_murfey[0].atlas_id is not None: - _transport_object.send( - _transport_object.feedback_queue, - { - "register": "atlas_update", - "atlas_id": dcg_murfey[0].atlas_id, - "atlas": dcg_params.atlas, - "sample": dcg_params.sample, - "atlas_pixel_size": dcg_params.atlas_pixel_size, - }, - ) - else: - atlas_id_response = _transport_object.do_insert_atlas( - Atlas( - dataCollectionGroupId=dcg_murfey[0].id, - atlasImage=dcg_params.atlas, - pixelSize=dcg_params.atlas_pixel_size, - cassetteSlot=dcg_params.sample, - ) - ) - dcg_murfey[0].atlas_id = atlas_id_response["return_value"] - db.add(dcg_murfey[0]) - db.commit() - else: - dcg_parameters = { - "start_time": str(datetime.datetime.now()), - "experiment_type": dcg_params.experiment_type, - "experiment_type_id": dcg_params.experiment_type_id, - "tag": dcg_params.tag, - "session_id": session_id, - "atlas": dcg_params.atlas, - "sample": dcg_params.sample, - "atlas_pixel_size": dcg_params.atlas_pixel_size, - } - - if _transport_object: - _transport_object.send( - _transport_object.feedback_queue, {"register": "data_collection_group", **dcg_parameters, "microscope": instrument_name, "proposal_code": ispyb_proposal_code, "proposal_number": ispyb_proposal_number, "visit_number": ispyb_visit_number} # type: ignore - ) - return dcg_params - - -@router.post("/visits/{visit_name}/{session_id}/start_data_collection") -def start_dc( - visit_name, session_id: MurfeySessionID, dc_params: DCParameters, db=murfey_db -): - ispyb_proposal_code = visit_name[:2] - ispyb_proposal_number = visit_name.split("-")[0][2:] - ispyb_visit_number = visit_name.split("-")[-1] - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)).one().instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - log.info( - f"Starting data collection on microscope {instrument_name!r} " - f"with basepath {sanitise(str(machine_config.rsync_basepath))} and directory {sanitise(dc_params.image_directory)}" - ) - dc_parameters = { - "visit": visit_name, - "image_directory": str( - machine_config.rsync_basepath / dc_params.image_directory - ), - "start_time": str(datetime.datetime.now()), - "voltage": dc_params.voltage, - "pixel_size": str(float(dc_params.pixel_size_on_image) * 1e9), - "image_suffix": dc_params.file_extension, - "experiment_type": dc_params.experiment_type, - "image_size_x": dc_params.image_size_x, - "image_size_y": dc_params.image_size_y, - "acquisition_software": dc_params.acquisition_software, - "tag": dc_params.tag, - "source": dc_params.source, - "magnification": dc_params.magnification, - "total_exposed_dose": dc_params.total_exposed_dose, - "c2aperture": dc_params.c2aperture, - "exposure_time": dc_params.exposure_time, - "slit_width": dc_params.slit_width, - "phase_plate": dc_params.phase_plate, - "session_id": session_id, - } - - if _transport_object: - _transport_object.send( - _transport_object.feedback_queue, - { - "register": "data_collection", - **dc_parameters, - "microscope": instrument_name, - "proposal_code": ispyb_proposal_code, - "proposal_number": ispyb_proposal_number, - "visit_number": ispyb_visit_number, - }, - ) - if dc_params.exposure_time: - prom.exposure_time.set(dc_params.exposure_time) - return dc_params - - -@router.post("/visits/{visit_name}/{session_id}/register_processing_job") -def register_proc( - visit_name: str, - session_id: MurfeySessionID, - proc_params: ProcessingJobParameters, - db=murfey_db, -): - proc_parameters: dict = { - "session_id": session_id, - "experiment_type": proc_params.experiment_type, - "recipe": proc_params.recipe, - "source": proc_params.source, - "tag": proc_params.tag, - "job_parameters": { - k: v for k, v in proc_params.parameters.items() if v not in (None, "None") - }, - } - - session_processing_parameters = db.exec( - select(SessionProcessingParameters).where( - SessionProcessingParameters.session_id == session_id - ) - ).all() - - if session_processing_parameters: - job_parameters: dict = proc_parameters["job_parameters"] - job_parameters.update( - { - "gain_ref": session_processing_parameters[0].gain_ref, - "dose_per_frame": session_processing_parameters[0].dose_per_frame, - "eer_fractionation_file": session_processing_parameters[ - 0 - ].eer_fractionation_file, - "symmetry": session_processing_parameters[0].symmetry, - } - ) - proc_parameters["job_parameters"] = job_parameters - - if _transport_object: - _transport_object.send( - _transport_object.feedback_queue, - {"register": "processing_job", **proc_parameters}, - ) - return proc_params - - -@router.post("/sessions/{session_id}/process_gain") -async def process_gain( - session_id: MurfeySessionID, gain_reference_params: GainReference, db=murfey_db -): - murfey_session = db.exec(select(Session).where(Session.id == session_id)).one() - visit_name = murfey_session.visit - instrument_name = murfey_session.instrument_name - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - camera = getattr(Camera, machine_config.camera) - if gain_reference_params.eer: - executables = machine_config.external_executables_eer - else: - executables = machine_config.external_executables - env = machine_config.external_environment - safe_path_name = secure_filename(gain_reference_params.gain_ref.name) - filepath = ( - Path(machine_config.rsync_basepath) - / str(datetime.datetime.now().year) - / secure_filename(visit_name) - / machine_config.gain_directory_name - ) - - # Check under previous year if the folder doesn't exist - if not filepath.exists(): - filepath_prev = filepath - filepath = ( - Path(machine_config.rsync_basepath) - / str(datetime.datetime.now().year - 1) - / secure_filename(visit_name) - / machine_config.gain_directory_name - ) - # If it's not in the previous year, it's a genuine error - if not filepath.exists(): - log_message = ( - "Unable to find gain reference directory under " - f"{str(filepath_prev)!r} or {str(filepath)}" - ) - log.error(log_message) - raise FileNotFoundError(log_message) - - if gain_reference_params.eer: - new_gain_ref, new_gain_ref_superres = await prepare_eer_gain( - filepath / safe_path_name, - executables, - env, - tag=gain_reference_params.tag, - ) - else: - new_gain_ref, new_gain_ref_superres = await prepare_gain( - camera, - filepath / safe_path_name, - executables, - env, - rescale=gain_reference_params.rescale, - tag=gain_reference_params.tag, - ) - if new_gain_ref and new_gain_ref_superres: - return { - "gain_ref": new_gain_ref.relative_to(Path(machine_config.rsync_basepath)), - "gain_ref_superres": new_gain_ref_superres.relative_to( - Path(machine_config.rsync_basepath) - ), - } - elif new_gain_ref: - return { - "gain_ref": new_gain_ref.relative_to(Path(machine_config.rsync_basepath)), - "gain_ref_superres": None, - } - else: - return {"gain_ref": str(filepath / safe_path_name), "gain_ref_superres": None} - - -@router.delete("/sessions/{session_id}") -def remove_session_by_id(session_id: MurfeySessionID, db=murfey_db): - session = db.exec(select(Session).where(Session.id == session_id)).one() - sessions_for_visit = db.exec( - select(Session).where(Session.visit == session.visit) - ).all() - # Don't remove prometheus metrics if there are other sessions using them - if len(sessions_for_visit) == 1: - safe_run( - prom.monitoring_switch.remove, - args=(session.visit,), - label="monitoring_switch", - ) - rsync_instances = db.exec( - select(RsyncInstance).where(RsyncInstance.session_id == session_id) - ).all() - for ri in rsync_instances: - safe_run( - prom.seen_files.remove, - args=(ri.source, session.visit), - label="seen_files", - ) - safe_run( - prom.transferred_files.remove, - args=(ri.source, session.visit), - label="transferred_files", - ) - safe_run( - prom.transferred_files_bytes.remove, - args=(ri.source, session.visit), - label="transferred_files_bytes", - ) - safe_run( - prom.seen_data_files.remove, - args=(ri.source, session.visit), - label="seen_data_files", - ) - safe_run( - prom.transferred_data_files.remove, - args=(ri.source, session.visit), - label="transferred_data_files", - ) - safe_run( - prom.transferred_data_files_bytes.remove, - args=(ri.source, session.visit), - label="transferred_data_file_bytes", - ) - collected_ids = db.exec( - select(DataCollectionGroup, DataCollection, ProcessingJob) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollection.dcg_id == DataCollectionGroup.id) - .where(ProcessingJob.dc_id == DataCollection.id) - ).all() - for c in collected_ids: - safe_run( - prom.preprocessed_movies.remove, - args=(c[2].id,), - label="preprocessed_movies", - ) - db.delete(session) - db.commit() - return - - -@router.post("/visits/{visit_name}/{session_id}/eer_fractionation_file") -async def write_eer_fractionation_file( - visit_name: str, - session_id: int, - fractionation_params: FractionationParameters, - db=murfey_db, -) -> dict: - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)).one().instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - if machine_config.eer_fractionation_file_template: - file_path = Path( - machine_config.eer_fractionation_file_template.format( - visit=secure_filename(visit_name), - year=str(datetime.datetime.now().year), - ) - ) / secure_filename(fractionation_params.fractionation_file_name) - else: - file_path = ( - Path(machine_config.rsync_basepath) - / str(datetime.datetime.now().year) - / secure_filename(visit_name) - / machine_config.gain_directory_name - / secure_filename(fractionation_params.fractionation_file_name) - ) - - session_parameters = db.exec( - select(SessionProcessingParameters).where( - SessionProcessingParameters.session_id == session_id - ) - ).all() - if session_parameters: - session_parameters[0].eer_fractionation_file = str(file_path) - db.add(session_parameters[0]) - db.commit() - - if file_path.is_file(): - return {"eer_fractionation_file": str(file_path)} - - if fractionation_params.num_frames: - num_eer_frames = fractionation_params.num_frames - elif ( - fractionation_params.eer_path - and secure_path(Path(fractionation_params.eer_path)).is_file() - ): - num_eer_frames = murfey.util.eer.num_frames(Path(fractionation_params.eer_path)) - else: - log.warning( - f"EER fractionation unable to find {secure_path(Path(fractionation_params.eer_path)) if fractionation_params.eer_path else None} " - f"or use {int(sanitise(str(fractionation_params.num_frames)))} frames" - ) - return {"eer_fractionation_file": None} - with open(file_path, "w") as frac_file: - frac_file.write( - f"{num_eer_frames} {fractionation_params.fractionation} {fractionation_params.dose_per_frame / fractionation_params.fractionation}" - ) - return {"eer_fractionation_file": str(file_path)} - - -@router.post("/visits/{year}/{visit_name}/{session_id}/make_milling_gif") -async def make_gif( - year: int, - visit_name: str, - session_id: int, - gif_params: MillingParameters, - db=murfey_db, -): - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)).one().instrument_name - ) - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - output_dir = ( - Path(machine_config.rsync_basepath) - / secure_filename(year) - / secure_filename(visit_name) - / "processed" - ) - output_dir.mkdir(exist_ok=True) - output_dir = output_dir / secure_filename(gif_params.raw_directory) - output_dir.mkdir(exist_ok=True) - output_path = output_dir / f"lamella_{gif_params.lamella_number}_milling.gif" - image_full_paths = [ - output_dir.parent / gif_params.raw_directory / i for i in gif_params.images - ] - images = [Image.open(f) for f in image_full_paths] - for im in images: - im.thumbnail((512, 512)) - images[0].save( - output_path, - format="GIF", - append_images=images[1:], - save_all=True, - duration=30, - loop=0, - ) - return {"output_gif": str(output_path)} - - -@router.get("/new_client_id/") -async def new_client_id(db=murfey_db): - clients = db.exec(select(ClientEnvironment)).all() - if not clients: - return {"new_id": 0} - sorted_ids = sorted([c.client_id for c in clients]) - return {"new_id": sorted_ids[-1] + 1} - - -@router.get("/clients") -async def get_clients(db=murfey_db): - clients = db.exec(select(ClientEnvironment)).all() - return clients - - -@router.get("/sessions") -async def get_sessions(db=murfey_db): - sessions = db.exec(select(Session)).all() - clients = db.exec(select(ClientEnvironment)).all() - res = [] - for sess in sessions: - r = {"session": sess, "clients": []} - for cl in clients: - if cl.session_id == sess.id: - r["clients"].append(cl) - res.append(r) - return res - - -@router.get("/instruments/{instrument_name}/visits/{visit_name}/sessions") -def get_sessions_with_visit( - instrument_name: str, visit_name: str, db=murfey_db -) -> List[Session]: - sessions = db.exec( - select(Session) - .where(Session.instrument_name == instrument_name) - .where(Session.visit == visit_name) - ).all() - return sessions - - -@router.get("/instruments/{instrument_name}/sessions") -async def get_sessions_by_instrument_name( - instrument_name: str, db=murfey_db -) -> List[Session]: - sessions = db.exec( - select(Session).where(Session.instrument_name == instrument_name) - ).all() - return sessions - - -@router.post("/instruments/{instrument_name}/clients/{client_id}/session") -def link_client_to_session( - instrument_name: str, client_id: int, sess: SessionInfo, db=murfey_db -): - sid = sess.session_id - if sid is None: - s = Session(name=sess.session_name, instrument_name=instrument_name) - db.add(s) - db.commit() - sid = s.id - client = db.exec( - select(ClientEnvironment).where(ClientEnvironment.client_id == client_id) - ).one() - client.session_id = sid - db.add(client) - db.commit() - db.close() - return sid - - -@router.post("/sessions/{session_id}/successful_processing") -def register_processing_success_in_ispyb( - session_id: MurfeySessionID, db=murfey.server.ispyb.DB, murfey_db=murfey_db -): - collected_ids = murfey_db.exec( - select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollection.dcg_id == DataCollectionGroup.id) - .where(ProcessingJob.dc_id == DataCollection.id) - .where(AutoProcProgram.pj_id == ProcessingJob.id) - ).all() - appids = [c[3].id for c in collected_ids] - if _transport_object: - apps = db.query(ISPyBAutoProcProgram).filter( - ISPyBAutoProcProgram.autoProcProgramId.in_(appids) - ) - for updated in apps: - updated.processingStatus = True - _transport_object.do_update_processing_status(updated) - - -@router.post("/visits/{visit_name}/monitoring/{on}") -def change_monitoring_status(visit_name: str, on: int): - prom.monitoring_switch.labels(visit=visit_name) - prom.monitoring_switch.labels(visit=visit_name).set(on) - - -@router.post("/instruments/{instrument_name}/failed_client_post") -def failed_client_post(instrument_name: str, post_info: PostInfo): - zocalo_message = { - "register": "failed_client_post", - "url": post_info.url, - "json": post_info.data, - } - if _transport_object: - _transport_object.send(_transport_object.feedback_queue, zocalo_message) - - -@router.get("/sessions/{session_id}/upstream_visits") -async def find_upstream_visits(session_id: MurfeySessionID, db=murfey_db): - murfey_session = db.exec(select(Session).where(Session.id == session_id)).one() - visit_name = murfey_session.visit - instrument_name = murfey_session.instrument_name - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - upstream_visits = {} - # Iterates through provided upstream directories - for p in machine_config.upstream_data_directories: - # Looks for visit name in file path - for v in Path(p).glob(f"{visit_name.split('-')[0]}-*"): - upstream_visits[v.name] = v / machine_config.processed_directory_name - return upstream_visits - - -def _get_upstream_tiff_dirs(visit_name: str, instrument_name: str) -> List[Path]: - tiff_dirs = [] - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - for directory_name in machine_config.upstream_data_tiff_locations: - for p in machine_config.upstream_data_directories: - if (Path(p) / secure_filename(visit_name)).is_dir(): - processed_dir = Path(p) / secure_filename(visit_name) / directory_name - tiff_dirs.append(processed_dir) - break - if not tiff_dirs: - log.warning( - f"No candidate directory found for upstream download from visit {sanitise(visit_name)}" - ) - return tiff_dirs - - -@router.get("/visits/{visit_name}/{session_id}/upstream_tiff_paths") -async def gather_upstream_tiffs(visit_name: str, session_id: int, db=murfey_db): - """ - Looks for TIFF files associated with the current session in the permitted storage - servers, and returns their relative file paths as a list. - """ - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)).one().instrument_name - ) - upstream_tiff_paths = [] - tiff_dirs = _get_upstream_tiff_dirs(visit_name, instrument_name) - if not tiff_dirs: - return None - for tiff_dir in tiff_dirs: - for f in tiff_dir.glob("**/*.tiff"): - upstream_tiff_paths.append(str(f.relative_to(tiff_dir))) - for f in tiff_dir.glob("**/*.tif"): - upstream_tiff_paths.append(str(f.relative_to(tiff_dir))) - return upstream_tiff_paths - - -@router.get("/visits/{visit_name}/{session_id}/upstream_tiff/{tiff_path:path}") -async def get_tiff(visit_name: str, session_id: int, tiff_path: str, db=murfey_db): - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)).one().instrument_name - ) - tiff_dirs = _get_upstream_tiff_dirs(visit_name, instrument_name) - if not tiff_dirs: - return None - - tiff_path = "/".join(secure_filename(p) for p in tiff_path.split("/")) - for tiff_dir in tiff_dirs: - test_path = tiff_dir / tiff_path - if test_path.is_file(): - break - else: - log.warning(f"TIFF {tiff_path} not found") - return None - - return FileResponse(path=test_path) - - -class VisitEndTime(BaseModel): - end_time: Optional[datetime.datetime] = None - - -@router.post("/instruments/{instrument_name}/visits/{visit}/session/{name}") -def create_session( - instrument_name: str, - visit: str, - name: str, - visit_end_time: VisitEndTime, - db=murfey_db, -) -> int: - s = Session( - name=name, - visit=visit, - instrument_name=instrument_name, - visit_end_time=visit_end_time.end_time, - ) - db.add(s) - db.commit() - sid = s.id - return sid - - -@router.post("/sessions/{session_id}") -def update_session( - session_id: MurfeySessionID, process: bool = True, db=murfey_db -) -> None: - session = db.exec(select(Session).where(Session.id == session_id)).one() - session.process = process - db.add(session) - db.commit() - return None - - -@router.put("/sessions/{session_id}/current_gain_ref") -def update_current_gain_ref( - session_id: MurfeySessionID, new_gain_ref: CurrentGainRef, db=murfey_db -): - session = db.exec(select(Session).where(Session.id == session_id)).one() - session.current_gain_ref = new_gain_ref.path - db.add(session) - db.commit() - - -@router.get("/prometheus/{metric_name}") -def inspect_prometheus_metrics( - metric_name: str, -): - """ - A debugging endpoint that returns the current contents of any Prometheus - gauges and counters that have been set up thus far. - """ - - # Extract the Prometheus metric defined in the Prometheus module - metric: Optional[Counter | Gauge] = getattr(prom, metric_name, None) - if metric is None or not isinstance(metric, (Counter, Gauge)): - raise LookupError("No matching metric was found") - - # Package contents into dict and return - results = {} - if hasattr(metric, "_metrics"): - for i, (label_tuple, sub_metric) in enumerate(metric._metrics.items()): - labels = dict(zip(metric._labelnames, label_tuple)) - labels["value"] = sub_metric._value.get() - results[i] = labels - return results - else: - value = metric._value.get() - return {"value": value} +templates = Jinja2Templates(template_files) diff --git a/src/murfey/server/api/auth.py b/src/murfey/server/api/auth.py index 1fb30b003..edcdc6589 100644 --- a/src/murfey/server/api/auth.py +++ b/src/murfey/server/api/auth.py @@ -16,8 +16,8 @@ from pydantic import BaseModel from sqlmodel import Session, create_engine, select -from murfey.server import sanitise from murfey.server.murfey_db import murfey_db, url +from murfey.util.api import url_path_for from murfey.util.config import get_security_config from murfey.util.db import MurfeyUser as User from murfey.util.db import Session as MurfeySession @@ -26,7 +26,7 @@ logger = getLogger("murfey.server.api.auth") # Set up router -router = APIRouter() +router = APIRouter(tags=["Authentication"]) class CookieScheme(HTTPBearer): @@ -170,7 +170,7 @@ async def validate_token(token: Annotated[str, Depends(oauth2_scheme)]): ) async with aiohttp.ClientSession(cookies=cookies) as session: async with session.get( - f"{auth_url}/validate_token", + f"{auth_url}{url_path_for('auth.router', 'simple_token_validation')}", headers=headers, ) as response: success = response.status == 200 @@ -225,7 +225,7 @@ def create_access_token(data: dict, token: str = "") -> str: # check the session ID is alphanumeric for security raise ValueError("Session ID was invalid (not alphanumeric)") minted_token_response = requests.get( - f"{auth_url}/sessions/{sanitise(str(session_id))}/token", + f"{auth_url}{url_path_for('auth.router', 'mint_session_token', session_id=session_id)}", headers={"Authorization": f"Bearer {token}"}, ) if minted_token_response.status_code != 200: @@ -257,7 +257,7 @@ async def generate_token( data.add_field("password", form_data.password) async with aiohttp.ClientSession() as session: async with session.post( - f"{auth_url}/token", + f"{auth_url}{url_path_for('auth.router', 'generate_token')}", data=data, ) as response: validated = response.status == 200 diff --git a/src/murfey/server/api/bootstrap.py b/src/murfey/server/api/bootstrap.py index cc4fb1542..879bbfb11 100644 --- a/src/murfey/server/api/bootstrap.py +++ b/src/murfey/server/api/bootstrap.py @@ -21,6 +21,7 @@ import re import zipfile from io import BytesIO +from typing import Any from urllib.parse import quote import packaging.version @@ -29,10 +30,11 @@ from fastapi.responses import FileResponse, HTMLResponse, StreamingResponse import murfey -from murfey.server import get_machine_config, respond_with_template +from murfey.server.api import templates +from murfey.util.config import get_hostname, get_machine_config, get_microscope tag = { - "name": "bootstrap", + "name": "Bootstrap", "description": __doc__, "externalDocs": { "description": "PEP 503", @@ -42,13 +44,13 @@ # Set up API endpoint groups # NOTE: Routers MUST have prefixes. prefix="" causes an error -version = APIRouter(prefix="/version", tags=["bootstrap"]) -bootstrap = APIRouter(prefix="/bootstrap", tags=["bootstrap"]) -cygwin = APIRouter(prefix="/cygwin", tags=["bootstrap"]) -msys2 = APIRouter(prefix="/msys2", tags=["bootstrap"]) -rust = APIRouter(prefix="/rust", tags=["bootstrap"]) -pypi = APIRouter(prefix="/pypi", tags=["bootstrap"]) -plugins = APIRouter(prefix="/plugins", tags=["bootstrap"]) +version = APIRouter(prefix="/version", tags=["Bootstrap"]) +bootstrap = APIRouter(prefix="/bootstrap", tags=["Bootstrap"]) +cygwin = APIRouter(prefix="/cygwin", tags=["Bootstrap"]) +msys2 = APIRouter(prefix="/msys2", tags=["Bootstrap"]) +rust = APIRouter(prefix="/rust", tags=["Bootstrap"]) +pypi = APIRouter(prefix="/pypi", tags=["Bootstrap"]) +plugins = APIRouter(prefix="/plugins", tags=["Bootstrap"]) logger = logging.getLogger("murfey.server.api.bootstrap") @@ -99,6 +101,24 @@ def get_version(client_version: str = ""): """ +def respond_with_template( + request: Request, filename: str, parameters: dict[str, Any] | None = None +): + template_parameters = { + "hostname": get_hostname(), + "microscope": get_microscope(), + "version": murfey.__version__, + # Extra parameters to reconstruct URLs for forwarded requests + "netloc": request.url.netloc, + "proxy_path": "", + } + if parameters: + template_parameters.update(parameters) + return templates.TemplateResponse( + request=request, name=filename, context=template_parameters + ) + + @bootstrap.get("/", response_class=HTMLResponse) def get_bootstrap_instructions(request: Request): """ diff --git a/src/murfey/server/api/clem.py b/src/murfey/server/api/clem.py index b4e64327d..1bc043fb8 100644 --- a/src/murfey/server/api/clem.py +++ b/src/murfey/server/api/clem.py @@ -31,7 +31,7 @@ logger = getLogger("murfey.server.api.clem") # Create APIRouter class object -router = APIRouter() +router = APIRouter(tags=["Workflows: CLEM"]) # Valid file types valid_file_types = ( diff --git a/src/murfey/server/api/display.py b/src/murfey/server/api/display.py index 8f51548af..181bc3cd2 100644 --- a/src/murfey/server/api/display.py +++ b/src/murfey/server/api/display.py @@ -9,10 +9,20 @@ from murfey.util.db import DataCollectionGroup, FoilHole, GridSquare # Create APIRouter class object -router = APIRouter(prefix="/display", tags=["display"]) +router = APIRouter(prefix="/display", tags=["Display"]) machine_config = get_machine_config() +@router.get("/instruments/{instrument_name}/instrument_name") +def get_instrument_display_name(instrument_name: str) -> str: + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + if machine_config: + return machine_config.display_name + return "" + + @router.get("/instruments/{instrument_name}/image/") def get_mic_image(instrument_name: str): if machine_config[instrument_name].image_path: diff --git a/src/murfey/server/api/file_manip.py b/src/murfey/server/api/file_manip.py new file mode 100644 index 000000000..01724c7eb --- /dev/null +++ b/src/murfey/server/api/file_manip.py @@ -0,0 +1,257 @@ +from datetime import datetime +from logging import getLogger +from pathlib import Path +from typing import Optional + +from fastapi import APIRouter, Depends +from pydantic import BaseModel +from sqlmodel import select +from werkzeug.utils import secure_filename + +from murfey.server.api.auth import MurfeySessionID, validate_token +from murfey.server.gain import Camera, prepare_eer_gain, prepare_gain +from murfey.server.murfey_db import murfey_db +from murfey.util import sanitise, secure_path +from murfey.util.config import get_machine_config +from murfey.util.db import Session, SessionProcessingParameters +from murfey.util.eer import num_frames + +logger = getLogger("murfey.server.api.file_manip") + +router = APIRouter( + prefix="/file_manipulation", + dependencies=[Depends(validate_token)], + tags=["File Manipulation"], +) + + +class SuggestedPathParameters(BaseModel): + base_path: Path + touch: bool = False + extra_directory: str = "" + + +@router.post("/visits/{visit_name}/{session_id}/suggested_path") +def suggest_path( + visit_name: str, session_id: int, params: SuggestedPathParameters, db=murfey_db +): + count: Optional[int] = None + secure_path_parts = [secure_filename(p) for p in params.base_path.parts] + base_path = "/".join(secure_path_parts) + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)).one().instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + if not machine_config: + raise ValueError( + "No machine configuration set when suggesting destination path" + ) + + # Construct the full path to where the dataset is to be saved + check_path = machine_config.rsync_basepath / base_path + + # Check previous year to account for the year rolling over during data collection + if not check_path.parent.exists(): + base_path_parts = base_path.split("/") + for part in base_path_parts: + # Find the path part corresponding to the year + if len(part) == 4 and part.isdigit(): + year_idx = base_path_parts.index(part) + base_path_parts[year_idx] = str(int(part) - 1) + base_path = "/".join(base_path_parts) + check_path_prev = check_path + check_path = machine_config.rsync_basepath / base_path + + # If it's not in the previous year either, it's a genuine error + if not check_path.parent.exists(): + log_message = ( + "Unable to find current visit folder under " + f"{str(check_path_prev.parent)!r} or {str(check_path.parent)!r}" + ) + logger.error(log_message) + raise FileNotFoundError(log_message) + + check_path_name = check_path.name + while check_path.exists(): + count = count + 1 if count else 2 + check_path = check_path.parent / f"{check_path_name}{count}" + if params.touch: + check_path.mkdir(mode=0o750) + if params.extra_directory: + (check_path / secure_filename(params.extra_directory)).mkdir(mode=0o750) + return {"suggested_path": check_path.relative_to(machine_config.rsync_basepath)} + + +class Dest(BaseModel): + destination: Path + + +@router.post("/sessions/{session_id}/make_rsyncer_destination") +def make_rsyncer_destination(session_id: int, destination: Dest, db=murfey_db): + secure_path_parts = [secure_filename(p) for p in destination.destination.parts] + destination_path = "/".join(secure_path_parts) + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)).one().instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + if not machine_config: + raise ValueError("No machine configuration set when making rsyncer destination") + full_destination_path = machine_config.rsync_basepath / destination_path + for parent_path in full_destination_path.parents: + parent_path.mkdir(mode=0o750, exist_ok=True) + return destination + + +class GainReference(BaseModel): + gain_ref: Path + rescale: bool = True + eer: bool = False + tag: str = "" + + +@router.post("/sessions/{session_id}/process_gain") +async def process_gain( + session_id: MurfeySessionID, gain_reference_params: GainReference, db=murfey_db +): + murfey_session = db.exec(select(Session).where(Session.id == session_id)).one() + visit_name = murfey_session.visit + instrument_name = murfey_session.instrument_name + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + camera = getattr(Camera, machine_config.camera) + if gain_reference_params.eer: + executables = machine_config.external_executables_eer + else: + executables = machine_config.external_executables + env = machine_config.external_environment + safe_path_name = secure_filename(gain_reference_params.gain_ref.name) + filepath = ( + Path(machine_config.rsync_basepath) + / str(datetime.now().year) + / secure_filename(visit_name) + / machine_config.gain_directory_name + ) + + # Check under previous year if the folder doesn't exist + if not filepath.exists(): + filepath_prev = filepath + filepath = ( + Path(machine_config.rsync_basepath) + / str(datetime.now().year - 1) + / secure_filename(visit_name) + / machine_config.gain_directory_name + ) + # If it's not in the previous year, it's a genuine error + if not filepath.exists(): + log_message = ( + "Unable to find gain reference directory under " + f"{str(filepath_prev)!r} or {str(filepath)}" + ) + logger.error(log_message) + raise FileNotFoundError(log_message) + + if gain_reference_params.eer: + new_gain_ref, new_gain_ref_superres = await prepare_eer_gain( + filepath / safe_path_name, + executables, + env, + tag=gain_reference_params.tag, + ) + else: + new_gain_ref, new_gain_ref_superres = await prepare_gain( + camera, + filepath / safe_path_name, + executables, + env, + rescale=gain_reference_params.rescale, + tag=gain_reference_params.tag, + ) + if new_gain_ref and new_gain_ref_superres: + return { + "gain_ref": new_gain_ref.relative_to(Path(machine_config.rsync_basepath)), + "gain_ref_superres": new_gain_ref_superres.relative_to( + Path(machine_config.rsync_basepath) + ), + } + elif new_gain_ref: + return { + "gain_ref": new_gain_ref.relative_to(Path(machine_config.rsync_basepath)), + "gain_ref_superres": None, + } + else: + return {"gain_ref": str(filepath / safe_path_name), "gain_ref_superres": None} + + +class FractionationParameters(BaseModel): + fractionation: int + dose_per_frame: float + num_frames: int = 0 + eer_path: Optional[str] = None + fractionation_file_name: str = "eer_fractionation.txt" + + +@router.post("/visits/{visit_name}/{session_id}/eer_fractionation_file") +async def write_eer_fractionation_file( + visit_name: str, + session_id: int, + fractionation_params: FractionationParameters, + db=murfey_db, +) -> dict: + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)).one().instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + if machine_config.eer_fractionation_file_template: + file_path = Path( + machine_config.eer_fractionation_file_template.format( + visit=secure_filename(visit_name), + year=str(datetime.now().year), + ) + ) / secure_filename(fractionation_params.fractionation_file_name) + else: + file_path = ( + Path(machine_config.rsync_basepath) + / str(datetime.now().year) + / secure_filename(visit_name) + / machine_config.gain_directory_name + / secure_filename(fractionation_params.fractionation_file_name) + ) + + session_parameters = db.exec( + select(SessionProcessingParameters).where( + SessionProcessingParameters.session_id == session_id + ) + ).all() + if session_parameters: + session_parameters[0].eer_fractionation_file = str(file_path) + db.add(session_parameters[0]) + db.commit() + + if file_path.is_file(): + return {"eer_fractionation_file": str(file_path)} + + if fractionation_params.num_frames: + num_eer_frames = fractionation_params.num_frames + elif ( + fractionation_params.eer_path + and secure_path(Path(fractionation_params.eer_path)).is_file() + ): + num_eer_frames = num_frames(Path(fractionation_params.eer_path)) + else: + logger.warning( + f"EER fractionation unable to find {secure_path(Path(fractionation_params.eer_path)) if fractionation_params.eer_path else None} " + f"or use {int(sanitise(str(fractionation_params.num_frames)))} frames" + ) + return {"eer_fractionation_file": None} + with open(file_path, "w") as frac_file: + frac_file.write( + f"{num_eer_frames} {fractionation_params.fractionation} {fractionation_params.dose_per_frame / fractionation_params.fractionation}" + ) + return {"eer_fractionation_file": str(file_path)} diff --git a/src/murfey/server/api/hub.py b/src/murfey/server/api/hub.py index cc9205712..f17db4a1c 100644 --- a/src/murfey/server/api/hub.py +++ b/src/murfey/server/api/hub.py @@ -11,7 +11,7 @@ config = get_machine_config() -router = APIRouter() +router = APIRouter(tags=["Murfey Hub"]) class InstrumentInfo(BaseModel): diff --git a/src/murfey/server/api/instrument.py b/src/murfey/server/api/instrument.py index 3f60e35d9..59665399a 100644 --- a/src/murfey/server/api/instrument.py +++ b/src/murfey/server/api/instrument.py @@ -12,22 +12,26 @@ from sqlmodel import select from werkzeug.utils import secure_filename -from murfey.server import sanitise -from murfey.server.api import MurfeySessionID from murfey.server.api.auth import ( + MurfeySessionID, create_access_token, instrument_server_tokens, oauth2_scheme, validate_token, ) from murfey.server.murfey_db import murfey_db -from murfey.util import secure_path +from murfey.util import sanitise, secure_path +from murfey.util.api import url_path_for from murfey.util.config import get_machine_config from murfey.util.db import RsyncInstance, Session, SessionProcessingParameters from murfey.util.models import File, MultigridWatcherSetup # Create APIRouter class object -router = APIRouter(dependencies=[Depends(validate_token)]) +router = APIRouter( + prefix="/instrument_server", + dependencies=[Depends(validate_token)], + tags=["Instrument Server"], +) log = logging.getLogger("murfey.server.instrument") @@ -62,7 +66,7 @@ async def activate_instrument_server_for_session( ] async with aiohttp.ClientSession() as clientsession: async with clientsession.post( - f"{machine_config.instrument_server_url}/sessions/{int(sanitise(str(session_id)))}/token", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'token_handshake_for_session', session_id=session_id)}", json={"access_token": token, "token_type": "bearer"}, ) as response: success = response.status == 200 @@ -85,7 +89,7 @@ async def check_if_session_is_active(instrument_name: str, session_id: int): instrument_name ] async with clientsession.get( - f"{machine_config.instrument_server_url}/sessions/{int(sanitise(str(session_id)))}/check_token", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'check_token', session_id=session_id)}", headers={ "Authorization": f"Bearer {instrument_server_tokens[session_id]['access_token']}" }, @@ -106,22 +110,12 @@ async def setup_multigrid_watcher( if machine_config.instrument_server_url: session = db.exec(select(Session).where(Session.id == session_id)).one() visit = session.visit - _config = { - "acquisition_software": machine_config.acquisition_software, - "calibrations": machine_config.calibrations, - "data_directories": [str(k) for k in machine_config.data_directories], - "create_directories": [str(k) for k in machine_config.create_directories], - "rsync_basepath": str(machine_config.rsync_basepath), - "visit": visit, - "default_model": str(machine_config.default_model), - } async with aiohttp.ClientSession() as clientsession: async with clientsession.post( - f"{machine_config.instrument_server_url}/sessions/{session_id}/multigrid_watcher", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'setup_multigrid_watcher', session_id=session_id)}", json={ "source": str(secure_path(watcher_spec.source / visit)), "visit": visit, - "configuration": _config, "label": visit, "instrument_name": instrument_name, "skip_existing_processing": watcher_spec.skip_existing_processing, @@ -155,7 +149,7 @@ async def start_multigrid_watcher(session_id: MurfeySessionID, db=murfey_db): ) async with aiohttp.ClientSession() as clientsession: async with clientsession.post( - f"{machine_config.instrument_server_url}/sessions/{session_id}/start_multigrid_watcher?process={'true' if process else 'false'}", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'start_multigrid_watcher', session_id=session_id)}?process={'true' if process else 'false'}", headers={ "Authorization": f"Bearer {instrument_server_tokens[session_id]['access_token']}" }, @@ -199,7 +193,7 @@ async def pass_proc_params_to_instrument_server( label = db.exec(select(Session).where(Session.id == session_id)).one().name async with aiohttp.ClientSession() as clientsession: async with clientsession.post( - f"{machine_config.instrument_server_url}/sessions/{session_id}/processing_parameters", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'register_processing_parameters', session_id=session_id)}", json={ "label": label, "params": { @@ -228,7 +222,7 @@ async def check_instrument_server(instrument_name: str): if machine_config.instrument_server_url: async with aiohttp.ClientSession() as clientsession: async with clientsession.get( - f"{machine_config.instrument_server_url}/health", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'health')}", ) as resp: data = await resp.json() return data @@ -249,7 +243,7 @@ async def get_possible_gain_references( token = instrument_server_tokens[session_id]["access_token"] async with aiohttp.ClientSession() as clientsession: async with clientsession.get( - f"{machine_config.instrument_server_url}/instruments/{sanitise(instrument_name)}/sessions/{sanitise(str(session_id))}/possible_gain_references", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'get_possible_gain_references', instrument_name=sanitise(instrument_name), session_id=session_id)}", headers={"Authorization": f"Bearer {token}"}, ) as resp: data = await resp.json() @@ -278,7 +272,7 @@ async def request_gain_reference_upload( if machine_config.instrument_server_url: async with aiohttp.ClientSession() as clientsession: async with clientsession.post( - f"{machine_config.instrument_server_url}/instruments/{instrument_name}/sessions/{session_id}/upload_gain_reference", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'upload_gain_reference', instrument_name=instrument_name, session_id=session_id)}", json={ "gain_path": str(gain_reference_request.gain_path), "visit_path": visit_path, @@ -311,7 +305,7 @@ async def request_upstream_tiff_data_download( if machine_config.instrument_server_url: async with aiohttp.ClientSession() as clientsession: async with clientsession.post( - f"{machine_config.instrument_server_url}/visits/{secure_filename(visit_name)}/sessions/{sanitise(str(session_id))}/upstream_tiff_data_request", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'gather_upstream_tiffs', visit_name=secure_filename(visit_name), session_id=session_id)}", json={"download_dir": download_dir}, headers={ "Authorization": f"Bearer {instrument_server_tokens[session_id]['access_token']}" @@ -339,7 +333,7 @@ async def stop_rsyncer( if machine_config.instrument_server_url: async with aiohttp.ClientSession() as clientsession: async with clientsession.post( - f"{machine_config.instrument_server_url}/sessions/{session_id}/stop_rsyncer", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'stop_rsyncer', session_id=session_id)}", json={ "label": session_id, "source": str(secure_path(Path(rsyncer_source.source))), @@ -366,7 +360,7 @@ async def finalise_rsyncer( if machine_config.instrument_server_url: async with aiohttp.ClientSession() as clientsession: async with clientsession.post( - f"{machine_config.instrument_server_url}/sessions/{session_id}/finalise_rsyncer", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'finalise_rsyncer', session_id=session_id)}", json={ "label": session_id, "source": str(secure_path(Path(rsyncer_source.source))), @@ -391,7 +385,7 @@ async def finalise_session(session_id: MurfeySessionID, db=murfey_db): if machine_config.instrument_server_url: async with aiohttp.ClientSession() as clientsession: async with clientsession.post( - f"{machine_config.instrument_server_url}/sessions/{session_id}/finalise_session", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'finalise_session', session_id=session_id)}", headers={ "Authorization": f"Bearer {instrument_server_tokens[session_id]['access_token']}" }, @@ -412,7 +406,7 @@ async def abandon_session(session_id: MurfeySessionID, db=murfey_db): if machine_config.instrument_server_url: async with aiohttp.ClientSession() as clientsession: async with clientsession.post( - f"{machine_config.instrument_server_url}/sessions/{session_id}/abandon_controller", + f"{machine_config.instrument_server_url}{url_path_for('api_router', 'abandon_controller', session_id=session_id)}", headers={ "Authorization": f"Bearer {instrument_server_tokens[session_id]['access_token']}" }, @@ -436,7 +430,7 @@ async def remove_rsyncer( if machine_config.instrument_server_url: async with aiohttp.ClientSession() as clientsession: async with clientsession.post( - f"{machine_config.instrument_server_url}/sessions/{session_id}/remove_rsyncer", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'remove_rsyncer', session_id=session_id)}", json={ "label": session_id, "source": str(secure_path(Path(rsyncer_source.source))), @@ -464,7 +458,7 @@ async def restart_rsyncer( if machine_config.instrument_server_url: async with aiohttp.ClientSession() as clientsession: async with clientsession.post( - f"{machine_config.instrument_server_url}/sessions/{session_id}/restart_rsyncer", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'restart_rsyncer', session_id=session_id)}", json={ "label": session_id, "source": str(secure_path(Path(rsyncer_source.source))), @@ -513,7 +507,7 @@ async def get_rsyncer_info( token = instrument_server_tokens[session_id]["access_token"] async with aiohttp.ClientSession() as clientsession: async with clientsession.get( - f"{machine_config.instrument_server_url}/sessions/{session_id}/rsyncer_info", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'get_rsyncer_info', session_id=session_id)}", headers={"Authorization": f"Bearer {token}"}, ) as resp: if resp.status == 200: @@ -533,7 +527,7 @@ async def get_rsyncer_info( token = instrument_server_tokens[session_id]["access_token"] async with aiohttp.ClientSession() as clientsession: async with clientsession.get( - f"{machine_config.instrument_server_url}/sessions/{session_id}/analyser_info", + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'get_analyser_info', session_id=session_id)}", headers={"Authorization": f"Bearer {token}"}, ) as resp: if resp.status == 200: diff --git a/src/murfey/server/api/mag_table.py b/src/murfey/server/api/mag_table.py new file mode 100644 index 000000000..f285b92bb --- /dev/null +++ b/src/murfey/server/api/mag_table.py @@ -0,0 +1,35 @@ +from typing import List + +from fastapi import APIRouter, Depends +from sqlmodel import select + +from murfey.server.api.auth import validate_token +from murfey.server.murfey_db import murfey_db +from murfey.util.db import MagnificationLookup + +router = APIRouter( + prefix="/mag_table", + dependencies=[Depends(validate_token)], + tags=["Magnification Table"], +) + + +@router.get("/mag_table/") +def get_mag_table(db=murfey_db) -> List[MagnificationLookup]: + return db.exec(select(MagnificationLookup)).all() + + +@router.post("/mag_table/") +def add_to_mag_table(rows: List[MagnificationLookup], db=murfey_db): + for r in rows: + db.add(r) + db.commit() + + +@router.delete("/mag_table/{mag}") +def remove_mag_table_row(mag: int, db=murfey_db): + row = db.exec( + select(MagnificationLookup).where(MagnificationLookup.magnification == mag) + ).one() + db.delete(row) + db.commit() diff --git a/src/murfey/server/api/processing_parameters.py b/src/murfey/server/api/processing_parameters.py index bf506d30f..11b9f57de 100644 --- a/src/murfey/server/api/processing_parameters.py +++ b/src/murfey/server/api/processing_parameters.py @@ -12,7 +12,11 @@ logger = getLogger("murfey.server.api.processing_parameters") -router = APIRouter(dependencies=[Depends(validate_token)]) +router = APIRouter( + prefix="/session_parameters", + dependencies=[Depends(validate_token)], + tags=["Processing Parameters"], +) class EditableSessionProcessingParameters(BaseModel): diff --git a/src/murfey/server/api/prometheus.py b/src/murfey/server/api/prometheus.py new file mode 100644 index 000000000..5e8d867ea --- /dev/null +++ b/src/murfey/server/api/prometheus.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +from logging import getLogger +from typing import Optional + +from fastapi import APIRouter, Depends +from prometheus_client import Counter, Gauge +from sqlmodel import select + +import murfey.server.prometheus as prom +from murfey.server.api.auth import validate_token +from murfey.server.murfey_db import murfey_db +from murfey.util import sanitise +from murfey.util.db import RsyncInstance +from murfey.util.models import RsyncerInfo + +logger = getLogger("murfey.server.api.prometheus") + +router = APIRouter( + prefix="/prometheus", + dependencies=[Depends(validate_token)], + tags=["Prometheus"], +) + + +@router.post("/visits/{visit_name}/increment_rsync_file_count") +def increment_rsync_file_count( + visit_name: str, rsyncer_info: RsyncerInfo, db=murfey_db +): + try: + rsync_instance = db.exec( + select(RsyncInstance).where( + RsyncInstance.source == rsyncer_info.source, + RsyncInstance.destination == rsyncer_info.destination, + RsyncInstance.session_id == rsyncer_info.session_id, + ) + ).one() + except Exception: + logger.error( + f"Failed to find rsync instance for visit {sanitise(visit_name)} " + "with the following properties: \n" + f"{rsyncer_info.dict()}", + exc_info=True, + ) + return None + rsync_instance.files_counted += rsyncer_info.increment_count + db.add(rsync_instance) + db.commit() + db.close() + prom.seen_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).inc( + rsyncer_info.increment_count + ) + prom.seen_data_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).inc( + rsyncer_info.increment_data_count + ) + + +@router.post("/visits/{visit_name}/increment_rsync_transferred_files") +def increment_rsync_transferred_files( + visit_name: str, rsyncer_info: RsyncerInfo, db=murfey_db +): + rsync_instance = db.exec( + select(RsyncInstance).where( + RsyncInstance.source == rsyncer_info.source, + RsyncInstance.destination == rsyncer_info.destination, + RsyncInstance.session_id == rsyncer_info.session_id, + ) + ).one() + rsync_instance.files_transferred += rsyncer_info.increment_count + db.add(rsync_instance) + db.commit() + db.close() + + +@router.post("/visits/{visit_name}/increment_rsync_transferred_files_prometheus") +def increment_rsync_transferred_files_prometheus( + visit_name: str, rsyncer_info: RsyncerInfo, db=murfey_db +): + prom.transferred_files.labels( + rsync_source=rsyncer_info.source, visit=visit_name + ).inc(rsyncer_info.increment_count) + prom.transferred_files_bytes.labels( + rsync_source=rsyncer_info.source, visit=visit_name + ).inc(rsyncer_info.bytes) + prom.transferred_data_files.labels( + rsync_source=rsyncer_info.source, visit=visit_name + ).inc(rsyncer_info.increment_data_count) + prom.transferred_data_files_bytes.labels( + rsync_source=rsyncer_info.source, visit=visit_name + ).inc(rsyncer_info.data_bytes) + + +@router.post("/visits/{visit_name}/monitoring/{on}") +def change_monitoring_status(visit_name: str, on: int): + prom.monitoring_switch.labels(visit=visit_name) + prom.monitoring_switch.labels(visit=visit_name).set(on) + + +@router.get("/metrics/{metric_name}") +def inspect_prometheus_metrics( + metric_name: str, +): + """ + A debugging endpoint that returns the current contents of any Prometheus + gauges and counters that have been set up thus far. + """ + + # Extract the Prometheus metric defined in the Prometheus module + metric: Optional[Counter | Gauge] = getattr(prom, metric_name, None) + if metric is None or not isinstance(metric, (Counter, Gauge)): + raise LookupError("No matching metric was found") + + # Package contents into dict and return + results = {} + if hasattr(metric, "_metrics"): + for i, (label_tuple, sub_metric) in enumerate(metric._metrics.items()): + labels = dict(zip(metric._labelnames, label_tuple)) + labels["value"] = sub_metric._value.get() + results[i] = labels + return results + else: + value = metric._value.get() + return {"value": value} diff --git a/src/murfey/server/api/session_control.py b/src/murfey/server/api/session_control.py new file mode 100644 index 000000000..418226f87 --- /dev/null +++ b/src/murfey/server/api/session_control.py @@ -0,0 +1,421 @@ +from datetime import datetime +from logging import getLogger +from pathlib import Path +from typing import Dict, List, Optional + +from fastapi import APIRouter, Depends +from fastapi.responses import FileResponse +from ispyb.sqlalchemy import AutoProcProgram as ISPyBAutoProcProgram +from pydantic import BaseModel +from sqlmodel import select +from werkzeug.utils import secure_filename + +import murfey.server.prometheus as prom +from murfey.server import _transport_object +from murfey.server.api.auth import MurfeySessionID, validate_token +from murfey.server.api.shared import get_foil_hole as _get_foil_hole +from murfey.server.api.shared import ( + get_foil_holes_from_grid_square as _get_foil_holes_from_grid_square, +) +from murfey.server.api.shared import get_grid_squares as _get_grid_squares +from murfey.server.api.shared import ( + get_grid_squares_from_dcg as _get_grid_squares_from_dcg, +) +from murfey.server.api.shared import ( + get_machine_config_for_instrument, + get_upstream_tiff_dirs, + remove_session_by_id, +) +from murfey.server.ispyb import DB as ispyb_db +from murfey.server.ispyb import get_all_ongoing_visits +from murfey.server.murfey_db import murfey_db +from murfey.util import sanitise +from murfey.util.config import MachineConfig, get_machine_config +from murfey.util.db import ( + AutoProcProgram, + ClientEnvironment, + DataCollection, + DataCollectionGroup, + FoilHole, + GridSquare, + ProcessingJob, + RsyncInstance, + Session, +) +from murfey.util.models import ( + ClientInfo, + FoilHoleParameters, + GridSquareParameters, + RsyncerInfo, + Visit, +) +from murfey.workflows.spa.flush_spa_preprocess import ( + register_foil_hole as _register_foil_hole, +) +from murfey.workflows.spa.flush_spa_preprocess import ( + register_grid_square as _register_grid_square, +) + +logger = getLogger("murfey.server.api.session_control") + +router = APIRouter( + prefix="/session_control", + dependencies=[Depends(validate_token)], + tags=["Session Control: General"], +) + + +@router.get("/time") +async def get_current_timestamp(): + return {"timestamp": datetime.now().timestamp()} + + +@router.get("/instruments/{instrument_name}/machine") +def machine_info_by_instrument(instrument_name: str) -> Optional[MachineConfig]: + return get_machine_config_for_instrument(instrument_name) + + +@router.get("/new_client_id/") +async def new_client_id(db=murfey_db): + clients = db.exec(select(ClientEnvironment)).all() + if not clients: + return {"new_id": 0} + sorted_ids = sorted([c.client_id for c in clients]) + return {"new_id": sorted_ids[-1] + 1} + + +@router.get("/instruments/{instrument_name}/visits_raw", response_model=List[Visit]) +def get_current_visits(instrument_name: str, db=ispyb_db): + logger.debug( + f"Received request to look up ongoing visits for {sanitise(instrument_name)}" + ) + return get_all_ongoing_visits(instrument_name, db) + + +class SessionInfo(BaseModel): + session_id: Optional[int] + session_name: str = "" + rescale: bool = True + + +@router.post("/instruments/{instrument_name}/clients/{client_id}/session") +def link_client_to_session( + instrument_name: str, client_id: int, sess: SessionInfo, db=murfey_db +): + sid = sess.session_id + if sid is None: + s = Session(name=sess.session_name, instrument_name=instrument_name) + db.add(s) + db.commit() + sid = s.id + client = db.exec( + select(ClientEnvironment).where(ClientEnvironment.client_id == client_id) + ).one() + client.session_id = sid + db.add(client) + db.commit() + db.close() + return sid + + +@router.post("/visits/{visit_name}") +def register_client_to_visit(visit_name: str, client_info: ClientInfo, db=murfey_db): + client_env = db.exec( + select(ClientEnvironment).where(ClientEnvironment.client_id == client_info.id) + ).one() + session = db.exec(select(Session).where(Session.id == client_env.session_id)).one() + if client_env: + client_env.visit = visit_name + db.add(client_env) + db.commit() + if session: + session.visit = visit_name + db.add(session) + db.commit() + db.close() + return client_info + + +@router.get("/sessions") +async def get_sessions(db=murfey_db): + sessions = db.exec(select(Session)).all() + clients = db.exec(select(ClientEnvironment)).all() + res = [] + for sess in sessions: + r = {"session": sess, "clients": []} + for cl in clients: + if cl.session_id == sess.id: + r["clients"].append(cl) + res.append(r) + return res + + +@router.delete("/sessions/{session_id}") +def remove_session(session_id: MurfeySessionID, db=murfey_db): + remove_session_by_id(session_id, db) + + +@router.post("/sessions/{session_id}/successful_processing") +def register_processing_success_in_ispyb( + session_id: MurfeySessionID, db=ispyb_db, murfey_db=murfey_db +): + collected_ids = murfey_db.exec( + select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram) + .where(DataCollectionGroup.session_id == session_id) + .where(DataCollection.dcg_id == DataCollectionGroup.id) + .where(ProcessingJob.dc_id == DataCollection.id) + .where(AutoProcProgram.pj_id == ProcessingJob.id) + ).all() + appids = [c[3].id for c in collected_ids] + if _transport_object: + if db is not None: + apps = db.query(ISPyBAutoProcProgram).filter( + ISPyBAutoProcProgram.autoProcProgramId.in_(appids) + ) + for updated in apps: + updated.processingStatus = True + _transport_object.do_update_processing_status(updated) + + +class PostInfo(BaseModel): + url: str + data: dict + + +@router.post("/instruments/{instrument_name}/failed_client_post") +def failed_client_post(instrument_name: str, post_info: PostInfo): + zocalo_message = { + "register": "failed_client_post", + "url": post_info.url, + "json": post_info.data, + } + if _transport_object: + _transport_object.send(_transport_object.feedback_queue, zocalo_message) + + +@router.post("/sessions/{session_id}/rsyncer") +def register_rsyncer(session_id: int, rsyncer_info: RsyncerInfo, db=murfey_db): + visit_name = db.exec(select(Session).where(Session.id == session_id)).one().visit + rsync_instance = RsyncInstance( + source=rsyncer_info.source, + session_id=rsyncer_info.session_id, + transferring=rsyncer_info.transferring, + destination=rsyncer_info.destination, + tag=rsyncer_info.tag, + ) + db.add(rsync_instance) + db.commit() + db.close() + prom.seen_files.labels(rsync_source=rsyncer_info.source, visit=visit_name) + prom.seen_data_files.labels(rsync_source=rsyncer_info.source, visit=visit_name) + prom.transferred_files.labels(rsync_source=rsyncer_info.source, visit=visit_name) + prom.transferred_files_bytes.labels( + rsync_source=rsyncer_info.source, visit=visit_name + ) + prom.transferred_data_files.labels( + rsync_source=rsyncer_info.source, visit=visit_name + ) + prom.transferred_data_files_bytes.labels( + rsync_source=rsyncer_info.source, visit=visit_name + ) + prom.seen_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).set(0) + prom.transferred_files.labels( + rsync_source=rsyncer_info.source, visit=visit_name + ).set(0) + prom.transferred_files_bytes.labels( + rsync_source=rsyncer_info.source, visit=visit_name + ).set(0) + prom.seen_data_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).set( + 0 + ) + prom.transferred_data_files.labels( + rsync_source=rsyncer_info.source, visit=visit_name + ).set(0) + prom.transferred_data_files_bytes.labels( + rsync_source=rsyncer_info.source, visit=visit_name + ).set(0) + return rsyncer_info + + +@router.get("/sessions/{session_id}/rsyncers", response_model=List[RsyncInstance]) +def get_rsyncers_for_session(session_id: MurfeySessionID, db=murfey_db): + rsync_instances = db.exec( + select(RsyncInstance).where(RsyncInstance.session_id == session_id) + ) + return rsync_instances.all() + + +class RsyncerSource(BaseModel): + source: str + + +@router.post("/sessions/{session_id}/rsyncer_stopped") +def register_stopped_rsyncer( + session_id: int, rsyncer_source: RsyncerSource, db=murfey_db +): + rsyncer = db.exec( + select(RsyncInstance) + .where(RsyncInstance.session_id == session_id) + .where(RsyncInstance.source == rsyncer_source.source) + ).one() + rsyncer.transferring = False + db.add(rsyncer) + db.commit() + + +@router.post("/sessions/{session_id}/rsyncer_started") +def register_restarted_rsyncer( + session_id: int, rsyncer_source: RsyncerSource, db=murfey_db +): + rsyncer = db.exec( + select(RsyncInstance) + .where(RsyncInstance.session_id == session_id) + .where(RsyncInstance.source == rsyncer_source.source) + ).one() + rsyncer.transferring = True + db.add(rsyncer) + db.commit() + + +@router.delete("/sessions/{session_id}/rsyncer") +def delete_rsyncer(session_id: int, source: Path, db=murfey_db): + try: + rsync_instance = db.exec( + select(RsyncInstance) + .where(RsyncInstance.session_id == session_id) + .where(RsyncInstance.source == str(source)) + ).one() + db.delete(rsync_instance) + db.commit() + except Exception: + logger.error( + f"Failed to delete rsyncer for source directory {sanitise(str(source))!r} " + f"in session {session_id}.", + exc_info=True, + ) + + +spa_router = APIRouter( + prefix="/session_control/spa", + dependencies=[Depends(validate_token)], + tags=["Session Control: SPA"], +) + + +@spa_router.get("/sessions/{session_id}/grid_squares") +def get_grid_squares(session_id: MurfeySessionID, db=murfey_db): + return _get_grid_squares(session_id, db) + + +@spa_router.get("/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares") +def get_grid_squares_from_dcg( + session_id: MurfeySessionID, dcgid: int, db=murfey_db +) -> List[GridSquare]: + return _get_grid_squares_from_dcg(session_id, dcgid, db) + + +@spa_router.get( + "/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares/{gsid}/foil_holes" +) +def get_foil_holes_from_grid_square( + session_id: MurfeySessionID, dcgid: int, gsid: int, db=murfey_db +) -> List[FoilHole]: + return _get_foil_holes_from_grid_square(session_id, dcgid, gsid, db) + + +@spa_router.get("/sessions/{session_id}/foil_hole/{fh_name}") +def get_foil_hole( + session_id: MurfeySessionID, fh_name: int, db=murfey_db +) -> Dict[str, int]: + return _get_foil_hole(session_id, fh_name, db) + + +@spa_router.post("/sessions/{session_id}/grid_square/{gsid}") +def register_grid_square( + session_id: MurfeySessionID, + gsid: int, + grid_square_params: GridSquareParameters, + db=murfey_db, +): + return _register_grid_square(session_id, gsid, grid_square_params, db) + + +@spa_router.post("/sessions/{session_id}/grid_square/{gs_name}/foil_hole") +def register_foil_hole( + session_id: MurfeySessionID, + gs_name: int, + foil_hole_params: FoilHoleParameters, + db=murfey_db, +): + logger.info( + f"Registering foil hole {foil_hole_params.name} with position {(foil_hole_params.x_location, foil_hole_params.y_location)}" + ) + return _register_foil_hole(session_id, gs_name, foil_hole_params, db) + + +correlative_router = APIRouter( + prefix="/session_control/correlative", + dependencies=[Depends(validate_token)], + tags=["Session Control: Correlative Imaging"], +) + + +@correlative_router.get("/sessions/{session_id}/upstream_visits") +async def find_upstream_visits(session_id: MurfeySessionID, db=murfey_db): + murfey_session = db.exec(select(Session).where(Session.id == session_id)).one() + visit_name = murfey_session.visit + instrument_name = murfey_session.instrument_name + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + upstream_visits = {} + # Iterates through provided upstream directories + for p in machine_config.upstream_data_directories: + # Looks for visit name in file path + for v in Path(p).glob(f"{visit_name.split('-')[0]}-*"): + upstream_visits[v.name] = v / machine_config.processed_directory_name + return upstream_visits + + +@correlative_router.get("/visits/{visit_name}/{session_id}/upstream_tiff_paths") +async def gather_upstream_tiffs(visit_name: str, session_id: int, db=murfey_db): + """ + Looks for TIFF files associated with the current session in the permitted storage + servers, and returns their relative file paths as a list. + """ + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)).one().instrument_name + ) + upstream_tiff_paths = [] + tiff_dirs = get_upstream_tiff_dirs(visit_name, instrument_name) + if not tiff_dirs: + return None + for tiff_dir in tiff_dirs: + for f in tiff_dir.glob("**/*.tiff"): + upstream_tiff_paths.append(str(f.relative_to(tiff_dir))) + for f in tiff_dir.glob("**/*.tif"): + upstream_tiff_paths.append(str(f.relative_to(tiff_dir))) + return upstream_tiff_paths + + +@correlative_router.get( + "/visits/{visit_name}/{session_id}/upstream_tiff/{tiff_path:path}" +) +async def get_tiff(visit_name: str, session_id: int, tiff_path: str, db=murfey_db): + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)).one().instrument_name + ) + tiff_dirs = get_upstream_tiff_dirs(visit_name, instrument_name) + if not tiff_dirs: + return None + + tiff_path = "/".join(secure_filename(p) for p in tiff_path.split("/")) + for tiff_dir in tiff_dirs: + test_path = tiff_dir / tiff_path + if test_path.is_file(): + break + else: + logger.warning(f"TIFF {tiff_path} not found") + return None + + return FileResponse(path=test_path) diff --git a/src/murfey/server/api/session_info.py b/src/murfey/server/api/session_info.py new file mode 100644 index 000000000..50e38ba2b --- /dev/null +++ b/src/murfey/server/api/session_info.py @@ -0,0 +1,486 @@ +from datetime import datetime +from logging import getLogger +from pathlib import Path +from typing import Dict, List, Optional + +from fastapi import APIRouter, Depends, Request +from fastapi.responses import FileResponse, HTMLResponse +from pydantic import BaseModel +from sqlalchemy import func +from sqlmodel import select +from werkzeug.utils import secure_filename + +import murfey +import murfey.server.api.websocket as ws +from murfey.server import _transport_object +from murfey.server.api import templates +from murfey.server.api.auth import MurfeySessionID, validate_token +from murfey.server.api.shared import get_foil_hole as _get_foil_hole +from murfey.server.api.shared import ( + get_foil_holes_from_grid_square as _get_foil_holes_from_grid_square, +) +from murfey.server.api.shared import get_grid_squares as _get_grid_squares +from murfey.server.api.shared import ( + get_grid_squares_from_dcg as _get_grid_squares_from_dcg, +) +from murfey.server.api.shared import ( + get_machine_config_for_instrument, + get_upstream_tiff_dirs, + remove_session_by_id, +) +from murfey.server.ispyb import DB as ispyb_db +from murfey.server.ispyb import get_all_ongoing_visits +from murfey.server.murfey_db import murfey_db +from murfey.util import sanitise +from murfey.util.config import ( + MachineConfig, + get_hostname, + get_machine_config, + get_microscope, +) +from murfey.util.db import ( + ClientEnvironment, + DataCollection, + DataCollectionGroup, + FoilHole, + GridSquare, + Movie, + ProcessingJob, + RsyncInstance, + Session, + SPAFeedbackParameters, + SPARelionParameters, + Tilt, + TiltSeries, +) +from murfey.util.models import Visit + +logger = getLogger("murfey.server.api.session_info") + +router = APIRouter( + prefix="/session_info", + dependencies=[Depends(validate_token)], + tags=["Session Info: General"], +) + + +# This will be the homepage for a given microscope. +@router.get("/", response_class=HTMLResponse) +async def root(request: Request): + return templates.TemplateResponse( + request=request, + name="home.html", + context={ + "hostname": get_hostname(), + "microscope": get_microscope(), + "version": murfey.__version__, + }, + ) + + +@router.get("/health/") +def health_check(db=ispyb_db): + conn = db.connection() + conn.close() + return { + "ispyb_connection": True, + "rabbitmq_connection": _transport_object.transport.is_connected(), + } + + +@router.get("/connections/") +def connections_check(): + return {"connections": list(ws.manager.active_connections.keys())} + + +@router.get("/instruments/{instrument_name}/machine") +def machine_info_by_instrument(instrument_name: str) -> Optional[MachineConfig]: + return get_machine_config_for_instrument(instrument_name) + + +@router.get("/instruments/{instrument_name}/visits_raw", response_model=List[Visit]) +def get_current_visits(instrument_name: str, db=ispyb_db): + logger.debug( + f"Received request to look up ongoing visits for {sanitise(instrument_name)}" + ) + return get_all_ongoing_visits(instrument_name, db) + + +@router.get("/instruments/{instrument_name}/visits/") +def all_visit_info(instrument_name: str, request: Request, db=ispyb_db): + visits = get_all_ongoing_visits(instrument_name, db) + + if visits: + return_query = [ + { + "Start date": visit.start, + "End date": visit.end, + "Visit name": visit.name, + "Time remaining": str(visit.end - datetime.now()), + } + for visit in visits + ] # "Proposal title": visit.proposal_title + logger.debug( + f"{len(visits)} visits active for {sanitise(instrument_name)=}: {', '.join(v.name for v in visits)}" + ) + return templates.TemplateResponse( + request=request, + name="activevisits.html", + context={"info": return_query, "microscope": instrument_name}, + ) + else: + logger.debug(f"No visits identified for {sanitise(instrument_name)=}") + return templates.TemplateResponse( + request=request, + name="activevisits.html", + context={"info": [], "microscope": instrument_name}, + ) + + +@router.get("/sessions/{session_id}/rsyncers", response_model=List[RsyncInstance]) +def get_rsyncers_for_client(session_id: MurfeySessionID, db=murfey_db): + rsync_instances = db.exec( + select(RsyncInstance).where(RsyncInstance.session_id == session_id) + ) + return rsync_instances.all() + + +class SessionClients(BaseModel): + session: Session + clients: List[ClientEnvironment] + + +@router.get("/session/{session_id}") +async def get_session(session_id: MurfeySessionID, db=murfey_db) -> SessionClients: + session = db.exec(select(Session).where(Session.id == session_id)).one() + clients = db.exec( + select(ClientEnvironment).where(ClientEnvironment.session_id == session_id) + ).all() + return SessionClients(session=session, clients=clients) + + +@router.get("/sessions") +async def get_sessions(db=murfey_db): + sessions = db.exec(select(Session)).all() + clients = db.exec(select(ClientEnvironment)).all() + res = [] + for sess in sessions: + r = {"session": sess, "clients": []} + for cl in clients: + if cl.session_id == sess.id: + r["clients"].append(cl) + res.append(r) + return res + + +class VisitEndTime(BaseModel): + end_time: Optional[datetime] = None + + +@router.post("/instruments/{instrument_name}/visits/{visit}/session/{name}") +def create_session( + instrument_name: str, + visit: str, + name: str, + visit_end_time: VisitEndTime, + db=murfey_db, +) -> int: + s = Session( + name=name, + visit=visit, + instrument_name=instrument_name, + visit_end_time=visit_end_time.end_time, + ) + db.add(s) + db.commit() + sid = s.id + return sid + + +@router.post("/sessions/{session_id}") +def update_session( + session_id: MurfeySessionID, process: bool = True, db=murfey_db +) -> None: + session = db.exec(select(Session).where(Session.id == session_id)).one() + session.process = process + db.add(session) + db.commit() + return None + + +@router.delete("/sessions/{session_id}") +def remove_session(session_id: MurfeySessionID, db=murfey_db): + remove_session_by_id(session_id, db) + + +@router.get("/instruments/{instrument_name}/visits/{visit_name}/sessions") +def get_sessions_with_visit( + instrument_name: str, visit_name: str, db=murfey_db +) -> List[Session]: + sessions = db.exec( + select(Session) + .where(Session.instrument_name == instrument_name) + .where(Session.visit == visit_name) + ).all() + return sessions + + +@router.get("/instruments/{instrument_name}/sessions") +async def get_sessions_by_instrument_name( + instrument_name: str, db=murfey_db +) -> List[Session]: + sessions = db.exec( + select(Session).where(Session.instrument_name == instrument_name) + ).all() + return sessions + + +@router.get("/sessions/{session_id}/data_collection_groups") +def get_dc_groups( + session_id: MurfeySessionID, db=murfey_db +) -> Dict[str, DataCollectionGroup]: + data_collection_groups = db.exec( + select(DataCollectionGroup).where(DataCollectionGroup.session_id == session_id) + ).all() + return {dcg.tag: dcg for dcg in data_collection_groups} + + +@router.get("/sessions/{session_id}/data_collection_groups/{dcgid}/data_collections") +def get_data_collections( + session_id: MurfeySessionID, dcgid: int, db=murfey_db +) -> List[DataCollection]: + data_collections = db.exec( + select(DataCollection).where(DataCollection.dcg_id == dcgid) + ).all() + return data_collections + + +@router.get("/clients") +async def get_clients(db=murfey_db): + clients = db.exec(select(ClientEnvironment)).all() + return clients + + +@router.get("/num_movies") +def count_number_of_movies(db=murfey_db) -> Dict[str, int]: + res = db.exec( + select(Movie.tag, func.count(Movie.murfey_id)).group_by(Movie.tag) + ).all() + return {r[0]: r[1] for r in res} + + +class CurrentGainRef(BaseModel): + path: str + + +@router.put("/sessions/{session_id}/current_gain_ref") +def update_current_gain_ref( + session_id: MurfeySessionID, new_gain_ref: CurrentGainRef, db=murfey_db +): + session = db.exec(select(Session).where(Session.id == session_id)).one() + session.current_gain_ref = new_gain_ref.path + db.add(session) + db.commit() + + +spa_router = APIRouter( + prefix="/session_info/spa", + dependencies=[Depends(validate_token)], + tags=["Session Info: SPA"], +) + + +class ProcessingDetails(BaseModel): + data_collection_group: DataCollectionGroup + data_collections: List[DataCollection] + processing_jobs: List[ProcessingJob] + relion_params: SPARelionParameters + feedback_params: SPAFeedbackParameters + + +@spa_router.get("/sessions/{session_id}/spa_processing_parameters") +def get_spa_proc_param_details( + session_id: MurfeySessionID, db=murfey_db +) -> Optional[List[ProcessingDetails]]: + params = db.exec( + select( + DataCollectionGroup, + DataCollection, + ProcessingJob, + SPARelionParameters, + SPAFeedbackParameters, + ) + .where(DataCollectionGroup.session_id == session_id) + .where(DataCollectionGroup.id == DataCollection.dcg_id) + .where(DataCollection.id == ProcessingJob.dc_id) + .where(SPARelionParameters.pj_id == ProcessingJob.id) + .where(SPAFeedbackParameters.pj_id == ProcessingJob.id) + ).all() + if not params: + return None + unique_dcg_indices = [] + dcg_ids = [] + for i, p in enumerate(params): + if p[0].id not in dcg_ids: + dcg_ids.append(p[0].id) + unique_dcg_indices.append(i) + + def _parse(ps, i, dcg_id): + res = [] + for p in ps: + if p[0].id == dcg_id: + if p[i] not in res: + res.append(p[i]) + return res + + return [ + ProcessingDetails( + data_collection_group=params[i][0], + data_collections=_parse(params, 1, d), + processing_jobs=_parse(params, 2, d), + relion_params=_parse(params, 3, d)[0], + feedback_params=_parse(params, 4, d)[0], + ) + for i, d in zip(unique_dcg_indices, dcg_ids) + ] + + +@spa_router.get( + "/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares/{gsid}/foil_holes/{fhid}/num_movies" +) +def get_number_of_movies_from_foil_hole( + session_id: int, dcgid: int, gsid: int, fhid: int, db=murfey_db +) -> int: + movies = db.exec( + select(Movie, FoilHole, GridSquare, DataCollectionGroup) + .where(Movie.foil_hole_id == FoilHole.id) + .where(FoilHole.name == fhid) + .where(FoilHole.grid_square_id == GridSquare.id) + .where(GridSquare.name == gsid) + .where(GridSquare.session_id == session_id) + .where(GridSquare.tag == DataCollectionGroup.tag) + .where(DataCollectionGroup.id == dcgid) + ).all() + return len(movies) + + +@spa_router.get("/sessions/{session_id}/grid_squares") +def get_grid_squares(session_id: MurfeySessionID, db=murfey_db): + return _get_grid_squares(session_id, db) + + +@spa_router.get("/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares") +def get_grid_squares_from_dcg( + session_id: MurfeySessionID, dcgid: int, db=murfey_db +) -> List[GridSquare]: + return _get_grid_squares_from_dcg(session_id, dcgid, db) + + +@spa_router.get( + "/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares/{gsid}/foil_holes" +) +def get_foil_holes_from_grid_square( + session_id: MurfeySessionID, dcgid: int, gsid: int, db=murfey_db +) -> List[FoilHole]: + return _get_foil_holes_from_grid_square(session_id, dcgid, gsid, db) + + +@spa_router.get("/sessions/{session_id}/foil_hole/{fh_name}") +def get_foil_hole( + session_id: MurfeySessionID, fh_name: int, db=murfey_db +) -> Dict[str, int]: + return _get_foil_hole(session_id, fh_name, db) + + +tomo_router = APIRouter( + prefix="/session_info/tomo", + dependencies=[Depends(validate_token)], + tags=["Session Info: CryoET"], +) + + +@tomo_router.get("/sessions/{session_id}/tilt_series/{tilt_series_tag}/tilts") +def get_tilts( + session_id: MurfeySessionID, tilt_series_tag: str, db=murfey_db +) -> Dict[str, List[str]]: + res = db.exec( + select(TiltSeries, Tilt) + .where(TiltSeries.tag == tilt_series_tag) + .where(TiltSeries.session_id == session_id) + .where(Tilt.tilt_series_id == TiltSeries.id) + ).all() + tilts: Dict[str, List[str]] = {} + for el in res: + if tilts.get(el[1].rsync_source): + tilts[el[1].rsync_source].append(el[2].movie_path) + else: + tilts[el[1].rsync_source] = [el[2].movie_path] + return tilts + + +correlative_router = APIRouter( + prefix="/session_info/correlative", + dependencies=[Depends(validate_token)], + tags=["Session Info: Correlative Imaging"], +) + + +@correlative_router.get("/sessions/{session_id}/upstream_visits") +async def find_upstream_visits(session_id: MurfeySessionID, db=murfey_db): + murfey_session = db.exec(select(Session).where(Session.id == session_id)).one() + visit_name = murfey_session.visit + instrument_name = murfey_session.instrument_name + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + upstream_visits = {} + # Iterates through provided upstream directories + for p in machine_config.upstream_data_directories: + # Looks for visit name in file path + for v in Path(p).glob(f"{visit_name.split('-')[0]}-*"): + upstream_visits[v.name] = v / machine_config.processed_directory_name + return upstream_visits + + +@correlative_router.get("/visits/{visit_name}/{session_id}/upstream_tiff_paths") +async def gather_upstream_tiffs(visit_name: str, session_id: int, db=murfey_db): + """ + Looks for TIFF files associated with the current session in the permitted storage + servers, and returns their relative file paths as a list. + """ + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)).one().instrument_name + ) + upstream_tiff_paths = [] + tiff_dirs = get_upstream_tiff_dirs(visit_name, instrument_name) + if not tiff_dirs: + return None + for tiff_dir in tiff_dirs: + for f in tiff_dir.glob("**/*.tiff"): + upstream_tiff_paths.append(str(f.relative_to(tiff_dir))) + for f in tiff_dir.glob("**/*.tif"): + upstream_tiff_paths.append(str(f.relative_to(tiff_dir))) + return upstream_tiff_paths + + +@correlative_router.get( + "/visits/{visit_name}/{session_id}/upstream_tiff/{tiff_path:path}" +) +async def get_tiff(visit_name: str, session_id: int, tiff_path: str, db=murfey_db): + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)).one().instrument_name + ) + tiff_dirs = get_upstream_tiff_dirs(visit_name, instrument_name) + if not tiff_dirs: + return None + + tiff_path = "/".join(secure_filename(p) for p in tiff_path.split("/")) + for tiff_dir in tiff_dirs: + test_path = tiff_dir / tiff_path + if test_path.is_file(): + break + else: + logger.warning(f"TIFF {tiff_path} not found") + return None + + return FileResponse(path=test_path) diff --git a/src/murfey/server/api/shared.py b/src/murfey/server/api/shared.py new file mode 100644 index 000000000..116b3b9d6 --- /dev/null +++ b/src/murfey/server/api/shared.py @@ -0,0 +1,157 @@ +import logging +from functools import lru_cache +from pathlib import Path +from typing import Dict, List, Optional + +from sqlmodel import select +from werkzeug.utils import secure_filename + +import murfey.server.prometheus as prom +from murfey.util import safe_run, sanitise +from murfey.util.config import MachineConfig, from_file, get_machine_config, settings +from murfey.util.db import ( + DataCollection, + DataCollectionGroup, + FoilHole, + GridSquare, + ProcessingJob, + RsyncInstance, + Session, +) + +logger = logging.getLogger("murfey.server.api.shared") + + +@lru_cache(maxsize=5) +def get_machine_config_for_instrument(instrument_name: str) -> Optional[MachineConfig]: + if settings.murfey_machine_configuration: + return from_file(Path(settings.murfey_machine_configuration), instrument_name)[ + instrument_name + ] + return None + + +def remove_session_by_id(session_id: int, db): + session = db.exec(select(Session).where(Session.id == session_id)).one() + sessions_for_visit = db.exec( + select(Session).where(Session.visit == session.visit) + ).all() + # Don't remove prometheus metrics if there are other sessions using them + if len(sessions_for_visit) == 1: + safe_run( + prom.monitoring_switch.remove, + args=(session.visit,), + label="monitoring_switch", + ) + rsync_instances = db.exec( + select(RsyncInstance).where(RsyncInstance.session_id == session_id) + ).all() + for ri in rsync_instances: + safe_run( + prom.seen_files.remove, + args=(ri.source, session.visit), + label="seen_files", + ) + safe_run( + prom.transferred_files.remove, + args=(ri.source, session.visit), + label="transferred_files", + ) + safe_run( + prom.transferred_files_bytes.remove, + args=(ri.source, session.visit), + label="transferred_files_bytes", + ) + safe_run( + prom.seen_data_files.remove, + args=(ri.source, session.visit), + label="seen_data_files", + ) + safe_run( + prom.transferred_data_files.remove, + args=(ri.source, session.visit), + label="transferred_data_files", + ) + safe_run( + prom.transferred_data_files_bytes.remove, + args=(ri.source, session.visit), + label="transferred_data_file_bytes", + ) + collected_ids = db.exec( + select(DataCollectionGroup, DataCollection, ProcessingJob) + .where(DataCollectionGroup.session_id == session_id) + .where(DataCollection.dcg_id == DataCollectionGroup.id) + .where(ProcessingJob.dc_id == DataCollection.id) + ).all() + for c in collected_ids: + safe_run( + prom.preprocessed_movies.remove, + args=(c[2].id,), + label="preprocessed_movies", + ) + db.delete(session) + db.commit() + return + + +def get_grid_squares(session_id: int, db): + grid_squares = db.exec( + select(GridSquare).where(GridSquare.session_id == session_id) + ).all() + tags = {gs.tag for gs in grid_squares} + res = {} + for t in tags: + res[t] = [gs for gs in grid_squares if gs.tag == t] + return res + + +def get_grid_squares_from_dcg(session_id: int, dcgid: int, db) -> List[GridSquare]: + grid_squares = db.exec( + select(GridSquare, DataCollectionGroup) + .where(GridSquare.session_id == session_id) + .where(GridSquare.tag == DataCollectionGroup.tag) + .where(DataCollectionGroup.id == dcgid) + ).all() + return [gs[0] for gs in grid_squares] + + +def get_foil_holes_from_grid_square( + session_id: int, dcgid: int, gsid: int, db +) -> List[FoilHole]: + foil_holes = db.exec( + select(FoilHole, GridSquare, DataCollectionGroup) + .where(FoilHole.grid_square_id == GridSquare.id) + .where(GridSquare.name == gsid) + .where(GridSquare.session_id == session_id) + .where(GridSquare.tag == DataCollectionGroup.tag) + .where(DataCollectionGroup.id == dcgid) + ).all() + return [fh[0] for fh in foil_holes] + + +def get_foil_hole(session_id: int, fh_name: int, db) -> Dict[str, int]: + foil_holes = db.exec( + select(FoilHole, GridSquare) + .where(FoilHole.name == fh_name) + .where(FoilHole.session_id == session_id) + .where(GridSquare.id == FoilHole.grid_square_id) + ).all() + return {f[1].tag: f[0].id for f in foil_holes} + + +def get_upstream_tiff_dirs(visit_name: str, instrument_name: str) -> List[Path]: + tiff_dirs = [] + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + for directory_name in machine_config.upstream_data_tiff_locations: + for p in machine_config.upstream_data_directories: + if (Path(p) / secure_filename(visit_name)).is_dir(): + processed_dir = Path(p) / secure_filename(visit_name) / directory_name + tiff_dirs.append(processed_dir) + break + if not tiff_dirs: + logger.warning( + f"No candidate directory found for upstream download from visit {sanitise(visit_name)}" + ) + return tiff_dirs diff --git a/src/murfey/server/api/spa.py b/src/murfey/server/api/spa.py index 5a2a2c737..c34a4f6db 100644 --- a/src/murfey/server/api/spa.py +++ b/src/murfey/server/api/spa.py @@ -12,7 +12,7 @@ from murfey.util.db import Session as MurfeySession # Create APIRouter class object -router = APIRouter() +router = APIRouter(tags=["Workflows: crYOLO Models"]) @lru_cache(maxsize=5) diff --git a/src/murfey/server/websocket.py b/src/murfey/server/api/websocket.py similarity index 99% rename from src/murfey/server/websocket.py rename to src/murfey/server/api/websocket.py index 13cbcc73e..9765dd31b 100644 --- a/src/murfey/server/websocket.py +++ b/src/murfey/server/api/websocket.py @@ -16,7 +16,7 @@ T = TypeVar("T") -ws = APIRouter(prefix="/ws", tags=["websocket"]) +ws = APIRouter(prefix="/ws", tags=["Websocket"]) log = logging.getLogger("murfey.server.websocket") diff --git a/src/murfey/server/api/workflow.py b/src/murfey/server/api/workflow.py new file mode 100644 index 000000000..df6da6705 --- /dev/null +++ b/src/murfey/server/api/workflow.py @@ -0,0 +1,1056 @@ +import asyncio +from datetime import datetime +from logging import getLogger +from pathlib import Path +from typing import Any, Dict, List, Optional + +import sqlalchemy +from fastapi import APIRouter, Depends +from ispyb.sqlalchemy import ( + Atlas, + BLSample, + BLSampleGroup, + BLSampleGroupHasBLSample, + BLSampleImage, + BLSubSample, +) +from pydantic import BaseModel +from sqlalchemy.exc import OperationalError +from sqlmodel import col, select +from werkzeug.utils import secure_filename + +try: + from PIL import Image +except ImportError: + Image = None + +import murfey.server.prometheus as prom +from murfey.server import _transport_object +from murfey.server.api.auth import MurfeySessionID, validate_token +from murfey.server.api.spa import _cryolo_model_path +from murfey.server.feedback import ( + _murfey_id, + check_tilt_series_mc, + get_all_tilts, + get_angle, + get_job_ids, + get_tomo_proc_params, +) +from murfey.server.ispyb import DB as ispyb_db +from murfey.server.ispyb import get_proposal_id +from murfey.server.murfey_db import murfey_db +from murfey.util import sanitise +from murfey.util.config import get_machine_config +from murfey.util.db import ( + AutoProcProgram, + DataCollection, + DataCollectionGroup, + FoilHole, + GridSquare, + Movie, + PreprocessStash, + ProcessingJob, + Session, + SessionProcessingParameters, + SPAFeedbackParameters, + SPARelionParameters, + Tilt, + TiltSeries, +) +from murfey.util.models import ProcessingParametersSPA, ProcessingParametersTomo +from murfey.util.processing_params import default_spa_parameters +from murfey.util.tomo import midpoint + +logger = getLogger("murfey.server.api.workflow") + +router = APIRouter( + prefix="/workflow", + dependencies=[Depends(validate_token)], + tags=["Workflows: General"], +) + + +class DCGroupParameters(BaseModel): + # DC = Data collection + experiment_type: str + experiment_type_id: int + tag: str + atlas: str = "" + sample: Optional[int] = None + atlas_pixel_size: int = 0 + + +@router.post("/visits/{visit_name}/{session_id}/register_data_collection_group") +def register_dc_group( + visit_name, session_id: MurfeySessionID, dcg_params: DCGroupParameters, db=murfey_db +): + ispyb_proposal_code = visit_name[:2] + ispyb_proposal_number = visit_name.split("-")[0][2:] + ispyb_visit_number = visit_name.split("-")[-1] + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)).one().instrument_name + ) + logger.info(f"Registering data collection group on microscope {instrument_name}") + if dcg_murfey := db.exec( + select(DataCollectionGroup) + .where(DataCollectionGroup.session_id == session_id) + .where(DataCollectionGroup.tag == dcg_params.tag) + ).all(): + dcg_murfey[0].atlas = dcg_params.atlas + dcg_murfey[0].sample = dcg_params.sample + dcg_murfey[0].atlas_pixel_size = dcg_params.atlas_pixel_size + + if _transport_object: + if dcg_murfey[0].atlas_id is not None: + _transport_object.send( + _transport_object.feedback_queue, + { + "register": "atlas_update", + "atlas_id": dcg_murfey[0].atlas_id, + "atlas": dcg_params.atlas, + "sample": dcg_params.sample, + "atlas_pixel_size": dcg_params.atlas_pixel_size, + }, + ) + else: + atlas_id_response = _transport_object.do_insert_atlas( + Atlas( + dataCollectionGroupId=dcg_murfey[0].id, + atlasImage=dcg_params.atlas, + pixelSize=dcg_params.atlas_pixel_size, + cassetteSlot=dcg_params.sample, + ) + ) + dcg_murfey[0].atlas_id = atlas_id_response["return_value"] + db.add(dcg_murfey[0]) + db.commit() + else: + dcg_parameters = { + "start_time": str(datetime.now()), + "experiment_type": dcg_params.experiment_type, + "experiment_type_id": dcg_params.experiment_type_id, + "tag": dcg_params.tag, + "session_id": session_id, + "atlas": dcg_params.atlas, + "sample": dcg_params.sample, + "atlas_pixel_size": dcg_params.atlas_pixel_size, + } + + if _transport_object: + _transport_object.send( + _transport_object.feedback_queue, + { + "register": "data_collection_group", + **dcg_parameters, + "microscope": instrument_name, + "proposal_code": ispyb_proposal_code, + "proposal_number": ispyb_proposal_number, + "visit_number": ispyb_visit_number, + }, + ) + return dcg_params + + +class DCParameters(BaseModel): + voltage: float + pixel_size_on_image: str + experiment_type: str + image_size_x: int + image_size_y: int + file_extension: str + acquisition_software: str + image_directory: str + tag: str + source: str + magnification: float + total_exposed_dose: Optional[float] = None + c2aperture: Optional[float] = None + exposure_time: Optional[float] = None + slit_width: Optional[float] = None + phase_plate: bool = False + data_collection_tag: str = "" + + +@router.post("/visits/{visit_name}/{session_id}/start_data_collection") +def start_dc( + visit_name, session_id: MurfeySessionID, dc_params: DCParameters, db=murfey_db +): + ispyb_proposal_code = visit_name[:2] + ispyb_proposal_number = visit_name.split("-")[0][2:] + ispyb_visit_number = visit_name.split("-")[-1] + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)).one().instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + logger.info( + f"Starting data collection on microscope {instrument_name!r} " + f"with basepath {sanitise(str(machine_config.rsync_basepath))} and directory {sanitise(dc_params.image_directory)}" + ) + dc_parameters = { + "visit": visit_name, + "image_directory": str( + machine_config.rsync_basepath / dc_params.image_directory + ), + "start_time": str(datetime.now()), + "voltage": dc_params.voltage, + "pixel_size": str(float(dc_params.pixel_size_on_image) * 1e9), + "image_suffix": dc_params.file_extension, + "experiment_type": dc_params.experiment_type, + "image_size_x": dc_params.image_size_x, + "image_size_y": dc_params.image_size_y, + "acquisition_software": dc_params.acquisition_software, + "tag": dc_params.tag, + "source": dc_params.source, + "magnification": dc_params.magnification, + "total_exposed_dose": dc_params.total_exposed_dose, + "c2aperture": dc_params.c2aperture, + "exposure_time": dc_params.exposure_time, + "slit_width": dc_params.slit_width, + "phase_plate": dc_params.phase_plate, + "session_id": session_id, + } + + if _transport_object: + _transport_object.send( + _transport_object.feedback_queue, + { + "register": "data_collection", + **dc_parameters, + "microscope": instrument_name, + "proposal_code": ispyb_proposal_code, + "proposal_number": ispyb_proposal_number, + "visit_number": ispyb_visit_number, + }, + ) + if dc_params.exposure_time: + prom.exposure_time.set(dc_params.exposure_time) + return dc_params + + +class ProcessingJobParameters(BaseModel): + tag: str + source: str + recipe: str + parameters: Dict[str, Any] = {} + experiment_type: str = "spa" + + +@router.post("/visits/{visit_name}/{session_id}/register_processing_job") +def register_proc( + visit_name: str, + session_id: MurfeySessionID, + proc_params: ProcessingJobParameters, + db=murfey_db, +): + proc_parameters: dict = { + "session_id": session_id, + "experiment_type": proc_params.experiment_type, + "recipe": proc_params.recipe, + "source": proc_params.source, + "tag": proc_params.tag, + "job_parameters": { + k: v for k, v in proc_params.parameters.items() if v not in (None, "None") + }, + } + + session_processing_parameters = db.exec( + select(SessionProcessingParameters).where( + SessionProcessingParameters.session_id == session_id + ) + ).all() + + if session_processing_parameters: + job_parameters: dict = proc_parameters["job_parameters"] + job_parameters.update( + { + "gain_ref": session_processing_parameters[0].gain_ref, + "dose_per_frame": session_processing_parameters[0].dose_per_frame, + "eer_fractionation_file": session_processing_parameters[ + 0 + ].eer_fractionation_file, + "symmetry": session_processing_parameters[0].symmetry, + } + ) + proc_parameters["job_parameters"] = job_parameters + + if _transport_object: + _transport_object.send( + _transport_object.feedback_queue, + {"register": "processing_job", **proc_parameters}, + ) + return proc_params + + +spa_router = APIRouter( + prefix="/workflow/spa", + dependencies=[Depends(validate_token)], + tags=["Workflows: SPA"], +) + + +@spa_router.post("/sessions/{session_id}/spa_processing_parameters") +def register_spa_proc_params( + session_id: MurfeySessionID, proc_params: ProcessingParametersSPA, db=murfey_db +): + session_processing_parameters = db.exec( + select(SessionProcessingParameters).where( + SessionProcessingParameters.session_id == session_id + ) + ).all() + if session_processing_parameters: + proc_params.gain_ref = session_processing_parameters[0].gain_ref + proc_params.dose_per_frame = session_processing_parameters[0].dose_per_frame + proc_params.eer_fractionation_file = session_processing_parameters[ + 0 + ].eer_fractionation_file + proc_params.symmetry = session_processing_parameters[0].symmetry + + zocalo_message = { + "register": "spa_processing_parameters", + **dict(proc_params), + "session_id": session_id, + } + if _transport_object: + _transport_object.send(_transport_object.feedback_queue, zocalo_message) + + +class Tag(BaseModel): + tag: str + + +@spa_router.post("/visits/{visit_name}/{session_id}/flush_spa_processing") +def flush_spa_processing( + visit_name: str, session_id: MurfeySessionID, tag: Tag, db=murfey_db +): + zocalo_message = { + "register": "spa.flush_spa_preprocess", + "session_id": session_id, + "tag": tag.tag, + } + if _transport_object: + _transport_object.send(_transport_object.feedback_queue, zocalo_message) + return + + +class SPAProcessFile(BaseModel): + tag: str + path: str + description: str + processing_job: Optional[int] + data_collection_id: Optional[int] + image_number: int + autoproc_program_id: Optional[int] + foil_hole_id: Optional[int] + pixel_size: Optional[float] + dose_per_frame: Optional[float] + mc_binning: Optional[int] = 1 + gain_ref: Optional[str] = None + extract_downscale: bool = True + eer_fractionation_file: Optional[str] = None + source: str = "" + + +@spa_router.post("/visits/{visit_name}/{session_id}/spa_preprocess") +async def request_spa_preprocessing( + visit_name: str, + session_id: MurfeySessionID, + proc_file: SPAProcessFile, + db=murfey_db, +): + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)).one().instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + parts = [secure_filename(p) for p in Path(proc_file.path).parts] + visit_idx = parts.index(visit_name) + core = Path("/") / Path(*parts[: visit_idx + 1]) + ppath = Path("/") / Path(*parts) + sub_dataset = ppath.relative_to(core).parts[0] + extra_path = machine_config.processed_extra_directory + for i, p in enumerate(ppath.parts): + if p.startswith("raw"): + movies_path_index = i + break + else: + raise ValueError(f"{proc_file.path} does not contain a raw directory") + mrc_out = ( + core + / machine_config.processed_directory_name + / sub_dataset + / extra_path + / "MotionCorr" + / "job002" + / "Movies" + / "/".join(ppath.parts[movies_path_index + 1 : -1]) + / str(ppath.stem + "_motion_corrected.mrc") + ) + try: + collected_ids = db.exec( + select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram) + .where(DataCollectionGroup.session_id == session_id) + .where(DataCollectionGroup.tag == proc_file.tag) + .where(DataCollection.dcg_id == DataCollectionGroup.id) + .where(ProcessingJob.dc_id == DataCollection.id) + .where(AutoProcProgram.pj_id == ProcessingJob.id) + .where(ProcessingJob.recipe == "em-spa-preprocess") + ).one() + params = db.exec( + select(SPARelionParameters, SPAFeedbackParameters) + .where(SPARelionParameters.pj_id == collected_ids[2].id) + .where(SPAFeedbackParameters.pj_id == SPARelionParameters.pj_id) + ).one() + proc_params: Optional[dict] = dict(params[0]) + feedback_params = params[1] + except sqlalchemy.exc.NoResultFound: + proc_params = None + try: + foil_hole_id = ( + db.exec( + select(FoilHole, GridSquare) + .where(FoilHole.name == proc_file.foil_hole_id) + .where(FoilHole.session_id == session_id) + .where(GridSquare.id == FoilHole.grid_square_id) + .where(GridSquare.tag == proc_file.tag) + ) + .one()[0] + .id + ) + except Exception as e: + logger.warning( + f"Foil hole ID not found for foil hole {sanitise(str(proc_file.foil_hole_id))}: {e}", + exc_info=True, + ) + foil_hole_id = None + if proc_params: + + detached_ids = [c.id for c in collected_ids] + + murfey_ids = _murfey_id(detached_ids[3], db, number=2, close=False) + + if feedback_params.picker_murfey_id is None: + feedback_params.picker_murfey_id = murfey_ids[1] + db.add(feedback_params) + movie = Movie( + murfey_id=murfey_ids[0], + path=proc_file.path, + image_number=proc_file.image_number, + tag=proc_file.tag, + foil_hole_id=foil_hole_id, + ) + db.add(movie) + db.commit() + db.close() + + if not mrc_out.parent.exists(): + Path(secure_filename(str(mrc_out))).parent.mkdir( + parents=True, exist_ok=True + ) + recipe_name = machine_config.recipes.get( + "em-spa-preprocess", "em-spa-preprocess" + ) + zocalo_message: dict = { + "recipes": [recipe_name], + "parameters": { + "node_creator_queue": machine_config.node_creator_queue, + "dcid": detached_ids[1], + "kv": proc_params["voltage"], + "autoproc_program_id": detached_ids[3], + "movie": proc_file.path, + "mrc_out": str(mrc_out), + "pixel_size": proc_params["angpix"], + "image_number": proc_file.image_number, + "microscope": instrument_name, + "mc_uuid": murfey_ids[0], + "foil_hole_id": foil_hole_id, + "ft_bin": proc_params["motion_corr_binning"], + "fm_dose": proc_params["dose_per_frame"], + "gain_ref": proc_params["gain_ref"], + "picker_uuid": murfey_ids[1], + "session_id": session_id, + "particle_diameter": proc_params["particle_diameter"] or 0, + "fm_int_file": ( + proc_params["eer_fractionation_file"] + if proc_params["eer_fractionation_file"] + else proc_file.eer_fractionation_file + ), + "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, + "cryolo_model_weights": str( + _cryolo_model_path(visit_name, instrument_name) + ), + }, + } + # log.info(f"Sending Zocalo message {zocalo_message}") + if _transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = _transport_object.feedback_queue + _transport_object.send("processing_recipe", zocalo_message) + else: + logger.error( + f"Pe-processing was requested for {sanitise(ppath.name)} but no Zocalo transport object was found" + ) + return proc_file + + else: + for_stash = PreprocessStash( + file_path=str(proc_file.path), + tag=proc_file.tag, + session_id=session_id, + image_number=proc_file.image_number, + mrc_out=str(mrc_out), + eer_fractionation_file=str(proc_file.eer_fractionation_file), + foil_hole_id=foil_hole_id, + ) + db.add(for_stash) + db.commit() + db.close() + + return proc_file + + +tomo_router = APIRouter( + prefix="/workflow/tomo", + dependencies=[Depends(validate_token)], + tags=["Workflows: CryoET"], +) + + +@tomo_router.post("/sessions/{session_id}/tomography_processing_parameters") +def register_tomo_proc_params( + session_id: MurfeySessionID, proc_params: ProcessingParametersTomo, db=murfey_db +): + session_processing_parameters = db.exec( + select(SessionProcessingParameters).where( + SessionProcessingParameters.session_id == session_id + ) + ).all() + if session_processing_parameters: + proc_params.gain_ref = session_processing_parameters[0].gain_ref + proc_params.dose_per_frame = session_processing_parameters[0].dose_per_frame + proc_params.eer_fractionation_file = session_processing_parameters[ + 0 + ].eer_fractionation_file + + zocalo_message = { + "register": "tomography_processing_parameters", + **dict(proc_params), + "session_id": session_id, + } + if _transport_object: + _transport_object.send(_transport_object.feedback_queue, zocalo_message) + + +class Source(BaseModel): + rsync_source: str + + +@tomo_router.post("/visits/{visit_name}/{session_id}/flush_tomography_processing") +def flush_tomography_processing( + visit_name: str, session_id: MurfeySessionID, rsync_source: Source, db=murfey_db +): + zocalo_message = { + "register": "flush_tomography_preprocess", + "session_id": session_id, + "visit_name": visit_name, + "data_collection_group_tag": rsync_source.rsync_source, + } + if _transport_object: + _transport_object.send(_transport_object.feedback_queue, zocalo_message) + return + + +class TiltSeriesInfo(BaseModel): + session_id: int + tag: str + source: str + + +@tomo_router.post("/visits/{visit_name}/tilt_series") +def register_tilt_series( + visit_name: str, tilt_series_info: TiltSeriesInfo, db=murfey_db +): + session_id = tilt_series_info.session_id + if db.exec( + select(TiltSeries) + .where(TiltSeries.session_id == session_id) + .where(TiltSeries.tag == tilt_series_info.tag) + .where(TiltSeries.rsync_source == tilt_series_info.source) + ).all(): + return + tilt_series = TiltSeries( + session_id=session_id, + tag=tilt_series_info.tag, + rsync_source=tilt_series_info.source, + ) + db.add(tilt_series) + db.commit() + + +class TiltSeriesGroupInfo(BaseModel): + tags: List[str] + source: str + tilt_series_lengths: List[int] + + +@tomo_router.post("/sessions/{session_id}/tilt_series_length") +def register_tilt_series_length( + session_id: int, + tilt_series_group: TiltSeriesGroupInfo, + db=murfey_db, +): + tilt_series_db = db.exec( + select(TiltSeries) + .where(col(TiltSeries.tag).in_(tilt_series_group.tags)) + .where(TiltSeries.session_id == session_id) + .where(TiltSeries.rsync_source == tilt_series_group.source) + ).all() + for ts in tilt_series_db: + ts_index = tilt_series_group.tags.index(ts.tag) + ts.tilt_series_length = tilt_series_group.tilt_series_lengths[ts_index] + db.add(ts) + db.commit() + + +class TomoProcessFile(BaseModel): + path: str + description: str + tag: str + image_number: int + pixel_size: float + dose_per_frame: Optional[float] + frame_count: int + tilt_axis: Optional[float] + mc_uuid: Optional[int] = None + voltage: float = 300 + mc_binning: int = 1 + gain_ref: Optional[str] = None + extract_downscale: int = 1 + eer_fractionation_file: Optional[str] = None + group_tag: Optional[str] = None + + +@tomo_router.post("/visits/{visit_name}/{session_id}/tomography_preprocess") +async def request_tomography_preprocessing( + visit_name: str, + session_id: MurfeySessionID, + proc_file: TomoProcessFile, + db=murfey_db, +): + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)).one().instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + visit_idx = Path(proc_file.path).parts.index(visit_name) + core = Path(*Path(proc_file.path).parts[: visit_idx + 1]) + ppath = Path("/".join(secure_filename(p) for p in Path(proc_file.path).parts)) + sub_dataset = "/".join(ppath.relative_to(core).parts[:-1]) + extra_path = machine_config.processed_extra_directory + mrc_out = ( + core + / machine_config.processed_directory_name + / sub_dataset + / extra_path + / "MotionCorr" + / "job002" + / "Movies" + / str(ppath.stem + "_motion_corrected.mrc") + ) + mrc_out = Path("/".join(secure_filename(p) for p in mrc_out.parts)) + + recipe_name = machine_config.recipes.get("em-tomo-preprocess", "em-tomo-preprocess") + + data_collection = db.exec( + select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram) + .where(DataCollectionGroup.session_id == session_id) + .where(DataCollectionGroup.tag == proc_file.group_tag) + .where(DataCollection.tag == proc_file.tag) + .where(DataCollection.dcg_id == DataCollectionGroup.id) + .where(ProcessingJob.dc_id == DataCollection.id) + .where(AutoProcProgram.pj_id == ProcessingJob.id) + .where(ProcessingJob.recipe == recipe_name) + ).all() + if data_collection: + if registered_tilts := db.exec( + select(Tilt).where(Tilt.movie_path == proc_file.path) + ).all(): + if len(registered_tilts) == 1: + if registered_tilts[0].motion_corrected: + return proc_file + dcid = data_collection[0][1].id + appid = data_collection[0][3].id + murfey_ids = _murfey_id(appid, db, number=1, close=False) + if not mrc_out.parent.exists(): + mrc_out.parent.mkdir(parents=True, exist_ok=True) + + session_processing_parameters = db.exec( + select(SessionProcessingParameters).where( + SessionProcessingParameters.session_id == session_id + ) + ).all() + if session_processing_parameters: + proc_file.gain_ref = session_processing_parameters[0].gain_ref + proc_file.dose_per_frame = session_processing_parameters[0].dose_per_frame + proc_file.eer_fractionation_file = session_processing_parameters[ + 0 + ].eer_fractionation_file + + zocalo_message: dict = { + "recipes": [recipe_name], + "parameters": { + "node_creator_queue": machine_config.node_creator_queue, + "dcid": dcid, + # "timestamp": datetime.datetime.now(), + "autoproc_program_id": appid, + "movie": proc_file.path, + "mrc_out": str(mrc_out), + "pixel_size": (proc_file.pixel_size) * 10**10, + "image_number": proc_file.image_number, + "kv": int(proc_file.voltage), + "microscope": instrument_name, + "mc_uuid": murfey_ids[0], + "ft_bin": proc_file.mc_binning, + "fm_dose": proc_file.dose_per_frame, + "frame_count": proc_file.frame_count, + "gain_ref": ( + str(machine_config.rsync_basepath / proc_file.gain_ref) + if proc_file.gain_ref and machine_config.data_transfer_enabled + else proc_file.gain_ref + ), + "fm_int_file": proc_file.eer_fractionation_file, + }, + } + if _transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = _transport_object.feedback_queue + _transport_object.send("processing_recipe", zocalo_message) + else: + logger.error( + f"Pe-processing was requested for {sanitise(ppath.name)} but no Zocalo transport object was found" + ) + return proc_file + else: + for_stash = PreprocessStash( + file_path=str(proc_file.path), + session_id=session_id, + image_number=proc_file.image_number, + mrc_out=str(mrc_out), + tag=proc_file.tag, + group_tag=proc_file.group_tag, + ) + db.add(for_stash) + db.commit() + db.close() + return proc_file + + +@tomo_router.post("/visits/{visit_name}/{session_id}/completed_tilt_series") +def register_completed_tilt_series( + visit_name: str, + session_id: MurfeySessionID, + tilt_series_group: TiltSeriesGroupInfo, + db=murfey_db, +): + tilt_series_db = db.exec( + select(TiltSeries) + .where(col(TiltSeries.tag).in_(tilt_series_group.tags)) + .where(TiltSeries.session_id == session_id) + .where(TiltSeries.rsync_source == tilt_series_group.source) + ).all() + for ts in tilt_series_db: + ts_index = tilt_series_group.tags.index(ts.tag) + ts.tilt_series_length = tilt_series_group.tilt_series_lengths[ts_index] + db.add(ts) + db.commit() + for ts in tilt_series_db: + if ( + check_tilt_series_mc(ts.id) + and not ts.processing_requested + and ts.tilt_series_length > 2 + ): + ts.processing_requested = True + db.add(ts) + + collected_ids = db.exec( + select( + DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram + ) + .where(DataCollectionGroup.session_id == session_id) + .where(DataCollectionGroup.tag == tilt_series_group.source) + .where(DataCollection.tag == ts.tag) + .where(DataCollection.dcg_id == DataCollectionGroup.id) + .where(ProcessingJob.dc_id == DataCollection.id) + .where(AutoProcProgram.pj_id == ProcessingJob.id) + .where(ProcessingJob.recipe == "em-tomo-align") + ).one() + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + tilts = get_all_tilts(ts.id) + ids = get_job_ids(ts.id, collected_ids[3].id) + preproc_params = get_tomo_proc_params(ids.dcgid) + + first_tilt = db.exec( + select(Tilt).where(Tilt.tilt_series_id == ts.id) + ).first() + parts = [secure_filename(p) for p in Path(first_tilt.movie_path).parts] + visit_idx = parts.index(visit_name) + core = Path(*Path(first_tilt.movie_path).parts[: visit_idx + 1]) + ppath = Path( + "/".join(secure_filename(p) for p in Path(first_tilt.movie_path).parts) + ) + sub_dataset = "/".join(ppath.relative_to(core).parts[:-1]) + extra_path = machine_config.processed_extra_directory + stack_file = ( + core + / machine_config.processed_directory_name + / sub_dataset + / extra_path + / "Tomograms" + / "job006" + / "tomograms" + / f"{ts.tag}_stack.mrc" + ) + if not stack_file.parent.exists(): + stack_file.parent.mkdir(parents=True) + tilt_offset = midpoint([float(get_angle(t)) for t in tilts]) + zocalo_message = { + "recipes": ["em-tomo-align"], + "parameters": { + "input_file_list": str([[t, str(get_angle(t))] for t in tilts]), + "path_pattern": "", # blank for now so that it works with the tomo_align service changes + "dcid": ids.dcid, + "appid": ids.appid, + "stack_file": str(stack_file), + "dose_per_frame": preproc_params.dose_per_frame, + "frame_count": preproc_params.frame_count, + "kv": preproc_params.voltage, + "tilt_axis": preproc_params.tilt_axis, + "pixel_size": preproc_params.pixel_size, + "manual_tilt_offset": -tilt_offset, + "node_creator_queue": machine_config.node_creator_queue, + }, + } + if _transport_object: + logger.info(f"Sending Zocalo message for processing: {zocalo_message}") + _transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + else: + logger.info( + f"No transport object found. Zocalo message would be {zocalo_message}" + ) + db.commit() + + +@tomo_router.post("/visits/{visit_name}/rerun_tilt_series") +def register_tilt_series_for_rerun( + visit_name: str, tilt_series_info: TiltSeriesInfo, db=murfey_db +): + """Set processing to false for cases where an extra tilt is found for a series""" + session_id = tilt_series_info.session_id + tilt_series_db = db.exec( + select(TiltSeries) + .where(TiltSeries.session_id == session_id) + .where(TiltSeries.tag == tilt_series_info.tag) + .where(TiltSeries.rsync_source == tilt_series_info.source) + ).all() + for ts in tilt_series_db: + ts.processing_requested = False + db.add(ts) + db.commit() + + +class TiltInfo(BaseModel): + tilt_series_tag: str + movie_path: str + source: str + + +@tomo_router.post("/visits/{visit_name}/{session_id}/tilt") +async def register_tilt( + visit_name: str, session_id: MurfeySessionID, tilt_info: TiltInfo, db=murfey_db +): + def _add_tilt(): + tilt_series_id = ( + db.exec( + select(TiltSeries) + .where(TiltSeries.tag == tilt_info.tilt_series_tag) + .where(TiltSeries.session_id == session_id) + .where(TiltSeries.rsync_source == tilt_info.source) + ) + .one() + .id + ) + if db.exec( + select(Tilt) + .where(Tilt.movie_path == tilt_info.movie_path) + .where(Tilt.tilt_series_id == tilt_series_id) + ).all(): + return + tilt = Tilt(movie_path=tilt_info.movie_path, tilt_series_id=tilt_series_id) + db.add(tilt) + db.commit() + + try: + _add_tilt() + except OperationalError: + await asyncio.sleep(30) + _add_tilt() + + +correlative_router = APIRouter( + prefix="/workflow/correlative", + dependencies=[Depends(validate_token)], + tags=["Workflows: Correlative Imaging"], +) + + +class Sample(BaseModel): + sample_group_id: int + sample_id: int + subsample_id: int + image_path: Optional[Path] + + +@correlative_router.get("/visit/{visit_name}/samples") +def get_samples(visit_name: str, db=ispyb_db) -> List[Sample]: + proposal_id = get_proposal_id(visit_name[:2], visit_name.split("-")[0][2:], db) + samples = ( + db.query(BLSampleGroup, BLSampleGroupHasBLSample, BLSample, BLSubSample) + .join(BLSample, BLSample.blSampleId == BLSampleGroupHasBLSample.blSampleId) + .join( + BLSampleGroup, + BLSampleGroup.blSampleGroupId == BLSampleGroupHasBLSample.blSampleGroupId, + ) + .join(BLSubSample, BLSubSample.blSampleId == BLSample.blSampleId) + .filter(BLSampleGroup.proposalId == proposal_id) + .all() + ) + res = [ + Sample( + sample_group_id=s[1].blSampleGroupId, + sample_id=s[2].blSampleId, + subsample_id=s[3].blSubSampleId, + image_path=s[3].imgFilePath, + ) + for s in samples + ] + return res + + +@correlative_router.post("/visit/{visit_name}/sample_group") +def register_sample_group(visit_name: str, db=ispyb_db) -> dict: + proposal_id = get_proposal_id(visit_name[:2], visit_name.split("-")[0][2:], db=db) + record = BLSampleGroup(proposalId=proposal_id) + if _transport_object: + return _transport_object.do_insert_sample_group(record) + return {"success": False} + + +class BLSampleParameters(BaseModel): + sample_group_id: int + + +@correlative_router.post("/visit/{visit_name}/sample") +def register_sample(visit_name: str, sample_params: BLSampleParameters) -> dict: + record = BLSample() + if _transport_object: + return _transport_object.do_insert_sample(record, sample_params.sample_group_id) + return {"success": False} + + +class BLSubSampleParameters(BaseModel): + sample_id: int + image_path: Optional[Path] = None + + +@correlative_router.post("/visit/{visit_name}/subsample") +def register_subsample( + visit_name: str, subsample_params: BLSubSampleParameters +) -> dict: + record = BLSubSample( + blSampleId=subsample_params.sample_id, imgFilePath=subsample_params.image_path + ) + if _transport_object: + return _transport_object.do_insert_subsample(record) + return {"success": False} + + +class BLSampleImageParameters(BaseModel): + sample_id: int + sample_path: Path + + +@correlative_router.post("/visit/{visit_name}/sample_image") +def register_sample_image( + visit_name: str, sample_image_params: BLSampleImageParameters +) -> dict: + record = BLSampleImage( + blSampleId=sample_image_params.sample_id, + imageFullPath=sample_image_params.image_path, + ) + if _transport_object: + return _transport_object.do_insert_sample_image(record) + return {"success": False} + + +class MillingParameters(BaseModel): + lamella_number: int + images: List[str] + raw_directory: str + + +@correlative_router.post("/visits/{year}/{visit_name}/{session_id}/make_milling_gif") +async def make_gif( + year: int, + visit_name: str, + session_id: int, + gif_params: MillingParameters, + db=murfey_db, +): + instrument_name = ( + db.exec(select(Session).where(Session.id == session_id)).one().instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + output_dir = ( + Path(machine_config.rsync_basepath) + / secure_filename(year) + / secure_filename(visit_name) + / "processed" + ) + output_dir.mkdir(exist_ok=True) + output_dir = output_dir / secure_filename(gif_params.raw_directory) + output_dir.mkdir(exist_ok=True) + output_path = output_dir / f"lamella_{gif_params.lamella_number}_milling.gif" + image_full_paths = [ + output_dir.parent / gif_params.raw_directory / i for i in gif_params.images + ] + if Image is not None: + images = [Image.open(f) for f in image_full_paths] + else: + images = [] + for im in images: + im.thumbnail((512, 512)) + images[0].save( + output_path, + format="GIF", + append_images=images[1:], + save_all=True, + duration=30, + loop=0, + ) + return {"output_gif": str(output_path)} diff --git a/src/murfey/server/demo_api.py b/src/murfey/server/demo_api.py index 01163ec05..c10d6a14f 100644 --- a/src/murfey/server/demo_api.py +++ b/src/murfey/server/demo_api.py @@ -3,32 +3,26 @@ import datetime import logging import os -import random from functools import lru_cache from itertools import count from pathlib import Path from typing import Dict, List, Optional import packaging.version -import sqlalchemy from fastapi import APIRouter, Depends, Request from fastapi.responses import FileResponse, HTMLResponse from ispyb.sqlalchemy import BLSession from PIL import Image from pydantic import BaseModel, BaseSettings from sqlalchemy import func -from sqlmodel import col, select +from sqlmodel import select from werkzeug.utils import secure_filename import murfey.server.api.bootstrap import murfey.server.prometheus as prom -import murfey.server.websocket as ws -import murfey.util.eer from murfey.server import ( _flush_grid_square_records, - _flush_tomography_preprocessing, _murfey_id, - feedback_callback, get_hostname, get_microscope, sanitise, @@ -36,8 +30,13 @@ ) from murfey.server import shutdown as _shutdown from murfey.server import templates -from murfey.server.api import MurfeySessionID -from murfey.server.api.auth import validate_token +from murfey.server.api.auth import MurfeySessionID, validate_token +from murfey.server.api.session_info import Visit +from murfey.server.api.workflow import ( + DCGroupParameters, + DCParameters, + ProcessingJobParameters, +) from murfey.server.murfey_db import murfey_db from murfey.util.config import MachineConfig, from_file, security_from_file from murfey.util.db import ( @@ -57,35 +56,14 @@ SPARelionParameters, Tilt, TiltSeries, - TomographyProcessingParameters, ) from murfey.util.models import ( ClientInfo, - CurrentGainRef, - DCGroupParameters, - DCParameters, FoilHoleParameters, - FractionationParameters, - GainReference, GridSquareParameters, - PostInfo, - ProcessingJobParameters, - ProcessingParametersSPA, - ProcessingParametersTomo, RegistrationMessage, - RsyncerInfo, - RsyncerSource, - SessionInfo, - SPAProcessFile, - SuggestedPathParameters, - TiltInfo, - TiltSeriesGroupInfo, - TiltSeriesInfo, - TomoProcessFile, - Visit, ) from murfey.util.processing_params import default_spa_parameters -from murfey.workflows.spa.picking import _register_picked_particles_use_diameter log = logging.getLogger("murfey.server.demo_api") @@ -217,31 +195,6 @@ def count_number_of_movies(db=murfey_db) -> Dict[str, int]: return {r[0]: r[1] for r in res} -@router.post("/sessions/{session_id}/rsyncer") -def register_rsyncer(session_id: int, rsyncer_info: RsyncerInfo, db=murfey_db): - log.info(f"Registering rsync instance {sanitise(rsyncer_info.source)}") - visit_name = db.exec(select(Session).where(Session.id == session_id)).one().visit - rsync_instance = RsyncInstance( - source=rsyncer_info.source, - session_id=rsyncer_info.session_id, - transferring=rsyncer_info.transferring, - destination=rsyncer_info.destination, - tag=rsyncer_info.tag, - ) - db.add(rsync_instance) - db.commit() - prom.seen_files.labels(rsync_source=rsyncer_info.source, visit=visit_name) - prom.transferred_files.labels(rsync_source=rsyncer_info.source, visit=visit_name) - prom.seen_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).set(0) - prom.transferred_files.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).set(0) - prom.transferred_files_bytes.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).set(0) - return rsyncer_info - - @router.delete("/sessions/{session_id}/rsyncer/{source:path}") def delete_rsyncer(session_id: int, source: str, db=murfey_db): rsync_instance = db.exec( @@ -253,34 +206,6 @@ def delete_rsyncer(session_id: int, source: str, db=murfey_db): db.commit() -@router.post("/sessions/{session_id}/rsyncer_stopped") -def register_stopped_rsyncer( - session_id: int, rsyncer_source: RsyncerSource, db=murfey_db -): - rsyncer = db.exec( - select(RsyncInstance) - .where(RsyncInstance.session_id == session_id) - .where(RsyncInstance.source == rsyncer_source.source) - ).one() - rsyncer.transferring = False - db.add(rsyncer) - db.commit() - - -@router.post("/sessions/{session_id}/rsyncer_started") -def register_restarted_rsyncer( - session_id: int, rsyncer_source: RsyncerSource, db=murfey_db -): - rsyncer = db.exec( - select(RsyncInstance) - .where(RsyncInstance.session_id == session_id) - .where(RsyncInstance.source == rsyncer_source.source) - ).one() - rsyncer.transferring = True - db.add(rsyncer) - db.commit() - - @router.get("/clients/{client_id}/rsyncers") def get_rsyncers_for_client(client_id: int, db=murfey_db): log.info("rsyncers requested") @@ -306,59 +231,6 @@ async def get_session(session_id: MurfeySessionID, db=murfey_db) -> SessionClien return SessionClients(session=session, clients=clients) -@router.post("/visits/{visit_name}/increment_rsync_file_count") -def increment_rsync_file_count( - visit_name: str, rsyncer_info: RsyncerInfo, db=murfey_db -): - rsync_instance = db.exec( - select(RsyncInstance).where( - RsyncInstance.source == rsyncer_info.source, - RsyncInstance.destination == rsyncer_info.destination, - RsyncInstance.session_id == rsyncer_info.session_id, - ) - ).one() - rsync_instance.files_counted += 1 - db.add(rsync_instance) - db.commit() - prom.seen_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).inc( - rsyncer_info.increment_count - ) - - -@router.post("/visits/{visit_name}/increment_rsync_transferred_files") -def increment_rsync_transferred_files( - visit_name: str, rsyncer_info: RsyncerInfo, db=murfey_db -): - rsync_instance = db.exec( - select(RsyncInstance).where( - RsyncInstance.source == rsyncer_info.source, - RsyncInstance.destination == rsyncer_info.destination, - RsyncInstance.session_id == rsyncer_info.session_id, - ) - ).one() - rsync_instance.files_transferred += 1 - db.add(rsync_instance) - db.commit() - - -@router.post("/visits/{visit_name}/increment_rsync_transferred_files_prometheus") -def increment_rsync_transferred_files_prometheus( - visit_name: str, rsyncer_info: RsyncerInfo, db=murfey_db -): - prom.transferred_files.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).inc(rsyncer_info.increment_count) - prom.transferred_files_bytes.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).inc(rsyncer_info.bytes) - prom.transferred_data_files.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).inc(rsyncer_info.increment_data_count) - prom.transferred_data_files_bytes.labels( - rsync_source=rsyncer_info.source, visit=visit_name - ).inc(rsyncer_info.data_bytes) - - class ProcessingDetails(BaseModel): data_collection_group: DataCollectionGroup data_collections: List[DataCollection] @@ -414,106 +286,6 @@ def _parse(ps, i, dcg_id): ] -@router.post("/sessions/{session_id}/spa_processing_parameters") -def register_spa_proc_params( - session_id: MurfeySessionID, proc_params: ProcessingParametersSPA, db=murfey_db -): - log.info( - f"Registration request for SPA processing parameters with data: {proc_params.json()}" - ) - try: - collected_ids = db.exec( - select( - DataCollectionGroup, - DataCollection, - ProcessingJob, - AutoProcProgram, - ) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollectionGroup.tag == proc_params.tag) - .where(DataCollection.dcg_id == DataCollectionGroup.id) - .where(ProcessingJob.dc_id == DataCollection.id) - .where(AutoProcProgram.pj_id == ProcessingJob.id) - .where(ProcessingJob.recipe == "em-spa-preprocess") - ).one() - current_gain_ref = ( - db.exec(select(Session).where(Session.id == session_id)) - .one() - .current_gain_ref - ) - params = SPARelionParameters( - pj_id=collected_ids[2].id, - angpix=proc_params.pixel_size_on_image, - dose_per_frame=proc_params.dose_per_frame, - gain_ref=current_gain_ref or proc_params.gain_ref, - voltage=proc_params.voltage, - motion_corr_binning=proc_params.motion_corr_binning, - eer_grouping=proc_params.eer_fractionation, - symmetry=proc_params.symmetry, - particle_diameter=proc_params.particle_diameter, - downscale=proc_params.downscale, - boxsize=proc_params.boxsize, - small_boxsize=proc_params.small_boxsize, - mask_diameter=proc_params.mask_diameter, - ) - feedback_params = SPAFeedbackParameters( - pj_id=collected_ids[2].id, - estimate_particle_diameter=proc_params.particle_diameter is None, - hold_class2d=False, - hold_class3d=False, - class_selection_score=0, - star_combination_job=0, - initial_model="", - next_job=0, - ) - except Exception as e: - log.warning(f"registration failed: {e}") - return - db.add(params) - db.add(feedback_params) - db.commit() - - -@router.post("/sessions/{session_id}/tomography_processing_parameters") -def register_tomo_proc_params( - session_id: MurfeySessionID, proc_params: ProcessingParametersTomo, db=murfey_db -): - log.info( - f"Registering tomography preprocessing parameters {sanitise(proc_params.tag)}, {sanitise(proc_params.tilt_series_tag)}" - ) - collected_ids = db.exec( - select( - DataCollectionGroup, - DataCollection, - ProcessingJob, - AutoProcProgram, - ) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollectionGroup.tag == proc_params.tag) - .where(DataCollection.tag == proc_params.tilt_series_tag) - .where(DataCollection.dcg_id == DataCollectionGroup.id) - .where(ProcessingJob.dc_id == DataCollection.id) - .where(AutoProcProgram.pj_id == ProcessingJob.id) - .where(ProcessingJob.recipe == "em-tomo-preprocess") - ).one() - if not db.exec( - select(func.count(TomographyProcessingParameters.dcg_id)).where( - TomographyProcessingParameters.dcg_id == collected_ids[0].id - ) - ).one(): - params = TomographyProcessingParameters( - dcg_id=collected_ids[0].id, - pixel_size=proc_params.pixel_size_on_image, - dose_per_frame=proc_params.dose_per_frame, - gain_ref=proc_params.gain_ref, - motion_corr_binning=proc_params.motion_corr_binning, - voltage=proc_params.voltage, - ) - db.add(params) - db.commit() - db.close() - - @router.get("/clients/{client_id}/spa_processing_parameters") def get_spa_proc_params(client_id: int, db=murfey_db) -> List[dict]: params = db.exec( @@ -698,54 +470,6 @@ def register_foil_hole( db.close() -@router.post("/visits/{visit_name}/tilt_series") -def register_tilt_series( - visit_name: str, tilt_series_info: TiltSeriesInfo, db=murfey_db -): - session_id = ( - db.exec( - select(ClientEnvironment).where( - ClientEnvironment.client_id == tilt_series_info.client_id - ) - ) - .one() - .session_id - ) - tilt_series = TiltSeries( - session_id=session_id, - tag=tilt_series_info.tag, - rsync_source=tilt_series_info.source, - ) - db.add(tilt_series) - db.commit() - - -@router.post("/visits/{visit_name}/{client_id}/completed_tilt_series") -def register_completed_tilt_series( - visit_name: str, - client_id: int, - tilt_series_group: TiltSeriesGroupInfo, - db=murfey_db, -): - session_id = ( - db.exec( - select(ClientEnvironment).where(ClientEnvironment.client_id == client_id) - ) - .one() - .session_id - ) - tilt_series_db = db.exec( - select(TiltSeries) - .where(col(TiltSeries.tag).in_(tilt_series_group.tags)) - .where(TiltSeries.session_id == session_id) - .where(TiltSeries.rsync_source == tilt_series_group.source) - ).all() - for ts in tilt_series_db: - ts.complete = True - db.add(ts) - db.commit() - - @router.get("/clients/{client_id}/tilt_series/{tilt_series_tag}/tilts") def get_tilts(client_id: int, tilt_series_tag: str, db=murfey_db): res = db.exec( @@ -764,52 +488,6 @@ def get_tilts(client_id: int, tilt_series_tag: str, db=murfey_db): return tilts -@router.post("/visits/{visit_name}/{client_id}/tilt") -def register_tilt(visit_name: str, client_id: int, tilt_info: TiltInfo, db=murfey_db): - session_id = ( - db.exec( - select(ClientEnvironment).where(ClientEnvironment.client_id == client_id) - ) - .one() - .session_id - ) - tilt_series_id = ( - db.exec( - select(TiltSeries) - .where(TiltSeries.tag == tilt_info.tilt_series_tag) - .where(TiltSeries.session_id == session_id) - .where(TiltSeries.rsync_source == tilt_info.source) - ) - .one() - .id - ) - tilt = Tilt(movie_path=tilt_info.movie_path, tilt_series_id=tilt_series_id) - db.add(tilt) - db.commit() - - -# @router.get("/instruments/{instrument_name}/visits_raw", response_model=List[Visit]) -# def get_current_visits(instrument_name: str): -# return [ -# Visit( -# start=datetime.datetime.now(), -# end=datetime.datetime.now() + datetime.timedelta(days=1), -# session_id=1, -# name="cm31111-2", -# beamline="m12", -# proposal_title="Nothing of importance", -# ), -# Visit( -# start=datetime.datetime.now(), -# end=datetime.datetime.now() + datetime.timedelta(days=1), -# session_id=1, -# name="cm31111-3", -# beamline="m12", -# proposal_title="Nothing of importance", -# ), -# ] - - @router.get("/instruments/{instrument_name}/visits_raw", response_model=List[Visit]) def get_current_visits(instrument_name: str, db=murfey.server.ispyb.DB): return murfey.server.ispyb.get_all_ongoing_visits(instrument_name, db) @@ -962,235 +640,10 @@ def flush_spa_processing( return -@router.post("/visits/{visit_name}/{session_id}/spa_preprocess") -async def request_spa_preprocessing( - visit_name: str, - session_id: MurfeySessionID, - proc_file: SPAProcessFile, - db=murfey_db, -): - parts = [secure_filename(p) for p in Path(proc_file.path).parts] - visit_idx = parts.index(visit_name) - core = Path("/") / Path(*parts[: visit_idx + 1]) - ppath = Path("/") / Path(*parts) - sub_dataset = ppath.relative_to(core).parts[0] - for i, p in enumerate(ppath.parts): - if p.startswith("raw"): - movies_path_index = i - break - else: - raise ValueError(f"{proc_file.path} does not contain a raw directory") - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)).one().instrument_name - ) - mrc_out = ( - core - / machine_config[instrument_name].processed_directory_name - / sub_dataset - / "MotionCorr" - / "job002" - / "Movies" - / "/".join(ppath.parts[movies_path_index + 1 : -1]) - / str(ppath.stem + "_motion_corrected.mrc") - ) - try: - collected_ids = db.exec( - select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollectionGroup.tag == proc_file.tag) - .where(DataCollection.dcg_id == DataCollectionGroup.id) - .where(ProcessingJob.dc_id == DataCollection.id) - .where(AutoProcProgram.pj_id == ProcessingJob.id) - .where(ProcessingJob.recipe == "em-spa-preprocess") - ).one() - params = db.exec( - select(SPARelionParameters, SPAFeedbackParameters) - .where(SPARelionParameters.pj_id == collected_ids[2].id) - .where(SPAFeedbackParameters.pj_id == SPARelionParameters.pj_id) - ).one() - proc_params: dict | None = dict(params[0]) - feedback_params = params[1] - except sqlalchemy.exc.NoResultFound: - proc_params = None - try: - foil_hole_id = ( - db.exec( - select(FoilHole, GridSquare) - .where(FoilHole.name == proc_file.foil_hole_id) - .where(FoilHole.session_id == session_id) - .where(GridSquare.id == FoilHole.grid_square_id) - .where(GridSquare.tag == proc_file.tag) - ) - .one()[0] - .id - ) - except Exception: - foil_hole_id = None - if proc_params: - collected_ids = db.exec( - select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram) - .where( - DataCollectionGroup.session_id == session_id - and DataCollectionGroup.tag == "spa" - ) - .where(DataCollectionGroup.tag == proc_file.tag) - .where(DataCollection.dcg_id == DataCollectionGroup.id) - .where(ProcessingJob.dc_id == DataCollection.id) - .where(AutoProcProgram.pj_id == ProcessingJob.id) - .where(ProcessingJob.recipe == "em-spa-preprocess") - ).one() - - detached_ids = [c.id for c in collected_ids] - - murfey_ids = _murfey_id(detached_ids[3], db, number=2) - - feedback_params.picker_murfey_id = murfey_ids[1] - db.add(feedback_params) - movie = Movie( - murfey_id=murfey_ids[0], - path=proc_file.path, - image_number=proc_file.image_number, - tag=proc_file.tag, - foil_hole_id=foil_hole_id, - ) - db.add(movie) - db.commit() - - if not mrc_out.parent.exists(): - Path(secure_filename(mrc_out)).parent.mkdir(parents=True) - log.info("Sending Zocalo message") - movie = db.exec(select(Movie).where(Movie.murfey_id == murfey_ids[0])).one() - movie.preprocessed = True - db.add(movie) - db.commit() - _register_picked_particles_use_diameter( - { - "session_id": session_id, - "extraction_parameters": { - "micrographs_file": "MotionCorr/job002/micrographs.star", - "coord_list_file": "AutoPick/job007/coords.star", - "extract_file": "Extract/job008/particles.star", - "ctf_values": { - "CtfMaxResolution": 4, - "CtfFigureOfMerit": 0.8, - "DefocusU": 1, - "DefocusV": 1, - "DefocusAngle": 10, - "CtfImage": "CtfFind/job006/ctf.mrc", - }, - }, - "particle_diameters": [random.randint(20, 30) for i in range(400)], - "program_id": detached_ids[3], - }, - _db=db, - demo=True, - ) - prom.preprocessed_movies.labels(processing_job=detached_ids[2]).inc() - - else: - for_stash = PreprocessStash( - file_path=str(proc_file.path), - tag=proc_file.tag, - session_id=session_id, - image_number=proc_file.image_number, - mrc_out=str(mrc_out), - foil_hole_id=foil_hole_id, - ) - db.add(for_stash) - db.commit() - - return proc_file - - class Source(BaseModel): rsync_source: str -@router.post("/visits/{visit_name}/{client_id}/flush_tomography_processing") -def flush_tomography_processing( - visit_name: str, client_id: int, rsync_source: Source, db=murfey_db -): - zocalo_message = { - "register": "flush_tomography_preprocess", - "client_id": client_id, - "visit_name": visit_name, - "data_collection_group_tag": rsync_source.rsync_source, - } - _flush_tomography_preprocessing(zocalo_message) - return - - -@router.post("/visits/{visit_name}/{client_id}/tomography_preprocess") -async def request_tomography_preprocessing( - visit_name: str, client_id: int, proc_file: TomoProcessFile, db=murfey_db -): - if not sanitise_path(Path(proc_file.path)).exists(): - log.warning( - f"{sanitise(str(proc_file.path))} has not been transferred before preprocessing" - ) - log.info(f"Tomo preprocesing requested for {sanitise(str(proc_file.path))}") - visit_idx = Path(proc_file.path).parts.index(visit_name) - core = Path(*Path(proc_file.path).parts[: visit_idx + 1]) - ppath = Path("/".join(secure_filename(p) for p in Path(proc_file.path).parts)) - sub_dataset = ( - ppath.relative_to(core).parts[0] - if len(ppath.relative_to(core).parts) > 1 - else "" - ) - mrc_out = ( - core - / "processed" - / sub_dataset - / "MotionCorr" - / str(ppath.stem + "_motion_corrected.mrc") - ) - mrc_out = Path("/".join(secure_filename(p) for p in mrc_out.parts)) - session_id = ( - db.exec( - select(ClientEnvironment).where(ClientEnvironment.client_id == client_id) - ) - .one() - .session_id - ) - data_collection = db.exec( - select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollectionGroup.id == DataCollection.dcg_id) - .where(DataCollection.tag == proc_file.tag) - .where(ProcessingJob.dc_id == DataCollection.id) - .where(AutoProcProgram.pj_id == ProcessingJob.id) - .where(ProcessingJob.recipe == "em-tomo-preprocess") - ).all() - if data_collection: - if not mrc_out.parent.exists(): - mrc_out.parent.mkdir(parents=True) - feedback_callback( - {}, - { - "register": "motion_corrected", - "movie": str(proc_file.path), - "mrc_out": str(mrc_out), - "movie_id": proc_file.mc_uuid, - "fm_int_file": proc_file.eer_fractionation_file, - "program_id": data_collection[0][3].id, - }, - ) - await ws.manager.broadcast(f"Pre-processing requested for {ppath.name}") - mrc_out.touch() - else: - for_stash = PreprocessStash( - file_path=str(proc_file.path), - session_id=session_id, - image_number=proc_file.image_number, - mrc_out=str(mrc_out), - tag=proc_file.tag, - ) - db.add(for_stash) - db.commit() - db.close() - return proc_file - - @router.get("/version") def get_version(client_version: str = ""): result = { @@ -1218,92 +671,6 @@ def shutdown(): return {"success": True} -@router.post("/visits/{visit_name}/{session_id}/suggested_path") -def suggest_path( - visit_name: str, session_id: int, params: SuggestedPathParameters, db=murfey_db -): - count: int | None = router.raw_count - instrument_name = ( - db.exec(select(Session).where(Session.id == session_id)).one().instrument_name - ) - rsync_basepath = ( - machine_config[instrument_name].rsync_basepath - if machine_config - else Path(f"/dls/{get_microscope()}") - ) - check_path = rsync_basepath / params.base_path - check_path = check_path.parent / f"{check_path.stem}{count}{check_path.suffix}" - check_path = check_path.resolve() - - # Check for path traversal attempt - if not str(check_path).startswith(str(rsync_basepath)): - raise Exception(f"Path traversal attempt detected: {str(check_path)!r}") - - # Check previous year to account for the year rolling over during data collection - if not sanitise_path(check_path).exists(): - base_path_parts = list(params.base_path.parts) - for part in base_path_parts: - # Find the path part corresponding to the year - if len(part) == 4 and part.isdigit(): - year_idx = base_path_parts.index(part) - base_path_parts[year_idx] = str(int(part) - 1) - base_path = "/".join(base_path_parts) - check_path_prev = check_path - check_path = rsync_basepath / base_path - check_path = check_path.parent / f"{check_path.stem}{count}{check_path.suffix}" - check_path = check_path.resolve() - - # Check for path traversal attempt - if not str(check_path).startswith(str(rsync_basepath)): - raise Exception(f"Path traversal attempt detected: {str(check_path)!r}") - - # If visit is not in the previous year either, it's a genuine error - if not check_path.exists(): - log_message = sanitise( - "Unable to find current visit folder under " - f"{str(check_path_prev)!r} or {str(check_path)!r}" - ) - log.error(log_message) - raise FileNotFoundError(log_message) - - check_path_name = check_path.name - while sanitise_path(check_path).exists(): - count = count + 1 if count else 2 - check_path = check_path.parent / f"{check_path_name}{count}" - router.raw_count += 1 - if params.touch: - sanitise_path(check_path).mkdir(mode=0o750) - if params.extra_directory: - (sanitise_path(check_path) / secure_filename(params.extra_directory)).mkdir( - mode=0o750 - ) - return { - "suggested_path": check_path.relative_to( - machine_config[instrument_name].rsync_basepath - ) - } - - -@router.get("/sessions/{session_id}/data_collection_groups") -def get_dc_groups( - session_id: MurfeySessionID, db=murfey_db -) -> Dict[str, DataCollectionGroup]: - data_collection_groups = db.exec( - select(DataCollectionGroup).where(DataCollectionGroup.session_id == session_id) - ).all() - return {dcg.tag: dcg for dcg in data_collection_groups} - - -@router.get("/sessions/{session_id}/data_collection_groups/{dcgid}/data_collections") -def get_data_collections( - session_id: MurfeySessionID, dcgid: int, db=murfey_db -) -> List[DataCollection]: - data_collections = db.exec( - select(DataCollection).where(DataCollection.dcg_id == dcgid) - ).all() - return data_collections - - @router.post("/visits/{visit_name}/{session_id}/register_data_collection_group") def register_dc_group( visit_name: str, @@ -1443,30 +810,6 @@ def register_proc( return proc_params -@router.post("/sessions/{session_id}/process_gain") -async def process_gain( - session_id: MurfeySessionID, gain_reference_params: GainReference, db=murfey_db -): - visit_name = db.exec(select(Session).where(Session.id == session_id)).one().visit - if machine_config.get("rsync_basepath"): - filepath = ( - Path(machine_config["rsync_basepath"]) - / str(datetime.datetime.now().year) - / visit_name - ) - else: - return {"gain_ref": None} - gain_ref_folder = machine_config.get("gain_directory_name", "processing") - gain_ref_out = ( - (filepath / gain_ref_folder / f"gain_{gain_reference_params.tag}.mrc") - if gain_reference_params.tag - else (filepath / gain_ref_folder / "gain.mrc") - ) - return { - "gain_ref": gain_ref_out.relative_to(Path(machine_config["rsync_basepath"])) - } - - @router.get("/new_client_id/") async def new_client_id(db=murfey_db): clients = db.exec(select(ClientEnvironment)).all() @@ -1518,25 +861,6 @@ async def get_sessions_by_instrument_name( return sessions -@router.post("/instruments/{instrument_name}/clients/{client_id}/session") -def link_client_to_session( - instrument_name: str, client_id: int, sess: SessionInfo, db=murfey_db -): - sid = sess.session_id - if sid is None: - s = Session(name=sess.session_name, instrument_name=instrument_name) - db.add(s) - db.commit() - sid = s.id - client = db.exec( - select(ClientEnvironment).where(ClientEnvironment.client_id == client_id) - ).one() - client.session_id = sid - db.add(client) - db.commit() - return sid - - @router.delete("/clients/{client_id}/session") def remove_session(client_id: int, db=murfey_db): client = db.exec( @@ -1595,40 +919,6 @@ def remove_session_by_id(session_id: MurfeySessionID, db=murfey_db): return -@router.post("/visits/{visit_name}/{session_id}/eer_fractionation_file") -async def write_eer_fractionation_file( - visit_name: str, session_id: int, fractionation_params: FractionationParameters -) -> dict: - file_path = ( - Path(machine_config["rsync_basepath"]) - / str(datetime.datetime.now().year) - / secure_filename(visit_name) - / secure_filename(fractionation_params.fractionation_file_name) - ) - log.info(f"EER fractionation file {file_path} creation requested") - if file_path.is_file(): - return {"eer_fractionation_file": str(file_path)} - - if fractionation_params.num_frames: - num_eer_frames = fractionation_params.num_frames - elif ( - fractionation_params.eer_path - and sanitise_path(Path(fractionation_params.eer_path)).is_file() - ): - num_eer_frames = murfey.util.eer.num_frames(Path(fractionation_params.eer_path)) - else: - log.warning( - f"EER fractionation unable to find {sanitise_path(Path(fractionation_params.eer_path)) if fractionation_params.eer_path else None} " - f"or use {sanitise(str(fractionation_params.num_frames))} frames" - ) - return {"eer_fractionation_file": None} - with open(file_path, "w") as frac_file: - frac_file.write( - f"{num_eer_frames} {fractionation_params.fractionation} {fractionation_params.dose_per_frame}" - ) - return {"eer_fractionation_file": str(file_path)} - - @router.post("/visits/{visit_name}/monitoring/{on}") def change_monitoring_status(visit_name: str, on: int): prom.monitoring_switch.labels(visit=visit_name) @@ -1700,12 +990,6 @@ async def get_tiff(visit_name: str, session_id: int, tiff_path: str, db=murfey_d return FileResponse(path=test_path) -@router.post("/failed_client_post") -def failed_client_post(post_info: PostInfo): - log.info("Post failed") - return - - @router.post("/instruments/{instrument_name}/visits/{visit}/session/{name}") def create_session(instrument_name: str, visit: str, name: str, db=murfey_db) -> int: s = Session(name=name, visit=visit, instrument_name=instrument_name) @@ -1713,13 +997,3 @@ def create_session(instrument_name: str, visit: str, name: str, db=murfey_db) -> db.commit() sid = s.id return sid - - -@router.put("/sessions/{session_id}/current_gain_ref") -def update_current_gain_ref( - session_id: MurfeySessionID, new_gain_ref: CurrentGainRef, db=murfey_db -): - session = db.exec(select(Session).where(Session.id == session_id)).one() - session.current_gain_ref = new_gain_ref.path - db.add(session) - db.commit() diff --git a/src/murfey/server/feedback.py b/src/murfey/server/feedback.py new file mode 100644 index 000000000..8ef95bf6f --- /dev/null +++ b/src/murfey/server/feedback.py @@ -0,0 +1,2534 @@ +""" +Contains functions related to how Murfey will interact with messages that it relays +and receives. These functions should eventually be refactored into workflows module +files, but are stored here for now to prevent unnecessary dependency chains due to +being written and stored in 'murfey.server.__init__'. +""" + +from __future__ import annotations + +import logging +import math +import subprocess +import time +from datetime import datetime +from functools import partial, singledispatch +from importlib.metadata import EntryPoint # For type hinting only +from pathlib import Path +from typing import Dict, List, NamedTuple, Tuple + +import mrcfile +import numpy as np +from backports.entry_points_selectable import entry_points +from ispyb.sqlalchemy._auto_db_schema import ( + Atlas, + AutoProcProgram, + Base, + DataCollection, + DataCollectionGroup, + ProcessingJob, + ProcessingJobParameter, +) +from sqlalchemy import func +from sqlalchemy.exc import ( + InvalidRequestError, + OperationalError, + PendingRollbackError, + SQLAlchemyError, +) +from sqlalchemy.orm.exc import ObjectDeletedError +from sqlmodel import Session, create_engine, select + +import murfey.server +import murfey.server.prometheus as prom +import murfey.util.db as db +from murfey.server.ispyb import ISPyBSession, get_session_id +from murfey.server.murfey_db import url # murfey_db +from murfey.util import sanitise +from murfey.util.config import ( + MachineConfig, + get_machine_config, + get_microscope, + get_security_config, +) +from murfey.util.processing_params import default_spa_parameters +from murfey.util.tomo import midpoint + +logger = logging.getLogger("murfey.server.feedback") + + +try: + _url = url(get_security_config()) + engine = create_engine(_url) + murfey_db = Session(engine, expire_on_commit=False) +except Exception: + murfey_db = None + + +class ExtendedRecord(NamedTuple): + record: Base # type: ignore + record_params: List[Base] # type: ignore + + +class JobIDs(NamedTuple): + dcgid: int + dcid: int + pid: int + appid: int + + +def get_angle(tilt_file_name: str) -> float: + for p in Path(tilt_file_name).name.split("_"): + if "." in p: + return float(p) + raise ValueError(f"Tilt angle not found for file {tilt_file_name}") + + +def check_tilt_series_mc(tilt_series_id: int) -> bool: + results = murfey_db.exec( + select(db.Tilt, db.TiltSeries) + .where(db.Tilt.tilt_series_id == db.TiltSeries.id) + .where(db.TiltSeries.id == tilt_series_id) + ).all() + return ( + all(r[0].motion_corrected for r in results) + and len(results) >= results[0][1].tilt_series_length + and results[0][1].tilt_series_length > 0 + ) + + +def get_all_tilts(tilt_series_id: int) -> List[str]: + complete_results = murfey_db.exec( + select(db.Tilt, db.TiltSeries, db.Session) + .where(db.Tilt.tilt_series_id == db.TiltSeries.id) + .where(db.TiltSeries.id == tilt_series_id) + .where(db.TiltSeries.session_id == db.Session.id) + ).all() + if not complete_results: + return [] + instrument_name = complete_results[0][2].instrument_name + results = [r[0] for r in complete_results] + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + + def _mc_path(mov_path: Path) -> str: + for p in mov_path.parts: + if "-" in p and p.startswith(("bi", "nr", "nt", "cm", "sw")): + visit_name = p + break + else: + raise ValueError(f"No visit found in {mov_path}") + visit_idx = Path(mov_path).parts.index(visit_name) + core = Path(*Path(mov_path).parts[: visit_idx + 1]) + ppath = Path(mov_path) + sub_dataset = "/".join(ppath.relative_to(core).parts[:-1]) + extra_path = machine_config.processed_extra_directory + mrc_out = ( + core + / machine_config.processed_directory_name + / sub_dataset + / extra_path + / "MotionCorr" + / "job002" + / "Movies" + / str(ppath.stem + "_motion_corrected.mrc") + ) + return str(mrc_out) + + return [_mc_path(Path(r.movie_path)) for r in results] + + +def get_job_ids(tilt_series_id: int, appid: int) -> JobIDs: + results = murfey_db.exec( + select( + db.TiltSeries, + db.AutoProcProgram, + db.ProcessingJob, + db.DataCollection, + db.DataCollectionGroup, + db.Session, + ) + .where(db.TiltSeries.id == tilt_series_id) + .where(db.DataCollection.tag == db.TiltSeries.tag) + .where(db.ProcessingJob.id == db.AutoProcProgram.pj_id) + .where(db.AutoProcProgram.id == appid) + .where(db.ProcessingJob.dc_id == db.DataCollection.id) + .where(db.DataCollectionGroup.id == db.DataCollection.dcg_id) + .where(db.Session.id == db.TiltSeries.session_id) + ).all() + return JobIDs( + dcgid=results[0][4].id, + dcid=results[0][3].id, + pid=results[0][2].id, + appid=results[0][1].id, + ) + + +def get_tomo_proc_params(dcg_id: int, *args) -> db.TomographyProcessingParameters: + results = murfey_db.exec( + select(db.TomographyProcessingParameters).where( + db.TomographyProcessingParameters.dcg_id == dcg_id + ) + ).one() + return results + + +def _murfey_id(app_id: int, _db, number: int = 1, close: bool = True) -> List[int]: + murfey_ledger = [db.MurfeyLedger(app_id=app_id) for _ in range(number)] + for ml in murfey_ledger: + _db.add(ml) + _db.commit() + # There is a race condition between the IDs being read back from the database + # after the insert and the insert being synchronised so allow multiple attempts + attempts = 0 + while attempts < 100: + try: + for m in murfey_ledger: + _db.refresh(m) + res = [m.id for m in murfey_ledger if m.id is not None] + break + except (ObjectDeletedError, InvalidRequestError): + pass + attempts += 1 + time.sleep(0.1) + else: + raise RuntimeError( + "Maximum number of attempts exceeded when producing new Murfey IDs" + ) + if close: + _db.close() + return res + + +def _murfey_class2ds( + murfey_ids: List[int], particles_file: str, app_id: int, _db, close: bool = False +): + pj_id = _pj_id(app_id, _db, recipe="em-spa-class2d") + class2ds = [ + db.Class2D( + class_number=i, + particles_file=particles_file, + pj_id=pj_id, + murfey_id=mid, + ) + for i, mid in enumerate(murfey_ids) + ] + for c in class2ds: + _db.add(c) + _db.commit() + if close: + _db.close() + + +def _murfey_class3ds(murfey_ids: List[int], particles_file: str, app_id: int, _db): + pj_id = _pj_id(app_id, _db, recipe="em-spa-class3d") + class3ds = [ + db.Class3D( + class_number=i, + particles_file=str(Path(particles_file).parent), + pj_id=pj_id, + murfey_id=mid, + ) + for i, mid in enumerate(murfey_ids) + ] + for c in class3ds: + _db.add(c) + _db.commit() + _db.close() + + +def _murfey_refine(murfey_id: int, refine_dir: str, tag: str, app_id: int, _db): + pj_id = _pj_id(app_id, _db, recipe="em-spa-refine") + refine3d = db.Refine3D( + tag=tag, + refine_dir=refine_dir, + pj_id=pj_id, + murfey_id=murfey_id, + ) + _db.add(refine3d) + _db.commit() + _db.close() + + +def _2d_class_murfey_ids(particles_file: str, app_id: int, _db) -> Dict[str, int]: + pj_id = ( + _db.exec(select(db.AutoProcProgram).where(db.AutoProcProgram.id == app_id)) + .one() + .pj_id + ) + classes = _db.exec( + select(db.Class2D).where( + db.Class2D.particles_file == particles_file and db.Class2D.pj_id == pj_id + ) + ).all() + return {str(cl.class_number): cl.murfey_id for cl in classes} + + +def _3d_class_murfey_ids(particles_file: str, app_id: int, _db) -> Dict[str, int]: + pj_id = ( + _db.exec(select(db.AutoProcProgram).where(db.AutoProcProgram.id == app_id)) + .one() + .pj_id + ) + classes = _db.exec( + select(db.Class3D).where( + db.Class3D.particles_file == str(Path(particles_file).parent) + and db.Class3D.pj_id == pj_id + ) + ).all() + return {str(cl.class_number): cl.murfey_id for cl in classes} + + +def _refine_murfey_id(refine_dir: str, tag: str, app_id: int, _db) -> Dict[str, int]: + pj_id = ( + _db.exec(select(db.AutoProcProgram).where(db.AutoProcProgram.id == app_id)) + .one() + .pj_id + ) + refined_class = _db.exec( + select(db.Refine3D) + .where(db.Refine3D.refine_dir == refine_dir) + .where(db.Refine3D.pj_id == pj_id) + .where(db.Refine3D.tag == tag) + ).one() + return refined_class.murfey_id + + +def _app_id(pj_id: int, _db) -> int: + return ( + _db.exec(select(db.AutoProcProgram).where(db.AutoProcProgram.pj_id == pj_id)) + .one() + .id + ) + + +def _pj_id(app_id: int, _db, recipe: str = "") -> int: + if recipe: + dc_id = ( + _db.exec( + select(db.AutoProcProgram, db.ProcessingJob) + .where(db.AutoProcProgram.id == app_id) + .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id) + ) + .one()[1] + .dc_id + ) + pj_id = ( + _db.exec( + select(db.ProcessingJob) + .where(db.ProcessingJob.dc_id == dc_id) + .where(db.ProcessingJob.recipe == recipe) + ) + .one() + .id + ) + else: + pj_id = ( + _db.exec(select(db.AutoProcProgram).where(db.AutoProcProgram.id == app_id)) + .one() + .pj_id + ) + return pj_id + + +def _get_spa_params( + app_id: int, _db +) -> Tuple[db.SPARelionParameters, db.SPAFeedbackParameters]: + pj_id = _pj_id(app_id, _db, recipe="em-spa-preprocess") + relion_params = _db.exec( + select(db.SPARelionParameters).where(db.SPARelionParameters.pj_id == pj_id) + ).one() + feedback_params = _db.exec( + select(db.SPAFeedbackParameters).where(db.SPAFeedbackParameters.pj_id == pj_id) + ).one() + _db.expunge(relion_params) + _db.expunge(feedback_params) + return relion_params, feedback_params + + +def _release_2d_hold(message: dict, _db=murfey_db): + relion_params, feedback_params = _get_spa_params(message["program_id"], _db) + if not feedback_params.star_combination_job: + feedback_params.star_combination_job = feedback_params.next_job + ( + 3 if default_spa_parameters.do_icebreaker_jobs else 2 + ) + pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-class2d") + if feedback_params.rerun_class2d: + first_class2d = _db.exec( + select(db.Class2DParameters).where(db.Class2DParameters.pj_id == pj_id) + ).first() + instrument_name = ( + _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + zocalo_message: dict = { + "parameters": { + "particles_file": first_class2d.particles_file, + "class2d_dir": message["job_dir"], + "batch_is_complete": first_class2d.complete, + "batch_size": first_class2d.batch_size, + "particle_diameter": relion_params.particle_diameter, + "mask_diameter": relion_params.mask_diameter or 0, + "combine_star_job_number": feedback_params.star_combination_job, + "autoselect_min_score": feedback_params.class_selection_score or 0, + "autoproc_program_id": message["program_id"], + "nr_iter": default_spa_parameters.nr_iter_2d, + "nr_classes": default_spa_parameters.nr_classes_2d, + "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, + "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, + "picker_id": feedback_params.picker_ispyb_id, + "class_uuids": _2d_class_murfey_ids( + first_class2d.particles_file, message["program_id"], _db + ), + "class2d_grp_uuid": _db.exec( + select(db.Class2DParameters) + .where( + db.Class2DParameters.particles_file + == first_class2d.particles_file + ) + .where(db.Class2DParameters.pj_id == pj_id) + ) + .one() + .murfey_id, + "session_id": message["session_id"], + "node_creator_queue": machine_config.node_creator_queue, + }, + "recipes": ["em-spa-class2d"], + } + if first_class2d.complete: + feedback_params.next_job += ( + 4 if default_spa_parameters.do_icebreaker_jobs else 3 + ) + feedback_params.rerun_class2d = False + _db.add(feedback_params) + if first_class2d.complete: + _db.delete(first_class2d) + _db.commit() + _db.close() + if murfey.server._transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = murfey.server._transport_object.feedback_queue + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + else: + feedback_params.hold_class2d = False + _db.add(feedback_params) + _db.commit() + _db.close() + + +def _release_3d_hold(message: dict, _db=murfey_db): + pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") + pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-class3d") + relion_params = _db.exec( + select(db.SPARelionParameters).where( + db.SPARelionParameters.pj_id == pj_id_params + ) + ).one() + feedback_params = _db.exec( + select(db.SPAFeedbackParameters).where( + db.SPAFeedbackParameters.pj_id == pj_id_params + ) + ).one() + class3d_params = _db.exec( + select(db.Class3DParameters).where(db.Class3DParameters.pj_id == pj_id) + ).one() + if class3d_params.run: + instrument_name = ( + _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + zocalo_message: dict = { + "parameters": { + "particles_file": class3d_params.particles_file, + "class3d_dir": class3d_params.class3d_dir, + "batch_size": class3d_params.batch_size, + "symmetry": relion_params.symmetry, + "particle_diameter": relion_params.particle_diameter, + "mask_diameter": relion_params.mask_diameter or 0, + "do_initial_model": False if feedback_params.initial_model else True, + "initial_model_file": feedback_params.initial_model, + "picker_id": feedback_params.picker_ispyb_id, + "class_uuids": _3d_class_murfey_ids( + class3d_params.particles_file, _app_id(pj_id, _db), _db + ), + "class3d_grp_uuid": _db.exec( + select(db.Class3DParameters) + .where( + db.Class3DParameters.particles_file + == class3d_params.particles_file + ) + .where(db.Class3DParameters.pj_id == pj_id) + ) + .one() + .murfey_id, + "nr_iter": default_spa_parameters.nr_iter_3d, + "initial_model_iterations": default_spa_parameters.nr_iter_ini_model, + "nr_classes": default_spa_parameters.nr_classes_3d, + "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, + "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, + "session_id": message["session_id"], + "autoproc_program_id": _app_id( + _pj_id(message["program_id"], _db, recipe="em-spa-class3d"), _db + ), + "node_creator_queue": machine_config.node_creator_queue, + }, + "recipes": ["em-spa-class3d"], + } + if murfey.server._transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = murfey.server._transport_object.feedback_queue + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + class3d_params.run = False + _db.add(class3d_params) + else: + feedback_params.hold_class3d = False + _db.add(feedback_params) + _db.commit() + _db.close() + + +def _release_refine_hold(message: dict, _db=murfey_db): + pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") + pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-refine") + relion_params = _db.exec( + select(db.SPARelionParameters).where( + db.SPARelionParameters.pj_id == pj_id_params + ) + ).one() + feedback_params = _db.exec( + select(db.SPAFeedbackParameters).where( + db.SPAFeedbackParameters.pj_id == pj_id_params + ) + ).one() + refine_params = _db.exec( + select(db.RefineParameters) + .where(db.RefineParameters.pj_id == pj_id) + .where(db.RefineParameters.tag == "first") + ).one() + symmetry_refine_params = _db.exec( + select(db.RefineParameters) + .where(db.RefineParameters.pj_id == pj_id) + .where(db.RefineParameters.tag == "symmetry") + ).one() + if refine_params.run: + instrument_name = ( + _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + zocalo_message: dict = { + "parameters": { + "refine_job_dir": refine_params.refine_dir, + "class3d_dir": refine_params.class3d_dir, + "class_number": refine_params.class_number, + "pixel_size": relion_params.angpix, + "particle_diameter": relion_params.particle_diameter, + "mask_diameter": relion_params.mask_diameter or 0, + "symmetry": relion_params.symmetry, + "node_creator_queue": machine_config.node_creator_queue, + "nr_iter": default_spa_parameters.nr_iter_3d, + "picker_id": feedback_params.picker_ispyb_id, + "refined_class_uuid": _refine_murfey_id( + refine_dir=refine_params.refine_dir, + tag=refine_params.tag, + app_id=_app_id(pj_id, _db), + _db=_db, + ), + "refined_grp_uuid": refine_params.murfey_id, + "symmetry_refined_class_uuid": _refine_murfey_id( + refine_dir=symmetry_refine_params.refine_dir, + tag=symmetry_refine_params.tag, + app_id=_app_id(pj_id, _db), + _db=_db, + ), + "symmetry_refined_grp_uuid": symmetry_refine_params.murfey_id, + "session_id": message["session_id"], + "autoproc_program_id": _app_id( + _pj_id(message["program_id"], _db, recipe="em-spa-refine"), _db + ), + }, + "recipes": ["em-spa-refine"], + } + if murfey.server._transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = murfey.server._transport_object.feedback_queue + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + refine_params.run = False + _db.add(refine_params) + else: + feedback_params.hold_refine = False + _db.add(feedback_params) + _db.commit() + _db.close() + + +def _register_incomplete_2d_batch(message: dict, _db=murfey_db, demo: bool = False): + """Received first batch from particle selection service""" + # the general parameters are stored using the preprocessing auto proc program ID + logger.info("Registering incomplete particle batch for 2D classification") + instrument_name = ( + _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") + pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-class2d") + relion_params = _db.exec( + select(db.SPARelionParameters).where( + db.SPARelionParameters.pj_id == pj_id_params + ) + ).one() + feedback_params = _db.exec( + select(db.SPAFeedbackParameters).where( + db.SPAFeedbackParameters.pj_id == pj_id_params + ) + ).one() + if feedback_params.hold_class2d: + feedback_params.rerun_class2d = True + _db.add(feedback_params) + _db.commit() + _db.close() + return + feedback_params.next_job = 10 if default_spa_parameters.do_icebreaker_jobs else 7 + feedback_params.hold_class2d = True + relion_options = dict(relion_params) + other_options = dict(feedback_params) + if other_options["picker_ispyb_id"] is None: + logger.info("No ISPyB particle picker ID yet") + feedback_params.hold_class2d = False + _db.add(feedback_params) + _db.commit() + _db.expunge(feedback_params) + return + _db.add(feedback_params) + _db.commit() + _db.expunge(feedback_params) + class2d_message = message.get("class2d_message") + assert isinstance(class2d_message, dict) + if not _db.exec( + select(func.count(db.Class2DParameters.particles_file)) + .where(db.Class2DParameters.particles_file == class2d_message["particles_file"]) + .where(db.Class2DParameters.pj_id == pj_id) + ).one(): + class2d_params = db.Class2DParameters( + pj_id=pj_id, + murfey_id=_murfey_id(message["program_id"], _db)[0], + particles_file=class2d_message["particles_file"], + class2d_dir=class2d_message["class2d_dir"], + batch_size=class2d_message["batch_size"], + complete=False, + ) + _db.add(class2d_params) + _db.commit() + murfey_ids = _murfey_id(message["program_id"], _db, number=50) + _murfey_class2ds( + murfey_ids, class2d_message["particles_file"], message["program_id"], _db + ) + zocalo_message: dict = { + "parameters": { + "particles_file": class2d_message["particles_file"], + "class2d_dir": f"{class2d_message['class2d_dir']}{other_options['next_job']:03}", + "batch_is_complete": False, + "particle_diameter": relion_options["particle_diameter"], + "combine_star_job_number": -1, + "picker_id": other_options["picker_ispyb_id"], + "nr_iter": default_spa_parameters.nr_iter_2d, + "batch_size": default_spa_parameters.batch_size_2d, + "nr_classes": default_spa_parameters.nr_classes_2d, + "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, + "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, + "mask_diameter": 0, + "class_uuids": _2d_class_murfey_ids( + class2d_message["particles_file"], _app_id(pj_id, _db), _db + ), + "class2d_grp_uuid": _db.exec( + select(db.Class2DParameters).where( + db.Class2DParameters.particles_file + == class2d_message["particles_file"] + and db.Class2DParameters.pj_id == pj_id + ) + ) + .one() + .murfey_id, + "session_id": message["session_id"], + "autoproc_program_id": _app_id( + _pj_id(message["program_id"], _db, recipe="em-spa-class2d"), _db + ), + "node_creator_queue": machine_config.node_creator_queue, + }, + "recipes": ["em-spa-class2d"], + } + if murfey.server._transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = murfey.server._transport_object.feedback_queue + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + logger.info("2D classification requested") + if demo: + logger.info("Incomplete 2D batch registered in demo mode") + if not _db.exec( + select(func.count(db.Class2DParameters.particles_file)).where( + db.Class2DParameters.particles_file == class2d_message["particles_file"] + and db.Class2DParameters.pj_id == pj_id + and db.Class2DParameters.complete + ) + ).one(): + _register_complete_2d_batch(message, _db=_db, demo=demo) + message["class2d_message"]["particles_file"] = ( + message["class2d_message"]["particles_file"] + "_new" + ) + _register_complete_2d_batch(message, _db=_db, demo=demo) + _db.close() + + +def _register_complete_2d_batch(message: dict, _db=murfey_db, demo: bool = False): + """Received full batch from particle selection service""" + instrument_name = ( + _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + class2d_message = message.get("class2d_message") + assert isinstance(class2d_message, dict) + pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") + pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-class2d") + relion_params = _db.exec( + select(db.SPARelionParameters).where( + db.SPARelionParameters.pj_id == pj_id_params + ) + ).one() + feedback_params = _db.exec( + select(db.SPAFeedbackParameters).where( + db.SPAFeedbackParameters.pj_id == pj_id_params + ) + ).one() + _db.expunge(relion_params) + _db.expunge(feedback_params) + if feedback_params.hold_class2d or feedback_params.picker_ispyb_id is None: + feedback_params.rerun_class2d = True + _db.add(feedback_params) + _db.commit() + # If waiting then save the message + if _db.exec( + select(func.count(db.Class2DParameters.particles_file)) + .where(db.Class2DParameters.pj_id == pj_id) + .where( + db.Class2DParameters.particles_file == class2d_message["particles_file"] + ) + ).one(): + class2d_params = _db.exec( + select(db.Class2DParameters) + .where(db.Class2DParameters.pj_id == pj_id) + .where( + db.Class2DParameters.particles_file + == class2d_message["particles_file"] + ) + ).one() + class2d_params.complete = True + _db.add(class2d_params) + _db.commit() + _db.close() + else: + class2d_params = db.Class2DParameters( + pj_id=pj_id, + murfey_id=_murfey_id(message["program_id"], _db)[0], + particles_file=class2d_message["particles_file"], + class2d_dir=class2d_message["class2d_dir"], + batch_size=class2d_message["batch_size"], + ) + _db.add(class2d_params) + _db.commit() + _db.close() + murfey_ids = _murfey_id(_app_id(pj_id, _db), _db, number=50) + _murfey_class2ds( + murfey_ids, class2d_message["particles_file"], _app_id(pj_id, _db), _db + ) + if demo: + _register_class_selection( + {"session_id": message["session_id"], "class_selection_score": 0.5}, + _db=_db, + demo=demo, + ) + elif not feedback_params.class_selection_score: + # For the first batch, start a container and set the database to wait + job_number_after_first_batch = ( + 10 if default_spa_parameters.do_icebreaker_jobs else 7 + ) + if ( + feedback_params.next_job is not None + and feedback_params.next_job < job_number_after_first_batch + ): + feedback_params.next_job = job_number_after_first_batch + if not feedback_params.star_combination_job: + feedback_params.star_combination_job = feedback_params.next_job + ( + 3 if default_spa_parameters.do_icebreaker_jobs else 2 + ) + if _db.exec( + select(func.count(db.Class2DParameters.particles_file)) + .where(db.Class2DParameters.pj_id == pj_id) + .where( + db.Class2DParameters.particles_file == class2d_message["particles_file"] + ) + ).one(): + class_uuids = _2d_class_murfey_ids( + class2d_message["particles_file"], _app_id(pj_id, _db), _db + ) + class2d_grp_uuid = ( + _db.exec( + select(db.Class2DParameters) + .where(db.Class2DParameters.pj_id == pj_id) + .where( + db.Class2DParameters.particles_file + == class2d_message["particles_file"] + ) + ) + .one() + .murfey_id + ) + else: + class_uuids = { + str(i + 1): m + for i, m in enumerate(_murfey_id(_app_id(pj_id, _db), _db, number=50)) + } + class2d_grp_uuid = _murfey_id(_app_id(pj_id, _db), _db)[0] + zocalo_message: dict = { + "parameters": { + "particles_file": class2d_message["particles_file"], + "class2d_dir": f"{class2d_message['class2d_dir']}{feedback_params.next_job:03}", + "batch_is_complete": True, + "particle_diameter": relion_params.particle_diameter, + "mask_diameter": relion_params.mask_diameter or 0, + "combine_star_job_number": feedback_params.star_combination_job, + "autoselect_min_score": 0, + "picker_id": feedback_params.picker_ispyb_id, + "class_uuids": class_uuids, + "class2d_grp_uuid": class2d_grp_uuid, + "nr_iter": default_spa_parameters.nr_iter_2d, + "batch_size": default_spa_parameters.batch_size_2d, + "nr_classes": default_spa_parameters.nr_classes_2d, + "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, + "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, + "session_id": message["session_id"], + "autoproc_program_id": _app_id( + _pj_id(message["program_id"], _db, recipe="em-spa-class2d"), _db + ), + "node_creator_queue": machine_config.node_creator_queue, + }, + "recipes": ["em-spa-class2d"], + } + if murfey.server._transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = murfey.server._transport_object.feedback_queue + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + feedback_params.hold_class2d = True + feedback_params.next_job += ( + 4 if default_spa_parameters.do_icebreaker_jobs else 3 + ) + _db.add(feedback_params) + _db.commit() + _db.close() + else: + # Send all other messages on to a container + if _db.exec( + select(func.count(db.Class2DParameters.particles_file)) + .where(db.Class2DParameters.pj_id == pj_id) + .where( + db.Class2DParameters.particles_file == class2d_message["particles_file"] + ) + ).one(): + class_uuids = _2d_class_murfey_ids( + class2d_message["particles_file"], _app_id(pj_id, _db), _db + ) + class2d_grp_uuid = ( + _db.exec( + select(db.Class2DParameters) + .where(db.Class2DParameters.pj_id == pj_id) + .where( + db.Class2DParameters.particles_file + == class2d_message["particles_file"] + ) + ) + .one() + .murfey_id + ) + else: + class_uuids = { + str(i + 1): m + for i, m in enumerate(_murfey_id(_app_id(pj_id, _db), _db, number=50)) + } + class2d_grp_uuid = _murfey_id(_app_id(pj_id, _db), _db)[0] + zocalo_message = { + "parameters": { + "particles_file": class2d_message["particles_file"], + "class2d_dir": f"{class2d_message['class2d_dir']}{feedback_params.next_job:03}", + "batch_is_complete": True, + "particle_diameter": relion_params.particle_diameter, + "mask_diameter": relion_params.mask_diameter or 0, + "combine_star_job_number": feedback_params.star_combination_job, + "autoselect_min_score": feedback_params.class_selection_score or 0, + "picker_id": feedback_params.picker_ispyb_id, + "class_uuids": class_uuids, + "class2d_grp_uuid": class2d_grp_uuid, + "nr_iter": default_spa_parameters.nr_iter_2d, + "batch_size": default_spa_parameters.batch_size_2d, + "nr_classes": default_spa_parameters.nr_classes_2d, + "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, + "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, + "session_id": message["session_id"], + "autoproc_program_id": _app_id( + _pj_id(message["program_id"], _db, recipe="em-spa-class2d"), _db + ), + "node_creator_queue": machine_config.node_creator_queue, + }, + "recipes": ["em-spa-class2d"], + } + if murfey.server._transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = murfey.server._transport_object.feedback_queue + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + feedback_params.next_job += ( + 3 if default_spa_parameters.do_icebreaker_jobs else 2 + ) + _db.add(feedback_params) + _db.commit() + _db.close() + + +def _flush_class2d( + session_id: int, + app_id: int, + _db, + relion_params: db.SPARelionParameters | None = None, + feedback_params: db.SPAFeedbackParameters | None = None, +): + instrument_name = ( + _db.exec(select(db.Session).where(db.Session.id == session_id)) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + if not relion_params or feedback_params: + pj_id_params = _pj_id(app_id, _db, recipe="em-spa-preprocess") + if not relion_params: + relion_params = _db.exec( + select(db.SPARelionParameters).where( + db.SPARelionParameters.pj_id == pj_id_params + ) + ).one() + _db.expunge(relion_params) + if not feedback_params: + feedback_params = _db.exec( + select(db.SPAFeedbackParameters).where( + db.SPAFeedbackParameters.pj_id == pj_id_params + ) + ).one() + _db.expunge(feedback_params) + if not relion_params or not feedback_params: + return + pj_id = _pj_id(app_id, _db, recipe="em-spa-class2d") + class2d_db = _db.exec( + select(db.Class2DParameters) + .where(db.Class2DParameters.pj_id == pj_id) + .where(db.Class2DParameters.complete) + ).all() + if not feedback_params.next_job: + feedback_params.next_job = ( + 10 if default_spa_parameters.do_icebreaker_jobs else 7 + ) + if not feedback_params.star_combination_job: + feedback_params.star_combination_job = feedback_params.next_job + ( + 3 if default_spa_parameters.do_icebreaker_jobs else 2 + ) + for saved_message in class2d_db: + # Send all held Class2D messages on with the selection score added + _db.expunge(saved_message) + zocalo_message: dict = { + "parameters": { + "particles_file": saved_message.particles_file, + "class2d_dir": f"{saved_message.class2d_dir}{feedback_params.next_job:03}", + "batch_is_complete": True, + "batch_size": saved_message.batch_size, + "particle_diameter": relion_params.particle_diameter, + "mask_diameter": relion_params.mask_diameter or 0, + "combine_star_job_number": feedback_params.star_combination_job, + "autoselect_min_score": feedback_params.class_selection_score or 0, + "picker_id": feedback_params.picker_ispyb_id, + "class_uuids": _2d_class_murfey_ids( + saved_message.particles_file, _app_id(pj_id, _db), _db + ), + "class2d_grp_uuid": saved_message.murfey_id, + "nr_iter": default_spa_parameters.nr_iter_2d, + "nr_classes": default_spa_parameters.nr_classes_2d, + "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, + "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, + "session_id": session_id, + "autoproc_program_id": _app_id(pj_id, _db), + "node_creator_queue": machine_config.node_creator_queue, + }, + "recipes": ["em-spa-class2d"], + } + if murfey.server._transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = murfey.server._transport_object.feedback_queue + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + feedback_params.next_job += ( + 3 if default_spa_parameters.do_icebreaker_jobs else 2 + ) + _db.delete(saved_message) + _db.add(feedback_params) + _db.commit() + + +def _register_class_selection(message: dict, _db=murfey_db, demo: bool = False): + """Received selection score from class selection service""" + pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") + pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-class2d") + relion_params = _db.exec( + select(db.SPARelionParameters).where( + db.SPARelionParameters.pj_id == pj_id_params + ) + ).one() + class2d_db = _db.exec( + select(db.Class2DParameters).where(db.Class2DParameters.pj_id == pj_id) + ).all() + # Add the class selection score to the database + feedback_params = _db.exec( + select(db.SPAFeedbackParameters).where( + db.SPAFeedbackParameters.pj_id == pj_id_params + ) + ).one() + _db.expunge(feedback_params) + + if feedback_params.picker_ispyb_id is None: + selection_stash = db.SelectionStash( + pj_id=pj_id, + class_selection_score=message["class_selection_score"] or 0, + ) + _db.add(selection_stash) + _db.commit() + _db.close() + return + + feedback_params.class_selection_score = message.get("class_selection_score") or 0 + feedback_params.hold_class2d = False + next_job = feedback_params.next_job + if demo: + for saved_message in class2d_db: + # Send all held Class2D messages on with the selection score added + _db.expunge(saved_message) + particles_file = saved_message.particles_file + logger.info("Complete 2D classification registered in demo mode") + _register_3d_batch( + { + "session_id": message["session_id"], + "class3d_message": { + "particles_file": particles_file, + "class3d_dir": "Class3D", + "batch_size": 50000, + }, + }, + _db=_db, + demo=demo, + ) + logger.info("3D classification registered in demo mode") + _register_3d_batch( + { + "session_id": message["session_id"], + "class3d_message": { + "particles_file": particles_file + "_new", + "class3d_dir": "Class3D", + "batch_size": 50000, + }, + }, + _db=_db, + demo=demo, + ) + _register_initial_model( + { + "session_id": message["session_id"], + "initial_model": "InitialModel/job015/model.mrc", + }, + _db=_db, + demo=demo, + ) + next_job += 3 if default_spa_parameters.do_icebreaker_jobs else 2 + feedback_params.next_job = next_job + _db.close() + else: + _flush_class2d( + message["session_id"], + message["program_id"], + _db, + relion_params=relion_params, + feedback_params=feedback_params, + ) + _db.add(feedback_params) + for sm in class2d_db: + _db.delete(sm) + _db.commit() + _db.close() + + +def _find_initial_model(visit: str, machine_config: MachineConfig) -> Path | None: + if machine_config.initial_model_search_directory: + visit_directory = ( + machine_config.rsync_basepath / str(datetime.now().year) / visit + ) + possible_models = [ + p + for p in ( + visit_directory / machine_config.initial_model_search_directory + ).glob("*.mrc") + if "rescaled" not in p.name + ] + if possible_models: + return sorted(possible_models, key=lambda x: x.stat().st_ctime)[-1] + return None + + +def _downscaled_box_size( + particle_diameter: int, pixel_size: float +) -> Tuple[int, float]: + box_size = int(math.ceil(1.2 * particle_diameter)) + box_size = box_size + box_size % 2 + for small_box_pix in ( + 48, + 64, + 96, + 128, + 160, + 192, + 256, + 288, + 300, + 320, + 360, + 384, + 400, + 420, + 450, + 480, + 512, + 640, + 768, + 896, + 1024, + ): + # Don't go larger than the original box + if small_box_pix > box_size: + return box_size, pixel_size + # If Nyquist freq. is better than 8.5 A, use this downscaled box, else step size + small_box_angpix = pixel_size * box_size / small_box_pix + if small_box_angpix < 4.25: + return small_box_pix, small_box_angpix + raise ValueError(f"Box size is too large: {box_size}") + + +def _resize_intial_model( + downscaled_box_size: int, + downscaled_pixel_size: float, + input_path: Path, + output_path: Path, + executables: Dict[str, str], + env: Dict[str, str], +) -> None: + if executables.get("relion_image_handler"): + comp_proc = subprocess.run( + [ + f"{executables['relion_image_handler']}", + "--i", + str(input_path), + "--new_box", + str(downscaled_box_size), + "--rescale_angpix", + str(downscaled_pixel_size), + "--o", + str(output_path), + ], + capture_output=True, + text=True, + env=env, + ) + with mrcfile.open(output_path) as rescaled_mrc: + rescaled_mrc.header.cella = ( + downscaled_pixel_size, + downscaled_pixel_size, + downscaled_pixel_size, + ) + if comp_proc.returncode: + logger.error( + f"Resizing initial model {input_path} failed \n {comp_proc.stdout}" + ) + return None + + +def _register_3d_batch(message: dict, _db=murfey_db, demo: bool = False): + """Received 3d batch from class selection service""" + class3d_message = message.get("class3d_message") + assert isinstance(class3d_message, dict) + instrument_name = ( + _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") + pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-class3d") + relion_params = _db.exec( + select(db.SPARelionParameters).where( + db.SPARelionParameters.pj_id == pj_id_params + ) + ).one() + relion_options = dict(relion_params) + feedback_params = _db.exec( + select(db.SPAFeedbackParameters).where( + db.SPAFeedbackParameters.pj_id == pj_id_params + ) + ).one() + other_options = dict(feedback_params) + + visit_name = ( + _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) + .one() + .visit + ) + + provided_initial_model = _find_initial_model(visit_name, machine_config) + if provided_initial_model and not feedback_params.initial_model: + rescaled_initial_model_path = ( + provided_initial_model.parent + / f"{provided_initial_model.stem}_rescaled_{pj_id}{provided_initial_model.suffix}" + ) + if not rescaled_initial_model_path.is_file(): + _resize_intial_model( + *_downscaled_box_size( + message["particle_diameter"], + relion_options["angpix"], + ), + provided_initial_model, + rescaled_initial_model_path, + machine_config.external_executables, + machine_config.external_environment, + ) + feedback_params.initial_model = str(rescaled_initial_model_path) + other_options["initial_model"] = str(rescaled_initial_model_path) + next_job = feedback_params.next_job + class3d_dir = ( + f"{class3d_message['class3d_dir']}{(feedback_params.next_job+1):03}" + ) + feedback_params.next_job += 1 + _db.add(feedback_params) + _db.commit() + + class3d_grp_uuid = _murfey_id(message["program_id"], _db)[0] + class_uuids = _murfey_id(message["program_id"], _db, number=4) + class3d_params = db.Class3DParameters( + pj_id=pj_id, + murfey_id=class3d_grp_uuid, + particles_file=class3d_message["particles_file"], + class3d_dir=class3d_dir, + batch_size=class3d_message["batch_size"], + ) + _db.add(class3d_params) + _db.commit() + _murfey_class3ds( + class_uuids, + class3d_message["particles_file"], + message["program_id"], + _db, + ) + + if feedback_params.hold_class3d: + # If waiting then save the message + class3d_params = _db.exec( + select(db.Class3DParameters).where(db.Class3DParameters.pj_id == pj_id) + ).one() + class3d_params.run = True + class3d_params.particles_file = class3d_message["particles_file"] + class3d_params.batch_size = class3d_message["batch_size"] + _db.add(class3d_params) + _db.commit() + _db.close() + elif not feedback_params.initial_model: + # For the first batch, start a container and set the database to wait + next_job = feedback_params.next_job + class3d_dir = ( + f"{class3d_message['class3d_dir']}{(feedback_params.next_job+1):03}" + ) + class3d_grp_uuid = _murfey_id(message["program_id"], _db)[0] + class_uuids = _murfey_id(message["program_id"], _db, number=4) + class3d_params = db.Class3DParameters( + pj_id=pj_id, + murfey_id=class3d_grp_uuid, + particles_file=class3d_message["particles_file"], + class3d_dir=class3d_dir, + batch_size=class3d_message["batch_size"], + ) + _db.add(class3d_params) + _db.commit() + _murfey_class3ds( + class_uuids, class3d_message["particles_file"], message["program_id"], _db + ) + + feedback_params.hold_class3d = True + next_job += 2 + feedback_params.next_job = next_job + zocalo_message: dict = { + "parameters": { + "particles_file": class3d_message["particles_file"], + "class3d_dir": class3d_dir, + "batch_size": class3d_message["batch_size"], + "symmetry": relion_options["symmetry"], + "particle_diameter": relion_options["particle_diameter"], + "mask_diameter": relion_options["mask_diameter"] or 0, + "do_initial_model": True, + "picker_id": other_options["picker_ispyb_id"], + "class_uuids": {i + 1: m for i, m in enumerate(class_uuids)}, + "class3d_grp_uuid": class3d_grp_uuid, + "nr_iter": default_spa_parameters.nr_iter_3d, + "initial_model_iterations": default_spa_parameters.nr_iter_ini_model, + "nr_classes": default_spa_parameters.nr_classes_3d, + "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, + "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, + "session_id": message["session_id"], + "autoproc_program_id": _app_id( + _pj_id(message["program_id"], _db, recipe="em-spa-class3d"), _db + ), + "node_creator_queue": machine_config.node_creator_queue, + }, + "recipes": ["em-spa-class3d"], + } + if murfey.server._transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = murfey.server._transport_object.feedback_queue + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + _db.add(feedback_params) + _db.commit() + _db.close() + else: + # Send all other messages on to a container + class3d_params = _db.exec( + select(db.Class3DParameters).where(db.Class3DParameters.pj_id == pj_id) + ).one() + zocalo_message = { + "parameters": { + "particles_file": class3d_message["particles_file"], + "class3d_dir": class3d_params.class3d_dir, + "batch_size": class3d_message["batch_size"], + "symmetry": relion_options["symmetry"], + "particle_diameter": relion_options["particle_diameter"], + "mask_diameter": relion_options["mask_diameter"] or 0, + "do_initial_model": False, + "initial_model_file": other_options["initial_model"], + "picker_id": other_options["picker_ispyb_id"], + "class_uuids": _3d_class_murfey_ids( + class3d_params.particles_file, _app_id(pj_id, _db), _db + ), + "class3d_grp_uuid": class3d_params.murfey_id, + "nr_iter": default_spa_parameters.nr_iter_3d, + "initial_model_iterations": default_spa_parameters.nr_iter_ini_model, + "nr_classes": default_spa_parameters.nr_classes_3d, + "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs, + "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d, + "session_id": message["session_id"], + "autoproc_program_id": _app_id( + _pj_id(message["program_id"], _db, recipe="em-spa-class3d"), _db + ), + "node_creator_queue": machine_config.node_creator_queue, + }, + "recipes": ["em-spa-class3d"], + } + if murfey.server._transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = murfey.server._transport_object.feedback_queue + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + feedback_params.hold_class3d = True + _db.add(feedback_params) + _db.commit() + _db.close() + + +def _register_initial_model(message: dict, _db=murfey_db, demo: bool = False): + """Received initial model from 3d classification service""" + pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") + # Add the initial model file to the database + feedback_params = _db.exec( + select(db.SPAFeedbackParameters).where( + db.SPAFeedbackParameters.pj_id == pj_id_params + ) + ).one() + feedback_params.initial_model = message.get("initial_model") + _db.add(feedback_params) + _db.commit() + _db.close() + + +def _flush_tomography_preprocessing(message: dict): + session_id = message["session_id"] + instrument_name = ( + murfey_db.exec(select(db.Session).where(db.Session.id == session_id)) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + stashed_files = murfey_db.exec( + select(db.PreprocessStash) + .where(db.PreprocessStash.session_id == session_id) + .where(db.PreprocessStash.group_tag == message["data_collection_group_tag"]) + ).all() + if not stashed_files: + return + collected_ids = murfey_db.exec( + select( + db.DataCollectionGroup, + ) + .where(db.DataCollectionGroup.session_id == session_id) + .where(db.DataCollectionGroup.tag == message["data_collection_group_tag"]) + ).first() + proc_params = get_tomo_proc_params(collected_ids.id) + if not proc_params: + visit_name = message["visit_name"].replace("\r\n", "").replace("\n", "") + logger.warning( + f"No tomography processing parameters found for Murfey session {sanitise(str(message['session_id']))} on visit {sanitise(visit_name)}" + ) + return + + recipe_name = machine_config.recipes.get("em-tomo-preprocess", "em-tomo-preprocess") + + for f in stashed_files: + collected_ids = murfey_db.exec( + select( + db.DataCollectionGroup, + db.DataCollection, + db.ProcessingJob, + db.AutoProcProgram, + ) + .where(db.DataCollectionGroup.session_id == session_id) + .where(db.DataCollectionGroup.tag == message["data_collection_group_tag"]) + .where(db.DataCollection.dcg_id == db.DataCollectionGroup.id) + .where(db.DataCollection.tag == f.tag) + .where(db.ProcessingJob.dc_id == db.DataCollection.id) + .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id) + .where(db.ProcessingJob.recipe == recipe_name) + ).one() + detached_ids = [c.id for c in collected_ids] + + murfey_ids = _murfey_id(detached_ids[3], murfey_db, number=1, close=False) + p = Path(f.mrc_out) + if not p.parent.exists(): + p.parent.mkdir(parents=True) + movie = db.Movie( + murfey_id=murfey_ids[0], + path=f.file_path, + image_number=f.image_number, + tag=f.tag, + ) + murfey_db.add(movie) + zocalo_message: dict = { + "recipes": [recipe_name], + "parameters": { + "node_creator_queue": machine_config.node_creator_queue, + "dcid": detached_ids[1], + "autoproc_program_id": detached_ids[3], + "movie": f.file_path, + "mrc_out": f.mrc_out, + "pixel_size": proc_params.pixel_size, + "kv": proc_params.voltage, + "image_number": f.image_number, + "microscope": get_microscope(), + "mc_uuid": murfey_ids[0], + "ft_bin": proc_params.motion_corr_binning, + "fm_dose": proc_params.dose_per_frame, + "frame_count": proc_params.frame_count, + "gain_ref": ( + str(machine_config.rsync_basepath / proc_params.gain_ref) + if proc_params.gain_ref + else proc_params.gain_ref + ), + "fm_int_file": proc_params.eer_fractionation_file or "", + }, + } + logger.info( + f"Launching tomography preprocessing with Zocalo message: {zocalo_message}" + ) + if murfey.server._transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = murfey.server._transport_object.feedback_queue + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + else: + feedback_callback( + {}, + { + "register": "motion_corrected", + "movie": f.file_path, + "mrc_out": f.mrc_out, + "movie_id": murfey_ids[0], + "program_id": detached_ids[3], + }, + ) + murfey_db.delete(f) + murfey_db.commit() + murfey_db.close() + + +def _flush_grid_square_records(message: dict, _db=murfey_db, demo: bool = False): + tag = message["tag"] + session_id = message["session_id"] + gs_ids = [] + for gs in _db.exec( + select(db.GridSquare) + .where(db.GridSquare.session_id == session_id) + .where(db.GridSquare.tag == tag) + ).all(): + gs_ids.append(gs.id) + if demo: + logger.info(f"Flushing grid square {gs.name}") + for i in gs_ids: + _flush_foil_hole_records(i, _db=_db, demo=demo) + + +def _flush_foil_hole_records(grid_square_id: int, _db=murfey_db, demo: bool = False): + for fh in _db.exec( + select(db.FoilHole).where(db.FoilHole.grid_square_id == grid_square_id) + ).all(): + if demo: + logger.info(f"Flushing foil hole: {fh.name}") + + +def _register_refinement(message: dict, _db=murfey_db, demo: bool = False): + """Received class to refine from 3D classification""" + instrument_name = ( + _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") + pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-refine") + relion_params = _db.exec( + select(db.SPARelionParameters).where( + db.SPARelionParameters.pj_id == pj_id_params + ) + ).one() + relion_options = dict(relion_params) + feedback_params = _db.exec( + select(db.SPAFeedbackParameters).where( + db.SPAFeedbackParameters.pj_id == pj_id_params + ) + ).one() + other_options = dict(feedback_params) + + if feedback_params.hold_refine: + # If waiting then save the message + refine_params = _db.exec( + select(db.RefineParameters) + .where(db.RefineParameters.pj_id == pj_id) + .where(db.RefineParameters.tag == "first") + ).one() + # refine_params.refine_dir is not set as it will be the same as before + refine_params.run = True + refine_params.class3d_dir = message["class3d_dir"] + refine_params.class_number = message["best_class"] + _db.add(refine_params) + _db.commit() + _db.close() + else: + # Send all other messages on to a container + try: + refine_params = _db.exec( + select(db.RefineParameters) + .where(db.RefineParameters.pj_id == pj_id) + .where(db.RefineParameters.tag == "first") + ).one() + symmetry_refine_params = _db.exec( + select(db.RefineParameters) + .where(db.RefineParameters.pj_id == pj_id) + .where(db.RefineParameters.tag == "symmetry") + ).one() + except SQLAlchemyError: + next_job = feedback_params.next_job + refine_dir = f"{message['refine_dir']}{(feedback_params.next_job + 2):03}" + refined_grp_uuid = _murfey_id(message["program_id"], _db)[0] + refined_class_uuid = _murfey_id(message["program_id"], _db)[0] + symmetry_refined_grp_uuid = _murfey_id(message["program_id"], _db)[0] + symmetry_refined_class_uuid = _murfey_id(message["program_id"], _db)[0] + + refine_params = db.RefineParameters( + tag="first", + pj_id=pj_id, + murfey_id=refined_grp_uuid, + refine_dir=refine_dir, + class3d_dir=message["class3d_dir"], + class_number=message["best_class"], + ) + symmetry_refine_params = db.RefineParameters( + tag="symmetry", + pj_id=pj_id, + murfey_id=symmetry_refined_grp_uuid, + refine_dir=refine_dir, + class3d_dir=message["class3d_dir"], + class_number=message["best_class"], + ) + _db.add(refine_params) + _db.add(symmetry_refine_params) + _db.commit() + _murfey_refine( + murfey_id=refined_class_uuid, + refine_dir=refine_dir, + tag="first", + app_id=message["program_id"], + _db=_db, + ) + _murfey_refine( + murfey_id=symmetry_refined_class_uuid, + refine_dir=refine_dir, + tag="symmetry", + app_id=message["program_id"], + _db=_db, + ) + + if relion_options["symmetry"] == "C1": + # Extra Refine, Mask, PostProcess beyond for determined symmetry + next_job += 8 + else: + # Select and Extract particles, then Refine, Mask, PostProcess + next_job += 5 + feedback_params.next_job = next_job + + zocalo_message: dict = { + "parameters": { + "refine_job_dir": refine_params.refine_dir, + "class3d_dir": message["class3d_dir"], + "class_number": message["best_class"], + "pixel_size": relion_options["angpix"], + "particle_diameter": relion_options["particle_diameter"], + "mask_diameter": relion_options["mask_diameter"] or 0, + "symmetry": relion_options["symmetry"], + "node_creator_queue": machine_config.node_creator_queue, + "nr_iter": default_spa_parameters.nr_iter_3d, + "picker_id": other_options["picker_ispyb_id"], + "refined_class_uuid": _refine_murfey_id( + refine_dir=refine_params.refine_dir, + tag=refine_params.tag, + app_id=_app_id(pj_id, _db), + _db=_db, + ), + "refined_grp_uuid": refine_params.murfey_id, + "symmetry_refined_class_uuid": _refine_murfey_id( + refine_dir=symmetry_refine_params.refine_dir, + tag=symmetry_refine_params.tag, + app_id=_app_id(pj_id, _db), + _db=_db, + ), + "symmetry_refined_grp_uuid": symmetry_refine_params.murfey_id, + "session_id": message["session_id"], + "autoproc_program_id": _app_id( + _pj_id(message["program_id"], _db, recipe="em-spa-refine"), _db + ), + }, + "recipes": ["em-spa-refine"], + } + if murfey.server._transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = murfey.server._transport_object.feedback_queue + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + feedback_params.hold_refine = True + _db.add(feedback_params) + _db.commit() + _db.close() + + +def _register_bfactors(message: dict, _db=murfey_db, demo: bool = False): + """Received refined class to calculate b-factor""" + instrument_name = ( + _db.exec(select(db.Session).where(db.Session.id == message["session_id"])) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + pj_id_params = _pj_id(message["program_id"], _db, recipe="em-spa-preprocess") + pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-refine") + relion_params = _db.exec( + select(db.SPARelionParameters).where( + db.SPARelionParameters.pj_id == pj_id_params + ) + ).one() + relion_options = dict(relion_params) + feedback_params = _db.exec( + select(db.SPAFeedbackParameters).where( + db.SPAFeedbackParameters.pj_id == pj_id_params + ) + ).one() + + if message["symmetry"] != relion_params.symmetry: + # Currently don't do anything with a symmetrised re-run of the refinement + logger.info( + f"Received symmetrised structure of {sanitise(message['symmetry'])}" + ) + return True + + if not feedback_params.hold_refine: + logger.warning("B-Factors requested but refine hold is off") + return False + + # Add b-factor for refinement run + bfactor_run = db.BFactors( + pj_id=pj_id, + bfactor_directory=f"{message['project_dir']}/Refine3D/bfactor_{message['number_of_particles']}", + number_of_particles=message["number_of_particles"], + resolution=message["resolution"], + ) + _db.add(bfactor_run) + _db.commit() + + # All messages should create b-factor jobs as the refine hold is on at this point + try: + bfactor_params = _db.exec( + select(db.BFactorParameters).where(db.BFactorParameters.pj_id == pj_id) + ).one() + except SQLAlchemyError: + bfactor_params = db.BFactorParameters( + pj_id=pj_id, + project_dir=message["project_dir"], + batch_size=message["number_of_particles"], + refined_grp_uuid=message["refined_grp_uuid"], + refined_class_uuid=message["refined_class_uuid"], + class_reference=message["class_reference"], + class_number=message["class_number"], + mask_file=message["mask_file"], + ) + _db.add(bfactor_params) + _db.commit() + + bfactor_particle_count = default_spa_parameters.bfactor_min_particles + while bfactor_particle_count < bfactor_params.batch_size: + bfactor_run_name = ( + f"{bfactor_params.project_dir}/BFactors/bfactor_{bfactor_particle_count}" + ) + try: + bfactor_run = _db.exec( + select(db.BFactors) + .where(db.BFactors.pj_id == pj_id) + .where(db.BFactors.bfactor_directory == bfactor_run_name) + ).one() + bfactor_run.resolution = 0 + except SQLAlchemyError: + bfactor_run = db.BFactors( + pj_id=pj_id, + bfactor_directory=bfactor_run_name, + number_of_particles=bfactor_particle_count, + resolution=0, + ) + _db.add(bfactor_run) + _db.commit() + + bfactor_particle_count *= 2 + + zocalo_message: dict = { + "parameters": { + "bfactor_directory": bfactor_run.bfactor_directory, + "class_reference": bfactor_params.class_reference, + "class_number": bfactor_params.class_number, + "number_of_particles": bfactor_run.number_of_particles, + "batch_size": bfactor_params.batch_size, + "pixel_size": message["pixel_size"], + "mask": bfactor_params.mask_file, + "particle_diameter": relion_options["particle_diameter"], + "mask_diameter": relion_options["mask_diameter"] or 0, + "node_creator_queue": machine_config.node_creator_queue, + "picker_id": feedback_params.picker_ispyb_id, + "refined_grp_uuid": bfactor_params.refined_grp_uuid, + "refined_class_uuid": bfactor_params.refined_class_uuid, + "session_id": message["session_id"], + "autoproc_program_id": _app_id( + _pj_id(message["program_id"], _db, recipe="em-spa-refine"), _db + ), + }, + "recipes": ["em-spa-bfactor"], + } + if murfey.server._transport_object: + zocalo_message["parameters"][ + "feedback_queue" + ] = murfey.server._transport_object.feedback_queue + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + _db.close() + return True + + +def _save_bfactor(message: dict, _db=murfey_db, demo: bool = False): + """Received b-factor from refinement run""" + pj_id = _pj_id(message["program_id"], _db, recipe="em-spa-refine") + bfactor_run = _db.exec( + select(db.BFactors) + .where(db.BFactors.pj_id == pj_id) + .where(db.BFactors.number_of_particles == message["number_of_particles"]) + ).one() + bfactor_run.resolution = message["resolution"] + _db.add(bfactor_run) + _db.commit() + + # Find all the resolutions in the b-factors table + all_bfactors = _db.exec(select(db.BFactors).where(db.BFactors.pj_id == pj_id)).all() + particle_counts = [bf.number_of_particles for bf in all_bfactors] + resolutions = [bf.resolution for bf in all_bfactors] + + if all(resolutions): + # Calculate b-factor and add to ispyb class table + bfactor_fitting = np.polyfit( + np.log(particle_counts), 1 / np.array(resolutions) ** 2, 1 + ) + refined_class_uuid = message["refined_class_uuid"] + + # Request an ispyb insert of the b-factor fitting parameters + if murfey.server._transport_object: + murfey.server._transport_object.send( + "ispyb_connector", + { + "ispyb_command": "buffer", + "buffer_lookup": { + "particle_classification_id": refined_class_uuid, + }, + "buffer_command": { + "ispyb_command": "insert_particle_classification" + }, + "program_id": message["program_id"], + "bfactor_fit_intercept": str(bfactor_fitting[1]), + "bfactor_fit_linear": str(bfactor_fitting[0]), + }, + new_connection=True, + ) + + # Clean up the b-factors table and release the hold + [_db.delete(bf) for bf in all_bfactors] + _db.commit() + _release_refine_hold(message) + _db.close() + + +def feedback_callback(header: dict, message: dict) -> None: + try: + record = None + if "environment" in message: + params = message["recipe"][str(message["recipe-pointer"])].get( + "parameters", {} + ) + message = message["payload"] + message.update(params) + if message["register"] == "motion_corrected": + collected_ids = murfey_db.exec( + select( + db.DataCollectionGroup, + db.DataCollection, + db.ProcessingJob, + db.AutoProcProgram, + ) + .where(db.DataCollection.dcg_id == db.DataCollectionGroup.id) + .where(db.ProcessingJob.dc_id == db.DataCollection.id) + .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id) + .where(db.AutoProcProgram.id == message["program_id"]) + ).one() + session_id = collected_ids[0].session_id + + # Find the autoprocprogram id for the alignment recipe + alignment_ids = murfey_db.exec( + select( + db.DataCollection, + db.ProcessingJob, + db.AutoProcProgram, + ) + .where(db.ProcessingJob.dc_id == db.DataCollection.id) + .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id) + .where(db.DataCollection.id == collected_ids[1].id) + .where(db.ProcessingJob.recipe == "em-tomo-align") + ).one() + + relevant_tilt_and_series = murfey_db.exec( + select(db.Tilt, db.TiltSeries) + .where(db.Tilt.movie_path == message.get("movie")) + .where(db.Tilt.tilt_series_id == db.TiltSeries.id) + .where(db.TiltSeries.session_id == session_id) + ).one() + relevant_tilt = relevant_tilt_and_series[0] + relevant_tilt_series = relevant_tilt_and_series[1] + relevant_tilt.motion_corrected = True + murfey_db.add(relevant_tilt) + murfey_db.commit() + if ( + check_tilt_series_mc(relevant_tilt_series.id) + and not relevant_tilt_series.processing_requested + and relevant_tilt_series.tilt_series_length > 2 + ): + relevant_tilt_series.processing_requested = True + murfey_db.add(relevant_tilt_series) + + instrument_name = ( + murfey_db.exec( + select(db.Session).where(db.Session.id == session_id) + ) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + tilts = get_all_tilts(relevant_tilt_series.id) + ids = get_job_ids(relevant_tilt_series.id, alignment_ids[2].id) + preproc_params = get_tomo_proc_params(ids.dcgid) + stack_file = ( + Path(message["mrc_out"]).parents[3] + / "Tomograms" + / "job006" + / "tomograms" + / f"{relevant_tilt_series.tag}_stack.mrc" + ) + if not stack_file.parent.exists(): + stack_file.parent.mkdir(parents=True) + tilt_offset = midpoint([float(get_angle(t)) for t in tilts]) + zocalo_message = { + "recipes": ["em-tomo-align"], + "parameters": { + "input_file_list": str([[t, str(get_angle(t))] for t in tilts]), + "path_pattern": "", # blank for now so that it works with the tomo_align service changes + "dcid": ids.dcid, + "appid": ids.appid, + "stack_file": str(stack_file), + "dose_per_frame": preproc_params.dose_per_frame, + "frame_count": preproc_params.frame_count, + "kv": preproc_params.voltage, + "tilt_axis": preproc_params.tilt_axis, + "pixel_size": preproc_params.pixel_size, + "manual_tilt_offset": -tilt_offset, + "node_creator_queue": machine_config.node_creator_queue, + }, + } + if murfey.server._transport_object: + logger.info( + f"Sending Zocalo message for processing: {zocalo_message}" + ) + murfey.server._transport_object.send( + "processing_recipe", zocalo_message, new_connection=True + ) + else: + logger.info( + f"No transport object found. Zocalo message would be {zocalo_message}" + ) + + prom.preprocessed_movies.labels(processing_job=collected_ids[2].id).inc() + murfey_db.commit() + murfey_db.close() + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "data_collection_group": + ispyb_session_id = get_session_id( + microscope=message["microscope"], + proposal_code=message["proposal_code"], + proposal_number=message["proposal_number"], + visit_number=message["visit_number"], + db=ISPyBSession(), + ) + if dcg_murfey := murfey_db.exec( + select(db.DataCollectionGroup) + .where(db.DataCollectionGroup.session_id == message["session_id"]) + .where(db.DataCollectionGroup.tag == message.get("tag")) + ).all(): + dcgid = dcg_murfey[0].id + else: + if ispyb_session_id is None: + murfey_dcg = db.DataCollectionGroup( + session_id=message["session_id"], + tag=message.get("tag"), + ) + else: + record = DataCollectionGroup( + sessionId=ispyb_session_id, + experimentType=message["experiment_type"], + experimentTypeId=message["experiment_type_id"], + ) + dcgid = _register(record, header) + atlas_record = Atlas( + dataCollectionGroupId=dcgid, + atlasImage=message.get("atlas", ""), + pixelSize=message.get("atlas_pixel_size", 0), + cassetteSlot=message.get("sample"), + ) + if murfey.server._transport_object: + atlas_id = murfey.server._transport_object.do_insert_atlas( + atlas_record + )["return_value"] + murfey_dcg = db.DataCollectionGroup( + id=dcgid, + atlas_id=atlas_id, + session_id=message["session_id"], + tag=message.get("tag"), + ) + murfey_db.add(murfey_dcg) + murfey_db.commit() + murfey_db.close() + if murfey.server._transport_object: + if dcgid is None: + time.sleep(2) + murfey.server._transport_object.transport.nack(header, requeue=True) + return None + murfey.server._transport_object.transport.ack(header) + if dcg_hooks := entry_points().select( + group="murfey.hooks", name="data_collection_group" + ): + try: + for hook in dcg_hooks: + hook.load()(dcgid, session_id=message["session_id"]) + except Exception: + logger.error( + "Call to data collection group hook failed", exc_info=True + ) + return None + elif message["register"] == "atlas_update": + if murfey.server._transport_object: + murfey.server._transport_object.do_update_atlas( + message["atlas_id"], + message["atlas"], + message["atlas_pixel_size"], + message["sample"], + ) + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "data_collection": + logger.debug( + "Received message named 'data_collection' containing the following items:\n" + f"{', '.join([f'{sanitise(key)}: {sanitise(str(value))}' for key, value in message.items()])}" + ) + murfey_session_id = message["session_id"] + ispyb_session_id = get_session_id( + microscope=message["microscope"], + proposal_code=message["proposal_code"], + proposal_number=message["proposal_number"], + visit_number=message["visit_number"], + db=ISPyBSession(), + ) + dcg = murfey_db.exec( + select(db.DataCollectionGroup) + .where(db.DataCollectionGroup.session_id == murfey_session_id) + .where(db.DataCollectionGroup.tag == message["source"]) + ).all() + if dcg: + dcgid = dcg[0].id + # flush_data_collections(message["source"], murfey_db) + else: + logger.warning( + "No data collection group ID was found for image directory " + f"{sanitise(message['image_directory'])} and source " + f"{sanitise(message['source'])}" + ) + if murfey.server._transport_object: + murfey.server._transport_object.transport.nack(header, requeue=True) + return None + if dc_murfey := murfey_db.exec( + select(db.DataCollection) + .where(db.DataCollection.tag == message.get("tag")) + .where(db.DataCollection.dcg_id == dcgid) + ).all(): + dcid = dc_murfey[0].id + else: + if ispyb_session_id is None: + murfey_dc = db.DataCollection( + tag=message.get("tag"), + dcg_id=dcgid, + ) + else: + record = DataCollection( + SESSIONID=ispyb_session_id, + experimenttype=message["experiment_type"], + imageDirectory=message["image_directory"], + imageSuffix=message["image_suffix"], + voltage=message["voltage"], + dataCollectionGroupId=dcgid, + pixelSizeOnImage=message["pixel_size"], + imageSizeX=message["image_size_x"], + imageSizeY=message["image_size_y"], + slitGapHorizontal=message.get("slit_width"), + magnification=message.get("magnification"), + exposureTime=message.get("exposure_time"), + totalExposedDose=message.get("total_exposed_dose"), + c2aperture=message.get("c2aperture"), + phasePlate=int(message.get("phase_plate", 0)), + ) + dcid = _register( + record, + header, + tag=( + message.get("tag") + if message["experiment_type"] == "tomography" + else "" + ), + ) + murfey_dc = db.DataCollection( + id=dcid, + tag=message.get("tag"), + dcg_id=dcgid, + ) + murfey_db.add(murfey_dc) + murfey_db.commit() + dcid = murfey_dc.id + murfey_db.close() + if dcid is None and murfey.server._transport_object: + murfey.server._transport_object.transport.nack(header, requeue=True) + return None + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "processing_job": + murfey_session_id = message["session_id"] + logger.info("registering processing job") + dc = murfey_db.exec( + select(db.DataCollection, db.DataCollectionGroup) + .where(db.DataCollection.dcg_id == db.DataCollectionGroup.id) + .where(db.DataCollectionGroup.session_id == murfey_session_id) + .where(db.DataCollectionGroup.tag == message["source"]) + .where(db.DataCollection.tag == message["tag"]) + ).all() + if dc: + _dcid = dc[0][0].id + else: + logger.warning( + f"No data collection ID found for {sanitise(message['tag'])}" + ) + if murfey.server._transport_object: + murfey.server._transport_object.transport.nack(header, requeue=True) + return None + if pj_murfey := murfey_db.exec( + select(db.ProcessingJob) + .where(db.ProcessingJob.recipe == message["recipe"]) + .where(db.ProcessingJob.dc_id == _dcid) + ).all(): + pid = pj_murfey[0].id + else: + if ISPyBSession() is None: + murfey_pj = db.ProcessingJob(recipe=message["recipe"], dc_id=_dcid) + else: + record = ProcessingJob( + dataCollectionId=_dcid, recipe=message["recipe"] + ) + run_parameters = message.get("parameters", {}) + assert isinstance(run_parameters, dict) + if message.get("job_parameters"): + job_parameters = [ + ProcessingJobParameter(parameterKey=k, parameterValue=v) + for k, v in message["job_parameters"].items() + ] + pid = _register(ExtendedRecord(record, job_parameters), header) + else: + pid = _register(record, header) + murfey_pj = db.ProcessingJob( + id=pid, recipe=message["recipe"], dc_id=_dcid + ) + murfey_db.add(murfey_pj) + murfey_db.commit() + pid = murfey_pj.id + murfey_db.close() + if pid is None and murfey.server._transport_object: + murfey.server._transport_object.transport.nack(header, requeue=True) + return None + prom.preprocessed_movies.labels(processing_job=pid) + if not murfey_db.exec( + select(db.AutoProcProgram).where(db.AutoProcProgram.pj_id == pid) + ).all(): + if ISPyBSession() is None: + murfey_app = db.AutoProcProgram(pj_id=pid) + else: + record = AutoProcProgram( + processingJobId=pid, processingStartTime=datetime.now() + ) + appid = _register(record, header) + if appid is None and murfey.server._transport_object: + murfey.server._transport_object.transport.nack( + header, requeue=True + ) + return None + murfey_app = db.AutoProcProgram(id=appid, pj_id=pid) + murfey_db.add(murfey_app) + murfey_db.commit() + murfey_db.close() + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "flush_tomography_preprocess": + _flush_tomography_preprocessing(message) + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "spa_processing_parameters": + session_id = message["session_id"] + collected_ids = murfey_db.exec( + select( + db.DataCollectionGroup, + db.DataCollection, + db.ProcessingJob, + db.AutoProcProgram, + ) + .where(db.DataCollectionGroup.session_id == session_id) + .where(db.DataCollectionGroup.tag == message["tag"]) + .where(db.DataCollection.dcg_id == db.DataCollectionGroup.id) + .where(db.ProcessingJob.dc_id == db.DataCollection.id) + .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id) + .where(db.ProcessingJob.recipe == "em-spa-preprocess") + ).one() + pj_id = collected_ids[2].id + if not murfey_db.exec( + select(db.SPARelionParameters).where( + db.SPARelionParameters.pj_id == pj_id + ) + ).all(): + instrument_name = ( + murfey_db.exec( + select(db.Session).where(db.Session.id == session_id) + ) + .one() + .instrument_name + ) + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + params = db.SPARelionParameters( + pj_id=collected_ids[2].id, + angpix=float(message["pixel_size_on_image"]) * 1e10, + dose_per_frame=message["dose_per_frame"], + gain_ref=( + str(machine_config.rsync_basepath / message["gain_ref"]) + if message["gain_ref"] and machine_config.data_transfer_enabled + else message["gain_ref"] + ), + voltage=message["voltage"], + motion_corr_binning=message["motion_corr_binning"], + eer_fractionation_file=message["eer_fractionation_file"], + symmetry=message["symmetry"], + particle_diameter=message["particle_diameter"], + downscale=message["downscale"], + boxsize=message["boxsize"], + small_boxsize=message["small_boxsize"], + mask_diameter=message["mask_diameter"], + ) + feedback_params = db.SPAFeedbackParameters( + pj_id=collected_ids[2].id, + estimate_particle_diameter=not bool(message["particle_diameter"]), + hold_class2d=False, + hold_class3d=False, + class_selection_score=0, + star_combination_job=0, + initial_model="", + next_job=0, + ) + murfey_db.add(params) + murfey_db.add(feedback_params) + murfey_db.commit() + logger.info( + f"SPA processing parameters registered for processing job {collected_ids[2].id}" + ) + murfey_db.close() + else: + logger.info( + f"SPA processing parameters already exist for processing job ID {pj_id}" + ) + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "tomography_processing_parameters": + session_id = message["session_id"] + collected_ids = murfey_db.exec( + select( + db.DataCollectionGroup, + db.DataCollection, + db.ProcessingJob, + db.AutoProcProgram, + ) + .where(db.DataCollectionGroup.session_id == session_id) + .where(db.DataCollectionGroup.tag == message["tag"]) + .where(db.DataCollection.dcg_id == db.DataCollectionGroup.id) + .where(db.DataCollection.tag == message["tilt_series_tag"]) + .where(db.ProcessingJob.dc_id == db.DataCollection.id) + .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id) + .where(db.ProcessingJob.recipe == "em-tomo-preprocess") + ).one() + if not murfey_db.exec( + select(func.count(db.TomographyProcessingParameters.dcg_id)).where( + db.TomographyProcessingParameters.dcg_id == collected_ids[0].id + ) + ).one(): + params = db.TomographyProcessingParameters( + dcg_id=collected_ids[0].id, + pixel_size=float(message["pixel_size_on_image"]) * 10**10, + voltage=message["voltage"], + dose_per_frame=message["dose_per_frame"], + frame_count=message["frame_count"], + tilt_axis=message["tilt_axis"], + motion_corr_binning=message["motion_corr_binning"], + gain_ref=message["gain_ref"], + eer_fractionation_file=message["eer_fractionation_file"], + ) + murfey_db.add(params) + murfey_db.commit() + murfey_db.close() + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "done_incomplete_2d_batch": + _release_2d_hold(message) + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "incomplete_particles_file": + _register_incomplete_2d_batch(message) + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "complete_particles_file": + _register_complete_2d_batch(message) + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "save_class_selection_score": + _register_class_selection(message) + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "done_3d_batch": + _release_3d_hold(message) + if message.get("do_refinement"): + _register_refinement(message) + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "run_class3d": + _register_3d_batch(message) + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "save_initial_model": + _register_initial_model(message) + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "done_particle_selection": + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "done_class_selection": + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "atlas_registered": + _flush_grid_square_records(message) + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif message["register"] == "done_refinement": + bfactors_registered = _register_bfactors(message) + if murfey.server._transport_object: + if bfactors_registered: + murfey.server._transport_object.transport.ack(header) + else: + murfey.server._transport_object.transport.nack(header) + return None + elif message["register"] == "done_bfactor": + _save_bfactor(message) + if murfey.server._transport_object: + murfey.server._transport_object.transport.ack(header) + return None + elif ( + message["register"] in entry_points().select(group="murfey.workflows").names + ): + # Search for corresponding workflow + workflows: list[EntryPoint] = list( + entry_points().select( + group="murfey.workflows", name=message["register"] + ) + ) # Returns either 1 item or empty list + if not workflows: + logger.error(f"No workflow found for {sanitise(message['register'])}") + if murfey.server._transport_object: + murfey.server._transport_object.transport.nack( + header, requeue=False + ) + return None + # Run the workflow if a match is found + workflow: EntryPoint = workflows[0] + result = workflow.load()( + message=message, + murfey_db=murfey_db, + ) + if murfey.server._transport_object: + if result: + murfey.server._transport_object.transport.ack(header) + else: + # Send it directly to DLQ without trying to rerun it + murfey.server._transport_object.transport.nack( + header, requeue=False + ) + if not result: + logger.error( + f"Workflow {sanitise(message['register'])} returned {result}" + ) + return None + logger.error(f"No workflow found for {sanitise(message['register'])}") + if murfey.server._transport_object: + murfey.server._transport_object.transport.nack(header, requeue=False) + return None + except PendingRollbackError: + murfey_db.rollback() + murfey_db.close() + logger.warning("Murfey database required a rollback") + if murfey.server._transport_object: + murfey.server._transport_object.transport.nack(header, requeue=True) + except OperationalError: + logger.warning("Murfey database error encountered", exc_info=True) + time.sleep(1) + if murfey.server._transport_object: + murfey.server._transport_object.transport.nack(header, requeue=True) + except Exception: + logger.warning( + "Exception encountered in server RabbitMQ callback", exc_info=True + ) + if murfey.server._transport_object: + murfey.server._transport_object.transport.nack(header, requeue=False) + return None + + +@singledispatch +def _register(record, header: dict, **kwargs): + raise NotImplementedError(f"Not method to register {record} or type {type(record)}") + + +@_register.register +def _(record: Base, header: dict, **kwargs): # type: ignore + if not murfey.server._transport_object: + logger.error( + f"No transport object found when processing record {record}. Message header: {header}" + ) + return None + try: + if isinstance(record, DataCollection): + return murfey.server._transport_object.do_insert_data_collection( + record, **kwargs + )["return_value"] + if isinstance(record, DataCollectionGroup): + return murfey.server._transport_object.do_insert_data_collection_group( + record + )["return_value"] + if isinstance(record, ProcessingJob): + return murfey.server._transport_object.do_create_ispyb_job(record)[ + "return_value" + ] + if isinstance(record, AutoProcProgram): + return murfey.server._transport_object.do_update_processing_status(record)[ + "return_value" + ] + # session = Session() + # session.add(record) + # session.commit() + # murfey.server._transport_object.transport.ack(header, requeue=False) + return getattr(record, record.__table__.primary_key.columns[0].name) + + except SQLAlchemyError as e: + logger.error(f"Murfey failed to insert ISPyB record {record}", e, exc_info=True) + # murfey.server._transport_object.transport.nack(header) + return None + except AttributeError as e: + logger.error( + f"Murfey could not find primary key when inserting record {record}", + e, + exc_info=True, + ) + return None + + +@_register.register +def _(extended_record: ExtendedRecord, header: dict, **kwargs): + if not murfey.server._transport_object: + raise ValueError( + "Transport object should not be None if a database record is being updated" + ) + return murfey.server._transport_object.do_create_ispyb_job( + extended_record.record, params=extended_record.record_params + )["return_value"] + + +def feedback_listen(): + if murfey.server._transport_object: + if not murfey.server._transport_object.feedback_queue: + murfey.server._transport_object.feedback_queue = ( + murfey.server._transport_object.transport._subscribe_temporary( + channel_hint="", callback=None, sub_id=None + ) + ) + murfey.server._transport_object._connection_callback = partial( + murfey.server._transport_object.transport.subscribe, + murfey.server._transport_object.feedback_queue, + feedback_callback, + acknowledgement=True, + ) + murfey.server._transport_object.transport.subscribe( + murfey.server._transport_object.feedback_queue, + feedback_callback, + acknowledgement=True, + ) diff --git a/src/murfey/server/ispyb.py b/src/murfey/server/ispyb.py index ddf55f87e..24f8094ce 100644 --- a/src/murfey/server/ispyb.py +++ b/src/murfey/server/ispyb.py @@ -26,12 +26,13 @@ ZcZocaloBuffer, url, ) +from pydantic import BaseModel from sqlalchemy import create_engine from sqlalchemy.orm import Session, sessionmaker from murfey.util import sanitise from murfey.util.config import get_security_config -from murfey.util.models import FoilHoleParameters, GridSquareParameters, Sample, Visit +from murfey.util.models import FoilHoleParameters, GridSquareParameters log = logging.getLogger("murfey.server.ispyb") security_config = get_security_config() @@ -50,6 +51,27 @@ ISPyBSession = lambda: None +class Visit(BaseModel): + start: datetime.datetime + end: datetime.datetime + session_id: int + name: str + beamline: str + proposal_title: str + + def __repr__(self) -> str: + return ( + "Visit(" + f"start='{self.start:%Y-%m-%d %H:%M}', " + f"end='{self.end:%Y-%m-%d %H:%M}', " + f"session_id='{self.session_id!r}'," + f"name={self.name!r}, " + f"beamline={self.beamline!r}, " + f"proposal_title={self.proposal_title!r}" + ")" + ) + + def _send_using_new_connection(transport_type: str, queue: str, message: dict) -> None: transport = workflows.transport.lookup(transport_type)() transport.connect() @@ -558,8 +580,8 @@ def _get_session() -> Generator[Optional[Session], None, None]: db.close() -DB = Depends(_get_session) # Shortcut to access the database in a FastAPI endpoint +DB = Depends(_get_session) def get_session_id( @@ -612,31 +634,6 @@ def get_proposal_id(proposal_code: str, proposal_number: str, db: Session) -> in return query[0].proposalId -def get_sub_samples_from_visit(visit: str, db: Session) -> List[Sample]: - proposal_id = get_proposal_id(visit[:2], visit.split("-")[0][2:], db) - samples = ( - db.query(BLSampleGroup, BLSampleGroupHasBLSample, BLSample, BLSubSample) - .join(BLSample, BLSample.blSampleId == BLSampleGroupHasBLSample.blSampleId) - .join( - BLSampleGroup, - BLSampleGroup.blSampleGroupId == BLSampleGroupHasBLSample.blSampleGroupId, - ) - .join(BLSubSample, BLSubSample.blSampleId == BLSample.blSampleId) - .filter(BLSampleGroup.proposalId == proposal_id) - .all() - ) - res = [ - Sample( - sample_group_id=s[1].blSampleGroupId, - sample_id=s[2].blSampleId, - subsample_id=s[3].blSubSampleId, - image_path=s[3].imgFilePath, - ) - for s in samples - ] - return res - - def get_all_ongoing_visits(microscope: str, db: Session | None) -> list[Visit]: if db is None: print("No database found") diff --git a/src/murfey/server/main.py b/src/murfey/server/main.py index 1dcde56f0..f567d0e9b 100644 --- a/src/murfey/server/main.py +++ b/src/murfey/server/main.py @@ -1,7 +1,6 @@ from __future__ import annotations import logging -import os from backports.entry_points_selectable import entry_points from fastapi import FastAPI @@ -15,22 +14,20 @@ import murfey.server.api.bootstrap import murfey.server.api.clem import murfey.server.api.display +import murfey.server.api.file_manip import murfey.server.api.hub import murfey.server.api.instrument +import murfey.server.api.mag_table import murfey.server.api.processing_parameters +import murfey.server.api.prometheus +import murfey.server.api.session_control +import murfey.server.api.session_info import murfey.server.api.spa -import murfey.server.websocket -import murfey.util.models +import murfey.server.api.websocket +import murfey.server.api.workflow from murfey.server import template_files from murfey.util.config import get_security_config -# Import Murfey server or demo server based on settings -if os.getenv("MURFEY_DEMO"): - from murfey.server.demo_api import router -else: - from murfey.server.api import router - - log = logging.getLogger("murfey.server.main") tags_metadata = [murfey.server.api.bootstrap.tag] @@ -61,7 +58,6 @@ class Settings(BaseSettings): app.mount("/images", StaticFiles(directory=template_files / "images"), name="images") # Add router endpoints to the API -app.include_router(router) app.include_router(murfey.server.api.bootstrap.version) app.include_router(murfey.server.api.bootstrap.bootstrap) app.include_router(murfey.server.api.bootstrap.cygwin) @@ -69,14 +65,37 @@ class Settings(BaseSettings): app.include_router(murfey.server.api.bootstrap.rust) app.include_router(murfey.server.api.bootstrap.pypi) app.include_router(murfey.server.api.bootstrap.plugins) -app.include_router(murfey.server.api.clem.router) -app.include_router(murfey.server.api.spa.router) + app.include_router(murfey.server.api.auth.router) -app.include_router(murfey.server.api.display.router) -app.include_router(murfey.server.api.instrument.router) + app.include_router(murfey.server.api.hub.router) +app.include_router(murfey.server.api.display.router) app.include_router(murfey.server.api.processing_parameters.router) -app.include_router(murfey.server.websocket.ws) + +app.include_router(murfey.server.api.file_manip.router) + +app.include_router(murfey.server.api.instrument.router) + +app.include_router(murfey.server.api.mag_table.router) + +app.include_router(murfey.server.api.session_control.router) +app.include_router(murfey.server.api.session_control.spa_router) + +app.include_router(murfey.server.api.session_info.router) +app.include_router(murfey.server.api.session_info.correlative_router) +app.include_router(murfey.server.api.session_info.spa_router) +app.include_router(murfey.server.api.session_info.tomo_router) + +app.include_router(murfey.server.api.workflow.router) +app.include_router(murfey.server.api.workflow.correlative_router) +app.include_router(murfey.server.api.workflow.spa_router) +app.include_router(murfey.server.api.workflow.tomo_router) +app.include_router(murfey.server.api.spa.router) +app.include_router(murfey.server.api.clem.router) + +app.include_router(murfey.server.api.prometheus.router) + +app.include_router(murfey.server.api.websocket.ws) # Search external packages for additional routers to include in Murfey for r in entry_points(group="murfey.routers"): diff --git a/src/murfey/server/run.py b/src/murfey/server/run.py index 56fee104c..26eb75bc0 100644 --- a/src/murfey/server/run.py +++ b/src/murfey/server/run.py @@ -1,5 +1,209 @@ from __future__ import annotations -from murfey.server import run +import argparse +import logging +import os +from threading import Thread +from typing import Literal -__all__ = ["run"] +import graypy +import uvicorn +from rich.logging import RichHandler +from workflows import Error as WorkflowsError +from workflows.transport.pika_transport import PikaTransport + +import murfey +import murfey.server +from murfey.server.feedback import feedback_listen +from murfey.server.ispyb import TransportManager +from murfey.util import LogFilter +from murfey.util.config import get_microscope, get_security_config + +logger = logging.getLogger("murfey.server.run") + + +def _set_up_logging(quiet: bool, verbosity: int): + rich_handler = RichHandler(enable_link_path=False) + if quiet: + rich_handler.setLevel(logging.INFO) + log_levels = { + "murfey": logging.INFO, + "uvicorn": logging.WARNING, + "fastapi": logging.INFO, + "starlette": logging.INFO, + "sqlalchemy": logging.WARNING, + } + elif verbosity <= 0: + rich_handler.setLevel(logging.INFO) + log_levels = { + "murfey": logging.DEBUG, + "uvicorn": logging.INFO, + "uvicorn.access": logging.WARNING, + "fastapi": logging.INFO, + "starlette": logging.INFO, + "sqlalchemy": logging.WARNING, + } + elif verbosity <= 1: + rich_handler.setLevel(logging.DEBUG) + log_levels = { + "": logging.INFO, + "murfey": logging.DEBUG, + "uvicorn": logging.INFO, + "fastapi": logging.INFO, + "starlette": logging.INFO, + "sqlalchemy": logging.WARNING, + } + elif verbosity <= 2: + rich_handler.setLevel(logging.DEBUG) + log_levels = { + "": logging.INFO, + "murfey": logging.DEBUG, + "uvicorn": logging.DEBUG, + "fastapi": logging.DEBUG, + "starlette": logging.DEBUG, + "sqlalchemy": logging.WARNING, + } + else: + rich_handler.setLevel(logging.DEBUG) + log_levels = { + "": logging.DEBUG, + "murfey": logging.DEBUG, + "uvicorn": logging.DEBUG, + "fastapi": logging.DEBUG, + "starlette": logging.DEBUG, + "sqlalchemy": logging.DEBUG, + } + + logging.getLogger().addHandler(rich_handler) + for logger_name, log_level in log_levels.items(): + logging.getLogger(logger_name).setLevel(log_level) + + +def _set_up_transport(transport_type: Literal["PikaTransport"]): + # Update the existing TransportManager object in 'murfey.server' + murfey.server._transport_object = TransportManager(transport_type) + + +def run(): + """ + Main function that starts up the Murfey server + """ + + # Set up argument parser + parser = argparse.ArgumentParser(description="Start the Murfey server") + parser.add_argument( + "--host", + help="Listen for incoming connections on a specific interface (IP address or hostname; default: all)", + default="0.0.0.0", + ) + parser.add_argument( + "--port", + help="Listen for incoming TCP connections on this port (default: 8000)", + type=int, + default=8000, + ) + parser.add_argument( + "--workers", help="Number of workers for Uvicorn server", type=int, default=2 + ) + parser.add_argument( + "--demo", + action="store_true", + ) + parser.add_argument( + "--feedback", + action="store_true", + ) + parser.add_argument( + "--temporary", + action="store_true", + ) + parser.add_argument( + "--root-path", + default="", + type=str, + help="Uvicorn root path for use in conjunction with a proxy", + ) + verbosity = parser.add_mutually_exclusive_group() + verbosity.add_argument( + "-q", + "--quiet", + action="store_true", + default=False, + help="Decrease logging output verbosity", + ) + verbosity.add_argument( + "-v", + "--verbose", + action="count", + help="Increase logging output verbosity", + default=0, + ) + # Parse and separate known and unknown args + args, unknown = parser.parse_known_args() + + # Load the security configuration + security_config = get_security_config() + + # Set up GrayLog handler if provided in the configuration + if security_config.graylog_host: + handler = graypy.GELFUDPHandler( + security_config.graylog_host, security_config.graylog_port, level_names=True + ) + root_logger = logging.getLogger() + root_logger.addHandler(handler) + # Install a log filter to all existing handlers. + LogFilter.install() + + if args.demo: + # Run in demo mode with no connections set up + os.environ["MURFEY_DEMO"] = "1" + else: + # Load RabbitMQ configuration and set up the connection + try: + PikaTransport().load_configuration_file( + security_config.rabbitmq_credentials + ) + _set_up_transport("PikaTransport") + logger.info("Set up message transport manager") + except WorkflowsError: + logger.error( + "Error encountered setting up RabbitMQ connection", + exc_info=True, + ) + + # Set up logging now that the desired verbosity is known + _set_up_logging(quiet=args.quiet, verbosity=args.verbose) + + if not args.temporary and murfey.server._transport_object: + murfey.server._transport_object.feedback_queue = security_config.feedback_queue + rabbit_thread = Thread( + target=feedback_listen, + daemon=True, + ) + logger.info("Starting Murfey RabbitMQ thread") + if args.feedback: + rabbit_thread.start() + + logger.info( + f"Starting Murfey server version {murfey.__version__} for beamline {get_microscope()}, listening on {args.host}:{args.port}" + ) + config = uvicorn.Config( + "murfey.server.main:app", + host=args.host, + port=args.port, + log_config=None, + ws_ping_interval=300, + ws_ping_timeout=300, + workers=args.workers, + root_path=args.root_path, + ) + + murfey.server._running_server = uvicorn.Server(config=config) + murfey.server._running_server.run() + logger.info("Server shutting down") + + +def shutdown(): + if murfey.server._running_server: + murfey.server._running_server.should_exit = True + murfey.server._running_server.force_exit = True diff --git a/src/murfey/util/__init__.py b/src/murfey/util/__init__.py index 71613424f..27aabfb22 100644 --- a/src/murfey/util/__init__.py +++ b/src/murfey/util/__init__.py @@ -16,6 +16,10 @@ def sanitise(in_string: str) -> str: return in_string.replace("\r\n", "").replace("\n", "") +def sanitise_path(in_path: Path) -> Path: + return Path("/".join(secure_filename(p) for p in in_path.parts)) + + def sanitise_nonpath(in_string: str) -> str: for c in ("\r\n", "\n", "/", "\\", ":", ";"): in_string = in_string.replace(c, "") diff --git a/src/murfey/util/api.py b/src/murfey/util/api.py new file mode 100644 index 000000000..c402d66a6 --- /dev/null +++ b/src/murfey/util/api.py @@ -0,0 +1,148 @@ +""" +Utility functions to help with URL path lookups using function names for our FastAPI +servers. This makes reference to the route_manifest.yaml file that is also saved in +this directory. This routes_manifest.yaml file should be regenerated whenver changes +are made to the API endpoints. This can be done using the 'generate_route_manifest' +CLI function. +""" + +from __future__ import annotations + +import re +from functools import lru_cache +from logging import getLogger +from pathlib import Path +from typing import Any + +import yaml + +import murfey.util + +logger = getLogger("murfey.util.api") + +route_manifest_file = Path(murfey.util.__path__[0]) / "route_manifest.yaml" + + +@lru_cache(maxsize=1) # Load the manifest once and reuse +def load_route_manifest( + file: Path = route_manifest_file, +): + with open(file, "r") as f: + return yaml.safe_load(f) + + +def find_unique_index( + pattern: str, + candidates: list[str], + exact: bool = False, # Allows partial matches +) -> int: + """ + Finds the index of a unique entry in a list. + """ + counter = 0 + matches = [] + index = 0 + for i, candidate in enumerate(candidates): + if (not exact and pattern in candidate) or (exact and pattern == candidate): + counter += 1 + matches.append(candidate) + index = i + if counter == 0: + message = f"No match found for {pattern!r}" + logger.error(message) + raise KeyError(message) + if counter > 1: + message = f"Ambiguous match for {pattern!r}: {matches}" + logger.error(message) + raise KeyError(message) + return index + + +def render_path(path_template: str, kwargs: dict) -> str: + """ + Replace all FastAPI-style {param[:converter]} path parameters with corresponding + values from kwargs. + """ + + pattern = re.compile(r"{([^}]+)}") # Look for all path params + + def replace(match): + raw_str = match.group(1) + param_name = raw_str.split(":")[0] # Ignore :converter in the field + if param_name not in kwargs: + message = ( + f"Error constructing URL for {path_template!r}; " + f"missing path parameter {param_name!r}" + ) + logger.error(message) + raise KeyError(message) + return str(kwargs[param_name]) + + return pattern.sub(replace, path_template) + + +def url_path_for( + router_name: str, # With logic for partial matches + function_name: str, # With logic for partial matches + **kwargs, # Takes any path param and matches it against curly bracket contents +): + """ + Utility function that takes the function name and API router name, along with all + necessary path parameters, retrieves the matching URL path template from the route + manifest, and returns a correctly populated instance of the URL path. + """ + # Use 'Any' first and slowly reveal types as it is unpacked + route_manifest: dict[str, list[Any]] = load_route_manifest() + + # Load the routes in the desired router + routers = list(route_manifest.keys()) + routes: list[dict[str, Any]] = route_manifest[ + routers[find_unique_index(router_name, routers, exact=False)] + ] + + # Search router for the function + route_info = routes[ + find_unique_index(function_name, [r["function"] for r in routes], exact=True) + ] + + # Unpack the dictionary + route_path: str = route_info["path"] + path_params: list[dict[str, str]] = route_info["path_params"] + + # Validate the stored path params against the ones provided + if path_params: + for path_param in path_params: + param_name = path_param["name"] + param_type = path_param["type"] + if param_name not in kwargs.keys(): + message = ( + f"Error validating parameters for {function_name!r}; " + f"path parameter {param_name!r} was not provided" + ) + logger.error(message) + raise KeyError(message) + # Skip complicated type resolution for now + if param_type.startswith("typing."): + continue + elif type(kwargs[param_name]).__name__ not in param_type: + message = ( + f"Error validating parameters for {function_name!r}; " + f"{param_name!r} must be {param_type!r}, " + f"received {type(kwargs[param_name]).__name__!r}" + ) + logger.error(message) + raise TypeError(message) + + # Render and return the path + return render_path(route_path, kwargs) + + +if __name__ == "__main__": + # Run test on some existing routes + url_path = url_path_for( + "workflow.tomo_router", + "register_tilt", + visit_name="nt15587-15", + session_id=2, + ) + print(url_path) diff --git a/src/murfey/util/client.py b/src/murfey/util/client.py index 0e9bd3c4c..2721baa0f 100644 --- a/src/murfey/util/client.py +++ b/src/murfey/util/client.py @@ -17,11 +17,11 @@ from functools import lru_cache, partial from pathlib import Path from typing import Awaitable, Callable, Optional, Union -from urllib.parse import ParseResult, urlparse, urlunparse +from urllib.parse import urlparse, urlunparse import requests -from murfey.util.models import Visit +from murfey.util.api import url_path_for logger = logging.getLogger("murfey.util.client") @@ -59,7 +59,9 @@ def get_machine_config_client( _instrument_name: Optional[str] = instrument_name or os.getenv("BEAMLINE") if not _instrument_name: return {} - return requests.get(f"{url}/instruments/{_instrument_name}/machine").json() + return requests.get( + f"{url}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=_instrument_name)}" + ).json() def authorised_requests() -> tuple[Callable, Callable, Callable, Callable]: @@ -74,17 +76,6 @@ def authorised_requests() -> tuple[Callable, Callable, Callable, Callable]: requests.get, requests.post, requests.put, requests.delete = authorised_requests() -def _get_visit_list(api_base: ParseResult, instrument_name: str): - proxy_path = api_base.path.rstrip("/") - get_visits_url = api_base._replace( - path=f"{proxy_path}/instruments/{instrument_name}/visits_raw" - ) - server_reply = requests.get(get_visits_url.geturl()) - if server_reply.status_code != 200: - raise ValueError(f"Server unreachable ({server_reply.status_code})") - return [Visit.parse_obj(v) for v in server_reply.json()] - - def capture_post(url: str, json: Union[dict, list] = {}) -> Optional[requests.Response]: try: response = requests.post(url, json=json) diff --git a/src/murfey/util/config.py b/src/murfey/util/config.py index 95f3544be..31bb380fb 100644 --- a/src/murfey/util/config.py +++ b/src/murfey/util/config.py @@ -4,46 +4,62 @@ import socket from functools import lru_cache from pathlib import Path -from typing import Dict, List, Literal, Optional, Union +from typing import Literal, Optional, Union import yaml from backports.entry_points_selectable import entry_points from pydantic import BaseModel, BaseSettings, Extra, validator -class MachineConfig(BaseModel, extra=Extra.allow): # type: ignore - acquisition_software: List[str] - calibrations: Dict[str, Dict[str, Union[dict, float]]] - data_directories: List[Path] - rsync_basepath: Path - default_model: Path +class MachineConfig(BaseModel): # type: ignore + """ + Keys that describe the type of workflow conducted on the client side, and how + Murfey will handle its data transfer and processing + """ + + # General info -------------------------------------------------------------------- display_name: str = "" instrument_name: str = "" image_path: Optional[Path] = None - software_versions: Dict[str, str] = {} - external_executables: Dict[str, str] = {} - external_executables_eer: Dict[str, str] = {} - external_environment: Dict[str, str] = {} - rsync_module: str = "" + machine_override: str = "" + + # Hardware and software ----------------------------------------------------------- + camera: str = "FALCON" + superres: bool = False + calibrations: dict[str, dict[str, Union[dict, float]]] + acquisition_software: list[str] + software_versions: dict[str, str] = {} + software_settings_output_directories: dict[str, list[str]] = {} + data_required_substrings: dict[str, dict[str, list[str]]] = {} + + # Client side directory setup ----------------------------------------------------- + data_directories: list[Path] create_directories: list[str] = ["atlas"] - analyse_created_directories: List[str] = [] + analyse_created_directories: list[str] = [] gain_reference_directory: Optional[Path] = None eer_fractionation_file_template: str = "" - processed_directory_name: str = "processed" - gain_directory_name: str = "processing" - node_creator_queue: str = "node_creator" - superres: bool = False - camera: str = "FALCON" - data_required_substrings: Dict[str, Dict[str, List[str]]] = {} - allow_removal: bool = False + + # Data transfer setup ------------------------------------------------------------- + # Rsync setup data_transfer_enabled: bool = True + rsync_url: str = "" + rsync_module: str = "" + rsync_basepath: Path + allow_removal: bool = False + + # Upstream data download setup + upstream_data_directories: list[Path] = [] # Previous sessions + upstream_data_download_directory: Optional[Path] = None # Set by microscope config + upstream_data_tiff_locations: list[str] = ["processed"] # Location of CLEM TIFFs + + # Data processing setup ----------------------------------------------------------- + # General processing setup processing_enabled: bool = True - machine_override: str = "" - processed_extra_directory: str = "" - plugin_packages: Dict[str, Path] = {} - software_settings_output_directories: Dict[str, List[str]] = {} process_by_default: bool = True - recipes: Dict[str, str] = { + gain_directory_name: str = "processing" + processed_directory_name: str = "processed" + processed_extra_directory: str = "" + recipes: dict[str, str] = { "em-spa-bfactor": "em-spa-bfactor", "em-spa-class2d": "em-spa-class2d", "em-spa-class3d": "em-spa-class3d", @@ -53,26 +69,41 @@ class MachineConfig(BaseModel, extra=Extra.allow): # type: ignore "em-tomo-align": "em-tomo-align", } - # Find and download upstream directories - upstream_data_directories: List[Path] = [] # Previous sessions - upstream_data_download_directory: Optional[Path] = None # Set by microscope config - upstream_data_tiff_locations: List[str] = ["processed"] # Location of CLEM TIFFs - + # Particle picking setup + default_model: Path model_search_directory: str = "processing" initial_model_search_directory: str = "processing/initial_model" - failure_queue: str = "" - instrument_server_url: str = "http://localhost:8001" - frontend_url: str = "http://localhost:3000" - murfey_url: str = "http://localhost:8000" - rsync_url: str = "" + # Data analysis plugins + external_executables: dict[str, str] = {} + external_executables_eer: dict[str, str] = {} + external_environment: dict[str, str] = {} + plugin_packages: dict[str, Path] = {} + # Server and network setup -------------------------------------------------------- + # Configurations and URLs security_configuration_path: Optional[Path] = None + murfey_url: str = "http://localhost:8000" + frontend_url: str = "http://localhost:3000" + instrument_server_url: str = "http://localhost:8001" + # Messaging queues + failure_queue: str = "" + node_creator_queue: str = "node_creator" notifications_queue: str = "pato_notification" + class Config: + """ + Inner class that defines this model's parsing and serialising behaviour + """ -def from_file(config_file_path: Path, instrument: str = "") -> Dict[str, MachineConfig]: + extra = Extra.allow + json_encoders = { + Path: str, + } + + +def from_file(config_file_path: Path, instrument: str = "") -> dict[str, MachineConfig]: with open(config_file_path, "r") as config_stream: config = yaml.safe_load(config_stream) return { @@ -83,22 +114,36 @@ def from_file(config_file_path: Path, instrument: str = "") -> Dict[str, Machine class Security(BaseModel): + # Murfey database settings murfey_db_credentials: Path crypto_key: str - auth_key: str = "" + sqlalchemy_pooling: bool = True + + # ISPyB settings + ispyb_credentials: Optional[Path] = None + + # Murfey server connection settings auth_algorithm: str = "" + auth_key: str = "" + auth_type: Literal["password", "cookie"] = "password" auth_url: str = "" - sqlalchemy_pooling: bool = True - allow_origins: List[str] = ["*"] + cookie_key: str = "" session_validation: str = "" session_token_timeout: Optional[int] = None - auth_type: Literal["password", "cookie"] = "password" - cookie_key: str = "" + allow_origins: list[str] = ["*"] + + # RabbitMQ settings rabbitmq_credentials: Path feedback_queue: str = "murfey_feedback" + + # Graylog settings graylog_host: str = "" graylog_port: Optional[int] = None - ispyb_credentials: Optional[Path] = None + + class Config: + json_encoders = { + Path: str, + } @validator("graylog_port") def check_port_present_if_host_is( @@ -158,7 +203,7 @@ def get_security_config() -> Security: @lru_cache(maxsize=1) -def get_machine_config(instrument_name: str = "") -> Dict[str, MachineConfig]: +def get_machine_config(instrument_name: str = "") -> dict[str, MachineConfig]: machine_config = { "": MachineConfig( acquisition_software=[], diff --git a/src/murfey/util/instrument_models.py b/src/murfey/util/instrument_models.py index ee45b5468..4b2b1ff90 100644 --- a/src/murfey/util/instrument_models.py +++ b/src/murfey/util/instrument_models.py @@ -4,12 +4,9 @@ from pydantic import BaseModel -from murfey.util.config import MachineConfig - class MultigridWatcherSpec(BaseModel): source: Path - configuration: MachineConfig label: str visit: str instrument_name: str diff --git a/src/murfey/util/models.py b/src/murfey/util/models.py index fac7bfe0a..d51bbda38 100644 --- a/src/murfey/util/models.py +++ b/src/murfey/util/models.py @@ -34,50 +34,6 @@ def __repr__(self) -> str: ) -class SuggestedPathParameters(BaseModel): - base_path: Path - touch: bool = False - extra_directory: str = "" - - -class DCGroupParameters(BaseModel): - # DC = Data collection - experiment_type: str - experiment_type_id: int - tag: str - atlas: str = "" - sample: Optional[int] = None - atlas_pixel_size: int = 0 - - -class DCParameters(BaseModel): - voltage: float - pixel_size_on_image: str - experiment_type: str - image_size_x: int - image_size_y: int - file_extension: str - acquisition_software: str - image_directory: str - tag: str - source: str - magnification: float - total_exposed_dose: Optional[float] = None - c2aperture: Optional[float] = None - exposure_time: Optional[float] = None - slit_width: Optional[float] = None - phase_plate: bool = False - data_collection_tag: str = "" - - -class ProcessingJobParameters(BaseModel): - tag: str - source: str - recipe: str - parameters: Dict[str, Any] = {} - experiment_type: str = "spa" - - class RegistrationMessage(BaseModel): registration: str params: Optional[Dict[str, Any]] = None @@ -96,20 +52,10 @@ class ConnectionFileParameters(BaseModel): destinations: List[str] -class SessionInfo(BaseModel): - session_id: Optional[int] - session_name: str = "" - rescale: bool = True - - class ClientInfo(BaseModel): id: int -class RsyncerSource(BaseModel): - source: str - - class RsyncerInfo(BaseModel): source: str destination: str @@ -122,55 +68,6 @@ class RsyncerInfo(BaseModel): tag: str = "" -class GainReference(BaseModel): - gain_ref: Path - rescale: bool = True - eer: bool = False - tag: str = "" - - -class FractionationParameters(BaseModel): - fractionation: int - dose_per_frame: float - num_frames: int = 0 - eer_path: Optional[str] = None - fractionation_file_name: str = "eer_fractionation.txt" - - -""" -FIB -=== -Models related to FIB, as part of correlative workflow with TEM. -""" - - -class Sample(BaseModel): - sample_group_id: int - sample_id: int - subsample_id: int - image_path: Optional[Path] - - -class BLSampleImageParameters(BaseModel): - sample_id: int - sample_path: Path - - -class BLSampleParameters(BaseModel): - sample_group_id: int - - -class BLSubSampleParameters(BaseModel): - sample_id: int - image_path: Optional[Path] = None - - -class MillingParameters(BaseModel): - lamella_number: int - images: List[str] - raw_directory: str - - """ Single Particle Analysis ======================== @@ -178,24 +75,6 @@ class MillingParameters(BaseModel): """ -class SPAProcessFile(BaseModel): - tag: str - path: str - description: str - processing_job: Optional[int] - data_collection_id: Optional[int] - image_number: int - autoproc_program_id: Optional[int] - foil_hole_id: Optional[int] - pixel_size: Optional[float] - dose_per_frame: Optional[float] - mc_binning: Optional[int] = 1 - gain_ref: Optional[str] = None - extract_downscale: bool = True - eer_fractionation_file: Optional[str] = None - source: str = "" - - class ProcessingParametersSPA(BaseModel): tag: str dose_per_frame: float @@ -268,11 +147,6 @@ class FoilHoleParameters(BaseModel): diameter: Optional[float] = None -class PostInfo(BaseModel): - url: str - data: dict - - class MultigridWatcherSetup(BaseModel): source: Path skip_existing_processing: bool = False @@ -280,10 +154,6 @@ class MultigridWatcherSetup(BaseModel): rsync_restarts: List[str] = [] -class CurrentGainRef(BaseModel): - path: str - - class Token(BaseModel): access_token: str token_type: str @@ -296,47 +166,6 @@ class Token(BaseModel): """ -class TomoProcessFile(BaseModel): - path: str - description: str - tag: str - image_number: int - pixel_size: float - dose_per_frame: Optional[float] - frame_count: int - tilt_axis: Optional[float] - mc_uuid: Optional[int] = None - voltage: float = 300 - mc_binning: int = 1 - gain_ref: Optional[str] = None - extract_downscale: int = 1 - eer_fractionation_file: Optional[str] = None - group_tag: Optional[str] = None - - -class TiltInfo(BaseModel): - tilt_series_tag: str - movie_path: str - source: str - - -class TiltSeriesInfo(BaseModel): - session_id: int - tag: str - source: str - - -class TiltSeriesGroupInfo(BaseModel): - tags: List[str] - source: str - tilt_series_lengths: List[int] - - -class CompletedTiltSeries(BaseModel): - tilt_series: List[str] - rsync_source: str - - class ProcessingParametersTomo(BaseModel): dose_per_frame: Optional[float] frame_count: int @@ -358,3 +187,8 @@ class Base(BaseModel): dose_per_frame: Optional[float] gain_ref: Optional[str] eer_fractionation: int + + +class CompletedTiltSeries(BaseModel): + tilt_series: List[str] + rsync_source: str diff --git a/src/murfey/util/route_manifest.yaml b/src/murfey/util/route_manifest.yaml new file mode 100644 index 000000000..3775c6b0e --- /dev/null +++ b/src/murfey/util/route_manifest.yaml @@ -0,0 +1,1161 @@ +murfey.instrument_server.api.router: + - path: /health + function: health + path_params: [] + methods: + - GET + - path: /token + function: token_handshake + path_params: [] + methods: + - POST + - path: /sessions/{session_id}/token + function: token_handshake_for_session + path_params: + - name: session_id + type: int + methods: + - POST + - path: /sessions/{session_id}/check_token + function: check_token + path_params: [] + methods: + - GET + - path: /sessions/{session_id}/multigrid_watcher + function: setup_multigrid_watcher + path_params: [] + methods: + - POST + - path: /sessions/{session_id}/start_multigrid_watcher + function: start_multigrid_watcher + path_params: [] + methods: + - POST + - path: /sessions/{session_id}/multigrid_watcher/{label} + function: stop_multigrid_watcher + path_params: + - name: label + type: str + methods: + - DELETE + - path: /sessions/{session_id}/stop_rsyncer + function: stop_rsyncer + path_params: [] + methods: + - POST + - path: /sessions/{session_id}/remove_rsyncer + function: remove_rsyncer + path_params: [] + methods: + - POST + - path: /sessions/{session_id}/abandon_controller + function: abandon_controller + path_params: [] + methods: + - POST + - path: /sessions/{session_id}/finalise_rsyncer + function: finalise_rsyncer + path_params: [] + methods: + - POST + - path: /sessions/{session_id}/finalise_session + function: finalise_session + path_params: [] + methods: + - POST + - path: /sessions/{session_id}/restart_rsyncer + function: restart_rsyncer + path_params: [] + methods: + - POST + - path: /sessions/{session_id}/rsyncer_info + function: get_rsyncer_info + path_params: [] + methods: + - GET + - path: /sessions/{session_id}/analyser_info + function: get_analyser_info + path_params: [] + methods: + - GET + - path: /sessions/{session_id}/processing_parameters + function: register_processing_parameters + path_params: [] + methods: + - POST + - path: /instruments/{instrument_name}/sessions/{session_id}/possible_gain_references + function: get_possible_gain_references + path_params: + - name: instrument_name + type: str + methods: + - GET + - path: /instruments/{instrument_name}/sessions/{session_id}/upload_gain_reference + function: upload_gain_reference + path_params: + - name: instrument_name + type: str + methods: + - POST + - path: /visits/{visit_name}/sessions/{session_id}/upstream_tiff_data_request + function: gather_upstream_tiffs + path_params: + - name: visit_name + type: str + methods: + - POST +murfey.server.api.auth.router: + - path: /token + function: generate_token + path_params: [] + methods: + - POST + - path: /sessions/{session_id}/token + function: mint_session_token + path_params: [] + methods: + - GET + - path: /validate_token + function: simple_token_validation + path_params: [] + methods: + - GET +murfey.server.api.bootstrap.bootstrap: + - path: /bootstrap/ + function: get_bootstrap_instructions + path_params: [] + methods: + - GET + - path: /bootstrap/pip.whl + function: get_pip_wheel + path_params: [] + methods: + - GET + - path: /bootstrap/murfey.whl + function: get_murfey_wheel + path_params: [] + methods: + - GET +murfey.server.api.bootstrap.cygwin: + - path: /cygwin/setup-x86_64.exe + function: get_cygwin_setup + path_params: [] + methods: + - GET + - path: /cygwin/{request_path:path} + function: parse_cygwin_request + path_params: + - name: request_path + type: str + methods: + - GET +murfey.server.api.bootstrap.msys2: + - path: /msys2/config/pacman.d.zip + function: get_pacman_mirrors + path_params: [] + methods: + - GET + - path: /msys2/repo/distrib/{setup_file} + function: get_msys2_setup + path_params: + - name: setup_file + type: str + methods: + - GET + - path: /msys2/repo/ + function: get_msys2_main_index + path_params: [] + methods: + - GET + - path: /msys2/repo/{system}/ + function: get_msys2_environment_index + path_params: + - name: system + type: str + methods: + - GET + - path: /msys2/repo/{system}/{environment}/ + function: get_msys2_package_index + path_params: + - name: system + type: str + - name: environment + type: str + methods: + - GET + - path: /msys2/repo/{system}/{environment}/{package} + function: get_msys2_package_file + path_params: + - name: system + type: str + - name: environment + type: str + - name: package + type: str + methods: + - GET +murfey.server.api.bootstrap.plugins: + - path: /plugins/instruments/{instrument_name}/{package} + function: get_plugin_wheel + path_params: + - name: instrument_name + type: str + - name: package + type: str + methods: + - GET +murfey.server.api.bootstrap.pypi: + - path: /pypi/ + function: get_pypi_index + path_params: [] + methods: + - GET + - path: /pypi/{package}/ + function: get_pypi_package_downloads_list + path_params: + - name: package + type: str + methods: + - GET + - path: /pypi/{package}/{filename} + function: get_pypi_file + path_params: + - name: package + type: str + - name: filename + type: str + methods: + - GET +murfey.server.api.bootstrap.rust: + - path: /rust/cargo/config.toml + function: get_cargo_config + path_params: [] + methods: + - GET + - path: /rust/index/ + function: get_index_page + path_params: [] + methods: + - GET + - path: /rust/index/config.json + function: get_index_config + path_params: [] + methods: + - GET + - path: /rust/index/{c1}/{c2}/{package} + function: get_index_package_metadata + path_params: + - name: c1 + type: str + - name: c2 + type: str + - name: package + type: str + methods: + - GET + - path: /rust/index/{n}/{package} + function: get_index_package_metadata_for_short_package_names + path_params: + - name: n + type: str + - name: package + type: str + methods: + - GET + - path: /rust/crates/{package}/{version}/download + function: get_rust_package_download + path_params: + - name: package + type: str + - name: version + type: str + methods: + - GET + - path: /rust/api/v1/crates + function: get_rust_api_package_index + path_params: [] + methods: + - GET + - path: /rust/api/v1/crates/{package} + function: get_rust_api_package_info + path_params: + - name: package + type: str + methods: + - GET + - path: /rust/api/v1/crates/{package}/versions + function: get_rust_api_package_versions + path_params: + - name: package + type: str + methods: + - GET + - path: /rust/api/v1/crates/{package}/{version}/download + function: get_rust_api_package_download + path_params: + - name: package + type: str + - name: version + type: str + methods: + - GET + - path: /rust/crates/{package}/{crate} + function: get_rust_package_crate + path_params: + - name: package + type: str + - name: crate + type: str + methods: + - GET +murfey.server.api.bootstrap.version: + - path: /version/ + function: get_version + path_params: [] + methods: + - GET +murfey.server.api.clem.router: + - path: /sessions/{session_id}/clem/lif_files + function: register_lif_file + path_params: + - name: session_id + type: int + methods: + - POST + - path: /sessions/{session_id}/clem/tiff_files + function: register_tiff_file + path_params: + - name: session_id + type: int + methods: + - POST + - path: /sessions/{session_id}/clem/metadata_files + function: register_clem_metadata + path_params: + - name: session_id + type: int + methods: + - POST + - path: /sessions/{session_id}/clem/image_series + function: register_image_series + path_params: + - name: session_id + type: int + methods: + - POST + - path: /sessions/{session_id}/clem/image_stacks + function: register_image_stack + path_params: + - name: session_id + type: int + methods: + - POST + - path: /sessions/{session_id}/clem/preprocessing/process_raw_lifs + function: process_raw_lifs + path_params: + - name: session_id + type: int + methods: + - POST + - path: /sessions/{session_id}/clem/preprocessing/process_raw_tiffs + function: process_raw_tiffs + path_params: + - name: session_id + type: int + methods: + - POST + - path: /sessions/{session_id}/clem/processing/align_and_merge_stacks + function: align_and_merge_stacks + path_params: + - name: session_id + type: int + methods: + - POST +murfey.server.api.display.router: + - path: /display/instruments/{instrument_name}/instrument_name + function: get_instrument_display_name + path_params: + - name: instrument_name + type: str + methods: + - GET + - path: /display/instruments/{instrument_name}/image/ + function: get_mic_image + path_params: + - name: instrument_name + type: str + methods: + - GET + - path: /display/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares/{grid_square_name}/image + function: get_grid_square_img + path_params: + - name: session_id + type: int + - name: dcgid + type: int + - name: grid_square_name + type: int + methods: + - GET + - path: /display/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares/{grid_square_name}/foil_holes/{foil_hole_name}/image + function: get_foil_hole_img + path_params: + - name: session_id + type: int + - name: dcgid + type: int + - name: grid_square_name + type: int + - name: foil_hole_name + type: int + methods: + - GET +murfey.server.api.file_manip.router: + - path: /file_manipulation/visits/{visit_name}/{session_id}/suggested_path + function: suggest_path + path_params: + - name: visit_name + type: str + - name: session_id + type: int + methods: + - POST + - path: /file_manipulation/sessions/{session_id}/make_rsyncer_destination + function: make_rsyncer_destination + path_params: + - name: session_id + type: int + methods: + - POST + - path: /file_manipulation/sessions/{session_id}/process_gain + function: process_gain + path_params: [] + methods: + - POST + - path: /file_manipulation/visits/{visit_name}/{session_id}/eer_fractionation_file + function: write_eer_fractionation_file + path_params: + - name: visit_name + type: str + - name: session_id + type: int + methods: + - POST +murfey.server.api.hub.router: + - path: /instruments + function: get_instrument_info + path_params: [] + methods: + - GET + - path: /instrument/{instrument_name}/image + function: get_instrument_image + path_params: + - name: instrument_name + type: str + methods: + - GET +murfey.server.api.instrument.router: + - path: /instrument_server/instruments/{instrument_name}/sessions/{session_id}/activate_instrument_server + function: activate_instrument_server_for_session + path_params: + - name: instrument_name + type: str + - name: session_id + type: int + methods: + - POST + - path: /instrument_server/instruments/{instrument_name}/sessions/{session_id}/active + function: check_if_session_is_active + path_params: + - name: instrument_name + type: str + - name: session_id + type: int + methods: + - GET + - path: /instrument_server/sessions/{session_id}/multigrid_watcher + function: setup_multigrid_watcher + path_params: [] + methods: + - POST + - path: /instrument_server/sessions/{session_id}/start_multigrid_watcher + function: start_multigrid_watcher + path_params: [] + methods: + - POST + - path: /instrument_server/sessions/{session_id}/provided_processing_parameters + function: pass_proc_params_to_instrument_server + path_params: [] + methods: + - POST + - path: /instrument_server/instruments/{instrument_name}/instrument_server + function: check_instrument_server + path_params: + - name: instrument_name + type: str + methods: + - GET + - path: /instrument_server/instruments/{instrument_name}/sessions/{session_id}/possible_gain_references + function: get_possible_gain_references + path_params: + - name: instrument_name + type: str + methods: + - GET + - path: /instrument_server/sessions/{session_id}/upload_gain_reference + function: request_gain_reference_upload + path_params: [] + methods: + - POST + - path: /instrument_server/visits/{visit_name}/{session_id}/upstream_tiff_data_request + function: request_upstream_tiff_data_download + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /instrument_server/sessions/{session_id}/stop_rsyncer + function: stop_rsyncer + path_params: [] + methods: + - POST + - path: /instrument_server/sessions/{session_id}/finalise_rsyncer + function: finalise_rsyncer + path_params: [] + methods: + - POST + - path: /instrument_server/sessions/{session_id}/finalise_session + function: finalise_session + path_params: [] + methods: + - POST + - path: /instrument_server/sessions/{session_id}/abandon_session + function: abandon_session + path_params: [] + methods: + - POST + - path: /instrument_server/sessions/{session_id}/remove_rsyncer + function: remove_rsyncer + path_params: [] + methods: + - POST + - path: /instrument_server/sessions/{session_id}/restart_rsyncer + function: restart_rsyncer + path_params: [] + methods: + - POST + - path: /instrument_server/instruments/{instrument_name}/sessions/{session_id}/rsyncer_info + function: get_rsyncer_info + path_params: + - name: instrument_name + type: str + methods: + - GET +murfey.server.api.mag_table.router: + - path: /mag_table/mag_table/ + function: get_mag_table + path_params: [] + methods: + - GET + - path: /mag_table/mag_table/ + function: add_to_mag_table + path_params: [] + methods: + - POST + - path: /mag_table/mag_table/{mag} + function: remove_mag_table_row + path_params: + - name: mag + type: int + methods: + - DELETE +murfey.server.api.processing_parameters.router: + - path: /session_parameters/sessions/{session_id}/session_processing_parameters + function: get_session_processing_parameters + path_params: [] + methods: + - GET + - path: /session_parameters/sessions/{session_id}/session_processing_parameters + function: set_session_processing_parameters + path_params: [] + methods: + - POST +murfey.server.api.prometheus.router: + - path: /prometheus/visits/{visit_name}/increment_rsync_file_count + function: increment_rsync_file_count + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /prometheus/visits/{visit_name}/increment_rsync_transferred_files + function: increment_rsync_transferred_files + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /prometheus/visits/{visit_name}/increment_rsync_transferred_files_prometheus + function: increment_rsync_transferred_files_prometheus + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /prometheus/visits/{visit_name}/monitoring/{on} + function: change_monitoring_status + path_params: + - name: visit_name + type: str + - name: "on" + type: int + methods: + - POST + - path: /prometheus/metrics/{metric_name} + function: inspect_prometheus_metrics + path_params: + - name: metric_name + type: str + methods: + - GET +murfey.server.api.session_control.correlative_router: + - path: /session_control/correlative/sessions/{session_id}/upstream_visits + function: find_upstream_visits + path_params: [] + methods: + - GET + - path: /session_control/correlative/visits/{visit_name}/{session_id}/upstream_tiff_paths + function: gather_upstream_tiffs + path_params: + - name: visit_name + type: str + - name: session_id + type: int + methods: + - GET + - path: /session_control/correlative/visits/{visit_name}/{session_id}/upstream_tiff/{tiff_path:path} + function: get_tiff + path_params: + - name: visit_name + type: str + - name: session_id + type: int + - name: tiff_path + type: str + methods: + - GET +murfey.server.api.session_control.router: + - path: /session_control/time + function: get_current_timestamp + path_params: [] + methods: + - GET + - path: /session_control/instruments/{instrument_name}/machine + function: machine_info_by_instrument + path_params: + - name: instrument_name + type: str + methods: + - GET + - path: /session_control/new_client_id/ + function: new_client_id + path_params: [] + methods: + - GET + - path: /session_control/instruments/{instrument_name}/visits_raw + function: get_current_visits + path_params: + - name: instrument_name + type: str + methods: + - GET + - path: /session_control/instruments/{instrument_name}/clients/{client_id}/session + function: link_client_to_session + path_params: + - name: instrument_name + type: str + - name: client_id + type: int + methods: + - POST + - path: /session_control/visits/{visit_name} + function: register_client_to_visit + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /session_control/sessions + function: get_sessions + path_params: [] + methods: + - GET + - path: /session_control/sessions/{session_id} + function: remove_session + path_params: [] + methods: + - DELETE + - path: /session_control/sessions/{session_id}/successful_processing + function: register_processing_success_in_ispyb + path_params: [] + methods: + - POST + - path: /session_control/instruments/{instrument_name}/failed_client_post + function: failed_client_post + path_params: + - name: instrument_name + type: str + methods: + - POST + - path: /session_control/sessions/{session_id}/rsyncer + function: register_rsyncer + path_params: + - name: session_id + type: int + methods: + - POST + - path: /session_control/sessions/{session_id}/rsyncers + function: get_rsyncers_for_session + path_params: [] + methods: + - GET + - path: /session_control/sessions/{session_id}/rsyncer_stopped + function: register_stopped_rsyncer + path_params: + - name: session_id + type: int + methods: + - POST + - path: /session_control/sessions/{session_id}/rsyncer_started + function: register_restarted_rsyncer + path_params: + - name: session_id + type: int + methods: + - POST + - path: /session_control/sessions/{session_id}/rsyncer + function: delete_rsyncer + path_params: + - name: session_id + type: int + methods: + - DELETE +murfey.server.api.session_control.spa_router: + - path: /session_control/spa/sessions/{session_id}/grid_squares + function: get_grid_squares + path_params: [] + methods: + - GET + - path: /session_control/spa/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares + function: get_grid_squares_from_dcg + path_params: + - name: dcgid + type: int + methods: + - GET + - path: /session_control/spa/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares/{gsid}/foil_holes + function: get_foil_holes_from_grid_square + path_params: + - name: dcgid + type: int + - name: gsid + type: int + methods: + - GET + - path: /session_control/spa/sessions/{session_id}/foil_hole/{fh_name} + function: get_foil_hole + path_params: + - name: fh_name + type: int + methods: + - GET + - path: /session_control/spa/sessions/{session_id}/grid_square/{gsid} + function: register_grid_square + path_params: + - name: gsid + type: int + methods: + - POST + - path: /session_control/spa/sessions/{session_id}/grid_square/{gs_name}/foil_hole + function: register_foil_hole + path_params: + - name: gs_name + type: int + methods: + - POST +murfey.server.api.session_info.correlative_router: + - path: /session_info/correlative/sessions/{session_id}/upstream_visits + function: find_upstream_visits + path_params: [] + methods: + - GET + - path: /session_info/correlative/visits/{visit_name}/{session_id}/upstream_tiff_paths + function: gather_upstream_tiffs + path_params: + - name: visit_name + type: str + - name: session_id + type: int + methods: + - GET + - path: /session_info/correlative/visits/{visit_name}/{session_id}/upstream_tiff/{tiff_path:path} + function: get_tiff + path_params: + - name: visit_name + type: str + - name: session_id + type: int + - name: tiff_path + type: str + methods: + - GET +murfey.server.api.session_info.router: + - path: /session_info/ + function: root + path_params: [] + methods: + - GET + - path: /session_info/health/ + function: health_check + path_params: [] + methods: + - GET + - path: /session_info/connections/ + function: connections_check + path_params: [] + methods: + - GET + - path: /session_info/instruments/{instrument_name}/machine + function: machine_info_by_instrument + path_params: + - name: instrument_name + type: str + methods: + - GET + - path: /session_info/instruments/{instrument_name}/visits_raw + function: get_current_visits + path_params: + - name: instrument_name + type: str + methods: + - GET + - path: /session_info/instruments/{instrument_name}/visits/ + function: all_visit_info + path_params: + - name: instrument_name + type: str + methods: + - GET + - path: /session_info/sessions/{session_id}/rsyncers + function: get_rsyncers_for_client + path_params: [] + methods: + - GET + - path: /session_info/session/{session_id} + function: get_session + path_params: [] + methods: + - GET + - path: /session_info/sessions + function: get_sessions + path_params: [] + methods: + - GET + - path: /session_info/instruments/{instrument_name}/visits/{visit}/session/{name} + function: create_session + path_params: + - name: instrument_name + type: str + - name: visit + type: str + - name: name + type: str + methods: + - POST + - path: /session_info/sessions/{session_id} + function: update_session + path_params: [] + methods: + - POST + - path: /session_info/sessions/{session_id} + function: remove_session + path_params: [] + methods: + - DELETE + - path: /session_info/instruments/{instrument_name}/visits/{visit_name}/sessions + function: get_sessions_with_visit + path_params: + - name: instrument_name + type: str + - name: visit_name + type: str + methods: + - GET + - path: /session_info/instruments/{instrument_name}/sessions + function: get_sessions_by_instrument_name + path_params: + - name: instrument_name + type: str + methods: + - GET + - path: /session_info/sessions/{session_id}/data_collection_groups + function: get_dc_groups + path_params: [] + methods: + - GET + - path: /session_info/sessions/{session_id}/data_collection_groups/{dcgid}/data_collections + function: get_data_collections + path_params: + - name: dcgid + type: int + methods: + - GET + - path: /session_info/clients + function: get_clients + path_params: [] + methods: + - GET + - path: /session_info/num_movies + function: count_number_of_movies + path_params: [] + methods: + - GET + - path: /session_info/sessions/{session_id}/current_gain_ref + function: update_current_gain_ref + path_params: [] + methods: + - PUT +murfey.server.api.session_info.spa_router: + - path: /session_info/spa/sessions/{session_id}/spa_processing_parameters + function: get_spa_proc_param_details + path_params: [] + methods: + - GET + - path: /session_info/spa/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares/{gsid}/foil_holes/{fhid}/num_movies + function: get_number_of_movies_from_foil_hole + path_params: + - name: session_id + type: int + - name: dcgid + type: int + - name: gsid + type: int + - name: fhid + type: int + methods: + - GET + - path: /session_info/spa/sessions/{session_id}/grid_squares + function: get_grid_squares + path_params: [] + methods: + - GET + - path: /session_info/spa/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares + function: get_grid_squares_from_dcg + path_params: + - name: dcgid + type: int + methods: + - GET + - path: /session_info/spa/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares/{gsid}/foil_holes + function: get_foil_holes_from_grid_square + path_params: + - name: dcgid + type: int + - name: gsid + type: int + methods: + - GET + - path: /session_info/spa/sessions/{session_id}/foil_hole/{fh_name} + function: get_foil_hole + path_params: + - name: fh_name + type: int + methods: + - GET +murfey.server.api.session_info.tomo_router: + - path: /session_info/tomo/sessions/{session_id}/tilt_series/{tilt_series_tag}/tilts + function: get_tilts + path_params: + - name: tilt_series_tag + type: str + methods: + - GET +murfey.server.api.spa.router: + - path: /sessions/{session_id}/cryolo_model + function: get_cryolo_model_path + path_params: + - name: session_id + type: int + methods: + - GET +murfey.server.api.websocket.ws: + - path: /ws/test/{client_id} + function: websocket_endpoint + path_params: + - name: client_id + type: int + methods: [] + - path: /ws/connect/{client_id} + function: websocket_connection_endpoint + path_params: + - name: client_id + type: typing.Union[int, str] + methods: [] + - path: /ws/test/{client_id} + function: close_ws_connection + path_params: + - name: client_id + type: int + methods: + - DELETE + - path: /ws/connect/{client_id} + function: close_unrecorded_ws_connection + path_params: + - name: client_id + type: typing.Union[int, str] + methods: + - DELETE +murfey.server.api.workflow.correlative_router: + - path: /workflow/correlative/visit/{visit_name}/samples + function: get_samples + path_params: + - name: visit_name + type: str + methods: + - GET + - path: /workflow/correlative/visit/{visit_name}/sample_group + function: register_sample_group + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /workflow/correlative/visit/{visit_name}/sample + function: register_sample + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /workflow/correlative/visit/{visit_name}/subsample + function: register_subsample + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /workflow/correlative/visit/{visit_name}/sample_image + function: register_sample_image + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /workflow/correlative/visits/{year}/{visit_name}/{session_id}/make_milling_gif + function: make_gif + path_params: + - name: year + type: int + - name: visit_name + type: str + - name: session_id + type: int + methods: + - POST +murfey.server.api.workflow.router: + - path: /workflow/visits/{visit_name}/{session_id}/register_data_collection_group + function: register_dc_group + path_params: + - name: visit_name + type: typing.Any + methods: + - POST + - path: /workflow/visits/{visit_name}/{session_id}/start_data_collection + function: start_dc + path_params: + - name: visit_name + type: typing.Any + methods: + - POST + - path: /workflow/visits/{visit_name}/{session_id}/register_processing_job + function: register_proc + path_params: + - name: visit_name + type: str + methods: + - POST +murfey.server.api.workflow.spa_router: + - path: /workflow/spa/sessions/{session_id}/spa_processing_parameters + function: register_spa_proc_params + path_params: [] + methods: + - POST + - path: /workflow/spa/visits/{visit_name}/{session_id}/flush_spa_processing + function: flush_spa_processing + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /workflow/spa/visits/{visit_name}/{session_id}/spa_preprocess + function: request_spa_preprocessing + path_params: + - name: visit_name + type: str + methods: + - POST +murfey.server.api.workflow.tomo_router: + - path: /workflow/tomo/sessions/{session_id}/tomography_processing_parameters + function: register_tomo_proc_params + path_params: [] + methods: + - POST + - path: /workflow/tomo/visits/{visit_name}/{session_id}/flush_tomography_processing + function: flush_tomography_processing + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /workflow/tomo/visits/{visit_name}/tilt_series + function: register_tilt_series + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /workflow/tomo/sessions/{session_id}/tilt_series_length + function: register_tilt_series_length + path_params: + - name: session_id + type: int + methods: + - POST + - path: /workflow/tomo/visits/{visit_name}/{session_id}/tomography_preprocess + function: request_tomography_preprocessing + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /workflow/tomo/visits/{visit_name}/{session_id}/completed_tilt_series + function: register_completed_tilt_series + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /workflow/tomo/visits/{visit_name}/rerun_tilt_series + function: register_tilt_series_for_rerun + path_params: + - name: visit_name + type: str + methods: + - POST + - path: /workflow/tomo/visits/{visit_name}/{session_id}/tilt + function: register_tilt + path_params: + - name: visit_name + type: str + methods: + - POST diff --git a/src/murfey/workflows/clem/process_raw_lifs.py b/src/murfey/workflows/clem/process_raw_lifs.py index 1d56bff68..22f50c5c0 100644 --- a/src/murfey/workflows/clem/process_raw_lifs.py +++ b/src/murfey/workflows/clem/process_raw_lifs.py @@ -3,6 +3,7 @@ The recipe referred to here is stored on GitLab. """ +from logging import getLogger from pathlib import Path from typing import Optional @@ -11,6 +12,8 @@ except AttributeError: pass # Ignore if ISPyB credentials environment variable not set +logger = getLogger("murfey.workflows.clem.process_raw_lifs") + def zocalo_cluster_request( file: Path, @@ -43,24 +46,28 @@ def zocalo_cluster_request( # Load machine config to get the feedback queue feedback_queue: str = messenger.feedback_queue - # Send the message - # The keys under "parameters" will populate all the matching fields in {} - # in the processing recipe - messenger.send( - "processing_recipe", - { - "recipes": ["clem-lif-to-stack"], - "parameters": { - # Job parameters - "lif_file": f"{str(file)}", - "root_folder": root_folder, - # Other recipe parameters - "session_dir": f"{str(session_dir)}", - "session_id": session_id, - "job_name": job_name, - "feedback_queue": feedback_queue, - }, + # Construct recipe and submit it for processing + recipe = { + "recipes": ["clem-lif-to-stack"], + "parameters": { + # Job parameters + "lif_file": f"{str(file)}", + "root_folder": root_folder, + # Other recipe parameters + "session_dir": f"{str(session_dir)}", + "session_id": session_id, + "job_name": job_name, + "feedback_queue": feedback_queue, }, + } + logger.debug( + f"Submitting LIF processing request to {messenger.feedback_queue!r} " + "with the following recipe: \n" + f"{recipe}" + ) + messenger.send( + queue="processing_recipe", + message=recipe, new_connection=True, ) else: diff --git a/src/murfey/workflows/clem/process_raw_tiffs.py b/src/murfey/workflows/clem/process_raw_tiffs.py index 52c371092..2c0c1a5b3 100644 --- a/src/murfey/workflows/clem/process_raw_tiffs.py +++ b/src/murfey/workflows/clem/process_raw_tiffs.py @@ -3,9 +3,12 @@ The recipe referred to here is stored on GitLab. """ +from logging import getLogger from pathlib import Path from typing import Optional +logger = getLogger("murfey.workflows.clem.process_raw_tiffs") + try: from murfey.server.ispyb import TransportManager # Session except AttributeError: @@ -50,23 +53,30 @@ def zocalo_cluster_request( # Load machine config to get the feedback queue feedback_queue: str = messenger.feedback_queue - messenger.send( - "processing_recipe", - { - "recipes": ["clem-tiff-to-stack"], - "parameters": { - # Job parameters - "tiff_list": "null", - "tiff_file": f"{str(tiff_list[0])}", - "root_folder": root_folder, - "metadata": f"{str(metadata)}", - # Other recipe parameters - "session_dir": f"{str(session_dir)}", - "session_id": session_id, - "job_name": job_name, - "feedback_queue": feedback_queue, - }, + # Construct recipe and submit it for processing + recipe = { + "recipes": ["clem-tiff-to-stack"], + "parameters": { + # Job parameters + "tiff_list": "null", + "tiff_file": f"{str(tiff_list[0])}", + "root_folder": root_folder, + "metadata": f"{str(metadata)}", + # Other recipe parameters + "session_dir": f"{str(session_dir)}", + "session_id": session_id, + "job_name": job_name, + "feedback_queue": feedback_queue, }, + } + logger.debug( + f"Submitting TIFF processing request to {messenger.feedback_queue!r} " + "with the following recipe: \n" + f"{recipe}" + ) + messenger.send( + queue="processing_recipe", + message=recipe, new_connection=True, ) else: diff --git a/src/murfey/workflows/spa/flush_spa_preprocess.py b/src/murfey/workflows/spa/flush_spa_preprocess.py index db933198d..27664aa41 100644 --- a/src/murfey/workflows/spa/flush_spa_preprocess.py +++ b/src/murfey/workflows/spa/flush_spa_preprocess.py @@ -6,9 +6,10 @@ from sqlalchemy.exc import NoResultFound from sqlmodel import Session, select -from murfey.server import _murfey_id, _transport_object, sanitise +from murfey.server import _transport_object from murfey.server.api.auth import MurfeySessionID -from murfey.util import secure_path +from murfey.server.feedback import _murfey_id +from murfey.util import sanitise, secure_path from murfey.util.config import get_machine_config, get_microscope from murfey.util.db import ( AutoProcProgram, diff --git a/src/murfey/workflows/spa/picking.py b/src/murfey/workflows/spa/picking.py index 900cf6769..e383f438f 100644 --- a/src/murfey/workflows/spa/picking.py +++ b/src/murfey/workflows/spa/picking.py @@ -6,13 +6,13 @@ from sqlmodel import Session, select import murfey.server.prometheus as prom -from murfey.server import ( +from murfey.server import _transport_object +from murfey.server.feedback import ( _app_id, _flush_class2d, _pj_id, _register_class_selection, _register_incomplete_2d_batch, - _transport_object, ) from murfey.util.config import get_machine_config from murfey.util.db import ( diff --git a/tests/client/test.py b/tests/client/test.py deleted file mode 100644 index 1ccba7389..000000000 --- a/tests/client/test.py +++ /dev/null @@ -1,14 +0,0 @@ -from __future__ import annotations - -import os - -os.environ["BEAMLINE"] = "m12" -import pytest - -import murfey.client.main as main - - -@pytest.mark.xfail -def test_get_visit_info(): - response = main.get_visit_info("cm31111-1") # Should be valid until end of 2022 - assert response.status_code == 200 diff --git a/tests/client/test_context.py b/tests/client/test_context.py index c6431c72e..2c71d70fc 100644 --- a/tests/client/test_context.py +++ b/tests/client/test_context.py @@ -25,6 +25,7 @@ def test_tomography_context_add_tomo_tilt(mock_post, mock_get, tmp_path): sources=[tmp_path], default_destinations={tmp_path: str(tmp_path)}, instrument_name="", + visit="test", ) context = TomographyContext("tomo", tmp_path) (tmp_path / "Position_1_001_[30.0]_date_time_fractions.tiff").touch() @@ -80,6 +81,7 @@ def test_tomography_context_add_tomo_tilt_out_of_order(mock_post, mock_get, tmp_ sources=[tmp_path], default_destinations={tmp_path: str(tmp_path)}, instrument_name="", + visit="test", ) context = TomographyContext("tomo", tmp_path) (tmp_path / "Position_1_001_[30.0]_date_time_fractions.tiff").touch() @@ -163,6 +165,7 @@ def test_tomography_context_add_tomo_tilt_delayed_tilt(mock_post, mock_get, tmp_ sources=[tmp_path], default_destinations={tmp_path: str(tmp_path)}, instrument_name="", + visit="test", ) context = TomographyContext("tomo", tmp_path) (tmp_path / "Position_1_001_[30.0]_date_time_fractions.tiff").touch() @@ -226,6 +229,7 @@ def test_setting_tilt_series_size_and_completion_from_mdoc_parsing( sources=[tmp_path], default_destinations={tmp_path: str(tmp_path)}, instrument_name="", + visit="test", ) context = TomographyContext("tomo", tmp_path) assert len(context._tilt_series_sizes) == 0 diff --git a/tests/instrument_server/test_api.py b/tests/instrument_server/test_api.py index b4fa9cd73..785f1e812 100644 --- a/tests/instrument_server/test_api.py +++ b/tests/instrument_server/test_api.py @@ -11,6 +11,7 @@ upload_gain_reference, ) from murfey.util import posix_path +from murfey.util.api import url_path_for test_get_murfey_url_params_matrix = ( # Server URL to use @@ -119,7 +120,7 @@ def test_upload_gain_reference( ) # Check that the machine config request was called - machine_config_url = f"{server_url}/instruments/{instrument_name}/machine" + machine_config_url = f"{server_url}{url_path_for('session_control.router', 'machine_info_by_instrument', instrument_name=instrument_name)}" mock_request.get.assert_called_once_with( machine_config_url, headers={"Authorization": ANY}, diff --git a/tests/server/api/test_movies.py b/tests/server/api/test_movies.py index 59f57bc1f..bc633016b 100644 --- a/tests/server/api/test_movies.py +++ b/tests/server/api/test_movies.py @@ -52,7 +52,9 @@ def login(test_user): @patch("murfey.server.api.auth.check_user", return_value=True) def test_movie_count(mock_check, test_user): token = login(test_user) - response = client.get("/num_movies", headers={"Authorization": f"Bearer {token}"}) + response = client.get( + "/session_info/num_movies", headers={"Authorization": f"Bearer {token}"} + ) assert mock_check.called_once() assert response.status_code == 200 assert len(mock_session.method_calls) == 2 diff --git a/tests/server/test_main.py b/tests/server/test_main.py index 5a957e67e..24aab534b 100644 --- a/tests/server/test_main.py +++ b/tests/server/test_main.py @@ -30,7 +30,7 @@ def login(test_user): @patch("murfey.server.api.auth.check_user", return_value=True) def test_read_main(mock_check, test_user): token = login(test_user) - response = client.get("/", headers={"Authorization": f"Bearer {token}"}) + response = client.get("/session_info", headers={"Authorization": f"Bearer {token}"}) assert mock_check.called_once() assert response.status_code == 200 assert "" in response.text.lower() diff --git a/tests/util/test_client.py b/tests/util/test_client.py index 98e920f88..6bbfb468a 100644 --- a/tests/util/test_client.py +++ b/tests/util/test_client.py @@ -6,11 +6,8 @@ from pytest import mark -from murfey.util.client import ( - _get_visit_list, - read_config, - set_default_acquisition_output, -) +from murfey.client import _get_visit_list +from murfey.util.client import read_config, set_default_acquisition_output from murfey.util.models import Visit test_read_config_params_matrix = ( @@ -73,7 +70,7 @@ def test_read_config( @mark.parametrize("test_params", test_get_visit_list_params_matrix) -@patch("murfey.util.client.requests") +@patch("murfey.client.requests") def test_get_visit_list( mock_request, test_params: tuple[str], @@ -112,7 +109,9 @@ def test_get_visit_list( visits = _get_visit_list(urlparse(server_url), instrument_name) # Check that request was sent with the correct URL - expected_url = f"{server_url}/instruments/{instrument_name}/visits_raw" + expected_url = ( + f"{server_url}/session_control/instruments/{instrument_name}/visits_raw" + ) mock_request.get.assert_called_once_with(expected_url) # Check that expected outputs are correct (order-sensitive) diff --git a/tests/workflows/clem/test_process_raw_lifs.py b/tests/workflows/clem/test_process_raw_lifs.py index 857b0979c..d6d7bdea8 100644 --- a/tests/workflows/clem/test_process_raw_lifs.py +++ b/tests/workflows/clem/test_process_raw_lifs.py @@ -70,7 +70,7 @@ def test_zocalo_cluster_request( # Check that it sends the expected recipe mock_transport.send.assert_called_once_with( - "processing_recipe", - sent_recipe, + queue="processing_recipe", + message=sent_recipe, new_connection=True, ) diff --git a/tests/workflows/clem/test_process_raw_tiffs.py b/tests/workflows/clem/test_process_raw_tiffs.py index 885aa69a6..c2972f478 100644 --- a/tests/workflows/clem/test_process_raw_tiffs.py +++ b/tests/workflows/clem/test_process_raw_tiffs.py @@ -97,7 +97,7 @@ def test_zocalo_cluster_request( # Check that it sends the expected recipe mock_transport.send.assert_called_once_with( - "processing_recipe", - sent_recipe, + queue="processing_recipe", + message=sent_recipe, new_connection=True, )