Custom Prefect Flow Decorator: Handling Multiple Deployments with Dynamic Entrypoint Assignment

[abhishek815]

I created a custom flow decorator that wraps around the Prefect @flow decorator and adds specific functionalities tailored to my use cases, such as deploying to Prefect cloud. One issue I am running into is creating multiple deployments for a single flow. Since I cannot stack multiple decorators on my method, the other way I thought of accomplishing this is by creating multiple variable references to my custom @flow wrapper, for example:

from toolkit import custom_flow
from prefect import serve, settings, flow

def test_flow():
    print("hello")

deploy1 = custom_flow(
    deployment_name="v1",
    image_name="custom.azurecr.io/app/project:v1"
)(test_flow)

deploy2 = custom_flow(
    deployment_name="v2",
    image_name="custom.azurecr.io/app/project:v2"
)(test_flow)

However, when I run the deployments (locally and in the cloud), I get the following error

prefect.exceptions.MissingFlowError: Function with name 'test_flow' is not a flow. Make sure that it is decorated with '@flow'.

(running prefect 2.20.10)

I suspect this is because the entrypoint for the deployment is set to be test_flow , but in reality it should be deploy1 and deploy2 (depending on the deployment we're running). I tested this theory by manually configuring the entrypoint to for each deployment runner and this worked!

However, to create a robust solution in my custom flow decorator, I thought I would add a custom to_deployment function to that dynamically sets the correct entrypoint based on the flow object name (e.g., deploy1 or deploy2) instead of the base function name (test_flow). Is this the right approach or should I be handling this differently? I know we can define the entrypoint in the deployment yaml config but I want to handle this programmatically within my custom decorator to make the solution more dynamic and reusable, especially when managing multiple deployments for the same flow.

[desertaxle]

Hey @abhishek815! I think you should be able to get this to work, but you'll need to ensure that the object returned by custom_flow is a Flow object (or at least passes an instance check). Can you share the code for your custom_flow decorator?

[abhishek815]

Here's a slimmed down version of custom_flow:

def custom_flow(
    deployment_name: str,
    image_name: str,
    image_tag: Optional[str] = "dev",
    description: Optional[str] = None,
    flow_run_name: Optional[str] = None,
    environment_variables: Optional[Dict[str, str]] = None,
    work_pool_name: str = "default",
    retries: Optional[bool] = None,
    retry_delay_seconds: Optional[int] = None,
    task_runner: Optional[BaseTaskRunner] = ConcurrentTaskRunner,
    tags: Optional[list[str]] = None,
    custom_entrypoint_object: Optional[str] = None,
    **kwargs,
):
    tags = tags + ["custom-flow"] if tags else ["custom-flow"]

    def decorator(flow_func):
        job_variables = {}

        nonlocal environment_variables
        if environment_variables is None:
            environment_variables = {}

        extra_loggers = environment_variables.get("PREFECT_LOGGING_EXTRA_LOGGERS", "")
        if extra_loggers:
            environment_variables["PREFECT_LOGGING_EXTRA_LOGGERS"] = (
                f"{extra_loggers},project"
            )
        else:
            environment_variables["PREFECT_LOGGING_EXTRA_LOGGERS"] = "project"

        if "ENV" not in environment_variables:
            environment_variables["ENV"] = os.getenv("ENV", "dev")


        job_variables["env"] = environment_variables

        deployment_params = {
            "name": deployment_name,
            "work_pool_name": work_pool_name,
            "image": DeploymentImage(name=image_name, tag=image_tag),
            "job_variables": job_variables,
            "build": False,
            "push": False,
            "tags": tags,
        }

        deployment_params.update(**kwargs)
        wrapped_flow = flow(
            name=_get_flow_name(flow_func),
            retries=retries,
            retry_delay_seconds=retry_delay_seconds,
            flow_run_name=flow_run_name,
            task_runner=task_runner,
            description=description,
            log_prints=True,
        )(flow_func)

        flow_deployment_parameters[wrapped_flow] = deployment_params

        @sync_compatible
        async def custom_to_deployment(self, name: str, **kwargs):
            deployment = await Flow.to_deployment(self, name=name, **kwargs)

            file_path = inspect.getfile(flow_func)
            relative_path = os.path.relpath(file_path, os.getcwd())
            module_path = relative_path.replace(os.sep, "/").replace(".py", "")
            deployment.entrypoint = f"{module_path}.py:{custom_entrypoint_object}"
            return deployment

        if custom_entrypoint_object:
            # Add custom to_deployment method to the flow object if custom_entrypoint_object is provided
            wrapped_flow.to_deployment = custom_to_deployment.__get__(
                wrapped_flow, Flow
            )

        LOGGER.debug(
            f"Registered {wrapped_flow.name} for deployment with params {flow_deployment_parameters[wrapped_flow]}"
        )
        return wrapped_flow

    return decorator

When checking the type of deploy1 and deploy2 (from the original post), I see that they are of type Flow, but the to_deployment function sets the entry point of the deployment to be test_flow, which is not of type Flow. As a workaround, I added the custom_entrypoint_object parameter to my custom_flow to set the entrypoint to be deploy1 and deploy2 (which worked!)