File size: 5,974 Bytes
d572bbd | 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 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | """
Cloud storage integration module for Google Colab and Google Drive.
Provides secure, authenticated storage operations for training results,
checkpoints, and logs. Supports automatic environment detection for
Colab vs local execution.
"""
from __future__ import annotations
import json
import logging
import os
import shutil
import time
from pathlib import Path
from typing import Optional
logger = logging.getLogger(__name__)
def is_colab_environment() -> bool:
try:
import google.colab
return True
except ImportError:
return False
def mount_google_drive(mount_point: str = "/content/drive") -> bool:
if not is_colab_environment():
logger.info("Not running in Colab, skipping Google Drive mount")
return False
try:
from google.colab import drive
drive.mount(mount_point)
logger.info("Google Drive mounted at %s", mount_point)
return True
except Exception as e:
logger.error("Failed to mount Google Drive: %s", e)
return False
def get_drive_path(base_path: str = "/content/drive/MyDrive/EasyTranslate") -> Optional[Path]:
if not os.path.exists("/content/drive"):
logger.warning("Google Drive not mounted, cannot resolve drive path")
return None
drive_path = Path(base_path)
drive_path.mkdir(parents=True, exist_ok=True)
return drive_path
def sync_checkpoints_to_drive(
local_checkpoint_dir: str | Path,
drive_base_path: str = "/content/drive/MyDrive/EasyTranslate",
max_retries: int = 3,
) -> bool:
local_dir = Path(local_checkpoint_dir)
if not local_dir.exists():
logger.warning("Local checkpoint directory does not exist: %s", local_dir)
return False
drive_path = get_drive_path(drive_base_path)
if drive_path is None:
return False
drive_checkpoint_dir = drive_path / "checkpoints"
drive_checkpoint_dir.mkdir(parents=True, exist_ok=True)
success = True
for ckpt_file in local_dir.glob("*.pt"):
dest = drive_checkpoint_dir / ckpt_file.name
for attempt in range(max_retries):
try:
shutil.copy2(ckpt_file, dest)
logger.info("Synced checkpoint to Drive: %s", dest)
break
except Exception as e:
logger.warning("Sync attempt %d/%d failed for %s: %s", attempt + 1, max_retries, ckpt_file.name, e)
if attempt == max_retries - 1:
success = False
time.sleep(2 ** attempt)
return success
def sync_logs_to_drive(
local_log_dir: str | Path,
drive_base_path: str = "/content/drive/MyDrive/EasyTranslate",
max_retries: int = 3,
) -> bool:
local_dir = Path(local_log_dir)
if not local_dir.exists():
logger.warning("Local log directory does not exist: %s", local_dir)
return False
drive_path = get_drive_path(drive_base_path)
if drive_path is None:
return False
drive_log_dir = drive_path / "logs"
drive_log_dir.mkdir(parents=True, exist_ok=True)
success = True
for log_file in local_dir.glob("*"):
if log_file.is_file():
dest = drive_log_dir / log_file.name
for attempt in range(max_retries):
try:
shutil.copy2(log_file, dest)
break
except Exception as e:
logger.warning("Log sync attempt %d/%d failed: %s", attempt + 1, max_retries, e)
if attempt == max_retries - 1:
success = False
time.sleep(2 ** attempt)
return success
def save_training_summary_to_drive(
summary: dict,
drive_base_path: str = "/content/drive/MyDrive/EasyTranslate",
) -> bool:
drive_path = get_drive_path(drive_base_path)
if drive_path is None:
return False
summary_path = drive_path / "training_summary.json"
try:
with open(summary_path, "w", encoding="utf-8") as f:
json.dump(summary, f, indent=2, ensure_ascii=False, default=str)
logger.info("Training summary saved to Drive: %s", summary_path)
return True
except Exception as e:
logger.error("Failed to save training summary to Drive: %s", e)
return False
def sync_all_to_drive(
checkpoint_dir: str | Path = "checkpoints",
log_dir: str | Path = "logs",
drive_base_path: str = "/content/drive/MyDrive/EasyTranslate",
) -> dict:
results = {
"checkpoints_synced": sync_checkpoints_to_drive(checkpoint_dir, drive_base_path),
"logs_synced": sync_logs_to_drive(log_dir, drive_base_path),
}
summary_path = Path(checkpoint_dir) / "training_summary.json"
if summary_path.exists():
with open(summary_path, "r", encoding="utf-8") as f:
summary = json.load(f)
results["summary_saved"] = save_training_summary_to_drive(summary, drive_base_path)
logger.info("Drive sync results: %s", results)
return results
def setup_colab_environment() -> dict:
env_info = {
"is_colab": is_colab_environment(),
"gpu_available": False,
"gpu_info": None,
"drive_mounted": False,
"drive_path": None,
}
if env_info["is_colab"]:
env_info["drive_mounted"] = mount_google_drive()
env_info["drive_path"] = str(get_drive_path()) if env_info["drive_mounted"] else None
try:
import torch
env_info["gpu_available"] = torch.cuda.is_available()
if env_info["gpu_available"]:
env_info["gpu_info"] = {
"device_count": torch.cuda.device_count(),
"device_name": torch.cuda.get_device_name(0),
"device_capability": torch.cuda.get_device_capability(0),
}
except ImportError:
pass
logger.info("Environment setup: %s", {k: v for k, v in env_info.items() if k != "gpu_info"})
return env_info |