[Feature Request] Support for Shielded TPU VMs

[christianepeters]

The ./cloud/tool provision script fails to create a TPU VM in Google Cloud projects where the constraints/compute.requireShieldedVm org policy is enforced.

Error Log

google.api_core.exceptions.InvalidArgument: 400 Cloud TPU received a bad request. Constraint constraints/compute.requireShieldedVm violated. Please create TPUs with Secure Boot enabled or remove the Organization Policy Constraint.

Proposed Solution: Provide an option to create Shielded TPU VMs

Disabling the org policy is an option, but not recommended, and in some environments won't be possible given org admin settings.

Instead adapt the tool script to provide an optional flag to create a policy-compliant TPU VM with Secure Boot enabled. The default behavior (creating a non-shielded VM) should remain unchanged to avoid impacting users who do not have this policy constraint.

This involves adding a --shielded flag. If the flag is passed, the script should enable Secure Boot on the TPU VM. If the flag is not provided, the script should behave as it does today.

1. Add an optional --shielded flag to the provision command:

# In class Tool, method provision(self):
         ...
         parser.add_argument(
            "--shielded",
            action="store_true",
            default=False, # The flag is off by default
            help="Enable Shielded VM options (Secure Boot) for the TPU.",
         )

2. Conditionally enable Secure Boot in cloud/provision.py:

This change uses the shielded_instance_config field in the TPU API's Node object.

# In cloud/provision.py, function provision(...):

    from google.cloud import tpu_v2

    # ...
    node = tpu_v2.Node(
        accelerator_type=accelerator_type,
        runtime_version=runtime_version,
        # ... other parameters
    )

    if shielded: # New optional parameter passed from the tool
        node.shielded_instance_config = tpu_v2.ShieldedInstanceConfig(
            enable_secure_boot=True
        )

    op = core.client.create_node(
        parent=core.parent,
        node_id=name,
        node=node,
    )
    ...

This change would allow the tool to function correctly in security-conscious environments by creating a VM with Secure Boot, a core feature of what Google Cloud describes as a Shielded VM.

Test

Toggle the org policy and then deploy with the new flag.