refactor(api): remove Drain RPC, consolidate proto files, trim EngineError

README.md

@@ -109,7 +109,7 @@ The canonical schema is organized by domain under
 | Structured output     | JSON Schema, JSON object, regex, EBNF grammar, structural tags, and fixed choices                                 |
 | Token information     | Prompt and output logprobs, ranks, candidate-token selection, per-token records, and streamed text deltas         |
 | Discovery             | Server identity, deployment capacity, model limits, topology, parsers, and inference capabilities                 |
-| Lifecycle             | Health checks, targeted or global abort, graceful drain, progress, and terminal failures                          |
+| Lifecycle             | Health checks and targeted or global abort                                                                        |
 | Disaggregated serving | Prefill/decode roles, decode-context parallel topology, KV handoff, connector discovery, and cache controls       |
 | KV-aware routing      | Typed KV event streams plus discovery of engine-native event sources                                              |
 | Model extensions      | Multimodal inputs and LoRA adapter lifecycle                                                                      |
@@ -169,7 +169,7 @@ depend on it. Expect direct schema refinement during this phase.
 
 The intended adoption path is incremental:
 
-1. Aggregated generation and non-generative tasks, discovery, health, abort, and drain.
+1. Aggregated generation, discovery, health, and abort.
 2. Prefill/decode roles, KV handoff, rank affinity, and KV event integration.
 3. Logprobs, guided decoding, LoRA, and multimodal input as needed.
 

docs/api.md

@@ -34,7 +34,6 @@ service Control {
   // Health and lifecycle.
   rpc Health(HealthRequest) returns (HealthResponse);
   rpc Abort(AbortRequest) returns (AbortResponse);
-  rpc Drain(DrainRequest) returns (stream DrainResponse);
 
   // LoRA lifecycle.
   rpc LoadLora(LoadLoraRequest) returns (LoadLoraResponse);
@@ -744,7 +743,7 @@ server closes the stream with gRPC `OK`.
 
 ---
 
-## Health, abort, and drain
+## Health and abort
 
 ```protobuf
 message HealthRequest {
@@ -800,29 +799,6 @@ enum AbortStatus {
   ABORT_STATUS_ABORTED = 1;
   ABORT_STATUS_ALREADY_FINISHED = 2;
 }
-
-message DrainRequest {
-  bool stop_accepting_new_requests = 1;
-  optional uint32 deadline_ms = 2;
-  bool abort_after_deadline = 3;
-}
-
-message DrainResponse {
-  oneof event {
-    DrainState state = 1;
-    EngineError error = 5;
-  }
-  optional uint32 in_flight_requests = 2;
-  optional uint32 open_kv_sessions = 3;
-  string message = 4;
-}
-
-enum DrainState {
-  DRAIN_STATE_UNSPECIFIED = 0;
-  DRAIN_STATE_STARTED = 1;
-  DRAIN_STATE_IN_PROGRESS = 2;
-  DRAIN_STATE_COMPLETE = 3;
-}
 ```
 
 Exactly one abort target must be set. Use `all_requests {}` to abort every
@@ -831,13 +807,6 @@ or KV session target returns gRPC `NOT_FOUND`; an engine that does not support
 abort returns gRPC `UNIMPLEMENTED`. `ALREADY_FINISHED` remains a successful
 idempotent outcome rather than an error.
 
-`STARTED` and `IN_PROGRESS` are progress events; `COMPLETE` is terminal. A
-failure after the drain is accepted is represented by one terminal
-`EngineError`, not by a failed drain state. An absent `deadline_ms` means no
-deadline; an explicit zero means the deadline is immediate. Absent progress
-counts are unknown, while present zero values report that no requests or KV
-sessions remain.
-
 ---
 
 ## Standard errors
@@ -847,8 +816,6 @@ message EngineError {
   ErrorCode code = 1;
   string message = 2;
   bool retryable = 3;
-  optional uint64 retry_after_ms = 4;
-  google.protobuf.Struct details = 5;
 }
 
 enum ErrorCode {
@@ -883,14 +850,8 @@ accepted stream phase and report failures with non-OK gRPC status.
 `GenerationFinished` is terminal for its `output_index`; other output indexes
 may continue. The last `GenerationFinished` ends a successful aggregated or
 decode stream, `PrefillReady` ends a successful prefill stream, and
-`EngineError` ends any failed generation stream. `DrainState.COMPLETE` and
-`EngineError` terminate a drain stream. An `EngineError` also terminates a
-KV-event or runtime-event subscription. No response may follow a terminal
-`EngineError`. Application failure is neither a `GenerationFinished` reason nor
-a failed drain state.
+`EngineError` ends any failed generation stream. An `EngineError` also
+terminates a KV-event subscription. No response may follow a terminal
+`EngineError`. Application failure is not a `GenerationFinished` reason.
 
 `retryable` states whether the unchanged operation can succeed on retry.
-`retry_after_ms` is present only for retryable errors and is the recommended
-minimum delay; an explicit zero permits immediate retry. `details` contains
-machine-readable error context. Stable detail keys are part of this API;
-engine-specific keys should be namespaced to avoid collisions.

proto/openengine/v1/README.md

@@ -11,14 +11,12 @@ share the same package and together define the API.
 | File | Area |
 | --- | --- |
 | [`openengine.proto`](openengine.proto) | `Inference` and `Control` service declarations |
-| [`input.proto`](input.proto) | Shared text, token, and multimodal inputs |
-| [`server.proto`](server.proto) | Server identity, deployment capacity, engine roles, and parallelism |
+| [`server.proto`](server.proto) | Server identity, deployment capacity, engine roles, parallelism, and load |
 | [`model.proto`](model.proto) | Model metadata and inference capabilities |
-| [`generation.proto`](generation.proto) | Generation requests, parameters, streamed events, and usage |
+| [`generation.proto`](generation.proto) | Generation inputs, requests, parameters, streamed events, and usage |
 | [`lora.proto`](lora.proto) | LoRA adapter lifecycle |
 | [`kv.proto`](kv.proto) | KV sessions, connector discovery, and cache events |
-| [`lifecycle.proto`](lifecycle.proto) | Health, abort, and drain operations |
-| [`observability.proto`](observability.proto) | Load snapshots |
+| [`lifecycle.proto`](lifecycle.proto) | Health and abort operations |
 | [`error.proto`](error.proto) | Terminal errors for accepted streaming requests |
 
 Generate bindings from every `.proto` file in this directory. Compiling only

proto/openengine/v1/error.proto

@@ -7,16 +7,12 @@ syntax = "proto3";
 
 package openengine.v1;
 
-import "google/protobuf/struct.proto";
-
 // Accepted failures emit one terminal EngineError and close with OK.
 // Validation and transport failures use non-OK gRPC status instead.
 message EngineError {
   ErrorCode code = 1;
   string message = 2;
   bool retryable = 3; // Retry may succeed without changing the request.
-  optional uint64 retry_after_ms = 4; // Zero permits immediate retry.
-  google.protobuf.Struct details = 5; // Machine-readable context.
 }
 
 enum ErrorCode {

proto/openengine/v1/generation.proto

@@ -1,17 +1,42 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-// Generation requests, streamed events, and usage.
+// Generation inputs, requests, streamed events, and usage.
 
 syntax = "proto3";
 
 package openengine.v1;
 
 import "google/protobuf/struct.proto";
 import "openengine/v1/error.proto";
-import "openengine/v1/input.proto";
 import "openengine/v1/kv.proto";
 
+message TokenIds {
+  repeated uint32 ids = 1;
+}
+
+// Multimodal modality discriminator. UNSPECIFIED means the sender left it unset.
+enum Modality {
+  MODALITY_UNSPECIFIED = 0;
+  MODALITY_IMAGE = 1;
+  MODALITY_VIDEO = 2;
+  MODALITY_AUDIO = 3;
+}
+
+// A single multimodal input. Exactly one `source` should be set. The engine
+// owns fetch, decode, and preprocessing, so pre-decoded or RDMA media
+// descriptors are not represented here.
+message MediaItem {
+  Modality modality = 1;
+  oneof source {
+    string url = 2; // http(s):// -- engine fetches
+    string data_uri = 3; // data:<mime>;base64,<...> -- engine decodes
+    bytes raw_bytes = 4; // pre-fetched bytes -- engine still preprocesses
+  }
+  string mime_type = 5; // optional, hints raw_bytes decode
+  string uuid = 6; // optional caller id / mm_hash
+}
+
 message GenerateRequest {
   string request_id = 1;
   string model = 2;

proto/openengine/v1/lifecycle.proto

@@ -1,13 +1,12 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-// Health, abort, and drain operations.
+// Health and abort operations.
 
 syntax = "proto3";
 
 package openengine.v1;
 
-import "openengine/v1/error.proto";
 import "openengine/v1/kv.proto";
 import "openengine/v1/server.proto";
 
@@ -64,26 +63,3 @@ enum AbortStatus {
   ABORT_STATUS_ABORTED = 1;
   ABORT_STATUS_ALREADY_FINISHED = 2;
 }
-
-message DrainRequest {
-  bool stop_accepting_new_requests = 1;
-  optional uint32 deadline_ms = 2; // Absent means no deadline; zero is immediate.
-  bool abort_after_deadline = 3;
-}
-
-message DrainResponse {
-  oneof event {
-    DrainState state = 1; // Progress or successful completion.
-    EngineError error = 5; // Terminal failure.
-  }
-  optional uint32 in_flight_requests = 2;
-  optional uint32 open_kv_sessions = 3;
-  string message = 4;
-}
-
-enum DrainState {
-  DRAIN_STATE_UNSPECIFIED = 0;
-  DRAIN_STATE_STARTED = 1;
-  DRAIN_STATE_IN_PROGRESS = 2;
-  DRAIN_STATE_COMPLETE = 3;
-}

proto/openengine/v1/openengine.proto

@@ -21,7 +21,6 @@ import "openengine/v1/kv.proto";
 import "openengine/v1/lifecycle.proto";
 import "openengine/v1/lora.proto";
 import "openengine/v1/model.proto";
-import "openengine/v1/observability.proto";
 import "openengine/v1/server.proto";
 
 // Standard ASCII gRPC request metadata:
@@ -50,7 +49,6 @@ service Control {
   // Health and lifecycle.
   rpc Health(HealthRequest) returns (HealthResponse);
   rpc Abort(AbortRequest) returns (AbortResponse);
-  rpc Drain(DrainRequest) returns (stream DrainResponse);
 
   // LoRA lifecycle.
   rpc LoadLora(LoadLoraRequest) returns (LoadLoraResponse);

proto/openengine/v1/server.proto

@@ -1,7 +1,7 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-// Server identity, deployment capacity, engine role, and parallelism metadata.
+// Server identity, deployment capacity, roles, parallelism, and load.
 
 syntax = "proto3";
 
@@ -55,3 +55,33 @@ message ParallelismInfo {
   optional uint32 data_parallel_start_rank = 5;
   optional uint32 decode_context_parallel_size = 6; // Ranks per decode-context group; at least 1.
 }
+
+message GetLoadRequest {
+  bool include_per_rank = 1;
+}
+
+message LoadInfo {
+  string instance_id = 1;
+  optional uint64 timestamp_unix_nanos = 2;
+  optional uint32 running_requests = 3;
+  optional uint32 queued_requests = 4;
+  optional uint32 active_kv_sessions = 5;
+  optional uint64 used_kv_blocks = 6;
+  optional uint64 total_kv_blocks = 7;
+  optional uint64 running_tokens = 8;
+  optional uint64 waiting_tokens = 9;
+  optional uint32 prefill_batch_size = 10;
+  optional uint32 decode_batch_size = 11;
+  repeated RankLoadInfo ranks = 20;
+  google.protobuf.Struct attributes = 30;
+}
+
+message RankLoadInfo {
+  optional uint32 data_parallel_rank = 1;
+  optional uint32 running_requests = 2;
+  optional uint32 queued_requests = 3;
+  optional uint64 used_kv_blocks = 4;
+  optional uint64 total_kv_blocks = 5;
+  optional uint32 prefill_batch_size = 6;
+  optional uint32 decode_batch_size = 7;
+}