nntool 1.3.0__cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

nntool/slurm/function.py

@@ -0,0 +1,173 @@
+import submitit
+import copy
+
+from submitit import Job
+from typing import Any, Callable, Literal, Tuple, Union, Dict, List, Optional
+from .config import SlurmConfig
+from .csrc import _SlurmFunction
+
+
+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,
+    ) -> "SlurmFunction":
+        """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).
+
+        :param submit_fn: function to be submitted to Slurm, defaults to None
+        :param default_submit_fn_args: default args for submit_fn, defaults to ()
+        :param default_submit_fn_kwargs: default known word args for submit_fn, defaults to {}
+        :return: the wrapped submit function with configured slurm paramters
+        """
+        self.engine = _SlurmFunction(
+            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.
+
+        :return: 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.
+
+        :return: 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
+
+        :param slurm_config: SlurmConfig, the slurm configuration dataclass, defaults to None
+        :param slurm_params_kwargs: extra slurm arguments for the slurm configuration, defaults to {}
+        :param slurm_submit_kwargs: extra slurm arguments for `srun` or `sbatch`, defaults to {}
+        :param slurm_task_kwargs: extra arguments for the setting of distributed task, defaults to {}
+        :param 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
+        :return: 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
+
+        :param slurm_config: SlurmConfig, the slurm configuration dataclass
+        :return: 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.
+
+        :raises Exception: if the submit_fn is not set up
+        :return: Slurm Job or the return value of the submit_fn
+        """
+        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__``.
+
+        :raises Exception: if the submit_fn is not set up
+        :return: 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) -> List[Job]:
+        """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.
+
+        :raises Exception: if the submit_fn is not set up
+        :return: 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.
+
+        :param jobs: dependent jobs
+        :param condition: run condition, defaults to "afterok"
+        :return: 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: Tuple[Job]) -> "SlurmFunction":
+        """Mark the function should be executed after the provided slurm jobs have been done.
+
+        :return: the function itself
+        """
+        return self.on_condition(jobs, "afterok")
+
+    def afterany(self, *jobs: Tuple[Job]) -> "SlurmFunction":
+        """Mark the function should be executed after any one of the provided slurm jobs has been done.
+
+        :return: the function itself
+        """
+        return self.on_condition(jobs, "afterany")
+
+    def afternotok(self, *jobs: Tuple[Job]) -> "SlurmFunction":
+        """Mark the function should be executed after any one of the provided slurm jobs has been failed.
+
+        :return: the function itself
+        """
+        return self.on_condition(jobs, "afternotok")

nntool/slurm/task.py

@@ -0,0 +1,231 @@
+import os
+import shlex
+import shutil
+import submitit
+
+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) -> 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: Union[Callable[[str, str], bool], Callable[[str], bool]],
+    exclude_fn: Union[Callable[[str, str], bool], Callable[[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: Union[
+        Callable[[str, str], bool], Callable[[str], bool]
+    ] = _is_py_or_dockerfile,
+    exclude_fn: Union[
+        Callable[[str, str], bool], Callable[[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:
+    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):
+        if not self.verbose:
+            return
+
+        print(msg)
+
+    def command(self) -> str:
+        raise NotImplementedError
+
+    def checkpoint(self):
+        print("checkpointing")
+        return submitit.helpers.DelayedSubmission(self)
+
+
+@dataclass
+class DistributedTaskConfig:
+    num_processes: Union[int, str] = "$nntool_num_processes"
+    num_machines: Union[int, str] = "$nntool_num_machines"
+    machine_rank: Union[int, str] = "$nntool_machine_rank"
+    main_process_ip: str = "$nntool_main_process_ip"
+    main_process_port: Union[int, str] = "$nntool_main_process_port"
+
+    def export_bash(self, output_folder: str):
+        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):
+    """Ref:
+    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()
+        self.dist_env = None
+
+    def dist_set_up(self):
+        self.log("running task on slurm")
+        self.log("exporting PyTorch distributed environment variables")
+
+        # prepare enviroment variables
+        dist_env = submitit.helpers.TorchDistributedEnvironment().export()
+
+        # other setup
+        env_setup = {
+            # "NCCL_DEBUG": "info",
+            # "CUDA_LAUNCH_BLOCKING": "1",
+        }
+
+        # set CUDA visible devices if slurm has scheduled GPUs otherwise use all GPUs (without setting
+        # CUDA_VISIBLE_DEVICES)
+        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)
+
+        # 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:
+        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.dist_set_up()
+
+        # job environment
+        job_env = submitit.helpers.JobEnvironment()
+
+        # concrete run command
+        cmd = self.command()
+
+        # export distributed environment variables
+        if self.dist_env.local_rank == 0:
+            print(f"running command: {cmd}")
+            if self.slurm_config.mode == "slurm":
+                try:
+                    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
+            else:
+                return os.system(cmd)
+
+        return 0

nntool/slurm/wrap.py

@@ -0,0 +1,210 @@
+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()``.
+
+    **Example**
+
+    Here's an example of how to use this function:
+
+    .. code-block:: python
+
+        @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
+
+    :param submit_fn: the function to be run on slurm
+    :return: the function to be run on slurm
+    """
+    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.
+
+    #### 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.
+
+    :param ArgsType: the experiment arguments type, which should be a dataclass (it
+                     mush have a slurm field defined by `slurm_key`)
+    :param slurm_key: the key of the slurm field in the ArgsType, defaults to "slurm"
+    :param parser: the parser for the arguments, defaults to "tyro"
+    :param slurm_config: SlurmConfig, the slurm configuration dataclass
+    :param slurm_params_kwargs: extra slurm arguments for the slurm configuration, defaults to {}
+    :param slurm_submit_kwargs: extra slurm arguments for `srun` or `sbatch`, defaults to {}
+    :param slurm_task_kwargs: extra arguments for the setting of distributed task, defaults to {}
+    :param extra_args: extra arguments for the parser
+    :param extra_kwargs: extra keyword arguments for the parser
+    :return: decorator function with main entry
+    """
+    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_distributed_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 the distributed job. This decorator should be used for the distributed job only and as the program entry. The decorated function is non-blocking in the mode of `slurm`, while other modes cause blocking.
+
+    #### 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.
+
+    :param ArgsType: the experiment arguments type, which should be a dataclass (it
+                     mush have a slurm field defined by `slurm_key`)
+    :param slurm_key: the key of the slurm field in the ArgsType, defaults to "slurm"
+    :param parser: the parser for the arguments, defaults to "tyro"
+    :param slurm_config: SlurmConfig, the slurm configuration dataclass
+    :param slurm_params_kwargs: extra slurm arguments for the slurm configuration, defaults to {}
+    :param slurm_submit_kwargs: extra slurm arguments for `srun` or `sbatch`, defaults to {}
+    :param slurm_task_kwargs: extra arguments for the setting of distributed task, defaults to {}
+    :param extra_args: extra arguments for the parser
+    :param extra_kwargs: extra keyword arguments for the parser
+    :return: decorator function with main entry
+    """
+    warn(
+        "`slurm_distributed_launcher` has been deprecated. Please use `slurm_launcher` instead, which supports both distributed and non-distributed job (controlled by `use_distributed_env` in slurm field).",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    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_distributed_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.
+    ```
+    @slurm_function
+    def run_in_slurm(*args, **kwargs):
+        pass
+
+    job = run_in_slurm(slurm_config)(*args, **kwargs)
+    ```
+    The decorated function `submit_fn` is non-blocking now. To block and get the return value, you can call `job.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:
+        """Update the slurm configuration for the slurm function.
+
+        #### 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.
+
+        :param slurm_config: SlurmConfig, the slurm configuration dataclass
+        :param slurm_params_kwargs: extra slurm arguments for the slurm configuration, defaults to {}
+        :param slurm_submit_kwargs: extra slurm arguments for `srun` or `sbatch`, defaults to {}
+        :param slurm_task_kwargs: extra arguments for the setting of distributed task, defaults to {}
+        :param 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
+        :return: the wrapped submit function with configured slurm paramters
+        """
+        warn(
+            "`slurm_function` has been deprecated. Please use `slurm_fn` instead, which supports both distributed and non-distributed job (controlled by `use_distributed_env` in slurm field).",
+            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

nntool/train/__init__.py

@@ -0,0 +1,25 @@
+import importlib.metadata as importlib_metadata
+import importlib.util
+
+
+def is_torch_available():
+    package_exists = importlib.util.find_spec("torch") is not None
+
+    # Check we're not importing a "torch" directory somewhere but the actual library by
+    # trying to grab the version
+    if package_exists:
+        try:
+            _ = importlib_metadata.metadata("torch")
+            return True
+        except importlib_metadata.PackageNotFoundError:
+            return False
+
+
+if is_torch_available():
+    from .trainer import BaseTrainer
+else:
+    # Inherits from a dummy `object` if torch is not available, so that python
+    # succeeds to import this file.
+    # BaseTrainer abstraction code will never inherit this dummy object as it checks if
+    # torch is available.
+    from builtins import object as BaseTrainer