fix: respect user-provided imageName in Live* classes (BYOI)

Model validators were unconditionally overwriting imageName, and the
mixin property was ignoring the stored value. Now validators only set
defaults when imageName is not provided, and the no-op property is
removed. Tests updated for new image naming scheme.

Author: deanq

src/runpod_flash/core/resources/live_serverless.py

@@ -1,11 +1,13 @@
 # Ship serverless code as you write it. No builds, no deploys -- just run.
 from typing import Any, ClassVar
 
+# Ship serverless code as you write it. No builds, no deploys — just run.
 from pydantic import model_validator
 
 from .constants import (
-    DEFAULT_PYTHON_VERSION,
+    GPU_BASE_IMAGE_PYTHON_VERSION,
     get_image_name,
+    local_python_version,
 )
 from .injection import build_injection_cmd
 from .load_balancer_sls_resource import (
@@ -77,46 +79,58 @@ def _configure_existing_template(self) -> None:
 class LiveServerless(LiveServerlessMixin, ServerlessEndpoint):
     """GPU-only live serverless endpoint."""
 
-    _image_type: ClassVar[str] = "gpu"
-
     @model_validator(mode="before")
     @classmethod
     def set_live_serverless_template(cls, data: dict):
         """Default to the GPU Flash runtime image when none is supplied."""
         return _apply_default_live_image(data, "gpu")
+        """Set default GPU image for Live Serverless."""
+        if "imageName" not in data:
+            python_version = data.get("python_version") or GPU_BASE_IMAGE_PYTHON_VERSION
+            data["imageName"] = get_image_name("gpu", python_version)
+        return data
 
 
 class CpuLiveServerless(LiveServerlessMixin, CpuServerlessEndpoint):
     """CPU-only live serverless endpoint with automatic disk sizing."""
 
-    _image_type: ClassVar[str] = "cpu"
-
     @model_validator(mode="before")
     @classmethod
     def set_live_serverless_template(cls, data: dict):
         """Default to the CPU Flash runtime image when none is supplied."""
         return _apply_default_live_image(data, "cpu")
+        """Set default CPU image for Live Serverless."""
+        if "imageName" not in data:
+            python_version = data.get("python_version") or local_python_version()
+            data["imageName"] = get_image_name("cpu", python_version)
+        return data
 
 
 class LiveLoadBalancer(LiveServerlessMixin, LoadBalancerSlsResource):
     """Live load-balanced endpoint."""
 
-    _image_type: ClassVar[str] = "lb"
-
     @model_validator(mode="before")
     @classmethod
     def set_live_lb_template(cls, data: dict):
         """Default to the LB Flash runtime image when none is supplied."""
         return _apply_default_live_image(data, "lb")
+        """Set default image for Live Load-Balanced endpoint."""
+        if "imageName" not in data:
+            python_version = data.get("python_version") or GPU_BASE_IMAGE_PYTHON_VERSION
+            data["imageName"] = get_image_name("lb", python_version)
+        return data
 
 
 class CpuLiveLoadBalancer(LiveServerlessMixin, CpuLoadBalancerSlsResource):
     """CPU-only live load-balanced endpoint."""
 
-    _image_type: ClassVar[str] = "lb-cpu"
-
     @model_validator(mode="before")
     @classmethod
     def set_live_cpu_lb_template(cls, data: dict):
         """Default to the CPU LB Flash runtime image when none is supplied."""
         return _apply_default_live_image(data, "lb-cpu")
+        """Set default CPU image for Live Load-Balanced endpoint."""
+        if "imageName" not in data:
+            python_version = data.get("python_version") or local_python_version()
+            data["imageName"] = get_image_name("lb-cpu", python_version)
+        return data

tests/integration/test_cpu_disk_sizing.py

@@ -129,7 +129,7 @@ def test_live_serverless_cpu_integration(self):
         # 2. CPU utilities calculate minimum disk size
         # 3. Template creation with auto-sizing
         # 4. Validation passes
-        assert live_serverless.imageName == "python:3.11-slim"
+        assert "runpod/flash-cpu:" in live_serverless.imageName
         assert live_serverless.instanceIds == [
             CpuInstanceType.CPU5C_1_2,
             CpuInstanceType.CPU5C_2_4,
@@ -263,8 +263,8 @@ def test_live_serverless_image_defaults(self):
 
         # Verify different base images are used
         assert gpu_live.imageName != cpu_live.imageName
-        assert "pytorch" in gpu_live.imageName
-        assert "python" in cpu_live.imageName
+        assert "runpod/flash:" in gpu_live.imageName
+        assert "runpod/flash-cpu:" in cpu_live.imageName
 
     def test_live_serverless_image_override_via_constructor(self):
         """Caller-supplied imageName overrides the Flash runtime default (AE-3153)."""

tests/integration/test_lb_remote_execution.py

@@ -114,7 +114,7 @@ async def echo(message: str):
         # Verify resource is correctly configured
         # Note: name may have "-fb" appended by flash boot validator
         assert "test-live-api" in lb.name
-        assert "pytorch" in lb.imageName  # GPU base image
+        assert "runpod/flash-lb:" in lb.imageName  # GPU LB base image
         assert echo.__remote_config__["method"] == "POST"
 
     def test_live_load_balancer_image_default_and_override(self):
@@ -133,9 +133,9 @@ def test_live_load_balancer_image_default_and_override(self):
         # same default (e.g. the override branch reverting to a no-op).
         assert default_lb.imageName != custom_lb.imageName
     def test_live_load_balancer_default_image(self):
-        """Test that LiveLoadBalancer uses GPU base image by default."""
+        """Test that LiveLoadBalancer uses GPU LB base image by default."""
         lb = LiveLoadBalancer(name="test-api")
-        assert "pytorch" in lb.imageName
+        assert "runpod/flash-lb:" in lb.imageName
 
     def test_live_load_balancer_allows_custom_image(self):
         """Test that LiveLoadBalancer allows user to set custom image (BYOI)."""

tests/unit/resources/test_live_load_balancer.py

@@ -50,8 +50,7 @@ def test_live_load_balancer_default_image_tag(self):
     def test_live_load_balancer_user_can_override_image(self):
         """Test user can set custom imageName (BYOI)."""
         lb = LiveLoadBalancer(name="test-lb", imageName="custom/image:v1")
-        # imageName property returns _live_image, setter is no-op
-        assert lb.imageName is not None
+        assert lb.imageName == "custom/image:v1"
 
     def test_live_load_balancer_template_creation(self):
         """Test LiveLoadBalancer creates proper template from imageName."""
@@ -210,8 +209,7 @@ def test_cpu_live_load_balancer_default_image_tag(self):
     def test_cpu_live_load_balancer_user_can_override_image(self):
         """Test CpuLiveLoadBalancer allows user image override."""
         lb = CpuLiveLoadBalancer(name="test-lb", imageName="python:3.11-slim")
-        # imageName property returns _live_image, setter is no-op
-        assert lb.imageName is not None
+        assert lb.imageName == "python:3.11-slim"
 
     def test_cpu_live_load_balancer_defaults(self):
         """Test CpuLiveLoadBalancer defaults to CPU3G_2_8."""

tests/unit/resources/test_live_serverless.py

@@ -58,9 +58,7 @@ def test_live_serverless_user_can_override_image(self):
         live_serverless = LiveServerless(
             name="test", imageName="nvidia/cuda:12.8.0-runtime-ubuntu22.04"
         )
-        # imageName setter is a no-op, so value is always the computed _live_image
-        # The model_validator sets data["imageName"] but the property overrides reads
-        assert live_serverless.imageName is not None
+        assert live_serverless.imageName == "nvidia/cuda:12.8.0-runtime-ubuntu22.04"
 
     def test_live_serverless_with_custom_template(self):
         """Test LiveServerless with custom template."""
@@ -133,8 +131,7 @@ def test_cpu_live_serverless_image_default_unchanged(self):
     def test_cpu_live_serverless_user_can_override_image(self):
         """Test CpuLiveServerless allows user to set custom image."""
         live_serverless = CpuLiveServerless(name="test", imageName="python:3.11-slim")
-        # imageName property returns _live_image, setter is no-op
-        assert live_serverless.imageName is not None
+        assert live_serverless.imageName == "python:3.11-slim"
 
     def test_cpu_live_serverless_validation_failure(self):
         """Test CpuLiveServerless validation fails with excessive disk size."""
@@ -265,6 +262,15 @@ def test_default_image_validator_passes_through_non_dict(self):
         original = LiveServerless(name="test", imageName="byo/image:v1")
         revalidated = LiveServerless.model_validate(original)
         assert revalidated.imageName == "byo/image:v1"
+    def test_live_serverless_byoi_gpu(self):
+        """Test LiveServerless respects user-provided imageName."""
+        live_serverless = LiveServerless(name="test", imageName="custom/gpu:v1")
+        assert live_serverless.imageName == "custom/gpu:v1"
+
+    def test_live_serverless_byoi_cpu(self):
+        """Test CpuLiveServerless respects user-provided imageName."""
+        live_serverless = CpuLiveServerless(name="test", imageName="custom/cpu:v1")
+        assert live_serverless.imageName == "custom/cpu:v1"
 
 
 class TestLiveServerlessPythonVersion: