feat: Add new gcloud commands, API clients, and third-party libraries across various services.

gcloud auth application-default login/google-cloud-sdk/lib/googlecloudsdk/api_lib/edgeml/edgeml.py

@@ -0,0 +1,160 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2019 Google LLC. All Rights Reserved.
+#
+# Licensed 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.
+"""Utilities for Edge ML API."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+import re
+
+from apitools.base.py import encoding_helper
+from googlecloudsdk.api_lib.edgeml import operations
+from googlecloudsdk.api_lib.edgeml import util
+from googlecloudsdk.core import properties
+
+
+def _CompileDestination(source):
+  """Converts model[.tflite] style filename to model_edgetpu.tflite."""
+  return re.sub(r'(\.tflite)?$', '_edgetpu.tflite', source, count=1)
+
+
+def _ConvertDestination(source):
+  """Converts model[/saved_model.(pb|pbtxt)] style filename to model.tflite."""
+  return re.sub(r'(/saved_model\.(pb|pbtxt))?$', '.tflite', source, count=1)
+
+
+class EdgeMlClient(object):
+  """Client for Edge ML service.
+
+  Attributes:
+    client: Generated Edge ML API client.
+    messages: Generated Edge ML API messages.
+  """
+
+  def __init__(self, client=None, messages=None):
+    self.client = client or util.GetClientInstance()
+    self.messages = messages or util.GetMessagesModule(client)
+    self._service = self.client.projects_models
+    self._operations_client = operations.OperationsClient(client, messages)
+
+  def Analyze(self, source):
+    """Analyzes a machine learning model.
+
+    The AnalyzeResponse will contain model's framework, and for TF models
+    it will also contain Edge TPU compiliability and input/output tensor
+    information.
+
+    Args:
+      source: str, GCS object URI to the model file or directory to analyze.
+
+    Returns:
+      AnalyzedResponse on the finish of analyze operation.
+
+    Raises:
+      LongrunningError: when long running operation fails.
+    """
+    project = 'projects/' + properties.VALUES.core.project.GetOrFail()
+    analyze_req = self.messages.EdgemlProjectsModelsAnalyzeRequest(
+        analyzeModelRequest=self.messages.AnalyzeModelRequest(
+            gcsSource=self.messages.GcsSource(inputUris=[source])),
+        project=project)
+    operation = self._service.Analyze(analyze_req)
+    result = self._operations_client.WaitForOperation(operation)
+    response = encoding_helper.JsonToMessage(
+        self.messages.AnalyzeModelResponse,
+        encoding_helper.MessageToJson(result))
+
+    return response
+
+  def Compile(self, source, destination=None):
+    """Optimizes a TFLite model for EdgeTPU.
+
+    Args:
+      source: str, GCS object URI to the model file to compile. Must be a
+        .tflite file.
+      destination: str, GCS URI to an output tflite object. If not provided,
+        for source filename `model[.tflite]` this will be set to
+        `model_edgetpu.tflite`.
+
+    Returns:
+      (CompileModelResponse, output object URI) on the finish of compilation.
+
+    Raises:
+      LongrunningError: when long running operation fails.
+    """
+    project = 'projects/' + properties.VALUES.core.project.GetOrFail()
+    if not destination:
+      destination = _CompileDestination(source)
+    compile_req_type = self.messages.CompileModelRequest
+    chip_type = compile_req_type.ChipTypeValueValuesEnum.EDGE_TPU_V1
+
+    input_config = self.messages.InputConfig(
+        gcsSource=self.messages.GcsSource(inputUris=[source]))
+    output_config = self.messages.OutputConfig(
+        gcsDestination=self.messages.GcsDestination(outputUri=destination))
+
+    compile_req = self.messages.EdgemlProjectsModelsCompileRequest(
+        compileModelRequest=compile_req_type(
+            chipType=chip_type,
+            inputConfig=input_config,
+            outputConfig=output_config),
+        project=project)
+
+    operation = self._service.Compile(compile_req)
+    result = self._operations_client.WaitForOperation(operation)
+
+    response = encoding_helper.JsonToMessage(
+        self.messages.CompileModelResponse,
+        encoding_helper.MessageToJson(result))
+    return response, destination
+
+  def Convert(self, source, destination=None):
+    """Converts Tensorflow SavedModel to TFLite model.
+
+    Args:
+      source: str, GCS URI to an input SavedModel archive
+      destination: str, GCS URI to an output tflite object. If not provided,
+        for source filename `model[/saved_model.(pb|pbtxt)]` this will be
+        set to `model.tflite`.
+
+    Returns:
+      (ConvertModelResponse, output object URI) on the finish of
+      convert operation.
+
+    Raises:
+      LongrunningError: when long running operation fails.
+    """
+    project = 'projects/' + properties.VALUES.core.project.GetOrFail()
+    if not destination:
+      destination = _ConvertDestination(source)
+
+    input_config = self.messages.InputConfig(
+        gcsSource=self.messages.GcsSource(inputUris=[source]))
+    output_config = self.messages.OutputConfig(
+        gcsDestination=self.messages.GcsDestination(outputUri=destination))
+
+    convert_req = self.messages.EdgemlProjectsModelsConvertRequest(
+        convertModelRequest=self.messages.ConvertModelRequest(
+            inputConfig=input_config, outputConfig=output_config),
+        project=project)
+
+    operation = self._service.Convert(convert_req)
+    result = self._operations_client.WaitForOperation(operation)
+
+    response = encoding_helper.JsonToMessage(
+        self.messages.ConvertModelResponse,
+        encoding_helper.MessageToJson(result))
+    return response, destination

gcloud auth application-default login/google-cloud-sdk/lib/googlecloudsdk/api_lib/edgeml/operations.py

@@ -0,0 +1,88 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2019 Google LLC. All Rights Reserved.
+#
+# Licensed 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.
+"""Utilities for Edge ML API long-running operations."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+from googlecloudsdk.api_lib.edgeml import util
+from googlecloudsdk.api_lib.util import waiter
+from googlecloudsdk.core import resources
+
+
+_MAX_WAIT_TIME_MS = 10 * 60 * 1000  # 10 minutes.
+
+_WAITING_MESSAGE = {
+    'analyze': 'Analyzing Edge machine learning model. [{}]',
+    'compile': 'Optimizing Edge machine learning model for Edge TPU. [{}]',
+    'convert': 'Converting SavedModel to TF Lite model. [{}]',
+}
+
+
+class OperationsClient(object):
+  """Client for the Edge ML operations API.
+
+  Attributes:
+    client: Generated Edge ML API client.
+    messages: Generated Edge ML API messages.
+  """
+
+  def __init__(self, client=None, messages=None):
+    self.client = client or util.GetClientInstance()
+    self.messages = messages or util.GetMessagesModule(client)
+    self._service = self.client.operations
+
+  def Get(self, operation_ref):
+    """Calls get method for long-running operations.
+
+    Args:
+      operation_ref: Ref to long-running operation.
+
+    Returns:
+      edgeml.Operation message.
+    """
+    return self.client.projects_operations.Get(
+        self.messages.EdgemlOperationsGetRequest(
+            name=operation_ref.RelativeName()))
+
+  def WaitForOperation(self, operation):
+    """Wait until the operation is complete or times out.
+
+    Args:
+      operation: The operation resource to wait on.
+
+    Returns:
+      The operation resource when it has completed.
+
+    Raises:
+      OperationTimeoutError: When the operation polling times out.
+      OperationError: When the operation completed with an error.
+    """
+    poller = waiter.CloudOperationPollerNoResources(self._service)
+    operation_ref = resources.REGISTRY.Parse(
+        operation.name, collection='edgeml.operations')
+    operation_type = operation_ref.Name().split('/')[0]
+    message = _WAITING_MESSAGE.get(operation_type, 'Waiting for operation [{}]')
+    return waiter.WaitFor(
+        poller,
+        operation_ref,
+        message,
+        pre_start_sleep_ms=1000,
+        max_wait_ms=_MAX_WAIT_TIME_MS,
+        exponential_sleep_multiplier=None,
+        jitter_ms=None,
+        wait_ceiling_ms=None,
+        sleep_ms=5000)

gcloud auth application-default login/google-cloud-sdk/lib/googlecloudsdk/api_lib/edgeml/util.py

@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2019 Google LLC. All Rights Reserved.
+#
+# Licensed 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.
+"""Utilities for Edge ML API."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+from googlecloudsdk.api_lib.util import apis
+
+_API_NAME = 'edgeml'
+_API_VERSION = 'v1beta1'
+
+
+def GetClientInstance(no_http=False):
+  return apis.GetClientInstance(_API_NAME, _API_VERSION, no_http=no_http)
+
+
+def GetMessagesModule(client=None):
+  client = client or GetClientInstance()
+  return client.MESSAGES_MODULE