Repository: tensorflow/tfx

Issue:
Support for deterministic job name generation in the Trainer component
Support for deterministic job name generation in the Trainer component is needed to smoothly enable [HP tuning on AIP Training](https://cloud.google.com/ml-engine/docs/using-hyperparameter-tuning). Indeed, [retrieving](https://cloud.google.com/ml-engine/docs/monitor-training#checking_job_status) the best run requires the job name. 


Discussion:

Resolving patch:
diff --git a/tfx/extensions/google_cloud_ai_platform/runner.py b/tfx/extensions/google_cloud_ai_platform/runner.py
--- a/tfx/extensions/google_cloud_ai_platform/runner.py
+++ b/tfx/extensions/google_cloud_ai_platform/runner.py
@@ -21,7 +21,7 @@
 import json
 import sys
 import time
-from typing import Any, Dict, List, Text
+from typing import Any, Dict, List, Optional, Text
 
 import absl
 from googleapiclient import discovery
@@ -82,9 +82,10 @@ def _get_caip_python_version() -> Text:
 
 def start_aip_training(input_dict: Dict[Text, List[types.Artifact]],
                        output_dict: Dict[Text, List[types.Artifact]],
-                       exec_properties: Dict[Text, Any],
-                       executor_class_path: Text, training_inputs: Dict[Text,
-                                                                        Any]):
+                       exec_properties: Dict[Text,
+                                             Any], executor_class_path: Text,
+                       training_inputs: Dict[Text,
+                                             Any], job_id: Optional[Text]):
   """Start a trainer job on AI Platform (AIP).
 
   This is done by forwarding the inputs/outputs/exec_properties to the
@@ -95,27 +96,25 @@ def start_aip_training(input_dict: Dict[Text, List[types.Artifact]],
     output_dict: Passthrough input dict for tfx.components.Trainer.executor.
     exec_properties: Passthrough input dict for tfx.components.Trainer.executor.
     executor_class_path: class path for TFX core default trainer.
-    training_inputs: Training input for AI Platform training job.
-      'pythonModule', 'pythonVersion' and 'runtimeVersion' will be inferred by
-      the runner. For the full set of parameters supported, refer to
+    training_inputs: Training input argment for AI Platform training job.
+      'pythonModule', 'pythonVersion' and 'runtimeVersion' will be inferred. For
+      the full set of parameters, refer to
       https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#TrainingInput
-
+    job_id: Job ID for AI Platform Training job. If not supplied,
+      system-determined unique ID is given. Refer to
+    https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#resource-job
   Returns:
     None
   Raises:
     RuntimeError: if the Google Cloud AI Platform training job failed.
   """
   training_inputs = training_inputs.copy()
-  # Remove aip_args from exec_properties so AIP trainer doesn't call itself
-  for gaip_training_key in ['ai_platform_training_args', 'gaip_training_args']:
-    if gaip_training_key in exec_properties.get('custom_config'):
-      exec_properties['custom_config'].pop(gaip_training_key)
 
   json_inputs = artifact_utils.jsonify_artifact_dict(input_dict)
   absl.logging.info('json_inputs=\'%s\'.', json_inputs)
   json_outputs = artifact_utils.jsonify_artifact_dict(output_dict)
   absl.logging.info('json_outputs=\'%s\'.', json_outputs)
-  json_exec_properties = json.dumps(exec_properties)
+  json_exec_properties = json.dumps(exec_properties, sort_keys=True)
   absl.logging.info('json_exec_properties=\'%s\'.', json_exec_properties)
 
   # Configure AI Platform training job
@@ -143,20 +142,21 @@ def start_aip_training(input_dict: Dict[Text, List[types.Artifact]],
   project = training_inputs.pop('project')
   project_id = 'projects/{}'.format(project)
 
-  job_name = 'tfx_' + datetime.datetime.now().strftime('%Y%m%d%H%M%S')
-  job_spec = {'jobId': job_name, 'trainingInput': training_inputs}
+  # 'tfx_YYYYmmddHHMMSS' is the default job ID if not explicitly specified.
+  job_id = job_id or 'tfx_%s' % datetime.datetime.now().strftime('%Y%m%d%H%M%S')
+  job_spec = {'jobId': job_id, 'trainingInput': training_inputs}
 
   # Submit job to AIP Training
   absl.logging.info(
       'Submitting job=\'{}\', project=\'{}\' to AI Platform.'.format(
-          job_name, project))
+          job_id, project))
   request = api_client.projects().jobs().create(
       body=job_spec, parent=project_id)
   request.execute()
 
   # Wait for AIP Training job to finish
-  job_id = '{}/jobs/{}'.format(project_id, job_name)
-  request = api_client.projects().jobs().get(name=job_id)
+  job_name = '{}/jobs/{}'.format(project_id, job_id)
+  request = api_client.projects().jobs().get(name=job_name)
   response = request.execute()
   while response['state'] not in ('SUCCEEDED', 'FAILED'):
     time.sleep(_POLLING_INTERVAL_IN_SECONDS)
diff --git a/tfx/extensions/google_cloud_ai_platform/trainer/executor.py b/tfx/extensions/google_cloud_ai_platform/trainer/executor.py
--- a/tfx/extensions/google_cloud_ai_platform/trainer/executor.py
+++ b/tfx/extensions/google_cloud_ai_platform/trainer/executor.py
@@ -42,8 +42,9 @@ def Do(self, input_dict: Dict[Text, List[types.Artifact]],
       output_dict: Passthrough input dict for tfx.components.Trainer.executor.
       exec_properties: Mostly a passthrough input dict for
         tfx.components.Trainer.executor. custom_config.ai_platform_training_args
-        is consumed by this class.  For the full set of parameters supported by
-        Google Cloud AI Platform, refer to
+        and custom_config.ai_platform_training_job_id are consumed by this
+        class.  For the full set of parameters supported by Google Cloud AI
+        Platform, refer to
         https://cloud.google.com/ml-engine/docs/tensorflow/training-jobs#configuring_the_job
 
     Returns:
@@ -55,15 +56,16 @@ def Do(self, input_dict: Dict[Text, List[types.Artifact]],
     """
     self._log_startup(input_dict, output_dict, exec_properties)
 
-    if not exec_properties.get('custom_config',
-                               {}).get('ai_platform_training_args'):
+    custom_config = exec_properties.get('custom_config', {})
+    training_inputs = custom_config.get('ai_platform_training_args')
+    if training_inputs is None:
       err_msg = '\'ai_platform_training_args\' not found in custom_config.'
       absl.logging.error(err_msg)
       raise ValueError(err_msg)
 
-    training_inputs = exec_properties.get('custom_config',
-                                          {}).pop('ai_platform_training_args')
+    job_id = custom_config.get('ai_platform_training_job_id')
     executor_class_path = '%s.%s' % (tfx_trainer_executor.Executor.__module__,
                                      tfx_trainer_executor.Executor.__name__)
     return runner.start_aip_training(input_dict, output_dict, exec_properties,
-                                     executor_class_path, training_inputs)
+                                     executor_class_path, training_inputs,
+                                     job_id)


Test patch:
diff --git a/tfx/extensions/google_cloud_ai_platform/runner_test.py b/tfx/extensions/google_cloud_ai_platform/runner_test.py
--- a/tfx/extensions/google_cloud_ai_platform/runner_test.py
+++ b/tfx/extensions/google_cloud_ai_platform/runner_test.py
@@ -30,6 +30,7 @@
 class RunnerTest(tf.test.TestCase):
 
   def setUp(self):
+    super(RunnerTest, self).setUp()
     self._output_data_dir = os.path.join(
         os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', self.get_temp_dir()),
         self._testMethodName)
@@ -40,9 +41,10 @@ def setUp(self):
     self._training_inputs = {
         'project': self._project_id,
     }
+    self._job_id = 'my_jobid'
     self._exec_properties = {
         'custom_config': {
-            'gaip_training_args': self._training_inputs
+            'ai_platform_training_args': self._training_inputs,
         },
     }
     self._ai_platform_serving_args = {
@@ -67,7 +69,7 @@ def testStartAIPTraining(self, mock_discovery):
 
     runner.start_aip_training(self._inputs, self._outputs,
                               self._exec_properties, class_path,
-                              self._training_inputs)
+                              self._training_inputs, None)
 
     mock_create.assert_called_with(
         body=mock.ANY, parent='projects/{}'.format(self._project_id))
@@ -82,7 +84,8 @@ def testStartAIPTraining(self, mock_discovery):
             },
             'args': [
                 '--executor_class_path', class_path, '--inputs', '{}',
-                '--outputs', '{}', '--exec-properties', '{"custom_config": {}}'
+                '--outputs', '{}', '--exec-properties', '{"custom_config": '
+                '{"ai_platform_training_args": {"project": "12345"}}}'
             ],
         }, body['trainingInput'])
     self.assertStartsWith(body['jobId'], 'tfx_')
@@ -104,9 +107,11 @@ def testStartAIPTrainingWithUserContainer(self, mock_discovery):
 
     class_path = 'foo.bar.class'
 
+    self._exec_properties['custom_config'][
+        'ai_platform_training_job_id'] = self._job_id
     runner.start_aip_training(self._inputs, self._outputs,
                               self._exec_properties, class_path,
-                              self._training_inputs)
+                              self._training_inputs, self._job_id)
 
     mock_create.assert_called_with(
         body=mock.ANY, parent='projects/{}'.format(self._project_id))
@@ -119,10 +124,14 @@ def testStartAIPTrainingWithUserContainer(self, mock_discovery):
             },
             'args': [
                 '--executor_class_path', class_path, '--inputs', '{}',
-                '--outputs', '{}', '--exec-properties', '{"custom_config": {}}'
+                '--outputs', '{}', '--exec-properties', '{"custom_config": '
+                '{"ai_platform_training_args": '
+                '{"masterConfig": {"imageUri": "my-custom-image"}, '
+                '"project": "12345"}, '
+                '"ai_platform_training_job_id": "my_jobid"}}'
             ],
         }, body['trainingInput'])
-    self.assertStartsWith(body['jobId'], 'tfx_')
+    self.assertEqual(body['jobId'], 'my_jobid')
     mock_get.execute.assert_called_with()
 
   @mock.patch(
diff --git a/tfx/extensions/google_cloud_ai_platform/trainer/executor_test.py b/tfx/extensions/google_cloud_ai_platform/trainer/executor_test.py
--- a/tfx/extensions/google_cloud_ai_platform/trainer/executor_test.py
+++ b/tfx/extensions/google_cloud_ai_platform/trainer/executor_test.py
@@ -30,6 +30,8 @@
 class ExecutorTest(tf.test.TestCase):
 
   def setUp(self):
+    super(ExecutorTest, self).setUp()
+
     self._output_data_dir = os.path.join(
         os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', self.get_temp_dir()),
         self._testMethodName)
@@ -45,6 +47,9 @@ def setUp(self):
             },
         },
     }
+    self._executor_class_path = '%s.%s' % (
+        tfx_trainer_executor.Executor.__module__,
+        tfx_trainer_executor.Executor.__name__)
 
   @mock.patch(
       'tfx.extensions.google_cloud_ai_platform.trainer.executor.runner'
@@ -52,18 +57,28 @@ def setUp(self):
   def testDo(self, mock_runner):
     executor = Executor()
     executor.Do(self._inputs, self._outputs, self._exec_properties)
-    executor_class_path = '%s.%s' % (tfx_trainer_executor.Executor.__module__,
-                                     tfx_trainer_executor.Executor.__name__)
     mock_runner.start_aip_training.assert_called_with(
-        self._inputs,
-        self._outputs,
-        self._exec_properties,
-        executor_class_path,
-        {
+        self._inputs, self._outputs, self._exec_properties,
+        self._executor_class_path, {
             'project': self._project_id,
             'jobDir': self._job_dir,
-        },
-    )
+        }, None)
+
+  @mock.patch(
+      'tfx.extensions.google_cloud_ai_platform.trainer.executor.runner'
+  )
+  def testDoWithJobIdOverride(self, mock_runner):
+    executor = Executor()
+    job_id = 'overridden_job_id'
+    self._exec_properties['custom_config'][
+        'ai_platform_training_job_id'] = job_id
+    executor.Do(self._inputs, self._outputs, self._exec_properties)
+    mock_runner.start_aip_training.assert_called_with(
+        self._inputs, self._outputs, self._exec_properties,
+        self._executor_class_path, {
+            'project': self._project_id,
+            'jobDir': self._job_dir,
+        }, job_id)
 
 
 if __name__ == '__main__':