# Copyright 2021 The FedLearner Authors. All Rights Reserved.
#
# 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.

# coding: utf-8
from google.protobuf.json_format import MessageToDict
from fedlearner_webconsole.proto.workflow_definition_pb2 import (
    WorkflowDefinition, JobDefinition, JobDependency
)
from fedlearner_webconsole.proto.common_pb2 import (
    Variable
)


def make_workflow_template():
    workflow = WorkflowDefinition(
        group_alias='test_template',
        is_left=True,
        variables=[
            Variable(
                name='image_version',
                value='v1.5-rc3',
                access_mode=Variable.PEER_READABLE),
            Variable(
                name='num_partitions',
                value='4',
                access_mode=Variable.PEER_WRITABLE),
        ],
        job_definitions=[
            JobDefinition(
                name='raw-data-job',
                job_type=JobDefinition.RAW_DATA,
                is_federated=False,
                variables=[
                    Variable(
                        name='input_dir',
                        value='/app/deploy/integrated_test/tfrecord_raw_data',
                        access_mode=Variable.PRIVATE),
                    Variable(
                        name='file_wildcard',
                        value='*.rd',
                        access_mode=Variable.PRIVATE),
                    Variable(
                        name='batch_size',
                        value='1024',
                        access_mode=Variable.PEER_WRITABLE),
                    Variable(
                        name='input_format',
                        value='TF_RECORD',
                        access_mode=Variable.PRIVATE),
                    Variable(
                        name='output_format',
                        value='TF_RECORD',
                        access_mode=Variable.PRIVATE),
                    Variable(
                        name='master_cpu',
                        value='2000m',
                        access_mode=Variable.PEER_WRITABLE),
                    Variable(
                        name='master_mem',
                        value='3Gi',
                        access_mode=Variable.PEER_WRITABLE),
                    Variable(
                        name='worker_cpu',
                        value='2000m',
                        access_mode=Variable.PEER_WRITABLE),
                    Variable(
                        name='worker_mem',
                        value='4Gi',
                        access_mode=Variable.PEER_WRITABLE),
                ],
                yaml_template='''{
    "apiVersion": "fedlearner.k8s.io/v1alpha1",
    "kind": "FLApp",
    "metadata": {
        "name": "${workflow.jobs.raw-data-job.name}",
        "namespace": "${project.variables.namespace}"
    },
    "spec": {
        "role": "Follower",
        "peerSpecs": {
            "Leader": {
                "peerURL": "",
                "authority": ""
            }
        },
        "cleanPodPolicy": "All",
        "flReplicaSpecs": {
            "Master": {
                "template": {
                    "spec": {
                        "containers": [
                            {
                                "resources": {
                                    "limits": {
                                        "cpu": "${workflow.jobs.raw-data-job.variables.master_cpu}",
                                        "memory": "${workflow.jobs.raw-data-job.variables.master_mem}"
                                    },
                                    "requests": {
                                        "cpu": "${workflow.jobs.raw-data-job.variables.master_cpu}",
                                        "memory": "${workflow.jobs.raw-data-job.variables.master_mem}"
                                    }
                                },
                                "image": "artifact.bytedance.com/fedlearner/fedlearner:${workflow.variables.image_version}",
                                "ports": [
                                    {
                                        "containerPort": 50051,
                                        "name": "flapp-port"
                                    }
                                ],
                                "command": [
                                    "/app/deploy/scripts/data_portal/run_data_portal_master.sh"
                                ],
                                "args": [],
                                "env": [
                                    ${system.basic_envs},
                                    {
                                        "name": "ETCD_NAME",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "ETCD_ADDR",
                                        "value": "fedlearner-stack-etcd.default.svc.cluster.local:2379"
                                    },
                                    {
                                        "name": "ETCD_BASE_DIR",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "EGRESS_URL",
                                        "value": "fedlearner-stack-ingress-nginx-controller.default.svc.cluster.local:80"
                                    },
                                    {
                                        "name": "EGRESS_HOST",
                                        "value": "${project.participants[0].egress_host}"
                                    },
                                    {
                                        "name": "EGRESS_DOMAIN",
                                        "value": "${project.participants[0].egress_domain}"
                                    },
                                    {
                                        "name": "STORAGE_ROOT_PATH",
                                        "value": "${project.variables.storage_root_dir}"
                                    },
                                    {
                                        "name": "POD_IP",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "status.podIP"
                                            }
                                        }
                                    },
                                    {
                                        "name": "POD_NAME",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "metadata.name"
                                            }
                                        }
                                    },
                                    {
                                        "name": "APPLICATION_ID",
                                        "value": "${workflow.jobs.raw-data-job.name}"
                                    },
                                    {
                                        "name": "DATA_PORTAL_NAME",
                                        "value": "${workflow.jobs.raw-data-job.name}"
                                    },
                                    {
                                        "name": "OUTPUT_PARTITION_NUM",
                                        "value": "${workflow.variables.num_partitions}"
                                    },
                                    {
                                        "name": "INPUT_BASE_DIR",
                                        "value": "${workflow.jobs.raw-data-job.variables.input_dir}"
                                    },
                                    {
                                        "name": "OUTPUT_BASE_DIR",
                                        "value": "${project.variables.storage_root_dir}/raw_data/${workflow.jobs.raw-data-job.name}"
                                    },
                                    {
                                        "name": "RAW_DATA_PUBLISH_DIR",
                                        "value": "portal_publish_dir/${workflow.jobs.raw-data-job.name}"
                                    },
                                    {
                                        "name": "DATA_PORTAL_TYPE",
                                        "value": "Streaming"
                                    },
                                    {
                                        "name": "FILE_WILDCARD",
                                        "value": "${workflow.jobs.raw-data-job.variables.file_wildcard}"
                                    }
                                ],
                                "volumeMounts": [
                                    {
                                        "mountPath": "/data",
                                        "name": "data"
                                    }
                                ],
                                "imagePullPolicy": "IfNotPresent",
                                "name": "tensorflow"
                            }
                        ],
                        "imagePullSecrets": [
                            {
                                "name": "regcred"
                            }
                        ],
                        "volumes": [
                            {
                                "persistentVolumeClaim": {
                                    "claimName": "pvc-fedlearner-default"
                                },
                                "name": "data"
                            }
                        ],
                        "restartPolicy": "Never"
                    }
                },
                "pair": false,
                "replicas": 1
            },
            "Worker": {
                "replicas": 4,
                "template": {
                    "spec": {
                        "containers": [
                            {
                                "resources": {
                                    "limits": {
                                        "cpu": "${workflow.jobs.raw-data-job.variables.worker_cpu}",
                                        "memory": "${workflow.jobs.raw-data-job.variables.worker_mem}"
                                    },
                                    "requests": {
                                        "cpu": "${workflow.jobs.raw-data-job.variables.worker_cpu}",
                                        "memory": "${workflow.jobs.raw-data-job.variables.worker_mem}"
                                    }
                                },
                                "image": "artifact.bytedance.com/fedlearner/fedlearner:${workflow.variables.image_version}",
                                "command": [
                                    "/app/deploy/scripts/data_portal/run_data_portal_worker.sh"
                                ],
                                "args": [],
                                "env": [
                                    {
                                        "name": "ETCD_NAME",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "ETCD_ADDR",
                                        "value": "fedlearner-stack-etcd.default.svc.cluster.local:2379"
                                    },
                                    {
                                        "name": "ETCD_BASE_DIR",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "EGRESS_URL",
                                        "value": "fedlearner-stack-ingress-nginx-controller.default.svc.cluster.local:80"
                                    },
                                    {
                                        "name": "EGRESS_HOST",
                                        "value": "${project.participants[0].egress_host}"
                                    },
                                    {
                                        "name": "EGRESS_DOMAIN",
                                        "value": "${project.participants[0].egress_domain}"
                                    },
                                    {
                                        "name": "STORAGE_ROOT_PATH",
                                        "value": "${project.variables.storage_root_dir}"
                                    },
                                    {
                                        "name": "POD_IP",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "status.podIP"
                                            }
                                        }
                                    },
                                    {
                                        "name": "POD_NAME",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "metadata.name"
                                            }
                                        }
                                    },
                                    ${system.basic_envs},
                                    {
                                        "name": "CPU_REQUEST",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "requests.cpu"
                                            }
                                        }
                                    },
                                    {
                                        "name": "MEM_REQUEST",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "requests.memory"
                                            }
                                        }
                                    },
                                    {
                                        "name": "CPU_LIMIT",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "limits.cpu"
                                            }
                                        }
                                    },
                                    {
                                        "name": "MEM_LIMIT",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "limits.memory"
                                            }
                                        }
                                    },
                                    {
                                        "name": "APPLICATION_ID",
                                        "value": "${workflow.jobs.raw-data-job.name}"
                                    },
                                    {
                                        "name": "BATCH_SIZE",
                                        "value": "${workflow.jobs.raw-data-job.variables.batch_size}"
                                    },
                                    {
                                        "name": "INPUT_DATA_FORMAT",
                                        "value": "${workflow.jobs.raw-data-job.variables.input_format}"
                                    },
                                    {
                                        "name": "COMPRESSED_TYPE",
                                        "value": ""
                                    },
                                    {
                                        "name": "OUTPUT_DATA_FORMAT",
                                        "value": "${workflow.jobs.raw-data-job.variables.output_format}"
                                    }
                                ],
                                "volumeMounts": [
                                    {
                                        "mountPath": "/data",
                                        "name": "data"
                                    }
                                ],
                                "imagePullPolicy": "IfNotPresent",
                                "name": "tensorflow"
                            }
                        ],
                        "imagePullSecrets": [
                            {
                                "name": "regcred"
                            }
                        ],
                        "volumes": [
                            {
                                "persistentVolumeClaim": {
                                    "claimName": "pvc-fedlearner-default"
                                },
                                "name": "data"
                            }
                        ],
                        "restartPolicy": "Never"
                    }
                },
                "pair": false
            }
        }
    }
}
                '''
            ),
            JobDefinition(
                name='data-join-job',
                job_type=JobDefinition.DATA_JOIN,
                is_federated=True,
                variables=[
                    Variable(
                        name='master_cpu',
                        value='2000m',
                        access_mode=Variable.PEER_WRITABLE),
                    Variable(
                        name='master_mem',
                        value='3Gi',
                        access_mode=Variable.PEER_WRITABLE),
                    Variable(
                        name='worker_cpu',
                        value='4000m',
                        access_mode=Variable.PEER_WRITABLE),
                    Variable(
                        name='worker_mem',
                        value='4Gi',
                        access_mode=Variable.PEER_WRITABLE),
                ],
                dependencies=[
                    JobDependency(source='raw-data-job')
                ],
                yaml_template='''
{
    "apiVersion": "fedlearner.k8s.io/v1alpha1",
    "kind": "FLApp",
    "metadata": {
        "name": "${workflow.jobs.data-join-job.name}",
        "namespace": "${project.variables.namespace}"
    },
    "spec": {
        "role": "Leader",
        "cleanPodPolicy": "All",
        "peerSpecs": {
            "Follower": {
                "peerURL": "fedlearner-stack-ingress-nginx-controller.default.svc.cluster.local:80",
                "authority": "${project.participants[0].egress_domain}",
                "extraHeaders": {
                    "x-host": "default.fedlearner.operator"
                }
            }
        },
        "flReplicaSpecs": {
            "Master": {
                "template": {
                    "spec": {
                        "restartPolicy": "Never",
                        "containers": [
                            {
                                "env": [
                                    ${system.basic_envs},
                                    {
                                        "name": "STORAGE_ROOT_PATH",
                                        "value": "${project.variables.storage_root_dir}"
                                    },
                                    {
                                        "name": "POD_IP",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "status.podIP"
                                            }
                                        }
                                    },
                                    {
                                        "name": "POD_NAME",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "metadata.name"
                                            }
                                        }
                                    },
                                    {
                                        "name": "ROLE",
                                        "value": "leader"
                                    },
                                    {
                                        "name": "APPLICATION_ID",
                                        "value": "${workflow.jobs.data-join-job.name}"
                                    },
                                    {
                                        "name": "OUTPUT_BASE_DIR",
                                        "value": "${project.variables.storage_root_dir}/data_source/${workflow.jobs.data-join-job.name}"
                                    },
                                    {
                                        "name": "CPU_REQUEST",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "requests.cpu"
                                            }
                                        }
                                    },
                                    {
                                        "name": "MEM_REQUEST",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "requests.memory"
                                            }
                                        }
                                    },
                                    {
                                        "name": "CPU_LIMIT",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "limits.cpu"
                                            }
                                        }
                                    },
                                    {
                                        "name": "MEM_LIMIT",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "limits.memory"
                                            }
                                        }
                                    },
                                    {
                                        "name": "ETCD_NAME",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "ETCD_ADDR",
                                        "value": "fedlearner-stack-etcd.default.svc.cluster.local:2379"
                                    },
                                    {
                                        "name": "ETCD_BASE_DIR",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "EGRESS_URL",
                                        "value": "fedlearner-stack-ingress-nginx-controller.default.svc.cluster.local:80"
                                    },
                                    {
                                        "name": "EGRESS_HOST",
                                        "value": "${project.participants[0].egress_host}"
                                    },
                                    {
                                        "name": "EGRESS_DOMAIN",
                                        "value": "${project.participants[0].egress_domain}"
                                    },
                                    {
                                        "name": "BATCH_MODE",
                                        "value": "--batch_mode"
                                    },
                                    {
                                        "name": "PARTITION_NUM",
                                        "value": "${workflow.variables.num_partitions}"
                                    },
                                    {
                                        "name": "START_TIME",
                                        "value": "0"
                                    },
                                    {
                                        "name": "END_TIME",
                                        "value": "999999999999"
                                    },
                                    {
                                        "name": "NEGATIVE_SAMPLING_RATE",
                                        "value": "1.0"
                                    },
                                    {
                                        "name": "RAW_DATA_SUB_DIR",
                                        "value": "portal_publish_dir/${workflow.jobs.raw-data-job.name}"
                                    }
                                ],
                                "imagePullPolicy": "IfNotPresent",
                                "name": "tensorflow",
                                "volumeMounts": [
                                    {
                                        "mountPath": "/data",
                                        "name": "data"
                                    }
                                ],
                                "image": "artifact.bytedance.com/fedlearner/fedlearner:${workflow.variables.image_version}",
                                "ports": [
                                    {
                                        "containerPort": 50051,
                                        "name": "flapp-port"
                                    }
                                ],
                                "command": [
                                    "/app/deploy/scripts/wait4pair_wrapper.sh"
                                ],
                                "args": [
                                    "/app/deploy/scripts/data_join/run_data_join_master.sh"
                                ],
                                "resources": {
                                      "limits": {
                                        "cpu": "${workflow.jobs.data-join-job.variables.master_cpu}",
                                        "memory": "${workflow.jobs.data-join-job.variables.master_mem}"
                                      },
                                      "requests": {
                                        "cpu": "${workflow.jobs.data-join-job.variables.master_cpu}",
                                        "memory": "${workflow.jobs.data-join-job.variables.master_mem}"
                                      }
                                }
                            }
                        ],
                        "imagePullSecrets": [
                            {
                                "name": "regcred"
                            }
                        ],
                        "volumes": [
                            {
                                "persistentVolumeClaim": {
                                    "claimName": "pvc-fedlearner-default"
                                },
                                "name": "data"
                            }
                        ]
                    }
                },
                "pair": true,
                "replicas": 1
            },
            "Worker": {
                "template": {
                    "spec": {
                        "restartPolicy": "Never",
                        "containers": [
                            {
                                "env": [
                                    ${system.basic_envs},
                                    {
                                        "name": "STORAGE_ROOT_PATH",
                                        "value": "${project.variables.storage_root_dir}"
                                    },
                                    {
                                        "name": "POD_IP",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "status.podIP"
                                            }
                                        }
                                    },
                                    {
                                        "name": "POD_NAME",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "metadata.name"
                                            }
                                        }
                                    },
                                    {
                                        "name": "ROLE",
                                        "value": "leader"
                                    },
                                    {
                                        "name": "APPLICATION_ID",
                                        "value": "${workflow.jobs.data-join-job.name}"
                                    },
                                    {
                                        "name": "OUTPUT_BASE_DIR",
                                        "value": "${project.variables.storage_root_dir}/data_source/${workflow.jobs.data-join-job.name}"
                                    },
                                    {
                                        "name": "CPU_REQUEST",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "requests.cpu"
                                            }
                                        }
                                    },
                                    {
                                        "name": "MEM_REQUEST",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "requests.memory"
                                            }
                                        }
                                    },
                                    {
                                        "name": "CPU_LIMIT",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "limits.cpu"
                                            }
                                        }
                                    },
                                    {
                                        "name": "MEM_LIMIT",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "limits.memory"
                                            }
                                        }
                                    },
                                    {
                                        "name": "ETCD_NAME",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "ETCD_ADDR",
                                        "value": "fedlearner-stack-etcd.default.svc.cluster.local:2379"
                                    },
                                    {
                                        "name": "ETCD_BASE_DIR",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "EGRESS_URL",
                                        "value": "fedlearner-stack-ingress-nginx-controller.default.svc.cluster.local:80"
                                    },
                                    {
                                        "name": "EGRESS_HOST",
                                        "value": "${project.participants[0].egress_host}"
                                    },
                                    {
                                        "name": "EGRESS_DOMAIN",
                                        "value": "${project.participants[0].egress_domain}"
                                    },
                                    {
                                        "name": "PARTITION_NUM",
                                        "value": "${workflow.variables.num_partitions}"
                                    },
                                    {
                                        "name": "RAW_DATA_SUB_DIR",
                                        "value": "portal_publish_dir/${workflow.jobs.raw-data-job.name}"
                                    },
                                    {
                                        "name": "DATA_BLOCK_DUMP_INTERVAL",
                                        "value": "600"
                                    },
                                    {
                                        "name": "DATA_BLOCK_DUMP_THRESHOLD",
                                        "value": "65536"
                                    },
                                    {
                                        "name": "EXAMPLE_ID_DUMP_INTERVAL",
                                        "value": "600"
                                    },
                                    {
                                        "name": "EXAMPLE_ID_DUMP_THRESHOLD",
                                        "value": "65536"
                                    },
                                    {
                                        "name": "EXAMPLE_ID_BATCH_SIZE",
                                        "value": "4096"
                                    },
                                    {
                                        "name": "MAX_FLYING_EXAMPLE_ID",
                                        "value": "307152"
                                    },
                                    {
                                        "name": "MIN_MATCHING_WINDOW",
                                        "value": "2048"
                                    },
                                    {
                                        "name": "MAX_MATCHING_WINDOW",
                                        "value": "8192"
                                    },
                                    {
                                        "name": "RAW_DATA_ITER",
                                        "value": "${workflow.jobs.raw-data-job.variables.output_format}"
                                    },
                                    {
                                        "name": "RAW_DATA_SUB_DIR",
                                        "value": "portal_publish_dir/${workflow.jobs.raw-data-job.name}"
                                    },
                                    {
                                        "name": "PARTITION_NUM",
                                        "value": "${workflow.variables.num_partitions}"
                                    }
                                ],
                                "imagePullPolicy": "IfNotPresent",
                                "name": "tensorflow",
                                "volumeMounts": [
                                    {
                                        "mountPath": "/data",
                                        "name": "data"
                                    }
                                ],
                                "image": "artifact.bytedance.com/fedlearner/fedlearner:${workflow.variables.image_version}",
                                "ports": [
                                    {
                                        "containerPort": 50051,
                                        "name": "flapp-port"
                                    }
                                ],
                                "command": [
                                    "/app/deploy/scripts/wait4pair_wrapper.sh"
                                ],
                                "args": [
                                    "/app/deploy/scripts/data_join/run_data_join_worker.sh"
                                ],
                                "resources": {
                                    "limits": {
                                        "cpu": "${workflow.jobs.data-join-job.variables.worker_cpu}",
                                        "memory": "${workflow.jobs.data-join-job.variables.worker_mem}"
                                    },
                                    "requests": {
                                        "cpu": "${workflow.jobs.data-join-job.variables.worker_cpu}",
                                        "memory": "${workflow.jobs.data-join-job.variables.worker_mem}"
                                    }
                                }
                            }
                        ],
                        "imagePullSecrets": [
                            {
                                "name": "regcred"
                            }
                        ],
                        "volumes": [
                            {
                                "persistentVolumeClaim": {
                                    "claimName": "pvc-fedlearner-default"
                                },
                                "name": "data"
                            }
                        ]
                    }
                },
                "pair": true,
                "replicas": ${workflow.variables.num_partitions}
            }
        }
    }
}
                '''
            ),
            JobDefinition(
                name='train-job',
                job_type=JobDefinition.NN_MODEL_TRANINING,
                is_federated=True,
                variables=[
                    Variable(
                        name='code_key',
                        value='/app/deploy/integrated_test/code_key/criteo-train-2.tar.gz',
                        access_mode=Variable.PRIVATE
                    )
                ],
                dependencies=[
                    JobDependency(source='data-join-job')
                ],
                yaml_template='''
                {
    "apiVersion": "fedlearner.k8s.io/v1alpha1",
    "kind": "FLApp",
    "metadata": {
        "name": "${workflow.jobs.train-job.name}",
        "namespace": "${project.variables.namespace}"
    },
    "spec": {
        "role": "Follower",
        "cleanPodPolicy": "All",
        "peerSpecs": {
            "Leader": {
                "peerURL": "fedlearner-stack-ingress-nginx-controller.default.svc.cluster.local:80",
                "authority": "${project.participants[0].egress_domain}",
                "extraHeaders": {
                    "x-host": "default.fedlearner.operator"
                }
            }
        },
        "flReplicaSpecs": {
            "Master": {
                "template": {
                    "spec": {
                        "restartPolicy": "Never",
                        "containers": [
                            {
                                "env": [
                                    ${system.basic_envs},
                                    {
                                        "name": "STORAGE_ROOT_PATH",
                                        "value":  "${project.variables.storage_root_dir}"
                                    },
                                    {
                                        "name": "POD_IP",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "status.podIP"
                                            }
                                        }
                                    },
                                    {
                                        "name": "POD_NAME",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "metadata.name"
                                            }
                                        }
                                    },
                                    {
                                        "name": "ROLE",
                                        "value": "follower"
                                    },
                                    {
                                        "name": "APPLICATION_ID",
                                        "value": "${workflow.jobs.train-job.name}"
                                    },
                                    {
                                        "name": "OUTPUT_BASE_DIR",
                                        "value": "${project.variables.storage_root_dir}/job_output/${workflow.jobs.train-job.name}"
                                    },
                                    {
                                        "name": "CPU_REQUEST",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "requests.cpu"
                                            }
                                        }
                                    },
                                    {
                                        "name": "MEM_REQUEST",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "requests.memory"
                                            }
                                        }
                                    },
                                    {
                                        "name": "CPU_LIMIT",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "limits.cpu"
                                            }
                                        }
                                    },
                                    {
                                        "name": "MEM_LIMIT",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "limits.memory"
                                            }
                                        }
                                    },
                                    {
                                        "name": "ETCD_NAME",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "ETCD_ADDR",
                                        "value": "fedlearner-stack-etcd.default.svc.cluster.local:2379"
                                    },
                                    {
                                        "name": "ETCD_BASE_DIR",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "EGRESS_URL",
                                        "value": "fedlearner-stack-ingress-nginx-controller.default.svc.cluster.local:80"
                                    },
                                    {
                                        "name": "EGRESS_HOST",
                                        "value": "${project.participants[0].egress_host}"
                                    },
                                    {
                                        "name": "EGRESS_DOMAIN",
                                        "value": "${project.participants[0].egress_domain}"
                                    },
                                    {
                                        "name": "DATA_SOURCE",
                                        "value": "${workflow.jobs.data-join-job.name}"
                                    }
                                ],
                                "imagePullPolicy": "IfNotPresent",
                                "name": "tensorflow",
                                "volumeMounts": [
                                    {
                                        "mountPath": "/data",
                                        "name": "data"
                                    }
                                ],
                                "image": "artifact.bytedance.com/fedlearner/fedlearner:${workflow.variables.image_version}",
                                "ports": [
                                    {
                                        "containerPort": 50051,
                                        "name": "flapp-port"
                                    }
                                ],
                                "command": [
                                    "/app/deploy/scripts/trainer/run_trainer_master.sh"
                                ],
                                "args": [],
                                "resources": {
                                      "limits": {
                                        "cpu": "2000m",
                                        "memory": "2Gi"
                                      },
                                      "requests": {
                                        "cpu": "1000m",
                                        "memory": "2Gi"
                                      }
                                }
                            }
                        ],
                        "imagePullSecrets": [
                            {
                                "name": "regcred"
                            }
                        ],
                        "volumes": [
                            {
                                "persistentVolumeClaim": {
                                    "claimName": "pvc-fedlearner-default"
                                },
                                "name": "data"
                            }
                        ]
                    }
                },
                "replicas": 1,
                "pair": false
            },
            "PS": {
                "template": {
                    "spec": {
                        "restartPolicy": "Never",
                        "containers": [
                            {
                                "env": [
                                    ${system.basic_envs},
                                    {
                                        "name": "STORAGE_ROOT_PATH",
                                        "value": "${project.variables.storage_root_dir}"
                                    },
                                    {
                                        "name": "POD_IP",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "status.podIP"
                                            }
                                        }
                                    },
                                    {
                                        "name": "POD_NAME",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "metadata.name"
                                            }
                                        }
                                    },
                                    {
                                        "name": "ROLE",
                                        "value": "follower"
                                    },
                                    {
                                        "name": "APPLICATION_ID",
                                        "value": "${workflow.jobs.train-job.name}"
                                    },
                                    {
                                        "name": "OUTPUT_BASE_DIR",
                                        "value": "${project.variables.storage_root_dir}/job_output/${workflow.jobs.train-job.name}"
                                    },
                                    {
                                        "name": "CPU_REQUEST",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "requests.cpu"
                                            }
                                        }
                                    },
                                    {
                                        "name": "MEM_REQUEST",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "requests.memory"
                                            }
                                        }
                                    },
                                    {
                                        "name": "CPU_LIMIT",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "limits.cpu"
                                            }
                                        }
                                    },
                                    {
                                        "name": "MEM_LIMIT",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "limits.memory"
                                            }
                                        }
                                    },
                                    {
                                        "name": "ETCD_NAME",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "ETCD_ADDR",
                                        "value": "fedlearner-stack-etcd.default.svc.cluster.local:2379"
                                    },
                                    {
                                        "name": "ETCD_BASE_DIR",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "EGRESS_URL",
                                        "value": "fedlearner-stack-ingress-nginx-controller.default.svc.cluster.local:80"
                                    },
                                    {
                                        "name": "EGRESS_HOST",
                                        "value": "${project.participants[0].egress_host}"
                                    },
                                    {
                                        "name": "EGRESS_DOMAIN",
                                        "value": "${project.participants[0].egress_domain}"
                                    },
                                    {
                                        "name": "DATA_SOURCE",
                                        "value": "${workflow.jobs.data-join-job.name}"
                                    },
                                    {
                                        "name": "DATA_SOURCE",
                                        "value": "${workflow.jobs.data-join-job.name}"
                                    }
                                ],
                                "imagePullPolicy": "IfNotPresent",
                                "name": "tensorflow",
                                "volumeMounts": [
                                    {
                                        "mountPath": "/data",
                                        "name": "data"
                                    }
                                ],
                                "image": "artifact.bytedance.com/fedlearner/fedlearner:${workflow.variables.image_version}",
                                "ports": [
                                    {
                                        "containerPort": 50051,
                                        "name": "flapp-port"
                                    }
                                ],
                                "command": [
                                    "/app/deploy/scripts/trainer/run_trainer_ps.sh"
                                ],
                                "args": [],
                                "resources": {
                                    "limits": {
                                        "cpu": "2000m",
                                        "memory": "4Gi"
                                    },
                                    "requests": {
                                        "cpu": "1000m",
                                        "memory": "2Gi"
                                    }
                                }
                            }
                        ],
                        "imagePullSecrets": [
                            {
                                "name": "regcred"
                            }
                        ],
                        "volumes": [
                            {
                                "persistentVolumeClaim": {
                                    "claimName": "pvc-fedlearner-default"
                                },
                                "name": "data"
                            }
                        ]
                    }
                },
                "pair": false,
                "replicas": 1
            },
            "Worker": {
                "template": {
                    "spec": {
                        "restartPolicy": "Never",
                        "containers": [
                            {
                                "env": [
                                    ${system.basic_envs},
                                    {
                                        "name": "STORAGE_ROOT_PATH",
                                        "value": "${project.variables.storage_root_dir}"
                                    },
                                    {
                                        "name": "POD_IP",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "status.podIP"
                                            }
                                        }
                                    },
                                    {
                                        "name": "POD_NAME",
                                        "valueFrom": {
                                            "fieldRef": {
                                                "fieldPath": "metadata.name"
                                            }
                                        }
                                    },
                                    {
                                        "name": "ROLE",
                                        "value": "follower"
                                    },
                                    {
                                        "name": "APPLICATION_ID",
                                        "value": "${workflow.jobs.train-job.name}"
                                    },
                                    {
                                        "name": "OUTPUT_BASE_DIR",
                                        "value": "${project.variables.storage_root_dir}/job_output/${workflow.jobs.train-job.name}"
                                    },
                                    {
                                        "name": "CPU_REQUEST",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "requests.cpu"
                                            }
                                        }
                                    },
                                    {
                                        "name": "MEM_REQUEST",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "requests.memory"
                                            }
                                        }
                                    },
                                    {
                                        "name": "CPU_LIMIT",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "limits.cpu"
                                            }
                                        }
                                    },
                                    {
                                        "name": "MEM_LIMIT",
                                        "valueFrom": {
                                            "resourceFieldRef": {
                                                "resource": "limits.memory"
                                            }
                                        }
                                    },
                                    {
                                        "name": "ETCD_NAME",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "ETCD_ADDR",
                                        "value": "fedlearner-stack-etcd.default.svc.cluster.local:2379"
                                    },
                                    {
                                        "name": "ETCD_BASE_DIR",
                                        "value": "fedlearner"
                                    },
                                    {
                                        "name": "EGRESS_URL",
                                        "value": "fedlearner-stack-ingress-nginx-controller.default.svc.cluster.local:80"
                                    },
                                    {
                                        "name": "EGRESS_HOST",
                                        "value": "${project.participants[0].egress_host}"
                                    },
                                    {
                                        "name": "EGRESS_DOMAIN",
                                        "value": "${project.participants[0].egress_domain}"
                                    },
                                    {
                                        "name": "CODE_KEY",
                                        "value": "${workflow.jobs.train-job.variables.code_key}"
                                    },
                                    {
                                        "name": "SAVE_CHECKPOINT_STEPS",
                                        "value": "1000"
                                    },
                                    {
                                        "name": "DATA_SOURCE",
                                        "value": "${workflow.jobs.data-join-job.name}"
                                    },
                                    {
                                        "name": "DATA_SOURCE",
                                        "value": "${workflow.jobs.data-join-job.name}"
                                    }
                                ],
                                "imagePullPolicy": "IfNotPresent",
                                "name": "tensorflow",
                                "volumeMounts": [
                                    {
                                        "mountPath": "/data",
                                        "name": "data"
                                    }
                                ],
                                "image": "artifact.bytedance.com/fedlearner/fedlearner:${workflow.variables.image_version}",
                                "ports": [
                                    {
                                        "containerPort": 50051,
                                        "name": "flapp-port"
                                    },
                                    {
                                        "containerPort": 50052,
                                        "name": "tf-port"
                                    }
                                ],
                                "command": [
                                    "/app/deploy/scripts/wait4pair_wrapper.sh"
                                ],
                                "args": [
                                    "/app/deploy/scripts/trainer/run_trainer_worker.sh"
                                ],
                                "resources": {
                                    "limits": {
                                        "cpu": "2000m",
                                        "memory": "4Gi"
                                    },
                                    "requests": {
                                        "cpu": "1000m",
                                        "memory": "2Gi"
                                    }
                                }
                            }
                        ],
                        "imagePullSecrets": [
                            {
                                "name": "regcred"
                            }
                        ],
                        "volumes": [
                            {
                                "persistentVolumeClaim": {
                                    "claimName": "pvc-fedlearner-default"
                                },
                                "name": "data"
                            }
                        ]
                    }
                },
                "pair": true,
                "replicas": ${workflow.variables.num_partitions}
            }
        }
    }
}
                '''
            )
        ])

    return workflow


import json

if __name__ == '__main__':
    print(json.dumps(MessageToDict(
        make_workflow_template(),
        preserving_proto_field_name=True,
        including_default_value_fields=True)))
