Skip to content

Commit 1ca9a05

Browse files
yinghsienwucopybara-github
authored andcommitted
feat: Support NFS for Ray cluster creation
PiperOrigin-RevId: 698972917
1 parent 653ba88 commit 1ca9a05

File tree

4 files changed

+45
-1
lines changed

4 files changed

+45
-1
lines changed

google/cloud/aiplatform/vertex_ray/cluster_init.py

+14-1
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
from google.cloud.aiplatform import utils
2626
from google.cloud.aiplatform.utils import resource_manager_utils
2727
from google.cloud.aiplatform_v1beta1.types import persistent_resource_service
28-
28+
from google.cloud.aiplatform_v1beta1.types.machine_resources import NfsMount
2929
from google.cloud.aiplatform_v1beta1.types.persistent_resource import (
3030
PersistentResource,
3131
RayLogsSpec,
@@ -64,6 +64,7 @@ def create_ray_cluster(
6464
enable_logging: Optional[bool] = True,
6565
psc_interface_config: Optional[resources.PscIConfig] = None,
6666
reserved_ip_ranges: Optional[List[str]] = None,
67+
nfs_mounts: Optional[List[resources.NfsMount]] = None,
6768
labels: Optional[Dict[str, str]] = None,
6869
) -> str:
6970
"""Create a ray cluster on the Vertex AI.
@@ -312,6 +313,17 @@ def create_ray_cluster(
312313
ray_metric_spec=ray_metric_spec,
313314
ray_logs_spec=ray_logs_spec,
314315
)
316+
if nfs_mounts:
317+
gapic_nfs_mounts = []
318+
for nfs_mount in nfs_mounts:
319+
gapic_nfs_mounts.append(
320+
NfsMount(
321+
server=nfs_mount.server,
322+
path=nfs_mount.path,
323+
mount_point=nfs_mount.mount_point,
324+
)
325+
)
326+
ray_spec.nfs_mounts = gapic_nfs_mounts
315327
if service_account:
316328
service_account_spec = ServiceAccountSpec(
317329
enable_custom_service_account=True,
@@ -329,6 +341,7 @@ def create_ray_cluster(
329341
)
330342
else:
331343
gapic_psc_interface_config = None
344+
332345
persistent_resource = PersistentResource(
333346
resource_pools=resource_pools,
334347
network=network,

google/cloud/aiplatform/vertex_ray/util/resources.py

+18
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,24 @@ class PscIConfig:
107107
network_attachment: str = None
108108

109109

110+
@dataclasses.dataclass
111+
class NfsMount:
112+
"""NFS mount.
113+
114+
Attributes:
115+
server: Required. IP address of the NFS server.
116+
path: Required. Source path exported from NFS server. Has to start
117+
with '/', and combined with the ip address, it indicates the
118+
source mount path in the form of ``server:path``.
119+
mount_point: Required. Destination mount path. The NFS will be mounted
120+
for the user under /mnt/nfs/<mount_point>.
121+
"""
122+
123+
server: str = None
124+
path: str = None
125+
mount_point: str = None
126+
127+
110128
@dataclasses.dataclass
111129
class Cluster:
112130
"""Ray cluster (output only).

tests/unit/vertex_ray/test_cluster_init.py

+1
Original file line numberDiff line numberDiff line change
@@ -351,6 +351,7 @@ def test_create_ray_cluster_1_pool_custom_image_success(
351351
network=tc.ProjectConstants.TEST_VPC_NETWORK,
352352
cluster_name=tc.ClusterConstants.TEST_VERTEX_RAY_PR_ID,
353353
custom_images=custom_images,
354+
nfs_mounts=[tc.ClusterConstants.TEST_NFS_MOUNT],
354355
)
355356

356357
assert tc.ClusterConstants.TEST_VERTEX_RAY_PR_ADDRESS == cluster_name

tests/unit/vertex_ray/test_constants.py

+12
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,14 @@
2121
from google.cloud.aiplatform.vertex_ray.util.resources import Cluster
2222
from google.cloud.aiplatform.vertex_ray.util.resources import (
2323
AutoscalingSpec,
24+
NfsMount,
2425
PscIConfig,
2526
Resources,
2627
)
2728
from google.cloud.aiplatform_v1beta1.types.machine_resources import DiskSpec
29+
from google.cloud.aiplatform_v1beta1.types.machine_resources import (
30+
NfsMount as GapicNfsMount,
31+
)
2832
from google.cloud.aiplatform_v1beta1.types.machine_resources import (
2933
MachineSpec,
3034
)
@@ -105,6 +109,12 @@ class ClusterConstants:
105109
TEST_VERTEX_RAY_PR_ADDRESS = (
106110
f"{ProjectConstants.TEST_PARENT}/persistentResources/" + TEST_VERTEX_RAY_PR_ID
107111
)
112+
TEST_NFS_MOUNT = NfsMount(
113+
server="10.10.10.10", path="nfs_path", mount_point="nfs_mount_point"
114+
)
115+
TEST_GAPIC_NFS_MOUNT = GapicNfsMount(
116+
server="10.10.10.10", path="nfs_path", mount_point="nfs_mount_point"
117+
)
108118
TEST_CPU_IMAGE_2_9 = "us-docker.pkg.dev/vertex-ai/training/ray-cpu.2-9.py310:latest"
109119
TEST_GPU_IMAGE_2_9 = "us-docker.pkg.dev/vertex-ai/training/ray-gpu.2-9.py310:latest"
110120
TEST_CPU_IMAGE_2_33 = (
@@ -177,6 +187,7 @@ class ClusterConstants:
177187
resource_pool_images={"head-node": TEST_CUSTOM_IMAGE},
178188
ray_metric_spec=RayMetricSpec(disabled=False),
179189
ray_logs_spec=RayLogsSpec(disabled=False),
190+
nfs_mounts=[TEST_GAPIC_NFS_MOUNT],
180191
),
181192
),
182193
psc_interface_config=None,
@@ -227,6 +238,7 @@ class ClusterConstants:
227238
ray_spec=RaySpec(
228239
resource_pool_images={"head-node": TEST_CUSTOM_IMAGE},
229240
ray_metric_spec=RayMetricSpec(disabled=False),
241+
nfs_mounts=[TEST_GAPIC_NFS_MOUNT],
230242
),
231243
),
232244
psc_interface_config=None,

0 commit comments

Comments
 (0)