44 "context"
55 "crypto/tls"
66 "encoding/json"
7+ "fmt"
8+ "log/slog"
79 "net"
810 "net/http"
911 "os"
@@ -14,29 +16,39 @@ import (
1416 "time"
1517
1618 "github.com/docker/model-runner/pkg/inference"
19+ "github.com/docker/model-runner/pkg/inference/backends/diffusers"
1720 "github.com/docker/model-runner/pkg/inference/backends/llamacpp"
1821 "github.com/docker/model-runner/pkg/inference/backends/sglang"
1922 "github.com/docker/model-runner/pkg/inference/config"
2023 "github.com/docker/model-runner/pkg/inference/models"
24+ "github.com/docker/model-runner/pkg/logging"
2125 "github.com/docker/model-runner/pkg/metrics"
2226 "github.com/docker/model-runner/pkg/routing"
2327 modeltls "github.com/docker/model-runner/pkg/tls"
24- "github.com/sirupsen/logrus"
2528)
2629
2730const (
2831 // DefaultTLSPort is the default TLS port for Moby
2932 DefaultTLSPort = "12444"
3033)
3134
32- var log = logrus .New ()
35+ // initLogger creates the application logger based on LOG_LEVEL env var.
36+ func initLogger () * slog.Logger {
37+ level := logging .ParseLevel (os .Getenv ("LOG_LEVEL" ))
38+ return logging .NewLogger (level )
39+ }
40+
41+ var log = initLogger ()
3342
3443// Log is the logger used by the application, exported for testing purposes.
3544var Log = log
3645
3746// testLog is a test-override logger used by createLlamaCppConfigFromEnv.
3847var testLog = log
3948
49+ // exitFunc is used for Fatal-like exits; overridden in tests.
50+ var exitFunc = func (code int ) { os .Exit (code ) }
51+
4052func main () {
4153 ctx , cancel := signal .NotifyContext (context .Background (), syscall .SIGINT , syscall .SIGTERM )
4254 defer cancel ()
@@ -48,7 +60,8 @@ func main() {
4860
4961 userHomeDir , err := os .UserHomeDir ()
5062 if err != nil {
51- log .Fatalf ("Failed to get user home directory: %v" , err )
63+ log .Error ("Failed to get user home directory" , "error" , err )
64+ exitFunc (1 )
5265 }
5366
5467 modelPath := os .Getenv ("MODELS_PATH" )
@@ -90,21 +103,21 @@ func main() {
90103 }
91104 baseTransport .Proxy = http .ProxyFromEnvironment
92105
93- log .Infof ("LLAMA_SERVER_PATH: %s " , llamaServerPath )
106+ log .Info ("LLAMA_SERVER_PATH" , "path "llamaServerPath )
94107 if vllmServerPath != "" {
95- log .Infof ("VLLM_SERVER_PATH: %s " , vllmServerPath )
108+ log .Info ("VLLM_SERVER_PATH" , "path "vllmServerPath )
96109 }
97110 if sglangServerPath != "" {
98- log .Infof ("SGLANG_SERVER_PATH: %s " , sglangServerPath )
111+ log .Info ("SGLANG_SERVER_PATH" , "path "sglangServerPath )
99112 }
100113 if mlxServerPath != "" {
101- log .Infof ("MLX_SERVER_PATH: %s " , mlxServerPath )
114+ log .Info ("MLX_SERVER_PATH" , "path "mlxServerPath )
102115 }
103116 if diffusersServerPath != "" {
104- log .Infof ("DIFFUSERS_SERVER_PATH: %s " , diffusersServerPath )
117+ log .Info ("DIFFUSERS_SERVER_PATH" , "path "diffusersServerPath )
105118 }
106119 if vllmMetalServerPath != "" {
107- log .Infof ("VLLM_METAL_SERVER_PATH: %s " , vllmMetalServerPath )
120+ log .Info ("VLLM_METAL_SERVER_PATH" , "path "vllmMetalServerPath )
108121 }
109122
110123 // Create llama.cpp configuration from environment variables
@@ -121,7 +134,7 @@ func main() {
121134 Log : log ,
122135 ClientConfig : models.ClientConfig {
123136 StoreRootPath : modelPath ,
124- Logger : log .WithFields (logrus. Fields { "component" : "model-manager" } ),
137+ Logger : log .With ( "component" , "model-manager" ),
125138 Transport : baseTransport ,
126139 },
127140 Backends : append (
@@ -139,17 +152,21 @@ func main() {
139152 DiffusersPath : diffusersServerPath ,
140153 }),
141154 routing.BackendDef {Name : sglang .Name , Init : func (mm * models.Manager ) (inference.Backend , error ) {
142- return sglang .New (log , mm , log .WithFields (logrus.Fields {"component" : sglang .Name }), nil , sglangServerPath )
155+ return sglang .New (log , mm , log .With ("component" , sglang .Name ), nil , sglangServerPath )
156+ }},
157+ routing.BackendDef {Name : diffusers .Name , Init : func (mm * models.Manager ) (inference.Backend , error ) {
158+ return diffusers .New (log , mm , log .With ("component" , diffusers .Name ), nil , diffusersServerPath )
143159 }},
144160 ),
145161 OnBackendError : func (name string , err error ) {
146- log .Fatalf ("unable to initialize %s backend: %v" , name , err )
162+ log .Error ("unable to initialize backend" , "backend" , name , "error" , err )
163+ exitFunc (1 )
147164 },
148165 DefaultBackendName : llamacpp .Name ,
149166 HTTPClient : http .DefaultClient ,
150167 MetricsTracker : metrics .NewTracker (
151168 http .DefaultClient ,
152- log .WithField ("component" , "metrics" ),
169+ log .With ("component" , "metrics" ),
153170 "" ,
154171 false ,
155172 ),
@@ -169,14 +186,14 @@ func main() {
169186 r .HandleFunc ("/version" , func (w http.ResponseWriter , req * http.Request ) {
170187 w .Header ().Set ("Content-Type" , "application/json" )
171188 if err := json .NewEncoder (w ).Encode (map [string ]string {"version" : Version }); err != nil {
172- log .Warnf ("failed to write version response: %v " , err )
189+ log .Warn ("failed to write version response" , "error "err )
173190 }
174191 })
175192
176193 // Metrics endpoint
177194 if os .Getenv ("DISABLE_METRICS" ) != "1" {
178195 metricsHandler := metrics .NewAggregatedMetricsHandler (
179- log .WithField ("component" , "metrics" ),
196+ log .With ("component" , "metrics" ),
180197 s .SchedulerHTTP ,
181198 )
182199 r .Handle ("/metrics" , metricsHandler )
@@ -187,7 +204,8 @@ func main() {
187204 },
188205 })
189206 if err != nil {
190- log .Fatalf ("failed to initialize service: %v" , err )
207+ log .Error ("failed to initialize service" , "error" , err )
208+ exitFunc (1 )
191209 }
192210
193211 server := & http.Server {
@@ -205,7 +223,7 @@ func main() {
205223 if tcpPort != "" {
206224 // Use TCP port
207225 addr := ":" + tcpPort
208- log .Infof ("Listening on TCP port %s " , tcpPort )
226+ log .Info ("Listening on TCP port" , "port "tcpPort )
209227 server .Addr = addr
210228 go func () {
211229 serverErrors <- server .ListenAndServe ()
@@ -214,12 +232,14 @@ func main() {
214232 // Use Unix socket
215233 if err := os .Remove (sockName ); err != nil {
216234 if ! os .IsNotExist (err ) {
217- log .Fatalf ("Failed to remove existing socket: %v" , err )
235+ log .Error ("Failed to remove existing socket" , "error" , err )
236+ exitFunc (1 )
218237 }
219238 }
220239 ln , err := net .ListenUnix ("unix" , & net.UnixAddr {Name : sockName , Net : "unix" })
221240 if err != nil {
222- log .Fatalf ("Failed to listen on socket: %v" , err )
241+ log .Error ("Failed to listen on socket" , "error" , err )
242+ exitFunc (1 )
223243 }
224244 go func () {
225245 serverErrors <- server .Serve (ln )
@@ -244,19 +264,22 @@ func main() {
244264 var err error
245265 certPath , keyPath , err = modeltls .EnsureCertificates ("" , "" )
246266 if err != nil {
247- log .Fatalf ("Failed to ensure TLS certificates: %v" , err )
267+ log .Error ("Failed to ensure TLS certificates" , "error" , err )
268+ exitFunc (1 )
248269 }
249- log .Infof ("Using TLS certificate: %s " , certPath )
250- log .Infof ("Using TLS key: %s " , keyPath )
270+ log .Info ("Using TLS certificate" , "cert "certPath )
271+ log .Info ("Using TLS key" , "key "keyPath )
251272 } else {
252- log .Fatal ("TLS enabled but no certificate provided and auto-cert is disabled" )
273+ log .Error ("TLS enabled but no certificate provided and auto-cert is disabled" )
274+ exitFunc (1 )
253275 }
254276 }
255277
256278 // Load TLS configuration
257279 tlsConfig , err := modeltls .LoadTLSConfig (certPath , keyPath )
258280 if err != nil {
259- log .Fatalf ("Failed to load TLS configuration: %v" , err )
281+ log .Error ("Failed to load TLS configuration" , "error" , err )
282+ exitFunc (1 )
260283 }
261284
262285 tlsServer = & http.Server {
@@ -266,7 +289,7 @@ func main() {
266289 ReadHeaderTimeout : 10 * time .Second ,
267290 }
268291
269- log .Infof ("Listening on TLS port %s " , tlsPort )
292+ log .Info ("Listening on TLS port" , "port "tlsPort )
270293 go func () {
271294 // Use ListenAndServeTLS with empty strings since TLSConfig already has the certs
272295 ln , err := tls .Listen ("tcp" , tlsServer .Addr , tlsConfig )
@@ -294,30 +317,30 @@ func main() {
294317 select {
295318 case err := <- serverErrors :
296319 if err != nil {
297- log .Errorf ("Server error: %v " , err )
320+ log .Error ("Server error" , "error "err )
298321 }
299322 case err := <- tlsServerErrorsChan :
300323 if err != nil {
301- log .Errorf ("TLS server error: %v " , err )
324+ log .Error ("TLS server error" , "error "err )
302325 }
303326 case <- ctx .Done ():
304- log .Infoln ("Shutdown signal received" )
305- log .Infoln ("Shutting down the server" )
327+ log .Info ("Shutdown signal received" )
328+ log .Info ("Shutting down the server" )
306329 if err := server .Close (); err != nil {
307- log .Errorf ("Server shutdown error: %v " , err )
330+ log .Error ("Server shutdown error" , "error "err )
308331 }
309332 if tlsServer != nil {
310- log .Infoln ("Shutting down the TLS server" )
333+ log .Info ("Shutting down the TLS server" )
311334 if err := tlsServer .Close (); err != nil {
312- log .Errorf ("TLS server shutdown error: %v " , err )
335+ log .Error ("TLS server shutdown error" , "error "err )
313336 }
314337 }
315- log .Infoln ("Waiting for the scheduler to stop" )
338+ log .Info ("Waiting for the scheduler to stop" )
316339 if err := <- schedulerErrors ; err != nil {
317- log .Errorf ("Scheduler error: %v " , err )
340+ log .Error ("Scheduler error" , "error "err )
318341 }
319342 }
320- log .Infoln ("Docker Model Runner stopped" )
343+ log .Info ("Docker Model Runner stopped" )
321344}
322345
323346// createLlamaCppConfigFromEnv creates a LlamaCppConfig from environment variables
@@ -338,12 +361,13 @@ func createLlamaCppConfigFromEnv() config.BackendConfig {
338361 for _ , arg := range args {
339362 for _ , disallowed := range disallowedArgs {
340363 if arg == disallowed {
341- testLog .Fatalf ("LLAMA_ARGS cannot override the %s argument as it is controlled by the model runner" , disallowed )
364+ testLog .Error (fmt .Sprintf ("LLAMA_ARGS cannot override the %s argument as it is controlled by the model runner" , disallowed ))
365+ exitFunc (1 )
342366 }
343367 }
344368 }
345369
346- testLog .Infof ("Using custom arguments: %v" , args )
370+ testLog .Info ("Using custom arguments" , "args" , fmt . Sprintf ( " %v"args ) )
347371 return & llamacpp.Config {
348372 Args : args ,
349373 }
0 commit comments