def registerWorkflow(self, request, context):
workflow = transform_workflow_meta(request.workflow)
response = self.store.register_workflow(name=workflow.name,
project_id=workflow.project_id,
properties=workflow.properties)
return _wrap_meta_response(
MetaToProto.workflow_meta_to_proto(response))
def getModelById(self, request, context):
model_relation = self.store.get_model_relation_by_id(request.id)
if model_relation is None:
model_detail = None
else:
model_detail = self.model_center_client.get_registered_model_detail(model_relation.name)
return _wrap_meta_response(MetaToProto.model_meta_to_proto(model_relation, model_detail))
def registerJob(self, request, context):
job = transform_job_meta(request.job)
response = self.store.register_job(name=job.name, workflow_execution_id=job.workflow_execution_id,
job_state=job.job_state, properties=job.properties, job_id=job.job_id,
start_time=job.start_time, end_time=job.end_time, log_uri=job.log_uri,
signature=job.signature)
return _wrap_meta_response(MetaToProto.job_meta_to_proto(response))
def updateProject(self, request, context):
properties = None if request.properties == {} else request.properties
project = self.store.update_project(
project_name=request.name,
uri=request.uri.value if request.HasField('uri') else None,
properties=properties)
return _wrap_meta_response(MetaToProto.project_meta_to_proto(project))
def updateDataset(self, request, context):
properties = None if request.properties == {} else request.properties
name_list = request.name_list
type_list = request.type_list
if not name_list:
name_list = None
if not type_list:
data_type_list = None
else:
data_type_list = []
for data_type in type_list:
data_type_list.append(DataType(DataTypeProto.Name(data_type)))
dataset_meta = self.store.update_dataset(dataset_name=request.name,
data_format=request.data_format.value if request.HasField(
'data_format') else None,
description=request.description.value if request.HasField(
'description') else None,
uri=request.uri.value if request.HasField('uri') else None,
properties=properties,
name_list=name_list,
type_list=data_type_list,
catalog_name=request.catalog_name.value if request.HasField(
'catalog_name') else None,
catalog_type=request.catalog_type.value if request.HasField(
'catalog_type') else None,
catalog_database=request.catalog_database.value if request.HasField(
'catalog_database') else None,
catalog_connection_uri=request.catalog_connection_uri.value \
if request.HasField('catalog_connection_uri') else None,
catalog_table=request.catalog_table.value if request.HasField(
'catalog_table') else None)
return _wrap_meta_response(
MetaToProto.dataset_meta_to_proto(dataset_meta))
def registerModel(self, request, context):
model = transform_model_meta(request.model)
model_detail = self.model_center_client.create_registered_model(
model.name, ModelType.Name(model.model_type), model.model_desc)
model_relation = self.store.register_model_relation(
name=model.name, project_id=model.project_id)
return _wrap_meta_response(
MetaToProto.model_meta_to_proto(model_relation, model_detail))
def updateWorkflow(self, request, context):
properties = None if request.properties == {} else request.properties
workflow = self.store.update_workflow(
workflow_name=request.workflow_name,
project_name=request.project_name,
properties=properties)
return _wrap_meta_response(
MetaToProto.workflow_meta_to_proto(workflow))
def registerArtifact(self, request, context):
artifact = transform_artifact_meta(request.artifact)
response = self.store.register_artifact(name=artifact.name, data_format=artifact.data_format,
description=artifact.description,
batch_uri=artifact.batch_uri, stream_uri=artifact.stream_uri,
create_time=artifact.create_time,
update_time=artifact.update_time, properties=artifact.properties)
return _wrap_meta_response(MetaToProto.artifact_meta_to_proto(response))
def getModelVersionByVersion(self, request, context):
model_version_relation = self.store.get_model_version_relation_by_version(request.name, request.model_id)
if model_version_relation is None:
model_version_detail = None
else:
model_relation = self.store.get_model_relation_by_id(model_version_relation.model_id)
model_version_detail = self.model_center_client.get_model_version_detail(model_relation.name, request.name)
return _wrap_meta_response(
MetaToProto.model_version_meta_to_proto(model_version_relation, model_version_detail))
def registerModelVersionRelation(self, request, context):
model_version = transform_model_version_relation_meta(
request.model_version_relation)
response = self.store.register_model_version_relation(
version=model_version.version,
model_id=model_version.model_id,
workflow_execution_id=model_version.workflow_execution_id)
return _wrap_meta_response(
MetaToProto.model_version_relation_meta_to_proto(response))
def registerArtifact(self, request, context):
artifact = transform_artifact_meta(request.artifact)
response = self.store.register_artifact(
name=artifact.name,
artifact_type=artifact.artifact_type,
description=artifact.description,
uri=artifact.uri,
properties=artifact.properties)
return _wrap_meta_response(
MetaToProto.artifact_meta_to_proto(response))
def registerProject(self, request, context):
project = transform_project_meta(request.project)
response = self.store.register_project(
name=project.name,
uri=project.uri,
properties=project.properties,
user=project.user,
password=project.password,
project_type=project.project_type)
return _wrap_meta_response(MetaToProto.project_meta_to_proto(response))
def getLatestGeneratedModelVersion(self, request, context):
model_version_detail = self.store.get_latest_generated_model_version(request.name)
if model_version_detail is None:
model_version_relation = None
else:
model_relation = self.store.get_model_relation_by_name(request.name)
model_version_relation = self.store.get_model_version_relation_by_version(
model_version_detail.model_version,
model_relation.uuid)
return _wrap_meta_response(
MetaToProto.model_version_store_to_proto(model_version_relation, model_version_detail))
def registerExampleWithCatalog(self, request, context):
example = transform_example_meta(request.example)
example_meta = self.store.register_example_with_catalog(name=example.name,
support_type=example.support_type,
catalog_name=example.catalog_name,
catalog_type=example.catalog_type,
catalog_database=example.catalog_database,
catalog_connection_uri=example.catalog_connection_uri,
catalog_version=example.catalog_version,
catalog_table=example.catalog_table)
return _wrap_meta_response(MetaToProto.example_meta_to_proto(example_meta))
def registerDatasetWithCatalog(self, request, context):
dataset = transform_dataset_meta(request.dataset)
dataset_meta = self.store.register_dataset_with_catalog(
name=dataset.name,
catalog_name=dataset.catalog_name,
catalog_type=dataset.catalog_type,
catalog_database=dataset.catalog_database,
catalog_connection_uri=dataset.catalog_connection_uri,
catalog_table=dataset.catalog_table)
return _wrap_meta_response(
MetaToProto.dataset_meta_to_proto(dataset_meta))
def register_datasets(self,
datasets: List[DatasetMeta]) -> List[DatasetMeta]:
"""
register multiple datasets in metadata store.
:param datasets: A list of datasets
:return: List of :py:class:`ai_flow.meta.dataset_meta.DatasetMeta` objects.
"""
request = metadata_service_pb2.RegisterDatasetsRequest(
datasets=MetaToProto.dataset_meta_list_to_proto(datasets))
response = self.metadata_store_stub.registerDatasets(request)
return _unwrap_dataset_list_response(response)
def updateJob(self, request, context):
properties = None if request.properties == {} else request.properties
job_state = None if request.job_state == 0 else State(StateProto.Name(request.job_state))
job = self.store.update_job(job_name=request.name, job_state=job_state,
properties=properties,
job_id=request.job_id.value if request.HasField('job_id') else None,
workflow_execution_id=request.workflow_execution_id.value
if request.HasField('workflow_execution_id') else None,
end_time=request.end_time.value if request.HasField('end_time') else None,
log_uri=request.log_uri.value if request.HasField('log_uri') else None,
signature=request.signature.value if request.HasField('signature') else None)
return _wrap_meta_response(MetaToProto.job_meta_to_proto(job))
def _warp_job_list_response(job_list):
if job_list is not None:
job_proto_list = MetaToProto.job_meta_list_to_proto(job_list)
return Response(return_code=str(SUCCESS),
return_msg=ReturnCode.Name(SUCCESS).lower(),
data=MessageToJson(JobListProto(jobs=job_proto_list),
preserving_proto_field_name=True))
else:
return Response(
return_code=str(RESOURCE_DOES_NOT_EXIST),
return_msg=ReturnCode.Name(RESOURCE_DOES_NOT_EXIST).lower(),
data=None)
def registerDataset(self, request, context):
dataset = transform_dataset_meta(request.dataset)
dataset_meta = self.store.register_dataset(
name=dataset.name,
data_format=dataset.data_format,
description=dataset.description,
uri=dataset.uri,
properties=dataset.properties,
name_list=dataset.schema.name_list,
type_list=dataset.schema.type_list)
return _wrap_meta_response(
MetaToProto.dataset_meta_to_proto(dataset_meta))
def updateArtifact(self, request, context):
properties = None if request.properties == {} else request.properties
artifact = self.store.update_artifact(
name=request.name,
artifact_type=request.artifact_type.value
if request.HasField('artifact_type') else None,
properties=properties,
description=request.description.value
if request.HasField('description') else None,
uri=request.uri.value if request.HasField('uri') else None,
)
return _wrap_meta_response(
MetaToProto.artifact_meta_to_proto(artifact))
def registerWorkFlowExecution(self, request, context):
workflow_execution = transform_workflow_execution_meta(request.workflow_execution)
response = self.store.register_workflow_execution(name=workflow_execution.name,
project_id=workflow_execution.project_id,
execution_state=workflow_execution.execution_state,
properties=workflow_execution.properties,
start_time=workflow_execution.start_time,
end_time=workflow_execution.end_time,
log_uri=workflow_execution.log_uri,
workflow_json=workflow_execution.workflow_json,
signature=workflow_execution.signature
)
return _wrap_meta_response(MetaToProto.workflow_execution_meta_to_proto(response))
def _warp_artifact_list_response(artifact_list):
if artifact_list is not None:
artifact_proto_list = MetaToProto.artifact_meta_list_to_proto(
artifact_list)
return Response(return_code=str(SUCCESS),
return_msg=ReturnCode.Name(SUCCESS).lower(),
data=MessageToJson(
ArtifactListProto(artifacts=artifact_proto_list),
preserving_proto_field_name=True))
else:
return Response(
return_code=str(RESOURCE_DOES_NOT_EXIST),
return_msg=ReturnCode.Name(RESOURCE_DOES_NOT_EXIST).lower(),
data=None)
def _warp_example_list_response(example_meta_list):
if example_meta_list is not None:
example_proto_list = MetaToProto.example_meta_list_to_proto(
example_meta_list)
return Response(return_code=str(SUCCESS),
return_msg=ReturnCode.Name(SUCCESS).lower(),
data=MessageToJson(
ExampleListProto(examples=example_proto_list),
preserving_proto_field_name=True))
else:
return Response(
return_code=str(RESOURCE_DOES_NOT_EXIST),
return_msg=ReturnCode.Name(RESOURCE_DOES_NOT_EXIST).lower(),
data=None)
def _warp_model_version_relation_list_response(model_version_relation_list):
if model_version_relation_list is not None:
model_version_proto_list = MetaToProto.model_version_relation_meta_list_to_proto(
model_version_relation_list)
return Response(return_code=str(SUCCESS),
return_msg=ReturnCode.Name(SUCCESS).lower(),
data=MessageToJson(ModelVersionRelationListProto(
model_versions=model_version_proto_list),
preserving_proto_field_name=True))
else:
return Response(
return_code=str(RESOURCE_DOES_NOT_EXIST),
return_msg=ReturnCode.Name(RESOURCE_DOES_NOT_EXIST).lower(),
data=None)
def _wrap_workflow_list_response(workflow_list):
if workflow_list is not None:
workflow_proto_list = MetaToProto.workflow_meta_list_to_proto(
workflow_list)
return Response(return_code=str(SUCCESS),
return_msg=ReturnCode.Name(SUCCESS).lower(),
data=MessageToJson(
WorkflowListProto(workflows=workflow_proto_list),
preserving_proto_field_name=True))
else:
return Response(
return_code=str(RESOURCE_DOES_NOT_EXIST),
return_msg=ReturnCode.Name(RESOURCE_DOES_NOT_EXIST).lower(),
data=None)
def registerModelVersion(self, request, context):
model_version = transform_model_version_meta(request.model_version)
model_relation = self.store.get_model_relation_by_id(model_version.model_id)
model_version_detail = self.model_center_client.create_model_version(model_relation.name,
model_version.model_path,
model_version.model_metric,
model_version.model_flavor,
model_version.version_desc,
request.model_version.current_stage)
model_version_relation = self.store.register_model_version_relation(version=model_version_detail.model_version,
model_id=model_version.model_id,
workflow_execution_id=
model_version.workflow_execution_id)
return _wrap_meta_response(
MetaToProto.model_version_meta_to_proto(model_version_relation, model_version_detail))
def registerExample(self, request, context):
example = transform_example_meta(request.example)
example_meta = self.store.register_example(name=example.name,
support_type=example.support_type,
data_type=example.data_type,
data_format=example.data_format,
description=example.description,
batch_uri=example.batch_uri,
stream_uri=example.stream_uri,
create_time=example.create_time,
update_time=example.update_time,
properties=example.properties,
name_list=example.schema.name_list,
type_list=example.schema.type_list)
return _wrap_meta_response(MetaToProto.example_meta_to_proto(example_meta))
def updateArtifact(self, request, context):
properties = None if request.properties == {} else request.properties
artifact = self.store.update_artifact(artifact_name=request.name,
data_format=request.data_format.value if request.HasField(
'data_format') else None,
properties=properties,
description=request.description.value if request.HasField(
'description') else None,
batch_uri=request.batch_uri.value if request.HasField(
'batch_uri') else None,
stream_uri=request.stream_uri.value if request.HasField(
'stream_uri') else None,
update_time=request.update_time.value if request.HasField(
'update_time') else None
)
return _wrap_meta_response(MetaToProto.artifact_meta_to_proto(artifact))
def updateWorkflowExecution(self, request, context):
properties = None if request.properties == {} else request.properties
execution_state = None if request.execution_state == 0 else State(StateProto.Name(request.execution_state))
workflow_execution = self.store.update_workflow_execution(execution_name=request.name,
execution_state=execution_state,
project_id=request.project_id.value if request.HasField(
'project_id') else None,
properties=properties,
end_time=request.end_time.value if request.HasField(
'end_time') else None,
log_uri=request.log_uri_value if request.HasField(
'log_uri') else None,
workflow_json=request.workjson.value if request.HasField(
'workflow_json') else None,
signature=request.signature.value if request.HasField(
'signature') else None)
return _wrap_meta_response(MetaToProto.workflow_execution_meta_to_proto(workflow_execution))
def updateExample(self, request, context):
support_type = None if request.support_type == 0 else ExampleSupportType(
ExampleSupportTypeProto.Name(request.support_type))
properties = None if request.properties == {} else request.properties
name_list = request.name_list
type_list = request.type_list
if not name_list:
name_list = None
if not type_list:
data_type_list = None
else:
data_type_list = []
for data_type in type_list:
data_type_list.append(DataType(DataTypeProto.Name(data_type)))
example_meta = self.store.update_example(example_name=request.name,
support_type=support_type,
data_type=request.data_type.value if request.HasField(
'data_type') else None,
data_format=request.data_format.value if request.HasField(
'data_format') else None,
description=request.description.value if request.HasField(
'description') else None,
batch_uri=request.batch_uri.value if request.HasField(
'batch_uri') else None,
stream_uri=request.stream_uri.value if request.HasField(
'stream_uri') else None,
update_time=request.update_time.value if request.HasField(
'update_time') else None,
properties=properties,
name_list=name_list,
type_list=data_type_list,
catalog_name=request.catalog_name.value if request.HasField(
'catalog_name') else None,
catalog_type=request.catalog_type.value if request.HasField(
'catalog_type') else None,
catalog_database=request.catalog_database.value if request.HasField(
'catalog_database') else None,
catalog_version=request.catalog_version.value if request.HasField(
'catalog_version') else None,
catalog_connection_uri=request.catalog_connection_uri.value \
if request.HasField('catalog_connection_uri') else None,
catalog_table=request.catalog_table.value if request.HasField(
'catalog_table') else None)
return _wrap_meta_response(
MetaToProto.example_meta_to_proto(example_meta))
