61
61
TagPropertiesClass ,
62
62
TimeStampClass ,
63
63
VersionPropertiesClass ,
64
- VersionSetKeyClass ,
65
64
VersionTagClass ,
66
65
_Aspect ,
67
66
)
@@ -498,7 +497,6 @@ def _get_ml_model_workunits(self) -> Iterable[MetadataWorkUnit]:
498
497
for registered_model in registered_models :
499
498
version_set_urn = self ._get_version_set_urn (registered_model )
500
499
yield self ._get_ml_group_workunit (registered_model )
501
- yield self ._get_version_set (version_set_urn )
502
500
model_versions = self ._get_mlflow_model_versions (registered_model )
503
501
for model_version in model_versions :
504
502
run = self ._get_mlflow_run (model_version )
@@ -521,22 +519,6 @@ def _get_version_set_urn(self, registered_model: RegisteredModel) -> VersionSetU
521
519
522
520
return version_set_urn
523
521
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
-
540
522
def _get_ml_model_version_properties_workunit (
541
523
self ,
542
524
model_version : ModelVersion ,
@@ -554,7 +536,7 @@ def _get_ml_model_version_properties_workunit(
554
536
),
555
537
),
556
538
versionSet = str (version_set_urn ),
557
- sortId = "" ,
539
+ sortId = str ( model_version . version ). zfill ( 10 ) ,
558
540
aliases = [
559
541
VersionTagClass (versionTag = alias ) for alias in model_version .aliases
560
542
],
0 commit comments