Skip to content

Commit 1e3ff7f

Browse files
Logging improvements (#216)
* Implement internal logging package with Trace level support - Create internal/logging package with embedded zap logger - Add Trace level (-8) below Debug level with full sugar support - Implement Python process protection to cap log levels at debug - Replace external github.com/replicate/go/logging dependency - Wire new logger throughout codebase (main.go, server, service, manager) - Create loggingtest package for test logger compatibility - Update all test files to use new test logger infrastructure - Add comprehensive unit tests for logging functionality - Configure golangci-yml with depguard to prevent test package misuse - Support environment variables: COG_LOG_LEVEL, LOG_LEVEL, LOG_FORMAT, LOG_FILE * Add custom TRACE level encoders and migrate verbose Debug logs - Add custom level encoders (uppercase, lowercase, color) to display "TRACE" instead of "LEVEL(-8)" - Implement test logger with proper TRACE level display for zaptest compatibility - Add SugaredLogger chaining support (With/Named) to maintain custom type - Migrate 25 verbose Debug logs to Trace level across runner, server, service, webhook - Keep important Debug logs for actual debugging (errors, state changes, subprocess management) - Verbose implementation details now at Trace: file operations, response watchers, shutdown flow * Migrate Info logs to Debug level for cleaner operational output - Move 16 Info logs to Debug: component initialization, runner management, prediction flow details - Keep critical Info logs: service lifecycle, prediction completion, force kills, webhooks - Move 4 Info logs to Trace: internal shutdown flow and cleanup state tracking - Maintain clean Info level for exceptional events only - Remove unused customLevelEncoder function to fix linter warnings
1 parent 3ebc411 commit 1e3ff7f

File tree

15 files changed

+706
-203
lines changed

15 files changed

+706
-203
lines changed

.golangci.yml

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ linters:
2323
- containedctx
2424
- contextcheck
2525
- copyloopvar
26+
- depguard
2627
- dogsled
2728
- errcheck
2829
- errchkjson
@@ -58,6 +59,7 @@ linters:
5859
- errcheck
5960
- forcetypeassert
6061
- gosec
62+
- depguard
6163
- path: 'cmd/(.+)/main\.go'
6264
linters:
6365
- forbidigo
@@ -77,6 +79,12 @@ linters:
7779
allow-no-explanation: []
7880
require-explanation: true
7981
require-specific: true
82+
depguard:
83+
rules:
84+
test-only:
85+
deny:
86+
- pkg: github.com/replicate/cog-runtime/internal/loggingtest
87+
desc: "loggingtest is test-only"
8088
revive:
8189
confidence: 0.8
8290
severity: warning

cmd/cog/main.go

Lines changed: 7 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,9 @@ import (
99
"time"
1010

1111
"github.com/alecthomas/kong"
12-
"github.com/replicate/go/logging"
13-
"go.uber.org/zap"
14-
"go.uber.org/zap/zapcore"
1512

1613
"github.com/replicate/cog-runtime/internal/config"
14+
"github.com/replicate/cog-runtime/internal/logging"
1715
"github.com/replicate/cog-runtime/internal/runner"
1816
"github.com/replicate/cog-runtime/internal/service"
1917
"github.com/replicate/cog-runtime/internal/version"
@@ -42,30 +40,15 @@ type CLI struct {
4240
Test TestCmd `cmd:"" help:"Run model tests to verify functionality"`
4341
}
4442

45-
// createBaseLogger creates a base logger with configurable level
46-
func createBaseLogger(name string) *zap.Logger {
47-
logLevel := os.Getenv("COG_LOG_LEVEL")
48-
if logLevel == "" {
49-
logLevel = "info"
50-
}
51-
_, err := zapcore.ParseLevel(logLevel)
52-
if err != nil {
53-
fmt.Printf("Failed to parse log level \"%s\": %s\n", logLevel, err) //nolint:forbidigo // logger setup error reporting
54-
}
55-
56-
return logging.New(name).WithOptions(zap.IncreaseLevel(zapcore.DebugLevel))
57-
}
58-
5943
// buildServiceConfig converts CLI ServerCmd to service configuration
6044
func buildServiceConfig(s *ServerCmd) (config.Config, error) {
61-
log := createBaseLogger("cog-config").Sugar()
45+
log := logging.New("cog-config").Sugar()
6246

6347
logLevel := log.Level()
64-
log.Infow("log level", "level", logLevel)
65-
log.Infow("env log level", "level", os.Getenv("COG_LOG_LEVEL"))
48+
log.Debugw("log level", "level", logLevel)
6649
// One-shot mode requires procedure mode
6750
if s.OneShot && !s.UseProcedureMode {
68-
log.Error("one-shot mode requires procedure mode")
51+
log.Fatal("one-shot mode requires procedure mode")
6952
return config.Config{}, fmt.Errorf("one-shot mode requires procedure mode, use --use-procedure-mode")
7053
}
7154

@@ -115,7 +98,7 @@ func buildServiceConfig(s *ServerCmd) (config.Config, error) {
11598

11699
func (s *ServerCmd) Run() error {
117100
// Create base logger
118-
baseLogger := createBaseLogger("cog")
101+
baseLogger := logging.New("cog")
119102
log := baseLogger.Sugar()
120103

121104
// Build service configuration
@@ -143,7 +126,7 @@ func (s *ServerCmd) Run() error {
143126
}
144127

145128
func (s *SchemaCmd) Run() error {
146-
log := createBaseLogger("cog-schema").Sugar()
129+
log := logging.New("cog-schema").Sugar()
147130

148131
wd, err := os.Getwd()
149132
if err != nil {
@@ -169,7 +152,7 @@ func (s *SchemaCmd) Run() error {
169152
}
170153

171154
func (t *TestCmd) Run() error {
172-
log := createBaseLogger("cog-test").Sugar()
155+
log := logging.New("cog-test").Sugar()
173156

174157
wd, err := os.Getwd()
175158
if err != nil {

internal/logging/logger.go

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
package logging
2+
3+
import (
4+
"fmt"
5+
"os"
6+
"strings"
7+
8+
"go.uber.org/zap"
9+
"go.uber.org/zap/zapcore"
10+
)
11+
12+
// Custom log levels with Trace below Debug
13+
const (
14+
TraceLevel = zapcore.Level(-8) // Below Debug (-4)
15+
)
16+
17+
// customLowercaseLevelEncoder handles our custom Trace level display (lowercase)
18+
func customLowercaseLevelEncoder(level zapcore.Level, enc zapcore.PrimitiveArrayEncoder) {
19+
switch level {
20+
case TraceLevel:
21+
enc.AppendString("trace")
22+
default:
23+
zapcore.LowercaseLevelEncoder(level, enc)
24+
}
25+
}
26+
27+
// customColorLevelEncoder handles our custom Trace level display (with colors)
28+
func customColorLevelEncoder(level zapcore.Level, enc zapcore.PrimitiveArrayEncoder) {
29+
switch level {
30+
case TraceLevel:
31+
enc.AppendString("\x1b[90mTRACE\x1b[0m") // Gray color for trace
32+
default:
33+
zapcore.CapitalColorLevelEncoder(level, enc)
34+
}
35+
}
36+
37+
// Logger embeds zap.Logger and adds Trace level support
38+
type Logger struct {
39+
*zap.Logger
40+
}
41+
42+
// SugaredLogger embeds zap.SugaredLogger and adds Trace level support
43+
type SugaredLogger struct {
44+
*zap.SugaredLogger
45+
}
46+
47+
// New creates a new logger with the given name
48+
func New(name string) *Logger {
49+
// Check if we should use development config (console format)
50+
logFormat := os.Getenv("LOG_FORMAT")
51+
isDevelopment := logFormat == "development" || logFormat == "console"
52+
53+
var cfg zap.Config
54+
if isDevelopment {
55+
cfg = zap.NewDevelopmentConfig()
56+
cfg.Level = zap.NewAtomicLevelAt(zapcore.DebugLevel)
57+
cfg.EncoderConfig.EncodeLevel = customColorLevelEncoder
58+
} else {
59+
cfg = zap.NewProductionConfig()
60+
cfg.Level = zap.NewAtomicLevelAt(zapcore.InfoLevel)
61+
cfg.EncoderConfig.EncodeLevel = customLowercaseLevelEncoder
62+
}
63+
64+
// Set log level from environment (COG_LOG_LEVEL takes precedence, fallback to LOG_LEVEL)
65+
logLevel := os.Getenv("COG_LOG_LEVEL")
66+
if logLevel == "" {
67+
logLevel = os.Getenv("LOG_LEVEL")
68+
}
69+
if logLevel != "" {
70+
level, err := parseLevel(logLevel)
71+
if err != nil {
72+
fmt.Printf("Failed to parse log level \"%s\": %s\n", logLevel, err) //nolint:forbidigo // logger setup error reporting
73+
} else {
74+
cfg.Level = zap.NewAtomicLevelAt(level)
75+
}
76+
}
77+
78+
// Set output file if LOG_FILE is specified
79+
logFile := os.Getenv("LOG_FILE")
80+
if logFile != "" {
81+
cfg.OutputPaths = []string{logFile}
82+
cfg.ErrorOutputPaths = []string{logFile}
83+
} else {
84+
cfg.OutputPaths = []string{"stdout"}
85+
cfg.ErrorOutputPaths = []string{"stderr"}
86+
}
87+
88+
// Common encoder config
89+
cfg.EncoderConfig.TimeKey = "timestamp"
90+
cfg.EncoderConfig.LevelKey = "severity"
91+
cfg.EncoderConfig.NameKey = "logger"
92+
cfg.EncoderConfig.CallerKey = "caller"
93+
cfg.EncoderConfig.MessageKey = "message"
94+
cfg.EncoderConfig.StacktraceKey = "stacktrace"
95+
cfg.EncoderConfig.EncodeTime = zapcore.ISO8601TimeEncoder
96+
cfg.EncoderConfig.EncodeDuration = zapcore.StringDurationEncoder
97+
cfg.EncoderConfig.EncodeCaller = zapcore.ShortCallerEncoder
98+
99+
// Disable sampling for now (can be re-enabled later if needed)
100+
cfg.Sampling = nil
101+
102+
zapLogger, err := cfg.Build()
103+
if err != nil {
104+
panic(fmt.Sprintf("Failed to build logger: %v", err))
105+
}
106+
107+
return &Logger{Logger: zapLogger.Named(name)}
108+
}
109+
110+
// parseLevel parses log level string including our custom "trace" level
111+
func parseLevel(level string) (zapcore.Level, error) {
112+
switch strings.ToLower(level) {
113+
case "trace":
114+
return TraceLevel, nil
115+
case "debug":
116+
return zapcore.DebugLevel, nil
117+
case "info":
118+
return zapcore.InfoLevel, nil
119+
case "warn", "warning":
120+
return zapcore.WarnLevel, nil
121+
case "error":
122+
return zapcore.ErrorLevel, nil
123+
default:
124+
return zapcore.InfoLevel, fmt.Errorf("unknown log level: %s", level)
125+
}
126+
}
127+
128+
// Override Sugar to return our custom SugaredLogger
129+
func (l *Logger) Sugar() *SugaredLogger {
130+
return &SugaredLogger{SugaredLogger: l.Logger.Sugar()}
131+
}
132+
133+
// Override Named to return our custom Logger
134+
func (l *Logger) Named(name string) *Logger {
135+
return &Logger{Logger: l.Logger.Named(name)}
136+
}
137+
138+
// Override With to return our custom Logger
139+
func (l *Logger) With(fields ...zap.Field) *Logger {
140+
return &Logger{Logger: l.Logger.With(fields...)}
141+
}
142+
143+
// Override WithOptions to return our custom Logger
144+
func (l *Logger) WithOptions(opts ...zap.Option) *Logger {
145+
return &Logger{Logger: l.Logger.WithOptions(opts...)}
146+
}
147+
148+
// Add Trace method to Logger
149+
func (l *Logger) Trace(msg string, fields ...zap.Field) {
150+
l.Log(TraceLevel, msg, fields...)
151+
}
152+
153+
// Add Trace method to SugaredLogger
154+
func (s *SugaredLogger) Trace(args ...any) {
155+
s.Log(TraceLevel, args...)
156+
}
157+
158+
// Add Tracew method to SugaredLogger
159+
func (s *SugaredLogger) Tracew(msg string, keysAndValues ...any) {
160+
s.Logw(TraceLevel, msg, keysAndValues...)
161+
}
162+
163+
// Override With to return our custom SugaredLogger
164+
func (s *SugaredLogger) With(args ...any) *SugaredLogger {
165+
return &SugaredLogger{SugaredLogger: s.SugaredLogger.With(args...)}
166+
}
167+
168+
// Override Named to return our custom SugaredLogger
169+
func (s *SugaredLogger) Named(name string) *SugaredLogger {
170+
return &SugaredLogger{SugaredLogger: s.SugaredLogger.Named(name)}
171+
}

0 commit comments

Comments
 (0)