16
16
import ml .comet .experiment .impl .rest .ExperimentModelResponse ;
17
17
import ml .comet .experiment .impl .rest .RegistryModelCountResponse ;
18
18
import ml .comet .experiment .impl .rest .RegistryModelCreateRequest ;
19
+ import ml .comet .experiment .impl .rest .RegistryModelDeleteRequest ;
19
20
import ml .comet .experiment .impl .rest .RegistryModelDetailsResponse ;
20
21
import ml .comet .experiment .impl .rest .RegistryModelItemCreateRequest ;
21
22
import ml .comet .experiment .impl .rest .RegistryModelNotesResponse ;
22
23
import ml .comet .experiment .impl .rest .RegistryModelNotesUpdateRequest ;
23
24
import ml .comet .experiment .impl .rest .RegistryModelOverviewListResponse ;
24
25
import ml .comet .experiment .impl .rest .RegistryModelUpdateItemRequest ;
25
26
import ml .comet .experiment .impl .rest .RegistryModelUpdateRequest ;
27
+ import ml .comet .experiment .impl .rest .RegistryModelVersionStageAddRequest ;
26
28
import ml .comet .experiment .impl .rest .RestApiResponse ;
27
29
import ml .comet .experiment .impl .utils .CometUtils ;
28
30
import ml .comet .experiment .impl .utils .ExceptionUtils ;
71
73
import static ml .comet .experiment .impl .resources .LogMessages .EXPERIMENT_HAS_NO_MODELS ;
72
74
import static ml .comet .experiment .impl .resources .LogMessages .EXPERIMENT_WITH_KEY_NOT_FOUND ;
73
75
import static ml .comet .experiment .impl .resources .LogMessages .EXTRACTED_N_REGISTRY_MODEL_FILES ;
76
+ import static ml .comet .experiment .impl .resources .LogMessages .FAILED_TO_ADD_REGISTRY_MODEL_VERSION_STAGE ;
74
77
import static ml .comet .experiment .impl .resources .LogMessages .FAILED_TO_DELETE_REGISTRY_MODEL ;
75
78
import static ml .comet .experiment .impl .resources .LogMessages .FAILED_TO_DELETE_REGISTRY_MODEL_VERSION ;
76
79
import static ml .comet .experiment .impl .resources .LogMessages .FAILED_TO_DOWNLOAD_REGISTRY_MODEL ;
@@ -505,11 +508,29 @@ public void updateRegistryModelVersion(String registryName, String workspace, St
505
508
this .updateRegistryModelVersion (registryName , workspace , version , null , null );
506
509
}
507
510
511
+ @ Override
512
+ public void addRegistryModelVersionStage (@ NonNull String registryName , @ NonNull String workspace ,
513
+ @ NonNull String version , @ NonNull String stage ) {
514
+ // get version details
515
+ Optional <ModelVersionOverview > versionOverviewOptional = this .getRegistryModelVersion (
516
+ registryName , workspace , version );
517
+ if (!versionOverviewOptional .isPresent ()) {
518
+ throw new ModelVersionNotFoundException (
519
+ getString (REGISTRY_MODEL_VERSION_NOT_FOUND , version , workspace , registryName ));
520
+ }
521
+
522
+ String errorMessage = getString (
523
+ FAILED_TO_ADD_REGISTRY_MODEL_VERSION_STAGE , stage , workspace , registryName , version );
524
+ this .executeSyncRequest (this .restApiClient ::addRegistryModelVersionStage ,
525
+ new RegistryModelVersionStageAddRequest (versionOverviewOptional .get ().getRegistryModelItemId (), stage ),
526
+ errorMessage );
527
+ }
528
+
508
529
@ Override
509
530
public void deleteRegistryModel (@ NonNull String registryName , @ NonNull String workspace ) {
510
- RestApiResponse response = this . restApiClient . deleteRegistryModel ( registryName , workspace )
511
- . blockingGet ();
512
- this . checkRestApiResponse ( response , getString ( FAILED_TO_DELETE_REGISTRY_MODEL , registryName , workspace ));
531
+ String errorMsg = getString ( FAILED_TO_DELETE_REGISTRY_MODEL , registryName , workspace );
532
+ this . executeSyncRequest ( this . restApiClient :: deleteRegistryModel ,
533
+ new RegistryModelDeleteRequest ( registryName , workspace ), errorMsg );
513
534
}
514
535
515
536
@ Override
@@ -523,9 +544,8 @@ public void deleteRegistryModelVersion(@NonNull String registryName, @NonNull St
523
544
getString (REGISTRY_MODEL_VERSION_NOT_FOUND , version , workspace , registryName ));
524
545
}
525
546
String errorMsg = getString (FAILED_TO_DELETE_REGISTRY_MODEL_VERSION , workspace , registryName , version );
526
- RestApiResponse response = this .executeSyncRequest (this .restApiClient ::deleteRegistryModelVersion ,
547
+ this .executeSyncRequest (this .restApiClient ::deleteRegistryModelVersion ,
527
548
versionOverviewOptional .get ().getRegistryModelItemId (), errorMsg );
528
- this .checkRestApiResponse (response , errorMsg );
529
549
}
530
550
531
551
/**
0 commit comments