prefect-client 3.7.1.dev6__tar.gz → 3.7.1.dev8__tar.gz

{prefect_client-3.7.1.dev6 → prefect_client-3.7.1.dev8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prefect-client
-Version: 3.7.1.dev6
+Version: 3.7.1.dev8
 Summary: Workflow orchestration and management.
 Project-URL: Changelog, https://github.com/PrefectHQ/prefect/releases
 Project-URL: Documentation, https://docs.prefect.io

prefect_client-3.7.1.dev8/src/prefect/_build_info.py

@@ -0,0 +1,5 @@
+# Generated by versioningit
+__version__ = "3.7.1.dev8"
+__build_date__ = "2026-05-14 09:02:56.534802+00:00"
+__git_commit__ = "7057b0df499b140407d3abf264538c21444724ce"
+__dirty__ = False

{prefect_client-3.7.1.dev6 → prefect_client-3.7.1.dev8}/src/prefect/concurrency/AGENTS.md

@@ -18,7 +18,7 @@ Does NOT define concurrency limits (server-side in `server/`). Does NOT handle t
 
 **`raise_on_lease_renewal_failure` (public parameter):** Controls lease renewal failure behavior independently from `strict`. When `None` (default), falls back to the value of `strict` for backward compatibility. Set to `False` to let long-running tasks continue even if a transient renewal error occurs; set to `True` to terminate immediately on renewal failure regardless of `strict`. This means `strict=True, raise_on_lease_renewal_failure=False` gives strict slot acquisition but non-fatal renewal failures, and vice versa.
 
-**Lease renewal:** `_leases.py` runs a background loop that renews immediately on entry, then sleeps for 75% of `lease_duration` between renewals. Each renewal call uses `@retry_async_fn(max_attempts=3)` for transient failures. If all 3 attempts fail, the background task raises and a done-callback either cancels execution or logs a warning, depending on `raise_on_lease_renewal_failure`.
+**Lease renewal:** `_leases.py` starts a daemon thread that renews immediately on entry, then waits for 75% of `lease_duration` between renewals. The renewal thread uses the standard `get_client(sync_client=True)` path and retries each renewal up to 3 times with exponential backoff. If all attempts fail, a shared failure handler either cancels execution with the existing sync/async cancel scope or logs a warning, depending on `raise_on_lease_renewal_failure`.
 
 **Sync/async lockstep invariant:** `asyncio.py`/`sync.py` and `_asyncio.py`/`_sync.py` are parallel implementations. Any behavior change to one must be mirrored in the other.
 
@@ -39,6 +39,6 @@ Layered — public → internal → services, with leases and events as cross-cu
 ## Pitfalls
 
 - **Service singleton is keyed on `frozenset(names)`.** Passing the same names in different order reuses the same singleton; passing a strict subset creates a different singleton. Each unique name-set gets its own queue.
-- **Lease renewal runs on the global event loop** (sync path). If the global loop is blocked or torn down, renewal silently fails — you'll see the renewal failure callback fire after `max_attempts` retries are exhausted.
+- **Lease renewal runs on a daemon thread.** The thread copies the caller's context variables for settings and logging context, then uses the standard `get_client(sync_client=True)` path instead of lease-specific async-client cloning. Sync callers join the thread briefly on exit; async callers only signal the thread to stop so the event loop is not blocked.
 - **Cancellation during acquire or release** — if a `CancelledError` is raised either during slot acquisition or during `release_concurrency_slots_with_lease` in the `finally` block, the lease ID is appended to `ConcurrencyContext.cleanup_lease_ids`. `ConcurrencyContext.__exit__` releases them via a sync client. If `ConcurrencyContext` is not active, those leases are abandoned until server-side expiry.
 - **`v1/` subdirectory** contains the legacy slot API (no lease model). New code should use the top-level `concurrency()` / `rate_limit()` APIs.

prefect_client-3.7.1.dev8/src/prefect/concurrency/_leases.py

@@ -0,0 +1,230 @@
+import contextvars
+import threading
+from contextlib import asynccontextmanager, contextmanager
+from typing import TYPE_CHECKING, AsyncGenerator, Generator
+from uuid import UUID
+
+from prefect._internal.concurrency.cancellation import (
+    AsyncCancelScope,
+    WatcherThreadCancelScope,
+)
+from prefect._internal.retries import exponential_backoff_with_jitter
+from prefect.client.orchestration import get_client
+from prefect.logging.loggers import get_logger, get_run_logger
+
+if TYPE_CHECKING:
+    from prefect.client.orchestration import SyncPrefectClient
+
+_RENEWAL_FRACTION = 0.75
+_RENEWAL_MAX_ATTEMPTS = 3
+_RENEWAL_RETRY_BASE_DELAY = 1
+_RENEWAL_RETRY_MAX_DELAY = 10
+
+
+def _renew_concurrency_lease_with_retries(
+    client: "SyncPrefectClient",
+    lease_id: UUID,
+    lease_duration: float,
+    stop_event: threading.Event,
+) -> bool:
+    """
+    Renew a concurrency lease, retrying transient failures.
+
+    Returns `False` when shutdown is requested during retry backoff.
+    """
+    logger = get_logger("concurrency")
+
+    for attempt in range(_RENEWAL_MAX_ATTEMPTS):
+        try:
+            client.renew_concurrency_lease(
+                lease_id=lease_id, lease_duration=lease_duration
+            )
+            return True
+        except Exception as exc:
+            if attempt == _RENEWAL_MAX_ATTEMPTS - 1:
+                logger.exception(
+                    "Function 'concurrency lease renewal' failed after %s attempts",
+                    _RENEWAL_MAX_ATTEMPTS,
+                )
+                raise
+
+            delay = exponential_backoff_with_jitter(
+                attempt,
+                _RENEWAL_RETRY_BASE_DELAY,
+                _RENEWAL_RETRY_MAX_DELAY,
+            )
+            logger.warning(
+                "Attempt %s of function 'concurrency lease renewal' failed with "
+                "%s: %s. Retrying in %.2f seconds...",
+                attempt + 1,
+                type(exc).__name__,
+                exc,
+                delay,
+            )
+            if stop_event.wait(delay):
+                return False
+
+    return False
+
+
+def _lease_renewal_loop(
+    lease_id: UUID,
+    lease_duration: float,
+    stop_event: threading.Event,
+) -> None:
+    """
+    Maintain a concurrency lease until stopped.
+
+    Args:
+        lease_id: The ID of the lease to maintain.
+        lease_duration: The duration of the lease in seconds.
+        stop_event: Event set by the owning context manager on exit.
+    """
+    with get_client(sync_client=True) as client:
+        while not stop_event.is_set():
+            renewed = _renew_concurrency_lease_with_retries(
+                client, lease_id, lease_duration, stop_event
+            )
+            if not renewed:
+                return
+            stop_event.wait(lease_duration * _RENEWAL_FRACTION)
+
+
+def _handle_lease_renewal_failure(
+    exc: BaseException,
+    lease_id: UUID,
+    raise_on_lease_renewal_failure: bool,
+    suppress_warnings: bool,
+    cancel_scope: AsyncCancelScope | WatcherThreadCancelScope,
+) -> None:
+    try:
+        # Use a run logger if available
+        logger = get_run_logger()
+    except Exception:
+        logger = get_logger("concurrency")
+
+    if raise_on_lease_renewal_failure:
+        logger.error(
+            "Concurrency lease renewal failed - slots are no longer reserved. "
+            "Terminating execution to prevent over-allocation. "
+            "Lease ID: %s, exception: %s",
+            lease_id,
+            exc,
+            exc_info=(type(exc), exc, exc.__traceback__),
+        )
+        cancel_scope.cancel()
+    elif suppress_warnings:
+        logger.debug(
+            "Concurrency lease renewal failed - slots are no longer reserved. "
+            "Execution will continue, but concurrency limits may be exceeded. "
+            "Lease ID: %s, exception: %s",
+            lease_id,
+            exc,
+            exc_info=(type(exc), exc, exc.__traceback__),
+        )
+    else:
+        logger.warning(
+            "Concurrency lease renewal failed - slots are no longer reserved. "
+            "Execution will continue, but concurrency limits may be exceeded. "
+            "Lease ID: %s, exception: %s",
+            lease_id,
+            exc,
+            exc_info=(type(exc), exc, exc.__traceback__),
+        )
+
+
+def _start_lease_renewal_thread(
+    lease_id: UUID,
+    lease_duration: float,
+    raise_on_lease_renewal_failure: bool,
+    suppress_warnings: bool,
+    cancel_scope: AsyncCancelScope | WatcherThreadCancelScope,
+) -> tuple[threading.Event, threading.Thread]:
+    stop_event = threading.Event()
+    renewal_ctx = contextvars.copy_context()
+
+    def renewal_loop() -> None:
+        try:
+            renewal_ctx.run(_lease_renewal_loop, lease_id, lease_duration, stop_event)
+        except Exception as exc:
+            if not stop_event.is_set():
+                _handle_lease_renewal_failure(
+                    exc,
+                    lease_id,
+                    raise_on_lease_renewal_failure,
+                    suppress_warnings,
+                    cancel_scope,
+                )
+
+    thread = threading.Thread(
+        target=renewal_loop,
+        daemon=True,
+        name=f"concurrency-lease-renewal-{lease_id}",
+    )
+    thread.start()
+    return stop_event, thread
+
+
+@contextmanager
+def maintain_concurrency_lease(
+    lease_id: UUID,
+    lease_duration: float,
+    raise_on_lease_renewal_failure: bool = False,
+    suppress_warnings: bool = False,
+) -> Generator[None, None, None]:
+    """
+    Maintain a concurrency lease for the given lease ID.
+
+    Args:
+        lease_id: The ID of the lease to maintain.
+        lease_duration: The duration of the lease in seconds.
+        raise_on_lease_renewal_failure: A boolean specifying whether to raise an error if the lease renewal fails.
+    """
+    with WatcherThreadCancelScope() as cancel_scope:
+        stop_event: threading.Event | None = None
+        thread: threading.Thread | None = None
+        try:
+            stop_event, thread = _start_lease_renewal_thread(
+                lease_id,
+                lease_duration,
+                raise_on_lease_renewal_failure,
+                suppress_warnings,
+                cancel_scope,
+            )
+            yield
+        finally:
+            if stop_event is not None:
+                stop_event.set()
+            if thread is not None:
+                thread.join(timeout=2)
+
+
+@asynccontextmanager
+async def amaintain_concurrency_lease(
+    lease_id: UUID,
+    lease_duration: float,
+    raise_on_lease_renewal_failure: bool = False,
+    suppress_warnings: bool = False,
+) -> AsyncGenerator[None, None]:
+    """
+    Maintain a concurrency lease for the given lease ID.
+
+    Args:
+        lease_id: The ID of the lease to maintain.
+        lease_duration: The duration of the lease in seconds.
+        raise_on_lease_renewal_failure: A boolean specifying whether to raise an error if the lease renewal fails.
+    """
+    with AsyncCancelScope() as cancel_scope:
+        stop_event: threading.Event | None = None
+        try:
+            stop_event, _ = _start_lease_renewal_thread(
+                lease_id,
+                lease_duration,
+                raise_on_lease_renewal_failure,
+                suppress_warnings,
+                cancel_scope,
+            )
+            yield
+        finally:
+            if stop_event is not None:
+                stop_event.set()

{prefect_client-3.7.1.dev6 → prefect_client-3.7.1.dev8}/src/prefect/events/AGENTS.md

@@ -7,6 +7,8 @@ Client-side event system for emitting, subscribing to, and defining automations
 - **Both client and server define their own event schemas.** The client-side schemas live in `schemas/` here; the server has its own parallel definitions in `server/events/schemas/`. They are structurally similar but independently maintained — the server does not import schemas from this module.
 - Events follow the CloudEvents-inspired schema: `Event` with `Resource` and `RelatedResource`.
 - Automations combine triggers (event, metric, compound, sequence) with actions.
+- **Metric triggers are Cloud-only.** OSS's `ServerTriggerTypes` excludes `MetricTrigger`, so `POST /automations` returns 422. For "alert if a flow takes too long" on OSS, use a proactive `EventTrigger` (`after={prefect.flow-run.Running}`, terminal states in `expect`, SLO in `within`).
+- **Event name prefixes differ by backend.** OSS emits `prefect.*`; Cloud emits `prefect-cloud.*` for automation/action lifecycle events (e.g., `prefect-cloud.automation.triggered`). Cross-backend tooling that filters events must match both.
 - `DeploymentTriggerTypes` are the subset of triggers usable in `prefect.yaml` deployment definitions.
 
 ## clients.py — Checkpointing Invariant

{prefect_client-3.7.1.dev6 → prefect_client-3.7.1.dev8}/src/prefect/server/api/server.py

@@ -257,7 +257,13 @@ def _build_ui_path(base_url: str, relative_path: str) -> str:
         return normalized_relative_path
 
     if normalized_relative_path == "/":
-        return normalized_base_url
+        # Preserve the trailing slash so the redirect target falls inside
+        # the bundle's mount. Starlette's `Mount` only routes requests
+        # whose path starts with `{mount}/`; a bare `{mount}` falls through
+        # to whatever else matches, which in practice is the V1 SPA mount
+        # at "/" — that returns V1's index.html under the `/v2` URL and
+        # the V1 router can't resolve the route.
+        return f"{normalized_base_url}/"
 
     return f"{normalized_base_url}{normalized_relative_path}"
 

prefect_client-3.7.1.dev8/src/prefect/workers/_cleanup_handlers.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import copy
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any
+from uuid import UUID
+
+from prefect.client.schemas.worker_channel import (
+    CANCELLING_TIMEOUT_TEARDOWN,
+    CancellingTimeoutCleanupMessagePayload,
+    CleanupMessagePayload,
+)
+from prefect.exceptions import (
+    InfrastructureNotAvailable,
+    InfrastructureNotFound,
+    ObjectNotFound,
+)
+from prefect.workers._cleanup import (
+    CleanupExecutionResult,
+    WorkerCleanupHandlerRegistry,
+)
+
+if TYPE_CHECKING:
+    from prefect.client.schemas.objects import FlowRun
+    from prefect.workers.base import BaseWorker
+
+
+class CancellingTimeoutTeardownHandler:
+    """Idempotent infrastructure teardown for `cancelling_timeout_teardown.v1`.
+
+    The cleanup message is enqueued after the server has already committed
+    the CANCELLING-timeout state outcome. Boundary contract:
+
+    - Does not propose or force flow-run state transitions.
+    - Does not require the flow run to currently be in `CANCELLING` and
+      does not skip teardown when the flow run has a `start_time`.
+    - Treats `InfrastructureNotFound` and an absent infrastructure handle
+      (when no actionable handle can be determined) as idempotent success.
+    - Maps `InfrastructureNotAvailable`, `NotImplementedError`, and missing
+      configuration context to stable release reasons rather than generic
+      errors so the executor can release with intent.
+    - Uses the cleanup payload's stable target identifiers; current worker
+      configuration is consulted only for provider access, not as proof of
+      submission-time configuration.
+    """
+
+    cleanup_kind = CANCELLING_TIMEOUT_TEARDOWN
+
+    def __init__(
+        self,
+        worker: "BaseWorker[Any, Any, Any]",
+        *,
+        grace_seconds: int = 30,
+    ) -> None:
+        self._worker = worker
+        self._grace_seconds = grace_seconds
+
+    async def cleanup(
+        self, message: CleanupMessagePayload
+    ) -> CleanupExecutionResult | None:
+        if not isinstance(message, CancellingTimeoutCleanupMessagePayload):
+            return CleanupExecutionResult.release("unexpected_payload_kind")
+
+        flow_run_id = getattr(message.target, "flow_run_id", None)
+        if flow_run_id is None:
+            return CleanupExecutionResult.release("invalid_payload")
+
+        flow_run = await self._read_flow_run(flow_run_id)
+        if flow_run is None:
+            return CleanupExecutionResult.success()
+
+        infrastructure_pid = message.target.infrastructure_pid
+        if not infrastructure_pid:
+            stored_pid = getattr(flow_run, "infrastructure_pid", None)
+            if not stored_pid:
+                return CleanupExecutionResult.success()
+            return CleanupExecutionResult.release("missing_infrastructure_handle")
+
+        # Freeze the worker's work-pool view so a mid-flight snapshot apply
+        # cannot retemplate the configuration we're handing to kill_infrastructure.
+        work_pool = copy.deepcopy(self._worker.work_pool)
+        try:
+            configuration = await self._worker.job_configuration.resolve_for_flow_run(
+                flow_run,
+                client=self._worker.client,
+                work_pool=work_pool,
+                worker_name=self._worker.name,
+                worker_id=self._worker.backend_id,
+            )
+        except ObjectNotFound:
+            return CleanupExecutionResult.release("configuration_context_unavailable")
+
+        try:
+            await self._worker.kill_infrastructure(
+                infrastructure_pid=infrastructure_pid,
+                configuration=configuration,
+                grace_seconds=self._grace_seconds,
+            )
+        except InfrastructureNotFound:
+            return CleanupExecutionResult.success()
+        except NotImplementedError:
+            return CleanupExecutionResult.release("unsupported_worker_type")
+        except InfrastructureNotAvailable:
+            return CleanupExecutionResult.release("infrastructure_not_available")
+
+        return CleanupExecutionResult.success()
+
+    async def _read_flow_run(self, flow_run_id: UUID) -> "FlowRun | None":
+        try:
+            return await self._worker.client.read_flow_run(flow_run_id)
+        except ObjectNotFound:
+            return None
+
+
+@lru_cache(maxsize=None)
+def _class_implements_kill_infrastructure(cls: type) -> bool:
+    return sum("kill_infrastructure" in vars(c) for c in cls.__mro__) > 1
+
+
+def build_cleanup_handler_registry(
+    worker: "BaseWorker[Any, Any, Any]",
+) -> "WorkerCleanupHandlerRegistry":
+    """Build the cleanup handler registry for a worker.
+
+    Class-level `cleanup_handlers` for a given cleanup kind take precedence
+    over the per-instance default registered here.
+    """
+    registry = WorkerCleanupHandlerRegistry(worker.__class__.cleanup_handlers)
+
+    if (
+        _class_implements_kill_infrastructure(type(worker))
+        and registry.get(CANCELLING_TIMEOUT_TEARDOWN) is None
+    ):
+        registry.register(CancellingTimeoutTeardownHandler(worker))
+
+    return registry
+
+
+__all__ = [
+    "CancellingTimeoutTeardownHandler",
+    "build_cleanup_handler_registry",
+]

{prefect_client-3.7.1.dev6 → prefect_client-3.7.1.dev8}/src/prefect/workers/_worker_channel/__init__.py

@@ -9,6 +9,7 @@ from prefect.workers._worker_channel._state import (
     WorkerChannelState,
     WorkerChannelStatus,
     WorkerChannelTerminalError,
+    WorkPoolSnapshotState,
 )
 from prefect.workers._worker_channel._sync import WorkPoolWorkerChannel
 
@@ -21,5 +22,6 @@ __all__ = [
     "WorkerChannelState",
     "WorkerChannelStatus",
     "WorkerChannelTerminalError",
+    "WorkPoolSnapshotState",
     "WorkPoolWorkerChannel",
 ]

{prefect_client-3.7.1.dev6 → prefect_client-3.7.1.dev8}/src/prefect/workers/_worker_channel/_protocol.py

@@ -27,6 +27,7 @@ from prefect.client.schemas.worker_channel import (
     WorkerHelloFrame,
     WorkerReadyFrame,
     WorkPoolSnapshotFrame,
+    WorkPoolSnapshotPayload,
     validate_worker_channel_frame,
 )
 from prefect.settings import get_current_settings
@@ -38,6 +39,7 @@ from prefect.workers._worker_channel._state import (
     WorkerChannelRetryableError,
     WorkerChannelSession,
     WorkerChannelTerminalError,
+    WorkPoolSnapshotState,
 )
 
 
@@ -77,6 +79,7 @@ class WorkerChannelProtocolHandler:
         self._current_session = CurrentWorkerChannelSession()
         self._worker_id: UUID | None = None
         self._worker_metadata_sent = False
+        self._work_pool_snapshots = WorkPoolSnapshotState()
 
     @property
     def worker_id(self) -> UUID | None:
@@ -86,6 +89,10 @@ class WorkerChannelProtocolHandler:
     def worker_metadata_sent(self) -> bool:
         return self._worker_metadata_sent
 
+    @property
+    def work_pool_snapshots_available(self) -> bool:
+        return self._work_pool_snapshots.snapshots_available
+
     async def handshake(
         self, websocket: websockets.asyncio.client.ClientConnection
     ) -> WorkerReadyFrame:
@@ -175,12 +182,26 @@ class WorkerChannelProtocolHandler:
             )
 
     def _handle_ready(self, ready: WorkerReadyFrame) -> None:
-        self.handle_work_pool_snapshot(ready.payload.initial_snapshot.work_pool)
+        self.reset_work_pool_snapshot_sequence()
+        self.handle_work_pool_snapshot(ready.payload.initial_snapshot)
         worker_id = ready.payload.worker_id
         if worker_id is not None:
             self.record_worker_id(worker_id)
 
-    def handle_work_pool_snapshot(self, work_pool: WorkPool) -> None:
+    def reset_work_pool_snapshot_sequence(self) -> None:
+        self._work_pool_snapshots.reset_connection_sequence()
+
+    def handle_work_pool_snapshot(self, snapshot: WorkPoolSnapshotPayload) -> bool:
+        work_pool = self._work_pool_snapshots.apply_snapshot(snapshot)
+        if work_pool is None:
+            return False
+
+        if self._on_work_pool_snapshot is not None:
+            self._on_work_pool_snapshot(work_pool)
+        return True
+
+    def record_rest_work_pool_snapshot(self, work_pool: WorkPool) -> None:
+        work_pool = self._work_pool_snapshots.replace_from_rest(work_pool)
         if self._on_work_pool_snapshot is not None:
             self._on_work_pool_snapshot(work_pool)
 
@@ -299,7 +320,7 @@ class WorkerChannelProtocolHandler:
                 ) from exc
 
             if isinstance(frame, WorkPoolSnapshotFrame):
-                self.handle_work_pool_snapshot(frame.payload.work_pool)
+                self.handle_work_pool_snapshot(frame.payload)
                 continue
 
             if isinstance(frame, CleanupMessageFrame):

{prefect_client-3.7.1.dev6 → prefect_client-3.7.1.dev8}/src/prefect/workers/_worker_channel/_state.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import copy
 from dataclasses import dataclass
 from enum import Enum
 from typing import Protocol
@@ -10,10 +11,12 @@ import websockets.asyncio.client
 import websockets.exceptions
 
 from prefect.client.orchestration import PrefectClient
+from prefect.client.schemas.objects import WorkPool
 from prefect.client.schemas.worker_channel import (
     WorkerChannelCapability,
     WorkerChannelFrame,
     WorkerReadyFrame,
+    WorkPoolSnapshotPayload,
 )
 
 
@@ -77,6 +80,34 @@ class WorkerChannelState:
 WorkerChannelFallbackState = WorkerChannelState
 
 
+@dataclass
+class WorkPoolSnapshotState:
+    work_pool: WorkPool | None = None
+    last_applied_sequence: int | None = None
+
+    @property
+    def snapshots_available(self) -> bool:
+        return self.last_applied_sequence is not None
+
+    def reset_connection_sequence(self) -> None:
+        self.last_applied_sequence = None
+
+    def apply_snapshot(self, snapshot: WorkPoolSnapshotPayload) -> WorkPool | None:
+        if (
+            self.last_applied_sequence is not None
+            and snapshot.snapshot_sequence <= self.last_applied_sequence
+        ):
+            return None
+
+        self.last_applied_sequence = snapshot.snapshot_sequence
+        self.work_pool = copy.deepcopy(snapshot.work_pool)
+        return self.work_pool
+
+    def replace_from_rest(self, work_pool: WorkPool) -> WorkPool:
+        self.work_pool = copy.deepcopy(work_pool)
+        return self.work_pool
+
+
 @dataclass
 class WorkerChannelSession:
     """One successfully negotiated worker-channel WebSocket session."""
@@ -146,6 +177,9 @@ class WorkerChannel(Protocol):
     @property
     def rest_fallback_enabled(self) -> bool: ...
 
+    @property
+    def snapshots_available(self) -> bool: ...
+
     def set_client(self, client: PrefectClient) -> None: ...
 
     async def sync(self, task_group: anyio.abc.TaskGroup | None) -> None: ...

{prefect_client-3.7.1.dev6 → prefect_client-3.7.1.dev8}/src/prefect/workers/_worker_channel/_sync.py

@@ -112,6 +112,10 @@ class WorkPoolWorkerChannel:
     def rest_fallback_enabled(self) -> bool:
         return self.state.rest_fallback_enabled
 
+    @property
+    def snapshots_available(self) -> bool:
+        return self._protocol.work_pool_snapshots_available
+
     @property
     def worker_id(self) -> UUID | None:
         return self._protocol.worker_id
@@ -128,7 +132,7 @@ class WorkPoolWorkerChannel:
             self._run_scope.cancel()
 
     async def sync(self, task_group: anyio.abc.TaskGroup | None) -> None:
-        if self.state.healthy:
+        if self.state.healthy and self.snapshots_available:
             return
 
         if task_group is not None:
@@ -138,7 +142,7 @@ class WorkPoolWorkerChannel:
             if self.url is None:
                 self.state.mark_terminal("endpoint_unavailable")
 
-        if channel_started:
+        if channel_started and self.snapshots_available:
             return
 
         await self._sync_rest_work_pool()
@@ -317,13 +321,27 @@ class WorkPoolWorkerChannel:
                     )
                 return
 
+        if self.state.healthy and self.snapshots_available:
+            self._logger.debug(
+                "Skipping REST work pool sync because the worker channel applied a "
+                "snapshot while REST sync was in flight."
+            )
+            return
+
         if not work_pool.base_job_template:
             await self._set_work_pool_template(
                 work_pool, self.default_base_job_template
             )
             work_pool.base_job_template = self.default_base_job_template
 
-        self._protocol.handle_work_pool_snapshot(work_pool)
+        if self.state.healthy and self.snapshots_available:
+            self._logger.debug(
+                "Skipping REST work pool snapshot because the worker channel applied "
+                "a snapshot while REST sync was in flight."
+            )
+            return
+
+        self._protocol.record_rest_work_pool_snapshot(work_pool)
 
     async def _set_work_pool_template(
         self, work_pool: WorkPool, job_template: dict[str, Any]