tfargo.yaml

apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
  generateName: tf-workflow-
spec:
  entrypoint: tests
  onExit: exit-handler
  # Parameters can be passed/overridden via the argo CLI.
  # To override the printed message, run `argo submit` with the -p option:
  # $ argo submit examples/arguments-parameters.yaml -p message="goodbye world"
  arguments:
    parameters:
    - name: datasets-image
      # default image
      value: gcr.io/constant-cubist-173123/tf_workflow:nan_workflow
    - name: tf-master # number of tf masters
      value: 1
    - name: tf-worker # number of tf workers
      value: 1
    - name: tf-ps # number of tf parameter servers
      value: 2
    - name: tf-image
      value: elsonrodriguez/mytfmodel:1.45
    - name: tf-server-image
      value: elsonrodriguez/mytfserver:1.6
    - name: job-name
      value: job21
    - name: namespace
      value: default
    - name: s3sourceURL
      value: s3://tfoperator/data/mnist/
    - name: s3resultURL
      value: s3://tfoperator/models
    - name: aws-endpoint-url
      value: https://s3.amazonaws.com
    - name: aws-secret
      value: aws-creds
    - name: gcsresultURL
      value: s3://superpuke/models
  volumes:
  - name: training-data
    emptyDir: {}
  - name: training-output
  templates:
  - name: tests
    steps:
      - - name: mount-datasets
          template: kvc
      - - name: get-volume
          template: get-volumemanager-info
      - - name: train
          template: tf-train
          arguments:
            parameters:
            - name: nodeaffinity
              value: "{{steps.get-volume.outputs.parameters.nodeaffinity}}"
            - name: hostpath
              value: "{{steps.get-volume.outputs.parameters.hostpath}}"
      - - name: export-model
          template: tf-export
      - - name: inference
          template: tf-inference
    # - - name: prepare-datasets
    #     template: download-and-convert-data
  - name: exit-handler
    steps:
      - - name: cleanup
          template: clean
  # - name: download-and-convert-data
  #   container:
  #     image: gcr.io/constant-cubist-173123/tf_workflow:nan_workflow
  #     imagePullPolicy: Always
  #     command: ["bash", "-c",  "python /notebooks/models/research/slim/download_and_convert_data.py --dataset_name=flowers --dataset_dir=/tmp/data"]
  #     outputs:
  #       artifacts:
  #       - name: training-data
  #         path: /tmp/data
  - name: kvc
    resource:
      action: apply
      successCondition: status.state == Running
      manifest: |
        apiVersion: aipg.intel.com/v1
        kind: VolumeManager
        metadata:
          name: mnist
          namespace: {{workflow.parameters.namespace}}
        spec:
          volumeConfigs:
            - id: "mnist-{{workflow.uid}}"
              replicas: 1
              sourceType: "S3"
              sourceURL: "{{workflow.parameters.s3sourceURL}}"
              AccessMode: "ReadWriteOnce"
              Capacity: 1Gi
              Labels:
                tfdata: mnist
              Options:
                awsCredentialsSecretName: {{workflow.parameters.aws-secret}}
  - name: get-volumemanager-info
    container:
      image: nervana/circleci:master
      imagePullPolicy: Always
      command: ["bash", "-c", "kubectl get volumemanager mnist -o json | jq -r '.status.volumes[].volumeSource.hostPath.path' | tee /tmp/hostpath; kubectl get volumemanager mnist -o json | jq '.status.volumes[].nodeAffinity' | tee /tmp/nodeaffinity"]
    outputs:
      parameters:
      - name: hostpath
        valueFrom:
          path: /tmp/hostpath
      - name: nodeaffinity
        valueFrom:
          path: /tmp/nodeaffinity
  - name: tf-train
    inputs:
      parameters:
      - name: nodeaffinity
      - name: hostpath
    resource:
      action: apply
      # NOTE: need to detect master node complete
      successCondition: status.state == Succeeded
      manifest: |
        apiVersion: "kubeflow.org/v1alpha1"
        kind: "TFJob"
        metadata:
          name: {{workflow.parameters.job-name}}
          namespace: {{workflow.parameters.namespace}}
        spec:
          replicaSpecs:
            - replicas: {{workflow.parameters.tf-master}}
              tfReplicaType: MASTER
              template:
                spec:
                  affinity:
                    nodeAffinity:
                      {{inputs.parameters.nodeaffinity}}
                  serviceAccountName: tf-job-operator
                  containers:
                    - image: {{workflow.parameters.tf-image}}
                      name: tensorflow
                      imagePullPolicy: Always
                      args: ["python", "/opt/model.py", "--data_dir=/tmp/data", "--train_dir=/tmp/train", "--download=true", "--sync_replicas=true"]
                      volumeMounts:
                      - name: training-result
                        mountPath: /tmp/train
                      - name: training-data
                        mountPath: /tmp/data
                    - image: nervana/circleci:jose_wait_for_master
                      name: upload
                      env:
                      - name: POD_NAME
                        valueFrom:
                          fieldRef:
                            apiVersion: v1
                            fieldPath: metadata.name
                      - name: POD_NAMESPACE
                        valueFrom:
                          fieldRef:
                            apiVersion: v1
                            fieldPath: metadata.namespace
                      - name: AWS_ACCESS_KEY_ID
                        valueFrom:
                          secretKeyRef:
                            name: {{workflow.parameters.aws-secret}}
                            key: awsAccessKeyID
                      - name: AWS_SECRET_ACCESS_KEY
                        valueFrom:
                          secretKeyRef:
                            name: {{workflow.parameters.aws-secret}}
                            key: awsSecretAccessKey
                      command: ['sh', '-c', "./bin/wait_for_master $POD_NAMESPACE $POD_NAME; aws --endpoint-url {{workflow.parameters.aws-endpoint-url}} s3 cp --recursive /tmp/train {{workflow.parameters.s3resultURL}}/unfinished/{{workflow.parameters.job-name}}/"]
                      volumeMounts:
                      - name: training-result
                        mountPath: /tmp/train
                  volumes:
                  - name: training-result
                    emptyDir: {}
                  - name: training-data
                    hostPath:
                      path: {{inputs.parameters.hostpath}}
                  restartPolicy: OnFailure
            - replicas: {{workflow.parameters.tf-worker}}
              tfReplicaType: WORKER
              template:
                spec:
                  affinity:
                    nodeAffinity:
                      {{inputs.parameters.nodeaffinity}}
                  serviceAccountName: tf-job-operator
                  containers:
                    - image: {{workflow.parameters.tf-server-image}}
                      name: tensorflow
                      imagePullPolicy: Always
                      volumeMounts:
                      - name: training-result
                        mountPath: /tmp/train
                      - name: training-data
                        mountPath: /tmp/data
                    - image: nervana/circleci:jose_wait_for_master
                      name: upload
                      env:
                      - name: POD_NAME
                        valueFrom:
                          fieldRef:
                            apiVersion: v1
                            fieldPath: metadata.name
                      - name: POD_NAMESPACE
                        valueFrom:
                          fieldRef:
                            apiVersion: v1
                            fieldPath: metadata.namespace
                      - name: AWS_ACCESS_KEY_ID
                        valueFrom:
                          secretKeyRef:
                            name: {{workflow.parameters.aws-secret}}
                            key: awsAccessKeyID
                      - name: AWS_SECRET_ACCESS_KEY
                        valueFrom:
                          secretKeyRef:
                            name: {{workflow.parameters.aws-secret}}
                            key: awsSecretAccessKey
                      command: ['sh', '-c', "./bin/wait_for_master $POD_NAMESPACE $POD_NAME; aws --endpoint-url {{workflow.parameters.aws-endpoint-url}} s3 cp --recursive /tmp/train {{workflow.parameters.s3resultURL}}/unfinished/{{workflow.parameters.job-name}}/"]
                      volumeMounts:
                      - name: training-result
                        mountPath: /tmp/train
                  volumes:
                  - name: training-result
                    emptyDir: {}
                  - name: training-data
                    hostPath:
                      path: {{inputs.parameters.hostpath}}
                  restartPolicy: OnFailure
            - replicas: {{workflow.parameters.tf-ps}}
              tfReplicaType: PS
              template:
                spec:
                  containers:
                    - image: {{workflow.parameters.tf-server-image}}
                      name: tensorflow
                      imagePullPolicy: Always
                  restartPolicy: OnFailure
  - name: tf-export
    resource:
      action: apply
      successCondition: status.phase == Succeeded
      failureCondition: status.phase == Failed
      manifest: |
        apiVersion: v1
        kind: Pod
        metadata:
          name: exporter-{{workflow.parameters.job-name}}
        spec:
          restartPolicy: Never
          containers:
          - name: s3-download-unfinished
            image: volumecontroller/aws-cli:latest
            env:
            - name: AWS_ACCESS_KEY_ID
              valueFrom:
                secretKeyRef:
                  name: {{workflow.parameters.aws-secret}}
                  key: awsAccessKeyID
            - name: AWS_SECRET_ACCESS_KEY
              valueFrom:
                secretKeyRef:
                  name: {{workflow.parameters.aws-secret}}
                  key: awsSecretAccessKey
            - name: AWS_ENDPOINT_URL
              value: {{workflow.parameters.aws-endpoint-url}}
            command: ['sh', '-c', "aws s3 cp --recursive {{workflow.parameters.s3resultURL}}/unfinished/{{workflow.parameters.job-name}}/ /tmp/train && touch /tmp/download-done/done"]
            volumeMounts:
            - name: training-result
              mountPath: /tmp/train
            - name: download-done
              mountPath: /tmp/download-done
          - image: elsonrodriguez/mytfmodel:1.45
            name: export
            imagePullPolicy: Always
            command: ['/bin/bash','-c', 'while [ ! -f /tmp/download-done/done ]; do sleep 5; done; ls -r /tmp/export/model; python /opt/export.py --checkpoint_dir=/tmp/train --output_dir=/tmp/export/model && touch /tmp/export-done/done']
            volumeMounts:
            - name: training-result
              mountPath: /tmp/train
            - name: download-done
              mountPath: /tmp/download-done
            - name: export-done
              mountPath: /tmp/export-done
            - name: export-result
              mountPath: /tmp/export
          - name: s3-upload-exported
            image: volumecontroller/aws-cli:latest
            env:
            - name: AWS_ACCESS_KEY_ID
              valueFrom:
                secretKeyRef:
                  name: gcs-creds
                  key: awsAccessKeyID
            - name: AWS_SECRET_ACCESS_KEY
              valueFrom:
                secretKeyRef:
                  name: gcs-creds
                  key: awsSecretAccessKey
            - name: AWS_ENDPOINT_URL
              value: https://storage.googleapis.com
            command: ['sh', '-c', 'while [ ! -f /tmp/export-done/done ]; do sleep 5; done; aws s3 cp --recursive /tmp/export/model/ {{workflow.parameters.gcsresultURL}}/exported/{{workflow.parameters.job-name}}/']
            volumeMounts:
            - name: export-result
              mountPath: /tmp/export
            - name: export-done
              mountPath: /tmp/export-done
          volumes:
          - name: training-result
            emptyDir: {}
          - name: export-result
            emptyDir: {}
          - name: download-done
            emptyDir: {}
          - name: export-done
            emptyDir: {}
  - name: tf-inference
    script:
      image: elsonrodriguez/ksonnet:0.8.0-test6
      command: ["/ksonnet-entrypoint.sh"]
      source: |
        MODEL_COMPONENT=serveMnist
        MODEL_PATH=gs://superpuke/models/exported/{{workflow.parameters.job-name}} #{{workflow.parameters.gcsresultURL}}/exported/{{workflow.parameters.job-name}}
        MODEL_SERVER_IMAGE=gcr.io/constant-cubist-173123/model-server:1.0

        ks init my-model-server
        cd my-model-server
        ks registry add kubeflow github.com/kubeflow/kubeflow/tree/master/kubeflow
        ks pkg install kubeflow/tf-serving
        ks env add  cloud
        # TODO change mnist name to be specific to a job. Right now mnist name is required to serve the model.
        ks generate tf-serving ${MODEL_COMPONENT} --name=mnist2 --namespace={{workflow.parameters.namespace}} --model_path=${MODEL_PATH} --model_server_image=${MODEL_SERVER_IMAGE}
        ks apply cloud -c ${MODEL_COMPONENT}
      env:
      - name: SERVICE_ACCOUNT
        value: argo
  - name: clean
    container:
      image: nervana/circleci:master
      imagePullPolicy: Always
      command: ["bash", "-c", "kubectl delete tfjob {{workflow.parameters.job-name}} || true; kubectl delete pod exporter-{{workflow.parameters.job-name}} || true"]

# NOTE: export mnodel
#   - name: tf_model_export
#                 spec:
#                   containers:
#                     - image: elsonrodriguez/mytfmodel:1.45
#                       name: tensorflow
#     command: python export.py  --checkpoint_dir=/combined_training_results/ --output_dir=/saved_model