<commit_before># Copyright 2020 Google LLC
#
# 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.

r'''
Simple script that launches GKE pods to act as client workers for a TPU pod.

This script strongly assumes that it is running in the context of another GKE
pod that had a TPU attached. As such, this script expects that you will provide
the current pod's name, UID, and TPU pod addresses via the downward API (the
TPU addresses are automatically given in $KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS).
This script does not clean up created resources. Instead, it sets
`metadata.ownerReferences` such that GKE's garbage collector will clean up
the created pods and services when the invoking pod is deleted.

Example:

```
python3 launch_k8s_workers.py \
    --name=pytorch-xla-pods \
    --image=gcr.io/xl-ml-test/pytorch-xla:nightly \
    --owner_name=$POD_NAME \
    --owner_uid=$POD_UID \
    --tpu=$KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS \
    --cpu=4 \
    --memory=4Gi \
    -- \
    python3 /pytorch/xla/test/test_train_mp_imagenet.py --fake_data
```
'''

import concurrent.futures
import os
import random
import string
import time

from absl import app
from absl import flags
from absl import logging
import kubernetes

FLAGS = flags.FLAGS

flags.DEFINE_string('name', None,
                    'Name of worker StatefulSet. Must be unique in `namespace`.')
flags.DEFINE_string('command', None, 'Command to run on each worker.')
flags.DEFINE_string('namespace', 'default',
                    'The namespace of the created StatefulSet.')
flags.DEFINE_string('tpu', os.getenv('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS', None),
                    'List of grpc:// addresses for the TPU. Defaults to '
                    '$KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS.')
flags.DEFINE_string('owner_name', None, 'Name of Pod that owns workers, if any.')
flags.DEFINE_string('owner_uid', None, 'UUID of Pod that owns workers, if any.')
flags.DEFINE_string('image', 'gcr.io/tpu-pytorch/xla:nightly',
                    'Docker image used for workers in created StatefulSet.')
flags.DEFINE_string('cpu', '4', 'CPU request for each worker.')
flags.DEFINE_string('memory', '4Gi', 'Memory request for each worker.')
flags.DEFINE_list('volumes', None,
                  'Comma-separated list of [PVC_NAME]:[MOUNT_DIR], where '
                  '[PVC_NAME] is the name of a Kubernetes PersistentVolumeClaim '
                  'and [MOUNT_PATH] is the directory where the PVC will be '
                  'mounted.')

def _format_env(envs):
  return [{'name': k, 'value': v} for k, v in envs.items()]

def main(argv):
  if FLAGS.command and len(argv) > 1:
    logging.warning('`--command` defined. Ignoring positional arguments.')
  elif not FLAGS.command and len(argv) > 1:
    FLAGS.command = ' '.join(argv[1:])
  elif not FLAGS.command:
    logging.error(
        'Must define `--command` or give command as a positional argument.')
    return 1

  logging.info('Command to distribute: `%s`', FLAGS.command)

  try:
    kubernetes.config.load_incluster_config()
  except:
    logging.warning('No Kubernetes cluster config. Using local kube config.')
    kubernetes.config.load_kube_config()

  k8s_client = kubernetes.client.CoreV1Api()

  random_suffix = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6))
  job_name = '{}-{}'.format(FLAGS.name, random_suffix)

  if FLAGS.owner_name:
    ownerReferences = [{
      'apiVersion': 'v1',
      'controller': True,
      'blockOwnerDeletion': True,
      'kind': 'Pod',
      'name': FLAGS.owner_name,
      'uid': FLAGS.owner_uid
    }]
  else:
    ownerReferences = None

  service_request = kubernetes.client.V1Service(**{
    'metadata': {
      'name': job_name,
      'ownerReferences': ownerReferences,
    },
    'spec': {
      'ports': [{
        'name': 'xrt-mesh',<commit_msg>Sanitize job name before launching jobs in launch_k8s_workers.py (#252)<commit_after># Copyright 2020 Google LLC
#
# 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.

r'''
Simple script that launches GKE pods to act as client workers for a TPU pod.

This script strongly assumes that it is running in the context of another GKE
pod that had a TPU attached. As such, this script expects that you will provide
the current pod's name, UID, and TPU pod addresses via the downward API (the
TPU addresses are automatically given in $KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS).
This script does not clean up created resources. Instead, it sets
`metadata.ownerReferences` such that GKE's garbage collector will clean up
the created pods and services when the invoking pod is deleted.

Example:

```
python3 launch_k8s_workers.py \
    --name=pytorch-xla-pods \
    --image=gcr.io/xl-ml-test/pytorch-xla:nightly \
    --owner_name=$POD_NAME \
    --owner_uid=$POD_UID \
    --tpu=$KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS \
    --cpu=4 \
    --memory=4Gi \
    -- \
    python3 /pytorch/xla/test/test_train_mp_imagenet.py --fake_data
```
'''

import concurrent.futures
import os
import random
import re
import string
import time

from absl import app
from absl import flags
from absl import logging
import kubernetes

FLAGS = flags.FLAGS

flags.DEFINE_string('name', None,
                    'Name of worker StatefulSet. Must be unique in `namespace`.')
flags.DEFINE_string('command', None, 'Command to run on each worker.')
flags.DEFINE_string('namespace', 'default',
                    'The namespace of the created StatefulSet.')
flags.DEFINE_string('tpu', os.getenv('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS', None),
                    'List of grpc:// addresses for the TPU. Defaults to '
                    '$KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS.')
flags.DEFINE_string('owner_name', None, 'Name of Pod that owns workers, if any.')
flags.DEFINE_string('owner_uid', None, 'UUID of Pod that owns workers, if any.')
flags.DEFINE_string('image', 'gcr.io/tpu-pytorch/xla:nightly',
                    'Docker image used for workers in created StatefulSet.')
flags.DEFINE_string('cpu', '4', 'CPU request for each worker.')
flags.DEFINE_string('memory', '4Gi', 'Memory request for each worker.')
flags.DEFINE_list('volumes', None,
                  'Comma-separated list of [PVC_NAME]:[MOUNT_DIR], where '
                  '[PVC_NAME] is the name of a Kubernetes PersistentVolumeClaim '
                  'and [MOUNT_PATH] is the directory where the PVC will be '
                  'mounted.')

def _format_env(envs):
  return [{'name': k, 'value': v} for k, v in envs.items()]

# Name must consist of lower case alphanumeric characters or '-'.
def _sanitize_job_name(name):
  return re.sub(r'[^a-z0-9\-]', '-', name.lower())

def main(argv):
  if FLAGS.command and len(argv) > 1:
    logging.warning('`--command` defined. Ignoring positional arguments.')
  elif not FLAGS.command and len(argv) > 1:
    FLAGS.command = ' '.join(argv[1:])
  elif not FLAGS.command:
    logging.error(
        'Must define `--command` or give command as a positional argument.')
    return 1

  logging.info('Command to distribute: `%s`', FLAGS.command)

  try:
    kubernetes.config.load_incluster_config()
  except:
    logging.warning('No Kubernetes cluster config. Using local kube config.')
    kubernetes.config.load_kube_config()

  k8s_client = kubernetes.client.CoreV1Api()

  random_suffix = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6))
  job_name = '{}-{}'.format(FLAGS.name, random_suffix)

  if FLAGS.owner_name:
    ownerReferences = [{
      'apiVersion': 'v1',
      'controller': True,
      'blockOwnerDeletion': True,
      'kind': 'Pod',
      'name': FLAGS.owner_name,
      'uid': FLAGS.owner_uid
    }]
  else:
    ownerReferences = None

  service_request = kubernetes.client.V1Service(**{
    'metadata': {
      'name': _sanitize_job_name(job_name),
      'ownerReferences': ownerReferences,
    },
    'spec': {
      'ports': [{
        'name': 'xrt-mesh',
{'fraction_non_alphanumeric': 0.11045854980799638, 'fraction_numerical': 0.006098938332956856, 'mean_word_length': 3.403281949278966, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4216819', 'n_tokens_mistral': 2847, 'n_tokens_neox': 2730, 'n_words': 887}