worker.py

"""
worker.py

Python FaaS worker process. Listens on a Unix Domain Socket and processes one
request at a time, simulating AI model inference via Pillow image operations.

CSCI 599: Network Systems for Cloud Computing
University of Southern California
"""

import socket
import sys
import os
import time
import base64
import io

try:
    from PIL import Image, ImageOps
except ImportError:
    print("[-] 请先安装 Pillow 库: pip install Pillow")
    sys.exit(1)

def process_image(base64_str):
    """真实业务负载：解码 Base64 -> 反转图像颜色 -> 重新编码 Base64"""
    # 移除可能存在的前缀 data:image/jpeg;base64,
    if "base64," in base64_str:
        base64_str = base64_str.split("base64,")[1]
        
    image_data = base64.b64decode(base64_str)
    image = Image.open(io.BytesIO(image_data)).convert('RGB')
    
    # 图像处理核心逻辑 (反转颜色)
    inverted_image = ImageOps.invert(image)
    
    buffered = io.BytesIO()
    inverted_image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode('utf-8')

def main():
    if len(sys.argv) != 2:
        print("Usage: python3 worker.py <socket_path>")
        sys.exit(1)
        
    socket_path = sys.argv[1]

    print("[Python Worker] 正在初始化冷启动环境...", flush=True)
    
    # 🔥 核心：物理模拟极度缓慢的冷启动 (例如加载 AI 模型)
    time.sleep(0.8) 
    print("[Python Worker] 核心库加载完毕，冷启动结束！准备就绪。", flush=True)

    if os.path.exists(socket_path):
        os.remove(socket_path)

    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    server.bind(socket_path)
    server.listen(5)
    print(f"[Python Worker] Listening on {socket_path}...", flush=True)

    while True:
        conn, addr = server.accept()
        data = b""
        
        # 接收数据
        while True:
            chunk = conn.recv(8192)
            if not chunk:
                break
            data += chunk
            if len(chunk) < 8192:
                break
                
        if not data:
            conn.close()
            continue
        print(f"[Python Worker] Get a data", flush=True)
        # Simulate AI model inference time (~500ms per request)
        time.sleep(0.5)
        try:
            req_text = data.decode('utf-8', errors='ignore')
            
            # 简单的 HTTP 报文解析，剥离请求头，只取 Body 部分
            if "\r\n\r\n" in req_text:
                body = req_text.split("\r\n\r\n", 1)[1].strip()
            else:
                body = req_text.strip()
            
            try:
                # 当作图片进行 CV 处理
                result_body = process_image(body)
            except Exception as e:
                # 如果传来的不是图片，做个容错处理
                result_body = f"Fallback Processed (Not an image): {body[::-1]}"
            
            # 包装成标准的 HTTP 响应发回
            http_response = (
                "HTTP/1.1 200 OK\r\n"
                "Content-Type: text/plain\r\n"
                f"Content-Length: {len(result_body)}\r\n"
                "Connection: close\r\n\r\n"
                f"{result_body}"
            )
            
            conn.sendall(http_response.encode('utf-8'))
        except Exception as e:
            print(f"[Worker Error] {e}")
        finally:
            conn.close()

if __name__ == "__main__":
    main()