# helm_charts_statefulset-ps.yaml
apiVersion: apps/v1
kind: StatefulSet
metadata:
  name: {{ .Release.Name }}-ps
  labels:
    app: {{ template "distributed-tensorflow.name" . }}
    chart: {{ template "distributed-tensorflow.chart" . }}
    release: {{ .Release.Name }}
    heritage: {{ .Release.Service }}
    role: ps
spec:
  selector:
    matchLabels:
      app: {{ template "distributed-tensorflow.name" . }}
      release: {{ .Release.Name }}
      role: ps
  serviceName: {{ template "distributed-tensorflow.fullname" . }}-ps
  podManagementPolicy: {{ .Values.ps.podManagementPolicy }}
  replicas: {{.Values.ps.number}}
  template:
    metadata:
      labels:
        app: {{ template "distributed-tensorflow.name" . }}
        chart: {{ template "distributed-tensorflow.chart" . }}
        release: {{ .Release.Name }}
        heritage: {{ .Release.Service }}
        role: ps
    spec:
{{- if .Values.volumes }}
      volumes:
{{ toYaml .Values.volumes | indent  6 }}
{{- end }}
      containers:
      - name: ps
        image: "{{ .Values.ps.image.repository }}:{{ .Values.ps.image.tag }}"
        imagePullPolicy: {{ .Values.ps.image.pullPolicy }}
        command:
        - "python"
        - "train_distributed.py"
{{- if .Values.hyperparams.learningrate }}
        - --learning_rate
        - "{{ .Values.hyperparams.learningrate }}"
{{- end }}
{{- if .Values.hyperparams.batchsize }}
        - --batch_size
        - "{{ .Values.hyperparams.batchsize }}"
{{- end }}
{{- if .Values.hyperparams.trainsteps }}
        - --train_steps
        - "{{ .Values.hyperparams.trainsteps }}"
{{- end }}
{{- if .Values.hyperparams.datadir }}
        - --data_dir
        - "{{ .Values.hyperparams.datadir }}"
{{- end }}
{{- if .Values.hyperparams.logdir }}
        - --log_dir
        - "{{ .Values.hyperparams.logdir }}"
{{- end }}
{{- if .Values.hyperparams.hiddenunits }}
        - --hidden_units
        - "{{ .Values.hyperparams.hiddenunits }}"
{{- end }}
        env:
        - name: WORKER_HOSTS
          valueFrom:
            configMapKeyRef:
              name: {{ template "distributed-tensorflow.fullname" . }}
              key: worker.hostList
        - name: PS_HOSTS
          valueFrom:
            configMapKeyRef:
              name: {{ template "distributed-tensorflow.fullname" . }}
              key: ps.hostList
        - name: POD_NAME
          valueFrom:
            fieldRef:
              fieldPath: metadata.name
        - name: JOB_NAME
          value: ps
       {{- if .Values.ps.env }}
       {{- range $key, $value := .Values.ps.env }}
        - name: "{{ $key }}"
          value: "{{ $value }}"
       {{- end }}
       {{- end }}
{{- if .Values.ps.privileged }}
        securityContext:
          privileged: true
{{- end }}
        ports:
        - containerPort: {{ .Values.ps.port }}
{{- if .Values.volumeMounts }}
        volumeMounts:
{{ toYaml .Values.volumeMounts | indent 8 }}
{{- end }}
{{- if .Values.ps.resources }}
        resources:
{{ toYaml .Values.ps.resources | indent 10 }}
{{- end }}