prefect-client 3.8.2.dev6__tar.gz → 3.8.3.dev1__tar.gz

{prefect_client-3.8.2.dev6 → prefect_client-3.8.3.dev1}/PKG-INFO

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

prefect_client-3.8.3.dev1/src/prefect/_build_info.py

@@ -0,0 +1,5 @@
+# Generated by versioningit
+__version__ = "3.8.3.dev1"
+__build_date__ = "2026-08-11 08:24:15.796562+00:00"
+__git_commit__ = "7d377adfd43019a26cc901d7f50491489a86377f"
+__dirty__ = False

{prefect_client-3.8.2.dev6 → prefect_client-3.8.3.dev1}/src/prefect/_internal/concurrency/services.py

@@ -408,10 +408,21 @@ class FutureQueueService(
             # If the request to the increment endpoint fails in a non-standard
             # way, we need to set the future's result so it'll be re-raised in
             # the context of the caller.
-            future.set_exception(exc)
+            if future.set_running_or_notify_cancel():
+                future.set_exception(exc)
             raise exc
         else:
-            future.set_result(response)
+            # Transitioning the future to running before setting its result closes
+            # the window where the caller cancels it while `acquire` is in flight:
+            # after this point `cancel()` fails, and if the caller got there first
+            # the acquired resource is nobody's and must be discarded.
+            if future.set_running_or_notify_cancel():
+                future.set_result(response)
+            else:
+                await self._discard(response)
+
+    async def _discard(self, response: R) -> None:
+        """Dispose of a result whose caller is no longer waiting for it."""
 
     @abc.abstractmethod
     async def acquire(self, *args: Unpack[Ts]) -> R:

{prefect_client-3.8.2.dev6 → prefect_client-3.8.3.dev1}/src/prefect/concurrency/_asyncio.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
 import logging
 from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, AsyncGenerator, Literal, Optional
@@ -150,6 +151,9 @@ async def aacquire_concurrency_slots_with_lease(
     )
     try:
         response = await asyncio.wrap_future(future)
+    except asyncio.CancelledError:
+        _release_lease_when_granted(service, future)
+        raise
     except TimeoutError as timeout:
         raise AcquireConcurrencySlotTimeoutError(
             f"Attempt to acquire concurrency slots timed out after {timeout_seconds} second(s)"
@@ -225,6 +229,28 @@ def _discard_cleanup_lease(lease_id: UUID) -> None:
         ctx.cleanup_lease_ids.remove(lease_id)
 
 
+def _release_lease_when_granted(
+    service: ConcurrencySlotAcquisitionWithLeaseService,
+    future: "concurrent.futures.Future[httpx.Response]",
+) -> None:
+    """Release the lease for an acquisition whose caller has been cancelled.
+
+    The acquisition may still be granted after the caller is gone, and the lease
+    id never reaches the `ConcurrencyContext`, so the slots would stay occupied
+    until the lease expires. The cancelled caller cannot do the release itself, so
+    the service that acquired the slots is asked to release them.
+    """
+
+    def _release(completed: "concurrent.futures.Future[httpx.Response]") -> None:
+        if completed.cancelled() or completed.exception() is not None:
+            # Nothing was granted, so there is nothing to release.
+            return
+
+        service.release_orphaned_lease(completed.result())
+
+    future.add_done_callback(_release)
+
+
 def _response_to_minimal_concurrency_limit_response(
     response: httpx.Response,
 ) -> list[MinimalConcurrencyLimitResponse]:

{prefect_client-3.8.2.dev6 → prefect_client-3.8.3.dev1}/src/prefect/concurrency/services.py

@@ -1,10 +1,11 @@
 import asyncio
+import concurrent.futures
 import json
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
 from threading import Lock
 from typing import TYPE_CHECKING, Literal, Optional
-from uuid import uuid4
+from uuid import UUID, uuid4
 
 import cachetools
 import httpx
@@ -104,12 +105,25 @@ class ConcurrencySlotAcquisitionWithLeaseService(
         super().__init__(concurrency_limit_names)
         self._client: PrefectClient
         self.concurrency_limit_names: list[str] = sorted(list(concurrency_limit_names))
+        self._pending_releases: set[concurrent.futures.Future[None]] = set()
 
     @asynccontextmanager
     async def _lifespan(self) -> AsyncGenerator[None, None]:
         async with get_client() as client:
             self._client = client
-            yield
+            try:
+                yield
+            finally:
+                # The client closes when this exits, so releases scheduled by
+                # cancelled callers have to finish first or their slots leak.
+                if self._pending_releases:
+                    await asyncio.gather(
+                        *(
+                            asyncio.wrap_future(release)
+                            for release in tuple(self._pending_releases)
+                        ),
+                        return_exceptions=True,
+                    )
 
     async def acquire(
         self,
@@ -185,6 +199,74 @@ class ConcurrencySlotAcquisitionWithLeaseService(
                     if max_retries is not None:
                         max_retries -= 1
 
+    async def _discard(self, response: httpx.Response) -> None:
+        """Release a lease granted to a caller that was cancelled while waiting.
+
+        Without this the slots stay occupied until the lease expires, because the
+        caller never learns the lease id and so cannot record it for cleanup.
+        """
+        lease_id = self._lease_id_from_response(response)
+        if lease_id is None:
+            return
+
+        try:
+            await self._client.release_concurrency_slots_with_lease(lease_id=lease_id)
+        except Exception:
+            logger.warning(
+                f"Failed to release concurrency lease {lease_id} for "
+                f"{self.concurrency_limit_names} after the acquiring caller was cancelled",
+                exc_info=True,
+            )
+
+    def release_orphaned_lease(
+        self, response: httpx.Response
+    ) -> "concurrent.futures.Future[None]":
+        """Release a lease that was delivered to a caller which is already gone.
+
+        The release is scheduled on the service's own loop so it reuses the client
+        that acquired the slots, keeping acquisition and cleanup on the same API, and
+        is tracked so that draining the service waits for it.
+        """
+        with self._lock:
+            loop = self._loop
+            if self._stopped or loop is None or loop.is_closed():
+                logger.warning(
+                    f"Unable to release concurrency lease for {self.concurrency_limit_names} "
+                    "after the acquiring caller was cancelled: the acquisition service is "
+                    "no longer running. The slots remain occupied until the lease expires."
+                )
+                stopped: concurrent.futures.Future[None] = concurrent.futures.Future()
+                stopped.set_result(None)
+                return stopped
+
+            # Registering under the lock `_stop` uses keeps the service from
+            # finishing its lifespan, and closing its client, before this release.
+            release = asyncio.run_coroutine_threadsafe(self._discard(response), loop)
+            self._pending_releases.add(release)
+
+        release.add_done_callback(self._pending_releases.discard)
+        return release
+
+    def _lease_id_from_response(self, response: httpx.Response) -> Optional[UUID]:
+        """Read the lease id out of a successful acquisition response.
+
+        Returns `None` when no slots were acquired, and logs when the response is
+        malformed, since that leaves a possibly active lease we cannot release.
+        """
+        try:
+            data = response.json()
+            if not data.get("limits"):
+                return None
+            return UUID(data["lease_id"])
+        except Exception:
+            logger.warning(
+                f"Unable to read the lease id from the acquisition response for "
+                f"{self.concurrency_limit_names}; if a lease was granted its slots "
+                "remain occupied until it expires.",
+                exc_info=True,
+            )
+            return None
+
 
 def _should_use_cache(
     names: list[str], holder: Optional["ConcurrencyLeaseHolder"]

{prefect_client-3.8.2.dev6 → prefect_client-3.8.3.dev1}/src/prefect/flows.py

@@ -3134,6 +3134,7 @@ async def aserve(
 
         if __name__ == "__main__":
             asyncio.run(main())
+        ```
     """
 
     from prefect.runner import Runner

prefect_client-3.8.2.dev6/src/prefect/_build_info.py

@@ -1,5 +0,0 @@
-# Generated by versioningit
-__version__ = "3.8.2.dev6"
-__build_date__ = "2026-08-07 08:27:53.656449+00:00"
-__git_commit__ = "eff3759c0edd62edb46e6edbe3a5e04854f5cf73"
-__dirty__ = False