VertexAI on Prefect 2

I noticed there is nascent and experiemental support for VertexAI in prefect 2 via GitHub - PrefectHQ/prefect-gcp: Prefect integrations with Google Cloud Platform. but it’s extremely light on documentation and I can’t figure out how to use it.

I’ve used vertex on prefect 1. Conceptually what I don’t get is if we still need to have a dedicated vertex runner (started via ‘prefect agent vertex start’) which takes in flow runs off a queue and schedules appropriate vertex jobs with it. That command doesn’t work with the prefect 2 agent, and the documentation doesn’t mention this at all.

Any help would be appreciated.

The agent start command is prefect agent start -p "work-pool-name" so there is not a specific agent start command in Prefect 2.

Thanks, sounds like a normal local agent can schedule vertex flows - That’s helpful!

I don’t think I can get to the finish line without an example. The one in the docs[1] shows how to “echo Hello World” but not how to, as far as I can tell, get a flow run to execute on a vertex machine.

def vertex_ai_job_flow():
    gcp_credentials = GcpCredentials.load("MY_BLOCK")
    job = VertexAICustomTrainingJob(
        command=["echo", "hello world"],


I may be missing something obvious here. In prefect 1, there was a concept of a VertexRun which would run an entirely flow on a vertex instance.

Maybe I am expecting this to exist in the prefect2 module, but it doesn’t, and instead the VertexAICustomTrainingJob object is simply a way to kick off a normal vertex training job, and not a flow run. Is that right?

If so, what I’m looking for isn’t possible (yet)

Thanks for reporting this! The docs will be updated soon to improve this. In the meantime, if you follow Infrastructure - Prefect 2 - Coordinating the world's dataflows and replace KubernetesJob with VertexAI, it should work!

I’ll look forward to more docs, thanks.

I tried as you said, and got further, but the GCS VM which I setup to spin up the vertex jobs fails when it tries, complaining that, “KeyError: “No class found for dispatch key ‘vertex-ai-custom-training-job’ in registry for type ‘Block’.””

Let me know if this works:

Actually if that doesn’t resolve the issue, upgrading to prefect==2.8.2 should work.

Ah, thank you! Yes a key issue was my prefect agent didn’t have prefect-gcp available in it’s environment. I was imagining it was going to use my docker container (which has these python packages in it) to launch the vertex job, but that was mistaken. If I had installed the python package, as well as the other python packages it depends on, it likely would have worked.

Now I opted for another solution: I started my google cloud VM from my worker Docker image (as it had all the python modules already) with a startup command like:

prefect cloud login -k <key> -w "<username>/<workspace>"
prefect agent start -p <pool>

I don’t know if that’s a best practice, but its presently getting me going!