fix: run on runpod

runpod · Dec 23, 2024 · d7a0cf2 · d7a0cf2
1 parent deb06de
commit d7a0cf2
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 1 deletion.
diff --git a/handler.py b/handler.py
@@ -1,6 +1,8 @@
 import runpod
+import logging
 
 def handler(job):
+    logging.info(job)
     return job.get("input", {})
 
 runpod.serverless.start({

diff --git a/main.go b/main.go
@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"os"
 	"strconv"
+	"time"
 
 	"github.com/gin-gonic/gin"
 	prettyconsole "github.com/thessem/zap-prettyconsole"
@@ -190,6 +191,39 @@ func (h *Handler) JobDone(c *gin.Context) {
 	})
 }
 
+// LoggerMiddleware creates a middleware for logging requests
+func LoggerMiddleware(logger *zap.Logger) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		start := time.Now()
+		path := c.Request.URL.Path
+		query := c.Request.URL.RawQuery
+
+		// Log request
+		logger.Info("Incoming request",
+			zap.String("path", path),
+			zap.String("query", query),
+			zap.String("method", c.Request.Method),
+			zap.String("client_ip", c.ClientIP()),
+			zap.String("user_agent", c.Request.UserAgent()),
+		)
+
+		// Process request
+		c.Next()
+
+		// Log response
+		latency := time.Since(start)
+		status := c.Writer.Status()
+
+		logger.Info("Request completed",
+			zap.String("path", path),
+			zap.Int("status", status),
+			zap.Duration("latency", latency),
+			zap.Int("body_size", c.Writer.Size()),
+			zap.String("errors", c.Errors.ByType(gin.ErrorTypePrivate).String()),
+		)
+	}
+}
+
 func main() {
 	defer log.Sync()
 
@@ -198,7 +232,10 @@ func main() {
 	gin.SetMode(gin.ReleaseMode)
 
 	r := gin.New()
+	// Add recovery middleware
 	r.Use(gin.Recovery())
+	// Add logging middleware
+	r.Use(LoggerMiddleware(log))
 
 	h := NewHandler(log)
 
@@ -215,6 +252,8 @@ func main() {
 		port = "19981"
 	}
 
+	log.Info("Server starting", zap.String("port", port))
+
 	// Start server
 	if err := r.Run(":" + port); err != nil {
 		log.Fatal("Failed to start server", zap.Error(err))

diff --git a/script.sh b/script.sh
@@ -5,4 +5,4 @@ set -e
 
 sleep 5
 
-python3 handler.py
+RUNPOD_ENDPOINT_BASE_URL=http://localhost:19981/v2 python3 handler.py
Original file line number	Diff line number	Diff line change
Expand Up		@@ -5,4 +5,4 @@ set -e

		sleep 5

		python3 handler.py
		RUNPOD_ENDPOINT_BASE_URL=http://localhost:19981/v2 python3 handler.py