slurmray 6.0.4__py3-none-any.whl

slurmray/backend/local.py

@@ -0,0 +1,124 @@
+import os
+import sys
+import time
+import subprocess
+import dill
+from typing import Any
+
+from slurmray.backend.base import ClusterBackend
+
+class LocalBackend(ClusterBackend):
+    """Backend for local execution (no scheduler)"""
+
+    def run(self, cancel_old_jobs: bool = True, wait: bool = True) -> Any:
+        """Run the job locally"""
+        self.logger.info("No cluster detected, running locally...")
+        
+        self._write_python_script()
+        
+        # Determine log file path
+        # Use job name logic similar to Slurm if possible, or just project name
+        job_name = f"{self.launcher.project_name}_local"
+        log_path = os.path.join(self.launcher.project_path, f"{job_name}.log")
+        
+        with open(log_path, "w") as log_file:
+            process = subprocess.Popen(
+                [sys.executable, os.path.join(self.launcher.project_path, "spython.py")],
+                stdout=log_file,
+                stderr=subprocess.STDOUT
+            )
+        
+        if not wait:
+            self.logger.info(f"Local job started asynchronously. Logs: {log_path} PID: {process.pid}")
+            return str(process.pid)
+
+        # Wait for result
+        while not os.path.exists(os.path.join(self.launcher.project_path, "result.pkl")):
+            time.sleep(0.25)
+            # Check if process died
+            if process.poll() is not None:
+                # Process finished but no result?
+                if not os.path.exists(os.path.join(self.launcher.project_path, "result.pkl")):
+                    # Check logs
+                    with open(log_path, "r") as f:
+                        print(f.read())
+                    raise RuntimeError("Local process succeeded but no result.pkl found (or failed).")
+
+        with open(os.path.join(self.launcher.project_path, "result.pkl"), "rb") as f:
+            result = dill.load(f)
+
+        return result
+
+    def get_result(self, job_id: str) -> Any:
+        """Get result from local execution (result.pkl)"""
+        # job_id is just project_name descriptor here
+        result_path = os.path.join(self.launcher.project_path, "result.pkl")
+        if os.path.exists(result_path):
+            try:
+                with open(result_path, "rb") as f:
+                     return dill.load(f)
+            except Exception:
+                return None
+        return None
+
+    def get_logs(self, job_id: str) -> Any:
+        """Get logs from local execution"""
+        log_path = os.path.join(self.launcher.project_path, f"{job_id}.log")
+        if os.path.exists(log_path):
+            with open(log_path, "r") as f:
+                # Yield lines
+                for line in f:
+                    yield line.strip()
+        else:
+            yield "Log file not found."
+
+    def cancel(self, job_id: str):
+        """Cancel a job (kill local process)"""
+        self.logger.info(f"Canceling local job {job_id}...")
+        try:
+            # Check if job_id looks like a PID (digits)
+            if job_id and job_id.isdigit():
+                pid = int(job_id)
+                self.logger.info(f"Killing process {pid}...")
+                os.kill(pid, signal.SIGTERM)
+            else:
+                # Fallback to pkill by pattern if not a PID (old behavior or synchronous fallback)
+                cmd = ["pkill", "-f", f"{self.launcher.project_path}/spython.py"]
+                subprocess.run(cmd, check=False)
+            
+            self.logger.info("Local process killed.")
+        except Exception as e:
+            self.logger.warning(f"Failed to kill local process: {e}")
+
+    def _write_python_script(self):
+        """Write the python script that will be executed by the job"""
+        self.logger.info("Writing python script...")
+
+        # Remove the old python script
+        for file in os.listdir(self.launcher.project_path):
+            if file.endswith(".py"):
+                os.remove(os.path.join(self.launcher.project_path, file))
+
+        # Write the python script
+        with open(
+            os.path.join(self.launcher.module_path, "assets", "spython_template.py"),
+            "r",
+        ) as f:
+            text = f.read()
+
+        text = text.replace("{{PROJECT_PATH}}", f'"{self.launcher.project_path}"')
+        # Local mode doesn't need special address usually, or 'auto' is fine.
+        # Original code:
+        # if self.cluster or self.server_run:
+        #    local_mode = ...
+        # else:
+        #    local_mode = "" (empty)
+        
+        local_mode = ""
+        text = text.replace(
+            "{{LOCAL_MODE}}",
+            local_mode,
+        )
+        with open(os.path.join(self.launcher.project_path, "spython.py"), "w") as f:
+            f.write(text)
+

slurmray/backend/remote.py

@@ -0,0 +1,191 @@
+import os
+import sys
+import time
+import paramiko
+import subprocess
+import logging
+from getpass import getpass
+from typing import Any, Optional, List
+
+from slurmray.backend.base import ClusterBackend
+
+
+class RemoteMixin(ClusterBackend):
+    """Mixin for remote execution via SSH"""
+
+    def __init__(self, launcher):
+        super().__init__(launcher)
+        self.ssh_client = None
+        self.job_id = None
+
+    def _connect(self):
+        """Establish SSH connection"""
+        if (
+            self.ssh_client
+            and self.ssh_client.get_transport()
+            and self.ssh_client.get_transport().is_active()
+        ):
+            return
+
+        connected = False
+        self.logger.info("Connecting to the remote server...")
+        self.ssh_client = paramiko.SSHClient()
+        self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+        while not connected:
+            try:
+                if self.launcher.server_password is None:
+                    # Password not provided: prompt interactively
+                    # (Credentials from .env or explicit parameters are already loaded in RayLauncher.__init__)
+                    try:
+                        self.launcher.server_password = getpass(
+                            "Enter your cluster password: "
+                        )
+                    except Exception:
+                        # Handle case where getpass fails (e.g. non-interactive terminal)
+                        pass
+
+                if self.launcher.server_password is None:
+                    raise ValueError(
+                        "No password provided and cannot prompt (non-interactive)"
+                    )
+
+                self.ssh_client.connect(
+                    hostname=self.launcher.server_ssh,
+                    username=self.launcher.server_username,
+                    password=self.launcher.server_password,
+                )
+                connected = True
+            except paramiko.ssh_exception.AuthenticationException:
+                self.launcher.server_password = None
+                self.logger.warning("Wrong password, please try again.")
+                # Only retry interactively if we failed
+                if not sys.stdin.isatty():
+                    raise  # Fail fast if non-interactive
+
+    def _push_file(
+        self, file_path: str, sftp: paramiko.SFTPClient, remote_base_dir: str
+    ):
+        """Push a file to the remote server"""
+        self.logger.info(
+            f"Pushing file {os.path.basename(file_path)} to the remote server..."
+        )
+
+        # Determine the path to the file
+        local_path = file_path
+        if not os.path.isabs(local_path):
+            local_path = os.path.join(self.launcher.pwd_path, local_path)
+        local_path = os.path.abspath(local_path)
+
+        local_path_from_pwd = os.path.relpath(local_path, self.launcher.pwd_path)
+        remote_path = os.path.join(remote_base_dir, local_path_from_pwd)
+
+        # Create the directory if not exists
+        stdin, stdout, stderr = self.ssh_client.exec_command(
+            f"mkdir -p '{os.path.dirname(remote_path)}'"
+        )
+        # Wait for command to finish
+        stdout.channel.recv_exit_status()
+
+        # Copy the file to the server
+        sftp.put(local_path, remote_path)
+
+    def _push_file_wrapper(
+        self,
+        rel_path: str,
+        sftp: paramiko.SFTPClient,
+        remote_base_dir: str,
+        ssh_client: paramiko.SSHClient = None,
+    ):
+        """
+        Wrapper to call _push_file with the correct signature for each backend.
+        """
+        import inspect
+
+        # Get the signature of _push_file for this backend
+        sig = inspect.signature(self._push_file)
+        param_count = len(sig.parameters)
+
+        if param_count == 4:  # SlurmBackend: (self, file_path, sftp, ssh_client)
+            if ssh_client is None:
+                ssh_client = self.ssh_client
+            self._push_file(rel_path, sftp, ssh_client)
+        else:  # RemoteMixin/DesiBackend: (self, file_path, sftp, remote_base_dir)
+            self._push_file(rel_path, sftp, remote_base_dir)
+
+    def _sync_local_files_incremental(
+        self,
+        sftp: paramiko.SFTPClient,
+        remote_base_dir: str,
+        local_files: List[str],
+        ssh_client: paramiko.SSHClient = None,
+    ):
+        """
+        Synchronize local files incrementally using hash comparison.
+        Only uploads files that have changed.
+        """
+        from slurmray.file_sync import FileHashManager, LocalFileSyncManager
+
+        # Use provided ssh_client or self.ssh_client
+        if ssh_client is None:
+            ssh_client = self.ssh_client
+
+        # Initialize sync manager
+        hash_manager = FileHashManager(self.launcher.pwd_path, self.logger)
+        sync_manager = LocalFileSyncManager(
+            self.launcher.pwd_path, hash_manager, self.logger
+        )
+
+        # Remote hash file path
+        remote_hash_file = os.path.join(
+            remote_base_dir, ".slogs", ".remote_file_hashes.json"
+        )
+
+        # Fetch remote hashes
+        remote_hashes = sync_manager.fetch_remote_hashes(ssh_client, remote_hash_file)
+
+        # Determine which files need uploading
+        files_to_upload = sync_manager.get_files_to_upload(local_files, remote_hashes)
+
+        if not files_to_upload:
+            self.logger.info("✅ All local files are up to date, no upload needed.")
+            return
+
+        self.logger.info(
+            f"📤 Uploading {len(files_to_upload)} modified/new file(s) out of {len(local_files)} total..."
+        )
+
+        # Upload files
+        uploaded_files = []
+        for rel_path in files_to_upload:
+            # Convert relative path to absolute
+            abs_path = os.path.join(self.launcher.pwd_path, rel_path)
+
+            # Handle directories (packages)
+            if os.path.isdir(abs_path):
+                # Upload all Python files in the directory recursively
+                for root, dirs, files in os.walk(abs_path):
+                    # Skip __pycache__
+                    dirs[:] = [d for d in dirs if d != "__pycache__"]
+                    for file in files:
+                        if file.endswith(".py"):
+                            file_path = os.path.join(root, file)
+                            file_rel = os.path.relpath(
+                                file_path, self.launcher.pwd_path
+                            )
+                            self._push_file_wrapper(
+                                file_rel, sftp, remote_base_dir, ssh_client
+                            )
+                            uploaded_files.append(file_rel)
+            else:
+                # Single file
+                self._push_file_wrapper(rel_path, sftp, remote_base_dir, ssh_client)
+                uploaded_files.append(rel_path)
+
+        # Update remote hashes after successful upload
+        sync_manager.update_remote_hashes(uploaded_files, remote_hashes)
+        sync_manager.save_remote_hashes_to_server(
+            ssh_client, remote_hash_file, remote_hashes
+        )
+
+        self.logger.info(f"✅ Successfully uploaded {len(uploaded_files)} file(s).")