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

prefect/deployments/runner.py

@@ -29,7 +29,6 @@ Example:
 
 """
 
-import enum
 import importlib
 import tempfile
 from datetime import datetime, timedelta
@@ -37,34 +36,32 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
 from uuid import UUID
 
-import pendulum
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    PrivateAttr,
+    model_validator,
+)
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn, track
 from rich.table import Table
 
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import BaseModel, Field, PrivateAttr, root_validator, validator
-else:
-    from pydantic import BaseModel, Field, PrivateAttr, root_validator, validator
-
 from prefect._internal.concurrency.api import create_call, from_async
 from prefect._internal.schemas.validators import (
     reconcile_paused_deployment,
     reconcile_schedules_runner,
-    validate_automation_names,
 )
 from prefect.client.orchestration import get_client
-from prefect.client.schemas.objects import MinimalDeploymentSchedule
+from prefect.client.schemas.actions import DeploymentScheduleCreate
 from prefect.client.schemas.schedules import (
     SCHEDULE_TYPES,
     construct_schedule,
 )
 from prefect.deployments.schedules import (
-    FlexibleScheduleList,
-    create_minimal_deployment_schedule,
+    create_deployment_schedule_create,
 )
+from prefect.docker.docker_image import DockerImage
 from prefect.events import DeploymentTriggerTypes, TriggerTypes
 from prefect.exceptions import (
     ObjectNotFound,
@@ -72,24 +69,19 @@ from prefect.exceptions import (
 )
 from prefect.runner.storage import RunnerStorage
 from prefect.settings import (
-    PREFECT_DEFAULT_DOCKER_BUILD_NAMESPACE,
     PREFECT_DEFAULT_WORK_POOL_NAME,
     PREFECT_UI_URL,
 )
+from prefect.types.entrypoint import EntrypointType
 from prefect.utilities.asyncutils import sync_compatible
 from prefect.utilities.callables import ParameterSchema, parameter_schema
 from prefect.utilities.collections import get_from_dict, isiterable
 from prefect.utilities.dockerutils import (
-    PushError,
-    build_image,
-    docker_client,
-    generate_default_dockerfile,
     parse_image_tag,
-    split_repository_path,
 )
-from prefect.utilities.slugify import slugify
 
 if TYPE_CHECKING:
+    from prefect.client.types.flexible_schedule_list import FlexibleScheduleList
     from prefect.flows import Flow
 
 __all__ = ["RunnerDeployment"]
@@ -101,18 +93,6 @@ class DeploymentApplyError(RuntimeError):
     """
 
 
-class EntrypointType(enum.Enum):
-    """
-    Enum representing a entrypoint type.
-
-    File path entrypoints are in the format: `path/to/file.py:function_name`.
-    Module path entrypoints are in the format: `path.to.module.function_name`.
-    """
-
-    FILE_PATH = "file_path"
-    MODULE_PATH = "module_path"
-
-
 class RunnerDeployment(BaseModel):
     """
     A Prefect RunnerDeployment definition, used for specifying and building deployments.
@@ -126,7 +106,6 @@ class RunnerDeployment(BaseModel):
             are used only for organizational purposes. For delegating work to agents,
             see `work_queue_name`.
         schedule: A schedule to run this deployment on, once registered
-        is_schedule_active: Whether or not the schedule is active
         parameters: A dictionary of parameter values to pass to runs created from this
             deployment
         path: The path to the working directory for the workflow, relative to remote
@@ -144,8 +123,7 @@ class RunnerDeployment(BaseModel):
             available settings.
     """
 
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     name: str = Field(..., description="The name of the deployment.")
     flow_name: Optional[str] = Field(
@@ -161,17 +139,17 @@ class RunnerDeployment(BaseModel):
         default_factory=list,
         description="One of more tags to apply to this deployment.",
     )
-    schedules: Optional[List[MinimalDeploymentSchedule]] = Field(
+    schedules: Optional[List[DeploymentScheduleCreate]] = Field(
         default=None,
         description="The schedules that should cause this deployment to run.",
     )
-    schedule: Optional[SCHEDULE_TYPES] = None
+    concurrency_limit: Optional[int] = Field(
+        default=None,
+        description="The maximum number of concurrent runs of this deployment.",
+    )
     paused: Optional[bool] = Field(
         default=None, description="Whether or not the deployment is paused."
     )
-    is_schedule_active: Optional[bool] = Field(
-        default=None, description="DEPRECATED: Whether or not the schedule is active."
-    )
     parameters: Dict[str, Any] = Field(default_factory=dict)
     entrypoint: Optional[str] = Field(
         default=None,
@@ -184,7 +162,7 @@ class RunnerDeployment(BaseModel):
         description="The triggers that should cause this deployment to run.",
     )
     enforce_parameter_schema: bool = Field(
-        default=False,
+        default=True,
         description=(
             "Whether or not the Prefect API should enforce the parameter schema for"
             " this deployment."
@@ -232,16 +210,22 @@ class RunnerDeployment(BaseModel):
     def entrypoint_type(self) -> EntrypointType:
         return self._entrypoint_type
 
-    @validator("triggers", allow_reuse=True)
-    def validate_automation_names(cls, field_value, values):
+    @model_validator(mode="after")
+    def validate_automation_names(self):
         """Ensure that each trigger has a name for its automation if none is provided."""
-        return validate_automation_names(field_value, values)
+        trigger: Union[DeploymentTriggerTypes, TriggerTypes]
+        for i, trigger in enumerate(self.triggers, start=1):
+            if trigger.name is None:
+                trigger.name = f"{self.name}__automation_{i}"
+        return self
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def reconcile_paused(cls, values):
         return reconcile_paused_deployment(values)
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def reconcile_schedules(cls, values):
         return reconcile_schedules_runner(values)
 
@@ -294,6 +278,7 @@ class RunnerDeployment(BaseModel):
                 version=self.version,
                 paused=self.paused,
                 schedules=self.schedules,
+                concurrency_limit=self.concurrency_limit,
                 parameters=self.parameters,
                 description=self.description,
                 tags=self.tags,
@@ -301,7 +286,9 @@ class RunnerDeployment(BaseModel):
                 entrypoint=self.entrypoint,
                 storage_document_id=None,
                 infrastructure_document_id=None,
-                parameter_openapi_schema=self._parameter_openapi_schema.dict(),
+                parameter_openapi_schema=self._parameter_openapi_schema.model_dump(
+                    exclude_unset=True
+                ),
                 enforce_parameter_schema=self.enforce_parameter_schema,
             )
 
@@ -325,26 +312,25 @@ class RunnerDeployment(BaseModel):
                     f"Error while applying deployment: {str(exc)}"
                 ) from exc
 
-            if client.server_type.supports_automations():
-                try:
-                    # The triggers defined in the deployment spec are, essentially,
-                    # anonymous and attempting truly sync them with cloud is not
-                    # feasible. Instead, we remove all automations that are owned
-                    # by the deployment, meaning that they were created via this
-                    # mechanism below, and then recreate them.
-                    await client.delete_resource_owned_automations(
-                        f"prefect.deployment.{deployment_id}"
-                    )
-                except PrefectHTTPStatusError as e:
-                    if e.response.status_code == 404:
-                        # This Prefect server does not support automations, so we can safely
-                        # ignore this 404 and move on.
-                        return deployment_id
-                    raise e
+            try:
+                # The triggers defined in the deployment spec are, essentially,
+                # anonymous and attempting truly sync them with cloud is not
+                # feasible. Instead, we remove all automations that are owned
+                # by the deployment, meaning that they were created via this
+                # mechanism below, and then recreate them.
+                await client.delete_resource_owned_automations(
+                    f"prefect.deployment.{deployment_id}"
+                )
+            except PrefectHTTPStatusError as e:
+                if e.response.status_code == 404:
+                    # This Prefect server does not support automations, so we can safely
+                    # ignore this 404 and move on.
+                    return deployment_id
+                raise e
 
-                for trigger in self.triggers:
-                    trigger.set_deployment_id(deployment_id)
-                    await client.create_automation(trigger.as_automation())
+            for trigger in self.triggers:
+                trigger.set_deployment_id(deployment_id)
+                await client.create_automation(trigger.as_automation())
 
             return deployment_id
 
@@ -358,8 +344,8 @@ class RunnerDeployment(BaseModel):
         rrule: Optional[Union[Iterable[str], str]] = None,
         timezone: Optional[str] = None,
         schedule: Optional[SCHEDULE_TYPES] = None,
-        schedules: Optional[FlexibleScheduleList] = None,
-    ) -> Union[List[MinimalDeploymentSchedule], FlexibleScheduleList]:
+        schedules: Optional["FlexibleScheduleList"] = None,
+    ) -> Union[List[DeploymentScheduleCreate], "FlexibleScheduleList"]:
         """
         Construct a schedule or schedules from the provided arguments.
 
@@ -397,7 +383,7 @@ class RunnerDeployment(BaseModel):
         )
         if num_schedules > 1:
             raise ValueError(
-                "Only one of interval, cron, rrule, schedule, or schedules can be provided."
+                "Only one of interval, cron, rrule, or schedules can be provided."
             )
         elif num_schedules == 0:
             return []
@@ -417,7 +403,7 @@ class RunnerDeployment(BaseModel):
                 value = [value]
 
             return [
-                create_minimal_deployment_schedule(
+                create_deployment_schedule_create(
                     construct_schedule(
                         **{
                             schedule_type: v,
@@ -429,7 +415,7 @@ class RunnerDeployment(BaseModel):
                 for v in value
             ]
         else:
-            return [create_minimal_deployment_schedule(schedule)]
+            return [create_deployment_schedule_create(schedule)]
 
     def _set_defaults_from_flow(self, flow: "Flow"):
         self._parameter_openapi_schema = parameter_schema(flow)
@@ -450,15 +436,14 @@ class RunnerDeployment(BaseModel):
         cron: Optional[Union[Iterable[str], str]] = None,
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
-        schedules: Optional[FlexibleScheduleList] = None,
-        schedule: Optional[SCHEDULE_TYPES] = None,
-        is_schedule_active: Optional[bool] = None,
+        schedules: Optional["FlexibleScheduleList"] = None,
+        concurrency_limit: Optional[int] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
         tags: Optional[List[str]] = None,
         version: Optional[str] = None,
-        enforce_parameter_schema: bool = False,
+        enforce_parameter_schema: bool = True,
         work_pool_name: Optional[str] = None,
         work_queue_name: Optional[str] = None,
         job_variables: Optional[Dict[str, Any]] = None,
@@ -477,11 +462,6 @@ class RunnerDeployment(BaseModel):
             paused: Whether or not to set this deployment as paused.
             schedules: A list of schedule objects defining when to execute runs of this deployment.
                 Used to define multiple schedules or additional scheduling options like `timezone`.
-            schedule: A schedule object of when to execute runs of this flow. Used for
-                advanced scheduling options like timezone.
-            is_schedule_active: Whether or not to set the schedule for this deployment as active. If
-                not provided when creating a deployment, the schedule will be set as active. If not
-                provided when updating a deployment, the schedule's activation will not be changed.
             triggers: A list of triggers that should kick of a run of this flow.
             parameters: A dictionary of default parameter values to pass to runs of this flow.
             description: A description for the created deployment. Defaults to the flow's
@@ -502,7 +482,6 @@ class RunnerDeployment(BaseModel):
             interval=interval,
             cron=cron,
             rrule=rrule,
-            schedule=schedule,
             schedules=schedules,
         )
 
@@ -511,9 +490,8 @@ class RunnerDeployment(BaseModel):
         deployment = cls(
             name=Path(name).stem,
             flow_name=flow.name,
-            schedule=schedule,
             schedules=constructed_schedules,
-            is_schedule_active=is_schedule_active,
+            concurrency_limit=concurrency_limit,
             paused=paused,
             tags=tags or [],
             triggers=triggers or [],
@@ -586,15 +564,14 @@ class RunnerDeployment(BaseModel):
         cron: Optional[Union[Iterable[str], str]] = None,
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
-        schedules: Optional[FlexibleScheduleList] = None,
-        schedule: Optional[SCHEDULE_TYPES] = None,
-        is_schedule_active: Optional[bool] = None,
+        schedules: Optional["FlexibleScheduleList"] = None,
+        concurrency_limit: Optional[int] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
         tags: Optional[List[str]] = None,
         version: Optional[str] = None,
-        enforce_parameter_schema: bool = False,
+        enforce_parameter_schema: bool = True,
         work_pool_name: Optional[str] = None,
         work_queue_name: Optional[str] = None,
         job_variables: Optional[Dict[str, Any]] = None,
@@ -613,11 +590,6 @@ class RunnerDeployment(BaseModel):
             paused: Whether or not to set this deployment as paused.
             schedules: A list of schedule objects defining when to execute runs of this deployment.
                 Used to define multiple schedules or additional scheduling options like `timezone`.
-            schedule: A schedule object of when to execute runs of this flow. Used for
-                advanced scheduling options like timezone.
-            is_schedule_active: Whether or not to set the schedule for this deployment as active. If
-                not provided when creating a deployment, the schedule will be set as active. If not
-                provided when updating a deployment, the schedule's activation will not be changed.
             triggers: A list of triggers that should kick of a run of this flow.
             parameters: A dictionary of default parameter values to pass to runs of this flow.
             description: A description for the created deployment. Defaults to the flow's
@@ -643,17 +615,15 @@ class RunnerDeployment(BaseModel):
             interval=interval,
             cron=cron,
             rrule=rrule,
-            schedule=schedule,
             schedules=schedules,
         )
 
         deployment = cls(
             name=Path(name).stem,
             flow_name=flow.name,
-            schedule=schedule,
             schedules=constructed_schedules,
+            concurrency_limit=concurrency_limit,
             paused=paused,
-            is_schedule_active=is_schedule_active,
             tags=tags or [],
             triggers=triggers or [],
             parameters=parameters or {},
@@ -684,15 +654,14 @@ class RunnerDeployment(BaseModel):
         cron: Optional[Union[Iterable[str], str]] = None,
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
-        schedules: Optional[FlexibleScheduleList] = None,
-        schedule: Optional[SCHEDULE_TYPES] = None,
-        is_schedule_active: Optional[bool] = None,
+        schedules: Optional["FlexibleScheduleList"] = None,
+        concurrency_limit: Optional[int] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
         tags: Optional[List[str]] = None,
         version: Optional[str] = None,
-        enforce_parameter_schema: bool = False,
+        enforce_parameter_schema: bool = True,
         work_pool_name: Optional[str] = None,
         work_queue_name: Optional[str] = None,
         job_variables: Optional[Dict[str, Any]] = None,
@@ -711,11 +680,6 @@ class RunnerDeployment(BaseModel):
                 or a timedelta object. If a number is given, it will be interpreted as seconds.
             cron: A cron schedule of when to execute runs of this flow.
             rrule: An rrule schedule of when to execute runs of this flow.
-            schedule: A schedule object of when to execute runs of this flow. Used for
-                advanced scheduling options like timezone.
-            is_schedule_active: Whether or not to set the schedule for this deployment as active. If
-                not provided when creating a deployment, the schedule will be set as active. If not
-                provided when updating a deployment, the schedule's activation will not be changed.
             triggers: A list of triggers that should kick of a run of this flow.
             parameters: A dictionary of default parameter values to pass to runs of this flow.
             description: A description for the created deployment. Defaults to the flow's
@@ -738,7 +702,6 @@ class RunnerDeployment(BaseModel):
             interval=interval,
             cron=cron,
             rrule=rrule,
-            schedule=schedule,
             schedules=schedules,
         )
 
@@ -756,10 +719,9 @@ class RunnerDeployment(BaseModel):
         deployment = cls(
             name=Path(name).stem,
             flow_name=flow.name,
-            schedule=schedule,
             schedules=constructed_schedules,
+            concurrency_limit=concurrency_limit,
             paused=paused,
-            is_schedule_active=is_schedule_active,
             tags=tags or [],
             triggers=triggers or [],
             parameters=parameters or {},
@@ -781,74 +743,11 @@ class RunnerDeployment(BaseModel):
         return deployment
 
 
-class DeploymentImage:
-    """
-    Configuration used to build and push a Docker image for a deployment.
-
-    Attributes:
-        name: The name of the Docker image to build, including the registry and
-            repository.
-        tag: The tag to apply to the built image.
-        dockerfile: The path to the Dockerfile to use for building the image. If
-            not provided, a default Dockerfile will be generated.
-        **build_kwargs: Additional keyword arguments to pass to the Docker build request.
-            See the [`docker-py` documentation](https://docker-py.readthedocs.io/en/stable/images.html#docker.models.images.ImageCollection.build)
-            for more information.
-
-    """
-
-    def __init__(self, name, tag=None, dockerfile="auto", **build_kwargs):
-        image_name, image_tag = parse_image_tag(name)
-        if tag and image_tag:
-            raise ValueError(
-                f"Only one tag can be provided - both {image_tag!r} and {tag!r} were"
-                " provided as tags."
-            )
-        namespace, repository = split_repository_path(image_name)
-        # if the provided image name does not include a namespace (registry URL or user/org name),
-        # use the default namespace
-        if not namespace:
-            namespace = PREFECT_DEFAULT_DOCKER_BUILD_NAMESPACE.value()
-        # join the namespace and repository to create the full image name
-        # ignore namespace if it is None
-        self.name = "/".join(filter(None, [namespace, repository]))
-        self.tag = tag or image_tag or slugify(pendulum.now("utc").isoformat())
-        self.dockerfile = dockerfile
-        self.build_kwargs = build_kwargs
-
-    @property
-    def reference(self):
-        return f"{self.name}:{self.tag}"
-
-    def build(self):
-        full_image_name = self.reference
-        build_kwargs = self.build_kwargs.copy()
-        build_kwargs["context"] = Path.cwd()
-        build_kwargs["tag"] = full_image_name
-        build_kwargs["pull"] = build_kwargs.get("pull", True)
-
-        if self.dockerfile == "auto":
-            with generate_default_dockerfile():
-                build_image(**build_kwargs)
-        else:
-            build_kwargs["dockerfile"] = self.dockerfile
-            build_image(**build_kwargs)
-
-    def push(self):
-        with docker_client() as client:
-            events = client.api.push(
-                repository=self.name, tag=self.tag, stream=True, decode=True
-            )
-            for event in events:
-                if "error" in event:
-                    raise PushError(event["error"])
-
-
 @sync_compatible
 async def deploy(
     *deployments: RunnerDeployment,
     work_pool_name: Optional[str] = None,
-    image: Optional[Union[str, DeploymentImage]] = None,
+    image: Optional[Union[str, DockerImage]] = None,
     build: bool = True,
     push: bool = True,
     print_next_steps_message: bool = True,
@@ -870,7 +769,7 @@ async def deploy(
         work_pool_name: The name of the work pool to use for these deployments. Defaults to
             the value of `PREFECT_DEFAULT_WORK_POOL_NAME`.
         image: The name of the Docker image to build, including the registry and
-            repository. Pass a DeploymentImage instance to customize the Dockerfile used
+            repository. Pass a DockerImage instance to customize the Dockerfile used
             and build arguments.
         build: Whether or not to build a new image for the flow. If False, the provided
             image will be used as-is and pulled at runtime.
@@ -925,7 +824,7 @@ async def deploy(
 
     if image and isinstance(image, str):
         image_name, image_tag = parse_image_tag(image)
-        image = DeploymentImage(name=image_name, tag=image_tag)
+        image = DockerImage(name=image_name, tag=image_tag)
 
     try:
         async with get_client() as client:

prefect/deployments/schedules.py

@@ -1,42 +1,44 @@
-from typing import List, Optional, Sequence, Union, get_args
+from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union
 
-from prefect.client.schemas.objects import MinimalDeploymentSchedule
-from prefect.client.schemas.schedules import SCHEDULE_TYPES
+from prefect.client.schemas.actions import DeploymentScheduleCreate
+from prefect.client.schemas.schedules import is_schedule_type
 
-try:
-    from prefect.server.schemas.schedules import SCHEDULE_TYPES as SERVER_SCHEDULE_TYPES
+if TYPE_CHECKING:
+    from prefect.client.schemas.schedules import SCHEDULE_TYPES
 
-    SERVER_SCHEDULE_TYPES = get_args(SERVER_SCHEDULE_TYPES)
-except ImportError:
-    # `prefect-client` does not have access to the server schemas.
-    SERVER_SCHEDULE_TYPES = ()
+FlexibleScheduleList = Sequence[
+    Union[DeploymentScheduleCreate, dict[str, Any], "SCHEDULE_TYPES"]
+]
 
-FlexibleScheduleList = Sequence[Union[MinimalDeploymentSchedule, dict, SCHEDULE_TYPES]]
 
-
-def create_minimal_deployment_schedule(
-    schedule: SCHEDULE_TYPES,
+def create_deployment_schedule_create(
+    schedule: "SCHEDULE_TYPES",
     active: Optional[bool] = True,
-) -> MinimalDeploymentSchedule:
-    return MinimalDeploymentSchedule(
+    max_active_runs: Optional[int] = None,
+    catchup: bool = False,
+) -> DeploymentScheduleCreate:
+    """Create a DeploymentScheduleCreate object from common schedule parameters."""
+    return DeploymentScheduleCreate(
         schedule=schedule,
         active=active if active is not None else True,
+        max_active_runs=max_active_runs,
+        catchup=catchup,
     )
 
 
-def normalize_to_minimal_deployment_schedules(
-    schedules: Optional[FlexibleScheduleList],
-) -> List[MinimalDeploymentSchedule]:
-    normalized = []
+def normalize_to_deployment_schedule_create(
+    schedules: Optional["FlexibleScheduleList"],
+) -> List[DeploymentScheduleCreate]:
+    normalized: list[DeploymentScheduleCreate] = []
     if schedules is not None:
         for obj in schedules:
-            if isinstance(obj, get_args(SCHEDULE_TYPES)):
-                normalized.append(create_minimal_deployment_schedule(obj))
+            if is_schedule_type(obj):
+                normalized.append(create_deployment_schedule_create(obj))
             elif isinstance(obj, dict):
-                normalized.append(create_minimal_deployment_schedule(**obj))
-            elif isinstance(obj, MinimalDeploymentSchedule):
+                normalized.append(create_deployment_schedule_create(**obj))
+            elif isinstance(obj, DeploymentScheduleCreate):
                 normalized.append(obj)
-            elif isinstance(obj, SERVER_SCHEDULE_TYPES):
+            elif _is_server_schema(obj):
                 raise ValueError(
                     "Server schema schedules are not supported. Please use "
                     "the schedule objects from `prefect.client.schemas.schedules`"
@@ -44,7 +46,11 @@ def normalize_to_minimal_deployment_schedules(
             else:
                 raise ValueError(
                     "Invalid schedule provided. Must be a schedule object, a dict,"
-                    " or a MinimalDeploymentSchedule."
+                    "or a `DeploymentScheduleCreate` object"
                 )
 
     return normalized
+
+
+def _is_server_schema(obj: Any):
+    return obj.__class__.__module__.startswith("prefect.server.schemas")

prefect/deployments/steps/__init__.py

@@ -1,7 +1,6 @@
 from .core import run_step
 from .pull import (
     git_clone,
-    git_clone_project,
     set_working_directory,
     pull_from_remote_storage,
     pull_with_block,

prefect/deployments/steps/core.py

@@ -10,6 +10,7 @@ Whenever a step is run, the following actions are taken:
 - The step's function is called with the resolved inputs
 - The step's output is returned and used to resolve inputs for subsequent steps
 """
+
 import os
 import re
 import subprocess
@@ -21,6 +22,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 
 from prefect._internal.compatibility.deprecated import PrefectDeprecationWarning
 from prefect._internal.concurrency.api import Call, from_async
+from prefect._internal.integrations import KNOWN_EXTRAS_FOR_PACKAGES
 from prefect.logging.loggers import get_logger
 from prefect.settings import PREFECT_DEBUG_MODE
 from prefect.utilities.importtools import import_object
@@ -83,6 +85,11 @@ def _get_function_for_step(
             raise
 
     try:
+        packages = [
+            KNOWN_EXTRAS_FOR_PACKAGES.get(package, package)
+            for package in packages
+            if package
+        ]
         subprocess.check_call([sys.executable, "-m", "pip", "install", *packages])
     except subprocess.CalledProcessError:
         get_logger("deployments.steps.core").warning(