Coverage for cli/files.py: 90.94%

1"""

2File system operations for GCO CLI.

4Provides functionality to interact with EFS and FSx for Lustre file systems

5attached to GCO regional stacks.

6"""

8from dataclasses import dataclass

9from datetime import UTC, datetime

10from typing import Any

12import boto3

13from botocore.exceptions import ClientError

15from ._image_uri import aws_partition, aws_url_suffix

16from .aws_client import get_aws_client

17from .config import GCOConfig, get_config

18from .kubectl_helpers import update_kubeconfig

21@dataclass

22class FileSystemInfo:

23 """Information about a file system."""

25 file_system_id: str

26 file_system_type: str # "efs" or "fsx"

27 region: str

28 dns_name: str

29 mount_target_ip: str | None = None

30 size_bytes: int | None = None

31 status: str = "available"

32 created_time: datetime | None = None

33 tags: dict[str, str] | None = None

35 def __post_init__(self) -> None:

36 if self.tags is None:

37 self.tags = {}

40@dataclass

41class FileInfo:

42 """Information about a file or directory."""

44 path: str

45 name: str

46 is_directory: bool

47 size_bytes: int = 0

48 modified_time: datetime | None = None

49 owner: str | None = None

52class FileSystemClient:

53 """

54 Client for interacting with GCO file systems.

56 Supports:

57 - Listing file systems (EFS/FSx) in GCO stacks

58 - Getting file system information and access points

59 - Downloading files from pods via kubectl cp

60 """

62 def __init__(self, config: GCOConfig | None = None):

63 self.config = config or get_config()

64 self._session = boto3.Session()

65 self._aws_client = get_aws_client(config)

67 def get_file_systems(self, region: str | None = None) -> list[FileSystemInfo]:

68 """

69 Get all file systems associated with GCO stacks.

71 Args:

72 region: Specific region to query (None for all regions)

74 Returns:

75 List of FileSystemInfo objects

76 """

77 file_systems = []

79 # Get regional stacks

80 stacks = self._aws_client.discover_regional_stacks()

82 if region:

83 stacks = {k: v for k, v in stacks.items() if k == region}

85 for stack_region, stack in stacks.items():

86 # Get EFS file systems

87 if stack.efs_file_system_id: 87 ↛ 93line 87 didn't jump to line 93 because the condition on line 87 was always true

88 efs_info = self._get_efs_info(stack.efs_file_system_id, stack_region)

89 if efs_info: 89 ↛ 93line 89 didn't jump to line 93 because the condition on line 89 was always true

90 file_systems.append(efs_info)

92 # Get FSx file systems

93 if stack.fsx_file_system_id:

94 fsx_info = self._get_fsx_info(stack.fsx_file_system_id, stack_region)

95 if fsx_info: 95 ↛ 85line 95 didn't jump to line 85 because the condition on line 95 was always true

96 file_systems.append(fsx_info)

98 return file_systems

100 def _get_efs_info(self, file_system_id: str, region: str) -> FileSystemInfo | None:

101 """Get information about an EFS file system."""

102 try:

103 efs = self._session.client("efs", region_name=region)

104

105 response = efs.describe_file_systems(FileSystemId=file_system_id)

106 if not response["FileSystems"]:

107 return None

108

109 fs = response["FileSystems"][0]

110

111 # Get mount targets for DNS name

112 mt_response = efs.describe_mount_targets(FileSystemId=file_system_id)

113 mount_target_ip = None

114 if mt_response["MountTargets"]:

115 mount_target_ip = mt_response["MountTargets"][0].get("IpAddress")

116

117 # Get tags

118 tags_response = efs.describe_tags(FileSystemId=file_system_id)

119 tags = {t["Key"]: t["Value"] for t in tags_response.get("Tags", [])}

120

121 return FileSystemInfo(

122 file_system_id=file_system_id,

123 file_system_type="efs",

124 region=region,

125 dns_name=f"{file_system_id}.efs.{region}.{aws_url_suffix(region)}",

126 mount_target_ip=mount_target_ip,

127 size_bytes=fs.get("SizeInBytes", {}).get("Value"),

128 status=fs["LifeCycleState"],

129 created_time=fs.get("CreationTime"),

130 tags=tags,

131 )

132 except ClientError:

133 return None

134

135 def _get_fsx_info(self, file_system_id: str, region: str) -> FileSystemInfo | None:

136 """Get information about an FSx for Lustre file system."""

137 try:

138 fsx = self._session.client("fsx", region_name=region)

139

140 response = fsx.describe_file_systems(FileSystemIds=[file_system_id])

141 if not response["FileSystems"]:

142 return None

143

144 fs = response["FileSystems"][0]

145

146 # Get DNS name from Lustre configuration

147 dns_name = fs.get("DNSName", "")

148

149 # Get tags

150 tags = {t["Key"]: t["Value"] for t in fs.get("Tags", [])}

151

152 return FileSystemInfo(

153 file_system_id=file_system_id,

154 file_system_type="fsx",

155 region=region,

156 dns_name=dns_name,

157 size_bytes=fs.get("StorageCapacity", 0) * 1024 * 1024 * 1024, # GB to bytes

158 status=fs["Lifecycle"],

159 created_time=fs.get("CreationTime"),

160 tags=tags,

161 )

162 except ClientError:

163 return None

164

165 def get_file_system_by_region(self, region: str, fs_type: str = "efs") -> FileSystemInfo | None:

166 """

167 Get file system for a specific region.

168

169 Args:

170 region: AWS region

171 fs_type: "efs" or "fsx"

172

173 Returns:

174 FileSystemInfo or None

175 """

176 file_systems = self.get_file_systems(region)

177 for fs in file_systems:

178 if fs.file_system_type == fs_type: 178 ↛ 177line 178 didn't jump to line 177 because the condition on line 178 was always true

179 return fs

180 return None

181

182 def create_datasync_download_task(

183 self,

184 file_system_id: str,

185 region: str,

186 source_path: str,

187 destination_bucket: str,

188 destination_prefix: str = "",

189 ) -> str:

190 """

191 Create a DataSync task to download files from EFS/FSx to S3.

192

193 This is useful for downloading large amounts of data from file systems

194 that aren't directly accessible.

195

196 Args:

197 file_system_id: EFS or FSx file system ID

198 region: AWS region

199 source_path: Path within the file system

200 destination_bucket: S3 bucket name

201 destination_prefix: S3 key prefix

202

203 Returns:

204 DataSync task ARN

205 """

206 datasync = self._session.client("datasync", region_name=region)

207 partition = aws_partition(region)

208

209 # Determine file system type

210 fs_info = None

211 for fs in self.get_file_systems(region):

212 if fs.file_system_id == file_system_id: 212 ↛ 211line 212 didn't jump to line 211 because the condition on line 212 was always true

213 fs_info = fs

214 break

215

216 if not fs_info:

217 raise ValueError(f"File system {file_system_id} not found in region {region}")

218

219 # Create source location

220 if fs_info.file_system_type == "efs":

221 source_location = datasync.create_location_efs(

222 EfsFilesystemArn=f"arn:{partition}:elasticfilesystem:{region}:{self._get_account_id()}:file-system/{file_system_id}",

223 Subdirectory=source_path,

224 Ec2Config={

225 "SubnetArn": self._get_subnet_arn(region),

226 "SecurityGroupArns": [self._get_security_group_arn(region)],

227 },

228 )

229 source_arn = source_location["LocationArn"]

230 else:

231 source_location = datasync.create_location_fsx_lustre(

232 FsxFilesystemArn=f"arn:{partition}:fsx:{region}:{self._get_account_id()}:file-system/{file_system_id}",

233 Subdirectory=source_path,

234 SecurityGroupArns=[self._get_security_group_arn(region)],

235 )

236 source_arn = source_location["LocationArn"]

237

238 # Create destination location (S3)

239 dest_location = datasync.create_location_s3(

240 S3BucketArn=f"arn:{partition}:s3:::{destination_bucket}",

241 Subdirectory=destination_prefix,

242 S3Config={"BucketAccessRoleArn": self._get_datasync_role_arn(region)},

243 )

244 dest_arn = dest_location["LocationArn"]

245

246 # Create task

247 task = datasync.create_task(

248 SourceLocationArn=source_arn,

249 DestinationLocationArn=dest_arn,

250 Name=f"gco-download-{datetime.now(UTC).strftime('%Y%m%d-%H%M%S')}",

251 Options={

252 "VerifyMode": "ONLY_FILES_TRANSFERRED",

253 "OverwriteMode": "ALWAYS",

254 "PreserveDeletedFiles": "REMOVE",

255 "TransferMode": "CHANGED",

256 },

257 )

258

259 return str(task["TaskArn"])

260

261 def _get_account_id(self) -> str:

262 """Get current AWS account ID."""

263 sts = self._session.client("sts")

264 return str(sts.get_caller_identity()["Account"])

265

266 def _get_subnet_arn(self, _region: str) -> str:

267 """Get a subnet ARN for DataSync in the given region."""

268 # This would need to be implemented based on your VPC setup

269 # For now, return a placeholder

270 raise NotImplementedError("Subnet ARN lookup not implemented - configure via stack outputs")

271

272 def _get_security_group_arn(self, _region: str) -> str:

273 """Get a security group ARN for DataSync in the given region."""

274 raise NotImplementedError(

275 "Security group ARN lookup not implemented - configure via stack outputs"

276 )

277

278 def _get_datasync_role_arn(self, region: str) -> str:

279 """Get the DataSync IAM role ARN."""

280 raise NotImplementedError(

281 "DataSync role ARN lookup not implemented - configure via stack outputs"

282 )

283

284 def get_access_point_info(self, file_system_id: str, region: str) -> list[dict[str, Any]]:

285 """

286 Get EFS access points for a file system.

287

288 Args:

289 file_system_id: EFS file system ID

290 region: AWS region

291

292 Returns:

293 List of access point information

294 """

295 try:

296 efs = self._session.client("efs", region_name=region)

297 response = efs.describe_access_points(FileSystemId=file_system_id)

298

299 return [

300 {

301 "access_point_id": ap["AccessPointId"],

302 "name": ap.get("Name", ""),

303 "path": ap.get("RootDirectory", {}).get("Path", "/"),

304 "posix_user": ap.get("PosixUser", {}),

305 "status": ap["LifeCycleState"],

306 }

307 for ap in response.get("AccessPoints", [])

308 ]

309 except ClientError:

310 return []

311

312 def download_from_pod(

313 self,

314 region: str,

315 pod_name: str,

316 remote_path: str,

317 local_path: str,

318 namespace: str = "gco-jobs",

319 container: str | None = None,

320 ) -> dict[str, Any]:

321 """

322 Download files from a pod using kubectl cp.

323

324 This uses kubectl port-forward internally to copy files from a pod's

325 mounted file system (EFS/FSx) to the local machine.

326

327 Args:

328 region: AWS region where the cluster is located

329 pod_name: Name of the pod to copy from

330 remote_path: Path inside the pod (e.g., /mnt/efs/outputs)

331 local_path: Local destination path

332 namespace: Kubernetes namespace (default: gco-jobs)

333 container: Container name (optional, for multi-container pods)

334

335 Returns:

336 Dict with download status and details

337 """

338 import os

339 import subprocess

340

341 # Update kubeconfig for the cluster

342 cluster_name = f"{self.config.project_name}-{region}"

343 update_kubeconfig(cluster_name, region)

344

345 # Build kubectl cp command

346 # Format: kubectl cp <namespace>/<pod>:<remote_path> <local_path>

347 source = f"{namespace}/{pod_name}:{remote_path}"

348 cmd = ["kubectl", "cp", source, local_path]

349

350 if container:

351 cmd.extend(["-c", container])

352

353 try:

354 subprocess.run(

355 cmd, check=True, capture_output=True, text=True

356 ) # nosemgrep: dangerous-subprocess-use-audit - cmd is a list ["kubectl","cp",source,local_path]; source is namespace/pod:path, local_path is caller-provided destination

357 if os.path.isfile(local_path):

358 size = os.path.getsize(local_path)

359 elif os.path.isdir(local_path): 359 ↛ 366line 359 didn't jump to line 366 because the condition on line 359 was always true

360 size = sum(

361 os.path.getsize(os.path.join(dirpath, filename))

362 for dirpath, _, filenames in os.walk(local_path)

363 for filename in filenames

364 )

365 else:

366 size = 0

367

368 return {

369 "status": "success",

370 "source": source,

371 "destination": local_path,

372 "size_bytes": size,

373 "message": "Download completed successfully",

374 }

375

376 except subprocess.CalledProcessError as e:

377 raise RuntimeError(f"kubectl cp failed: {e.stderr}") from e

378 except FileNotFoundError as e:

379 raise RuntimeError(

380 "kubectl not found. Please install kubectl and ensure it's in your PATH."

381 ) from e

382

383 def list_storage_contents(

384 self,

385 region: str,

386 remote_path: str = "/",

387 storage_type: str = "efs",

388 namespace: str = "gco-jobs",

389 pvc_name: str | None = None,

390 ) -> dict[str, Any]:

391 """

392 List contents of EFS/FSx storage using a temporary helper pod.

393

394 This creates a temporary pod that mounts the storage, lists contents,

395 then cleans up. Useful for discovering what directories/files exist.

396

397 Args:

398 region: AWS region where the cluster is located

399 remote_path: Path inside the storage to list (default: root)

400 storage_type: "efs" or "fsx" (default: efs)

401 namespace: Kubernetes namespace (default: gco-jobs)

402 pvc_name: PVC name to mount (default: gco-shared-storage for EFS,

403 gco-fsx-storage for FSx)

404

405 Returns:

406 Dict with listing status and contents

407 """

408 import subprocess

409 import time

410 import uuid

411

412 # Determine PVC name based on storage type

413 if pvc_name is None: 413 ↛ 417line 413 didn't jump to line 417 because the condition on line 413 was always true

414 pvc_name = "gco-shared-storage" if storage_type == "efs" else "gco-fsx-storage"

415

416 # Determine mount path based on storage type

417 mount_path = "/efs" if storage_type == "efs" else "/fsx"

418

419 # Generate unique pod name

420 helper_pod_name = f"gco-list-helper-{uuid.uuid4().hex[:8]}"

421

422 # Update kubeconfig for the cluster

423 cluster_name = f"{self.config.project_name}-{region}"

424 update_kubeconfig(cluster_name, region)

425

426 # Create helper pod manifest.

427 #

428 # We set an explicit ``resources`` block with CPU + memory but no GPU

429 # so the gco-jobs LimitRange admission plugin does not substitute its

430 # ``max`` value as an implicit request. Without this, a namespace

431 # that already has all 32 GPUs in use (typical during a demo or

432 # heavy workload burst) would reject the helper pod — even though

433 # listing files doesn't need a GPU — because K8s quota admission

434 # would attribute the LimitRange's ``max.nvidia.com/gpu`` to the

435 # pod's request.

436 pod_manifest = f"""

437apiVersion: v1

438kind: Pod

439metadata:

440 name: {helper_pod_name}

441 namespace: {namespace}

442 labels:

443 app: gco-list-helper

444spec:

445 restartPolicy: Never

446 containers:

447 - name: helper

448 image: busybox:1.38.0

449 command: ["sleep", "300"]

450 resources:

451 requests:

452 cpu: "50m"

453 memory: "64Mi"

454 limits:

455 cpu: "200m"

456 memory: "256Mi"

457 volumeMounts:

458 - name: storage

459 mountPath: {mount_path}

460 volumes:

461 - name: storage

462 persistentVolumeClaim:

463 claimName: {pvc_name}

464"""

465

466 try:

467 # Create the helper pod

468 subprocess.run(

469 ["kubectl", "apply", "-f", "-"],

470 input=pod_manifest,

471 capture_output=True,

472 text=True,

473 check=True,

474 )

475

476 # Wait for pod to be ready

477 max_wait = 60

478 waited = 0

479 while waited < max_wait: 479 ↛ 499line 479 didn't jump to line 499 because the condition on line 479 was always true

480 status_result = subprocess.run(

481 [

482 "kubectl",

483 "get",

484 "pod",

485 helper_pod_name,

486 "-n",

487 namespace,

488 "-o",

489 "jsonpath={.status.phase}",

490 ],

491 capture_output=True,

492 text=True,

493 )

494 if status_result.stdout.strip() == "Running": 494 ↛ 496line 494 didn't jump to line 496 because the condition on line 494 was always true

495 break

496 time.sleep(2) # nosemgrep: arbitrary-sleep

497 waited += 2

498

499 if waited >= max_wait: 499 ↛ 500line 499 didn't jump to line 500 because the condition on line 499 was never true

500 raise RuntimeError("Helper pod did not become ready in time")

501

502 # Build the full path inside the pod

503 full_remote_path = f"{mount_path}/{remote_path.lstrip('/')}"

504

505 # List contents using kubectl exec

506 list_result = subprocess.run(

507 [

508 "kubectl",

509 "exec",

510 helper_pod_name,

511 "-n",

512 namespace,

513 "--",

514 "ls",

515 "-la",

516 full_remote_path,

517 ],

518 capture_output=True,

519 text=True,

520 )

521

522 if list_result.returncode != 0:

523 return {

524 "status": "error",

525 "path": remote_path,

526 "storage_type": storage_type,

527 "contents": [],

528 "message": f"Path not found or empty: {list_result.stderr.strip()}",

529 }

530

531 # Parse ls output

532 contents = []

533 for line in list_result.stdout.strip().split("\n"):

534 if line.startswith("total") or not line.strip():

535 continue

536 parts = line.split()

537 if len(parts) >= 9: 537 ↛ 533line 537 didn't jump to line 533 because the condition on line 537 was always true

538 name = " ".join(parts[8:])

539 is_dir = line.startswith("d")

540 size = int(parts[4]) if parts[4].isdigit() else 0

541 contents.append(

542 {

543 "name": name,

544 "is_directory": is_dir,

545 "size_bytes": size,

546 "permissions": parts[0],

547 }

548 )

549

550 return {

551 "status": "success",

552 "path": remote_path,

553 "storage_type": storage_type,

554 "contents": contents,

555 "message": f"Found {len(contents)} items",

556 }

557

558 except subprocess.CalledProcessError as e:

559 error_msg = e.stderr if e.stderr else str(e)

560 raise RuntimeError(f"List failed: {error_msg}") from e

561 except FileNotFoundError as e:

562 raise RuntimeError(

563 "kubectl not found. Please install kubectl and ensure it's in your PATH."

564 ) from e

565 finally:

566 # Always clean up the helper pod

567 import contextlib

568

569 with contextlib.suppress(Exception):

570 subprocess.run(

571 [

572 "kubectl",

573 "delete",

574 "pod",

575 helper_pod_name,

576 "-n",

577 namespace,

578 "--ignore-not-found",

579 ],

580 capture_output=True,

581 text=True,

582 )

583

584 def download_from_storage(

585 self,

586 region: str,

587 remote_path: str,

588 local_path: str,

589 storage_type: str = "efs",

590 namespace: str = "gco-jobs",

591 pvc_name: str | None = None,

592 ) -> dict[str, Any]:

593 """

594 Download files from EFS/FSx storage using a temporary helper pod.

595

596 This creates a temporary pod that mounts the storage, copies files via

597 kubectl cp, then cleans up. Works even after the original job pod is gone.

598

599 Args:

600 region: AWS region where the cluster is located

601 remote_path: Path inside the storage (e.g., /efs-output-example/results.json)

602 local_path: Local destination path

603 storage_type: "efs" or "fsx" (default: efs)

604 namespace: Kubernetes namespace (default: gco-jobs)

605 pvc_name: PVC name to mount (default: gco-shared-storage for EFS,

606 gco-fsx-storage for FSx)

607

608 Returns:

609 Dict with download status and details

610 """

611 import os

612 import subprocess

613 import time

614 import uuid

615

616 # Determine PVC name based on storage type

617 if pvc_name is None:

618 pvc_name = "gco-shared-storage" if storage_type == "efs" else "gco-fsx-storage"

619

620 # Determine mount path based on storage type

621 mount_path = "/efs" if storage_type == "efs" else "/fsx"

622

623 # Generate unique pod name

624 helper_pod_name = f"gco-download-helper-{uuid.uuid4().hex[:8]}"

625

626 # Update kubeconfig for the cluster

627 cluster_name = f"{self.config.project_name}-{region}"

628 update_kubeconfig(cluster_name, region)

629

630 # Create helper pod manifest.

631 #

632 # Explicit ``resources`` block avoids the LimitRange admission plugin

633 # substituting ``max.nvidia.com/gpu`` as an implicit request — see the

634 # ``ls`` helper above for the full rationale.

635 pod_manifest = f"""

636apiVersion: v1

637kind: Pod

638metadata:

639 name: {helper_pod_name}

640 namespace: {namespace}

641 labels:

642 app: gco-download-helper

643spec:

644 restartPolicy: Never

645 containers:

646 - name: helper

647 image: busybox:1.38.0

648 command: ["sleep", "300"]

649 resources:

650 requests:

651 cpu: "50m"

652 memory: "64Mi"

653 limits:

654 cpu: "200m"

655 memory: "256Mi"

656 volumeMounts:

657 - name: storage

658 mountPath: {mount_path}

659 volumes:

660 - name: storage

661 persistentVolumeClaim:

662 claimName: {pvc_name}

663"""

664

665 try:

666 # Create the helper pod

667 subprocess.run(

668 ["kubectl", "apply", "-f", "-"],

669 input=pod_manifest,

670 capture_output=True,

671 text=True,

672 check=True,

673 )

674

675 # Wait for pod to be ready

676 max_wait = 60

677 waited = 0

678 while waited < max_wait:

679 status_result = subprocess.run(

680 [

681 "kubectl",

682 "get",

683 "pod",

684 helper_pod_name,

685 "-n",

686 namespace,

687 "-o",

688 "jsonpath={.status.phase}",

689 ],

690 capture_output=True,

691 text=True,

692 )

693 if status_result.stdout.strip() == "Running":

694 break

695 time.sleep(2) # nosemgrep: arbitrary-sleep

696 waited += 2

697

698 if waited >= max_wait:

699 raise RuntimeError("Helper pod did not become ready in time")

700

701 # Build the full path inside the pod

702 full_remote_path = f"{mount_path}/{remote_path.lstrip('/')}"

703

704 # Copy files from the helper pod

705 source = f"{namespace}/{helper_pod_name}:{full_remote_path}"

706 cmd = ["kubectl", "cp", source, local_path]

707

708 subprocess.run(

709 cmd, check=True, capture_output=True, text=True

710 ) # nosemgrep: dangerous-subprocess-use-audit - cmd is a list ["kubectl","cp",source,local_path]; source is namespace/pod:path, local_path is caller-provided destination

711

712 # Get file info

713 if os.path.isfile(local_path): 713 ↛ 715line 713 didn't jump to line 715 because the condition on line 713 was always true

714 size = os.path.getsize(local_path)

715 elif os.path.isdir(local_path):

716 size = sum(

717 os.path.getsize(os.path.join(dirpath, filename))

718 for dirpath, _, filenames in os.walk(local_path)

719 for filename in filenames

720 )

721 else:

722 size = 0

723

724 return {

725 "status": "success",

726 "source": f"{storage_type}:{remote_path}",

727 "destination": local_path,

728 "size_bytes": size,

729 "storage_type": storage_type,

730 "message": "Download completed successfully",

731 }

732

733 except subprocess.CalledProcessError as e:

734 error_msg = e.stderr if e.stderr else str(e)

735 raise RuntimeError(f"Download failed: {error_msg}") from e

736 except FileNotFoundError as e:

737 raise RuntimeError(

738 "kubectl not found. Please install kubectl and ensure it's in your PATH."

739 ) from e

740 finally:

741 # Always clean up the helper pod

742 import contextlib

743

744 with contextlib.suppress(Exception):

745 subprocess.run(

746 [

747 "kubectl",

748 "delete",

749 "pod",

750 helper_pod_name,

751 "-n",

752 namespace,

753 "--ignore-not-found",

754 ],

755 capture_output=True,

756 text=True,

757 )

758

759

760def get_file_system_client(config: GCOConfig | None = None) -> FileSystemClient:

761 """Get a configured file system client instance."""

762 return FileSystemClient(config)