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..f63cd76261 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 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 bd4045ae56..e513601782 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -12,6 +12,7 @@ DockerImageBody, HealthcheckResponse, PullResponse, + StopBody, SubmitBody, ) @@ -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 = StopBody(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()