From 931144d48f53f076eb9c42fe496c82230e7d4c57 Mon Sep 17 00:00:00 2001 From: Hongquan Li Date: Sun, 20 Sep 2026 14:22:48 -0700 Subject: [PATCH] feat(server): pre-flight disk check for the TCP run commands The "Not Enough Disk Space" check lived only in the two GUI acquisition widgets, so an acquisition started over TCP/MCP (run_acquisition, run_acquisition_from_yaml) was never checked: one that did not fit started anyway and filled the disk. Both commands now run the same check, after the controller is fully configured and before anything is created on disk. A run the save disk cannot hold is refused with an error that names the opt-in ('large_acquisition_mode: true' under 'acquisition:' in the YAML, or the Large Acquisitions setting); with the mode on it starts and the disk-space guard pauses it while data is offloaded. Runs that save nothing (skip_saving) are not checked, as in the GUI. The decision is shared: control.core.disk_space.preflight_disk_check (Qt-free; measures the nearest existing parent when the save folder does not exist yet) now backs the GUI dialog too, so both paths use the same estimate, safety factor and message. Co-Authored-By: Claude Fable 5.1 --- software/control/core/disk_space.py | 46 +++++++++++ software/control/microscope_control_server.py | 47 ++++++++++- software/control/widgets.py | 15 ++-- software/docs/mcp_integration.md | 5 ++ .../tests/control/core/test_disk_space.py | 47 +++++++++++ .../control/test_microscope_control_server.py | 82 +++++++++++++++++++ .../test_preflight_disk_check_controller.py | 33 ++++++++ 7 files changed, 265 insertions(+), 10 deletions(-) create mode 100644 software/tests/control/test_preflight_disk_check_controller.py diff --git a/software/control/core/disk_space.py b/software/control/core/disk_space.py index 02f513f29..fdddb4ebc 100644 --- a/software/control/core/disk_space.py +++ b/software/control/core/disk_space.py @@ -199,3 +199,49 @@ def _directory_usage(self) -> int: if not self._directory.is_dir(): return 0 return control.utils.get_directory_disk_usage(self._directory) + + +PREFLIGHT_SAFETY_FACTOR = 1.03 + + +@dataclass(frozen=True) +class PreflightDiskCheck: + """Whether the configured acquisition fits on the save disk, as checked before it starts.""" + + required_bytes: float + available_bytes: int + image_count: int + save_directory: str + + @property + def fits(self) -> bool: + return self.required_bytes <= self.available_bytes + + def describe(self) -> str: + megabytes_required = int(self.required_bytes / 1024 / 1024) + megabytes_available = int(self.available_bytes / 1024 / 1024) + return ( + f"This acquisition will capture {self.image_count:,} images, which will" + f" require {megabytes_required:,} [MB], but '{self.save_directory}' only has" + f" {megabytes_available:,} [MB] available." + ) + + +def preflight_disk_check( + multi_point_controller, save_directory: str, safety_factor: float = PREFLIGHT_SAFETY_FACTOR +) -> PreflightDiskCheck: + """Compare the fully configured controller's storage estimate with the free space where it will save. + + Shared by the GUI's "Not Enough Disk Space" dialog and the TCP run commands so both decide alike. + ``save_directory`` may not exist yet (the TCP commands check before creating anything); free space is + then measured on its nearest existing parent, which is the disk it will be created on. + """ + existing = pathlib.Path(save_directory) + while not existing.exists() and existing != existing.parent: + existing = existing.parent + return PreflightDiskCheck( + required_bytes=safety_factor * multi_point_controller.get_estimated_acquisition_disk_storage(), + available_bytes=control.utils.get_available_disk_space(existing), + image_count=multi_point_controller.get_acquisition_image_count(), + save_directory=str(save_directory), + ) diff --git a/software/control/microscope_control_server.py b/software/control/microscope_control_server.py index 3f504a8d2..052fbb164 100644 --- a/software/control/microscope_control_server.py +++ b/software/control/microscope_control_server.py @@ -21,6 +21,7 @@ from control.core.acquisition_settings import apply_acquisition_settings, parse_wells import control._def # Module import for runtime access to MCP-modifiable settings +from control.core.disk_space import preflight_disk_check # Qt imports for thread-safe GUI operations try: @@ -693,7 +694,11 @@ def _cmd_run_acquisition( le=99, ), ) -> Dict[str, Any]: - """Run a multi-point acquisition across wells using the MultiPointController.""" + """Run a multi-point acquisition across wells using the MultiPointController. + + An acquisition the save disk cannot hold is refused with an error before anything is created, + unless the Large Acquisitions setting is on. + """ import os import control._def @@ -781,6 +786,10 @@ def _cmd_run_acquisition( # Set the base path and start new experiment self.multipoint_controller.set_base_path(base_path) + self._refuse_if_acquisition_will_not_fit( + base_path, + opt_in_hint="enable Settings > Acquisition > Large Acquisitions (large_acquisition_mode = True in the INI)", + ) self.multipoint_controller.start_new_experiment(experiment_id) # Calculate total FOVs for status reporting @@ -1066,6 +1075,11 @@ def _cmd_run_acquisition_from_yaml( and region coordinates - wellplate or flexible), updates the GUI to reflect these settings, and starts the acquisition. Regions and every controller setting come from the file through control.core.acquisition_settings (shared with the fluidics protocol runner). + + An acquisition the save disk cannot hold is refused with an error before anything is created, + unless large acquisition mode is on for it ('large_acquisition_mode: true' under 'acquisition:' + in the YAML, or the Large Acquisitions setting); it then pauses when space runs low and resumes + as finished data is offloaded. """ import os @@ -1150,6 +1164,13 @@ def _cmd_run_acquisition_from_yaml( # Set the base path and start new experiment self.multipoint_controller.set_base_path(base_path) + self._refuse_if_acquisition_will_not_fit( + base_path, + opt_in_hint=( + "add 'large_acquisition_mode: true' under 'acquisition:' in the YAML, or enable " + "Settings > Acquisition > Large Acquisitions" + ), + ) self.multipoint_controller.start_new_experiment(experiment_id) # Calculate total FOVs for status reporting @@ -1187,6 +1208,30 @@ def _cmd_run_acquisition_from_yaml( self._log.error(traceback.format_exc()) raise RuntimeError(f"Failed to start acquisition: {str(e)}") from e + def _refuse_if_acquisition_will_not_fit(self, base_path: str, opt_in_hint: str) -> None: + """The GUI's "Not Enough Disk Space" check for TCP clients, who never see that dialog. + + Call with the controller fully configured and before anything is created on disk. A run the save + disk cannot hold is refused unless large acquisition mode is on for it (then the disk-space guard + pauses it while data is offloaded). Runs that save nothing are not checked. + """ + controller = self.multipoint_controller + if controller.skip_saving: + return + check = preflight_disk_check(controller, base_path) + if check.fits: + return + if controller.large_acquisition_mode or control._def.LARGE_ACQUISITION_MODE: + self._log.warning( + f"{check.describe()} Large acquisition mode is on: the run will pause when the disk runs low " + "and resume as data is offloaded." + ) + return + raise RuntimeError( + f"{check.describe()} Free up disk space or reduce the acquisition. To run an acquisition larger " + f"than the disk (it pauses when space runs low while you offload finished data), {opt_in_hint}." + ) + @schema_method def _cmd_set_performance_mode( self, diff --git a/software/control/widgets.py b/software/control/widgets.py index 4e197b504..ad17df883 100644 --- a/software/control/widgets.py +++ b/software/control/widgets.py @@ -21,6 +21,7 @@ from control.core.multi_point_controller import MultiPointController from control.core.mosaic_utils import format_well_id from control.core.geometry_utils import get_effective_well_size, calculate_well_coverage +from control.core.disk_space import preflight_disk_check from control.microcontroller import Microcontroller from control.piezo import PiezoStage from control.channel_sequence import enable_channel_sequence @@ -126,9 +127,10 @@ def check_space_available_with_error_dialog( # To check how much disk space is required, we need to have the MultiPointController all configured. That is # a precondition of this function. save_directory = multi_point_controller.base_path - available_disk_space = utils.get_available_disk_space(save_directory) - space_required = factor_of_safecty * multi_point_controller.get_estimated_acquisition_disk_storage() - image_count = multi_point_controller.get_acquisition_image_count() + check = preflight_disk_check(multi_point_controller, save_directory, safety_factor=factor_of_safecty) + available_disk_space = check.available_bytes + space_required = check.required_bytes + image_count = check.image_count logger.info( f"Checking space available: {space_required=}, {available_disk_space=}, {image_count=}, {save_directory=}" @@ -137,12 +139,7 @@ def check_space_available_with_error_dialog( if not options.needs_dialog: return True - megabytes_required = int(space_required / 1024 / 1024) - megabytes_available = int(available_disk_space / 1024 / 1024) - error_message = ( - f"This acquisition will capture {image_count:,} images, which will" - f" require {megabytes_required:,} [MB], but '{save_directory}' only has {megabytes_available:,} [MB] available." - ) + error_message = check.describe() # Module-level lookup (not a default argument) so tests can monkeypatch the dialog. action = preflight_disk_dialog(f"{error_message}\n\n{LARGE_ACQUISITION_MODE_EXPLANATION}", options) if action == "enable": diff --git a/software/docs/mcp_integration.md b/software/docs/mcp_integration.md index 0b69373bf..22315ffa6 100644 --- a/software/docs/mcp_integration.md +++ b/software/docs/mcp_integration.md @@ -152,6 +152,11 @@ The `python_exec` command is disabled by default for security. To enable it: | `get_acquisition_status` | - | Check acquisition progress | | `abort_acquisition` | - | Stop running acquisition | +> Both run commands check the save disk first, like the GUI's "Not Enough Disk Space" dialog: an acquisition that +> will not fit is refused with an error unless large acquisition mode is on for it (`large_acquisition_mode: true` +> under `acquisition:` in the YAML, or Settings > Acquisition > Large Acquisitions). With the mode on, the run +> pauses when space runs low and resumes as finished data is offloaded. +> > For scripted automation, see [Automation](automation.md). ### Performance diff --git a/software/tests/control/core/test_disk_space.py b/software/tests/control/core/test_disk_space.py index 9d14660fa..1b58c5556 100644 --- a/software/tests/control/core/test_disk_space.py +++ b/software/tests/control/core/test_disk_space.py @@ -6,6 +6,7 @@ import pytest import control._def +import control.utils from control.core.disk_space import DiskSpaceGuard, DiskStatus from control.core.pause_gate import PauseGate @@ -373,3 +374,49 @@ def test_simulated_capacity_poll_interval_defaults_to_setting(tmp_path, monkeypa assert guard.free_bytes() == 100_000 _write_file(tmp_path / "img.bin", 40_000) assert guard.free_bytes() == 60_000 + + +# --- pre-flight check shared by the GUI dialog and the TCP run commands ---------------------------- + + +class _EstimateOnlyController: + def __init__(self, estimate_bytes, image_count=12): + self._estimate = estimate_bytes + self._count = image_count + + def get_estimated_acquisition_disk_storage(self): + return self._estimate + + def get_acquisition_image_count(self): + return self._count + + +def test_preflight_check_applies_the_safety_factor_and_reports_fit(tmp_path, monkeypatch): + from control.core.disk_space import preflight_disk_check + + monkeypatch.setattr(control.utils, "get_available_disk_space", lambda d: 1030) + fits = preflight_disk_check(_EstimateOnlyController(1000), str(tmp_path)) + assert fits.fits and fits.required_bytes == 1030 and fits.available_bytes == 1030 and fits.image_count == 12 + + too_big = preflight_disk_check(_EstimateOnlyController(1001), str(tmp_path)) + assert not too_big.fits + + +def test_preflight_check_measures_the_nearest_existing_folder(tmp_path, monkeypatch): + from control.core.disk_space import preflight_disk_check + + asked = [] + monkeypatch.setattr(control.utils, "get_available_disk_space", lambda d: asked.append(str(d)) or 10**9) + # The TCP commands check before the experiment folder (or even the base path) has been created. + check = preflight_disk_check(_EstimateOnlyController(1), str(tmp_path / "not" / "created" / "yet")) + assert asked == [str(tmp_path)] + assert check.save_directory == str(tmp_path / "not" / "created" / "yet") + + +def test_preflight_check_describes_the_shortfall(tmp_path, monkeypatch): + from control.core.disk_space import preflight_disk_check + + monkeypatch.setattr(control.utils, "get_available_disk_space", lambda d: 5 * 1024 * 1024) + check = preflight_disk_check(_EstimateOnlyController(2000 * 1024 * 1024, image_count=1234), str(tmp_path)) + text = check.describe() + assert "1,234 images" in text and "2,060 [MB]" in text and "5 [MB] available" in text and str(tmp_path) in text diff --git a/software/tests/control/test_microscope_control_server.py b/software/tests/control/test_microscope_control_server.py index 12f841ee8..9331c6d7b 100644 --- a/software/tests/control/test_microscope_control_server.py +++ b/software/tests/control/test_microscope_control_server.py @@ -35,6 +35,15 @@ def create_mock_server(objective: str = "20x", channels: list = None): mock_multipoint = MagicMock() mock_multipoint.acquisition_in_progress.return_value = False mock_multipoint.experiment_ID = "test_experiment" + # Pre-flight disk check inputs: a tiny acquisition that saves, with large acquisition mode off. + mock_multipoint.get_estimated_acquisition_disk_storage.return_value = 1000 + mock_multipoint.get_acquisition_image_count.return_value = 12 + mock_multipoint.skip_saving = False + mock_multipoint.large_acquisition_mode = False + mock_multipoint.set_large_acquisition_mode.side_effect = lambda on: setattr( + mock_multipoint, "large_acquisition_mode", on + ) + mock_multipoint.set_skip_saving.side_effect = lambda on: setattr(mock_multipoint, "skip_saving", on) # What set_selected_configurations would have resolved for these channels mock_multipoint.selected_configurations = mock_channels @@ -641,3 +650,76 @@ def test_schemas_document_the_commands(self, real_server): assert schemas[name]["parameters"] == {} assert schemas[name]["required"] == [] assert len(schemas[name]["description"]) > 20 + + +class TestTcpDiskPreflight: + """Both TCP run commands refuse an acquisition the save disk cannot hold unless large acquisition mode is + on, mirroring the GUI's "Not Enough Disk Space" check (which a TCP client never sees).""" + + @pytest.fixture + def yaml_file(self, tmp_path): + path = tmp_path / "acq.yaml" + path.write_text(SAMPLE_WELLPLATE_YAML) + return str(path) + + @pytest.fixture + def opt_in_yaml_file(self, tmp_path): + path = tmp_path / "acq_large.yaml" + path.write_text( + SAMPLE_WELLPLATE_YAML.replace("acquisition:\n", "acquisition:\n large_acquisition_mode: true\n", 1) + ) + assert "large_acquisition_mode: true" in path.read_text() + return str(path) + + @pytest.fixture + def tiny_disk(self, monkeypatch): + import control._def + import control.utils + + monkeypatch.setattr(control._def, "LARGE_ACQUISITION_MODE", False) + monkeypatch.setattr(control.utils, "get_available_disk_space", lambda d: 500) # the mock run needs 1030 + + def test_yaml_run_that_does_not_fit_is_refused_before_anything_is_created(self, tiny_disk, yaml_file, tmp_path): + server = create_mock_server() + with pytest.raises(RuntimeError) as err: + server._cmd_run_acquisition_from_yaml(yaml_path=yaml_file, base_path=str(tmp_path)) + message = str(err.value) + assert "large_acquisition_mode: true" in message, "the error must name the opt-in" + assert "12 images" in message and str(tmp_path) in message + server.multipoint_controller.start_new_experiment.assert_not_called() + server.multipoint_controller.run_acquisition.assert_not_called() + + def test_yaml_opt_in_lets_the_run_start(self, tiny_disk, opt_in_yaml_file, tmp_path): + server = create_mock_server() + result = server._cmd_run_acquisition_from_yaml(yaml_path=opt_in_yaml_file, base_path=str(tmp_path)) + assert result["started"] is True + server.multipoint_controller.run_acquisition.assert_called_once() + + def test_global_setting_lets_the_run_start(self, tiny_disk, yaml_file, tmp_path, monkeypatch): + import control._def + + monkeypatch.setattr(control._def, "LARGE_ACQUISITION_MODE", True) + server = create_mock_server() + assert server._cmd_run_acquisition_from_yaml(yaml_path=yaml_file, base_path=str(tmp_path))["started"] is True + + def test_runs_that_save_nothing_are_not_checked(self, tiny_disk, tmp_path): + path = tmp_path / "acq_skip.yaml" + path.write_text(SAMPLE_WELLPLATE_YAML.replace("acquisition:\n", "acquisition:\n skip_saving: true\n", 1)) + server = create_mock_server() + assert server._cmd_run_acquisition_from_yaml(yaml_path=str(path), base_path=str(tmp_path))["started"] is True + + def test_a_run_that_fits_starts_as_before(self, yaml_file, tmp_path, monkeypatch): + import control.utils + + monkeypatch.setattr(control.utils, "get_available_disk_space", lambda d: 10**9) + server = create_mock_server() + assert server._cmd_run_acquisition_from_yaml(yaml_path=yaml_file, base_path=str(tmp_path))["started"] is True + + def test_the_wells_command_is_refused_too_and_names_the_setting(self, tiny_disk, tmp_path): + server = create_mock_server() + # This command validates channels against the live controller rather than the config repo. + server.microscope.live_controller.get_channels.return_value = [create_mock_channel("BF LED matrix full")] + with pytest.raises(RuntimeError) as err: + server._cmd_run_acquisition(wells="B6", channels=["BF LED matrix full"], base_path=str(tmp_path)) + assert "Large Acquisitions" in str(err.value) + server.multipoint_controller.run_acquisition.assert_not_called() diff --git a/software/tests/control/test_preflight_disk_check_controller.py b/software/tests/control/test_preflight_disk_check_controller.py new file mode 100644 index 000000000..fcf4aa0a6 --- /dev/null +++ b/software/tests/control/test_preflight_disk_check_controller.py @@ -0,0 +1,33 @@ +"""The Qt-free pre-flight disk check (control.core.disk_space.preflight_disk_check) against a real simulated +controller, as the TCP run commands use it: before the experiment folder exists.""" + + +def test_shared_check_works_on_a_real_controller_before_the_experiment_folder_exists(tmp_path, monkeypatch): + """The TCP run commands check before start_new_experiment(), so the estimate must not need the folder.""" + import control.microscope + import control.utils + import tests.control.test_stubs as ts + from control.core.disk_space import PREFLIGHT_SAFETY_FACTOR, preflight_disk_check + from tests.control.test_MultiPointController import select_some_configs + + scope = control.microscope.Microscope.build_from_global_config(True) + mpc = ts.get_test_multi_point_controller(microscope=scope) + cfg = mpc.stage.get_config() + mpc.scanCoordinates.add_single_fov_region( + "region_1", + center_x=cfg.X_AXIS.MIN_POSITION + 1.0, + center_y=cfg.Y_AXIS.MIN_POSITION + 1.0, + center_z=cfg.Z_AXIS.MIN_POSITION + 1.0, + ) + select_some_configs(mpc, scope.objective_store.current_objective) + mpc.set_Nt(3) + save_dir = tmp_path / "base" / "not_created_yet" + mpc.set_base_path(str(save_dir)) + + monkeypatch.setattr(control.utils, "get_available_disk_space", lambda d: 1) + check = preflight_disk_check(mpc, str(save_dir)) + + assert check.image_count == mpc.get_acquisition_image_count() > 0 + assert check.required_bytes == PREFLIGHT_SAFETY_FACTOR * mpc.get_estimated_acquisition_disk_storage() > 1 + assert not check.fits + assert not save_dir.exists(), "checking must not create anything"