feat: add `kill_infrastructure` implementation for vertex ai worker #213

parkedwards · 2023-09-21T22:12:46Z

Closes #212

Followup to #211

Example

Screenshots

Checklist

References any related issue by including "Closes #" or "Closes ".
- If no issue exists and your change is not a small fix, please create an issue first.
Includes tests or only affects documentation.
Passes pre-commit checks.
- Run pre-commit install && pre-commit run --all locally for formatting and linting.
Includes screenshots of documentation updates.
- Run mkdocs serve view documentation locally.
Summarizes PR's changes in CHANGELOG.md

parkedwards · 2023-09-22T03:03:33Z

prefect_gcp/workers/vertex.py

+
+        # Vertex will include an error message upon valid
+        # flow cancellations, so we'll avoid raising an error in that case
+        if error_msg and "CANCELED" not in error_msg:


had to exclude cancellation messages, as this will throw the runtime error

parkedwards · 2023-09-22T03:05:03Z

tests/test_vertex_worker.py

+        mock = job_config.credentials.job_service_client.create_custom_job
+        # the CancelCustomJobRequest class seems to reject a MagicMock value
+        # so here, we'll use a SimpleNamespace as the mocked return values
+        mock.return_value = SimpleNamespace(


im not fully understanding why, but returning a plain MagicMock here with the correct name and state keys would cause the test to fail, somewhere in the CancelCustomJobRequest class

self = <[AttributeError('Unknown field for CancelCustomJobRequest: _pb') raised in repr()] CancelCustomJobRequest object at 0x168bb4490> mapping = {'name': <MagicMock name='mock_name.name' id='5749201248'>}, ignore_unknown_fields = False, kwargs = {'name': <MagicMock name='mock_name.name' id='5749201248'>} params = {'name': <MagicMock name='mock_name.name' id='5749201248'>}, marshal = <proto.marshal.marshal.Marshal object at 0x106f4c790>, key = 'name' value = <MagicMock name='mock_name.name' id='5749201248'>, pb_value = <MagicMock name='mock_name.name' id='5749201248'>

so, here i just use a SimpleNamespace to return the right key values that can be access via dot notation when it's being passed into CancelCustomJobRequest downstream

parkedwards · 2023-09-22T14:38:28Z

prefect_gcp/workers/vertex.py

        title="Accelerator Count",
        description=(
            "The number of accelerators to attach to the machine. "
            "See https://cloud.google.com/vertex-ai/docs/reference/rest/v1/MachineSpec"
        ),
-        default=0,
+        example=1,


this is unrelated to the kill infra path, but something i found while testing this further

both the accelerator_* params are optional, and they have an XOR relationship - both have to exist, or both have to be null. so, i swapped these to be optional + removed the keys from the job body validation

For these to be optional in the UI, they'll need a default value of None.

im cool to add the default None, but with the current configuration, the UI looks like this (the fields are empty btw, the screenshot shows the example text):

Ok gotcha. Do you get errors in the worker when you don't fill in those fields? I'd expect Pydantic to complain that no value is provided if there's no default value.

that's so weird, i would expect that as well. but i just created a new work pool and started the worker, and:

$ prefect work-pool create 'a-new-pool' --type vertex-ai Created work pool 'a-new-pool'. $ prefect worker start -p a-new-pool Discovered type 'vertex-ai' for work pool 'a-new-pool'. Worker 'VertexAIWorker b1b27bbe-1070-4f5a-9ccc-f63a6e7adad1' started!

here are the work pool variables in the UI

You wouldn't see the error until a flow run is prepared to be run because that's when a configuration object is created.

yeah i didn't see it at all during my e2e testing, where i created a bunch of vertex jobs. does it make a difference that these properties are on the variables class, not the configuration?

btw i added the default None values in, this thread is more for my understanding / sanity

🤦‍♂️ I didn't realize this is the variables class. No objects are ever created from this class, so you wouldn't see that error. It all makes sense now.

ok whew, sanity restored

desertaxle

Overall, looks great! I have one small comment and I'm going to push some changelog updates to this branch so that we can release once this is merged!

desertaxle · 2023-09-22T14:51:38Z

prefect_gcp/workers/vertex.py

        title="Accelerator Count",
        description=(
            "The number of accelerators to attach to the machine. "
            "See https://cloud.google.com/vertex-ai/docs/reference/rest/v1/MachineSpec"
        ),
-        default=0,
+        example=1,


For these to be optional in the UI, they'll need a default value of None.

desertaxle

LGTM! 🚀

parkedwards added 2 commits September 21, 2023 14:17

feat: add kill_infrastructure implementation for vertex

d48b6eb

add unit tests

c11c3f6

parkedwards requested a review from a team as a code owner September 21, 2023 22:12

fix tests

35996c4

parkedwards commented Sep 22, 2023

View reviewed changes

parkedwards requested a review from desertaxle September 22, 2023 03:05

maybe assert like this

9d34c01

parkedwards commented Sep 22, 2023

View reviewed changes

desertaxle reviewed Sep 22, 2023

View reviewed changes

desertaxle and others added 2 commits September 22, 2023 10:05

Updates changelog for upcoming release

e231cda

add default Nones to optional props

0c5e648

desertaxle approved these changes Sep 22, 2023

View reviewed changes

parkedwards merged commit a21aca4 into main Sep 22, 2023
10 checks passed

parkedwards deleted the feat/vertex-worker-kill-infra branch September 22, 2023 16:36

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: add `kill_infrastructure` implementation for vertex ai worker #213

feat: add `kill_infrastructure` implementation for vertex ai worker #213

parkedwards commented Sep 21, 2023 •

edited

Loading

parkedwards Sep 22, 2023

parkedwards Sep 22, 2023 •

edited

Loading

parkedwards Sep 22, 2023

desertaxle Sep 22, 2023

parkedwards Sep 22, 2023 •

edited

Loading

desertaxle Sep 22, 2023

parkedwards Sep 22, 2023

desertaxle Sep 22, 2023

parkedwards Sep 22, 2023

parkedwards Sep 22, 2023

desertaxle Sep 22, 2023

parkedwards Sep 22, 2023

desertaxle left a comment

desertaxle Sep 22, 2023

desertaxle left a comment

feat: add kill_infrastructure implementation for vertex ai worker #213

feat: add kill_infrastructure implementation for vertex ai worker #213

Conversation

parkedwards commented Sep 21, 2023 • edited Loading

Example

Screenshots

Checklist

Choose a reason for hiding this comment

parkedwards Sep 22, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

parkedwards Sep 22, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

desertaxle left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

desertaxle left a comment

Choose a reason for hiding this comment

feat: add `kill_infrastructure` implementation for vertex ai worker #213

feat: add `kill_infrastructure` implementation for vertex ai worker #213

parkedwards commented Sep 21, 2023 •

edited

Loading

parkedwards Sep 22, 2023 •

edited

Loading

parkedwards Sep 22, 2023 •

edited

Loading