syntax = "proto3";

package api;

import "google/api/annotations.proto";

/* For each RPC service, we define mapping to HTTP REST API method.
 * The mapping includes the URL path, query parameters and request body.
 * https://cloud.google.com/service-infrastructure/docs/service-management/reference/rpc/google.api#http */
service Manager {
    rpc CreateStudy(CreateStudyRequest) returns (CreateStudyReply){
        option (google.api.http) = {
            post: "/api/Manager/CreateStudy"
            body: "study_config"
        };
    };
    rpc GetStudy(GetStudyRequest) returns (GetStudyReply){
        option (google.api.http) = {
            get: "/api/Manager/GetStudy/{study_id}"
        };
    };
    rpc GetStudyList(GetStudyListRequest) returns (GetStudyListReply){
        option (google.api.http) = {
            get: "/api/Manager/GetStudyList"
        };
    };
    rpc CreateTrial(CreateTrialRequest) returns (CreateTrialReply){
        option (google.api.http) = {
            post: "/api/Manager/CreateTrial"
            body: "trial"
        };
    };
    rpc GetTrials(GetTrialsRequest) returns (GetTrialsReply){
        option (google.api.http) = {
            get: "/api/Manager/GetTrials/{study_id}"
        };
    };
    rpc RegisterWorker(RegisterWorkerRequest) returns (RegisterWorkerReply){
        option (google.api.http) = {
            post: "/api/Manager/RegisterWorker"
            body: "worker"
        };
    };
    rpc GetWorkers(GetWorkersRequest) returns (GetWorkersReply){
        option (google.api.http) = {
            get: "/api/Manager/GetWorkers"
        };
    };
    rpc UpdateWorkerState(UpdateWorkerStateRequest) returns (UpdateWorkerStateReply){
        option (google.api.http) = {
            put: "/api/Manager/UpdateWorkerState",
            body: "*"
        };
    };
    rpc GetSuggestions(GetSuggestionsRequest) returns (GetSuggestionsReply){
        option (google.api.http) = {
            post: "/api/Manager/GetSuggestions"
            body: "*"
        };
    };
    rpc GetShouldStopWorkers(GetShouldStopWorkersRequest) returns (GetShouldStopWorkersReply){
        option (google.api.http) = {
            post: "/api/Manager/GetTrials/GetShouldStopWorkers"
            body: "*"
        };
    };
    rpc GetMetrics(GetMetricsRequest) returns (GetMetricsReply){
        option (google.api.http) = {
            post: "/api/Manager/GetMetrics"
            body: "*"
        };
    };
    rpc SetSuggestionParameters(SetSuggestionParametersRequest) returns (SetSuggestionParametersReply){
        option (google.api.http) = {
            post: "/api/Manager/SetSuggestionParameters"
            body: "*"
        };
    };
    rpc GetSuggestionParameters(GetSuggestionParametersRequest) returns (GetSuggestionParametersReply){
        option (google.api.http) = {
            get: "/api/Manager/GetSuggestionParameters/{param_id}"
        };
    };
    rpc GetSuggestionParameterList(GetSuggestionParameterListRequest) returns (GetSuggestionParameterListReply){
        option (google.api.http) = {
            get: "/api/Manager/GetSuggestionParameterList/{study_id}"
        };
    };
    rpc SetEarlyStoppingParameters(SetEarlyStoppingParametersRequest) returns (SetEarlyStoppingParametersReply){
        option (google.api.http) = {
            post: "/api/Manager/SetEarlyStoppingParameters"
            body: "*"
        };
    };
    rpc GetEarlyStoppingParameters(GetEarlyStoppingParametersRequest) returns (GetEarlyStoppingParametersReply){
        option (google.api.http) = {
            get: "/api/Manager/GetEarlyStoppingParameters/{param_id}"
        };
    };
    rpc GetEarlyStoppingParameterList(GetEarlyStoppingParameterListRequest) returns (GetEarlyStoppingParameterListReply){
        option (google.api.http) = {
            get: "/api/Manager/GetEarlyStoppingParameterList/{study_id}"
        };
    };
    rpc SaveStudy(SaveStudyRequest) returns(SaveStudyReply){
        option (google.api.http) = {
            post: "/api/Manager/SaveStudy"
            body: "*"
        };
    };
    rpc SaveModel(SaveModelRequest) returns(SaveModelReply){
        option (google.api.http) = {
            post: "/api/Manager/SaveModel"
            body: "*"
        };
    };
    rpc ReportMetricsLogs(ReportMetricsLogsRequest) returns (ReportMetricsLogsReply){
        option (google.api.http) = {
            post: "/api/Manager/ReportMetricsLogs"
            body: "*"
        };
    };
    rpc GetSavedStudies(GetSavedStudiesRequest) returns(GetSavedStudiesReply){
        option (google.api.http) = {
            get: "/api/Manager/GetSavedStudies"
        };
    };
    rpc GetSavedModels(GetSavedModelsRequest) returns(GetSavedModelsReply){
        option (google.api.http) = {
            get: "/api/Manager/GetSavedModels"
        };
    };
}

service Suggestion {
    rpc GetSuggestions(GetSuggestionsRequest) returns (GetSuggestionsReply);
}

service EarlyStopping {
    rpc GetShouldStopWorkers(GetShouldStopWorkersRequest) returns (GetShouldStopWorkersReply);
}

enum ParameterType {
    // Not used
    UNKNOWN_TYPE = 0;

    DOUBLE = 1;
    INT = 2;
    DISCRETE = 3;
    CATEGORICAL = 4;
}

enum OptimizationType {
    // Not used
    UNKNOWN_OPTIMIZATION = 0;

    MINIMIZE = 1;
    MAXIMIZE = 2;
}

message FeasibleSpace {
    string max = 1;
    string min = 2;
    repeated string list = 3;
}

message ParameterConfig {
    string name = 1;
    ParameterType parameter_type = 2;
    // The following values defines a feasible parameter space.
    FeasibleSpace feasible = 3;
}

message Parameter {
    string name = 1;
    ParameterType parameter_type = 2;
    string value = 3;
}

// This value is stored as TINYINT in MySQL.
enum State {
    PENDING = 0;
    RUNNING = 1;
    COMPLETED = 2;
    KILLED = 3;
    ERROR = 120;
}

message MetricsLogSet {
    string worker_id = 1;
    repeated MetricsLog metrics_logs= 2;
    State worker_status = 3;
}

message Metrics {
    string name = 1;
    string value = 2;
}

message MetricsValueTime {
    string time = 1;
    string value = 2;
}

message MetricsLog {
    string name = 1;
    repeated MetricsValueTime values = 2;
}

message SuggestionParameter {
    string name = 1;
    string value = 2;
}

message EarlyStoppingParameter {
    string name = 1;
    string value = 2;
}

message Tag {
    string name = 1;
    string value = 2;
}

message StudyOverview {
    string name = 1;
    string owner  = 2;
    string id = 3;
    string description = 4;
}

message Trial {
    string trial_id = 1;
    string study_id = 2;
    repeated Parameter parameter_set = 3;
    string objective_value = 4;
    repeated Tag tags = 5;
}

message VolumeConf {
   string name = 1;
   string pvc_name = 2;
   string mount_path = 3;
}

message Worker {
    string worker_id = 1;
    string study_id = 2;
    string trial_id = 3;
    string Type = 4;
    State status = 5;
    string TemplatePath = 6;
    repeated Tag tags = 7;
}

message StudyConfig {
    message ParameterConfigs {
        repeated ParameterConfig configs = 1;
    }
    string name = 1;
    string owner = 2;
    OptimizationType optimization_type = 3;
    double optimization_goal = 4;
    ParameterConfigs parameter_configs = 5;
    repeated string access_permissions = 6;
    repeated Tag tags = 7;
    string objective_value_name = 8;
    repeated string metrics = 9;
    string jobId = 10;
}

message CreateStudyRequest {
    StudyConfig study_config = 1;
}

message CreateStudyReply {
    string study_id = 1;
}

message StopStudyReply {
}

message GetStudyRequest {
    string study_id = 1;
}

message GetStudyReply {
    StudyConfig study_config= 1;
}

message GetStudyListRequest {
}

message GetStudyListReply {
    repeated StudyOverview study_overviews = 1;
}

message CreateTrialRequest {
    Trial trial = 1;
}

message CreateTrialReply {
    string trial_id = 1;
}

message GetTrialsRequest {
    string study_id = 1;
}

message GetTrialsReply {
    repeated Trial trials = 1;
}

message RegisterWorkerRequest {
    Worker worker = 1;
}

message RegisterWorkerReply {
    string worker_id = 1;
}

message StopWorkersRequest {
    string study_id = 1;
    repeated string worker_ids = 2;
    bool is_complete = 3;
}

message StopWorkersReply {
}

message GetWorkersRequest {
    string study_id = 1;
    string trial_id = 2;
    string worker_id = 3;
}

message GetWorkersReply {
    repeated Worker workers = 1;
}

message UpdateWorkerStateRequest{
    string workerId = 1;
    State status = 2;
}

message UpdateWorkerStateReply{
}

message GetSuggestionsRequest {
    string study_id = 1;
    string suggestion_algorithm = 2;
    int32 request_number = 3;
    repeated string log_worker_ids = 4;
    string param_id = 5;
}

message GetSuggestionsReply {
    repeated Trial trials = 1;
}

message GetShouldStopWorkersRequest {
    string study_id = 1;
    string early_stopping_algorithm = 2;
    string param_id = 5;
}

message GetShouldStopWorkersReply {
    repeated string should_stop_worker_ids = 1;
}

message GetMetricsRequest {
    string study_id = 1;
    repeated string worker_ids = 2;
    repeated string metrics_names = 3;
}

message GetMetricsReply {
    repeated MetricsLogSet metrics_log_sets = 1;
}

message ReportMetricsLogsRequest {
    string study_id = 1;
    repeated MetricsLogSet metrics_log_sets = 3;
}

message ReportMetricsLogsReply {
}

message ModelInfo {
    string study_name = 1;  
    string worker_id = 2; 
    repeated Parameter parameters = 3; 
    repeated Metrics metrics = 4;
    string model_path = 5; 
}

message DataSetInfo {
    string name = 1;
    string path = 2;
}

message SaveStudyRequest {
    string study_name = 1;
    string owner = 2;
    string description = 3;
}

message SaveStudyReply {
}

message SaveModelRequest {
    ModelInfo model = 1;
    DataSetInfo data_set = 2;
    bool tensor_board = 3;
}

message SaveModelReply {
}

message GetSavedStudiesRequest {
}

message GetSavedStudiesReply {
    repeated StudyOverview studies = 1;
}

message GetSavedModelsRequest {
    string study_name = 1;
}

message GetSavedModelsReply {
    repeated ModelInfo models = 1;
}

message GetSavedModelRequest {
    string study_name = 1;
    string worker_id = 2;
}

message GetSavedModelReply {
    ModelInfo model = 1;
}

message SetSuggestionParametersRequest {
    string study_id = 1;
    string suggestion_algorithm = 2;
    string param_id = 3;
    repeated SuggestionParameter suggestion_parameters = 4;
}

message SetSuggestionParametersReply {
    string param_id = 1;
}

message GetSuggestionParametersRequest {
    string param_id = 1;
}

message GetSuggestionParametersReply {
    repeated SuggestionParameter suggestion_parameters = 1;
}

message GetSuggestionParameterListRequest {
    string study_id = 1;
}

message SuggestionParameterSet {
    string param_id = 1;
    string suggestion_algorithm = 2;
    repeated SuggestionParameter suggestion_parameters = 3;
}

message GetSuggestionParameterListReply {

    repeated SuggestionParameterSet suggestion_parameter_sets = 1;
}

message StopSuggestionRequest {
    string study_id = 1;
}

message StopSuggestionReply {
}

message SetEarlyStoppingParametersRequest {
    string study_id = 1;
    string early_stopping_algorithm = 2;
    string param_id = 3;
    repeated EarlyStoppingParameter early_stopping_parameters = 4;
}

message SetEarlyStoppingParametersReply {
    string param_id = 1;
}

message GetEarlyStoppingParametersRequest {
    string param_id = 1;
}

message GetEarlyStoppingParametersReply {
    repeated EarlyStoppingParameter early_stopping_parameters = 1;
}

message GetEarlyStoppingParameterListRequest {
    string study_id = 1;
}

message EarlyStoppingParameterSet {
    string param_id = 1;
    string early_stopping_algorithm = 2;
    repeated EarlyStoppingParameter early_stopping_parameters = 3;
}

message GetEarlyStoppingParameterListReply {
    repeated EarlyStoppingParameterSet early_stopping_parameter_sets = 1;
}