text: # Copyright 2021 The Flax Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  # Script for creating VM on Google cloud and run a Flax example inside.
  # See ./README.md for instructions.
  
  import datetime
  import os
  import re
  import subprocess
  import time
  from typing import Sequence
  
  from absl import app
  from absl import flags
  from absl import logging
  
  # General options.
  flags.DEFINE_bool(
      'dry_run',
      False,
      help='If set, then the command to launch the GCE instance will only be '
      'printed to stdout.')
  flags.DEFINE_bool(
      'connect',
      False,
      help='Same as --wait, but directly connect to VM once it is ready.')
  flags.DEFINE_bool(
      'wait', False,
      help='If set, then the script will wait until VM is ready. If VM_READY_CMD '
      'is set in environment, then that command will be executed once the VM '
      'is ready. Useful for sending a notification, e.g. "osascript"')
  
  # Machine configuration.
  flags.DEFINE_string('project', None, help='Name of the Google Cloud project.')
  flags.DEFINE_string('zone', None, help='Zone in which the VM will be created.')
  flags.DEFINE_string(
      'machine_type',
      None,
      help='Machine type to use for VM. See "gcloud compute machine-types list".')
  flags.DEFINE_string(
      'accelerator_type',
      '',
      help='Type of accelerator to use, or empty. '
      'See "gcloud compute accelerator-types list".'
  )
  flags.DEFINE_integer(
      'accelerator_count', 8, help='Number of accelerators to use.')
  
  # GCS configuration.
  flags.DEFINE_string(
      'gcs_workdir_base',
      None,
      help='GCS base directory for model output. The --workdir argument will be '
      'constructed from {gcs_workdir_base}/{example}/{name}/{timestamp} .')
  flags.DEFINE_string(
      'tfds_data_dir',
      '',
      help='Optional tfds data directory. This can be useful to prepare datasets '
      'on GCS and then point the jobs to this preloaded directory. Dataset will '
      'be downloaded from the web if not specified.')
  
  # Repo configuration.
  flags.DEFINE_string(
      'repo', 'https://github.com/google/flax', help='Git repository')
  flags.DEFINE_string('branch', 'master', help='Git repository')
  
  # Example configuration.
  flags.DEFINE_string(
      'example', None, help='Name of Flax example (e.g. "imagenet").')
  flags.DEFINE_string(
      'args',
      '',
      help='Any additional command line arguments for {example}_main.py, like '
      'for example --config. Note that --workdir will be provided by the '
      'script.')
  
  # Run configuration.
  flags.DEFINE_string(
      'name',
      None,
      help='Name of the experiment. Note that the provided name will be '
      'extended to {example}/{name}/{timestamp}')
  
  FLAGS = flags.FLAGS
  flags.mark_flags_as_required(
      ['project', 'zone', 'machine_type', 'gcs_workdir_base', 'example', 'name'])
  
  timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
  
  
  def generate_startup_file(vm_name: str) -> str:
    directory = os.path.dirname(os.path.abspath(__file__))
    startup_script_src = os.path.join(directory, 'startup_script.sh')
    startup_script_dst = os.path.join(directory, f'{vm_name}-startup_script.sh')
    assert not os.path.exists(startup_script_dst)
    with open(startup_script_src, encoding='utf8') as f:
      startup_script_content = f.read()
    for from_str, to_str in (
        ('__REPO__', FLAGS.repo),
        ('__BRANCH__', FLAGS.branch),
        ('__EXAMPLE__', FLAGS.example),
        ('__TIMESTAMP__', timestamp),
        ('__NAME__', FLAGS.name),
        ('__ARGS__', FLAGS.args),
        ('__GCS_WORKDIR_BASE__', FLAGS.gcs_workdir_base),
        ('__TFDS_DATA_DIR__', FLAGS.tfds_data_dir),
        ('__ACCELERATOR_TYPE__', FLAGS.accelerator_type),
    ):
      startup_script_content = startup_script_content.replace(from_str, to_str)
    with open(startup_script_dst, 'w', encoding='utf8') as f:
      f.write(startup_script_content)
    return startup_script_dst
  
  
  def launch_gce(*, vm_name: str, startup_script: str):
    # Note : Use `gcloud compute images list --project ml-images` to get a list
    # of available VM images.
    args = [
        'gcloud', 'compute', 'instances', 'create', vm_name,
        f'--project={FLAGS.project}', f'--zone={FLAGS.zone}',
        '--image=c1-deeplearning-tf-2-4-cu110-v20210329-debian-10',
        '--image-project=ml-images', f'--machine-type={FLAGS.machine_type}',
        '--scopes=cloud-platform,storage-full', '--boot-disk-size=256GB',
        '--boot-disk-type=pd-ssd', '--metadata=install-nvidia-driver=True',
        f'--metadata-from-file=startup-script={startup_script}'
    ]
    if FLAGS.accelerator_type and FLAGS.accelerator_count:
      args.extend([
          '--maintenance-policy=TERMINATE',
          f'--accelerator=type={FLAGS.accelerator_type},count={FLAGS.accelerator_count}',
      ])
  
    if FLAGS.dry_run:
      print()
      print('Would run the following command without --dry-run:')
      print()
      print(' \\\n    '.join(args))
      print()
      return
  
    print()
    print('Creating instance on GCE... This will take some minutes...')
    print()
    result = subprocess.run(args)
    if result.returncode:
      raise RuntimeError('Could not create VM!')
  
  
  def print_howto(login_args: Sequence[str]):
    print(f'''
  You can start/stop the instace via the web UI:
  https://console.cloud.google.com/compute/instances?project={FLAGS.project}
  
  Once the VM has started, you can login and connect to the training session:
  
  {' '.join(login_args)}
  
  Note that you can disconnect from the tmux session without stopping the training
  with the keystrokes 'CTRL-B A'. See "man tmux" for help about tmux.
  
  To observe the training via Tensorboard, simply run in your local computer:
  
  $ tensorboard --logdir={FLAGS.gcs_workdir_base}
  
  ''')
  
  
  def main(_):
    for name in ('repo', 'branch', 'example', 'name', 'gcs_workdir_base'):
      value = getattr(FLAGS, name)
      if re.match(r'[^\w:/_-]', value):
        raise ValueError(f'Invalid flag value: --{name}="{value}"')
    example_base_directory = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        os.path.pardir,
    )
    if not os.path.isdir(os.path.join(example_base_directory, FLAGS.example)):
      raise ValueError(f'Could not find --example={FLAGS.example}')
    if FLAGS.connect and FLAGS.dry_run:
      raise ValueError('Cannot --connect to VM with --dry_run')
  
    vm_name = '-'.join([
        'flax',
        FLAGS.example,
        timestamp,
    ])
    vm_name = re.sub(r'[^a-z0-9-]', '-', vm_name)
  
    startup_script = generate_startup_file(vm_name)
    launch_gce(vm_name=vm_name, startup_script=startup_script)
  
    login_args = [
        'gcloud',
        'compute',
        'ssh',
        '--project',
        FLAGS.project,
        '--zone',
        FLAGS.zone,
        vm_name,
        '--',
        '/sudo_tmux_a.sh',
    ]
  
    print('Your instance is being started...')
  
    print_howto(login_args)
  
    if FLAGS.connect or FLAGS.wait:
      login_true_args = login_args[:-1] + ['true']
      while True:
        try:
          result = subprocess.run(
              login_true_args, timeout=10, stderr=subprocess.PIPE)
          if result.returncode == 0:
            break
          stderr = result.stderr.decode('utf8')
          if 'connection refused' in stderr.lower():
            print('(Connection refused - waiting a little longer...)')
            time.sleep(20)
          else:
            raise ValueError(f'Unknown error: {stderr}')
        except subprocess.TimeoutExpired:
          print('(Timeout - waiting a little longer...)')
          time.sleep(20)
      if 'VM_READY_CMD' in os.environ:
        os.system(os.environ['VM_READY_CMD'])
      if FLAGS.connect:
        result = subprocess.run(login_args)
        # SSH session has cleared previous message, print it again.
        print_howto(login_args)
  
  
  if __name__ == '__main__':
    app.run(main)
  