Skip to content

Commit 6bb1fbb

Browse files
authored
Merge pull request #77 from comet-ml/CM-2344-add-registry-model-version-stage
[CM-2344]: Implement equivalent of API.add_registry_model_version_stage
2 parents a0333e2 + 6208324 commit 6bb1fbb

24 files changed

+181
-50
lines changed

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ public class RegistryModelExample {
4444
static final String SOME_MODEL_VERSION = "1.0.0";
4545
static final String SOME_MODEL_VERSION_UP = "1.0.1";
4646
static final String STAGE_PRODUCTION = "production";
47+
static final String STAGE_STAGING = "staging";
4748
static final String SOME_NOTES = "some model notes";
4849

4950
/**
@@ -147,6 +148,13 @@ record = api.registerModel(updatedModel, experiment.getExperimentKey());
147148
System.out.printf("Overview of the model '%s' not found\n", registryName);
148149
}
149150

151+
// add stage to the model
152+
//
153+
System.out.printf("Adding stage `%s' to the registered model version '%s:%s'\n",
154+
STAGE_STAGING, registryName, SOME_MODEL_VERSION_UP);
155+
api.addRegistryModelVersionStage(registryName, experiment.getWorkspaceName(),
156+
SOME_MODEL_VERSION_UP, STAGE_STAGING);
157+
150158
// get details about model version
151159
//
152160
System.out.printf("Retrieving details of the model version '%s:%s'\n", registryName, SOME_MODEL_VERSION_UP);

comet-examples/src/main/java/ml/comet/examples/mnist/MnistExperimentExample.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ public final class MnistExperimentExample {
4848
* The number of epochs to perform.
4949
*/
5050
@Parameter(names = {"--epochs", "-e"}, description = "number of epochs to perform")
51+
final
5152
int numEpochs = 2;
5253

5354
/**

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,4 +206,14 @@ void updateRegistryModelVersion(String registryName, String workspace, String ve
206206
* @param version the version of the registered model to be deleted.
207207
*/
208208
void deleteRegistryModelVersion(String registryName, String workspace, String version);
209+
210+
/**
211+
* Adds a stage to a registered model version.
212+
*
213+
* @param registryName the name of the model.
214+
* @param workspace the name of the model's workspace.
215+
* @param version the version of the registered model to be updated.
216+
* @param stage the name of the stage to be added.
217+
*/
218+
void addRegistryModelVersionStage(String registryName, String workspace, String version, String stage);
209219
}

comet-java-client/src/main/java/ml/comet/experiment/builder/OnlineExperimentBuilder.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ public interface OnlineExperimentBuilder extends BaseCometBuilder<OnlineExperime
3434
/**
3535
* Set the URL of your comet installation.
3636
*
37-
* @param urlOverride full url of comet installation. Default is https://www.comet.ml
37+
* @param urlOverride full url of comet installation. Default is https://www.comet.com
3838
* @return the builder configured with specified URL of the Comet installation.
3939
*/
4040
OnlineExperimentBuilder withUrlOverride(String urlOverride);

comet-java-client/src/main/java/ml/comet/experiment/impl/BaseExperimentAsync.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@
7575
* using asynchronous networking.
7676
*/
7777
abstract class BaseExperimentAsync extends BaseExperiment {
78-
ExperimentContext baseContext;
78+
final ExperimentContext baseContext;
7979

8080
BaseExperimentAsync(@NonNull final String apiKey,
8181
@NonNull final String baseUrl,

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

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,15 @@
1616
import ml.comet.experiment.impl.rest.ExperimentModelResponse;
1717
import ml.comet.experiment.impl.rest.RegistryModelCountResponse;
1818
import ml.comet.experiment.impl.rest.RegistryModelCreateRequest;
19+
import ml.comet.experiment.impl.rest.RegistryModelDeleteRequest;
1920
import ml.comet.experiment.impl.rest.RegistryModelDetailsResponse;
2021
import ml.comet.experiment.impl.rest.RegistryModelItemCreateRequest;
2122
import ml.comet.experiment.impl.rest.RegistryModelNotesResponse;
2223
import ml.comet.experiment.impl.rest.RegistryModelNotesUpdateRequest;
2324
import ml.comet.experiment.impl.rest.RegistryModelOverviewListResponse;
2425
import ml.comet.experiment.impl.rest.RegistryModelUpdateItemRequest;
2526
import ml.comet.experiment.impl.rest.RegistryModelUpdateRequest;
27+
import ml.comet.experiment.impl.rest.RegistryModelVersionStageAddRequest;
2628
import ml.comet.experiment.impl.rest.RestApiResponse;
2729
import ml.comet.experiment.impl.utils.CometUtils;
2830
import ml.comet.experiment.impl.utils.ExceptionUtils;
@@ -71,6 +73,7 @@
7173
import static ml.comet.experiment.impl.resources.LogMessages.EXPERIMENT_HAS_NO_MODELS;
7274
import static ml.comet.experiment.impl.resources.LogMessages.EXPERIMENT_WITH_KEY_NOT_FOUND;
7375
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;
7477
import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_DELETE_REGISTRY_MODEL;
7578
import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_DELETE_REGISTRY_MODEL_VERSION;
7679
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
505508
this.updateRegistryModelVersion(registryName, workspace, version, null, null);
506509
}
507510

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+
508529
@Override
509530
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);
513534
}
514535

515536
@Override
@@ -523,9 +544,8 @@ public void deleteRegistryModelVersion(@NonNull String registryName, @NonNull St
523544
getString(REGISTRY_MODEL_VERSION_NOT_FOUND, version, workspace, registryName));
524545
}
525546
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,
527548
versionOverviewOptional.get().getRegistryModelItemId(), errorMsg);
528-
this.checkRestApiResponse(response, errorMsg);
529549
}
530550

531551
/**

comet-java-client/src/main/java/ml/comet/experiment/impl/OnlineExperimentImpl.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -755,7 +755,7 @@ Optional<Action> getLogAssetOnCompleteAction() {
755755
* The runnable to be invoked to send periodic heartbeat ping to mark this experiment as still running.
756756
*/
757757
static class HeartbeatPing implements Runnable {
758-
OnlineExperimentImpl onlineExperiment;
758+
final OnlineExperimentImpl onlineExperiment;
759759

760760
HeartbeatPing(OnlineExperimentImpl onlineExperiment) {
761761
this.onlineExperiment = onlineExperiment;

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

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
import ml.comet.experiment.impl.rest.RegistryModelCountResponse;
4545
import ml.comet.experiment.impl.rest.RegistryModelCreateRequest;
4646
import ml.comet.experiment.impl.rest.RegistryModelCreateResponse;
47+
import ml.comet.experiment.impl.rest.RegistryModelDeleteRequest;
4748
import ml.comet.experiment.impl.rest.RegistryModelDetailsResponse;
4849
import ml.comet.experiment.impl.rest.RegistryModelItemCreateRequest;
4950
import ml.comet.experiment.impl.rest.RegistryModelItemCreateResponse;
@@ -52,6 +53,7 @@
5253
import ml.comet.experiment.impl.rest.RegistryModelOverviewListResponse;
5354
import ml.comet.experiment.impl.rest.RegistryModelUpdateItemRequest;
5455
import ml.comet.experiment.impl.rest.RegistryModelUpdateRequest;
56+
import ml.comet.experiment.impl.rest.RegistryModelVersionStageAddRequest;
5557
import ml.comet.experiment.impl.rest.RestApiResponse;
5658
import ml.comet.experiment.impl.rest.SetSystemDetailsRequest;
5759
import ml.comet.experiment.impl.rest.TagsResponse;
@@ -74,6 +76,7 @@
7476
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_METRIC;
7577
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_OUTPUT;
7678
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_PARAMETER;
79+
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_REGISTRY_MODEL_VERSION_STAGE;
7780
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_START_END_TIME;
7881
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_TAG;
7982
import static ml.comet.experiment.impl.constants.ApiEndpoints.CREATE_REGISTRY_MODEL;
@@ -118,6 +121,7 @@
118121
import static ml.comet.experiment.impl.constants.QueryParamName.MODEL_NAME;
119122
import static ml.comet.experiment.impl.constants.QueryParamName.PROJECT_ID;
120123
import static ml.comet.experiment.impl.constants.QueryParamName.PROJECT_NAME;
124+
import static ml.comet.experiment.impl.constants.QueryParamName.STAGE;
121125
import static ml.comet.experiment.impl.constants.QueryParamName.TYPE;
122126
import static ml.comet.experiment.impl.constants.QueryParamName.WORKSPACE_NAME;
123127
import static ml.comet.experiment.impl.http.ConnectionUtils.checkResponseStatus;
@@ -383,6 +387,13 @@ Single<RegistryModelCountResponse> getRegistryModelsCount(String workspaceName)
383387
RegistryModelCountResponse.class);
384388
}
385389

390+
Single<RestApiResponse> addRegistryModelVersionStage(RegistryModelVersionStageAddRequest request) {
391+
Map<QueryParamName, String> queryParams = new HashMap<>();
392+
queryParams.put(MODEL_ITEM_ID, request.getRegistryModelItemId());
393+
queryParams.put(STAGE, request.getStage());
394+
return this.singleFromSyncGetWithRetries(ADD_REGISTRY_MODEL_VERSION_STAGE, queryParams);
395+
}
396+
386397
Single<RestApiResponse> downloadRegistryModel(
387398
final OutputStream output, String workspace, String registryName, final DownloadModelOptions options) {
388399
Map<QueryParamName, String> queryParams = downloadModelParams(workspace, registryName, options);
@@ -409,10 +420,10 @@ Single<RestApiResponse> updateRegistryModelVersion(RegistryModelUpdateItemReques
409420
return singleFromAsyncPost(request, UPDATE_REGISTRY_MODEL_VERSION);
410421
}
411422

412-
Single<RestApiResponse> deleteRegistryModel(String modelName, String workspaceName) {
423+
Single<RestApiResponse> deleteRegistryModel(RegistryModelDeleteRequest request) {
413424
Map<QueryParamName, String> queryParams = new HashMap<>();
414-
queryParams.put(WORKSPACE_NAME, workspaceName);
415-
queryParams.put(MODEL_NAME, modelName);
425+
queryParams.put(WORKSPACE_NAME, request.getWorkspace());
426+
queryParams.put(MODEL_NAME, request.getRegistryModelName());
416427
return singleFromSyncGetWithRetries(DELETE_REGISTRY_MODEL, queryParams);
417428
}
418429

comet-java-client/src/main/java/ml/comet/experiment/impl/asset/DownloadArtifactAssetOptions.java

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package ml.comet.experiment.impl.asset;
22

3-
import lombok.AllArgsConstructor;
43
import lombok.Data;
54
import lombok.EqualsAndHashCode;
65
import lombok.NoArgsConstructor;

comet-java-client/src/main/java/ml/comet/experiment/impl/asset/RemoteAssetImpl.java

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import lombok.EqualsAndHashCode;
44
import lombok.Getter;
5-
import lombok.NoArgsConstructor;
65
import lombok.Setter;
76
import ml.comet.experiment.asset.RemoteAsset;
87

0 commit comments

Comments
 (0)