# 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":
  "labels":
    "accelerator": "v3-8"
    "benchmarkId": "pt-r1.7-dlrm-convergence-conv-v3-8"
    "frameworkVersion": "pt-r1.7"
    "mode": "conv"
    "model": "dlrm-convergence"
  "name": "pt-r1.7-dlrm-convergence-conv-v3-8"
  "namespace": "automated"
"spec":
  "concurrencyPolicy": "Forbid"
  "jobTemplate":
    "metadata":
      "annotations":
        "ml-testing-accelerators/gcs-subdir": "pt-r1.7/dlrm-convergence/conv/v3-8"
        "ml-testing-accelerators/metric-config": |
          {
            "sources": [
              {
                "literals": {
                  "assertions": {
                    "duration": {
                      "inclusive_bounds": false,
                      "std_devs_from_mean": {
                        "comparison": "LESS",
                        "std_devs": 5
                      },
                      "wait_for_n_data_points": 10
                    }
                  }
                }
              },
              {
                "tensorboard": {
                  "aggregate_assertions": [
                    {
                      "assertion": {
                        "inclusive_bounds": false,
                        "std_devs_from_mean": {
                          "comparison": "LESS",
                          "std_devs": 5
                        },
                        "wait_for_n_data_points": 10
                      },
                      "strategy": "FINAL",
                      "tag": "ExecuteTime__Percentile_99_sec"
                    },
                    {
                      "assertion": {
                        "inclusive_bounds": true,
                        "std_devs_from_mean": {
                          "comparison": "LESS",
                          "std_devs": 0
                        },
                        "wait_for_n_data_points": 0
                      },
                      "strategy": "FINAL",
                      "tag": "aten_ops_sum"
                    }
                  ],
                  "exclude_tags": [
                    "LearningRate"
                  ],
                  "include_tags": [
                    {
                      "strategies": [
                        "FINAL"
                      ],
                      "tag_pattern": "*"
                    }
                  ],
                  "merge_runs": true
                }
              }
            ]
          }
      "labels":
        "accelerator": "v3-8"
        "benchmarkId": "pt-r1.7-dlrm-convergence-conv-v3-8"
        "frameworkVersion": "pt-r1.7"
        "mode": "conv"
        "model": "dlrm-convergence"
    "spec":
      "activeDeadlineSeconds": 21600
      "backoffLimit": 1
      "template":
        "metadata":
          "annotations":
            "reserved.cloud-tpus.google.com": "false"
            "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
              
              apt-get install -y bc
              pip install onnx
              git clone --recursive https://github.com/pytorch-tpu/examples.git
              python examples/deps/dlrm/dlrm_tpu_runner.py \
                  --arch-sparse-feature-size=16 \
                  --arch-mlp-bot="13-512-256-64-16" \
                  --arch-mlp-top="512-256-1" \
                  --data-generation=dataset \
                  --data-set=kaggle \
                  --raw-data-file=/datasets/criteo-kaggle/train.txt \
                  --processed-data-file=/datasets/criteo-kaggle/kaggleAdDisplayChallenge_processed.npz \
                  --loss-function=bce \
                  --round-targets=True \
                  --learning-rate=0.1 \
                  --mini-batch-size=128 \
                  --print-freq=1024 \
                  --print-time \
                  --test-mini-batch-size=16384 \
                  --test-num-workers=4 \
                  --test-freq=101376 \
                      --use-tpu \
                      --num-indices-per-lookup=1 \
                      --num-indices-per-lookup-fixed \
                      --tpu-model-parallel-group-len 8 \
                      --tpu-metrics-debug \
                      --tpu-cores=8 |& tee dlrm_logs.txt
              acc=`grep Testing dlrm_logs.txt | tail -1 | grep -oP 'best \K[+-]?([0-9]*[.])?[0-9]+'`
              echo 'Accuracy is' $acc
              test $(echo $acc'>'78.75 | bc -l) -eq 1  # assert cls acc higher than 78.75
              
            "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/dlrm-convergence/conv/v3-8/$(JOB_NAME)"
            - "name": "XLA_USE_BF16"
              "value": "0"
            "envFrom":
            - "configMapRef":
                "name": "gcs-buckets"
            "image": "gcr.io/xl-ml-test/pytorch-xla:r1.7"
            "imagePullPolicy": "Always"
            "name": "train"
            "resources":
              "limits":
                "cloud-tpus.google.com/v3": 8
              "requests":
                "cpu": "40.0"
                "memory": "500Gi"
            "volumeMounts":
            - "mountPath": "/dev/shm"
              "name": "dshm"
              "readOnly": false
            - "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/dlrm-convergence/conv/v3-8/$(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",
                   "total_wall_time",
                   "Accuracy/test_final",
                   "aten_ops_sum_final"
                  ],
                  "metric_success_conditions": {
                   "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
                    }
                   },
                   "total_wall_time": {
                    "comparison": "less",
                    "success_threshold": {
                     "stddevs_from_mean": 5
                    },
                    "wait_for_n_points_of_history": 10
                   }
                  }
                 },
                 "test_name": "pt-r1.7-dlrm-convergence-conv-v3-8"
                }
            "envFrom":
            - "configMapRef":
                "name": "gcs-buckets"
            "image": "gcr.io/xl-ml-test/publisher:stable"
            "imagePullPolicy": "Always"
            "name": "publisher"
          "nodeSelector":
            "tpu-available": "true"
          "restartPolicy": "Never"
          "volumes":
          - "emptyDir":
              "medium": "Memory"
            "name": "dshm"
          - "name": "pytorch-datasets-claim"
            "persistentVolumeClaim":
              "claimName": "pytorch-datasets-claim"
  "schedule": "0 0 * * 0,2,4"
  "successfulJobsHistoryLimit": 1