Batch inference with Token Factory¶
This notebook uploads a JSONL request file, creates a batch, checks that same batch, and downloads its output. Set NEBIUS_BATCH_INPUT_PATH to your JSONL file. Output is written to NEBIUS_BATCH_OUTPUT_PATH or a portable temporary-directory default.
Batch access may require project enablement and a configured batch endpoint. Confirm the current Token Factory batch documentation before running this notebook.
In [ ]:
import json
import os
import tempfile
from pathlib import Path
from openai import OpenAI
client = OpenAI(
base_url="https://api.tokenfactory.nebius.com/v1",
api_key=os.environ["NEBIUS_API_KEY"],
)
In [ ]:
batch_input_path = Path(
os.environ.get("NEBIUS_BATCH_INPUT_PATH", "batch.jsonl")
).expanduser()
if not batch_input_path.is_file():
raise FileNotFoundError(
f"Batch input not found at {batch_input_path}. "
"Set NEBIUS_BATCH_INPUT_PATH to a JSONL request file."
)
with batch_input_path.open("rb") as batch_input:
uploaded_file = client.files.create(file=batch_input, purpose="batch")
print({"input_file_id": uploaded_file.id, "filename": batch_input_path.name})
In [ ]:
batch = client.batches.create(
input_file_id=uploaded_file.id,
endpoint="/v1/chat/completions",
completion_window="24h",
metadata={"description": "nightly eval job"},
)
print({"batch_id": batch.id, "status": batch.status})
In [ ]:
# Re-run this cell until the batch reaches a terminal state.
status = client.batches.retrieve(batch.id)
print(
{
"batch_id": status.id,
"status": status.status,
"request_counts": status.request_counts,
}
)
In [ ]:
if not status.output_file_id:
raise RuntimeError(
f"Batch {status.id} has no output file yet (status: {status.status})."
)
file_response = client.files.content(status.output_file_id)
default_output_path = Path(tempfile.gettempdir()) / "nebius-batch-output.jsonl"
batch_output_path = Path(
os.environ.get("NEBIUS_BATCH_OUTPUT_PATH", str(default_output_path))
).expanduser()
batch_output_path.parent.mkdir(parents=True, exist_ok=True)
batch_output_path.write_bytes(file_response.content)
with batch_output_path.open(encoding="utf-8") as output_file:
result_count = sum(1 for line in output_file if line.strip() and json.loads(line))
print(f"Saved {result_count} batch result(s) to {batch_output_path}")