prefect-client 2.14.8__py3-none-any.whl → 2.14.10__py3-none-any.whl

prefect/events/filters.py

@@ -0,0 +1,214 @@
+from typing import List, Optional, Tuple, cast
+from uuid import UUID
+
+import pendulum
+
+from prefect._internal.pydantic import HAS_PYDANTIC_V2
+from prefect.events.schemas import Event, Resource, ResourceSpecification
+from prefect.server.utilities.schemas import DateTimeTZ, PrefectBaseModel
+
+if HAS_PYDANTIC_V2:
+    from pydantic.v1 import Field
+else:
+    from pydantic import Field
+
+
+class EventDataFilter(PrefectBaseModel):
+    """A base class for filtering event data."""
+
+    class Config:
+        extra = "forbid"
+
+    def get_filters(self) -> List["EventDataFilter"]:
+        return [
+            filter
+            for filter in [
+                getattr(self, name)
+                for name, field in self.__fields__.items()
+                if issubclass(field.type_, EventDataFilter)
+            ]
+            if filter
+        ]
+
+    def includes(self, event: Event) -> bool:
+        """Does the given event match the criteria of this filter?"""
+        return all(filter.includes(event) for filter in self.get_filters())
+
+    def excludes(self, event: Event) -> bool:
+        """Would the given filter exclude this event?"""
+        return not self.includes(event)
+
+
+class EventOccurredFilter(EventDataFilter):
+    since: DateTimeTZ = Field(
+        default_factory=lambda: cast(
+            DateTimeTZ,
+            pendulum.now("UTC").start_of("day").subtract(days=180),
+        ),
+        description="Only include events after this time (inclusive)",
+    )
+    until: DateTimeTZ = Field(
+        default_factory=lambda: cast(DateTimeTZ, pendulum.now("UTC")),
+        description="Only include events prior to this time (inclusive)",
+    )
+
+    def includes(self, event: Event) -> bool:
+        return self.since <= event.occurred <= self.until
+
+
+class EventNameFilter(EventDataFilter):
+    prefix: Optional[List[str]] = Field(
+        None, description="Only include events matching one of these prefixes"
+    )
+    exclude_prefix: Optional[List[str]] = Field(
+        None, description="Exclude events matching one of these prefixes"
+    )
+
+    name: Optional[List[str]] = Field(
+        None, description="Only include events matching one of these names exactly"
+    )
+    exclude_name: Optional[List[str]] = Field(
+        None, description="Exclude events matching one of these names exactly"
+    )
+
+    def includes(self, event: Event) -> bool:
+        if self.prefix:
+            if not any(event.event.startswith(prefix) for prefix in self.prefix):
+                return False
+
+        if self.exclude_prefix:
+            if any(event.event.startswith(prefix) for prefix in self.exclude_prefix):
+                return False
+
+        if self.name:
+            if not any(event.event == name for name in self.name):
+                return False
+
+        if self.exclude_name:
+            if any(event.event == name for name in self.exclude_name):
+                return False
+
+        return True
+
+
+class EventResourceFilter(EventDataFilter):
+    id: Optional[List[str]] = Field(
+        None, description="Only include events for resources with these IDs"
+    )
+    id_prefix: Optional[List[str]] = Field(
+        None,
+        description=(
+            "Only include events for resources with IDs starting with these prefixes."
+        ),
+    )
+    labels: Optional[ResourceSpecification] = Field(
+        None, description="Only include events for resources with these labels"
+    )
+
+    def includes(self, event: Event) -> bool:
+        if self.id:
+            if not any(event.resource.id == resource_id for resource_id in self.id):
+                return False
+
+        if self.id_prefix:
+            if not any(
+                event.resource.id.startswith(prefix) for prefix in self.id_prefix
+            ):
+                return False
+
+        if self.labels:
+            if not self.labels.matches(event.resource):
+                return False
+
+        return True
+
+
+class EventRelatedFilter(EventDataFilter):
+    id: Optional[List[str]] = Field(
+        None, description="Only include events for related resources with these IDs"
+    )
+    role: Optional[List[str]] = Field(
+        None, description="Only include events for related resources in these roles"
+    )
+    resources_in_roles: Optional[List[Tuple[str, str]]] = Field(
+        None,
+        description=(
+            "Only include events with specific related resources in specific roles"
+        ),
+    )
+    labels: Optional[ResourceSpecification] = Field(
+        None, description="Only include events for related resources with these labels"
+    )
+
+
+class EventAnyResourceFilter(EventDataFilter):
+    id: Optional[List[str]] = Field(
+        None, description="Only include events for resources with these IDs"
+    )
+    id_prefix: Optional[List[str]] = Field(
+        None,
+        description=(
+            "Only include events for resources with IDs starting with these prefixes"
+        ),
+    )
+    labels: Optional[ResourceSpecification] = Field(
+        None, description="Only include events for related resources with these labels"
+    )
+
+    def includes(self, event: Event) -> bool:
+        resources = [event.resource] + event.related
+        if not any(self._includes(resource) for resource in resources):
+            return False
+        return True
+
+    def _includes(self, resource: Resource) -> bool:
+        if self.id:
+            if not any(resource.id == resource_id for resource_id in self.id):
+                return False
+
+        if self.id_prefix:
+            if not any(resource.id.startswith(prefix) for prefix in self.id_prefix):
+                return False
+
+        if self.labels:
+            if not self.labels.matches(resource):
+                return False
+
+        return True
+
+
+class EventIDFilter(EventDataFilter):
+    id: Optional[List[UUID]] = Field(
+        None, description="Only include events with one of these IDs"
+    )
+
+    def includes(self, event: Event) -> bool:
+        if self.id:
+            if not any(event.id == id for id in self.id):
+                return False
+
+        return True
+
+
+class EventFilter(EventDataFilter):
+    occurred: EventOccurredFilter = Field(
+        default_factory=EventOccurredFilter,
+        description="Filter criteria for when the events occurred",
+    )
+    event: Optional[EventNameFilter] = Field(
+        None,
+        description="Filter criteria for the event name",
+    )
+    any_resource: Optional[EventAnyResourceFilter] = Field(
+        None, description="Filter criteria for any resource involved in the event"
+    )
+    resource: Optional[EventResourceFilter] = Field(
+        None, description="Filter criteria for the resource of the event"
+    )
+    related: Optional[EventRelatedFilter] = Field(
+        None, description="Filter criteria for the related resources of the event"
+    )
+    id: EventIDFilter = Field(
+        default_factory=EventIDFilter,
+        description="Filter criteria for the events' ID",
+    )

prefect/exceptions.py

@@ -296,6 +296,10 @@ class Pause(PrefectSignal):
     Raised when a flow run is PAUSED and needs to exit for resubmission.
     """
 
+    def __init__(self, *args, state=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.state = state
+
 
 class ExternalSignal(BaseException):
     """

prefect/infrastructure/base.py

@@ -11,22 +11,27 @@ from prefect._internal.compatibility.experimental import (
     experiment_enabled,
 )
 from prefect._internal.pydantic import HAS_PYDANTIC_V2
+from prefect.client.schemas.actions import WorkPoolCreate
+from prefect.exceptions import ObjectAlreadyExists
 
 if HAS_PYDANTIC_V2:
     import pydantic.v1 as pydantic
 else:
     import pydantic
 
+from rich.console import Console
 from typing_extensions import Self
 
 import prefect
-from prefect.blocks.core import Block
+from prefect.blocks.core import Block, BlockNotSavedError
 from prefect.logging import get_logger
 from prefect.settings import (
     PREFECT_EXPERIMENTAL_WARN,
     PREFECT_EXPERIMENTAL_WARN_ENHANCED_CANCELLATION,
+    PREFECT_UI_URL,
     get_current_settings,
 )
+from prefect.utilities.asyncutils import sync_compatible
 
 MIN_COMPAT_PREFECT_VERSION = "2.0b12"
 
@@ -66,6 +71,106 @@ class Infrastructure(Block, abc.ABC):
         description="The command to run in the infrastructure.",
     )
 
+    async def generate_work_pool_base_job_template(self):
+        if self._block_document_id is None:
+            raise BlockNotSavedError(
+                "Cannot publish as work pool, block has not been saved. Please call"
+                " `.save()` on your block before publishing."
+            )
+
+        block_schema = self.__class__.schema()
+        return {
+            "job_configuration": {"block": "{{ block }}"},
+            "variables": {
+                "type": "object",
+                "properties": {
+                    "block": {
+                        "title": "Block",
+                        "description": (
+                            "The infrastructure block to use for job creation."
+                        ),
+                        "allOf": [{"$ref": f"#/definitions/{self.__class__.__name__}"}],
+                        "default": {
+                            "$ref": {"block_document_id": str(self._block_document_id)}
+                        },
+                    }
+                },
+                "required": ["block"],
+                "definitions": {self.__class__.__name__: block_schema},
+            },
+        }
+
+    def get_corresponding_worker_type(self):
+        return "block"
+
+    @sync_compatible
+    async def publish_as_work_pool(self, work_pool_name: Optional[str] = None):
+        """
+        Creates a work pool configured to use the given block as the job creator.
+
+        Used to migrate from a agents setup to a worker setup.
+
+        Args:
+            work_pool_name: The name to give to the created work pool. If not provided, the name of the current
+                block will be used.
+        """
+
+        base_job_template = await self.generate_work_pool_base_job_template()
+        work_pool_name = work_pool_name or self._block_document_name
+
+        if work_pool_name is None:
+            raise ValueError(
+                "`work_pool_name` must be provided if the block has not been saved."
+            )
+
+        console = Console()
+
+        try:
+            async with prefect.get_client() as client:
+                work_pool = await client.create_work_pool(
+                    work_pool=WorkPoolCreate(
+                        name=work_pool_name,
+                        type=self.get_corresponding_worker_type(),
+                        base_job_template=base_job_template,
+                    )
+                )
+        except ObjectAlreadyExists:
+            console.print(
+                (
+                    f"Work pool with name {work_pool_name!r} already exists, please use"
+                    " a different name."
+                ),
+                style="red",
+            )
+            return
+
+        console.print(
+            f"Work pool {work_pool.name} created!",
+            style="green",
+        )
+        if PREFECT_UI_URL:
+            console.print(
+                "You see your new work pool in the UI at"
+                f" {PREFECT_UI_URL.value()}/work-pools/work-pool/{work_pool.name}"
+            )
+
+        deploy_script = (
+            "my_flow.deploy(work_pool_name='{work_pool.name}', image='my_image:tag')"
+        )
+        if not hasattr(self, "image"):
+            deploy_script = (
+                "my_flow.from_source(source='https://github.com/org/repo.git',"
+                f" entrypoint='flow.py:my_flow').deploy(work_pool_name='{work_pool.name}')"
+            )
+        console.print(
+            "\nYou can deploy a flow to this work pool by calling"
+            f" [blue].deploy[/]:\n\n\t{deploy_script}\n"
+        )
+        console.print(
+            "\nTo start a worker to execute flow runs in this work pool run:\n"
+        )
+        console.print(f"\t[blue]prefect worker start --pool {work_pool.name}[/]\n")
+
     @abc.abstractmethod
     async def run(
         self,

prefect/infrastructure/container.py

@@ -1,5 +1,6 @@
 import json
 import re
+import shlex
 import sys
 import urllib.parse
 import warnings
@@ -385,6 +386,57 @@ class DockerContainer(Infrastructure):
         finally:
             docker_client.close()
 
+    async def generate_work_pool_base_job_template(self):
+        from prefect.workers.utilities import (
+            get_default_base_job_template_for_infrastructure_type,
+        )
+
+        base_job_template = await get_default_base_job_template_for_infrastructure_type(
+            self.get_corresponding_worker_type()
+        )
+        if base_job_template is None:
+            return await super().generate_work_pool_base_job_template()
+        for key, value in self.dict(exclude_unset=True, exclude_defaults=True).items():
+            if key == "command":
+                base_job_template["variables"]["properties"]["command"]["default"] = (
+                    shlex.join(value)
+                )
+            elif key == "image_registry":
+                self.logger.warning(
+                    "Image registry blocks are not supported by Docker"
+                    " work pools. Please authenticate to your registry using"
+                    " the `docker login` command on your worker instances."
+                )
+            elif key in [
+                "type",
+                "block_type_slug",
+                "_block_document_id",
+                "_block_document_name",
+                "_is_anonymous",
+            ]:
+                continue
+            elif key == "image_pull_policy":
+                new_value = None
+                if value == ImagePullPolicy.ALWAYS:
+                    new_value = "Always"
+                elif value == ImagePullPolicy.NEVER:
+                    new_value = "Never"
+                elif value == ImagePullPolicy.IF_NOT_PRESENT:
+                    new_value = "IfNotPresent"
+
+                base_job_template["variables"]["properties"][key]["default"] = new_value
+            elif key in base_job_template["variables"]["properties"]:
+                base_job_template["variables"]["properties"][key]["default"] = value
+            else:
+                self.logger.warning(
+                    f"Variable {key!r} is not supported by Docker work pools. Skipping."
+                )
+
+        return base_job_template
+
+    def get_corresponding_worker_type(self):
+        return "docker"
+
     def _get_infrastructure_pid(self, container_id: str) -> str:
         """Generates a Docker infrastructure_pid string in the form of
         `<docker_host_base_url>:<container_id>`.

prefect/infrastructure/process.py

@@ -1,6 +1,7 @@
 import asyncio
 import contextlib
 import os
+import shlex
 import signal
 import socket
 import subprocess
@@ -247,6 +248,43 @@ class Process(Infrastructure):
     def _base_flow_run_command(self):
         return [get_sys_executable(), "-m", "prefect.engine"]
 
+    def get_corresponding_worker_type(self):
+        return "process"
+
+    async def generate_work_pool_base_job_template(self):
+        from prefect.workers.utilities import (
+            get_default_base_job_template_for_infrastructure_type,
+        )
+
+        base_job_template = await get_default_base_job_template_for_infrastructure_type(
+            self.get_corresponding_worker_type(),
+        )
+        assert (
+            base_job_template is not None
+        ), "Failed to generate default base job template for Process worker."
+        for key, value in self.dict(exclude_unset=True, exclude_defaults=True).items():
+            if key == "command":
+                base_job_template["variables"]["properties"]["command"]["default"] = (
+                    shlex.join(value)
+                )
+            elif key in [
+                "type",
+                "block_type_slug",
+                "_block_document_id",
+                "_block_document_name",
+                "_is_anonymous",
+            ]:
+                continue
+            elif key in base_job_template["variables"]["properties"]:
+                base_job_template["variables"]["properties"][key]["default"] = value
+            else:
+                self.logger.warning(
+                    f"Variable {key!r} is not supported by Process work pools."
+                    " Skipping."
+                )
+
+        return base_job_template
+
 
 class ProcessResult(InfrastructureResult):
     """Contains information about the final state of a completed process"""

prefect/infrastructure/provisioners/__init__.py

@@ -1,11 +1,15 @@
-from typing import Any, Dict, Optional, Protocol
+from typing import Any, Dict, Optional, Protocol, Type
 
 from prefect.client.orchestration import PrefectClient
 from .cloud_run import CloudRunPushProvisioner
+from .container_instance import ContainerInstancePushProvisioner
+from .ecs import ElasticContainerServicePushProvisioner
 import rich.console
 
 _provisioners = {
     "cloud-run:push": CloudRunPushProvisioner,
+    "azure-container-instance:push": ContainerInstancePushProvisioner,
+    "ecs:push": ElasticContainerServicePushProvisioner,
 }
 
 
@@ -29,7 +33,7 @@ class Provisioner(Protocol):
 
 def get_infrastructure_provisioner_for_work_pool_type(
     work_pool_type: str,
-) -> Provisioner:
+) -> Type[Provisioner]:
     """
     Retrieve an instance of the infrastructure provisioner for the given work pool type.
 

prefect/infrastructure/provisioners/cloud_run.py

@@ -259,6 +259,12 @@ class CloudRunPushProvisioner:
             ] = {"$ref": {"block_document_id": str(block_doc_id)}}
             progress.advance(task)
 
-        self._console.print("Infrastructure successfully provisioned!", style="green")
+        self._console.print(
+            (
+                f"Infrastructure successfully provisioned for '{work_pool_name}' work"
+                " pool!"
+            ),
+            style="green",
+        )
 
         return base_job_template_copy