dstack/runner/internal/shim/models.go at 22bed96c96d1f65b447a9502dd1ca6fd8d58dbd8 · dstackai/dstack · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
package shim

import (
	"github.com/docker/docker/api/types/mount"
)

type DockerParameters interface {
	DockerPrivileged() bool
	DockerShellCommands(authorizedKeys []string, runnerHttpAddress string) []string
	DockerMounts(string) ([]mount.Mount, error)
	DockerPorts() []int
	RunnerHTTPPort() int
	MakeRunnerDir(name string) (string, error)
	DockerPJRTDevice() string
}

type CLIArgs struct {
	Shim struct {
		HTTPPort   int
		HomeDir    string
		BinaryPath string
		LogLevel   int
	}

	Runner struct {
		HTTPPort    int
		SSHPort     int
		DownloadURL string
		BinaryPath  string
		LogLevel    int
	}

	DCGMExporter struct {
		HTTPPort int
		Interval int // milliseconds
	}

	DCGM struct {
		Address string
	}

	Docker struct {
		Privileged bool
		PJRTDevice string
	}
}

type NetworkMode string

const (
	NetworkModeHost   = "host"
	NetworkModeBridge = "bridge"
)

type VolumeMountPoint struct {
	Name string `json:"name"`
	Path string `json:"path"`
}

type InstanceMountPoint struct {
	InstancePath string `json:"instance_path"`
	Path         string `json:"path"`
}

type VolumeInfo struct {
	Backend    string `json:"backend"`
	Name       string `json:"name"`
	VolumeId   string `json:"volume_id"`
	InitFs     bool   `json:"init_fs"`
	DeviceName string `json:"device_name"`
}

type PortMapping struct {
	Host      int `json:"host"`
	Container int `json:"container"`
}

type GPUDevice struct {
	PathOnHost      string `json:"path_on_host"`
	PathInContainer string `json:"path_in_container"`
}

type TaskConfig struct {
	ID               string               `json:"id"`
	Name             string               `json:"name"`
	RegistryUsername string               `json:"registry_username"`
	RegistryPassword string               `json:"registry_password"`
	ImageName        string               `json:"image_name"`
	ContainerUser    string               `json:"container_user"`
	Privileged       bool                 `json:"privileged"`
	GPU              int                  `json:"gpu"`      // -1 = all available, even if zero; 0 = zero, ...
	CPU              float64              `json:"cpu"`      // 0.0 = all available; 0.5 = a half of CPU, ...
	Memory           int64                `json:"memory"`   // bytes; 0 = all avaliable
	ShmSize          int64                `json:"shm_size"` // bytes; 0 = default (64MiB)
	NetworkMode      NetworkMode          `json:"network_mode"`
	Volumes          []VolumeInfo         `json:"volumes"`
	VolumeMounts     []VolumeMountPoint   `json:"volume_mounts"`
	InstanceMounts   []InstanceMountPoint `json:"instance_mounts"`
	// GPUDevices allows the server to set gpu devices instead of relying on the runner default logic.
	// E.g. passing nvidia devices directly instead of using nvidia-container-toolkit.
	GPUDevices       []GPUDevice `json:"gpu_devices"`
	HostSshUser      string      `json:"host_ssh_user"`
	HostSshKeys      []string    `json:"host_ssh_keys"`
	ContainerSshKeys []string    `json:"container_ssh_keys"`
}

type TaskListItem struct {
	ID     string     `json:"id"`
	Status TaskStatus `json:"status"`
}

type TaskInfo struct {
	ID                 string
	Status             TaskStatus
	TerminationReason  string
	TerminationMessage string
	Ports              []PortMapping
	ContainerName      string
	ContainerID        string
	GpuIDs             []string
}