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 Instance(models.Resource):
54  """Represent a Vertex AI Workbench user-managed notebook instance
55
56  https://cloud.google.com/vertex-ai/docs/workbench/reference/rest/v1/projects.locations.instances#resource:-instance
57  """
58
59  _resource_data: dict
60
61  def __init__(self, project_id, resource_data):
62    super().__init__(project_id=project_id)
63    self._resource_data = resource_data
64
65  @property
66  def full_path(self) -> str:
67    """
68    The 'name' of the instance is already in the full path form
69    projects/{project}/locations/{location}/instances/{instance}.
70    """
71    return self._resource_data[NAME_KEY]
72
73  @property
74  def short_path(self) -> str:
75    path = self.full_path
76    path = re.sub(r'^projects/', '', path)
77    path = re.sub(r'/locations/', '/', path)
78    path = re.sub(r'/instances/', '/', path)
79    return path
80
81  @property
82  def metadata(self) -> dict:
83    return self._resource_data.get('metadata', {})
84
85  @property
86  def name(self) -> str:
87    logging.info(self._resource_data)
88    return self._resource_data[NAME_KEY]

 91class Runtime(models.Resource):
 92  """Represent a Vertex AI Workbench runtime for a managed notebook instance
 93
 94  https://cloud.google.com/vertex-ai/docs/workbench/reference/rest/v1/projects.locations.runtimes#resource:-runtime
 95  """
 96
 97  _resource_data: dict
 98
 99  def __init__(self, project_id, resource_data):
100    super().__init__(project_id=project_id)
101    self._resource_data = resource_data
102
103  @property
104  def full_path(self) -> str:
105    """
106    The 'name' of the runtime is already in the full path form
107    projects/{project}/locations/{location}/runtimes/{instance}.
108    """
109    return self._resource_data[NAME_KEY]
110
111  @property
112  def short_path(self) -> str:
113    path = self.full_path
114    path = re.sub(r'^projects/', '', path)
115    path = re.sub(r'/locations/', '/', path)
116    path = re.sub(r'/runtimes/', '/', path)
117    return path
118
119  @property
120  def metadata(self) -> dict:
121    return self._resource_data.get('metadata', {})
122
123  @property
124  def name(self) -> str:
125    logging.info(self._resource_data)
126    return self._resource_data[NAME_KEY]
127
128  @property
129  def software_config(self) -> dict:
130    return self._resource_data.get('softwareConfig', {})
131
132  @property
133  def idle_shutdown(self) -> bool:
134    return self.software_config.get('idleShutdown', False)
135
136  @property
137  def is_upgradeable(self) -> bool:
138    return self.software_config.get('upgradeable', False)
139
140  @property
141  def version(self) -> str:
142    return self.software_config.get('version', '')
143
144  @property
145  def health_state(self) -> HealthStateEnum:
146    return self._resource_data.get(HEALTH_STATE_KEY,
147                                   HealthStateEnum.HEALTH_STATE_UNSPECIFIED)

150@caching.cached_api_call
151def get_instances(context: models.Context) -> Mapping[str, Instance]:
152  instances: Dict[str, Instance] = {}
153  if not apis.is_enabled(context.project_id, 'notebooks'):
154    return instances
155  logging.info(
156      'fetching list of Vertex AI Workbench notebook instances in project %s',
157      context.project_id)
158  notebooks_api = apis.get_api('notebooks', 'v1', context.project_id)
159  query = notebooks_api.projects().locations().instances().list(
160      parent=f'projects/{context.project_id}/locations/-'
161  )  #'-' (wildcard) all regions
162  try:
163    resp = query.execute(num_retries=config.API_RETRIES)
164    if INSTANCES_KEY not in resp:
165      return instances
166    for i in resp[INSTANCES_KEY]:
167      # verify that we have some minimal data that we expect
168      if NAME_KEY not in i:
169        raise RuntimeError(
170            'missing instance name in projects.locations.instances.list response'
171        )
172      # projects/{projectId}/locations/{location}/instances/{instanceId}
173      result = re.match(r'projects/[^/]+/locations/([^/]+)/instances/([^/]+)',
174                        i['name'])
175      if not result:
176        logging.error('invalid notebook instances data: %s', i['name'])
177        continue
178
179      if not context.match_project_resource(location=result.group(1),
180                                            resource=result.group(2),
181                                            labels=i.get('labels', {})):
182        continue
183      instances[i[NAME_KEY]] = Instance(project_id=context.project_id,
184                                        resource_data=i)
185  except googleapiclient.errors.HttpError as err:
186    raise utils.GcpApiError(err) from err
187  return instances

201def get_instance_health_info(context: models.Context, name: str) -> dict:
202  try:
203    return _get_instance_health(context, name).get(HEALTH_INFO_KEY, {})
204  except googleapiclient.errors.HttpError as err:
205    raise utils.GcpApiError(err) from err

208def get_instance_health_state(context: models.Context,
209                              name: str) -> HealthStateEnum:
210  instance_health_state = HealthStateEnum('HEALTH_STATE_UNSPECIFIED')
211
212  try:
213    resp = _get_instance_health(context, name)
214    if HEALTH_STATE_KEY not in resp:
215      raise RuntimeError(
216          'missing instance health state in projects.locations.instances:getInstanceHealth response'
217      )
218    instance_health_state = HealthStateEnum(resp[HEALTH_STATE_KEY])
219    return instance_health_state
220  except googleapiclient.errors.HttpError as err:
221    raise utils.GcpApiError(err) from err
222  return instance_health_state

225@caching.cached_api_call
226def instance_is_upgradeable(
227    context: models.Context,
228    notebook_instance: str) -> Dict[str, Union[str, bool]]:
229  is_upgradeable: Dict[str, Union[str, bool]] = {}
230  if not apis.is_enabled(context.project_id, 'notebooks'):
231    logging.error('Notebooks API is not enabled')
232    return is_upgradeable
233  if not notebook_instance:
234    logging.error('notebookInstance not provided')
235    return is_upgradeable
236  logging.info(
237      'fetching Vertex AI user-managed notebook instance is upgradeable in project %s',
238      context.project_id)
239  notebooks_api = apis.get_api('notebooks', 'v1', context.project_id)
240  query = notebooks_api.projects().locations().instances().isUpgradeable(
241      notebookInstance=notebook_instance)
242  try:
243    resp = query.execute(num_retries=config.API_RETRIES)
244    if 'upgradeable' not in resp:
245      raise RuntimeError(
246          'missing instance upgradeable data in projects.locations.instances:isUpgradeable response'
247      )
248    is_upgradeable = resp
249    return is_upgradeable
250  except googleapiclient.errors.HttpError as err:
251    raise utils.GcpApiError(err) from err
252  return is_upgradeable

255@caching.cached_api_call
256def get_runtimes(context: models.Context) -> Mapping[str, Runtime]:
257  runtimes: Dict[str, Runtime] = {}
258  if not apis.is_enabled(context.project_id, 'notebooks'):
259    return runtimes
260  logging.info(
261      'fetching list of Vertex AI Workbench managed notebook runtimes in project %s',
262      context.project_id)
263  notebooks_api = apis.get_api('notebooks', 'v1', context.project_id)
264  query = notebooks_api.projects().locations().runtimes().list(
265      parent=f'projects/{context.project_id}/locations/-'
266  )  #'-' (wildcard) all regions
267  try:
268    resp = query.execute(num_retries=config.API_RETRIES)
269    if RUNTIMES_KEY not in resp:
270      return runtimes
271    for i in resp[RUNTIMES_KEY]:
272      # verify that we have some minimal data that we expect
273      if NAME_KEY not in i:
274        raise RuntimeError(
275            'missing runtime name in projects.locations.runtimes.list response')
276      # projects/{projectId}/locations/{location}/runtimes/{runtimeId}
277      result = re.match(r'projects/[^/]+/locations/([^/]+)/runtimes/([^/]+)',
278                        i['name'])
279      if not result:
280        logging.error('invalid notebook runtimes data: %s', i['name'])
281        continue
282
283      if not context.match_project_resource(location=result.group(1),
284                                            resource=result.group(2),
285                                            labels=i.get('labels', {})):
286        continue
287      runtimes[i[NAME_KEY]] = Runtime(project_id=context.project_id,
288                                      resource_data=i)
289  except googleapiclient.errors.HttpError as err:
290    raise utils.GcpApiError(err) from err
291  return runtimes