/**
 * Katib API
 */
syntax = "proto3";

package api;

import "google/api/annotations.proto";

/**
 * Service for Main API for Katib
 * 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 {
    /**
     * Create a Study from Study Config.
     * Generate a unique ID and store the Study to DB.
     */
    rpc CreateStudy(CreateStudyRequest) returns (CreateStudyReply){
        option (google.api.http) = {
            post: "/api/Manager/CreateStudy"
            body: "study_config"
        };
    };
    /** 
     * Get a Study Config from DB by ID of Study.
     */
    rpc GetStudy(GetStudyRequest) returns (GetStudyReply){
        option (google.api.http) = {
            get: "/api/Manager/GetStudy/{study_id}"
        };
    };
    /** 
     * Delete a Study from DB by Study ID.
     */
    rpc DeleteStudy(DeleteStudyRequest) returns (DeleteStudyReply){
        option (google.api.http) = {
            get: "/api/Manager/DeleteStudy/{study_id}"
        };
    };
    /**
     * Get all Study Configs from DB.
     */
    rpc GetStudyList(GetStudyListRequest) returns (GetStudyListReply){
        option (google.api.http) = {
            get: "/api/Manager/GetStudyList"
        };
    };
    /**
     * Create a Trial from Trial Config.
     * Generate a unique ID and store the Trial to DB.
     */
    rpc CreateTrial(CreateTrialRequest) returns (CreateTrialReply){
        option (google.api.http) = {
            post: "/api/Manager/CreateTrial"
            body: "trial"
        };
    };
    /** 
     * Get a Trial Configs from DB by ID of Study.
     */
    rpc GetTrials(GetTrialsRequest) returns (GetTrialsReply){
        option (google.api.http) = {
            get: "/api/Manager/GetTrials/{study_id}"
        };
    };
    /**
     * Get a Trial Configuration from DB by ID of Trial.
     */
    rpc GetTrial(GetTrialRequest) returns (GetTrialReply){
        option (google.api.http) = {
            get: "/api/Manager/GetTrial/{trial_id}"
        };
    };
    /**
     * Create a Worker from Worker Config.
     * Generate a unique ID and store the Worker to DB.
     */
    rpc RegisterWorker(RegisterWorkerRequest) returns (RegisterWorkerReply){
        option (google.api.http) = {
            post: "/api/Manager/RegisterWorker"
            body: "worker"
        };
    };
    /** 
     * Get a Worker Configs and Status from DB by ID of Study, Trial or Worker.
     */
    rpc GetWorkers(GetWorkersRequest) returns (GetWorkersReply){
        option (google.api.http) = {
            get: "/api/Manager/GetWorkers"
        };
    };
    /** 
     * Update a Status of Worker.
     */
    rpc UpdateWorkerState(UpdateWorkerStateRequest) returns (UpdateWorkerStateReply){
        option (google.api.http) = {
            put: "/api/Manager/UpdateWorkerState",
            body: "*"
        };
    };
    /** 
     * Get full information related to specified Workers.
     * It includes Worker Config, HyperParameters and Metrics Logs.
     */
    rpc GetWorkerFullInfo(GetWorkerFullInfoRequest) returns (GetWorkerFullInfoReply){
        option (google.api.http) = {
            get: "/api/Manager/GetWorkerFullInfo"
        };
    };
    /** 
     * Get Suggestions from a Suggestion service.
     */
    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: "*"
        };
    };
    /**
     * Get metrics of workers.
     * You can get all logs of metrics since start of the worker.
     */
    rpc GetMetrics(GetMetricsRequest) returns (GetMetricsReply){
        option (google.api.http) = {
            post: "/api/Manager/GetMetrics"
            body: "*"
        };
    };
    /**
     * Create or Update parameter set for a suggestion service.
     * If you specify an ID of parameter set, it will update the parameter set by your request.
     * If you don't specify an ID, it will create a new parameter set for corresponding study and suggestion service.
     * The parameters are key-value format.
     */
    rpc SetSuggestionParameters(SetSuggestionParametersRequest) returns (SetSuggestionParametersReply){
        option (google.api.http) = {
            post: "/api/Manager/SetSuggestionParameters"
            body: "*"
        };
    };
    /**
     * Get suggestion parameter set from DB specified.
     */
    rpc GetSuggestionParameters(GetSuggestionParametersRequest) returns (GetSuggestionParametersReply){
        option (google.api.http) = {
            get: "/api/Manager/GetSuggestionParameters/{param_id}"
        };
    };
    /**
     * Get all suggestion parameter sets from DB.
     */
    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: "*"
        };
    };
    /**
     * Report a logs of metrics for workers.
     * The logs for each worker must have timestamp and must be ordered in time series.
     * When the log you reported are already reported before, it will be dismissed and get no error.
     */
    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);
}

/**
 * Types of value for HyperParameter.
 */
enum ParameterType {
    UNKNOWN_TYPE = 0; /// Undefined type and not used.

    DOUBLE = 1; /// Double float type. Use "Max/Min".
    INT = 2; /// Int type. Use "Max/Min".
    DISCRETE = 3; /// Discrete number type. Use "List" as float.
    CATEGORICAL = 4; /// Categorical type. Use "List" as string.
}

/**
 * Direction of optimization. Minimize or Maximize.
 */
enum OptimizationType {
    UNKNOWN_OPTIMIZATION = 0; /// Undefined type and not used.

    MINIMIZE = 1; /// Minimize
    MAXIMIZE = 2; /// Maximize
}

/**
 * Feasible space for optimization.
 * Int and Double type use Max/Min.
 * Discrete and Categorical type use List.
 */
message FeasibleSpace {
    string max = 1; /// Max Value
    string min = 2; /// Minimum Value
    repeated string list = 3; /// List of Values.
}

/**
 * Config for a Hyper parameter.
 * Katib will create each Hyper parameter from this config.
 */
message ParameterConfig {
    string name = 1; /// Name of the parameter.
    ParameterType parameter_type = 2; /// Type of the parameter.
    FeasibleSpace feasible = 3; /// FeasibleSpace for the parameter.
}

/**
 * Value of a Hyper parameter.
 * This will be created from a correcponding Config.
 */
message Parameter {
    string name = 1; /// Name of the parameter.
    ParameterType parameter_type = 2; /// Type of the parameter.
    string value = 3; /// Value of the parameter.
}

/**
 * Status code for worker.
 * This value is stored as TINYINT in MySQL.
 */
enum State {
    PENDING = 0; /// Pending. Created but not running.
    RUNNING = 1; /// Running. 
    COMPLETED = 2; /// Completed.
    KILLED = 3; /// Killed. Not failed.
    ERROR = 120; /// Error.
}

/**
 * Logs of metrics for a worker.
 */
message MetricsLogSet {
    string worker_id = 1; /// ID of the corresponding worker.
    repeated MetricsLog metrics_logs= 2; /// Logs of metrics.
    State worker_status = 3; /// Status of the corresponding worker.
}

/**
 * Metrics of a worker
 */
message Metrics {
    string name = 1; /// Name of metrics.
    string value = 2; /// Value of metrics. Double float.
}

/**
 * Metrics of a worker with timestamp
 */
message MetricsValueTime {
    string time = 1; /// Timestamp RFC3339 format.
    string value = 2; /// Value of metrics. Double float.
}

/**
 * Metrics logs of a worker
 */
message MetricsLog {
    string name = 1;  /// Name of metrics.
    repeated MetricsValueTime values = 2; /// Log of metrics. Ordered by time series.
}

/**
 * Parameter for Suggestion service. Key-value format.
 */
message SuggestionParameter {
    string name = 1; /// Name of Parameter.
    string value = 2; /// Value of Parameter.
}

/**
 * Parameter for EarlyStopping service. Key-value format.
 */
message EarlyStoppingParameter {
    string name = 1; /// Name of Parameter.
    string value = 2; /// Value of Parameter.
}

/**
 * Tag for each resource.
 */
message Tag {
    string name = 1; /// Name of tag.
    string value = 2; /// Value of tag.
}

/**
 * Overview of a study. For UI.
 */
message StudyOverview {
    string name = 1; /// Name of Study.
    string owner  = 2; /// Owner of Study.
    string id = 3; /// Study ID.
    string description = 4; /// Discretption of Study.
}

/**
 * A set of Hyperparameter.
 * In a study, multiple trials are evaluated by workers.
 * Suggestion service will generate next trials.
 */
message Trial {
    string trial_id = 1; /// Trial ID.
    string study_id = 2; /// Study ID.
    repeated Parameter parameter_set = 3; /// Hyperparameter set
    string objective_value = 4; /// Objective Value
    repeated Tag tags = 5; /// Tags of Trial.
}

/** 
 * A process of evaluation for a trial.
 * Types of worker supported by Katib are k8s Job, TF-Job, and Pytorch-Job.
 */
message Worker {
    string worker_id = 1; /// Worker ID.
    string study_id = 2; /// Study ID.
    string trial_id = 3; /// Trial ID.
    string Type = 4; /// Type of Worker
    State status = 5; /// Status of Worker.
    string TemplatePath = 6; /// Path for the manufest template of Worker.
    repeated Tag tags = 7; /// Tags of Worker.
}

/**
 * Config of a Study. Study represents a single optimization run over a feasible space. 
 * Each Study contains a configuration describing the feasible space, as well as a set of Trials.
 * It is assumed that objective function f(x) does not change in the course of a Study.
 */
message StudyConfig {
    /** 
     * List of ParameterConfig
     */
    message ParameterConfigs {
        repeated ParameterConfig configs = 1;
    }
    string name = 1; /// Name of Study.
    string owner = 2; /// Owner of Study.
    OptimizationType optimization_type = 3; /// Optimization type.
    double optimization_goal = 4; /// Goal of optimization value.
    ParameterConfigs parameter_configs = 5; /// List of ParameterConfig
    repeated string access_permissions = 6; /// Access Permission
    repeated Tag tags = 7; /// Tag for Study
    string objective_value_name = 8; /// Name of objective value.
    repeated string metrics = 9; /// List of metrics name.
    string jobId = 10; /// ID of studyjob that is created from this config.
}

/** 
 * Create a Study from Study Config.
 * Generate an unique ID and store the Study to DB.
 */
message CreateStudyRequest {
    StudyConfig study_config = 1;
}

/**
 * Return generated StudyID.
 */
message CreateStudyReply {
    string study_id = 1;
}

/** 
 * Delete a Study from DB by Study ID.
 */
message DeleteStudyRequest {
    string study_id = 1;
}

/**
 * Return deleted Study ID.
 */
message DeleteStudyReply {
    string study_id = 1;
}

/**
 * Get a Study Config from DB by ID of Study.
 */
message GetStudyRequest {
    string study_id = 1;
}

/**
 * Return a config of specified Study.
 */
message GetStudyReply {
    StudyConfig study_config= 1;
}

/** 
 * Get all Study Configs from DB.
 */
message GetStudyListRequest {
}

/** 
 * Return a overview list of Studies.
 */
message GetStudyListReply {
    repeated StudyOverview study_overviews = 1;
}

/**
 * Create a Trial from Trial Config.
 * Generate an unique ID and store the Trial to DB.
 */
message CreateTrialRequest {
    Trial trial = 1;
}

/**
 * Return generated TrialID.
 */
message CreateTrialReply {
    string trial_id = 1;
}

/** 
 * Get a Trial Configs from DB by ID of Study.
 */
message GetTrialsRequest {
    string study_id = 1;
}

/**
 * Return a trial list in specified Study.
 */
message GetTrialsReply {
    repeated Trial trials = 1;
}

/**
 * Get a trial configuration from DB by trial ID
 */
message GetTrialRequest {
    string trial_id = 1;
}

/**
 * Return a trial configuration by specified trial ID
 */
message GetTrialReply {
    Trial trial = 1;
}

/**
 * Create a Worker from Worker Config.
 * Generate an unique ID and store the Worker to DB.
 */
message RegisterWorkerRequest {
    Worker worker = 1;
}

/**
 * Return generated WorkerID.
 */
message RegisterWorkerReply {
    string worker_id = 1;
}

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

message StopWorkersReply {
}

/** 
 * Get a configs and status of a Worker from DB by ID of Study, Trial or Worker.
 */
message GetWorkersRequest {
    string study_id = 1;
    string trial_id = 2;
    string worker_id = 3;
}

/**
 * Return a Worker list by specified condition.
 */
message GetWorkersReply {
    repeated Worker workers = 1;
}

/** 
 * Update a Status of Worker.
 */
message UpdateWorkerStateRequest {
    string worker_id = 1;
    State status = 2;
}

message UpdateWorkerStateReply{
}

/** 
 * Get a full information related to specified Workers.
 * It includes Worker Config, HyperParameters and Metrics Logs.
 */
message GetWorkerFullInfoRequest {
    string study_id = 1;
    string trial_id = 2;
    string worker_id = 3;
    bool only_latest_log = 4;
}

message WorkerFullInfo {
    Worker Worker = 1;
    repeated Parameter parameter_set = 2;
    repeated MetricsLog metrics_logs= 3;
}

message GetWorkerFullInfoReply {
    repeated WorkerFullInfo worker_full_infos =1;
}

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;
}