sl-shared-assets 1.0.0rc14__py3-none-any.whl → 1.0.0rc15__py3-none-any.whl

sl_shared_assets/server/job.py

@@ -0,0 +1,140 @@
+"""This module provides the core Job class, used as the starting point for all SLURM-managed job executed on lab compute
+server(s). Specifically, the Job class acts as a wrapper around the SLURM configuration and specific logic of each
+job. During runtime, Server class interacts with input job objects to manage their transfer and execution on the
+remote servers."""
+
+# noinspection PyProtectedMember
+from pathlib import Path
+import datetime
+
+from simple_slurm import Slurm  # type: ignore
+
+
+class Job:
+    """Aggregates the data of a single SLURM-managed job to be executed on the Sun lab BioHPC cluster.
+
+    This class provides the API for constructing any server-side job in the Sun lab. Internally, it wraps an instance
+    of a Slurm class to package the job data into the format expected by the SLURM job manager. All jobs managed by this
+    class instance should be submitted to an initialized Server class 'submit_job' method to be executed on the server.
+
+    Notes:
+        The initialization method of the class contains the arguments for configuring the SLURM and Conda environments
+        used by the job. Do not submit additional SLURM or Conda commands via the 'add_command' method, as this may
+        produce unexpected behavior.
+
+        Each job can be conceptualized as a sequence of shell instructions to execute on the remote compute server. For
+        the lab, that means that the bulk of the command consists of calling various CLIs exposed by data processing or
+        analysis pipelines, installed in the Conda environment on the server. Other than that, the job contains commands
+        for activating the target conda environment and, in some cases, doing other preparatory or cleanup work. The
+        source code of a 'remote' job is typically identical to what a human operator would type in a 'local' terminal
+        to run the same job on their PC.
+
+        A key feature of server-side jobs is that they are executed on virtual machines managed by SLURM. Since the
+        server has a lot more compute and memory resources than likely needed by individual jobs, each job typically
+        requests a subset of these resources. Upon being executed, SLURM creates an isolated environment with the
+        requested resources and runs the job in that environment.
+
+        Since all jobs are expected to use the CLIs from python packages (pre)installed on the BioHPC server, make sure
+        that the target environment is installed and configured before submitting jobs to the server. See notes in
+        ReadMe to learn more about configuring server-side conda environments.
+
+    Args:
+        job_name: The descriptive name of the SLURM job to be created. Primarily, this name is used in terminal
+            printouts to identify the job to human operators.
+        output_log: The absolute path to the .txt file on the processing server, where to store the standard output
+            data of the job.
+        error_log: The absolute path to the .txt file on the processing server, where to store the standard error
+            data of the job.
+        working_directory: The absolute path to the directory where temporary job files will be stored. During runtime,
+            classes from this library use that directory to store files such as the job's shell script. All such files
+            are automatically removed from the directory at the end of a non-errors runtime.
+        conda_environment: The name of the conda environment to activate on the server before running the job logic. The
+            environment should contain the necessary Python packages and CLIs to support running the job's logic.
+        cpus_to_use: The number of CPUs to use for the job.
+        ram_gb: The amount of RAM to allocate for the job, in Gigabytes.
+        time_limit: The maximum time limit for the job, in minutes. If the job is still running at the end of this time
+            period, it will be forcibly terminated. It is highly advised to always set adequate maximum runtime limits
+            to prevent jobs from hogging the server in case of runtime or algorithm errors.
+
+    Attributes:
+        remote_script_path: Stores the path to the script file relative to the root of the remote server that runs the
+            command.
+        job_id: Stores the unique job identifier assigned by the SLURM manager to this job, when it is accepted for
+            execution. This field initialized to None and is overwritten by the Server class that submits the job.
+        job_name: Stores the descriptive name of the SLURM job.
+        _command: Stores the managed SLURM command object.
+    """
+
+    def __init__(
+        self,
+        job_name: str,
+        output_log: Path,
+        error_log: Path,
+        working_directory: Path,
+        conda_environment: str,
+        cpus_to_use: int = 10,
+        ram_gb: int = 10,
+        time_limit: int = 60,
+    ) -> None:
+        # Resolves the paths to the remote (server-side) .sh script file. This is the path where the job script
+        # will be stored on the server, once it is transferred by the Server class instance.
+        self.remote_script_path = str(working_directory.joinpath(f"{job_name}.sh"))
+
+        # Defines additional arguments used by the Server class that executed the job.
+        self.job_id: str | None = None  # This is set by the Server that submits the job.
+        self.job_name: str = job_name  # Also stores the job name to support more informative terminal prints
+
+        # Builds the slurm command object filled with configuration information
+        self._command: Slurm = Slurm(
+            cpus_per_task=cpus_to_use,
+            job_name=job_name,
+            output=str(output_log),
+            error=str(error_log),
+            mem=f"{ram_gb}G",
+            time=datetime.timedelta(minutes=time_limit),
+        )
+
+        # Conda shell initialization commands
+        self._command.add_cmd("eval $(conda shell.bash hook)")
+        self._command.add_cmd("conda init bash")
+
+        # Activates the target conda environment for the command.
+        self._command.add_cmd(f"conda activate {conda_environment}")
+
+    def __repr__(self) -> str:
+        """Returns the string representation of the Job instance."""
+        return f"Job(name={self.job_name}, id={self.job_id})"
+
+    def add_command(self, command: str) -> None:
+        """Adds the input command string to the end of the managed SLURM job command list.
+
+        This method is a wrapper around simple_slurm's 'add_cmd' method. It is used to iteratively build the shell
+        command sequence of the job.
+
+        Args:
+            command: The command string to add to the command list, e.g.: 'python main.py --input 1'.
+        """
+
+        self._command.add_cmd(command)
+
+    @property
+    def command_script(self) -> str:
+        """Translates the managed job data into a shell-script-writable string and returns it to caller.
+
+        This method is used by the Server class to translate the job into the format that can be submitted to and
+        executed on the remote compute server. Do not call this method manually unless you know what you are doing.
+        The returned string is safe to dump into a .sh (shell script) file and move to the BioHPC server for execution.
+        """
+
+        # Appends the command to clean up (remove) the temporary script file after processing runtime is over
+        self._command.add_cmd(f"rm -f {self.remote_script_path}")
+
+        # Translates the command to string format
+        script_content = str(self._command)
+
+        # Replaces escaped $ (/$) with $. This is essential, as without this correction, things like conda
+        # initialization would not work as expected.
+        fixed_script_content = script_content.replace("\\$", "$")
+
+        # Returns the script content to caller as a string
+        return fixed_script_content

sl_shared_assets/server/job.pyi

@@ -0,0 +1,94 @@
+from pathlib import Path
+
+from _typeshed import Incomplete
+from simple_slurm import Slurm
+
+class Job:
+    """Aggregates the data of a single SLURM-managed job to be executed on the Sun lab BioHPC cluster.
+
+    This class provides the API for constructing any server-side job in the Sun lab. Internally, it wraps an instance
+    of a Slurm class to package the job data into the format expected by the SLURM job manager. All jobs managed by this
+    class instance should be submitted to an initialized Server class 'submit_job' method to be executed on the server.
+
+    Notes:
+        The initialization method of the class contains the arguments for configuring the SLURM and Conda environments
+        used by the job. Do not submit additional SLURM or Conda commands via the 'add_command' method, as this may
+        produce unexpected behavior.
+
+        Each job can be conceptualized as a sequence of shell instructions to execute on the remote compute server. For
+        the lab, that means that the bulk of the command consists of calling various CLIs exposed by data processing or
+        analysis pipelines, installed in the Conda environment on the server. Other than that, the job contains commands
+        for activating the target conda environment and, in some cases, doing other preparatory or cleanup work. The
+        source code of a 'remote' job is typically identical to what a human operator would type in a 'local' terminal
+        to run the same job on their PC.
+
+        A key feature of server-side jobs is that they are executed on virtual machines managed by SLURM. Since the
+        server has a lot more compute and memory resources than likely needed by individual jobs, each job typically
+        requests a subset of these resources. Upon being executed, SLURM creates an isolated environment with the
+        requested resources and runs the job in that environment.
+
+        Since all jobs are expected to use the CLIs from python packages (pre)installed on the BioHPC server, make sure
+        that the target environment is installed and configured before submitting jobs to the server. See notes in
+        ReadMe to learn more about configuring server-side conda environments.
+
+    Args:
+        job_name: The descriptive name of the SLURM job to be created. Primarily, this name is used in terminal
+            printouts to identify the job to human operators.
+        output_log: The absolute path to the .txt file on the processing server, where to store the standard output
+            data of the job.
+        error_log: The absolute path to the .txt file on the processing server, where to store the standard error
+            data of the job.
+        working_directory: The absolute path to the directory where temporary job files will be stored. During runtime,
+            classes from this library use that directory to store files such as the job's shell script. All such files
+            are automatically removed from the directory at the end of a non-errors runtime.
+        conda_environment: The name of the conda environment to activate on the server before running the job logic. The
+            environment should contain the necessary Python packages and CLIs to support running the job's logic.
+        cpus_to_use: The number of CPUs to use for the job.
+        ram_gb: The amount of RAM to allocate for the job, in Gigabytes.
+        time_limit: The maximum time limit for the job, in minutes. If the job is still running at the end of this time
+            period, it will be forcibly terminated. It is highly advised to always set adequate maximum runtime limits
+            to prevent jobs from hogging the server in case of runtime or algorithm errors.
+
+    Attributes:
+        remote_script_path: Stores the path to the script file relative to the root of the remote server that runs the
+            command.
+        job_id: Stores the unique job identifier assigned by the SLURM manager to this job, when it is accepted for
+            execution. This field initialized to None and is overwritten by the Server class that submits the job.
+        job_name: Stores the descriptive name of the SLURM job.
+        _command: Stores the managed SLURM command object.
+    """
+
+    remote_script_path: Incomplete
+    job_id: str | None
+    job_name: str
+    _command: Slurm
+    def __init__(
+        self,
+        job_name: str,
+        output_log: Path,
+        error_log: Path,
+        working_directory: Path,
+        conda_environment: str,
+        cpus_to_use: int = 10,
+        ram_gb: int = 10,
+        time_limit: int = 60,
+    ) -> None: ...
+    def __repr__(self) -> str:
+        """Returns the string representation of the Job instance."""
+    def add_command(self, command: str) -> None:
+        """Adds the input command string to the end of the managed SLURM job command list.
+
+        This method is a wrapper around simple_slurm's 'add_cmd' method. It is used to iteratively build the shell
+        command sequence of the job.
+
+        Args:
+            command: The command string to add to the command list, e.g.: 'python main.py --input 1'.
+        """
+    @property
+    def command_script(self) -> str:
+        """Translates the managed job data into a shell-script-writable string and returns it to caller.
+
+        This method is used by the Server class to translate the job into the format that can be submitted to and
+        executed on the remote compute server. Do not call this method manually unless you know what you are doing.
+        The returned string is safe to dump into a .sh (shell script) file and move to the BioHPC server for execution.
+        """

sl_shared_assets/server/server.py

@@ -0,0 +1,213 @@
+"""This module provides the tools for working with the Sun lab BioHPC cluster. Specifically, the classes from this
+module establish an API for submitting jobs to the shared data processing cluster (managed via SLURM) and monitoring
+the running job status. All lab processing and analysis pipelines use this interface for accessing shared compute
+resources.
+"""
+
+import time
+from pathlib import Path
+import tempfile
+from dataclasses import dataclass
+
+import paramiko
+
+# noinspection PyProtectedMember
+from simple_slurm import Slurm  # type: ignore
+from paramiko.client import SSHClient
+from ataraxis_base_utilities import LogLevel, console
+from ataraxis_data_structures import YamlConfig
+
+from .job import Job
+
+
+def generate_server_credentials(
+    output_directory: Path, username: str, password: str, host: str = "cbsuwsun.biohpc.cornell.edu"
+) -> None:
+    """Generates a new server_credentials.yaml file under the specified directory, using input information.
+
+    This function provides a convenience interface for generating new BioHPC server credential files. Generally, this is
+    only used when setting up new host-computers in the lab.
+    """
+    ServerCredentials(username=username, password=password, host=host).to_yaml(
+        file_path=output_directory.joinpath("server_credentials.yaml")
+    )
+
+
+@dataclass()
+class ServerCredentials(YamlConfig):
+    """This class stores the hostname and credentials used to log into the BioHPC cluster to run Sun lab processing
+    pipelines.
+
+    Primarily, this is used as part of the sl-experiment library runtime to start data processing once it is
+    transferred to the BioHPC server during preprocessing. However, the same file can be used together with the Server
+    class API to run any computation jobs on the lab's BioHPC server.
+    """
+
+    username: str = "YourNetID"
+    """The username to use for server authentication."""
+    password: str = "YourPassword"
+    """The password to use for server authentication."""
+    host: str = "cbsuwsun.biohpc.cornell.edu"
+    """The hostname or IP address of the server to connect to."""
+
+
+class Server:
+    """Encapsulates access to the Sun lab BioHPC processing server.
+
+    This class provides the API that allows accessing the BioHPC server to create and submit various SLURM-managed jobs
+    to the server. It functions as the central interface used by all processing pipelines in the lab to execute costly
+    data processing on the server.
+
+    Notes:
+        All lab processing pipelines expect the data to be stored on the server and all processing logic to be packaged
+        and installed into dedicated conda environments on the server.
+
+        This class assumes that the target server has SLURM job manager installed and accessible to the user whose
+        credentials are used to connect to the server as part of this class instantiation.
+
+    Args:
+        credentials_path: The path to the locally stored .yaml file that contains the server hostname and access
+            credentials.
+
+    Attributes:
+        _open: Tracks whether the connection to the server is open or not.
+        _client: Stores the initialized SSHClient instance used to interface with the server.
+    """
+
+    def __init__(self, credentials_path: Path) -> None:
+        # Tracker used to prevent __del__ from classing stop() for a partially initialized class.
+        self._open: bool = False
+
+        # Loads the credentials from the provided .yaml file
+        self._credentials: ServerCredentials = ServerCredentials.from_yaml(credentials_path)  # type: ignore
+
+        # Establishes the SSH connection to the specified processing server. At most, attempts to connect to the server
+        # 30 times before terminating with an error
+        attempt = 0
+        while True:
+            console.echo(
+                f"Trying to connect to {self._credentials.host} (attempt {attempt}/30)...", level=LogLevel.INFO
+            )
+            try:
+                self._client: SSHClient = paramiko.SSHClient()
+                self._client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+                self._client.connect(
+                    self._credentials.host, username=self._credentials.username, password=self._credentials.password
+                )
+                console.echo(f"Connected to {self._credentials.host}", level=LogLevel.SUCCESS)
+                break
+            except paramiko.AuthenticationException:
+                message = (
+                    f"Authentication failed when connecting to {self._credentials.host} using "
+                    f"{self._credentials.username} user."
+                )
+                console.error(message, RuntimeError)
+                raise RuntimeError
+            except:
+                if attempt == 30:
+                    message = f"Could not connect to {self._credentials.host} after 30 attempts. Aborting runtime."
+                    console.error(message, RuntimeError)
+                    raise RuntimeError
+
+                console.echo(
+                    f"Could not SSH to {self._credentials.host}, retrying after a 2-second delay...",
+                    level=LogLevel.WARNING,
+                )
+                attempt += 1
+                time.sleep(2)
+
+    def __del__(self) -> None:
+        """If the instance is connected to the server, terminates the connection before the instance is destroyed."""
+        self.close()
+
+    def submit_job(self, job: Job) -> Job:
+        """Submits the input job to the managed BioHPC server via SLURM job manager.
+
+        This method submits various jobs for execution via SLURM-managed BioHPC cluster. As part of its runtime, the
+        method translates the Job object into the shell script, moves the script to the target working directory on
+        the server, and instructs the server to execute the shell script (via SLURM).
+
+        Args:
+            job: The Job object that contains all job data.
+
+        Returns:
+            The job object whose 'job_id' attribute had been modified with the job ID, if the job was successfully
+            submitted.
+
+        Raises:
+            RuntimeError: If job submission to the server fails.
+        """
+
+        # Generates a temporary shell script on the local machine. Uses tempfile to automatically remove the
+        # local script as soon as it is uploaded to the server.
+        with tempfile.TemporaryDirectory() as temp_dir:
+            local_script_path = Path(temp_dir).joinpath(f"{job.job_name}.sh")
+            fixed_script_content = job.command_script
+
+            # Creates a temporary script file locally and dumps translated command data into the file
+            with open(local_script_path, "w") as f:
+                f.write(fixed_script_content)
+
+            # Uploads the command script to the server
+            sftp = self._client.open_sftp()
+            sftp.put(localpath=local_script_path, remotepath=job.remote_script_path)
+            sftp.close()
+
+        # Makes the server-side script executable
+        self._client.exec_command(f"chmod +x {job.remote_script_path}")
+
+        # Submits the job to SLURM with sbatch and verifies submission state
+        job_output = self._client.exec_command(f"sbatch {job.remote_script_path}")[1].read().strip().decode()
+
+        # If batch_job is not in the output received from SLURM in response to issuing the submission command, raises an
+        # error.
+        if "Submitted batch job" not in job_output:
+            message = f"Failed to submit the {job.job_name} job to the BioHPC cluster."
+            console.error(message, RuntimeError)
+
+            # Fallback to appease mypy, should not be reachable
+            raise RuntimeError(message)
+
+        # Otherwise, extracts the job id assigned to the job by SLURM from the response and writes it to the processed
+        # Job object
+        job_id = job_output.split()[-1]
+        job.job_id = job_id
+        return job
+
+    def job_complete(self, job: Job) -> bool:
+        """Returns True if the job managed by the input Job instance has been completed or terminated its runtime due
+        to an error.
+
+        If the job is still running or is waiting inside the execution queue, returns False.
+
+        Args:
+            job: The Job object whose status needs to be checked.
+
+        Raises:
+            ValueError: If the input Job object does not contain a valid job_id, suggesting that it has not been
+                submitted to the server.
+        """
+
+        if job.job_id is None:
+            message = (
+                f"The input Job object for the job {job.job_name} does not contain a valid job_id. This indicates that "
+                f"the job has not been submitted to the server."
+            )
+            console.error(message, ValueError)
+
+            # This is here to appease mypy, it should not be reachable
+            raise ValueError(message)
+
+        if job.job_id not in self._client.exec_command(f"squeue -j {job.job_id}")[1].read().decode().strip():
+            return True
+        else:
+            return False
+
+    def close(self) -> None:
+        """Closes the SSH connection to the server.
+
+        This method has to be called before destroying the class instance to ensure proper resource cleanup.
+        """
+        # Prevents closing already closed connections
+        if self._open:
+            self._client.close()

sl_shared_assets/server/server.pyi

@@ -0,0 +1,95 @@
+from pathlib import Path
+from dataclasses import dataclass
+
+from simple_slurm import Slurm as Slurm
+from paramiko.client import SSHClient as SSHClient
+from ataraxis_data_structures import YamlConfig
+
+from .job import Job as Job
+
+def generate_server_credentials(
+    output_directory: Path, username: str, password: str, host: str = "cbsuwsun.biohpc.cornell.edu"
+) -> None:
+    """Generates a new server_credentials.yaml file under the specified directory, using input information.
+
+    This function provides a convenience interface for generating new BioHPC server credential files. Generally, this is
+    only used when setting up new host-computers in the lab.
+    """
+@dataclass()
+class ServerCredentials(YamlConfig):
+    """This class stores the hostname and credentials used to log into the BioHPC cluster to run Sun lab processing
+    pipelines.
+
+    Primarily, this is used as part of the sl-experiment library runtime to start data processing once it is
+    transferred to the BioHPC server during preprocessing. However, the same file can be used together with the Server
+    class API to run any computation jobs on the lab's BioHPC server.
+    """
+
+    username: str = ...
+    password: str = ...
+    host: str = ...
+
+class Server:
+    """Encapsulates access to the Sun lab BioHPC processing server.
+
+    This class provides the API that allows accessing the BioHPC server to create and submit various SLURM-managed jobs
+    to the server. It functions as the central interface used by all processing pipelines in the lab to execute costly
+    data processing on the server.
+
+    Notes:
+        All lab processing pipelines expect the data to be stored on the server and all processing logic to be packaged
+        and installed into dedicated conda environments on the server.
+
+        This class assumes that the target server has SLURM job manager installed and accessible to the user whose
+        credentials are used to connect to the server as part of this class instantiation.
+
+    Args:
+        credentials_path: The path to the locally stored .yaml file that contains the server hostname and access
+            credentials.
+
+    Attributes:
+        _open: Tracks whether the connection to the server is open or not.
+        _client: Stores the initialized SSHClient instance used to interface with the server.
+    """
+
+    _open: bool
+    _credentials: ServerCredentials
+    _client: SSHClient
+    def __init__(self, credentials_path: Path) -> None: ...
+    def __del__(self) -> None:
+        """If the instance is connected to the server, terminates the connection before the instance is destroyed."""
+    def submit_job(self, job: Job) -> Job:
+        """Submits the input job to the managed BioHPC server via SLURM job manager.
+
+        This method submits various jobs for execution via SLURM-managed BioHPC cluster. As part of its runtime, the
+        method translates the Job object into the shell script, moves the script to the target working directory on
+        the server, and instructs the server to execute the shell script (via SLURM).
+
+        Args:
+            job: The Job object that contains all job data.
+
+        Returns:
+            The job object whose 'job_id' attribute had been modified with the job ID, if the job was successfully
+            submitted.
+
+        Raises:
+            RuntimeError: If job submission to the server fails.
+        """
+    def job_complete(self, job: Job) -> bool:
+        """Returns True if the job managed by the input Job instance has been completed or terminated its runtime due
+        to an error.
+
+        If the job is still running or is waiting inside the execution queue, returns False.
+
+        Args:
+            job: The Job object whose status needs to be checked.
+
+        Raises:
+            ValueError: If the input Job object does not contain a valid job_id, suggesting that it has not been
+                submitted to the server.
+        """
+    def close(self) -> None:
+        """Closes the SSH connection to the server.
+
+        This method has to be called before destroying the class instance to ensure proper resource cleanup.
+        """

sl_shared_assets/suite2p/__init__.py

@@ -0,0 +1,8 @@
+"""This package provides the configuration classes used by the Sun lab maintained version of the suite2p library
+(sl-suite2p package, https://github.com/Sun-Lab-NBB/suite2p) to process brain activity data within and across sessions
+(days)."""
+
+from .multi_day import MultiDayS2PConfiguration
+from .single_day import SingleDayS2PConfiguration
+
+__all__ = ["MultiDayS2PConfiguration", "SingleDayS2PConfiguration"]

sl_shared_assets/suite2p/__init__.pyi

@@ -0,0 +1,4 @@
+from .multi_day import MultiDayS2PConfiguration as MultiDayS2PConfiguration
+from .single_day import SingleDayS2PConfiguration as SingleDayS2PConfiguration
+
+__all__ = ["MultiDayS2PConfiguration", "SingleDayS2PConfiguration"]