diff --git a/configs/config_example.toml b/configs/config_example.toml index a0682f6..4b057dc 100644 --- a/configs/config_example.toml +++ b/configs/config_example.toml @@ -7,6 +7,10 @@ router_name = "router1" router_address = "xx.xx.xx.xx" # Replace with actual IP address router_port = 5555 +# Address of the peer (follower) node used for two-player protocols (CHSH, QKD/SSM). +# The health check probes this address on startup and disables CHSH/SSM if unreachable. +follower_node_address = "xx.xx.xx.xx:8000" # Replace with actual follower node address + # Rotary encoder serial adress rotary_encoder_address = "/dev/tty.usbmodem1101" diff --git a/pyproject.toml b/pyproject.toml index 9929745..0b9b4b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ authors = [ { name = "Soroush Hoseini", email = "soroushhoseini0@gmail.com" }, ] dependencies = [ - "pqn-hardware @ git+https://github.com/PublicQuantumNetwork/pqn-hardware.git@6440be2830fddccb09b461caea023c84b7c3bea1", + "pqn-hardware @ git+https://github.com/PublicQuantumNetwork/pqn-hardware.git@master", "fastapi[standard]>=0.115.14", "httpx>=0.28.1", "pydantic>=2.0", diff --git a/src/pqn_node/api/deps.py b/src/pqn_node/api/deps.py index 3ef982e..fda86d1 100644 --- a/src/pqn_node/api/deps.py +++ b/src/pqn_node/api/deps.py @@ -24,7 +24,9 @@ async def get_http_client() -> AsyncGenerator[httpx.AsyncClient, None]: async def get_instrument_client() -> AsyncGenerator[Client, None]: - async with Client(host=settings.router_address, port=settings.router_port) as client: + async with Client( + host=settings.router_address, port=settings.router_port, router_name=settings.router_name + ) as client: yield client diff --git a/src/pqn_node/api/main.py b/src/pqn_node/api/main.py index 63c2383..60ad59e 100644 --- a/src/pqn_node/api/main.py +++ b/src/pqn_node/api/main.py @@ -1,14 +1,22 @@ from fastapi import APIRouter +from pydantic import BaseModel from pqn_node.api.routes import chsh from pqn_node.api.routes import coordination from pqn_node.api.routes import debug -from pqn_node.api.routes import games from pqn_node.api.routes import health from pqn_node.api.routes import qkd from pqn_node.api.routes import rng from pqn_node.api.routes import serial from pqn_node.api.routes import timetagger +from pqn_node.core.config import GamesAvailability +from pqn_node.core.config import get_settings +from pqn_node.core.config import settings + + +class NodeConfig(BaseModel): + follower_node_address: str | None + api_router = APIRouter() api_router.include_router(chsh.router) @@ -18,5 +26,14 @@ api_router.include_router(serial.router) api_router.include_router(coordination.router) api_router.include_router(debug.router) -api_router.include_router(games.router) api_router.include_router(health.router) + + +@api_router.get("/games/availability", tags=["games"]) +def get_availability() -> GamesAvailability: + return get_settings().games_availability + + +@api_router.get("/node/config", tags=["node"]) +def get_node_config() -> NodeConfig: + return NodeConfig(follower_node_address=settings.follower_node_address) diff --git a/src/pqn_node/api/routes/chsh.py b/src/pqn_node/api/routes/chsh.py index 3784d51..69ced58 100644 --- a/src/pqn_node/api/routes/chsh.py +++ b/src/pqn_node/api/routes/chsh.py @@ -77,11 +77,15 @@ async def event_generator() -> AsyncGenerator[str, None]: async def _chsh( # Complexity is high due to the nature of the CHSH experiment. basis: tuple[float, float], - follower_node_address: str, http_client: ClientDep, timetagger_address: str, state: StateDep, ) -> ChshResult: + follower_node_address = settings.follower_node_address + if follower_node_address is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="follower_node_address not configured" + ) logger.debug("Starting CHSH") # Initialize progress tracking @@ -91,7 +95,9 @@ async def _chsh( # Complexity is high due to the nature of the CHSH experiment. chsh_progress_event.set() logger.debug("Instantiating client") - client = Client(host=settings.router_address, port=settings.router_port, timeout=600_000) + client = Client( + host=settings.router_address, port=settings.router_port, router_name=settings.router_name, timeout=600_000 + ) # TODO: Check if settings.chsh_settings.hwp is set before even trying to get the device. hwp = cast("RotatorInstrument", client.get_device(settings.chsh_settings.hwp[0], settings.chsh_settings.hwp[1])) @@ -186,18 +192,19 @@ async def _chsh( # Complexity is high due to the nature of the CHSH experiment. @router.post("/") async def chsh( basis: tuple[float, float], - follower_node_address: str, http_client: ClientDep, timetagger_address: str, state: StateDep, ) -> ChshResult: logger.info("Starting CHSH experiment with basis: %s", basis) - return await _chsh(basis, follower_node_address, http_client, timetagger_address, state) + return await _chsh(basis, http_client, timetagger_address, state) @router.post("/request-angle-by-basis") async def request_angle_by_basis(index: int, state: StateDep, *, perp: bool = False) -> bool: - client = Client(host=settings.router_address, port=settings.router_port, timeout=600_000) + client = Client( + host=settings.router_address, port=settings.router_port, router_name=settings.router_name, timeout=600_000 + ) hwp = cast( "RotatorInstrument", client.get_device(settings.chsh_settings.request_hwp[0], settings.chsh_settings.request_hwp[1]), diff --git a/src/pqn_node/api/routes/coordination.py b/src/pqn_node/api/routes/coordination.py index c7b8367..0e60acf 100644 --- a/src/pqn_node/api/routes/coordination.py +++ b/src/pqn_node/api/routes/coordination.py @@ -113,9 +113,7 @@ async def protocol_cancelled( @router.post("/collect_follower") -async def collect_follower( - request: Request, address: str, state: StateDep, http_client: ClientDep -) -> CollectFollowerResponse: +async def collect_follower(request: Request, state: StateDep, http_client: ClientDep) -> CollectFollowerResponse: """ Endpoint called by a leader node (this one) to request a follower node (other node) to follow it. @@ -123,6 +121,11 @@ async def collect_follower( ------- CollectFollowerResponse indicating if the follower accepted the request. """ + address = settings.follower_node_address + if address is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="follower_node_address not configured" + ) logger.info("Requesting client at %s to follow", address) # Get the port this server is listening on diff --git a/src/pqn_node/api/routes/games.py b/src/pqn_node/api/routes/games.py deleted file mode 100644 index 23f103b..0000000 --- a/src/pqn_node/api/routes/games.py +++ /dev/null @@ -1,11 +0,0 @@ -from fastapi import APIRouter - -from pqn_node.core.config import GamesAvailability -from pqn_node.core.config import get_settings - -router = APIRouter(prefix="/games", tags=["games"]) - - -@router.get("/availability") -def get_availability() -> GamesAvailability: - return get_settings().games_availability diff --git a/src/pqn_node/api/routes/health.py b/src/pqn_node/api/routes/health.py index b623050..dcfdd6e 100644 --- a/src/pqn_node/api/routes/health.py +++ b/src/pqn_node/api/routes/health.py @@ -1,11 +1,11 @@ +import concurrent.futures import logging import time -from typing import Annotated +from collections.abc import Callable import httpx import serial from fastapi import APIRouter -from fastapi import Query from pqn_hardware.network.client import Client from pydantic import BaseModel from pydantic import Field @@ -17,7 +17,20 @@ router = APIRouter(prefix="/health", tags=["health"]) _ROUTER_TIMEOUT_MS = 5000 +_ROUTER_WALL_TIMEOUT_S = 6.0 + +_DEVICES_WALL_TIMEOUT_S = 6.0 + +_SERIAL_WALL_TIMEOUT_S = 3.0 + _FOLLOWER_TIMEOUT_S = 5.0 +_FOLLOWER_WALL_TIMEOUT_S = 6.0 + +_probe_executor = concurrent.futures.ThreadPoolExecutor(max_workers=4, thread_name_prefix="health-probe") + + +def _run_with_timeout[T](fn: Callable[[], T], timeout_s: float) -> T: + return _probe_executor.submit(fn).result(timeout=timeout_s) class ComponentStatus(BaseModel): @@ -158,15 +171,23 @@ def _probe_follower(follower_node_address: str) -> ComponentStatus: @router.get("/") -def health( - follower_node_address: Annotated[str | None, Query()] = None, -) -> HealthStatus: +def health() -> HealthStatus: """Probe router, configured devices, rotary encoder, and optional follower node.""" - router_status, client = _connect_router() + follower_node_address = settings.follower_node_address + try: + router_status, client = _run_with_timeout(_connect_router, _ROUTER_WALL_TIMEOUT_S) + except concurrent.futures.TimeoutError: + router_status = ComponentStatus(reachable=False, error="timeout") + client = None if client is not None: try: - devices = _probe_devices(client) + devices = _run_with_timeout(lambda: _probe_devices(client), _DEVICES_WALL_TIMEOUT_S) + except concurrent.futures.TimeoutError: + devices = [ + DeviceStatus(provider=provider, name=name, purpose=purpose, reachable=False, error="timeout") + for provider, name, purpose in _configured_devices() + ] finally: client.disconnect() else: @@ -175,8 +196,23 @@ def health( for provider, name, purpose in _configured_devices() ] - rotary_encoder = _probe_rotary_encoder() - follower_node = _probe_follower(follower_node_address) if follower_node_address else None + if not settings.virtual_rotator: + try: + rotary_encoder = _run_with_timeout(_probe_rotary_encoder, _SERIAL_WALL_TIMEOUT_S) + except concurrent.futures.TimeoutError: + rotary_encoder = ComponentStatus(reachable=False, error="timeout") + else: + rotary_encoder = _probe_rotary_encoder() + + if follower_node_address: + try: + follower_node = _run_with_timeout(lambda: _probe_follower(follower_node_address), _FOLLOWER_WALL_TIMEOUT_S) + except concurrent.futures.TimeoutError: + follower_node = ComponentStatus(reachable=False, error="timeout") + else: + follower_node = None + + _apply_games_override(router_status, follower_node) return HealthStatus( router=router_status, @@ -184,3 +220,14 @@ def health( rotary_encoder=rotary_encoder, follower_node=follower_node, ) + + +def _apply_games_override(router_status: ComponentStatus, follower_node: ComponentStatus | None) -> None: + ga = settings.games_availability + if not router_status.reachable: + ga.chsh = False + ga.qf = False + ga.ssm = False + elif follower_node is not None and not follower_node.reachable: + ga.chsh = False + ga.ssm = False diff --git a/src/pqn_node/api/routes/qkd.py b/src/pqn_node/api/routes/qkd.py index 858c3f8..f60b68d 100644 --- a/src/pqn_node/api/routes/qkd.py +++ b/src/pqn_node/api/routes/qkd.py @@ -38,13 +38,21 @@ class QKDResult(BaseModel): async def _qkd( - follower_node_address: str, http_client: ClientDep, state: StateDep, timetagger_address: str | None = None, + follower_node_address: str | None = None, ) -> list[int]: + if follower_node_address is None: + follower_node_address = settings.follower_node_address + if follower_node_address is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="follower_node_address not configured" + ) logger.debug("Starting QKD") - client = Client(host=settings.router_address, port=settings.router_port, timeout=600_000) + client = Client( + host=settings.router_address, port=settings.router_port, router_name=settings.router_name, timeout=600_000 + ) hwp = cast("RotatorInstrument", client.get_device(settings.qkd_settings.hwp[0], settings.qkd_settings.hwp[1])) if hwp is None: @@ -126,7 +134,6 @@ def get_outcome(state: int, basis: int, choice: int, counts: int) -> int: @router.post("") async def qkd( - follower_node_address: str, http_client: ClientDep, state: StateDep, timetagger_address: str | None = None, @@ -139,12 +146,14 @@ async def qkd( detail="QKD basis list is empty", ) - return await _qkd(follower_node_address, http_client, state, timetagger_address) + return await _qkd(http_client, state, timetagger_address) @router.post("/single_bit") async def request_qkd_single_pass(state: StateDep) -> bool: - client = Client(host=settings.router_address, port=settings.router_port, timeout=600_000) + client = Client( + host=settings.router_address, port=settings.router_port, router_name=settings.router_name, timeout=600_000 + ) hwp = cast( "RotatorInstrument", client.get_device(settings.qkd_settings.request_hwp[0], settings.qkd_settings.request_hwp[1]), @@ -345,7 +354,7 @@ async def _submit_basis_list_leader( state.qkd_leader_basis_list = basis_list await _wait_for_follower_ready(state, http_client) - ret = await _qkd(state.followers_address, http_client, state, timetagger_address) + ret = await _qkd(http_client, state, timetagger_address, follower_node_address=state.followers_address) logger.info("Final QKD bits: %s", str(ret)) # Assemble QKDResult object diff --git a/src/pqn_node/api/routes/timetagger.py b/src/pqn_node/api/routes/timetagger.py index 22eb740..845cd24 100644 --- a/src/pqn_node/api/routes/timetagger.py +++ b/src/pqn_node/api/routes/timetagger.py @@ -40,7 +40,9 @@ async def measure_correlation( channel1=channel1, channel2=channel2, ) - client = Client(host=settings.router_address, port=settings.router_port, timeout=600_000) + client = Client( + host=settings.router_address, port=settings.router_port, router_name=settings.router_name, timeout=600_000 + ) tagger = cast("TimeTaggerInstrument", client.get_device(settings.timetagger[0], settings.timetagger[1])) if tagger is None: logger.error("Could not find time tagger device") @@ -73,7 +75,9 @@ async def count_singles( detail="No timetagger configured", ) - client = Client(host=settings.router_address, port=settings.router_port, timeout=600_000) + client = Client( + host=settings.router_address, port=settings.router_port, router_name=settings.router_name, timeout=600_000 + ) tagger = cast("TimeTaggerInstrument", client.get_device(settings.timetagger[0], settings.timetagger[1])) if tagger is None: logger.error("Could not find time tagger device") diff --git a/src/pqn_node/core/config.py b/src/pqn_node/core/config.py index 46642b5..1def27c 100644 --- a/src/pqn_node/core/config.py +++ b/src/pqn_node/core/config.py @@ -70,6 +70,7 @@ class Settings(BaseSettings): rotary_encoder_address: str = "/dev/ttyACM0" virtual_rotator: bool = False # If True, use terminal input instead of hardware rotary encoder games_availability: GamesAvailability = Field(default_factory=GamesAvailability) + follower_node_address: str | None = None model_config = SettingsConfigDict( toml_file="./config.toml", diff --git a/src/pqn_node/main.py b/src/pqn_node/main.py index ca19942..1998c02 100644 --- a/src/pqn_node/main.py +++ b/src/pqn_node/main.py @@ -1,15 +1,27 @@ import logging +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from pqn_node.api.main import api_router +from pqn_node.api.routes.health import health logging.basicConfig(level=logging.DEBUG) logger = logging.getLogger(__name__) + +@asynccontextmanager +async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]: + logger.info("Running startup health check") + health() + yield + + app = FastAPI( title="Public Quantum Network", + lifespan=lifespan, ) # Add CORS middleware to allow all origins diff --git a/uv.lock b/uv.lock index c012ee6..297f598 100644 --- a/uv.lock +++ b/uv.lock @@ -727,7 +727,7 @@ wheels = [ [[package]] name = "pqn-hardware" version = "0.2.0" -source = { git = "https://github.com/PublicQuantumNetwork/pqn-hardware.git?rev=6440be2830fddccb09b461caea023c84b7c3bea1#6440be2830fddccb09b461caea023c84b7c3bea1" } +source = { git = "https://github.com/PublicQuantumNetwork/pqn-hardware.git?rev=master#d9f46be7df481d11da7acc9a02e9e7eed914f6c6" } dependencies = [ { name = "numpy" }, { name = "pydantic" }, @@ -765,7 +765,7 @@ dev = [ requires-dist = [ { name = "fastapi", extras = ["standard"], specifier = ">=0.115.14" }, { name = "httpx", specifier = ">=0.28.1" }, - { name = "pqn-hardware", git = "https://github.com/PublicQuantumNetwork/pqn-hardware.git?rev=6440be2830fddccb09b461caea023c84b7c3bea1" }, + { name = "pqn-hardware", git = "https://github.com/PublicQuantumNetwork/pqn-hardware.git?rev=master" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pydantic-settings", specifier = ">=2.10.1" }, { name = "pyserial", specifier = ">=3.5" },