Skip to content

lablink_allocator_service.utils.aws_utils

Bases: RuntimeError

IMDSv2 unreachable (running outside EC2).

Check if a given EC2 instance type supports NVIDIA GPUs. Args: machine_type (str): The EC2 instance type to check. Returns: bool: True if the instance type supports NVIDIA GPUs, False otherwise.

Source code in packages/allocator/src/lablink_allocator_service/utils/aws_utils.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def check_support_nvidia(machine_type) -> bool:
    """Check if a given EC2 instance type supports NVIDIA GPUs.
    Args:
        machine_type (str): The EC2 instance type to check.
    Returns:
        bool: True if the instance type supports NVIDIA GPUs, False otherwise.
    """
    kwargs = {
        "region_name": os.getenv("AWS_REGION", "us-west-2"),
        "aws_access_key_id": os.getenv("AWS_ACCESS_KEY_ID"),
        "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
    }
    if os.getenv("AWS_SESSION_TOKEN"):
        kwargs["aws_session_token"] = os.getenv("AWS_SESSION_TOKEN")

    ec2 = boto3.client("ec2", **kwargs)
    try:
        response = ec2.describe_instance_types(InstanceTypes=[machine_type])
        gpu_info = response["InstanceTypes"][0].get("GpuInfo", {})

        # Check if GPU is present
        if not gpu_info:
            logger.debug(f"No GPU info found for instance type {machine_type}.")
            return False

        # Check if any GPU supports NVIDIA
        for gpu in gpu_info.get("Gpus", []):
            if "NVIDIA" in gpu.get("Manufacturer", ""):
                logger.info(f"Instance type {machine_type} supports NVIDIA GPUs.")
                return True

        logger.debug(f"Instance type {machine_type} does not support NVIDIA GPUs.")
    except ClientError as e:
        logger.error(f"Error checking NVIDIA support for {machine_type}: {e}")
    return False

Return the primary SG ID of the EC2 this process runs on.

Parameters:

Name Type Description Default
region str

AWS region. Caller should pass cfg.app.region — boto3 won't auto-detect inside the allocator container even though the instance metadata is reachable, because the container's env doesn't have AWS_REGION set.

 'us-west-2'

Raises NotOnEC2Error if IMDSv2 is unreachable.

Source code in packages/allocator/src/lablink_allocator_service/utils/aws_utils.py 
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
def current_instance_security_group(region: str = "us-west-2") -> str:
    """Return the primary SG ID of the EC2 this process runs on.

    Args:
        region: AWS region. Caller should pass cfg.app.region — boto3
            won't auto-detect inside the allocator container even though
            the instance metadata is reachable, because the container's
            env doesn't have AWS_REGION set.

    Raises NotOnEC2Error if IMDSv2 is unreachable.
    """
    try:
        token = _imds_token()
        instance_id = _imds_get("meta-data/instance-id", token)
    except (ConnectionError, _RequestException) as exc:
        raise NotOnEC2Error(str(exc)) from exc
    ec2 = boto3.client("ec2", region_name=region)
    resp = ec2.describe_instances(InstanceIds=[instance_id])
    sgs = resp["Reservations"][0]["Instances"][0]["SecurityGroups"]
    if not sgs:
        raise NotOnEC2Error(f"instance {instance_id} has no SGs")
    return sgs[0]["GroupId"]

Look up an EC2 instance ID by its Name tag.

Parameters:

Name Type Description Default
name str

The Name tag value of the EC2 instance.

 required
region str

The AWS region to search in.

 'us-west-2'

Returns:

Type Description
Optional[str]

The instance ID if found, None otherwise.
Source code in packages/allocator/src/lablink_allocator_service/utils/aws_utils.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def get_instance_id_by_name(name: str, region: str = "us-west-2") -> Optional[str]:
    """Look up an EC2 instance ID by its Name tag.

    Args:
        name: The Name tag value of the EC2 instance.
        region: The AWS region to search in.

    Returns:
        The instance ID if found, None otherwise.
    """
    kwargs = {
        "region_name": region,
        "aws_access_key_id": os.getenv("AWS_ACCESS_KEY_ID"),
        "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
    }
    if os.getenv("AWS_SESSION_TOKEN"):
        kwargs["aws_session_token"] = os.getenv("AWS_SESSION_TOKEN")

    ec2 = boto3.client("ec2", **kwargs)
    try:
        response = ec2.describe_instances(
            Filters=[
                {"Name": "tag:Name", "Values": [name]},
                {
                    "Name": "instance-state-name",
                    "Values": ["running", "stopped", "pending"],
                },
            ]
        )
        for reservation in response.get("Reservations", []):
            for instance in reservation.get("Instances", []):
                return instance["InstanceId"]
    except ClientError as e:
        logger.error(f"Error looking up instance by name '{name}': {e}")
    return None

Get the private IP address of an EC2 instance.

Parameters:

Name Type Description Default
instance_id str

The EC2 instance ID.

 required
region str

The AWS region where the instance is located.

 'us-west-2'

Returns:

Type Description
Optional[str]

The private IP address if available, None otherwise.
Source code in packages/allocator/src/lablink_allocator_service/utils/aws_utils.py 
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def get_instance_private_ip(
    instance_id: str, region: str = "us-west-2"
) -> Optional[str]:
    """Get the private IP address of an EC2 instance.

    Args:
        instance_id: The EC2 instance ID.
        region: The AWS region where the instance is located.

    Returns:
        The private IP address if available, None otherwise.
    """
    kwargs = {
        "region_name": region,
        "aws_access_key_id": os.getenv("AWS_ACCESS_KEY_ID"),
        "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
    }
    if os.getenv("AWS_SESSION_TOKEN"):
        kwargs["aws_session_token"] = os.getenv("AWS_SESSION_TOKEN")

    ec2 = boto3.client("ec2", **kwargs)
    try:
        response = ec2.describe_instances(InstanceIds=[instance_id])
        for reservation in response.get("Reservations", []):
            for instance in reservation.get("Instances", []):
                ip = instance.get("PrivateIpAddress")
                if ip:
                    return ip
    except ClientError as e:
        logger.error(f"Error getting private IP for instance {instance_id}: {e}")
    return None

Get the public IP address of an EC2 instance.

Parameters:

Name Type Description Default
instance_id str

The EC2 instance ID.

 required
region str

The AWS region where the instance is located.

 'us-west-2'

Returns:

Type Description
Optional[str]

The public IP address if available, None otherwise.
Source code in packages/allocator/src/lablink_allocator_service/utils/aws_utils.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
def get_instance_public_ip(
    instance_id: str, region: str = "us-west-2"
) -> Optional[str]:
    """Get the public IP address of an EC2 instance.

    Args:
        instance_id: The EC2 instance ID.
        region: The AWS region where the instance is located.

    Returns:
        The public IP address if available, None otherwise.
    """
    kwargs = {
        "region_name": region,
        "aws_access_key_id": os.getenv("AWS_ACCESS_KEY_ID"),
        "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
    }
    if os.getenv("AWS_SESSION_TOKEN"):
        kwargs["aws_session_token"] = os.getenv("AWS_SESSION_TOKEN")

    ec2 = boto3.client("ec2", **kwargs)
    try:
        response = ec2.describe_instances(InstanceIds=[instance_id])
        for reservation in response.get("Reservations", []):
            for instance in reservation.get("Instances", []):
                ip = instance.get("PublicIpAddress")
                if ip:
                    return ip
    except ClientError as e:
        logger.error(f"Error getting public IP for instance {instance_id}: {e}")
    return None

Stop and start an EC2 instance (last-resort fallback).

Calls stop_instances, waits for stopped state, then starts again. Note: This does NOT clear cloud-init state, so user_data.sh will not re-run. This is a best-effort restart for cases where SSH is unreachable (e.g., hung processes, OOM).

Parameters:

Name Type Description Default
instance_id str

The EC2 instance ID.

 required
region str

The AWS region where the instance is located.

 'us-west-2'

Returns:

Type Description
bool

True if stop/start completed successfully, False otherwise.
Source code in packages/allocator/src/lablink_allocator_service/utils/aws_utils.py 
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
def stop_start_ec2_instance(instance_id: str, region: str = "us-west-2") -> bool:
    """Stop and start an EC2 instance (last-resort fallback).

    Calls stop_instances, waits for stopped state, then starts again.
    Note: This does NOT clear cloud-init state, so user_data.sh will
    not re-run. This is a best-effort restart for cases where SSH is
    unreachable (e.g., hung processes, OOM).

    Args:
        instance_id: The EC2 instance ID.
        region: The AWS region where the instance is located.

    Returns:
        True if stop/start completed successfully, False otherwise.
    """
    kwargs = {
        "region_name": region,
        "aws_access_key_id": os.getenv("AWS_ACCESS_KEY_ID"),
        "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
    }
    if os.getenv("AWS_SESSION_TOKEN"):
        kwargs["aws_session_token"] = os.getenv("AWS_SESSION_TOKEN")

    ec2 = boto3.client("ec2", **kwargs)
    try:
        logger.info(f"Stopping instance {instance_id}...")
        ec2.stop_instances(InstanceIds=[instance_id])
        waiter = ec2.get_waiter("instance_stopped")
        waiter.wait(
            InstanceIds=[instance_id],
            WaiterConfig={"Delay": 10, "MaxAttempts": 40},
        )
        logger.info(f"Instance {instance_id} stopped, starting...")
        ec2.start_instances(InstanceIds=[instance_id])
        logger.info(f"Start initiated for instance {instance_id}")
        return True
    except Exception as e:
        logger.error(f"Failed to stop/start instance {instance_id}: {e}")
        return False

Uploads a file to an S3 bucket.

Parameters:

Name Type Description Default
local_path Path

The local file path to upload.

 required
env str

The environment (e.g., dev, test, prod) for the upload.

 required
bucket_name str

The name of the S3 bucket to upload to.

 required
region str

The AWS region where the S3 bucket is located.

 required
deployment_name str

The deployment name for S3 key scoping. Defaults to "lablink".

 'lablink'
kms_key_id Optional[str]

The KMS key ID for server-side encryption. Defaults to None.

 None
Source code in packages/allocator/src/lablink_allocator_service/utils/aws_utils.py 
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
def upload_to_s3(
    local_path: Path,
    env: str,
    bucket_name: str,
    region: str,
    deployment_name: str = "lablink",
    kms_key_id: Optional[str] = None,
) -> None:
    """Uploads a file to an S3 bucket.

    Args:
        local_path (Path): The local file path to upload.
        env (str): The environment (e.g., dev, test, prod) for the upload.
        bucket_name (str): The name of the S3 bucket to upload to.
        region (str): The AWS region where the S3 bucket is located.
        deployment_name (str): The deployment name for S3 key scoping.
            Defaults to "lablink".
        kms_key_id (Optional[str], optional): The KMS key ID for server-side encryption.
            Defaults to None.
    """
    s3 = boto3.client("s3", region_name=region)
    key = f"{deployment_name}/{env}/client/{local_path.name}"
    extra = {"ContentType": "text/plain"}
    if kms_key_id:
        extra.update({"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": kms_key_id})

    # Upload the variable file
    s3.upload_file(local_path, bucket_name, key, ExtraArgs=extra)