Merge pull request #149 from awslabs/fix/aws-resource-cleanup-leaks

fix: resolve AWS resource and launch template cleanup leaks

Author: fgogolli

config/default_config.json

@@ -101,6 +101,17 @@
           "spot_fleet_request_expiry": 30,
           "percent_on_demand": 0,
           "root_device_volume_size": 20
+        },
+        "cleanup": {
+          "enabled": true,
+          "delete_launch_template": true,
+          "dry_run": false,
+          "resources": {
+            "asg": true,
+            "ec2_fleet": true,
+            "spot_fleet": true,
+            "run_instances": true
+          }
         }
       }
     }
@@ -393,16 +404,5 @@
       "asg": "",
       "tag": ""
     }
-  },
-  "cleanup": {
-    "enabled": true,
-    "delete_launch_template": true,
-    "dry_run": false,
-    "resources": {
-      "asg": true,
-      "ec2_fleet": true,
-      "spot_fleet": true,
-      "run_instances": true
-    }
   }
 }

pyproject.toml

@@ -304,6 +304,8 @@ markers = [
     "provider_contract: part of the provider contract test suite",
     "simulator_limitation: assertion weakened due to known simulator gap",
     "requires_real_provider: skipped unless --run-real-provider flag passed",
+    "moto: marks tests that use moto-mocked AWS",
+    "cli: marks CLI integration tests",
 ]
 env = [
     "AWS_DEFAULT_REGION=us-east-1",
@@ -416,8 +418,16 @@ combine-as-imports = true
 [tool.pyright]
 pythonVersion = "3.12"
 pythonPlatform = "Linux"
-include = ["src/orb"]  # keep in sync with .project.yml build.package_root
-exclude = [".venv", "build", "dist"]
+include = ["src/orb", "tests/shared", "tests/onmoto"]  # keep in sync with .project.yml build.package_root
+exclude = [
+    ".venv", "build", "dist",
+    # onmoto files with pre-existing type errors, excluded until resolved
+    "tests/onmoto/test_config_driven_provision.py",
+    "tests/onmoto/test_cqrs_control_loop.py",
+    "tests/onmoto/test_hf_contract.py",
+    "tests/onmoto/test_provision_lifecycle.py",
+]
+extraPaths = ["."]
 typeCheckingMode = "basic"
 reportMissingImports = true
 reportMissingTypeStubs = false

src/orb/config/schemas/app_schema.py

@@ -5,7 +5,6 @@
 
 from pydantic import BaseModel, Field, field_validator, model_validator
 
-from .cleanup_schema import CleanupConfig
 from .common_schema import (
     DatabaseConfig,
     EventsConfig,
@@ -43,7 +42,6 @@ class AppConfig(BaseModel):
     performance: PerformanceConfig = Field(default_factory=lambda: PerformanceConfig())  # type: ignore[call-arg]
     server: ServerConfig = Field(default_factory=lambda: ServerConfig())  # type: ignore[call-arg]
     native_spec: NativeSpecConfig = Field(default_factory=lambda: NativeSpecConfig())  # type: ignore[call-arg]
-    cleanup: CleanupConfig = Field(default_factory=lambda: CleanupConfig())  # type: ignore[call-arg]
     environment: str = Field("development", description="Environment")
     debug: bool = Field(False, description="Debug mode")
     request_timeout: int = Field(300, description="Request timeout in seconds")

src/orb/config/schemas/provider_strategy_schema.py

@@ -6,6 +6,7 @@
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 from .base_config import BaseCircuitBreakerConfig
+from .cleanup_schema import CleanupConfig
 
 
 class HandlerConfig(BaseModel):
@@ -38,8 +39,9 @@ class ProviderDefaults(BaseModel):
         default_factory=dict, description="Template defaults for this provider type"
     )
     extensions: Optional[dict[str, Any]] = Field(
-        None, description="Provider-specific extensions configuration"
+        default=None, description="Provider-specific extensions configuration"
     )
+    cleanup: Optional[CleanupConfig] = Field(default=None)
 
 
 class ProviderMode(str, Enum):
@@ -95,7 +97,6 @@ class ProviderInstanceConfig(BaseModel):
     enabled: bool = Field(True, description="Whether this provider is enabled")
     priority: int = Field(0, description="Provider priority (lower = higher priority)")
     weight: int = Field(100, description="Provider weight for load balancing")
-    # Keep dict for backward compatibility
     config: dict[str, Any] = Field(
         default_factory=dict, description="Provider-specific configuration"
     )

src/orb/domain/base/ports/configuration_port.py

@@ -93,10 +93,6 @@ def set_configuration_value(self, key: str, value: Any) -> None:
     def get_resource_prefix(self, resource_type: str) -> str:
         """Get resource naming prefix for the given resource type."""
 
-    @abstractmethod
-    def get_cleanup_config(self) -> dict[str, Any]:
-        """Get cleanup configuration."""
-
     @abstractmethod
     def get_config_file_path(self) -> str:
         """Get the config file path."""

src/orb/infrastructure/adapters/configuration_adapter.py

@@ -334,37 +334,11 @@ def get_resource_prefix(self, resource_type: str) -> str:
             if hasattr(resource_config.prefixes, resource_type):
                 return getattr(resource_config.prefixes, resource_type)
             return resource_config.default_prefix
-        except Exception:
-            return ""
-
-    def get_cleanup_config(self) -> dict[str, Any]:
-        """Get cleanup configuration."""
-        try:
-            cleanup = self._config_manager.app_config.cleanup
-            return {
-                "enabled": cleanup.enabled,
-                "delete_launch_template": cleanup.delete_launch_template,
-                "dry_run": cleanup.dry_run,
-                "resources": {
-                    "asg": cleanup.resources.asg,
-                    "ec2_fleet": cleanup.resources.ec2_fleet,
-                    "spot_fleet": cleanup.resources.spot_fleet,
-                    "run_instances": cleanup.resources.run_instances,
-                },
-            }
         except Exception as e:
-            _logger.warning("Failed to load cleanup config, using defaults: %s", e)
-            return {
-                "enabled": True,
-                "delete_launch_template": True,
-                "dry_run": False,
-                "resources": {
-                    "asg": True,
-                    "ec2_fleet": True,
-                    "spot_fleet": True,
-                    "run_instances": True,
-                },
-            }
+            _logger.warning(
+                "Failed to get resource prefix for '%s', using empty prefix: %s", resource_type, e
+            )
+            return ""
 
     def get_active_provider_override(self) -> str | None:
         """Get current provider override from CLI."""

src/orb/infrastructure/scheduler/hostfactory/field_mappings.py

@@ -31,6 +31,7 @@ class HostFactoryFieldMappings:
             "instanceTags": "tags",  # Will be parsed from string format
             # Template metadata
             "name": "name",
+            "requestId": "request_id",
             "providerName": "provider_name",
             "providerApi": "provider_api",
             "providerType": "provider_type",

src/orb/interface/mcp/server/core.py

@@ -58,6 +58,7 @@ def _register_core_tools(self) -> None:
         from orb.interface.request_command_handlers import (
             handle_get_request_status,
             handle_get_return_requests,
+            handle_list_requests,
             handle_request_machines,
             handle_request_return_machines,
         )
@@ -86,6 +87,7 @@ def _register_core_tools(self) -> None:
 
         # Request tools
         self.tools["get_request_status"] = handle_get_request_status
+        self.tools["list_requests"] = handle_list_requests
         self.tools["request_machines"] = handle_request_machines
         self.tools["list_return_requests"] = handle_get_return_requests
         self.tools["return_machines"] = handle_request_return_machines

src/orb/interface/request_command_handlers.py

@@ -424,6 +424,21 @@ async def handle_request_return_machines(args: "argparse.Namespace") -> dict[str
     return {"requestId": request_id, "message": "Return request created successfully"}
 
 
+@handle_interface_exceptions(context="list_requests", interface_type="cli")
+async def handle_list_requests(args: "argparse.Namespace") -> dict[str, Any]:
+    """List all active provisioning requests."""
+    container = get_container()
+    query_bus = container.get(QueryBus)
+    scheduler_strategy = container.get(SchedulerPort)
+
+    from orb.application.dto.queries import ListActiveRequestsQuery
+
+    query = ListActiveRequestsQuery(all_resources=True)
+    request_dtos = await query_bus.execute(query)
+
+    return scheduler_strategy.format_request_status_response(request_dtos)
+
+
 @handle_interface_exceptions(context="cancel_request", interface_type="cli")
 async def handle_cancel_request(args: "argparse.Namespace") -> dict[str, Any]:
     """Handle cancel request operations."""

src/orb/providers/aws/infrastructure/handlers/asg/capacity_manager.py

@@ -138,7 +138,24 @@ def release_instances(
 
         if not asg_details:
             self._logger.warning(
-                "ASG details missing for %s, terminating instances without ASG operations",
+                "ASG details missing for %s, attempting direct describe before cleanup",
+                asg_name,
+            )
+            try:
+                response = self._retry_with_backoff(
+                    self._aws_client.autoscaling_client.describe_auto_scaling_groups,
+                    operation_type="read_only",
+                    AutoScalingGroupNames=[asg_name],
+                )
+                groups = response.get("AutoScalingGroups", [])
+                if groups:
+                    asg_details = groups[0]
+            except Exception as exc:
+                self._logger.warning("Retry describe for ASG %s also failed: %s", asg_name, exc)
+
+        if not asg_details:
+            self._logger.warning(
+                "ASG details unavailable for %s, terminating instances and proceeding with cleanup",
                 asg_name,
             )
             self._aws_ops.terminate_instances_with_fallback(
@@ -151,6 +168,8 @@ def release_instances(
                 asg_name,
                 instance_ids,
             )
+            self._call_delete_asg(asg_name)
+            self._cleanup_on_zero_capacity("asg", asg_name)
             return
 
         # Detach instances (API limit: 50 per call; use 20 for safety)