20
20
from google .cloud .aiplatform .preview .vertex_ray .util .resources import Cluster
21
21
from google .cloud .aiplatform .preview .vertex_ray .util .resources import (
22
22
Resources ,
23
+ NodeImages ,
23
24
)
24
25
from google .cloud .aiplatform_v1beta1 .types .machine_resources import DiskSpec
25
26
from google .cloud .aiplatform_v1beta1 .types .machine_resources import (
@@ -82,6 +83,7 @@ class ClusterConstants:
82
83
)
83
84
_TEST_CPU_IMAGE = "us-docker.pkg.dev/vertex-ai/training/ray-cpu.2-4.py310:latest"
84
85
_TEST_GPU_IMAGE = "us-docker.pkg.dev/vertex-ai/training/ray-gpu.2-4.py310:latest"
86
+ _TEST_CUSTOM_IMAGE = "us-docker.pkg.dev/my-project/ray-custom.2-4.py310:latest"
85
87
# RUNNING Persistent Cluster w/o Ray
86
88
_TEST_RESPONSE_NO_RAY_RUNNING = PersistentResource (
87
89
name = _TEST_VERTEX_RAY_PR_ADDRESS ,
@@ -127,6 +129,13 @@ class ClusterConstants:
127
129
network = ProjectConstants ._TEST_VPC_NETWORK ,
128
130
labels = _TEST_LABELS ,
129
131
)
132
+ _TEST_REQUEST_RUNNING_1_POOL_CUSTOM_IMAGES = PersistentResource (
133
+ resource_pools = [_TEST_RESOURCE_POOL_0 ],
134
+ resource_runtime_spec = ResourceRuntimeSpec (
135
+ ray_spec = RaySpec (resource_pool_images = {"head-node" : _TEST_CUSTOM_IMAGE }),
136
+ ),
137
+ network = ProjectConstants ._TEST_VPC_NETWORK ,
138
+ )
130
139
# Get response has generated name, and URIs
131
140
_TEST_RESPONSE_RUNNING_1_POOL = PersistentResource (
132
141
name = _TEST_VERTEX_RAY_PR_ADDRESS ,
@@ -143,6 +152,22 @@ class ClusterConstants:
143
152
),
144
153
state = "RUNNING" ,
145
154
)
155
+ # Get response has generated name, and URIs
156
+ _TEST_RESPONSE_RUNNING_1_POOL_CUSTOM_IMAGES = PersistentResource (
157
+ name = _TEST_VERTEX_RAY_PR_ADDRESS ,
158
+ resource_pools = [_TEST_RESOURCE_POOL_0 ],
159
+ resource_runtime_spec = ResourceRuntimeSpec (
160
+ ray_spec = RaySpec (resource_pool_images = {"head-node" : _TEST_CUSTOM_IMAGE }),
161
+ ),
162
+ network = ProjectConstants ._TEST_VPC_NETWORK ,
163
+ resource_runtime = ResourceRuntime (
164
+ access_uris = {
165
+ "RAY_DASHBOARD_URI" : _TEST_VERTEX_RAY_DASHBOARD_ADDRESS ,
166
+ "RAY_HEAD_NODE_INTERNAL_IP" : _TEST_VERTEX_RAY_HEAD_NODE_IP ,
167
+ }
168
+ ),
169
+ state = "RUNNING" ,
170
+ )
146
171
# 2_POOL: worker_node_types and head_node_type have different MachineSpecs
147
172
_TEST_HEAD_NODE_TYPE_2_POOLS = Resources ()
148
173
_TEST_WORKER_NODE_TYPES_2_POOLS = [
@@ -213,6 +238,7 @@ class ClusterConstants:
213
238
head_node_type = _TEST_HEAD_NODE_TYPE_1_POOL ,
214
239
worker_node_types = _TEST_WORKER_NODE_TYPES_1_POOL ,
215
240
dashboard_address = _TEST_VERTEX_RAY_DASHBOARD_ADDRESS ,
241
+ node_images = NodeImages (head = _TEST_CPU_IMAGE , worker = _TEST_CPU_IMAGE ),
216
242
)
217
243
_TEST_CLUSTER_2 = Cluster (
218
244
cluster_resource_name = _TEST_VERTEX_RAY_PR_ADDRESS ,
@@ -223,6 +249,7 @@ class ClusterConstants:
223
249
head_node_type = _TEST_HEAD_NODE_TYPE_2_POOLS ,
224
250
worker_node_types = _TEST_WORKER_NODE_TYPES_2_POOLS ,
225
251
dashboard_address = _TEST_VERTEX_RAY_DASHBOARD_ADDRESS ,
252
+ node_images = NodeImages (head = _TEST_CPU_IMAGE , worker = _TEST_GPU_IMAGE ),
226
253
)
227
254
_TEST_BEARER_TOKEN = "test-bearer-token"
228
255
_TEST_HEADERS = {
0 commit comments