slurmray 6.0.4__py3-none-any.whl

slurmray/__main__.py

@@ -0,0 +1,5 @@
+from slurmray.cli import main
+
+if __name__ == "__main__":
+    main()
+

slurmray/assets/cleanup_old_projects.py

@@ -0,0 +1,171 @@
+#!/usr/bin/env python3
+"""
+Cleanup script for old SlurmRay projects on cluster.
+
+This script is designed to be run as a cron job to automatically remove
+old project files and venv that have exceeded their retention period.
+
+Usage:
+    python cleanup_old_projects.py [--cluster-type slurm|desi] [--base-path /path/to/users]
+
+The script will:
+- Scan user directories for slurmray-server/{project_name}/ directories
+- Read .retention_timestamp and .retention_days files
+- Delete projects that have exceeded their retention period
+"""
+
+import os
+import sys
+import time
+import argparse
+import logging
+from pathlib import Path
+
+
+def read_retention_files(project_dir):
+    """
+    Read retention timestamp and days from project directory.
+    
+    Args:
+        project_dir: Path to project directory
+        
+    Returns:
+        tuple: (timestamp, retention_days) or (None, None) if files don't exist
+    """
+    timestamp_file = os.path.join(project_dir, ".retention_timestamp")
+    retention_days_file = os.path.join(project_dir, ".retention_days")
+    
+    if not os.path.exists(timestamp_file) or not os.path.exists(retention_days_file):
+        return None, None
+    
+    try:
+        with open(timestamp_file, 'r') as f:
+            timestamp_str = f.read().strip()
+            timestamp = int(timestamp_str)
+        
+        with open(retention_days_file, 'r') as f:
+            retention_days_str = f.read().strip()
+            retention_days = int(retention_days_str)
+        
+        return timestamp, retention_days
+    except (ValueError, IOError):
+        return None, None
+
+
+def cleanup_user_projects(username, base_path, cluster_type, logger):
+    """
+    Cleanup old projects for a specific user.
+    
+    Args:
+        username: Username
+        base_path: Base path to user directories (e.g., /users or /home)
+        cluster_type: Type of cluster ('slurm' or 'desi')
+        logger: Logger instance
+    """
+    user_dir = os.path.join(base_path, username)
+    if not os.path.exists(user_dir):
+        return
+    
+    slurmray_base = os.path.join(user_dir, "slurmray-server")
+    if not os.path.exists(slurmray_base):
+        return
+    
+    # Scan for project directories
+    if not os.path.isdir(slurmray_base):
+        return
+    
+    for item in os.listdir(slurmray_base):
+        project_dir = os.path.join(slurmray_base, item)
+        if not os.path.isdir(project_dir):
+            continue
+        
+        # Read retention files
+        timestamp, retention_days = read_retention_files(project_dir)
+        
+        if timestamp is None or retention_days is None:
+            # No retention files, skip this project
+            continue
+        
+        # Calculate age in days
+        current_time = int(time.time())
+        age_seconds = current_time - timestamp
+        age_days = age_seconds / (24 * 3600)
+        
+        # Check if project should be deleted
+        if age_days > retention_days:
+            logger.info(f"Deleting project {username}/{item}: age={age_days:.1f} days, retention={retention_days} days")
+            try:
+                import shutil
+                shutil.rmtree(project_dir)
+                logger.info(f"Successfully deleted project {username}/{item}")
+            except Exception as e:
+                logger.error(f"Failed to delete project {username}/{item}: {e}")
+                raise  # Fail-fast on errors
+
+
+def main():
+    """Main entry point for cleanup script"""
+    parser = argparse.ArgumentParser(description="Cleanup old SlurmRay projects on cluster")
+    parser.add_argument(
+        "--cluster-type",
+        choices=["slurm", "desi"],
+        default="slurm",
+        help="Type of cluster (default: slurm)"
+    )
+    parser.add_argument(
+        "--base-path",
+        help="Base path to user directories (default: /users for slurm, /home for desi)"
+    )
+    parser.add_argument(
+        "--log-file",
+        help="Log file path (default: /tmp/slurmray_cleanup.log)"
+    )
+    
+    args = parser.parse_args()
+    
+    # Setup logging
+    log_file = args.log_file or "/tmp/slurmray_cleanup.log"
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(levelname)s - %(message)s',
+        handlers=[
+            logging.FileHandler(log_file),
+            logging.StreamHandler(sys.stdout)
+        ]
+    )
+    logger = logging.getLogger(__name__)
+    
+    # Determine base path
+    if args.base_path:
+        base_path = args.base_path
+    else:
+        if args.cluster_type == "slurm":
+            base_path = "/users"
+        else:  # desi
+            base_path = "/home"
+    
+    logger.info(f"Starting cleanup for cluster type: {args.cluster_type}, base path: {base_path}")
+    
+    if not os.path.exists(base_path):
+        logger.error(f"Base path does not exist: {base_path}")
+        sys.exit(1)
+    
+    # Scan all user directories
+    cleaned_count = 0
+    for username in os.listdir(base_path):
+        user_path = os.path.join(base_path, username)
+        if not os.path.isdir(user_path):
+            continue
+        
+        try:
+            cleanup_user_projects(username, base_path, args.cluster_type, logger)
+        except Exception as e:
+            logger.error(f"Error processing user {username}: {e}")
+            raise  # Fail-fast on errors
+    
+    logger.info("Cleanup completed")
+
+
+if __name__ == "__main__":
+    main()
+

slurmray/assets/sbatch_template.sh

@@ -0,0 +1,67 @@
+#!/bin/bash
+
+# THIS FILE IS GENERATED BY AUTOMATION SCRIPT! PLEASE REFER TO ORIGINAL SCRIPT!
+# THIS FILE IS A TEMPLATE AND IT SHOULD NOT BE DEPLOYED TO PRODUCTION!
+
+#SBATCH --partition={{PARTITION_NAME}}
+#SBATCH --mem {{MEMORY}}G
+#SBATCH --job-name={{JOB_NAME}}
+#SBATCH --output={{JOB_NAME}}.log
+#SBATCH --time {{RUNNING_TIME}}
+{{GIVEN_NODE}}
+
+### This script works for any number of nodes, Ray will find and manage all resources
+#SBATCH --nodes={{NUM_NODES}}
+
+### Give all resources to a single Ray task, ray can manage the resources internally
+#SBATCH --ntasks-per-node=1
+{{PARTITION_SPECIFICS}}
+
+# Load modules or your own conda environment here
+{{LOAD_ENV}}
+
+################# DON NOT CHANGE THINGS HERE UNLESS YOU KNOW WHAT YOU ARE DOING ###############
+# This script is a modification to the implementation suggest by gregSchwartz18 here:
+# https://github.com/ray-project/ray/issues/826#issuecomment-522116599
+redis_password=$(uuidgen)
+export redis_password
+
+nodes=$(scontrol show hostnames $SLURM_JOB_NODELIST) # Getting the node names
+nodes_array=($nodes)
+
+node_1=${nodes_array[0]}
+ip=$(srun --nodes=1 --ntasks=1 -w $node_1 hostname --ip-address) # making redis-address
+
+if [[ $ip == *" "* ]]; then
+  IFS=' ' read -ra ADDR <<<"$ip"
+  if [[ ${#ADDR[0]} > 16 ]]; then
+    ip=${ADDR[1]}
+  else
+    ip=${ADDR[0]}
+  fi
+  echo "We detect space in ip! You are using IPV6 address. We split the IPV4 address as $ip"
+fi
+
+port=6379
+ip_head=$ip:$port
+export ip_head
+echo "IP Head: $ip_head"
+
+echo "STARTING HEAD at $node_1"
+# srun --nodes=1 --ntasks=1 -w $node_1 start-head.sh $ip $redis_password &
+srun --nodes=1 --ntasks=1 -w $node_1 \
+  ray start --head --include-dashboard=true --dashboard-host 0.0.0.0 --dashboard-port=8265 --node-ip-address=$ip --port=6379 --redis-password=$redis_password --block &
+sleep 30
+
+worker_num=$(($SLURM_JOB_NUM_NODES - 1)) #number of nodes other than the head node
+for ((i = 1; i <= $worker_num; i++)); do
+  node_i=${nodes_array[$i]}
+  echo "STARTING WORKER $i at $node_i"
+  srun --nodes=1 --ntasks=1 -w $node_i ray start --address $ip_head --redis-password=$redis_password --block &
+  sleep 5
+done
+
+##############################################################################################
+
+#### call your code below
+{{COMMAND_PLACEHOLDER}} {{COMMAND_SUFFIX}}

slurmray/assets/slurmray_server.sh

@@ -0,0 +1,145 @@
+#!/bin/sh
+
+echo "Installing slurmray server"
+
+# Copy files
+mv -t slurmray-server requirements.txt slurmray_server.py
+mv -t slurmray-server/.slogs/server func.pkl args.pkl 
+cd slurmray-server
+
+# Load modules
+# Using specific versions for Curnagl compatibility (SLURM 24.05.3)
+# gcc/13.2.0: Latest GCC version
+# python/3.12.1: Latest Python version on Curnagl
+# cuda/12.6.2: Latest CUDA version
+# cudnn/9.2.0.82-12: Compatible with cuda/12.6.2
+module load gcc/13.2.0 rust python/3.12.1 cuda/12.6.2 cudnn/9.2.0.82-12
+
+# Create venv if it doesn't exist (hash check is done in Python before file upload)
+# If venv needs recreation, it has already been removed by Python
+# Check for force reinstall flag
+if [ -f ".force_reinstall" ]; then
+    echo "Force reinstall flag detected: removing existing virtualenv..."
+    rm -rf .venv
+    rm -f .force_reinstall
+fi
+
+if [ ! -d ".venv" ]; then
+    echo "Creating virtualenv..."
+    python3 -m venv .venv
+else
+    echo "Using existing virtualenv (requirements unchanged)..."
+    VENV_EXISTED=true
+fi
+
+source .venv/bin/activate
+
+# Install requirements if file exists and is not empty
+if [ -f requirements.txt ]; then
+    # Check if requirements.txt is empty (only whitespace)
+    if [ -s requirements.txt ]; then
+        echo "📥 Installing dependencies from requirements.txt..."
+        
+        # Get installed packages once (fast, single command) - create lookup file
+        uv pip list --format=freeze 2>/dev/null | sed 's/==/ /' | awk '{print $1" "$2}' > /tmp/installed_packages.txt || touch /tmp/installed_packages.txt
+        
+        # Process requirements: filter duplicates and check what needs installation
+        INSTALL_ERRORS=0
+        SKIPPED_COUNT=0
+        > /tmp/to_install.txt  # Clear file
+        
+        while IFS= read -r line || [ -n "$line" ]; do
+            # Skip empty lines and comments
+            line=$(echo "$line" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
+            if [ -z "$line" ] || [ "${line#\#}" != "$line" ]; then
+                continue
+            fi
+            
+            # Extract package name (remove version specifiers and extras)
+            pkg_name=$(echo "$line" | sed 's/[<>=!].*//' | sed 's/\[.*\]//' | sed 's/[[:space:]]*//' | tr '[:upper:]' '[:lower:]')
+            if [ -z "$pkg_name" ]; then
+                continue
+            fi
+            
+            # Skip duplicates (check if we've already processed this package)
+            if grep -qi "^$pkg_name$" /tmp/seen_packages.txt 2>/dev/null; then
+                continue
+            fi
+            echo "$pkg_name" >> /tmp/seen_packages.txt
+            
+            # Extract required version if present
+            required_version=""
+            if echo "$line" | grep -q "=="; then
+                required_version=$(echo "$line" | sed 's/.*==\([^;]*\).*/\1/' | sed 's/[[:space:]]*//')
+            fi
+            
+            # Check if package is already installed with correct version
+            installed_version=$(grep -i "^$pkg_name " /tmp/installed_packages.txt 2>/dev/null | awk '{print $2}' | head -1)
+            
+            if [ -n "$installed_version" ]; then
+                if [ -z "$required_version" ] || [ "$installed_version" = "$required_version" ]; then
+                    echo "  ⏭️  $pkg_name==$installed_version (already installed)"
+                    SKIPPED_COUNT=$((SKIPPED_COUNT + 1))
+                    continue
+                fi
+            fi
+            
+            # Package not installed or version mismatch, add to install list
+            echo "$line" >> /tmp/to_install.txt
+        done < requirements.txt
+        
+        # Install packages that need installation
+        if [ -s /tmp/to_install.txt ]; then
+            > /tmp/install_errors.txt  # Track errors
+            while IFS= read -r line; do
+                pkg_name=$(echo "$line" | sed 's/[<>=!].*//' | sed 's/\[.*\]//' | sed 's/[[:space:]]*//')
+                if uv pip install --quiet "$line" >/dev/null 2>&1; then
+                    echo "  ✅ $pkg_name"
+                else
+                    echo "  ❌ $pkg_name"
+                    echo "1" >> /tmp/install_errors.txt
+                    # Show error details
+                    uv pip install "$line" 2>&1 | grep -E "(error|Error|ERROR|failed|Failed|FAILED)" | head -3 | sed 's/^/      /' || true
+                fi
+            done < /tmp/to_install.txt
+            INSTALL_ERRORS=$(wc -l < /tmp/install_errors.txt 2>/dev/null | tr -d ' ' || echo "0")
+            rm -f /tmp/install_errors.txt
+        fi
+        
+        # Count newly installed packages before cleanup
+        NEWLY_INSTALLED=0
+        if [ -s /tmp/to_install.txt ]; then
+            NEWLY_INSTALLED=$(wc -l < /tmp/to_install.txt 2>/dev/null | tr -d ' ' || echo "0")
+        fi
+        
+        # Cleanup temp files
+        rm -f /tmp/installed_packages.txt /tmp/seen_packages.txt /tmp/to_install.txt
+        
+        if [ $INSTALL_ERRORS -eq 0 ]; then
+            if [ $SKIPPED_COUNT -gt 0 ]; then
+                echo "✅ All dependencies up to date ($SKIPPED_COUNT already installed, $NEWLY_INSTALLED newly installed)"
+            else
+                echo "✅ All dependencies installed successfully"
+            fi
+        else
+            echo "❌ Failed to install $INSTALL_ERRORS package(s)" >&2
+            exit 1
+        fi
+    else
+        if [ "$VENV_EXISTED" = "true" ]; then
+            echo "✅ All dependencies already installed (requirements.txt is empty)"
+        else
+            echo "⚠️  requirements.txt is empty, skipping dependency installation"
+        fi
+    fi
+else
+    echo "⚠️  No requirements.txt found, skipping dependency installation"
+fi
+
+# Fix torch bug (https://github.com/pytorch/pytorch/issues/111469)
+PYTHON_VERSION=$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')
+export LD_LIBRARY_PATH=$HOME/slurmray-server/.venv/lib/python$PYTHON_VERSION/site-packages/nvidia/nvjitlink/lib:$LD_LIBRARY_PATH
+
+
+# Run server
+python -u slurmray_server.py

slurmray/assets/slurmray_server_template.py

@@ -0,0 +1,28 @@
+from slurmray.RayLauncher import RayLauncher
+
+if __name__ == "__main__":
+    # Note: This template creates a RayLauncher instance that runs on the cluster.
+    # The function execution is handled by spython.py which loads the serialized function.
+    # With the refactored API, we don't need to call cluster() here because:
+    # 1. The function is already serialized (func_source.py, func.pkl, args.pkl)
+    # 2. The actual execution is done via sbatch -> spython.py
+    # The RayLauncher instance here is just used to detect cluster mode and submit the job.
+    
+    cluster = RayLauncher(
+        project_name="server",
+        modules={{MODULES}},
+        node_nbr={{NODE_NBR}},
+        use_gpu={{USE_GPU}},
+        memory={{MEMORY}},
+        max_running_time={{MAX_RUNNING_TIME}},
+        server_run=False,
+        server_ssh=None,
+        server_username=None,
+    )
+    
+    # The job execution is handled by the backend when cluster.run() is called.
+    # Since we're on the cluster, the backend will detect cluster=True and submit via sbatch.
+    # The actual function execution happens in spython.py.
+    # Note: With the refactored API, this would need a dummy function to work.
+    # However, this template should be refactored to not need this call at all.
+    # For now, we skip execution here as it's handled by the backend infrastructure.

slurmray/assets/spython_template.py

@@ -0,0 +1,113 @@
+import ray
+import dill
+import os
+import sys
+
+PROJECT_PATH = {{PROJECT_PATH}}
+
+# Add the project path to the python path
+sys.path.append(PROJECT_PATH)
+
+# Add editable package source directories to sys.path
+# This handles packages uploaded from editable installs (e.g., Poetry projects)
+# Check for src/ directory (common Poetry src/ layout)
+# For flat layout, PROJECT_PATH is already in sys.path, so packages at root are importable
+src_path = os.path.join(PROJECT_PATH, "src")
+if os.path.exists(src_path) and os.path.isdir(src_path):
+    if src_path not in sys.path:
+        sys.path.insert(0, src_path)  # Insert at beginning for priority
+
+# Suppress Ray FutureWarning about accelerator visible devices
+os.environ.setdefault("RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO", "0")
+
+# Start the ray cluster
+ray.init({{LOCAL_MODE}})
+
+# Load the function
+# Read the serialization method used
+serialization_method = "dill_pickle"  # Default
+method_file = os.path.join(PROJECT_PATH, "serialization_method.txt")
+if os.path.exists(method_file):
+    with open(method_file, "r") as f:
+        serialization_method = f.read().strip()
+
+if serialization_method == "source_extraction":
+    # Load from source extraction
+    try:
+        func_source_path = os.path.join(PROJECT_PATH, "func_source.py")
+        func_name_path = os.path.join(PROJECT_PATH, "func_name.txt")
+
+        if not os.path.exists(func_source_path) or not os.path.exists(func_name_path):
+            raise FileNotFoundError(
+                "Source files missing: func_source.py or func_name.txt not found"
+            )
+
+        import importlib.util
+
+        spec = importlib.util.spec_from_file_location("func_module", func_source_path)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+
+        with open(func_name_path, "r") as f:
+            func_name = f.read().strip()
+
+        func = getattr(module, func_name)
+        print(f"Loaded function '{func_name}' from source extraction.")
+    except Exception as e:
+        print(f"❌ Error loading function from source extraction: {e}")
+        print("Falling back to dill pickle...")
+        import traceback
+
+        traceback.print_exc()
+        with open(os.path.join(PROJECT_PATH, "func.pkl"), "rb") as f:
+            func = dill.load(f)
+        print("✅ Loaded function from dill pickle (fallback).")
+else:
+    # Load from dill pickle (default method)
+    try:
+        func_pickle_path = os.path.join(PROJECT_PATH, "func.pkl")
+        if not os.path.exists(func_pickle_path):
+            raise FileNotFoundError(f"Pickle file not found: {func_pickle_path}")
+
+        with open(func_pickle_path, "rb") as f:
+            func = dill.load(f)
+        print(f"✅ Loaded function from dill pickle.")
+    except Exception as e:
+        print(f"❌ Error loading function from dill pickle: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+# Load the arguments
+with open(os.path.join(PROJECT_PATH, "args.pkl"), "rb") as f:
+    args = dill.load(f)
+
+# Run the function
+try:
+    result = func(**args)
+except Exception as e:
+    print(f"Error executing function: {e}")
+    import traceback
+
+    traceback.print_exc()
+    sys.exit(1)
+
+# Write the result
+result_path = os.path.join(PROJECT_PATH, "result.pkl")
+try:
+    with open(result_path, "wb") as f:
+        dill.dump(result, f)
+    print(f"Result written to {result_path}")
+except Exception as e:
+    print(f"Error writing result: {e}")
+    import traceback
+
+    traceback.print_exc()
+    sys.exit(1)
+
+# Stop ray
+try:
+    ray.shutdown()
+except Exception as e:
+    print(f"Warning: Error shutting down Ray: {e}")