Unable to set Vertex runner disk size (but was able to in prefect 1)

[acgourley]

Expectation / Proposal

In prefect v1 we could specify the machine parameters (like boot disk size) of the vertex runner. This functionality seems missing from this v2 version. It's critical for my use case which deals with large data processing.

I am open to any work around which allows me to edit the "workerPoolSpecs" ultimately sent to GCS. I've tried editing my local aiplatform.py to ask for a 500gb disk as follows

from google.cloud.aiplatform_v1.types.machine_resources import DiskSpec

container_spec=container_spec, machine_spec=machine_spec, replica_count=1, disk_spec=DiskSpec(boot_disk_type="pd-ssd", boot_disk_size_gb=500)

But even if I deploy the VertexAICustomTrainingJob with save("name", overwrite=True) and re-rerun the deploy, new jobs on vertex still show this in the run parameters: "bootDiskSizeGb": 100 so it didn't seem to work.

• I would like to help contribute a pull request to resolve this!

[acgourley]

I've realized my hack fix works, I was just confused as to how to deploy it. I needed to edit that file in the environment of the prefect agent who was receiving the jobs and sending them to vertex.

So, for anyone looking for a work around, you can edit aiplatform.py locally and then also run the agent locally.

[acgourley]

As for a proper fix, would this be the approach?

1. Update VertexAICustomTrainingJob to have two more Fields

    boot_disk_type: str = Field(
        default="pd-ssd",
        title="Boot Disk Type",
        description="The machine boot disk type."
    )
   boot_disk_size_gb: str = Field(
        default="100",
        title="Boot Disk Size",
        description="The machine boot disk size in gigabytes."
    )

2. Update _build_job_spec with a line like:

container_spec=container_spec, machine_spec=machine_spec, replica_count=1, disk_spec=DiskSpec(boot_disk_type= boot_disk_type, boot_disk_size_gb=boot_disk_size_gb)

I am not clear on if you can just add new Fields to infrastructure blocks and if that will bounce through prefect cloud correctly.

[zanieb]

@acgourley you can add things to infrastructure blocks, your local version will be registered with the API when it is saved. You just need the agent to have the same version.

I'd be happy to review a pull request that adds the fields as you've described!

[acgourley]

Sure I'll throw a pull request together after I've tested it. Thanks for confirming the approach should work.