gcpdiag.queries.notebooks

37class HealthStateEnum(enum.Enum):
38  """Vertex AI Workbench user-managed notebooks instance health states
39
40  https://cloud.google.com/vertex-ai/docs/workbench/reference/rest/v1/projects.locations.instances/getInstanceHealth#healthstate
41  """
42
43  HEALTH_STATE_UNSPECIFIED = 'HEALTH_STATE_UNSPECIFIED'
44  HEALTHY = 'HEALTHY'
45  UNHEALTHY = 'UNHEALTHY'
46  AGENT_NOT_INSTALLED = 'AGENT_NOT_INSTALLED'
47  AGENT_NOT_RUNNING = 'AGENT_NOT_RUNNING'
48
49  def __str__(self):
50    return str(self.value)

53class StateEnum(enum.Enum):
54  """Vertex AI Workbench instance states
55
56  https://cloud.google.com/vertex-ai/docs/workbench/reference/rest/v2/projects.locations.instances#state
57  """
58
59  STATE_UNSPECIFIED = 'STATE_UNSPECIFIED'
60  STARTING = 'STARTING'
61  PROVISIONING = 'PROVISIONING'
62  ACTIVE = 'ACTIVE'
63  STOPPING = 'STOPPING'
64  STOPPED = 'STOPPED'
65  DELETED = 'DELETED'
66  UPGRADING = 'UPGRADING'
67  INITIALIZING = 'INITIALIZING'
68  SUSPENDING = 'SUSPENDING'
69  SUSPENDED = 'SUSPENDED'
70
71  def __str__(self):
72    return str(self.value)

 75class Instance(models.Resource):
 76  """Represent a Vertex AI Workbench user-managed notebook instance
 77
 78  https://cloud.google.com/vertex-ai/docs/workbench/reference/rest/v1/projects.locations.instances#resource:-instance
 79  """
 80
 81  _resource_data: dict
 82
 83  def __init__(self, project_id, resource_data):
 84    super().__init__(project_id=project_id)
 85    self._resource_data = resource_data
 86
 87  @property
 88  def full_path(self) -> str:
 89    """
 90    The 'name' of the instance is already in the full path form
 91    projects/{project}/locations/{location}/instances/{instance}.
 92    """
 93    return self._resource_data[NAME_KEY]
 94
 95  @property
 96  def short_path(self) -> str:
 97    path = self.full_path
 98    path = re.sub(r'^projects/', '', path)
 99    path = re.sub(r'/locations/', '/', path)
100    path = re.sub(r'/instances/', '/', path)
101    return path
102
103  @property
104  def metadata(self) -> dict:
105    return self._resource_data.get('metadata', {})
106
107  @property
108  def name(self) -> str:
109    logging.info(self._resource_data)
110    return self._resource_data[NAME_KEY]

113class Runtime(models.Resource):
114  """Represent a Vertex AI Workbench runtime for a managed notebook instance
115
116  https://cloud.google.com/vertex-ai/docs/workbench/reference/rest/v1/projects.locations.runtimes#resource:-runtime
117  """
118
119  _resource_data: dict
120
121  def __init__(self, project_id, resource_data):
122    super().__init__(project_id=project_id)
123    self._resource_data = resource_data
124
125  @property
126  def full_path(self) -> str:
127    """
128    The 'name' of the runtime is already in the full path form
129    projects/{project}/locations/{location}/runtimes/{instance}.
130    """
131    return self._resource_data[NAME_KEY]
132
133  @property
134  def short_path(self) -> str:
135    path = self.full_path
136    path = re.sub(r'^projects/', '', path)
137    path = re.sub(r'/locations/', '/', path)
138    path = re.sub(r'/runtimes/', '/', path)
139    return path
140
141  @property
142  def metadata(self) -> dict:
143    return self._resource_data.get('metadata', {})
144
145  @property
146  def name(self) -> str:
147    logging.info(self._resource_data)
148    return self._resource_data[NAME_KEY]
149
150  @property
151  def software_config(self) -> dict:
152    return self._resource_data.get('softwareConfig', {})
153
154  @property
155  def idle_shutdown(self) -> bool:
156    return self.software_config.get('idleShutdown', False)
157
158  @property
159  def is_upgradeable(self) -> bool:
160    return self.software_config.get('upgradeable', False)
161
162  @property
163  def version(self) -> str:
164    return self.software_config.get('version', '')
165
166  @property
167  def health_state(self) -> HealthStateEnum:
168    return self._resource_data.get(HEALTH_STATE_KEY,
169                                   HealthStateEnum.HEALTH_STATE_UNSPECIFIED)

172class WorkbenchInstance(Instance):
173  """Represent a Vertex AI Workbench Instance
174
175  https://cloud.google.com/vertex-ai/docs/workbench/reference/rest/v2/projects.locations.instances#Instance
176  """
177
178  _resource_data: dict
179
180  def __init__(self, project_id, resource_data):
181    super().__init__(project_id=project_id, resource_data=resource_data)
182    self._resource_data = resource_data
183
184  @property
185  def state(self) -> StateEnum:
186    return StateEnum[self._resource_data.get('state', 'STATE_UNSPECIFIED')]
187
188  @property
189  def gce_setup(self) -> dict:
190    return self._resource_data.get('gceSetup', {})
191
192  @property
193  def gce_service_account_email(self) -> str:
194    gce_setup_service_accounts = self.gce_setup.get('serviceAccounts', [])
195    return gce_setup_service_accounts[0].get(
196        'email', '') if len(gce_setup_service_accounts) > 0 else ''
197
198  @property
199  def network(self) -> str:
200    gce_setup_network_interfaces = self.gce_setup.get('networkInterfaces', [])
201    return gce_setup_network_interfaces[0].get(
202        'network', '') if len(gce_setup_network_interfaces) > 0 else ''
203
204  @property
205  def subnet(self) -> str:
206    gce_setup_network_interfaces = self.gce_setup.get('networkInterfaces', [])
207    return gce_setup_network_interfaces[0].get(
208        'subnet', '') if len(gce_setup_network_interfaces) > 0 else ''
209
210  @property
211  def disable_public_ip(self) -> bool:
212    return self.gce_setup.get('disablePublicIp', False)
213
214  @property
215  def metadata(self) -> dict:
216    # https://cloud.google.com/vertex-ai/docs/workbench/instances/manage-metadata#keys
217    return self.gce_setup.get('metadata', {})
218
219  @property
220  def environment_version(self) -> int:
221    return int(self.metadata.get('version', '0'))
222
223  @property
224  def disable_mixer(self) -> bool:
225    return self.metadata.get('disable-mixer', '').lower() == 'true'
226
227  @property
228  def serial_port_logging_enabled(self) -> bool:
229    return self.metadata.get('serial-port-logging-enable', '').lower() == 'true'
230
231  @property
232  def report_event_health(self) -> bool:
233    return self.metadata.get('report-event-health', '').lower() == 'true'
234
235  @property
236  def post_startup_script(self) -> str:
237    return self.metadata.get('post-startup-script', '')
238
239  @property
240  def startup_script(self) -> str:
241    return self.metadata.get('startup-script', '')
242
243  @property
244  def startup_script_url(self) -> str:
245    return self.metadata.get('startup-script-url', '')
246
247  @property
248  def health_state(self) -> HealthStateEnum:
249    return self._resource_data.get(HEALTH_STATE_KEY,
250                                   HealthStateEnum.HEALTH_STATE_UNSPECIFIED)
251
252  @property
253  def health_info(self) -> dict:
254    return self._resource_data.get(HEALTH_INFO_KEY, {})
255
256  @property
257  def is_jupyterlab_status_healthy(self) -> bool:
258    return self.health_info.get('jupyterlab_status', '') == '1'
259
260  @property
261  def is_jupyterlab_api_status_healthy(self) -> bool:
262    return self.health_info.get('jupyterlab_api_status', '') == '1'
263
264  @property
265  def is_notebooks_api_dns_healthy(self) -> bool:
266    return self.health_info.get('notebooks_api_dns', '') == '1'
267
268  @property
269  def is_proxy_registration_dns_healthy(self) -> bool:
270    return self.health_info.get('proxy_registration_dns', '') == '1'
271
272  @property
273  def is_system_healthy(self) -> bool:
274    return self.health_info.get('system_health', '') == '1'
275
276  @property
277  def is_docker_status_healthy(self) -> bool:
278    return self.health_info.get('docker_status', '') == '1'
279
280  @property
281  def is_docker_proxy_agent_status_healthy(self) -> bool:
282    return self.metadata.get('docker_proxy_agent_status', '') == '1'

285@caching.cached_api_call
286def get_instances(context: models.Context) -> Mapping[str, Instance]:
287  instances: Dict[str, Instance] = {}
288  if not apis.is_enabled(context.project_id, 'notebooks'):
289    return instances
290  logging.info(
291      'fetching list of Vertex AI Workbench notebook instances in project %s',
292      context.project_id)
293  notebooks_api = apis.get_api('notebooks', 'v1', context.project_id)
294  query = notebooks_api.projects().locations().instances().list(
295      parent=f'projects/{context.project_id}/locations/-'
296  )  #'-' (wildcard) all regions
297  try:
298    resp = query.execute(num_retries=config.API_RETRIES)
299    if INSTANCES_KEY not in resp:
300      return instances
301    for i in resp[INSTANCES_KEY]:
302      # verify that we have some minimal data that we expect
303      if NAME_KEY not in i:
304        raise RuntimeError(
305            'missing instance name in projects.locations.instances.list response'
306        )
307      # projects/{projectId}/locations/{location}/instances/{instanceId}
308      result = re.match(r'projects/[^/]+/locations/([^/]+)/instances/([^/]+)',
309                        i['name'])
310      if not result:
311        logging.error('invalid notebook instances data: %s', i['name'])
312        continue
313
314      if not context.match_project_resource(location=result.group(1),
315                                            resource=result.group(2),
316                                            labels=i.get('labels', {})):
317        continue
318      instances[i[NAME_KEY]] = Instance(project_id=context.project_id,
319                                        resource_data=i)
320  except googleapiclient.errors.HttpError as err:
321    raise utils.GcpApiError(err) from err
322  return instances

336def get_instance_health_info(context: models.Context, name: str) -> dict:
337  try:
338    return _get_instance_health(context, name).get(HEALTH_INFO_KEY, {})
339  except googleapiclient.errors.HttpError as err:
340    raise utils.GcpApiError(err) from err

343def get_instance_health_state(context: models.Context,
344                              name: str) -> HealthStateEnum:
345  instance_health_state = HealthStateEnum('HEALTH_STATE_UNSPECIFIED')
346
347  try:
348    resp = _get_instance_health(context, name)
349    if HEALTH_STATE_KEY not in resp:
350      raise RuntimeError(
351          'missing instance health state in projects.locations.instances:getInstanceHealth response'
352      )
353    instance_health_state = HealthStateEnum(resp[HEALTH_STATE_KEY])
354    return instance_health_state
355  except googleapiclient.errors.HttpError as err:
356    raise utils.GcpApiError(err) from err
357  return instance_health_state

360@caching.cached_api_call
361def instance_is_upgradeable(
362    context: models.Context,
363    notebook_instance: str) -> Dict[str, Union[str, bool]]:
364  is_upgradeable: Dict[str, Union[str, bool]] = {}
365  if not apis.is_enabled(context.project_id, 'notebooks'):
366    logging.error('Notebooks API is not enabled')
367    return is_upgradeable
368  if not notebook_instance:
369    logging.error('notebookInstance not provided')
370    return is_upgradeable
371  logging.info(
372      'fetching Vertex AI user-managed notebook instance is upgradeable in project %s',
373      context.project_id)
374  notebooks_api = apis.get_api('notebooks', 'v1', context.project_id)
375  query = notebooks_api.projects().locations().instances().isUpgradeable(
376      notebookInstance=notebook_instance)
377  try:
378    resp = query.execute(num_retries=config.API_RETRIES)
379    if 'upgradeable' not in resp:
380      raise RuntimeError(
381          'missing instance upgradeable data in projects.locations.instances:isUpgradeable response'
382      )
383    is_upgradeable = resp
384    return is_upgradeable
385  except googleapiclient.errors.HttpError as err:
386    raise utils.GcpApiError(err) from err
387  return is_upgradeable

390@caching.cached_api_call
391def get_runtimes(context: models.Context) -> Mapping[str, Runtime]:
392  runtimes: Dict[str, Runtime] = {}
393  if not apis.is_enabled(context.project_id, 'notebooks'):
394    return runtimes
395  logging.info(
396      'fetching list of Vertex AI Workbench managed notebook runtimes in project %s',
397      context.project_id)
398  notebooks_api = apis.get_api('notebooks', 'v1', context.project_id)
399  query = notebooks_api.projects().locations().runtimes().list(
400      parent=f'projects/{context.project_id}/locations/-'
401  )  #'-' (wildcard) all regions
402  try:
403    resp = query.execute(num_retries=config.API_RETRIES)
404    if RUNTIMES_KEY not in resp:
405      return runtimes
406    for i in resp[RUNTIMES_KEY]:
407      # verify that we have some minimal data that we expect
408      if NAME_KEY not in i:
409        raise RuntimeError(
410            'missing runtime name in projects.locations.runtimes.list response')
411      # projects/{projectId}/locations/{location}/runtimes/{runtimeId}
412      result = re.match(r'projects/[^/]+/locations/([^/]+)/runtimes/([^/]+)',
413                        i['name'])
414      if not result:
415        logging.error('invalid notebook runtimes data: %s', i['name'])
416        continue
417
418      if not context.match_project_resource(location=result.group(1),
419                                            resource=result.group(2),
420                                            labels=i.get('labels', {})):
421        continue
422      runtimes[i[NAME_KEY]] = Runtime(project_id=context.project_id,
423                                      resource_data=i)
424  except googleapiclient.errors.HttpError as err:
425    raise utils.GcpApiError(err) from err
426  return runtimes

429@caching.cached_api_call
430def get_workbench_instance(project_id: str, zone: str,
431                           instance_name: str) -> Instance:
432  """Returns workbench instance object matching instance name and zone
433  https://cloud.google.com/vertex-ai/docs/workbench/reference/rest/v2/projects.locations.instances/get
434  """
435  workbench_instance: WorkbenchInstance = WorkbenchInstance(
436      project_id=project_id, resource_data={})
437  if not apis.is_enabled(project_id, 'notebooks'):
438    return workbench_instance
439  notebooks_api = apis.get_api('notebooks', 'v2', project_id)
440  name = f'projects/{project_id}/locations/{zone}/instances/{instance_name}'
441  query = notebooks_api.projects().locations().instances().get(name=name)
442  try:
443    response = query.execute(num_retries=config.API_RETRIES)
444    workbench_instance = WorkbenchInstance(project_id=project_id,
445                                           resource_data=response)
446  except googleapiclient.errors.HttpError as err:
447    raise utils.GcpApiError(err) from err
448  return workbench_instance

451@caching.cached_api_call
452def workbench_instance_check_upgradability(
453    project_id: str,
454    workbench_instance_name: str) -> Dict[str, Union[str, bool]]:
455  """Returns if workbench instance is upgradable and upgrade details
456  https://cloud.google.com/vertex-ai/docs/workbench/reference/rest/v2/projects.locations.instances/checkUpgradability"""
457  check_upgradability: Dict[str, Union[str, bool]] = {}
458  if not apis.is_enabled(project_id, 'notebooks'):
459    logging.error('Notebooks API is not enabled')
460    return check_upgradability
461  if not workbench_instance_name:
462    logging.error('Workbench Instance name not provided')
463    return check_upgradability
464  logging.info(
465      'fetching Vertex AI Workbench Instance is upgradeable in project %s',
466      project_id)
467  notebooks_api = apis.get_api('notebooks', 'v2', project_id)
468  query = notebooks_api.projects().locations().instances().checkUpgradability(
469      notebookInstance=workbench_instance_name)
470  try:
471    response = query.execute(num_retries=config.API_RETRIES)
472    return response
473  except googleapiclient.errors.HttpError as err:
474    raise utils.GcpApiError(err) from err
475  return check_upgradability