/
clusters.py
168 lines (142 loc) · 4.77 KB
/
clusters.py
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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
# -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function
from marshmallow import fields
from polyaxon.schemas.base import BaseConfig, BaseSchema
from polyaxon.schemas.fields import UUID
class NodeGPUSchema(BaseSchema):
index = fields.Int()
name = fields.Str()
uuid = UUID()
memory = fields.Int()
serial = fields.Str()
cluster_node = UUID()
@staticmethod
def schema_config():
return NodeGPUConfig
class NodeGPUConfig(BaseConfig):
"""
Node gpu config.
Args:
index: `int`. The index of the gpu during the discovery.
name: `str`. The name of gpu.
uuid: `UUID`. The uuid of gpu.
memory: `int`. The memory size of the gpu.
serial: `str`. The serial of the gpu.
cluster_node: `UUID`. the uuid of the cluster node.
"""
SCHEMA = NodeGPUSchema
IDENTIFIER = "NodeGPU"
DEFAULT_EXCLUDE_ATTRIBUTES = ["uuid", "cluster_node"]
def __init__(self, index, name, uuid, memory, serial, cluster_node):
self.uuid = uuid
self.serial = serial
self.name = name
self.index = index
self.memory = memory
self.cluster_node = cluster_node
class ClusterNodeSchema(BaseSchema):
sequence = fields.Int(allow_none=True)
name = fields.Str(allow_none=True)
uuid = UUID()
status = fields.Str(allow_none=True)
hostname = fields.Str(allow_none=True)
role = fields.Str(allow_none=True)
memory = fields.Int(allow_none=True)
cpu = fields.Float(allow_none=True)
n_gpus = fields.Int(allow_none=True)
kubelet_version = fields.Str(allow_none=True)
docker_version = fields.Str(allow_none=True)
os_image = fields.Str(allow_none=True)
kernel_version = fields.Str(allow_none=True)
schedulable_taints = fields.Bool(allow_none=True)
schedulable_state = fields.Bool(allow_none=True)
gpus = fields.Nested(NodeGPUSchema, many=True, allow_none=True)
@staticmethod
def schema_config():
return ClusterNodeConfig
class ClusterNodeConfig(BaseConfig):
"""
Node gpu config.
Args:
uuid: `UUID`. the uuid of the cluster node.
sequence: `int`. The sequence of the node in the cluster.
name: `str`. The name of node.
hostname: `str`. The node hostname.
role: `str`. The role of the node.
docker_version: `str`. The docker version used in the node.
kubelet_version: `str`. The kubelet version used in the node.
os_image: `str`. The os image used of the node.
kernel_version: `str`. The kernel version of the node.
schedulable_taints: `bool`. The schedulable taints of the node.
schedulable_state: `bool`. The schedulable state of the node.
memory: `int`. The memory size of the node.
cpu: `float`. The cpu of the node.
n_gpus: `int`. The number of gpus in the node.
status: `str`. The status of the node (ready or ...)
gpus: `list(NodeGPUConfig)`. The node gpus.
"""
SCHEMA = ClusterNodeSchema
IDENTIFIER = "ClusterNode"
DEFAULT_INCLUDE_ATTRIBUTES = [
"sequence",
"name",
"hostname",
"role",
"memory",
"cpu",
"n_gpus",
"status",
]
def __init__(
self,
uuid,
sequence=None,
name=None,
hostname=None,
role=None,
docker_version=None,
kubelet_version=None,
os_image=None,
kernel_version=None,
schedulable_taints=None,
schedulable_state=None,
memory=None,
cpu=None,
n_gpus=None,
status=None,
gpus=None,
):
self.uuid = uuid
self.sequence = sequence
self.name = name
self.hostname = hostname
self.role = role
self.docker_version = docker_version
self.kubelet_version = kubelet_version
self.os_image = os_image
self.kernel_version = kernel_version
self.schedulable_taints = schedulable_taints
self.schedulable_state = schedulable_state
self.memory = memory
self.cpu = cpu
self.n_gpus = n_gpus
self.status = status
self.gpus = gpus
class PolyaxonClusterSchema(BaseSchema):
version_api = fields.Dict()
nodes = fields.Nested(ClusterNodeSchema, many=True, allow_none=True)
@staticmethod
def schema_config():
return PolyaxonClusterConfig
class PolyaxonClusterConfig(BaseConfig):
"""
Polyaxon cluster config.
Args:
version_api: `dict`. The cluster's version api.
nodes: list(ClusterNodeConfig). The nodes in the cluster.
"""
SCHEMA = PolyaxonClusterSchema
IDENTIFIER = "PolyaxonCluster"
def __init__(self, version_api, nodes=None):
self.version_api = version_api
self.nodes = nodes