nntool 2.0.0rc0__py3-none-any.whl

nntool/slurm/function.py

@@ -0,0 +1,209 @@
+import submitit
+import copy
+
+from submitit import Job
+from typing import Any, Callable, Literal, Tuple, Union, Dict, List, Optional
+from .config import SlurmConfig
+from .core import SlurmBackend
+
+
+class SlurmFunction:
+    """The function for the slurm job, which can be used for distributed or non-distributed job (controlled by `use_distributed_env` in the slurm dataclass)."""
+
+    def __init__(
+        self,
+        submit_fn: Callable[..., Any],
+        default_submit_fn_args: Optional[Tuple[Any]] = None,
+        default_submit_fn_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """A slurm function for the slurm job, which can be used for distributed or non-distributed job (controlled by `use_distributed_env` in the slurm dataclass).
+
+        Args:
+            submit_fn: function to be submitted to Slurm, defaults to None
+            default_submit_fn_args: default args for submit_fn, defaults to ()
+            default_submit_fn_kwargs: default known word args for submit_fn, defaults to {}
+
+        Returns:
+            the wrapped submit function with configured slurm paramters
+        """
+        self.engine = SlurmBackend(submit_fn, default_submit_fn_args, default_submit_fn_kwargs)
+
+    def __create_copy(self) -> "SlurmFunction":
+        return copy.copy(self)
+
+    def is_configured(self) -> bool:
+        """Whether the slurm function has been configured.
+
+        Returns:
+            True if the slurm function has been configured, False otherwise
+        """
+        return self.engine.is_configured()
+
+    def is_distributed(self) -> bool:
+        """Whether the slurm function is distributed.
+
+        Returns:
+            True if the slurm function is distributed, False otherwise
+        """
+        return self.engine.is_distributed()
+
+    def get_executor(
+        self,
+    ) -> submitit.AutoExecutor:
+        return self.engine.get_executor()
+
+    def configure(
+        self,
+        slurm_config: SlurmConfig,
+        slurm_params_kwargs: Optional[Dict[str, str]] = None,
+        slurm_submit_kwargs: Optional[Dict[str, str]] = None,
+        slurm_task_kwargs: Optional[Dict[str, str]] = None,
+        system_argv: Optional[List[str]] = None,
+        pack_code_include_fn: Optional[Callable[[str, str], bool]] = None,
+        pack_code_exclude_fn: Optional[Callable[[str, str], bool]] = None,
+    ) -> "SlurmFunction":
+        """Update the slurm configuration for the slurm function. A slurm function for the slurm job, which can be used for distributed or non-distributed job (controlled by `use_distributed_env` in the slurm dataclass).
+
+        **Exported Distributed Enviroment Variables**
+
+        - ``NNTOOL_SLURM_HAS_BEEN_SET_UP`` is a special environment variable to indicate that the slurm has been set up.
+        - After the set up, the distributed job will be launched and the following variables are exported:
+            - ``num_processes``: int
+            - ``num_machines``: int
+            - ``machine_rank``: int
+            - ``main_process_ip``: str
+            - ``main_process_port``: int
+
+        Args:
+            slurm_config: SlurmConfig, the slurm configuration dataclass, defaults to None
+            slurm_params_kwargs: extra slurm arguments for the slurm configuration, defaults to {}
+            slurm_submit_kwargs: extra slurm arguments for `srun` or `sbatch`, defaults to {}
+            slurm_task_kwargs: extra arguments for the setting of distributed task, defaults to {}
+            system_argv: the system arguments for the second launch in the distributed task (by default it will use the current system arguments `sys.argv[1:]`), defaults to None
+
+        Returns:
+            a new copy with configured slurm parameters
+        """
+        configured_slurm_function = self.__create_copy()
+        configured_slurm_function.engine = self.engine.configure(
+            slurm_config,
+            slurm_params_kwargs,
+            slurm_submit_kwargs,
+            slurm_task_kwargs,
+            system_argv,
+            pack_code_include_fn,
+            pack_code_exclude_fn,
+        )
+        return configured_slurm_function
+
+    def __getitem__(self, slurm_config: Union[Dict[str, Any], Tuple[Any], Any]) -> "SlurmFunction":
+        """Instantiate the slurm configuration for the slurm function. A slurm function for the slurm job, which can be used for distributed or non-distributed job (controlled by `use_distributed_env` in the slurm dataclass).
+
+        **Exported Distributed Enviroment Variables**
+
+        - ``NNTOOL_SLURM_HAS_BEEN_SET_UP`` is a special environment variable to indicate that the slurm has been set up.
+        - After the set up, the distributed job will be launched and the following variables are exported:
+            - ``num_processes``: int
+            - ``num_machines``: int
+            - ``machine_rank``: int
+            - ``main_process_ip``: str
+            - ``main_process_port``: int
+
+        Args:
+            slurm_config: SlurmConfig, the slurm configuration dataclass
+
+        Returns:
+            the wrapped submit function with configured slurm paramters
+        """
+        configured_slurm_function = self.__create_copy()
+        configured_slurm_function.engine = self.engine[slurm_config]
+        return configured_slurm_function
+
+    def __call__(self, *submit_fn_args, **submit_fn_kwargs) -> Union[Job, Any]:
+        """Run the submit_fn with the given arguments and keyword arguments. The function is non-blocking in the mode of `slurm`, while other modes cause blocking. If there is no given arguments or keyword arguments, the default arguments and keyword arguments will be used.
+
+        Args:
+            submit_fn_args: arguments for the submit_fn
+            submit_fn_kwargs: keyword arguments for the submit_fn
+
+        Returns:
+            Slurm Job or the return value of the submit_fn, depends on the submit mode
+
+        Raises:
+            Exception: if the submit_fn is not set up
+        """
+        return self.engine(*submit_fn_args, **submit_fn_kwargs)
+
+    def submit(self, *submit_fn_args, **submit_fn_kwargs) -> Union[Job, Any]:
+        """An alias function to ``__call__``.
+
+        Args:
+            submit_fn_args: arguments for the submit_fn
+            submit_fn_kwargs: keyword arguments for the submit_fn
+
+        Raises:
+            Exception: if the submit_fn is not set up
+
+        Returns:
+            Slurm Job or the return value of the submit_fn
+        """
+        return self(*submit_fn_args, **submit_fn_kwargs)
+
+    def map_array(
+        self, *submit_fn_args, **submit_fn_kwargs
+    ) -> Union[Job[Any], List[Job[Any]], Any]:
+        """Run the submit_fn with the given arguments and keyword arguments. The function is non-blocking in the mode of `slurm`, while other modes cause blocking. If there is no given arguments or keyword arguments, the default arguments and keyword arguments will be used.
+
+        Args:
+            submit_fn_args: arguments for the submit_fn
+            submit_fn_kwargs: keyword arguments for the submit_fn
+
+        Raises:
+            Exception: if the submit_fn is not set up
+
+        Returns:
+            Slurm Job or the return value of the submit_fn
+        """
+        return self.engine.map_array(*submit_fn_args, **submit_fn_kwargs)
+
+    def on_condition(
+        self,
+        jobs: Union[Job, List[Job], Tuple[Job]],
+        condition: Literal["afterany", "afterok", "afternotok"] = "afterok",
+    ) -> "SlurmFunction":
+        """Mark this job should be executed after the provided slurm jobs have been done. This function allows combining different conditions by multiple calling.
+
+        Args:
+            jobs: dependent jobs
+            condition: run condition, defaults to "afterok"
+
+        Returns:
+            the function itself
+        """
+        configured_slurm_function = self.__create_copy()
+        configured_slurm_function.engine = self.engine.on_condition(jobs, condition)
+        return configured_slurm_function
+
+    def afterok(self, *jobs: Job) -> "SlurmFunction":
+        """Mark the function should be executed after the provided slurm jobs have been done.
+
+        Returns:
+            the new slurm function with the condition
+        """
+        return self.on_condition(list(jobs), "afterok")
+
+    def afterany(self, *jobs: Job) -> "SlurmFunction":
+        """Mark the function should be executed after any one of the provided slurm jobs has been done.
+
+        Returns:
+            the new slurm function with the condition
+        """
+        return self.on_condition(list(jobs), "afterany")
+
+    def afternotok(self, *jobs: Job) -> "SlurmFunction":
+        """Mark the function should be executed after any one of the provided slurm jobs has been failed.
+
+        Returns:
+            the new slurm function with the condition
+        """
+        return self.on_condition(list(jobs), "afternotok")

nntool/slurm/parser/__init__.py

@@ -0,0 +1,6 @@
+from .parse import parse_from_cli
+
+
+__all__ = [
+    "parse_from_cli",
+]

nntool/slurm/parser/parse.py

@@ -0,0 +1,22 @@
+import tyro
+
+from typing import Union, Callable, Type, Any
+
+
+def parse_from_cli(
+    ArgsType: Type[Any],
+    parser: Union[str, Callable] = "tyro",
+    *args,
+    **kwargs,
+) -> Any:
+    # parse with built-in parser or custom parser function
+    parser_fn = None
+    if isinstance(parser, str):
+        if parser == "tyro":
+            parser_fn = tyro.cli
+        else:
+            raise ValueError(f"Parser `{parser}` is not supported.")
+    else:
+        parser_fn = parser
+    args: ArgsType = parser_fn(ArgsType, *args, **kwargs)
+    return args

nntool/slurm/task.py

@@ -0,0 +1,300 @@
+import os
+import shlex
+import shutil
+import submitit
+import subprocess
+
+from pathlib import Path
+from typing import Union, Generator, Callable
+from dataclasses import dataclass
+from .config import SlurmConfig
+from .accelerator.utils import nvidia_smi_gpu_memory_stats_str
+
+WANDB_DIRS = ("wandb", ".wandb")
+
+
+def _is_py_or_dockerfile(path: str, root: str) -> bool:
+    file = os.path.basename(path)
+    return file.endswith(".py") or file.startswith("Dockerfile")
+
+
+def include_code_files(path: str, root: str, code_ext: list[str]):
+    file = os.path.basename(path)
+    return any(file.endswith(ext) for ext in code_ext) or file.startswith("Dockerfile")
+
+
+def exclude_code_folders(path: str, root: str, code_folders: list[str]):
+    return any(
+        os.path.relpath(path, root).startswith(code_folders + os.sep)
+        for code_folders in code_folders
+    )
+
+
+def exclude_wandb_fn(path: str, root: str) -> bool:
+    return any(
+        os.path.relpath(path, root).startswith(wandb_dir + os.sep) for wandb_dir in WANDB_DIRS
+    )
+
+
+def filtered_dir(
+    root: str,
+    include_fn: Callable[[str, str], bool],
+    exclude_fn: Callable[[str, str], bool],
+) -> Generator[str, None, None]:
+    """Simple generator to walk a directory."""
+
+    for dirpath, _, files in os.walk(root):
+        for fname in files:
+            file_path = os.path.join(dirpath, fname)
+            if include_fn(file_path, root) and not exclude_fn(file_path, root):
+                yield file_path
+
+
+def pack_code_files(
+    root: str,
+    target_root: str,
+    include_fn: Callable[[str, str], bool] = _is_py_or_dockerfile,
+    exclude_fn: Callable[[str, str], bool] = exclude_wandb_fn,
+):
+    root = os.path.abspath(root)
+    code_root = Path(os.path.abspath(root))
+    code_target = Path(os.path.abspath(target_root)) / "code"
+    if not code_root.exists():
+        raise ValueError(f"Code root {code_root} does not exist.")
+    if not code_target.exists():
+        code_target.mkdir(parents=True)
+
+    for file_path in filtered_dir(root, include_fn, exclude_fn):
+        save_name = os.path.relpath(file_path, root)
+        sub_file_path, file_name = os.path.split(save_name)
+        sub_file_full_path = code_target / sub_file_path
+        if not sub_file_full_path.exists():
+            sub_file_full_path.mkdir(parents=True)
+        shutil.copy(file_path, sub_file_full_path / file_name)
+
+    return code_target
+
+
+def reconstruct_command_line(argv):
+    # Quote each argument that needs special handling (like spaces or shell characters)
+    # and join them with spaces to form the command line
+    return " ".join(shlex.quote(arg) for arg in argv)
+
+
+class Task:
+    """The base class for all tasks that will be run on Slurm. Especially useful for
+    distributed tasks that need to set up the distributed environment variables.
+
+    Args:
+        argv (list[str]): the command line arguments to run the task. This will be passed to the command method to reconstruct the command line.
+        slurm_config (SlurmConfig): the Slurm configuration to use for the task.
+        verbose (bool, optional): whether to print verbose output. Defaults to False.
+    """
+
+    def __init__(self, argv: list[str], slurm_config: SlurmConfig, verbose: bool = False):
+        self.argv = argv
+        self.slurm_config = slurm_config
+        self.verbose = verbose
+
+    def log(self, msg: str):
+        """Log a message to the console if verbose is enabled.
+
+        Args:
+            msg (str): the message to log.
+        """
+        if not self.verbose:
+            return
+        print(msg)
+
+    def command(self) -> str:
+        """Return the command to run the task. This method should be implemented by
+        subclasses to return the actual command line to run the task.
+
+        Raises:
+            NotImplementedError: If the method is not implemented by the subclass.
+
+        Returns:
+            str: the command to run the task.
+        """
+        raise NotImplementedError
+
+    def checkpoint(self):
+        """Return a checkpoint for the task. This is used to save the state of the task."""
+        return submitit.helpers.DelayedSubmission(self)
+
+
+@dataclass
+class DistributedTaskConfig:
+    """Configuration for distributed tasks. This is used to set up the distributed environment
+    variables for PyTorch distributed training.
+
+    Args:
+        num_processes (int): The total number of processes to run across all machines.
+        num_machines (int): The number of machines to run the task on.
+        machine_rank (int): The rank of the current machine in the distributed setup.
+        main_process_ip (str): The IP address of the main process (rank 0)
+            in the distributed setup.
+        main_process_port (int): The port of the main process (rank 0)
+            in the distributed setup.
+    """
+
+    # The number of processes to run in total across all machines.
+    num_processes: Union[int, str] = "$nntool_num_processes"
+
+    # The number of machines to run the task on.
+    num_machines: Union[int, str] = "$nntool_num_machines"
+
+    # The rank of the current machine in the distributed setup.
+    machine_rank: Union[int, str] = "$nntool_machine_rank"
+
+    # The IP address of the main process (rank 0) in the distributed setup.
+    main_process_ip: str = "$nntool_main_process_ip"
+
+    # The port of the main process (rank 0) in the distributed setup.
+    main_process_port: Union[int, str] = "$nntool_main_process_port"
+
+    def export_bash(self, output_folder: str):
+        """Export the distributed environment variables to a bash script.
+        This script can be sourced to set the environment variables for the distributed task.
+
+        Args:
+            output_folder (str): the folder to save the bash script to.
+        """
+        lines = ["#!/bin/bash"]
+        for k, v in self.__dict__.items():
+            lines.append(f"export nntool_{k}={v}")
+        with open(os.path.join(output_folder, "nntool_distributed_env.sh"), "w") as f:
+            f.write("\n".join(lines))
+
+
+class PyTorchDistributedTask(Task):
+    """A task that runs on Slurm and sets up the PyTorch distributed environment variables. It runs the command locally
+    if in other modes.
+
+    Args:
+        launch_cmd (str): The command to launch the task.
+        argv (list[str]): The command line arguments for the task.
+        slurm_config (SlurmConfig): The Slurm configuration to use for the task.
+        verbose (bool, optional): _description_. Defaults to False.
+
+    References:
+        https://github.com/huggingface/accelerate/issues/1239
+        https://github.com/yuvalkirstain/PickScore/blob/main/trainer/slurm_scripts/slurm_train.py
+        https://github.com/facebookincubator/submitit/pull/1703
+    """
+
+    def __init__(
+        self,
+        launch_cmd: str,
+        argv: list[str],
+        slurm_config: SlurmConfig,
+        verbose: bool = False,
+        **env_setup_kwargs,
+    ):
+        super().__init__(argv, slurm_config, verbose)
+        self.launch_cmd = launch_cmd
+        self.env_setup_kwargs = env_setup_kwargs
+
+        # to be set up in the dist_set_up method
+        self.dist_args: DistributedTaskConfig = DistributedTaskConfig()
+        self.dist_env: Union[None, submitit.helpers.TorchDistributedEnvironment] = None
+
+    def set_up_dist_env(self):
+        """Set up the distributed environment variables for PyTorch distributed training."""
+
+        self.log("running task on slurm")
+        self.log("exporting PyTorch distributed environment variables")
+
+        # prepare enviroment variables
+        dist_env = submitit.helpers.TorchDistributedEnvironment().export(
+            set_cuda_visible_devices=False
+        )
+
+        # other setup
+        env_setup = {}
+
+        # set CUDA visible devices if slurm has scheduled GPUs otherwise use all GPUs (without setting
+        # CUDA_VISIBLE_DEVICES)
+        if self.slurm_config.mode == "slurm":
+            env_setup.update(
+                {"CUDA_VISIBLE_DEVICES": os.environ["SLURM_JOB_GPUS"]}
+                if "SLURM_JOB_GPUS" in os.environ
+                else {}
+            )
+
+        # other environment variables set by the user
+        env_setup.update(self.env_setup_kwargs)
+        self.log(f"Env setup: {env_setup}")
+
+        # update environment variables
+        os.environ.update(**env_setup)
+
+        self.log(nvidia_smi_gpu_memory_stats_str())
+        self.log(f"Master: {dist_env.master_addr}:{dist_env.master_port}")
+        self.log(f"Rank: {dist_env.rank}")
+        self.log(f"World size: {dist_env.world_size}")
+        self.log(f"Local rank: {dist_env.local_rank}")
+        self.log(f"Local world size: {dist_env.local_world_size}")
+        self.log(
+            f"Local rank {dist_env.local_rank}: CUDA_VISIBLE_DEVICES {os.environ.get('CUDA_VISIBLE_DEVICES', 'all')}"
+        )
+
+        # set distributed arguments
+        num_processes = (
+            self.slurm_config.tasks_per_node
+            * self.slurm_config.processes_per_task
+            * self.slurm_config.num_of_node
+        )
+        machine_rank = dist_env.rank // self.slurm_config.tasks_per_node
+        self.dist_args = DistributedTaskConfig(
+            num_processes=num_processes,
+            num_machines=self.slurm_config.num_of_node,
+            machine_rank=machine_rank,
+            main_process_ip=dist_env.master_addr,
+            main_process_port=dist_env.master_port,
+        )
+        self.dist_env = dist_env
+
+        return self.dist_args, self.dist_env
+
+    def command(self) -> str:
+        """Return the command to run the task. This method should be implemented by
+        subclasses to return the actual command line to run the task.
+
+        Returns:
+            str: the command to run the task.
+        """
+        cmd = self.launch_cmd.format(**self.dist_args.__dict__)
+        cmd += " " + reconstruct_command_line(self.argv)
+        return cmd
+
+    def __call__(self):
+        # Set up distributed environment
+        self.set_up_dist_env()
+
+        # Job environment
+        job_env = submitit.helpers.JobEnvironment()
+
+        # Concrete run command
+        cmd = self.command()
+
+        # Export distributed environment variables only the global rank 0 process will run the command
+        if self.dist_env.rank == 0:
+            print(f"running command: {cmd}")
+            if self.slurm_config.mode == "slurm":
+                try:
+                    # Export distributed environment variables to a bash script
+                    # and the fn will be launched after the job is scheduled
+                    self.dist_args.export_bash(shlex.quote(str(job_env.paths.folder)))
+                except Exception as e:
+                    print(f"failed to export distributed environment variables: {e}")
+                    return -1
+            elif self.slurm_config.mode == "local":
+                cmd_list = shlex.split(cmd)
+                return subprocess.Popen(cmd_list)
+            else:
+                # If not on slurm mode, we can just run the command directly
+                # This is useful for local testing or when running on a single machine
+                return os.system(cmd)
+
+        return 0

nntool/slurm/wrap.py

@@ -0,0 +1,148 @@
+import sys
+
+from warnings import warn
+from typing import Any, Callable, Type, Union, Dict, List
+
+from .config import SlurmConfig
+from .function import SlurmFunction
+from .parser import parse_from_cli
+
+
+def slurm_fn(
+    submit_fn: Callable,
+) -> SlurmFunction:
+    """A decorator to wrap a function to be run on slurm. The function decorated by this decorator should be launched on the way below. The decorated function `submit_fn` is non-blocking now. To block and get the return value, you can call ``job.result()``.
+
+    Args:
+        submit_fn: the function to be run on slurm
+
+    Returns:
+        the function to be run on slurm
+
+    Example:
+        >>> @slurm_fn
+        ... def run_on_slurm(a, b):
+        ...     return a + b
+        >>> slurm_config = SlurmConfig(
+        ...     mode="slurm",
+        ...     partition="PARTITION",
+        ...     job_name="EXAMPLE",
+        ...     tasks_per_node=1,
+        ...     cpus_per_task=8,
+        ...     mem="1GB",
+        ... )
+        >>> job = run_on_slurm[slurm_config](1, b=2)
+        >>> result = job.result()  # block and get the result
+    """
+    slurm_fn = SlurmFunction(submit_fn=submit_fn)
+
+    return slurm_fn
+
+
+def slurm_launcher(
+    ArgsType: Type[Any],
+    parser: Union[str, Callable] = "tyro",
+    slurm_key: str = "slurm",
+    slurm_params_kwargs: dict = {},
+    slurm_submit_kwargs: dict = {},
+    slurm_task_kwargs: dict = {},
+    *extra_args,
+    **extra_kwargs,
+) -> Callable[[Callable[..., Any]], SlurmFunction]:
+    """A slurm launcher decorator for distributed or non-distributed job (controlled by `use_distributed_env` in slurm field). This decorator should be used as the program entry. The decorated function is non-blocking in the mode of `slurm`, while other modes cause blocking.
+
+    Args:
+        ArgsType: the experiment arguments type, which should be a dataclass (it
+                         mush have a slurm field defined by `slurm_key`)
+        slurm_key: the key of the slurm field in the ArgsType, defaults to "slurm"
+        parser: the parser for the arguments, defaults to "tyro"
+        slurm_config: SlurmConfig, the slurm configuration dataclass
+        slurm_params_kwargs: extra slurm arguments for the slurm configuration, defaults to {}
+        slurm_submit_kwargs: extra slurm arguments for `srun` or `sbatch`, defaults to {}
+        slurm_task_kwargs: extra arguments for the setting of distributed task, defaults to {}
+        extra_args: extra arguments for the parser
+        extra_kwargs: extra keyword arguments for the parser
+
+    Returns:
+        decorator function with main entry
+
+    Exported Distributed Enviroment Variables:
+        1. NNTOOL_SLURM_HAS_BEEN_SET_UP is a special environment variable to indicate that the slurm has been set up.
+        2. After the set up, the distributed job will be launched and the following variables are exported: num_processes: int, num_machines: int, machine_rank: int, main_process_ip: str, main_process_port: int.
+    """
+    argv = list(sys.argv[1:])
+    args = parse_from_cli(ArgsType, parser, *extra_args, **extra_kwargs)
+
+    # check if args have slurm field
+    if not hasattr(args, slurm_key):
+        raise ValueError(
+            f"ArgsType should have a field named `{slurm_key}` to use `slurm_launcher` decorator."
+        )
+    slurm_config: SlurmConfig = getattr(args, slurm_key)
+
+    def decorator(
+        submit_fn: Callable[..., Any],
+    ) -> SlurmFunction:
+        return SlurmFunction(
+            submit_fn=submit_fn,
+            default_submit_fn_args=(args,),
+        ).configure(
+            slurm_config,
+            slurm_params_kwargs,
+            slurm_submit_kwargs,
+            slurm_task_kwargs,
+            system_argv=argv,
+        )
+
+    return decorator
+
+
+def slurm_function(
+    submit_fn: Callable,
+):
+    """
+    A decorator to annoate a function to be run in slurm. The function decorated by this decorator should be launched in the way below.
+
+    Deprecated:
+        This function is deprecated and will be removed in future versions. Please use `slurm_fn` instead.
+
+    Example:
+        >>> @slurm_function
+        ... def run_on_slurm(a, b):
+        ...     return a + b
+        >>> slurm_config = SlurmConfig(
+        ...     mode="slurm",
+        ...     partition="PARTITION",
+        ...     job_name="EXAMPLE",
+        ...     tasks_per_node=1,
+        ...     cpus_per_task=8,
+        ...     mem="1GB",
+        ... )
+        >>> job = run_on_slurm(slurm_config)(1, b=2)
+        >>> result = job.result()  # block and get the result
+    """
+
+    def wrapper(
+        slurm_config: SlurmConfig,
+        slurm_params_kwargs: Dict[str, Any] = {},
+        slurm_submit_kwargs: Dict[str, Any] = {},
+        slurm_task_kwargs: Dict[str, Any] = {},
+        system_argv: Union[List[str], None] = None,
+    ) -> SlurmFunction:
+        warn(
+            "`slurm_function` has been deprecated. Please use `slurm_fn` instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        slurm_fn = SlurmFunction(
+            submit_fn=submit_fn,
+        ).configure(
+            slurm_config,
+            slurm_params_kwargs,
+            slurm_submit_kwargs,
+            slurm_task_kwargs,
+            system_argv,
+        )
+        return slurm_fn
+
+    return wrapper