[AIRFLOW-1273] Add Google Cloud ML version and model operators

airflow/contrib/hooks/__init__.py

@@ -40,6 +40,7 @@
     'qubole_hook': ['QuboleHook'],
     'gcs_hook': ['GoogleCloudStorageHook'],
     'datastore_hook': ['DatastoreHook'],
+    'gcp_cloudml_hook': ['CloudMLHook'],
     'gcp_dataproc_hook': ['DataProcHook'],
     'gcp_dataflow_hook': ['DataFlowHook'],
     'spark_submit_operator': ['SparkSubmitOperator'],

airflow/contrib/hooks/gcp_cloudml_hook.py

@@ -0,0 +1,167 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+import random
+import time
+from airflow import settings
+from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook
+from apiclient.discovery import build
+from apiclient import errors
+from oauth2client.client import GoogleCredentials
+
+logging.getLogger('GoogleCloudML').setLevel(settings.LOGGING_LEVEL)
+
+
+def _poll_with_exponential_delay(request, max_n, is_done_func, is_error_func):
+
+    for i in range(0, max_n):
+        try:
+            response = request.execute()
+            if is_error_func(response):
+                raise ValueError('The response contained an error: {}'.format(response))
+            elif is_done_func(response):
+                logging.info('Operation is done: {}'.format(response))
+                return response
+            else:
+                time.sleep((2**i) + (random.randint(0, 1000) / 1000))
+        except errors.HttpError as e:
+            if e.resp.status != 429:
+                logging.info('Something went wrong. Not retrying: {}'.format(e))
+                raise e
+            else:
+                time.sleep((2**i) + (random.randint(0, 1000) / 1000))
+
+
+class CloudMLHook(GoogleCloudBaseHook):
+
+    def __init__(self, gcp_conn_id='google_cloud_default', delegate_to=None):
+        super(CloudMLHook, self).__init__(gcp_conn_id, delegate_to)
+        self._cloudml = self.get_conn()
+
+    def get_conn(self):
+        """
+        Returns a Google CloudML service object.
+        """
+        credentials = GoogleCredentials.get_application_default()
+        return build('ml', 'v1', credentials=credentials)
+
+    def create_version(self, project_name, model_name, version_spec):
+        """
+        Creates the Version on Cloud ML.
+
+        Returns the operation if the version was created successfully and raises
+        an error otherwise.
+        """
+        parent_name = 'projects/{}/models/{}'.format(project_name, model_name)
+        create_request = self._cloudml.projects().models().versions().create(
+            parent=parent_name, body=version_spec)
+        response = create_request.execute()
+        get_request = self._cloudml.projects().operations().get(
+            name=response['name'])
+
+        return _poll_with_exponential_delay(
+            request=get_request,
+            max_n=9,
+            is_done_func=lambda resp: resp.get('done', False),
+            is_error_func=lambda resp: resp.get('error', None) is not None)
+
+    def set_default_version(self, project_name, model_name, version_name):
+        """
+        Sets a version to be the default. Blocks until finished.
+        """
+        full_version_name = 'projects/{}/models/{}/versions/{}'.format(
+            project_name, model_name, version_name)
+        request = self._cloudml.projects().models().versions().setDefault(
+            name=full_version_name, body={})
+
+        try:
+            response = request.execute()
+            logging.info('Successfully set version: {} to default'.format(response))
+            return response
+        except errors.HttpError as e:
+            logging.error('Something went wrong: {}'.format(e))
+            raise e
+
+    def list_versions(self, project_name, model_name):
+        """
+        Lists all available versions of a model. Blocks until finished.
+        """
+        result = []
+        full_parent_name = 'projects/{}/models/{}'.format(
+            project_name, model_name)
+        request = self._cloudml.projects().models().versions().list(
+            parent=full_parent_name, pageSize=100)
+
+        response = request.execute()
+        next_page_token = response.get('nextPageToken', None)
+        result.extend(response.get('versions', []))
+        while next_page_token is not None:
+            next_request = self._cloudml.projects().models().versions().list(
+                parent=full_parent_name,
+                pageToken=next_page_token,
+                pageSize=100)
+            response = next_request.execute()
+            next_page_token = response.get('nextPageToken', None)
+            result.extend(response.get('versions', []))
+            time.sleep(5)
+        return result
+
+    def delete_version(self, project_name, model_name, version_name):
+        """
+        Deletes the given version of a model. Blocks until finished.
+        """
+        full_name = 'projects/{}/models/{}/versions/{}'.format(
+            project_name, model_name, version_name)
+        delete_request = self._cloudml.projects().models().versions().delete(
+            name=full_name)
+        response = delete_request.execute()
+        get_request = self._cloudml.projects().operations().get(
+            name=response['name'])
+
+        return _poll_with_exponential_delay(
+            request=get_request,
+            max_n=9,
+            is_done_func=lambda resp: resp.get('done', False),
+            is_error_func=lambda resp: resp.get('error', None) is not None)
+
+    def create_model(self, project_name, model):
+        """
+        Create a Model. Blocks until finished.
+        """
+        assert model['name'] is not None and model['name'] is not ''
+        project = 'projects/{}'.format(project_name)
+
+        request = self._cloudml.projects().models().create(
+            parent=project, body=model)
+        return request.execute()
+
+    def get_model(self, project_name, model_name):
+        """
+        Gets a Model. Blocks until finished.
+        """
+        assert model_name is not None and model_name is not ''
+        full_model_name = 'projects/{}/models/{}'.format(
+            project_name, model_name)
+        request = self._cloudml.projects().models().get(name=full_model_name)
+        try:
+            return request.execute()
+        except errors.HttpError as e:
+            if e.resp.status == 404:
+                logging.error('Model was not found: {}'.format(e))
+                return None
+            raise e

airflow/contrib/operators/cloudml_operator.py

@@ -0,0 +1,178 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the 'License'); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+from airflow import settings
+from airflow.contrib.hooks.gcp_cloudml_hook import CloudMLHook
+from airflow.operators import BaseOperator
+from airflow.utils.decorators import apply_defaults
+
+logging.getLogger('GoogleCloudML').setLevel(settings.LOGGING_LEVEL)
+
+
+class CloudMLVersionOperator(BaseOperator):
+    """
+    Operator for managing a Google Cloud ML version.
+
+    :param model_name: The name of the Google Cloud ML model that the version
+        belongs to.
+    :type model_name: string
+
+    :param project_name: The Google Cloud project name to which CloudML
+        model belongs.
+    :type project_name: string
+
+    :param version: A dictionary containing the information about the version.
+        If the `operation` is `create`, `version` should contain all the
+        information about this version such as name, and deploymentUrl.
+        If the `operation` is `get` or `delete`, the `version` parameter
+        should contain the `name` of the version.
+        If it is None, the only `operation` possible would be `list`.
+    :type version: dict
+
+    :param gcp_conn_id: The connection ID to use when fetching connection info.
+    :type gcp_conn_id: string
+
+    :param operation: The operation to perform. Available operations are:
+        'create': Creates a new version in the model specified by `model_name`,
+            in which case the `version` parameter should contain all the
+            information to create that version
+            (e.g. `name`, `deploymentUrl`).
+        'get': Gets full information of a particular version in the model
+            specified by `model_name`.
+            The name of the version should be specified in the `version`
+            parameter.
+
+        'list': Lists all available versions of the model specified
+            by `model_name`.
+
+        'delete': Deletes the version specified in `version` parameter from the
+            model specified by `model_name`).
+            The name of the version should be specified in the `version`
+            parameter.
+     :type operation: string
+
+    :param delegate_to: The account to impersonate, if any.
+        For this to work, the service account making the request must have
+        domain-wide delegation enabled.
+    :type delegate_to: string
+    """
+
+
+    template_fields = [
+        '_model_name',
+        '_version',
+    ]
+
+    @apply_defaults
+    def __init__(self,
+                 model_name,
+                 project_name,
+                 version=None,
+                 gcp_conn_id='google_cloud_default',
+                 operation='create',
+                 delegate_to=None,
+                 *args,
+                 **kwargs):
+
+        super(CloudMLVersionOperator, self).__init__(*args, **kwargs)
+        self._model_name = model_name
+        self._version = version
+        self._gcp_conn_id = gcp_conn_id
+        self._delegate_to = delegate_to
+        self._project_name = project_name
+        self._operation = operation
+
+    def execute(self, context):
+        hook = CloudMLHook(
+            gcp_conn_id=self._gcp_conn_id, delegate_to=self._delegate_to)
+
+        if self._operation == 'create':
+            assert self._version is not None
+            return hook.create_version(self._project_name, self._model_name,
+                                       self._version)
+        elif self._operation == 'set_default':
+            return hook.set_default_version(
+                self._project_name, self._model_name,
+                self._version['name'])
+        elif self._operation == 'list':
+            return hook.list_versions(self._project_name, self._model_name)
+        elif self._operation == 'delete':
+            return hook.delete_version(self._project_name, self._model_name,
+                                       self._version['name'])
+        else:
+            raise ValueError('Unknown operation: {}'.format(self._operation))
+
+
+class CloudMLModelOperator(BaseOperator):
+    """
+    Operator for managing a Google Cloud ML model.
+
+    :param model: A dictionary containing the information about the model.
+        If the `operation` is `create`, then the `model` parameter should
+        contain all the information about this model such as `name`.
+
+        If the `operation` is `get`, the `model` parameter
+        should contain the `name` of the model.
+    :type model: dict
+
+    :param project_name: The Google Cloud project name to which CloudML
+        model belongs.
+    :type project_name: string
+
+    :param gcp_conn_id: The connection ID to use when fetching connection info.
+    :type gcp_conn_id: string
+
+    :param operation: The operation to perform. Available operations are:
+        'create': Creates a new model as provided by the `model` parameter.
+        'get': Gets a particular model where the name is specified in `model`.
+
+    :param delegate_to: The account to impersonate, if any.
+        For this to work, the service account making the request must have
+        domain-wide delegation enabled.
+    :type delegate_to: string
+    """
+
+    template_fields = [
+        '_model',
+    ]
+
+    @apply_defaults
+    def __init__(self,
+                 model,
+                 project_name,
+                 gcp_conn_id='google_cloud_default',
+                 operation='create',
+                 delegate_to=None,
+                 *args,
+                 **kwargs):
+        super(CloudMLModelOperator, self).__init__(*args, **kwargs)
+        self._model = model
+        self._operation = operation
+        self._gcp_conn_id = gcp_conn_id
+        self._delegate_to = delegate_to
+        self._project_name = project_name
+
+    def execute(self, context):
+        hook = CloudMLHook(
+            gcp_conn_id=self._gcp_conn_id, delegate_to=self._delegate_to)
+        if self._operation == 'create':
+            hook.create_model(self._project_name, self._model)
+        elif self._operation == 'get':
+            hook.get_model(self._project_name, self._model['name'])
+        else:
+            raise ValueError('Unknown operation: {}'.format(self._operation))

airflow/utils/db.py

@@ -128,6 +128,10 @@ def initdb():
             conn_id='presto_default', conn_type='presto',
             host='localhost',
             schema='hive', port=3400))
+    merge_conn(
+        models.Connection(
+            conn_id='google_cloud_default', conn_type='google_cloud_platform',
+            schema='default',))
     merge_conn(
         models.Connection(
             conn_id='hive_cli_default', conn_type='hive_cli',