fix bug with endpoint metrics

Author: aarora79

src/fmbench/3_run_inference.ipynb

@@ -70,6 +70,7 @@
     "from fmbench.globals import * \n",
     "from datetime import datetime\n",
     "from datetime import timezone\n",
+    "from datetime import timedelta\n",
     "from transformers import AutoTokenizer\n",
     "from sagemaker.predictor import Predictor\n",
     "import importlib.resources as pkg_resources\n",
@@ -494,7 +495,15 @@
     "    sys.modules[module_name] = inference_module\n",
     "    spec.loader.exec_module(inference_module)\n",
     "    # create a predictor from each endpoint in experiments\n",
-    "    return inference_module.create_predictor(ep_name, inference_spec)"
+    "    metadata: Optional[Dict] = None\n",
+    "    if ep_info[0].get('endpoint'):\n",
+    "        production_variants = ep_info[0].get('endpoint').get(\"ProductionVariants\")\n",
+    "        if production_variants is not None:\n",
+    "            variant_name = production_variants[0].get(\"VariantName\")\n",
+    "            metadata = dict(variant_name=variant_name)\n",
+    "            logger.info(f\"ep_name={ep_name}, variant_name={variant_name}\")\n",
+    "    logger.info(f\"ep_name={ep_name}, metadata={metadata}\")\n",
+    "    return inference_module.create_predictor(ep_name, inference_spec, metadata)"
    ]
   },
   {
@@ -611,6 +620,9 @@
     "\n",
     "# dataframe list to hold metrics for each endpoint\n",
     "df_ep_metrics_list = []\n",
+    "# list for holding predictors and run start and end timestamp\n",
+    "# because cloud watch metrics are available after a 1-minute delay\n",
+    "predictors_and_metrics_timestamp_list = []\n",
     "\n",
     "for e_idx, experiment in enumerate(config['experiments']):\n",
     "    # Start timer for the experiment   \n",
@@ -626,7 +638,8 @@
     "    prompt_tokens_total: int = 0\n",
     "    completion_tokens_total: int = 0\n",
     "    for concurrency, payload_file, split_payload in combination_data:\n",
-    "        experiment_at_concurrency_start_dttm = datetime.now()\n",
+    "        # track time at minute boundaries\n",
+    "        experiment_at_concurrency_start_dttm = datetime.utcnow().replace(second=0, microsecond=0)\n",
     "        for chunk_index, chunk in enumerate(split_payload):\n",
     "            logger.info(f\"experiment_index={e_idx+1}/{num_experiments}, \"\n",
     "                        f\"concurrency={concurrency}, payload_file={payload_file}, \"\n",
@@ -666,17 +679,18 @@
     "                                METRICS_PER_INFERENCE_DIR,\n",
     "                                response_file_name)\n",
     "        # save endpoint metrics\n",
-    "        df_ep_metrics = predictor.get_metrics(experiment_at_concurrency_start_dttm,\n",
-    "                                              datetime.now())\n",
-    "        if df_ep_metrics is not None:\n",
-    "            # we want concurrency after timestamp, endpoint name\n",
-    "            df_ep_metrics.insert(loc=2,\n",
-    "                                 column='instance_type',\n",
-    "                                 value=experiment['instance_type'])\n",
-    "            df_ep_metrics.insert(loc=3,\n",
-    "                                 column='concurrency',\n",
-    "                                 value=concurrency)\n",
-    "            df_ep_metrics_list.append(df_ep_metrics)\n",
+    "        experiment_at_concurrency_end_dttm = datetime.utcnow().replace(second=0, microsecond=0)\n",
+    "        # if the endtime and start time are in the same minute then move the endtime to the next\n",
+    "        # minute otherwise cloudwatch would return an empty resonse\n",
+    "        time_delta_in_seconds = (experiment_at_concurrency_end_dttm - experiment_at_concurrency_start_dttm).seconds\n",
+    "        if time_delta_in_seconds < 60:\n",
+    "            experiment_at_concurrency_end_dttm += timedelta(seconds=60)\n",
+    "\n",
+    "        predictors_and_metrics_timestamp_list.append((predictor,\n",
+    "                                                      experiment_at_concurrency_start_dttm,\n",
+    "                                                      experiment_at_concurrency_end_dttm,\n",
+    "                                                      concurrency,\n",
+    "                                                      experiment['instance_type']))\n",
     "\n",
     "    # Experiment done, stopping the timer for this given experiment\n",
     "    experiment_end_time = time.perf_counter()\n",
@@ -715,6 +729,37 @@
     "                f\"duration={experiment_duration:.6f} seconds, exp_cost={exp_cost:.6f}, done\")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# add a 1-minute sleep to be able to grab the CW metrics from the last run\n",
+    "sleep_time: int = 60\n",
+    "logger.info(f\"going to sleep for {sleep_time}s before querying metrics from the endpoint\")\n",
+    "time.sleep(sleep_time)\n",
+    "logger.info(f\"after sleep for {sleep_time}s before querying metrics from the endpoint\")\n",
+    "\n",
+    "for predictor, \\\n",
+    "    experiment_at_concurrency_start_dttm, \\\n",
+    "    experiment_at_concurrency_end_dttm, \\\n",
+    "    concurrency, \\\n",
+    "    instance_type in predictors_and_metrics_timestamp_list:\n",
+    "    # save endpoint metrics\n",
+    "    df_ep_metrics = predictor.get_metrics(experiment_at_concurrency_start_dttm,\n",
+    "                                          experiment_at_concurrency_end_dttm)\n",
+    "    if df_ep_metrics is not None:\n",
+    "        # we want concurrency after timestamp, endpoint name\n",
+    "        df_ep_metrics.insert(loc=2,\n",
+    "                             column='instance_type',\n",
+    "                             value=instance_type)\n",
+    "        df_ep_metrics.insert(loc=3,\n",
+    "                             column='concurrency',\n",
+    "                             value=concurrency)\n",
+    "        df_ep_metrics_list.append(df_ep_metrics)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

src/fmbench/configs/llama3/8b/config-llama3-8b-instruct-p5-djl-lmi-dist.yml

@@ -45,11 +45,11 @@ s3_read_data:
 ## section that enables container to run notebooks and python scripts automatically 
 run_steps:
     0_setup.ipynb: yes
-    1_generate_data.ipynb: no
+    1_generate_data.ipynb: yes
     2_deploy_model.ipynb: yes
     3_run_inference.ipynb: yes
     4_model_metric_analysis.ipynb: yes
-    5_cleanup.ipynb: no
+    5_cleanup.ipynb: yes
 
 
 datasets:

src/fmbench/scripts/bedrock_predictor.py

@@ -29,7 +29,8 @@ class BedrockPredictor(FMBenchPredictor):
     # overriding abstract method
     def __init__(self,
                  endpoint_name: str,
-                 inference_spec: Optional[Dict]):
+                 inference_spec: Optional[Dict],
+                 metadata: Optional[Dict]):
         try:
             # initialize private member variables
             self._endpoint_name = endpoint_name
@@ -221,9 +222,9 @@ def inference_parameters(self) -> Dict:
                     top_p=self._top_p)
 
 
-def create_predictor(endpoint_name: str, inference_spec: Optional[Dict]):
+def create_predictor(endpoint_name: str, inference_spec: Optional[Dict], metadata: Optional[Dict]):
     if endpoint_name in EMBEDDING_MODELS:
         logger.error(f"embeddings models not supported for now")
         return None
     else:
-        return BedrockPredictor(endpoint_name, inference_spec)
+        return BedrockPredictor(endpoint_name, inference_spec, metadata)

src/fmbench/scripts/fmbench_predictor.py

@@ -9,7 +9,8 @@ class FMBenchPredictor(ABC):
     @abstractmethod
     def __init__(self,
                  endpoint_name: str,
-                 inference_spec: Optional[Dict]):
+                 inference_spec: Optional[Dict],
+                 metadata: Optional[Dict]):
         pass
 
     @abstractmethod

src/fmbench/scripts/rest_predictor.py

@@ -20,7 +20,8 @@ class RESTPredictor(FMBenchPredictor):
     # overriding abstract method
     def __init__(self,
                  endpoint_name: str,
-                 inference_spec: Optional[Dict]):
+                 inference_spec: Optional[Dict],
+                 metadata: Optional[Dict]):
         try:
             self._endpoint_name: str = endpoint_name
             self._inference_spec: Dict = inference_spec 
@@ -123,5 +124,5 @@ def inference_parameters(self) -> Dict:
         """The inference parameters property."""
         return self._inference_spec.get("parameters")
 
-def create_predictor(endpoint_name: str, inference_spec: Optional[Dict]):
-    return RESTPredictor(endpoint_name, inference_spec)
+def create_predictor(endpoint_name: str, inference_spec: Optional[Dict], metadata: Optional[Dict]):
+    return RESTPredictor(endpoint_name, inference_spec, metadata)

src/fmbench/scripts/sagemaker_metrics.py

@@ -14,6 +14,7 @@
 logger = logging.getLogger(__name__)
 
 def _get_endpoint_utilization_metrics(endpoint_name: str,
+                                      variant_name: str,
                                       start_time: datetime,
                                       end_time: datetime,
                                       period : int = 60) -> pd.DataFrame:
@@ -42,6 +43,8 @@ def _get_endpoint_utilization_metrics(endpoint_name: str,
     namespace = "/aws/sagemaker/Endpoints"
 
     for metric_name in metrics:
+        logger.debug(f"_get_endpoint_utilization_metrics, endpoint_name={endpoint_name}, variant_name={variant_name}, "
+                     f"metric_name={metric_name}, start_time={start_time}, end_time={end_time}")
         response = client.get_metric_statistics(
             Namespace=namespace,
             MetricName=metric_name,
@@ -52,15 +55,15 @@ def _get_endpoint_utilization_metrics(endpoint_name: str,
                 },
                 {
                     'Name': 'VariantName',
-                    'Value': 'AllTraffic'
+                    'Value': variant_name
                 }
             ],
             StartTime=start_time,
             EndTime=end_time,
             Period=period,
             Statistics=['Average']  # You can also use 'Sum', 'Minimum', 'Maximum', 'SampleCount'
         )
-
+        logger.debug(response)
         for datapoint in response['Datapoints']:
             data.append({
                 'EndpointName': endpoint_name, 
@@ -82,6 +85,7 @@ def _get_endpoint_utilization_metrics(endpoint_name: str,
 
 
 def _get_endpoint_invocation_metrics(endpoint_name: str,
+                                     variant_name: str,
                                      start_time: datetime,
                                      end_time: datetime,
                                      period : int = 60):
@@ -114,6 +118,8 @@ def _get_endpoint_invocation_metrics(endpoint_name: str,
             stat = 'Average'
         else:
             stat = 'Sum'
+        logger.debug(f"_get_endpoint_invocation_metrics, endpoint_name={endpoint_name}, variant_name={variant_name}, "
+                     f"metric_name={metric_name}, start_time={start_time}, end_time={end_time}")
         # Get metric data for the specified metric
         response = client.get_metric_data(
             MetricDataQueries=[
@@ -130,7 +136,7 @@ def _get_endpoint_invocation_metrics(endpoint_name: str,
                                 },
                                 {
                                     'Name': 'VariantName',
-                                    'Value': 'AllTraffic'
+                                    'Value': variant_name
                                 }
                             ]
                         },
@@ -143,7 +149,7 @@ def _get_endpoint_invocation_metrics(endpoint_name: str,
             StartTime=start_time,
             EndTime=end_time
         )
-        
+        logger.debug(response)
         # Extract the data points from the response
         timestamps = response['MetricDataResults'][0]['Timestamps']
         values = response['MetricDataResults'][0]['Values']
@@ -169,6 +175,7 @@ def _get_endpoint_invocation_metrics(endpoint_name: str,
 
 
 def get_endpoint_metrics(endpoint_name: str,
+                         variant_name: str,
                          start_time: datetime,
                          end_time: datetime,
                          period: int = 60):
@@ -188,26 +195,31 @@ def get_endpoint_metrics(endpoint_name: str,
     endpoint_metrics_df: Optional[pd.DataFrame] = None
     try:
         logger.info(f"get_endpoint_metrics, going to retrieve endpoint utlization metrics for "
-                    f"endpoint={endpoint_name}")
+                    f"endpoint={endpoint_name}, variant_name={variant_name}, start_time={start_time}, "
+                    f"end_time={end_time}, period={period}")
         utilization_metrics_df = _get_endpoint_utilization_metrics(endpoint_name=endpoint_name,
+                                                                   variant_name=variant_name,
                                                                    start_time=start_time,
                                                                    end_time=end_time,
                                                                    period=period)
-        logger.info(f"get_endpoint_metrics, going to retrieve invocation metrics for "
-                    f"endpoint={endpoint_name}")
-        invocation_metrics_df = _get_endpoint_invocation_metrics(endpoint_name=endpoint_name, 
-                                                                start_time=start_time,
-                                                                end_time=end_time,
-                                                                period=period)
+        logger.info(f"get_endpoint_metrics, going to retrieve endpoint invocation metrics for "
+                    f"endpoint={endpoint_name}, variant_name={variant_name}, start_time={start_time}, "
+                    f"end_time={end_time}, period={period}")
+        invocation_metrics_df = _get_endpoint_invocation_metrics(endpoint_name=endpoint_name,
+                                                                 variant_name=variant_name,
+                                                                 start_time=start_time,
+                                                                 end_time=end_time,
+                                                                 period=period)
 
         endpoint_metrics_df = pd.merge(utilization_metrics_df,
                                        invocation_metrics_df,
                                        on=['Timestamp', 'EndpointName'],
                                        how='outer')
         logger.info(f"get_endpoint_metrics, shape of invocation and utilization metrics for "
                     f"endpoint={endpoint_name} is {endpoint_metrics_df.shape}")
+        logger.info(f"get_endpoint_metrics, endpoint_metrics_df={endpoint_metrics_df.head()}")
     except Exception as e:
         logger.error(f"get_endpoint_metrics, exception occured while retrieving metrics for {endpoint_name}, "
                      f"exception={e}")
-        
+
     return endpoint_metrics_df

src/fmbench/scripts/sagemaker_predictor.py

@@ -21,10 +21,15 @@ class SageMakerPredictor(FMBenchPredictor):
     # overriding abstract method
     def __init__(self,
                  endpoint_name: str,
-                 inference_spec: Optional[Dict]):
+                 inference_spec: Optional[Dict],
+                 metadata: Optional[Dict]):
         self._predictor: Optional[sagemaker.base_predictor.Predictor] = None
         self._endpoint_name: str = endpoint_name
         self._inference_spec: Dict = inference_spec
+        self._variant_name: Optional[str] = None
+        if metadata is not None:
+            self._variant_name = metadata.get("variant_name")
+
         try:
             # Create a SageMaker Predictor object
             self._predictor = Predictor(
@@ -115,13 +120,13 @@ def get_metrics(self,
                     start_time: datetime,
                     end_time: datetime,
                     period: int = 60) -> pd.DataFrame:
-        return get_endpoint_metrics(self._endpoint_name, start_time, end_time)
+        return get_endpoint_metrics(self._endpoint_name, self._variant_name, start_time, end_time)
 
     @property
     def inference_parameters(self) -> Dict:
         """The inference parameters property."""
         return self._inference_spec.get("parameters")
 
 
-def create_predictor(endpoint_name: str, inference_spec: Optional[Dict]):
-    return SageMakerPredictor(endpoint_name, inference_spec)
+def create_predictor(endpoint_name: str, inference_spec: Optional[Dict], metadata: Optional[Dict]):
+    return SageMakerPredictor(endpoint_name, inference_spec, metadata)