workbench 0.8.168__py3-none-any.whl → 0.8.192__py3-none-any.whl

workbench/core/artifacts/data_capture_core.py

@@ -0,0 +1,355 @@
+"""DataCaptureCore class for managing SageMaker endpoint data capture"""
+
+import logging
+import re
+import time
+from datetime import datetime
+from typing import Tuple
+import pandas as pd
+from sagemaker import Predictor
+from sagemaker.model_monitor import DataCaptureConfig
+import awswrangler as wr
+
+# Workbench Imports
+from workbench.core.artifacts.endpoint_core import EndpointCore
+from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
+from workbench.utils.monitor_utils import process_data_capture
+
+# Setup logging
+log = logging.getLogger("workbench")
+
+
+class DataCaptureCore:
+    """Manages data capture configuration and retrieval for SageMaker endpoints"""
+
+    def __init__(self, endpoint_name: str):
+        """DataCaptureCore Class
+
+        Args:
+            endpoint_name (str): Name of the endpoint to manage data capture for
+        """
+        self.log = logging.getLogger("workbench")
+        self.endpoint_name = endpoint_name
+        self.endpoint = EndpointCore(self.endpoint_name)
+
+        # Initialize Class Attributes
+        self.sagemaker_session = self.endpoint.sm_session
+        self.sagemaker_client = self.endpoint.sm_client
+        self.data_capture_path = self.endpoint.endpoint_data_capture_path
+        self.workbench_role_arn = AWSAccountClamp().aws_session.get_workbench_execution_role_arn()
+
+    def summary(self) -> dict:
+        """Return the summary of data capture configuration
+
+        Returns:
+            dict: Summary of data capture status
+        """
+        if self.endpoint.is_serverless():
+            return {"endpoint_type": "serverless", "data_capture": "not supported"}
+        else:
+            return {
+                "endpoint_type": "realtime",
+                "data_capture_enabled": self.is_enabled(),
+                "capture_percentage": self.capture_percentage(),
+                "capture_modes": self.capture_modes() if self.is_enabled() else [],
+                "data_capture_path": self.data_capture_path if self.is_enabled() else None,
+            }
+
+    def enable(self, capture_percentage=100, capture_options=None, force_redeploy=False):
+        """
+        Enable data capture for the SageMaker endpoint.
+
+        Args:
+            capture_percentage (int): Percentage of data to capture. Defaults to 100.
+            capture_options (list): List of what to capture - ["REQUEST"], ["RESPONSE"], or ["REQUEST", "RESPONSE"].
+                                    Defaults to ["REQUEST", "RESPONSE"] to capture both.
+            force_redeploy (bool): If True, force redeployment even if data capture is already enabled.
+        """
+        # Early returns for cases where we can't/don't need to add data capture
+        if self.endpoint.is_serverless():
+            self.log.warning("Data capture is not supported for serverless endpoints.")
+            return
+
+        # Default to capturing both if not specified
+        if capture_options is None:
+            capture_options = ["REQUEST", "RESPONSE"]
+
+        # Validate capture_options
+        valid_options = {"REQUEST", "RESPONSE"}
+        if not all(opt in valid_options for opt in capture_options):
+            self.log.error("Invalid capture_options. Must be a list containing 'REQUEST' and/or 'RESPONSE'")
+            return
+
+        if self.is_enabled() and not force_redeploy:
+            self.log.important(f"Data capture already configured for {self.endpoint_name}.")
+            return
+
+        # Get the current endpoint configuration name for later deletion
+        current_endpoint_config_name = self.endpoint.endpoint_config_name()
+
+        # Log the data capture operation
+        self.log.important(f"Enabling Data Capture for {self.endpoint_name} --> {self.data_capture_path}")
+        self.log.important(f"Capturing: {', '.join(capture_options)} at {capture_percentage}% sampling")
+        self.log.important("This will redeploy the endpoint...")
+
+        # Create and apply the data capture configuration
+        data_capture_config = DataCaptureConfig(
+            enable_capture=True,
+            sampling_percentage=capture_percentage,
+            destination_s3_uri=self.data_capture_path,
+            capture_options=capture_options,
+        )
+
+        # Update endpoint with the new capture configuration
+        Predictor(self.endpoint_name, sagemaker_session=self.sagemaker_session).update_data_capture_config(
+            data_capture_config=data_capture_config
+        )
+
+        # Clean up old endpoint configuration
+        try:
+            self.sagemaker_client.delete_endpoint_config(EndpointConfigName=current_endpoint_config_name)
+            self.log.info(f"Deleted old endpoint configuration: {current_endpoint_config_name}")
+        except Exception as e:
+            self.log.warning(f"Could not delete old endpoint configuration {current_endpoint_config_name}: {e}")
+
+    def disable(self):
+        """
+        Disable data capture for the SageMaker endpoint.
+        """
+        # Early return if data capture isn't configured
+        if not self.is_enabled():
+            self.log.important(f"Data capture is not currently enabled for {self.endpoint_name}.")
+            return
+
+        # Get the current endpoint configuration name for later deletion
+        current_endpoint_config_name = self.endpoint.endpoint_config_name()
+
+        # Log the operation
+        self.log.important(f"Disabling Data Capture for {self.endpoint_name}")
+        self.log.important("This normally redeploys the endpoint...")
+
+        # Create a configuration with capture disabled
+        data_capture_config = DataCaptureConfig(enable_capture=False, destination_s3_uri=self.data_capture_path)
+
+        # Update endpoint with the new configuration
+        Predictor(self.endpoint_name, sagemaker_session=self.sagemaker_session).update_data_capture_config(
+            data_capture_config=data_capture_config
+        )
+
+        # Clean up old endpoint configuration
+        self.sagemaker_client.delete_endpoint_config(EndpointConfigName=current_endpoint_config_name)
+
+    def is_enabled(self) -> bool:
+        """
+        Check if data capture is enabled on the endpoint.
+
+        Returns:
+            bool: True if data capture is enabled, False otherwise.
+        """
+        try:
+            endpoint_config_name = self.endpoint.endpoint_config_name()
+            endpoint_config = self.sagemaker_client.describe_endpoint_config(EndpointConfigName=endpoint_config_name)
+            data_capture_config = endpoint_config.get("DataCaptureConfig", {})
+
+            # Check if data capture is enabled
+            is_enabled = data_capture_config.get("EnableCapture", False)
+            return is_enabled
+        except Exception as e:
+            self.log.error(f"Error checking data capture configuration: {e}")
+            return False
+
+    def capture_percentage(self) -> int:
+        """
+        Get the data capture percentage from the endpoint configuration.
+
+        Returns:
+            int: Data capture percentage if enabled, None otherwise.
+        """
+        try:
+            endpoint_config_name = self.endpoint.endpoint_config_name()
+            endpoint_config = self.sagemaker_client.describe_endpoint_config(EndpointConfigName=endpoint_config_name)
+            data_capture_config = endpoint_config.get("DataCaptureConfig", {})
+
+            # Check if data capture is enabled and return the percentage
+            if data_capture_config.get("EnableCapture", False):
+                return data_capture_config.get("InitialSamplingPercentage", 0)
+            else:
+                return None
+        except Exception as e:
+            self.log.error(f"Error checking data capture percentage: {e}")
+            return None
+
+    def get_config(self) -> dict:
+        """
+        Returns the complete data capture configuration from the endpoint config.
+
+        Returns:
+            dict: Complete DataCaptureConfig from AWS, or None if not configured
+        """
+        config_name = self.endpoint.endpoint_config_name()
+        response = self.sagemaker_client.describe_endpoint_config(EndpointConfigName=config_name)
+        data_capture_config = response.get("DataCaptureConfig")
+        if not data_capture_config:
+            self.log.error(f"No data capture configuration found for endpoint config {config_name}")
+            return None
+        return data_capture_config
+
+    def capture_modes(self) -> list:
+        """Get the current capture modes (REQUEST/RESPONSE)"""
+        if not self.is_enabled():
+            return []
+
+        config = self.get_config()
+        if not config:
+            return []
+
+        capture_options = config.get("CaptureOptions", [])
+        modes = [opt.get("CaptureMode") for opt in capture_options]
+        return ["REQUEST" if m == "Input" else "RESPONSE" for m in modes if m]
+
+    def get_captured_data(self, from_date: str = None, add_timestamp: bool = True) -> Tuple[pd.DataFrame, pd.DataFrame]:
+        """
+        Read and process captured data from S3.
+
+        Args:
+            from_date (str, optional): Only process files from this date onwards (YYYY-MM-DD format).
+                                       Defaults to None to process all files.
+            add_timestamp (bool, optional): Whether to add a timestamp column to the DataFrame.
+
+        Returns:
+            Tuple[pd.DataFrame, pd.DataFrame]: Processed input and output DataFrames.
+        """
+        files = wr.s3.list_objects(self.data_capture_path)
+        if not files:
+            self.log.warning(f"No data capture files found in {self.data_capture_path}.")
+            return pd.DataFrame(), pd.DataFrame()
+
+        # Filter by date if specified
+        if from_date:
+            from_date_obj = datetime.strptime(from_date, "%Y-%m-%d").date()
+            files = [f for f in files if self._file_date_filter(f, from_date_obj)]
+            self.log.info(f"Processing {len(files)} files from {from_date} onwards.")
+        else:
+            self.log.info(f"Processing all {len(files)} files...")
+
+        # Check if any files remain after filtering
+        if not files:
+            self.log.info("No files to process after date filtering.")
+            return pd.DataFrame(), pd.DataFrame()
+
+        # Sort files by name (assumed to include timestamp)
+        files.sort()
+
+        # Get all timestamps in one batch if needed
+        timestamps = {}
+        if add_timestamp:
+            # Batch describe operation - much more efficient than per-file calls
+            timestamps = wr.s3.describe_objects(path=files)
+
+        # Process files using concurrent.futures
+        start_time = time.time()
+
+        def process_single_file(file_path):
+            """Process a single file and return input/output DataFrames."""
+            try:
+                log.debug(f"Processing file: {file_path}...")
+                df = wr.s3.read_json(path=file_path, lines=True)
+                if not df.empty:
+                    input_df, output_df = process_data_capture(df)
+                    if add_timestamp and file_path in timestamps:
+                        output_df["timestamp"] = timestamps[file_path]["LastModified"]
+                    return input_df, output_df
+                return pd.DataFrame(), pd.DataFrame()
+            except Exception as e:
+                self.log.warning(f"Error processing {file_path}: {e}")
+                return pd.DataFrame(), pd.DataFrame()
+
+        # Use ThreadPoolExecutor for I/O-bound operations
+        from concurrent.futures import ThreadPoolExecutor
+
+        max_workers = min(32, len(files))  # Cap at 32 threads or number of files
+
+        all_input_dfs, all_output_dfs = [], []
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = [executor.submit(process_single_file, file_path) for file_path in files]
+            for future in futures:
+                input_df, output_df = future.result()
+                if not input_df.empty:
+                    all_input_dfs.append(input_df)
+                if not output_df.empty:
+                    all_output_dfs.append(output_df)
+
+        if not all_input_dfs:
+            self.log.warning("No valid data was processed.")
+            return pd.DataFrame(), pd.DataFrame()
+
+        input_df = pd.concat(all_input_dfs, ignore_index=True)
+        output_df = pd.concat(all_output_dfs, ignore_index=True)
+
+        elapsed_time = time.time() - start_time
+        self.log.info(f"Processed {len(files)} files in {elapsed_time:.2f} seconds.")
+        return input_df, output_df
+
+    def _file_date_filter(self, file_path, from_date_obj):
+        """Extract date from S3 path and compare with from_date."""
+        try:
+            # Match YYYY/MM/DD pattern in the path
+            date_match = re.search(r"/(\d{4})/(\d{2})/(\d{2})/", file_path)
+            if date_match:
+                year, month, day = date_match.groups()
+                file_date = datetime(int(year), int(month), int(day)).date()
+                return file_date >= from_date_obj
+            return False  # No date pattern found
+        except ValueError:
+            return False
+
+    def __repr__(self) -> str:
+        """String representation of this DataCaptureCore object
+
+        Returns:
+            str: String representation of this DataCaptureCore object
+        """
+        summary_dict = self.summary()
+        summary_items = [f"  {repr(key)}: {repr(value)}" for key, value in summary_dict.items()]
+        summary_str = f"{self.__class__.__name__}: {self.endpoint_name}\n" + ",\n".join(summary_items)
+        return summary_str
+
+
+# Test function for the class
+if __name__ == "__main__":
+    """Exercise the MonitorCore class"""
+    from pprint import pprint
+
+    # Set options for actually seeing the dataframe
+    pd.set_option("display.max_columns", None)
+    pd.set_option("display.width", None)
+
+    # Create the Class and test it out
+    endpoint_name = "abalone-regression-rt"
+    my_endpoint = EndpointCore(endpoint_name)
+    if not my_endpoint.exists():
+        print(f"Endpoint {endpoint_name} does not exist.")
+        exit(1)
+    dc = my_endpoint.data_capture()
+
+    # Check the summary of the data capture class
+    pprint(dc.summary())
+
+    # Enable data capture on the endpoint
+    # dc.enable(force_redeploy=True)
+    my_endpoint.enable_data_capture()
+
+    # Test the data capture by running some predictions
+    # pred_df = my_endpoint.auto_inference()
+    # print(pred_df.head())
+
+    # Check that data capture is working
+    input_df, output_df = dc.get_captured_data(from_date="2025-09-01")
+    if input_df.empty and output_df.empty:
+        print("No data capture files found, for a new endpoint it may take a few minutes to start capturing data")
+    else:
+        print("Found data capture files")
+        print("Input")
+        print(input_df.head())
+        print("Output")
+        print(output_df.head())

workbench/core/artifacts/endpoint_core.py

@@ -8,7 +8,7 @@ import pandas as pd
 import numpy as np
 from io import StringIO
 import awswrangler as wr
-from typing import Union, Optional
+from typing import Union, Optional, Tuple
 import hashlib
 
 # Model Performance Scores
@@ -32,11 +32,11 @@ from sagemaker import Predictor
 from workbench.core.artifacts.artifact import Artifact
 from workbench.core.artifacts import FeatureSetCore, ModelCore, ModelType
 from workbench.utils.endpoint_metrics import EndpointMetrics
-from workbench.utils.fast_inference import fast_inference
 from workbench.utils.cache import Cache
 from workbench.utils.s3_utils import compute_s3_object_hash
 from workbench.utils.model_utils import uq_metrics
 from workbench.utils.xgboost_model_utils import cross_fold_inference
+from workbench_bridges.endpoints.fast_inference import fast_inference
 
 
 class EndpointCore(Artifact):
@@ -164,11 +164,17 @@ class EndpointCore(Artifact):
         """
         return "Serverless" in self.endpoint_meta["InstanceType"]
 
-    def add_data_capture(self):
+    def data_capture(self):
+        """Get the MonitorCore class for this endpoint"""
+        from workbench.core.artifacts.data_capture_core import DataCaptureCore
+
+        return DataCaptureCore(self.endpoint_name)
+
+    def enable_data_capture(self):
         """Add data capture to the endpoint"""
-        self.get_monitor().add_data_capture()
+        self.data_capture().enable()
 
-    def get_monitor(self):
+    def monitor(self):
         """Get the MonitorCore class for this endpoint"""
         from workbench.core.artifacts.monitor_core import MonitorCore
 
@@ -350,7 +356,7 @@ class EndpointCore(Artifact):
             return pd.DataFrame()
 
         # Grab the evaluation data from the FeatureSet
-        table = fs.view("training").table
+        table = model.training_view().table
         eval_df = fs.query(f'SELECT * FROM "{table}" where training = FALSE')
         capture_name = "auto_inference" if capture else None
         return self.inference(eval_df, capture_name, id_column=fs.id_column)
@@ -414,8 +420,12 @@ class EndpointCore(Artifact):
 
             # Capture the inference results and metrics
             if capture_name is not None:
+
+                # If we don't have an id_column, we'll pull it from the model's FeatureSet
+                if id_column is None:
+                    fs = FeatureSetCore(model.get_input())
+                    id_column = fs.id_column
                 description = capture_name.replace("_", " ").title()
-                features = model.features()
                 self._capture_inference_results(
                     capture_name, prediction_df, target_column, model_type, metrics, description, features, id_column
                 )
@@ -423,31 +433,81 @@ class EndpointCore(Artifact):
                 # For UQ Models we also capture the uncertainty metrics
                 if model_type in [ModelType.UQ_REGRESSOR]:
                     metrics = uq_metrics(prediction_df, target_column)
-
-                    # Now put into the Parameter Store Model Inference Namespace
                     self.param_store.upsert(f"/workbench/models/{model.name}/inference/{capture_name}", metrics)
 
         # Return the prediction DataFrame
         return prediction_df
 
-    def cross_fold_inference(self, nfolds: int = 5) -> dict:
+    def cross_fold_inference(self, nfolds: int = 5) -> Tuple[dict, pd.DataFrame]:
         """Run cross-fold inference (only works for XGBoost models)
 
         Args:
             nfolds (int): Number of folds to use for cross-fold (default: 5)
 
         Returns:
-            dict: Dictionary with the cross-fold inference results
+            Tuple[dict, pd.DataFrame]: Tuple of (cross_fold_metrics, out_of_fold_df)
         """
 
         # Grab our model
         model = ModelCore(self.model_name)
 
         # Compute CrossFold Metrics
-        cross_fold_metrics = cross_fold_inference(model, nfolds=nfolds)
+        cross_fold_metrics, out_of_fold_df = cross_fold_inference(model, nfolds=nfolds)
         if cross_fold_metrics:
             self.param_store.upsert(f"/workbench/models/{model.name}/inference/cross_fold", cross_fold_metrics)
-        return cross_fold_metrics
+
+        # Capture the results
+        capture_name = "full_cross_fold"
+        description = capture_name.replace("_", " ").title()
+        target_column = model.target()
+        model_type = model.model_type
+
+        # Get the id_column from the model's FeatureSet
+        fs = FeatureSetCore(model.get_input())
+        id_column = fs.id_column
+
+        # Is this a UQ Model? If so, run full inference and merge the results
+        additional_columns = []
+        if model_type == ModelType.UQ_REGRESSOR:
+            self.log.important("UQ Regressor detected, running full inference to get uncertainty estimates...")
+
+            # Get the training view dataframe for inference
+            training_df = model.training_view().pull_dataframe()
+
+            # Run inference on the endpoint to get UQ outputs
+            uq_df = self.inference(training_df)
+
+            # Identify UQ-specific columns (quantiles and prediction_std)
+            uq_columns = [col for col in uq_df.columns if col.startswith("q_") or col == "prediction_std"]
+
+            # Merge UQ columns with out-of-fold predictions
+            if uq_columns:
+                # Keep id_column and UQ columns, drop 'prediction' to avoid conflict when merging
+                uq_df = uq_df[[id_column] + uq_columns]
+
+                # Drop duplicates in uq_df based on id_column
+                uq_df = uq_df.drop_duplicates(subset=[id_column])
+
+                # Merge UQ columns into out_of_fold_df
+                out_of_fold_df = pd.merge(out_of_fold_df, uq_df, on=id_column, how="left")
+                additional_columns = uq_columns
+                self.log.info(f"Added UQ columns: {', '.join(additional_columns)}")
+
+                # Also compute UQ metrics
+                metrics = uq_metrics(out_of_fold_df, target_column)
+                self.param_store.upsert(f"/workbench/models/{model.name}/inference/{capture_name}", metrics)
+
+        self._capture_inference_results(
+            capture_name,
+            out_of_fold_df,
+            target_column,
+            model_type,
+            pd.DataFrame([cross_fold_metrics["summary_metrics"]]),
+            description,
+            features=additional_columns,
+            id_column=id_column,
+        )
+        return cross_fold_metrics, out_of_fold_df
 
     def fast_inference(self, eval_df: pd.DataFrame, threads: int = 4) -> pd.DataFrame:
         """Run inference on the Endpoint using the provided DataFrame
@@ -642,6 +702,10 @@ class EndpointCore(Artifact):
     @staticmethod
     def _hash_dataframe(df: pd.DataFrame, hash_length: int = 8):
         # Internal: Compute a data hash for the dataframe
+        if df.empty:
+            return "--hash--"
+
+        # Sort the dataframe by columns to ensure consistent ordering
         df = df.copy()
         df = df.sort_values(by=sorted(df.columns.tolist()))
         row_hashes = pd.util.hash_pandas_object(df, index=False)
@@ -696,8 +760,8 @@ class EndpointCore(Artifact):
         wr.s3.to_csv(metrics, f"{inference_capture_path}/inference_metrics.csv", index=False)
 
         # Grab the target column, prediction column, any _proba columns, and the ID column (if present)
-        prediction_col = "prediction" if "prediction" in pred_results_df.columns else "predictions"
-        output_columns = [target_column, prediction_col]
+        output_columns = [target_column]
+        output_columns += [col for col in pred_results_df.columns if "prediction" in col]
 
         # Add any _proba columns to the output columns
         output_columns += [col for col in pred_results_df.columns if col.endswith("_proba")]
@@ -707,7 +771,7 @@ class EndpointCore(Artifact):
 
         # Add the ID column
         if id_column and id_column in pred_results_df.columns:
-            output_columns.append(id_column)
+            output_columns.insert(0, id_column)
 
         # Write the predictions to our S3 Model Inference Folder
         self.log.info(f"Writing predictions to {inference_capture_path}/inference_predictions.csv")
@@ -929,9 +993,9 @@ class EndpointCore(Artifact):
         self.upsert_workbench_meta({"workbench_input": input})
 
     def delete(self):
-        """ "Delete an existing Endpoint: Underlying Models, Configuration, and Endpoint"""
+        """Delete an existing Endpoint: Underlying Models, Configuration, and Endpoint"""
         if not self.exists():
-            self.log.warning(f"Trying to delete an Model that doesn't exist: {self.name}")
+            self.log.warning(f"Trying to delete an Endpoint that doesn't exist: {self.name}")
 
         # Remove this endpoint from the list of registered endpoints
         self.log.info(f"Removing {self.name} from the list of registered endpoints...")
@@ -972,12 +1036,23 @@ class EndpointCore(Artifact):
             cls.log.info(f"Deleting Monitoring Schedule {schedule['MonitoringScheduleName']}...")
             cls.sm_client.delete_monitoring_schedule(MonitoringScheduleName=schedule["MonitoringScheduleName"])
 
-        # Recursively delete all endpoint S3 artifacts (inference, data capture, monitoring, etc)
+        # Recursively delete all endpoint S3 artifacts (inference, etc)
+        # Note: We do not want to delete the data_capture/ files since these
+        #       might be used for collection and data drift analysis
         base_endpoint_path = f"{cls.endpoints_s3_path}/{endpoint_name}"
-        s3_objects = wr.s3.list_objects(base_endpoint_path, boto3_session=cls.boto3_session)
-        cls.log.info(f"Deleting S3 Objects at {base_endpoint_path}...")
-        cls.log.info(f"{s3_objects}")
-        wr.s3.delete_objects(s3_objects, boto3_session=cls.boto3_session)
+        all_s3_objects = wr.s3.list_objects(base_endpoint_path, boto3_session=cls.boto3_session)
+
+        # Filter out objects that contain 'data_capture/' in their path
+        s3_objects_to_delete = [obj for obj in all_s3_objects if "/data_capture/" not in obj]
+        cls.log.info(f"Found {len(all_s3_objects)} total objects at {base_endpoint_path}")
+        cls.log.info(f"Filtering out data_capture files, will delete {len(s3_objects_to_delete)} objects...")
+        cls.log.info(f"Objects to delete: {s3_objects_to_delete}")
+
+        if s3_objects_to_delete:
+            wr.s3.delete_objects(s3_objects_to_delete, boto3_session=cls.boto3_session)
+            cls.log.info(f"Successfully deleted {len(s3_objects_to_delete)} objects")
+        else:
+            cls.log.info("No objects to delete (only data_capture files found)")
 
         # Delete any dataframes that were stored in the Dataframe Cache
         cls.log.info("Deleting Dataframe Cache...")
@@ -1028,7 +1103,7 @@ class EndpointCore(Artifact):
 if __name__ == "__main__":
     """Exercise the Endpoint Class"""
     from workbench.api import FeatureSet
-    from workbench.utils.endpoint_utils import fs_evaluation_data
+    from workbench.utils.endpoint_utils import get_evaluation_data
     import random
 
     # Grab an EndpointCore object and pull some information from it
@@ -1036,7 +1111,7 @@ if __name__ == "__main__":
 
     # Test various error conditions (set row 42 length to pd.NA)
     # Note: This test should return ALL rows
-    my_eval_df = fs_evaluation_data(my_endpoint)
+    my_eval_df = get_evaluation_data(my_endpoint)
     my_eval_df.at[42, "length"] = pd.NA
     pred_results = my_endpoint.inference(my_eval_df, drop_error_rows=True)
     print(f"Sent rows: {len(my_eval_df)}")
@@ -1044,6 +1119,9 @@ if __name__ == "__main__":
     assert len(pred_results) == len(my_eval_df), "Predictions should match the number of sent rows"
 
     # Now we put in an invalid value
+    print("*" * 80)
+    print("NOW TESTING ERROR CONDITIONS...")
+    print("*" * 80)
     my_eval_df.at[42, "length"] = "invalid_value"
     pred_results = my_endpoint.inference(my_eval_df, drop_error_rows=True)
     print(f"Sent rows: {len(my_eval_df)}")
@@ -1104,16 +1182,20 @@ if __name__ == "__main__":
     # Run Inference where we provide the data
     # Note: This dataframe could be from a FeatureSet or any other source
     print("Running Inference...")
-    my_eval_df = fs_evaluation_data(my_endpoint)
+    my_eval_df = get_evaluation_data(my_endpoint)
     pred_results = my_endpoint.inference(my_eval_df)
 
     # Now set capture=True to save inference results and metrics
-    my_eval_df = fs_evaluation_data(my_endpoint)
+    my_eval_df = get_evaluation_data(my_endpoint)
     pred_results = my_endpoint.inference(my_eval_df, capture_name="holdout_xyz")
 
     # Run predictions using the fast_inference method
     fast_results = my_endpoint.fast_inference(my_eval_df)
 
+    # Test the cross_fold_inference method
+    print("Running Cross-Fold Inference...")
+    metrics, all_results = my_endpoint.cross_fold_inference()
+
     # Run Inference and metrics for a Classification Endpoint
     class_endpoint = EndpointCore("wine-classification")
     auto_predictions = class_endpoint.auto_inference()
@@ -1122,6 +1204,10 @@ if __name__ == "__main__":
     target = "wine_class"
     print(class_endpoint.generate_confusion_matrix(target, auto_predictions))
 
+    # Test the cross_fold_inference method
+    print("Running Cross-Fold Inference...")
+    metrics, all_results = class_endpoint.cross_fold_inference()
+
     # Test the class method delete (commented out for now)
     # from workbench.api import Model
     # model = Model("abalone-regression")