Add submit stop

runner/internal/repo/diff_test.go

@@ -10,7 +10,6 @@ import (
 )
 
 func TestApplyDiff(t *testing.T) {
-
 	cases := []struct {
 		name    string
 		diff    string

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
 }

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", "/api/healthcheck", nil)
 	responseRecorder := httptest.NewRecorder()
 

runner/internal/runner/api/submit_test.go

@@ -13,7 +13,6 @@ import (
 )
 
 func TestSubmit(t *testing.T) {
-
 	request := httptest.NewRequest("POST", "/api/submit", strings.NewReader("{\"image_name\":\"ubuntu\"}"))
 	responseRecorder := httptest.NewRecorder()
 

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
+}

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,

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("/api/submit", api.JSONResponseHandler("POST", s.SubmitPostHandler))
 	mux.HandleFunc("/api/healthcheck", api.JSONResponseHandler("GET", s.HealthcheckGetHandler))
 	mux.HandleFunc("/api/pull", api.JSONResponseHandler("GET", s.PullGetHandler))
+	mux.HandleFunc("/api/stop", api.JSONResponseHandler("POST", s.StopPostHandler))
 	return s
 }

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
 }
 

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)
 	}

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):

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)

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

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("/api/stop"), json=body.dict())
+        resp.raise_for_status()
+
     def pull(self):  # TODO return
         resp = requests.get(self._url("/api/pull"))
         resp.raise_for_status()