prefect-client 2.20.2__py3-none-any.whl → 3.0.0__py3-none-any.whl

prefect/infrastructure/__init__.py

@@ -1,27 +1,6 @@
-from prefect.infrastructure.base import Infrastructure, InfrastructureResult
-from prefect.infrastructure.container import DockerContainer, DockerContainerResult
-from prefect.infrastructure.kubernetes import (
-    KubernetesClusterConfig,
-    KubernetesImagePullPolicy,
-    KubernetesJob,
-    KubernetesJobResult,
-    KubernetesManifest,
-    KubernetesRestartPolicy,
-)
-from prefect.infrastructure.process import Process, ProcessResult
+"""
+2024-06-27: This surfaces an actionable error message for moved or removed objects in Prefect 3.0 upgrade.
+"""
+from prefect._internal.compatibility.migration import getattr_migration
 
-# Declare API
-__all__ = [
-    "DockerContainer",
-    "DockerContainerResult",
-    "Infrastructure",
-    "InfrastructureResult",
-    "KubernetesClusterConfig",
-    "KubernetesImagePullPolicy",
-    "KubernetesJob",
-    "KubernetesJobResult",
-    "KubernetesManifest",
-    "KubernetesRestartPolicy",
-    "Process",
-    "ProcessResult",
-]
+__getattr__ = getattr_migration(__name__)

prefect/infrastructure/base.py

@@ -1,323 +1,6 @@
 """
-DEPRECATION WARNING:
-
-This module is deprecated as of March 2024 and will not be available after September 2024.
-Infrastructure blocks have been replaced by workers, which offer enhanced functionality and better performance.
-
-For upgrade instructions, see https://docs.prefect.io/latest/guides/upgrade-guide-agents-to-workers/.
+2024-06-27: This surfaces an actionable error message for moved or removed objects in Prefect 3.0 upgrade.
 """
-import abc
-import shlex
-import warnings
-from typing import TYPE_CHECKING, Dict, List, Optional
-
-import anyio.abc
-
-from prefect._internal.compatibility.deprecated import deprecated_class
-from prefect._internal.compatibility.experimental import (
-    EXPERIMENTAL_WARNING,
-    ExperimentalFeature,
-    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, 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"
-
-
-if TYPE_CHECKING:
-    from prefect.client.schemas.objects import Deployment, Flow, FlowRun
-
-
-class InfrastructureResult(pydantic.BaseModel, abc.ABC):
-    identifier: str
-    status_code: int
-
-    def __bool__(self):
-        return self.status_code == 0
-
-
-@deprecated_class(
-    start_date="Mar 2024",
-    help="Use the `BaseWorker` class to create custom infrastructure integrations instead."
-    " Refer to the upgrade guide for more information:"
-    " https://docs.prefect.io/latest/guides/upgrade-guide-agents-to-workers/.",
-)
-class Infrastructure(Block, abc.ABC):
-    _block_schema_capabilities = ["run-infrastructure"]
-
-    type: str
-
-    env: Dict[str, Optional[str]] = pydantic.Field(
-        default_factory=dict,
-        title="Environment",
-        description="Environment variables to set in the configured infrastructure.",
-    )
-    labels: Dict[str, str] = pydantic.Field(
-        default_factory=dict,
-        description="Labels applied to the infrastructure for metadata purposes.",
-    )
-    name: Optional[str] = pydantic.Field(
-        default=None,
-        description="Name applied to the infrastructure for identification.",
-    )
-    command: Optional[List[str]] = pydantic.Field(
-        default=None,
-        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,
-        task_status: anyio.abc.TaskStatus = None,
-    ) -> InfrastructureResult:
-        """
-        Run the infrastructure.
-
-        If provided a `task_status`, the status will be reported as started when the
-        infrastructure is successfully created. The status return value will be an
-        identifier for the infrastructure.
-
-        The call will then monitor the created infrastructure, returning a result at
-        the end containing a status code indicating if the infrastructure exited cleanly
-        or encountered an error.
-        """
-        # Note: implementations should include `sync_compatible`
-
-    @abc.abstractmethod
-    def preview(self) -> str:
-        """
-        View a preview of the infrastructure that would be run.
-        """
-
-    @property
-    def logger(self):
-        return get_logger(f"prefect.infrastructure.{self.type}")
-
-    @property
-    def is_using_a_runner(self):
-        return self.command is not None and "prefect flow-run execute" in shlex.join(
-            self.command
-        )
-
-    @classmethod
-    def _base_environment(cls) -> Dict[str, str]:
-        """
-        Environment variables that should be passed to all created infrastructure.
-
-        These values should be overridable with the `env` field.
-        """
-        return get_current_settings().to_environment_variables(exclude_unset=True)
-
-    def prepare_for_flow_run(
-        self: Self,
-        flow_run: "FlowRun",
-        deployment: Optional["Deployment"] = None,
-        flow: Optional["Flow"] = None,
-    ) -> Self:
-        """
-        Return an infrastructure block that is prepared to execute a flow run.
-        """
-        if deployment is not None:
-            deployment_labels = self._base_deployment_labels(deployment)
-        else:
-            deployment_labels = {}
-
-        if flow is not None:
-            flow_labels = self._base_flow_labels(flow)
-        else:
-            flow_labels = {}
-
-        return self.copy(
-            update={
-                "env": {**self._base_flow_run_environment(flow_run), **self.env},
-                "labels": {
-                    **self._base_flow_run_labels(flow_run),
-                    **deployment_labels,
-                    **flow_labels,
-                    **self.labels,
-                },
-                "name": self.name or flow_run.name,
-                "command": self.command or self._base_flow_run_command(),
-            }
-        )
-
-    @staticmethod
-    def _base_flow_run_command() -> List[str]:
-        """
-        Generate a command for a flow run job.
-        """
-        if experiment_enabled("enhanced_cancellation"):
-            if (
-                PREFECT_EXPERIMENTAL_WARN
-                and PREFECT_EXPERIMENTAL_WARN_ENHANCED_CANCELLATION
-            ):
-                warnings.warn(
-                    EXPERIMENTAL_WARNING.format(
-                        feature="Enhanced flow run cancellation",
-                        group="enhanced_cancellation",
-                        help="",
-                    ),
-                    ExperimentalFeature,
-                    stacklevel=3,
-                )
-            return ["prefect", "flow-run", "execute"]
-
-        return ["python", "-m", "prefect.engine"]
-
-    @staticmethod
-    def _base_flow_run_labels(flow_run: "FlowRun") -> Dict[str, str]:
-        """
-        Generate a dictionary of labels for a flow run job.
-        """
-        return {
-            "prefect.io/flow-run-id": str(flow_run.id),
-            "prefect.io/flow-run-name": flow_run.name,
-            "prefect.io/version": prefect.__version__,
-        }
-
-    @staticmethod
-    def _base_flow_run_environment(flow_run: "FlowRun") -> Dict[str, str]:
-        """
-        Generate a dictionary of environment variables for a flow run job.
-        """
-        environment = {}
-        environment["PREFECT__FLOW_RUN_ID"] = str(flow_run.id)
-        return environment
-
-    @staticmethod
-    def _base_deployment_labels(deployment: "Deployment") -> Dict[str, str]:
-        labels = {
-            "prefect.io/deployment-name": deployment.name,
-        }
-        if deployment.updated is not None:
-            labels["prefect.io/deployment-updated"] = deployment.updated.in_timezone(
-                "utc"
-            ).to_iso8601_string()
-        return labels
+from prefect._internal.compatibility.migration import getattr_migration
 
-    @staticmethod
-    def _base_flow_labels(flow: "Flow") -> Dict[str, str]:
-        return {
-            "prefect.io/flow-name": flow.name,
-        }
+__getattr__ = getattr_migration(__name__)

prefect/infrastructure/provisioners/__init__.py

@@ -1,12 +1,14 @@
-from typing import Any, Dict, Optional, Protocol, Type
+from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol, Type
 
-from prefect.client.orchestration import PrefectClient
 from prefect.infrastructure.provisioners.modal import ModalPushProvisioner
 from .cloud_run import CloudRunPushProvisioner
 from .container_instance import ContainerInstancePushProvisioner
 from .ecs import ElasticContainerServicePushProvisioner
 import rich.console
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
 _provisioners = {
     "cloud-run:push": CloudRunPushProvisioner,
     "cloud-run-v2:push": CloudRunPushProvisioner,
@@ -29,7 +31,7 @@ class Provisioner(Protocol):
         self,
         work_pool_name: str,
         base_job_template: Dict[str, Any],
-        client: Optional[PrefectClient] = None,
+        client: Optional["PrefectClient"] = None,
     ) -> Dict[str, Any]:
         ...
 

prefect/infrastructure/provisioners/cloud_run.py

@@ -5,7 +5,7 @@ import tempfile
 from copy import deepcopy
 from pathlib import Path
 from textwrap import dedent
-from typing import Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional
 from uuid import UUID
 
 from anyio import run_process
@@ -17,7 +17,7 @@ from rich.prompt import Confirm
 from rich.syntax import Syntax
 
 from prefect.cli._prompts import prompt, prompt_select_from_table
-from prefect.client.orchestration import PrefectClient, ServerType
+from prefect.client.orchestration import ServerType
 from prefect.client.schemas.actions import BlockDocumentCreate
 from prefect.client.utilities import inject_client
 from prefect.exceptions import ObjectAlreadyExists
@@ -27,6 +27,9 @@ from prefect.settings import (
     update_current_profile,
 )
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
 
 class CloudRunPushProvisioner:
     def __init__(self):
@@ -215,7 +218,7 @@ class CloudRunPushProvisioner:
             ) from e
 
     async def _create_gcp_credentials_block(
-        self, block_document_name: str, key: dict, client: PrefectClient
+        self, block_document_name: str, key: dict, client: "PrefectClient"
     ) -> UUID:
         credentials_block_type = await client.read_block_type_by_slug("gcp-credentials")
 
@@ -242,7 +245,9 @@ class CloudRunPushProvisioner:
             )
             return block_doc.id
 
-    async def _create_provision_table(self, work_pool_name: str, client: PrefectClient):
+    async def _create_provision_table(
+        self, work_pool_name: str, client: "PrefectClient"
+    ):
         return Panel(
             dedent(
                 f"""\
@@ -268,7 +273,7 @@ class CloudRunPushProvisioner:
         )
 
     async def _customize_resource_names(
-        self, work_pool_name: str, client: PrefectClient
+        self, work_pool_name: str, client: "PrefectClient"
     ) -> bool:
         self._service_account_name = prompt(
             "Please enter a name for the service account",
@@ -294,7 +299,7 @@ class CloudRunPushProvisioner:
         self,
         work_pool_name: str,
         base_job_template: dict,
-        client: Optional[PrefectClient] = None,
+        client: Optional["PrefectClient"] = None,
     ) -> Dict[str, Any]:
         assert client, "Client injection failed"
         await self._verify_gcloud_ready()
@@ -399,7 +404,7 @@ class CloudRunPushProvisioner:
                         dedent(
                             f"""\
                         from prefect import flow
-                        from prefect.deployments import DeploymentImage
+                        from prefect.docker import DockerImage
 
 
                         @flow(log_prints=True)
@@ -411,7 +416,7 @@ class CloudRunPushProvisioner:
                             my_flow.deploy(
                                 name="my-deployment",
                                 work_pool_name="{work_pool_name}",
-                                image=DeploymentImage(
+                                image=DockerImage(
                                     name="my-image:latest",
                                     platform="linux/amd64",
                                 )

prefect/infrastructure/provisioners/container_instance.py

@@ -10,6 +10,7 @@ Classes:
     ContainerInstancePushProvisioner: A class for provisioning infrastructure using Azure Container Instances.
 
 """
+
 import json
 import random
 import shlex
@@ -18,7 +19,7 @@ import subprocess
 import time
 from copy import deepcopy
 from textwrap import dedent
-from typing import Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional
 from uuid import UUID
 
 from anyio import run_process
@@ -29,7 +30,6 @@ from rich.prompt import Confirm
 from rich.syntax import Syntax
 
 from prefect.cli._prompts import prompt, prompt_select_from_table
-from prefect.client.orchestration import PrefectClient
 from prefect.client.schemas.actions import BlockDocumentCreate
 from prefect.client.utilities import inject_client
 from prefect.exceptions import ObjectAlreadyExists, ObjectNotFound
@@ -38,6 +38,9 @@ from prefect.settings import (
     update_current_profile,
 )
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
 
 class AzureCLI:
     """
@@ -684,7 +687,7 @@ class ContainerInstancePushProvisioner:
         client_id: str,
         tenant_id: str,
         client_secret: str,
-        client: PrefectClient,
+        client: "PrefectClient",
     ) -> UUID:
         """
         Creates a credentials block for Azure Container Instance.
@@ -755,7 +758,7 @@ class ContainerInstancePushProvisioner:
             raise e
 
     async def _aci_credentials_block_exists(
-        self, block_name: str, client: PrefectClient
+        self, block_name: str, client: "PrefectClient"
     ) -> bool:
         """
         Checks if an ACI credentials block with the given name already exists.
@@ -782,7 +785,9 @@ class ContainerInstancePushProvisioner:
         else:
             return False
 
-    async def _create_provision_table(self, work_pool_name: str, client: PrefectClient):
+    async def _create_provision_table(
+        self, work_pool_name: str, client: "PrefectClient"
+    ):
         return Panel(
             dedent(
                 f"""\
@@ -809,7 +814,7 @@ class ContainerInstancePushProvisioner:
         )
 
     async def _customize_resource_names(
-        self, work_pool_name: str, client: PrefectClient
+        self, work_pool_name: str, client: "PrefectClient"
     ) -> bool:
         self._resource_group_name = prompt(
             "Please enter a name for the resource group",
@@ -851,7 +856,7 @@ class ContainerInstancePushProvisioner:
         self,
         work_pool_name: str,
         base_job_template: Dict[str, Any],
-        client: Optional[PrefectClient] = None,
+        client: Optional["PrefectClient"] = None,
     ) -> Dict[str, Any]:
         """
         Orchestrates the provisioning of Azure resources and setup for the push work pool.
@@ -1037,7 +1042,7 @@ class ContainerInstancePushProvisioner:
                     dedent(
                         f"""\
                         from prefect import flow
-                        from prefect.deployments import DeploymentImage
+                        from prefect.docker import DockerImage
 
 
                         @flow(log_prints=True)
@@ -1049,7 +1054,7 @@ class ContainerInstancePushProvisioner:
                             my_flow.deploy(
                                 name="my-deployment",
                                 work_pool_name="{work_pool_name}",
-                                image=DeploymentImage(
+                                image=DockerImage(
                                     name="my-image:latest",
                                     platform="linux/amd64",
                                 )

prefect/infrastructure/provisioners/ecs.py

@@ -9,7 +9,7 @@ import sys
 from copy import deepcopy
 from functools import partial
 from textwrap import dedent
-from typing import Any, Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
 
 import anyio
 from anyio import run_process
@@ -20,7 +20,6 @@ from rich.prompt import Confirm
 from rich.syntax import Syntax
 
 from prefect.cli._prompts import prompt
-from prefect.client.orchestration import PrefectClient
 from prefect.client.schemas.actions import BlockDocumentCreate
 from prefect.client.utilities import inject_client
 from prefect.exceptions import ObjectNotFound
@@ -31,6 +30,9 @@ from prefect.settings import (
 from prefect.utilities.collections import get_from_dict
 from prefect.utilities.importtools import lazy_import
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
 boto3 = lazy_import("boto3")
 
 current_console = contextvars.ContextVar("console", default=Console())
@@ -250,7 +252,7 @@ class CredentialsBlockResource:
 
     @inject_client
     async def requires_provisioning(
-        self, client: Optional[PrefectClient] = None
+        self, client: Optional["PrefectClient"] = None
     ) -> bool:
         if self._requires_provisioning is None:
             try:
@@ -280,7 +282,7 @@ class CredentialsBlockResource:
         self,
         base_job_template: Dict[str, Any],
         advance: Callable[[], None],
-        client: Optional[PrefectClient] = None,
+        client: Optional["PrefectClient"] = None,
     ):
         """
         Provisions an AWS credentials block.
@@ -365,7 +367,7 @@ class AuthenticationResource:
         work_pool_name: str,
         user_name: str = "prefect-ecs-user",
         policy_name: str = "prefect-ecs-policy",
-        credentials_block_name: str = None,
+        credentials_block_name: Optional[str] = None,
     ):
         self._user_name = user_name
         self._credentials_block_name = (
@@ -948,7 +950,7 @@ class ContainerRepositoryResource:
                         dedent(
                             f"""\
                                 from prefect import flow
-                                from prefect.deployments import DeploymentImage
+                                from prefect.docker import DockerImage
 
 
                                 @flow(log_prints=True)
@@ -960,7 +962,7 @@ class ContainerRepositoryResource:
                                     my_flow.deploy(
                                         name="my-deployment",
                                         work_pool_name="{self._work_pool_name}",
-                                        image=DeploymentImage(
+                                        image=DockerImage(
                                             name="{self._repository_name}:latest",
                                             platform="linux/amd64",
                                         )
@@ -1128,7 +1130,7 @@ class ElasticContainerServicePushProvisioner:
         work_pool_name: str,
         user_name: str = "prefect-ecs-user",
         policy_name: str = "prefect-ecs-policy",
-        credentials_block_name: str = None,
+        credentials_block_name: Optional[str] = None,
         cluster_name: str = "prefect-ecs-cluster",
         vpc_name: str = "prefect-ecs-vpc",
         ecs_security_group_name: str = "prefect-ecs-security-group",

prefect/infrastructure/provisioners/modal.py

@@ -2,20 +2,23 @@ import importlib
 import shlex
 import sys
 from copy import deepcopy
-from typing import Any, Dict, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
 
 from anyio import run_process
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn
 from rich.prompt import Confirm
 
-from prefect.client.orchestration import PrefectClient
 from prefect.client.schemas.actions import BlockDocumentCreate
 from prefect.client.schemas.objects import BlockDocument
 from prefect.client.utilities import inject_client
 from prefect.exceptions import ObjectNotFound
 from prefect.utilities.importtools import lazy_import
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
+
 modal = lazy_import("modal")
 
 
@@ -24,7 +27,7 @@ class ModalPushProvisioner:
     A infrastructure provisioner for Modal push work pools.
     """
 
-    def __init__(self, client: Optional[PrefectClient] = None):
+    def __init__(self, client: Optional["PrefectClient"] = None):
         self._console = Console()
 
     @property
@@ -91,7 +94,7 @@ class ModalPushProvisioner:
         block_document_name: str,
         modal_token_id: str,
         modal_token_secret: str,
-        client: PrefectClient,
+        client: "PrefectClient",
     ) -> BlockDocument:
         """
         Creates a ModalCredentials block containing the provided token ID and secret.
@@ -142,7 +145,7 @@ class ModalPushProvisioner:
         self,
         work_pool_name: str,
         base_job_template: Dict[str, Any],
-        client: Optional[PrefectClient] = None,
+        client: Optional["PrefectClient"] = None,
     ) -> Dict[str, Any]:
         """
         Provisions resources necessary for a Modal push work pool.