nshtrainer 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

nshtrainer/_submit/session/unified.py

@@ -1,350 +0,0 @@
-import copy
-import logging
-import os
-import signal
-import subprocess
-from collections.abc import Callable, Mapping, Sequence
-from datetime import timedelta
-from pathlib import Path
-from typing import Any, Literal
-
-from typing_extensions import (
-    TypeAlias,
-    TypedDict,
-    TypeVar,
-    TypeVarTuple,
-    Unpack,
-    assert_never,
-)
-
-from . import lsf, slurm
-from ._output import SubmitOutput
-
-TArgs = TypeVarTuple("TArgs")
-_Path: TypeAlias = str | Path | os.PathLike
-
-log = logging.getLogger(__name__)
-
-
-class GenericJobKwargs(TypedDict, total=False):
-    name: str
-    """The name of the job."""
-
-    partition: str | Sequence[str]
-    """The partition or queue to submit the job to. Same as `queue`."""
-
-    queue: str | Sequence[str]
-    """The queue to submit the job to. Same as `partition`."""
-
-    qos: str
-    """
-    The quality of service to submit the job to.
-
-    This corresponds to the "--qos" option in sbatch (only for Slurm).
-    """
-
-    account: str
-    """The account (or project) to charge the job to. Same as `project`."""
-
-    project: str
-    """The project (or account) to charge the job to. Same as `account`."""
-
-    output_file: _Path
-    """
-    The file to write the job output to.
-
-    This corresponds to the "-o" option in bsub. If not specified, the output will be written to the default output file.
-    """
-
-    error_file: _Path
-    """
-    The file to write the job errors to.
-
-    This corresponds to the "-e" option in bsub. If not specified, the errors will be written to the default error file.
-    """
-
-    nodes: int
-    """The number of nodes to request."""
-
-    tasks_per_node: int
-    """The number of tasks to request per node."""
-
-    cpus_per_task: int
-    """The number of CPUs to request per task."""
-
-    gpus_per_task: int
-    """The number of GPUs to request per task."""
-
-    memory_mb: int
-    """The maximum memory for the job in MB."""
-
-    walltime: timedelta
-    """The maximum walltime for the job."""
-
-    email: str
-    """The email address to send notifications to."""
-
-    notifications: set[Literal["begin", "end"]]
-    """The notifications to send via email."""
-
-    setup_commands: Sequence[str]
-    """
-    The setup commands to run before the job.
-
-    These commands will be executed prior to everything else in the job script.
-    """
-
-    environment: Mapping[str, str]
-    """
-    The environment variables to set for the job.
-
-    These variables will be set prior to executing any commands in the job script.
-    """
-
-    command_prefix: str
-    """
-    A command to prefix the job command with.
-
-    This is used to add commands like `srun` or `jsrun` to the job command.
-    """
-
-    constraint: str | Sequence[str]
-    """
-    The constraint to request for the job. For SLRUM, this corresponds to the `--constraint` option. For LSF, this is unused.
-    """
-
-    signal: signal.Signals
-    """The signal that will be sent to the job when it is time to stop it."""
-
-    command_template: str
-    """
-    The template for the command to execute the helper script.
-
-    Default: `bash {script}`.
-    """
-
-    requeue_on_preempt: bool
-    """
-    Whether to requeue the job if it is preempted.
-
-    This corresponds to the "--requeue" option in sbatch (only for Slurm).
-    """
-
-    slurm_options: slurm.SlurmJobKwargs
-    """Additional keyword arguments for Slurm jobs."""
-
-    lsf_options: lsf.LSFJobKwargs
-    """Additional keyword arguments for LSF jobs."""
-
-
-Scheduler: TypeAlias = Literal["slurm", "lsf"]
-
-
-T = TypeVar("T", infer_variance=True)
-
-
-def _one_of(*fns: Callable[[], T | None]) -> T | None:
-    values = [value for fn in fns if (value := fn()) is not None]
-
-    # Only one (or zero) value should be set. If not, raise an error.
-    if len(set(values)) > 1:
-        raise ValueError(f"Multiple values set: {values}")
-
-    return next((value for value in values if value is not None), None)
-
-
-def _to_slurm(kwargs: GenericJobKwargs) -> slurm.SlurmJobKwargs:
-    slurm_kwargs: slurm.SlurmJobKwargs = {}
-    if (name := kwargs.get("name")) is not None:
-        slurm_kwargs["name"] = name
-    if (
-        account := _one_of(
-            lambda: kwargs.get("account"),
-            lambda: kwargs.get("project"),
-        )
-    ) is not None:
-        slurm_kwargs["account"] = account
-    if (
-        partition := _one_of(
-            lambda: kwargs.get("partition"),
-            lambda: kwargs.get("queue"),
-        )
-    ) is not None:
-        slurm_kwargs["partition"] = partition
-    if (qos := kwargs.get("qos")) is not None:
-        slurm_kwargs["qos"] = qos
-    if (output_file := kwargs.get("output_file")) is not None:
-        slurm_kwargs["output_file"] = output_file
-    if (error_file := kwargs.get("error_file")) is not None:
-        slurm_kwargs["error_file"] = error_file
-    if (walltime := kwargs.get("walltime")) is not None:
-        slurm_kwargs["time"] = walltime
-    if (memory_mb := kwargs.get("memory_mb")) is not None:
-        slurm_kwargs["memory_mb"] = memory_mb
-    if (nodes := kwargs.get("nodes")) is not None:
-        slurm_kwargs["nodes"] = nodes
-    if (tasks_per_node := kwargs.get("tasks_per_node")) is not None:
-        slurm_kwargs["ntasks_per_node"] = tasks_per_node
-    if (cpus_per_task := kwargs.get("cpus_per_task")) is not None:
-        slurm_kwargs["cpus_per_task"] = cpus_per_task
-    if (gpus_per_task := kwargs.get("gpus_per_task")) is not None:
-        slurm_kwargs["gpus_per_task"] = gpus_per_task
-    if (constraint := kwargs.get("constraint")) is not None:
-        slurm_kwargs["constraint"] = constraint
-    if (signal := kwargs.get("signal")) is not None:
-        slurm_kwargs["signal"] = signal
-    if (email := kwargs.get("email")) is not None:
-        slurm_kwargs["mail_user"] = email
-    if (notifications := kwargs.get("notifications")) is not None:
-        mail_type: list[slurm.MailType] = []
-        for notification in notifications:
-            match notification:
-                case "begin":
-                    mail_type.append("BEGIN")
-                case "end":
-                    mail_type.append("END")
-                case _:
-                    raise ValueError(f"Unknown notification type: {notification}")
-        slurm_kwargs["mail_type"] = mail_type
-    if (setup_commands := kwargs.get("setup_commands")) is not None:
-        slurm_kwargs["setup_commands"] = setup_commands
-    if (environment := kwargs.get("environment")) is not None:
-        slurm_kwargs["environment"] = environment
-    if (command_prefix := kwargs.get("command_prefix")) is not None:
-        slurm_kwargs["command_prefix"] = command_prefix
-    if (requeue_on_preempt := kwargs.get("requeue_on_preempt")) is not None:
-        slurm_kwargs["requeue"] = requeue_on_preempt
-    if (additional_kwargs := kwargs.get("slurm_options")) is not None:
-        slurm_kwargs.update(additional_kwargs)
-
-    return slurm_kwargs
-
-
-def _to_lsf(kwargs: GenericJobKwargs) -> lsf.LSFJobKwargs:
-    lsf_kwargs: lsf.LSFJobKwargs = {}
-    if (name := kwargs.get("name")) is not None:
-        lsf_kwargs["name"] = name
-    if (
-        account := _one_of(
-            lambda: kwargs.get("account"),
-            lambda: kwargs.get("project"),
-        )
-    ) is not None:
-        lsf_kwargs["project"] = account
-    if (
-        partition := _one_of(
-            lambda: kwargs.get("partition"),
-            lambda: kwargs.get("queue"),
-        )
-    ) is not None:
-        lsf_kwargs["queue"] = partition
-    if (output_file := kwargs.get("output_file")) is not None:
-        lsf_kwargs["output_file"] = output_file
-    if (error_file := kwargs.get("error_file")) is not None:
-        lsf_kwargs["error_file"] = error_file
-    if (walltime := kwargs.get("walltime")) is not None:
-        lsf_kwargs["walltime"] = walltime
-    if (memory_mb := kwargs.get("memory_mb")) is not None:
-        lsf_kwargs["memory_mb"] = memory_mb
-    if (nodes := kwargs.get("nodes")) is not None:
-        lsf_kwargs["nodes"] = nodes
-    if (tasks_per_node := kwargs.get("tasks_per_node")) is not None:
-        lsf_kwargs["rs_per_node"] = tasks_per_node
-    if (cpus_per_task := kwargs.get("cpus_per_task")) is not None:
-        lsf_kwargs["cpus_per_rs"] = cpus_per_task
-    if (gpus_per_task := kwargs.get("gpus_per_task")) is not None:
-        lsf_kwargs["gpus_per_rs"] = gpus_per_task
-    if (constraint := kwargs.get("constraint")) is not None:
-        log.warning(f'LSF does not support constraints, ignoring "{constraint=}".')
-    if (email := kwargs.get("email")) is not None:
-        lsf_kwargs["email"] = email
-    if (notifications := kwargs.get("notifications")) is not None:
-        if "begin" in notifications:
-            lsf_kwargs["notify_begin"] = True
-        if "end" in notifications:
-            lsf_kwargs["notify_end"] = True
-    if (setup_commands := kwargs.get("setup_commands")) is not None:
-        lsf_kwargs["setup_commands"] = setup_commands
-    if (environment := kwargs.get("environment")) is not None:
-        lsf_kwargs["environment"] = environment
-    if (command_prefix := kwargs.get("command_prefix")) is not None:
-        lsf_kwargs["command_prefix"] = command_prefix
-    if (signal := kwargs.get("signal")) is not None:
-        lsf_kwargs["signal"] = signal
-    if (requeue_on_preempt := kwargs.get("requeue_on_preempt")) is not None:
-        log.warning(
-            f'LSF does not support requeueing, ignoring "{requeue_on_preempt=}".'
-        )
-    if (additional_kwargs := kwargs.get("lsf_options")) is not None:
-        lsf_kwargs.update(additional_kwargs)
-
-    return lsf_kwargs
-
-
-def validate_kwargs(scheduler: Scheduler, kwargs: GenericJobKwargs) -> None:
-    match scheduler:
-        case "slurm":
-            _to_slurm(copy.deepcopy(kwargs))
-        case "lsf":
-            _to_lsf(copy.deepcopy(kwargs))
-        case _:
-            assert_never(scheduler)
-
-
-def to_array_batch_script(
-    scheduler: Scheduler,
-    dest: Path,
-    callable: Callable[[Unpack[TArgs]], Any],
-    args_list: Sequence[tuple[Unpack[TArgs]]],
-    /,
-    job_index_variable: str | None = None,
-    print_environment_info: bool = False,
-    python_command_prefix: str | None = None,
-    **kwargs: Unpack[GenericJobKwargs],
-) -> SubmitOutput:
-    job_index_variable_kwargs = {}
-    if job_index_variable is not None:
-        job_index_variable_kwargs["job_index_variable"] = job_index_variable
-    match scheduler:
-        case "slurm":
-            slurm_kwargs = _to_slurm(kwargs)
-            return slurm.to_array_batch_script(
-                dest,
-                callable,
-                args_list,
-                **job_index_variable_kwargs,
-                print_environment_info=print_environment_info,
-                python_command_prefix=python_command_prefix,
-                **slurm_kwargs,
-            )
-        case "lsf":
-            lsf_kwargs = _to_lsf(kwargs)
-            return lsf.to_array_batch_script(
-                dest,
-                callable,
-                args_list,
-                **job_index_variable_kwargs,
-                print_environment_info=print_environment_info,
-                python_command_prefix=python_command_prefix,
-                **lsf_kwargs,
-            )
-        case _:
-            assert_never(scheduler)
-
-
-def infer_current_scheduler() -> Scheduler:
-    # First, we check for `bsub` as it's much less common than `sbatch`.
-    try:
-        subprocess.check_output(["bsub", "-V"])
-        return "lsf"
-    except BaseException:
-        pass
-
-    # Next, we check for `sbatch` as it's the most common scheduler.
-    try:
-        subprocess.check_output(["sbatch", "--version"])
-        return "slurm"
-    except BaseException:
-        pass
-
-    raise RuntimeError("Could not determine the current scheduler.")

nshtrainer/config.py

@@ -1,289 +0,0 @@
-from collections.abc import Mapping, MutableMapping
-from typing import TYPE_CHECKING, Any, ClassVar
-
-from pydantic import BaseModel, ConfigDict
-from pydantic import Field as Field
-from pydantic import PrivateAttr as PrivateAttr
-from typing_extensions import deprecated, override
-
-from ._config.missing import MISSING, validate_no_missing_values
-from ._config.missing import AllowMissing as AllowMissing
-from ._config.missing import MissingField as MissingField
-
-_MutableMappingBase = MutableMapping[str, Any]
-if TYPE_CHECKING:
-    _MutableMappingBase = object
-
-
-_DraftConfigContextSentinel = object()
-
-
-class TypedConfig(BaseModel, _MutableMappingBase):
-    _is_draft_config: bool = PrivateAttr(default=False)
-    """
-    Whether this config is a draft config or not.
-
-    Draft configs are configs that are not yet fully validated.
-    They allow for a nicer API when creating configs, e.g.:
-
-        ```python
-        config = MyConfig.draft()
-
-        # Set some values
-        config.a = 10
-        config.b = "hello"
-
-        # Finalize the config
-        config = config.finalize()
-        ```
-    """
-
-    repr_diff_only: ClassVar[bool] = True
-    """
-    If `True`, the repr methods will only show values for fields that are different from the default.
-    """
-
-    MISSING: ClassVar[Any] = MISSING
-    """
-    Alias for the `MISSING` constant.
-    """
-
-    model_config: ClassVar[ConfigDict] = ConfigDict(
-        # By default, Pydantic will throw a warning if a field starts with "model_",
-        # so we need to disable that warning (beacuse "model_" is a popular prefix for ML).
-        protected_namespaces=(),
-        validate_assignment=True,
-        validate_return=True,
-        validate_default=True,
-        strict=True,
-        revalidate_instances="always",
-        arbitrary_types_allowed=True,
-        extra="ignore",
-        validation_error_cause=True,
-        use_attribute_docstrings=True,
-    )
-
-    def __draft_pre_init__(self):
-        """Called right before a draft config is finalized."""
-        pass
-
-    def __post_init__(self):
-        """Called after the final config is validated."""
-        pass
-
-    @classmethod
-    @deprecated("Use `model_validate` instead.")
-    def from_dict(cls, model_dict: Mapping[str, Any]):
-        return cls.model_validate(model_dict)
-
-    def model_deep_validate(self, strict: bool = True):
-        """
-        Validate the config and all of its sub-configs.
-
-        Args:
-            config: The config to validate.
-            strict: Whether to validate the config strictly.
-        """
-        config_dict = self.model_dump(round_trip=True)
-        config = self.model_validate(config_dict, strict=strict)
-
-        # Make sure that this is not a draft config
-        if config._is_draft_config:
-            raise ValueError("Draft configs are not valid. Call `finalize` first.")
-
-        return config
-
-    @classmethod
-    def draft(cls, **kwargs):
-        config = cls.model_construct_draft(**kwargs)
-        return config
-
-    def finalize(self, strict: bool = True):
-        # This must be a draft config, otherwise we raise an error
-        if not self._is_draft_config:
-            raise ValueError("Finalize can only be called on drafts.")
-
-        # First, we call `__draft_pre_init__` to allow the config to modify itself a final time
-        self.__draft_pre_init__()
-
-        # Then, we dump the config to a dict and then re-validate it
-        return self.model_deep_validate(strict=strict)
-
-    @override
-    def model_post_init(self, __context: Any) -> None:
-        super().model_post_init(__context)
-
-        # Call the `__post_init__` method if this is not a draft config
-        if __context is _DraftConfigContextSentinel:
-            return
-
-        self.__post_init__()
-
-        # After `_post_init__` is called, we perform the final round of validation
-        self.model_post_init_validate()
-
-    def model_post_init_validate(self):
-        validate_no_missing_values(self)
-
-    @classmethod
-    def model_construct_draft(cls, _fields_set: set[str] | None = None, **values: Any):
-        """
-        NOTE: This is a copy of the `model_construct` method from Pydantic's `Model` class,
-            with the following changes:
-            - The `model_post_init` method is called with the `_DraftConfigContext` context.
-            - The `_is_draft_config` attribute is set to `True` in the `values` dict.
-
-        Creates a new instance of the `Model` class with validated data.
-
-        Creates a new model setting `__dict__` and `__pydantic_fields_set__` from trusted or pre-validated data.
-        Default values are respected, but no other validation is performed.
-
-        !!! note
-            `model_construct()` generally respects the `model_config.extra` setting on the provided model.
-            That is, if `model_config.extra == 'allow'`, then all extra passed values are added to the model instance's `__dict__`
-            and `__pydantic_extra__` fields. If `model_config.extra == 'ignore'` (the default), then all extra passed values are ignored.
-            Because no validation is performed with a call to `model_construct()`, having `model_config.extra == 'forbid'` does not result in
-            an error if extra values are passed, but they will be ignored.
-
-        Args:
-            _fields_set: The set of field names accepted for the Model instance.
-            values: Trusted or pre-validated data dictionary.
-
-        Returns:
-            A new instance of the `Model` class with validated data.
-        """
-
-        values["_is_draft_config"] = True
-
-        m = cls.__new__(cls)
-        fields_values: dict[str, Any] = {}
-        fields_set = set()
-
-        for name, field in cls.model_fields.items():
-            if field.alias and field.alias in values:
-                fields_values[name] = values.pop(field.alias)
-                fields_set.add(name)
-            elif name in values:
-                fields_values[name] = values.pop(name)
-                fields_set.add(name)
-            elif not field.is_required():
-                fields_values[name] = field.get_default(call_default_factory=True)
-        if _fields_set is None:
-            _fields_set = fields_set
-
-        _extra: dict[str, Any] | None = None
-        if cls.model_config.get("extra") == "allow":
-            _extra = {}
-            for k, v in values.items():
-                _extra[k] = v
-        object.__setattr__(m, "__dict__", fields_values)
-        object.__setattr__(m, "__pydantic_fields_set__", _fields_set)
-        if not cls.__pydantic_root_model__:
-            object.__setattr__(m, "__pydantic_extra__", _extra)
-
-        if cls.__pydantic_post_init__:
-            m.model_post_init(_DraftConfigContextSentinel)
-            # update private attributes with values set
-            if (
-                hasattr(m, "__pydantic_private__")
-                and m.__pydantic_private__ is not None
-            ):
-                for k, v in values.items():
-                    if k in m.__private_attributes__:
-                        m.__pydantic_private__[k] = v
-
-        elif not cls.__pydantic_root_model__:
-            # Note: if there are any private attributes, cls.__pydantic_post_init__ would exist
-            # Since it doesn't, that means that `__pydantic_private__` should be set to None
-            object.__setattr__(m, "__pydantic_private__", None)
-
-        return m
-
-    @override
-    def __repr_args__(self):
-        # If `repr_diff_only` is `True`, we only show the fields that are different from the default.
-        if not self.repr_diff_only:
-            yield from super().__repr_args__()
-            return
-
-        # First, we get the default values for all fields.
-        default_values = self.model_construct_draft()
-
-        # Then, we compare the default values with the current values.
-        for k, v in super().__repr_args__():
-            if k is None:
-                yield k, v
-                continue
-
-            # If there is no default value or the value is different from the default, we yield it.
-            if not hasattr(default_values, k) or getattr(default_values, k) != v:
-                yield k, v
-                continue
-
-            # Otherwise, we can skip this field.
-
-    # region MutableMapping implementation
-    if not TYPE_CHECKING:
-        # This is mainly so the config can be used with lightning's hparams
-        #   transparently and without any issues.
-
-        @property
-        def _ll_dict(self):
-            return self.model_dump()
-
-        # We need to make sure every config class
-        #   is a MutableMapping[str, Any] so that it can be used
-        #   with lightning's hparams.
-        @override
-        def __getitem__(self, key: str):
-            # Key can be of the format "a.b.c"
-            #   so we need to split it into a list of keys.
-            [first_key, *rest_keys] = key.split(".")
-            value = self._ll_dict[first_key]
-
-            for key in rest_keys:
-                if isinstance(value, Mapping):
-                    value = value[key]
-                else:
-                    value = getattr(value, key)
-
-            return value
-
-        @override
-        def __setitem__(self, key: str, value: Any):
-            # Key can be of the format "a.b.c"
-            #   so we need to split it into a list of keys.
-            [first_key, *rest_keys] = key.split(".")
-            if len(rest_keys) == 0:
-                self._ll_dict[first_key] = value
-                return
-
-            # We need to traverse the keys until we reach the last key
-            #   and then set the value
-            current_value = self._ll_dict[first_key]
-            for key in rest_keys[:-1]:
-                if isinstance(current_value, Mapping):
-                    current_value = current_value[key]
-                else:
-                    current_value = getattr(current_value, key)
-
-            # Set the value
-            if isinstance(current_value, MutableMapping):
-                current_value[rest_keys[-1]] = value
-            else:
-                setattr(current_value, rest_keys[-1], value)
-
-        @override
-        def __delitem__(self, key: str):
-            # This is unsupported for this class
-            raise NotImplementedError
-
-        @override
-        def __iter__(self):
-            return iter(self._ll_dict)
-
-        @override
-        def __len__(self):
-            return len(self._ll_dict)
-
-    # endregion