fix: Implement Persistent OS Login Authentication for TPU SSH (GoogleCloudPlatform#1172)

This change transitions the TPU SSH connection mechanism from ephemeral key injection to a persistent OS Login architecture. By leveraging long-lived SSH keys stored in the Service Account's OS Login profile, we eliminate the race conditions (409 Conflict) frequently encountered when running multiple concurrent TPU tasks in Airflow.

Author: AidenYu1673

xlml/utils/tpu.py

@@ -374,6 +374,16 @@ def ssh_tpu(
       client.get_node(name=os.path.join(node.parent, 'nodes', node.node_id))
       for node in queued_resource.tpu.node_spec
   ]
+  node_metadata = nodes[0].metadata
+  is_oslogin_enabled = node_metadata.get('enable-oslogin', '') == 'TRUE'
+
+  user = 'ml-auto-solutions'
+  if is_oslogin_enabled:
+    logging.info('Auto-detected OS Login enabled on node {nodes[0].name}..')
+    # get private key from  Airflow Variable
+    user = Variable.get('os-login-ssh-user')
+    ssh_keys.private = Variable.get('os-login-ssh-private-key')
+    ssh_keys.public = Variable.get('os-login-ssh-public-key')
 
   if all_workers:
     endpoints = itertools.chain.from_iterable(
@@ -397,7 +407,7 @@ def ssh_tpu(
       *ip_addresses,
       connect_kwargs={
           'auth_strategy': paramiko.auth_strategy.InMemoryPrivateKey(
-              'ml-auto-solutions', pkey
+              user, pkey
           ),
           # See https://stackoverflow.com/a/59453832
           'banner_timeout': 200,