Use latest AMI for cluster (#180)

Author: prateekdesai04

src/autogluon/cloud/cluster/ray_aws_cluster_config_generator.py

@@ -1,6 +1,7 @@
 import os
 from typing import Any, Dict, List, Optional, Union
 
+from ..utils.aws_utils import get_latest_amazon_linux_ami
 from ..utils.utils import get_utc_timestamp_now
 from .cluster_config_generator import DEFAULT_CONFIG_LOCATION
 from .constants import (
@@ -36,6 +37,7 @@ def __init__(
         config: Optional[Union[str, Dict[str, Any]]] = None,
         cluster_name: Optional[str] = None,
         region: str = "us-east-1",
+        use_latest_ami: bool = True,
         **kwargs,
     ) -> None:
         """
@@ -57,6 +59,13 @@ def __init__(
             self._update_cluster_name(cluster_name)
             self._update_region(region)
 
+            if use_latest_ami:
+                try:
+                    latest_ami = get_latest_amazon_linux_ami()
+                    self._update_ami(latest_ami)
+                except Exception as e:
+                    print(f"Warning: Could not fetch latest Amazon Linux AMI : {e}")
+
     def _update_config(
         self,
         instance_type: Optional[str] = None,

src/autogluon/cloud/utils/aws_utils.py

@@ -5,6 +5,21 @@
 from botocore.config import Config
 
 
+def get_latest_amazon_linux_ami(region="us-east-1", version="al2023"):
+    ec2_client = boto3.client("ec2", region_name=region)
+    filters = [
+        {"Name": "owner-alias", "Values": ["amazon"]},
+        {"Name": "name", "Values": [f"{version}-ami-*"]},  # Amazon Linux 2 or AL2023
+        {"Name": "state", "Values": ["available"]},
+    ]
+    response = ec2_client.describe_images(Filters=filters, Owners=["amazon"])
+    if not response["Images"]:
+        raise ValueError("No Amazon Linux AMI found!")
+
+    latest_ami = sorted(response["Images"], key=lambda x: x["CreationDate"], reverse=True)[0]
+    return latest_ami["ImageId"]
+
+
 def setup_sagemaker_session(
     config: Optional[Config] = None,
     connect_timeout: int = 60,

tests/unittests/cluster/test_cluster_config_generator.py

@@ -24,6 +24,7 @@
     SSH_PRIVATE_KEY,
     VOLUME_SIZE,
 )
+from autogluon.cloud.utils.aws_utils import get_latest_amazon_linux_ami
 
 
 def _create_config_file(config):
@@ -47,10 +48,12 @@ def test_generate_config(config_generator, config):
 
 @pytest.mark.parametrize("config", [None, {"foo": "bar"}, "dummy.yaml"])
 def test_update_ray_aws_cluster_config(config):
+    latest_ami = get_latest_amazon_linux_ami()
+
     with tempfile.TemporaryDirectory() as temp_dir:
         os.chdir(temp_dir)
         _create_config_file(config)
-        config_generator = RayAWSClusterConfigGenerator(config)
+        config_generator = RayAWSClusterConfigGenerator(config, use_latest_ami=False)
         # Test update
         dummy_config = {"cluster_name": "foo"}
         config_generator.update_config(dummy_config)
@@ -60,7 +63,7 @@ def test_update_ray_aws_cluster_config(config):
             instance_type="foo",
             instance_count=2,
             volumes_size=2,
-            ami="dummy_ami",
+            ami=latest_ami,
             custom_image_uri="bar",
             ssh_key_path="dummy.pem",
             initialization_commands=["abc"],
@@ -72,7 +75,7 @@ def test_update_ray_aws_cluster_config(config):
         assert node_config[BLOCK_DEVICE_MAPPINGS][0][EBS][VOLUME_SIZE] == 2
         assert config_generator.config[DOCKER][IMAGE] == "bar"
         assert config_generator.config[AUTH][SSH_PRIVATE_KEY] == os.path.abspath("dummy.pem")
-        assert node_config[IMAGE_ID] == "dummy_ami"
+        assert node_config[IMAGE_ID] == latest_ami
         assert node_config[KEY_NAME] == "dummy"
         assert config_generator.config[INITIALIZATION_COMMANDS] == ["abc"]
         config = config_generator.config