kostyl-toolkit 0.1.35__py3-none-any.whl → 0.1.37__py3-none-any.whl

kostyl/utils/logging.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import inspect
+import os
 import sys
 import uuid
 from collections import namedtuple
@@ -18,32 +19,18 @@ from loguru import logger as _base_logger
 if TYPE_CHECKING:
     from loguru import Logger
 
-    class CustomLogger(Logger):  # noqa: D101
+    class KostylLogger(Logger):  # noqa: D101
         def log_once(self, level: str, message: str, *args, **kwargs) -> None: ...  # noqa: ANN003, D102
         def warning_once(self, message: str, *args, **kwargs) -> None: ...  # noqa: ANN003, D102
 else:
-    CustomLogger = type(_base_logger)
+    KostylLogger = type(_base_logger)
 
 try:
-    import torch.distributed as dist
     from torch.nn.modules.module import (
         _IncompatibleKeys,  # pyright: ignore[reportAssignmentType]
     )
 except Exception:
 
-    class _Dummy:
-        @staticmethod
-        def is_available() -> bool:
-            return False
-
-        @staticmethod
-        def is_initialized() -> bool:
-            return False
-
-        @staticmethod
-        def get_rank() -> int:
-            return 0
-
     class _IncompatibleKeys(
         namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"]),
     ):
@@ -56,14 +43,13 @@ except Exception:
 
         __str__ = __repr__
 
-    dist = _Dummy()
     _IncompatibleKeys = _IncompatibleKeys
 
 _once_lock = Lock()
 _once_keys: set[tuple[str, str]] = set()
 
 
-def _log_once(self: CustomLogger, level: str, message: str, *args, **kwargs) -> None:  # noqa: ANN003
+def _log_once(self: KostylLogger, level: str, message: str, *args, **kwargs) -> None:  # noqa: ANN003
     key = (message, level)
 
     with _once_lock:
@@ -75,7 +61,7 @@ def _log_once(self: CustomLogger, level: str, message: str, *args, **kwargs) ->
     return
 
 
-_base_logger = cast(CustomLogger, _base_logger)
+_base_logger = cast(KostylLogger, _base_logger)
 _base_logger.log_once = _log_once  # pyright: ignore[reportAttributeAccessIssue]
 _base_logger.warning_once = partialmethod(_log_once, "WARNING")  # pyright: ignore[reportAttributeAccessIssue]
 
@@ -91,44 +77,83 @@ _DEFAULT_FMT = "<level>{level: <8}</level> {time:HH:mm:ss.SSS} [{extra[channel]}
 _ONLY_MESSAGE_FMT = "<level>{message}</level>"
 _PRESETS = {"default": _DEFAULT_FMT, "only_message": _ONLY_MESSAGE_FMT}
 
+KOSTYL_LOG_LEVEL = os.getenv("KOSTYL_LOG_LEVEL", "INFO")
+
 
 def setup_logger(
     name: str | None = None,
-    fmt: Literal["default", "only_message"] | str = "default",
-    level: str = "INFO",
-    add_rank: bool | None = None,
+    fmt: Literal["default", "only_message"] | str = "only_message",
+    level: str | None = None,
     sink=sys.stdout,
     colorize: bool = True,
     serialize: bool = False,
-) -> CustomLogger:
+) -> KostylLogger:
     """
-    Returns a bound logger with its own sink and formatting.
+    Creates and configures a logger with custom formatting and output.
+
+    The function automatically removes the default sink on first call and creates
+    an isolated logger with a unique identifier for message filtering.
+
+    Args:
+        name (str | None, optional): Logger channel name. If None, automatically
+            uses the calling function's filename. Defaults to None.
+        fmt (Literal["default", "only_message"] | str, optional): Log message format.
+            Available presets:
+            - "default": includes level, time, and channel
+            - "only_message": outputs only the message itself
+            Custom format strings are also supported. Defaults to "only_message".
+        level (str | None, optional): Logging level (TRACE, DEBUG, INFO, SUCCESS,
+            WARNING, ERROR, CRITICAL). If None, uses the KOSTYL_LOG_LEVEL environment
+            variable or "INFO" by default. Defaults to None.
+        sink: Output object for logs (file, sys.stdout, sys.stderr, etc.).
+            Defaults to sys.stdout.
+        colorize (bool, optional): Enable colored output formatting.
+            Defaults to True.
+        serialize (bool, optional): Serialize logs to JSON format.
+            Defaults to False.
+
+    Returns:
+        CustomLogger: Configured logger instance with additional methods
+            log_once() and warning_once().
+
+    Example:
+        >>> # Basic usage with automatic name detection
+        >>> logger = setup_logger()
+        >>> logger.info("Hello World")
 
-    Note: If name=None, the caller's filename (similar to __file__) is used automatically.
+        >>> # With custom name and level
+        >>> logger = setup_logger(name="MyApp", level="DEBUG")
+
+        >>> # With custom format
+        >>> logger = setup_logger(
+        ...     name="API",
+        ...     fmt="{level} | {time:YYYY-MM-DD HH:mm:ss} | {message}"
+        ... )
 
-    Format example: "{level} {time:MM-DD HH:mm:ss} [{extra[channel]}] {message}"
     """
     global _DEFAULT_SINK_REMOVED
     if not _DEFAULT_SINK_REMOVED:
         _base_logger.remove()
         _DEFAULT_SINK_REMOVED = True
 
-    if name is None:
-        base = _caller_filename()
-    else:
-        base = name
+    if level is None:
+        if KOSTYL_LOG_LEVEL not in {
+            "TRACE",
+            "DEBUG",
+            "INFO",
+            "SUCCESS",
+            "WARNING",
+            "ERROR",
+            "CRITICAL",
+        }:
+            level = "INFO"
+        else:
+            level = KOSTYL_LOG_LEVEL
 
-    if (add_rank is None) or add_rank:
-        try:
-            add_rank = dist.is_available() and dist.is_initialized()
-        except Exception:
-            add_rank = False
-
-    if add_rank:
-        rank = dist.get_rank()
-        channel = f"rank:{rank} - {base}"
+    if name is None:
+        channel = _caller_filename()
     else:
-        channel = base
+        channel = name
 
     if fmt in _PRESETS:
         fmt = _PRESETS[fmt]
@@ -146,7 +171,7 @@ def setup_logger(
         filter=lambda r: r["extra"].get("logger_id") == logger_id,
     )
     logger = _base_logger.bind(logger_id=logger_id, channel=channel)
-    return cast(CustomLogger, logger)
+    return cast(KostylLogger, logger)
 
 
 def log_incompatible_keys(
@@ -154,22 +179,12 @@ def log_incompatible_keys(
     incompatible_keys: _IncompatibleKeys
     | tuple[list[str], list[str]]
     | dict[str, list[str]],
-    model_specific_msg: str = "",
+    postfix_msg: str = "",
 ) -> None:
     """
     Logs warnings for incompatible keys encountered during model loading or state dict operations.
 
     Note: If incompatible_keys is of an unsupported type, an error message is logged and the function returns early.
-
-    Args:
-        logger (Logger): The logger instance used to output warning messages.
-        incompatible_keys (_IncompatibleKeys | tuple[list[str], list[str]] | dict[str, list[str]]): An object containing lists of missing and unexpected keys.
-        model_specific_msg (str, optional): A custom message to append to the log output, typically
-            indicating the model or context. Defaults to an empty string.
-
-    Returns:
-        None
-
     """
     incompatible_keys_: dict[str, list[str]] = {}
     match incompatible_keys:
@@ -192,5 +207,5 @@ def log_incompatible_keys(
             return
 
     for name, keys in incompatible_keys_.items():
-        logger.warning(f"{name} {model_specific_msg}: {', '.join(keys)}")
+        logger.warning(f"{name} {postfix_msg}: {', '.join(keys)}")
     return

{kostyl_toolkit-0.1.35.dist-info → kostyl_toolkit-0.1.37.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.35
+Version: 0.1.37
 Summary: Kickass Orchestration System for Training, Yielding & Logging 
 Requires-Dist: case-converter>=1.2.0
 Requires-Dist: loguru>=0.7.3

{kostyl_toolkit-0.1.35.dist-info → kostyl_toolkit-0.1.37.dist-info}/RECORD

@@ -6,32 +6,33 @@ kostyl/ml/clearml/logging_utils.py,sha256=GBjIIZbH_itd5sj7XpvxjkyZwxxGOpEcQ3BiWa
 kostyl/ml/clearml/pulling_utils.py,sha256=jMlVXcYRumwWnPlELRlgEdfq5L6Wir_EcfTmOoWBLTA,4077
 kostyl/ml/configs/__init__.py,sha256=IetcivbqYGutowLqxdKp7QR4tkXKBr4m8t4Zkk9jHZU,911
 kostyl/ml/configs/base_model.py,sha256=Eofn14J9RsjpVx_J4rp6C19pDDCANU4hr3JtX-d0FpQ,4820
-kostyl/ml/configs/hyperparams.py,sha256=2S_VEZ07RWquNFSWjHBb3OUpBlTznbUpFSchzMpSBOc,2879
-kostyl/ml/configs/training_settings.py,sha256=Sq2tiRuwkbmi9zKDG2JghZLXo5DDt_eQqN_KYJSdcTY,2509
+kostyl/ml/configs/hyperparams.py,sha256=lvtbvOFEoTBAJug7FR35xMQdPLgDQjRoP2fyDP-jD7E,3305
+kostyl/ml/configs/training_settings.py,sha256=wT9CHuLaKrLwonsc87Ee421EyFis_c9fqOgn9bSClm8,2747
 kostyl/ml/data_processing_utils.py,sha256=jjEjV0S0wREgZkzg27ip0LpI8cQqkwe2QwATmAqm9-g,3832
-kostyl/ml/dist_utils.py,sha256=Onf0KHVLA8oeUgZTcTdmR9qiM22f2uYLoNwgLbMGJWk,3495
+kostyl/ml/dist_utils.py,sha256=lK9_aAh9L1SvvXWzcWiBoFjczfDiKzEpcno5csImAYQ,4635
 kostyl/ml/lightning/__init__.py,sha256=R36PImjVvzBF9t_z9u6RYVnUFJJ-sNDUOdboWUojHmM,173
 kostyl/ml/lightning/callbacks/__init__.py,sha256=EnKkNwwNDZnEqKRlpY4FVrqP88ECPF6nlT2bSLUIKRk,194
-kostyl/ml/lightning/callbacks/checkpoint.py,sha256=sZ9OqudO-gXp7FqtWaOH46TXVpeCJxV-EowyRPN836k,18983
+kostyl/ml/lightning/callbacks/checkpoint.py,sha256=HI17gu-GxnfXUchflWBTwly7cCYnlpKcshuR-TgD6s4,19066
 kostyl/ml/lightning/callbacks/early_stopping.py,sha256=D5nyjktCJ9XYAf28-kgXG8jORvXLl1N3nbDQnvValPM,615
 kostyl/ml/lightning/extensions/__init__.py,sha256=OY6QGv1agYgqqKf1xJBrxgp_i8FunVfPzYezfaRrGXU,182
-kostyl/ml/lightning/extensions/custom_module.py,sha256=iQrnPz-WTmRfvLo94C5fQc2Qwa1IpHtUh1sCpVwTSFM,6602
-kostyl/ml/lightning/extensions/pretrained_model.py,sha256=eRfQBzAjVernHl9A4PP5uTLvjjmcNKPdTu7ABFLq7HI,5196
+kostyl/ml/lightning/extensions/custom_module.py,sha256=qYffgPwIB_ePwK_MIaRruuDxPKJZb42kg2yy996eGwY,6415
+kostyl/ml/lightning/extensions/pretrained_model.py,sha256=hVIsIUu6Iryrz6S7GQTqog9vNq8LQyjJd2aoJ5Ws6KU,5253
 kostyl/ml/lightning/loggers/__init__.py,sha256=e51dszaoJbuzwBkbdugmuDsPldoSO4yaRgmZUg1Bdy0,71
-kostyl/ml/lightning/loggers/tb_logger.py,sha256=j02HK5ue8yzXXV8FWKmmXyHkFlIxgHx-ahHWk_rFCZs,893
-kostyl/ml/lightning/training_utils.py,sha256=u7X9ysF9Gqy8CdwacdcDlNQNsbagYAhslbv-1WLJ45k,9052
+kostyl/ml/lightning/loggers/tb_logger.py,sha256=CpjlcEIT187cJXJgRYafqfzvcnwPgPaVZ0vLUflIr7k,899
+kostyl/ml/lightning/utils.py,sha256=DhLy_3JA5VyMQkB1v6xxRxDNHfisjXFYVjuIKPpO81M,1967
 kostyl/ml/metrics_formatting.py,sha256=U6vdNENZLvp2dT1L3HqFKtXrHwGKoDXN93hvamPGHjM,1341
 kostyl/ml/params_groups.py,sha256=nUyw5d06Pvy9QPiYtZzLYR87xwXqJLxbHthgQH8oSCM,3583
-kostyl/ml/registry_uploader.py,sha256=W90TYo_WKv2oBE6nqEJl4hecYmJyyuKwQJ9_uUPGnJQ,3346
-kostyl/ml/schedulers/__init__.py,sha256=bxXbsU_WYnVbhvNNnuI7cOAh2Axz7D25TaleBTZhYfc,197
-kostyl/ml/schedulers/base.py,sha256=9M2iOoOVSRojR_liPX1qo3Nn4iMXSM5ZJuAFWZTulUk,1327
+kostyl/ml/registry_uploader.py,sha256=BbyLXvF8AL145k7g6MRkJ7gf_3Um53p3Pn5280vVD9U,4384
+kostyl/ml/schedulers/__init__.py,sha256=_EtZu8DwTCSv4-eR84kRstEZblHylVqda7WQUOXIKfw,534
+kostyl/ml/schedulers/base.py,sha256=bjmwgdZpnSqpCnHPnKC6MEiRO79cwxMJpZq-eQVNs2M,1353
 kostyl/ml/schedulers/composite.py,sha256=ee4xlMDMMtjKPkbTF2ue9GTr9DuGCGjZWf11mHbi6aE,2387
-kostyl/ml/schedulers/cosine.py,sha256=t74_ByT22L5NQKpnBVU9UGzBVx1ZM2GTylb9ct3_PVg,7627
-kostyl/ml/schedulers/linear.py,sha256=7HPkVWcPa0lbaZywutXSDdVLLSihAyWk5XIE2Dzj_5Q,5168
+kostyl/ml/schedulers/cosine.py,sha256=y8ylrgVOkVcr2-ExoqqNW--tdDX88TBYPQCOppIf2_M,8685
+kostyl/ml/schedulers/cosine_with_plateu.py,sha256=0-X6wl3HgsTiLIbISb9lOxIVWXHDEND7rILitMWtIiM,10195
+kostyl/ml/schedulers/linear.py,sha256=RnnnblRuRXP3LT03QVIHUaK2kNsiMP1AedrMoeyh3qk,5843
 kostyl/utils/__init__.py,sha256=hkpmB6c5pr4Ti5BshOROebb7cvjDZfNCw83qZ_FFKMM,240
 kostyl/utils/dict_manipulations.py,sha256=e3vBicID74nYP8lHkVTQc4-IQwoJimrbFELy5uSF6Gk,1073
 kostyl/utils/fs.py,sha256=gAQNIU4R_2DhwjgzOS8BOMe0gZymtY1eZwmdgOdDgqo,510
-kostyl/utils/logging.py,sha256=Vye0u4-yeOSUc-f03gpQbxSktTbFiilTWLEVr00ZHvc,5796
-kostyl_toolkit-0.1.35.dist-info/WHEEL,sha256=ZyFSCYkV2BrxH6-HRVRg3R9Fo7MALzer9KiPYqNxSbo,79
-kostyl_toolkit-0.1.35.dist-info/METADATA,sha256=KL4-Z421DpchI6KUZ6tVATy99urk1OP2OY4Uf5r9R3U,4269
-kostyl_toolkit-0.1.35.dist-info/RECORD,,
+kostyl/utils/logging.py,sha256=CgNFNogcK0hoZmygvBWlTcq5A3m2Pfv9eOAP_gwx0pM,6633
+kostyl_toolkit-0.1.37.dist-info/WHEEL,sha256=eycQt0QpYmJMLKpE3X9iDk8R04v2ZF0x82ogq-zP6bQ,79
+kostyl_toolkit-0.1.37.dist-info/METADATA,sha256=yHPgSAhPnm5tDQjvDIfs213-bsVX6vMfVsUbX9GboGU,4269
+kostyl_toolkit-0.1.37.dist-info/RECORD,,

{kostyl_toolkit-0.1.35.dist-info → kostyl_toolkit-0.1.37.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: uv 0.9.18
+Generator: uv 0.9.24
 Root-Is-Purelib: true
 Tag: py3-none-any

kostyl/ml/lightning/training_utils.py

@@ -1,241 +0,0 @@
-from dataclasses import dataclass
-from dataclasses import fields
-from pathlib import Path
-from typing import Literal
-from typing import cast
-
-import lightning as L
-import torch
-import torch.distributed as dist
-from clearml import OutputModel
-from clearml import Task
-from lightning.pytorch.callbacks import Callback
-from lightning.pytorch.callbacks import EarlyStopping
-from lightning.pytorch.callbacks import LearningRateMonitor
-from lightning.pytorch.callbacks import ModelCheckpoint
-from lightning.pytorch.loggers import TensorBoardLogger
-from lightning.pytorch.strategies import DDPStrategy
-from lightning.pytorch.strategies import FSDPStrategy
-from torch.distributed import ProcessGroup
-from torch.distributed.fsdp import MixedPrecision
-from torch.nn import Module
-
-from kostyl.ml.configs import CheckpointConfig
-from kostyl.ml.configs import DDPStrategyConfig
-from kostyl.ml.configs import EarlyStoppingConfig
-from kostyl.ml.configs import FSDP1StrategyConfig
-from kostyl.ml.configs import SingleDeviceStrategyConfig
-from kostyl.ml.lightning.callbacks import setup_checkpoint_callback
-from kostyl.ml.lightning.callbacks import setup_early_stopping_callback
-from kostyl.ml.lightning.loggers import setup_tb_logger
-from kostyl.ml.registry_uploader import ClearMLRegistryUploaderCallback
-from kostyl.utils.logging import setup_logger
-
-
-TRAINING_STRATEGIES = (
-    FSDP1StrategyConfig | DDPStrategyConfig | SingleDeviceStrategyConfig
-)
-
-logger = setup_logger(add_rank=True)
-
-
-def estimate_total_steps(
-    trainer: L.Trainer, process_group: ProcessGroup | None = None
-) -> int:
-    """
-    Estimates the total number of training steps based on the
-    dataloader length, accumulation steps, and distributed world size.
-    """  # noqa: D205
-    if dist.is_initialized():
-        world_size = dist.get_world_size(process_group)
-    else:
-        world_size = 1
-
-    datamodule = trainer.datamodule  # type: ignore
-    if datamodule is None:
-        raise ValueError("Trainer must have a datamodule to estimate total steps.")
-    datamodule = cast(L.LightningDataModule, datamodule)
-
-    logger.info("Loading `train_dataloader` to estimate number of stepping batches.")
-    datamodule.setup("fit")
-
-    dataloader_len = len(datamodule.train_dataloader())
-    steps_per_epoch = dataloader_len // trainer.accumulate_grad_batches // world_size
-
-    if trainer.max_epochs is None:
-        raise ValueError("Trainer must have `max_epochs` set to estimate total steps.")
-    total_steps = steps_per_epoch * trainer.max_epochs
-
-    logger.info(
-        f"Total steps: {total_steps} (per-epoch: {steps_per_epoch})\n"
-        f"-> Dataloader len: {dataloader_len}\n"
-        f"-> Accumulate grad batches: {trainer.accumulate_grad_batches}\n"
-        f"-> Epochs: {trainer.max_epochs}\n "
-        f"-> World size: {world_size}"
-    )
-    return total_steps
-
-
-@dataclass
-class Callbacks:
-    """Dataclass to hold PyTorch Lightning callbacks."""
-
-    checkpoint: ModelCheckpoint
-    lr_monitor: LearningRateMonitor
-    early_stopping: EarlyStopping | None = None
-
-    def to_list(self) -> list[Callback]:
-        """Convert dataclass fields to a list of Callbacks. None values are omitted."""
-        callbacks: list[Callback] = [
-            getattr(self, field.name)
-            for field in fields(self)
-            if getattr(self, field.name) is not None
-        ]
-        return callbacks
-
-
-def setup_callbacks(
-    task: Task,
-    root_path: Path,
-    checkpoint_cfg: CheckpointConfig,
-    early_stopping_cfg: EarlyStoppingConfig | None,
-    output_model: OutputModel,
-    checkpoint_upload_strategy: Literal["only-best", "every-checkpoint"],
-    config_dict: dict[str, str] | None = None,
-    enable_tag_versioning: bool = False,
-) -> Callbacks:
-    """
-    Set up PyTorch Lightning callbacks for training.
-
-    Creates and configures a set of callbacks including checkpoint saving,
-    learning rate monitoring, model registry uploading, and optional early stopping.
-
-    Args:
-        task: ClearML task for organizing checkpoints by task name and ID.
-        root_path: Root directory for saving checkpoints.
-        checkpoint_cfg: Configuration for checkpoint saving behavior.
-        checkpoint_upload_strategy: Model upload strategy:
-            - `"only-best"`: Upload only the best checkpoint based on monitored metric.
-            - `"every-checkpoint"`: Upload every saved checkpoint.
-        output_model: ClearML OutputModel instance for model registry integration.
-        early_stopping_cfg: Configuration for early stopping. If None, early stopping
-            is disabled.
-        config_dict: Optional configuration dictionary to store with the model
-            in the registry.
-        enable_tag_versioning: Whether to auto-increment version tags (e.g., "v1.0")
-            on the uploaded model.
-
-    Returns:
-        Callbacks dataclass containing configured checkpoint, lr_monitor,
-        and optionally early_stopping callbacks.
-
-    """
-    lr_monitor = LearningRateMonitor(
-        logging_interval="step", log_weight_decay=True, log_momentum=False
-    )
-    model_uploader = ClearMLRegistryUploaderCallback(
-        output_model=output_model,
-        config_dict=config_dict,
-        verbose=True,
-        enable_tag_versioning=enable_tag_versioning,
-    )
-    checkpoint_callback = setup_checkpoint_callback(
-        root_path / "checkpoints" / task.name / task.id,
-        checkpoint_cfg,
-        registry_uploader_callback=model_uploader,
-        uploading_strategy=checkpoint_upload_strategy,
-    )
-    if early_stopping_cfg is not None:
-        early_stopping_callback = setup_early_stopping_callback(early_stopping_cfg)
-    else:
-        early_stopping_callback = None
-
-    callbacks = Callbacks(
-        checkpoint=checkpoint_callback,
-        lr_monitor=lr_monitor,
-        early_stopping=early_stopping_callback,
-    )
-    return callbacks
-
-
-def setup_loggers(task: Task, root_path: Path) -> list[TensorBoardLogger]:
-    """
-    Set up PyTorch Lightning loggers for training.
-
-    Args:
-        task: ClearML task used to organize log directories by task name and ID.
-        root_path: Root directory for storing TensorBoard logs.
-
-    Returns:
-        List of configured TensorBoard loggers.
-
-    """
-    loggers = [
-        setup_tb_logger(root_path / "runs" / task.name / task.id),
-    ]
-    return loggers
-
-
-def setup_strategy(
-    strategy_settings: TRAINING_STRATEGIES,
-    devices: list[int] | int,
-    auto_wrap_policy: set[type[Module]] | None = None,
-) -> Literal["auto"] | FSDPStrategy | DDPStrategy:
-    """
-    Configure and return a PyTorch Lightning training strategy.
-
-    Args:
-        strategy_settings: Strategy configuration object. Must be one of:
-            - `FSDP1StrategyConfig`: Fully Sharded Data Parallel strategy (requires 2+ devices).
-            - `DDPStrategyConfig`: Distributed Data Parallel strategy (requires 2+ devices).
-            - `SingleDeviceStrategyConfig`: Single device training (requires exactly 1 device).
-        devices: Device(s) to use for training. Either a list of device IDs or
-            a single integer representing the number of devices.
-        auto_wrap_policy: Set of module types that should be wrapped for FSDP.
-            Required when using `FSDP1StrategyConfig`, ignored otherwise.
-
-    Returns:
-        Configured strategy: `FSDPStrategy`, `DDPStrategy`, or `"auto"` for single device.
-
-    Raises:
-        ValueError: If device count doesn't match strategy requirements or
-            if `auto_wrap_policy` is missing for FSDP.
-
-    """
-    if isinstance(devices, list):
-        num_devices = len(devices)
-    else:
-        num_devices = devices
-
-    match strategy_settings:
-        case FSDP1StrategyConfig():
-            if num_devices < 2:
-                raise ValueError("FSDP strategy requires multiple devices.")
-
-            if auto_wrap_policy is None:
-                raise ValueError("auto_wrap_policy must be provided for FSDP strategy.")
-
-            mixed_precision_config = MixedPrecision(
-                param_dtype=getattr(torch, strategy_settings.param_dtype),
-                reduce_dtype=getattr(torch, strategy_settings.reduce_dtype),
-                buffer_dtype=getattr(torch, strategy_settings.buffer_dtype),
-            )
-            strategy = FSDPStrategy(
-                auto_wrap_policy=auto_wrap_policy,
-                mixed_precision=mixed_precision_config,
-            )
-        case DDPStrategyConfig():
-            if num_devices < 2:
-                raise ValueError("DDP strategy requires at least two devices.")
-            strategy = DDPStrategy(
-                find_unused_parameters=strategy_settings.find_unused_parameters
-            )
-        case SingleDeviceStrategyConfig():
-            if num_devices != 1:
-                raise ValueError("SingleDevice strategy requires exactly one device.")
-            strategy = "auto"
-        case _:
-            raise ValueError(
-                f"Unsupported strategy type: {type(strategy_settings.trainer.strategy)}"
-            )
-    return strategy