apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
  generateName: tf-workflow-
spec:
  entrypoint: deploy-model 
  # 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: workflow
      value: workflow-name
  templates:
  - name: deploy-model
    steps:
      - - name: get-workflow-info
          template: get-workflow-info
      - - name: serve-model
          template: tf-inference
          arguments:
            parameters:
            - name: s3-model-url
              value: "{{steps.get-workflow-info.outputs.parameters.s3-model-url}}"
            - name: s3-exported-url
              value: "{{steps.get-workflow-info.outputs.parameters.s3-exported-url}}"
            - name: aws-secret 
              value: "{{steps.get-workflow-info.outputs.parameters.aws-secret}}"
            - name: namespace
              value: "{{steps.get-workflow-info.outputs.parameters.namespace}}"
            - name: aws-region 
              value: "{{steps.get-workflow-info.outputs.parameters.aws-region}}"
            - name: s3-endpoint
              value: "{{steps.get-workflow-info.outputs.parameters.s3-endpoint}}"
            - name: s3-use-https
              value: "{{steps.get-workflow-info.outputs.parameters.s3-use-https}}"
            - name: s3-verify-ssl
              value: "{{steps.get-workflow-info.outputs.parameters.s3-verify-ssl}}"
            - name: job-name
              value: "{{steps.get-workflow-info.outputs.parameters.job-name}}"
            - name: tf-serving-image 
              value: "{{steps.get-workflow-info.outputs.parameters.tf-serving-image}}"
            - name: model-serving-servicetype
              value: "{{steps.get-workflow-info.outputs.parameters.model-serving-servicetype}}"
            - name: model-serving-ks-url
              value: "{{steps.get-workflow-info.outputs.parameters.model-serving-ks-url}}"
            - name: model-serving-ks-tag
              value: "{{steps.get-workflow-info.outputs.parameters.model-serving-ks-tag}}"
            - name: model-name
              value: "{{steps.get-workflow-info.outputs.parameters.model-name}}"
  - name: get-workflow-info
    container:
      image: nervana/circleci:master
      imagePullPolicy: Always
      command: ["bash", "-c", "-x", 'for var in s3-model-url s3-exported-url; do kubectl get workflow {{workflow.parameters.workflow}} -o json | jq -r ".status.nodes[] | select(.name|contains(\"get-workflow-info\")) | .outputs.parameters[] | select(.name == \"${var}\") | .value" > /tmp/${var} ; done; for var in job-name namespace aws-secret aws-region s3-endpoint s3-use-https s3-verify-ssl tf-serving-image model-serving-servicetype model-serving-ks-url model-serving-ks-tag model-name; do kubectl get workflow {{workflow.parameters.workflow}} -o jsonpath="{.spec.arguments.parameters[?(.name==\"${var}\")].value}" > /tmp/${var}; done']
    outputs:
      parameters:
      - name: s3-model-url
        valueFrom:
          path: /tmp/s3-model-url
      - name: s3-exported-url
        valueFrom:
          path: /tmp/s3-exported-url
      - name: aws-secret
        valueFrom:
          path: /tmp/aws-secret
      - name: namespace
        valueFrom:
          path: /tmp/namespace
      - name: aws-region
        valueFrom:
          path: /tmp/aws-region
      - name: s3-endpoint
        valueFrom:
          path: /tmp/s3-endpoint
      - name: s3-use-https
        valueFrom:
          path: /tmp/s3-use-https
      - name: s3-verify-ssl
        valueFrom:
          path: /tmp/s3-verify-ssl
      - name: job-name
        valueFrom:
          path: /tmp/job-name
      - name: tf-serving-image
        valueFrom:
          path: /tmp/tf-serving-image
      - name: model-serving-servicetype
        valueFrom:
          path: /tmp/model-serving-servicetype
      - name: model-serving-ks-url
        valueFrom:
          path: /tmp/model-serving-ks-url
      - name: model-serving-ks-tag
        valueFrom:
          path: /tmp/model-serving-ks-tag
      - name: model-name
        valueFrom:
          path: /tmp/model-name
  - name: tf-inference
    inputs:
      parameters:
      - name: s3-model-url
      - name: s3-exported-url
      - name: aws-secret
      - name: namespace
      - name: aws-region
      - name: s3-endpoint
      - name: s3-use-https
      - name: s3-verify-ssl
      - name: job-name
      - name: tf-serving-image
      - name: model-serving-servicetype
      - name: model-serving-ks-url
      - name: model-serving-ks-tag
      - name: model-name
    script:
      image: elsonrodriguez/ksonnet:0.10.1
      command: ["/ksonnet-entrypoint.sh"]
      source: |
        ks init my-model-server
        cd my-model-server
        ks registry add kubeflow {{inputs.parameters.model-serving-ks-url}}
        ks pkg install kubeflow/tf-serving@{{inputs.parameters.model-serving-ks-tag}}
        ks env add default
        # TODO change mnist name to be specific to a job. Right now mnist name is required to serve the model.
        ks generate tf-serving {{inputs.parameters.model-name}} --name=mnist-{{inputs.parameters.job-name}} --namespace={{inputs.parameters.namespace}} --model_path={{inputs.parameters.s3-exported-url}}
        ks param set {{inputs.parameters.model-name}} model_server_image {{inputs.parameters.tf-serving-image}}
        ks param set {{inputs.parameters.model-name}} model_name {{inputs.parameters.model-name}}
        ks param set {{inputs.parameters.model-name}} namespace {{inputs.parameters.namespace}}
        ks param set {{inputs.parameters.model-name}} service_type {{inputs.parameters.model-serving-servicetype}}
        ks param set {{inputs.parameters.model-name}} s3_create_secret false
        ks param set {{inputs.parameters.model-name}} s3_secret_name {{inputs.parameters.aws-secret}}
        ks param set {{inputs.parameters.model-name}} s3_secret_accesskeyid_key_name awsAccessKeyID
        ks param set {{inputs.parameters.model-name}} s3_secret_secretaccesskey_key_name awsSecretAccessKey
        ks param set {{inputs.parameters.model-name}} s3_aws_region {{inputs.parameters.aws-region}}
        ks param set {{inputs.parameters.model-name}} s3_endpoint {{inputs.parameters.s3-endpoint}}
        ks param set {{inputs.parameters.model-name}} s3_use_https {{inputs.parameters.s3-use-https}} --as-string
        ks param set {{inputs.parameters.model-name}} s3_verify_ssl {{inputs.parameters.s3-verify-ssl}} --as-string
        ks apply default -c {{inputs.parameters.model-name}}
      #FIXME This doesn't actually work in the current version of argo. We're using a default of `tf-user` in the container entrypoint for now.
      env:
      - name: SERVICE_ACCOUNT
        value: tf-user
