feat: process injection via LiveServerlessMixin

src/runpod_flash/core/resources/constants.py

@@ -119,6 +119,20 @@ def get_image_name(
     f"runpod/flash-lb-cpu:py{DEFAULT_PYTHON_VERSION}-{FLASH_IMAGE_TAG}",
 )
 
+# Base images for process injection (no flash-worker baked in)
+FLASH_GPU_BASE_IMAGE = os.environ.get(
+    "FLASH_GPU_BASE_IMAGE", "pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime"
+)
+FLASH_CPU_BASE_IMAGE = os.environ.get("FLASH_CPU_BASE_IMAGE", "python:3.11-slim")
+
+# Worker tarball for process injection
+FLASH_WORKER_VERSION = os.environ.get("FLASH_WORKER_VERSION", "1.1.1")
+FLASH_WORKER_TARBALL_URL_TEMPLATE = os.environ.get(
+    "FLASH_WORKER_TARBALL_URL",
+    "https://github.com/runpod-workers/flash/releases/download/"
+    "v{version}/flash-worker-v{version}-py3.11-linux-x86_64.tar.gz",
+)
+
 # Worker configuration defaults
 DEFAULT_WORKERS_MIN = 0
 DEFAULT_WORKERS_MAX = 1

src/runpod_flash/core/resources/injection.py

@@ -0,0 +1,54 @@
+"""Process injection utilities for flash-worker tarball delivery."""
+
+from .constants import FLASH_WORKER_TARBALL_URL_TEMPLATE, FLASH_WORKER_VERSION
+
+
+def build_injection_cmd(
+    worker_version: str = FLASH_WORKER_VERSION,
+    tarball_url: str | None = None,
+) -> str:
+    """Build the dockerArgs command that downloads, extracts, and runs flash-worker.
+
+    Supports remote URLs (curl/wget) and local file paths (file://) for testing.
+    Includes version-based caching to skip re-extraction on warm workers.
+    Network volume caching stores extracted tarball at /runpod-volume/.flash-worker/v{version}.
+    """
+    if tarball_url is None:
+        tarball_url = FLASH_WORKER_TARBALL_URL_TEMPLATE.format(version=worker_version)
+
+    if tarball_url.startswith("file://"):
+        local_path = tarball_url[7:]
+        return (
+            "bash -c '"
+            "set -e; FW_DIR=/opt/flash-worker; "
+            "mkdir -p $FW_DIR; "
+            f"tar xzf {local_path} -C $FW_DIR --strip-components=1; "
+            "exec $FW_DIR/bootstrap.sh'"
+        )
+
+    return (
+        "bash -c '"
+        f"set -e; FW_DIR=/opt/flash-worker; FW_VER={worker_version}; "
+        # Network volume cache check
+        'NV_CACHE="/runpod-volume/.flash-worker/v$FW_VER"; '
+        'if [ -d "$NV_CACHE" ] && [ -f "$NV_CACHE/.version" ]; then '
+        'cp -r "$NV_CACHE" "$FW_DIR"; '
+        # Local cache check (container disk persistence between restarts)
+        'elif [ -f "$FW_DIR/.version" ] && [ "$(cat $FW_DIR/.version)" = "$FW_VER" ]; then '
+        "true; "
+        "else "
+        "mkdir -p $FW_DIR; "
+        f'DL_URL="{tarball_url}"; '
+        "dl() { "
+        '(command -v curl >/dev/null 2>&1 && curl -sSL "$1" || '
+        'command -v wget >/dev/null 2>&1 && wget -qO- "$1" || '
+        'python3 -c "import urllib.request,sys;sys.stdout.buffer.write(urllib.request.urlopen(sys.argv[1]).read())" "$1"); '
+        "}; "
+        'dl "$DL_URL" '
+        "| tar xz -C $FW_DIR --strip-components=1; "
+        # Cache to network volume if available
+        "if [ -d /runpod-volume ]; then "
+        'mkdir -p "$NV_CACHE" && cp -r "$FW_DIR"/* "$NV_CACHE/" 2>/dev/null || true; fi; '
+        "fi; "
+        "exec $FW_DIR/bootstrap.sh'"
+    )

src/runpod_flash/core/resources/live_serverless.py

@@ -1,42 +1,46 @@
 # Ship serverless code as you write it. No builds, no deploys -- just run.
-from typing import Any, ClassVar
+from typing import Any
 
 from pydantic import model_validator
 
-from .constants import (
-    DEFAULT_PYTHON_VERSION,
-    get_image_name,
-)
+from .constants import DEFAULT_PYTHON_VERSION, get_image_name
+from .injection import build_injection_cmd
 from .load_balancer_sls_resource import (
     CpuLoadBalancerSlsResource,
     LoadBalancerSlsResource,
 )
 from .serverless import ServerlessEndpoint
 from .serverless_cpu import CpuServerlessEndpoint
+from .template import PodTemplate
 
 
 class LiveServerlessMixin:
     """Common mixin for live serverless endpoints.
 
-    Treats the Flash runtime image as a *default*: if the caller passes an
-    ``imageName`` (e.g. via ``Endpoint(image=...)`` in client mode), that
-    value wins. Otherwise the Flash runtime image for this resource type is
-    used so decorator-mode workloads continue to deploy the Flash wrapper.
-
-    The default is applied via the ``@model_validator(mode="before")`` on each
-    concrete subclass (see ``_apply_default_live_image``); reads and writes of
-    ``imageName`` go through the normal Pydantic field machinery so model
-    serialization, drift detection, and ``setattr`` all stay consistent.
+    Configures process injection via ``dockerArgs`` for any base image, and
+    treats the Flash runtime image as a *default*: if the caller passes an
+    ``imageName`` (e.g. via ``Endpoint(image=...)`` in client mode), that value
+    wins. Otherwise the Flash runtime image for this resource type is applied by
+    the ``@model_validator(mode="before")`` on each concrete subclass (see
+    ``_apply_default_live_image``), so decorator-mode workloads continue to
+    deploy the Flash wrapper.
+
+    The injection ``dockerArgs`` download, extract, and run the flash-worker
+    tarball at container start; QB vs LB mode is determined by the
+    ``FLASH_ENDPOINT_TYPE`` env var at runtime, not by the Docker image.
     """
 
-    _image_type: ClassVar[str] = (
-        ""  # override in subclasses: 'gpu', 'cpu', 'lb', 'lb-cpu'
-    )
+    def _create_new_template(self) -> PodTemplate:
+        """Create template with dockerArgs for process injection."""
+        template = super()._create_new_template()  # type: ignore[misc]
+        template.dockerArgs = build_injection_cmd()
+        return template
 
-    @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)
+    def _configure_existing_template(self) -> None:
+        """Configure existing template, adding dockerArgs for injection if not user-set."""
+        super()._configure_existing_template()  # type: ignore[misc]
+        if self.template is not None and not self.template.dockerArgs:  # type: ignore[attr-defined]
+            self.template.dockerArgs = build_injection_cmd()  # type: ignore[attr-defined]
 
 
 def _apply_default_live_image(data: Any, image_type: str):
@@ -57,8 +61,6 @@ def _apply_default_live_image(data: Any, image_type: str):
 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):
@@ -69,8 +71,6 @@ def set_live_serverless_template(cls, data: dict):
 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):
@@ -81,8 +81,6 @@ def set_live_serverless_template(cls, data: dict):
 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):
@@ -93,8 +91,6 @@ def set_live_lb_template(cls, data: dict):
 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):

tests/integration/test_cpu_disk_sizing.py

@@ -125,11 +125,11 @@ def test_live_serverless_cpu_integration(self):
         )
 
         # Verify integration:
-        # 1. Uses CPU image (locked)
+        # 1. Uses CPU base image (default)
         # 2. CPU utilities calculate minimum disk size
         # 3. Template creation with auto-sizing
         # 4. Validation passes
-        assert "flash-cpu:" in live_serverless.imageName
+        assert "runpod/flash-cpu:" in live_serverless.imageName
         assert live_serverless.instanceIds == [
             CpuInstanceType.CPU5C_1_2,
             CpuInstanceType.CPU5C_2_4,
@@ -244,18 +244,18 @@ def test_mixed_cpu_generations_integration(self):
         assert "cpu5c-1-2: max 15GB" in error_msg
 
 
-class TestLiveServerlessImageIntegration:
-    """Test image default + override behavior in live serverless variants (AE-3153)."""
+class TestLiveServerlessImageDefaultsIntegration:
+    """Test image defaults in live serverless variants."""
 
-    def test_live_serverless_image_consistency(self):
-        """LiveServerless variants default to distinct Flash runtime images."""
+    def test_live_serverless_image_defaults(self):
+        """Test that LiveServerless variants use correct base images."""
         gpu_live = LiveServerless(name="gpu-live")
         cpu_live = CpuLiveServerless(name="cpu-live")
 
-        # Verify different default images are used per resource type.
+        # Verify different base images are used
         assert gpu_live.imageName != cpu_live.imageName
-        assert "flash:" in gpu_live.imageName
-        assert "flash-cpu:" 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)."""
@@ -266,6 +266,11 @@ def test_live_serverless_image_override_via_constructor(self):
 
         assert gpu_live.imageName == "custom/image:latest"
         assert cpu_live.imageName == "custom/cpu-image:latest"
+        # Verify images can be overridden (BYOI)
+        custom_gpu = LiveServerless(
+            name="custom-gpu", imageName="nvidia/cuda:12.8.0-runtime"
+        )
+        assert custom_gpu.imageName == "nvidia/cuda:12.8.0-runtime"
 
     def test_live_serverless_template_integration(self):
         """Test live serverless template integration with disk sizing."""

tests/integration/test_lb_remote_execution.py

@@ -112,8 +112,9 @@ async def echo(message: str):
             return {"echo": message}
 
         # Verify resource is correctly configured
-        assert lb.name == "test-live-api"
-        assert "flash-lb" in lb.imageName
+        # Note: name may have "-fb" appended by flash boot validator
+        assert "test-live-api" in lb.name
+        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):
@@ -132,6 +133,16 @@ 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 LB base image by default."""
+        lb = LiveLoadBalancer(name="test-api")
+        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)."""
+        lb = LiveLoadBalancer(name="test-api", imageName="custom-image:latest")
+        assert lb.imageName == "custom-image:latest"
+
     def test_load_balancer_vs_queue_based_endpoints(self):
         """Test that LB and QB endpoints have different characteristics."""
         from runpod_flash import ServerlessEndpoint
@@ -160,48 +171,50 @@ async def qb_func():
         assert qb_func.__remote_config__["path"] is None
 
     def test_scanner_discovers_load_balancer_resources(self):
-        """Test that scanner can discover LiveLoadBalancer and LoadBalancerSlsResource."""
+        """Test that scanner discovers @Endpoint load-balanced route handlers."""
         from runpod_flash.cli.commands.build_utils.scanner import RuntimeScanner
         from pathlib import Path
         import tempfile
 
-        # Create temporary Python file with LoadBalancer resource
+        # Uses the current @Endpoint API (remote is deprecated). The worker file
+        # must NOT be named test_*.py — the scanner skips test files, so a
+        # test-prefixed fixture would yield zero discovered functions.
         code = """
-from runpod_flash import LiveLoadBalancer, LoadBalancerSlsResource, remote
+from runpod_flash.endpoint import Endpoint
+from runpod_flash.core.resources.gpu import GpuGroup
 
-# Test LiveLoadBalancer discovery
-api = LiveLoadBalancer(name="test-api")
+api = Endpoint(name="test-api", gpu=GpuGroup.AMPERE_16)
 
-@remote(api, method="POST", path="/api/process")
+@api.post("/api/process")
 async def process_data(x: int):
     return {"result": x}
 
-# Test LoadBalancerSlsResource discovery
-deployed = LoadBalancerSlsResource(name="deployed-api", imageName="test:latest")
+status_api = Endpoint(name="deployed-api", gpu=GpuGroup.AMPERE_16)
 
-@remote(deployed, method="GET", path="/api/status")
+@status_api.get("/api/status")
 def get_status():
     return {"status": "ok"}
 """
 
         with tempfile.TemporaryDirectory() as tmpdir:
             project_dir = Path(tmpdir)
-            py_file = project_dir / "api_worker.py"
+            py_file = project_dir / "worker_api.py"
             py_file.write_text(code)
 
             scanner = RuntimeScanner(project_dir)
             functions = scanner.discover_remote_functions()
 
-            # Verify both resources were discovered
+            # Both LB route handlers are discovered.
             assert len(functions) == 2
 
-            # Verify resource types are correctly identified
-            resource_types = {f.resource_type for f in functions}
-            assert "LiveLoadBalancer" in resource_types
-            assert "LoadBalancerSlsResource" in resource_types
+            # Endpoint(gpu=...) with route handlers resolves to LiveLoadBalancer.
+            assert {f.resource_type for f in functions} == {"LiveLoadBalancer"}
+            assert all(f.is_load_balanced for f in functions)
+
+            # HTTP method and path are captured per route.
+            routes = {(f.http_method, f.http_path) for f in functions}
+            assert routes == {("POST", "/api/process"), ("GET", "/api/status")}
 
-            # Verify resource configs were extracted
-            assert "test-api" in scanner.resource_types
+            # Resource configs are tracked by name.
             assert scanner.resource_types["test-api"] == "LiveLoadBalancer"
-            assert "deployed-api" in scanner.resource_types
-            assert scanner.resource_types["deployed-api"] == "LoadBalancerSlsResource"
+            assert scanner.resource_types["deployed-api"] == "LiveLoadBalancer"