Support local app serving (#628)

# Summary
This PR implements local app serving for AppEnvironment, enabling
developers to serve Flyte apps on localhost for rapid local development
and testing. It introduces a symmetry with the existing
`flyte.with_runcontext(mode="local") ` pattern by adding
`flyte.with_servecontext(mode="local")` for apps.

Key changes:
- New `_LocalApp` class and local serving infrastructure in `_serve.py`
(server-decorator, command/subprocess, and FastAPI-based apps)
- `AppEnvironment.endpoint` now transparently resolves to localhost when
the app is served locally
`ephemeral_ctx` / `ephemeral_ctx_sync` context managers on both
_LocalApp and remote App for scoped activate/deactivate lifecycle
- `_LocalApp` `activate` and `deactivate` methods mirror the methods in
remote `App` object.
- `activate`: if the server is not running locally, starts it, and
optionally waits for the server to be active.
  - `deactivate`: stop the server locally
- `--local` flag on the flyte serve CLI command
- Signal handler and atexit cleanup to ensure child processes are torn
down on `Ctrl-C` or process exit
- 6 example scripts and comprehensive unit tests

---------

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

Author: cosmicBboy

examples/apps/local_app_calling_app.py

@@ -0,0 +1,111 @@
+"""Local serving example where one FastAPI app calls another FastAPI app.
+
+This example demonstrates how to serve two FastAPI apps locally and have one
+app proxy requests to the other.  Both apps are started via
+``flyte.with_servecontext(mode="local")`` and communicate over HTTP.
+
+Usage (SDK):
+    python examples/apps/local_app_calling_app.py
+"""
+
+import httpx
+from fastapi import FastAPI
+
+import flyte
+from flyte.app.extras import FastAPIAppEnvironment
+
+image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("fastapi", "uvicorn", "httpx")
+
+# ---------------------------------------------------------------------------
+# App 1 - a simple "square" service
+# ---------------------------------------------------------------------------
+
+app1 = FastAPI(
+    title="Square Service",
+    description="A FastAPI app that squares a number",
+    version="1.0.0",
+)
+
+app1_env = FastAPIAppEnvironment(
+    name="local-square-service",
+    app=app1,
+    description="Squares a number",
+    image=image,
+    resources=flyte.Resources(cpu=1, memory="512Mi"),
+    port=8091,
+    requires_auth=False,
+)
+
+
+@app1.get("/")
+async def square(x: int) -> dict[str, int]:
+    """Return x squared."""
+    return {"result": x * x}
+
+
+@app1.get("/health")
+async def app1_health() -> dict[str, str]:
+    return {"status": "healthy"}
+
+
+# ---------------------------------------------------------------------------
+# App 2 - proxies to App 1 and adds its own logic
+# ---------------------------------------------------------------------------
+
+app2 = FastAPI(
+    title="Square-Plus-One Service",
+    description="A FastAPI app that calls the square service and adds one",
+    version="1.0.0",
+)
+
+app2_env = FastAPIAppEnvironment(
+    name="local-square-plus-one",
+    app=app2,
+    description="Calls the square service and adds one",
+    image=image,
+    resources=flyte.Resources(cpu=1, memory="512Mi"),
+    port=8092,
+    requires_auth=False,
+    depends_on=[app1_env],
+)
+
+
+@app2.get("/")
+async def square_plus_one(x: int) -> dict[str, int]:
+    """Call app1 to square x, then add one."""
+    async with httpx.AsyncClient() as client:
+        response = await client.get(app1_env.endpoint, params={"x": x})
+        response.raise_for_status()
+        squared = response.json()["result"]
+    return {"result": squared + 1}
+
+
+@app2.get("/health")
+async def app2_health() -> dict[str, str]:
+    return {"status": "healthy"}
+
+
+if __name__ == "__main__":
+    serve_ctx = flyte.with_servecontext(mode="local")
+
+    # Serve app1 first (app2 depends on it)
+    local_app1 = serve_ctx.serve(app1_env)
+    local_app1.activate(wait=True)
+    print(f"App 1 (square) is ready at {local_app1.endpoint}")
+
+    # Serve app2
+    local_app2 = serve_ctx.serve(app2_env)
+    local_app2.activate(wait=True)
+    print(f"App 2 (square+1) is ready at {local_app2.endpoint}")
+
+    # Call app2 which internally calls app1
+    response = httpx.get(f"{local_app2.endpoint}", params={"x": 5})
+    response.raise_for_status()
+    data = response.json()
+    print(f"Response: {data}")
+    assert data["result"] == 26  # 5^2 + 1 = 26
+
+    # Shut down both apps
+    local_app2.deactivate(wait=True)
+    local_app1.deactivate(wait=True)
+    print("Done!")

examples/apps/local_app_calling_task.py

@@ -0,0 +1,75 @@
+"""Local serving example where a FastAPI app calls a Flyte task.
+
+This example demonstrates how to serve a FastAPI app locally that invokes a
+Flyte task from one of its endpoints.  The task is executed via
+``flyte.with_runcontext(mode="local").run(...)`` so it runs in the same process.
+
+Usage (SDK):
+    python examples/apps/local_app_calling_task.py
+"""
+
+import httpx
+from fastapi import FastAPI
+
+import flyte
+from flyte.app.extras import FastAPIAppEnvironment
+
+app = FastAPI(
+    title="App Calling Task",
+    description="A FastAPI app that delegates computation to a Flyte task",
+    version="1.0.0",
+)
+
+app_env = FastAPIAppEnvironment(
+    name="local-app-calling-task",
+    app=app,
+    description="App that calls a Flyte task to double a number",
+    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("fastapi", "uvicorn", "httpx"),
+    resources=flyte.Resources(cpu=1, memory="512Mi"),
+    requires_auth=False,
+)
+
+task_env = flyte.TaskEnvironment(
+    name="local-doubler-task-env",
+    image=flyte.Image.from_debian_base(python_version=(3, 12)),
+    resources=flyte.Resources(cpu=1, memory="512Mi"),
+)
+
+
+@task_env.task
+async def double(x: int) -> int:
+    """A simple task that doubles the input."""
+    return x * 2
+
+
+@app.get("/")
+async def double_endpoint(x: int) -> dict[str, int]:
+    """Endpoint that invokes the ``double`` task and returns the result."""
+    result = flyte.with_runcontext(mode="local").run(double, x=x)
+    return {"result": result.outputs()[0]}
+
+
+@app.get("/health")
+async def health() -> dict[str, str]:
+    """Health check endpoint."""
+    return {"status": "healthy"}
+
+
+if __name__ == "__main__":
+    # Serve the app locally (non-blocking)
+    local_app = flyte.with_servecontext(mode="local").serve(app_env)
+
+    # Wait for the app to be ready
+    local_app.activate(wait=True)
+    print(f"App is ready at {local_app.endpoint}")
+
+    # Call the app endpoint which internally runs the task
+    response = httpx.get(f"{local_app.endpoint}", params={"x": 21})
+    response.raise_for_status()
+    data = response.json()
+    print(f"Response: {data}")
+    assert data["result"] == 42
+
+    # Shut down the local app
+    local_app.deactivate(wait=True)
+    print("Done!")

examples/apps/local_app_ephemeral_server.py

@@ -0,0 +1,109 @@
+"""Local serving example using ephemeral_server context manager.
+
+This example demonstrates how to serve a FastAPI app locally using the
+ephemeral_server() context manager, which ensures the app is properly
+activated and deactivated. The FastAPI lifespan simulates loading a model
+on startup.
+
+Usage (SDK):
+    python examples/apps/local_app_ephemeral_server.py
+
+Usage (CLI):
+    flyte serve --local examples/apps/local_app_ephemeral_server.py app_env
+"""
+
+import asyncio
+from contextlib import asynccontextmanager
+
+import httpx
+from fastapi import FastAPI
+
+import flyte
+from flyte.app.extras import FastAPIAppEnvironment
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Lifespan function that simulates loading a model on startup."""
+    print("Simulate loading model...")
+    await asyncio.sleep(3)
+    app.state.model = {"m": 5, "b": 11}
+    print("Model loaded!")
+    yield
+    print("Shutting down, releasing model resources...")
+
+
+app = FastAPI(
+    title="Local Linear Regression",
+    description="A local FastAPI app that performs linear regression",
+    version="1.0.0",
+    lifespan=lifespan,
+)
+
+app_env = FastAPIAppEnvironment(
+    name="local-linear-regression",
+    app=app,
+    description="Performs linear regression",
+    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("fastapi", "uvicorn", "httpx"),
+    resources=flyte.Resources(cpu=1, memory="512Mi"),
+    requires_auth=False,
+)
+
+task_env = flyte.TaskEnvironment(
+    name="local-linear-regression-task-env",
+    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("httpx"),
+    resources=flyte.Resources(cpu=1, memory="512Mi"),
+)
+
+
+@app.post("/predict")
+async def predict(x: int) -> dict[str, int]:
+    """Perform linear regression."""
+    result = app.state.model["m"] * x + app.state.model["b"]
+    return {"result": result}
+
+
+@app.get("/health")
+async def health() -> dict[str, str]:
+    """Health check endpoint."""
+    return {"status": "healthy"}
+
+
+@task_env.task
+async def predict_task(data: list[int]) -> list[int]:
+    """Task that calls the local app endpoint."""
+    print(f"Calling app at {app_env.endpoint}")
+    async with httpx.AsyncClient() as client:
+        results = []
+        for x in data:
+            response = await client.post(f"{app_env.endpoint}/predict", params={"x": x})
+            response.raise_for_status()
+            results.append(response.json()["result"])
+        return results
+
+
+if __name__ == "__main__":
+    # Serve the app locally (non-blocking)
+    local_app = flyte.with_servecontext(mode="local").serve(app_env)
+
+    # Use ephemeral_server to ensure the app is activated and deactivated
+    async def main():
+        async with local_app.ephemeral_ctx():
+            print(f"App is ready at {local_app.endpoint}")
+
+            # Call the app endpoint directly
+            async with httpx.AsyncClient() as client:
+                response = await client.post(f"{local_app.endpoint}/predict", params={"x": 5})
+                response.raise_for_status()
+                print(f"Direct call result: {response.json()}")
+
+            # Run a task that calls the local app
+            result = await flyte.with_runcontext(mode="local").run.aio(predict_task, data=[5, 10, 15])
+            print(f"Task result: {result.outputs()[0]}")
+            assert result.outputs()[0] == [36, 61, 86]
+
+        assert local_app.is_deactivated()
+
+    asyncio.run(main())
+    # App is automatically deactivated after exiting the context
+    print("Done!")

examples/apps/local_command_app.py

@@ -0,0 +1,45 @@
+"""Local serving example using command-based AppEnvironment.
+
+This example demonstrates how to serve an app locally using the ``command``
+specification pattern. The app is run as a subprocess.
+
+Usage (CLI):
+    flyte serve --local examples/apps/local_command_app.py app_env
+
+Note: This example uses ``python -m http.server`` as the command, which only
+serves static files. For a real application, you would use a proper web server
+command like ``uvicorn main:app --host 0.0.0.0 --port 8082``.
+"""
+
+import flyte
+from flyte.app import AppEnvironment
+
+app_env = AppEnvironment(
+    name="local-static-server",
+    image=flyte.Image.from_debian_base(python_version=(3, 12)),
+    resources=flyte.Resources(cpu=1, memory="512Mi"),
+    command="python -m http.server 8082",
+    port=8082,
+    requires_auth=False,
+)
+
+
+if __name__ == "__main__":
+    import httpx
+
+    flyte.init()
+
+    # Serve the app locally (non-blocking)
+    local_app = flyte.with_servecontext(mode="local", health_check_path="/").serve(app_env)
+
+    # Wait for the app to be ready
+    local_app.activate(wait=True)
+    print(f"App is ready at {local_app.endpoint}")
+
+    # Test the endpoint
+    response = httpx.get(local_app.endpoint)
+    print(f"Response status: {response.status_code}")
+
+    # Shut down the local app
+    local_app.deactivate(wait=True)
+    print("Done!")