# 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
#
#     https://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.

"apiVersion": "batch/v1beta1"
"kind": "CronJob"
"metadata":
  "name": "pt-r1.7-roberta-pre-func-v3-32"
  "namespace": "automated"
"spec":
  "concurrencyPolicy": "Forbid"
  "jobTemplate":
    "spec":
      "activeDeadlineSeconds": 3600
      "backoffLimit": 0
      "template":
        "metadata":
          "annotations":
            "tf-version.cloud-tpus.google.com": "pytorch-1.7"
        "spec":
          "containers":
          - "env":
            - "name": "POD_NAME"
              "valueFrom":
                "fieldRef":
                  "fieldPath": "metadata.name"
            - "name": "POD_NAMESPACE"
              "valueFrom":
                "fieldRef":
                  "fieldPath": "metadata.namespace"
            "image": "gcr.io/xl-ml-test/health-monitor:stable"
            "imagePullPolicy": "Always"
            "name": "monitor"
          - "args":
            - "/bin/bash"
            - "-c"
            - |
              set -u
              set -e
              set -x
              
              pip install --editable tpu-examples/deps/fairseq
              python3 \
                /tpu-examples/deps/fairseq/train.py \
                /datasets/wikitext-103 \
                --task=masked_lm --criterion=masked_lm \
                --arch=roberta_base --sample-break-mode=complete \
                --tokens-per-sample=512 \
                --optimizer=adam \
                --adam-betas='(0.9,0.98)' \
                --adam-eps=1e-6 \
                --clip-norm=0.0 \
                --lr-scheduler=polynomial_decay \
                --lr=0.0005 \
                --warmup-updates=10000 \
                --dropout=0.1 \
                --attention-dropout=0.1 \
                --weight-decay=0.01 \
                --update-freq=16 \
                --log-format=simple \
                --train-subset=train \
                --valid-subset=valid \
                --num_cores=8 \
                --metrics_debug \
                --save-dir=checkpoints \
                --log_steps=30 \
                --skip-invalid-size-inputs-valid-test \
                --suppress_loss_report \
                --input_shapes 16x512 18x480 21x384 \
                --max-epoch=1
              
            "command":
            - "python3"
            - "launch_k8s_workers.py"
            - "--name=$(JOB_NAME)"
            - "--image=gcr.io/xl-ml-test/pytorch-xla:r1.7"
            - "--owner_name=$(POD_NAME)"
            - "--owner_uid=$(POD_UID)"
            - "--namespace=$(POD_NAMESPACE)"
            - "--tpu=$(KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS)"
            - "--cpu=4"
            - "--memory=4Gi"
            - "--volumes=pytorch-datasets-claim:/datasets"
            - "--"
            "env":
            - "name": "POD_NAME"
              "valueFrom":
                "fieldRef":
                  "fieldPath": "metadata.name"
            - "name": "POD_UID"
              "valueFrom":
                "fieldRef":
                  "fieldPath": "metadata.uid"
            - "name": "POD_NAMESPACE"
              "valueFrom":
                "fieldRef":
                  "fieldPath": "metadata.namespace"
            - "name": "JOB_NAME"
              "valueFrom":
                "fieldRef":
                  "fieldPath": "metadata.labels['job-name']"
            - "name": "MODEL_DIR"
              "value": "$(OUTPUT_BUCKET)/pt-r1.7/roberta-pre/func/v3-32/$(JOB_NAME)"
            "envFrom":
            - "configMapRef":
                "name": "gcs-buckets"
            "image": "gcr.io/xl-ml-test/pytorch-pods:nightly"
            "imagePullPolicy": "Always"
            "name": "train"
            "resources":
              "limits":
                "cloud-tpus.google.com/v3": 32
              "requests":
                "cpu": "9.0"
                "ephemeral-storage": "10Gi"
                "memory": "30Gi"
            "volumeMounts":
            - "mountPath": "/datasets"
              "name": "pytorch-datasets-claim"
              "readOnly": true
          "initContainers":
          - "env":
            - "name": "POD_NAME"
              "valueFrom":
                "fieldRef":
                  "fieldPath": "metadata.name"
            - "name": "POD_UID"
              "valueFrom":
                "fieldRef":
                  "fieldPath": "metadata.uid"
            - "name": "POD_NAMESPACE"
              "valueFrom":
                "fieldRef":
                  "fieldPath": "metadata.namespace"
            - "name": "JOB_NAME"
              "valueFrom":
                "fieldRef":
                  "fieldPath": "metadata.labels['job-name']"
            - "name": "MODEL_DIR"
              "value": "$(OUTPUT_BUCKET)/pt-r1.7/roberta-pre/func/v3-32/$(JOB_NAME)"
            - "name": "METRIC_CONFIG"
              "value": |
                {
                 "metric_collection_config": {
                  "default_aggregation_strategies": [
                   "final"
                  ],
                  "tags_to_ignore": [
                   "LearningRate"
                  ],
                  "write_to_bigquery": true
                 },
                 "regression_test_config": null,
                 "test_name": "pt-r1.7-roberta-pre-func-v3-32"
                }
            "envFrom":
            - "configMapRef":
                "name": "gcs-buckets"
            "image": "gcr.io/xl-ml-test/publisher:stable"
            "imagePullPolicy": "Always"
            "name": "publisher"
          "nodeSelector":
            "tpu-available": "true"
          "restartPolicy": "Never"
          "serviceAccountName": "pytorch-xla-pods"
          "volumes":
          - "name": "pytorch-datasets-claim"
            "persistentVolumeClaim":
              "claimName": "pytorch-datasets-claim"
  "schedule": "0 22 * * *"
  "successfulJobsHistoryLimit": 1