template.yaml.jinja

{%- set name = "criteo-sample-kubernetes" -%}
{%- set has_eval = False -%}
{%- set has_tensorboard = False -%}
{%- set script = "trainer/trainer.py" -%}
{%- set port = 5000 -%}
{% set cmdline_arg_list = cmdline_args.split(" ") %}

{%- set replicas = {"worker": num_workers,
                    "chief": num_chief,
                    "ps": num_ps,
                    "evaluator": has_eval|int,
                    "tensorboard": has_tensorboard|int} -%}

{%- macro chief_host() -%}
    \"{{ name }}-chief-0:{{ port }}\"
{%- endmacro -%}

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

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

{%- macro tf_config(task_type, task_id) -%}
{
  \"cluster\": {
    {%- if num_chief > 0 -%}
    \"chief\": [{{ chief_host() }}]
    {%- endif -%}
    {%- if num_workers > 0 -%},
    \"worker\": [{{ worker_hosts() }}]
    {%- endif -%}
    {%- if num_ps > 0 -%}, \"ps\": [{{ ps_hosts() }}]{%- endif -%}
    {%- if has_eval -%},
    \"evaluator\": [\"{{ name }}-evaluator-0:{{ port }}\"]{%- endif -%}
  },
  \"task\": {
    \"type\":  \"{{ task_type }}\",
    \"index\": {{ task_id }}
  },
  \"environment\": \"cloud\"
}
{%- endmacro -%}

{% for job in ["worker", "ps", "evaluator", "tensorboard", "chief"] -%}
{%- for i in range(replicas[job]) -%}
kind: Service
apiVersion: v1
metadata:
  name: {{ name }}-{{ job }}-{{ i }}
spec:
  type: ClusterIP
  clusterIP: None
  selector:
    job: {{ name }}-{{ job }}-{{ i }}
  ports:
  - port: {{ port }}
---
kind: Pod
apiVersion: v1
metadata:
  name: {{ name }}-{{ job }}-{{ i }}
  labels:
    job: {{ name }}-{{ job }}-{{ i }}
  annotations:
    tf-version.cloud-tpus.google.com: "2.1"
spec:
  hostNetwork: true
  dnsPolicy: ClusterFirstWithHostNet
  restartPolicy: Never
{% if job == "tensorboard" %}
  containers:
  - name: tensorflow
    image: tensorflow/tensorflow
{% else %}
  containers:
  - name: tensorflow
    image: {{ image }}
    imagePullPolicy: Always
    securityContext:
      privileged: true
    resources:
      limits:
{% if num_tpus > 0 %}
        cloud-tpus.google.com/v2: {{ num_tpus }}
{% endif %}
        cpu: 3
{% if job == "ps" %}
        memory: 12G
{% else %}
        memory: 12G
        nvidia.com/gpu: {{ num_gpus_per_worker }}
{% endif %}
{% endif %}
    env:
{% if job != "tensorboard" %}
    - name: TF_CONFIG
      value: "{{ tf_config(job, i) }}"
{% endif %}
    ports:
    - containerPort: {{ port }}
{% if job == "tensorboard" %}
    command:
    - "tensorboard"
    args:
    - "--logdir={{ train_dir }}"
    - "--port={{ port }}"
{% else %}
    command:
    - "python"
    - "{{ script }}"
    {%- for cmdline_arg in cmdline_arg_list %}
    - "{{ cmdline_arg }}"
    {%- endfor -%}

{% endif %}
---
{% endfor %}
{%- endfor -%}
apiVersion: policy/v1beta1
kind: PodSecurityPolicy
metadata:
  name: privileged
  annotations:
    seccomp.security.alpha.kubernetes.io/allowedProfileNames: '*'
spec:
  privileged: true
  allowPrivilegeEscalation: true
  allowedCapabilities:
  - '*'
  volumes:
  - '*'
  hostNetwork: true
  hostPorts:
  - min: 0
    max: 65535
  hostIPC: true
  hostPID: true
  runAsUser:
    rule: 'RunAsAny'
  seLinux:
    rule: 'RunAsAny'
  supplementalGroups:
    rule: 'RunAsAny'
  fsGroup:
    rule: 'RunAsAny'