Add GPU indexing to the live service

.dockerignore

@@ -13,6 +13,7 @@ Makefile
 *.tar
 .*_cache
 ENV/
+*.so
 # This is better, but doesn't work with podman. sigh.
 # *
 # !**/CMakeLists.txt
@@ -24,3 +25,4 @@ ENV/
 # !**/*.cc
 # !**/*.cpp
 # !**/*.h
+

CMakeLists.txt

@@ -43,10 +43,9 @@ include(AlwaysColourCompilation)
 
 include_directories(include)
 
-# Find Python
-find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
-include_directories(${Python3_INCLUDE_DIRS})
+find_package(Python COMPONENTS Interpreter Development.Module NumPy REQUIRED)
 
+# Dependency fetching
 # Find packages that are in conda environment
 find_package(fmt REQUIRED)
 find_package(spdlog REQUIRED)
@@ -72,6 +71,21 @@ FetchContent_Declare(
 )
 FetchContent_MakeAvailable(argparse)
 
+# Required for Python module build of ffbidx
+set(BUILD_FAST_INDEXER ON CACHE BOOL "" FORCE)
+set(PYTHON_MODULE ON CACHE BOOL "" FORCE)
+set(PYTHON_MODULE_RPATH ON CACHE BOOL "" FORCE)
+
+# ffbidx uses this to find site-packages
+set(PYTHON_MODULE_PATH ${Python_SITEARCH} CACHE PATH "ffbidx python site-packages path" FORCE)
+
+FetchContent_Declare(
+  ffbidx
+  GIT_REPOSITORY https://github.com/paulscherrerinstitute/fast-feedback-indexer.git
+  GIT_TAG main
+)
+FetchContent_MakeAvailable(ffbidx)
+
 # GTest - try find_package first, fallback to FetchContent
 find_package(GTest QUIET)
 if(NOT GTest_FOUND AND NOT TARGET GTest::gtest_main)
@@ -191,5 +205,3 @@ add_subdirectory(h5read)
 add_subdirectory(baseline)
 add_subdirectory(spotfinder)
 add_subdirectory(integrator)
-
-configure_file(${CMAKE_SOURCE_DIR}/scripts/ssx_index ${CMAKE_BINARY_DIR}/bin/ssx_index COPYONLY)

Dockerfile

@@ -24,29 +24,30 @@ RUN micromamba create -y -f /opt/runtime-environment.yml -p /opt/ffs
 
 # Copy source
 COPY . /opt/ffs_src
+ENV CMAKE_GENERATOR Ninja
 
 # Build the C++/CUDA backend
 WORKDIR /opt/build
-RUN cmake /opt/ffs_src -G Ninja -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/opt/ffs -DHDF5_ROOT=/opt/ffs -DCUDA_ARCH=80
-RUN cmake --build . --target spotfinder --target spotfinder32
+RUN cmake /opt/ffs_src \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCMAKE_INSTALL_PREFIX=/opt/ffs \
+    -DHDF5_ROOT=/opt/ffs \
+    -DPython3_ROOT_DIR=/opt/ffs \
+    -DCUDA_ARCH=80
 
-RUN cmake --install . --component Runtime
+RUN cmake --build .
 
-# Build and install dx2 submodule
-WORKDIR /opt/build_dx2
-RUN cmake /opt/ffs_src/dx2 -G Ninja -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/opt/ffs
-RUN cmake --build . && cmake --install .
+RUN cmake --install .
 
 # Install Python package
-WORKDIR /opt/build
 RUN SETUPTOOLS_SCM_PRETEND_VERSION_FOR_FFS=1.0 /opt/ffs/bin/pip3 install /opt/ffs_src
 
 # Now copy this into an isolated runtime container
 FROM nvcr.io/nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu24.04
 
 LABEL org.opencontainers.image.title="fast-feedback-service" \
       org.opencontainers.image.description="GPU-accelerated fast-feedback X-ray diffraction analysis service" \
-      org.opencontainers.image.authors="Nicholas Devenish <nicholas.devenish@diamond.ac.uk>, Dimitrios Vlachos <dimitrios.vlachos@diamond.ac.uk>" \
+      org.opencontainers.image.authors="Nicholas Devenish <nicholas.devenish@diamond.ac.uk>, Dimitrios Vlachos <dimitrios.vlachos@diamond.ac.uk>, James Beilsten-Edmands <james.beilsten-edmands@diamond.ac.uk>" \
       org.opencontainers.image.source="https://github.com/DiamondLightSource/fast-feedback-service" \
       org.opencontainers.image.licenses="BSD-3-Clause"
 

baseline/indexer/CMakeLists.txt

@@ -60,9 +60,11 @@ target_link_libraries(baseline_indexer
 # Create Python module
 nanobind_add_module(index index_module.cpp)
 target_include_directories(index PRIVATE ${CMAKE_SOURCE_DIR}/baseline/predictor)
-target_link_libraries(index PRIVATE Eigen3::Eigen dx2 ${Python3_LIBRARIES})
+target_link_libraries(index PRIVATE Eigen3::Eigen dx2 Python::Module)
 
-set(PYTHON_MODULE_DIR ${CMAKE_SOURCE_DIR}/src/ffs)
-add_custom_command(TARGET index POST_BUILD
-    COMMAND ${CMAKE_COMMAND} -E copy $<TARGET_FILE:index> ${PYTHON_MODULE_DIR}
+
+# Install into the runtime Python sitearch
+install(TARGETS index
+    LIBRARY DESTINATION ${Python_SITEARCH}/ffs
+    COMPONENT Runtime
 )

baseline/indexer/index_module.cpp

@@ -1,11 +1,13 @@
 #include <nanobind/eigen/dense.h>
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/array.h>
+#include <nanobind/stl/string.h>
 #include <nanobind/stl/tuple.h>
 #include <nanobind/stl/vector.h>
 
 #include <dx2/crystal.hpp>
 #include <dx2/detector.hpp>
+#include <dx2/detector_attenuations.hpp>
 #include <dx2/reflection.hpp>
 #include <experimental/mdspan>
 #include <tuple>
@@ -220,6 +222,9 @@ NB_MODULE(index, m) {
       .def_prop_ro("delpsi", [](const IndexingResult &r) { return r.delpsi; })
       .def_prop_ro("rmsds", [](const IndexingResult &r) { return r.rmsds; });
     m.def("make_panel", &make_panel, "Create a configured Panel object");
+    m.def("calculate_mu_for_material_at_wavelength",
+          &calculate_mu_for_material_at_wavelength,
+          "Calculate the absorption coefficient from material and wavelength");
     m.def("ssx_xyz_to_rlp",
           &ssx_xyz_to_rlp,
           nb::arg("xyzobs_px"),

build.sh

@@ -161,6 +161,17 @@ build_directory() {
             make -j"$JOBS"
         fi
     )
+
+    # Install
+    print_status "Installing $description..."
+    (
+        cd "$build_dir"
+        if [[ "$BUILD_CMD" == "ninja" ]]; then
+            ninja "install"
+        else
+            make "install" -j"$JOBS"
+        fi
+    )
 
     print_success "Successfully built $description"
 }

chart/values.yaml

@@ -27,4 +27,4 @@ podSecurityContext:
 # Container-level security context
 containerSecurityContext:
   # k8s-i24-beamline user
-  runAsUser: 36145
+  runAsUser: 36145

environment.yml

@@ -15,6 +15,7 @@ dependencies:
   - hdf5=1.14.6
   - hdf5-external-filter-plugins=0.1.0
   - nlohmann_json=3.12.0
+  - numpy=2.3.4
   - pre-commit=4.3.0
   - pydantic=2.12.3
   - pytest=8.4.2

pyproject.toml

@@ -10,6 +10,7 @@ readme = "README.md"
 authors = [
     { name = "Nicholas Devenish", email = "ndevenish@gmail.com" },
     { name = "Dimitri Vlachos", email = "dimitrios.vlachos@diamond.ac.uk" },
+    { name = "James Beilsten-Edmands", email = "james.beilsten-edmands@diamond.ac.uk"},
 ]
 
 requires-python = ">=3.12"
@@ -22,6 +23,9 @@ dependencies = [
     "zocalo>=1.2.0",
 ]
 
+[project.scripts]
+ssx_index = "ffs.ssx_index:run"
+
 [dependency-groups]
 dev = [
     "pytest",

runtime-environment.yml

@@ -7,10 +7,12 @@ dependencies:
   - gemmi=0.7.3
   - hdf5=1.14.6
   - hdf5-external-filter-plugins=0.1.0
+  - h5py
   - python=3.13.7
   - pip
   - zocalo
   - workflows
   - pydantic
   - rich
   - spdlog=1.15.3
+  - numpy=2.3.4