@@ -68,8 +68,9 @@ class InvalidInstanceException(Exception):
6868
6969 _MESSAGE = (
7070 'The specified instance, {}, does not appear '
71- 'to have been created by the `datalab` tool, and '
72- 'so cannot be managed by it.' )
71+ 'to have been created by the `datalab` tool, or '
72+ 'from any GCE Deeplearning images. Therefore it '
73+ 'cannot be managed by `datalab` tool.' )
7374
7475 def __init__ (self , instance_name ):
7576 super (InvalidInstanceException , self ).__init__ (
@@ -240,22 +241,35 @@ def flatten_metadata(metadata):
240241 return result
241242
242243
243- def _check_datalab_tag (instance , tags ):
244+ def _check_instance_allowed (instance , status_tags_and_metadata ):
244245 """Check that the given "tags" object contains `datalab`.
245246
246247 This is used to verify that a VM was created by the `datalab create`
247- command.
248+ command or was from GCE Deeplearning images, by checking if the VM
249+ description contains a tag of 'datalab' or includes a c2d-tensorflow
250+ licence string.
248251
249252 Args:
250253 instance: The name of the instance to check
251- tags: An object with an 'items' field that is a list of tags.
254+ status_tags_and_metadata: An object containing the result of GCE
255+ VM instance description.
252256 Raises:
253257 InvalidInstanceException: If the check fails.
254258 """
259+ tags = status_tags_and_metadata .get ('tags' , {})
255260 items = tags .get ('items' , [])
256- if 'datalab' not in items :
257- raise InvalidInstanceException (instance )
258- return
261+
262+ if 'datalab' in items :
263+ return
264+ else :
265+ _license = ('https://www.googleapis.com/compute/v1/projects/'
266+ 'click-to-deploy-images/global/licenses/c2d-tensorflow' )
267+ disks = status_tags_and_metadata .get ('disks' , [])
268+ for disk in disks :
269+ if _license in disk .get ('licenses' , []):
270+ return
271+
272+ raise InvalidInstanceException (instance )
259273
260274
261275def describe_instance (args , gcloud_compute , instance ):
@@ -282,16 +296,16 @@ def describe_instance(args, gcloud_compute, instance):
282296 if args .zone :
283297 get_cmd .extend (['--zone' , args .zone ])
284298 get_cmd .extend (
285- ['--format' , 'json(status,tags.items,metadata.items)' , instance ])
299+ ['--format' , 'json(status,tags.items,metadata.items,disks[].licenses)' ,
300+ instance ])
286301 with tempfile .TemporaryFile () as stdout , \
287302 tempfile .TemporaryFile () as stderr :
288303 try :
289304 gcloud_compute (args , get_cmd , stdout = stdout , stderr = stderr )
290305 stdout .seek (0 )
291306 json_result = stdout .read ().decode ('utf-8' ).strip ()
292307 status_tags_and_metadata = json .loads (json_result )
293- tags = status_tags_and_metadata .get ('tags' , {})
294- _check_datalab_tag (instance , tags )
308+ _check_instance_allowed (instance , status_tags_and_metadata )
295309
296310 status = status_tags_and_metadata .get ('status' , 'UNKNOWN' )
297311 metadata = status_tags_and_metadata .get ('metadata' , {})
0 commit comments