# 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-1.5-fs-transformer-convergence-v3-8"
  "namespace": "automated"
"spec":
  "concurrencyPolicy": "Forbid"
  "jobTemplate":
    "spec":
      "activeDeadlineSeconds": 18000
      "backoffLimit": 2
      "template":
        "metadata":
          "annotations":
            "tf-version.cloud-tpus.google.com": "pytorch-1.5"
        "spec":
          "containers":
          - "args":
            - "python3"
            - "/tpu-examples/deps/fairseq/train.py"
            - "/datasets/wmt18_en_de_bpej32k"
            - "--tensorboard-logdir=$(MODEL_DIR)"
            - "--metrics_debug"
            - "--arch=transformer_vaswani_wmt_en_de_big"
            - "--max-target-positions=64"
            - "--attention-dropout=0.1"
            - "--no-progress-bar"
            - "--no-save"
            - "--save-interval=1"
            - "--criterion=label_smoothed_cross_entropy"
            - "--source-lang=en"
            - "--lr-scheduler=inverse_sqrt"
            - "--min-lr=1e-09"
            - "--skip-invalid-size-inputs-valid-test"
            - "--target-lang=de"
            - "--label-smoothing=0.1"
            - "--update-freq=1"
            - "--optimizer=adam"
            - "--adam-betas=(0.9,0.98)"
            - "--warmup-init-lr=1e-07"
            - "--lr=0.0005"
            - "--warmup-updates=4000"
            - "--share-all-embeddings"
            - "--dropout=0.3"
            - "--weight-decay=0.0"
            - "--valid-subset=valid"
            - "--num_cores=8"
            - "--max-epoch=3"
            - "--log_steps=100"
            - "--train-subset=train"
            - "--input_shapes"
            - "64x64"
            - "128x32"
            - "256x16"
            "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": "METRIC_CONFIG"
              "value": |
                {
                 "metric_collection_config": {
                  "default_aggregation_strategies": [
                   "final"
                  ],
                  "tags_to_ignore": [
                   "LearningRate"
                  ],
                  "write_to_bigquery": true
                 },
                 "regression_test_config": {
                  "metric_subset_to_alert": [
                   "ExecuteTime__Percentile_99_sec_final",
                   "CompileTime__Percentile_99_sec_final",
                   "total_wall_time",
                   "Accuracy/test_final",
                   "aten_ops_sum_final"
                  ],
                  "metric_success_conditions": {
                   "CompileTime__Percentile_99_sec_final": {
                    "comparison": "less",
                    "success_threshold": {
                     "stddevs_from_mean": 5
                    },
                    "wait_for_n_points_of_history": 10
                   },
                   "ExecuteTime__Percentile_99_sec_final": {
                    "comparison": "less",
                    "success_threshold": {
                     "stddevs_from_mean": 5
                    },
                    "wait_for_n_points_of_history": 10
                   },
                   "aten_ops_sum_final": {
                    "comparison": "less_or_equal",
                    "success_threshold": {
                     "stddevs_from_mean": 0
                    }
                   },
                   "loss_final": {
                    "comparison": "less",
                    "success_threshold": {
                     "fixed_value": 4.5
                    }
                   },
                   "total_wall_time": {
                    "comparison": "less",
                    "success_threshold": {
                     "stddevs_from_mean": 5
                    },
                    "wait_for_n_points_of_history": 10
                   }
                  },
                  "write_to_error_reporting": true
                 },
                 "test_name": "pt-1.5-fs-transformer-convergence-v3-8"
                }
            - "name": "MODEL_DIR"
              "value": "gs://xl-ml-test-us-central1/k8s/fs-transformer/convergence/v3-8/$(JOB_NAME)"
            - "name": "XLA_USE_BF16"
              "value": "0"
            "image": "gcr.io/xl-ml-test/pytorch-xla:r1.5"
            "imagePullPolicy": "Always"
            "name": "train"
            "resources":
              "limits":
                "cloud-tpus.google.com/v3": 8
              "requests":
                "cpu": "9.0"
                "memory": "30Gi"
            "volumeMounts":
            - "mountPath": "/dev/shm"
              "name": "dshm"
            - "mountPath": "/datasets"
              "name": "wmt18-pd"
              "readOnly": true
          "restartPolicy": "Never"
          "volumes":
          - "emptyDir":
              "medium": "Memory"
            "name": "dshm"
          - "gcePersistentDisk":
              "fsType": "ext4"
              "pdName": "wmt18-en-de-pd-central1-b"
              "readOnly": true
            "name": "wmt18-pd"
  "schedule": "0 20 * * 1,6"