diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml index 4117fb5f..6a154cdb 100644 --- a/.github/workflows/lint.yaml +++ b/.github/workflows/lint.yaml @@ -12,6 +12,6 @@ jobs: timeout-minutes: 5 steps: - uses: actions/checkout@v6 - - uses: astral-sh/ruff-action@v3 + - uses: astral-sh/ruff-action@v4.0.0 - run: ruff check - run: ruff format --check diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml new file mode 100644 index 00000000..e909ade3 --- /dev/null +++ b/.github/workflows/test.yaml @@ -0,0 +1,28 @@ +--- +name: Python tests + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + linux-test: + runs-on: ubuntu-latest + timeout-minutes: 30 + strategy: + matrix: + python-version: ["3.10", "3.11", "3.12", "3.13"] + steps: + - uses: actions/checkout@v6 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + pipx install uv + uv sync --extra server --extra client + - name: Full Python tests + run: uv run pytest -v -m "not container" diff --git a/pyproject.toml b/pyproject.toml index 89728777..b0b08ab8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -144,6 +144,9 @@ ignore_missing_imports = false [tool.pytest.ini_options] asyncio_default_fixture_loop_scope = "session" +markers = [ + "container: marks tests which need a container runtime", +] # [tool.uv.sources] # tiqi-zedboard = { path = "../tiqi-zedboard", editable = true } diff --git a/src/icon/config/config.py b/src/icon/config/config.py index afd9d682..6fe60a61 100644 --- a/src/icon/config/config.py +++ b/src/icon/config/config.py @@ -1,17 +1,20 @@ import logging +import os from pathlib import Path, PosixPath import yaml from confz import BaseConfig, FileSource -from icon.config import v1 -from icon.config import v2 as latest -from icon.config.config_path import get_config_path +from icon.config import latest, v1 from icon.config.migrations import migration_by_version +_ENV_KEY = "ICON_CONFIG" + logger = logging.getLogger("config") -VERSIONS: dict[int, BaseConfig] = {1: v1.ServiceConfigV1, 2: latest.ServiceConfig} +VERSIONS: dict[int, BaseConfig] = { + cfg.__version__: cfg.ServiceConfig for cfg in (v1, latest) +} # https://github.com/yaml/pyyaml/issues/617#issuecomment-1039273397 @@ -27,11 +30,11 @@ def get_config() -> latest.ServiceConfig: if not source.is_file(): source.parent.mkdir(parents=True, exist_ok=True) - with source.open("a") as file: - file.write(yaml.dump(latest.ServiceConfig().model_dump())) + with source.open("a") as f: + f.write(yaml.dump(latest.ServiceConfig().model_dump())) - with source.open("r") as file: - content = yaml.safe_load(file) + with source.open("r") as f: + content = yaml.safe_load(f) config_version = content.get("version", None) schema = VERSIONS.get(config_version) @@ -51,3 +54,32 @@ def get_config() -> latest.ServiceConfig: file.write(yaml.dump(config.model_dump())) return config + + +def save_config(config: latest.ServiceConfig) -> None: + """Save the configuration to the source YAML file. + + Serializes the updated configuration and writes it back to the file. + + Args: + config: + The validated configuration instance. + """ + with get_config_path().open("w") as f: + f.write(yaml.dump(config.model_dump())) + + +def _normalize(p: str | Path) -> Path: + return Path(p).expanduser().resolve() + + +def set_config_path(p: Path) -> None: + """Set once at startup; children inherit via environment.""" + os.environ[_ENV_KEY] = str(_normalize(p)) + + +def get_config_path() -> Path: + """Read from env, else default.""" + if env := os.environ.get(_ENV_KEY): + return _normalize(env) + return _normalize(Path.home() / ".config/icon/config.yaml") diff --git a/src/icon/config/config_path.py b/src/icon/config/config_path.py deleted file mode 100644 index f4b6e00f..00000000 --- a/src/icon/config/config_path.py +++ /dev/null @@ -1,22 +0,0 @@ -from __future__ import annotations - -import os -from pathlib import Path - -_ENV_KEY = "ICON_CONFIG" - - -def _normalize(p: str | Path) -> Path: - return Path(p).expanduser().resolve() - - -def set_config_path(p: Path) -> None: - """Set once at startup; children inherit via environment.""" - os.environ[_ENV_KEY] = str(_normalize(p)) - - -def get_config_path() -> Path: - """Read from env, else default.""" - if env := os.environ.get(_ENV_KEY): - return _normalize(env) - return _normalize(Path.home() / ".config/icon/config.yaml") diff --git a/src/icon/config/latest.py b/src/icon/config/latest.py new file mode 100644 index 00000000..0145f7b4 --- /dev/null +++ b/src/icon/config/latest.py @@ -0,0 +1,73 @@ +from pathlib import Path +from typing import Any + +from confz import BaseConfig +from pydantic import BaseModel + +__version__ = 2 + + +class HealthCheckConfig(BaseModel): + interval_seconds: float = 10.0 + + +class DataConfiguration(BaseModel): + results_dir: str = str(Path.cwd() / "output") + + +class ExperimentLibraryConfig(BaseModel): + module: str = "icon.server.data_access.pycrystal_experiment_library_client" + client_class: str = "AsyncPyCrystalClient" + client_args: dict[str, Any] = {} + update_interval: int = 30 + + +class InfluxDBv1Config(BaseModel): + host: str = "localhost" + port: int = 8087 + username: str = "admin" + password: str = "admin" # noqa: S105 + database: str = "testing" + measurement: str = "Experiment Parameters" + ssl: bool = False + verify_ssl: bool = False + headers: dict[str, str] = {} + + +class SQLiteConfig(BaseModel): + file: str = str(Path.cwd() / "icon.db") + + +class DatabaseConfig(BaseModel): + influxdbv1: InfluxDBv1Config = InfluxDBv1Config() + sqlite: SQLiteConfig = SQLiteConfig() + + +class DateConfig(BaseModel): + timezone: str = "Europe/Zurich" + + +class PreProcessingConfig(BaseModel): + workers: int = 2 + + +class ServerConfig(BaseModel): + port: int = 8004 + host: str = "0.0.0.0" + pre_processing: PreProcessingConfig = PreProcessingConfig() + + +class HardwareConfig(BaseModel): + host: str = "localhost" + port: int = 6007 + + +class ServiceConfig(BaseConfig): # type: ignore[misc] + version: int = __version__ + experiment_library: ExperimentLibraryConfig = ExperimentLibraryConfig() + databases: DatabaseConfig = DatabaseConfig() + date: DateConfig = DateConfig() + server: ServerConfig = ServerConfig() + hardware: HardwareConfig = HardwareConfig() + health_check: HealthCheckConfig = HealthCheckConfig() + data: DataConfiguration = DataConfiguration() diff --git a/src/icon/config/migrations.py b/src/icon/config/migrations.py index b0669b8d..189580ca 100644 --- a/src/icon/config/migrations.py +++ b/src/icon/config/migrations.py @@ -3,9 +3,8 @@ from collections.abc import Callable from typing import Any -from icon.config.v1 import ServiceConfigV1 -from icon.config.v2 import ExperimentLibraryConfig -from icon.config.v2 import ServiceConfig as ServiceConfigV2 +from icon.config import latest as v2 +from icon.config import v1 migration_by_version = {} @@ -19,10 +18,10 @@ def decorate(f: Callable[[Any], Any]) -> Callable[[Any], Any]: @migration(version=1) -def migrate_v1_to_v2(old_config: ServiceConfigV1) -> ServiceConfigV2: +def migrate_v1_to_v2(old_config: v1.ServiceConfig) -> v2.ServiceConfig: exp_lib = old_config.experiment_library - return ServiceConfigV2( - experiment_library=ExperimentLibraryConfig( + return v2.ServiceConfig( + experiment_library=v2.ExperimentLibraryConfig( update_interval=exp_lib.update_interval, client_args={ "checkout_path": exp_lib.dir, diff --git a/src/icon/config/reloader.py b/src/icon/config/reloader.py new file mode 100644 index 00000000..2bab60d2 --- /dev/null +++ b/src/icon/config/reloader.py @@ -0,0 +1,44 @@ +import logging +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +from icon.config.config import get_config + +if TYPE_CHECKING: + from icon.config.latest import ServiceConfig + +logger = logging.getLogger(__name__) + + +class ReloadError(Exception): + pass + + +class Reloader: + def __init__( + self, + obj_factory: Callable[..., Any], + fallback_obj: Any, + subconfig: "Callable[[ServiceConfig], dict[str, Any]]", + ) -> None: + self.current_config: dict[str, Any] = {} + self.obj_factory = staticmethod(obj_factory) + self.obj = fallback_obj + self.fallback_obj = fallback_obj + self.subconfig = staticmethod(subconfig) + + def reload(self) -> Any: + new_config = self.subconfig(get_config()) + if new_config == self.current_config: + return self.obj + try: + self.obj = self.obj_factory(**new_config) + except ReloadError as e: + logger.warning(format(e)) + self.obj = self.fallback_obj + + self.current_config = new_config + return self.obj + + def is_configured(self) -> bool: + return self.obj is not self.fallback_obj diff --git a/src/icon/config/v1.py b/src/icon/config/v1.py index a29d6f04..3d1e4974 100644 --- a/src/icon/config/v1.py +++ b/src/icon/config/v1.py @@ -1,68 +1,27 @@ -from pathlib import Path - from confz import BaseConfig from pydantic import BaseModel -__version__ = 1 - - -class HealthCheckConfig(BaseModel): - interval_seconds: float = 10.0 - +from icon.config.latest import ( + DatabaseConfig, + DataConfiguration, + DateConfig, + HardwareConfig, + HealthCheckConfig, + ServerConfig, +) -class DataConfiguration(BaseModel): - results_dir: str = str(Path.cwd() / "output") +__version__ = 1 -class ExperimentLibraryConfigV1(BaseModel): +class ExperimentLibraryConfig(BaseModel): dir: str | None = None git_repository: str = "https://..." update_interval: int = 30 -class InfluxDBv1Config(BaseModel): - host: str = "localhost" - port: int = 8087 - username: str = "admin" - password: str = "admin" # noqa: S105 - database: str = "testing" - measurement: str = "Experiment Parameters" - ssl: bool = False - verify_ssl: bool = False - headers: dict[str, str] = {} - - -class SQLiteConfig(BaseModel): - file: str = str(Path.cwd() / "icon.db") - - -class DatabaseConfig(BaseModel): - influxdbv1: InfluxDBv1Config = InfluxDBv1Config() - sqlite: SQLiteConfig = SQLiteConfig() - - -class DateConfig(BaseModel): - timezone: str = "Europe/Zurich" - - -class PreProcessingConfig(BaseModel): - workers: int = 2 - - -class ServerConfig(BaseModel): - port: int = 8004 - host: str = "0.0.0.0" - pre_processing: PreProcessingConfig = PreProcessingConfig() - - -class HardwareConfig(BaseModel): - host: str = "localhost" - port: int = 6007 - - -class ServiceConfigV1(BaseConfig): # type: ignore[misc] +class ServiceConfig(BaseConfig): # type: ignore[misc] version: int = __version__ - experiment_library: ExperimentLibraryConfigV1 = ExperimentLibraryConfigV1() + experiment_library: ExperimentLibraryConfig = ExperimentLibraryConfig() databases: DatabaseConfig = DatabaseConfig() date: DateConfig = DateConfig() server: ServerConfig = ServerConfig() diff --git a/src/icon/config/v2.py b/src/icon/config/v2.py deleted file mode 100644 index a3417f98..00000000 --- a/src/icon/config/v2.py +++ /dev/null @@ -1,33 +0,0 @@ -from typing import Any - -from confz import BaseConfig -from pydantic import BaseModel - -from icon.config.v1 import ( - DatabaseConfig, - DataConfiguration, - DateConfig, - HardwareConfig, - HealthCheckConfig, - ServerConfig, -) - -__version__ = 2 - - -class ExperimentLibraryConfig(BaseModel): - module: str = "icon.server.data_access.pycrystal_experiment_library_client" - client_class: str = "AsyncPyCrystalClient" - client_args: dict[str, Any] = {} - update_interval: int = 30 - - -class ServiceConfig(BaseConfig): # type: ignore[misc] - version: int = __version__ - experiment_library: ExperimentLibraryConfig = ExperimentLibraryConfig() - databases: DatabaseConfig = DatabaseConfig() - date: DateConfig = DateConfig() - server: ServerConfig = ServerConfig() - hardware: HardwareConfig = HardwareConfig() - health_check: HealthCheckConfig = HealthCheckConfig() - data: DataConfiguration = DataConfiguration() diff --git a/src/icon/server/__main__.py b/src/icon/server/__main__.py index 8aca699c..728f2f96 100644 --- a/src/icon/server/__main__.py +++ b/src/icon/server/__main__.py @@ -7,11 +7,12 @@ import click -from icon.config.config_path import set_config_path +from icon.config.config import set_config_path from icon.logging import setup_logging from icon.server.data_access.reconfigurable_experiment_library_client import ( ReconfigurableExperimentLibraryClient, ) +from icon.server.hardware_processing.zedboard_controller import ZedboardController from icon.server.shared_resource_manager import SRM if TYPE_CHECKING: @@ -91,6 +92,7 @@ def start_server() -> None: hardware_processing_queue=SRM.hardware_processing_queue, post_processing_queue=post_processing_queue, manager=SRM, + hardware_controller=ZedboardController(), ) hardware_processing_worker.start() @@ -103,6 +105,7 @@ def start_server() -> None: APIService( experiment_library_client=exp_lib_client, pre_processing_event_queues=pre_processing_update_queues, + hardware_controller=ZedboardController(connect=False), ), host=get_config().server.host, web_port=get_config().server.port, @@ -135,7 +138,10 @@ def main(*, version: bool, verbose: int, quiet: int, config: pathlib.Path) -> No setup_logging(level) set_config_path(config or pathlib.Path.home() / ".config/icon/config.yaml") - start_server() + try: + start_server() + except KeyboardInterrupt: + logging.getLogger(__name__).info("Exit") if __name__ == "__main__": diff --git a/src/icon/server/api/api_service.py b/src/icon/server/api/api_service.py index a36e724d..cdf1398f 100644 --- a/src/icon/server/api/api_service.py +++ b/src/icon/server/api/api_service.py @@ -26,9 +26,10 @@ if TYPE_CHECKING: import multiprocessing - from icon.server.data_access.experiment_library_client import ( - ExperimentLibraryClient, + from icon.server.data_access.reconfigurable_experiment_library_client import ( + ReconfigurableExperimentLibraryClient, ) + from icon.server.hardware_processing.hardware_controller import HardwareController from icon.server.utils.types import UpdateQueue logger = logging.getLogger(__name__) @@ -51,7 +52,8 @@ class APIService(pydase.DataService): def __init__( self, pre_processing_event_queues: list[multiprocessing.Queue[UpdateQueue]], - experiment_library_client: ExperimentLibraryClient, + experiment_library_client: ReconfigurableExperimentLibraryClient, + hardware_controller: HardwareController, ) -> None: """Create a new APIService. @@ -59,6 +61,7 @@ def __init__( pre_processing_event_queues: Queues used by `ScansController` to notify pre-processing workers. experiment_library_client: Client for an experiment library + hardware_controller: Controller for the hardware """ super().__init__() @@ -82,7 +85,7 @@ def __init__( ) """Controller for triggering update events for jobs across multiple worker processes.""" - self.status = StatusController() + self.status = StatusController(hardware_controller) """Controller for system status monitoring.""" self._experiment_library_client = experiment_library_client diff --git a/src/icon/server/api/configuration_controller.py b/src/icon/server/api/configuration_controller.py index a404a2d6..c5dd57e6 100644 --- a/src/icon/server/api/configuration_controller.py +++ b/src/icon/server/api/configuration_controller.py @@ -2,12 +2,10 @@ from typing import Any import pydase -import yaml from confz import DataSource -from icon.config.config import get_config -from icon.config.config_path import get_config_path -from icon.config.v2 import ServiceConfig +from icon.config.config import get_config, save_config +from icon.config.latest import ServiceConfig from icon.server.web_server.socketio_emit_queue import emit_queue logger = logging.getLogger(__name__) @@ -48,7 +46,7 @@ def update_config_option(self, key: str, value: Any) -> bool: updated_config = ServiceConfig(config_sources=DataSource(current_config)) # Save the updated configuration back to the file - self._save_configuration(updated_config) + save_config(updated_config) emit_queue.put( {"event": "config.update", "data": updated_config.model_dump()} ) @@ -57,28 +55,37 @@ def update_config_option(self, key: str, value: Any) -> bool: return False return True - def _save_configuration(self, new_config: ServiceConfig) -> None: - """Save the updated configuration to the source YAML file. - - Serializes the updated configuration and writes it back to the file. - - Args: - new_config: - The validated configuration instance. - """ - with get_config_path().open("w") as file: - file.write(yaml.dump(new_config.model_dump())) - def set_nested(config: dict[str, Any], nested_key: str, value: Any) -> None: """Set a value in a nested dict.""" - current = config - *fields, last_field = nested_key.split(".") + current: dict[str, Any] | list[Any] = config + *fields, last_field = parse_config_key(nested_key) # Traverse to the nested key for field in fields: - if field not in current: + if isinstance(current, dict) and ( + not isinstance(field, str) or field not in current + ): raise KeyError(f"Key {nested_key!r} not found in configuration.") - current = current[field] + if isinstance(current, list) and ( + not isinstance(field, int) or field >= len(current) + ): + raise IndexError( + f"Configuration error: Index out of range: {field} in {nested_key!r}" + ) + current = current[field] # type: ignore[index] # Update the value - current[last_field] = value + current[last_field] = value # type: ignore[index] + + +def parse_config_key(nested_key: str) -> list[str | int]: + components = nested_key.split(".") + + def split_index(key: str) -> tuple[str | int, ...]: + try: + key, index_str = key.removesuffix("]").split("[", 1) + return key, int(index_str) + except ValueError: + return (key,) + + return [c for group in components for c in split_index(group)] diff --git a/src/icon/server/api/status_controller.py b/src/icon/server/api/status_controller.py index bbac1312..5d48841c 100644 --- a/src/icon/server/api/status_controller.py +++ b/src/icon/server/api/status_controller.py @@ -16,9 +16,9 @@ class StatusController(pydase.DataService): via the Socket.IO queue. """ - def __init__(self) -> None: + def __init__(self, hardware_controller: HardwareController) -> None: super().__init__() - self.__hardware_controller = HardwareController(connect=False) + self.__hardware_controller = hardware_controller self._influxdb_available = False self._hardware_available = False diff --git a/src/icon/server/data_access/experiment_library_client.py b/src/icon/server/data_access/experiment_library_client.py index 1d4a24c4..9698858d 100644 --- a/src/icon/server/data_access/experiment_library_client.py +++ b/src/icon/server/data_access/experiment_library_client.py @@ -1,7 +1,7 @@ """Abstraction over experiment library clients.""" from contextlib import AbstractContextManager, nullcontext -from typing import TYPE_CHECKING, TypedDict +from typing import TYPE_CHECKING, Any, TypedDict if TYPE_CHECKING: from icon.server.api.models.experiment_dict import ( @@ -93,7 +93,7 @@ async def get_experiment_readout_metadata( """ raise NotImplementedError("Must be implemented by a subclass") - async def get_setup_hardware_description(self) -> dict[str, dict]: + async def get_setup_hardware_description(self) -> dict[str, dict[str, Any]]: """Fetch hardware description from experiment library. Returns: @@ -160,7 +160,7 @@ async def get_experiment_readout_metadata( "vector_channel_windows": [], } - async def get_setup_hardware_description(self) -> dict[str, dict]: + async def get_setup_hardware_description(self) -> dict[str, dict[str, Any]]: """Fetch hardware description from experiment library. Returns: diff --git a/src/icon/server/data_access/pycrystal_experiment_library_client.py b/src/icon/server/data_access/pycrystal_experiment_library_client.py index 2cae2754..98cd7145 100644 --- a/src/icon/server/data_access/pycrystal_experiment_library_client.py +++ b/src/icon/server/data_access/pycrystal_experiment_library_client.py @@ -169,7 +169,7 @@ def plot_window_metadata(data: Any) -> "PlotWindowMetadata": ], } - def get_setup_hardware_description(self) -> dict[str, dict]: + def get_setup_hardware_description(self) -> dict[str, dict[str, Any]]: """Fetch hardware description from experiment library. Returns: diff --git a/src/icon/server/data_access/reconfigurable_experiment_library_client.py b/src/icon/server/data_access/reconfigurable_experiment_library_client.py index 45482c81..02cf1c6e 100644 --- a/src/icon/server/data_access/reconfigurable_experiment_library_client.py +++ b/src/icon/server/data_access/reconfigurable_experiment_library_client.py @@ -5,7 +5,7 @@ from contextlib import AbstractContextManager from typing import TYPE_CHECKING, Any -from icon.config.config import get_config +from icon.config.reloader import Reloader, ReloadError from icon.server.data_access.experiment_library_client import ( ExperimentLibraryClient, FallbackExperimentLibraryClient, @@ -21,45 +21,46 @@ ReadoutMetadata, ) + logger = logging.getLogger(__name__) +def load_client( + module: str, client_class: str, client_args: dict[str, Any] +) -> ExperimentLibraryClient: + try: + exp_lib_client_module = importlib.import_module(module) + exp_lib_client_class = getattr(exp_lib_client_module, client_class) + logger.info("Using experiment library client %s.%s", module, client_class) + return exp_lib_client_class(**client_args) + except (ValueError, ImportError, AttributeError) as e: + raise ReloadError( + "Experiment library client is misconfigured.\n" + f"configured module: {module}\n" + f"configured class: {client_class}\n" + f"Error message: {e}\n" + "Please reconfigure!" + ) from None + + class ReconfigurableExperimentLibraryClient(ExperimentLibraryClient): """Wrapper reconfiguring an underlying client, whenever relevant config changes.""" def __init__(self) -> None: self.client = FallbackExperimentLibraryClient() - self.current_config: dict[str, Any] = {} - self.reload() - - def reload(self) -> None: - config = get_config().experiment_library - new_config = config.model_dump() - if new_config == self.current_config: - return - logger.info( - "Using experiment library client %s.%s", config.module, config.client_class + self.reloader = Reloader( + load_client, + fallback_obj=self.client, + subconfig=lambda config: { + key: val + for key, val in config.experiment_library.model_dump().items() + if key != "update_interval" + }, ) - self.current_config = new_config - try: - exp_lib_client_module = importlib.import_module(config.module) - exp_lib_client_class = getattr(exp_lib_client_module, config.client_class) - self.client = exp_lib_client_class(**config.client_args) - except (ValueError, ImportError, AttributeError) as e: - logger.warning( - "Experiment library client is misconfigured.\n" - "configured module: %s\n" - "configured class: %s\n" - "Error message: %s\n" - "Please reconfigure!", - config.module, - config.client_class, - e, - ) + self.client = self.reloader.reload() def is_configured(self) -> bool: - self.reload() - return not isinstance(self.client, FallbackExperimentLibraryClient) + return self.reloader.is_configured() def checkout_revision(self, revision: str | None) -> str | None: """Restore a state of the library defined by `revision`. @@ -68,7 +69,7 @@ def checkout_revision(self, revision: str | None) -> str | None: Should be implemented by experiment library clients based on a git repository. """ - self.reload() + self.client = self.reloader.reload() return self.client.checkout_revision(revision) def isolated(self) -> AbstractContextManager[ExperimentLibraryClient]: @@ -77,7 +78,7 @@ def isolated(self) -> AbstractContextManager[ExperimentLibraryClient]: By default isolation is not implemented and only a reference to the original library is returned. """ - self.reload() + self.client = self.reloader.reload() return self.client.isolated() async def load_metadata(self) -> "tuple[ExperimentDict, ParameterMetadataDict]": @@ -86,7 +87,7 @@ async def load_metadata(self) -> "tuple[ExperimentDict, ParameterMetadataDict]": To support hot-reloading of user data modules, this is a method and not static data. """ - self.reload() + self.client = self.reloader.reload() return await self.client.load_metadata() async def generate_json_sequence( @@ -108,7 +109,7 @@ async def generate_json_sequence( Returns: JSON string containing the generated sequence. """ - self.reload() + self.client = self.reloader.reload() return await self.client.generate_json_sequence( exp_module_name=exp_module_name, exp_instance_name=exp_instance_name, @@ -133,18 +134,18 @@ async def get_experiment_readout_metadata( Returns: Dictionary containing readout metadata for the experiment. """ - self.reload() + self.client = self.reloader.reload() return await self.client.get_experiment_readout_metadata( exp_module_name=exp_module_name, exp_instance_name=exp_instance_name, parameter_dict=parameter_dict, ) - async def get_setup_hardware_description(self) -> dict[str, dict]: + async def get_setup_hardware_description(self) -> dict[str, dict[str, Any]]: """Fetch hardware description from experiment library. Returns: Dictionary containing a description of the experiment setup. """ - self.reload() + self.client = self.reloader.reload() return await self.client.get_setup_hardware_description() diff --git a/src/icon/server/data_access/repositories/experiment_data_repository.py b/src/icon/server/data_access/repositories/experiment_data_repository.py index 51a77b26..91dd20c5 100644 --- a/src/icon/server/data_access/repositories/experiment_data_repository.py +++ b/src/icon/server/data_access/repositories/experiment_data_repository.py @@ -29,7 +29,8 @@ logger = logging.getLogger(__name__) -class ResultDict(TypedDict): +@dataclass +class ResultDict: """Scalar/vector/shot readouts for a single data point.""" result_channels: dict[str, float] @@ -40,6 +41,7 @@ class ResultDict(TypedDict): """Mapping from shot channel name to per-shot integers.""" +@dataclass class ExperimentDataPoint(ResultDict): """A single data point with its context.""" @@ -361,6 +363,9 @@ def update_metadata_by_job_id( local_parameter_timestamp: Optional timestamp for local parameters. parameters: Scan parameters. """ + if parameters is None: + parameters = [] + filename = get_filename_by_job_id(job_id) h5_path = Path(get_config().data.results_dir) / filename @@ -463,54 +468,54 @@ def write_experiment_data_by_job_id( write_scan_parameters_and_timestamp_to_dataset( h5file=h5file, - data_point_index=data_point["index"], - scan_params=data_point["scan_params"], - timestamp=data_point["timestamp"], + data_point_index=data_point.index, + scan_params=data_point.scan_params, + timestamp=data_point.timestamp, number_of_data_points=number_of_data_points, ) write_results_to_dataset( h5file=h5file, - data_point_index=data_point["index"], - result_channels=data_point["result_channels"], + data_point_index=data_point.index, + result_channels=data_point.result_channels, number_of_data_points=number_of_data_points, ) write_shot_channels_to_datasets( h5file=h5file, - data_point_index=data_point["index"], - shot_channels=data_point["shot_channels"], + data_point_index=data_point.index, + shot_channels=data_point.shot_channels, number_of_data_points=number_of_data_points, number_of_shots=number_of_shots, ) write_vector_channels_to_datasets( h5file=h5file, - data_point_index=data_point["index"], - vector_channels=data_point["vector_channels"], + data_point_index=data_point.index, + vector_channels=data_point.vector_channels, ) write_sequence_json_to_dataset( h5file=h5file, - data_point_index=data_point["index"], - sequence_json=data_point["sequence_json"], + data_point_index=data_point.index, + sequence_json=data_point.sequence_json, ) - if data_point["index"] >= number_of_data_points: - h5file.attrs["number_of_data_points"] = data_point["index"] + 1 + if data_point.index >= number_of_data_points: + h5file.attrs["number_of_data_points"] = data_point.index + 1 logger.debug("Appended data to %s", h5_path) emit_queue.put( { "event": f"experiment_{job_id}", - "data": data_point, + "data": asdict(data_point), } ) emit_queue.put( { "event": "last_experiment_sequence", - "data": data_point["sequence_json"], + "data": data_point.sequence_json, } ) diff --git a/src/icon/server/data_access/venv_experiment_library_client.py b/src/icon/server/data_access/venv_experiment_library_client.py index 5cb1b3f5..d2649f59 100644 --- a/src/icon/server/data_access/venv_experiment_library_client.py +++ b/src/icon/server/data_access/venv_experiment_library_client.py @@ -1,7 +1,7 @@ """Access isolated experiment libraries.""" import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from icon.server.data_access.experiment_library_client import ExperimentLibraryClient from icon.server.data_access.venv_exec import VirtualEnvironment @@ -72,7 +72,7 @@ def get_experiment_readout_metadata( """ raise NotImplementedError("Must be implemented by a subclass") - def get_setup_hardware_description(self) -> dict[str, dict]: + def get_setup_hardware_description(self) -> dict[str, dict[str, Any]]: """Fetch hardware description from experiment library. Returns: @@ -153,7 +153,7 @@ async def get_experiment_readout_metadata( logger=venv_logger, ) - async def get_setup_hardware_description(self) -> dict[str, dict]: + async def get_setup_hardware_description(self) -> dict[str, dict[str, Any]]: """Fetch hardware description from experiment library. Returns: diff --git a/src/icon/server/hardware_processing/hardware_controller.py b/src/icon/server/hardware_processing/hardware_controller.py index d972a57d..db34f188 100644 --- a/src/icon/server/hardware_processing/hardware_controller.py +++ b/src/icon/server/hardware_processing/hardware_controller.py @@ -1,71 +1,34 @@ import logging +from enum import Enum, auto +from typing import Any -try: - import tiqi_zedboard.zedboard # type: ignore - - HAS_TIQI_ZEDBOARD = True -except ImportError: - HAS_TIQI_ZEDBOARD = False - -from icon.config.config import get_config from icon.server.data_access.repositories.experiment_data_repository import ResultDict -from icon.server.utils.sockets import is_socket_closed logger = logging.getLogger(__name__) -class HardwareController: - def __init__(self, *, connect: bool = True) -> None: - self._host = get_config().hardware.host - self._port = get_config().hardware.port - self._zedboard: tiqi_zedboard.zedboard.Zedboard | None = None - if connect: - self.connect() +class StatusFlag(Enum): + SUCCESS = auto() + ERROR = auto() + UNKNOWN = auto() + +class HardwareController: def connect(self) -> None: - logger.info("Connecting to the Zedboard") - self._host = get_config().hardware.host - self._port = get_config().hardware.port - if HAS_TIQI_ZEDBOARD: - self._zedboard = tiqi_zedboard.zedboard.Zedboard( - hostname=self._host, port=self._port - ) - if not self.connected: - logger.warning("Failed to connect to the Zedboard") + raise NotImplementedError("Must be implemented by a derived class") @property def connected(self) -> bool: - return ( - self._zedboard is not None - and hasattr(self._zedboard, "_client") - and self._zedboard._client is not None - and not is_socket_closed(self._zedboard._client._socket) - ) - - def _update_zedboard_sequence(self, *, sequence: str) -> None: - if self._zedboard is not None: - self._zedboard.sequence_JSON_parser.Sequence_JSON = sequence # type: ignore + raise NotImplementedError("Must be implemented by a derived class") - def run(self, *, sequence: str) -> ResultDict: - if not self.connected: - self.connect() + def send(self, data: Any) -> None: + raise NotImplementedError("Must be implemented by a derived class") - if not self.connected: - if HAS_TIQI_ZEDBOARD: - raise RuntimeError("Could not connect to the Zedboard") - raise RuntimeError( - "Tiqi zedboard package is not available. " - "Please use 'uv sync --all-extras' to install all dependencies" - ) + def run(self) -> None: + raise NotImplementedError("Must be implemented by a derived class") - self._update_zedboard_sequence(sequence=sequence) - self._zedboard.sequence_JSON_parser.Parse_JSON_Header() # type: ignore - results: tiqi_zedboard.zedboard.Result = self._zedboard.sequence_JSON_parser() # type: ignore + def status(self) -> tuple[StatusFlag, str, Any]: + raise NotImplementedError("Must be implemented by a derived class") - return { - "result_channels": results.result_channels, - "vector_channels": results.vector_channels - if results.vector_channels is not None - else {}, - "shot_channels": results.shot_channels, - } + def receive(self) -> ResultDict: + raise NotImplementedError("Must be implemented by a derived class") diff --git a/src/icon/server/hardware_processing/worker.py b/src/icon/server/hardware_processing/worker.py index e4bf424c..c3c040e7 100644 --- a/src/icon/server/hardware_processing/worker.py +++ b/src/icon/server/hardware_processing/worker.py @@ -14,22 +14,23 @@ from icon.config.config import get_config from icon.server.data_access.models.enums import DeviceStatus, JobRunStatus from icon.server.data_access.repositories.device_repository import DeviceRepository +from icon.server.data_access.repositories.experiment_data_repository import ( + ExperimentDataPoint, +) from icon.server.data_access.repositories.job_run_repository import ( JobRunRepository, job_run_cancelled_or_failed, ) -from icon.server.hardware_processing.hardware_controller import HardwareController from icon.server.hardware_processing.utils import extract_hardware_error_message from icon.server.post_processing.task import PostProcessingTask +from icon.server.utils.handle_keyboard_interrupt import handle_keyboard_interrupt if TYPE_CHECKING: import queue from icon.server.data_access.db_context.influxdb_v1 import DatabaseValueType from icon.server.data_access.models.sqlite.device import Device - from icon.server.data_access.repositories.experiment_data_repository import ( - ExperimentDataPoint, - ) + from icon.server.hardware_processing.hardware_controller import HardwareController from icon.server.hardware_processing.task import HardwareProcessingTask from icon.server.shared_resource_manager import SharedResourceManager @@ -70,6 +71,7 @@ def __init__( hardware_processing_queue: queue.PriorityQueue[HardwareProcessingTask], post_processing_queue: multiprocessing.Queue[PostProcessingTask], manager: SharedResourceManager, + hardware_controller: HardwareController, ) -> None: super().__init__() self._queue = hardware_processing_queue @@ -77,7 +79,7 @@ def __init__( self._manager = manager self._pydase_clients: dict[str, pydase.Client] = {} - self._hardware_controller = HardwareController() + self._hardware_controller = hardware_controller def _update_pydase_service_parameter( self, device: Device, access_path: str, new_value: DatabaseValueType @@ -142,6 +144,7 @@ def _set_pydase_service_values( new_value=value, ) + @handle_keyboard_interrupt(logger) def run(self) -> None: self._pydase_clients = { device.name: pydase.Client( @@ -173,17 +176,19 @@ def run(self) -> None: self._set_pydase_service_values(scanned_params=task.scanned_params) timestamp = datetime.now(timezone) - result = self._hardware_controller.run(sequence=task.sequence_json) - - experiment_data_point: ExperimentDataPoint = { - "index": task.data_point_index, - "scan_params": task.scanned_params, - "result_channels": result["result_channels"], - "shot_channels": result["shot_channels"], - "vector_channels": result["vector_channels"], - "timestamp": timestamp.isoformat(), - "sequence_json": task.sequence_json, - } + self._hardware_controller.send(data=task.sequence_json.encode("utf-8")) + self._hardware_controller.run() + result = self._hardware_controller.receive() + + experiment_data_point = ExperimentDataPoint( + index=task.data_point_index, + scan_params=task.scanned_params, + result_channels=result.result_channels, + shot_channels=result.shot_channels, + vector_channels=result.vector_channels, + timestamp=timestamp.isoformat(), + sequence_json=task.sequence_json, + ) post_processing_task = PostProcessingTask( priority=task.priority, diff --git a/src/icon/server/hardware_processing/zedboard_controller.py b/src/icon/server/hardware_processing/zedboard_controller.py new file mode 100644 index 00000000..36572bde --- /dev/null +++ b/src/icon/server/hardware_processing/zedboard_controller.py @@ -0,0 +1,76 @@ +import logging +from typing import Any + +try: + import tiqi_zedboard.zedboard +except ImportError: + raise ImportError( + "Tiqi zedboard package is not available. Please enable the `zedboard` extra." + ) from None + + +from icon.config.config import get_config +from icon.server.data_access.repositories.experiment_data_repository import ResultDict +from icon.server.hardware_processing.hardware_controller import ( + HardwareController, + StatusFlag, +) +from icon.server.utils.sockets import is_socket_closed + +logger = logging.getLogger(__name__) + + +class ZedboardController(HardwareController): + def __init__(self, *, connect: bool = True) -> None: + self._host = get_config().hardware.host + self._port = get_config().hardware.port + self._zedboard: tiqi_zedboard.zedboard.Zedboard | None = None + if connect: + self.connect() + + def connect(self) -> None: + logger.info("Connecting to the Zedboard") + self._host = get_config().hardware.host + self._port = get_config().hardware.port + self._zedboard = tiqi_zedboard.zedboard.Zedboard( + hostname=self._host, port=self._port + ) + if not self.connected: + logger.warning("Failed to connect to the Zedboard") + + @property + def connected(self) -> bool: + return ( + self._zedboard is not None + and hasattr(self._zedboard, "_client") + and self._zedboard._client is not None + and not is_socket_closed(self._zedboard._client._socket) + ) + + def _update_zedboard_sequence(self, *, sequence: str) -> None: + if self._zedboard is not None: + self._zedboard.sequence_JSON_parser.Sequence_JSON = sequence # type: ignore + + def send(self, data: bytes) -> None: + if not self.connected: + self.connect() + if not self.connected: + raise RuntimeError("Could not connect to the Zedboard") + self._update_zedboard_sequence(sequence=data.decode()) + + def run(self) -> None: + self._zedboard.sequence_JSON_parser.Parse_JSON_Header() # type: ignore + + def receive(self) -> ResultDict: + results: tiqi_zedboard.zedboard.Result = self._zedboard.sequence_JSON_parser() # type: ignore + + return ResultDict( + result_channels=results.result_channels, + vector_channels=results.vector_channels + if results.vector_channels is not None + else {}, + shot_channels=results.shot_channels, + ) + + def status(self) -> tuple[StatusFlag, str, Any]: + return (StatusFlag.UNKNOWN, "", None) diff --git a/src/icon/server/post_processing/worker.py b/src/icon/server/post_processing/worker.py index 1e0a8f42..9538ce2e 100644 --- a/src/icon/server/post_processing/worker.py +++ b/src/icon/server/post_processing/worker.py @@ -10,6 +10,7 @@ from icon.server.data_access.repositories.job_run_repository import ( job_run_cancelled_or_failed, ) +from icon.server.utils.handle_keyboard_interrupt import handle_keyboard_interrupt if TYPE_CHECKING: from icon.server.post_processing.task import PostProcessingTask @@ -25,6 +26,7 @@ def __init__( super().__init__() self._post_processing_queue = post_processing_queue + @handle_keyboard_interrupt(logger) def run(self) -> None: logger.info("Pre-processing worker started") diff --git a/src/icon/server/pre_processing/worker.py b/src/icon/server/pre_processing/worker.py index f93a38c2..886f0655 100644 --- a/src/icon/server/pre_processing/worker.py +++ b/src/icon/server/pre_processing/worker.py @@ -35,6 +35,7 @@ ) from icon.server.fitting.auto_fit import try_auto_fit from icon.server.hardware_processing.task import HardwareProcessingTask +from icon.server.utils.handle_keyboard_interrupt import handle_keyboard_interrupt if TYPE_CHECKING: from collections.abc import Iterable, Iterator @@ -183,6 +184,7 @@ def __init__( ) self._experiment_library_client = experiment_library_client + @handle_keyboard_interrupt(logger) def run(self) -> None: with self._experiment_library_client.isolated() as isolated_lib_client: logger.debug( diff --git a/src/icon/server/scheduler/scheduler.py b/src/icon/server/scheduler/scheduler.py index cac28972..a8127a21 100644 --- a/src/icon/server/scheduler/scheduler.py +++ b/src/icon/server/scheduler/scheduler.py @@ -15,6 +15,7 @@ JobRunRepository, ) from icon.server.pre_processing.task import PreProcessingTask +from icon.server.utils.handle_keyboard_interrupt import handle_keyboard_interrupt logger = logging.getLogger(__name__) @@ -56,6 +57,7 @@ def __init__( self.kwargs = kwargs self._pre_processing_queue = pre_processing_queue + @handle_keyboard_interrupt(logger) def run(self) -> None: initialise_job_tables() while not should_exit(): diff --git a/src/icon/server/shared_resource_manager.py b/src/icon/server/shared_resource_manager.py index 1e0bb539..9774abdb 100644 --- a/src/icon/server/shared_resource_manager.py +++ b/src/icon/server/shared_resource_manager.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from icon.server.data_access.db_context.influxdb_v1 import DatabaseValueType + from icon.server.data_access.experiment_data import DatabaseValueType from icon.server.hardware_processing.task import HardwareProcessingTask from icon.server.pre_processing.task import PreProcessingTask diff --git a/src/icon/server/utils/handle_keyboard_interrupt.py b/src/icon/server/utils/handle_keyboard_interrupt.py new file mode 100644 index 00000000..c45d8f5e --- /dev/null +++ b/src/icon/server/utils/handle_keyboard_interrupt.py @@ -0,0 +1,18 @@ +import logging +from collections.abc import Callable +from typing import Any + + +def handle_keyboard_interrupt( + logger: logging.Logger, +) -> Callable[[Callable[..., None]], Callable[..., None]]: + def wrapper(f: Callable[..., None]) -> Callable[..., None]: + def handled_f(*args: Any, **kwargs: Any) -> None: + try: + f(*args, **kwargs) + except KeyboardInterrupt: + logger.info("Forced Exit") + + return handled_f + + return wrapper diff --git a/tests/__init__.py b/tests/__init__.py index 9fe046e2..ab7fcbf6 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,5 +1,5 @@ from pathlib import Path -from icon.config.config_path import set_config_path +from icon.config.config import set_config_path set_config_path(Path(__file__).parent / "config.yaml") diff --git a/tests/server/__init__.py b/tests/server/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/server/api/test_configuration_controller.py b/tests/server/api/test_configuration_controller.py new file mode 100644 index 00000000..fd9c3675 --- /dev/null +++ b/tests/server/api/test_configuration_controller.py @@ -0,0 +1,19 @@ +from icon.server.api.configuration_controller import ( + ConfigurationController, + parse_config_key, +) + + +def test_update_config_option() -> None: + controller = ConfigurationController() + original_port = controller.get_config()["server"]["port"] + target_port = original_port + 1 + controller.update_config_option("server.port", target_port) + updated_port = controller.get_config()["server"]["port"] + assert updated_port == target_port + + +def test_parse_config_key() -> None: + nested_key = "hardware.devices[1].id" + path = parse_config_key(nested_key) + assert path == ["hardware", "devices", 1, "id"] diff --git a/tests/server/data_access/db_context/test_influxdbv1.py b/tests/server/data_access/db_context/test_influxdbv1.py index 40ad8659..c3509e22 100644 --- a/tests/server/data_access/db_context/test_influxdbv1.py +++ b/tests/server/data_access/db_context/test_influxdbv1.py @@ -52,6 +52,7 @@ def check_port(port: int) -> bool: yield +@pytest.mark.container def test_InfluxDBv1Session(influxdbv1_service: None) -> None: # noqa: N802, ARG001 test_value = 1337