From dc94f95fd3d0995bc0ecdd637c58ec7815638267 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 16 Oct 2026 22:36:24 -0700 Subject: [PATCH] Fix Flex-start TPU lease expiry lookup --- baseline/gpt2_small/speedrun.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/baseline/gpt2_small/speedrun.py b/baseline/gpt2_small/speedrun.py index 8d7f1444..54298ede 100644 --- a/baseline/gpt2_small/speedrun.py +++ b/baseline/gpt2_small/speedrun.py @@ -35,7 +35,8 @@ def lease(a, required_seconds=SUITE_SECONDS+1800): node=json.loads(run(['gcloud','compute','tpus','tpu-vm','describe',a.node,*flags],capture_output=True).stdout) queue_name=node.get('queuedResource','').rsplit('/',1)[-1] if not queue_name: raise RuntimeError('No queued-resource lease found') - queue=json.loads(run(['gcloud','compute','tpus','queued-resources','describe',queue_name,*flags],capture_output=True).stdout) + # Flex-start terminationTimestamp is exposed on the alpha queued-resource surface. + queue=json.loads(run(['gcloud','alpha','compute','tpus','queued-resources','describe',queue_name,*flags],capture_output=True).stdout) if node.get('state')!='READY' or queue.get('state',{}).get('state')!='ACTIVE' or node.get('acceleratorType')!='v5litepod-8': raise RuntimeError('Requires an active single-host v5e-8 allocation') expiries=[node.get('schedulingConfig',{}).get('terminationTimestamp'),queue.get('runDuration',{}).get('terminationTime')]