Enable docker image feature with xpk

[RissyRan]

New PR.

Description

Enable docker image feature with xpk:

• Add this as an GKE example, and plan to add more examples for quick start
• Add TpuGkeTask to execute run_model and post_process task groups
• Add cluster configs under configs/cluster folder, so Airflow has the access to the created cluster with TPUs. It seems certificate-authority-data is safe to share based on this stackoverflow thread.
• Update TpuTask to TpuGceTask to distinguish TpuGkeTask, and similar for test config

A few TODO once b/311073979 is approved:

• Update GCP and cluster configs, and host clusters in cloud-ml-auto-solutions project
• Will have a separate PR to document how to provision cluster using XPK. It is quite straightforward with 1-2 commands

Tests

Please describe the tests that you ran on Cloud VM to verify changes.

Instruction and/or command lines to reproduce your tests: ...

• Upload codes to GCS bucket and test

List links for your tests (use go/shortn-gen for any internal link): ...

• Airflow test: link

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run one-shot tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[will-cromar]

If this test is actually going to be used, should it be in the example directory? If it is just an example, I would go even simpler and just print jax.device_count() to show the configuration is correct.

[RissyRan]

This is a good point. Yes, currently I just want to show an example. Since we are testing models, I was thinking probably it's better to have a hello world E2E example. Do you feel this is a complex one?

[will-cromar]

This is fine if the script prints the topology or number of devices. IMO the main part we want to confirm with the example is that we created the right machine and any metadata needed for initialization is plumbed through correctly to the runtime.