Skip to content

Commit 5d67288

Browse files
committed
delete versionset and add sortid
1 parent c0100cd commit 5d67288

File tree

1 file changed

+1
-19
lines changed
  • metadata-ingestion/src/datahub/ingestion/source

1 file changed

+1
-19
lines changed

metadata-ingestion/src/datahub/ingestion/source/mlflow.py

+1-19
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,6 @@
6161
TagPropertiesClass,
6262
TimeStampClass,
6363
VersionPropertiesClass,
64-
VersionSetKeyClass,
6564
VersionTagClass,
6665
_Aspect,
6766
)
@@ -498,7 +497,6 @@ def _get_ml_model_workunits(self) -> Iterable[MetadataWorkUnit]:
498497
for registered_model in registered_models:
499498
version_set_urn = self._get_version_set_urn(registered_model)
500499
yield self._get_ml_group_workunit(registered_model)
501-
yield self._get_version_set(version_set_urn)
502500
model_versions = self._get_mlflow_model_versions(registered_model)
503501
for model_version in model_versions:
504502
run = self._get_mlflow_run(model_version)
@@ -521,22 +519,6 @@ def _get_version_set_urn(self, registered_model: RegisteredModel) -> VersionSetU
521519

522520
return version_set_urn
523521

524-
def _get_version_set(
525-
self,
526-
version_set_urn: VersionSetUrn,
527-
) -> MetadataWorkUnit:
528-
version_set_key = VersionSetKeyClass(
529-
id=version_set_urn.id,
530-
entityType="mlModel",
531-
)
532-
533-
wu = MetadataChangeProposalWrapper(
534-
entityUrn=str(version_set_urn),
535-
aspect=version_set_key,
536-
).as_workunit()
537-
538-
return wu
539-
540522
def _get_ml_model_version_properties_workunit(
541523
self,
542524
model_version: ModelVersion,
@@ -554,7 +536,7 @@ def _get_ml_model_version_properties_workunit(
554536
),
555537
),
556538
versionSet=str(version_set_urn),
557-
sortId="",
539+
sortId=str(model_version.version).zfill(10),
558540
aliases=[
559541
VersionTagClass(versionTag=alias) for alias in model_version.aliases
560542
],

0 commit comments

Comments
 (0)