def execute(self, context: 'Context'):
hook = EndpointServiceHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
impersonation_chain=self.impersonation_chain,
)
try:
self.log.info("Get endpoint: %s", self.endpoint_id)
endpoint_obj = hook.get_endpoint(
project_id=self.project_id,
region=self.region,
endpoint=self.endpoint_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
VertexAIEndpointLink.persist(context=context,
task_instance=self,
endpoint_id=self.endpoint_id)
self.log.info("Endpoint was gotten.")
return Endpoint.to_dict(endpoint_obj)
except NotFound:
self.log.info("The Endpoint ID %s does not exist.",
self.endpoint_id)
def execute(self, context: 'Context'):
hook = EndpointServiceHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
impersonation_chain=self.impersonation_chain,
)
self.log.info("Creating endpoint")
operation = hook.create_endpoint(
project_id=self.project_id,
region=self.region,
endpoint=self.endpoint,
endpoint_id=self.endpoint_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
result = hook.wait_for_operation(timeout=self.timeout,
operation=operation)
endpoint = Endpoint.to_dict(result)
endpoint_id = hook.extract_endpoint_id(endpoint)
self.log.info("Endpoint was created. Endpoint ID: %s", endpoint_id)
self.xcom_push(context, key="endpoint_id", value=endpoint_id)
VertexAIEndpointLink.persist(context=context,
task_instance=self,
endpoint_id=endpoint_id)
return endpoint
def execute(self, context: 'Context'):
hook = EndpointServiceHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
impersonation_chain=self.impersonation_chain,
)
self.log.info("Updating endpoint: %s", self.endpoint_id)
result = hook.update_endpoint(
project_id=self.project_id,
region=self.region,
endpoint_id=self.endpoint_id,
endpoint=self.endpoint,
update_mask=self.update_mask,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
self.log.info("Endpoint was updated")
VertexAIEndpointLink.persist(context=context,
task_instance=self,
endpoint_id=self.endpoint_id)
return Endpoint.to_dict(result)