gunicorn/examples/dirty_example/wsgi_app.py

#
# This file is part of gunicorn released under the MIT license.
# See the NOTICE for more information.

"""
Example WSGI Application that uses Dirty Workers

This demonstrates how HTTP workers can call dirty workers
for heavy operations like ML inference.

Run with:
    cd examples/dirty_example
    gunicorn wsgi_app:app -c gunicorn_conf.py
"""

import json
import os
from urllib.parse import parse_qs


def get_dirty_client():
    """Get the dirty client, with fallback for when dirty workers aren't enabled."""
    try:
        from gunicorn.dirty import get_dirty_client as _get_dirty_client
        return _get_dirty_client()
    except Exception as e:
        return None


def app(environ, start_response):
    """WSGI application that demonstrates dirty worker integration."""
    path = environ.get('PATH_INFO', '/')
    method = environ.get('REQUEST_METHOD', 'GET')

    # Parse query string
    query = parse_qs(environ.get('QUERY_STRING', ''))

    # Get dirty client
    client = get_dirty_client()

    try:
        if path == '/':
            result = {
                "message": "Dirty Workers Demo",
                "dirty_enabled": client is not None,
                "pid": os.getpid(),
                "endpoints": {
                    "/models": "List loaded models",
                    "/load?name=MODEL": "Load a model",
                    "/inference?model=NAME&data=INPUT": "Run inference",
                    "/unload?name=MODEL": "Unload a model",
                    "/fibonacci?n=NUMBER": "Compute fibonacci",
                    "/prime?n=NUMBER": "Check if prime",
                    "/stats": "Get dirty worker stats",
                    "/session/login?user_id=ID&name=NAME": "Login user (stash demo)",
                    "/session/get?user_id=ID": "Get session (stash demo)",
                    "/session/list": "List all sessions (stash demo)",
                    "/session/logout?user_id=ID": "Logout user (stash demo)",
                    "/session/stats": "Get stash stats (stash demo)",
                }
            }

        elif path == '/models':
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                result = client.execute(
                    "examples.dirty_example.dirty_app:MLApp",
                    "list_models"
                )

        elif path == '/load':
            name = query.get('name', ['model1'])[0]
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                result = client.execute(
                    "examples.dirty_example.dirty_app:MLApp",
                    "load_model",
                    name
                )

        elif path == '/inference':
            model = query.get('model', ['default'])[0]
            data = query.get('data', ['test input'])[0]
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                result = client.execute(
                    "examples.dirty_example.dirty_app:MLApp",
                    "inference",
                    model,
                    data
                )

        elif path == '/unload':
            name = query.get('name', ['model1'])[0]
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                result = client.execute(
                    "examples.dirty_example.dirty_app:MLApp",
                    "unload_model",
                    name
                )

        elif path == '/fibonacci':
            n = int(query.get('n', ['10'])[0])
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                result = client.execute(
                    "examples.dirty_example.dirty_app:ComputeApp",
                    "fibonacci",
                    n
                )

        elif path == '/prime':
            n = int(query.get('n', ['17'])[0])
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                result = client.execute(
                    "examples.dirty_example.dirty_app:ComputeApp",
                    "prime_check",
                    n
                )

        elif path == '/stats':
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                ml_stats = client.execute(
                    "examples.dirty_example.dirty_app:MLApp",
                    "list_models"
                )
                compute_stats = client.execute(
                    "examples.dirty_example.dirty_app:ComputeApp",
                    "stats"
                )
                result = {
                    "ml_app": ml_stats,
                    "compute_app": compute_stats,
                    "http_worker_pid": os.getpid(),
                }

        # =====================================================================
        # Session endpoints (stash demo)
        # =====================================================================
        elif path == '/session/login':
            user_id = query.get('user_id', ['1'])[0]
            name = query.get('name', ['Anonymous'])[0]
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                result = client.execute(
                    "examples.dirty_example.dirty_app:SessionApp",
                    "login",
                    user_id=user_id,
                    user_data={"name": name}
                )

        elif path == '/session/get':
            user_id = query.get('user_id', ['1'])[0]
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                result = client.execute(
                    "examples.dirty_example.dirty_app:SessionApp",
                    "get_session",
                    user_id=user_id
                )

        elif path == '/session/list':
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                result = client.execute(
                    "examples.dirty_example.dirty_app:SessionApp",
                    "list_sessions"
                )

        elif path == '/session/logout':
            user_id = query.get('user_id', ['1'])[0]
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                result = client.execute(
                    "examples.dirty_example.dirty_app:SessionApp",
                    "logout",
                    user_id=user_id
                )

        elif path == '/session/stats':
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                result = client.execute(
                    "examples.dirty_example.dirty_app:SessionApp",
                    "get_stats"
                )

        elif path == '/session/clear':
            if client is None:
                result = {"error": "Dirty workers not enabled"}
            else:
                result = client.execute(
                    "examples.dirty_example.dirty_app:SessionApp",
                    "clear_all"
                )

        else:
            start_response('404 Not Found', [('Content-Type', 'application/json')])
            return [json.dumps({"error": "Not found"}).encode()]

        # Success response
        start_response('200 OK', [('Content-Type', 'application/json')])
        return [json.dumps(result, indent=2).encode()]

    except Exception as e:
        start_response('500 Internal Server Error', [('Content-Type', 'application/json')])
        return [json.dumps({
            "error": str(e),
            "type": type(e).__name__
        }).encode()]