diff --git a/changelog/added/marching-marmot.md b/changelog/added/marching-marmot.md new file mode 100644 index 0000000000..a30a60b081 --- /dev/null +++ b/changelog/added/marching-marmot.md @@ -0,0 +1 @@ +Add `temporalio.contrib.modal.worker` for running versioned Temporal Workers within Modal Function invocations with conservative worker defaults and deadline-aware graceful shutdown. diff --git a/temporalio/contrib/modal/__init__.py b/temporalio/contrib/modal/__init__.py new file mode 100644 index 0000000000..bd0303ec32 --- /dev/null +++ b/temporalio/contrib/modal/__init__.py @@ -0,0 +1 @@ +"""Modal integrations for Temporal.""" diff --git a/temporalio/contrib/modal/worker/README.md b/temporalio/contrib/modal/worker/README.md new file mode 100644 index 0000000000..1d4cde629a --- /dev/null +++ b/temporalio/contrib/modal/worker/README.md @@ -0,0 +1,159 @@ +# Modal worker + +Run a Temporal Worker for the lifetime of one +[Modal Function](https://modal.com/docs/guide/functions) invocation. The +integration creates a fresh Temporal client and worker, applies conservative +serverless defaults, polls for tasks, and shuts the worker down before the +Function timeout. + +## Quick start + +Define an Activity in `activities.py`: + +```python +from temporalio import activity + + +@activity.defn +async def greet(name: str) -> str: + return f"Hello, {name}!" +``` + +Define a Workflow in `workflows.py`: + +```python +from datetime import timedelta + +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from activities import greet + + +@workflow.defn +class GreetingWorkflow: + @workflow.run + async def run(self, name: str) -> str: + return await workflow.execute_activity( + greet, + name, + start_to_close_timeout=timedelta(seconds=30), + ) +``` + +Create the Modal Function in `worker.py`: + +```python +from datetime import timedelta + +import modal + +from temporalio.common import WorkerDeploymentVersion +from temporalio.contrib.modal.worker import ModalWorkerConfig, run_worker + +from activities import greet +from workflows import GreetingWorkflow + +FUNCTION_TIMEOUT = timedelta(minutes=5) +VERSION = WorkerDeploymentVersion( + deployment_name="my-service", + build_id="v1.0", +) + +app = modal.App("temporal-worker") + + +def configure(config: ModalWorkerConfig) -> None: + config.worker_config["task_queue"] = "my-task-queue" + config.worker_config["workflows"] = [GreetingWorkflow] + config.worker_config["activities"] = [greet] + + +@app.function(timeout=int(FUNCTION_TIMEOUT.total_seconds())) +async def temporal_worker() -> None: + await run_worker( + VERSION, + configure, + modal_function_timeout=FUNCTION_TIMEOUT, + ) +``` + +`run_worker` is async because Modal provides the invocation's event loop. Call +it once from a top-level Modal Function. Do not enable concurrent inputs for +that Function: every input owns one Temporal client and worker lifecycle. +Sequential container reuse is supported because every invocation constructs +fresh SDK objects. + +## Timeout and shutdown + +[Modal Functions default to a five-minute timeout](https://modal.com/docs/guide/timeouts). +Pass the Function's configured timeout to `run_worker` as +`modal_function_timeout`, as shown above. The two values must match because +Modal does not expose the remaining invocation time to running Function code. + +The SDK reserves 15 seconds of internal safety headroom. By default, the rest +is the execution budget, which includes configuration, connection, polling, +worker shutdown, and shutdown hooks. + +With a five-minute Function timeout and the default five-second graceful +shutdown timeout, the lifecycle is approximately: + +| Time after `run_worker` starts | Event | +| --- | --- | +| 0:00 | Configuration and connection begin. | +| 4:38 | Polling stops and worker shutdown begins. | +| 4:43 | The default graceful Activity drain period ends. | +| 4:45 | The SDK targets completion after reserving two seconds for shutdown hooks. | +| 5:00 | Modal's Function timeout is reached. | + +Setup and connection time consume the polling budget, so the worker may poll +for less than 4 minutes 38 seconds. When shutdown begins, the worker stops +polling for new tasks and lets in-flight Activities finish for up to +`graceful_shutdown_timeout`. It then runs the configured shutdown hooks. + +The default `shutdown_deadline_buffer` is calculated after `configure` as +`graceful_shutdown_timeout + 2 seconds`. Increasing +`graceful_shutdown_timeout` therefore moves the stop-polling time earlier +without requiring a second setting. Set `shutdown_deadline_buffer` explicitly +only when shutdown hooks or other drain work require a different reserve. +Modal's Function timeout remains the final deadline. + +Worker errors and Modal cancellation propagate after shutdown hooks run, so +Modal can record or retry an interrupted input. + +## Configuration + +Client settings are loaded using `temporalio.envconfig`. The optional config +file is resolved in this order: + +1. `TEMPORAL_CONFIG_FILE`, when set. +2. `temporal.toml` in the current working directory. + +`TEMPORAL_TASK_QUEUE` pre-populates the worker task queue. The configure +callback may override it and any client or worker option. The callback runs once +per invocation and may be synchronous, async, an async generator, or an async +context manager. + +`execution_timeout` defaults to the Modal Function timeout minus the SDK's +internal safety headroom. The configure callback may shorten it when an +invocation should stop earlier, but cannot extend it into that headroom. + +The integration always enables Worker Deployment Versioning and always disables +eager activity execution. Its conservative defaults are: + +| Setting | Default | +| --- | --- | +| `max_concurrent_activities` | 2 | +| `max_concurrent_workflow_tasks` | 10 | +| `max_concurrent_local_activities` | 2 | +| `max_concurrent_nexus_tasks` | 5 | +| `workflow_task_poller_behavior` | `SimpleMaximum(2)` | +| `activity_task_poller_behavior` | `SimpleMaximum(1)` | +| `nexus_task_poller_behavior` | `SimpleMaximum(1)` | +| `graceful_shutdown_timeout` | 5 seconds | +| `max_cached_workflows` | 30 | +| `disable_eager_activity_execution` | Always `True` | + +The default Temporal worker identity is +`@`. A user-supplied client +`identity` takes precedence. diff --git a/temporalio/contrib/modal/worker/__init__.py b/temporalio/contrib/modal/worker/__init__.py new file mode 100644 index 0000000000..aa9254978d --- /dev/null +++ b/temporalio/contrib/modal/worker/__init__.py @@ -0,0 +1,12 @@ +"""Run Temporal workers inside Modal Functions. + +See the package README for a complete Workflow and Activity example. +""" + +from temporalio.contrib.modal.worker._configure import ModalWorkerConfig +from temporalio.contrib.modal.worker._run_worker import run_worker + +__all__ = [ + "ModalWorkerConfig", + "run_worker", +] diff --git a/temporalio/contrib/modal/worker/_configure.py b/temporalio/contrib/modal/worker/_configure.py new file mode 100644 index 0000000000..c24ce869ae --- /dev/null +++ b/temporalio/contrib/modal/worker/_configure.py @@ -0,0 +1,62 @@ +"""Configuration for a Temporal worker running in a Modal Function.""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from datetime import timedelta + +from temporalio.client import ClientConnectConfig +from temporalio.worker import WorkerConfig + +logger = logging.getLogger(__name__) + + +@dataclass +class ModalWorkerConfig: + """Configuration passed to :py:func:`run_worker`'s callback. + + Attributes: + execution_timeout: Total SDK lifecycle budget for this invocation, + including configuration and connection setup. It defaults to the + Modal Function timeout minus the SDK's internal safety margin. The + callback may shorten it. + client_connect_config: Keyword arguments passed to + :py:meth:`temporalio.client.Client.connect`. Values are loaded from + Temporal environment configuration before the callback runs. + worker_config: Keyword arguments passed to + :py:class:`temporalio.worker.Worker`. The ``client`` key is managed + internally. Conservative concurrency and polling defaults are + applied, and ``task_queue`` is populated from + ``TEMPORAL_TASK_QUEUE`` when present. + shutdown_deadline_buffer: Time reserved at the end of the invocation + for worker drain and shutdown hooks. When unset, it is derived + after configuration as ``graceful_shutdown_timeout + 2s``. + shutdown_hooks: Sync or async functions called in order after the + worker stops, including when worker execution fails or is cancelled. + """ + + execution_timeout: timedelta + client_connect_config: ClientConnectConfig = field( + default_factory=ClientConnectConfig + ) + worker_config: WorkerConfig = field(default_factory=WorkerConfig) + shutdown_deadline_buffer: timedelta | None = None + shutdown_hooks: list[Callable[[], Awaitable[None] | None]] = field( + default_factory=list + ) + + +async def _run_shutdown_hooks( # type: ignore[reportUnusedFunction] + config: ModalWorkerConfig, +) -> None: + """Run registered shutdown hooks in order and log individual failures.""" + for fn in config.shutdown_hooks: + try: + result = fn() + if asyncio.iscoroutine(result): + await result + except Exception as err: + logger.error("shutdown hook error: %s", err) diff --git a/temporalio/contrib/modal/worker/_defaults.py b/temporalio/contrib/modal/worker/_defaults.py new file mode 100644 index 0000000000..34cf8f4944 --- /dev/null +++ b/temporalio/contrib/modal/worker/_defaults.py @@ -0,0 +1,71 @@ +"""Defaults for Temporal workers running in Modal Functions.""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from datetime import timedelta +from pathlib import Path + +from temporalio.worker import PollerBehaviorSimpleMaximum, WorkerConfig + +DEFAULT_MAX_CONCURRENT_ACTIVITIES: int = 2 +DEFAULT_MAX_CONCURRENT_WORKFLOW_TASKS: int = 10 +DEFAULT_MAX_CONCURRENT_LOCAL_ACTIVITIES: int = 2 +DEFAULT_MAX_CONCURRENT_NEXUS_TASKS: int = 5 +DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT: timedelta = timedelta(seconds=5) +DEFAULT_SHUTDOWN_HOOK_BUFFER: timedelta = timedelta(seconds=2) +DEFAULT_MAX_CACHED_WORKFLOWS: int = 30 +_FUNCTION_TIMEOUT_HEADROOM: timedelta = timedelta(seconds=15) + +DEFAULT_WORKFLOW_TASK_POLLER_BEHAVIOR = PollerBehaviorSimpleMaximum(maximum=2) +DEFAULT_ACTIVITY_TASK_POLLER_BEHAVIOR = PollerBehaviorSimpleMaximum(maximum=1) +DEFAULT_NEXUS_TASK_POLLER_BEHAVIOR = PollerBehaviorSimpleMaximum(maximum=1) + +ENV_TASK_QUEUE = "TEMPORAL_TASK_QUEUE" +ENV_MODAL_TASK_ID = "MODAL_TASK_ID" +ENV_CONFIG_FILE = "TEMPORAL_CONFIG_FILE" +DEFAULT_CONFIG_FILE = "temporal.toml" + + +def apply_modal_worker_defaults(config: WorkerConfig) -> None: + """Apply conservative defaults suitable for a Modal Function. + + Values already present in *config* are preserved, except eager activity + execution, which is always disabled for an ephemeral worker. + """ + config.setdefault("max_concurrent_activities", DEFAULT_MAX_CONCURRENT_ACTIVITIES) + config.setdefault( + "max_concurrent_workflow_tasks", DEFAULT_MAX_CONCURRENT_WORKFLOW_TASKS + ) + config.setdefault( + "max_concurrent_local_activities", DEFAULT_MAX_CONCURRENT_LOCAL_ACTIVITIES + ) + config.setdefault("max_concurrent_nexus_tasks", DEFAULT_MAX_CONCURRENT_NEXUS_TASKS) + config.setdefault("graceful_shutdown_timeout", DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT) + config.setdefault("max_cached_workflows", DEFAULT_MAX_CACHED_WORKFLOWS) + config.setdefault( + "workflow_task_poller_behavior", DEFAULT_WORKFLOW_TASK_POLLER_BEHAVIOR + ) + config.setdefault( + "activity_task_poller_behavior", DEFAULT_ACTIVITY_TASK_POLLER_BEHAVIOR + ) + config.setdefault("nexus_task_poller_behavior", DEFAULT_NEXUS_TASK_POLLER_BEHAVIOR) + config["disable_eager_activity_execution"] = True + + +def build_modal_identity(function_call_id: str, container_id: str) -> str: + """Build ``@`` for a worker invocation.""" + return f"{function_call_id or 'unknown'}@{container_id or 'unknown'}" + + +def modal_default_config_file_path( + getenv: Callable[[str], str | None] = os.environ.get, +) -> Path: + """Return the Temporal config path for a Modal Function. + + ``TEMPORAL_CONFIG_FILE`` takes precedence over ``temporal.toml`` in the + current working directory. + """ + config_file = getenv(ENV_CONFIG_FILE) + return Path(config_file) if config_file else Path(DEFAULT_CONFIG_FILE) diff --git a/temporalio/contrib/modal/worker/_run_worker.py b/temporalio/contrib/modal/worker/_run_worker.py new file mode 100644 index 0000000000..4010f00eaf --- /dev/null +++ b/temporalio/contrib/modal/worker/_run_worker.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +import asyncio +import importlib +import inspect +import logging +import os +import time +from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from dataclasses import dataclass, field +from datetime import timedelta +from typing import Any, TypeAlias + +import temporalio.client +import temporalio.worker +from temporalio.client import ClientConnectConfig +from temporalio.common import WorkerDeploymentVersion +from temporalio.contrib.modal.worker._configure import ( + ModalWorkerConfig, + _run_shutdown_hooks, +) +from temporalio.contrib.modal.worker._defaults import ( + _FUNCTION_TIMEOUT_HEADROOM, + DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT, + DEFAULT_SHUTDOWN_HOOK_BUFFER, + ENV_MODAL_TASK_ID, + ENV_TASK_QUEUE, + apply_modal_worker_defaults, + build_modal_identity, + modal_default_config_file_path, +) +from temporalio.envconfig import ClientConfigProfile +from temporalio.worker import WorkerConfig, WorkerDeploymentConfig + +logger = logging.getLogger(__name__) + +ConfigureCallback: TypeAlias = Callable[ + [ModalWorkerConfig], + None + | Awaitable[None] + | AsyncGenerator[None, None] + | AbstractAsyncContextManager[None], +] + + +@dataclass +class _WorkerDeps: + """External dependencies injected for testability.""" + + connect: Callable[..., Awaitable[temporalio.client.Client]] = field( + default_factory=lambda: temporalio.client.Client.connect + ) + create_worker: Callable[..., temporalio.worker.Worker] = field( + default_factory=lambda: temporalio.worker.Worker + ) + load_config: Callable[[], ClientConfigProfile] | None = None + getenv: Callable[[str], str | None] = field(default_factory=lambda: os.environ.get) + get_function_call_id: Callable[[], str] | None = None + monotonic: Callable[[], float] = time.monotonic + + +def _default_load_config(getenv: Callable[[str], str | None]) -> ClientConfigProfile: + config_path = modal_default_config_file_path(getenv) + return ClientConfigProfile.load(config_source=config_path) + + +def _default_get_function_call_id() -> str: + try: + modal = importlib.import_module("modal") + except ImportError as err: + raise RuntimeError( + "the Modal Python package is required to run a Modal worker" + ) from err + return modal.current_function_call_id() + + +def _validate_version(version: WorkerDeploymentVersion) -> None: + if not version.deployment_name or not version.build_id: + raise ValueError( + "version is required (deployment_name and build_id must be set)" + ) + + +def _validate_task_queue(config: ModalWorkerConfig) -> None: + if not config.worker_config.get("task_queue"): + raise ValueError( + "task queue not configured: set " + 'worker_config["task_queue"] or the ' + "TEMPORAL_TASK_QUEUE environment variable" + ) + + +def _work_time( + config: ModalWorkerConfig, + shutdown_deadline_buffer: timedelta, + invocation_started: float, + monotonic: Callable[[], float], +) -> timedelta: + if config.execution_timeout <= timedelta(0): + raise ValueError("execution_timeout must be greater than zero") + if shutdown_deadline_buffer < timedelta(0): + raise ValueError("shutdown_deadline_buffer must not be negative") + elapsed = timedelta(seconds=monotonic() - invocation_started) + return config.execution_timeout - elapsed - shutdown_deadline_buffer + + +def _default_execution_timeout(modal_function_timeout: timedelta) -> timedelta: + if modal_function_timeout <= timedelta(0): + raise ValueError("modal_function_timeout must be greater than zero") + if modal_function_timeout <= _FUNCTION_TIMEOUT_HEADROOM: + raise ValueError( + "modal_function_timeout must be greater than the SDK's " + f"{_FUNCTION_TIMEOUT_HEADROOM.total_seconds():.0f}s safety margin" + ) + return modal_function_timeout - _FUNCTION_TIMEOUT_HEADROOM + + +def _resolve_shutdown_deadline_buffer(config: ModalWorkerConfig) -> timedelta: + shutdown_buffer = config.shutdown_deadline_buffer + if shutdown_buffer is None: + graceful_timeout = config.worker_config.get( + "graceful_shutdown_timeout", DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT + ) + shutdown_buffer = graceful_timeout + DEFAULT_SHUTDOWN_HOOK_BUFFER + config.shutdown_deadline_buffer = shutdown_buffer + return shutdown_buffer + + +def _validate_function_timeout( + config: ModalWorkerConfig, modal_function_timeout: timedelta +) -> None: + minimum_function_timeout = config.execution_timeout + _FUNCTION_TIMEOUT_HEADROOM + if modal_function_timeout < minimum_function_timeout: + raise ValueError( + "modal_function_timeout must be at least execution_timeout + " + f"{_FUNCTION_TIMEOUT_HEADROOM.total_seconds():.0f}s " + f"({minimum_function_timeout.total_seconds():.1f}s required, " + f"{modal_function_timeout.total_seconds():.1f}s configured)" + ) + + +async def _run_worker_until_timeout( + worker: temporalio.worker.Worker, timeout: float +) -> None: + """Run a worker, suppressing only the timeout imposed by this function.""" + worker_raised_timeout = False + + async def tracked_run() -> None: + nonlocal worker_raised_timeout + try: + await worker.run() + except asyncio.TimeoutError: + # Python 3.11 aliases asyncio.TimeoutError to TimeoutError, so track + # whether the worker or wait_for raised it. + worker_raised_timeout = True + raise + + try: + await asyncio.wait_for(tracked_run(), timeout=timeout) + except asyncio.TimeoutError: + if worker_raised_timeout: + raise + + +async def run_worker( + version: WorkerDeploymentVersion, + configure: ConfigureCallback, + *, + modal_function_timeout: timedelta, +) -> None: + """Run a Temporal worker for one Modal Function invocation. + + Call this from a top-level async function decorated with + ``@app.function``. A fresh Temporal client and worker are created for every + invocation. The worker polls until its execution budget approaches the + configured shutdown buffer, then drains and runs shutdown hooks. + + The callback may be synchronous, async, an async generator, or an async + context manager. Async-generator setup runs before the Temporal client + connects; its teardown runs after the worker and shutdown hooks stop. + + Worker Deployment Versioning is always enabled. Pass the timeout configured + on ``@app.function`` as *modal_function_timeout*; the SDK reserves an + internal safety margin from it. + + Args: + version: Required worker deployment version. + configure: Per-invocation callback that registers workflows and + activities and may override client or worker options. + modal_function_timeout: Timeout configured on the enclosing Modal + Function. The deployment and this argument must use the same value. + """ + await _run_worker_internal( + version, + configure, + modal_function_timeout, + _WorkerDeps(), + ) + + +async def _run_worker_internal( + version: WorkerDeploymentVersion, + configure: ConfigureCallback, + modal_function_timeout: timedelta, + deps: _WorkerDeps, +) -> None: + """Run one invocation with injected dependencies.""" + invocation_started = deps.monotonic() + _validate_version(version) + execution_timeout = _default_execution_timeout(modal_function_timeout) + + load_config = deps.load_config or (lambda: _default_load_config(deps.getenv)) + profile = load_config() + base_connect_config: ClientConnectConfig = {**profile.to_client_connect_config()} + + base_worker_config: WorkerConfig = {} + apply_modal_worker_defaults(base_worker_config) + base_worker_config["deployment_config"] = WorkerDeploymentConfig( + version=version, + use_worker_versioning=True, + ) + + config = ModalWorkerConfig( + execution_timeout=execution_timeout, + client_connect_config=base_connect_config, + worker_config=base_worker_config, + ) + env_task_queue = deps.getenv(ENV_TASK_QUEUE) + if env_task_queue: + config.worker_config["task_queue"] = env_task_queue + + async with _invocation_config_scope(configure, config): + try: + _validate_function_timeout(config, modal_function_timeout) + shutdown_buffer = _resolve_shutdown_deadline_buffer(config) + work_time = _work_time( + config, + shutdown_buffer, + invocation_started, + deps.monotonic, + ) + if work_time <= timedelta(seconds=1): + raise RuntimeError( + "Modal execution timeout is too short: " + f"{config.execution_timeout.total_seconds():.1f}s configured but " + f"{shutdown_buffer.total_seconds():.1f}s is " + "reserved for shutdown, leaving no time for work after setup. " + "Increase the Modal Function timeout, increase execution_timeout " + "if it was shortened, or decrease shutdown_deadline_buffer" + ) + if work_time < timedelta(seconds=5): + logger.warning( + "Modal execution timeout leaves less than 5s for work after setup " + "and the shutdown buffer; consider increasing the Modal Function " + "timeout, increasing execution_timeout if it was shortened, or " + "decreasing shutdown_deadline_buffer " + "(work_time=%s, shutdown_buffer=%s)", + work_time, + shutdown_buffer, + ) + + invocation_connect_config: ClientConnectConfig = { + **config.client_connect_config + } + if "identity" not in invocation_connect_config: + get_function_call_id = ( + deps.get_function_call_id or _default_get_function_call_id + ) + invocation_connect_config["identity"] = build_modal_identity( + get_function_call_id(), deps.getenv(ENV_MODAL_TASK_ID) or "" + ) + + client = await deps.connect(**invocation_connect_config) + worker = deps.create_worker(client, **config.worker_config) + + work_time = _work_time( + config, + shutdown_buffer, + invocation_started, + deps.monotonic, + ) + if work_time <= timedelta(0): + raise RuntimeError( + "Modal execution budget was exhausted before the worker started" + ) + await _run_worker_until_timeout(worker, work_time.total_seconds()) + finally: + await _run_shutdown_hooks(config) + + +@asynccontextmanager +async def _invocation_config_scope( + configure: ConfigureCallback, + config: ModalWorkerConfig, +) -> AsyncGenerator[ModalWorkerConfig, None]: + """Run configure setup and bracket any async teardown around the worker.""" + if inspect.isasyncgenfunction(configure): + cm: Any = asynccontextmanager(configure)(config) + else: + result = configure(config) + if inspect.isawaitable(result): + await result + cm = result if result is not None and hasattr(result, "__aenter__") else None + + if cm is not None: + async with cm: + _validate_task_queue(config) + yield config + else: + _validate_task_queue(config) + yield config diff --git a/tests/contrib/modal/__init__.py b/tests/contrib/modal/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/contrib/modal/worker/__init__.py b/tests/contrib/modal/worker/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/contrib/modal/worker/test_modal_worker.py b/tests/contrib/modal/worker/test_modal_worker.py new file mode 100644 index 0000000000..b1bb589e48 --- /dev/null +++ b/tests/contrib/modal/worker/test_modal_worker.py @@ -0,0 +1,593 @@ +"""Tests for temporalio.contrib.modal.worker.""" + +from __future__ import annotations + +import asyncio +import dataclasses +from collections.abc import Awaitable +from contextlib import asynccontextmanager +from datetime import timedelta +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from temporalio.common import WorkerDeploymentVersion +from temporalio.contrib.modal.worker._configure import ( + ModalWorkerConfig, + _run_shutdown_hooks, +) +from temporalio.contrib.modal.worker._defaults import ( + _FUNCTION_TIMEOUT_HEADROOM, + DEFAULT_MAX_CACHED_WORKFLOWS, + DEFAULT_MAX_CONCURRENT_ACTIVITIES, + DEFAULT_MAX_CONCURRENT_LOCAL_ACTIVITIES, + DEFAULT_MAX_CONCURRENT_NEXUS_TASKS, + DEFAULT_MAX_CONCURRENT_WORKFLOW_TASKS, + apply_modal_worker_defaults, + build_modal_identity, + modal_default_config_file_path, +) +from temporalio.contrib.modal.worker._run_worker import ( + _run_worker_internal, + _run_worker_until_timeout, + _WorkerDeps, +) +from temporalio.envconfig import ClientConfigProfile +from temporalio.worker import WorkerConfig + +TEST_VERSION = WorkerDeploymentVersion( + deployment_name="test-deployment", + build_id="test-build", +) +TEST_MODAL_FUNCTION_TIMEOUT = timedelta(minutes=5) +TEST_EXECUTION_TIMEOUT = TEST_MODAL_FUNCTION_TIMEOUT - _FUNCTION_TIMEOUT_HEADROOM + + +def _make_test_deps( + *, + connect_kwargs_capture: list[dict[str, Any]] | None = None, + worker_kwargs_capture: list[dict[str, Any]] | None = None, +) -> _WorkerDeps: + mock_client = MagicMock() + mock_worker = MagicMock() + mock_worker.run = AsyncMock() + + async def fake_connect(**kwargs: Any) -> Any: + if connect_kwargs_capture is not None: + connect_kwargs_capture.append(kwargs) + return mock_client + + def fake_create_worker(_client: Any, **kwargs: Any) -> Any: + if worker_kwargs_capture is not None: + worker_kwargs_capture.append(kwargs) + return mock_worker + + env = { + "TEMPORAL_TASK_QUEUE": "test-queue", + "MODAL_TASK_ID": "ta-container", + } + return _WorkerDeps( + connect=fake_connect, + create_worker=fake_create_worker, + load_config=lambda: ClientConfigProfile(), + getenv=env.get, + get_function_call_id=lambda: "fc-call", + monotonic=lambda: 0, + ) + + +class TestModalWorkerConfig: + def test_fields_are_modifiable(self) -> None: + config = ModalWorkerConfig(execution_timeout=TEST_EXECUTION_TIMEOUT) + config.client_connect_config["namespace"] = "custom-ns" + config.worker_config["task_queue"] = "my-queue" + config.execution_timeout = timedelta(minutes=10) + config.shutdown_deadline_buffer = timedelta(seconds=10) + + assert config.client_connect_config["namespace"] == "custom-ns" + assert config.worker_config["task_queue"] == "my-queue" + assert config.execution_timeout == timedelta(minutes=10) + assert config.shutdown_deadline_buffer == timedelta(seconds=10) + + def test_mutable_fields_are_independent(self) -> None: + first = ModalWorkerConfig(execution_timeout=TEST_EXECUTION_TIMEOUT) + second = ModalWorkerConfig(execution_timeout=TEST_EXECUTION_TIMEOUT) + first.worker_config["max_concurrent_activities"] = 99 + first.shutdown_hooks.append(lambda: None) + + assert "max_concurrent_activities" not in second.worker_config + assert not second.shutdown_hooks + assert dataclasses.is_dataclass(ModalWorkerConfig) + + @pytest.mark.asyncio + async def test_shutdown_hooks_run_in_order_and_continue_after_error( + self, + ) -> None: + config = ModalWorkerConfig(execution_timeout=TEST_EXECUTION_TIMEOUT) + calls: list[str] = [] + + async def async_hook() -> None: + calls.append("async") + + def failing_hook() -> None: + calls.append("failing") + raise RuntimeError("flush failed") + + config.shutdown_hooks.extend( + [lambda: calls.append("sync"), failing_hook, async_hook] + ) + await _run_shutdown_hooks(config) + + assert calls == ["sync", "failing", "async"] + + +class TestDefaults: + def test_apply_modal_worker_defaults(self) -> None: + config: WorkerConfig = {} + apply_modal_worker_defaults(config) + + assert ( + config.get("max_concurrent_activities") == DEFAULT_MAX_CONCURRENT_ACTIVITIES + ) + assert ( + config.get("max_concurrent_workflow_tasks") + == DEFAULT_MAX_CONCURRENT_WORKFLOW_TASKS + ) + assert ( + config.get("max_concurrent_local_activities") + == DEFAULT_MAX_CONCURRENT_LOCAL_ACTIVITIES + ) + assert ( + config.get("max_concurrent_nexus_tasks") + == DEFAULT_MAX_CONCURRENT_NEXUS_TASKS + ) + assert config.get("max_cached_workflows") == DEFAULT_MAX_CACHED_WORKFLOWS + assert config.get("disable_eager_activity_execution") is True + + def test_defaults_preserve_overrides_but_always_disable_eager(self) -> None: + config: WorkerConfig = { + "max_concurrent_activities": 50, + "disable_eager_activity_execution": False, + } + apply_modal_worker_defaults(config) + + assert config.get("max_concurrent_activities") == 50 + assert config.get("disable_eager_activity_execution") is True + + def test_build_modal_identity(self) -> None: + assert build_modal_identity("fc-123", "ta-456") == "fc-123@ta-456" + assert build_modal_identity("", "") == "unknown@unknown" + + def test_default_config_file_path(self) -> None: + empty_env: dict[str, str] = {} + configured_env = {"TEMPORAL_CONFIG_FILE": "/config/temporal.toml"} + assert modal_default_config_file_path(empty_env.get) == Path("temporal.toml") + assert modal_default_config_file_path(configured_env.get) == Path( + "/config/temporal.toml" + ) + + +class TestRunWorkerUntilTimeout: + @pytest.mark.asyncio + async def test_budget_timeout_is_suppressed(self) -> None: + worker = MagicMock() + worker_cancelled = False + + async def run() -> None: + nonlocal worker_cancelled + try: + await asyncio.Event().wait() + finally: + worker_cancelled = True + + worker.run = run + + await _run_worker_until_timeout(worker, 0.001) + + assert worker_cancelled + + @pytest.mark.asyncio + async def test_worker_timeout_is_propagated(self) -> None: + worker = MagicMock() + worker_error = TimeoutError("worker failed") + + async def run() -> None: + raise worker_error + + worker.run = run + + with pytest.raises(TimeoutError) as err: + await _run_worker_until_timeout(worker, 1) + + assert err.value is worker_error + + +class TestRunWorkerInternal: + @pytest.mark.asyncio + async def test_success_applies_defaults_and_versioning(self) -> None: + worker_capture: list[dict[str, Any]] = [] + deps = _make_test_deps(worker_kwargs_capture=worker_capture) + + await _run_worker_internal( + TEST_VERSION, + lambda config: None, + TEST_MODAL_FUNCTION_TIMEOUT, + deps, + ) + + worker_config = worker_capture[0] + assert worker_config["task_queue"] == "test-queue" + assert ( + worker_config["max_concurrent_activities"] + == DEFAULT_MAX_CONCURRENT_ACTIVITIES + ) + assert worker_config["disable_eager_activity_execution"] is True + deployment_config = worker_config["deployment_config"] + assert deployment_config.use_worker_versioning is True + assert deployment_config.version == TEST_VERSION + + @pytest.mark.asyncio + async def test_config_receives_defaults_and_can_override_them(self) -> None: + connect_capture: list[dict[str, Any]] = [] + worker_capture: list[dict[str, Any]] = [] + deps = _make_test_deps( + connect_kwargs_capture=connect_capture, + worker_kwargs_capture=worker_capture, + ) + + def configure(config: ModalWorkerConfig) -> None: + assert config.execution_timeout == TEST_EXECUTION_TIMEOUT + assert ( + config.worker_config.get("max_concurrent_activities") + == DEFAULT_MAX_CONCURRENT_ACTIVITIES + ) + config.client_connect_config["namespace"] = "custom-ns" + config.worker_config["task_queue"] = "custom-queue" + config.worker_config["max_concurrent_activities"] = 42 + + await _run_worker_internal( + TEST_VERSION, configure, TEST_MODAL_FUNCTION_TIMEOUT, deps + ) + + assert connect_capture[0]["namespace"] == "custom-ns" + assert worker_capture[0]["task_queue"] == "custom-queue" + assert worker_capture[0]["max_concurrent_activities"] == 42 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("modal_function_timeout", "expected_execution_timeout"), + [ + (timedelta(minutes=5), timedelta(minutes=4, seconds=45)), + (timedelta(minutes=15), timedelta(minutes=14, seconds=45)), + ], + ) + async def test_execution_timeout_defaults_from_function_timeout( + self, + modal_function_timeout: timedelta, + expected_execution_timeout: timedelta, + ) -> None: + deps = _make_test_deps() + execution_timeouts: list[timedelta] = [] + + def configure(config: ModalWorkerConfig) -> None: + execution_timeouts.append(config.execution_timeout) + + await _run_worker_internal( + TEST_VERSION, + configure, + modal_function_timeout, + deps, + ) + + assert execution_timeouts == [expected_execution_timeout] + + @pytest.mark.asyncio + async def test_identity_uses_function_call_and_container(self) -> None: + connect_capture: list[dict[str, Any]] = [] + deps = _make_test_deps(connect_kwargs_capture=connect_capture) + + await _run_worker_internal( + TEST_VERSION, + lambda config: None, + TEST_MODAL_FUNCTION_TIMEOUT, + deps, + ) + + assert connect_capture[0]["identity"] == "fc-call@ta-container" + + @pytest.mark.asyncio + async def test_user_identity_override_does_not_read_modal_context(self) -> None: + connect_capture: list[dict[str, Any]] = [] + deps = _make_test_deps(connect_kwargs_capture=connect_capture) + get_function_call_id = MagicMock(side_effect=AssertionError("not expected")) + deps.get_function_call_id = get_function_call_id + + def configure(config: ModalWorkerConfig) -> None: + config.client_connect_config["identity"] = "custom-worker" + + await _run_worker_internal( + TEST_VERSION, configure, TEST_MODAL_FUNCTION_TIMEOUT, deps + ) + + assert connect_capture[0]["identity"] == "custom-worker" + get_function_call_id.assert_not_called() + + @pytest.mark.asyncio + async def test_missing_task_queue(self) -> None: + deps = _make_test_deps() + deps.getenv = lambda _name: None + + with pytest.raises(ValueError, match="task queue not configured"): + await _run_worker_internal( + TEST_VERSION, + lambda config: None, + TEST_MODAL_FUNCTION_TIMEOUT, + deps, + ) + + @pytest.mark.asyncio + async def test_missing_version(self) -> None: + deps = _make_test_deps() + with pytest.raises(ValueError, match="version is required"): + await _run_worker_internal( + WorkerDeploymentVersion(deployment_name="", build_id=""), + lambda config: None, + TEST_MODAL_FUNCTION_TIMEOUT, + deps, + ) + + @pytest.mark.asyncio + async def test_function_timeout_requires_headroom(self) -> None: + connect_capture: list[dict[str, Any]] = [] + worker_capture: list[dict[str, Any]] = [] + deps = _make_test_deps( + connect_kwargs_capture=connect_capture, + worker_kwargs_capture=worker_capture, + ) + + def configure(config: ModalWorkerConfig) -> None: + config.execution_timeout = TEST_MODAL_FUNCTION_TIMEOUT + + with pytest.raises( + ValueError, + match=("modal_function_timeout must be at least execution_timeout \\+ 15s"), + ): + await _run_worker_internal( + TEST_VERSION, + configure, + TEST_MODAL_FUNCTION_TIMEOUT, + deps, + ) + + assert not connect_capture + assert not worker_capture + + @pytest.mark.asyncio + async def test_function_timeout_uses_configured_execution_timeout(self) -> None: + deps = _make_test_deps() + + def configure(config: ModalWorkerConfig) -> None: + config.execution_timeout = TEST_EXECUTION_TIMEOUT + timedelta(seconds=1) + + with pytest.raises(ValueError, match="301.0s required, 300.0s configured"): + await _run_worker_internal( + TEST_VERSION, + configure, + TEST_MODAL_FUNCTION_TIMEOUT, + deps, + ) + + @pytest.mark.asyncio + async def test_function_timeout_must_exceed_internal_headroom(self) -> None: + deps = _make_test_deps() + + with pytest.raises(ValueError, match="greater than the SDK's 15s safety"): + await _run_worker_internal( + TEST_VERSION, + lambda config: None, + _FUNCTION_TIMEOUT_HEADROOM, + deps, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("configured_buffer", "expected_buffer", "expected_work_time"), + [ + (None, timedelta(seconds=32), 28.0), + (timedelta(seconds=10), timedelta(seconds=10), 50.0), + ], + ) + async def test_shutdown_buffer_follows_graceful_timeout_unless_overridden( + self, + configured_buffer: timedelta | None, + expected_buffer: timedelta, + expected_work_time: float, + ) -> None: + deps = _make_test_deps() + captured_config: list[ModalWorkerConfig] = [] + captured_timeout: list[float] = [] + + async def fake_wait_for(awaitable: Awaitable[None], timeout: float) -> None: + captured_timeout.append(timeout) + await awaitable + + def configure(config: ModalWorkerConfig) -> None: + config.execution_timeout = timedelta(seconds=60) + config.worker_config["graceful_shutdown_timeout"] = timedelta(seconds=30) + config.shutdown_deadline_buffer = configured_buffer + captured_config.append(config) + + with patch( + "temporalio.contrib.modal.worker._run_worker.asyncio.wait_for", + fake_wait_for, + ): + await _run_worker_internal( + TEST_VERSION, + configure, + TEST_MODAL_FUNCTION_TIMEOUT, + deps, + ) + + assert captured_config[0].shutdown_deadline_buffer == expected_buffer + assert captured_timeout == [expected_work_time] + + @pytest.mark.asyncio + async def test_timeout_too_short_after_setup(self) -> None: + deps = _make_test_deps() + shutdown_called = False + + def configure(config: ModalWorkerConfig) -> None: + config.execution_timeout = timedelta(seconds=8) + config.shutdown_deadline_buffer = timedelta(seconds=7) + config.shutdown_hooks.append(_mark_shutdown) + + def _mark_shutdown() -> None: + nonlocal shutdown_called + shutdown_called = True + + with pytest.raises(RuntimeError, match="execution timeout is too short"): + await _run_worker_internal( + TEST_VERSION, configure, TEST_MODAL_FUNCTION_TIMEOUT, deps + ) + assert shutdown_called + + @pytest.mark.asyncio + async def test_short_work_time_logs_warning(self) -> None: + deps = _make_test_deps() + + def configure(config: ModalWorkerConfig) -> None: + config.execution_timeout = timedelta(seconds=10) + config.shutdown_deadline_buffer = timedelta(seconds=7) + + with patch("temporalio.contrib.modal.worker._run_worker.logger") as mock_logger: + await _run_worker_internal( + TEST_VERSION, configure, TEST_MODAL_FUNCTION_TIMEOUT, deps + ) + + mock_logger.warning.assert_called_once() + assert "less than 5s" in mock_logger.warning.call_args[0][0] + + @pytest.mark.asyncio + async def test_budget_includes_configuration_and_connection(self) -> None: + deps = _make_test_deps() + times = iter([100.0, 102.0, 105.0]) + deps.monotonic = lambda: next(times) + captured_timeout: list[float] = [] + + async def fake_wait_for(awaitable: Awaitable[None], timeout: float) -> None: + captured_timeout.append(timeout) + await awaitable + + def configure(config: ModalWorkerConfig) -> None: + config.execution_timeout = timedelta(seconds=20) + config.shutdown_deadline_buffer = timedelta(seconds=7) + + with patch( + "temporalio.contrib.modal.worker._run_worker.asyncio.wait_for", + fake_wait_for, + ): + await _run_worker_internal( + TEST_VERSION, configure, TEST_MODAL_FUNCTION_TIMEOUT, deps + ) + + assert captured_timeout == [8.0] + + @pytest.mark.asyncio + async def test_shutdown_hooks_run_on_worker_error(self) -> None: + deps = _make_test_deps() + shutdown_called = False + + async def failing_run() -> None: + raise RuntimeError("worker failed") + + def create_worker(_client: Any, **_kwargs: Any) -> Any: + worker = MagicMock() + worker.run = failing_run + return worker + + deps.create_worker = create_worker + + def configure(config: ModalWorkerConfig) -> None: + def shutdown() -> None: + nonlocal shutdown_called + shutdown_called = True + + config.shutdown_hooks.append(shutdown) + + with pytest.raises(RuntimeError, match="worker failed"): + await _run_worker_internal( + TEST_VERSION, configure, TEST_MODAL_FUNCTION_TIMEOUT, deps + ) + assert shutdown_called + + @pytest.mark.asyncio + async def test_modal_cancellation_propagates_after_shutdown_hooks(self) -> None: + deps = _make_test_deps() + shutdown_called = False + + async def cancelled_run() -> None: + raise asyncio.CancelledError + + def create_worker(_client: Any, **_kwargs: Any) -> Any: + worker = MagicMock() + worker.run = cancelled_run + return worker + + deps.create_worker = create_worker + + def configure(config: ModalWorkerConfig) -> None: + def shutdown() -> None: + nonlocal shutdown_called + shutdown_called = True + + config.shutdown_hooks.append(shutdown) + + with pytest.raises(asyncio.CancelledError): + await _run_worker_internal( + TEST_VERSION, configure, TEST_MODAL_FUNCTION_TIMEOUT, deps + ) + assert shutdown_called + + @pytest.mark.asyncio + async def test_async_generator_brackets_worker_and_hooks(self) -> None: + events: list[str] = [] + deps = _make_test_deps() + + original_connect = deps.connect + + async def connect(**kwargs: Any) -> Any: + events.append("connect") + return await original_connect(**kwargs) + + deps.connect = connect + + async def configure(config: ModalWorkerConfig): + events.append("setup") + config.shutdown_hooks.append(lambda: events.append("shutdown")) + yield + events.append("teardown") + + await _run_worker_internal( + TEST_VERSION, configure, TEST_MODAL_FUNCTION_TIMEOUT, deps + ) + + assert events == ["setup", "connect", "shutdown", "teardown"] + + @pytest.mark.asyncio + async def test_async_context_manager_configure(self) -> None: + events: list[str] = [] + deps = _make_test_deps() + + @asynccontextmanager + async def configure(_config: ModalWorkerConfig): + events.append("setup") + yield + events.append("teardown") + + await _run_worker_internal( + TEST_VERSION, configure, TEST_MODAL_FUNCTION_TIMEOUT, deps + ) + + assert events == ["setup", "teardown"]