from typing import Any, Dict, Callable, Mapping, Type, Sequence
import os
import importlib
from dataclasses import dataclass, field
from interpretune.adapter_registry import ADAPTER_REGISTRY
from interpretune.config import ITDataModuleConfig, ITConfig, ITSerializableCfg, ITSharedConfig
from interpretune.protocol import Adapter, DataModuleInitable, ModuleSteppable, ITModuleProtocol, ITDataModuleProtocol
from interpretune.utils import unexpected_state_msg_suffix, rank_zero_warn
from interpretune.metadata import ITClassMetadata
class NamedWrapper:
def __repr__(self) -> str:
orig_module = getattr(self, "_orig_module_name", "Original module attribute not set, instantiation incomplete.")
composed_classes = getattr(self, "_composed_classes", "N/A")
enriched_mod_str = f"Original module: {orig_module} {os.linesep}"
enriched_mod_str += f"Now {self.__class__.__name__} composing {orig_module} with: {os.linesep} - "
composed_mod_lines = [c.__name__ for c in composed_classes] if not isinstance(composed_classes, str) else "N/A"
enriched_mod_str += f"{os.linesep} - ".join(composed_mod_lines) + f"{os.linesep}"
return enriched_mod_str + super().__repr__()
class ITMeta(type):
supported_component_keys = ("datamodule", "dm", "module", "m")
def __new__(mcs, name, bases, classdict, **kwargs):
component, input_cls, ctx = mcs._validate_build_ctx(kwargs)
# TODO: add runtime checks for adherence to IT protocol here?
composition_classes = mcs._map_composition_target(component, ctx)
new_bases: tuple[type, ...] = (NamedWrapper, input_cls, *composition_classes) # type: ignore[misc]
built_class = super().__new__(mcs, name, new_bases, classdict)
built_class._orig_module_name = input_cls.__qualname__ # type: ignore[attr-defined] # dynamic attribute for session tracking
built_class._composed_classes = composition_classes # type: ignore[attr-defined] # dynamic attribute for session tracking
return built_class
@staticmethod
def _validate_build_ctx(kwargs: dict) -> tuple[str, Callable, tuple]:
required_kwargs = ("component", "input", "ctx")
for kwarg in required_kwargs:
if kwarg not in kwargs:
raise ValueError(f"{kwarg} must be provided")
if (component := kwargs.get("component")) not in ITMeta.supported_component_keys:
raise ValueError(f"Specified component was {component}, should be either 'module' or 'datamodule'")
if not callable(input := kwargs.get("input")):
raise ValueError(f"Specified input {input} is not a callable, it should be the class to be enriched.")
if not isinstance(ctx := kwargs.get("ctx"), tuple):
raise ValueError(f"Specified ctx {ctx} must be a tuple specifying the desired class enrichment")
return component, input, ctx
@staticmethod
def _map_composition_target(component, ctx):
component_key = None
match component:
case "datamodule" | "dm":
component_key = "datamodule"
case "module" | "m":
component_key = "module"
assert component_key is not None, f"invalid component, should be in: {ITMeta.supported_component_keys}"
return ADAPTER_REGISTRY.get((component_key, *ctx))
[docs]
@dataclass(kw_only=True)
class UnencapsulatedArgs(ITSerializableCfg):
# Most use cases will encapsulate datamodule/module config by subclassing the relevant dataclasses for a given
# experiment/application but we also allow unencapsulated args/kwargs to be passed to the datamodule and module
dm_args: tuple = ()
dm_kwargs: dict[str, Any] = field(default_factory=dict)
module_args: tuple = ()
module_kwargs: dict[str, Any] = field(default_factory=dict)
[docs]
@dataclass(kw_only=True)
class ITSessionConfig(UnencapsulatedArgs):
adapter_ctx: Sequence[Adapter | str] = (Adapter.core,)
datamodule_cfg: ITDataModuleConfig
module_cfg: ITConfig
shared_cfg: ITSharedConfig | Dict | None = None
datamodule_cls: Type[DataModuleInitable] | str | None = None
module_cls: Type[ModuleSteppable] | str | None = None
datamodule: ITDataModuleProtocol | None = None
module: ITModuleProtocol | None = None
def __post_init__(self):
self.adapter_ctx = ADAPTER_REGISTRY.canonicalize_composition(self.adapter_ctx)
# TODO: add checks validating supported adapter combinations are provided
# N.B. we must defer validating data/module_cls against their respective protocols until after composing
# TODO: add warnings/errors here in case provided classes cannot be resolved or are invalid
for attr_k in ["datamodule_cls", "module_cls"]:
if isinstance(getattr(self, attr_k), str):
module, module_cls = getattr(self, attr_k).rsplit(".", 1)
mod = importlib.import_module(module)
setattr(self, attr_k, getattr(mod, module_cls, None))
if self.shared_cfg:
if isinstance(self.shared_cfg, dict):
self.shared_cfg = ITSharedConfig(**self.shared_cfg)
for attr, value in self.shared_cfg.__dict__.items():
for cfg in [self.datamodule_cfg, self.module_cfg]:
if getattr(cfg, attr, None) not in (None, {}, "", value):
rank_zero_warn(
f"Overriding `{attr}` from `{cfg.__class__.__name__}` with value from `shared_cfg`"
)
setattr(cfg, attr, value)
cfg._validate_on_session_cfg_init()
[docs]
class ITSession(Mapping):
# Consolidated class-level metadata to reduce attribute clutter
_it_cls_metadata = ITClassMetadata(
base_attrs={Adapter.core: ("datamodule", "module"), Adapter.lightning: ("datamodule", "model")},
ready_attrs=("datamodule", "module"),
composition_target_attrs=("datamodule_cls", "module_cls"),
ready_protocols=(ITDataModuleProtocol, ITModuleProtocol),
)
def __init__(self, session_cfg: ITSessionConfig, *args, **kwargs):
super().__init__(*args, **kwargs)
self.datamodule = None
self.module = None
# to improve usability, run a datamodule cross-validation hook to allow auto-reconfiguration prior to
# session instantiation
session_cfg.datamodule_cfg._cross_validate(session_cfg.module_cfg)
self.compose_interpretunable(session_cfg)
# TODO: filter adapter_ctx to find first adapter w custom BASE_ATTRS in the future (only lightning right now)
base_attrs = type(self)._it_cls_metadata.base_attrs
self._ctx = (
base_attrs[Adapter.lightning] if Adapter.lightning in session_cfg.adapter_ctx else base_attrs[Adapter.core]
)
def __getitem__(self, key):
return self.to_dict()[key]
def __iter__(self):
return iter(self.to_dict())
def __len__(self):
return len(self.to_dict())
def to_dict(self) -> dict[str, Any]:
return {self._ctx[0]: self.datamodule, self._ctx[1]: self.module}
def __repr__(self) -> str:
try:
dm = getattr(self, "datamodule", None)
m = getattr(self, "module", None)
parts = []
parts.append(f"datamodule={dm.__class__.__name__ if dm is not None else None}")
parts.append(f"module={m.__class__.__name__ if m is not None else None}")
# include module._it_state summary if available
if m is not None and hasattr(m, "_it_state"):
try:
parts.append(m._it_state.to_summary())
except Exception:
parts.append("ITState(<unavailable>)")
return f"ITSession({', '.join(parts)})"
except Exception:
return super().__repr__()
def _check_ready(self, session_cfg: ITSessionConfig) -> None:
meta = type(self)._it_cls_metadata
for ready, tocompose, ready_type in zip(meta.ready_attrs, meta.composition_target_attrs, meta.ready_protocols):
if ready_mod := getattr(session_cfg, ready): # if a module is ready, ensure we don't try to transform it
if not isinstance(ready_mod, ready_type):
raise ValueError(
f"{ready_mod} is not a {ready_type}, {ready_mod} should either be left unset (to"
f" enable auto-composition of {getattr(session_cfg, tocompose)}) or be a"
f" {ready_type}."
)
setattr(session_cfg, tocompose, None)
def compose_interpretunable(self, session_cfg: ITSessionConfig) -> None:
dm_cls, m_cls = None, None
# 1. We first check to see if the datamodule and module classes are already defined and adhere to the relevant
# protocol.
self._check_ready(session_cfg)
# 1.1. If so, we can skip composition
meta = type(self)._it_cls_metadata
for ready, tocompose in zip(meta.ready_attrs, meta.composition_target_attrs):
if getattr(session_cfg, tocompose) is not None:
continue
setattr(self, ready, getattr(session_cfg, ready))
# 2. For the components (i.e. datamodule and module) that aren't ready, we compose the provided input component
# with the relevant classes based on the specified execution context and instantiate our enriched components!
if session_cfg.datamodule_cls: # 2.1 Compose datamodule if necessary
dm_cls = ITMeta(
"InterpretunableDataModule",
(),
{},
component="dm",
input=session_cfg.datamodule_cls,
ctx=session_cfg.adapter_ctx,
)
self.datamodule = dm_cls(itdm_cfg=session_cfg.datamodule_cfg, *session_cfg.dm_args, **session_cfg.dm_kwargs) # type: ignore[call-arg]
self._set_dm_handles_for_instantiation(session_cfg)
if session_cfg.module_cls: # 2.2 Compose module if necessary
m_cls = ITMeta(
"InterpretunableModule",
(),
{},
component="m",
input=session_cfg.module_cls,
ctx=session_cfg.adapter_ctx,
)
self.module = m_cls(it_cfg=session_cfg.module_cfg, *session_cfg.module_args, **session_cfg.module_kwargs) # type: ignore[call-arg]
self._set_model_handles_for_instantiation()
self._validate_session(dm_cls, m_cls, session_cfg)
def _set_dm_handles_for_instantiation(self, session_cfg: ITSessionConfig):
# some datamodule handles may be required for module init, we update the module_cfg to provide them here
supported_dm_handles_for_module = {"tokenizer": "tokenizer"}
for m_attr, dm_handle in supported_dm_handles_for_module.items():
# attach directly to module if it's ready, otherwise pass the handles to our module_cfg
attr_target_obj = getattr(self, "module", None) or session_cfg.module_cfg
setattr(attr_target_obj, m_attr, getattr(self.datamodule, dm_handle))
def _set_model_handles_for_instantiation(self):
# having access to a model handle may be useful in some pre-setup hook steps (e.g. signature inspection in
# `prepare_data`) we provide early access to the model handle for datamodule
# TODO: for Lightning, since we're setting _module before trainer.model is set, we double check attr coherency
# TODO: wrt Lightning compatibility, since we're providing access to the model handle before `setup_environment`
# or `configure_model` hooks are executed, need to evaluate the the validity/need to update this reference
# after subsequent model hooks are called
supported_module_handles_for_datamodule = {"_module": self.module}
for dm_attr, m_handle in supported_module_handles_for_datamodule.items():
setattr(self.datamodule, dm_attr, m_handle)
def _issue_warnings(self, dm_composed, m_composed, session_cfg, warn_datamodule, warn_module):
if warn_datamodule:
dm_warn_name = getattr(self.datamodule, "_orig_module_name", self.datamodule.__class__.__name__)
dm_warn_base = f"{dm_warn_name} is not {ITDataModuleProtocol}"
dm_warn_composed = (
f" even after composing {session_cfg.datamodule_cls} with {dm_composed}. {unexpected_state_msg_suffix}"
)
not_ready_dm_warn = dm_warn_base + dm_warn_composed if dm_composed else dm_warn_base
rank_zero_warn(not_ready_dm_warn)
if warn_module:
m_warn_name = getattr(self.module, "_orig_module_name", self.module.__class__.__name__)
m_warn_base = f"{m_warn_name} is not {ITModuleProtocol}"
m_warn_composed = (
f" even after composing {session_cfg.module_cls} with {m_composed}. {unexpected_state_msg_suffix}"
)
not_ready_m_warn = m_warn_base + m_warn_composed if m_composed else m_warn_base
rank_zero_warn(not_ready_m_warn)
def _validate_session(self, dm_composed, m_composed, session_cfg):
warn_datamodule = not isinstance(self.datamodule, ITDataModuleProtocol)
warn_module = not isinstance(self.module, ITModuleProtocol)
if any([warn_datamodule, warn_module]):
self._issue_warnings(dm_composed, m_composed, session_cfg, warn_datamodule, warn_module)
