From c6050c7b7660c129978e140f9cb7a90c6f737db7 Mon Sep 17 00:00:00 2001 From: Sam Shi Date: Mon, 14 Jan 2019 17:48:43 -0800 Subject: [PATCH] =?UTF-8?q?Support=20datalab=20connect()=20to=20connect=20?= =?UTF-8?q?a=20GCE=20instance=20created=20from=20Deep=E2=80=A6=20(#2112)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * support datalab connect() to connect a GCE instance created from Deeplearning images * adding soupsieve licence --- containers/base/third_party_licenses.csv | 1 + tools/cli/commands/utils.py | 36 ++++++++++++++++-------- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/containers/base/third_party_licenses.csv b/containers/base/third_party_licenses.csv index b9690c076..ee665c2bf 100644 --- a/containers/base/third_party_licenses.csv +++ b/containers/base/third_party_licenses.csv @@ -251,3 +251,4 @@ xgboost,https://raw.githubusercontent.com/dmlc/xgboost/master/LICENSE,Apache 2.0 zict,https://raw.githubusercontent.com/dask/zict/master/LICENSE.txt,BSD zope.deprecation,https://raw.githubusercontent.com/zopefoundation/zope.deprecation/master/LICENSE.txt,ZPL google-cloud-dataflow,https://raw.githubusercontent.com/apache/beam/master/LICENSE,Apache 2.0 +soupsieve,https://raw.githubusercontent.com/facelessuser/soupsieve/master/LICENSE.md,MIT diff --git a/tools/cli/commands/utils.py b/tools/cli/commands/utils.py index f3a047889..d76937fa6 100644 --- a/tools/cli/commands/utils.py +++ b/tools/cli/commands/utils.py @@ -68,8 +68,9 @@ class InvalidInstanceException(Exception): _MESSAGE = ( 'The specified instance, {}, does not appear ' - 'to have been created by the `datalab` tool, and ' - 'so cannot be managed by it.') + 'to have been created by the `datalab` tool, or ' + 'from any GCE Deeplearning images. Therefore it ' + 'cannot be managed by `datalab` tool.') def __init__(self, instance_name): super(InvalidInstanceException, self).__init__( @@ -240,22 +241,35 @@ def flatten_metadata(metadata): return result -def _check_datalab_tag(instance, tags): +def _check_instance_allowed(instance, status_tags_and_metadata): """Check that the given "tags" object contains `datalab`. This is used to verify that a VM was created by the `datalab create` - command. + command or was from GCE Deeplearning images, by checking if the VM + description contains a tag of 'datalab' or includes a c2d-tensorflow + licence string. Args: instance: The name of the instance to check - tags: An object with an 'items' field that is a list of tags. + status_tags_and_metadata: An object containing the result of GCE + VM instance description. Raises: InvalidInstanceException: If the check fails. """ + tags = status_tags_and_metadata.get('tags', {}) items = tags.get('items', []) - if 'datalab' not in items: - raise InvalidInstanceException(instance) - return + + if 'datalab' in items: + return + else: + _license = ('https://www.googleapis.com/compute/v1/projects/' + 'click-to-deploy-images/global/licenses/c2d-tensorflow') + disks = status_tags_and_metadata.get('disks', []) + for disk in disks: + if _license in disk.get('licenses', []): + return + + raise InvalidInstanceException(instance) def describe_instance(args, gcloud_compute, instance): @@ -282,7 +296,8 @@ def describe_instance(args, gcloud_compute, instance): if args.zone: get_cmd.extend(['--zone', args.zone]) get_cmd.extend( - ['--format', 'json(status,tags.items,metadata.items)', instance]) + ['--format', 'json(status,tags.items,metadata.items,disks[].licenses)', + instance]) with tempfile.TemporaryFile() as stdout, \ tempfile.TemporaryFile() as stderr: try: @@ -290,8 +305,7 @@ def describe_instance(args, gcloud_compute, instance): stdout.seek(0) json_result = stdout.read().decode('utf-8').strip() status_tags_and_metadata = json.loads(json_result) - tags = status_tags_and_metadata.get('tags', {}) - _check_datalab_tag(instance, tags) + _check_instance_allowed(instance, status_tags_and_metadata) status = status_tags_and_metadata.get('status', 'UNKNOWN') metadata = status_tags_and_metadata.get('metadata', {})