refactor: integrate LoadBalancerSls with dynamic GraphQL deployment

Author: pandyamarut

src/tetra_rp/core/resources/load_balancer_sls/integration.py

@@ -9,7 +9,7 @@
 import logging
 from typing import List, Optional
 
-from tetra_rp.core.resources import ServerlessResource
+from tetra_rp.core.resources import ServerlessResource, ResourceManager
 from .client import LoadBalancerSls
 
 log = logging.getLogger(__name__)
@@ -25,12 +25,12 @@ def create_load_balancer_sls_class(
     """
     Create a LoadBalancerSls-enabled class following existing deployment patterns.
 
-    This function follows the same synchronous pattern as create_remote_class but uses
-    LoadBalancerSls for execution. The deployment will happen lazily when methods are called.
+    This function creates a wrapper that deploys the LoadBalancerSls endpoint dynamically
+    using GraphQL and uses the deployed endpoint URL for LoadBalancerSls execution.
 
     Args:
         cls: The class to be wrapped for LoadBalancerSls execution
-        resource_config: Configuration object specifying the serverless resource
+        resource_config: Configuration object specifying the serverless resource (with type="LB")
         dependencies: List of pip packages to install
         system_dependencies: List of system packages to install
         extra: Additional parameters for execution
@@ -43,18 +43,75 @@ def create_load_balancer_sls_class(
     # Follow the same deployment pattern as existing system
     log.info(f"Creating LoadBalancerSls class for {cls.__name__}")
 
-    # For now, use the hardcoded URL as requested
-    # TODO: Integrate with resource_config to get actual deployed endpoint URL
-    endpoint_url = "https://9ttr6h4l3f17w3.api.runpod.ai"
+    # Verify this is a LoadBalancer resource
+    if getattr(resource_config, 'type', None) != 'LB':
+        raise ValueError(f"Expected LoadBalancer resource with type='LB', got type='{getattr(resource_config, 'type', None)}'")
 
-    # Note: resource_config parameter will be used in future for actual deployment
+    # Create a deployment wrapper that handles async deployment
+    class LoadBalancerSlsClassWrapper:
+        def __init__(self):
+            self._deployed_endpoint_url = None
+            self._deployment_lock = None
+            
+        async def _ensure_deployed(self):
+            """Ensure the endpoint is deployed and get the URL."""
+            if self._deployed_endpoint_url:
+                return self._deployed_endpoint_url
+                
+            # Use ResourceManager to handle caching properly
+            resource_manager = ResourceManager()
+            deployed_resource = await resource_manager.get_or_deploy_resource(resource_config)
 
-    log.info(f"Using LoadBalancerSls endpoint: {endpoint_url}")
+            # Construct the endpoint URL from the deployed endpoint ID
+            # Format: https://ENDPOINT_ID.api.runpod.ai
+            self._deployed_endpoint_url = f"https://{deployed_resource.id}.api.runpod.ai"
+            
+            log.info(f"LoadBalancerSls endpoint ready: {self._deployed_endpoint_url}")
+            return self._deployed_endpoint_url
+        
+        def __call__(self, *args, **kwargs):
+            """Create an instance that handles dynamic deployment."""
+            return LoadBalancerSlsInstanceWrapper(
+                cls, self, dependencies, system_dependencies, args, kwargs
+            )
+    
+    return LoadBalancerSlsClassWrapper()
 
-    # Create LoadBalancerSls instance following existing patterns
-    runtime = LoadBalancerSls(endpoint_url=endpoint_url)
 
-    # Return the wrapped class using LoadBalancerSls
-    return runtime.remote_class(
-        dependencies=dependencies, system_dependencies=system_dependencies
-    )(cls)
+class LoadBalancerSlsInstanceWrapper:
+    """Instance wrapper that handles async deployment and method routing."""
+    
+    def __init__(self, original_cls, class_wrapper, dependencies, system_dependencies, args, kwargs):
+        self._original_cls = original_cls
+        self._class_wrapper = class_wrapper
+        self._dependencies = dependencies
+        self._system_dependencies = system_dependencies
+        self._args = args
+        self._kwargs = kwargs
+        self._runtime = None
+        
+    async def _get_runtime(self):
+        """Get or create the LoadBalancerSls runtime with deployed endpoint."""
+        if not self._runtime:
+            endpoint_url = await self._class_wrapper._ensure_deployed()
+            self._runtime = LoadBalancerSls(endpoint_url=endpoint_url)
+            
+            # Apply the remote_class decorator
+            self._wrapped_class = self._runtime.remote_class(
+                dependencies=self._dependencies, 
+                system_dependencies=self._system_dependencies
+            )(self._original_cls)
+            
+            # Create the actual instance
+            self._instance = self._wrapped_class(*self._args, **self._kwargs)
+            
+        return self._runtime, self._instance
+    
+    def __getattr__(self, name):
+        """Route method calls through the deployed runtime."""
+        async def async_method_wrapper(*args, **kwargs):
+            runtime, instance = await self._get_runtime()
+            method = getattr(instance, name)
+            return await method(*args, **kwargs)
+        
+        return async_method_wrapper

src/tetra_rp/core/resources/load_balancer_sls_resource.py

@@ -8,6 +8,8 @@
 import os
 from pydantic import model_validator
 from .serverless import ServerlessResource
+from .template import PodTemplate, KeyValuePair
+from .serverless import get_env_vars
 
 TETRA_IMAGE_TAG = os.environ.get("TETRA_IMAGE_TAG", "latest")
 TETRA_GPU_IMAGE = os.environ.get(
@@ -34,9 +36,11 @@ def set_load_balancer_defaults(cls, data: dict):
         data["type"] = "LB"
 
         # Set default image based on instanceIds presence
-        data["imageName"] = (
-            TETRA_CPU_IMAGE if data.get("instanceIds") else TETRA_GPU_IMAGE
-        )
+        # This ensures imageName is available for template creation
+        if not data.get("imageName"):
+            data["imageName"] = (
+                TETRA_CPU_IMAGE if data.get("instanceIds") else TETRA_GPU_IMAGE
+            )
 
         return data
 
@@ -47,7 +51,32 @@ def imageName(self):
             TETRA_CPU_IMAGE if getattr(self, "instanceIds", None) else TETRA_GPU_IMAGE
         )
 
-    @property
+    @property 
     def type(self):
         """Always return 'LB' for LoadBalancerSls resources."""
-        return "LB"
+        return "LB"
+
+    @model_validator(mode="after")
+    def ensure_template_creation(self):
+        """Ensure template is created for LoadBalancerSls deployment."""
+        # Call the parent class template creation logic
+        if not self.templateId and not self.template and self.imageName:
+            self.template = PodTemplate(
+                name=self.resource_id,
+                imageName=self.imageName,
+                env=KeyValuePair.from_dict(self.env or get_env_vars()),
+            )
+        elif self.template:
+            self.template.name = f"{self.resource_id}__{self.template.resource_id}"
+            if self.imageName:
+                self.template.imageName = self.imageName
+            if self.env:
+                self.template.env = KeyValuePair.from_dict(self.env)
+        
+        return self
+    
+    def model_dump(self, **kwargs):
+        """Override model_dump to ensure type='LB' is included."""
+        data = super().model_dump(**kwargs)
+        data["type"] = "LB"  # Ensure type is always included in serialization
+        return data