diff --git a/src/wf_config/__init__.py b/src/wf_config/__init__.py index 720af7ad..87306291 100644 --- a/src/wf_config/__init__.py +++ b/src/wf_config/__init__.py @@ -12,6 +12,7 @@ from .models import ( PythonSourceConfig, RpcHttpTargetConfig, RpcHttpTransportConfig, + SchedulerConfig, ServerConfig, ServerStoresConfig, SourceConfigOwnership, @@ -32,6 +33,7 @@ __all__ = [ "PythonSourceConfig", "RpcHttpTargetConfig", "RpcHttpTransportConfig", + "SchedulerConfig", "ServerConfig", "ServerStoresConfig", "SourceConfigOwnership", diff --git a/src/wf_config/models.py b/src/wf_config/models.py index e1606be6..edc3d1f6 100644 --- a/src/wf_config/models.py +++ b/src/wf_config/models.py @@ -162,6 +162,20 @@ SourceConfig = Annotated[ ] +class SchedulerConfig(WorkflowConfigModel): + """Opt-in same-server scheduler tuning (T12). + + No filesystem paths live in this section, so the config loader leaves + it untouched. ``enabled`` defaults off: scheduling never runs unless + the config section or the server CLI flag turns it on. + """ + + enabled: bool = False + poll_interval_s: float = Field(default=1.0, gt=0) + max_concurrent_runs: int = Field(default=4, ge=1) + drain_grace_s: float = Field(default=30.0, ge=0) + + class ServerStoresConfig(WorkflowConfigModel): """Optional role-specific store overrides. @@ -181,6 +195,7 @@ class ServerConfig(WorkflowConfigModel): stores: ServerStoresConfig = Field(default_factory=ServerStoresConfig) transports: list[ServerTransportConfig] = Field(default_factory=list) sources: list[SourceConfig] = Field(default_factory=list) + scheduler: SchedulerConfig | None = None @model_validator(mode="after") def validate_unique_source_ids(self) -> ServerConfig: diff --git a/src/wf_server/cli.py b/src/wf_server/cli.py index 54e344b6..62e63faa 100644 --- a/src/wf_server/cli.py +++ b/src/wf_server/cli.py @@ -9,6 +9,7 @@ from dotenv import load_dotenv from wf_config import ( FilesystemStoreConfig, RpcHttpTransportConfig, + WorkflowConfigFile, load_workflow_config, ) from wf_server.config import ( @@ -16,6 +17,7 @@ from wf_server.config import ( build_workflow_server_from_workflow_config, ) from wf_server.context import build_local_static_workflow_server +from wf_server.scheduling import scheduler_lifespan, server_scheduler_config from wf_transport_rpc_http import create_rpc_app app = typer.Typer(add_completion=False) @@ -50,6 +52,11 @@ def serve( max=65535, help="Override RPC bind port; defaults to config or 8765.", ), + enable_scheduler: bool = typer.Option( + False, + "--enable-scheduler", + help="Enable the opt-in same-server scheduler (or set server.scheduler.enabled).", + ), ) -> None: """Serve WorkflowApi over JSON-RPC HTTP.""" if mcp_config is not None and store_root is not None: @@ -61,6 +68,7 @@ def serve( resolved_rpc_path = "/rpc" server = None + workflow_config: WorkflowConfigFile | None = None if mcp_config is not None: server = build_workflow_server_from_legacy_mcp_config(mcp_config) @@ -109,7 +117,15 @@ def serve( # endpoint; opt into its persistence explicitly at this boundary. server = build_local_static_workflow_server(resolved_store_root, drafts=True) - rpc_app = create_rpc_app(server, rpc_path=resolved_rpc_path, drafts=True) + sched_config = server_scheduler_config(workflow_config, enable_scheduler) + rpc_app = create_rpc_app( + server, + rpc_path=resolved_rpc_path, + drafts=True, + lifespan=scheduler_lifespan(server, sched_config) + if sched_config is not None + else None, + ) uvicorn.run( rpc_app, host=resolved_host or "127.0.0.1", diff --git a/src/wf_server/scheduling.py b/src/wf_server/scheduling.py new file mode 100644 index 00000000..fab873ea --- /dev/null +++ b/src/wf_server/scheduling.py @@ -0,0 +1,84 @@ +"""Opt-in same-server scheduler wiring (T12). + +Maps neutral config plus the server CLI flag to a +:class:`~wf_scheduling.lifecycle.SchedulerService` over the server's own +stores, and exposes the transport lifespan that starts/stops the service. +Scheduling stays fully off unless enabled: :func:`server_scheduler_config` +returns ``None`` and the transport runs exactly as before. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +from wf_config import WorkflowConfigFile +from wf_scheduling.lifecycle import SchedulerService, SchedulerServiceConfig +from wf_scheduling.ownership import SchedulerOwnership +from wf_scheduling.store import FileScheduleStore + +from .context import WorkflowServer + +SCHEDULER_OWNER = "wf-server-scheduler" + + +def build_scheduler_service( + server: WorkflowServer, config: SchedulerServiceConfig +) -> SchedulerService: + """Build the scheduler service over the server's own stores. + + The schedule store is a sibling of the run store under the same + composition root (server run data lives at ``/runs``, + schedules at ``/schedules``), so one lock at + ``/scheduler.lock`` covers both stores. + """ + return SchedulerService( + schedule_store=FileScheduleStore(server.config.store_root), + run_store=server.stores.run_store, + runtime=server.context.runtime, + artifact_store=server.stores.artifact_store, + ownership=SchedulerOwnership(server.config.store_root, owner=SCHEDULER_OWNER), + config=config, + ) + + +def server_scheduler_config( + workflow_config: WorkflowConfigFile | None, + cli_enabled: bool = False, +) -> SchedulerServiceConfig | None: + """Resolve the server scheduler config, or ``None`` when disabled. + + The file section's ``enabled`` flag OR the CLI flag enables + scheduling; file values map ``poll_interval_s`` / + ``capacity=max_concurrent_runs`` / ``drain_grace_s``. The server + always ticks on its own event loop, so ``auto_tick`` is always True. + """ + section = workflow_config.server.scheduler if workflow_config is not None else None + enabled = cli_enabled or (section.enabled if section is not None else False) + if not enabled: + return None + if section is None: + return SchedulerServiceConfig(auto_tick=True) + return SchedulerServiceConfig( + poll_interval_s=section.poll_interval_s, + capacity=section.max_concurrent_runs, + drain_grace_s=section.drain_grace_s, + auto_tick=True, + ) + + +@asynccontextmanager +async def scheduler_lifespan( + server: WorkflowServer, config: SchedulerServiceConfig +) -> AsyncIterator[SchedulerService]: + """Transport lifespan owning the scheduler service. + + A failed ``start`` releases the lock (per the service contract) and + propagates before yielding, so the server never runs unprotected. + """ + service = build_scheduler_service(server, config) + await service.start() + try: + yield service + finally: + await service.stop() diff --git a/src/wf_transport_rpc_http/app.py b/src/wf_transport_rpc_http/app.py index 804a232e..26335011 100644 --- a/src/wf_transport_rpc_http/app.py +++ b/src/wf_transport_rpc_http/app.py @@ -4,6 +4,8 @@ Return annotations stay eagerly evaluated because fastapi-jsonrpc captures them while registering nested handlers for response validation and OpenRPC output. """ +from typing import Any + import fastapi_jsonrpc as jsonrpc from wf_api.models import HealthResult @@ -29,17 +31,19 @@ def create_rpc_app( *, rpc_path: str = "/rpc", drafts: bool = False, + lifespan: Any = None, ) -> jsonrpc.API: """Build a JSON-RPC HTTP app over an existing WorkflowServer. Transport code owns only JSON-RPC envelope handling. Workflow semantics stay behind server.api, so this package remains swappable with WebSocket/MCP - transports later. + transports later. ``lifespan`` (e.g. the opt-in scheduler lifespan) is + passed through to the ASGI app; ``None`` preserves existing behavior. """ if not rpc_path.startswith("/"): raise ValueError("rpc_path must start with '/'") - app = jsonrpc.API() + app = jsonrpc.API(lifespan=lifespan) entrypoint = jsonrpc.Entrypoint(rpc_path) @app.get("/healthz") diff --git a/tests/wf_config/test_config_models.py b/tests/wf_config/test_config_models.py index 0b9417f6..11f7a4cc 100644 --- a/tests/wf_config/test_config_models.py +++ b/tests/wf_config/test_config_models.py @@ -13,6 +13,7 @@ from wf_config import ( McpSourceConfig, RpcHttpTargetConfig, RpcHttpTransportConfig, + SchedulerConfig, StdioSourceTransportConfig, StdlibSourceConfig, WorkflowConfigFile, @@ -490,3 +491,50 @@ def test_workflow_config_parses_oauth_provider_profile() -> None: "access_type": "offline", "prompt": "consent", } + + +def test_server_config_scheduler_absent_is_none() -> None: + config = WorkflowConfigFile.model_validate({"version": 1}) + + assert config.server.scheduler is None + + +def test_server_config_parses_scheduler_section_with_values() -> None: + config = WorkflowConfigFile.model_validate( + { + "version": 1, + "server": { + "scheduler": { + "enabled": True, + "poll_interval_s": 2.5, + "max_concurrent_runs": 8, + "drain_grace_s": 60.0, + }, + }, + } + ) + + assert isinstance(config.server.scheduler, SchedulerConfig) + assert config.server.scheduler.enabled is True + assert config.server.scheduler.poll_interval_s == 2.5 + assert config.server.scheduler.max_concurrent_runs == 8 + assert config.server.scheduler.drain_grace_s == 60.0 + + +def test_server_config_scheduler_defaults_to_disabled() -> None: + config = WorkflowConfigFile.model_validate( + {"version": 1, "server": {"scheduler": {}}} + ) + + assert config.server.scheduler is not None + assert config.server.scheduler.enabled is False + assert config.server.scheduler.poll_interval_s == 1.0 + assert config.server.scheduler.max_concurrent_runs == 4 + assert config.server.scheduler.drain_grace_s == 30.0 + + +def test_server_config_rejects_unknown_scheduler_key() -> None: + with pytest.raises(ValidationError): + WorkflowConfigFile.model_validate( + {"version": 1, "server": {"scheduler": {"bogus_key": 1}}} + ) diff --git a/tests/wf_server/test_cli.py b/tests/wf_server/test_cli.py index 93a35869..2b1c5248 100644 --- a/tests/wf_server/test_cli.py +++ b/tests/wf_server/test_cli.py @@ -88,7 +88,7 @@ def test_rpc_server_cli_uses_configured_store_and_transport( captured["drafts"] = drafts return object() - def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False): + def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False, lifespan=None): captured["server"] = server captured["rpc_path"] = rpc_path captured["drafts"] = drafts @@ -135,7 +135,7 @@ def test_rpc_server_cli_uses_mcp_config_server(monkeypatch, tmp_path) -> None: captured["mcp_config_path"] = path return object() - def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False): + def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False, lifespan=None): captured["server"] = server captured["rpc_path"] = rpc_path captured["drafts"] = drafts @@ -211,7 +211,7 @@ def test_rpc_server_cli_mcp_config_builds_registry_capable_server( ) captured: dict[str, object] = {} - def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False): + def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False, lifespan=None): captured["source_registry_admin"] = server.source_registry_admin captured["rpc_path"] = rpc_path captured["drafts"] = drafts @@ -268,7 +268,7 @@ def test_rpc_server_cli_mcp_config_with_config_uses_transport_settings( ) captured: dict[str, object] = {} - def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False): + def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False, lifespan=None): captured["server"] = server captured["rpc_path"] = rpc_path captured["drafts"] = drafts @@ -310,7 +310,7 @@ def test_rpc_server_cli_config_with_mcp_source_uses_mcp_builder( captured["build_drafts"] = drafts return object() - def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False): + def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False, lifespan=None): captured["server"] = server captured["rpc_path"] = rpc_path captured["drafts"] = drafts @@ -422,7 +422,7 @@ def test_rpc_server_cli_config_uses_workflow_store_override( captured["build_drafts"] = drafts return object() - def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False): + def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False, lifespan=None): captured["server"] = server captured["rpc_path"] = rpc_path captured["drafts"] = drafts @@ -442,3 +442,126 @@ def test_rpc_server_cli_config_uses_workflow_store_override( assert result.exit_code == 0, result.output assert captured["workflow_store_root"] == (tmp_path / ".workflow").resolve() + + +def test_rpc_server_cli_help_mentions_enable_scheduler() -> None: + result = CliRunner().invoke(app, ["--help"]) + + assert result.exit_code == 0 + assert "--enable-scheduler" in result.output + + +def test_rpc_server_cli_scheduler_disabled_by_default(monkeypatch, tmp_path) -> None: + captured: dict[str, object] = {} + + def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False, lifespan=None): + captured["lifespan"] = lifespan + return object() + + def fake_uvicorn_run(app_obj, *, host, port, access_log): + captured["app"] = app_obj + + monkeypatch.setattr("wf_server.cli.create_rpc_app", fake_create_rpc_app) + monkeypatch.setattr("wf_server.cli.uvicorn.run", fake_uvicorn_run) + + result = CliRunner().invoke(app, ["--store-root", str(tmp_path / "store")]) + + assert result.exit_code == 0, result.output + assert captured["app"] is not None + assert captured["lifespan"] is None + + +def test_rpc_server_cli_enable_scheduler_with_store_root_builds_app( + monkeypatch, tmp_path +) -> None: + captured: dict[str, object] = {} + + def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False, lifespan=None): + captured["server"] = server + captured["lifespan"] = lifespan + return object() + + def fake_uvicorn_run(app_obj, *, host, port, access_log): + captured["app"] = app_obj + + monkeypatch.setattr("wf_server.cli.create_rpc_app", fake_create_rpc_app) + monkeypatch.setattr("wf_server.cli.uvicorn.run", fake_uvicorn_run) + + result = CliRunner().invoke( + app, ["--store-root", str(tmp_path / "store"), "--enable-scheduler"] + ) + + assert result.exit_code == 0, result.output + assert captured["server"] is not None + assert captured["app"] is not None + assert captured["lifespan"] is not None + + +def test_rpc_server_cli_config_scheduler_section_enables_without_flag( + monkeypatch, tmp_path +) -> None: + config_path = tmp_path / "wf.json" + config_path.write_text( + json.dumps( + { + "version": 1, + "server": { + "store": {"kind": "filesystem", "root": ".wf_store"}, + "scheduler": {"enabled": True}, + }, + } + ), + encoding="utf-8", + ) + captured: dict[str, object] = {} + + def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False, lifespan=None): + captured["lifespan"] = lifespan + return object() + + def fake_uvicorn_run(app_obj, *, host, port, access_log): + captured["app"] = app_obj + + monkeypatch.setattr("wf_server.cli.create_rpc_app", fake_create_rpc_app) + monkeypatch.setattr("wf_server.cli.uvicorn.run", fake_uvicorn_run) + + result = CliRunner().invoke(app, ["--config", str(config_path)]) + + assert result.exit_code == 0, result.output + assert captured["lifespan"] is not None + + +def test_rpc_server_cli_flag_overrides_disabled_config_scheduler( + monkeypatch, tmp_path +) -> None: + config_path = tmp_path / "wf.json" + config_path.write_text( + json.dumps( + { + "version": 1, + "server": { + "store": {"kind": "filesystem", "root": ".wf_store"}, + "scheduler": {"enabled": False}, + }, + } + ), + encoding="utf-8", + ) + captured: dict[str, object] = {} + + def fake_create_rpc_app(server, *, rpc_path="/rpc", drafts=False, lifespan=None): + captured["lifespan"] = lifespan + return object() + + def fake_uvicorn_run(app_obj, *, host, port, access_log): + captured["app"] = app_obj + + monkeypatch.setattr("wf_server.cli.create_rpc_app", fake_create_rpc_app) + monkeypatch.setattr("wf_server.cli.uvicorn.run", fake_uvicorn_run) + + result = CliRunner().invoke( + app, ["--config", str(config_path), "--enable-scheduler"] + ) + + assert result.exit_code == 0, result.output + assert captured["lifespan"] is not None diff --git a/tests/wf_server/test_scheduler_composition.py b/tests/wf_server/test_scheduler_composition.py new file mode 100644 index 00000000..bc30ee01 --- /dev/null +++ b/tests/wf_server/test_scheduler_composition.py @@ -0,0 +1,139 @@ +"""Server scheduler composition (T12): config mapping and store wiring. + +Scheduling stays off unless the config section or the CLI flag enables +it. When enabled, the service runs over the server's own stores behind +one composition lock, and the transport lifespan owns start/stop. +""" + +from __future__ import annotations + +from pathlib import Path + +from wf_config import WorkflowConfigFile +from wf_scheduling.lifecycle import DrainReport, SchedulerServiceConfig +from wf_scheduling.ownership import SchedulerOwnership +from wf_server.context import build_local_static_workflow_server +from wf_server.scheduling import ( + build_scheduler_service, + scheduler_lifespan, + server_scheduler_config, +) + + +def test_server_scheduler_config_disabled_by_default() -> None: + assert server_scheduler_config(None, False) is None + + bare = WorkflowConfigFile.model_validate({"version": 1}) + assert server_scheduler_config(bare, False) is None + + file_disabled = WorkflowConfigFile.model_validate( + {"version": 1, "server": {"scheduler": {"enabled": False}}} + ) + assert server_scheduler_config(file_disabled, False) is None + + +def test_server_scheduler_config_flag_enables_defaults_without_config() -> None: + resolved = server_scheduler_config(None, True) + + assert resolved is not None + assert resolved.poll_interval_s == 1.0 + assert resolved.capacity == 4 + assert resolved.drain_grace_s == 30.0 + assert resolved.auto_tick is True + + +def test_server_scheduler_config_maps_file_values() -> None: + config = WorkflowConfigFile.model_validate( + { + "version": 1, + "server": { + "scheduler": { + "enabled": True, + "poll_interval_s": 2.5, + "max_concurrent_runs": 8, + "drain_grace_s": 60.0, + }, + }, + } + ) + + resolved = server_scheduler_config(config, False) + + assert resolved is not None + assert resolved.poll_interval_s == 2.5 + assert resolved.capacity == 8 + assert resolved.drain_grace_s == 60.0 + assert resolved.auto_tick is True + + +def test_server_scheduler_config_flag_overrides_disabled_section() -> None: + config = WorkflowConfigFile.model_validate( + { + "version": 1, + "server": { + "scheduler": {"enabled": False, "max_concurrent_runs": 2}, + }, + } + ) + + resolved = server_scheduler_config(config, True) + + assert resolved is not None + assert resolved.capacity == 2 + assert resolved.auto_tick is True + + +def test_build_scheduler_service_wires_server_stores(tmp_path: Path) -> None: + server = build_local_static_workflow_server(tmp_path) + service = build_scheduler_service(server, SchedulerServiceConfig(auto_tick=False)) + + assert service.schedule_store.root == server.config.store_root + assert ( + service.schedule_store.schedules_dir == server.config.store_root / "schedules" + ) + assert service.run_store is server.stores.run_store + assert service.artifact_store is server.stores.artifact_store + assert service.runtime is server.context.runtime + assert service.ownership.lock_path == server.config.store_root / "scheduler.lock" + + +async def test_scheduler_service_start_stop_on_server_stores( + tmp_path: Path, +) -> None: + server = build_local_static_workflow_server(tmp_path) + service = build_scheduler_service(server, SchedulerServiceConfig(auto_tick=False)) + + await service.start() + try: + assert service.running is True + assert service.ownership.covers( + service.schedule_store.root, service.run_store.root + ) + finally: + report = await service.stop() + + assert isinstance(report, DrainReport) + assert service.running is False + # The lock is released: a fresh owner can acquire the same composition. + probe = SchedulerOwnership(tmp_path, owner="probe") + probe.acquire() + try: + assert probe.held is True + finally: + probe.release() + + +async def test_scheduler_lifespan_releases_lock_on_exit(tmp_path: Path) -> None: + server = build_local_static_workflow_server(tmp_path) + resolved = server_scheduler_config(None, True) + assert resolved is not None + + async with scheduler_lifespan(server, resolved) as service: + assert service.running is True + + probe = SchedulerOwnership(tmp_path, owner="probe") + probe.acquire() + try: + assert probe.held is True + finally: + probe.release()