import modal from suno_utils.worker.settings import s3_client from suno_utils.worker.modal_base import get_modal_base_image import os from pathlib import Path import argparse import fnmatch import sys import platform from suno_utils.worker.modal_model_volume import ( GPT_PATH_DICT, DIFFUSION_PATH_DICT, CODEC_PATH_DICT, GPT_CKPT_PATH, TOKENIZER_PATH, SEMANTIC_CKPT_PATH, SEMANTIC_CENTROIDS_PATH, CODEC_CKPT_PATH, HOOT_CKPT_PATH, HOOT_TOKENIZER_PATH, FASTTEXT_CKPT_PATH, TOKENIZER_FILEPATH, SEMANTIC_MODEL_FILEPATH, SEMANTIC_CLUSTERS_FILEPATH, CODEC_FILEPATH, DIT_MODEL_FILEPATH, ) # Before sending to Modal local_os = platform.system() print(f"Local OS: {local_os}") # -------------------------------------- # Runbook: https://www.notion.so/suno-ai/Model-Manager-1bbb01573ccf80aea280fea5084764fb?pvs=4 # -------------------------------------- APP_NAME = "model-store" SECRETS = [ (modal.Secret.from_name("studio-aws")), modal.Secret.from_name("redis-test"), ] volume = modal.Volume.from_name("model-store") disk_prefix = "/app2/suno/modal/models" ssd_prefix = "/mnt/localdisk/models" local_path = "/local/models" volume_path = "/volume/models" base_image = get_modal_base_image().add_local_python_source("suno_utils") app = modal.App(APP_NAME, image=base_image) # Example usage in the cron function @app.function( volumes={"/volume": volume}, secrets=SECRETS, cloud="aws", region="us-east", ) def model_store(): print("Start inspect volume") # Recursively list all files in the volume with sizes import os import datetime from datetime import timezone def get_size_format(num, suffix="B"): """Convert bytes to human-readable format""" for unit in ["", "K", "M", "G", "T", "P", "E", "Z"]: if abs(num) < 1024.0: return f"{num:.2f} {unit}{suffix}" num /= 1024.0 return f"{num:.2f} Y{suffix}" volume_path = "/volume" total_size = 0 file_count = 0 dir_count = 0 print(f"Contents of volume mounted at {volume_path}:") print("-" * 100) print(f"{'Size':>10} | {'Modified (UTC)':>16} | {'Path':<65}") print("-" * 100) for root, dirs, files in os.walk(volume_path): dir_count += len(dirs) for file in files: file_path = os.path.join(root, file) try: # Get file size file_size = os.path.getsize(file_path) total_size += file_size # Get modification time in UTC mod_time = os.path.getmtime(file_path) mod_time_dt = datetime.datetime.fromtimestamp(mod_time, tz=timezone.utc) # Format time to minutes precision time_str = mod_time_dt.strftime("%Y-%m-%d %H:%M") file_count += 1 rel_path = os.path.relpath(file_path, volume_path) if len(rel_path) > 65: rel_path = "..." + rel_path[-62:] print(f"{get_size_format(file_size):>10} | {time_str:>16} | {rel_path:<65}") except Exception as e: print(f"Error accessing {file_path}: {e}") print("-" * 100) print(f"Total: {file_count} files in {dir_count} directories") print(f"Total size: {get_size_format(total_size)}") print("End inspect volume") @app.function(volumes={"/volume": volume}) def overwrite_timestamp(filepaths, access_times, mod_times): for filepath, access_time, mod_time in zip(filepaths, access_times, mod_times): filepath = "/volume" + filepath print(f"overwriting {filepath} with access time {access_time} and mod time {mod_time}") os.utime(filepath, (access_time, mod_time)) print(Path(filepath).stat().st_mtime) # -------------------------------------- # Below this is backfill script to download models from S3 to local, this should be run on a ubuntu machine in Oracle Cloud because it requires a lot of space # -------------------------------------- # Backfill helper function def download_s3_file(bucket_name: str, file_path: str = "", output_prefix: str = disk_prefix): if os.path.exists(os.path.join(output_prefix, file_path)): print(f"File {file_path} already exists in {output_prefix}") return os.path.join(output_prefix, file_path) directory_path = os.path.dirname(file_path) local_path = os.path.join(output_prefix, directory_path.lstrip("/")) # Create output directory if it doesn't exist Path(local_path).mkdir(parents=True, exist_ok=True) # List objects in the bucket with the given prefix response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=directory_path.lstrip("/")) # Download the file try: s3_client.download_file( Bucket=bucket_name, Key=file_path.lstrip("/"), Filename=os.path.join(output_prefix, file_path.lstrip("/")), ) return local_path except Exception as e: print(f"Error downloading {file_path}: {e}") return None # Backfill helper function def is_ubuntu(): """ Checks if the current system is Ubuntu. Returns: bool: True if running on Ubuntu, False otherwise """ if os.path.exists("/etc/os-release"): with open("/etc/os-release", "r") as f: os_release = f.read().lower() if "ubuntu" in os_release: return True return False def download_model_from_path(model: str): print(f"Downloading {model} from S3") local_path = os.path.join(prefix, model.lstrip("/")) print(local_path) # Check if file already exists locally if os.path.exists(local_path): print(f"Skipping {model}: Local file already exists at {local_path}") return # Create directory if it doesn't exist os.makedirs(os.path.dirname(local_path), exist_ok=True) # Download file try: local_path = download_s3_file("suno-data", model, prefix) print(f"Downloaded {model} to {local_path}") except Exception as e: print(f"Error downloading {model}: {e}") # Backfill helper function def download_models_in_dict(path_dict: dict): for model in path_dict: print(f"Downloading {model} from S3") local_path = os.path.join(prefix, path_dict[model].lstrip("/")) print(local_path) # Check if file already exists locally if os.path.exists(local_path): print(f"Skipping {model}: Local file already exists at {local_path}") continue # Create directory if it doesn't exist os.makedirs(os.path.dirname(local_path), exist_ok=True) # Download file try: local_path = download_s3_file("suno-data", path_dict[model], prefix) print(f"Downloaded {model} to {local_path}") except Exception as e: print(f"Error downloading {model}: {e}") if __name__ == "__main__": """backfill script to download models from S3 to local""" if not is_ubuntu(): print( "This script is only supported on Ubuntu, your computer likely does not have sufficient space to download models" ) exit(1) parser = argparse.ArgumentParser(description="Model store management tool") parser.add_argument( "action", choices=["download", "upload", "sync"], help="Action to perform: 'download' models from S3 to local, or 'upload' models to Modal volume", ) parser.add_argument( "--sync-mode", choices=["soft", "hard"], default="soft", help="For 'sync' action: 'soft' only copies directory structure, 'hard' copies directories and all files", ) parser.add_argument( "--path", choices=["disk", "ssd"], default="disk", help="For 'download': local path to store downloaded models; For 'upload': path to models to upload", ) parser.add_argument( "--force", nargs="*", default=[], help="For 'upload': specify file patterns to force upload even if already exist in volume. Use '*' for all files.", ) args = parser.parse_args() print(f"Action: {args.action}") print(f"Path: {args.path}") if args.path == "disk": prefix = disk_prefix elif args.path == "ssd": prefix = ssd_prefix else: raise ValueError(f"Invalid path: {args.path}") if args.action == "download": print("Downloading models from remote storage...") download_models_in_dict(GPT_PATH_DICT) download_models_in_dict(DIFFUSION_PATH_DICT) download_models_in_dict(CODEC_PATH_DICT) download_model_from_path(GPT_CKPT_PATH) download_model_from_path(TOKENIZER_PATH) download_model_from_path(SEMANTIC_CKPT_PATH) download_model_from_path(SEMANTIC_CENTROIDS_PATH) download_model_from_path(CODEC_CKPT_PATH) download_model_from_path(HOOT_CKPT_PATH) download_model_from_path(HOOT_TOKENIZER_PATH) download_model_from_path(FASTTEXT_CKPT_PATH) download_model_from_path(TOKENIZER_FILEPATH) download_model_from_path(SEMANTIC_MODEL_FILEPATH) download_model_from_path(SEMANTIC_CLUSTERS_FILEPATH) download_model_from_path(CODEC_FILEPATH) download_model_from_path(DIT_MODEL_FILEPATH) elif args.action == "upload": should_force = "--force" in sys.argv force_all = len(args.force) == 0 and should_force print(f"Copying files from {prefix} to volume...") existing_files = set() existing_file_sizes = {} existing_file_timestamps = {} res = volume.listdir(path="/", recursive=True) # for i in res: # existing_files.add("/" + i.path) try: # Get list of files in the volume using Modal's volume API res = volume.listdir(path="/", recursive=True) for i in res: file_path = "/" + i.path existing_files.add(file_path) if hasattr(i, "size"): existing_file_sizes[file_path] = i.size if hasattr(i, "mtime"): existing_file_timestamps[file_path] = i.mtime print(f"Found {len(existing_files)} existing files in volume") except Exception as e: print(f"Error listing volume contents: {e}") print("Will attempt to upload all files") local_file_paths = [] volume_file_paths = [] access_times = [] mod_times = [] with volume.batch_upload(force=True) as batch: # Walk through all files in the source directory for root, dirs, files in os.walk(prefix): # Get the relative path to maintain directory structure rel_path = os.path.relpath(root, prefix) # Handle the root directory case if rel_path == ".": rel_path = "" # Copy each file while maintaining the structure for file in files: local_file_path = os.path.join(root, file) # Ensure paths use forward slashes for consistency rel_file_path = os.path.join(rel_path, file).replace("\\", "/") volume_file_path = os.path.join("/models", rel_file_path).replace("\\", "/") # Check if we should force upload this file based on patterns should_force = force_all or any( fnmatch.fnmatch(file, pattern) for pattern in args.force ) # Get local file size local_size = os.path.getsize(local_file_path) # Check if file exists and sizes are different (if we have size info) size_mismatch = False if ( volume_file_path in existing_file_sizes and existing_file_sizes[volume_file_path] != local_size ): size_mismatch = True print(f"Size mismatch for {rel_file_path}") timestamp_mismatch = False if volume_file_path in existing_file_timestamps and existing_file_timestamps[ volume_file_path ] != int(os.path.getmtime(local_file_path)): timestamp_mismatch = True print(f"Timestamp mismatch for {rel_file_path}") # Skip if file exists with same size and not forced if ( volume_file_path in existing_files and not size_mismatch and not timestamp_mismatch and not should_force ): print(f"Skipping {volume_file_path}: File already exists in volume") continue print(f"Uploading {local_file_path} to {volume_file_path}") batch.put_file(local_file_path, volume_file_path) local_file_paths.append(local_file_path) volume_file_paths.append(volume_file_path) access_times.append(os.path.getatime(local_file_path)) mod_times.append(os.path.getmtime(local_file_path)) with modal.enable_output(): with app.run(): overwrite_timestamp.remote( volume_file_paths, access_times, mod_times, ) if args.path == "ssd": print("Starting background job to copy SSD files to disk storage...") # Function to launch background process that copies specific files from SSD to disk def start_background_copy(file_paths): import subprocess import datetime import tempfile import os # Create a timestamp for the log file timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") log_file = f"/tmp/ssd_to_disk_copy_{timestamp}.log" # Create a temp shell script for the copy operation with tempfile.NamedTemporaryFile(mode="w", suffix=".sh", delete=False) as script: script_path = script.name script.write("#!/bin/bash\n\n") script.write( f"echo 'Starting copy of {len(file_paths)} files from {ssd_prefix} to {disk_prefix} at $(date)' > {log_file}\n" ) # Copy each file while preserving the directory structure for src_path in file_paths: # Get the relative path from SSD prefix rel_path = os.path.relpath(src_path, ssd_prefix) # Construct destination path dst_path = os.path.join(disk_prefix, rel_path) # Ensure destination directory exists dst_dir = os.path.dirname(dst_path) script.write(f"sudo mkdir -p '{dst_dir}'\n") # Copy the file preserving attributes (-p) script.write(f"sudo cp -p '{src_path}' '{dst_path}' >> {log_file} 2>&1\n") script.write(f"echo 'Copied: {rel_path}' >> {log_file}\n") script.write(f"echo 'Copy completed at $(date)' >> {log_file}\n") # Make the script executable subprocess.run(["chmod", "+x", script_path]) # Execute the script in background print( f"Starting background copy of {len(file_paths)} files from {ssd_prefix} to {disk_prefix}" ) print(f"Progress will be logged to {log_file}") subprocess.Popen(f"nohup {script_path} > /dev/null 2>&1 &", shell=True) # Start the background copy with the list of files we uploaded try: start_background_copy(local_file_paths) print("Background copy process started. You can safely exit this script.") except Exception as e: print(f"Failed to start background copy: {e}") print("You may need to manually copy files from SSD to disk.") print("All files copied to volume successfully.") exit(0) elif args.action == "sync": print(f"Syncing models from disk to ssd (Mode: {args.sync_mode})...") # Ensure both paths exist if not os.path.exists(disk_prefix): print(f"Error: Source path {disk_prefix} does not exist") exit(1) for root, dirs, files in os.walk(disk_prefix): # Get relative path to maintain directory structure rel_path = os.path.relpath(root, disk_prefix) # Create corresponding directory on SSD if rel_path != ".": ssd_dir = os.path.join(ssd_prefix, rel_path) os.makedirs(ssd_dir, exist_ok=True) if args.sync_mode == "soft": print("Soft sync mode: Only creating directories on SSD") exit(0) # Create SSD directory if it doesn't exist os.makedirs(ssd_prefix, exist_ok=True) # Walk through all files on disk copied_count = 0 skipped_count = 0 error_count = 0 total_size_copied = 0 print(f"Copying files from {disk_prefix} to {ssd_prefix}...") print("-" * 80) for root, dirs, files in os.walk(disk_prefix): # Get relative path to maintain directory structure rel_path = os.path.relpath(root, disk_prefix) # Create corresponding directory on SSD if rel_path != ".": ssd_dir = os.path.join(ssd_prefix, rel_path) os.makedirs(ssd_dir, exist_ok=True) # Copy each file for file in files: disk_file = os.path.join(root, file) ssd_file = os.path.join(ssd_prefix, rel_path, file) # Skip if destination file exists and has same size if os.path.exists(ssd_file) and os.path.getsize(disk_file) == os.path.getsize(ssd_file): print(f"Skipping (already exists): {os.path.join(rel_path, file)}") skipped_count += 1 continue # Copy the file try: print(f"Copying: {os.path.join(rel_path, file)}") # Ensure directory exists (for the "." case) os.makedirs(os.path.dirname(ssd_file), exist_ok=True) # Faster file copying with larger buffer def fast_copy(src, dst, buffer_size=10 * 1024 * 1024): # 10MB buffer with open(src, "rb") as fsrc, open(dst, "wb") as fdst: while True: buf = fsrc.read(buffer_size) if not buf: break fdst.write(buf) # Still preserve metadata like copy2 import shutil shutil.copystat(src, dst) # Try system commands first (much faster), fall back to Python implementation import subprocess import platform try: if platform.system() == "Linux" or platform.system() == "Darwin": # Use cp on Linux/Mac (much faster than Python's copy) result = subprocess.run(["cp", disk_file, ssd_file], check=True) # Apply metadata separately import shutil shutil.copystat(disk_file, ssd_file) else: # Fall back to fast Python implementation fast_copy(disk_file, ssd_file) except Exception as e: # If system commands fail, use the Python implementation print(f" System copy failed ({str(e)}), using Python implementation") fast_copy(disk_file, ssd_file) file_size = os.path.getsize(disk_file) total_size_copied += file_size copied_count += 1 except Exception as e: print(f"Error copying {disk_file}: {e}") error_count += 1 # Format total size copied def get_size_format(bytes): for unit in ["B", "KB", "MB", "GB", "TB"]: if bytes < 1024.0: return f"{bytes:.2f} {unit}" bytes /= 1024.0 return f"{bytes:.2f} PB" print("-" * 80) print("Sync complete:") print(f" - Files copied: {copied_count}") print(f" - Files skipped (already exist): {skipped_count}") print(f" - Errors: {error_count}") print(f" - Total data copied: {get_size_format(total_size_copied)}") if error_count > 0: print("\nWarning: Some files could not be copied. Check the errors above.") else: print("Invalid action. Please specify 'download' or 'upload'.") exit(1)