Skip to content

Commit

Permalink
Merge branch 'main' into 1020_hutch-shutter-i19
Browse files Browse the repository at this point in the history
  • Loading branch information
@noemifrisina
noemifrisina committed Jan 29, 2025
2 parents 14a092f + 388db6a commit d52f517
Show file tree
Hide file tree
Showing 2 changed files with 91 additions and 37 deletions.
25 changes: 21 additions & 4 deletions src/dodal/devices/zocalo/zocalo_results.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,12 @@ class ZocaloResults(StandardReadable, Triggerable):
prefix (str): EPICS PV prefix for the device
use_cpu_and_gpu (bool): When True, ZocaloResults will wait for results from the CPU and the GPU, compare them, and provide a warning if the results differ. When False, ZocaloResults will only use results from the CPU
use_cpu_and_gpu (bool): When True, ZocaloResults will wait for results from the
CPU and the GPU, compare them, and provide a warning if the results differ. When
False, ZocaloResults will only use results from the CPU
use_gpu (bool): When True, ZocaloResults will take the first set of
results that it receives (which are likely the GPU results)
"""

Expand All @@ -142,6 +147,7 @@ def __init__(
timeout_s: float = DEFAULT_TIMEOUT,
prefix: str = "",
use_cpu_and_gpu: bool = False,
use_gpu: bool = False,
) -> None:
self.zocalo_environment = zocalo_environment
self.sort_key = SortKeys[sort_key]
Expand All @@ -151,6 +157,7 @@ def __init__(
self._raw_results_received: Queue = Queue()
self.transport: CommonTransport | None = None
self.use_cpu_and_gpu = use_cpu_and_gpu
self.use_gpu = use_gpu

self.centre_of_mass, self._com_setter = soft_signal_r_and_setter(
Array1D[np.float64], name="centre_of_mass"
Expand Down Expand Up @@ -213,6 +220,11 @@ async def stage(self):
clearing the queue. Plans using this device should wait on ZOCALO_STAGE_GROUP
before triggering processing for the experiment"""

if self.use_cpu_and_gpu and self.use_gpu:
raise ValueError(
"Cannot compare GPU and CPU results and use GPU results at the same time."
)

LOGGER.info("Subscribing to results queue")
try:
self._subscribe_to_results()
Expand Down Expand Up @@ -253,8 +265,13 @@ async def trigger(self):
raw_results = self._raw_results_received.get(timeout=self.timeout_s)
source_of_first_results = source_from_results(raw_results)

# Wait for results from CPU and GPU, warn and continue if only GPU times out. Error if CPU times out
if self.use_gpu and source_of_first_results == ZocaloSource.CPU:
LOGGER.warning(
"Configured to use GPU results but CPU came first, using CPU results."
)

if self.use_cpu_and_gpu:
# Wait for results from CPU and GPU, warn and continue if only GPU times out. Error if CPU times out
if source_of_first_results == ZocaloSource.CPU:
LOGGER.warning("Received zocalo results from CPU before GPU")
raw_results_two_sources = [raw_results]
Expand Down Expand Up @@ -303,7 +320,7 @@ async def trigger(self):
raise err

LOGGER.info(
f"Zocalo results from {ZocaloSource.CPU.value} processing: found {len(raw_results['results'])} crystals."
f"Zocalo results from {source_from_results(raw_results)} processing: found {len(raw_results['results'])} crystals."
)
# Sort from strongest to weakest in case of multiple crystals
await self._put_results(
Expand Down Expand Up @@ -335,7 +352,7 @@ def _receive_result(

results = message.get("results", [])

if self.use_cpu_and_gpu:
if self.use_cpu_and_gpu or self.use_gpu:
self._raw_results_received.put(
{"results": results, "recipe_parameters": recipe_parameters}
)
Expand Down
103 changes: 70 additions & 33 deletions tests/devices/unit_tests/test_zocalo_results.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -257,20 +257,7 @@ def zocalo_plan():
recipe_wrapper,
{},
{
"results": [
{
"centre_of_mass": [
2.207133058984911,
1.4175240054869684,
13.317215363511659,
],
"max_voxel": [2, 1, 13],
"max_count": 702.0,
"n_voxels": 12,
"total_count": 5832.0,
"bounding_box": [[1, 0, 12], [4, 3, 15]],
}
],
"results": [TEST_RESULTS[0]],
"status": "success",
"type": "3d",
},
Expand Down Expand Up @@ -313,6 +300,18 @@ async def test_if_use_cpu_and_gpu_zocalos_then_wait_twice_for_results(
assert zocalo_results._raw_results_received.get.call_count == 2


async def test_if_use_gpu_then_only_use_first_result(
zocalo_results: ZocaloResults, RE: RunEngine
):
zocalo_results.use_gpu = True
await zocalo_results.stage()
zocalo_results._raw_results_received.put([])
zocalo_results._raw_results_received.put([])
zocalo_results._raw_results_received.get = MagicMock()
RE(bps.trigger(zocalo_results, wait=False))
assert zocalo_results._raw_results_received.get.call_count == 1


@patch("dodal.devices.zocalo.zocalo_results.LOGGER")
async def test_source_of_zocalo_results_correctly_identified(
mock_logger, zocalo_results: ZocaloResults, RE: RunEngine
Expand Down Expand Up @@ -351,23 +350,44 @@ async def test_if_zocalo_results_timeout_from_gpu_then_warn(
)


async def test_if_zocalo_results_from_gpu_but_not_cpu_then_error(
async def test_given_comparing_results_if_zocalo_results_from_gpu_but_not_cpu_then_error(
zocalo_results: ZocaloResults, RE: RunEngine
):
zocalo_results.use_cpu_and_gpu = True
await zocalo_results.stage()
zocalo_results._raw_results_received.get = MagicMock(
side_effect=[
{"recipe_parameters": {"test": 0, "gpu": True}, "results": []},
{
"recipe_parameters": {"test": 0, "gpu": True},
"results": [TEST_RESULTS[0]],
},
Empty,
]
)
with pytest.raises(NoResultsFromZocalo):
await zocalo_results.trigger()


async def test_given_using_gpu_results_if_zocalo_results_from_gpu_but_not_cpu_then_uses_gpu(
zocalo_results: ZocaloResults, RE: RunEngine
):
zocalo_results.use_gpu = True
await zocalo_results.stage()
zocalo_results._raw_results_received.get = MagicMock(
side_effect=[
{
"recipe_parameters": {"dcgid": 0, "dcid": 0, "gpu": True},
"results": [TEST_RESULTS[0]],
},
Empty,
]
)
await zocalo_results.trigger()
assert len(await zocalo_results.centre_of_mass.get_value())


@patch("dodal.devices.zocalo.zocalo_results.LOGGER")
async def test_if_cpu_results_arrive_before_gpu_then_warn(
async def test_given_comparing_results_if_cpu_results_arrive_before_gpu_then_warn(
mock_logger, zocalo_results: ZocaloResults, RE: RunEngine
):
zocalo_results.use_cpu_and_gpu = True
Expand Down Expand Up @@ -418,7 +438,7 @@ async def test_if_cpu_results_arrive_before_gpu_then_warn(
],
)
@patch("dodal.devices.zocalo.zocalo_results.LOGGER")
async def test_warning_if_results_are_different(
async def test_given_comparing_results_then_warning_if_results_are_different(
mock_logger, zocalo_results: ZocaloResults, RE: RunEngine, dict1, dict2, output
):
zocalo_results.use_cpu_and_gpu = True
Expand Down Expand Up @@ -476,20 +496,7 @@ def zocalo_plan():
recipe_wrapper,
{},
{
"results": [
{
"centre_of_mass": [
2.207133058984911,
1.4175240054869684,
13.317215363511659,
],
"max_voxel": [2, 1, 13],
"max_count": 702.0,
"n_voxels": 12,
"total_count": 5832.0,
"bounding_box": [[1, 0, 12], [4, 3, 15]],
}
],
"results": [TEST_RESULTS[0]],
"status": "success",
"type": "3d",
},
Expand All @@ -507,7 +514,7 @@ def zocalo_plan():
RE(zocalo_plan())


async def test_given_gpu_enabled_when_no_results_found_then_returns_no_results(
async def test_given_comparing_results_when_no_results_found_then_returns_no_results(
zocalo_results: ZocaloResults,
):
zocalo_results.use_cpu_and_gpu = True
Expand All @@ -520,3 +527,33 @@ async def test_given_gpu_enabled_when_no_results_found_then_returns_no_results(
)
await zocalo_results.trigger()
assert len(await zocalo_results.centre_of_mass.get_value()) == 0


@patch("dodal.devices.zocalo.zocalo_results.LOGGER")
async def test_given_using_gpu_results_if_results_from_cpu_first_then_warn_and_use(
mock_logger: MagicMock, zocalo_results: ZocaloResults
):
zocalo_results.use_gpu = True
await zocalo_results.stage()
zocalo_results._raw_results_received.get = MagicMock(
side_effect=[
{
"recipe_parameters": {"dcgid": 0, "dcid": 0},
"results": [TEST_RESULTS[0]],
},
]
)
await zocalo_results.trigger()
assert len(await zocalo_results.centre_of_mass.get_value())
mock_logger.warning.assert_called_with(
"Configured to use GPU results but CPU came first, using CPU results."
)


async def test_given_using_gpu_results_and_comparing_results_both_on_then_error_when_staged(
zocalo_results: ZocaloResults,
):
zocalo_results.use_gpu = True
zocalo_results.use_cpu_and_gpu = True
with pytest.raises(ValueError):
await zocalo_results.stage()

0 comments on commit d52f517

Please sign in to comment.