ml.comet.experiment.impl.RestApiClient Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of comet-java-client Show documentation
There is a newer version: 1.1.14
package ml.comet.experiment.impl;

import io.reactivex.rxjava3.core.Single;
import io.reactivex.rxjava3.disposables.Disposable;
import lombok.NonNull;
import ml.comet.experiment.exception.CometApiException;
import ml.comet.experiment.impl.asset.Asset;
import ml.comet.experiment.impl.asset.AssetType;
import ml.comet.experiment.impl.asset.RemoteAsset;
import ml.comet.experiment.impl.constants.FormParamName;
import ml.comet.experiment.impl.constants.QueryParamName;
import ml.comet.experiment.impl.http.Connection;
import ml.comet.experiment.impl.utils.AssetUtils;
import ml.comet.experiment.impl.utils.JsonUtils;
import ml.comet.experiment.model.AddExperimentTagsRest;
import ml.comet.experiment.model.AddGraphRest;
import ml.comet.experiment.model.CreateExperimentRequest;
import ml.comet.experiment.model.CreateExperimentResponse;
import ml.comet.experiment.model.ExperimentAssetListResponse;
import ml.comet.experiment.model.ExperimentMetadataRest;
import ml.comet.experiment.model.ExperimentStatusResponse;
import ml.comet.experiment.model.ExperimentTimeRequest;
import ml.comet.experiment.model.GetExperimentsResponse;
import ml.comet.experiment.model.GetGraphResponse;
import ml.comet.experiment.model.GetHtmlResponse;
import ml.comet.experiment.model.GetOutputResponse;
import ml.comet.experiment.model.GetProjectsResponse;
import ml.comet.experiment.model.GetWorkspacesResponse;
import ml.comet.experiment.model.GitMetadata;
import ml.comet.experiment.model.GitMetadataRest;
import ml.comet.experiment.model.HtmlRest;
import ml.comet.experiment.model.LogDataResponse;
import ml.comet.experiment.model.LogOtherRest;
import ml.comet.experiment.model.MetricRest;
import ml.comet.experiment.model.MinMaxResponse;
import ml.comet.experiment.model.OutputUpdate;
import ml.comet.experiment.model.ParameterRest;
import ml.comet.experiment.model.TagsResponse;

import java.io.File;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_ASSET;
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_GIT_METADATA;
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_GRAPH;
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_HTML;
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_LOG_OTHER;
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_METRIC;
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_OUTPUT;
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_PARAMETER;
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_START_END_TIME;
import static ml.comet.experiment.impl.constants.ApiEndpoints.ADD_TAG;
import static ml.comet.experiment.impl.constants.ApiEndpoints.EXPERIMENTS;
import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_ASSET_INFO;
import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_GIT_METADATA;
import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_GRAPH;
import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_HTML;
import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_LOG_OTHER;
import static ml.comet.experiment.impl.constants.ApiEndpoints.GET_METADATA;
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_TAGS;
import static ml.comet.experiment.impl.constants.ApiEndpoints.NEW_EXPERIMENT;
import static ml.comet.experiment.impl.constants.ApiEndpoints.PROJECTS;
import static ml.comet.experiment.impl.constants.ApiEndpoints.SET_EXPERIMENT_STATUS;
import static ml.comet.experiment.impl.constants.ApiEndpoints.WORKSPACES;
import static ml.comet.experiment.impl.constants.FormParamName.LINK;
import static ml.comet.experiment.impl.constants.QueryParamName.EXPERIMENT_KEY;
import static ml.comet.experiment.impl.constants.QueryParamName.IS_REMOTE;
import static ml.comet.experiment.impl.constants.QueryParamName.PROJECT_ID;
import static ml.comet.experiment.impl.constants.QueryParamName.TYPE;
import static ml.comet.experiment.impl.constants.QueryParamName.WORKSPACE_NAME;

/**
 * Represents Comet REST API client providing access to all exposed REST endpoints.
 */
final class RestApiClient implements Disposable {
    private Connection connection;
    private boolean disposed;

    static final IllegalStateException ALREADY_DISPOSED = new IllegalStateException("REST API client already disposed");

    RestApiClient(Connection connection) {
        this.connection = connection;
    }

    Single getAllWorkspaces() {
        return singleFromSyncGet(WORKSPACES, Collections.emptyMap(), GetWorkspacesResponse.class);
    }

    Single getAllProjects(String workspaceName) {
        return singleFromSyncGet(
                PROJECTS, Collections.singletonMap(WORKSPACE_NAME, workspaceName), GetProjectsResponse.class);
    }

    Single getAllExperiments(String projectId) {
        return singleFromSyncGet(
                EXPERIMENTS, Collections.singletonMap(PROJECT_ID, projectId), GetExperimentsResponse.class);
    }

    Single getMetadata(String experimentKey) {
        return singleFromSyncGet(GET_METADATA, experimentKey, ExperimentMetadataRest.class);
    }

    Single getGitMetadata(String experimentKey) {
        return singleFromSyncGet(GET_GIT_METADATA, experimentKey, GitMetadataRest.class);
    }

    Single getHtml(String experimentKey) {
        return singleFromSyncGet(GET_HTML, experimentKey, GetHtmlResponse.class);
    }

    Single getOutput(String experimentKey) {
        return singleFromSyncGet(GET_OUTPUT, experimentKey, GetOutputResponse.class);
    }

    Single getGraph(String experimentKey) {
        return singleFromSyncGet(GET_GRAPH, experimentKey, GetGraphResponse.class);
    }

    Single getParameters(String experimentKey) {
        return singleFromSyncGet(GET_PARAMETERS, experimentKey, MinMaxResponse.class);
    }

    Single getMetrics(String experimentKey) {
        return singleFromSyncGet(GET_METRICS, experimentKey, MinMaxResponse.class);
    }

    Single getLogOther(String experimentKey) {
        return singleFromSyncGet(GET_LOG_OTHER, experimentKey, MinMaxResponse.class);
    }

    Single getTags(String experimentKey) {
        return singleFromSyncGet(GET_TAGS, experimentKey, TagsResponse.class);
    }

    Single getAssetList(String experimentKey, AssetType type) {
        HashMap params = new HashMap() {{
            put(EXPERIMENT_KEY, experimentKey);
            put(TYPE, type.type());
        }};
        return singleFromSyncGet(GET_ASSET_INFO, params, ExperimentAssetListResponse.class);
    }

    Single sendExperimentStatus(String experimentKey) {
        return singleFromSyncGet(SET_EXPERIMENT_STATUS, experimentKey, ExperimentStatusResponse.class);
    }

    Single logMetric(final MetricRest request, String experimentKey) {
        request.setExperimentKey(experimentKey);
        return singleFromAsyncPost(request, ADD_METRIC, LogDataResponse.class);
    }

    Single logParameter(final ParameterRest request, String experimentKey) {
        request.setExperimentKey(experimentKey);
        return singleFromAsyncPost(request, ADD_PARAMETER, LogDataResponse.class);
    }

    Single logOutputLine(final OutputUpdate request, String experimentKey) {
        request.setExperimentKey(experimentKey);
        return singleFromAsyncPost(request, ADD_OUTPUT, LogDataResponse.class);
    }

    Single logHtml(final HtmlRest request, String experimentKey) {
        request.setExperimentKey(experimentKey);
        return singleFromAsyncPost(request, ADD_HTML, LogDataResponse.class);
    }

    Single logOther(final LogOtherRest request, String experimentKey) {
        request.setExperimentKey(experimentKey);
        return singleFromAsyncPost(request, ADD_LOG_OTHER, LogDataResponse.class);
    }

    Single addTag(final AddExperimentTagsRest request, String experimentKey) {
        request.setExperimentKey(experimentKey);
        return singleFromAsyncPost(request, ADD_TAG, LogDataResponse.class);
    }

    Single logGraph(final AddGraphRest request, String experimentKey) {
        request.setExperimentKey(experimentKey);
        return singleFromAsyncPost(request, ADD_GRAPH, LogDataResponse.class);
    }

    Single logStartEndTime(final ExperimentTimeRequest request, String experimentKey) {
        request.setExperimentKey(experimentKey);
        return singleFromAsyncPost(request, ADD_START_END_TIME, LogDataResponse.class);
    }

    Single logGitMetadata(final GitMetadata request, String experimentKey) {
        request.setExperimentKey(experimentKey);
        return singleFromAsyncPost(request, ADD_GIT_METADATA, LogDataResponse.class);
    }

    Single registerExperiment(final CreateExperimentRequest request) {
        return singleFromSyncPost(request, NEW_EXPERIMENT, true, CreateExperimentResponse.class);
    }

    Single logAsset(final Asset asset, String experimentKey) {
        Map queryParams = AssetUtils.assetQueryParameters(asset, experimentKey);
        Map formParams = AssetUtils.assetFormParameters(asset);

        // call appropriate send method
        if (asset.getFile() != null) {
            return singleFromAsyncPost(asset.getFile(), ADD_ASSET, queryParams,
                    formParams, LogDataResponse.class);
        } else if (asset.getFileLikeData() != null) {
            return singleFromAsyncPost(asset.getFileLikeData(), ADD_ASSET, queryParams,
                    formParams, LogDataResponse.class);
        }

        // no data response
        LogDataResponse response = new LogDataResponse();
        response.setMsg("asset has no data");
        response.setCode(-1);
        return Single.just(response);
    }

    Single logRemoteAsset(final RemoteAsset asset, String experimentKey) {
        Map queryParams = AssetUtils.assetQueryParameters(asset, experimentKey);
        queryParams.put(IS_REMOTE, Boolean.valueOf(true).toString());

        Map formParams = AssetUtils.assetFormParameters(asset);
        formParams.put(LINK, asset.getLink().toASCIIString());

        return singleFromAsyncPost(ADD_ASSET, queryParams, formParams, LogDataResponse.class);
    }

    private  Single singleFromAsyncPost(@NonNull String endpoint,
                                              @NonNull Map queryParams,
                                              @NonNull Map formParams,
                                              @NonNull Class clazz) {
        if (isDisposed()) {
            return Single.error(ALREADY_DISPOSED);
        }
        return Single.fromFuture(this.connection.sendPostAsync(endpoint, queryParams, formParams))
                .onTerminateDetach()
                .map(response -> JsonUtils.fromJson(response.getResponseBody(), clazz));
    }

    private  Single singleFromAsyncPost(
            byte[] fileLikeData, @NonNull String endpoint,
            @NonNull Map queryParams, Map formParams,
            @NonNull Class clazz) {
        if (isDisposed()) {
            return Single.error(ALREADY_DISPOSED);
        }

        return Single.fromFuture(this.connection.sendPostAsync(fileLikeData, endpoint, queryParams, formParams))
                .onTerminateDetach()
                .map(response -> JsonUtils.fromJson(response.getResponseBody(), clazz));
    }

    private  Single singleFromAsyncPost(
            @NonNull File file, @NonNull String endpoint,
            @NonNull Map queryParams, Map formParams,
            @NonNull Class clazz) {
        if (isDisposed()) {
            return Single.error(ALREADY_DISPOSED);
        }

        return Single.fromFuture(this.connection.sendPostAsync(file, endpoint, queryParams, formParams))
                .onTerminateDetach()
                .map(response -> JsonUtils.fromJson(response.getResponseBody(), clazz));
    }

    private  Single singleFromAsyncPost(
            @NonNull Object payload, @NonNull String endpoint, @NonNull Class clazz) {
        if (isDisposed()) {
            return Single.error(ALREADY_DISPOSED);
        }

        return Single.fromFuture(this.connection.sendPostAsync(JsonUtils.toJson(payload), endpoint))
                .onTerminateDetach()
                .map(response -> JsonUtils.fromJson(response.getResponseBody(), clazz));
    }

    private  Single singleFromSyncPost(@NonNull Object payload,
                                             @NonNull String endpoint,
                                             boolean throwOnFailure,
                                             @NonNull Class clazz) {
        if (isDisposed()) {
            return Single.error(ALREADY_DISPOSED);
        }

        String request = JsonUtils.toJson(payload);
        return this.connection.sendPostWithRetries(request, endpoint, throwOnFailure)
                .map(body -> Single.just(JsonUtils.fromJson(body, clazz)))
                .orElse(Single.error(new CometApiException(
                        String.format("No response was returned by endpoint: %s", endpoint))));
    }

    private  Single singleFromSyncGet(@NonNull String endpoint,
                                            @NonNull String experimentKey,
                                            @NonNull Class clazz) {
        return singleFromSyncGet(endpoint, Collections.singletonMap(EXPERIMENT_KEY, experimentKey), clazz);
    }

    private  Single singleFromSyncGet(@NonNull String endpoint,
                                            @NonNull Map params,
                                            @NonNull Class clazz) {
        if (isDisposed()) {
            return Single.error(ALREADY_DISPOSED);
        }
        return this.connection.sendGetWithRetries(endpoint, params)
                .map(body -> Single.just(JsonUtils.fromJson(body, clazz)))
                .orElse(Single.error(new CometApiException(
                        String.format("No response was returned by endpoint: %s", endpoint))));
    }

    @Override
    public void dispose() {
        this.disposed = true;
        // release reference to the connection
        this.connection = null;
    }

    @Override
    public boolean isDisposed() {
        return this.disposed;
    }
}