template.yaml.jinja

{%- set name = "" -%}
{%- set image = "" -%}
{%- set worker_replicas = 2 -%}
{%- set ps_replicas = 1 -%}
{%- set script = "" -%}
{%- set data_dir = "" -%}
{%- set train_dir = "" -%}
{%- set credential_secret_name = "" -%}
{%- set credential_secret_key = "" -%}

{%- set port = 5000 -%}
{%- set replicas = {"worker": worker_replicas, "ps": ps_replicas} -%}

{%- macro worker_hosts() -%}
  {%- for i in range(worker_replicas) -%}
    {%- if not loop.first -%},{%- endif -%}
    {{ name }}-worker-{{ i }}:{{ port }}
  {%- endfor -%}
{%- endmacro -%}

{%- macro ps_hosts() -%}
  {%- for i in range(ps_replicas) -%}
    {%- if not loop.first -%},{%- endif -%}
    {{ name }}-ps-{{ i }}:{{ port }}
  {%- endfor -%}
{%- endmacro -%}

{%- for job in ["worker", "ps"] -%}
{%- for i in range(replicas[job]) -%}
kind: Service
apiVersion: v1
metadata:
  name: {{ name }}-{{ job }}-{{ i }}
spec:
  selector:
    name: {{ name }}
    job: {{ job }}
    task: "{{ i }}"
  ports:
  - port: {{ port }}
---
kind: ReplicaSet
apiVersion: extensions/v1beta1
metadata:
  name: {{ name }}-{{ job }}-{{ i }}
spec:
  replicas: 1
  template:
    metadata:
      labels:
        name: {{ name }}
        job: {{ job }}
        task: "{{ i }}"
    spec:
      containers:
      - name: tensorflow
        image: {{ image }}
{% if credential_secret_name != "" %}
        env:
        - name: GOOGLE_APPLICATION_CREDENTIALS
          value: "/etc/credential/{{ credential_secret_key }}"
{% endif %}
        ports:
        - containerPort: {{ port }}
{% if job == "tensorboard" %}
        command:
        - "tensorboard"
        args:
        - "--logdir={{ train_dir }}"
{% else %}
        command:
        - "/usr/bin/python"
        - "{{ script }}"
        args:
        - "--data_dir={{ data_dir }}"
        - "--train_dir={{ train_dir }}"
        - "--task_index={{ i }}"
        - "--job_name={{ job }}"
        - "--worker_hosts={{ worker_hosts() }}"
        - "--ps_hosts={{ ps_hosts() }}"
{% endif %}
{% if credential_secret_name != "" %}
        volumeMounts:
        - name: credential
          mountPath: /etc/credential
          readOnly: true
      volumes:
      - name: credential
        secret:
          secretName: {{ credential_secret_name }}
{% endif %}
---
{% endfor %}
{%- endfor -%}