def execute(self, context: "Context"):
hook = ModelServiceHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
impersonation_chain=self.impersonation_chain,
)
self.log.info("Upload model")
operation = hook.upload_model(
project_id=self.project_id,
region=self.region,
model=self.model,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
result = hook.wait_for_operation(timeout=self.timeout,
operation=operation)
model_resp = model_service.UploadModelResponse.to_dict(result)
model_id = hook.extract_model_id(model_resp)
self.log.info("Model was uploaded. Model ID: %s", model_id)
self.xcom_push(context, key="model_id", value=model_id)
VertexAIModelLink.persist(context=context,
task_instance=self,
model_id=model_id)
return model_resp
def execute(self, context: "Context"):
hook = ModelServiceHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
impersonation_chain=self.impersonation_chain,
)
results = hook.list_models(
project_id=self.project_id,
region=self.region,
filter=self.filter,
page_size=self.page_size,
page_token=self.page_token,
read_mask=self.read_mask,
order_by=self.order_by,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
VertexAIModelListLink.persist(context=context, task_instance=self)
return [Model.to_dict(result) for result in results]
def execute(self, context: "Context"):
hook = ModelServiceHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
impersonation_chain=self.impersonation_chain,
)
try:
self.log.info("Deleting model: %s", self.model_id)
operation = hook.delete_model(
project_id=self.project_id,
region=self.region,
model=self.model_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
hook.wait_for_operation(timeout=self.timeout, operation=operation)
self.log.info("Model was deleted.")
except NotFound:
self.log.info("The Model ID %s does not exist.", self.model_id)
def execute(self, context: "Context"):
hook = ModelServiceHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
impersonation_chain=self.impersonation_chain,
)
try:
self.log.info("Exporting model: %s", self.model_id)
operation = hook.export_model(
project_id=self.project_id,
region=self.region,
model=self.model_id,
output_config=self.output_config,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
hook.wait_for_operation(timeout=self.timeout, operation=operation)
VertexAIModelExportLink.persist(context=context,
task_instance=self)
self.log.info("Model was exported.")
except NotFound:
self.log.info("The Model ID %s does not exist.", self.model_id)