Babysit command doesn't recover preemptied tpu

[jsrimr]

I created tpu named 'incremental-scale-width-1-0' and ran 'pu babysit incremental-scale-width-1-0'.

At first, everything seemed okay but some hours passed after without feeding tpu anything,
babysit command stopped displaying status view and threw

ValueError: No TPUs matched id 'incremental-scale-width-1-0'

Maybe is there any policy change in TPU preemtation? As I remember, even after TPU is preemptied, we could track the preemptied tpu's name but now it seems impossible and I think this is causing the problem.

[shawwn]

There was a change in policy some time ago: if you leave your TPU idle for four hours or so, they'll delete it.

Unfortunately, I wrote tpunicorn before that change, so it wasn't designed to handle deletions. The library assumes that you can always "figure out how to recreate the specified TPU." But clearly, if you delete a TPU and run pu recreate <that_tpu>, then tpunicorn has no way of knowing how to recreate it. It never occurred to me that I needed to handle random deletions.

There's a way to solve this, I think.

At the top of the babysit function,

tpunicorn/tpunicorn/program.py

Lines 399 to 410 in 58e51c8

	def babysit(ctx, tpu, zone, project, dry_run, interval, command):
	"""Checks TPU every INTERVAL seconds. Recreates the TPU if (and only if) the tpu has preempted."""
	# cmd = cli.get_command(ctx, 'babysit')
	# ctx = click.Context(cmd, parent=ctx, ignore_unknown_options=True)
	ctx.forward(recreate, yes=True, preempted=True)
	while True:
	time.sleep(interval)
	try:
	ctx.forward(recreate, yes=True, preempted=True)
	except:
	import traceback
	traceback.print_exc()

after ctx.forward(recreate, yes=True, preempted=True), write something like:

   tpu = tpunicorn.get_tpu(tpu=tpu, zone=zone, project=project)
   create = tpunicorn.create_tpu_command(tpu, zone=zone, project=project)

create contains a string that you can pass to os.system() to recreate the TPU. (It's the gcloud shell command you'd normally run from the command line.)

Then, inside the while True:, you could put:

    try:
      ...
    except ValueError as e:
      if 'No TPUs matched' in str(e):
        if 0 != do_step("The TPU was deleted by someone. Recreating...", create):
          click.echo("TPU failed to create (is the region out of capacity?)", err=True)
    else:
      ...

(do_step("foo", cmd) is equivalent to click.echo("foo"); os.system(cmd).)

I think that would do the trick.

Putting it all together, it might look like:

def babysit(ctx, tpu, zone, project, dry_run, interval, command):
  """Checks TPU every INTERVAL seconds. Recreates the TPU if (and only if) the tpu has preempted."""
  ctx.forward(recreate, yes=True, preempted=True)
   tpu = tpunicorn.get_tpu(tpu=tpu, zone=zone, project=project)
   create = tpunicorn.create_tpu_command(tpu, zone=zone, project=project)
  while True:
    time.sleep(interval)
    try:
      ctx.forward(recreate, yes=True, preempted=True)
    except ValueError as e:
      if 'No TPUs matched' in str(e):
        if 0 != do_step("The TPU was deleted by someone. Recreating...", create):
          click.echo("TPU failed to create (is the region out of capacity?)", err=True)
    except:
      import traceback
      traceback.print_exc()

Would you be interested in making these changes and testing it out? If not, no worries.