VertexAI on Prefect 2

Help
,
1

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.

2

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

3

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.


@flow
def vertex_ai_job_flow():
    gcp_credentials = GcpCredentials.load("MY_BLOCK")
    job = VertexAICustomTrainingJob(
        command=["echo", "hello world"],
        region="us-east1",
        image="us-docker.pkg.dev/cloudrun/container/job:latest",
        gcp_credentials=gcp_credentials,
    )
    job.run()

vertex_ai_job_flow()
4

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)

5

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!

6

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’.””

7

Let me know if this works:

8

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

9

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!