gcpdiag.queries.notebooks

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

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

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

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

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

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

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

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

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

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

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

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