Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 35 additions & 1 deletion runner/docs/shim.openapi.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ openapi: 3.1.2

info:
title: dstack-shim API
version: v2/0.20.1
version: v2/0.20.30
x-logo:
url: https://avatars.githubusercontent.com/u/54146142?s=260
description: >
Expand Down Expand Up @@ -86,6 +86,22 @@ paths:
schema:
$ref: "#/components/schemas/InstanceHealthResponse"

/instance/info:
get:
summary: Get instance info
description: >
(since [0.20.30](https://github.com/dstackai/dstack/releases/tag/0.20.30))
Returns facts about the host observed by shim, e.g., the GPU driver version.
Unlike `/components`, the reported entities are not managed by shim.
tags: [Instance]
responses:
"200":
description: ""
content:
application/json:
schema:
$ref: "#/components/schemas/InstanceInfoResponse"

/components:
get:
summary: Get components
Expand Down Expand Up @@ -504,6 +520,24 @@ components:
$ref: "#/components/schemas/DCGMHealth"
additionalProperties: false

InstanceInfoResponse:
title: shim.api.InstanceInfoResponse
type: object
properties:
gpu_vendor:
description: Host GPU vendor. Omitted on hosts without GPUs.
type: string
examples:
- nvidia
gpu_driver_version:
description: >
Host GPU driver version. Omitted on hosts without GPUs
or if detection failed.
type: string
examples:
- 570.86.15
additionalProperties: false

ComponentListResponse:
title: shim.api.ComponentListResponse
type: object
Expand Down
6 changes: 6 additions & 0 deletions runner/internal/shim/api/api_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,12 @@ import (
"sync"

"github.com/dstackai/dstack/runner/internal/shim"
"github.com/dstackai/dstack/runner/internal/shim/host"
)

type DummyRunner struct {
tasks map[string]bool
gpus []host.GpuInfo
mu sync.Mutex
}

Expand Down Expand Up @@ -46,6 +48,10 @@ func (ds *DummyRunner) Resources(context.Context) shim.Resources {
return shim.Resources{}
}

func (ds *DummyRunner) Gpus(context.Context) []host.GpuInfo {
return ds.gpus
}

func NewDummyRunner() *DummyRunner {
return &DummyRunner{
tasks: map[string]bool{},
Expand Down
12 changes: 12 additions & 0 deletions runner/internal/shim/api/handlers.go
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,18 @@ func (s *ShimServer) InstanceHealthHandler(w http.ResponseWriter, r *http.Reques
return &response, nil
}

func (s *ShimServer) InstanceInfoHandler(w http.ResponseWriter, r *http.Request) (interface{}, error) {
response := InstanceInfoResponse{}
// GPUs are detected once on startup, so this is not an expensive call.
// The driver is a host-wide property, hence any GPU can be used as the source.
if gpus := s.runner.Gpus(r.Context()); len(gpus) > 0 {
response.GpuVendor = string(gpus[0].Vendor)
response.GpuDriverVersion = gpus[0].DriverVersion
}

return &response, nil
}

func (s *ShimServer) TaskListHandler(w http.ResponseWriter, r *http.Request) (interface{}, error) {
tasks := s.runner.TaskList()
return &TaskListResponse{tasks}, nil
Expand Down
46 changes: 46 additions & 0 deletions runner/internal/shim/api/handlers_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@ import (
"testing"

commonapi "github.com/dstackai/dstack/runner/internal/common/api"
"github.com/dstackai/dstack/runner/internal/common/gpu"
"github.com/dstackai/dstack/runner/internal/shim/host"
)

func TestHealthcheck(t *testing.T) {
Expand All @@ -29,6 +31,50 @@ func TestHealthcheck(t *testing.T) {
}
}

// TestInstanceInfo goes through the router to also cover the endpoint registration
func TestInstanceInfo(t *testing.T) {
request := httptest.NewRequest("GET", "/api/instance/info", nil)
responseRecorder := httptest.NewRecorder()

runner := NewDummyRunner()
runner.gpus = []host.GpuInfo{
{Vendor: gpu.GpuVendorNvidia, Name: "T4", Vram: 16384, DriverVersion: "570.86.15"},
}
server := NewShimServer(context.Background(), ":12346", "0.0.1.dev2", runner, nil, nil, nil, nil)

server.httpServer.Handler.ServeHTTP(responseRecorder, request)

if responseRecorder.Code != 200 {
t.Errorf("Want status '%d', got '%d'", 200, responseRecorder.Code)
}

expected := `{"gpu_vendor":"nvidia","gpu_driver_version":"570.86.15"}`

if strings.TrimSpace(responseRecorder.Body.String()) != expected {
t.Errorf("Want '%s', got '%s'", expected, responseRecorder.Body.String())
}
}

func TestInstanceInfoWithoutGpus(t *testing.T) {
request := httptest.NewRequest("GET", "/api/instance/info", nil)
responseRecorder := httptest.NewRecorder()

server := NewShimServer(context.Background(), ":12347", "0.0.1.dev2", NewDummyRunner(), nil, nil, nil, nil)

f := commonapi.JSONResponseHandler(server.InstanceInfoHandler)
f(responseRecorder, request)

if responseRecorder.Code != 200 {
t.Errorf("Want status '%d', got '%d'", 200, responseRecorder.Code)
}

expected := "{}"

if strings.TrimSpace(responseRecorder.Body.String()) != expected {
t.Errorf("Want '%s', got '%s'", expected, responseRecorder.Body.String())
}
}

func TestTaskSubmit(t *testing.T) {
server := NewShimServer(context.Background(), ":12340", "0.0.1.dev2", NewDummyRunner(), nil, nil, nil, nil)
requestBody := `{
Expand Down
7 changes: 7 additions & 0 deletions runner/internal/shim/api/schemas.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,13 @@ type InstanceHealthResponse struct {
DCGM *dcgm.Health `json:"dcgm"`
}

// InstanceInfoResponse reports facts about the host observed by shim. Fields are
// omitted if the corresponding fact is not applicable or could not be detected.
type InstanceInfoResponse struct {
GpuVendor string `json:"gpu_vendor,omitempty"`
GpuDriverVersion string `json:"gpu_driver_version,omitempty"`
}

type TaskListResponse struct {
Tasks []*shim.TaskListItem `json:"tasks"`
}
Expand Down
3 changes: 3 additions & 0 deletions runner/internal/shim/api/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import (
"github.com/dstackai/dstack/runner/internal/shim"
"github.com/dstackai/dstack/runner/internal/shim/components"
"github.com/dstackai/dstack/runner/internal/shim/dcgm"
"github.com/dstackai/dstack/runner/internal/shim/host"
)

type TaskRunner interface {
Expand All @@ -22,6 +23,7 @@ type TaskRunner interface {
Remove(ctx context.Context, taskID string) error

Resources(context.Context) shim.Resources
Gpus(context.Context) []host.GpuInfo
TaskList() []*shim.TaskListItem
TaskInfo(taskID string) shim.TaskInfo
}
Expand Down Expand Up @@ -85,6 +87,7 @@ func NewShimServer(
r.AddHandler("GET", "/api/healthcheck", s.HealthcheckHandler)
r.AddHandler("POST", "/api/shutdown", s.ShutdownHandler)
r.AddHandler("GET", "/api/instance/health", s.InstanceHealthHandler)
r.AddHandler("GET", "/api/instance/info", s.InstanceInfoHandler)
r.AddHandler("GET", "/api/components", s.ComponentListHandler)
r.AddHandler("POST", "/api/components/install", s.ComponentInstallHandler)
r.AddHandler("GET", "/api/tasks", s.TaskListHandler)
Expand Down
6 changes: 6 additions & 0 deletions runner/internal/shim/docker.go
Original file line number Diff line number Diff line change
Expand Up @@ -324,6 +324,12 @@ func (d *DockerRunner) Resources(ctx context.Context) Resources {
}
}

// Gpus returns the GPUs detected at startup without collecting other host
// resources, making it suitable for frequently called paths.
func (d *DockerRunner) Gpus(ctx context.Context) []host.GpuInfo {
return d.gpus
}

func (d *DockerRunner) TaskList() []*TaskListItem {
tasks := d.tasks.List()
result := make([]*TaskListItem, 0, len(tasks))
Expand Down
Loading
Loading