Skip to content
This repository was archived by the owner on Apr 15, 2026. It is now read-only.

Commit f0875dc

Browse files
committed
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
1 parent 0ec1e75 commit f0875dc

15 files changed

Lines changed: 566 additions & 145 deletions

File tree

.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: 5 additions & 22 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,27 +40,12 @@ 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()
6448
log.Infow("log level", "level", logLevel)
65-
log.Infow("env log level", "level", os.Getenv("COG_LOG_LEVEL"))
6649
// One-shot mode requires procedure mode
6750
if s.OneShot && !s.UseProcedureMode {
6851
log.Error("one-shot mode requires procedure mode")
@@ -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: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
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+
// Logger embeds zap.Logger and adds Trace level support
18+
type Logger struct {
19+
*zap.Logger
20+
}
21+
22+
// SugaredLogger embeds zap.SugaredLogger and adds Trace level support
23+
type SugaredLogger struct {
24+
*zap.SugaredLogger
25+
}
26+
27+
// New creates a new logger with the given name
28+
func New(name string) *Logger {
29+
// Check if we should use development config (console format)
30+
logFormat := os.Getenv("LOG_FORMAT")
31+
isDevelopment := logFormat == "development" || logFormat == "console"
32+
33+
var cfg zap.Config
34+
if isDevelopment {
35+
cfg = zap.NewDevelopmentConfig()
36+
cfg.Level = zap.NewAtomicLevelAt(zapcore.DebugLevel)
37+
cfg.EncoderConfig.EncodeLevel = zapcore.CapitalColorLevelEncoder
38+
} else {
39+
cfg = zap.NewProductionConfig()
40+
cfg.Level = zap.NewAtomicLevelAt(zapcore.InfoLevel)
41+
cfg.EncoderConfig.EncodeLevel = zapcore.LowercaseLevelEncoder
42+
}
43+
44+
// Set log level from environment (COG_LOG_LEVEL takes precedence, fallback to LOG_LEVEL)
45+
logLevel := os.Getenv("COG_LOG_LEVEL")
46+
if logLevel == "" {
47+
logLevel = os.Getenv("LOG_LEVEL")
48+
}
49+
if logLevel != "" {
50+
level, err := parseLevel(logLevel)
51+
if err != nil {
52+
fmt.Printf("Failed to parse log level \"%s\": %s\n", logLevel, err) //nolint:forbidigo // logger setup error reporting
53+
} else {
54+
cfg.Level = zap.NewAtomicLevelAt(level)
55+
}
56+
}
57+
58+
// Set output file if LOG_FILE is specified
59+
logFile := os.Getenv("LOG_FILE")
60+
if logFile != "" {
61+
cfg.OutputPaths = []string{logFile}
62+
cfg.ErrorOutputPaths = []string{logFile}
63+
} else {
64+
cfg.OutputPaths = []string{"stdout"}
65+
cfg.ErrorOutputPaths = []string{"stderr"}
66+
}
67+
68+
// Common encoder config
69+
cfg.EncoderConfig.TimeKey = "timestamp"
70+
cfg.EncoderConfig.LevelKey = "severity"
71+
cfg.EncoderConfig.NameKey = "logger"
72+
cfg.EncoderConfig.CallerKey = "caller"
73+
cfg.EncoderConfig.MessageKey = "message"
74+
cfg.EncoderConfig.StacktraceKey = "stacktrace"
75+
cfg.EncoderConfig.EncodeTime = zapcore.ISO8601TimeEncoder
76+
cfg.EncoderConfig.EncodeDuration = zapcore.StringDurationEncoder
77+
cfg.EncoderConfig.EncodeCaller = zapcore.ShortCallerEncoder
78+
79+
// Disable sampling for now (can be re-enabled later if needed)
80+
cfg.Sampling = nil
81+
82+
zapLogger, err := cfg.Build()
83+
if err != nil {
84+
panic(fmt.Sprintf("Failed to build logger: %v", err))
85+
}
86+
87+
return &Logger{Logger: zapLogger.Named(name)}
88+
}
89+
90+
// parseLevel parses log level string including our custom "trace" level
91+
func parseLevel(level string) (zapcore.Level, error) {
92+
switch strings.ToLower(level) {
93+
case "trace":
94+
return TraceLevel, nil
95+
case "debug":
96+
return zapcore.DebugLevel, nil
97+
case "info":
98+
return zapcore.InfoLevel, nil
99+
case "warn", "warning":
100+
return zapcore.WarnLevel, nil
101+
case "error":
102+
return zapcore.ErrorLevel, nil
103+
default:
104+
return zapcore.InfoLevel, fmt.Errorf("unknown log level: %s", level)
105+
}
106+
}
107+
108+
// Override Sugar to return our custom SugaredLogger
109+
func (l *Logger) Sugar() *SugaredLogger {
110+
return &SugaredLogger{SugaredLogger: l.Logger.Sugar()}
111+
}
112+
113+
// Override Named to return our custom Logger
114+
func (l *Logger) Named(name string) *Logger {
115+
return &Logger{Logger: l.Logger.Named(name)}
116+
}
117+
118+
// Override With to return our custom Logger
119+
func (l *Logger) With(fields ...zap.Field) *Logger {
120+
return &Logger{Logger: l.Logger.With(fields...)}
121+
}
122+
123+
// Override WithOptions to return our custom Logger
124+
func (l *Logger) WithOptions(opts ...zap.Option) *Logger {
125+
return &Logger{Logger: l.Logger.WithOptions(opts...)}
126+
}
127+
128+
// Add Trace method to Logger
129+
func (l *Logger) Trace(msg string, fields ...zap.Field) {
130+
l.Log(TraceLevel, msg, fields...)
131+
}
132+
133+
// Add Trace method to SugaredLogger
134+
func (s *SugaredLogger) Trace(args ...any) {
135+
s.Log(TraceLevel, args...)
136+
}
137+
138+
// Add Tracew method to SugaredLogger
139+
func (s *SugaredLogger) Tracew(msg string, keysAndValues ...any) {
140+
s.Logw(TraceLevel, msg, keysAndValues...)
141+
}

0 commit comments

Comments
 (0)