Remove ctpu tool usage as it's deprecated. (#521)

Author: sganeshb

perfzero/lib/perfzero/utils.py

@@ -398,22 +398,6 @@ def get_gpu_info():
   return _get_amd_gpu_info() if shutil.which("rocm-smi") \
     else _get_nvidia_gpu_info()
 
-
-def _install_tpu_tool():
-  """Installs the ctpu tool to managing cloud TPUs.
-
-  Follows the instructions here:
-  https://github.com/tensorflow/tpu/tree/master/tools/ctpu
-  """
-  if not os.path.exists('ctpu'):
-    logging.info('Installing TPU tool')
-    commands = [
-        'wget https://dl.google.com/cloud_tpu/ctpu/latest/linux/ctpu',
-        'chmod a+x ctpu',
-    ]
-    run_commands(commands)
-
-
 def setup_tpu(parameters):
   """Sets up a TPU with a given set of parameters.
 
@@ -424,26 +408,23 @@ def setup_tpu(parameters):
     True if an error occurs during setup.
   """
   try:
-    _install_tpu_tool()
-
+    base_cmd = 'gcloud compute tpus execution-groups create'
     args = [
+        '--tpu-only',
         '--name={}'.format(parameters.get('name')),
         '--project={}'.format(parameters.get('project')),
         '--zone={}'.format(parameters.get('zone')),
-        '--tpu-size={}'.format(parameters.get('size')),
+        '--accelerator-type={}'.format(parameters.get('size')),
         '--tf-version={}'.format(parameters.get('version')),
-        '--tpu-only',
-        '-noconf',
     ]
-    command = './ctpu up {}'.format(' '.join(args))
+    command = '{} {}'.format(base_cmd, ' '.join(args))
     logging.info('Setting up TPU: %s', command)
     exit_code, output = run_command(command)
     if exit_code != 0:
       logging.error('Error in setup with output: %s', output)
     return exit_code != 0
   except Exception:
     logging.error('Unable to setup TPU')
-    run_command('rm -f ctpu')
     sys.exit(1)
 
 
@@ -456,16 +437,15 @@ def cleanup_tpu(parameters):
   Returns:
     True if an error occurs during cleanup.
   """
-  _install_tpu_tool()
+  
+  base_cmd = 'gcloud compute tpus execution-groups delete'
 
   args = [
-      '--name={}'.format(parameters.get('name')),
+      '{}'.format(parameters.get('name')),
       '--project={}'.format(parameters.get('project')),
       '--zone={}'.format(parameters.get('zone')),
-      '--tpu-only',
-      '-noconf',
   ]
-  command = './ctpu delete {}'.format(' '.join(args))
+  command = '{} {}'.format(base_cmd, ' '.join(args))
   logging.info('Cleaning up TPU: %s', command)
   exit_code, output = run_command(command)
   if exit_code != 0: