From e81c7b95d7e1101943862655e38abcb103477a75 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Thu, 30 Jul 2026 11:41:44 +0200 Subject: [PATCH] Show host GPU driver on fleet instances Currently, the only way to know which GPU driver a fleet instance runs is to provision it and run a task there. Now shim detects the host GPU driver version on start (`nvidia-smi` for NVIDIA, `amd-smi` for AMD, `/sys/module/tenstorrent/version` for Tenstorrent) and reports it via the new `GET /api/instance/info` endpoint. The server stores it in `JobProvisioningData`, so that no migration is needed, and exposes it as `Instance.gpu_driver` and as a DRIVER column in `dstack fleet -v`. Notes: - Instance info reports facts observed by shim, as opposed to `/api/components`, which reports software managed by shim. - GPU hosts are asked on every instance check, as the facts change when shim restarts, e.g., after a driver upgrade, which the server does not necessarily observe. The provisioning data is only updated if the driver changed. This also fills the driver for instances provisioned before the server upgrade. - The driver stays unknown on hosts without GPUs, if detection fails, and on backends where shim does not run (runpod, vastai, kubernetes). Co-Authored-By: Claude Opus 5 (1M context) --- runner/docs/shim.openapi.yaml | 36 +++- runner/internal/shim/api/api_test.go | 6 + runner/internal/shim/api/handlers.go | 12 ++ runner/internal/shim/api/handlers_test.go | 46 +++++ runner/internal/shim/api/schemas.go | 7 + runner/internal/shim/api/server.go | 3 + runner/internal/shim/docker.go | 6 + runner/internal/shim/host/gpu.go | 156 +++++++++++--- runner/internal/shim/host/gpu_test.go | 117 +++++++++++ src/dstack/_internal/cli/utils/fleet.py | 2 + .../core/backends/kubernetes/compute.py | 4 + .../_internal/core/backends/vastai/compute.py | 1 + src/dstack/_internal/core/models/instances.py | 8 + src/dstack/_internal/core/models/runs.py | 7 + .../pipeline_tasks/instances/check.py | 52 ++++- .../pipeline_tasks/instances/common.py | 22 ++ .../_internal/server/schemas/instances.py | 2 + src/dstack/_internal/server/schemas/runner.py | 8 + .../_internal/server/services/instances.py | 1 + .../server/services/runner/client.py | 41 +++- .../test_instances/test_check.py | 195 +++++++++++++++++- .../_internal/server/routers/test_fleets.py | 4 + .../server/services/runner/test_client.py | 66 ++++++ .../server/services/test_instances.py | 5 +- 24 files changed, 771 insertions(+), 36 deletions(-) diff --git a/runner/docs/shim.openapi.yaml b/runner/docs/shim.openapi.yaml index e375e4e9d3..25f853c435 100644 --- a/runner/docs/shim.openapi.yaml +++ b/runner/docs/shim.openapi.yaml @@ -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: > @@ -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 @@ -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 diff --git a/runner/internal/shim/api/api_test.go b/runner/internal/shim/api/api_test.go index b6879187af..777c8e67a6 100644 --- a/runner/internal/shim/api/api_test.go +++ b/runner/internal/shim/api/api_test.go @@ -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 } @@ -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{}, diff --git a/runner/internal/shim/api/handlers.go b/runner/internal/shim/api/handlers.go index b3382d0f26..b40ed1c8e0 100644 --- a/runner/internal/shim/api/handlers.go +++ b/runner/internal/shim/api/handlers.go @@ -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 diff --git a/runner/internal/shim/api/handlers_test.go b/runner/internal/shim/api/handlers_test.go index bb19ebbf1b..f76ca8714a 100644 --- a/runner/internal/shim/api/handlers_test.go +++ b/runner/internal/shim/api/handlers_test.go @@ -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) { @@ -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 := `{ diff --git a/runner/internal/shim/api/schemas.go b/runner/internal/shim/api/schemas.go index 0e96028a5b..d035bdb176 100644 --- a/runner/internal/shim/api/schemas.go +++ b/runner/internal/shim/api/schemas.go @@ -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"` } diff --git a/runner/internal/shim/api/server.go b/runner/internal/shim/api/server.go index 9008aa2efe..db9ae7a7b1 100644 --- a/runner/internal/shim/api/server.go +++ b/runner/internal/shim/api/server.go @@ -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 { @@ -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 } @@ -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) diff --git a/runner/internal/shim/docker.go b/runner/internal/shim/docker.go index b149516a6c..7fbda1458d 100644 --- a/runner/internal/shim/docker.go +++ b/runner/internal/shim/docker.go @@ -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)) diff --git a/runner/internal/shim/host/gpu.go b/runner/internal/shim/host/gpu.go index eff57f2e00..a9e7f97058 100644 --- a/runner/internal/shim/host/gpu.go +++ b/runner/internal/shim/host/gpu.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "io" + "os" "path/filepath" "strconv" "strings" @@ -41,6 +42,10 @@ type GpuInfo struct { // AMD: empty string // Intel: accelerator index: ("0", "1", ...), as reported by `hl-smi -Q index` Index string + // Version of the installed host driver, e.g., "570.86.15" (NVIDIA), + // "6.10.5" (AMD amdgpu), "2.0.0" (Tenstorrent TT-KMD). + // Empty string if detection failed. All GPUs on a host share the same driver. + DriverVersion string } func GetGpuInfo(ctx context.Context) []GpuInfo { @@ -59,12 +64,23 @@ func GetGpuInfo(ctx context.Context) []GpuInfo { return []GpuInfo{} } +// normalizeDriverVersion filters out placeholder values SMI tools emit when a +// query field is not available, e.g., "N/A" or "[Not Supported]". +func normalizeDriverVersion(value string) string { + value = strings.TrimSpace(value) + switch strings.ToUpper(value) { + case "N/A", "[N/A]", "UNKNOWN", "[UNKNOWN]", "[NOT SUPPORTED]", "[NOT AVAILABLE]": + return "" + } + return value +} + func getNvidiaGpuInfo(ctx context.Context) []GpuInfo { gpus := []GpuInfo{} cmd := execute.ExecTask{ Command: "nvidia-smi", - Args: []string{"--query-gpu=name,memory.total,uuid", "--format=csv,noheader,nounits"}, + Args: []string{"--query-gpu=name,memory.total,uuid,driver_version", "--format=csv,noheader,nounits"}, StreamStdio: false, } res, err := cmd.Execute(ctx) @@ -90,8 +106,8 @@ func getNvidiaGpuInfo(ctx context.Context) []GpuInfo { log.Error(ctx, "cannot read csv", "err", err) return gpus } - if len(record) != 3 { - log.Error(ctx, "3 csv fields expected", "len", len(record)) + if len(record) != 4 { + log.Error(ctx, "4 csv fields expected", "len", len(record)) return gpus } vram, err := strconv.Atoi(strings.TrimSpace(record[1])) @@ -100,19 +116,43 @@ func getNvidiaGpuInfo(ctx context.Context) []GpuInfo { vram = 0 } gpus = append(gpus, GpuInfo{ - Vendor: gpu.GpuVendorNvidia, - Name: strings.TrimSpace(record[0]), - Vram: vram, - ID: strings.TrimSpace(record[2]), + Vendor: gpu.GpuVendorNvidia, + Name: strings.TrimSpace(record[0]), + Vram: vram, + ID: strings.TrimSpace(record[2]), + DriverVersion: normalizeDriverVersion(record[3]), }) } return gpus } type amdGpu struct { - Asic amdAsic `json:"asic"` - Vram amdVram `json:"vram"` - Bus amdBus `json:"bus"` + Asic amdAsic `json:"asic"` + Vram amdVram `json:"vram"` + Bus amdBus `json:"bus"` + Driver amdDriver `json:"driver"` +} + +// amdDriver is the `driver` section of `amd-smi static --driver`. +// Key names and value shapes differ between amd-smi versions, so it is parsed +// defensively: an unexpected format leaves Version empty instead of failing +// the whole GPU detection. Key matching is case-insensitive (encoding/json). +type amdDriver struct { + Version string +} + +func (d *amdDriver) UnmarshalJSON(data []byte) error { + var section struct { + Version string `json:"version"` + DriverVersion string `json:"driver_version"` + } + // The error is ignored deliberately: an unexpected shape leaves Version empty. + _ = json.Unmarshal(data, §ion) + d.Version = normalizeDriverVersion(section.Version) + if d.Version == "" { + d.Version = normalizeDriverVersion(section.DriverVersion) + } + return nil } // amd-smi >= 7.x wraps the array in {"gpu_data": [...]} @@ -151,38 +191,73 @@ func parseAmdSmiOutput(data []byte) ([]amdGpu, error) { return wrapped.GpuData, nil } -func getAmdGpuInfo(ctx context.Context) []GpuInfo { - gpus := []GpuInfo{} - +func execAmdSmiStatic(ctx context.Context, withDriver bool) (string, error) { ctx, cancel := context.WithTimeout(ctx, 2*time.Minute) defer cancel() + args := []string{ + "run", + "--rm", + "--device", "/dev/kfd", + "--device", "/dev/dri", + amdSmiImage, + "static", "--json", "--asic", "--vram", "--bus", + } + if withDriver { + args = append(args, "--driver") + } cmd := execute.ExecTask{ - Command: "docker", - Args: []string{ - "run", - "--rm", - "--device", "/dev/kfd", - "--device", "/dev/dri", - amdSmiImage, - "static", "--json", "--asic", "--vram", "--bus", - }, + Command: "docker", + Args: args, StreamStdio: false, } res, err := cmd.Execute(ctx) if err != nil { - log.Error(ctx, "failed to execute amd-smi", "err", err) - return gpus + return "", err } if res.ExitCode != 0 { - log.Error( - ctx, "failed to execute amd-smi", - "exitcode", res.ExitCode, "stdout", res.Stdout, "stderr", res.Stderr, - ) + return "", &amdSmiExitError{ + ExitCode: res.ExitCode, + Stdout: res.Stdout, + Stderr: res.Stderr, + } + } + return res.Stdout, nil +} + +// amdSmiExitError means amd-smi ran and exited with a non-zero code, e.g., because +// the installed version does not support one of the requested options. Failures to +// run amd-smi at all, including timeouts, are reported as other error types. +type amdSmiExitError struct { + ExitCode int + Stdout string + Stderr string +} + +func (e *amdSmiExitError) Error() string { + return fmt.Sprintf( + "exitcode: %d, stdout: %s, stderr: %s", e.ExitCode, e.Stdout, e.Stderr, + ) +} + +func getAmdGpuInfo(ctx context.Context) []GpuInfo { + gpus := []GpuInfo{} + + stdout, err := execAmdSmiStatic(ctx, true) + // Only an exited-with-error amd-smi may not support --driver. Retrying after + // a timeout or a docker failure would only double the wait, delaying shim startup. + var exitErr *amdSmiExitError + if err != nil && errors.As(err, &exitErr) { + // Fall back for amd-smi versions without the --driver option. + log.Error(ctx, "failed to execute amd-smi with --driver, retrying without", "err", err) + stdout, err = execAmdSmiStatic(ctx, false) + } + if err != nil { + log.Error(ctx, "failed to execute amd-smi", "err", err) return gpus } - amdGpus, err := parseAmdSmiOutput([]byte(res.Stdout)) + amdGpus, err := parseAmdSmiOutput([]byte(stdout)) if err != nil { log.Error(ctx, "cannot read json", "err", err) return gpus @@ -198,6 +273,7 @@ func getAmdGpuInfo(ctx context.Context) []GpuInfo { Name: amdGpu.Asic.Name, Vram: amdGpu.Vram.Size.Value, RenderNodePath: renderNodePath, + DriverVersion: amdGpu.Driver.Version, }) } return gpus @@ -413,6 +489,20 @@ func getGpusFromTtSmiSnapshot(snapshot *ttSmiSnapshot) []GpuInfo { return gpus } +// tenstorrentDriverVersionPath is the TT-KMD version file; it is what tt-smi +// itself reads to report the driver version. It is a variable so tests can +// override it. +var tenstorrentDriverVersionPath = "/sys/module/tenstorrent/version" + +func getTenstorrentDriverVersion(ctx context.Context) string { + data, err := os.ReadFile(tenstorrentDriverVersionPath) + if err != nil { + log.Error(ctx, "failed to read tenstorrent driver version", "err", err) + return "" + } + return normalizeDriverVersion(string(data)) +} + func getTenstorrentGpuInfo(ctx context.Context) []GpuInfo { gpus := []GpuInfo{} @@ -447,7 +537,13 @@ func getTenstorrentGpuInfo(ctx context.Context) []GpuInfo { return gpus } - return getGpusFromTtSmiSnapshot(ttSmiSnapshot) + gpus = getGpusFromTtSmiSnapshot(ttSmiSnapshot) + if driverVersion := getTenstorrentDriverVersion(ctx); driverVersion != "" { + for i := range gpus { + gpus[i].DriverVersion = driverVersion + } + } + return gpus } func getAmdRenderNodePath(bdf string) (string, error) { diff --git a/runner/internal/shim/host/gpu_test.go b/runner/internal/shim/host/gpu_test.go index 4110816003..f036a99fae 100644 --- a/runner/internal/shim/host/gpu_test.go +++ b/runner/internal/shim/host/gpu_test.go @@ -15,6 +15,123 @@ func loadTestData(filename string) ([]byte, error) { return os.ReadFile(path) } +func TestParseAmdSmiOutputWithDriver(t *testing.T) { + tests := []struct { + name string + data string + wantName string + wantDriver string + }{ + { + name: "rocm 6.x flat array with driver", + data: `[{"gpu": 0, "asic": {"market_name": "MI300X"}, "vram": {"size": {"value": 196592, "unit": "MB"}},` + + ` "bus": {"bdf": "0000:05:00.0"}, "driver": {"name": "amdgpu", "version": "6.10.5"}}]`, + wantName: "MI300X", + wantDriver: "6.10.5", + }, + { + name: "version preferred over driver_version when both present", + data: `[{"gpu": 0, "asic": {"market_name": "MI300X"}, "vram": {"size": {"value": 196592}},` + + ` "bus": {"bdf": "0000:05:00.0"}, "driver": {"version": "6.10.5", "driver_version": "6.8.5"}}]`, + wantName: "MI300X", + wantDriver: "6.10.5", + }, + { + name: "rocm 7.x wrapped with uppercase driver keys", + data: `{"gpu_data": [{"gpu": 0, "asic": {"market_name": "MI300X"}, "vram": {"size": {"value": 196592}},` + + ` "bus": {"bdf": "0000:05:00.0"}, "driver": {"NAME": "amdgpu", "VERSION": "6.12.12"}}]}`, + wantName: "MI300X", + wantDriver: "6.12.12", + }, + { + name: "driver_version key variant", + data: `[{"gpu": 0, "asic": {"market_name": "MI300X"}, "vram": {"size": {"value": 196592}},` + + ` "bus": {"bdf": "0000:05:00.0"}, "driver": {"driver_name": "amdgpu", "driver_version": "6.8.5"}}]`, + wantName: "MI300X", + wantDriver: "6.8.5", + }, + { + name: "no driver section", + data: `[{"gpu": 0, "asic": {"market_name": "MI300X"}, "vram": {"size": {"value": 196592}},` + + ` "bus": {"bdf": "0000:05:00.0"}}]`, + wantName: "MI300X", + wantDriver: "", + }, + { + name: "unexpected driver section shape does not fail parsing", + data: `[{"gpu": 0, "asic": {"market_name": "MI300X"}, "vram": {"size": {"value": 196592}},` + + ` "bus": {"bdf": "0000:05:00.0"}, "driver": "amdgpu 6.8.5"}]`, + wantName: "MI300X", + wantDriver: "", + }, + { + // amd-smi reports N/A if it fails to read the driver info + name: "placeholder driver version is treated as unknown", + data: `[{"gpu": 0, "asic": {"market_name": "MI300X"}, "vram": {"size": {"value": 196592}},` + + ` "bus": {"bdf": "0000:05:00.0"}, "driver": {"name": "N/A", "version": "N/A"}}]`, + wantName: "MI300X", + wantDriver: "", + }, + { + name: "driver_version is used if version is a placeholder", + data: `[{"gpu": 0, "asic": {"market_name": "MI300X"}, "vram": {"size": {"value": 196592}},` + + ` "bus": {"bdf": "0000:05:00.0"}, "driver": {"version": "N/A", "driver_version": "6.8.5"}}]`, + wantName: "MI300X", + wantDriver: "6.8.5", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + amdGpus, err := parseAmdSmiOutput([]byte(tt.data)) + if err != nil { + t.Fatalf("parseAmdSmiOutput() error = %v", err) + } + if len(amdGpus) != 1 { + t.Fatalf("parseAmdSmiOutput() returned %d GPUs, want 1", len(amdGpus)) + } + if amdGpus[0].Asic.Name != tt.wantName { + t.Errorf("name = %q, want %q", amdGpus[0].Asic.Name, tt.wantName) + } + if amdGpus[0].Driver.Version != tt.wantDriver { + t.Errorf("driver version = %q, want %q", amdGpus[0].Driver.Version, tt.wantDriver) + } + }) + } +} + +func TestNormalizeDriverVersion(t *testing.T) { + for input, want := range map[string]string{ + " 570.86.15 ": "570.86.15", + "N/A": "", + "[Not Supported]": "", + "Unknown": "", + } { + if got := normalizeDriverVersion(input); got != want { + t.Errorf("normalizeDriverVersion(%q) = %q, want %q", input, got, want) + } + } +} + +func TestGetTenstorrentDriverVersion(t *testing.T) { + versionFile := filepath.Join(t.TempDir(), "version") + if err := os.WriteFile(versionFile, []byte("2.0.0\n"), 0o644); err != nil { + t.Fatalf("failed to write version file: %v", err) + } + origPath := tenstorrentDriverVersionPath + tenstorrentDriverVersionPath = versionFile + defer func() { tenstorrentDriverVersionPath = origPath }() + + if got := getTenstorrentDriverVersion(t.Context()); got != "2.0.0" { + t.Errorf("getTenstorrentDriverVersion() = %q, want %q", got, "2.0.0") + } + + tenstorrentDriverVersionPath = filepath.Join(t.TempDir(), "nonexistent") + if got := getTenstorrentDriverVersion(t.Context()); got != "" { + t.Errorf("getTenstorrentDriverVersion() = %q, want empty string", got) + } +} + func TestUnmarshalTtSmiSnapshot(t *testing.T) { tests := []struct { name string diff --git a/src/dstack/_internal/cli/utils/fleet.py b/src/dstack/_internal/cli/utils/fleet.py index ccb2400857..875e2ef436 100644 --- a/src/dstack/_internal/cli/utils/fleet.py +++ b/src/dstack/_internal/cli/utils/fleet.py @@ -33,6 +33,7 @@ def get_fleets_table( table.add_column("NODES") if verbose: table.add_column("RESOURCES") + table.add_column("DRIVER") else: table.add_column("GPU") table.add_column("SPOT") @@ -123,6 +124,7 @@ def get_fleets_table( "RESOURCES": _format_instance_resources(instance), "GPU": _format_instance_gpu(instance), "BACKEND": backend_with_region, + "DRIVER": instance.gpu_driver.version if instance.gpu_driver else "-", "PRICE": instance_price, "SPOT": instance_spot, "STATUS": _format_instance_status(instance), diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index 7b3990714b..3340aeee5e 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -402,6 +402,10 @@ def update_provisioning_data( provisioning_data.hostname = get_or_error(service_spec.cluster_ip) pod_spec = get_or_error(pod.spec) node = api.read_node(name=get_or_error(pod_spec.node_name)) + # TODO: Set provisioning_data.gpu_driver from the node labels: + # nvidia.com/cuda.driver-version.full set by GPU Feature Discovery + # (or nvidia.com/cuda.driver.{major,minor,rev} set by older GFD versions), + # amd.com/gpu.driver-version set by the AMD GPU Operator. instance_offer = get_instance_offer_from_node(node=node, region=cluster.region) if instance_offer is not None: resource_requirements = get_or_error(pod_spec.containers[0].resources) diff --git a/src/dstack/_internal/core/backends/vastai/compute.py b/src/dstack/_internal/core/backends/vastai/compute.py index d51e40d297..bf85fc315a 100644 --- a/src/dstack/_internal/core/backends/vastai/compute.py +++ b/src/dstack/_internal/core/backends/vastai/compute.py @@ -200,6 +200,7 @@ def update_provisioning_data( provisioning_data.ssh_port = int( resp["ports"][f"{DSTACK_RUNNER_SSH_PORT}/tcp"][0]["HostPort"] ) + # TODO: Set provisioning_data.gpu_driver from resp["driver_version"] if ( resp["actual_status"] == "created" and ": OCI runtime create failed:" in resp["status_msg"] diff --git a/src/dstack/_internal/core/models/instances.py b/src/dstack/_internal/core/models/instances.py index a7623bd27a..cb7656f15c 100644 --- a/src/dstack/_internal/core/models/instances.py +++ b/src/dstack/_internal/core/models/instances.py @@ -46,6 +46,12 @@ def validate_name_and_vendor(cls, values): return values +class GpuDriverInfo(CoreModel): + vendor: Optional[gpuhunt.AcceleratorVendor] = None + """`vendor` is not set on hosts where shim could not detect it.""" + version: str + + class Disk(CoreModel): size_mib: int """`size_mib=0` has a special meaning -- size is unknown""" @@ -325,3 +331,5 @@ class Instance(CoreModel): price: Optional[float] = None total_blocks: Optional[int] = None busy_blocks: int = 0 + gpu_driver: Optional[GpuDriverInfo] = None + """`gpu_driver` is the accelerator driver installed on the host, when known.""" diff --git a/src/dstack/_internal/core/models/runs.py b/src/dstack/_internal/core/models/runs.py index d373fb23bc..611ba87118 100644 --- a/src/dstack/_internal/core/models/runs.py +++ b/src/dstack/_internal/core/models/runs.py @@ -30,6 +30,7 @@ ) from dstack._internal.core.models.files import FileArchiveMapping from dstack._internal.core.models.instances import ( + GpuDriverInfo, InstanceOfferWithAvailability, InstanceType, SSHConnectionParams, @@ -336,6 +337,12 @@ class JobProvisioningData(CoreModel): ssh_proxy: Optional[SSHConnectionParams] = None backend_data: Optional[str] = None """`backend_data` stores backend-specific data in JSON.""" + gpu_driver: Optional[GpuDriverInfo] = None + """`gpu_driver` is the accelerator driver installed on the host, when known. + Detected via the shim for VM-based backends and SSH fleets; taken from the + provider API or node labels for some container-based backends. May be set + after provisioning. + """ def get_base_backend(self) -> BackendType: if self.base_backend is not None: diff --git a/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py b/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py index 486c83dbf6..779d83435a 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py @@ -21,7 +21,11 @@ from dstack._internal.core.errors import ProvisioningError from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.health import HealthStatus -from dstack._internal.core.models.instances import InstanceStatus, InstanceTerminationReason +from dstack._internal.core.models.instances import ( + GpuDriverInfo, + InstanceStatus, + InstanceTerminationReason, +) from dstack._internal.core.models.profiles import TerminationPolicy from dstack._internal.core.models.runs import JobProvisioningData from dstack._internal.server import settings as server_settings @@ -32,6 +36,7 @@ can_terminate_fleet_instances_on_idle_duration, get_instance_idle_duration, get_provisioning_deadline, + set_gpu_driver_update, set_health_update, set_status_update, set_unreachable_update, @@ -145,6 +150,7 @@ async def check_instance(instance_model: InstanceModel) -> ProcessResult: instance_model=instance_model, job_provisioning_data=job_provisioning_data, check_instance_health=check_instance_health, + check_instance_info=_should_check_instance_info(job_provisioning_data), ) health_status = _get_health_status_for_instance_check( instance_model=instance_model, @@ -181,6 +187,11 @@ async def check_instance(instance_model: InstanceModel) -> ProcessResult: if instance_check.reachable: result.instance_update_map["termination_deadline"] = None + set_gpu_driver_update( + update_map=result.instance_update_map, + job_provisioning_data=job_provisioning_data, + gpu_driver=instance_check.gpu_driver, + ) if instance_model.status == InstanceStatus.PROVISIONING: set_status_update( update_map=result.instance_update_map, @@ -235,10 +246,22 @@ async def _should_check_instance_health(instance_id) -> bool: return res.scalar_one() == 0 +def _should_check_instance_info(job_provisioning_data: JobProvisioningData) -> bool: + """ + Instance info reports host facts that shim detects on start, e.g., the GPU driver + version, so they change after shim is restarted, which is required if the host GPUs + or their driver change. Such a restart is not necessarily observed by the server, + hence the facts are requested on every check, and only stored if they changed. + Hosts without GPUs report nothing, hence are never asked. + """ + return bool(job_provisioning_data.instance_type.resources.gpus) + + async def _run_instance_check( instance_model: InstanceModel, job_provisioning_data: JobProvisioningData, check_instance_health: bool, + check_instance_info: bool, ) -> InstanceCheck: ssh_private_keys = get_instance_ssh_private_keys(instance_model) instance_check = await run_async( @@ -248,6 +271,7 @@ async def _run_instance_check( None, instance=instance_model, check_instance_health=check_instance_health, + check_instance_info=check_instance_info, ) # May return False if fails to establish ssh connection. if instance_check is False: @@ -380,6 +404,7 @@ def _check_instance_inner( *, instance: InstanceModel, check_instance_health: bool = False, + check_instance_info: bool = False, ) -> InstanceCheck: instance_health_response: Optional[InstanceHealthResponse] = None shim_client = runner_client.ShimClient.from_address(addresses[DSTACK_SHIM_HTTP_PORT]) @@ -400,6 +425,8 @@ def _check_instance_inner( logger.exception(template, *args) return InstanceCheck(reachable=False, message=template % args) + gpu_driver = _get_gpu_driver(instance, shim_client) if check_instance_info else None + try: remove_dangling_tasks_from_instance(shim_client, instance) except Exception as exc: @@ -410,9 +437,32 @@ def _check_instance_inner( return runner_client.healthcheck_response_to_instance_check( healthcheck_response, instance_health_response, + gpu_driver, ) +def _get_gpu_driver( + instance_model: InstanceModel, + shim_client: runner_client.ShimClient, +) -> Optional[GpuDriverInfo]: + """ + Returns the host GPU driver reported by shim, or `None` if it cannot be retrieved. + The driver is optional metadata, so errors are not propagated to the instance check. + """ + try: + instance_info = shim_client.get_instance_info() + except requests.RequestException as exc: + logger.warning( + "Instance %s: shim.get_instance_info(): request error: %s", instance_model.name, exc + ) + return None + try: + return runner_client.instance_info_response_to_gpu_driver(instance_info) + except ValueError as exc: + logger.warning("Instance %s: unexpected instance info: %s", instance_model.name, exc) + return None + + def _maybe_install_components( instance_model: InstanceModel, shim_client: runner_client.ShimClient, diff --git a/src/dstack/_internal/server/background/pipeline_tasks/instances/common.py b/src/dstack/_internal/server/background/pipeline_tasks/instances/common.py index a386960478..7ae5ad852a 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/instances/common.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/instances/common.py @@ -11,6 +11,7 @@ from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.health import HealthStatus from dstack._internal.core.models.instances import ( + GpuDriverInfo, InstanceStatus, InstanceTerminationReason, SSHKey, @@ -176,3 +177,24 @@ def set_unreachable_update( return False update_map["unreachable"] = unreachable return True + + +def set_gpu_driver_update( + update_map: InstanceUpdateMap, + job_provisioning_data: JobProvisioningData, + gpu_driver: Optional[GpuDriverInfo], +) -> bool: + """ + Stores the shim-reported GPU driver in the instance provisioning data if it changed, + e.g., because the host driver was upgraded. Also fills it for instances provisioned + before the server upgrade. A GPU driver that could not be detected keeps the stored + one, as an undetectable driver is less informative than a possibly outdated one. + """ + if gpu_driver is None: + return False + current = job_provisioning_data.gpu_driver + if current is not None and current.dict() == gpu_driver.dict(): + return False + job_provisioning_data.gpu_driver = gpu_driver + update_map["job_provisioning_data"] = job_provisioning_data.json() + return True diff --git a/src/dstack/_internal/server/schemas/instances.py b/src/dstack/_internal/server/schemas/instances.py index 8f87935b92..8d0850983b 100644 --- a/src/dstack/_internal/server/schemas/instances.py +++ b/src/dstack/_internal/server/schemas/instances.py @@ -4,6 +4,7 @@ from dstack._internal.core.models.common import CoreModel from dstack._internal.core.models.health import HealthCheck, HealthStatus +from dstack._internal.core.models.instances import GpuDriverInfo from dstack._internal.server.schemas.runner import InstanceHealthResponse @@ -26,6 +27,7 @@ class InstanceCheck(CoreModel): reachable: bool message: Optional[str] = None health_response: Optional[InstanceHealthResponse] = None + gpu_driver: Optional[GpuDriverInfo] = None def get_health_status(self) -> HealthStatus: if self.health_response is None: diff --git a/src/dstack/_internal/server/schemas/runner.py b/src/dstack/_internal/server/schemas/runner.py index 3d366e6574..8363443171 100644 --- a/src/dstack/_internal/server/schemas/runner.py +++ b/src/dstack/_internal/server/schemas/runner.py @@ -131,6 +131,14 @@ class InstanceHealthResponse(CoreModel): dcgm: Optional[DCGMHealthResponse] = None +class InstanceInfoResponse(CoreModel): + gpu_vendor: Optional[str] = None + """`gpu_vendor` is not set on hosts without GPUs.""" + gpu_driver_version: Optional[str] = None + """`gpu_driver_version` is not set on hosts without GPUs + and when driver detection fails.""" + + class ShutdownRequest(CoreModel): force: bool diff --git a/src/dstack/_internal/server/services/instances.py b/src/dstack/_internal/server/services/instances.py index 913d3c9f44..71cd3aac2d 100644 --- a/src/dstack/_internal/server/services/instances.py +++ b/src/dstack/_internal/server/services/instances.py @@ -259,6 +259,7 @@ def instance_model_to_instance(instance_model: InstanceModel) -> Instance: instance.instance_type = jpd.instance_type instance.hostname = jpd.hostname instance.availability_zone = jpd.availability_zone + instance.gpu_driver = jpd.gpu_driver return instance diff --git a/src/dstack/_internal/server/services/runner/client.py b/src/dstack/_internal/server/services/runner/client.py index 93cf2f7238..1ba5e556ee 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -15,6 +15,7 @@ from dstack._internal.core.errors import DstackError from dstack._internal.core.models.common import CoreModel, NetworkMode from dstack._internal.core.models.envs import Env +from dstack._internal.core.models.instances import GpuDriverInfo from dstack._internal.core.models.repos.remote import RemoteRepoCreds from dstack._internal.core.models.resources import Memory from dstack._internal.core.models.runs import ClusterInfo, Job, Run @@ -29,6 +30,7 @@ GPUDevice, HealthcheckResponse, InstanceHealthResponse, + InstanceInfoResponse, JobInfoResponse, LegacyPullResponse, LegacyStopBody, @@ -308,6 +310,9 @@ class ShimClient: # `/api/instance/health` _INSTANCE_HEALTH_MIN_SHIM_VERSION = (0, 19, 22) + # `/api/instance/info` + _INSTANCE_INFO_MIN_SHIM_VERSION = (0, 20, 30) + # `/api/components` _COMPONENTS_MIN_SHIM_VERSION = (0, 20, 0) @@ -361,6 +366,14 @@ def is_instance_health_supported(self) -> bool: or self._shim_version_tuple >= self._INSTANCE_HEALTH_MIN_SHIM_VERSION ) + def is_instance_info_supported(self) -> bool: + if not self._negotiated: + self._negotiate() + return ( + self._shim_version_tuple is None + or self._shim_version_tuple >= self._INSTANCE_INFO_MIN_SHIM_VERSION + ) + def are_components_supported(self) -> bool: if not self._negotiated: self._negotiate() @@ -407,6 +420,18 @@ def get_instance_health(self) -> Optional[InstanceHealthResponse]: self._raise_for_status(resp) return self._response(InstanceHealthResponse, resp) + def get_instance_info(self) -> Optional[InstanceInfoResponse]: + if not self.is_instance_info_supported(): + logger.debug("instance info is not supported: %s", self._shim_version_string) + return None + resp = self._request("GET", "/api/instance/info") + if resp.status_code == HTTPStatus.NOT_FOUND: + # Old dev build of shim + logger.debug("instance info is not supported: %s", self._shim_version_string) + return None + self._raise_for_status(resp) + return self._response(InstanceInfoResponse, resp) + def shutdown(self, *, force: bool) -> bool: if not self.is_shutdown_supported(): logger.debug("shim shutdown is not supported: %s", self._shim_version_string) @@ -675,6 +700,7 @@ def _make_session_and_base_url( def healthcheck_response_to_instance_check( response: HealthcheckResponse, instance_health_response: Optional[InstanceHealthResponse] = None, + gpu_driver: Optional[GpuDriverInfo] = None, ) -> InstanceCheck: if response.service == "dstack-shim": message: Optional[str] = None @@ -685,7 +711,10 @@ def healthcheck_response_to_instance_check( ): message = instance_health_response.dcgm.incidents[0].error_message return InstanceCheck( - reachable=True, health_response=instance_health_response, message=message + reachable=True, + health_response=instance_health_response, + message=message, + gpu_driver=gpu_driver, ) return InstanceCheck( reachable=False, @@ -694,6 +723,16 @@ def healthcheck_response_to_instance_check( ) +def instance_info_response_to_gpu_driver( + response: Optional[InstanceInfoResponse], +) -> Optional[GpuDriverInfo]: + if response is None or not response.gpu_driver_version: + return None + return GpuDriverInfo.parse_obj( + {"vendor": response.gpu_vendor, "version": response.gpu_driver_version} + ) + + def _volume_to_shim_volume_info(volume: Volume, instance_id: str) -> ShimVolumeInfo: device_name = None attachment_data = volume.get_attachment_data_for_instance(instance_id) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_check.py b/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_check.py index 33e57df016..90bb351f96 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_check.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_check.py @@ -1,19 +1,30 @@ import datetime as dt import logging +from typing import Optional from unittest.mock import Mock import pytest import pytest_asyncio +import requests +from gpuhunt import AcceleratorVendor from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal.core.models.fleets import FleetNodesSpec from dstack._internal.core.models.health import HealthStatus -from dstack._internal.core.models.instances import InstanceStatus, InstanceTerminationReason +from dstack._internal.core.models.instances import ( + GpuDriverInfo, + InstanceStatus, + InstanceTerminationReason, +) from dstack._internal.core.models.profiles import TerminationPolicy -from dstack._internal.core.models.runs import JobStatus +from dstack._internal.core.models.runs import JobProvisioningData, JobStatus from dstack._internal.server.background.pipeline_tasks.instances import InstanceWorker from dstack._internal.server.background.pipeline_tasks.instances import check as instances_check +from dstack._internal.server.background.pipeline_tasks.instances.common import ( + InstanceUpdateMap, + set_gpu_driver_update, +) from dstack._internal.server.models import InstanceHealthCheckModel, InstanceModel from dstack._internal.server.schemas.health.dcgm import DCGMHealthResponse, DCGMHealthResult from dstack._internal.server.schemas.instances import InstanceCheck @@ -23,6 +34,7 @@ ComponentStatus, HealthcheckResponse, InstanceHealthResponse, + InstanceInfoResponse, TaskListResponse, ) from dstack._internal.server.services.runner.client import ComponentList, ShimClient @@ -36,6 +48,7 @@ create_user, get_fleet_configuration, get_fleet_spec, + get_job_provisioning_data, get_remote_connection_info, list_events, ) @@ -146,6 +159,75 @@ async def test_check_shim_transitions_provisioning_on_busy( assert instance.termination_deadline is None assert job.instance == instance + @pytest.mark.parametrize( + "stored_version", + [ + pytest.param(None, id="driver-not-yet-known"), + pytest.param("550.90.07", id="driver-upgraded-on-the-host"), + ], + ) + async def test_check_shim_stores_gpu_driver( + self, + test_db, + session: AsyncSession, + worker: InstanceWorker, + monkeypatch: pytest.MonkeyPatch, + stored_version: Optional[str], + ): + project = await create_project(session=session) + job_provisioning_data = get_job_provisioning_data(dockerized=True, gpu_count=1) + if stored_version is not None: + job_provisioning_data.gpu_driver = GpuDriverInfo( + vendor=AcceleratorVendor.NVIDIA, version=stored_version + ) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.IDLE, + job_provisioning_data=job_provisioning_data, + ) + await session.commit() + + check_instance_inner_mock = Mock( + return_value=InstanceCheck( + reachable=True, + gpu_driver=GpuDriverInfo(vendor=AcceleratorVendor.NVIDIA, version="570.86.15"), + ) + ) + monkeypatch.setattr(instances_check, "_check_instance_inner", check_instance_inner_mock) + await process_instance(session, worker, instance) + + await session.refresh(instance) + + assert check_instance_inner_mock.call_args.kwargs["check_instance_info"] + assert instance.job_provisioning_data is not None + jpd = JobProvisioningData.__response__.parse_raw(instance.job_provisioning_data) + assert jpd.gpu_driver is not None + assert jpd.gpu_driver.vendor == AcceleratorVendor.NVIDIA + assert jpd.gpu_driver.version == "570.86.15" + + async def test_check_shim_does_not_request_instance_info_without_gpus( + self, + test_db, + session: AsyncSession, + worker: InstanceWorker, + monkeypatch: pytest.MonkeyPatch, + ): + project = await create_project(session=session) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.IDLE, + job_provisioning_data=get_job_provisioning_data(dockerized=True, gpu_count=0), + ) + await session.commit() + + check_instance_inner_mock = Mock(return_value=InstanceCheck(reachable=True)) + monkeypatch.setattr(instances_check, "_check_instance_inner", check_instance_inner_mock) + await process_instance(session, worker, instance) + + assert not check_instance_inner_mock.call_args.kwargs["check_instance_info"] + async def test_check_shim_start_termination_deadline( self, test_db, @@ -539,6 +621,7 @@ def shim_client_mock( version=self.EXPECTED_VERSION, ) mock.get_instance_health.return_value = InstanceHealthResponse() + mock.get_instance_info.return_value = None mock.get_components.return_value = component_list mock.list_tasks.return_value = TaskListResponse(tasks=[]) mock.is_safe_to_restart.return_value = False @@ -942,3 +1025,111 @@ async def test_outdated_but_shim_installation_requested( shim_client_mock.get_components.assert_called_once() shim_client_mock.shutdown.assert_not_called() + + +class TestSetGpuDriverUpdate: + def test_noop_without_driver(self): + jpd = get_job_provisioning_data(dockerized=True) + update_map = InstanceUpdateMap() + assert not set_gpu_driver_update( + update_map=update_map, + job_provisioning_data=jpd, + gpu_driver=None, + ) + assert update_map == {} + + def test_noop_when_driver_unchanged(self): + jpd = get_job_provisioning_data(dockerized=True) + jpd.gpu_driver = GpuDriverInfo(vendor=AcceleratorVendor.NVIDIA, version="570.86.15") + update_map = InstanceUpdateMap() + assert not set_gpu_driver_update( + update_map=update_map, + job_provisioning_data=jpd, + gpu_driver=GpuDriverInfo(vendor=AcceleratorVendor.NVIDIA, version="570.86.15"), + ) + assert update_map == {} + + @pytest.mark.parametrize("current_version", [None, "550.90.07"]) + def test_sets_new_or_changed_driver(self, current_version): + jpd = get_job_provisioning_data(dockerized=True) + if current_version is not None: + jpd.gpu_driver = GpuDriverInfo( + vendor=AcceleratorVendor.NVIDIA, version=current_version + ) + update_map = InstanceUpdateMap() + assert set_gpu_driver_update( + update_map=update_map, + job_provisioning_data=jpd, + gpu_driver=GpuDriverInfo(vendor=AcceleratorVendor.NVIDIA, version="570.86.15"), + ) + assert "job_provisioning_data" in update_map + parsed = JobProvisioningData.__response__.parse_raw(update_map["job_provisioning_data"]) + assert parsed.gpu_driver is not None + assert parsed.gpu_driver.version == "570.86.15" + + +class TestShouldCheckInstanceInfo: + def _should_check(self, gpu_count: int, driver_known: bool) -> bool: + jpd = get_job_provisioning_data(dockerized=True, gpu_count=gpu_count) + if driver_known: + jpd.gpu_driver = GpuDriverInfo(vendor=AcceleratorVendor.NVIDIA, version="570.86.15") + return instances_check._should_check_instance_info(jpd) + + def test_no_gpus(self): + assert not self._should_check(gpu_count=0, driver_known=False) + + def test_gpus_without_known_driver(self): + assert self._should_check(gpu_count=1, driver_known=False) + + def test_gpus_with_known_driver(self): + # Shim restarts after a driver upgrade, which the server may not observe, + # so a known driver is re-requested to detect the change + assert self._should_check(gpu_count=1, driver_known=True) + + +class TestGetGpuDriver(BaseTestMaybeInstallComponents): + async def test_returns_reported_driver( + self, + test_db, + instance: InstanceModel, + shim_client_mock: Mock, + ): + shim_client_mock.get_instance_info.return_value = InstanceInfoResponse( + gpu_vendor="nvidia", gpu_driver_version="570.86.15" + ) + + gpu_driver = instances_check._get_gpu_driver(instance, shim_client_mock) + + assert gpu_driver == GpuDriverInfo(vendor=AcceleratorVendor.NVIDIA, version="570.86.15") + + async def test_returns_none_if_driver_not_detected( + self, + test_db, + instance: InstanceModel, + shim_client_mock: Mock, + ): + shim_client_mock.get_instance_info.return_value = InstanceInfoResponse() + + assert instances_check._get_gpu_driver(instance, shim_client_mock) is None + + async def test_returns_none_on_request_error( + self, + test_db, + instance: InstanceModel, + shim_client_mock: Mock, + ): + shim_client_mock.get_instance_info.side_effect = requests.RequestException("boom") + + assert instances_check._get_gpu_driver(instance, shim_client_mock) is None + + async def test_returns_none_on_unknown_gpu_vendor( + self, + test_db, + instance: InstanceModel, + shim_client_mock: Mock, + ): + shim_client_mock.get_instance_info.return_value = InstanceInfoResponse( + gpu_vendor="quantumx", gpu_driver_version="1.2.3" + ) + + assert instances_check._get_gpu_driver(instance, shim_client_mock) is None diff --git a/src/tests/_internal/server/routers/test_fleets.py b/src/tests/_internal/server/routers/test_fleets.py index 04d0145dfe..b7a48b54df 100644 --- a/src/tests/_internal/server/routers/test_fleets.py +++ b/src/tests/_internal/server/routers/test_fleets.py @@ -1006,6 +1006,7 @@ async def test_creates_fleet(self, test_db, session: AsyncSession, client: Async "price": None, "total_blocks": 1, "busy_blocks": 0, + "gpu_driver": None, } ], } @@ -1138,6 +1139,7 @@ async def test_creates_ssh_fleet(self, test_db, session: AsyncSession, client: A "price": 0.0, "total_blocks": 1, "busy_blocks": 0, + "gpu_driver": None, } ], } @@ -1358,6 +1360,7 @@ async def test_updates_ssh_fleet(self, test_db, session: AsyncSession, client: A "price": 0.0, "total_blocks": 1, "busy_blocks": 0, + "gpu_driver": None, }, { "id": SomeUUID4Str(), @@ -1393,6 +1396,7 @@ async def test_updates_ssh_fleet(self, test_db, session: AsyncSession, client: A "price": 0.0, "total_blocks": 1, "busy_blocks": 0, + "gpu_driver": None, }, ], } diff --git a/src/tests/_internal/server/services/runner/test_client.py b/src/tests/_internal/server/services/runner/test_client.py index e7db24226b..e776cd14b6 100644 --- a/src/tests/_internal/server/services/runner/test_client.py +++ b/src/tests/_internal/server/services/runner/test_client.py @@ -5,11 +5,13 @@ import pytest import requests_mock +from gpuhunt import AcceleratorVendor from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT, DSTACK_SHIM_HTTP_PORT from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import NetworkMode from dstack._internal.core.models.configurations import TaskConfiguration +from dstack._internal.core.models.instances import GpuDriverInfo from dstack._internal.core.models.resources import Memory from dstack._internal.core.models.runs import ClusterInfo, Job, JobSpec, JobSubmission, Run from dstack._internal.core.models.volumes import ( @@ -21,6 +23,7 @@ ) from dstack._internal.server.schemas.runner import ( HealthcheckResponse, + InstanceInfoResponse, JobResult, LegacyPullResponse, PortMapping, @@ -32,6 +35,8 @@ ShimClient, ShimHTTPError, _parse_version, + healthcheck_response_to_instance_check, + instance_info_response_to_gpu_driver, ) from dstack._internal.server.testing.common import ( get_run_spec, @@ -598,3 +603,64 @@ def test_valid_major_only(self, value: str): @pytest.mark.parametrize("value", ["", "foo", "1.12.3-next.20241231"]) def test_invalid(self, value: str): assert _parse_version(value) is None + + +class TestHealthcheckResponseToInstanceCheck: + def test_reachable(self): + response = HealthcheckResponse(service="dstack-shim", version="0.19.0") + check = healthcheck_response_to_instance_check(response) + assert check.reachable + assert check.gpu_driver is None + + def test_unexpected_service(self): + response = HealthcheckResponse(service="not-dstack-shim", version="0.19.0") + check = healthcheck_response_to_instance_check(response) + assert not check.reachable + assert check.gpu_driver is None + + +class TestInstanceInfoResponseToGpuDriver: + def test_none_response(self): + assert instance_info_response_to_gpu_driver(None) is None + + def test_no_gpus(self): + assert instance_info_response_to_gpu_driver(InstanceInfoResponse()) is None + + def test_vendor_without_version(self): + response = InstanceInfoResponse(gpu_vendor="nvidia") + assert instance_info_response_to_gpu_driver(response) is None + + def test_gpu_driver(self): + response = InstanceInfoResponse(gpu_vendor="nvidia", gpu_driver_version="570.86.15") + assert instance_info_response_to_gpu_driver(response) == GpuDriverInfo( + vendor=AcceleratorVendor.NVIDIA, version="570.86.15" + ) + + +class TestShimClientGetInstanceInfo(BaseShimClientTest): + @pytest.mark.shim_version("0.20.30") + def test_returns_instance_info(self, adapter: requests_mock.Adapter, client: ShimClient): + adapter.register_uri( + "GET", + "/api/instance/info", + json={"gpu_vendor": "nvidia", "gpu_driver_version": "570.86.15"}, + ) + resp = client.get_instance_info() + assert resp is not None + assert resp.gpu_vendor == "nvidia" + assert resp.gpu_driver_version == "570.86.15" + self.assert_request(adapter, 1, "GET", "/api/instance/info") + + @pytest.mark.shim_version("0.20.29") + def test_returns_none_if_not_supported( + self, adapter: requests_mock.Adapter, client: ShimClient + ): + assert client.get_instance_info() is None + assert len(adapter.request_history) == 1 # healthcheck only + + @pytest.mark.shim_version("latest") + def test_returns_none_if_not_found(self, adapter: requests_mock.Adapter, client: ShimClient): + adapter.register_uri("GET", "/api/instance/info", status_code=404, text="Not Found") + assert client.get_instance_info() is None + # An unknown shim version is assumed to support the endpoint, so it is requested + self.assert_request(adapter, 1, "GET", "/api/instance/info") diff --git a/src/tests/_internal/server/services/test_instances.py b/src/tests/_internal/server/services/test_instances.py index cba11c67ec..7a9606d5f9 100644 --- a/src/tests/_internal/server/services/test_instances.py +++ b/src/tests/_internal/server/services/test_instances.py @@ -2,12 +2,14 @@ from unittest.mock import Mock, call import pytest +from gpuhunt import AcceleratorVendor from sqlalchemy.ext.asyncio import AsyncSession import dstack._internal.server.services.instances as instances_services from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.health import HealthStatus from dstack._internal.core.models.instances import ( + GpuDriverInfo, Instance, InstanceStatus, InstanceTerminationReason, @@ -502,6 +504,7 @@ async def test_converts_instance(self, test_db, session: AsyncSession): price=1.0, total_blocks=1, busy_blocks=0, + gpu_driver=GpuDriverInfo(vendor=AcceleratorVendor.NVIDIA, version="570.86.15"), ) im = InstanceModel( id=instance_id, @@ -512,7 +515,7 @@ async def test_converts_instance(self, test_db, session: AsyncSession): unreachable=False, health=HealthStatus.WARNING, project=project, - job_provisioning_data='{"ssh_proxy":null, "backend":"aws","hostname":"hostname_test","region":"eu-west","price":1.0,"username":"user1","ssh_port":12345,"dockerized":false,"instance_id":"test_instance","instance_type": {"name": "instance", "resources": {"cpus": 1, "memory_mib": 512, "gpus": [], "spot": false, "disk": {"size_mib": 102400}, "description":""}}}', + job_provisioning_data='{"ssh_proxy":null, "backend":"aws","hostname":"hostname_test","region":"eu-west","price":1.0,"username":"user1","ssh_port":12345,"dockerized":false,"instance_id":"test_instance","gpu_driver":{"vendor":"nvidia","version":"570.86.15"},"instance_type": {"name": "instance", "resources": {"cpus": 1, "memory_mib": 512, "gpus": [], "spot": false, "disk": {"size_mib": 102400}, "description":""}}}', offer='{"price":1.0, "backend":"aws", "region":"eu-west-1", "availability":"available","instance": {"name": "instance", "resources": {"cpus": 1, "memory_mib": 512, "gpus": [], "spot": false, "disk": {"size_mib": 102400}, "description":""}}}', total_blocks=1, busy_blocks=0,