[AIRFLOW-6915] Add AI Platform Console Link for MLEngineStartTrainingJobOperator (#7535)

Author: turbaszek

airflow/providers/google/cloud/operators/mlengine.py

@@ -24,7 +24,8 @@
 from typing import List, Optional
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
+from airflow.models import BaseOperator, BaseOperatorLink
+from airflow.models.taskinstance import TaskInstance
 from airflow.providers.google.cloud.hooks.mlengine import MLEngineHook
 from airflow.utils.decorators import apply_defaults
 
@@ -852,6 +853,23 @@ def execute(self, context):
         )
 
 
+class AIPlatformConsoleLink(BaseOperatorLink):
+    """
+    Helper class for constructing AI Platform Console link.
+    """
+    name = "AI Platform Console"
+
+    def get_link(self, operator, dttm):
+        task_instance = TaskInstance(task=operator, execution_date=dttm)
+        gcp_metadata_dict = task_instance.xcom_pull(task_ids=operator.task_id, key="gcp_metadata")
+        if not gcp_metadata_dict:
+            return ''
+        job_id = gcp_metadata_dict['job_id']
+        project_id = gcp_metadata_dict['project_id']
+        console_link = f"https://console.cloud.google.com/ai-platform/jobs/{job_id}?project={project_id}"
+        return console_link
+
+
 class MLEngineStartTrainingJobOperator(BaseOperator):
     """
     Operator for launching a MLEngine training job.
@@ -915,6 +933,10 @@ class MLEngineStartTrainingJobOperator(BaseOperator):
         '_job_dir'
     ]
 
+    operator_extra_links = (
+        AIPlatformConsoleLink(),
+    )
+
     @apply_defaults
     def __init__(self,  # pylint: disable=too-many-arguments
                  job_id: str,
@@ -1016,6 +1038,12 @@ def check_existing_job(existing_job):
             self.log.error('MLEngine training job failed: %s', str(finished_training_job))
             raise RuntimeError(finished_training_job['errorMessage'])
 
+        gcp_metadata = {
+            "job_id": job_id,
+            "project_id": self._project_id,
+        }
+        context['task_instance'].xcom_push("gcp_metadata", gcp_metadata)
+
 
 class MLEngineTrainingJobFailureOperator(BaseOperator):
 

airflow/serialization/serialized_objects.py

@@ -42,6 +42,7 @@
 BUILTIN_OPERATOR_EXTRA_LINKS: List[str] = [
     "airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleLink",
     "airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink",
+    "airflow.providers.google.cloud.operators.mlengine.AIPlatformConsoleLink",
     "airflow.providers.qubole.operators.qubole.QDSLink"
 ]
 

tests/providers/google/cloud/operators/test_mlengine.py

@@ -18,23 +18,27 @@
 import copy
 import datetime
 import unittest
-from unittest.mock import ANY, patch
+from unittest.mock import ANY, MagicMock, patch
 
 import httplib2
 from googleapiclient.errors import HttpError
 
 from airflow.exceptions import AirflowException
+from airflow.models import TaskInstance
 from airflow.models.dag import DAG
 from airflow.providers.google.cloud.operators.mlengine import (
-    MLEngineCreateModelOperator, MLEngineCreateVersionOperator, MLEngineDeleteModelOperator,
-    MLEngineDeleteVersionOperator, MLEngineGetModelOperator, MLEngineListVersionsOperator,
-    MLEngineManageModelOperator, MLEngineManageVersionOperator, MLEngineSetDefaultVersionOperator,
-    MLEngineStartBatchPredictionJobOperator, MLEngineStartTrainingJobOperator,
-    MLEngineTrainingJobFailureOperator,
+    AIPlatformConsoleLink, MLEngineCreateModelOperator, MLEngineCreateVersionOperator,
+    MLEngineDeleteModelOperator, MLEngineDeleteVersionOperator, MLEngineGetModelOperator,
+    MLEngineListVersionsOperator, MLEngineManageModelOperator, MLEngineManageVersionOperator,
+    MLEngineSetDefaultVersionOperator, MLEngineStartBatchPredictionJobOperator,
+    MLEngineStartTrainingJobOperator, MLEngineTrainingJobFailureOperator,
 )
+from airflow.serialization.serialized_objects import SerializedDAG
+from airflow.utils.dates import days_ago
 
 DEFAULT_DATE = datetime.datetime(2017, 6, 6)
 
+TEST_DAG_ID = "test-mlengine-operators"
 TEST_PROJECT_ID = "test-project-id"
 TEST_MODEL_NAME = "test-model-name"
 TEST_VERSION_NAME = "test-version"
@@ -304,7 +308,8 @@ class TestMLEngineTrainingOperator(unittest.TestCase):
         'training_args': '--some_arg=\'aaa\'',
         'region': 'us-east1',
         'scale_tier': 'STANDARD_1',
-        'task_id': 'test-training'
+        'task_id': 'test-training',
+        'start_date': days_ago(1)
     }
     TRAINING_INPUT = {
         'jobId': 'test_training',
@@ -317,6 +322,9 @@ class TestMLEngineTrainingOperator(unittest.TestCase):
         }
     }
 
+    def setUp(self):
+        self.dag = DAG(TEST_DAG_ID, default_args=self.TRAINING_DEFAULT_ARGS)
+
     @patch('airflow.providers.google.cloud.operators.mlengine.MLEngineHook')
     def test_success_create_training_job(self, mock_hook):
         success_response = self.TRAINING_INPUT.copy()
@@ -326,7 +334,7 @@ def test_success_create_training_job(self, mock_hook):
 
         training_op = MLEngineStartTrainingJobOperator(
             **self.TRAINING_DEFAULT_ARGS)
-        training_op.execute(None)
+        training_op.execute(MagicMock())
 
         mock_hook.assert_called_once_with(
             gcp_conn_id='google_cloud_default', delegate_to=None)
@@ -352,7 +360,7 @@ def test_success_create_training_job_with_optional_args(self, mock_hook):
             python_version='3.5',
             job_dir='gs://some-bucket/jobs/test_training',
             **self.TRAINING_DEFAULT_ARGS)
-        training_op.execute(None)
+        training_op.execute(MagicMock())
 
         mock_hook.assert_called_once_with(gcp_conn_id='google_cloud_default', delegate_to=None)
         # Make sure only 'create_job' is invoked on hook instance
@@ -404,6 +412,73 @@ def test_failed_job_error(self, mock_hook):
             project_id='test-project', job=self.TRAINING_INPUT, use_existing_job_fn=ANY)
         self.assertEqual('A failure message', str(context.exception))
 
+    @patch('airflow.providers.google.cloud.operators.mlengine.MLEngineHook')
+    def test_console_extra_link(self, mock_hook):
+        training_op = MLEngineStartTrainingJobOperator(
+            **self.TRAINING_DEFAULT_ARGS)
+
+        ti = TaskInstance(
+            task=training_op,
+            execution_date=DEFAULT_DATE,
+        )
+
+        job_id = self.TRAINING_DEFAULT_ARGS['job_id']
+        project_id = self.TRAINING_DEFAULT_ARGS['project_id']
+        gcp_metadata = {
+            "job_id": job_id,
+            "project_id": project_id,
+        }
+        ti.xcom_push(key='gcp_metadata', value=gcp_metadata)
+
+        self.assertEqual(
+            f"https://console.cloud.google.com/ai-platform/jobs/{job_id}?project={project_id}",
+            training_op.get_extra_links(DEFAULT_DATE, AIPlatformConsoleLink.name),
+        )
+
+        self.assertEqual(
+            '',
+            training_op.get_extra_links(datetime.datetime(2019, 1, 1), AIPlatformConsoleLink.name),
+        )
+
+    def test_console_extra_link_serialized_field(self):
+        with self.dag:
+            training_op = MLEngineStartTrainingJobOperator(**self.TRAINING_DEFAULT_ARGS)
+        serialized_dag = SerializedDAG.to_dict(self.dag)
+        dag = SerializedDAG.from_dict(serialized_dag)
+        simple_task = dag.task_dict[self.TRAINING_DEFAULT_ARGS['task_id']]
+
+        # Check Serialized version of operator link
+        self.assertEqual(
+            serialized_dag["dag"]["tasks"][0]["_operator_extra_links"],
+            [{"airflow.providers.google.cloud.operators.mlengine.AIPlatformConsoleLink": {}}]
+        )
+
+        # Check DeSerialized version of operator link
+        self.assertIsInstance(list(simple_task.operator_extra_links)[0], AIPlatformConsoleLink)
+
+        job_id = self.TRAINING_DEFAULT_ARGS['job_id']
+        project_id = self.TRAINING_DEFAULT_ARGS['project_id']
+        gcp_metadata = {
+            "job_id": job_id,
+            "project_id": project_id,
+        }
+
+        ti = TaskInstance(
+            task=training_op,
+            execution_date=DEFAULT_DATE,
+        )
+        ti.xcom_push(key='gcp_metadata', value=gcp_metadata)
+
+        self.assertEqual(
+            f"https://console.cloud.google.com/ai-platform/jobs/{job_id}?project={project_id}",
+            simple_task.get_extra_links(DEFAULT_DATE, AIPlatformConsoleLink.name),
+        )
+
+        self.assertEqual(
+            '',
+            simple_task.get_extra_links(datetime.datetime(2019, 1, 1), AIPlatformConsoleLink.name),
+        )
+
 
 class TestMLEngineTrainingJobFailureOperator(unittest.TestCase):