Add cmd_stop configuration to better support docker (#35)

Add `cmd_stop` to model configuration to run a command instead of sending a SIGTERM to shutdown a process before swapping.
This commit is contained in:
Benson Wong
2025-01-30 16:59:57 -08:00
committed by GitHub
parent 2833517eef
commit baeb0c4e7f
6 changed files with 125 additions and 22 deletions

View File

@@ -11,6 +11,7 @@ import (
type ModelConfig struct {
Cmd string `yaml:"cmd"`
CmdStop string `yaml:"cmd_stop"`
Proxy string `yaml:"proxy"`
Aliases []string `yaml:"aliases"`
Env []string `yaml:"env"`
@@ -22,6 +23,9 @@ type ModelConfig struct {
func (m *ModelConfig) SanitizedCommand() ([]string, error) {
return SanitizeCommand(m.Cmd)
}
func (m *ModelConfig) SanitizeCommandStop() ([]string, error) {
return SanitizeCommand(m.CmdStop)
}
type Config struct {
HealthCheckTimeout int `yaml:"healthCheckTimeout"`

View File

@@ -35,6 +35,11 @@ models:
aliases:
- "m2"
checkEndpoint: "/"
docker:
cmd: docker run -p 9999:8080 --name "my_container"
cmd_stop: docker stop my_container
proxy: "http://localhost:9999"
checkEndpoint: "/health"
healthCheckTimeout: 15
profiles:
test:
@@ -56,6 +61,7 @@ profiles:
Models: map[string]ModelConfig{
"model1": {
Cmd: "path/to/cmd --arg1 one",
CmdStop: "",
Proxy: "http://localhost:8080",
Aliases: []string{"m1", "model-one"},
Env: []string{"VAR1=value1", "VAR2=value2"},
@@ -63,11 +69,19 @@ profiles:
},
"model2": {
Cmd: "path/to/cmd --arg1 one",
CmdStop: "",
Proxy: "http://localhost:8081",
Aliases: []string{"m2"},
Env: nil,
CheckEndpoint: "/",
},
"docker": {
Cmd: `docker run -p 9999:8080 --name "my_container"`,
CmdStop: "docker stop my_container",
Proxy: "http://localhost:9999",
Env: nil,
CheckEndpoint: "/health",
},
},
HealthCheckTimeout: 15,
Profiles: map[string][]string{
@@ -99,6 +113,18 @@ func TestConfig_ModelConfigSanitizedCommand(t *testing.T) {
assert.Equal(t, []string{"python", "model1.py", "--arg1", "value1", "--arg2", "value2"}, args)
}
func TestConfig_ModelConfigSanitizedCommandStop(t *testing.T) {
config := &ModelConfig{
CmdStop: `docker stop my_container \
--arg1 1
--arg2 2`,
}
args, err := config.SanitizeCommandStop()
assert.NoError(t, err)
assert.Equal(t, []string{"docker", "stop", "my_container", "--arg1", "1", "--arg2", "2"}, args)
}
func TestConfig_FindConfig(t *testing.T) {
// TODO?

View File

@@ -153,12 +153,13 @@ func (p *Process) Stop() {
defer p.stateMutex.Unlock()
if p.state != StateReady {
fmt.Fprintf(p.logMonitor, "!!! Stop() called but Process State is not READY\n")
return
}
if p.cmd == nil || p.cmd.Process == nil {
// this situation should never happen... but if it does just update the state
fmt.Fprintf(p.logMonitor, "!!! State is Ready but Command is nil.")
fmt.Fprintf(p.logMonitor, "!!! State is Ready but Command is nil.\n")
p.state = StateStopped
return
}
@@ -166,30 +167,59 @@ func (p *Process) Stop() {
// Pretty sure this stopping code needs some work for windows and
// will be a source of pain in the future.
sigtermTimeout, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
sigtermNormal := make(chan error, 1)
go func() {
sigtermNormal <- p.cmd.Wait()
}()
p.cmd.Process.Signal(syscall.SIGTERM)
select {
case <-sigtermTimeout.Done():
fmt.Fprintf(p.logMonitor, "XXX Process for %s timed out waiting to stop, sending SIGKILL to PID: %d\n", p.ID, p.cmd.Process.Pid)
p.cmd.Process.Kill()
p.cmd.Wait()
case err := <-sigtermNormal:
if p.config.CmdStop != "" {
// for issue #35 to do things like `docker stop`
args, err := p.config.SanitizeCommandStop()
if err != nil {
if err.Error() != "wait: no child processes" {
// possible that simple-responder for testing is just not
// existing right, so suppress those errors.
fmt.Fprintf(p.logMonitor, "!!! process for %s stopped with error > %v\n", p.ID, err)
fmt.Fprintf(p.logMonitor, "!!! Error sanitizing stop command: %v\n", err)
// leave the state as it is?
return
}
fmt.Fprintf(p.logMonitor, "!!! Running stop command: %s\n", strings.Join(args, " "))
cmd := exec.Command(args[0], args[1:]...)
cmd.Stdout = p.logMonitor
cmd.Stderr = p.logMonitor
err = cmd.Start()
if err != nil {
fmt.Fprintf(p.logMonitor, "!!! Error running stop command: %v\n", err)
// leave the state as it is?
return
}
err = cmd.Wait()
if err != nil {
fmt.Fprintf(p.logMonitor, "!!! WARNING error waiting for stop command to complete: %v\n", err)
}
} else {
sigtermTimeout, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
sigtermNormal := make(chan error, 1)
go func() {
sigtermNormal <- p.cmd.Wait()
}()
p.cmd.Process.Signal(syscall.SIGTERM)
select {
case <-sigtermTimeout.Done():
fmt.Fprintf(p.logMonitor, "XXX Process for %s timed out waiting to stop, sending SIGKILL to PID: %d\n", p.ID, p.cmd.Process.Pid)
p.cmd.Process.Kill()
p.cmd.Wait()
case err := <-sigtermNormal:
if err != nil {
if err.Error() != "wait: no child processes" {
// possible that simple-responder for testing is just not
// existing right, so suppress those errors.
fmt.Fprintf(p.logMonitor, "!!! process for %s stopped with error > %v\n", p.ID, err)
}
}
}
}
p.state = StateStopped
}