fix: codestyle, type hints, license, and docstrings

Author: Uemit Yoldas

src/sagemaker/amtviz/__init__.py

@@ -1,17 +1,17 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: MIT-0
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy of this
-# software and associated documentation files (the "Software"), to deal in the Software
-# without restriction, including without limitation the rights to use, copy, modify,
-# merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
-# permit persons to whom the Software is furnished to do so.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
-# INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
-# PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
-# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
-# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Placeholder docstring"""
+from __future__ import absolute_import
 
 from sagemaker.amtviz.visualization import visualize_tuning_job
-__all__ = ['visualize_tuning_job']
+__all__ = ['visualize_tuning_job']

src/sagemaker/amtviz/job_metrics.py

@@ -1,29 +1,28 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: MIT-0
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
-# IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Helper functions to retrieve job metrics from CloudWatch."""
+from __future__ import absolute_import
 
 from datetime import datetime, timedelta
 from typing import Callable, List, Optional, Tuple, Dict, Any
 import hashlib
 import os
 from pathlib import Path
 
+import logging
 import pandas as pd
 import numpy as np
 import boto3
-import logging
 
 logger = logging.getLogger(__name__)
 
@@ -58,16 +57,16 @@ def inner(*args: Any, **kwargs: Any) -> pd.DataFrame:
                 logger.debug("H", end="")
                 df["ts"] = pd.to_datetime(df["ts"])
                 df["ts"] = df["ts"].dt.tz_localize(None)
-                df["rel_ts"] = pd.to_datetime(df["rel_ts"]) # pyright: ignore [reportIndexIssue, reportOptionalSubscript]
+                # pyright: ignore [reportIndexIssue, reportOptionalSubscript]
+                df["rel_ts"] = pd.to_datetime(df["rel_ts"])
                 df["rel_ts"] = df["rel_ts"].dt.tz_localize(None)
                 return df
             except KeyError:
                 # Empty file leads to empty df, hence no df['ts'] possible
                 pass
             # nosec b110 - doesn't matter why we could not load it.
             except BaseException as e:
-                logger.error("\nException", type(e), e)
-                pass  # continue with calling the outer function
+                logger.error("\nException: %s - %s", type(e), e)
 
         logger.debug("M", end="")
         df = outer(*args, **kwargs)
@@ -82,6 +81,7 @@ def inner(*args: Any, **kwargs: Any) -> pd.DataFrame:
 
 
 def _metric_data_query_tpl(metric_name: str, dim_name: str, dim_value: str) -> Dict[str, Any]:
+    """Returns a CloudWatch metric data query template."""
     return {
         "Id": metric_name.lower().replace(":", "_").replace("-", "_"),
         "MetricStat": {
@@ -100,18 +100,19 @@ def _metric_data_query_tpl(metric_name: str, dim_name: str, dim_value: str) -> D
 
 
 def _get_metric_data(
-    queries: List[Dict[str, Any]], 
-    start_time: datetime, 
+    queries: List[Dict[str, Any]],
+    start_time: datetime,
     end_time: datetime
 ) -> pd.DataFrame:
+    """Fetches CloudWatch metrics between timestamps and returns a DataFrame with selected columns."""
     start_time = start_time - timedelta(hours=1)
     end_time = end_time + timedelta(hours=1)
     response = cw.get_metric_data(MetricDataQueries=queries, StartTime=start_time, EndTime=end_time)
 
     df = pd.DataFrame()
     if "MetricDataResults" not in response:
         return df
-    
+
     for metric_data in response["MetricDataResults"]:
         values = metric_data["Values"]
         ts = np.array(metric_data["Timestamps"], dtype=np.datetime64)
@@ -130,11 +131,11 @@ def _get_metric_data(
 
 @disk_cache
 def _collect_metrics(
-    dimensions: List[Tuple[str, str]], 
-    start_time: datetime, 
+    dimensions: List[Tuple[str, str]],
+    start_time: datetime,
     end_time: Optional[datetime]
 ) -> pd.DataFrame:
-
+    """Collects SageMaker training job metrics from CloudWatch based on given dimensions and time range."""
     df = pd.DataFrame()
     for dim_name, dim_value in dimensions:
         response = cw.list_metrics(
@@ -158,8 +159,8 @@ def _collect_metrics(
 
 
 def get_cw_job_metrics(
-    job_name: str, 
-    start_time: Optional[datetime] = None, 
+    job_name: str,
+    start_time: Optional[datetime] = None,
     end_time: Optional[datetime] = None
 ) -> pd.DataFrame:
     """Retrieves CloudWatch metrics for a SageMaker training job.
@@ -182,4 +183,4 @@ def get_cw_job_metrics(
     # If not given, use reasonable defaults for start and end time
     start_time = start_time or datetime.now() - timedelta(hours=4)
     end_time = end_time or start_time + timedelta(hours=4)
-    return _collect_metrics(dimensions, start_time, end_time)
+    return _collect_metrics(dimensions, start_time, end_time)