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,10 @@
-# Ship serverless code as you write it. No builds, no deploys -- just run.
-from typing import 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 (
@@ -26,23 +25,6 @@ class LiveServerlessMixin:
     runtime, not by the Docker image.
     """
 
-    _image_type: ClassVar[str] = (
-        ""  # override in subclasses: 'gpu', 'cpu', 'lb', 'lb-cpu'
-    )
-
-    @property
-    def _live_image(self) -> str:
-        python_version = getattr(self, "python_version", None) or DEFAULT_PYTHON_VERSION
-        return get_image_name(self._image_type, python_version)
-
-    @property
-    def imageName(self):
-        return self._live_image
-
-    @imageName.setter
-    def imageName(self, value):
-        pass
-
     def _create_new_template(self) -> PodTemplate:
         """Create template with dockerArgs for process injection."""
         template = super()._create_new_template()  # type: ignore[misc]
@@ -59,54 +41,50 @@ 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):
         """Set default GPU image for Live Serverless."""
-        python_version = data.get("python_version") or DEFAULT_PYTHON_VERSION
-        data["imageName"] = get_image_name("gpu", python_version)
+        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):
         """Set default CPU image for Live Serverless."""
-        python_version = data.get("python_version") or DEFAULT_PYTHON_VERSION
-        data["imageName"] = get_image_name("cpu", python_version)
+        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):
         """Set default image for Live Load-Balanced endpoint."""
-        python_version = data.get("python_version") or DEFAULT_PYTHON_VERSION
-        data["imageName"] = get_image_name("lb", python_version)
+        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):
         """Set default CPU image for Live Load-Balanced endpoint."""
-        python_version = data.get("python_version") or DEFAULT_PYTHON_VERSION
-        data["imageName"] = get_image_name("lb-cpu", python_version)
+        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,
@@ -254,8 +254,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
 
         # Verify images can be overridden (BYOI)
         custom_gpu = LiveServerless(

tests/integration/test_lb_remote_execution.py

@@ -114,13 +114,13 @@ 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_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

@@ -44,9 +44,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."""
@@ -102,8 +100,7 @@ def test_cpu_live_serverless_multiple_instances(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."""
@@ -200,23 +197,15 @@ def test_cpu_live_load_balancer_defaults(self):
         assert lb.template is not None
         assert lb.template.dockerArgs
 
-    def test_image_name_setter_ignored_gpu(self):
-        """Test LiveServerless imageName setter is ignored."""
-        live_serverless = LiveServerless(name="test")
-        original_image = live_serverless.imageName
-
-        live_serverless.imageName = "should-be-ignored"
-
-        assert live_serverless.imageName == original_image
-
-    def test_image_name_setter_ignored_cpu(self):
-        """Test CpuLiveServerless imageName setter is ignored."""
-        live_serverless = CpuLiveServerless(name="test")
-        original_image = live_serverless.imageName
-
-        live_serverless.imageName = "should-be-ignored"
+    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"
 
-        assert live_serverless.imageName == original_image
+    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: