Merge pull request #76 from comet-ml/CM-2347-registry-model-count

[CM-2347]: Implement equivalent of API.get_registry_model_count

Author: yaricom

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

@@ -203,6 +203,17 @@ record = api.registerModel(updatedModel, experiment.getExperimentKey());
                     newComment, newStages);
             System.out.printf("Model version details was successfully updated for: %s\n", SOME_MODEL_VERSION_UP);
 
+            // read the number of registered models
+            //
+            Optional<Integer> modelsCount = api.getRegistryModelsCount(experiment.getWorkspaceName());
+            if (modelsCount.isPresent()) {
+                System.out.printf("The number of registered models under workspace '%s' is: %d\n",
+                        experiment.getWorkspaceName(), modelsCount.get());
+            } else {
+                System.out.printf("Failed to get registered models count. The workspace '%s' doesn't exists.\n",
+                        experiment.getWorkspaceName());
+            }
+
             // delete registry model version
             //
             System.out.printf("Deleting registry model version '%s/%s:%s'.\n",

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

@@ -129,6 +129,15 @@ ModelDownloadInfo downloadRegistryModel(Path outputPath, String registryName, St
      */
     List<String> getRegistryModelVersions(String registryName, String workspace);
 
+    /**
+     * Returns the number of models registered in the specified workspace that is available to the current user.
+     *
+     * @param workspace the name of the workspace.
+     * @return the {@code Optional} number of models registered in the specified workspace.
+     * If workspace doesn't exist the empty {@code Optional} returned.
+     */
+    Optional<Integer> getRegistryModelsCount(String workspace);
+
     /**
      * Adds or updates notes associated with the registry model.
      *

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

@@ -14,6 +14,7 @@
 import ml.comet.experiment.impl.http.ConnectionInitializer;
 import ml.comet.experiment.impl.rest.ExperimentModelListResponse;
 import ml.comet.experiment.impl.rest.ExperimentModelResponse;
+import ml.comet.experiment.impl.rest.RegistryModelCountResponse;
 import ml.comet.experiment.impl.rest.RegistryModelCreateRequest;
 import ml.comet.experiment.impl.rest.RegistryModelDetailsResponse;
 import ml.comet.experiment.impl.rest.RegistryModelItemCreateRequest;
@@ -74,6 +75,7 @@
 import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_DELETE_REGISTRY_MODEL_VERSION;
 import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_DOWNLOAD_REGISTRY_MODEL;
 import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_FIND_EXPERIMENT_MODEL_BY_NAME;
+import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_GET_REGISTRY_MODELS_COUNT;
 import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_GET_REGISTRY_MODEL_DETAILS;
 import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_GET_REGISTRY_MODEL_NOTES;
 import static ml.comet.experiment.impl.resources.LogMessages.FAILED_TO_GET_REGISTRY_MODEL_VERSIONS;
@@ -395,6 +397,22 @@ public List<String> getRegistryModelVersions(@NonNull String registryName, @NonN
                 ArrayList::addAll);
     }
 
+    @Override
+    public Optional<Integer> getRegistryModelsCount(@NonNull String workspace) {
+        try {
+            RegistryModelCountResponse countResponse = this.restApiClient
+                    .getRegistryModelsCount(workspace)
+                    .blockingGet();
+            return Optional.of(countResponse.getRegistryModelCount());
+        } catch (CometApiException ex) {
+            this.logger.error(getString(FAILED_TO_GET_REGISTRY_MODELS_COUNT, workspace), ex);
+            if (ex.getStatusCode() == 400) {
+                return Optional.empty();
+            }
+            throw ex;
+        }
+    }
+
     @Override
     public void updateRegistryModelNotes(
             @NonNull String notes, @NonNull String registryName, @NonNull String workspace) {

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

@@ -41,6 +41,7 @@
 import ml.comet.experiment.impl.rest.MinMaxResponse;
 import ml.comet.experiment.impl.rest.OutputUpdate;
 import ml.comet.experiment.impl.rest.ParameterRest;
+import ml.comet.experiment.impl.rest.RegistryModelCountResponse;
 import ml.comet.experiment.impl.rest.RegistryModelCreateRequest;
 import ml.comet.experiment.impl.rest.RegistryModelCreateResponse;
 import ml.comet.experiment.impl.rest.RegistryModelDetailsResponse;
@@ -94,6 +95,7 @@
 import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_METRICS;
 import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_OUTPUT;
 import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_PARAMETERS;
+import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_REGISTRY_MODELS_COUNT;
 import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_REGISTRY_MODEL_DETAILS;
 import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_REGISTRY_MODEL_LIST;
 import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_REGISTRY_MODEL_NOTES;
@@ -375,6 +377,12 @@ Single<RegistryModelDetailsResponse> getRegistryModelDetails(String modelName, S
                 RegistryModelDetailsResponse.class);
     }
 
+    Single<RegistryModelCountResponse> getRegistryModelsCount(String workspaceName) {
+        return this.singleFromSyncGetWithRetries(GET_REGISTRY_MODELS_COUNT,
+                Collections.singletonMap(WORKSPACE_NAME, workspaceName), true,
+                RegistryModelCountResponse.class);
+    }
+
     Single<RestApiResponse> downloadRegistryModel(
             final OutputStream output, String workspace, String registryName, final DownloadModelOptions options) {
         Map<QueryParamName, String> queryParams = downloadModelParams(workspace, registryName, options);

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

@@ -53,4 +53,6 @@ public final class ApiEndpoints {
     public static final String DOWNLOAD_REGISTRY_MODEL = READ_API_URL + "/registry-model/item/download";
     public static final String GET_REGISTRY_MODEL_DETAILS = READ_API_URL + "/registry-model/details";
     public static final String GET_REGISTRY_MODEL_NOTES = READ_API_URL + "/registry-model/notes";
+
+    public static final String GET_REGISTRY_MODELS_COUNT = READ_API_URL + "/registry-model/count";
 }

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

@@ -95,6 +95,7 @@ public class LogMessages {
     public static final String FAILED_TO_GET_REGISTRY_MODEL_DETAILS = "FAILED_TO_GET_REGISTRY_MODEL_DETAILS";
     public static final String FAILED_TO_GET_REGISTRY_MODEL_VERSIONS = "FAILED_TO_GET_REGISTRY_MODEL_VERSIONS";
     public static final String FAILED_TO_GET_REGISTRY_MODEL_NOTES = "FAILED_TO_GET_REGISTRY_MODEL_NOTES";
+    public static final String FAILED_TO_GET_REGISTRY_MODELS_COUNT = "FAILED_TO_GET_REGISTRY_MODELS_COUNT";
     public static final String FAILED_TO_UPDATE_REGISTRY_MODEL_NOTES = "FAILED_TO_UPDATE_REGISTRY_MODEL_NOTES";
     public static final String FAILED_TO_UPDATE_REGISTRY_MODEL = "FAILED_TO_UPDATE_REGISTRY_MODEL";
     public static final String FAILED_TO_UPDATE_REGISTRY_MODEL_VERSION = "FAILED_TO_UPDATE_REGISTRY_MODEL_VERSION";

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

@@ -0,0 +1,16 @@
+package ml.comet.experiment.impl.rest;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import lombok.AllArgsConstructor;
+import lombok.Data;
+import lombok.NoArgsConstructor;
+
+@JsonIgnoreProperties(ignoreUnknown = true)
+@JsonInclude(JsonInclude.Include.NON_NULL)
+@NoArgsConstructor
+@AllArgsConstructor
+@Data
+public class RegistryModelCountResponse {
+    private int registryModelCount;
+}

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

@@ -80,6 +80,7 @@ FAILED_TO_DOWNLOAD_REGISTRY_MODEL=Failed to download registry model, reason: '%s
 FAILED_TO_GET_REGISTRY_MODEL_DETAILS=Failed to get details of the registry model '%s/%s', reason: '%s', sdk error code: '%s'.
 FAILED_TO_GET_REGISTRY_MODEL_VERSIONS=Failed to get versions of the registry model '%s/%s'.
 FAILED_TO_GET_REGISTRY_MODEL_NOTES=Failed to get notes of the registry model '%s/%s', reason: '%s', sdk error code: '%s'.
+FAILED_TO_GET_REGISTRY_MODELS_COUNT=Failed to get number of registry models under workspace '%s'.
 FAILED_TO_UPDATE_REGISTRY_MODEL_NOTES=Failed to update notes of the registry model '%s/%s'.
 FAILED_TO_UPDATE_REGISTRY_MODEL=Failed to update registry model '%s/%s' with data: '%s'.
 FAILED_TO_UPDATE_REGISTRY_MODEL_VERSION=Failed to update registry model's version '%s/%s:%s' with data: '%s'.

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

@@ -477,6 +477,37 @@ public void testUpdateRegistryModelNotes_model_doesnt_exists() {
                 COMET_API.updateRegistryModelNotes(SOME_NOTES, modelName, SHARED_EXPERIMENT.getWorkspaceName()));
     }
 
+    @Test
+    public void testGetRegistryModelsCount() {
+        String modelName = String.format("%s-%d", SOME_MODEL_NAME, System.currentTimeMillis());
+
+        // register model with defaults
+        //
+        registerModelWithDefaults(modelName);
+
+        // wait for registry model to be processed by backend
+        //
+        Awaitility.await("failed to get registry model")
+                .pollInterval(1, TimeUnit.SECONDS)
+                .atMost(60, TimeUnit.SECONDS)
+                .until(() -> COMET_API.getRegistryModelDetails(modelName,
+                        SHARED_EXPERIMENT.getWorkspaceName()).isPresent());
+
+        // get models count and do assert
+        //
+        Optional<Integer> modelsCount = COMET_API.getRegistryModelsCount(SHARED_EXPERIMENT.getWorkspaceName());
+        assertTrue(modelsCount.isPresent());
+        // there could be more than one model registered from other tests
+        assertTrue(modelsCount.get() > 1);
+    }
+
+    @Test
+    public void testGetRegistryModelsCount_not_existing_workspace() {
+        String notExistingWorkspace = "not existing workspace";
+        Optional<Integer> modelsCount = COMET_API.getRegistryModelsCount(notExistingWorkspace);
+        assertFalse(modelsCount.isPresent());
+    }
+
     @Test
     public void testUpdateRegistryModelNotes() {
         String modelName = String.format("%s-%d", SOME_MODEL_NAME, System.currentTimeMillis());