daggerml.util

  1import json
  2import logging
  3import os
  4import time
  5import urllib.request
  6from dataclasses import dataclass
  7from random import randint
  8
  9import boto3
 10from botocore.client import Config
 11from botocore.exceptions import BotoCoreError, NoRegionError
 12
 13logger = logging.getLogger(__name__)
 14
 15
 16def current_time_millis():
 17    return round(time.time() * 1000)
 18
 19
 20@dataclass
 21class BackoffWithJitter:
 22    min: int = 100
 23    max: int = 10000
 24    k: int = 2
 25    state: int = 0
 26
 27    def __call__(self):
 28        if self.state == 0:
 29            self.state = randint(self.min, self.min * 2)
 30            return self.state
 31        self.state = min(self.max, randint(self.min, max(self.min, self.state) * self.k))
 32        return self.state
 33
 34
 35def _get_region_from_metadata():
 36    """
 37    Attempts to retrieve the AWS region from ECS or EC2 metadata.
 38
 39    Returns
 40    -------
 41    Optional[str]
 42        The AWS region string if found, otherwise None.
 43    """
 44    # ECS (used in AWS Batch)
 45    metadata_uri = os.environ.get("ECS_CONTAINER_METADATA_URI_V4") or os.environ.get("ECS_CONTAINER_METADATA_URI")
 46    if metadata_uri:
 47        try:
 48            with urllib.request.urlopen(metadata_uri, timeout=2) as response:
 49                metadata = json.load(response)
 50            cluster_label = metadata.get("Labels", {}).get("com.amazonaws.ecs.cluster", "")
 51            region = cluster_label.split(":")[0] if ":" in cluster_label else None
 52            if region:
 53                return region
 54        except Exception as e:
 55            logger.warning("Failed to get region from ECS metadata: %s", e)
 56    # EC2 fallback: use IMDSv2
 57    try:
 58        token_req = urllib.request.Request(
 59            "http://169.254.169.254/latest/api/token",
 60            method="PUT",
 61            headers={"X-aws-ec2-metadata-token-ttl-seconds": "60"},
 62        )
 63        with urllib.request.urlopen(token_req, timeout=2) as token_response:
 64            token = token_response.read().decode()
 65        region_req = urllib.request.Request(
 66            "http://169.254.169.254/latest/dynamic/instance-identity/document",
 67            headers={"X-aws-ec2-metadata-token": token},
 68        )
 69        with urllib.request.urlopen(region_req, timeout=2) as region_response:
 70            identity_doc = json.load(region_response)
 71        return identity_doc.get("region")
 72    except Exception as e:
 73        logger.warning("Failed to get region from EC2 metadata: %s", e)
 74    return
 75
 76
 77def get_client(
 78    name,
 79    region=None,
 80    default_region="us-east-1",
 81    *,
 82    connection_timeout=5,
 83    read_timeout=60,
 84    max_attempts=5,
 85    retry_mode="adaptive",
 86    max_pool_connections=20,
 87):
 88    """
 89    Creates a robust boto3 client, determining the AWS region in the following order:
 90        1. Explicit argument
 91        2. AWS_REGION / AWS_DEFAULT_REGION environment variables
 92        3. boto3/botocore session
 93        4. ECS/EC2 metadata
 94        5. Fallback default region (us-east-1)
 95
 96    Parameters
 97    ----------
 98    name : str
 99        The name of the AWS service client.
100    region : Optional[str], default=None
101        The AWS region to use.
102    default_region : str, default="us-east-1"
103        The fallback AWS region.
104    connection_timeout : int | float, default=5
105        The connection timeout in seconds.
106    read_timeout : int | float, default=60
107        The read timeout in seconds.
108    max_attempts : int, default=5
109        The maximum number of retry attempts for retriable AWS requests.
110    retry_mode : str, default="adaptive"
111        The botocore retry mode.
112    max_pool_connections : int, default=20
113        The maximum number of pooled HTTP connections.
114
115    Returns
116    -------
117    boto3.client
118        A boto3 client for the specified service.
119    """
120    # Step 1–3: Try common boto3 config methods
121    region = region or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
122    if not region:
123        try:
124            region = boto3.Session().region_name
125        except (BotoCoreError, NoRegionError, ConnectionRefusedError):
126            logger.debug("could not instantiate boto client...")
127            pass
128    # Step 4: Metadata if still no region
129    if not region:
130        logger.debug("inferring aws region from metadata")
131        region = _get_region_from_metadata()
132    # Step 5: Fallback default
133    if not region:
134        logger.warning(f"falling back to default region '{default_region}'")
135        region = default_region
136    config = Config(
137        region_name=region,
138        connect_timeout=connection_timeout,
139        read_timeout=read_timeout,
140        retries={"max_attempts": max_attempts, "mode": retry_mode},
141        max_pool_connections=max_pool_connections,
142    )
143    return boto3.client(name, config=config)