diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index 0169adf1..91224688 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -1219,6 +1219,9 @@ def _admit_sequences_from_message(self, msg: dict) -> None: seq.batchgen_debug = entry.get("batchgen_debug") seq.priority = entry.get("priority", 0) seq.sampling_params = entry.get("sampling_params") + # Per-request ignore_eos (vendor extension via extra_body). Honored + # per-sequence at the decode EOS check; OR'd with the global flag. + seq.ignore_eos = bool(entry.get("ignore_eos", False)) self.global_batch.add_sequence(seq) new_uuids.append(seq.uuid) @@ -4744,7 +4747,11 @@ def _check_and_handle_completions( decoded_lens[i] = seq.decoded_length max_lens[i] = seq.max_decode_length ctx_lens[i] = seq.current_context_length - eos_flags[i] = seq.eos_reached and not ignore_eos + # Honor ignore_eos per-sequence (vendor extension via extra_body), + # OR'd with the server-global flag. + eos_flags[i] = seq.eos_reached and not ( + ignore_eos or getattr(seq, "ignore_eos", False) + ) # Variable-length N-gram repetition detection at decision boundary # Catches repeating patterns of length 2-100 tokens (32 repetitions required) diff --git a/batchgen/sequence.py b/batchgen/sequence.py index db5ba6bc..e1b7dd67 100644 --- a/batchgen/sequence.py +++ b/batchgen/sequence.py @@ -97,6 +97,7 @@ class SequenceEntry: 'pool_slot_index', # Index in SchedulingPool's pre-allocated QueryBook 'priority', # 0=NORMAL, 1=HIGH (inherited from batch) 'sampling_params', # Per-request sampling params for this sequence + 'ignore_eos', # Per-request ignore_eos (vendor extension via extra_body) # Lifespan monitoring (BATCHGEN_SEQ_LIFESPAN=1) '_lifespan_log', # List[SeqEventRecord], ring buffer '_lifespan_idx', # int, next write position @@ -171,6 +172,7 @@ def __init__( self.pool_slot_index: int = -1 self.priority: int = 0 # 0=NORMAL, 1=HIGH self.sampling_params: Optional[Dict] = None + self.ignore_eos: bool = False # per-request vendor extension (extra_body) # Lifespan monitoring self._lifespan_log: list = [] diff --git a/batchgen/server/batch_scheduler.py b/batchgen/server/batch_scheduler.py index f2400285..3ff6d1c7 100644 --- a/batchgen/server/batch_scheduler.py +++ b/batchgen/server/batch_scheduler.py @@ -294,7 +294,7 @@ async def _process_batch(self, batch_id: str) -> None: prompts, None, # max_input_len: dynamically determined from prompts max_tokens, - False, # ignore_eos + bool(getattr(batch, "ignore_eos", False)), # ignore_eos (batch-level; per-request honored in pool mode) None, # temperature: handled via per-request sampling_params None, # top_p: handled via per-request sampling_params max_context_length=batch.max_context_length, @@ -885,7 +885,14 @@ async def _process_batch_pool_mode( # Build IntakeEntry objects and push to IntakePool entries = [] + batch_ignore_eos = bool(getattr(batch, "ignore_eos", False)) for idx, req in enumerate(requests): + # Per-request ignore_eos (vendor extension via extra_body), falling + # back to the batch-level CreateBatchRequest.ignore_eos. + req_ignore_eos = getattr(req.body, "ignore_eos", None) + eff_ignore_eos = ( + req_ignore_eos if req_ignore_eos is not None else batch_ignore_eos + ) entries.append(IntakeEntry( request_id=req.custom_id or f"{batch_id}_req_{idx}", batch_id=batch_id, @@ -894,6 +901,7 @@ async def _process_batch_pool_mode( "max_tokens": per_request_max_tokens[idx], "priority": 0, # TODO: support per-batch priority from API "sampling_params": sampling_params[idx] if sampling_params else {}, + "ignore_eos": eff_ignore_eos, "batchgen_debug": batch.batchgen_debug or {}, }, priority=Priority.NORMAL, @@ -1042,6 +1050,7 @@ async def _drain_intake_to_worker(self) -> None: "batch_id": entry.batch_id, "priority": entry.priority.value, "sampling_params": entry.raw_request.get("sampling_params", {}), + "ignore_eos": entry.raw_request.get("ignore_eos", False), "batchgen_debug": entry.raw_request.get("batchgen_debug", {}), }) diff --git a/batchgen/server/io_struct.py b/batchgen/server/io_struct.py index e534600a..11fb4425 100644 --- a/batchgen/server/io_struct.py +++ b/batchgen/server/io_struct.py @@ -67,6 +67,13 @@ class ChatCompletionRequest(BaseModel): default=None, description="Preserve prior assistant reasoning_content in Kimi-style chat templates", ) + # Vendor extension — pass via the OpenAI SDK `extra_body` (or as a top-level + # body key in a Batch JSONL line). Per-request; falls back to the batch-level + # CreateBatchRequest.ignore_eos when None. + ignore_eos: Optional[bool] = Field( + default=None, + description="If true, ignore EOS and decode to the max output length (vendor extension via extra_body).", + ) @validator("stream") def validate_stream(cls, value: Optional[bool]) -> Optional[bool]: @@ -94,6 +101,11 @@ class CompletionRequest(BaseModel): presence_penalty: Optional[float] = Field(default=0, ge=-2, le=2) frequency_penalty: Optional[float] = Field(default=0, ge=-2, le=2) user: Optional[str] = None + # Vendor extension (see ChatCompletionRequest.ignore_eos). + ignore_eos: Optional[bool] = Field( + default=None, + description="If true, ignore EOS and decode to the max output length (vendor extension via extra_body).", + ) @validator("stream") def validate_stream(cls, value: Optional[bool]) -> Optional[bool]: diff --git a/batchgen/worker/completion.py b/batchgen/worker/completion.py index 0b72099f..1ce1a176 100644 --- a/batchgen/worker/completion.py +++ b/batchgen/worker/completion.py @@ -68,7 +68,11 @@ def is_sequence_completed(ctx: CompletionContext, seq: "SequenceEntry") -> bool: return True if seq.current_context_length >= ctx.model_context_length: return True - if seq.eos_reached and not ctx.ignore_eos: + # ignore_eos honored per-sequence (vendor extension via extra_body), + # OR'd with the server-global flag. + if seq.eos_reached and not ( + ctx.ignore_eos or getattr(seq, "ignore_eos", False) + ): return True if seq._rep_detected: return True @@ -100,7 +104,10 @@ def get_finish_reason(ctx: CompletionContext, seq: "SequenceEntry") -> str: elif seq.current_context_length >= ctx.model_context_length: finish = "length" # Real EOS only — the token at seq.decoded_length-1 matches an EOS id - elif seq.eos_reached and not ctx.ignore_eos: + # (ignore_eos honored per-sequence OR globally). + elif seq.eos_reached and not ( + ctx.ignore_eos or getattr(seq, "ignore_eos", False) + ): finish = "stop" else: finish = "length" diff --git a/docs/input-format.md b/docs/input-format.md index 6f4b1872..df669c89 100644 --- a/docs/input-format.md +++ b/docs/input-format.md @@ -40,10 +40,16 @@ Reference for the JSONL input file format used by BatchGen's `/v1/batches` API. | `max_tokens` | int | No | Maximum output tokens to generate per request (legacy alias for `max_completion_tokens`) | | `temperature` | float | No | Sampling temperature (see [Sampling Parameters](#sampling-parameters)) | | `top_p` | float | No | Nucleus sampling threshold (see [Sampling Parameters](#sampling-parameters)) | -| `top_k` | int | No | Top-k filtering threshold (see [Sampling Parameters](#sampling-parameters)) | +| `top_k` | int | No | Top-k filtering threshold (vendor extension — not an OpenAI field; see [Sampling Parameters](#sampling-parameters)) | +| `ignore_eos` | bool | No | Ignore EOS and decode to the max output length (vendor extension; falls back to the batch-level `ignore_eos`) | *One of `messages` or `prompt` is required depending on the endpoint. +> **Vendor extensions (`extra_body`).** Some body fields above are not part of the OpenAI API +> (`top_k`, `ignore_eos`, `enable_thinking`/`thinking`, `preserve_thinking`, `reasoning_effort`). +> With the OpenAI SDK, pass them via `extra_body={...}`; in Batch JSONL, put them directly in `body`. +> See the full list in `batchgen_design/server/openai_extended_kwargs.md`. + --- ## Sampling Parameters diff --git a/tests/worker/test_completion.py b/tests/worker/test_completion.py index 382807da..ccf4835d 100644 --- a/tests/worker/test_completion.py +++ b/tests/worker/test_completion.py @@ -181,3 +181,26 @@ def test_handler_is_stateless(ctx_strict): # seq fields we read (not written) unchanged assert seq.decoded_length == 4 assert seq.eos_reached is False + + +def test_is_sequence_completed_honors_per_sequence_ignore_eos(ctx_strict): + """Per-sequence ignore_eos (vendor extension via extra_body) overrides a real + EOS even when the global ctx.ignore_eos is False; length limits still apply.""" + seq = _make_seq(eos_reached=True) + # Baseline under strict (global ignore_eos=False) ctx: real EOS completes. + assert CompletionHandler.is_sequence_completed(ctx_strict, seq) is True + # Per-sequence override: not completed by EOS. + seq.ignore_eos = True + assert CompletionHandler.is_sequence_completed(ctx_strict, seq) is False + # Length limit still completes regardless of ignore_eos. + seq.decoded_length = seq.max_decode_length + assert CompletionHandler.is_sequence_completed(ctx_strict, seq) is True + + +def test_get_finish_reason_per_sequence_ignore_eos_becomes_length(ctx_strict): + """A real EOS reports finish_reason 'stop' under strict ctx, but 'length' when + the sequence sets ignore_eos (per-request), without touching the global flag.""" + seq = _make_seq(eos_reached=True, decoded_length=4, max_decode_length=16) + assert CompletionHandler.get_finish_reason(ctx_strict, seq) == "stop" + seq.ignore_eos = True + assert CompletionHandler.get_finish_reason(ctx_strict, seq) == "length"