Add VertexAIService.launch_graph_store_job

[kmontemayor2-sc]

Scope of work done

Add ability to launch heterogeneous VAI clusters via new VertexAIService.launch_graph_store_job API

Where is the documentation for this feature?: In doc comments

Did you add automated tests or write a test plan? Added unit tests

Updated Changelog.md? NO

Ready for code review?: YES

[kmontemayor2-sc]

/unit_test

[kmontemayor2-sc]

/integration_tests

[kmontemayor2-sc]

/e2e_test

[github-actions]

GiGL Automation

@ 17:22:07UTC : 🔄 Unit Test started.

@ 17:28:21UTC : ❌ Workflow failed.
Please check the logs for more details.

[github-actions]

GiGL Automation

@ 17:22:12UTC : 🔄 Integration Test started.

@ 18:31:11UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 17:22:21UTC : 🔄 E2E Test started.

@ 18:42:28UTC : ✅ Workflow completed successfully.

[kmontemayor2-sc]

/unit_test

[github-actions]

GiGL Automation

@ 18:54:41UTC : 🔄 Unit Test started.

@ 19:46:25UTC : ✅ Workflow completed successfully.

[kmontemayor2-sc]

/unit_test

[kmontemayor2-sc]

/integration_test

[github-actions]

GiGL Automation

@ 16:30:50UTC : 🔄 Unit Test started.

@ 17:23:00UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 16:30:58UTC : 🔄 Integration Test started.

@ 17:32:29UTC : ✅ Workflow completed successfully.

[mkolodner-sc]

Thanks Kyle, generally LGTM -- just would prefer us to use existing VAI environment vars that are autopopulated instead of creating our own here

[kmontemayor2-sc]

Seeing as this lives in python/gigl/common/services/vertex_ai.py, I don't think we'd be kicking ourselves in the foot by relying on VAI environment variables in this file specifically.

The reliance isn't on this side, it's on the "application code" e.g. whatever code we run in these jobs). If we rely on the VAI CLUSTER_SPEC then any of that code is going to be harder to update in the future. We could expose this all through some utility but I'm not really sure what the downside to adding new env variables that we are in control of is here.

[mkolodner-sc]

Gotcha, thanks for the clarification. I don't feel too strongly here as that makes sense to me as well -- will defer to @svij-sc here. Approving to unblock on my end

[kmontemayor2-sc]

/unit_test

[kmontemayor2-sc]

/integration_test

[github-actions]

GiGL Automation

@ 20:50:34UTC : 🔄 Unit Test started.

@ 21:41:16UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 20:50:39UTC : 🔄 Integration Test started.

@ 20:56:51UTC : ✅ Workflow completed successfully.

[svij-sc]

Sorry for another round of comments, but seems minimal changes- althought test cycles for worker pool 2 might / switching around compute / storage cluster might take a little time.

I will go ahead and pre-emptively approve given that the rest of the comments will be address pre-merge.

[kmontemayor2-sc]

Could we try worker pools index 0, 1 continuing to be the main training/inference processes?
With either 2, or 3 being the storage pools? (Ideally 2; the the graph store is already the “server” lane \ and maps cleanly to paradigms off Parameter Servers / Reduction Server.

Sure, we can put the compute pool first.

For whatever reason, VAI rejects it when I put anything (gpu or not) into worker pool 2. idk why.

[kmontemayor2-sc]

For whatever reason, VAI rejects it when I put anything (gpu or not) into worker pool 2. idk why.

oh lol it randomly works now...