Skip to content

Commit bac421f

Browse files
yaricomyaricom
authored
Merge pull request #61 from comet-ml/CM-2372-update-registry-model-version
[CM-2372]: Implement equivalent of API.update_registry_model_version
2 parents 1c66049 + 3b98d76 commit bac421f

File tree

10 files changed

+190
-0
lines changed

10 files changed

+190
-0
lines changed

‎comet-examples/src/main/java/ml/comet/examples/RegistryModelExample.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,16 @@ record = api.registerModel(updatedModel, experiment.getExperimentKey());
193193
api.updateRegistryModel(registryName, experiment.getWorkspaceName(), newModelName, newDescription);
194194
System.out.println("Model was successfully updated");
195195

196+
// update registry model version details
197+
//
198+
System.out.printf("Updating version details of the registry model '%s/%s:%s'.\n",
199+
experiment.getWorkspaceName(), newModelName, SOME_MODEL_VERSION_UP);
200+
String newComment = "updated version comment";
201+
List<String> newStages = Collections.singletonList("updated stage");
202+
api.updateRegistryModelVersion(newModelName, experiment.getWorkspaceName(), SOME_MODEL_VERSION_UP,
203+
newComment, newStages);
204+
System.out.printf("Model version details was successfully updated for: %s\n", SOME_MODEL_VERSION_UP);
205+
196206
} finally {
197207
PathUtils.deleteDirectory(modelTmpDir);
198208
}

‎comet-java-client/src/main/java/ml/comet/experiment/CometApi.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,4 +143,18 @@ void updateRegistryModel(String registryName, String workspace, String newRegist
143143

144144
void updateRegistryModel(String registryName, String workspace, String newRegistryName)
145145
throws ModelNotFoundException;
146+
147+
/**
148+
* Updates the comments and stages of particular version of the registered model.
149+
*
150+
* @param registryName the name of the model.
151+
* @param workspace the name of the model's workspace.
152+
* @param version the version of the registered model to be updated.
153+
* @param comments the comment to associate with new version.
154+
* @param stages the stages to associate with new version.
155+
*/
156+
void updateRegistryModelVersion(String registryName, String workspace, String version,
157+
String comments, List<String> stages);
158+
159+
void updateRegistryModelVersion(String registryName, String workspace, String version, String comments);
146160
}

‎comet-java-client/src/main/java/ml/comet/experiment/impl/CometApiImpl.java‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import ml.comet.experiment.impl.rest.RegistryModelNotesResponse;
2020
import ml.comet.experiment.impl.rest.RegistryModelNotesUpdateRequest;
2121
import ml.comet.experiment.impl.rest.RegistryModelOverviewListResponse;
22+
import ml.comet.experiment.impl.rest.RegistryModelUpdateItemRequest;
2223
import ml.comet.experiment.impl.rest.RegistryModelUpdateRequest;
2324
import ml.comet.experiment.impl.rest.RestApiResponse;
2425
import ml.comet.experiment.impl.utils.CometUtils;
@@ -34,6 +35,7 @@
3435
import ml.comet.experiment.registrymodel.ModelNotFoundException;
3536
import ml.comet.experiment.registrymodel.ModelOverview;
3637
import ml.comet.experiment.registrymodel.ModelRegistryRecord;
38+
import ml.comet.experiment.registrymodel.ModelVersionNotFoundException;
3739
import ml.comet.experiment.registrymodel.ModelVersionOverview;
3840
import org.apache.commons.lang3.StringUtils;
3941
import org.slf4j.Logger;
@@ -72,9 +74,11 @@
7274
import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_GET_REGISTRY_MODEL_VERSIONS;
7375
import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_UPDATE_REGISTRY_MODEL;
7476
import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_UPDATE_REGISTRY_MODEL_NOTES;
77+
import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_UPDATE_REGISTRY_MODEL_VERSION;
7578
import static ml.comet.experiment.impl.resources.LogMessages.MODEL_REGISTERED_IN_WORKSPACE;
7679
import static ml.comet.experiment.impl.resources.LogMessages.MODEL_VERSION_CREATED_IN_WORKSPACE;
7780
import static ml.comet.experiment.impl.resources.LogMessages.REGISTRY_MODEL_NOT_FOUND;
81+
import static ml.comet.experiment.impl.resources.LogMessages.REGISTRY_MODEL_VERSION_NOT_FOUND;
7882
import static ml.comet.experiment.impl.resources.LogMessages.UPDATE_REGISTRY_MODEL_DESCRIPTION_IGNORED;
7983
import static ml.comet.experiment.impl.resources.LogMessages.UPDATE_REGISTRY_MODEL_IS_PUBLIC_IGNORED;
8084
import static ml.comet.experiment.impl.resources.LogMessages.WORKSPACE_HAS_NO_REGISTRY_MODELS;
@@ -380,6 +384,46 @@ private void updateRegistryModel(@NonNull String registryName, @NonNull String w
380384
this.checkRestApiResponse(response, errorMsg);
381385
}
382386

387+
@Override
388+
public void updateRegistryModelVersion(@NonNull String registryName, @NonNull String workspace,
389+
@NonNull String version, String comments, List<String> stages) {
390+
// get registry model details
391+
Optional<ModelOverview> overviewOptional = this.getRegistryModelDetails(registryName, workspace);
392+
if (!overviewOptional.isPresent()) {
393+
throw new ModelNotFoundException(getString(REGISTRY_MODEL_NOT_FOUND, workspace, registryName));
394+
}
395+
396+
ModelOverview modelOverview = overviewOptional.get();
397+
if (modelOverview.getVersions() == null || modelOverview.getVersions().size() == 0) {
398+
throw new ModelNotFoundException(getString(FAILED_TO_GET_REGISTRY_MODEL_VERSIONS, workspace, registryName));
399+
}
400+
401+
// get registry model's item ID which match version
402+
Optional<ModelVersionOverview> versionOverviewOptional = modelOverview.getVersions()
403+
.stream()
404+
.filter(modelVersionOverview -> Objects.equals(modelVersionOverview.getVersion(), version))
405+
.findFirst();
406+
if (!versionOverviewOptional.isPresent()) {
407+
throw new ModelVersionNotFoundException(
408+
getString(REGISTRY_MODEL_VERSION_NOT_FOUND, version, workspace, registryName));
409+
}
410+
411+
// update model version details
412+
RegistryModelUpdateItemRequest request = new RegistryModelUpdateItemRequest();
413+
request.setRegistryModelItemId(versionOverviewOptional.get().getRegistryModelItemId());
414+
request.setComment(comments);
415+
request.setStages(stages);
416+
String errorMsg = getString(FAILED_TO_UPDATE_REGISTRY_MODEL_VERSION, workspace, registryName, version, request);
417+
RestApiResponse response = this.executeSyncRequest(
418+
this.restApiClient::updateRegistryModelVersion, request, errorMsg);
419+
this.checkRestApiResponse(response, errorMsg);
420+
}
421+
422+
@Override
423+
public void updateRegistryModelVersion(String registryName, String workspace, String version, String comments) {
424+
this.updateRegistryModelVersion(registryName, workspace, version, null, null);
425+
}
426+
383427
/**
384428
* Release all resources hold by this instance, such as connection to the Comet server.
385429
*

‎comet-java-client/src/main/java/ml/comet/experiment/impl/RestApiClient.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
import ml.comet.experiment.impl.rest.RegistryModelNotesResponse;
5050
import ml.comet.experiment.impl.rest.RegistryModelNotesUpdateRequest;
5151
import ml.comet.experiment.impl.rest.RegistryModelOverviewListResponse;
52+
import ml.comet.experiment.impl.rest.RegistryModelUpdateItemRequest;
5253
import ml.comet.experiment.impl.rest.RegistryModelUpdateRequest;
5354
import ml.comet.experiment.impl.rest.RestApiResponse;
5455
import ml.comet.experiment.impl.rest.SetSystemDetailsRequest;
@@ -102,6 +103,7 @@
102103
import static ml.comet.experiment.impl.constants.ApiEndpoints.UPDATE_ARTIFACT_STATE;
103104
import static ml.comet.experiment.impl.constants.ApiEndpoints.UPDATE_REGISTRY_MODEL;
104105
import static ml.comet.experiment.impl.constants.ApiEndpoints.UPDATE_REGISTRY_MODEL_NOTES;
106+
import static ml.comet.experiment.impl.constants.ApiEndpoints.UPDATE_REGISTRY_MODEL_VERSION;
105107
import static ml.comet.experiment.impl.constants.ApiEndpoints.UPSERT_ARTIFACT;
106108
import static ml.comet.experiment.impl.constants.ApiEndpoints.WORKSPACES;
107109
import static ml.comet.experiment.impl.constants.FormParamName.LINK;
@@ -379,6 +381,10 @@ Single<RestApiResponse> updateRegistryModel(RegistryModelUpdateRequest request)
379381
return singleFromAsyncPost(request, UPDATE_REGISTRY_MODEL);
380382
}
381383

384+
Single<RestApiResponse> updateRegistryModelVersion(RegistryModelUpdateItemRequest request) {
385+
return singleFromAsyncPost(request, UPDATE_REGISTRY_MODEL_VERSION);
386+
}
387+
382388
private Single<RestApiResponse> singleFromAsyncDownload(@NonNull OutputStream output,
383389
@NonNull String endpoint,
384390
@NonNull Map<QueryParamName, String> queryParams) {

‎comet-java-client/src/main/java/ml/comet/experiment/impl/constants/ApiEndpoints.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ public final class ApiEndpoints {
2727
public static final String CREATE_REGISTRY_MODEL_ITEM = CREATE_REGISTRY_MODEL + "/item";
2828
public static final String UPDATE_REGISTRY_MODEL_NOTES = UPDATE_API_URL + "/registry-model/notes";
2929
public static final String UPDATE_REGISTRY_MODEL = UPDATE_API_URL + "/registry-model/update";
30+
public static final String UPDATE_REGISTRY_MODEL_VERSION = UPDATE_API_URL + "/registry-model/item/update";
3031

3132
public static final String READ_API_URL = "/api/rest/v2";
3233
public static final String GET_ASSETS_LIST = READ_API_URL + "/experiment/asset/list";

‎comet-java-client/src/main/java/ml/comet/experiment/impl/resources/LogMessages.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,8 @@ public class LogMessages {
9595
public static final String FAILED_TO_GET_REGISTRY_MODEL_NOTES = "FAILED_TO_GET_REGISTRY_MODEL_NOTES";
9696
public static final String FAILED_TO_UPDATE_REGISTRY_MODEL_NOTES = "FAILED_TO_UPDATE_REGISTRY_MODEL_NOTES";
9797
public static final String FAILED_TO_UPDATE_REGISTRY_MODEL = "FAILED_TO_UPDATE_REGISTRY_MODEL";
98+
public static final String FAILED_TO_UPDATE_REGISTRY_MODEL_VERSION = "FAILED_TO_UPDATE_REGISTRY_MODEL_VERSION";
99+
public static final String REGISTRY_MODEL_VERSION_NOT_FOUND = "REGISTRY_MODEL_VERSION_NOT_FOUND";
98100

99101

100102
/**

‎comet-java-client/src/main/java/ml/comet/experiment/impl/rest/RegistryModelUpdateItemRequest.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
package ml.comet.experiment.impl.rest;
2+
3+
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
4+
import com.fasterxml.jackson.annotation.JsonInclude;
5+
import lombok.AllArgsConstructor;
6+
import lombok.Data;
7+
import lombok.NoArgsConstructor;
8+
9+
import java.util.List;
10+
11+
@JsonIgnoreProperties(ignoreUnknown = true)
12+
@JsonInclude(JsonInclude.Include.NON_NULL)
13+
@Data
14+
@NoArgsConstructor
15+
@AllArgsConstructor
16+
public class RegistryModelUpdateItemRequest {
17+
private String registryModelItemId;
18+
private String comment;
19+
private List<String> stages;
20+
}

‎comet-java-client/src/main/java/ml/comet/experiment/registrymodel/ModelVersionNotFoundException.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
package ml.comet.experiment.registrymodel;
2+
3+
import ml.comet.experiment.exception.CometApiException;
4+
5+
/**
6+
* Exception to be thrown if specific version of the registered Comet model was not found.
7+
*/
8+
public class ModelVersionNotFoundException extends CometApiException {
9+
/**
10+
* Constructs a new runtime exception with the specified detail message.
11+
* The cause is not initialized, and may subsequently be initialized by a
12+
* call to {@link #initCause}.
13+
*
14+
* @param message the detail message. The detail message is saved for
15+
* later retrieval by the {@link #getMessage()} method.
16+
*/
17+
public ModelVersionNotFoundException(String message) {
18+
super(message);
19+
}
20+
}

‎comet-java-client/src/main/resources/messages.properties‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,3 +80,5 @@ FAILED_TO_GET_REGISTRY_MODEL_VERSIONS=Failed to get versions of the registry mod
8080
FAILED_TO_GET_REGISTRY_MODEL_NOTES=Failed to get notes of the registry model '%s/%s', reason: '%s', sdk error code: '%s'.
8181
FAILED_TO_UPDATE_REGISTRY_MODEL_NOTES=Failed to update notes of the registry model '%s/%s'.
8282
FAILED_TO_UPDATE_REGISTRY_MODEL=Failed to update registry model '%s/%s' with data: '%s'.
83+
FAILED_TO_UPDATE_REGISTRY_MODEL_VERSION=Failed to update registry model's version '%s/%s:%s' with data: '%s'.
84+
REGISTRY_MODEL_VERSION_NOT_FOUND=Version '%s' of the registry model '%s/%s' is not found.

‎comet-java-client/src/test/java/ml/comet/experiment/impl/CometApiTest.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import ml.comet.experiment.registrymodel.ModelNotFoundException;
1818
import ml.comet.experiment.registrymodel.ModelOverview;
1919
import ml.comet.experiment.registrymodel.ModelRegistryRecord;
20+
import ml.comet.experiment.registrymodel.ModelVersionNotFoundException;
2021
import ml.comet.experiment.registrymodel.ModelVersionOverview;
2122
import org.apache.commons.io.file.Counters;
2223
import org.apache.commons.io.file.PathUtils;
@@ -479,6 +480,76 @@ public void testUpdateRegistryModel() {
479480
assertFalse(modelOverview.isPublic(), "wrong visibility status");
480481
}
481482

483+
@Test
484+
public void testUpdateRegistryModelVersion_model_doesnt_exists() {
485+
// try to update not existing model
486+
//
487+
String modelName = "not existing model";
488+
assertThrows(CometApiException.class, () ->
489+
COMET_API.updateRegistryModelVersion(
490+
modelName, SHARED_EXPERIMENT.getWorkspaceName(), "1.0.1", "some comment"));
491+
}
492+
493+
@Test
494+
public void testUpdateRegistryModelVersion_version_doesnt_exists() {
495+
String modelName = String.format("%s-%d", SOME_MODEL_NAME, System.currentTimeMillis());
496+
497+
// register model with defaults
498+
//
499+
registerModelWithDefaults(modelName);
500+
501+
// try to update not existing model
502+
//
503+
assertThrows(ModelVersionNotFoundException.class, () ->
504+
COMET_API.updateRegistryModelVersion(
505+
modelName, SHARED_EXPERIMENT.getWorkspaceName(), "1.0.1", "some comment"));
506+
}
507+
508+
@Test
509+
public void testUpdateRegistryModelVersion() {
510+
String modelName = String.format("%s-%d", SOME_MODEL_NAME, System.currentTimeMillis());
511+
512+
// register model with defaults
513+
//
514+
registerModelWithDefaults(modelName);
515+
516+
// wait for registry model to be processed by backend
517+
//
518+
Awaitility.await("failed to get registry model")
519+
.pollInterval(1, TimeUnit.SECONDS)
520+
.atMost(60, TimeUnit.SECONDS)
521+
.until(() -> COMET_API.getRegistryModelDetails(modelName,
522+
SHARED_EXPERIMENT.getWorkspaceName()).isPresent());
523+
524+
// update model's version
525+
//
526+
String comment = "testUpdateRegistryModelVersion comment";
527+
List<String> stages = Collections.singletonList("testUpdateRegistryModelVersion");
528+
COMET_API.updateRegistryModelVersion(
529+
modelName, SHARED_EXPERIMENT.getWorkspaceName(), DEFAULT_MODEL_VERSION, comment, stages);
530+
531+
// get registry model and check that it was updated
532+
//
533+
Awaitility.await("failed to get registry model")
534+
.pollInterval(1, TimeUnit.SECONDS)
535+
.atMost(60, TimeUnit.SECONDS)
536+
.until(() -> COMET_API.getRegistryModelDetails(modelName,
537+
SHARED_EXPERIMENT.getWorkspaceName()).isPresent());
538+
539+
Optional<ModelOverview> overviewOptional = COMET_API.getRegistryModelDetails(modelName,
540+
SHARED_EXPERIMENT.getWorkspaceName());
541+
542+
assertTrue(overviewOptional.isPresent(), "model expected");
543+
ModelOverview modelOverview = overviewOptional.get();
544+
545+
assertNotNull(modelOverview.getVersions(), "versions list expected");
546+
assertEquals(1, modelOverview.getVersions().size(), "wrong versions list size");
547+
548+
ModelVersionOverview versionOverview = modelOverview.getVersions().get(0);
549+
assertEquals(comment, versionOverview.getComment(), "wrong comment");
550+
assertEquals(stages, versionOverview.getStages(), "wrong stages");
551+
}
552+
482553
private static String registerModelWithDefaults(String modelName) {
483554
// log model folder
484555
//

0 commit comments

Comments
 (0)