// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

syntax = "proto3";

package google.cloud.aiplatform.v1beta1;

import "google/api/field_behavior.proto";
import "google/api/resource.proto";
import "google/cloud/aiplatform/v1beta1/custom_job.proto";
import "google/cloud/aiplatform/v1beta1/encryption_spec.proto";
import "google/cloud/aiplatform/v1beta1/job_state.proto";
import "google/cloud/aiplatform/v1beta1/study.proto";
import "google/protobuf/timestamp.proto";
import "google/rpc/status.proto";

option csharp_namespace = "Google.Cloud.AIPlatform.V1Beta1";
option go_package = "cloud.google.com/go/aiplatform/apiv1beta1/aiplatformpb;aiplatformpb";
option java_multiple_files = true;
option java_outer_classname = "NasJobProto";
option java_package = "com.google.cloud.aiplatform.v1beta1";
option php_namespace = "Google\\Cloud\\AIPlatform\\V1beta1";
option ruby_package = "Google::Cloud::AIPlatform::V1beta1";

// Represents a Neural Architecture Search (NAS) job.
message NasJob {
  option (google.api.resource) = {
    type: "aiplatform.googleapis.com/NasJob"
    pattern: "projects/{project}/locations/{location}/nasJobs/{nas_job}"
  };

  // Output only. Resource name of the NasJob.
  string name = 1 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Required. The display name of the NasJob.
  // The name can be up to 128 characters long and can consist of any UTF-8
  // characters.
  string display_name = 2 [(google.api.field_behavior) = REQUIRED];

  // Required. The specification of a NasJob.
  NasJobSpec nas_job_spec = 4 [(google.api.field_behavior) = REQUIRED];

  // Output only. Output of the NasJob.
  NasJobOutput nas_job_output = 5 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. The detailed state of the job.
  JobState state = 6 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Time when the NasJob was created.
  google.protobuf.Timestamp create_time = 7
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Time when the NasJob for the first time entered the
  // `JOB_STATE_RUNNING` state.
  google.protobuf.Timestamp start_time = 8
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Time when the NasJob entered any of the following states:
  // `JOB_STATE_SUCCEEDED`, `JOB_STATE_FAILED`, `JOB_STATE_CANCELLED`.
  google.protobuf.Timestamp end_time = 9
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Time when the NasJob was most recently updated.
  google.protobuf.Timestamp update_time = 10
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Only populated when job's state is JOB_STATE_FAILED or
  // JOB_STATE_CANCELLED.
  google.rpc.Status error = 11 [(google.api.field_behavior) = OUTPUT_ONLY];

  // The labels with user-defined metadata to organize NasJobs.
  //
  // Label keys and values can be no longer than 64 characters
  // (Unicode codepoints), can only contain lowercase letters, numeric
  // characters, underscores and dashes. International characters are allowed.
  //
  // See https://goo.gl/xmQnxf for more information and examples of labels.
  map<string, string> labels = 12;

  // Customer-managed encryption key options for a NasJob.
  // If this is set, then all resources created by the NasJob
  // will be encrypted with the provided encryption key.
  EncryptionSpec encryption_spec = 13;

  // Optional. Enable a separation of Custom model training
  // and restricted image training for tenant project.
  bool enable_restricted_image_training = 14
      [deprecated = true, (google.api.field_behavior) = OPTIONAL];

  // Output only. Reserved for future use.
  bool satisfies_pzs = 15 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Reserved for future use.
  bool satisfies_pzi = 16 [(google.api.field_behavior) = OUTPUT_ONLY];
}

// Represents a NasTrial details along with its parameters. If there is a
// corresponding train NasTrial, the train NasTrial is also returned.
message NasTrialDetail {
  option (google.api.resource) = {
    type: "aiplatform.googleapis.com/NasTrialDetail"
    pattern: "projects/{project}/locations/{location}/nasJobs/{nas_job}/nasTrialDetails/{nas_trial_detail}"
  };

  // Output only. Resource name of the NasTrialDetail.
  string name = 1 [(google.api.field_behavior) = OUTPUT_ONLY];

  // The parameters for the NasJob NasTrial.
  string parameters = 2;

  // The requested search NasTrial.
  NasTrial search_trial = 3;

  // The train NasTrial corresponding to
  // [search_trial][google.cloud.aiplatform.v1beta1.NasTrialDetail.search_trial].
  // Only populated if
  // [search_trial][google.cloud.aiplatform.v1beta1.NasTrialDetail.search_trial]
  // is used for training.
  NasTrial train_trial = 4;
}

// Represents the spec of a NasJob.
message NasJobSpec {
  // The spec of multi-trial Neural Architecture Search (NAS).
  message MultiTrialAlgorithmSpec {
    // Represents a metric to optimize.
    message MetricSpec {
      // The available types of optimization goals.
      enum GoalType {
        // Goal Type will default to maximize.
        GOAL_TYPE_UNSPECIFIED = 0;

        // Maximize the goal metric.
        MAXIMIZE = 1;

        // Minimize the goal metric.
        MINIMIZE = 2;
      }

      // Required. The ID of the metric. Must not contain whitespaces.
      string metric_id = 1 [(google.api.field_behavior) = REQUIRED];

      // Required. The optimization goal of the metric.
      GoalType goal = 2 [(google.api.field_behavior) = REQUIRED];
    }

    // Represent spec for search trials.
    message SearchTrialSpec {
      // Required. The spec of a search trial job. The same spec applies to
      // all search trials.
      CustomJobSpec search_trial_job_spec = 1
          [(google.api.field_behavior) = REQUIRED];

      // Required. The maximum number of Neural Architecture Search (NAS) trials
      // to run.
      int32 max_trial_count = 2 [(google.api.field_behavior) = REQUIRED];

      // Required. The maximum number of trials to run in parallel.
      int32 max_parallel_trial_count = 3
          [(google.api.field_behavior) = REQUIRED];

      // The number of failed trials that need to be seen before failing
      // the NasJob.
      //
      // If set to 0, Vertex AI decides how many trials must fail
      // before the whole job fails.
      int32 max_failed_trial_count = 4;
    }

    // Represent spec for train trials.
    message TrainTrialSpec {
      // Required. The spec of a train trial job. The same spec applies to
      // all train trials.
      CustomJobSpec train_trial_job_spec = 1
          [(google.api.field_behavior) = REQUIRED];

      // Required. The maximum number of trials to run in parallel.
      int32 max_parallel_trial_count = 2
          [(google.api.field_behavior) = REQUIRED];

      // Required. Frequency of search trials to start train stage. Top N
      // [TrainTrialSpec.max_parallel_trial_count]
      // search trials will be trained for every M
      // [TrainTrialSpec.frequency] trials searched.
      int32 frequency = 3 [(google.api.field_behavior) = REQUIRED];
    }

    // The available types of multi-trial algorithms.
    enum MultiTrialAlgorithm {
      // Defaults to `REINFORCEMENT_LEARNING`.
      MULTI_TRIAL_ALGORITHM_UNSPECIFIED = 0;

      // The Reinforcement Learning Algorithm for Multi-trial Neural
      // Architecture Search (NAS).
      REINFORCEMENT_LEARNING = 1;

      // The Grid Search Algorithm for Multi-trial Neural
      // Architecture Search (NAS).
      GRID_SEARCH = 2;
    }

    // The multi-trial Neural Architecture Search (NAS) algorithm
    // type. Defaults to `REINFORCEMENT_LEARNING`.
    MultiTrialAlgorithm multi_trial_algorithm = 1;

    // Metric specs for the NAS job.
    // Validation for this field is done at `multi_trial_algorithm_spec` field.
    MetricSpec metric = 2;

    // Required. Spec for search trials.
    SearchTrialSpec search_trial_spec = 3
        [(google.api.field_behavior) = REQUIRED];

    // Spec for train trials. Top N [TrainTrialSpec.max_parallel_trial_count]
    // search trials will be trained for every M
    // [TrainTrialSpec.frequency] trials searched.
    TrainTrialSpec train_trial_spec = 4;
  }

  // The Neural Architecture Search (NAS) algorithm specification.
  oneof nas_algorithm_spec {
    // The spec of multi-trial algorithms.
    MultiTrialAlgorithmSpec multi_trial_algorithm_spec = 2;
  }

  // The ID of the existing NasJob in the same Project and Location
  // which will be used to resume search. search_space_spec and
  // nas_algorithm_spec are obtained from previous NasJob hence should not
  // provide them again for this NasJob.
  string resume_nas_job_id = 3;

  // It defines the search space for Neural Architecture Search (NAS).
  string search_space_spec = 1;
}

// Represents a uCAIP NasJob output.
message NasJobOutput {
  // The output of a multi-trial Neural Architecture Search (NAS) jobs.
  message MultiTrialJobOutput {
    // Output only. List of NasTrials that were started as part of search stage.
    repeated NasTrial search_trials = 1
        [(google.api.field_behavior) = OUTPUT_ONLY];

    // Output only. List of NasTrials that were started as part of train stage.
    repeated NasTrial train_trials = 2
        [(google.api.field_behavior) = OUTPUT_ONLY];
  }

  // The output of this Neural Architecture Search (NAS) job.
  oneof output {
    // Output only. The output of this multi-trial Neural Architecture Search
    // (NAS) job.
    MultiTrialJobOutput multi_trial_job_output = 1
        [(google.api.field_behavior) = OUTPUT_ONLY];
  }
}

// Represents a uCAIP NasJob trial.
message NasTrial {
  // Describes a NasTrial state.
  enum State {
    // The NasTrial state is unspecified.
    STATE_UNSPECIFIED = 0;

    // Indicates that a specific NasTrial has been requested, but it has not yet
    // been suggested by the service.
    REQUESTED = 1;

    // Indicates that the NasTrial has been suggested.
    ACTIVE = 2;

    // Indicates that the NasTrial should stop according to the service.
    STOPPING = 3;

    // Indicates that the NasTrial is completed successfully.
    SUCCEEDED = 4;

    // Indicates that the NasTrial should not be attempted again.
    // The service will set a NasTrial to INFEASIBLE when it's done but missing
    // the final_measurement.
    INFEASIBLE = 5;
  }

  // Output only. The identifier of the NasTrial assigned by the service.
  string id = 1 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. The detailed state of the NasTrial.
  State state = 2 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. The final measurement containing the objective value.
  Measurement final_measurement = 3 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Time when the NasTrial was started.
  google.protobuf.Timestamp start_time = 4
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Time when the NasTrial's status changed to `SUCCEEDED` or
  // `INFEASIBLE`.
  google.protobuf.Timestamp end_time = 5
      [(google.api.field_behavior) = OUTPUT_ONLY];
}