@@ -4,22 +4,94 @@ import (
44 "context"
55 "fmt"
66 "io"
7- _ "net/http/pprof"
87 "os"
98 "path/filepath"
109
1110 "github.com/sirupsen/logrus"
11+ "github.com/urfave/cli/v3"
1212
1313 "github.com/dstackai/dstack/runner/consts"
1414 "github.com/dstackai/dstack/runner/internal/log"
1515 "github.com/dstackai/dstack/runner/internal/runner/api"
1616)
1717
18+ // Version is a build-time variable. The value is overridden by ldflags.
19+ var Version string
20+
1821func main () {
19- App ()
22+ os .Exit (mainInner ())
23+ }
24+
25+ func mainInner () int {
26+ var tempDir string
27+ var homeDir string
28+ var httpPort int
29+ var sshPort int
30+ var logLevel int
31+
32+ cmd := & cli.Command {
33+ Name : "dstack-runner" ,
34+ Usage : "configure and start dstack-runner" ,
35+ Version : Version ,
36+ Flags : []cli.Flag {
37+ & cli.IntFlag {
38+ Name : "log-level" ,
39+ Value : 2 ,
40+ DefaultText : "4 (Info)" ,
41+ Usage : "log verbosity level: 2 (Error), 3 (Warning), 4 (Info), 5 (Debug), 6 (Trace)" ,
42+ Destination : & logLevel ,
43+ },
44+ },
45+ Commands : []* cli.Command {
46+ {
47+ Name : "start" ,
48+ Usage : "Start dstack-runner" ,
49+ Flags : []cli.Flag {
50+ & cli.StringFlag {
51+ Name : "temp-dir" ,
52+ Usage : "Temporary directory for logs and other files" ,
53+ Value : consts .RunnerTempDir ,
54+ Destination : & tempDir ,
55+ TakesFile : true ,
56+ },
57+ & cli.StringFlag {
58+ Name : "home-dir" ,
59+ Usage : "HomeDir directory for credentials and $HOME" ,
60+ Value : consts .RunnerHomeDir ,
61+ Destination : & homeDir ,
62+ TakesFile : true ,
63+ },
64+ & cli.IntFlag {
65+ Name : "http-port" ,
66+ Usage : "Set a http port" ,
67+ Value : consts .RunnerHTTPPort ,
68+ Destination : & httpPort ,
69+ },
70+ & cli.IntFlag {
71+ Name : "ssh-port" ,
72+ Usage : "Set the ssh port" ,
73+ Value : consts .RunnerSSHPort ,
74+ Destination : & sshPort ,
75+ },
76+ },
77+ Action : func (cxt context.Context , cmd * cli.Command ) error {
78+ return start (cxt , tempDir , homeDir , httpPort , sshPort , logLevel , Version )
79+ },
80+ },
81+ },
82+ }
83+
84+ ctx := context .Background ()
85+
86+ if err := cmd .Run (ctx , os .Args ); err != nil {
87+ log .Error (ctx , err .Error ())
88+ return 1
89+ }
90+
91+ return 0
2092}
2193
22- func start (tempDir string , homeDir string , httpPort int , sshPort int , logLevel int , version string ) error {
94+ func start (ctx context. Context , tempDir string , homeDir string , httpPort int , sshPort int , logLevel int , version string ) error {
2395 if err := os .MkdirAll (tempDir , 0o755 ); err != nil {
2496 return fmt .Errorf ("create temp directory: %w" , err )
2597 }
@@ -31,20 +103,20 @@ func start(tempDir string, homeDir string, httpPort int, sshPort int, logLevel i
31103 defer func () {
32104 closeErr := defaultLogFile .Close ()
33105 if closeErr != nil {
34- log .Error (context . TODO () , "Failed to close default log file" , "err" , closeErr )
106+ log .Error (ctx , "Failed to close default log file" , "err" , closeErr )
35107 }
36108 }()
37109
38110 log .DefaultEntry .Logger .SetOutput (io .MultiWriter (os .Stdout , defaultLogFile ))
39111 log .DefaultEntry .Logger .SetLevel (logrus .Level (logLevel ))
40112
41- server , err := api .NewServer (context . TODO () , tempDir , homeDir , fmt .Sprintf (":%d" , httpPort ), sshPort , version )
113+ server , err := api .NewServer (ctx , tempDir , homeDir , fmt .Sprintf (":%d" , httpPort ), sshPort , version )
42114 if err != nil {
43115 return fmt .Errorf ("create server: %w" , err )
44116 }
45117
46- log .Trace (context . TODO () , "Starting API server" , "port" , httpPort )
47- if err := server .Run (); err != nil {
118+ log .Trace (ctx , "Starting API server" , "port" , httpPort )
119+ if err := server .Run (ctx ); err != nil {
48120 return fmt .Errorf ("server failed: %w" , err )
49121 }
50122
0 commit comments