-
Notifications
You must be signed in to change notification settings - Fork 13
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge branch 'master' into dependabot/go_modules/github.com/golang-jw…
…t/jwt/v4-4.5.1
- Loading branch information
Showing
51 changed files
with
1,459 additions
and
73 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,11 @@ | ||
[[source]] | ||
url = "https://pypi.org/simple" | ||
verify_ssl = true | ||
name = "pypi" | ||
|
||
[packages] | ||
|
||
[dev-packages] | ||
|
||
[requires] | ||
python_version = "3.13" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
*.env | ||
gpu_names_cache.json |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,277 @@ | ||
package main | ||
|
||
import ( | ||
"context" | ||
"encoding/json" | ||
"flag" | ||
"fmt" | ||
"log" | ||
"log/slog" | ||
"net/http" | ||
"net/http/httputil" | ||
"net/url" | ||
"os" | ||
"sync" | ||
"time" | ||
|
||
"github.com/google/uuid" | ||
"github.com/joho/godotenv" | ||
"within.website/x/internal" | ||
"within.website/x/web/vastai/vastaicli" | ||
) | ||
|
||
var ( | ||
bind = flag.String("bind", ":3238", "HTTP port to bind to") | ||
diskSizeGB = flag.Int("vastai-disk-size-gb", 32, "amount of disk we need from vast.ai") | ||
dockerImage = flag.String("docker-image", "reg.xeiaso.net/runner/sdxl-tigris:latest", "docker image to start") | ||
onstartCmd = flag.String("onstart-cmd", "python -m cog.server.http", "onstart command to run in vast.ai") | ||
vastaiPort = flag.Int("vastai-port", 5000, "port that the guest will use in vast.ai") | ||
vastaiFilters = flag.String("vastai-filters", "verified=False cuda_max_good>=12.1 gpu_ram>=12 num_gpus=1 inet_down>=850", "vast.ai search filters") | ||
) | ||
|
||
func main() { | ||
internal.HandleStartup() | ||
ctx, cancel := context.WithCancel(context.Background()) | ||
defer cancel() | ||
|
||
if flag.NArg() != 1 { | ||
fmt.Println("usage: orodyagzou [flags] <whatever.env>") | ||
os.Exit(2) | ||
} | ||
|
||
fname := flag.Arg(0) | ||
slog.Debug("using env file", "fname", fname) | ||
|
||
env, err := godotenv.Read(fname) | ||
if err != nil { | ||
slog.Error("can't read env file", "fname", fname, "err", err) | ||
os.Exit(1) | ||
} | ||
|
||
var cfg vastaicli.InstanceConfig | ||
|
||
cfg.DiskSizeGB = *diskSizeGB | ||
cfg.Environment = env | ||
cfg.DockerImage = *dockerImage | ||
cfg.OnStartCMD = *onstartCmd | ||
cfg.Ports = append(cfg.Ports, *vastaiPort) | ||
|
||
images := &ScaleToZeroProxy{ | ||
cfg: cfg, | ||
} | ||
|
||
go images.slayLoop(ctx) | ||
|
||
mux := http.NewServeMux() | ||
mux.Handle("/v1/images", images) | ||
|
||
fmt.Printf("http://localhost%s\n", *bind) | ||
log.Fatal(http.ListenAndServe(*bind, mux)) | ||
} | ||
|
||
type ScaleToZeroProxy struct { | ||
cfg vastaicli.InstanceConfig | ||
|
||
// locked fields | ||
lock sync.RWMutex | ||
endpointURL string | ||
instanceID int | ||
ready bool | ||
lastUsed time.Time | ||
} | ||
|
||
func (s *ScaleToZeroProxy) slayLoop(ctx context.Context) { | ||
t := time.NewTicker(time.Minute) | ||
defer t.Stop() | ||
|
||
for { | ||
select { | ||
case <-ctx.Done(): | ||
slog.Error("context canceled", "err", ctx.Err()) | ||
return | ||
case <-t.C: | ||
s.lock.RLock() | ||
ready := s.ready | ||
lastUsed := s.lastUsed | ||
s.lock.RUnlock() | ||
|
||
if !ready { | ||
continue | ||
} | ||
|
||
if lastUsed.Add(5 * time.Minute).Before(time.Now()) { | ||
if err := s.slay(ctx); err != nil { | ||
slog.Error("can't slay instance", "err", err) | ||
} | ||
} | ||
} | ||
} | ||
} | ||
|
||
func (s *ScaleToZeroProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { | ||
s.lock.RLock() | ||
ready := s.ready | ||
s.lock.RUnlock() | ||
|
||
if !ready { | ||
if err := s.mint(r.Context()); err != nil { | ||
slog.Error("can't mint", "err", err) | ||
http.Error(w, "can't mint", http.StatusInternalServerError) | ||
return | ||
} | ||
} | ||
|
||
s.lock.RLock() | ||
endpointURL := s.endpointURL | ||
s.lock.RUnlock() | ||
|
||
u, err := url.Parse(endpointURL) | ||
if err != nil { | ||
slog.Error("can't url parse", "err", err, "url", s.endpointURL) | ||
http.Error(w, "can't url parse", http.StatusInternalServerError) | ||
return | ||
} | ||
|
||
next := httputil.NewSingleHostReverseProxy(u) | ||
od := next.Director | ||
next.Director = func(r *http.Request) { | ||
od(r) | ||
r.URL.Path = "/predictions/" + uuid.NewString() | ||
} | ||
next.ServeHTTP(w, r) | ||
|
||
s.lock.Lock() | ||
s.lastUsed = time.Now() | ||
s.lock.Unlock() | ||
} | ||
|
||
func (s *ScaleToZeroProxy) mint(ctx context.Context) error { | ||
s.lock.Lock() | ||
defer s.lock.Unlock() | ||
|
||
candidates, err := vastaicli.Search(ctx, *vastaiFilters, "dph+") | ||
if err != nil { | ||
return err | ||
} | ||
|
||
candidate := candidates[0] | ||
slog.Info("found instance", "costDPH", candidate.DphTotal, "gpuName", candidate.GpuName) | ||
|
||
instanceData, err := vastaicli.Mint(ctx, candidate.AskContractID, s.cfg) | ||
if err != nil { | ||
return err | ||
} | ||
|
||
slog.Info("created instance, waiting for things to settle", "id", instanceData.NewContract) | ||
|
||
instance, err := s.delayUntilRunning(ctx, instanceData.NewContract) | ||
if err != nil { | ||
return err | ||
} | ||
|
||
addr, ok := instance.AddrFor(s.cfg.Ports[0]) | ||
if !ok { | ||
return fmt.Errorf("somehow can't get port %d for instance %d, is god dead?", s.cfg.Ports[0], instance.ID) | ||
} | ||
|
||
s.endpointURL = "http://" + addr + "/" | ||
s.ready = true | ||
s.instanceID = instance.ID | ||
s.lastUsed = time.Now().Add(5 * time.Minute) | ||
|
||
if err := s.delayUntilReady(ctx, s.endpointURL); err != nil { | ||
return fmt.Errorf("can't do healthcheck: %w", err) | ||
} | ||
|
||
slog.Info("ready", "endpointURL", s.endpointURL, "instanceID", s.instanceID) | ||
|
||
return nil | ||
} | ||
|
||
func (s *ScaleToZeroProxy) slay(ctx context.Context) error { | ||
s.lock.Lock() | ||
defer s.lock.Unlock() | ||
|
||
if err := vastaicli.Slay(ctx, s.instanceID); err != nil { | ||
return err | ||
} | ||
|
||
s.endpointURL = "" | ||
s.ready = false | ||
s.lastUsed = time.Now() | ||
s.instanceID = 0 | ||
|
||
slog.Info("instance slayed", "docker_image", s.cfg.DockerImage) | ||
|
||
return nil | ||
} | ||
|
||
func (s *ScaleToZeroProxy) delayUntilReady(ctx context.Context, endpointURL string) error { | ||
type cogHealthCheck struct { | ||
Status string `json:"status"` | ||
} | ||
|
||
u, err := url.Parse(endpointURL) | ||
if err != nil { | ||
return fmt.Errorf("[unexpected] can't parse endpoint url %q: %w", endpointURL, err) | ||
} | ||
|
||
u.Path = "/health-check" | ||
|
||
t := time.NewTicker(time.Second) | ||
defer t.Stop() | ||
|
||
for { | ||
select { | ||
case <-ctx.Done(): | ||
return ctx.Err() | ||
case <-t.C: | ||
resp, err := http.Get(u.String()) | ||
if err != nil { | ||
return fmt.Errorf("can't fetch health check: %w", err) | ||
} | ||
|
||
var status cogHealthCheck | ||
if err := json.NewDecoder(resp.Body).Decode(&status); err != nil { | ||
return fmt.Errorf("can't parse health check response: %w", err) | ||
} | ||
|
||
if status.Status == "READY" { | ||
slog.Info("health check passed") | ||
return nil | ||
} | ||
} | ||
} | ||
} | ||
|
||
func (s *ScaleToZeroProxy) delayUntilRunning(ctx context.Context, instanceID int) (*vastaicli.Instance, error) { | ||
t := time.NewTicker(10 * time.Second) | ||
defer t.Stop() | ||
|
||
var instance *vastaicli.Instance | ||
var err error | ||
|
||
for { | ||
select { | ||
case <-ctx.Done(): | ||
return nil, ctx.Err() | ||
case <-t.C: | ||
instance, err = vastaicli.GetInstance(ctx, instanceID) | ||
if err != nil { | ||
return nil, err | ||
} | ||
|
||
slog.Debug("instance is cooking", "curr", instance.ActualStatus, "next", instance.NextState, "status", instance.StatusMsg) | ||
|
||
if instance.ActualStatus == "running" { | ||
_, ok := instance.AddrFor(s.cfg.Ports[0]) | ||
if !ok { | ||
slog.Info("no addr", "ports", s.cfg.Ports) | ||
continue | ||
} | ||
|
||
return instance, nil | ||
} | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
.terraform | ||
*.tfvars |
Oops, something went wrong.