workbench 0.8.162__py3-none-any.whl → 0.8.220__py3-none-any.whl

workbench/utils/cloudwatch_utils.py

@@ -0,0 +1,137 @@
+"""AWS CloudWatch utility functions for Workbench."""
+
+import time
+import logging
+from datetime import datetime, timezone
+from typing import List, Optional, Dict, Generator
+from urllib.parse import quote
+from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
+
+log = logging.getLogger("workbench")
+
+
+def get_cloudwatch_client():
+    """Get the CloudWatch Logs client using the Workbench assumed role session."""
+    session = AWSAccountClamp().boto3_session
+    return session.client("logs")
+
+
+def get_cloudwatch_logs_url(log_group: str, log_stream: str) -> Optional[str]:
+    """
+    Generate CloudWatch logs URL for the specified log group and stream.
+
+    Args:
+        log_group: Log group name (e.g., '/aws/batch/job')
+        log_stream: Log stream name
+
+    Returns:
+        CloudWatch console URL or None if unable to generate
+    """
+    try:
+        region = AWSAccountClamp().region
+
+        # URL encode the log group and stream
+        encoded_group = quote(log_group, safe="")
+        encoded_stream = quote(log_stream, safe="")
+
+        return (
+            f"https://{region}.console.aws.amazon.com/cloudwatch/home?"
+            f"region={region}#logsV2:log-groups/log-group/{encoded_group}"
+            f"/log-events/{encoded_stream}"
+        )
+    except Exception as e:  # noqa: BLE001
+        log.warning(f"Failed to generate CloudWatch logs URL: {e}")
+        return None
+
+
+def get_active_log_streams(
+    log_group_name: str, start_time_ms: int, stream_filter: Optional[str] = None, client=None
+) -> List[str]:
+    """Retrieve log streams that have events after the specified start time."""
+    if not client:
+        client = get_cloudwatch_client()
+    active_streams = []
+    stream_params = {
+        "logGroupName": log_group_name,
+        "orderBy": "LastEventTime",
+        "descending": True,
+    }
+    while True:
+        response = client.describe_log_streams(**stream_params)
+        log_streams = response.get("logStreams", [])
+        for log_stream in log_streams:
+            log_stream_name = log_stream["logStreamName"]
+            last_event_timestamp = log_stream.get("lastEventTimestamp", 0)
+            if last_event_timestamp >= start_time_ms:
+                active_streams.append(log_stream_name)
+            else:
+                break
+        if "nextToken" in response:
+            stream_params["nextToken"] = response["nextToken"]
+        else:
+            break
+    # Sort and filter streams
+    active_streams.sort()
+    if stream_filter and active_streams:
+        active_streams = [stream for stream in active_streams if stream_filter in stream]
+    return active_streams
+
+
+def stream_log_events(
+    log_group_name: str,
+    log_stream_name: str,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+    follow: bool = False,
+    client=None,
+) -> Generator[Dict, None, None]:
+    """
+    Stream log events from a specific log stream.
+    Yields:
+        Log events as dictionaries
+    """
+    if not client:
+        client = get_cloudwatch_client()
+    params = {"logGroupName": log_group_name, "logStreamName": log_stream_name, "startFromHead": True}
+    if start_time:
+        params["startTime"] = int(start_time.timestamp() * 1000)
+    if end_time:
+        params["endTime"] = int(end_time.timestamp() * 1000)
+    next_token = None
+    while True:
+        if next_token:
+            params["nextToken"] = next_token
+            params.pop("startTime", None)
+        try:
+            response = client.get_log_events(**params)
+            events = response.get("events", [])
+            for event in events:
+                event["logStreamName"] = log_stream_name
+                yield event
+            next_token = response.get("nextForwardToken")
+            # Break if no more events or same token
+            if not next_token or next_token == params.get("nextToken"):
+                if not follow:
+                    break
+                time.sleep(2)
+        except client.exceptions.ResourceNotFoundException:
+            if not follow:
+                break
+            time.sleep(2)
+
+
+def print_log_event(
+    event: dict, show_stream: bool = True, local_time: bool = True, custom_format: Optional[str] = None
+):
+    """Print a formatted log event."""
+    timestamp = datetime.fromtimestamp(event["timestamp"] / 1000, tz=timezone.utc)
+    if local_time:
+        timestamp = timestamp.astimezone()
+    message = event["message"].rstrip()
+    if custom_format:
+        # Allow custom formatting
+        print(custom_format.format(stream=event.get("logStreamName", ""), time=timestamp, message=message))
+    elif show_stream and "logStreamName" in event:
+        print(f"[{event['logStreamName']}] [{timestamp:%Y-%m-%d %I:%M%p}] {message}")
+    else:
+        print(f"[{timestamp:%H:%M:%S}] {message}")

workbench/utils/config_manager.py

@@ -4,15 +4,12 @@ import os
 import sys
 import platform
 import logging
-import importlib.resources as resources  # noqa: F401 Python 3.9 compatibility
 from typing import Any, Dict
+from importlib.resources import files, as_file
 
 # Workbench imports
 from workbench.utils.license_manager import LicenseManager
-from workbench.utils.execution_environment import running_as_service
-
-# Python 3.9 compatibility
-from workbench.utils.resource_utils import get_resource_path
+from workbench_bridges.utils.execution_environment import running_as_service
 
 
 class FatalConfigError(Exception):
@@ -172,8 +169,7 @@ class ConfigManager:
         Returns:
             str: The open source API key.
         """
-        # Python 3.9 compatibility
-        with get_resource_path("workbench.resources", "open_source_api.key") as open_source_key_path:
+        with as_file(files("workbench.resources").joinpath("open_source_api.key")) as open_source_key_path:
             with open(open_source_key_path, "r") as key_file:
                 return key_file.read().strip()
 

workbench/utils/endpoint_utils.py

@@ -7,9 +7,7 @@ from typing import Union, Optional
 import pandas as pd
 
 # Workbench Imports
-from workbench.api.feature_set import FeatureSet
-from workbench.api.model import Model
-from workbench.api.endpoint import Endpoint
+from workbench.api import FeatureSet, Model, Endpoint
 
 # Set up the log
 log = logging.getLogger("workbench")
@@ -77,7 +75,7 @@ def internal_model_data_url(endpoint_config_name: str, session: boto3.Session) -
         return None
 
 
-def fs_training_data(end: Endpoint) -> pd.DataFrame:
+def get_training_data(end: Endpoint) -> pd.DataFrame:
     """Code to get the training data from the FeatureSet used to train the Model
 
     Args:
@@ -100,7 +98,7 @@ def fs_training_data(end: Endpoint) -> pd.DataFrame:
     return train_df
 
 
-def fs_evaluation_data(end: Endpoint) -> pd.DataFrame:
+def get_evaluation_data(end: Endpoint) -> pd.DataFrame:
     """Code to get the evaluation data from the FeatureSet NOT used for training
 
     Args:
@@ -178,11 +176,11 @@ if __name__ == "__main__":
     print(model_data_url)
 
     # Get the training data
-    my_train_df = fs_training_data(my_endpoint)
+    my_train_df = get_training_data(my_endpoint)
     print(my_train_df)
 
     # Get the evaluation data
-    my_eval_df = fs_evaluation_data(my_endpoint)
+    my_eval_df = get_evaluation_data(my_endpoint)
     print(my_eval_df)
 
     # Backtrack to the FeatureSet

workbench/utils/license_manager.py

@@ -6,15 +6,12 @@ import json
 import logging
 import requests
 from typing import Union
-import importlib.resources as resources  # noqa: F401 Python 3.9 compatibility
 from datetime import datetime
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.asymmetric import padding
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.backends import default_backend
-
-# Python 3.9 compatibility
-from workbench.utils.resource_utils import get_resource_path
+from importlib.resources import files, as_file
 
 
 class FatalLicenseError(Exception):
@@ -140,8 +137,7 @@ class LicenseManager:
         Returns:
             The public key as an object.
         """
-        # Python 3.9 compatibility
-        with get_resource_path("workbench.resources", "signature_verify_pub.pem") as public_key_path:
+        with as_file(files("workbench.resources").joinpath("signature_verify_pub.pem")) as public_key_path:
             with open(public_key_path, "rb") as key_file:
                 public_key_data = key_file.read()