feat: replace per-conversation ZMQ fallback with dataset client metadata request

src/aiperf/common/enums/enums.py

@@ -213,6 +213,8 @@ class MessageType(CaseInsensitiveStrEnum):
     CREDIT_PHASE_START = "credit_phase_start"
     CREDIT_PHASES_CONFIGURED = "credit_phases_configured"
     CREDITS_COMPLETE = "credits_complete"
+    DATASET_CLIENT_REQUEST = "dataset_client_request"
+    DATASET_CLIENT_RESPONSE = "dataset_client_response"
     DATASET_CONFIGURED_NOTIFICATION = "dataset_configured_notification"
     ERROR = "error"
     HEARTBEAT = "heartbeat"

src/aiperf/common/messages/__init__.py

@@ -33,6 +33,8 @@
     ConversationResponseMessage,
     ConversationTurnRequestMessage,
     ConversationTurnResponseMessage,
+    DatasetClientRequestMessage,
+    DatasetClientResponseMessage,
     DatasetConfiguredNotification,
 )
 from aiperf.common.messages.inference_messages import (
@@ -87,6 +89,8 @@
     "ConversationResponseMessage",
     "ConversationTurnRequestMessage",
     "ConversationTurnResponseMessage",
+    "DatasetClientRequestMessage",
+    "DatasetClientResponseMessage",
     "DatasetConfiguredNotification",
     "ErrorMessage",
     "HeartbeatMessage",

src/aiperf/common/messages/dataset_messages.py

@@ -16,6 +16,31 @@
 from aiperf.common.types import MessageTypeT
 
 
+class DatasetClientRequestMessage(BaseServiceMessage):
+    """Request the dataset client metadata from the DatasetManager."""
+
+    message_type: MessageTypeT = MessageType.DATASET_CLIENT_REQUEST
+
+
+class DatasetClientResponseMessage(BaseServiceMessage):
+    """Response containing the dataset client metadata for connection."""
+
+    message_type: MessageTypeT = MessageType.DATASET_CLIENT_RESPONSE
+
+    client_metadata: SerializeAsAny[DatasetClientMetadata] = Field(
+        ...,
+        description="Client access metadata (e.g., mmap file paths) for reading the dataset.",
+    )
+
+    @field_validator("client_metadata", mode="before")
+    @classmethod
+    def route_client_metadata(cls, v: Any) -> DatasetClientMetadata:
+        """Route nested AutoRoutedModel field to correct subclass."""
+        if isinstance(v, dict):
+            return DatasetClientMetadata.from_json(v)
+        return v
+
+
 class ConversationRequestMessage(BaseServiceMessage):
     """Message to request a full conversation by ID."""
 

src/aiperf/dataset/dataset_manager.py

@@ -25,6 +25,8 @@
     ConversationResponseMessage,
     ConversationTurnRequestMessage,
     ConversationTurnResponseMessage,
+    DatasetClientRequestMessage,
+    DatasetClientResponseMessage,
     DatasetConfiguredNotification,
     ProfileConfigureCommand,
 )
@@ -464,6 +466,22 @@ async def _handle_conversation_turn_request(
             turn=turn,
         )
 
+    @on_request(MessageType.DATASET_CLIENT_REQUEST)
+    async def _handle_dataset_client_request(
+        self, message: DatasetClientRequestMessage
+    ) -> DatasetClientResponseMessage:
+        """Return dataset client metadata so the requester can initialize its own client."""
+        self.debug(lambda: f"Handling dataset client request from {message.service_id}")
+
+        await self._wait_for_dataset_configuration()
+
+        client_metadata = self._backing_store.get_client_metadata()
+        return DatasetClientResponseMessage(
+            service_id=self.service_id,
+            request_id=message.request_id,
+            client_metadata=client_metadata,
+        )
+
     async def _wait_for_dataset_configuration(self) -> None:
         """Wait for the dataset to be configured if it is not already."""
         if not self.dataset_configured.is_set():

src/aiperf/workers/worker.py

@@ -26,12 +26,13 @@
     WorkerHealthMessage,
 )
 from aiperf.common.messages.dataset_messages import (
-    ConversationRequestMessage,
-    ConversationResponseMessage,
+    DatasetClientRequestMessage,
+    DatasetClientResponseMessage,
 )
 from aiperf.common.mixins import ProcessHealthMixin
 from aiperf.common.models import (
     Conversation,
+    DatasetClientMetadata,
     ErrorDetails,
     ModelEndpointInfo,
     ReasoningResponseData,
@@ -197,10 +198,9 @@ def __init__(
             or self.user_config.loadgen.warmup_prefill_concurrency is not None
         )
 
-        # Only used as a fallback when dataset client is not initialized
-        # or was not available when the credit was dropped. Must be created here
-        # so it can be attached to the worker lifecycle.
-        self.conversation_request_client: RequestClientProtocol = (
+        # Used to request dataset client metadata from DatasetManager when
+        # the dataset client is not yet initialized (race condition at startup).
+        self._dataset_client_request_client: RequestClientProtocol = (
             self.comms.create_request_client(
                 address=CommAddress.DATASET_MANAGER_PROXY_FRONTEND,
                 bind=False,
@@ -221,24 +221,30 @@ async def _send_worker_ready_message(self) -> None:
 
     @on_message(MessageType.DATASET_CONFIGURED_NOTIFICATION)
     async def _on_dataset_configured(self, msg: DatasetConfiguredNotification) -> None:
-        """Initialize dataset client when configuration is received.
+        """Initialize dataset client when configuration is received."""
+        await self._initialize_dataset_client(msg.client_metadata)
+        self.debug(
+            lambda: (
+                f"Dataset client initialized: type={msg.client_metadata.client_type}"
+            )
+        )
+
+    async def _initialize_dataset_client(
+        self, client_metadata: DatasetClientMetadata
+    ) -> None:
+        """Create and initialize the dataset client from metadata.
 
         Uses factory pattern to dynamically create the appropriate client.
         The factory auto-extracts client_type from client_metadata, leveraging
         the discriminated union pattern for type-safe routing. This allows new
         storage backends (S3, Redis, etc.) to work without modifying Worker code.
         """
         ClientStoreClass = plugins.get_class(
-            PluginType.DATASET_CLIENT_STORE, msg.client_metadata.client_type
+            PluginType.DATASET_CLIENT_STORE, client_metadata.client_type
         )
-        self._dataset_client = ClientStoreClass(client_metadata=msg.client_metadata)
+        self._dataset_client = ClientStoreClass(client_metadata=client_metadata)
         await self._dataset_client.initialize()
         self._dataset_configured_event.set()
-        self.debug(
-            lambda: (
-                f"Dataset client initialized: type={msg.client_metadata.client_type}"
-            )
-        )
 
     @on_stop
     async def _send_worker_shutdown_message(self) -> None:
@@ -576,52 +582,32 @@ async def _retrieve_conversation(
         elif self.stop_requested:
             raise asyncio.CancelledError("Stop requested while retrieving conversation")
 
-        return await self._request_conversation_from_dataset_manager(
-            conversation_id, credit_context
-        )
+        await self._request_dataset_client_from_dataset_manager()
+        return await self._dataset_client.get_conversation(conversation_id)
 
-    async def _request_conversation_from_dataset_manager(
-        self, conversation_id: str, credit_context: CreditContext
-    ) -> Conversation:
-        """Fallback: Request from DatasetManager via ZMQ"""
-        conversation_response: (
-            ConversationResponseMessage | ErrorMessage
-        ) = await self.conversation_request_client.request(
-            ConversationRequestMessage(
-                service_id=self.service_id,
-                conversation_id=conversation_id,
-                credit_phase=credit_context.credit.phase,
-            )
+    async def _request_dataset_client_from_dataset_manager(self) -> None:
+        """Fallback: Request dataset client metadata from DatasetManager and initialize client."""
+        self.info(
+            "Dataset client not available, requesting metadata from DatasetManager"
         )
-        if self.is_trace_enabled:
-            self.trace(f"Received response message: {conversation_response}")
-
-        # Check for error in conversation response
-        if isinstance(conversation_response, ErrorMessage):
-            error = conversation_response.error
-            await self._send_inference_result_message(
-                RequestRecord(
-                    request_info=RequestInfo(
-                        model_endpoint=self.model_endpoint,
-                        conversation_id=conversation_id,
-                        turn_index=0,
-                        turns=[],
-                        credit_num=credit_context.credit.id,
-                        credit_phase=credit_context.credit.phase,
-                        x_request_id=str(uuid.uuid4()),
-                        x_correlation_id=credit_context.credit.x_correlation_id,
-                        drop_perf_ns=credit_context.drop_perf_ns,
-                    ),
-                    model_name=self.model_endpoint.primary_model_name,
-                    timestamp_ns=time.time_ns(),
-                    start_perf_ns=time.perf_counter_ns(),
-                    end_perf_ns=time.perf_counter_ns(),
-                    error=error,
-                )
+        response: (
+            DatasetClientResponseMessage | ErrorMessage
+        ) = await self._dataset_client_request_client.request(
+            DatasetClientRequestMessage(service_id=self.service_id)
+        )
+
+        if isinstance(response, ErrorMessage):
+            raise ValueError(
+                f"Failed to retrieve dataset client metadata: {response.error}"
             )
-            raise ValueError(f"Failed to retrieve conversation response: {error}")
 
-        return conversation_response.conversation
+        await self._initialize_dataset_client(response.client_metadata)
+        self.info(
+            lambda: (
+                f"Dataset client initialized via fallback request: "
+                f"type={response.client_metadata.client_type}"
+            )
+        )
 
     async def _process_response(self, record: RequestRecord) -> Turn | None:
         """Extract assistant response from RequestRecord and convert to Turn for session.

tests/unit/workers/test_worker.py

@@ -299,12 +299,19 @@ async def test_raises_cancelled_error_when_stop_requested_and_no_client(
     async def test_falls_back_to_dataset_manager_when_no_client_and_not_stopping(
         self, monkeypatch, mock_worker, sample_credit_context
     ):
-        """When _dataset_client is None and not stopping, should request from DatasetManager."""
+        """When _dataset_client is None and not stopping, should request client metadata from DatasetManager."""
         mock_worker._dataset_client = None
         expected_conversation = Conversation(session_id="test-conv-123", turns=[])
-        mock_fallback = AsyncMock(return_value=expected_conversation)
+
+        async def mock_request_client(self_worker):
+            mock_client = AsyncMock()
+            mock_client.get_conversation = AsyncMock(return_value=expected_conversation)
+            self_worker._dataset_client = mock_client
+
         monkeypatch.setattr(
-            mock_worker, "_request_conversation_from_dataset_manager", mock_fallback
+            mock_worker,
+            "_request_dataset_client_from_dataset_manager",
+            lambda: mock_request_client(mock_worker),
         )
 
         result = await mock_worker._retrieve_conversation(
@@ -313,4 +320,3 @@ async def test_falls_back_to_dataset_manager_when_no_client_and_not_stopping(
         )
 
         assert result == expected_conversation
-        mock_fallback.assert_called_once_with("test-conv-123", sample_credit_context)