minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] improve: improve realtime webui playback pacing (#27148)

Browse Source
This commit is contained in:
Mick
2026-06-04 00:33:56 +08:00
committed by GitHub
parent 45a66f4088
commit b0f78bef97
19 changed files with 1980 additions and 310 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/apps/realtime_webui/app.js
+769 -204
View File
File diff suppressed because it is too large Load Diff
python/sglang/multimodal_gen/apps/realtime_webui/decoder_worker.js
+29 -19
View File
@@ -80,24 +80,36 @@ function rawFramesToRgbaBuffers(header, payload) {
return buffers; return buffers;
} }
async function encodedFrameToRgbaBuffers(header, payload) { function splitEncodedPayload(header, payload) {
if (typeof createImageBitmap === "undefined" || typeof OffscreenCanvas === "undefined") { const bytes = payload instanceof Uint8Array ? payload : new Uint8Array(payload);
const lengths = Array.isArray(header.payload_lengths) && header.payload_lengths.length
? header.payload_lengths.map(Number)
: [bytes.byteLength];
const payloads = [];
let offset = 0;
for (const length of lengths) {
payloads.push(bytes.buffer.slice(
bytes.byteOffset + offset,
bytes.byteOffset + offset + length,
));
offset += length;
}
return payloads;
}
async function encodedFramesToImageBitmaps(header, payload) {
if (typeof createImageBitmap === "undefined") {
throw new Error("This browser does not support worker image decoding"); throw new Error("This browser does not support worker image decoding");
} }
const blob = new Blob([payload], { type: header.content_type }); const frames = await Promise.all(splitEncodedPayload(header, payload).map((framePayload) => (
const bitmap = await createImageBitmap(blob); createImageBitmap(new Blob([framePayload], { type: header.content_type }))
const width = bitmap.width; )));
const height = bitmap.height;
const canvas = new OffscreenCanvas(width, height);
const ctx = canvas.getContext("2d", { alpha: false });
ctx.drawImage(bitmap, 0, 0);
const image = ctx.getImageData(0, 0, width, height);
bitmap.close?.();
return { return {
width, width: frames[0]?.width || 0,
height, height: frames[0]?.height || 0,
frames: [image.data.buffer], frame_type: "bitmap",
frames,
}; };
} }
@@ -107,12 +119,13 @@ async function decode(header, payload) {
header.content_type === WEBP_FRAME_CONTENT_TYPE || header.content_type === WEBP_FRAME_CONTENT_TYPE ||
header.content_type === JPEG_FRAME_CONTENT_TYPE header.content_type === JPEG_FRAME_CONTENT_TYPE
) { ) {
const decoded = await encodedFrameToRgbaBuffers(header, payload); const decoded = await encodedFramesToImageBitmaps(header, payload);
return { return {
id: header.__decode_id, id: header.__decode_id,
width: decoded.width, width: decoded.width,
height: decoded.height, height: decoded.height,
chunk: Number(header.chunk_index), chunk: Number(header.chunk_index),
frame_type: decoded.frame_type,
frames: decoded.frames, frames: decoded.frames,
}; };
} else if (header.content_type === RAW_RGB_CONTENT_TYPE) { } else if (header.content_type === RAW_RGB_CONTENT_TYPE) {
@@ -148,10 +161,7 @@ self.onmessage = async (event) => {
return; return;
} }
const result = await decode(message.header, message.payload); const result = await decode(message.header, message.payload);
self.postMessage( self.postMessage({ type: "decoded", ...result }, result.frames);
{ type: "decoded", ...result },
result.frames,
);
} catch (error) { } catch (error) {
self.postMessage({ self.postMessage({
type: "error", type: "error",
python/sglang/multimodal_gen/apps/realtime_webui/index.html
+23 -5
View File
@@ -4,7 +4,7 @@
<meta charset="utf-8" /> <meta charset="utf-8" />
<meta name="viewport" content="width=device-width, initial-scale=1" /> <meta name="viewport" content="width=device-width, initial-scale=1" />
<title>sglang-diffusion Realtime Studio</title> <title>sglang-diffusion Realtime Studio</title>
<link rel="stylesheet" href="./styles.css?v=realtime-sr-v38" /> <link rel="stylesheet" href="./styles.css?v=realtime-record-v49" />
</head> </head>
<body> <body>
<main class="shell"> <main class="shell">
@@ -60,6 +60,18 @@
</select> </select>
</label> </label>
</div> </div>
<label>SR model
<select id="upscalingModel">
<option value="">Quality x2</option>
<option
value="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"
selected
>
Fast general
</option>
<option value="/scratch/realesr-animevideov3.pth">Fast anime</option>
</select>
</label>
<label class="toggle-row"><input id="frameInterpolation" type="checkbox" />Smooth 2x frames</label> <label class="toggle-row"><input id="frameInterpolation" type="checkbox" />Smooth 2x frames</label>
<label class="toggle-row"><input id="continuous" type="checkbox" checked />Continuous session</label> <label class="toggle-row"><input id="continuous" type="checkbox" checked />Continuous session</label>
<div class="actions"> <div class="actions">
@@ -75,6 +87,11 @@
<span id="statusDot" class="dot"></span> <span id="statusDot" class="dot"></span>
<span id="statusText">Idle</span> <span id="statusText">Idle</span>
<span id="chunkText">chunk -</span> <span id="chunkText">chunk -</span>
<button id="recordBtn" class="record-button" type="button" aria-pressed="false" title="Record preview">
<span class="record-button-icon" aria-hidden="true"></span>
<span id="recordLabel">Record</span>
<span id="recordDuration" class="record-button-duration">00:00</span>
</button>
<span class="topbar-spacer"></span> <span class="topbar-spacer"></span>
<label class="preview-scale-control">Preview <label class="preview-scale-control">Preview
<input id="previewScale" type="range" min="80" max="170" value="120" /> <input id="previewScale" type="range" min="80" max="170" value="120" />
@@ -82,8 +99,8 @@
</label> </label>
<span class="stage-stat">output <b id="outputSizeText">832x480</b></span> <span class="stage-stat">output <b id="outputSizeText">832x480</b></span>
<span class="stage-stat">render <b id="renderFps">0</b> fps</span> <span class="stage-stat">render <b id="renderFps">0</b> fps</span>
<span class="stage-stat">theoretical <b id="theoreticalFpsText">-</b></span> <span class="stage-stat">source <b id="theoreticalFpsText">-</b></span>
<span class="stage-stat">wait <b id="stageLatencyText">-</b></span> <span class="stage-stat">buffer <b id="stageLatencyText">-</b></span>
</div> </div>
<div class="preview-frame"> <div class="preview-frame">
<canvas id="viewport" width="1280" height="720"></canvas> <canvas id="viewport" width="1280" height="720"></canvas>
@@ -145,6 +162,7 @@
</section> </section>
</section> </section>
</main> </main>
<script src="./app.js?v=realtime-sr-v38"></script> <script src="./playback_controller.js?v=realtime-playback-v13"></script>
</body> <script src="./app.js?v=realtime-record-v72"></script>
</body>
</html> </html>
python/sglang/multimodal_gen/apps/realtime_webui/playback_controller.js
+517
View File
@@ -0,0 +1,517 @@
(function attachRealtimePlaybackController(global) {
const DEFAULT_CONFIG = {
targetFps: 25,
minSourceFps: 1,
serverFpsAlphaUp: 0.28,
serverFpsAlphaDown: 0.2,
deliveryFpsAlphaUp: 0.08,
deliveryFpsAlphaDown: 0.55,
targetLeadChunkRatio: 1.5,
minTargetLeadMs: 1500,
maxTargetLeadMs: 2600,
maxLeadExtraChunkRatio: 8.0,
startLeadChunkRatio: 1.85,
minStartLeadMs: 1700,
resumeLeadChunkRatio: 2.5,
minResumeLeadMs: 1000,
maxResumeLeadMs: 1800,
rebufferLeadBoostMs: 250,
rebufferLeadBoostDecayMsPerSecond: 120,
deliveryLeadBoostDecayMsPerSecond: 80,
maxDeliveryLeadBoostMs: 2000,
deliveryStallExpectedMultiplier: 1.25,
receiveStallPlaybackRateMin: 0.65,
receiveStallPlaybackRateSlewPerSecond: 0.5,
lowWaterRatio: 0.4,
playbackRateGain: 0.14,
playbackRateMin: 0.92,
playbackRateMax: 1.08,
emergencyPlaybackRateMin: 0.9,
emergencyPlaybackRateMax: 1.12,
playbackRateSlewPerSecond: 0.08,
eventCutoverMaxMs: 420,
eventCutoverMaxFrames: 10,
settleEventCutoverMaxMs: 720,
settleEventCutoverMaxFrames: 18,
startupWarmupMinMs: 1500,
startupWarmupExpectedMultiplier: 3,
};
function clamp(value, min, max) {
return Math.min(max, Math.max(min, value));
}
function finitePositive(value) {
return Number.isFinite(value) && value > 0;
}
class RealtimePlaybackController {
constructor(config = {}) {
this.config = { ...DEFAULT_CONFIG, ...config };
this.reset({ targetFps: this.config.targetFps });
}
reset({ targetFps } = {}) {
this.targetFps = Math.max(1, Number(targetFps || this.config.targetFps));
this.sourceFps = this.targetFps;
this.serverFps = this.targetFps;
this.deliveryFps = this.targetFps;
this.hasServerSample = false;
this.hasDeliverySample = false;
this.latestChunkDurationMs = 1000 / this.targetFps;
this.latestChunkFrames = 1;
this.playbackRate = 1;
this.renderFps = this.targetFps;
this.queue = [];
this.lastDrawAt = 0;
this.lastRateUpdateAt = 0;
this.renderedFrames = 0;
this.droppedFrames = 0;
this.buffering = true;
this.pendingEventId = 0;
this.pendingEventSentAt = 0;
this.pendingEventCutoverMode = "motion";
this.lastDropReason = "";
this.lastDropAt = 0;
this.lastDropCount = 0;
this.rebufferLeadBoostMs = 0;
this.deliveryLeadBoostMs = 0;
this.chunkReceives = new Map();
this.serverStatChunks = new Set();
this.lastFinalReceiveAt = 0;
this.receiveStalled = false;
}
setTargetFps(targetFps) {
const nextTargetFps = Math.max(1, Number(targetFps || this.config.targetFps));
this.targetFps = nextTargetFps;
if (!this.hasServerSample && !this.hasDeliverySample) {
this.serverFps = nextTargetFps;
this.deliveryFps = nextTargetFps;
this.sourceFps = nextTargetFps;
this.renderFps = nextTargetFps;
} else {
this.serverFps = clamp(this.serverFps, this.config.minSourceFps, nextTargetFps);
this.deliveryFps = clamp(this.deliveryFps, this.config.minSourceFps, nextTargetFps);
this.sourceFps = clamp(this.sourceFps, this.config.minSourceFps, nextTargetFps);
this.renderFps = this.sourceFps * this.playbackRate;
}
this.latestChunkDurationMs = Math.max(this.latestChunkDurationMs, 1000 / this.targetFps);
}
clear() {
const frames = this.queue.splice(0);
this.lastDrawAt = 0;
this.buffering = true;
return frames;
}
noteInputEvent(eventId, now, { cutoverMode = "motion" } = {}) {
this.pendingEventId = Number(eventId || 0);
this.pendingEventSentAt = Number(now || 0);
this.pendingEventCutoverMode = cutoverMode;
}
observeServerStats(stats, now) {
const chunkIndex = Number(stats.chunk_index || 0);
const numFrames = Number(stats.num_frames || 0);
const chunkTotalMs = Number(stats.chunk_total_ms || 0);
if (numFrames > 0 && chunkTotalMs > 0) {
this.serverStatChunks.add(chunkIndex);
if (this.serverStatChunks.size > 128) {
this.serverStatChunks.delete(this.serverStatChunks.values().next().value);
}
const expectedMs = numFrames / Math.max(1, this.targetFps) * 1000;
const isStartupWarmup =
chunkIndex === 0 &&
chunkTotalMs > Math.max(
this.config.startupWarmupMinMs,
expectedMs * this.config.startupWarmupExpectedMultiplier,
);
if (isStartupWarmup) return this.snapshot();
this.#observeFpsSample("server", {
fps: numFrames / (chunkTotalMs / 1000),
frameCount: numFrames,
durationMs: chunkTotalMs,
now,
});
}
return this.snapshot();
}
enqueueDecodedFrames(header, frames, now) {
const chunkIndex = Number(header.chunk_index || 0);
const eventId = Number(header.event_id || 0);
const receivedAt = Number(header.__received_at || now);
const preparedFrames = frames.map((frame) => ({
...frame,
chunk: Number(frame.chunk ?? chunkIndex),
chunkIndex,
eventId,
}));
const droppedFrames = [];
let cutover = null;
if (this.pendingEventId && eventId >= this.pendingEventId) {
const oldEventFrameCount = this.#oldEventFrameCount(eventId);
const graceFrames = this.#eventGraceFrames();
const dropCount = Math.max(0, oldEventFrameCount - graceFrames);
if (dropCount > 0) {
droppedFrames.push(...this.queue.splice(graceFrames, dropCount));
this.#recordDrop(dropCount, "event cutover", now);
}
cutover = {
eventId,
latencyMs: this.pendingEventSentAt ? now - this.pendingEventSentAt : 0,
};
this.pendingEventId = 0;
this.pendingEventSentAt = 0;
this.pendingEventCutoverMode = "motion";
}
this.queue.push(...preparedFrames);
this.#observeChunkArrival(header, preparedFrames.length, receivedAt, now);
droppedFrames.push(...this.#trimBacklog(now));
return { droppedFrames, cutover, snapshot: this.snapshot() };
}
render(now, { hasPendingInput = true } = {}) {
this.#decayRebufferBoost(now);
this.#updateReceiveStallGuard(now);
const droppedFrames = this.#trimBacklog(now);
if (!this.queue.length) {
if (this.renderedFrames && hasPendingInput && !this.buffering) {
this.buffering = true;
this.rebufferLeadBoostMs = Math.max(
this.rebufferLeadBoostMs,
this.config.rebufferLeadBoostMs,
);
}
return { action: "hold", droppedFrames, snapshot: this.snapshot() };
}
const bufferMs = this.bufferDurationMs;
if (
hasPendingInput &&
this.receiveStalled &&
this.renderedFrames &&
bufferMs < this.targetLeadMs
) {
this.buffering = true;
this.lastDrawAt = 0;
return { action: "hold", droppedFrames, snapshot: this.snapshot() };
}
if (
hasPendingInput &&
this.buffering &&
bufferMs < (this.renderedFrames ? this.#resumeLeadMs() : this.#startLeadMs())
) {
this.buffering = true;
this.lastDrawAt = 0;
return { action: "hold", droppedFrames, snapshot: this.snapshot() };
}
if (this.buffering) {
this.buffering = false;
this.lastDrawAt = 0;
}
this.#updatePlaybackRate(now);
const targetMs = 1000 / Math.max(1, this.renderFps);
const elapsedMs = this.lastDrawAt ? now - this.lastDrawAt : targetMs;
if (elapsedMs < targetMs) {
return { action: "wait", droppedFrames, snapshot: this.snapshot() };
}
const frame = this.queue.shift();
this.renderedFrames += 1;
this.lastDrawAt = !this.lastDrawAt || elapsedMs > targetMs * 4
? now
: now - (elapsedMs % targetMs);
return { action: "draw", frame, droppedFrames, snapshot: this.snapshot() };
}
get queuedFrames() {
return this.queue.length;
}
get bufferDurationMs() {
return this.queue.length / Math.max(1, this.sourceFps) * 1000;
}
get targetLeadMs() {
const base = clamp(
this.latestChunkDurationMs * this.config.targetLeadChunkRatio,
this.config.minTargetLeadMs,
this.config.maxTargetLeadMs,
);
return clamp(
base + this.rebufferLeadBoostMs + this.deliveryLeadBoostMs,
this.config.minTargetLeadMs,
this.config.maxTargetLeadMs +
this.config.rebufferLeadBoostMs +
this.config.maxDeliveryLeadBoostMs,
);
}
get maxLeadMs() {
return this.targetLeadMs + this.latestChunkDurationMs * this.config.maxLeadExtraChunkRatio;
}
snapshot() {
return {
queueFrames: this.queue.length,
bufferMs: this.bufferDurationMs,
targetLeadMs: this.targetLeadMs,
maxLeadMs: this.maxLeadMs,
sourceFps: this.sourceFps,
serverFps: this.serverFps,
deliveryFps: this.deliveryFps,
targetFps: this.targetFps,
renderFps: this.renderFps,
playbackRate: this.playbackRate,
droppedFrames: this.droppedFrames,
lastDropAt: this.lastDropAt,
lastDropCount: this.lastDropCount,
buffering: this.buffering,
lastDropReason: this.lastDropReason,
};
}
#observeFpsSample(kind, { fps, frameCount, durationMs, now }) {
if (!finitePositive(fps)) return;
const cappedFps = clamp(fps, this.config.minSourceFps, this.targetFps);
const isDelivery = kind === "delivery";
const currentFps = isDelivery ? this.deliveryFps : this.serverFps;
const hasSample = isDelivery ? this.hasDeliverySample : this.hasServerSample;
let nextFps;
if (!hasSample) {
nextFps = cappedFps;
} else {
const alpha = cappedFps > currentFps
? (isDelivery ? this.config.deliveryFpsAlphaUp : this.config.serverFpsAlphaUp)
: (isDelivery ? this.config.deliveryFpsAlphaDown : this.config.serverFpsAlphaDown);
nextFps = currentFps * (1 - alpha) + cappedFps * alpha;
}
if (isDelivery) {
this.deliveryFps = nextFps;
this.hasDeliverySample = true;
this.#observeDeliveryJitter(frameCount, durationMs);
} else {
this.serverFps = nextFps;
this.hasServerSample = true;
}
const effectiveFps = this.hasServerSample
? this.serverFps
: (this.hasDeliverySample ? this.deliveryFps : this.targetFps);
this.sourceFps = clamp(effectiveFps, this.config.minSourceFps, this.targetFps);
if (!isDelivery || !this.hasServerSample) {
this.latestChunkFrames = Math.max(1, Number(frameCount || this.latestChunkFrames));
this.latestChunkDurationMs = clamp(
Number(durationMs || (this.latestChunkFrames / Math.max(1, this.sourceFps) * 1000)),
1000 / Math.max(1, this.targetFps),
2500,
);
}
this.#updatePlaybackRate(now);
}
#observeDeliveryJitter(frameCount, durationMs) {
if (!this.hasServerSample || !finitePositive(durationMs)) return;
const expectedMs = Number(frameCount || 0) / Math.max(1, this.serverFps) * 1000;
if (expectedMs <= 0) return;
if (durationMs <= expectedMs * this.config.deliveryStallExpectedMultiplier) return;
const boostMs = clamp(
durationMs - expectedMs,
0,
this.config.maxDeliveryLeadBoostMs,
);
this.deliveryLeadBoostMs = Math.max(this.deliveryLeadBoostMs, boostMs);
}
#updateReceiveStallGuard(now) {
this.receiveStalled = false;
if (!this.lastFinalReceiveAt || !this.hasServerSample) return;
const elapsedMs = now - this.lastFinalReceiveAt;
const expectedMs = Math.max(
this.latestChunkDurationMs,
this.latestChunkFrames / Math.max(1, this.serverFps) * 1000,
);
if (elapsedMs <= expectedMs * this.config.deliveryStallExpectedMultiplier) return;
this.receiveStalled = true;
this.deliveryLeadBoostMs = Math.max(
this.deliveryLeadBoostMs,
clamp(elapsedMs - expectedMs, 0, this.config.maxDeliveryLeadBoostMs),
);
}
#observeChunkArrival(header, frameCount, receivedAt, now) {
const chunkIndex = Number(header.chunk_index || 0);
const state = this.chunkReceives.get(chunkIndex) || {
firstAt: receivedAt,
frames: 0,
};
state.frames += Number(frameCount || 0);
state.lastAt = receivedAt;
this.chunkReceives.set(chunkIndex, state);
const frameBatchIndex = Number(header.frame_batch_index || 0);
const numFrameBatches = Number(header.num_frame_batches || 1);
const isFinalFrameBatch =
Boolean(header.is_final_frame_batch) ||
frameBatchIndex + 1 >= numFrameBatches;
if (!isFinalFrameBatch) return;
const durationMs = this.lastFinalReceiveAt
? receivedAt - this.lastFinalReceiveAt
: 0;
this.lastFinalReceiveAt = receivedAt;
if (state.frames > 0 && durationMs > 0) {
this.#observeFpsSample("delivery", {
fps: state.frames / (durationMs / 1000),
frameCount: state.frames,
durationMs,
now,
});
}
this.chunkReceives.delete(chunkIndex);
}
#updatePlaybackRate(now) {
const bufferMs = this.bufferDurationMs;
const targetLeadMs = Math.max(1, this.targetLeadMs);
const error = (bufferMs - targetLeadMs) / targetLeadMs;
const emergency =
bufferMs > this.maxLeadMs ||
bufferMs < targetLeadMs * this.config.lowWaterRatio ||
(this.receiveStalled && bufferMs < targetLeadMs);
const minRate = emergency
? (
this.receiveStalled
? this.config.receiveStallPlaybackRateMin
: this.config.emergencyPlaybackRateMin
)
: this.config.playbackRateMin;
const maxRate = this.receiveStalled && bufferMs < targetLeadMs
? 1
: emergency
? this.config.emergencyPlaybackRateMax
: this.config.playbackRateMax;
const desiredRate = clamp(
1 + error * this.config.playbackRateGain,
minRate,
maxRate,
);
if (!this.lastRateUpdateAt) {
this.playbackRate = desiredRate;
} else {
const dtSeconds = Math.max(0.001, (now - this.lastRateUpdateAt) / 1000);
const slewPerSecond = this.receiveStalled
? this.config.receiveStallPlaybackRateSlewPerSecond
: this.config.playbackRateSlewPerSecond;
const maxDelta = slewPerSecond * dtSeconds;
this.playbackRate = clamp(
desiredRate,
this.playbackRate - maxDelta,
this.playbackRate + maxDelta,
);
}
this.lastRateUpdateAt = now;
this.renderFps = clamp(
this.sourceFps * this.playbackRate,
this.config.minSourceFps,
this.targetFps * this.config.emergencyPlaybackRateMax,
);
}
#trimBacklog(now) {
const droppedFrames = [];
while (this.queue.length && this.bufferDurationMs > this.maxLeadMs) {
const firstChunk = this.queue[0].chunkIndex;
let dropCount = 0;
while (
dropCount < this.queue.length &&
this.queue[dropCount].chunkIndex === firstChunk
) {
dropCount += 1;
}
if (!dropCount || dropCount >= this.queue.length) break;
droppedFrames.push(...this.queue.splice(0, dropCount));
this.#recordDrop(dropCount, "backlog", now);
}
return droppedFrames;
}
#oldEventFrameCount(nextEventId) {
let count = 0;
while (count < this.queue.length && this.queue[count].eventId < nextEventId) {
count += 1;
}
return count;
}
#eventGraceFrames() {
const byTime = Math.max(
1,
Math.round(this.sourceFps * this.#eventCutoverMaxMs() / 1000),
);
const byChunkRatio = this.pendingEventCutoverMode === "settle" ? 1.5 : 0.85;
const byChunk = Math.max(1, Math.round(this.latestChunkFrames * byChunkRatio));
return Math.min(this.#eventCutoverMaxFrames(), byTime, byChunk);
}
#eventCutoverMaxMs() {
return this.pendingEventCutoverMode === "settle"
? this.config.settleEventCutoverMaxMs
: this.config.eventCutoverMaxMs;
}
#eventCutoverMaxFrames() {
return this.pendingEventCutoverMode === "settle"
? this.config.settleEventCutoverMaxFrames
: this.config.eventCutoverMaxFrames;
}
#startLeadMs() {
return Math.max(
this.config.minStartLeadMs,
this.latestChunkDurationMs * this.config.startLeadChunkRatio,
this.targetLeadMs,
);
}
#resumeLeadMs() {
return clamp(
this.latestChunkDurationMs * this.config.resumeLeadChunkRatio,
this.config.minResumeLeadMs,
this.config.maxResumeLeadMs,
);
}
#decayRebufferBoost(now) {
if ((!this.rebufferLeadBoostMs && !this.deliveryLeadBoostMs) || !this.lastRateUpdateAt) return;
const dtSeconds = Math.max(0, (now - this.lastRateUpdateAt) / 1000);
this.rebufferLeadBoostMs = Math.max(
0,
this.rebufferLeadBoostMs - dtSeconds * this.config.rebufferLeadBoostDecayMsPerSecond,
);
this.deliveryLeadBoostMs = Math.max(
0,
this.deliveryLeadBoostMs - dtSeconds * this.config.deliveryLeadBoostDecayMsPerSecond,
);
}
#recordDrop(count, reason, now) {
this.droppedFrames += count;
this.lastDropAt = Number(now || 0);
this.lastDropCount = count;
this.lastDropReason = reason;
}
}
global.RealtimePlaybackController = RealtimePlaybackController;
if (typeof module !== "undefined" && module.exports) {
module.exports = { RealtimePlaybackController };
}
})(typeof globalThis !== "undefined" ? globalThis : window);
python/sglang/multimodal_gen/apps/realtime_webui/playback_controller_test.js
+116
View File
@@ -0,0 +1,116 @@
const assert = require("node:assert/strict");
const { RealtimePlaybackController } = require("./playback_controller.js");
function frames(count, chunk) {
return Array.from({ length: count }, (_, index) => ({
image: { close() {} },
chunk,
index,
}));
}
function enqueueChunk(controller, {
chunk,
eventId = 0,
frameCount = 12,
durationMs = 480,
now,
}) {
controller.observeServerStats({
chunk_index: chunk,
num_frames: frameCount,
chunk_total_ms: durationMs,
}, now);
return controller.enqueueDecodedFrames({
chunk_index: chunk,
event_id: eventId,
num_frames: frameCount,
__received_at: now,
is_final_frame_batch: true,
}, frames(frameCount, chunk), now);
}
function renderFor(controller, startMs, durationMs) {
let rendered = 0;
for (let now = startMs; now <= startMs + durationMs; now += 16.67) {
const decision = controller.render(now, { hasPendingInput: true });
if (decision.action === "draw") rendered += 1;
}
return rendered;
}
function stableSourceDoesNotDrop() {
const controller = new RealtimePlaybackController({ targetFps: 25 });
let now = 0;
for (let chunk = 0; chunk < 8; chunk += 1) {
now += 480;
enqueueChunk(controller, { chunk, now });
renderFor(controller, now, 480);
}
const snapshot = controller.snapshot();
assert.equal(snapshot.droppedFrames, 0);
assert.ok(snapshot.sourceFps > 24 && snapshot.sourceFps <= 25);
}
function slowServerCapsRenderFps() {
const controller = new RealtimePlaybackController({ targetFps: 25 });
let now = 0;
for (let chunk = 0; chunk < 8; chunk += 1) {
now += 1360;
enqueueChunk(controller, { chunk, durationMs: 1360, now });
renderFor(controller, now, 1360);
}
const snapshot = controller.snapshot();
assert.ok(snapshot.sourceFps > 8 && snapshot.sourceFps < 10);
assert.ok(snapshot.renderFps <= 10);
}
function backlogDropsContiguousOldFrames() {
const controller = new RealtimePlaybackController({ targetFps: 25 });
let now = 100;
for (let chunk = 0; chunk < 13; chunk += 1) {
enqueueChunk(controller, { chunk, now, durationMs: 480 });
now += 20;
}
const snapshot = controller.snapshot();
assert.ok(snapshot.droppedFrames > 0);
assert.equal(snapshot.lastDropReason, "backlog");
}
function eventCutoverKeepsShortGrace() {
const controller = new RealtimePlaybackController({ targetFps: 25 });
enqueueChunk(controller, { chunk: 1, frameCount: 24, durationMs: 960, now: 1000 });
controller.noteInputEvent(5, 1050);
const result = enqueueChunk(controller, {
chunk: 2,
eventId: 5,
frameCount: 12,
durationMs: 480,
now: 1150,
});
assert.ok(result.cutover);
assert.ok(result.droppedFrames.length >= 14);
assert.equal(controller.queue[0].chunk, 1);
assert.equal(controller.queue[0].index, 0);
}
function settleEventCutoverKeepsWiderGrace() {
const controller = new RealtimePlaybackController({ targetFps: 25 });
enqueueChunk(controller, { chunk: 1, frameCount: 24, durationMs: 960, now: 1000 });
controller.noteInputEvent(5, 1050, { cutoverMode: "settle" });
const result = enqueueChunk(controller, {
chunk: 2,
eventId: 5,
frameCount: 12,
durationMs: 480,
now: 1150,
});
assert.ok(result.cutover);
assert.ok(result.droppedFrames.length <= 12);
}
stableSourceDoesNotDrop();
slowServerCapsRenderFps();
backlogDropsContiguousOldFrames();
eventCutoverKeepsShortGrace();
settleEventCutoverKeepsWiderGrace();
python/sglang/multimodal_gen/apps/realtime_webui/styles.css
+100 -4
View File
@@ -17,6 +17,7 @@
body { body {
margin: 0; margin: 0;
overflow-x: hidden;
min-height: 100vh; min-height: 100vh;
background: background:
linear-gradient(90deg, rgba(23, 26, 22, 0.035) 1px, transparent 1px), linear-gradient(90deg, rgba(23, 26, 22, 0.035) 1px, transparent 1px),
@@ -32,8 +33,10 @@ button:disabled { cursor: wait; opacity: 0.64; transform: none; }
.shell { .shell {
display: grid; display: grid;
grid-template-columns: minmax(260px, 320px) minmax(560px, 1fr); grid-template-columns: minmax(260px, 320px) minmax(0, 1fr);
gap: 18px; gap: 18px;
width: 100%;
max-width: 100vw;
min-height: 100vh; min-height: 100vh;
padding: 18px; padding: 18px;
} }
@@ -219,6 +222,7 @@ button:focus-visible {
display: grid; display: grid;
gap: 18px; gap: 18px;
min-width: 0; min-width: 0;
max-width: 100%;
} }
.stage { .stage {
@@ -227,6 +231,7 @@ button:focus-visible {
grid-template-rows: auto auto auto auto auto; grid-template-rows: auto auto auto auto auto;
align-self: start; align-self: start;
justify-self: center; justify-self: center;
min-width: 0;
width: 100%; width: 100%;
max-width: min(1500px, 100%); max-width: min(1500px, 100%);
overflow: hidden; overflow: hidden;
@@ -273,12 +278,15 @@ button:focus-visible {
.topbar, .timeline { .topbar, .timeline {
display: flex; display: flex;
align-items: center; align-items: center;
gap: 14px; gap: 10px;
min-width: 0;
height: 44px; height: 44px;
padding: 0 14px; padding: 0 14px;
color: #e8eadf; color: #e8eadf;
background: rgba(17, 20, 15, 0.88); background: rgba(17, 20, 15, 0.88);
font-size: 12px; font-size: 12px;
font-variant-numeric: tabular-nums;
white-space: nowrap;
} }
.topbar > * { .topbar > * {
@@ -287,11 +295,76 @@ button:focus-visible {
} }
.topbar-spacer { flex: 1; } .topbar-spacer { flex: 1; }
.record-button {
display: inline-flex;
align-items: center;
justify-content: center;
gap: 6px;
flex: 0 0 118px;
width: 118px;
min-height: 28px;
height: 28px;
padding: 0 9px;
border-color: rgba(232, 234, 223, 0.22);
background: rgba(238, 241, 236, 0.08);
color: #e8eadf;
font-size: 11px;
font-variant-numeric: tabular-nums;
}
.record-button:hover:not(:disabled) {
border-color: rgba(232, 234, 223, 0.44);
background: rgba(238, 241, 236, 0.14);
box-shadow: none;
transform: none;
}
.record-button:active:not(:disabled),
.record-button.is-pressed:not(:disabled),
.record-button:focus-visible {
transform: none;
}
.record-button.is-recording {
border-color: color-mix(in oklch, var(--accent), white 18%);
background: var(--accent);
color: #fffdf7;
}
.record-button.is-saving {
cursor: wait;
opacity: 0.76;
}
#recordLabel {
flex: 0 0 36px;
text-align: left;
}
.record-button-icon {
flex: 0 0 9px;
width: 9px;
height: 9px;
min-width: 9px;
border-radius: 50%;
background: var(--accent);
box-shadow: 0 0 0 3px rgba(185, 84, 60, 0.16);
}
.record-button.is-recording .record-button-icon {
border-radius: 2px;
background: #fffdf7;
box-shadow: none;
}
.record-button-duration {
display: inline-block;
flex: 0 0 34px;
min-width: 34px;
text-align: right;
color: rgba(232, 234, 223, 0.7);
}
.record-button.is-recording .record-button-duration {
color: rgba(255, 253, 247, 0.86);
}
.preview-scale-control { .preview-scale-control {
display: inline-flex; display: inline-flex;
align-items: center; align-items: center;
gap: 8px; gap: 8px;
min-width: 190px; flex: 0 0 170px;
min-width: 170px;
margin: 0; margin: 0;
color: rgba(232, 234, 223, 0.72); color: rgba(232, 234, 223, 0.72);
font-size: 11px; font-size: 11px;
@@ -324,13 +397,36 @@ button:focus-visible {
display: inline-flex; display: inline-flex;
align-items: center; align-items: center;
gap: 5px; gap: 5px;
min-width: max-content; flex: 0 1 auto;
min-width: 0;
color: rgba(232, 234, 223, 0.72); color: rgba(232, 234, 223, 0.72);
line-height: 1; line-height: 1;
} }
.stage-stat b { .stage-stat b {
display: inline-block;
color: #fffdf7; color: #fffdf7;
font-weight: 650; font-weight: 650;
font-variant-numeric: tabular-nums;
}
#outputSizeText { min-width: 206px; }
#renderFps { min-width: 2ch; text-align: right; }
#theoreticalFpsText { min-width: 116px; }
#stageLatencyText { min-width: 120px; }
@media (max-width: 1180px) {
.topbar {
flex-wrap: wrap;
height: auto;
min-height: 44px;
padding: 8px 14px;
row-gap: 7px;
}
.topbar-spacer { display: none; }
.preview-scale-control { flex-basis: 170px; min-width: 170px; }
#outputSizeText { min-width: 156px; }
#theoreticalFpsText { min-width: 100px; }
#stageLatencyText { min-width: 108px; }
} }
.timeline { justify-content: flex-end; border-top: 1px solid rgba(232, 234, 223, 0.12); } .timeline { justify-content: flex-end; border-top: 1px solid rgba(232, 234, 223, 0.12); }
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
+5 -5
View File
@@ -43,8 +43,10 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import (
log_batch_completion, log_batch_completion,
log_generation_timer, log_generation_timer,
) )
from sglang.multimodal_gen.runtime.utils.trace_wrapper import trace_req from sglang.multimodal_gen.runtime.utils.trace_wrapper import (
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info init_diffusion_tracing,
trace_req,
)
logger = init_logger(__name__) logger = init_logger(__name__)
@@ -123,9 +125,7 @@ class DiffGenerator:
instance = cls( instance = cls(
server_args=server_args, server_args=server_args,
) )
if server_args.enable_trace: init_diffusion_tracing(server_args, "DiffGenerator")
process_tracing_init(server_args.otlp_traces_endpoint, "sglang-diffusion")
trace_set_thread_info("DiffGenerator")
logger.info(f"Local mode: {local_mode}") logger.info(f"Local mode: {local_mode}")
if local_mode: if local_mode:
python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py
+2
View File
@@ -151,6 +151,8 @@ class RealtimeVideoGenerationsRequest(VideoGenerationsRequest):
num_profiled_timesteps: Optional[int] = None num_profiled_timesteps: Optional[int] = None
profile_all_stages: Optional[bool] = False profile_all_stages: Optional[bool] = False
realtime_output_format: Optional[Literal["raw", "webp", "jpeg"]] = None realtime_output_format: Optional[Literal["raw", "webp", "jpeg"]] = None
realtime_preview_max_width: Optional[int] = None
realtime_output_pacing: Optional[bool] = False
realtime_causal_sink_size: Optional[int] = None realtime_causal_sink_size: Optional[int] = None
realtime_causal_kv_cache_num_frames: Optional[int] = None realtime_causal_kv_cache_num_frames: Optional[int] = None
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py
+4
View File
@@ -430,6 +430,10 @@ class LingBotWorldRealtimeAdapter(RealtimeModelAdapter):
batch.realtime_event_id = self._state(session).latest_sampled_event_id batch.realtime_event_id = self._state(session).latest_sampled_event_id
if session.request is not None: if session.request is not None:
batch.realtime_output_format = session.request.realtime_output_format batch.realtime_output_format = session.request.realtime_output_format
batch.realtime_preview_max_width = (
session.request.realtime_preview_max_width
)
batch.realtime_output_pacing = bool(session.request.realtime_output_pacing)
batch.realtime_causal_sink_size = session.request.realtime_causal_sink_size batch.realtime_causal_sink_size = session.request.realtime_causal_sink_size
batch.realtime_causal_kv_cache_num_frames = ( batch.realtime_causal_kv_cache_num_frames = (
session.request.realtime_causal_kv_cache_num_frames session.request.realtime_causal_kv_cache_num_frames
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/generate_session.py
+4
View File
@@ -36,6 +36,8 @@ class GenerateSession:
self.realtime_session = RealtimeSession() self.realtime_session = RealtimeSession()
self.adapter: RealtimeModelAdapter | None = None self.adapter: RealtimeModelAdapter | None = None
self.adapter_state: Any = None self.adapter_state: Any = None
self.output_pace_next_send_at: float | None = None
self.output_pace_last_event_id: int | None = None
def set_adapter(self, adapter: RealtimeModelAdapter): def set_adapter(self, adapter: RealtimeModelAdapter):
self.adapter = adapter self.adapter = adapter
@@ -53,6 +55,8 @@ class GenerateSession:
self.current_chunk = None self.current_chunk = None
self.adapter = None self.adapter = None
self.adapter_state = None self.adapter_state = None
self.output_pace_next_send_at = None
self.output_pace_last_event_id = None
self.realtime_session.dispose() self.realtime_session.dispose()
def new_chunk(self) -> RealtimeChunkContext: def new_chunk(self) -> RealtimeChunkContext:
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_output_adapter.py
+161 -22
View File
@@ -48,6 +48,7 @@ class RealtimeFrameBatchHeader(TypedDict, total=False):
raw_size: int raw_size: int
encoding: str encoding: str
delta_reference: str delta_reference: str
payload_lengths: list[int]
event_id: int event_id: int
frame_batch_index: int frame_batch_index: int
num_frame_batches: int num_frame_batches: int
@@ -64,6 +65,7 @@ class RealtimeFrameSendStats(TypedDict):
raw_payload_build_ms: float raw_payload_build_ms: float
raw_write_ms: float raw_write_ms: float
ws_write_ms: float ws_write_ms: float
pace_wait_ms: float
raw_bytes: int raw_bytes: int
ws_payload_bytes: int ws_payload_bytes: int
num_frames: int num_frames: int
@@ -79,6 +81,7 @@ def empty_frame_send_stats(content_type: str = "") -> RealtimeFrameSendStats:
"raw_payload_build_ms": 0.0, "raw_payload_build_ms": 0.0,
"raw_write_ms": 0.0, "raw_write_ms": 0.0,
"ws_write_ms": 0.0, "ws_write_ms": 0.0,
"pace_wait_ms": 0.0,
"raw_bytes": 0, "raw_bytes": 0,
"ws_payload_bytes": 0, "ws_payload_bytes": 0,
"num_frames": 0, "num_frames": 0,
@@ -135,7 +138,7 @@ ENCODED_PREVIEW_FORMATS = {"webp", "jpeg"}
class _TransportPayload: class _TransportPayload:
content_type: str content_type: str
payload: bytes payload: bytes
metadata: dict[str, int | str | bool] metadata: dict[str, int | str | bool | list[int]]
last_raw_rgb_frame: bytes | None = None last_raw_rgb_frame: bytes | None = None
last_event_id: int | None = None last_event_id: int | None = None
@@ -155,9 +158,14 @@ def _encode_rgb_frame_to_webp(
width: int, width: int,
height: int, height: int,
quality: int, quality: int,
preview_max_width: int | None,
) -> bytes: ) -> bytes:
buffer = io.BytesIO() buffer = io.BytesIO()
Image.frombytes("RGB", (width, height), frame).save( image = _resize_preview_image(
Image.frombytes("RGB", (width, height), frame),
preview_max_width=preview_max_width,
)
image.save(
buffer, buffer,
format="WEBP", format="WEBP",
quality=quality, quality=quality,
@@ -172,9 +180,14 @@ def _encode_rgb_frame_to_jpeg(
width: int, width: int,
height: int, height: int,
quality: int, quality: int,
preview_max_width: int | None,
) -> bytes: ) -> bytes:
buffer = io.BytesIO() buffer = io.BytesIO()
Image.frombytes("RGB", (width, height), frame).save( image = _resize_preview_image(
Image.frombytes("RGB", (width, height), frame),
preview_max_width=preview_max_width,
)
image.save(
buffer, buffer,
format="JPEG", format="JPEG",
quality=quality, quality=quality,
@@ -183,6 +196,39 @@ def _encode_rgb_frame_to_jpeg(
return buffer.getvalue() return buffer.getvalue()
def _preview_dimensions(
*,
width: int,
height: int,
preview_max_width: int | None,
) -> tuple[int, int]:
if (
preview_max_width is None
or preview_max_width <= 0
or width <= preview_max_width
):
return width, height
preview_width = int(preview_max_width)
preview_height = max(1, round(height * preview_width / width))
return preview_width, preview_height
def _resize_preview_image(
image: Image.Image,
*,
preview_max_width: int | None,
) -> Image.Image:
width, height = image.size
preview_width, preview_height = _preview_dimensions(
width=width,
height=height,
preview_max_width=preview_max_width,
)
if (preview_width, preview_height) == image.size:
return image
return image.resize((preview_width, preview_height), Image.Resampling.BICUBIC)
def _pack_frame_batch_message( def _pack_frame_batch_message(
header: RealtimeFrameBatchHeader, header: RealtimeFrameBatchHeader,
payload: bytes, payload: bytes,
@@ -202,11 +248,12 @@ def _build_transport_payload(
metadata: dict[str, int | str], metadata: dict[str, int | str],
output_format: str | None, output_format: str | None,
transport_quality: int | None, transport_quality: int | None,
preview_max_width: int | None,
reference_frame: bytes | None, reference_frame: bytes | None,
event_id: int | None, event_id: int | None,
) -> _TransportPayload: ) -> _TransportPayload:
payload_content_type = content_type payload_content_type = content_type
payload_metadata: dict[str, int | str | bool] = {} payload_metadata: dict[str, int | str | bool | list[int]] = {}
raw_payload = b"" raw_payload = b""
if ( if (
@@ -215,24 +262,45 @@ def _build_transport_payload(
and transport_frames and transport_frames
): ):
if output_format == "webp": if output_format == "webp":
raw_payload = _encode_rgb_frame_to_webp( encoded_frames = [
transport_frames[0], _encode_rgb_frame_to_webp(
width=int(metadata["width"]), frame,
height=int(metadata["height"]), width=int(metadata["width"]),
quality=int(transport_quality or WEBP_DEFAULT_QUALITY), height=int(metadata["height"]),
) quality=int(transport_quality or WEBP_DEFAULT_QUALITY),
preview_max_width=preview_max_width,
)
for frame in transport_frames
]
payload_content_type = WEBP_FRAME_CONTENT_TYPE payload_content_type = WEBP_FRAME_CONTENT_TYPE
else: else:
raw_payload = _encode_rgb_frame_to_jpeg( encoded_frames = [
transport_frames[0], _encode_rgb_frame_to_jpeg(
width=int(metadata["width"]), frame,
height=int(metadata["height"]), width=int(metadata["width"]),
quality=int(transport_quality or JPEG_DEFAULT_QUALITY), height=int(metadata["height"]),
) quality=int(transport_quality or JPEG_DEFAULT_QUALITY),
preview_max_width=preview_max_width,
)
for frame in transport_frames
]
payload_content_type = JPEG_FRAME_CONTENT_TYPE payload_content_type = JPEG_FRAME_CONTENT_TYPE
raw_payload = b"".join(encoded_frames)
preview_width, preview_height = _preview_dimensions(
width=int(metadata["width"]),
height=int(metadata["height"]),
preview_max_width=preview_max_width,
)
payload_metadata = { payload_metadata = {
"format": output_format, "format": output_format,
"encoding": output_format, "encoding": output_format,
"source_width": int(metadata["width"]),
"source_height": int(metadata["height"]),
"preview_width": preview_width,
"preview_height": preview_height,
"width": preview_width,
"height": preview_height,
"payload_lengths": [len(frame) for frame in encoded_frames],
} }
elif ( elif (
output_format == RAW_LOSSLESS_OUTPUT_FORMAT output_format == RAW_LOSSLESS_OUTPUT_FORMAT
@@ -302,20 +370,18 @@ async def _build_encoded_preview_payloads(
metadata: dict[str, int | str], metadata: dict[str, int | str],
output_format: str, output_format: str,
transport_quality: int | None, transport_quality: int | None,
preview_max_width: int | None,
event_id: int | None, event_id: int | None,
) -> list[_TransportPayload]: ) -> list[_TransportPayload]:
return list( return list(
await asyncio.gather( await asyncio.gather(
*( *(
asyncio.to_thread( _build_encoded_preview_payload(
_build_transport_payload,
transport_frames, transport_frames,
content_type=content_type,
metadata=metadata, metadata=metadata,
output_format=output_format, output_format=output_format,
transport_quality=transport_quality, transport_quality=transport_quality,
reference_frame=None, preview_max_width=preview_max_width,
event_id=event_id,
) )
for transport_frames in split_batches for transport_frames in split_batches
) )
@@ -323,6 +389,73 @@ async def _build_encoded_preview_payloads(
) )
async def _build_encoded_preview_payload(
transport_frames: list[bytes],
*,
metadata: dict[str, int | str],
output_format: str,
transport_quality: int | None,
preview_max_width: int | None,
) -> _TransportPayload:
width = int(metadata["width"])
height = int(metadata["height"])
if output_format == "webp":
encoded_frames = list(
await asyncio.gather(
*(
asyncio.to_thread(
_encode_rgb_frame_to_webp,
frame,
width=width,
height=height,
quality=int(transport_quality or WEBP_DEFAULT_QUALITY),
preview_max_width=preview_max_width,
)
for frame in transport_frames
)
)
)
payload_content_type = WEBP_FRAME_CONTENT_TYPE
else:
encoded_frames = list(
await asyncio.gather(
*(
asyncio.to_thread(
_encode_rgb_frame_to_jpeg,
frame,
width=width,
height=height,
quality=int(transport_quality or JPEG_DEFAULT_QUALITY),
preview_max_width=preview_max_width,
)
for frame in transport_frames
)
)
)
payload_content_type = JPEG_FRAME_CONTENT_TYPE
preview_width, preview_height = _preview_dimensions(
width=width,
height=height,
preview_max_width=preview_max_width,
)
return _TransportPayload(
content_type=payload_content_type,
payload=b"".join(encoded_frames),
metadata={
"format": output_format,
"encoding": output_format,
"source_width": width,
"source_height": height,
"preview_width": preview_width,
"preview_height": preview_height,
"width": preview_width,
"height": preview_height,
"payload_lengths": [len(frame) for frame in encoded_frames],
},
)
class RawRGBRealtimeOutputAdapter: class RawRGBRealtimeOutputAdapter:
"""send raw RGB over WebSocket using lossless transport compression""" """send raw RGB over WebSocket using lossless transport compression"""
@@ -354,6 +487,7 @@ class RawRGBRealtimeOutputAdapter:
else {} else {}
) )
output_format = getattr(batch, "realtime_output_format", None) output_format = getattr(batch, "realtime_output_format", None)
preview_max_width = getattr(batch, "realtime_preview_max_width", None)
stats = await self._send_frame_batches( stats = await self._send_frame_batches(
ws, ws,
result.raw_frame_batches, result.raw_frame_batches,
@@ -364,6 +498,7 @@ class RawRGBRealtimeOutputAdapter:
frame_metadata=frame_metadata, frame_metadata=frame_metadata,
output_format=output_format, output_format=output_format,
transport_quality=getattr(batch, "output_compression", None), transport_quality=getattr(batch, "output_compression", None),
preview_max_width=preview_max_width,
) )
stats["frame_shape"] = _frame_shape_from_metadata(frame_metadata) stats["frame_shape"] = _frame_shape_from_metadata(frame_metadata)
return stats return stats
@@ -380,13 +515,14 @@ class RawRGBRealtimeOutputAdapter:
frame_metadata: dict[str, int | str] | None = None, frame_metadata: dict[str, int | str] | None = None,
output_format: str | None = None, output_format: str | None = None,
transport_quality: int | None = None, transport_quality: int | None = None,
preview_max_width: int | None = None,
) -> RealtimeFrameSendStats: ) -> RealtimeFrameSendStats:
chunk_index = chunk_index_start chunk_index = chunk_index_start
metadata = frame_metadata or {} metadata = frame_metadata or {}
stats = empty_frame_send_stats(content_type) stats = empty_frame_send_stats(content_type)
for frames in frame_batches: for frames in frame_batches:
split_batches = ( split_batches = (
[[frame] for frame in frames] [frames]
if _is_encoded_preview_transport( if _is_encoded_preview_transport(
content_type=content_type, content_type=content_type,
output_format=output_format, output_format=output_format,
@@ -410,6 +546,7 @@ class RawRGBRealtimeOutputAdapter:
metadata=metadata, metadata=metadata,
output_format=output_format, output_format=output_format,
transport_quality=transport_quality, transport_quality=transport_quality,
preview_max_width=preview_max_width,
event_id=event_id, event_id=event_id,
) )
stats["raw_payload_build_ms"] += timer.mark_ms() stats["raw_payload_build_ms"] += timer.mark_ms()
@@ -434,6 +571,7 @@ class RawRGBRealtimeOutputAdapter:
metadata=metadata, metadata=metadata,
output_format=output_format, output_format=output_format,
transport_quality=transport_quality, transport_quality=transport_quality,
preview_max_width=preview_max_width,
reference_frame=reference_frame, reference_frame=reference_frame,
event_id=event_id, event_id=event_id,
) )
@@ -444,6 +582,7 @@ class RawRGBRealtimeOutputAdapter:
metadata=metadata, metadata=metadata,
output_format=output_format, output_format=output_format,
transport_quality=transport_quality, transport_quality=transport_quality,
preview_max_width=preview_max_width,
reference_frame=reference_frame, reference_frame=reference_frame,
event_id=event_id, event_id=event_id,
) )
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_video_api.py
+70 -4
View File
@@ -41,7 +41,7 @@ if TYPE_CHECKING:
logger = init_logger(__name__) logger = init_logger(__name__)
router = APIRouter(prefix="/v1/realtime_video", tags=["realtime"]) router = APIRouter(prefix="/v1/realtime_video", tags=["realtime"])
_ACTIVE_SESSION_IDS: set[str] = set() _ACTIVE_SESSION_IDS: set[str] = set()
_ACTIVE_SESSION_WAIT_SECONDS = 15.0 _ACTIVE_SESSION_WAIT_SECONDS = 1.0
_ACTIVE_SESSION_WAIT_INTERVAL_SECONDS = 0.1 _ACTIVE_SESSION_WAIT_INTERVAL_SECONDS = 0.1
@@ -73,6 +73,7 @@ def _log_realtime_chunk_timing(
"realtime chunk timing: session_id=%s request_id=%s " "realtime chunk timing: session_id=%s request_id=%s "
"chunk_idx=%s event_id=%s condition_kinds=%s " "chunk_idx=%s event_id=%s condition_kinds=%s "
"request_prepare=%.2fms scheduler_forward=%.2fms " "request_prepare=%.2fms scheduler_forward=%.2fms "
"output_pace=%.2fms "
"header_pack=%.2fms " "header_pack=%.2fms "
"header_write=%.2fms raw_payload_build=%.2fms raw_write=%.2fms " "header_write=%.2fms raw_payload_build=%.2fms raw_write=%.2fms "
"ws_write=%.2fms chunk_total=%.2fms batches=%d frames=%d " "ws_write=%.2fms chunk_total=%.2fms batches=%d frames=%d "
@@ -84,6 +85,7 @@ def _log_realtime_chunk_timing(
sorted(batch.condition_inputs) if batch.condition_inputs else [], sorted(batch.condition_inputs) if batch.condition_inputs else [],
request_prepare_ms, request_prepare_ms,
scheduler_forward_ms, scheduler_forward_ms,
send_stats["pace_wait_ms"],
send_stats["header_pack_ms"], send_stats["header_pack_ms"],
send_stats["header_write_ms"], send_stats["header_write_ms"],
send_stats["raw_payload_build_ms"], send_stats["raw_payload_build_ms"],
@@ -119,6 +121,7 @@ async def _send_realtime_chunk_stats(
"event_id": getattr(batch, "realtime_event_id", None), "event_id": getattr(batch, "realtime_event_id", None),
"request_prepare_ms": _transport_ms(request_prepare_ms), "request_prepare_ms": _transport_ms(request_prepare_ms),
"scheduler_forward_ms": _transport_ms(scheduler_forward_ms), "scheduler_forward_ms": _transport_ms(scheduler_forward_ms),
"pace_wait_ms": _transport_ms(send_stats["pace_wait_ms"]),
"header_write_ms": _transport_ms(send_stats["header_write_ms"]), "header_write_ms": _transport_ms(send_stats["header_write_ms"]),
"raw_payload_build_ms": _transport_ms( "raw_payload_build_ms": _transport_ms(
send_stats["raw_payload_build_ms"] send_stats["raw_payload_build_ms"]
@@ -178,8 +181,8 @@ async def _generate_loop(ws: WebSocket, session: GenerateSession):
adapter.on_chunk_complete(session, result) adapter.on_chunk_complete(session, result)
if pending_send_task is not None: if pending_send_task is not None:
await pending_send_task await pending_send_task
pending_send_task = asyncio.create_task( if batch.realtime_output_pacing:
_send_output_and_log( await _send_output_and_log(
ws, ws,
session, session,
chunk, chunk,
@@ -189,7 +192,20 @@ async def _generate_loop(ws: WebSocket, session: GenerateSession):
scheduler_forward_ms, scheduler_forward_ms,
chunk_started, chunk_started,
) )
) pending_send_task = None
else:
pending_send_task = asyncio.create_task(
_send_output_and_log(
ws,
session,
chunk,
batch,
result,
request_prepare_ms,
scheduler_forward_ms,
chunk_started,
)
)
except asyncio.CancelledError: except asyncio.CancelledError:
if pending_send_task is not None: if pending_send_task is not None:
@@ -241,12 +257,14 @@ async def _send_output_and_log(
) -> RealtimeFrameSendStats: ) -> RealtimeFrameSendStats:
if session.adapter is None: if session.adapter is None:
raise ValueError("realtime adapter is not initialized") raise ValueError("realtime adapter is not initialized")
pace_wait_ms = await _wait_for_realtime_output_slot(session, batch, result)
send_stats = await session.adapter.send_output( send_stats = await session.adapter.send_output(
ws, ws,
session, session,
result, result,
batch, batch,
) )
send_stats["pace_wait_ms"] = pace_wait_ms
chunk_total_ms = (time.perf_counter() - chunk_started) * 1000 chunk_total_ms = (time.perf_counter() - chunk_started) * 1000
_log_realtime_chunk_timing( _log_realtime_chunk_timing(
session, session,
@@ -270,6 +288,54 @@ async def _send_output_and_log(
return send_stats return send_stats
def _result_num_frames(result) -> int:
if result.raw_frame_batches is None:
return 0
return sum(len(frames) for frames in result.raw_frame_batches)
def _output_pacing_fps(batch: "Req") -> float:
fps = float(batch.fps or 0)
if batch.enable_frame_interpolation:
fps *= 2 ** int(batch.frame_interpolation_exp or 1)
return fps
async def _wait_for_realtime_output_slot(
session: GenerateSession,
batch: "Req",
result,
) -> float:
if not batch.realtime_output_pacing:
return 0.0
frame_count = _result_num_frames(result)
output_fps = _output_pacing_fps(batch)
if frame_count <= 0 or output_fps <= 0:
return 0.0
now = time.perf_counter()
next_send_at = session.output_pace_next_send_at
if next_send_at is None:
next_send_at = now
if (
batch.realtime_event_id is not None
and batch.realtime_event_id != session.output_pace_last_event_id
):
next_send_at = min(next_send_at, now)
session.output_pace_last_event_id = batch.realtime_event_id
wait_s = max(0.0, next_send_at - now)
if wait_s > 0:
await asyncio.sleep(wait_s)
send_started_at = time.perf_counter()
session.output_pace_next_send_at = (
max(next_send_at, send_started_at) + frame_count / output_fps
)
return wait_s * 1000
async def _await_realtime_task(task: asyncio.Task | None) -> None: async def _await_realtime_task(task: asyncio.Task | None) -> None:
if task is None: if task is None:
return return
python/sglang/multimodal_gen/runtime/launch_server.py
+2 -4
View File
@@ -23,7 +23,7 @@ from sglang.multimodal_gen.runtime.server_args import (
) )
from sglang.multimodal_gen.runtime.utils.common import is_port_available from sglang.multimodal_gen.runtime.utils.common import is_port_available
from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger, logger from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger, logger
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info from sglang.multimodal_gen.runtime.utils.trace_wrapper import init_diffusion_tracing
def _find_available_port( def _find_available_port(
@@ -443,9 +443,7 @@ def _run_disagg_role_process(
def launch_http_server_only(server_args): def launch_http_server_only(server_args):
if server_args.enable_trace: init_diffusion_tracing(server_args, "DiffHTTPServer")
process_tracing_init(server_args.otlp_traces_endpoint, "sglang-diffusion")
trace_set_thread_info("DiffHTTPServer")
# set for endpoints to access global_server_args # set for endpoints to access global_server_args
set_global_server_args(server_args) set_global_server_args(server_args)
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
+6 -5
View File
@@ -72,8 +72,11 @@ from sglang.multimodal_gen.runtime.utils.realtime_video import (
RAW_RGB_CONTENT_TYPE, RAW_RGB_CONTENT_TYPE,
build_raw_rgb_frame_batches, build_raw_rgb_frame_batches,
) )
from sglang.multimodal_gen.runtime.utils.trace_wrapper import DiffStage, trace_slice from sglang.multimodal_gen.runtime.utils.trace_wrapper import (
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info DiffStage,
init_diffusion_tracing,
trace_slice,
)
from sglang.srt.utils.network import NetworkAddress from sglang.srt.utils.network import NetworkAddress
logger = init_logger(__name__) logger = init_logger(__name__)
@@ -989,9 +992,7 @@ def run_scheduler_process(
elif current_platform.is_musa(): elif current_platform.is_musa():
set_musa_arch() set_musa_arch()
if server_args.enable_trace: init_diffusion_tracing(server_args, f"DiffWorker_rank{rank}")
process_tracing_init(server_args.otlp_traces_endpoint, "sglang-diffusion")
trace_set_thread_info(f"DiffWorker_rank{rank}")
port_args = PortArgs.from_server_args(server_args) port_args = PortArgs.from_server_args(server_args)
python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py
+2
View File
@@ -211,6 +211,8 @@ class Req:
realtime_chunk_size: int | None = None realtime_chunk_size: int | None = None
realtime_event_id: int | None = None realtime_event_id: int | None = None
realtime_output_format: str | None = None realtime_output_format: str | None = None
realtime_preview_max_width: int | None = None
realtime_output_pacing: bool = False
realtime_causal_sink_size: int | None = None realtime_causal_sink_size: int | None = None
realtime_causal_kv_cache_num_frames: int | None = None realtime_causal_kv_cache_num_frames: int | None = None
# return websocket-friendly raw RGB frame bytes instead of rwa tensors # return websocket-friendly raw RGB frame bytes instead of rwa tensors
python/sglang/multimodal_gen/runtime/postprocess/realesrgan_upscaler.py
+22
View File
@@ -12,7 +12,9 @@ The ImageUpscaler wrapper and integration code are original work.
import math import math
import os import os
import time import time
from hashlib import sha256
from typing import Optional from typing import Optional
from urllib.parse import unquote, urlparse
import numpy as np import numpy as np
import torch import torch
@@ -689,6 +691,7 @@ def _resolve_model_path(model_path: str) -> str:
Accepts: Accepts:
- An existing local file path (pass-through). - An existing local file path (pass-through).
- An http(s) URL to a .pth file, downloaded into the local cache.
- A HuggingFace ``repo_id`` → downloads the default weight file - A HuggingFace ``repo_id`` → downloads the default weight file
(``RealESRGAN_x4.pth``). (``RealESRGAN_x4.pth``).
- A HuggingFace ``repo_id:filename`` → downloads *filename* from *repo_id*, - A HuggingFace ``repo_id:filename`` → downloads *filename* from *repo_id*,
@@ -702,6 +705,25 @@ def _resolve_model_path(model_path: str) -> str:
_RESOLVED_MODEL_PATH_CACHE[model_path] = model_path _RESOLVED_MODEL_PATH_CACHE[model_path] = model_path
return model_path return model_path
parsed_url = urlparse(model_path)
if parsed_url.scheme in ("http", "https"):
filename = (
os.path.basename(unquote(parsed_url.path)) or _DEFAULT_REALESRGAN_FILENAME
)
cache_dir = os.path.join(
os.path.expanduser("~"), ".cache", "sglang", "realesrgan"
)
os.makedirs(cache_dir, exist_ok=True)
cache_key = sha256(model_path.encode("utf-8")).hexdigest()[:12]
local_path = os.path.join(cache_dir, f"{cache_key}-{filename}")
if not os.path.isfile(local_path):
tmp_path = f"{local_path}.tmp"
logger.info("Downloading Real-ESRGAN weights from URL %s", model_path)
torch.hub.download_url_to_file(model_path, tmp_path, progress=False)
os.replace(tmp_path, local_path)
_RESOLVED_MODEL_PATH_CACHE[model_path] = local_path
return local_path
# Parse optional "repo_id:filename" syntax; fall back to default filename. # Parse optional "repo_id:filename" syntax; fall back to default filename.
if ":" in model_path and not model_path.startswith("/"): if ":" in model_path and not model_path.startswith("/"):
repo_id, filename = model_path.split(":", 1) repo_id, filename = model_path.split(":", 1)
python/sglang/multimodal_gen/runtime/postprocess/rife_interpolator.py
+55
View File
@@ -11,6 +11,7 @@ The FrameInterpolator wrapper and integration code are original work.
""" """
import os import os
import time
from typing import Optional from typing import Optional
import numpy as np import numpy as np
@@ -28,6 +29,7 @@ _DEFAULT_RIFE_HF_REPO = "elfgum/RIFE-4.22.lite"
# Module-level cache: model_path -> Model instance # Module-level cache: model_path -> Model instance
_MODEL_CACHE: dict[str, "Model"] = {} _MODEL_CACHE: dict[str, "Model"] = {}
_MAX_RIFE_BATCH_PAIRS = 16
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -392,6 +394,20 @@ class FrameInterpolator:
arr = t.squeeze(0).permute(1, 2, 0).clamp(0.0, 1.0).cpu().numpy() arr = t.squeeze(0).permute(1, 2, 0).clamp(0.0, 1.0).cpu().numpy()
return (arr * 255.0).astype(np.uint8) return (arr * 255.0).astype(np.uint8)
@staticmethod
def _frames_to_tensor(
frames: list[np.ndarray], device: torch.device
) -> torch.Tensor:
t = torch.from_numpy(np.stack(frames, axis=0))
t = t.permute(0, 3, 1, 2).contiguous().float() / 255.0
return t.to(device, non_blocking=True)
@staticmethod
def _tensor_to_frames(t: torch.Tensor) -> list[np.ndarray]:
arr = t.permute(0, 2, 3, 1).clamp(0.0, 1.0).cpu().numpy()
arr = (arr * 255.0).astype(np.uint8)
return [arr[i] for i in range(arr.shape[0])]
def _make_inference( def _make_inference(
self, model: Model, I0: torch.Tensor, I1: torch.Tensor, n: int, scale: float self, model: Model, I0: torch.Tensor, I1: torch.Tensor, n: int, scale: float
) -> list[torch.Tensor]: ) -> list[torch.Tensor]:
@@ -409,6 +425,30 @@ class FrameInterpolator:
+ self._make_inference(model, mid, I1, n // 2, scale) + self._make_inference(model, mid, I1, n // 2, scale)
) )
def _interpolate_2x_batched(
self, model: Model, frames: list[np.ndarray], scale: float
) -> list[np.ndarray]:
device = model.device()
source = self._frames_to_tensor(frames, device)
intermediate_frames: list[np.ndarray] = []
with torch.inference_mode():
for start in range(0, len(frames) - 1, _MAX_RIFE_BATCH_PAIRS):
end = min(start + _MAX_RIFE_BATCH_PAIRS, len(frames) - 1)
mids = model.inference(
source[start:end],
source[start + 1 : end + 1],
scale=scale,
)
intermediate_frames.extend(self._tensor_to_frames(mids))
result: list[np.ndarray] = []
for i, mid in enumerate(intermediate_frames):
result.append(frames[i])
result.append(mid)
result.append(frames[-1])
return result
def interpolate( def interpolate(
self, self,
frames: list[np.ndarray], frames: list[np.ndarray],
@@ -436,6 +476,16 @@ class FrameInterpolator:
device = model.device() device = model.device()
n_intermediate = 2**exp // 2 # intermediates per adjacent pair n_intermediate = 2**exp // 2 # intermediates per adjacent pair
start_time = time.perf_counter()
if n_intermediate == 1:
result = self._interpolate_2x_batched(model, frames, scale)
logger.info(
"RIFE batched interpolation completed in %.3f seconds for %d frames",
time.perf_counter() - start_time,
len(frames),
)
return result, 2**exp
result: list[np.ndarray] = [] result: list[np.ndarray] = []
for i in range(len(frames) - 1): for i in range(len(frames) - 1):
@@ -452,6 +502,11 @@ class FrameInterpolator:
result.append(frames[-1]) result.append(frames[-1])
multiplier = 2**exp multiplier = 2**exp
logger.info(
"RIFE interpolation completed in %.3f seconds for %d frames",
time.perf_counter() - start_time,
len(frames),
)
return result, multiplier return result, multiplier
python/sglang/multimodal_gen/runtime/utils/trace_wrapper.py
+35
View File
@@ -10,6 +10,8 @@ from __future__ import annotations
from contextlib import contextmanager from contextlib import contextmanager
from dataclasses import dataclass from dataclasses import dataclass
DIFFUSION_TRACE_MODULE = "diffusion"
@dataclass(frozen=True) @dataclass(frozen=True)
class DiffStageConfig: class DiffStageConfig:
@@ -26,6 +28,39 @@ class DiffStage:
GPU_FORWARD = DiffStageConfig("gpu_forward", level=2) GPU_FORWARD = DiffStageConfig("gpu_forward", level=2)
def init_diffusion_tracing(server_args, thread_label: str):
if not server_args.enable_trace:
return
from types import SimpleNamespace
from sglang.srt import server_args as srt_server_args_module
from sglang.srt.observability.trace import (
process_tracing_init,
trace_set_thread_info,
)
from sglang.srt.server_args import set_global_server_args_for_scheduler
# srt owns TraceReqContext and filters spans through its global trace_modules
try:
srt_server_args = srt_server_args_module.get_global_server_args()
except ValueError:
srt_server_args = SimpleNamespace(trace_modules=DIFFUSION_TRACE_MODULE)
set_global_server_args_for_scheduler(srt_server_args)
trace_modules = [
module.strip()
for module in getattr(srt_server_args, "trace_modules", "").split(",")
if module.strip()
]
if DIFFUSION_TRACE_MODULE not in trace_modules:
trace_modules.append(DIFFUSION_TRACE_MODULE)
srt_server_args.trace_modules = ",".join(trace_modules)
process_tracing_init(server_args.otlp_traces_endpoint, "sglang-diffusion")
trace_set_thread_info(thread_label)
@contextmanager @contextmanager
def trace_req(trace_ctx): def trace_req(trace_ctx):
"""Ensure ``trace_req_finish()`` is called when a request scope exits. """Ensure ``trace_req_finish()`` is called when a request scope exits.
python/sglang/multimodal_gen/test/unit/test_disagg_trace.py
+58 -38
View File
@@ -19,6 +19,8 @@ from __future__ import annotations
import json import json
import unittest import unittest
from contextlib import contextmanager
from types import SimpleNamespace
import torch import torch
@@ -28,8 +30,10 @@ from sglang.multimodal_gen.runtime.disaggregation.scheduler_mixin import (
) )
from sglang.multimodal_gen.runtime.disaggregation.transport.codec import pack_tensors from sglang.multimodal_gen.runtime.disaggregation.transport.codec import pack_tensors
from sglang.multimodal_gen.runtime.pipelines_core import Req from sglang.multimodal_gen.runtime.pipelines_core import Req
from sglang.srt import server_args as srt_server_args_module
from sglang.srt.observability import trace as srt_trace from sglang.srt.observability import trace as srt_trace
from sglang.srt.observability.trace import TraceNullContext, TraceReqContext from sglang.srt.observability.trace import TraceNullContext, TraceReqContext
from sglang.srt.server_args import set_global_server_args_for_scheduler
try: try:
from opentelemetry import propagate as otel_propagate from opentelemetry import propagate as otel_propagate
@@ -56,6 +60,16 @@ def _enable_minimal_otel() -> None:
srt_trace.trace_set_thread_info("TestThread") srt_trace.trace_set_thread_info("TestThread")
@contextmanager
def _srt_trace_server_args():
prev_server_args = srt_server_args_module._global_server_args
set_global_server_args_for_scheduler(SimpleNamespace(trace_modules="request"))
try:
yield
finally:
srt_server_args_module._global_server_args = prev_server_args
def _traceparent_from(ctx) -> str | None: def _traceparent_from(ctx) -> str | None:
"""Re-inject a W3C carrier from an OTel Context and return the traceparent. """Re-inject a W3C carrier from an OTel Context and return the traceparent.
@@ -133,35 +147,37 @@ class TestDisaggTracePropagation(unittest.TestCase):
TraceReqContext with an OTel Context (not the raw dict).""" TraceReqContext with an OTel Context (not the raw dict)."""
_enable_minimal_otel() _enable_minimal_otel()
ctx = TraceReqContext(rid="test-on", role="server", module_name="request") with _srt_trace_server_args():
ctx.trace_req_start() ctx = TraceReqContext(rid="test-on", role="server", module_name="request")
self.assertTrue(ctx.tracing_enable) ctx.trace_req_start()
self.assertFalse(ctx.is_copy) self.assertTrue(ctx.tracing_enable)
self.assertFalse(ctx.is_copy)
req = Req(request_id="test-on", prompt="x") req = Req(request_id="test-on", prompt="x")
req.trace_ctx = ctx req.trace_ctx = ctx
_, scalar_fields = extract_transfer_fields(req) _, scalar_fields = extract_transfer_fields(req)
self.assertNotIn("trace_ctx", scalar_fields) self.assertNotIn("trace_ctx", scalar_fields)
self.assertIn("_trace_state", scalar_fields) self.assertIn("_trace_state", scalar_fields)
state = scalar_fields["_trace_state"] state = scalar_fields["_trace_state"]
self.assertTrue(state.get("tracing_enable")) self.assertTrue(state.get("tracing_enable"))
# W3C carrier must be present so downstream roles can nest spans. # W3C carrier must be present so downstream roles can nest spans.
self.assertIn("traceparent", state.get("root_span_context", {})) self.assertIn("traceparent", state.get("root_span_context", {}))
decoded = _roundtrip_scalar_fields(scalar_fields) decoded = _roundtrip_scalar_fields(scalar_fields)
self.assertEqual(decoded["_trace_state"], state) self.assertEqual(decoded["_trace_state"], state)
rebuilt = object.__new__(TraceReqContext) rebuilt = object.__new__(TraceReqContext)
rebuilt.__setstate__(decoded["_trace_state"]) rebuilt.__setstate__(decoded["_trace_state"])
self.assertTrue(rebuilt.tracing_enable) self.assertTrue(rebuilt.tracing_enable)
self.assertTrue(rebuilt.is_copy) self.assertTrue(rebuilt.is_copy)
# The sender's traceparent must survive into the rebuilt Context so # The sender's traceparent must survive into the rebuilt Context so
# downstream role spans nest under the original trace_id. # downstream role spans nest under the original trace_id.
self.assertEqual( self.assertEqual(
_traceparent_from(rebuilt.root_span_context), _traceparent_from(rebuilt.root_span_context),
state["root_span_context"]["traceparent"], state["root_span_context"]["traceparent"],
) )
ctx.trace_req_finish()
@unittest.skipUnless(_OTEL_AVAILABLE, "opentelemetry SDK not installed") @unittest.skipUnless(_OTEL_AVAILABLE, "opentelemetry SDK not installed")
def test_build_disagg_req_installs_rebuilt_ctx(self): def test_build_disagg_req_installs_rebuilt_ctx(self):
@@ -170,22 +186,26 @@ class TestDisaggTracePropagation(unittest.TestCase):
the Req as a stray attribute.""" the Req as a stray attribute."""
_enable_minimal_otel() _enable_minimal_otel()
ctx = TraceReqContext(rid="test-brq", role="server", module_name="request") with _srt_trace_server_args():
ctx.trace_req_start() ctx = TraceReqContext(rid="test-brq", role="server", module_name="request")
ctx.trace_req_start()
req = Req(request_id="test-brq", prompt="x") req = Req(request_id="test-brq", prompt="x")
req.trace_ctx = ctx req.trace_ctx = ctx
_, scalar_fields = extract_transfer_fields(req) _, scalar_fields = extract_transfer_fields(req)
self.assertIn("_trace_state", scalar_fields) self.assertIn("_trace_state", scalar_fields)
# _build_disagg_req is an instance method but its body does not touch # _build_disagg_req is an instance method but its body does not
# ``self``; call via __func__ to avoid needing a real Scheduler. # touch ``self``; call via __func__ to avoid needing a real Scheduler.
rebuilt = SchedulerDisaggMixin._build_disagg_req(None, dict(scalar_fields), {}) rebuilt = SchedulerDisaggMixin._build_disagg_req(
None, dict(scalar_fields), {}
)
self.assertIsInstance(rebuilt.trace_ctx, TraceReqContext) self.assertIsInstance(rebuilt.trace_ctx, TraceReqContext)
self.assertTrue(rebuilt.trace_ctx.tracing_enable) self.assertTrue(rebuilt.trace_ctx.tracing_enable)
self.assertTrue(rebuilt.trace_ctx.is_copy) self.assertTrue(rebuilt.trace_ctx.is_copy)
self.assertFalse(hasattr(rebuilt, "_trace_state")) self.assertFalse(hasattr(rebuilt, "_trace_state"))
ctx.trace_req_finish()
@unittest.skipUnless(_OTEL_AVAILABLE, "opentelemetry SDK not installed") @unittest.skipUnless(_OTEL_AVAILABLE, "opentelemetry SDK not installed")
def test_build_disagg_req_falls_back_when_tracing_off(self): def test_build_disagg_req_falls_back_when_tracing_off(self):