From 95e45fb913ef76e7b1594c92e63c9b9cefa4a9d3 Mon Sep 17 00:00:00 2001 From: Sergey Mezentsev Date: Mon, 12 Feb 2024 13:15:42 +0300 Subject: [PATCH 1/7] Add submit stop --- runner/internal/runner/api/http_test.go | 2 + runner/internal/shim/api/http.go | 17 ++++++++ runner/internal/shim/api/schemas.go | 8 ++++ runner/internal/shim/api/server.go | 2 + runner/internal/shim/docker.go | 42 +++++++++++++++++-- src/dstack/_internal/core/models/runs.py | 2 +- .../background/tasks/process_finished_jobs.py | 33 ++++++++++++++- src/dstack/_internal/server/schemas/runner.py | 4 ++ .../server/services/runner/client.py | 6 +++ 9 files changed, 111 insertions(+), 5 deletions(-) diff --git a/runner/internal/runner/api/http_test.go b/runner/internal/runner/api/http_test.go index df5ef85751..7ce9b12e72 100644 --- a/runner/internal/runner/api/http_test.go +++ b/runner/internal/runner/api/http_test.go @@ -23,6 +23,8 @@ func (ds DummyRunner) Run(context.Context, shim.DockerImageConfig) error { return nil } +func (ds DummyRunner) SubmitStop(force bool) {} + func TestHealthcheck(t *testing.T) { request := httptest.NewRequest("GET", "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/api/healthcheck", nil) diff --git a/runner/internal/shim/api/http.go b/runner/internal/shim/api/http.go index 5508dfebff..95fc8d9661 100644 --- a/runner/internal/shim/api/http.go +++ b/runner/internal/shim/api/http.go @@ -51,3 +51,20 @@ func (s *ShimServer) PullGetHandler(w http.ResponseWriter, r *http.Request) (int State: string(s.runner.GetState()), }, nil } + +func (s *ShimServer) SubmitStopPostHandler(w http.ResponseWriter, r *http.Request) (interface{}, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + var body SubmitStopBody + if err := api.DecodeJSONBody(w, r, &body, true); err != nil { + log.Println("Failed to decode submit stop body", "err", err) + return nil, err + } + + s.runner.SubmitStop(body.Force) + + return &SubmitStopResponse{ + State: string(s.runner.GetState()), + }, nil +} diff --git a/runner/internal/shim/api/schemas.go b/runner/internal/shim/api/schemas.go index 8f86b8a21e..3b20f6600f 100644 --- a/runner/internal/shim/api/schemas.go +++ b/runner/internal/shim/api/schemas.go @@ -8,6 +8,10 @@ type DockerTaskBody struct { ImageName string `json:"image_name"` } +type SubmitStopBody struct { + Force bool `json:"force"` +} + type HealthcheckResponse struct { Service string `json:"service"` Version string `json:"version"` @@ -17,6 +21,10 @@ type PullResponse struct { State string `json:"state"` } +type SubmitStopResponse struct { + State string `json:"state"` +} + func (ra DockerTaskBody) TaskParams() shim.DockerImageConfig { res := shim.DockerImageConfig{ ImageName: ra.ImageName, diff --git a/runner/internal/shim/api/server.go b/runner/internal/shim/api/server.go index ce1bcfd591..ff19155959 100644 --- a/runner/internal/shim/api/server.go +++ b/runner/internal/shim/api/server.go @@ -12,6 +12,7 @@ import ( type TaskRunner interface { Run(context.Context, shim.DockerImageConfig) error GetState() shim.RunnerStatus + SubmitStop(bool) } type ShimServer struct { @@ -38,5 +39,6 @@ func NewShimServer(address string, runner TaskRunner, version string) *ShimServe mux.HandleFunc("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/api/submit", api.JSONResponseHandler("POST", s.SubmitPostHandler)) mux.HandleFunc("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/api/healthcheck", api.JSONResponseHandler("GET", s.HealthcheckGetHandler)) mux.HandleFunc("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/api/pull", api.JSONResponseHandler("GET", s.PullGetHandler)) + mux.HandleFunc("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/api/submit_stop", api.JSONResponseHandler("POST", s.SubmitStopPostHandler)) return s } diff --git a/runner/internal/shim/docker.go b/runner/internal/shim/docker.go index b70170cb3e..9bfc9228dd 100644 --- a/runner/internal/shim/docker.go +++ b/runner/internal/shim/docker.go @@ -23,9 +23,10 @@ import ( ) type DockerRunner struct { - client *docker.Client - dockerParams DockerParameters - state RunnerStatus + client *docker.Client + dockerParams DockerParameters + currentContainer string + state RunnerStatus } func NewDockerRunner(dockerParams DockerParameters) (*DockerRunner, error) { @@ -62,6 +63,8 @@ func (d *DockerRunner) Run(ctx context.Context, cfg DockerImageConfig) error { return err } + d.currentContainer = containerID + if !d.dockerParams.DockerKeepContainer() { defer func() { log.Println("Deleting container") @@ -82,10 +85,43 @@ func (d *DockerRunner) Run(ctx context.Context, cfg DockerImageConfig) error { log.Printf("Container finished successfully, id=%s\n", containerID) + d.currentContainer = "" d.state = Pending return nil } +func (d *DockerRunner) SubmitStop(force bool) { + if d.currentContainer == "" { + return + } + + ctx := context.Background() + + stopOptions := container.StopOptions{} + if force { + timeout := int(0) + stopOptions.Timeout = &timeout + } + + err := d.client.ContainerStop(ctx, d.currentContainer, stopOptions) + if err != nil { + log.Printf("Failed to stop container: %s", err) + } + + if !d.dockerParams.DockerKeepContainer() { + removeOptions := types.ContainerRemoveOptions{ + RemoveVolumes: true, + Force: true, + } + + err = d.client.ContainerRemove(ctx, d.currentContainer, removeOptions) + if err != nil { + log.Printf("Failed to remove container: %s", err) + } + } + +} + func (d DockerRunner) GetState() RunnerStatus { return d.state } diff --git a/src/dstack/_internal/core/models/runs.py b/src/dstack/_internal/core/models/runs.py index d14ebe9c4d..ac218b509f 100644 --- a/src/dstack/_internal/core/models/runs.py +++ b/src/dstack/_internal/core/models/runs.py @@ -120,7 +120,7 @@ class JobSpec(BaseModel): requirements: Requirements retry_policy: RetryPolicy working_dir: str - pool_name: Optional[str] + pool_name: Optional[str] # TODO: remove pool_name from JobSpec class JobProvisioningData(BaseModel): diff --git a/src/dstack/_internal/server/background/tasks/process_finished_jobs.py b/src/dstack/_internal/server/background/tasks/process_finished_jobs.py index 18cea5460e..a72f1e4643 100644 --- a/src/dstack/_internal/server/background/tasks/process_finished_jobs.py +++ b/src/dstack/_internal/server/background/tasks/process_finished_jobs.py @@ -1,7 +1,15 @@ +from typing import Dict + +from pydantic import parse_raw_as from sqlalchemy import or_, select from sqlalchemy.orm import joinedload -from dstack._internal.core.models.runs import InstanceStatus, JobSpec, JobStatus +from dstack._internal.core.models.runs import ( + InstanceStatus, + JobProvisioningData, + JobSpec, + JobStatus, +) from dstack._internal.server.db import get_session_ctx from dstack._internal.server.models import GatewayModel, JobModel from dstack._internal.server.services.gateways import gateway_connections_pool @@ -10,6 +18,8 @@ TERMINATING_PROCESSING_JOBS_LOCK, ) from dstack._internal.server.services.logging import job_log +from dstack._internal.server.services.runner import client +from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel from dstack._internal.server.utils.common import run_async from dstack._internal.utils.common import get_current_datetime from dstack._internal.utils.logging import get_logger @@ -90,7 +100,28 @@ async def _process_job(job_id): job_model.instance.last_job_processed_at = get_current_datetime() job_model.instance = None + server_ssh_private_key = job_model.project.ssh_private_key + job_provisioning_data = parse_raw_as(JobProvisioningData, job_model.job_provisioning_data) + await run_async( + submit_stop, + server_ssh_private_key, + job_provisioning_data, + job_model, + ) + job_model.removed = True job_model.last_processed_at = get_current_datetime() await session.commit() logger.info(*job_log("marked as removed", job_model)) + + +@runner_ssh_tunnel(ports=[client.REMOTE_SHIM_PORT, client.REMOTE_RUNNER_PORT]) +def submit_stop(job_model: JobModel, ports: Dict[int, int]): + shim_client = client.ShimClient(port=ports[client.REMOTE_SHIM_PORT]) + + resp = shim_client.healthcheck() + if resp is None: + logger.debug(*job_log("shim is not available yet", job_model)) + return False # shim is not available yet + + shim_client.submit_stop(force=False) diff --git a/src/dstack/_internal/server/schemas/runner.py b/src/dstack/_internal/server/schemas/runner.py index 97672ca406..84df329af1 100644 --- a/src/dstack/_internal/server/schemas/runner.py +++ b/src/dstack/_internal/server/schemas/runner.py @@ -70,3 +70,7 @@ class DockerImageBody(BaseModel): username: str password: str image_name: str + + +class SubmitStopBody(BaseModel): + force: bool = False diff --git a/src/dstack/_internal/server/services/runner/client.py b/src/dstack/_internal/server/services/runner/client.py index bd4045ae56..3d58c2ef61 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -13,6 +13,7 @@ HealthcheckResponse, PullResponse, SubmitBody, + SubmitStopBody, ) REMOTE_SHIM_PORT = 10998 @@ -109,6 +110,11 @@ def submit(self, username: str, password: str, image_name: str): ) resp.raise_for_status() + def submit_stop(self, force: bool = False): + body = SubmitStopBody(force=force) + resp = requests.post(self._url("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/api/submit_stop"), json=body.dict()) + resp.raise_for_status() + def pull(self): # TODO return resp = requests.get(self._url("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/api/pull")) resp.raise_for_status() From 8356226d15d9971f3f11e5146dcc49c52912204d Mon Sep 17 00:00:00 2001 From: Sergey Mezentsev Date: Wed, 14 Feb 2024 07:00:06 +0300 Subject: [PATCH 2/7] fixup! Add submit stop --- runner/internal/repo/diff_test.go | 1 - runner/internal/repo/manager.go | 2 ++ runner/internal/runner/api/http_test.go | 3 +-- runner/internal/runner/api/submit_test.go | 1 - runner/internal/shim/api/http.go | 4 ++-- runner/internal/shim/api/server.go | 4 ++-- runner/internal/shim/docker.go | 17 +++-------------- runner/internal/shim/runner.go | 1 - .../background/tasks/process_finished_jobs.py | 2 +- 9 files changed, 11 insertions(+), 24 deletions(-) diff --git a/runner/internal/repo/diff_test.go b/runner/internal/repo/diff_test.go index 2bf54eef66..a9de2a3c7d 100644 --- a/runner/internal/repo/diff_test.go +++ b/runner/internal/repo/diff_test.go @@ -10,7 +10,6 @@ import ( ) func TestApplyDiff(t *testing.T) { - cases := []struct { name string diff string diff --git a/runner/internal/repo/manager.go b/runner/internal/repo/manager.go index 3d2ad42d69..9a60ca8a1b 100644 --- a/runner/internal/repo/manager.go +++ b/runner/internal/repo/manager.go @@ -102,6 +102,7 @@ func (m *Manager) Checkout() error { return nil } + func (m *Manager) CheckoutBranch(branch string) error { log.Info(m.ctx, "git checkout", "auth", fmt.Sprintf("%T", (&m.clo).Auth)) ref, err := git.PlainClone(m.localPath, false, &m.clo) @@ -137,6 +138,7 @@ func (m *Manager) CheckoutMaster() error { return nil } + func (m *Manager) URL() string { return m.clo.URL } diff --git a/runner/internal/runner/api/http_test.go b/runner/internal/runner/api/http_test.go index 7ce9b12e72..0fc5b85bcf 100644 --- a/runner/internal/runner/api/http_test.go +++ b/runner/internal/runner/api/http_test.go @@ -23,10 +23,9 @@ func (ds DummyRunner) Run(context.Context, shim.DockerImageConfig) error { return nil } -func (ds DummyRunner) SubmitStop(force bool) {} +func (ds DummyRunner) Stop(force bool) {} func TestHealthcheck(t *testing.T) { - request := httptest.NewRequest("GET", "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/api/healthcheck", nil) responseRecorder := httptest.NewRecorder() diff --git a/runner/internal/runner/api/submit_test.go b/runner/internal/runner/api/submit_test.go index 9ff24efd06..b170c81be3 100644 --- a/runner/internal/runner/api/submit_test.go +++ b/runner/internal/runner/api/submit_test.go @@ -13,7 +13,6 @@ import ( ) func TestSubmit(t *testing.T) { - request := httptest.NewRequest("POST", "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/api/submit", strings.NewReader("{\"image_name\":\"ubuntu\"}")) responseRecorder := httptest.NewRecorder() diff --git a/runner/internal/shim/api/http.go b/runner/internal/shim/api/http.go index 95fc8d9661..7312512957 100644 --- a/runner/internal/shim/api/http.go +++ b/runner/internal/shim/api/http.go @@ -52,7 +52,7 @@ func (s *ShimServer) PullGetHandler(w http.ResponseWriter, r *http.Request) (int }, nil } -func (s *ShimServer) SubmitStopPostHandler(w http.ResponseWriter, r *http.Request) (interface{}, error) { +func (s *ShimServer) StopPostHandler(w http.ResponseWriter, r *http.Request) (interface{}, error) { s.mu.RLock() defer s.mu.RUnlock() @@ -62,7 +62,7 @@ func (s *ShimServer) SubmitStopPostHandler(w http.ResponseWriter, r *http.Reques return nil, err } - s.runner.SubmitStop(body.Force) + s.runner.Stop(body.Force) return &SubmitStopResponse{ State: string(s.runner.GetState()), diff --git a/runner/internal/shim/api/server.go b/runner/internal/shim/api/server.go index ff19155959..2e352faa7e 100644 --- a/runner/internal/shim/api/server.go +++ b/runner/internal/shim/api/server.go @@ -12,7 +12,7 @@ import ( type TaskRunner interface { Run(context.Context, shim.DockerImageConfig) error GetState() shim.RunnerStatus - SubmitStop(bool) + Stop(bool) } type ShimServer struct { @@ -39,6 +39,6 @@ func NewShimServer(address string, runner TaskRunner, version string) *ShimServe mux.HandleFunc("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/api/submit", api.JSONResponseHandler("POST", s.SubmitPostHandler)) mux.HandleFunc("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/api/healthcheck", api.JSONResponseHandler("GET", s.HealthcheckGetHandler)) mux.HandleFunc("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/api/pull", api.JSONResponseHandler("GET", s.PullGetHandler)) - mux.HandleFunc("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/api/submit_stop", api.JSONResponseHandler("POST", s.SubmitStopPostHandler)) + mux.HandleFunc("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/api/submit_stop", api.JSONResponseHandler("POST", s.StopPostHandler)) return s } diff --git a/runner/internal/shim/docker.go b/runner/internal/shim/docker.go index 9bfc9228dd..ae058bc576 100644 --- a/runner/internal/shim/docker.go +++ b/runner/internal/shim/docker.go @@ -90,7 +90,7 @@ func (d *DockerRunner) Run(ctx context.Context, cfg DockerImageConfig) error { return nil } -func (d *DockerRunner) SubmitStop(force bool) { +func (d *DockerRunner) Stop(force bool) { if d.currentContainer == "" { return } @@ -107,19 +107,6 @@ func (d *DockerRunner) SubmitStop(force bool) { if err != nil { log.Printf("Failed to stop container: %s", err) } - - if !d.dockerParams.DockerKeepContainer() { - removeOptions := types.ContainerRemoveOptions{ - RemoveVolumes: true, - Force: true, - } - - err = d.client.ContainerRemove(ctx, d.currentContainer, removeOptions) - if err != nil { - log.Printf("Failed to remove container: %s", err) - } - } - } func (d DockerRunner) GetState() RunnerStatus { @@ -202,12 +189,14 @@ func runContainer(ctx context.Context, client docker.APIClient, containerID stri if err := client.ContainerStart(ctx, containerID, types.ContainerStartOptions{}); err != nil { return tracerr.Wrap(err) } + waitCh, errorCh := client.ContainerWait(ctx, containerID, "") select { case <-waitCh: case err := <-errorCh: return tracerr.Wrap(err) } + return nil } diff --git a/runner/internal/shim/runner.go b/runner/internal/shim/runner.go index 1388c4b851..8ded5c6a7a 100644 --- a/runner/internal/shim/runner.go +++ b/runner/internal/shim/runner.go @@ -87,7 +87,6 @@ func downloadRunner(url string) (string, error) { } resp, err := http.DefaultClient.Do(req) - if err != nil { return "", gerrors.Wrap(err) } diff --git a/src/dstack/_internal/server/background/tasks/process_finished_jobs.py b/src/dstack/_internal/server/background/tasks/process_finished_jobs.py index a72f1e4643..0dfab60dc7 100644 --- a/src/dstack/_internal/server/background/tasks/process_finished_jobs.py +++ b/src/dstack/_internal/server/background/tasks/process_finished_jobs.py @@ -115,7 +115,7 @@ async def _process_job(job_id): logger.info(*job_log("marked as removed", job_model)) -@runner_ssh_tunnel(ports=[client.REMOTE_SHIM_PORT, client.REMOTE_RUNNER_PORT]) +@runner_ssh_tunnel(ports=[client.REMOTE_SHIM_PORT]) def submit_stop(job_model: JobModel, ports: Dict[int, int]): shim_client = client.ShimClient(port=ports[client.REMOTE_SHIM_PORT]) From a47660a8711a4fa1e0125af6b78a6faf1152ad7e Mon Sep 17 00:00:00 2001 From: Sergey Mezentsev Date: Wed, 14 Feb 2024 07:04:37 +0300 Subject: [PATCH 3/7] Enable golangci-lint presets: format and comment --- .pre-commit-config.yaml | 2 +- runner/consts/consts.go | 36 ++++++++++++++++++++++-------------- runner/internal/log/log.go | 2 +- 3 files changed, 24 insertions(+), 16 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ee577a901d..9a1ef5f096 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -10,5 +10,5 @@ repos: rev: v1.56.1 hooks: - id: golangci-lint-full - entry: bash -c 'cd runner && golangci-lint run -D depguard --presets import,module,unused "$@"' + entry: bash -c 'cd runner && golangci-lint run -D depguard,godox --presets import,module,unused,format,comment "$@"' stages: [manual] diff --git a/runner/consts/consts.go b/runner/consts/consts.go index 0fabfed538..456c610ac8 100644 --- a/runner/consts/consts.go +++ b/runner/consts/consts.go @@ -2,21 +2,25 @@ package consts import "time" -const DSTACK_DIR_PATH string = ".dstack" -const CONFIG_FILE_NAME string = "backend.yaml" -const RUNNER_FILE_NAME string = "runner.yaml" -const TMP_DIR_PATH = DSTACK_DIR_PATH + "/tmp" -const USER_ARTIFACTS_DIR = "artifacts" -const FUSE_DIR = "fuse" -const RUNS_DIR = "runs" +const ( + DSTACK_DIR_PATH string = ".dstack" + CONFIG_FILE_NAME string = "backend.yaml" + RUNNER_FILE_NAME string = "runner.yaml" + TMP_DIR_PATH = DSTACK_DIR_PATH + "/tmp" + USER_ARTIFACTS_DIR = "artifacts" + FUSE_DIR = "fuse" + RUNS_DIR = "runs" +) -const FILE_LOCK_FULL_DOWNLOAD = ".lock.full" -const SETUP_COMPLETED_FILE_NAME = ".setup_completed" +const ( + FILE_LOCK_FULL_DOWNLOAD = ".lock.full" + SETUP_COMPLETED_FILE_NAME = ".setup_completed" +) // ServerUrl A default build-time variable. The value is overridden via ldflags. var ServerUrl = "https://api.stgn.dstack.ai" -//var ServerUrl = "https://api.dstack.ai" +// var ServerUrl = "https://api.dstack.ai" // GPU constants const ( @@ -32,13 +36,17 @@ const ( EXPOSE_PORT_END = 4000 ) -const MAX_ATTEMPTS = 10 -const DELAY_TRY = 6 * time.Second +const ( + MAX_ATTEMPTS = 10 + DELAY_TRY = 6 * time.Second +) const DELAY_READ_STATUS = 5 * time.Second -const REPO_HTTPS_URL = "https://%s/%s/%s.git" -const REPO_GIT_URL = "git@%s:%s/%s.git" +const ( + REPO_HTTPS_URL = "https://%s/%s/%s.git" + REPO_GIT_URL = "git@%s:%s/%s.git" +) const ( TERMINATE_POLICY = "terminate" diff --git a/runner/internal/log/log.go b/runner/internal/log/log.go index 99478a8f9a..72d49e561e 100644 --- a/runner/internal/log/log.go +++ b/runner/internal/log/log.go @@ -100,7 +100,7 @@ func GetLogger(ctx context.Context) *logrus.Entry { } func CreateAppendFile(path string) (*os.File, error) { - f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644) + f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0o644) if err != nil { return nil, gerrors.Wrap(err) } From d824871dfca1aa62535cff293dea2b1f9c56dcc1 Mon Sep 17 00:00:00 2001 From: Sergey Mezentsev Date: Thu, 15 Feb 2024 11:11:05 +0300 Subject: [PATCH 4/7] fixup! Add submit stop --- runner/internal/shim/api/http.go | 26 ++++++++++++++++--- runner/internal/shim/api/server.go | 5 ++-- .../background/tasks/process_finished_jobs.py | 2 +- .../server/services/runner/client.py | 4 +-- 4 files changed, 29 insertions(+), 8 deletions(-) diff --git a/runner/internal/shim/api/http.go b/runner/internal/shim/api/http.go index 7312512957..e0e17f48b4 100644 --- a/runner/internal/shim/api/http.go +++ b/runner/internal/shim/api/http.go @@ -33,12 +33,18 @@ func (s *ShimServer) SubmitPostHandler(w http.ResponseWriter, r *http.Request) ( return nil, err } - go func(taskParams shim.DockerImageConfig) { - err := s.runner.Run(context.TODO(), taskParams) + ctx := context.Background() + ctx, cancel := context.WithCancel(ctx) + + s.runnerRunCancelFunc = &cancel + + go func(taskParams shim.DockerImageConfig, cancelFunc *context.CancelFunc) { + err := s.runner.Run(ctx, taskParams) if err != nil { fmt.Printf("failed Run %v", err) } - }(body.TaskParams()) + *cancelFunc = nil + }(body.TaskParams(), s.runnerRunCancelFunc) return nil, nil } @@ -56,6 +62,20 @@ func (s *ShimServer) StopPostHandler(w http.ResponseWriter, r *http.Request) (in s.mu.RLock() defer s.mu.RUnlock() + if s.runner.GetState() == shim.Pending { + return &SubmitStopResponse{ + State: string(s.runner.GetState()), + }, nil + } + + // if s.runner.GetState() == shim.Pulling { + // (*s.runnerRunCancelFunc)() + // // TODO: add barrier + // return &SubmitStopResponse{ + // State: string(s.runner.GetState()), + // }, nil + // } + var body SubmitStopBody if err := api.DecodeJSONBody(w, r, &body, true); err != nil { log.Println("Failed to decode submit stop body", "err", err) diff --git a/runner/internal/shim/api/server.go b/runner/internal/shim/api/server.go index 2e352faa7e..eab024a9b0 100644 --- a/runner/internal/shim/api/server.go +++ b/runner/internal/shim/api/server.go @@ -19,7 +19,8 @@ type ShimServer struct { HttpServer *http.Server mu sync.RWMutex - runner TaskRunner + runner TaskRunner + runnerRunCancelFunc *context.CancelFunc version string } @@ -39,6 +40,6 @@ func NewShimServer(address string, runner TaskRunner, version string) *ShimServe mux.HandleFunc("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/api/submit", api.JSONResponseHandler("POST", s.SubmitPostHandler)) mux.HandleFunc("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/api/healthcheck", api.JSONResponseHandler("GET", s.HealthcheckGetHandler)) mux.HandleFunc("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/api/pull", api.JSONResponseHandler("GET", s.PullGetHandler)) - mux.HandleFunc("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/api/submit_stop", api.JSONResponseHandler("POST", s.StopPostHandler)) + mux.HandleFunc("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/api/stop", api.JSONResponseHandler("POST", s.StopPostHandler)) return s } diff --git a/src/dstack/_internal/server/background/tasks/process_finished_jobs.py b/src/dstack/_internal/server/background/tasks/process_finished_jobs.py index 0dfab60dc7..5385acd263 100644 --- a/src/dstack/_internal/server/background/tasks/process_finished_jobs.py +++ b/src/dstack/_internal/server/background/tasks/process_finished_jobs.py @@ -124,4 +124,4 @@ def submit_stop(job_model: JobModel, ports: Dict[int, int]): logger.debug(*job_log("shim is not available yet", job_model)) return False # shim is not available yet - shim_client.submit_stop(force=False) + shim_client.stop(force=False) diff --git a/src/dstack/_internal/server/services/runner/client.py b/src/dstack/_internal/server/services/runner/client.py index 3d58c2ef61..f897198460 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -110,9 +110,9 @@ def submit(self, username: str, password: str, image_name: str): ) resp.raise_for_status() - def submit_stop(self, force: bool = False): + def stop(self, force: bool = False): body = SubmitStopBody(force=force) - resp = requests.post(self._url("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/api/submit_stop"), json=body.dict()) + resp = requests.post(self._url("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/api/stop"), json=body.dict()) resp.raise_for_status() def pull(self): # TODO return From 35081a96801a32a07c3cac3cbc3210f3e7e0b71f Mon Sep 17 00:00:00 2001 From: Sergey Mezentsev Date: Mon, 12 Feb 2024 13:15:42 +0300 Subject: [PATCH 5/7] Add submit stop --- runner/internal/repo/diff_test.go | 1 - runner/internal/repo/manager.go | 2 + runner/internal/runner/api/http_test.go | 3 +- runner/internal/runner/api/submit_test.go | 1 - runner/internal/shim/api/http.go | 25 +++++++++++- runner/internal/shim/api/schemas.go | 8 ++++ runner/internal/shim/api/server.go | 2 + runner/internal/shim/docker.go | 39 +++++++++++++++++-- runner/internal/shim/runner.go | 1 - src/dstack/_internal/core/models/runs.py | 2 +- .../background/tasks/process_finished_jobs.py | 33 +++++++++++++++- src/dstack/_internal/server/schemas/runner.py | 4 ++ .../server/services/runner/client.py | 6 +++ 13 files changed, 117 insertions(+), 10 deletions(-) diff --git a/runner/internal/repo/diff_test.go b/runner/internal/repo/diff_test.go index 2bf54eef66..a9de2a3c7d 100644 --- a/runner/internal/repo/diff_test.go +++ b/runner/internal/repo/diff_test.go @@ -10,7 +10,6 @@ import ( ) func TestApplyDiff(t *testing.T) { - cases := []struct { name string diff string diff --git a/runner/internal/repo/manager.go b/runner/internal/repo/manager.go index 3d2ad42d69..9a60ca8a1b 100644 --- a/runner/internal/repo/manager.go +++ b/runner/internal/repo/manager.go @@ -102,6 +102,7 @@ func (m *Manager) Checkout() error { return nil } + func (m *Manager) CheckoutBranch(branch string) error { log.Info(m.ctx, "git checkout", "auth", fmt.Sprintf("%T", (&m.clo).Auth)) ref, err := git.PlainClone(m.localPath, false, &m.clo) @@ -137,6 +138,7 @@ func (m *Manager) CheckoutMaster() error { return nil } + func (m *Manager) URL() string { return m.clo.URL } diff --git a/runner/internal/runner/api/http_test.go b/runner/internal/runner/api/http_test.go index df5ef85751..0fc5b85bcf 100644 --- a/runner/internal/runner/api/http_test.go +++ b/runner/internal/runner/api/http_test.go @@ -23,8 +23,9 @@ func (ds DummyRunner) Run(context.Context, shim.DockerImageConfig) error { return nil } -func TestHealthcheck(t *testing.T) { +func (ds DummyRunner) Stop(force bool) {} +func TestHealthcheck(t *testing.T) { request := httptest.NewRequest("GET", "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/api/healthcheck", nil) responseRecorder := httptest.NewRecorder() diff --git a/runner/internal/runner/api/submit_test.go b/runner/internal/runner/api/submit_test.go index 9ff24efd06..b170c81be3 100644 --- a/runner/internal/runner/api/submit_test.go +++ b/runner/internal/runner/api/submit_test.go @@ -13,7 +13,6 @@ import ( ) func TestSubmit(t *testing.T) { - request := httptest.NewRequest("POST", "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/api/submit", strings.NewReader("{\"image_name\":\"ubuntu\"}")) responseRecorder := httptest.NewRecorder() diff --git a/runner/internal/shim/api/http.go b/runner/internal/shim/api/http.go index 5508dfebff..ef259cd30a 100644 --- a/runner/internal/shim/api/http.go +++ b/runner/internal/shim/api/http.go @@ -34,7 +34,7 @@ func (s *ShimServer) SubmitPostHandler(w http.ResponseWriter, r *http.Request) ( } go func(taskParams shim.DockerImageConfig) { - err := s.runner.Run(context.TODO(), taskParams) + err := s.runner.Run(context.Background(), taskParams) if err != nil { fmt.Printf("failed Run %v", err) } @@ -51,3 +51,26 @@ func (s *ShimServer) PullGetHandler(w http.ResponseWriter, r *http.Request) (int State: string(s.runner.GetState()), }, nil } + +func (s *ShimServer) StopPostHandler(w http.ResponseWriter, r *http.Request) (interface{}, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + if s.runner.GetState() == shim.Pending { + return &StopResponse{ + State: string(s.runner.GetState()), + }, nil + } + + var body StopBody + if err := api.DecodeJSONBody(w, r, &body, true); err != nil { + log.Println("Failed to decode submit stop body", "err", err) + return nil, err + } + + s.runner.Stop(body.Force) + + return &StopResponse{ + State: string(s.runner.GetState()), + }, nil +} diff --git a/runner/internal/shim/api/schemas.go b/runner/internal/shim/api/schemas.go index 8f86b8a21e..65d3bd284d 100644 --- a/runner/internal/shim/api/schemas.go +++ b/runner/internal/shim/api/schemas.go @@ -8,6 +8,10 @@ type DockerTaskBody struct { ImageName string `json:"image_name"` } +type StopBody struct { + Force bool `json:"force"` +} + type HealthcheckResponse struct { Service string `json:"service"` Version string `json:"version"` @@ -17,6 +21,10 @@ type PullResponse struct { State string `json:"state"` } +type StopResponse struct { + State string `json:"state"` +} + func (ra DockerTaskBody) TaskParams() shim.DockerImageConfig { res := shim.DockerImageConfig{ ImageName: ra.ImageName, diff --git a/runner/internal/shim/api/server.go b/runner/internal/shim/api/server.go index ce1bcfd591..3b5393a97d 100644 --- a/runner/internal/shim/api/server.go +++ b/runner/internal/shim/api/server.go @@ -12,6 +12,7 @@ import ( type TaskRunner interface { Run(context.Context, shim.DockerImageConfig) error GetState() shim.RunnerStatus + Stop(bool) } type ShimServer struct { @@ -38,5 +39,6 @@ func NewShimServer(address string, runner TaskRunner, version string) *ShimServe mux.HandleFunc("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/api/submit", api.JSONResponseHandler("POST", s.SubmitPostHandler)) mux.HandleFunc("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/api/healthcheck", api.JSONResponseHandler("GET", s.HealthcheckGetHandler)) mux.HandleFunc("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/api/pull", api.JSONResponseHandler("GET", s.PullGetHandler)) + mux.HandleFunc("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/api/stop", api.JSONResponseHandler("POST", s.StopPostHandler)) return s } diff --git a/runner/internal/shim/docker.go b/runner/internal/shim/docker.go index b70170cb3e..45e26c9e38 100644 --- a/runner/internal/shim/docker.go +++ b/runner/internal/shim/docker.go @@ -23,9 +23,12 @@ import ( ) type DockerRunner struct { - client *docker.Client - dockerParams DockerParameters - state RunnerStatus + client *docker.Client + dockerParams DockerParameters + currentContainer string + state RunnerStatus + + cancelRun context.CancelFunc } func NewDockerRunner(dockerParams DockerParameters) (*DockerRunner, error) { @@ -45,6 +48,11 @@ func NewDockerRunner(dockerParams DockerParameters) (*DockerRunner, error) { func (d *DockerRunner) Run(ctx context.Context, cfg DockerImageConfig) error { var err error + ctx, cancel := context.WithTimeout(ctx, 10*time.Minute) + + d.cancelRun = cancel + defer cancel() + log.Println("Pulling image") d.state = Pulling if err = pullImage(ctx, d.client, cfg); err != nil { @@ -62,6 +70,8 @@ func (d *DockerRunner) Run(ctx context.Context, cfg DockerImageConfig) error { return err } + d.currentContainer = containerID + if !d.dockerParams.DockerKeepContainer() { defer func() { log.Println("Deleting container") @@ -82,10 +92,31 @@ func (d *DockerRunner) Run(ctx context.Context, cfg DockerImageConfig) error { log.Printf("Container finished successfully, id=%s\n", containerID) + d.currentContainer = "" d.state = Pending return nil } +func (d *DockerRunner) Stop(force bool) { + if d.currentContainer == "" && d.state == Pulling { + d.cancelRun() + return + } + + ctx := context.Background() + + stopOptions := container.StopOptions{} + if force { + timeout := int(0) + stopOptions.Timeout = &timeout + } + + err := d.client.ContainerStop(ctx, d.currentContainer, stopOptions) + if err != nil { + log.Printf("Failed to stop container: %s", err) + } +} + func (d DockerRunner) GetState() RunnerStatus { return d.state } @@ -166,12 +197,14 @@ func runContainer(ctx context.Context, client docker.APIClient, containerID stri if err := client.ContainerStart(ctx, containerID, types.ContainerStartOptions{}); err != nil { return tracerr.Wrap(err) } + waitCh, errorCh := client.ContainerWait(ctx, containerID, "") select { case <-waitCh: case err := <-errorCh: return tracerr.Wrap(err) } + return nil } diff --git a/runner/internal/shim/runner.go b/runner/internal/shim/runner.go index 1388c4b851..8ded5c6a7a 100644 --- a/runner/internal/shim/runner.go +++ b/runner/internal/shim/runner.go @@ -87,7 +87,6 @@ func downloadRunner(url string) (string, error) { } resp, err := http.DefaultClient.Do(req) - if err != nil { return "", gerrors.Wrap(err) } diff --git a/src/dstack/_internal/core/models/runs.py b/src/dstack/_internal/core/models/runs.py index d14ebe9c4d..ac218b509f 100644 --- a/src/dstack/_internal/core/models/runs.py +++ b/src/dstack/_internal/core/models/runs.py @@ -120,7 +120,7 @@ class JobSpec(BaseModel): requirements: Requirements retry_policy: RetryPolicy working_dir: str - pool_name: Optional[str] + pool_name: Optional[str] # TODO: remove pool_name from JobSpec class JobProvisioningData(BaseModel): diff --git a/src/dstack/_internal/server/background/tasks/process_finished_jobs.py b/src/dstack/_internal/server/background/tasks/process_finished_jobs.py index 18cea5460e..5385acd263 100644 --- a/src/dstack/_internal/server/background/tasks/process_finished_jobs.py +++ b/src/dstack/_internal/server/background/tasks/process_finished_jobs.py @@ -1,7 +1,15 @@ +from typing import Dict + +from pydantic import parse_raw_as from sqlalchemy import or_, select from sqlalchemy.orm import joinedload -from dstack._internal.core.models.runs import InstanceStatus, JobSpec, JobStatus +from dstack._internal.core.models.runs import ( + InstanceStatus, + JobProvisioningData, + JobSpec, + JobStatus, +) from dstack._internal.server.db import get_session_ctx from dstack._internal.server.models import GatewayModel, JobModel from dstack._internal.server.services.gateways import gateway_connections_pool @@ -10,6 +18,8 @@ TERMINATING_PROCESSING_JOBS_LOCK, ) from dstack._internal.server.services.logging import job_log +from dstack._internal.server.services.runner import client +from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel from dstack._internal.server.utils.common import run_async from dstack._internal.utils.common import get_current_datetime from dstack._internal.utils.logging import get_logger @@ -90,7 +100,28 @@ async def _process_job(job_id): job_model.instance.last_job_processed_at = get_current_datetime() job_model.instance = None + server_ssh_private_key = job_model.project.ssh_private_key + job_provisioning_data = parse_raw_as(JobProvisioningData, job_model.job_provisioning_data) + await run_async( + submit_stop, + server_ssh_private_key, + job_provisioning_data, + job_model, + ) + job_model.removed = True job_model.last_processed_at = get_current_datetime() await session.commit() logger.info(*job_log("marked as removed", job_model)) + + +@runner_ssh_tunnel(ports=[client.REMOTE_SHIM_PORT]) +def submit_stop(job_model: JobModel, ports: Dict[int, int]): + shim_client = client.ShimClient(port=ports[client.REMOTE_SHIM_PORT]) + + resp = shim_client.healthcheck() + if resp is None: + logger.debug(*job_log("shim is not available yet", job_model)) + return False # shim is not available yet + + shim_client.stop(force=False) diff --git a/src/dstack/_internal/server/schemas/runner.py b/src/dstack/_internal/server/schemas/runner.py index 97672ca406..84df329af1 100644 --- a/src/dstack/_internal/server/schemas/runner.py +++ b/src/dstack/_internal/server/schemas/runner.py @@ -70,3 +70,7 @@ class DockerImageBody(BaseModel): username: str password: str image_name: str + + +class SubmitStopBody(BaseModel): + force: bool = False diff --git a/src/dstack/_internal/server/services/runner/client.py b/src/dstack/_internal/server/services/runner/client.py index bd4045ae56..f897198460 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -13,6 +13,7 @@ HealthcheckResponse, PullResponse, SubmitBody, + SubmitStopBody, ) REMOTE_SHIM_PORT = 10998 @@ -109,6 +110,11 @@ def submit(self, username: str, password: str, image_name: str): ) resp.raise_for_status() + def stop(self, force: bool = False): + body = SubmitStopBody(force=force) + resp = requests.post(self._url("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/api/stop"), json=body.dict()) + resp.raise_for_status() + def pull(self): # TODO return resp = requests.get(self._url("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/api/pull")) resp.raise_for_status() From 2bc0b2e872f619c3f1c311843935c3ed685546d5 Mon Sep 17 00:00:00 2001 From: Sergey Mezentsev Date: Thu, 15 Feb 2024 16:00:03 +0300 Subject: [PATCH 6/7] Revert "Enable golangci-lint presets: format and comment" This reverts commit a47660a8711a4fa1e0125af6b78a6faf1152ad7e. --- .pre-commit-config.yaml | 2 +- runner/consts/consts.go | 36 ++++++++++++++---------------------- runner/internal/log/log.go | 2 +- 3 files changed, 16 insertions(+), 24 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9a1ef5f096..ee577a901d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -10,5 +10,5 @@ repos: rev: v1.56.1 hooks: - id: golangci-lint-full - entry: bash -c 'cd runner && golangci-lint run -D depguard,godox --presets import,module,unused,format,comment "$@"' + entry: bash -c 'cd runner && golangci-lint run -D depguard --presets import,module,unused "$@"' stages: [manual] diff --git a/runner/consts/consts.go b/runner/consts/consts.go index 456c610ac8..0fabfed538 100644 --- a/runner/consts/consts.go +++ b/runner/consts/consts.go @@ -2,25 +2,21 @@ package consts import "time" -const ( - DSTACK_DIR_PATH string = ".dstack" - CONFIG_FILE_NAME string = "backend.yaml" - RUNNER_FILE_NAME string = "runner.yaml" - TMP_DIR_PATH = DSTACK_DIR_PATH + "/tmp" - USER_ARTIFACTS_DIR = "artifacts" - FUSE_DIR = "fuse" - RUNS_DIR = "runs" -) +const DSTACK_DIR_PATH string = ".dstack" +const CONFIG_FILE_NAME string = "backend.yaml" +const RUNNER_FILE_NAME string = "runner.yaml" +const TMP_DIR_PATH = DSTACK_DIR_PATH + "/tmp" +const USER_ARTIFACTS_DIR = "artifacts" +const FUSE_DIR = "fuse" +const RUNS_DIR = "runs" -const ( - FILE_LOCK_FULL_DOWNLOAD = ".lock.full" - SETUP_COMPLETED_FILE_NAME = ".setup_completed" -) +const FILE_LOCK_FULL_DOWNLOAD = ".lock.full" +const SETUP_COMPLETED_FILE_NAME = ".setup_completed" // ServerUrl A default build-time variable. The value is overridden via ldflags. var ServerUrl = "https://api.stgn.dstack.ai" -// var ServerUrl = "https://api.dstack.ai" +//var ServerUrl = "https://api.dstack.ai" // GPU constants const ( @@ -36,17 +32,13 @@ const ( EXPOSE_PORT_END = 4000 ) -const ( - MAX_ATTEMPTS = 10 - DELAY_TRY = 6 * time.Second -) +const MAX_ATTEMPTS = 10 +const DELAY_TRY = 6 * time.Second const DELAY_READ_STATUS = 5 * time.Second -const ( - REPO_HTTPS_URL = "https://%s/%s/%s.git" - REPO_GIT_URL = "git@%s:%s/%s.git" -) +const REPO_HTTPS_URL = "https://%s/%s/%s.git" +const REPO_GIT_URL = "git@%s:%s/%s.git" const ( TERMINATE_POLICY = "terminate" diff --git a/runner/internal/log/log.go b/runner/internal/log/log.go index 72d49e561e..99478a8f9a 100644 --- a/runner/internal/log/log.go +++ b/runner/internal/log/log.go @@ -100,7 +100,7 @@ func GetLogger(ctx context.Context) *logrus.Entry { } func CreateAppendFile(path string) (*os.File, error) { - f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0o644) + f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644) if err != nil { return nil, gerrors.Wrap(err) } From 5b48ff064979146c54f4e77b5a08d9d62c14175a Mon Sep 17 00:00:00 2001 From: Sergey Mezentsev Date: Thu, 15 Feb 2024 16:03:27 +0300 Subject: [PATCH 7/7] SubmitStopBody -> SubmitBody --- src/dstack/_internal/server/schemas/runner.py | 2 +- src/dstack/_internal/server/services/runner/client.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/dstack/_internal/server/schemas/runner.py b/src/dstack/_internal/server/schemas/runner.py index 84df329af1..f63cd76261 100644 --- a/src/dstack/_internal/server/schemas/runner.py +++ b/src/dstack/_internal/server/schemas/runner.py @@ -72,5 +72,5 @@ class DockerImageBody(BaseModel): image_name: str -class SubmitStopBody(BaseModel): +class StopBody(BaseModel): force: bool = False diff --git a/src/dstack/_internal/server/services/runner/client.py b/src/dstack/_internal/server/services/runner/client.py index f897198460..e513601782 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -12,8 +12,8 @@ DockerImageBody, HealthcheckResponse, PullResponse, + StopBody, SubmitBody, - SubmitStopBody, ) REMOTE_SHIM_PORT = 10998 @@ -111,7 +111,7 @@ def submit(self, username: str, password: str, image_name: str): resp.raise_for_status() def stop(self, force: bool = False): - body = SubmitStopBody(force=force) + body = StopBody(force=force) resp = requests.post(self._url("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/api/stop"), json=body.dict()) resp.raise_for_status()