Ilker ManapandClaude Opus 5.5
3387f97f6a
CI / test (push) Successful in 9m40s
slurm.conf sets SrunPortRange=60001-63000, so a site's firewall can let srun's step connections in. A node given a part in Slurm adds 6817-6818 and that range to llmgw's nftables table, open to the cluster's networks and closed to everyone else, as the engines' ports are. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
85 lines
2.7 KiB
Go
85 lines
2.7 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"log/slog"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/slurmnode"
|
|
)
|
|
|
|
// fakeShell answers the firewall package's scripts and records what was
|
|
// installed or removed.
|
|
type fakeShell struct {
|
|
present bool
|
|
installed string
|
|
deleted bool
|
|
}
|
|
|
|
func (f *fakeShell) Run(_ context.Context, script string, stdin []byte) ([]byte, error) {
|
|
switch {
|
|
case strings.HasPrefix(script, "ip -o route get 10.0.0.140"):
|
|
return []byte("10.0.0.140 dev eth0 src 10.0.0.141 uid 0\n2: eth0 inet 10.0.0.141/24 brd 10.0.0.255 scope global eth0\n"), nil
|
|
case strings.HasPrefix(script, "sudo -n /usr/sbin/nft list table"):
|
|
if f.present {
|
|
return nil, nil
|
|
}
|
|
return nil, errors.New("no table")
|
|
case strings.Contains(script, "nft delete table") && !strings.Contains(script, "cat > "):
|
|
f.deleted, f.present = true, false
|
|
return nil, nil
|
|
case strings.HasPrefix(script, "sudo cat "):
|
|
return nil, nil
|
|
default:
|
|
f.installed, f.present = string(stdin), true
|
|
return []byte("changed\n"), nil
|
|
}
|
|
}
|
|
|
|
func TestANodeKeepsOrDropsItsFirewallAsTheClusterSays(t *testing.T) {
|
|
sh := &fakeShell{}
|
|
a := &Agent{log: slog.New(slog.DiscardHandler), fwRunner: sh, serverList: []string{"https://10.0.0.140:7443"}}
|
|
|
|
a.applyFirewall(&llmgwv1.Firewall{Enabled: true, Networks: []string{"192.168.5.0/24", "10.9.9.9"},
|
|
EnginePortsFrom: 20000, EnginePortsTo: 20999})
|
|
for _, want := range []string{"10.0.0.0/24", "192.168.5.0/24", "10.9.9.9/32", "20000-20999", "9100"} {
|
|
if !strings.Contains(sh.installed, want) {
|
|
t.Errorf("the rules lack %s:\n%s", want, sh.installed)
|
|
}
|
|
}
|
|
|
|
a.applyFirewall(&llmgwv1.Firewall{Enabled: false})
|
|
if !sh.deleted {
|
|
t.Error("turned off, the table stayed")
|
|
}
|
|
}
|
|
|
|
func TestANodeGivenAPartInSlurmLetsSlurmInFromTheClusterOnly(t *testing.T) {
|
|
sh := &fakeShell{}
|
|
a := &Agent{log: slog.New(slog.DiscardHandler), fwRunner: sh, serverList: []string{"https://10.0.0.140:7443"},
|
|
slurm: slurmnode.New(nil, nil, "", slog.New(slog.DiscardHandler))}
|
|
a.applyFirewall(&llmgwv1.Firewall{Enabled: true, EnginePortsFrom: 20000, EnginePortsTo: 20999})
|
|
if strings.Contains(sh.installed, "6817") {
|
|
t.Fatalf("Slurm's ports before the node has a part in it:\n%s", sh.installed)
|
|
}
|
|
|
|
a.slurmMessage(&llmgwv1.Slurm{Name: "node05", Node: true, ConfServer: "10.0.0.148:6817"})
|
|
deadline := time.Now().Add(5 * time.Second)
|
|
for {
|
|
a.fwMu.Lock()
|
|
got := sh.installed
|
|
a.fwMu.Unlock()
|
|
if strings.Contains(got, "6817-6818") && strings.Contains(got, "60001-63000") {
|
|
break
|
|
}
|
|
if time.Now().After(deadline) {
|
|
t.Fatalf("given a part in Slurm, the rules still lack its ports:\n%s", got)
|
|
}
|
|
time.Sleep(10 * time.Millisecond)
|
|
}
|
|
}
|