Merge branch 'main' into refactor/fc-outcome-from-error

Author: LukeAVanDrie

pkg/kvevents/engineadapter/common.go

@@ -33,6 +33,13 @@ const (
 	eventTagAllBlocksCleared = "AllBlocksCleared"
 )
 
+type msgpackEventBatch struct {
+	_                struct{} `msgpack:",array"`
+	TS               float64
+	Events           []msgpack.RawMessage
+	DataParallelRank *int `msgpack:",omitempty"`
+}
+
 // parseTopic extracts pod ID and model name from the topic format "kv@<pod-id>@<model-name>".
 //
 //nolint:gocritic // unnamedResult: named returns conflict with nonamedreturns linter

pkg/kvevents/engineadapter/sglang_adapter.go

@@ -71,7 +71,7 @@ func (s *SGLangAdapter) ShardingKey(msg *kvevents.RawMessage) string {
 func (s *SGLangAdapter) ParseMessage(msg *kvevents.RawMessage) (string, string, kvevents.EventBatch, error) {
 	podID, modelName := parseTopic(msg.Topic)
 
-	var batch msgpackSGLangEventBatch
+	var batch msgpackEventBatch
 	if err := msgpack.Unmarshal(msg.Payload, &batch); err != nil {
 		return "", "", kvevents.EventBatch{}, fmt.Errorf("failed to decode SGLang event batch: %w", err)
 	}
@@ -86,22 +86,14 @@ func (s *SGLangAdapter) ParseMessage(msg *kvevents.RawMessage) (string, string,
 	}
 
 	eventBatch := kvevents.EventBatch{
-		Timestamp: batch.TS,
-		Events:    genericEvents,
+		Timestamp:        batch.TS,
+		Events:           genericEvents,
+		DataParallelRank: batch.DataParallelRank,
 	}
 
 	return podID, modelName, eventBatch, nil
 }
 
-// SGLang msgpack event structures.
-// These match the vLLM wire format (SGLang uses the same positional encoding).
-type msgpackSGLangEventBatch struct {
-	_                struct{} `msgpack:",array"`
-	TS               float64
-	Events           []msgpack.RawMessage
-	DataParallelRank *int `msgpack:",omitempty"`
-}
-
 type msgpackSGLangBlockStoredEvent struct {
 	_               struct{} `msgpack:",array"`
 	Tag             string

pkg/kvevents/engineadapter/sglang_adapter_test.go

@@ -50,7 +50,7 @@ func TestSGLangParseMessage_Valid(t *testing.T) {
 	batch := []any{
 		1234567890.0,
 		[]any{blockStoredEvent},
-		nil,
+		3,
 	}
 	payload, err := msgpack.Marshal(batch)
 	require.NoError(t, err)
@@ -65,6 +65,8 @@ func TestSGLangParseMessage_Valid(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, "pod-1", podID)
 	assert.Equal(t, "llama-2-7b", modelName)
+	require.NotNil(t, eventBatch.DataParallelRank)
+	assert.Equal(t, 3, *eventBatch.DataParallelRank)
 	assert.Len(t, eventBatch.Events, 1)
 
 	blockStored, ok := eventBatch.Events[0].(*kvevents.BlockStoredEvent)

pkg/kvevents/engineadapter/vllm_adapter.go

@@ -64,7 +64,7 @@ func (v *VLLMAdapter) ShardingKey(msg *kvevents.RawMessage) string {
 func (v *VLLMAdapter) ParseMessage(msg *kvevents.RawMessage) (string, string, kvevents.EventBatch, error) {
 	podID, modelName := parseTopic(msg.Topic)
 
-	var vllmBatch msgpackVLLMEventBatch
+	var vllmBatch msgpackEventBatch
 	if err := msgpack.Unmarshal(msg.Payload, &vllmBatch); err != nil {
 		return "", "", kvevents.EventBatch{}, fmt.Errorf("failed to decode vLLM event batch: %w", err)
 	}
@@ -79,22 +79,14 @@ func (v *VLLMAdapter) ParseMessage(msg *kvevents.RawMessage) (string, string, kv
 	}
 
 	batch := kvevents.EventBatch{
-		Timestamp: vllmBatch.TS,
-		Events:    genericEvents,
+		Timestamp:        vllmBatch.TS,
+		Events:           genericEvents,
+		DataParallelRank: vllmBatch.DataParallelRank,
 	}
 
 	return podID, modelName, batch, nil
 }
 
-// vLLM msgpack event batch structure.
-// This struct uses array encoding to match vLLM's msgspec array_like=True format.
-type msgpackVLLMEventBatch struct {
-	_                struct{} `msgpack:",array"`
-	TS               float64
-	Events           []msgpack.RawMessage
-	DataParallelRank *int `msgpack:",omitempty"`
-}
-
 // decodeVLLMEvent decodes a single vLLM event from msgpack bytes and dispatches
 // it to the matching converter. Map-encoded events are first normalized to the
 // positional []any layout the converters consume.

pkg/kvevents/engineadapter/vllm_adapter_test.go

@@ -51,7 +51,7 @@ func TestVLLMParseMessage_Valid(t *testing.T) {
 	batch := []any{
 		1234567890.0,
 		[]any{blockStoredEvent},
-		nil,
+		3,
 	}
 	payload, err := msgpack.Marshal(batch)
 	require.NoError(t, err)
@@ -66,6 +66,8 @@ func TestVLLMParseMessage_Valid(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, "pod-1", podID)
 	assert.Equal(t, "llama-2-7b", modelName)
+	require.NotNil(t, eventBatch.DataParallelRank)
+	assert.Equal(t, 3, *eventBatch.DataParallelRank)
 	assert.Len(t, eventBatch.Events, 1)
 
 	blockStored, ok := eventBatch.Events[0].(*kvevents.BlockStoredEvent)

pkg/kvevents/events.go

@@ -50,8 +50,9 @@ type GenericEvent interface {
 
 // EventBatch represents a batch of generic events from an inference engine.
 type EventBatch struct {
-	Timestamp float64
-	Events    []GenericEvent
+	Timestamp        float64
+	Events           []GenericEvent
+	DataParallelRank *int
 }
 
 // RawMessage holds the raw transport-level data from a received pub/sub message.