// Copyright 2025 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.hypercomputecluster.v1alpha;

import "google/api/annotations.proto";
import "google/api/client.proto";
import "google/api/field_behavior.proto";
import "google/api/resource.proto";
import "google/cloud/hypercomputecluster/v1alpha/operation_metadata.proto";
import "google/longrunning/operations.proto";
import "google/protobuf/duration.proto";
import "google/protobuf/empty.proto";
import "google/protobuf/field_mask.proto";
import "google/protobuf/timestamp.proto";

option csharp_namespace = "Google.Cloud.HypercomputeCluster.V1Alpha";
option go_package = "cloud.google.com/go/hypercomputecluster/apiv1alpha/hypercomputeclusterpb;hypercomputeclusterpb";
option java_multiple_files = true;
option java_outer_classname = "MachineLearningRunsProto";
option java_package = "com.google.cloud.hypercomputecluster.v1alpha";
option php_namespace = "Google\\Cloud\\HypercomputeCluster\\V1alpha";
option ruby_package = "Google::Cloud::HypercomputeCluster::V1alpha";

// Service for managing Machine Learning Runs.
service MachineLearningRuns {
  option (google.api.default_host) = "hypercomputecluster.googleapis.com";
  option (google.api.oauth_scopes) =
      "https://www.googleapis.com/auth/cloud-platform";

  // Creates a Machine Learning Run.
  rpc CreateMachineLearningRun(CreateMachineLearningRunRequest)
      returns (google.longrunning.Operation) {
    option (google.api.http) = {
      post: "/v1alpha/{parent=projects/*/locations/*}/machineLearningRuns"
      body: "machine_learning_run"
    };
    option (google.api.method_signature) =
        "parent,machine_learning_run,machine_learning_run_id";
    option (google.longrunning.operation_info) = {
      response_type: "MachineLearningRun"
      metadata_type: "OperationMetadata"
    };
  }

  // Updates an Machine Learning Run.
  rpc UpdateMachineLearningRun(UpdateMachineLearningRunRequest)
      returns (google.longrunning.Operation) {
    option (google.api.http) = {
      patch: "/v1alpha/{machine_learning_run.name=projects/*/locations/*/machineLearningRuns/*}"
      body: "machine_learning_run"
    };
    option (google.api.method_signature) = "machine_learning_run,update_mask";
    option (google.longrunning.operation_info) = {
      response_type: "MachineLearningRun"
      metadata_type: "OperationMetadata"
    };
  }

  // Deletes a Machine Learning Run.
  rpc DeleteMachineLearningRun(DeleteMachineLearningRunRequest)
      returns (google.longrunning.Operation) {
    option (google.api.http) = {
      delete: "/v1alpha/{name=projects/*/locations/*/machineLearningRuns/*}"
    };
    option (google.api.method_signature) = "name";
    option (google.longrunning.operation_info) = {
      response_type: "google.protobuf.Empty"
      metadata_type: "OperationMetadata"
    };
  }

  // Gets a Machine Learning Run.
  rpc GetMachineLearningRun(GetMachineLearningRunRequest)
      returns (MachineLearningRun) {
    option (google.api.http) = {
      get: "/v1alpha/{name=projects/*/locations/*/machineLearningRuns/*}"
    };
    option (google.api.method_signature) = "name";
  }

  // Lists Machine Learning Runs.
  rpc ListMachineLearningRuns(ListMachineLearningRunsRequest)
      returns (ListMachineLearningRunsResponse) {
    option (google.api.http) = {
      get: "/v1alpha/{parent=projects/*/locations/*}/machineLearningRuns"
    };
    option (google.api.method_signature) = "parent";
  }

  // Gets a Profile Session.
  rpc GetProfileSession(GetProfileSessionRequest) returns (ProfileSession) {
    option (google.api.http) = {
      get: "/v1alpha/{name=projects/*/locations/*/machineLearningRuns/*/profileSessions/*}"
    };
    option (google.api.method_signature) = "name";
  }

  // Lists Profile Sessions for a Machine Learning Run.
  rpc ListProfileSessions(ListProfileSessionsRequest)
      returns (ListProfileSessionsResponse) {
    option (google.api.http) = {
      get: "/v1alpha/{parent=projects/*/locations/*/machineLearningRuns/*}/profileSessions"
    };
    option (google.api.method_signature) = "parent";
  }
}

// State of the Machine Learning run.
enum State {
  // State is not specified.
  STATE_UNSPECIFIED = 0;

  // Run is active.
  STATE_ACTIVE = 1;

  // Run is completed.
  STATE_COMPLETED = 2;

  // Run is failed.
  STATE_FAILED = 3;
}

// A Machine Learning run.
message MachineLearningRun {
  option (google.api.resource) = {
    type: "hypercomputecluster.googleapis.com/MachineLearningRun"
    pattern: "projects/{project}/locations/{location}/machineLearningRuns/{machine_learning_run}"
    plural: "machineLearningRuns"
    singular: "machineLearningRun"
  };

  // Orchestrator for the ML run like Google Compute Engine, Google
  // Kubernetes Engine, Slurm etc.
  enum Orchestrator {
    // Orchestrator type is not specified.
    ORCHESTRATOR_UNSPECIFIED = 0;

    // Google Compute Engine orchestrator.
    GCE = 1;

    // Google Kubernetes Engine orchestrator.
    GKE = 2;

    // Slurm cluster orchestrator.
    SLURM = 3;
  }

  // RunPhase defines the phase of the Machine Learning run.
  enum RunPhase {
    // State is not specified.
    RUN_PHASE_UNSPECIFIED = 0;

    // Run is active.
    ACTIVE = 1;

    // Run is completed.
    COMPLETED = 2;

    // Run is failed.
    FAILED = 3;
  }

  // Identifier. The name of the Machine Learning run.
  string name = 1 [(google.api.field_behavior) = IDENTIFIER];

  // Optional. Display name for the run.
  string display_name = 2 [(google.api.field_behavior) = OPTIONAL];

  // Required. Allows grouping of similar runs.
  //  * Helps improving UI rendering performance.
  //  * Allows comparing similar runs via fast filters.
  string run_set = 3 [(google.api.field_behavior) = REQUIRED];

  // Optional. ETag for the run.
  // It must be provided for update/delete operations and must match the
  // server's etag.
  string etag = 4 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Tracks configuration for this run,
  // example: batch_size, jax_version, tpu_generation etc.
  Configs configs = 5 [(google.api.field_behavior) = OPTIONAL];

  // Required. List of tools enabled for this run
  // example: XProf, NSys
  repeated Tool tools = 6 [(google.api.field_behavior) = REQUIRED];

  // Optional. Metrics for the run.
  Metrics metrics = 7 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Any custom labels for this run
  // Example: type:workload, type:simulation etc.
  map<string, string> labels = 8 [(google.api.field_behavior) = OPTIONAL];

  // Output only. State of the run.
  State state = 9 [(google.api.field_behavior) = OUTPUT_ONLY];

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

  // Output only. Time when the run was last updated.
  google.protobuf.Timestamp update_time = 11
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Optional. Artifacts for the run.
  Artifacts artifacts = 12 [(google.api.field_behavior) = OPTIONAL];

  // Required. The orchestrator used for the run.
  Orchestrator orchestrator = 13 [(google.api.field_behavior) = REQUIRED];

  // Optional. The metadata for the workload associated with the run.
  WorkloadDetails workload_details = 14
      [(google.api.field_behavior) = OPTIONAL];

  // Optional. RunPhase defines the phase of the run.
  RunPhase run_phase = 15 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Error details for the run.
  // This field is set when the run is failed.
  string error_details = 16 [(google.api.field_behavior) = OPTIONAL];

  // Output only. Time when the run was completed.
  // This field is set when the run is completed or failed.
  google.protobuf.Timestamp end_time = 17
      [(google.api.field_behavior) = OUTPUT_ONLY];
}

// Request message for MachineLearningRuns.CreateMachineLearningRun.
message CreateMachineLearningRunRequest {
  // Required. Parent format: projects/{project}/locations/{location}
  string parent = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      child_type: "hypercomputecluster.googleapis.com/MachineLearningRun"
    }
  ];

  // Required. The Machine Learning Run to create.
  MachineLearningRun machine_learning_run = 2
      [(google.api.field_behavior) = REQUIRED];

  // Optional. The ID of the Machine Learning Run to create.
  string machine_learning_run_id = 3 [(google.api.field_behavior) = OPTIONAL];
}

// Request message for MachineLearningRuns.UpdateMachineLearningRun.
message UpdateMachineLearningRunRequest {
  // Required. The Machine Learning Run to update.
  MachineLearningRun machine_learning_run = 1
      [(google.api.field_behavior) = REQUIRED];

  // Optional. Used to identify fields to be updated. Allowed fields is *, we
  // don't support partial updates.
  google.protobuf.FieldMask update_mask = 2
      [(google.api.field_behavior) = OPTIONAL];
}

// Request message for MachineLearningRuns.DeleteMachineLearningRun.
message DeleteMachineLearningRunRequest {
  // Required.
  // projects/{project}/locations/{location}/machineLearningRuns/{machineLearningRun}
  string name = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      type: "hypercomputecluster.googleapis.com/MachineLearningRun"
    }
  ];

  // Optional. Etag for the run. Etag gets updated with every update/create
  // operation. If provided, it must match the server's etag.
  string etag = 2 [(google.api.field_behavior) = OPTIONAL];
}

// Request message for MachineLearningRuns.GetMachineLearningRun.
message GetMachineLearningRunRequest {
  // Required.
  // projects/{project}/locations/{location}/machineLearningRuns/{machineLearningRun}
  string name = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      type: "hypercomputecluster.googleapis.com/MachineLearningRun"
    }
  ];
}

// Request message for MachineLearningRuns.ListMachineLearningRuns.
message ListMachineLearningRunsRequest {
  // Required. projects/{project}/locations/{location}
  string parent = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      child_type: "hypercomputecluster.googleapis.com/MachineLearningRun"
    }
  ];

  // Optional. The maximum number of Machine Learning Runs to return. The
  // service may return fewer than this value. If unspecified, at most 10
  // Machine Learning Runs will be returned. The maximum value is 1000; values
  // above 1000 will be coerced to 1000. Please refer to
  // https://google.aip.dev/158 for more details.
  int32 page_size = 2 [(google.api.field_behavior) = OPTIONAL];

  // Optional. A page token, received from a previous `ListMachineLearningRuns`
  // call. Provide this to retrieve the subsequent page. Please refer to
  // https://google.aip.dev/158 for more details.
  string page_token = 3 [(google.api.field_behavior) = OPTIONAL];

  // Optional. An expression for filtering the results of the request.
  //
  // Supported filter fields are:
  // - run_id (string): The ID of the run (e.g., "run-1").
  // - run_set (string): The run set (e.g., "run-set-1").
  // - display_name (string): The display name of the run (e.g.,
  // "display-name-1").
  // - state (enum): The state of the run (e.g., "STATE_COMPLETED").
  // - region (string): The region of the run (e.g., "us-central1").
  // - orchestrator (enum): The orchestrator used for the run (e.g.,
  // "ORCHESTRATOR_GKE").
  // - create_time (timestamp): The creation timestamp of the run.
  // - update_time (timestamp): The last update timestamp of the run.
  // - etag (string): The ETag of the run.
  //
  // Filter expressions can be combined using "AND" and "OR" operators.
  //
  // Examples:
  //
  // To filter by a single field:
  //
  //   `run_set = "run-set-1"`
  //   `state = "STATE_ACTIVE"`
  //
  // To filter by multiple fields:
  //
  //   `run_set = "run-set-1" AND state = "STATE_ACTIVE"`
  //   `region = "us-central1" AND orchestrator = "ORCHESTRATOR_GKE"`
  //
  // To filter by timestamp:
  //
  //   `create_time > "2025-05-26T10:00:00Z"`
  string filter = 4 [(google.api.field_behavior) = OPTIONAL];

  // Optional. A comma-separated list of fields to order by, sorted in ascending
  // order. Use "desc" after a field name for descending.
  //
  // Supported fields:
  //
  //   * `run_id`
  //   * `run_set`
  //   * `display_name`
  //   * `region`
  //   * `orchestrator`
  //   * `state`
  //   * `create_time`
  //   * `update_time`
  //
  // Example: `update_time desc, create_time`.
  string order_by = 5 [(google.api.field_behavior) = OPTIONAL];
}

// Response message for MachineLearningRuns.ListMachineLearningRuns.
message ListMachineLearningRunsResponse {
  // The Machine Learning Runs.
  repeated MachineLearningRun machine_learning_runs = 1;

  // A token, which can be sent as `page_token` to retrieve the next page.
  // If this field is omitted, there are no subsequent pages.
  string next_page_token = 2;
}

// Request message for MachineLearningRuns.ListProfileSessions.
message ListProfileSessionsRequest {
  // Required. Parent format:
  // projects/{project}/locations/{location}/machineLearningRuns/{machineLearningRun}
  string parent = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      child_type: "hypercomputecluster.googleapis.com/ProfileSession"
    }
  ];

  // Optional. The maximum number of Profile Sessions to return. The service may
  // return fewer than this value.
  // If unspecified, at most 50 Profile Sessions will be returned.
  // The maximum value is 1000; values above 1000 will be coerced to 1000.
  int32 page_size = 2 [(google.api.field_behavior) = OPTIONAL];

  // Optional. A page token, received from a previous `ListProfileSessions`
  // call. Provide this to retrieve the subsequent page.
  string page_token = 3 [(google.api.field_behavior) = OPTIONAL];
}

// Response message for MachineLearningRuns.ListProfileSessions.
message ListProfileSessionsResponse {
  // A list of profile sessions.
  repeated ProfileSession profile_sessions = 1;

  // A token, which can be sent as `page_token` to retrieve the next page.
  // If this field is omitted, there are no subsequent pages.
  string next_page_token = 2;
}

// Request message for MachineLearningRuns.GetProfileSession.
message GetProfileSessionRequest {
  // Required.
  // projects/{project}/locations/{location}/machineLearningRuns/{machineLearningRun}/profileSessions/{profileSession}
  string name = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      type: "hypercomputecluster.googleapis.com/ProfileSession"
    }
  ];
}

// Represents a single profiling session.
message ProfileSession {
  option (google.api.resource) = {
    type: "hypercomputecluster.googleapis.com/ProfileSession"
    pattern: "projects/{project}/locations/{location}/machineLearningRuns/{machine_learning_run}/profileSessions/{profile_session}"
    plural: "profileSessions"
    singular: "profileSession"
  };

  // Identifier. The name of the profile session.
  // Format:
  // projects/{project}/locations/{location}/machineLearningRuns/{machine_learning_run}/profileSessions/{profile_session}
  string name = 1 [(google.api.field_behavior) = IDENTIFIER];

  // Output only. The cloud storage path of the session. Example:
  // `gs://my-bucket/my-run-directory/session-1`.
  string gcs_path = 2 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. The creation time of the session.
  google.protobuf.Timestamp create_time = 3
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. The URL to dashboard to see session specific data. Not
  // specified if URL is not ready yet. Currently represents only Xprof.
  //
  // Form https://<backend>?<params(gsc_bucket|etc)>
  // Could contains user information like bucket name etc.
  string dashboard_url = 4 [(google.api.field_behavior) = OUTPUT_ONLY];
}

// Configuration for a Machine Learning run.
message Configs {
  // Optional. User defined configs.
  map<string, string> user_configs = 1 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Software configs.
  map<string, string> software_configs = 2
      [(google.api.field_behavior) = OPTIONAL];

  // Optional. Hardware configs.
  map<string, string> hardware_configs = 3
      [(google.api.field_behavior) = OPTIONAL];
}

// Metrics for a Machine Learning run.
message Metrics {
  // Output only. Average step time.
  google.protobuf.Duration avg_step = 1
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Average Model Flops Utilization (MFU).
  float avg_mfu = 2 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Average throughput. Value depends on the kind of workload ran
  // on TPU.
  float avg_throughput = 3 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Average latency.
  google.protobuf.Duration avg_latency = 4
      [(google.api.field_behavior) = OUTPUT_ONLY];
}

// Represents information about the artifacts of the Machine Learning Run.
message Artifacts {
  // Required. The Cloud Storage path where the artifacts of the run are stored.
  // Example: `gs://my-bucket/my-run-directory`.
  string gcs_path = 1 [(google.api.field_behavior) = REQUIRED];
}

// A tool for a Machine Learning run.
message Tool {
  // tool - one of the diagnostic tool.
  oneof tool {
    // Required. XProf related metadata
    Xprof xprof = 1 [(google.api.field_behavior) = REQUIRED];
  }
}

// Workload details associated for the Machine Learning Run.
// Workload have different metadata based on the orchestrator like
// GKE cluster, Slurm cluster, Google Compute Engine instance etc.
message WorkloadDetails {
  // The metadata for the workload for different orchestrators.
  oneof metadata {
    // GKE Workload metadata.
    GKEWorkloadDetails gke = 1;
  }
}

// Workload details for the GKE orchestrator.
message GKEWorkloadDetails {
  // Required. The identifier of the workload. Example - jobset-abcd
  string id = 1 [(google.api.field_behavior) = REQUIRED];

  // Required. The kind of the workload. Example - JobSet
  string kind = 2 [(google.api.field_behavior) = REQUIRED];

  // Required. The cluster of the workload.
  // Example - /projects/<project id>/locations/<location>/clusters/<cluster
  // name>
  string cluster = 3 [(google.api.field_behavior) = REQUIRED];

  // Required. The namespace of the workload. Example - default
  string namespace = 4 [(google.api.field_behavior) = REQUIRED];

  // Optional. The parent workload. Example - JobSet for a Job, RayJob,
  // LeaderWorkerSet etc. If a workload is a part of a RunSet, this field is
  // populate as `RunSet` and if it was ray job, it will be set to `RayJob` etc.
  // This will be an additional field to representing a workload.
  string parent_workload = 5 [(google.api.field_behavior) = OPTIONAL];

  // Optional. labels for the workload.
  // Example: {"type": "workload", "app": "simulation"}.
  map<string, string> labels = 6 [(google.api.field_behavior) = OPTIONAL];
}

// XProf related metadata
message Xprof {
  // Required. XProf session id
  string session_id = 1 [(google.api.field_behavior) = REQUIRED];
}
