Update Watson training component to output model location UID (#1587)

* Update wml-train.py * update watson training components to output model location * update variables
kubeflow · Jul 3, 2019 · f64782c · f64782c
1 parent f0c8432
commit f64782c
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 5 deletions.
diff --git a/components/ibm-components/watson/train/component.yaml b/components/ibm-components/watson/train/component.yaml
@@ -28,6 +28,7 @@ inputs:
   - {name: author_name,       description: 'Name of this training job author', default: 'default-author'}
 outputs:
   - {name: run_uid,           description: 'UID for the Watson Machine Learning training-runs'}
+  - {name: training_uid,      description: 'Training Location UID for the Watson Machine Learning training-runs'}
 implementation:
   container:
     image: docker.io/aipipeline/wml-train:latest
@@ -47,3 +48,4 @@ implementation:
     ]
     fileOutputs:
       run_uid: /tmp/run_uid
+      training_uid: /tmp/training_uid
diff --git a/components/ibm-components/watson/train/src/wml-train.py b/components/ibm-components/watson/train/src/wml-train.py
@@ -133,6 +133,13 @@ def train(args):
         status = client.training.get_status( run_uid )
     print(status)
 
+    # Get training details
+    training_details = client.training.get_details(run_uid)
+    with open("/tmp/training_uid", "w") as f:
+        training_uid = training_details['entity']['training_results_reference']['location']['model_location']
+        f.write(training_uid)
+    f.close()
+
 if __name__ == "__main__":
     import argparse
     parser = argparse.ArgumentParser()

diff --git a/samples/ibm-samples/watson/watson_train_serve_pipeline.py b/samples/ibm-samples/watson/watson_train_serve_pipeline.py
@@ -19,10 +19,10 @@
 import ai_pipeline_params as params
 
 secret_name = 'kfp-creds'
-configuration_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/56134b8bf3eac9ea124d07c818724e2a64789c11/components/ibm-components/commons/config/component.yaml')
-train_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/56134b8bf3eac9ea124d07c818724e2a64789c11/components/ibm-components/watson/train/component.yaml')
-store_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/56134b8bf3eac9ea124d07c818724e2a64789c11/components/ibm-components/watson/store/component.yaml')
-deploy_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/56134b8bf3eac9ea124d07c818724e2a64789c11/components/ibm-components/watson/deploy/component.yaml')
+configuration_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/master/components/ibm-components/commons/config/component.yaml')
+train_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/master/components/ibm-components/watson/train/component.yaml')
+store_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/master/components/ibm-components/watson/store/component.yaml')
+deploy_op = components.load_component_from_url('https://raw.githubusercontent.com/kubeflow/pipelines/master/components/ibm-components/watson/deploy/component.yaml')
 
 # create pipelines
 
@@ -66,7 +66,7 @@ def kfp_wml_pipeline(
 
     # op3 - this operation stores the model trained above
     wml_store = store_op(
-                   wml_train.output,
+                   wml_train.outputs['run_uid'],
                    model_name
                   ).apply(params.use_ai_pipeline_params(secret_name))