feat(evaluation): add DatasetClient and dataset management service provider (#491)

* feat(evaluation): add DatasetClient and ServiceDatasetProvider

Add Dataset Management SDK support with:
- DatasetClient: pass-through client for all 11 dataset APIs
  (create/get/list/update/delete datasets, versions, examples)
  with 6 _and_wait helpers for async operations
- ServiceDatasetProvider: fetches datasets from the service and
  returns SDK Dataset objects compatible with OnDemandEvaluationDatasetRunner
  and BatchEvaluationRunner
- Unit tests (20 tests) and integration tests (17 tests, verified against
  live AWS)

* refactor(evaluation): use downloadUrl for ServiceDatasetProvider

Switch ServiceDatasetProvider from list_dataset_examples pagination
to downloading the JSONL file via the presigned downloadUrl from
GetDataset. Single HTTP request is simpler and faster for large datasets.

* test(evaluation): add runner integ test helpers and skipped runner tests

- helpers.py: get_or_create_agent_runtime(), make_agent_invoker()
  with retry logic and warmup for cold start handling
- test_runners_with_service_dataset.py: OnDemandRunner + ServiceDatasetProvider
  end-to-end test (skipped until a working deployed agent is available —
  current account has 30s init timeout that prevents cold starts)

* test(evaluation): runner integ test passes with real agent

Update test_runners_with_service_dataset.py to use env-var config:
- INTEG_AGENT_RUNTIME_ARN: skips if not set
- BEDROCK_TEST_REGION: region matching the agent
- Verified end-to-end: ServiceDatasetProvider → OnDemandRunner → real agent invocation → COMPLETED

* refactor: extract region from agent ARN, no separate BEDROCK_TEST_REGION needed

* fix(evaluation): address PR review feedback

- ServiceDatasetProvider: accept client in __init__ (eliminates region_name)
- ServiceDatasetProvider: validate schemaType against supported runner schemas
- ServiceDatasetProvider: proper error message on download failure
- Remove helpers.py (not needed)
- Add unit tests for unsupported schema and download failure cases

* fix(evaluation): address second round of PR feedback

- ServiceDatasetProvider: import DatasetClient at top, default in __init__
- ScenarioExecutor: add schema_type field, override in Predefined/Simulated
- ServiceDatasetProvider: collect supported schemas dynamically from executors
- delete_dataset_and_wait: add DELETE_FAILED as failed status

* style: fix lint and formatting

* style: fix remaining line-too-long issues

* fix: remove unused variable assignment (F841)

* refactor: move schema_type to Scenario classes, _parse_scenario to module level

- Add schema_type field to Scenario base, PredefinedScenario, SimulatedScenario
- Remove schema_type from ScenarioExecutor (doesn't belong there)
- Move _parse_scenario from FileDatasetProvider to module-level function
- SUPPORTED_SCHEMA_TYPES derived from Scenario classes directly

* fix: remove unused _DATASET_FAILED_STATUSES constant

* fix(evaluation): address reviewer findings

- Add timeout=60 to requests.get() for presigned URL download (#1)
- Use r.content.decode("utf-8") instead of r.text for explicit encoding (#3)
- Replace model_fields introspection with plain constant set (#8)
- Guard __getattr__ against recursion when _cp_client not initialized (#5)
- Remove dead _mock_client function from tests (#11)

* fix(evaluation): stream JSONL download and consolidate test fixtures

- Stream JSONL via iter_lines() instead of loading entire file into memory (#2)
- Consolidate repetitive test mock setup with pytest fixtures (#12)

* refactor(evaluation): rename ServiceDatasetProvider to DatasetManagementServiceProvider

Addresses PR review feedback: the name makes explicit which service
the provider loads datasets from (Dataset Management Service).

* fix: Removed delete_dataset_version api from allowlist

* feat(evaluation): support JSONL files in FileDatasetProvider

Dispatch on file extension: paths ending in .jsonl are read line-by-line
(one scenario per line). All other paths keep the existing
{"scenarios": [...]} JSON shape.

Adds 8 unit tests covering predefined/simulated/mixed JSONL content,
blank-line tolerance, malformed lines, and extension dispatch.

* fix(evaluation): handle version-specific deletes in delete_dataset_and_wait

A version-specific delete (datasetVersion provided) does not remove the
dataset itself — it transitions the dataset to UPDATING and back to ACTIVE.
The previous waiter polled for ResourceNotFoundException and timed out.

Branch on whether datasetVersion is passed:
- Without datasetVersion: poll until ResourceNotFoundException (DELETE_FAILED)
- With datasetVersion: poll until ACTIVE (UPDATE_FAILED), return dataset dict

Add unit tests for both version-delete paths (success + UPDATE_FAILED)
and an integ test that creates two versions, deletes the oldest via
delete_dataset_and_wait, and verifies the dataset stays ACTIVE.

Author: jariy17

src/bedrock_agentcore/evaluation/__init__.py

@@ -6,6 +6,7 @@
     EvaluatorOutput,
     custom_code_based_evaluator,
 )
+from bedrock_agentcore.evaluation.dataset_client import DatasetClient
 from bedrock_agentcore.evaluation.runner.batch.batch_evaluation_models import (
     BatchEvaluationResult,
     BatchEvaluationRunConfig,
@@ -21,6 +22,7 @@
     BatchEvaluationRunner,
 )
 from bedrock_agentcore.evaluation.runner.dataset_providers import (
+    DatasetManagementServiceProvider,
     DatasetProvider,
     FileDatasetProvider,
 )
@@ -77,6 +79,7 @@
     "AgentInvokerOutput",
     "CloudWatchAgentSpanCollector",
     "Dataset",
+    "DatasetClient",
     "DatasetProvider",
     "EvaluationClient",
     "EvaluationResult",
@@ -86,6 +89,7 @@
     "EvaluatorOutput",
     "EvaluatorResult",
     "FileDatasetProvider",
+    "DatasetManagementServiceProvider",
     "Input",
     "OnDemandEvaluationDatasetRunner",
     "ReferenceInputs",

src/bedrock_agentcore/evaluation/dataset_client.py

@@ -0,0 +1,277 @@
+"""DatasetClient for managing evaluation datasets."""
+
+import logging
+from typing import Any, Dict, Optional
+
+import boto3
+from botocore.config import Config
+
+from bedrock_agentcore._utils.config import WaitConfig
+from bedrock_agentcore._utils.polling import wait_until, wait_until_deleted
+from bedrock_agentcore._utils.snake_case import accept_snake_case_kwargs, convert_kwargs
+from bedrock_agentcore._utils.user_agent import build_user_agent_suffix
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetClient:
+    """Client for managing evaluation datasets.
+
+    Provides pass-through access to all dataset management APIs on the
+    bedrock-agentcore-control service, plus *_and_wait helpers for async operations.
+
+    Example::
+
+        client = DatasetClient(region_name="us-west-2")
+
+        # Create a dataset and wait for ACTIVE
+        dataset = client.create_dataset_and_wait(
+            datasetName="my-dataset",
+            schemaType="AGENTCORE_EVALUATION_PREDEFINED_V1",
+            source={"inlineExamples": {"examples": [...]}},
+        )
+
+        # Pass-through to any dataset API
+        client.list_datasets(maxResults=10)
+        client.add_dataset_examples(datasetId="ds-123", examples=[...])
+    """
+
+    _ALLOWED_CP_METHODS = {
+        # Dataset CRUD
+        "create_dataset",
+        "get_dataset",
+        "list_datasets",
+        "update_dataset",
+        "delete_dataset",
+        # Version management
+        "create_dataset_version",
+        "list_dataset_versions",
+        # Examples management
+        "add_dataset_examples",
+        "update_dataset_examples",
+        "delete_dataset_examples",
+        "list_dataset_examples",
+    }
+
+    def __init__(
+        self,
+        region_name: Optional[str] = None,
+        integration_source: Optional[str] = None,
+        boto3_session: Optional[boto3.Session] = None,
+    ):
+        """Initialize the DatasetClient.
+
+        Args:
+            region_name: AWS region. Falls back to boto3 session region or us-west-2.
+            integration_source: Optional integration framework identifier for telemetry.
+            boto3_session: Optional boto3 Session. If not provided, a default is created.
+        """
+        session = boto3_session if boto3_session else boto3.Session()
+        self.region_name = region_name or session.region_name or "us-west-2"
+        self.integration_source = integration_source
+
+        user_agent_extra = build_user_agent_suffix(integration_source)
+        client_config = Config(user_agent_extra=user_agent_extra)
+
+        self._cp_client = session.client(
+            "bedrock-agentcore-control",
+            region_name=self.region_name,
+            config=client_config,
+        )
+
+        logger.info("Initialized DatasetClient in region %s", self.region_name)
+
+    def __getattr__(self, name: str):
+        """Dynamically forward allowlisted method calls to the boto3 client."""
+        if "_cp_client" not in self.__dict__:
+            raise AttributeError(name)
+
+        if name in self._ALLOWED_CP_METHODS and hasattr(self._cp_client, name):
+            method = getattr(self._cp_client, name)
+            logger.debug("Forwarding method '%s' to _cp_client", name)
+            return accept_snake_case_kwargs(method)
+
+        raise AttributeError(
+            f"'{self.__class__.__name__}' object has no attribute '{name}'. "
+            f"Method not found on control plane client. "
+            f"Available methods can be found in the boto3 documentation for "
+            f"'bedrock-agentcore-control' service."
+        )
+
+    # *_and_wait methods
+    # -------------------------------------------------------------------------
+
+    def create_dataset_and_wait(
+        self,
+        wait_config: Optional[WaitConfig] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Create a dataset and wait for it to reach ACTIVE status.
+
+        Args:
+            wait_config: Optional WaitConfig for polling behavior.
+            **kwargs: Arguments forwarded to the create_dataset API.
+
+        Returns:
+            Dataset details when ACTIVE.
+
+        Raises:
+            RuntimeError: If the dataset reaches CREATE_FAILED status.
+            TimeoutError: If the dataset doesn't become ACTIVE within max_wait.
+        """
+        response = self._cp_client.create_dataset(**convert_kwargs(kwargs))
+        dataset_id = response["datasetId"]
+        return wait_until(
+            lambda: self._cp_client.get_dataset(datasetId=dataset_id),
+            "ACTIVE",
+            {"CREATE_FAILED"},
+            wait_config,
+        )
+
+    def delete_dataset_and_wait(
+        self,
+        wait_config: Optional[WaitConfig] = None,
+        **kwargs,
+    ) -> Optional[Dict[str, Any]]:
+        """Delete a dataset (or a single version) and wait for completion.
+
+        - Full delete (no ``datasetVersion``): polls until ``GetDataset``
+          raises ``ResourceNotFoundException``. Fails on ``DELETE_FAILED``.
+        - Version-specific delete (``datasetVersion`` provided): the dataset
+          itself is not removed. Polls ``GetDataset`` until status returns to
+          ``ACTIVE``. Fails on ``UPDATE_FAILED``. Returns the dataset details.
+
+        Args:
+            wait_config: Optional WaitConfig for polling behavior.
+            **kwargs: Arguments forwarded to the delete_dataset API.
+
+        Raises:
+            RuntimeError: On ``DELETE_FAILED`` or ``UPDATE_FAILED``.
+            TimeoutError: If the operation does not finish within ``max_wait``.
+        """
+        converted = convert_kwargs(kwargs)
+        response = self._cp_client.delete_dataset(**converted)
+        dataset_id = response["datasetId"]
+
+        if "datasetVersion" in converted:
+            return wait_until(
+                lambda: self._cp_client.get_dataset(datasetId=dataset_id),
+                "ACTIVE",
+                {"UPDATE_FAILED"},
+                wait_config,
+            )
+
+        wait_until_deleted(
+            lambda: self._cp_client.get_dataset(datasetId=dataset_id),
+            failed={"DELETE_FAILED"},
+            wait_config=wait_config,
+        )
+        return None
+
+    def create_dataset_version_and_wait(
+        self,
+        wait_config: Optional[WaitConfig] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Create a dataset version and wait for the dataset to reach ACTIVE status.
+
+        Args:
+            wait_config: Optional WaitConfig for polling behavior.
+            **kwargs: Arguments forwarded to the create_dataset_version API.
+
+        Returns:
+            Dataset details when ACTIVE.
+
+        Raises:
+            RuntimeError: If the dataset reaches UPDATE_FAILED status.
+            TimeoutError: If the dataset doesn't become ACTIVE within max_wait.
+        """
+        response = self._cp_client.create_dataset_version(**convert_kwargs(kwargs))
+        dataset_id = response["datasetId"]
+        return wait_until(
+            lambda: self._cp_client.get_dataset(datasetId=dataset_id),
+            "ACTIVE",
+            {"UPDATE_FAILED"},
+            wait_config,
+        )
+
+    def add_examples_and_wait(
+        self,
+        wait_config: Optional[WaitConfig] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Add examples to a dataset and wait for ACTIVE status.
+
+        Args:
+            wait_config: Optional WaitConfig for polling behavior.
+            **kwargs: Arguments forwarded to the add_dataset_examples API.
+
+        Returns:
+            Dataset details when ACTIVE.
+
+        Raises:
+            RuntimeError: If the dataset reaches UPDATE_FAILED status.
+            TimeoutError: If the dataset doesn't become ACTIVE within max_wait.
+        """
+        response = self._cp_client.add_dataset_examples(**convert_kwargs(kwargs))
+        dataset_id = response["datasetId"]
+        return wait_until(
+            lambda: self._cp_client.get_dataset(datasetId=dataset_id),
+            "ACTIVE",
+            {"UPDATE_FAILED"},
+            wait_config,
+        )
+
+    def update_examples_and_wait(
+        self,
+        wait_config: Optional[WaitConfig] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Update examples in a dataset and wait for ACTIVE status.
+
+        Args:
+            wait_config: Optional WaitConfig for polling behavior.
+            **kwargs: Arguments forwarded to the update_dataset_examples API.
+
+        Returns:
+            Dataset details when ACTIVE.
+
+        Raises:
+            RuntimeError: If the dataset reaches UPDATE_FAILED status.
+            TimeoutError: If the dataset doesn't become ACTIVE within max_wait.
+        """
+        response = self._cp_client.update_dataset_examples(**convert_kwargs(kwargs))
+        dataset_id = response["datasetId"]
+        return wait_until(
+            lambda: self._cp_client.get_dataset(datasetId=dataset_id),
+            "ACTIVE",
+            {"UPDATE_FAILED"},
+            wait_config,
+        )
+
+    def delete_examples_and_wait(
+        self,
+        wait_config: Optional[WaitConfig] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Delete examples from a dataset and wait for ACTIVE status.
+
+        Args:
+            wait_config: Optional WaitConfig for polling behavior.
+            **kwargs: Arguments forwarded to the delete_dataset_examples API.
+
+        Returns:
+            Dataset details when ACTIVE.
+
+        Raises:
+            RuntimeError: If the dataset reaches UPDATE_FAILED status.
+            TimeoutError: If the dataset doesn't become ACTIVE within max_wait.
+        """
+        response = self._cp_client.delete_dataset_examples(**convert_kwargs(kwargs))
+        dataset_id = response["datasetId"]
+        return wait_until(
+            lambda: self._cp_client.get_dataset(datasetId=dataset_id),
+            "ACTIVE",
+            {"UPDATE_FAILED"},
+            wait_config,
+        )