workbench 0.8.167__tar.gz → 0.8.168__tar.gz

{workbench-0.8.167 → workbench-0.8.168}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.167
+Version: 0.8.168
 Summary: Workbench: A Dashboard and Python API for creating and deploying AWS SageMaker Model Pipelines
 Author-email: SuperCowPowers LLC <support@supercowpowers.com>
 License-Expression: MIT

workbench-0.8.168/aws_setup/workbench_core/app.py

@@ -0,0 +1,70 @@
+import os
+import boto3
+import logging
+import aws_cdk as cdk
+from workbench_core.workbench_core_stack import WorkbenchCoreStack, WorkbenchCoreStackProps
+
+# Initialize ConfigManager with error handling
+try:
+    # Temporarily disable logging
+    # logging.disable(logging.CRITICAL)
+    from workbench.utils.config_manager import ConfigManager
+
+    cm = ConfigManager()
+except Exception as e:
+    print(f"Workbench ConfigManager initialization failed: {e}")
+    print("Falling back to environment variables only")
+    cm = None
+finally:
+    # Re-enable logging
+    logging.disable(logging.NOTSET)
+
+# Grab the account and region using boto3
+session = boto3.session.Session()
+aws_account = session.client("sts").get_caller_identity().get("Account")
+aws_region = session.region_name
+print(f"Account: {aws_account}")
+print(f"Region: {aws_region}")
+
+# Get configurations for Workbench Bucket, SSO Group, and Additional Buckets
+workbench_bucket = (cm and cm.get_config("WORKBENCH_BUCKET")) or os.getenv("WORKBENCH_BUCKET")
+if not workbench_bucket:
+    print("Error: WORKBENCH_BUCKET is required but not found in config or environment variables.")
+    exit(1)
+sso_group = (cm and cm.get_config("WORKBENCH_SSO_GROUP")) or os.getenv("WORKBENCH_SSO_GROUP")
+additional_buckets_str = (cm and cm.get_config("WORKBENCH_ADDITIONAL_BUCKETS")) or os.getenv(
+    "WORKBENCH_ADDITIONAL_BUCKETS", ""
+)
+additional_buckets = (
+    [bucket.strip() for bucket in additional_buckets_str.split(",") if bucket.strip()] if additional_buckets_str else []
+)
+existing_vpc_id = (cm and cm.get_config("WORKBENCH_VPC_ID")) or os.getenv("WORKBENCH_VPC_ID")
+subnet_ids = (cm and cm.get_config("WORKBENCH_SUBNET_IDS")) or os.getenv("WORKBENCH_SUBNET_IDS", "")
+
+# Log the configuration for transparency
+print("Configuration:")
+print(f"  WORKBENCH_BUCKET: {workbench_bucket}")
+print(f"  WORKBENCH_SSO_GROUP: {sso_group}")
+print(f"  WORKBENCH_ADDITIONAL_BUCKETS: {additional_buckets}")
+print(f"  WORKBENCH_VPC_ID: {existing_vpc_id}")
+print(f"  WORKBENCH_SUBNET_IDS: {subnet_ids}")
+
+# Our CDK App and Environment
+app = cdk.App()
+env = cdk.Environment(account=aws_account, region=aws_region)
+
+# Create the Workbench Core Stack
+sandbox_stack = WorkbenchCoreStack(
+    app,
+    "WorkbenchCore",
+    env=env,
+    props=WorkbenchCoreStackProps(
+        workbench_bucket=workbench_bucket,
+        sso_group=sso_group,
+        additional_buckets=additional_buckets,
+        existing_vpc_id=existing_vpc_id,
+        subnet_ids=subnet_ids,
+    ),
+)
+
+app.synth()

workbench-0.8.168/aws_setup/workbench_core/requirements.txt

@@ -0,0 +1,3 @@
+aws-cdk-lib>=2.211.0
+constructs>=10.4.2,<11.0.0
+boto3>=1.34.0

{workbench-0.8.167 → workbench-0.8.168}/aws_setup/workbench_core/workbench_core/workbench_core_stack.py

@@ -17,6 +17,8 @@ class WorkbenchCoreStackProps:
     workbench_bucket: str
     sso_group: str
     additional_buckets: List[str] = field(default_factory=list)
+    existing_vpc_id: str = None
+    subnet_ids: List[str] = field(default_factory=list)  # Optional subnets ids for the Batch compute environment
 
 
 class WorkbenchCoreStack(Stack):
@@ -39,6 +41,8 @@ class WorkbenchCoreStack(Stack):
         self.workbench_bucket = props.workbench_bucket
         self.sso_group = props.sso_group
         self.additional_buckets = props.additional_buckets
+        self.existing_vpc_id = props.existing_vpc_id
+        self.subnet_ids = props.subnet_ids
 
         # Workbench Role Names
         self.execution_role_name = "Workbench-ExecutionRole"  # Main role
@@ -377,12 +381,31 @@ class WorkbenchCoreStack(Stack):
     #   Batch Compute   #
     #####################
     def create_batch_compute_environment(self) -> batch.FargateComputeEnvironment:
-        """Create the Batch compute environment - super simple with Fargate."""
+        """Create the Batch compute environment with Fargate."""
+
+        # Do we have an existing VPC we want to use? Otherwise use a default
+        if self.existing_vpc_id:
+            vpc = ec2.Vpc.from_lookup(self, "ImportedVPC", vpc_id=self.existing_vpc_id)
+        else:
+            vpc = ec2.Vpc.from_lookup(self, "DefaultVpc", is_default=True)
+
+        # Use specific subnets if provided, otherwise let CDK choose
+        vpc_subnets = None
+        if self.subnet_ids:
+            vpc_subnets = ec2.SubnetSelection(
+                subnets=[
+                    ec2.Subnet.from_subnet_id(self, f"BatchSubnet{i}", subnet_id)
+                    for i, subnet_id in enumerate(self.subnet_ids)
+                ]
+            )
+
         return batch.FargateComputeEnvironment(
             self,
             "WorkbenchBatchComputeEnvironment",
             compute_environment_name="workbench-compute-env",
-            vpc=ec2.Vpc.from_lookup(self, "DefaultVpc", is_default=True),
+            vpc=vpc,
+            vpc_subnets=vpc_subnets,
+            replace_compute_environment=True,
         )
 
     def create_batch_job_queue(self) -> batch.JobQueue:
@@ -1380,7 +1403,6 @@ class WorkbenchCoreStack(Stack):
 
         # Add policies for the Batch Role
         batch_role.add_to_policy(self.batch_job_logs())
-        # batch_role.add_to_policy(self.ecr_policy_statement())  # Not sure we need this
         batch_role.add_to_policy(self.parameter_store_full())
         batch_role.add_to_policy(self.dataframe_store_full())
         batch_role.add_managed_policy(self.datasource_policy)

{workbench-0.8.167 → workbench-0.8.168}/model_docker_images/ml_pipelines/Dockerfile

@@ -14,12 +14,22 @@ RUN pip install --no-cache-dir -r /tmp/requirements.txt
 # Final runtime image
 FROM python:3.12-slim
 
+# Install required X11 libraries for RDKit Draw module
+RUN apt-get update && apt-get install -y \
+    libxrender1 \
+    libxext6 \
+    libsm6 \
+    libxinerama1 \
+    libfontconfig1 \
+    libcairo2 \
+    && rm -rf /var/lib/apt/lists/*
+
 # Copy installed packages from builder
 COPY --from=builder /usr/local/lib/python3.12/site-packages /usr/local/lib/python3.12/site-packages
 COPY --from=builder /usr/local/bin /usr/local/bin
 
 # Install workbench only
-RUN pip install --no-cache-dir workbench==0.8.166
+RUN pip install --no-cache-dir workbench==0.8.167
 
 # Add the script runner
 COPY ml_pipeline_runner.py /app/ml_pipeline_runner.py

{workbench-0.8.167 → workbench-0.8.168}/model_docker_images/ml_pipelines/ml_pipeline_runner.py

@@ -15,43 +15,38 @@ def download_ml_pipeline_from_s3(s3_path: str, local_path: str):
     parsed = urlparse(s3_path)
     bucket = parsed.netloc
     key = parsed.path.lstrip("/")
-
     log.info(f"Downloading {s3_path} to {local_path}")
-
     s3_client = boto3.client("s3")
     s3_client.download_file(bucket, key, local_path)
 
-    # Make it executable
-    os.chmod(local_path, 0o755)
-    log.info(f"Downloaded and made executable: {local_path}")
-
 
 def run_ml_pipeline(script_path: str):
     """Execute the ML pipeline script."""
     log.info(f"Executing ML pipeline: {script_path}")
-
     try:
         # Run the script with python
-        result = subprocess.run([sys.executable, script_path], capture_output=True, text=True, check=True)
-
-        log.info("ML pipeline completed successfully")
-        log.info(f"STDOUT: {result.stdout}")
-        return 0
-
-    except subprocess.CalledProcessError as e:
-        log.error(f"ML pipeline failed with exit code {e.returncode}")
-        log.error(f"STDOUT: {e.stdout}")
-        log.error(f"STDERR: {e.stderr}")
-        return e.returncode
+        result = subprocess.run(
+            [sys.executable, script_path], check=False, text=True  # Don't raise exception on non-zero exit
+        )
+        if result.returncode == 0:
+            log.info("ML pipeline completed successfully")
+        else:
+            log.error(f"ML pipeline failed with exit code {result.returncode}")
+
+        return result.returncode
+    except Exception as e:
+        log.error(f"ML pipeline execution error: {e}")
+        return 1
 
 
 def main():
+    """Main entry point for the ML pipeline runner."""
+
     # Report the version of the workbench package
     log.info(f"Workbench version: {workbench.__version__}")
 
     # Get the ML Pipeline S3 path from environment variable
     ml_pipeline_s3_path = os.environ.get("ML_PIPELINE_S3_PATH")
-
     if not ml_pipeline_s3_path:
         log.error("ML_PIPELINE_S3_PATH environment variable not set")
         sys.exit(1)
@@ -59,19 +54,14 @@ def main():
     # Extract script name for local file
     script_name = os.path.basename(ml_pipeline_s3_path)
     local_script_path = f"/tmp/{script_name}"
-
     try:
         # Download the ML pipeline script from S3
         download_ml_pipeline_from_s3(ml_pipeline_s3_path, local_script_path)
-
         # Execute the ML pipeline
         exit_code = run_ml_pipeline(local_script_path)
-
         # Clean up
         os.remove(local_script_path)
-
         sys.exit(exit_code)
-
     except Exception as e:
         log.error(f"Error in run_script: {e}")
         sys.exit(1)

{workbench-0.8.167 → workbench-0.8.168}/src/workbench/scripts/ml_pipeline_launcher.py

@@ -8,6 +8,7 @@ from pathlib import Path
 from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
 from workbench.utils.config_manager import ConfigManager
 from workbench.utils.s3_utils import upload_content_to_s3
+from workbench.utils.cloudwatch_utils import get_cloudwatch_logs_url
 
 log = logging.getLogger("workbench")
 cm = ConfigManager()
@@ -27,10 +28,9 @@ def get_batch_role_arn() -> str:
 
 
 def ensure_job_definition():
-    """Ensure the job definition exists with network configuration."""
+    """Register or update the Batch job definition for ML pipeline runner."""
     batch = AWSAccountClamp().boto3_session.client("batch")
     name = "workbench-ml-pipeline-runner"
-
     response = batch.register_job_definition(
         jobDefinitionName=name,
         type="container",
@@ -40,27 +40,42 @@ def ensure_job_definition():
             "resourceRequirements": [{"type": "VCPU", "value": "2"}, {"type": "MEMORY", "value": "4096"}],
             "jobRoleArn": get_batch_role_arn(),
             "executionRoleArn": get_batch_role_arn(),
-            "environment": [{"name": "WORKBENCH_BUCKET", "value": workbench_bucket}],
-            "networkConfiguration": {"assignPublicIp": "ENABLED"},  # This is required so the ECR image can be pulled
+            "environment": [
+                {"name": "WORKBENCH_BUCKET", "value": workbench_bucket},
+                {"name": "PYTHONUNBUFFERED", "value": "1"},
+            ],
+            # "networkConfiguration": {"assignPublicIp": "ENABLED"},  # Required for ECR Image Pull (when not in VPC)
         },
         timeout={"attemptDurationSeconds": 10800},  # 3 hours
     )
-
     log.info(f"Job definition ready: {name} (revision {response['revision']})")
     return name
 
 
 def run_batch_job(script_path: str) -> int:
-    """Upload script, submit job, and track to completion."""
+    """
+    Submit and monitor an AWS Batch job for ML pipeline execution.
+    This function:
+    1. Uploads the ML pipeline script to S3
+    2. Submits a Batch job to run the script in a container
+    3. Monitors job status until completion
+    4. Returns the job's exit code
+
+    Args:
+        script_path: Local path to the ML pipeline script
+
+    Returns:
+        Exit code from the batch job (0 for success, non-zero for failure)
+    """
     batch = AWSAccountClamp().boto3_session.client("batch")
     script_name = Path(script_path).stem
 
-    # Upload script to S3
+    # Upload script to S3 for the container to download
     s3_path = f"s3://{workbench_bucket}/batch-jobs/{Path(script_path).name}"
     log.info(f"Uploading script to {s3_path}")
     upload_content_to_s3(Path(script_path).read_text(), s3_path)
 
-    # Submit job
+    # Submit the Batch job
     job_name = f"workbench_{script_name}_{datetime.now():%Y%m%d_%H%M%S}"
     response = batch.submit_job(
         jobName=job_name,
@@ -68,35 +83,50 @@ def run_batch_job(script_path: str) -> int:
         jobDefinition=ensure_job_definition(),
         containerOverrides={
             "environment": [
-                {"name": "SCRIPT_S3_PATH", "value": s3_path},
+                {"name": "ML_PIPELINE_S3_PATH", "value": s3_path},
                 {"name": "WORKBENCH_BUCKET", "value": workbench_bucket},
             ]
         },
     )
-
     job_id = response["jobId"]
     log.info(f"Submitted job: {job_name} ({job_id})")
 
-    # Track job to completion
+    # Monitor job execution
+    last_status = None
     while True:
+        # Check job status
         job = batch.describe_jobs(jobs=[job_id])["jobs"][0]
         status = job["status"]
-        log.info(f"Job status: {status}")
+        if status != last_status:
+            log.info(f"Job status: {status}")
+            last_status = status
 
+        # Check if job completed
         if status in ["SUCCEEDED", "FAILED"]:
             exit_code = job.get("attempts", [{}])[-1].get("exitCode", 1)
             if status == "FAILED":
                 log.error(f"Job failed: {job.get('statusReason', 'Unknown reason')}")
+            else:
+                log.info("Job completed successfully")
+
+            # Get CloudWatch logs URL
+            log_stream_name = job.get("container", {}).get("logStreamName")
+            logs_url = get_cloudwatch_logs_url(log_group="/aws/batch/job", log_stream=log_stream_name)
+            if logs_url:
+                # OSC 8 hyperlink format for modern terminals
+                clickable_url = f"\033]8;;{logs_url}\033\\{logs_url}\033]8;;\033\\"
+                log.info(f"View logs: {clickable_url}")
             return exit_code
 
-        time.sleep(30)
+        # Sleep a bit before next status check
+        time.sleep(10)
 
 
 def main():
+    """CLI entry point for running ML pipelines on AWS Batch."""
     parser = argparse.ArgumentParser(description="Run ML pipeline script on AWS Batch")
     parser.add_argument("script_file", help="Local path to ML pipeline script")
     args = parser.parse_args()
-
     try:
         exit_code = run_batch_job(args.script_file)
         exit(exit_code)

{workbench-0.8.167 → workbench-0.8.168}/src/workbench/scripts/monitor_cloud_watch.py

@@ -4,8 +4,10 @@ import sys
 import time
 import argparse
 from datetime import datetime, timedelta, timezone
-from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
+
+# Workbench Imports
 from workbench.utils.repl_utils import cprint, Spinner
+from workbench.utils.cloudwatch_utils import get_cloudwatch_client, get_active_log_streams, stream_log_events
 
 # Define the log levels to include all log levels above the specified level
 log_level_map = {
@@ -33,64 +35,6 @@ def date_display(dt):
         return dt.strftime("%Y-%m-%d %I:%M%p") + "(UTC)"
 
 
-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_active_log_streams(client, log_group_name, start_time_ms, stream_filter=None):
-    """Retrieve log streams that have events after the specified start time."""
-
-    # Get all the streams in the log group
-    active_streams = []
-    stream_params = {
-        "logGroupName": log_group_name,
-        "orderBy": "LastEventTime",
-        "descending": True,
-    }
-
-    # Loop to retrieve all log streams (maximum 50 per call)
-    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")
-
-            # Include streams with events since the specified start time
-            # Note: There's some issue where the last event timestamp is 'off'
-            #       so we're going to add 60 minutes from the last event timestamp
-            last_event_timestamp += 60 * 60 * 1000
-            if last_event_timestamp >= start_time_ms:
-                active_streams.append(log_stream_name)
-            else:
-                break  # Stop if we reach streams older than the start time
-
-        # Check if there are more streams to retrieve
-        if "nextToken" in response:
-            stream_params["nextToken"] = response["nextToken"]
-        else:
-            break
-
-    # Sort and report the active log streams
-    active_streams.sort()
-    if active_streams:
-        print("Active log streams:", len(active_streams))
-
-    # Filter the active streams by a substring if provided
-    if stream_filter and active_streams:
-        print(f"Filtering active log streams by '{stream_filter}'...")
-        active_streams = [stream for stream in active_streams if stream_filter in stream]
-
-    for stream in active_streams:
-        print(f"\t - {stream}")
-
-    # Return the active log streams
-    return active_streams
-
-
 def get_latest_log_events(client, log_group_name, start_time, end_time=None, stream_filter=None):
     """Retrieve the latest log events from the active/filtered log streams in a CloudWatch Logs group."""
 
@@ -99,11 +43,15 @@ def get_latest_log_events(client, log_group_name, start_time, end_time=None, str
         get_latest_log_events.first_run = True
 
     log_events = []
-    start_time_ms = int(start_time.timestamp() * 1000)  # Convert start_time to milliseconds
+    start_time_ms = int(start_time.timestamp() * 1000)
+
+    # Use the util function to get active streams
+    active_streams = get_active_log_streams(log_group_name, start_time_ms, stream_filter, client)
 
-    # Get the active log streams with events since start_time
-    active_streams = get_active_log_streams(client, log_group_name, start_time_ms, stream_filter)
     if active_streams:
+        print(f"Active log streams: {len(active_streams)}")
+        for stream in active_streams:
+            print(f"\t - {stream}")
         print(f"Processing log events from {date_display(start_time)} on {len(active_streams)} active log streams...")
         get_latest_log_events.first_run = False
     else:
@@ -114,50 +62,22 @@ def get_latest_log_events(client, log_group_name, start_time, end_time=None, str
             print("Monitoring for new events...")
         return log_events
 
-    # Iterate over the active streams and fetch log events
+    # Use the util function to stream events from each log stream
     for log_stream_name in active_streams:
-        params = {
-            "logGroupName": log_group_name,
-            "logStreamName": log_stream_name,
-            "startTime": start_time_ms,  # Use start_time in milliseconds
-            "startFromHead": True,  # Start from the nearest event to start_time
-        }
-        next_event_token = None
-        if end_time is not None:
-            params["endTime"] = int(end_time.timestamp() * 1000)
-
-        # Process the log events from this log stream
         spinner = Spinner("lightpurple", f"Pulling events from {log_stream_name}:")
         spinner.start()
         log_stream_events = 0
 
-        # Get the log events for the active log stream
-        while True:
-            if next_event_token:
-                params["nextToken"] = next_event_token
-                params.pop("startTime", None)  # Remove startTime when using nextToken
+        # Stream events using the util function
+        for event in stream_log_events(
+            log_group_name, log_stream_name, start_time, end_time, follow=False, client=client
+        ):
+            log_stream_events += 1
+            log_events.append(event)
 
-            # Fetch the log events (this call takes a while: optimize if we can)
-            events_response = client.get_log_events(**params)
+        spinner.stop()
+        print(f"Processed {log_stream_events} events from {log_stream_name} (Total: {len(log_events)})")
 
-            events = events_response.get("events", [])
-            for event in events:
-                event["logStreamName"] = log_stream_name
-
-            # Add the log stream events to our list of all log events
-            log_stream_events += len(events)
-            log_events.extend(events)
-
-            # Handle pagination for log events
-            next_event_token = events_response.get("nextForwardToken")
-
-            # Break the loop if there are no more events to fetch
-            if not next_event_token or next_event_token == params.get("nextToken"):
-                spinner.stop()
-                print(f"Processed {log_stream_events} events from {log_stream_name} (Total: {len(log_events)})")
-                break
-
-    # Return the log events
     return log_events
 
 
@@ -206,6 +126,7 @@ def monitor_log_group(
     print(f"Monitoring log group: {log_group_name} from {date_display(start_time)}")
     print(f"Log levels: {log_levels}")
     print(f"Search terms: {search_terms}")
+
     while True:
         # Get the latest log events with stream filtering if provided
         all_log_events = get_latest_log_events(client, log_group_name, start_time, end_time, stream_filter)
@@ -218,7 +139,6 @@ def monitor_log_group(
 
                 # Check the search terms
                 if not search_terms or any(term in event["message"].lower() for term in search_terms):
-
                     # Calculate the start and end index for this match
                     start_index = max(i - before, 0)
                     end_index = min(i + after, len(all_log_events) - 1)

workbench-0.8.168/src/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}")