Fix pickling of VanillaDataManager (#2129)

* Fix pickling of VanillaDataManager

* Fix vanilladatamanager init

nerfstudio/data/datamanagers/base_datamanager.py

@@ -18,11 +18,11 @@ Datamanager.
 
 from __future__ import annotations
 
-import functools
 from abc import abstractmethod
 from collections import defaultdict
 from dataclasses import dataclass, field
 from pathlib import Path
+from functools import cached_property
 from typing import (
     Any,
     Callable,
@@ -35,6 +35,9 @@ from typing import (
     Type,
     Union,
     cast,
+    ForwardRef,
+    get_origin,
+    get_args,
 )
 
 import torch
@@ -66,6 +69,7 @@ from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttrib
 from nerfstudio.model_components.ray_generators import RayGenerator
 from nerfstudio.utils.misc import IterableWrapper
 from nerfstudio.utils.rich_utils import CONSOLE
+from nerfstudio.utils.misc import get_orig_class
 
 
 def variable_res_collate(batch: List[Dict]) -> Dict:
@@ -376,7 +380,6 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
         local_rank: int = 0,
         **kwargs,
     ):
-        self.dataset_type: Type[TDataset] = kwargs.get("_dataset_type", getattr(TDataset, "__default__"))
         self.config = config
         self.device = device
         self.world_size = world_size
@@ -409,15 +412,32 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
                         CONSOLE.print("Variable resolution, using variable_res_collate")
                         self.config.collate_fn = variable_res_collate
                         break
-
         super().__init__()
 
-    def __class_getitem__(cls, item):
-        return type(
-            cls.__name__,
-            (cls,),
-            {"__module__": cls.__module__, "__init__": functools.partialmethod(cls.__init__, _dataset_type=item)},
-        )
+    @cached_property
+    def dataset_type(self) -> Type[TDataset]:
+        """Returns the dataset type passed as the generic argument"""
+        default: Type[TDataset] = cast(TDataset, TDataset.__default__)  # type: ignore
+        orig_class: Type[VanillaDataManager] = get_orig_class(self, default=None)  # type: ignore
+        if type(self) is VanillaDataManager and orig_class is None:
+            return default
+        if orig_class is not None and get_origin(orig_class) is VanillaDataManager:
+            return get_args(orig_class)[0]
+
+        # For inherited classes, we need to find the correct type to instantiate
+        for base in getattr(self, "__orig_bases__", []):
+            if get_origin(base) is VanillaDataManager:
+                for value in get_args(base):
+                    if isinstance(value, ForwardRef):
+                        if value.__forward_evaluated__:
+                            value = value.__forward_value__
+                        elif value.__forward_module__ is None:
+                            value.__forward_module__ = type(self).__module__
+                            value = getattr(value, "_evaluate")(None, None, set())
+                    assert isinstance(value, type)
+                    if issubclass(value, InputDataset):
+                        return cast(Type[TDataset], value)
+        return default
 
     def create_train_dataset(self) -> TDataset:
         """Sets up the data loaders for training"""

nerfstudio/utils/misc.py

@@ -17,6 +17,8 @@ Miscellaneous helper code.
 """
 
 
+from inspect import currentframe
+import typing
 import platform
 from typing import Any, Callable, Dict, List, Optional, TypeVar, Union
 import warnings
@@ -182,3 +184,32 @@ def torch_compile(*args, **kwargs):
         return lambda x: x
     else:
         return torch.compile(*args, **kwargs)
+
+
+def get_orig_class(obj, default=None):
+    """Returns the __orig_class__ class of `obj` even when it is not initialized in __init__ (Python>=3.8).
+
+    Workaround for https://github.com/python/typing/issues/658.
+    Inspired by https://github.com/Stewori/pytypes/pull/53.
+    """
+    try:
+        return object.__getattribute__(obj, "__orig_class__")
+    except AttributeError:
+        cls = object.__getattribute__(obj, "__class__")
+        try:
+            is_type_generic = isinstance(cls, typing.GenericMeta)  # type: ignore
+        except AttributeError:  # Python 3.8
+            is_type_generic = issubclass(cls, typing.Generic)
+        if is_type_generic:
+            frame = currentframe().f_back.f_back  # type: ignore
+            try:
+                while frame:
+                    try:
+                        res = frame.f_locals["self"]
+                        if res.__origin__ is cls:
+                            return res
+                    except (KeyError, AttributeError):
+                        frame = frame.f_back
+            finally:
+                del frame
+        return default

tests/data/test_datamanager.py

@@ -1,3 +1,4 @@
+import pickle
 from pathlib import Path
 from typing import Any
 
@@ -52,18 +53,49 @@ def config():
 def test_data_manager_type_inference(config):
     # Mock for a faster test
 
-    assert VanillaDataManager(config).dataset_type is InputDataset
     assert VanillaDataManager[DepthDataset](config).dataset_type is DepthDataset
+    assert VanillaDataManager(config).dataset_type is InputDataset
+
+    class tmp2(VanillaDataManager[DepthDataset]):
+        pass
+
+    assert tmp2(config).dataset_type is DepthDataset
 
     class tmp(VanillaDataManager):
         pass
 
     assert tmp(config).dataset_type is InputDataset
 
-    class tmp2(VanillaDataManager[DepthDataset]):
-        pass
 
-    assert tmp2(config).dataset_type is DepthDataset
+class _pickle_enabled_tmp(VanillaDataManager):
+    pass
+
+
+class _pickle_enabled_tmp2(VanillaDataManager[DepthDataset]):
+    pass
+
+
+def test_data_manager_type_can_be_pickled(config):
+    # Mock for a faster test
+    assert VanillaDataManager[DepthDataset](config).dataset_type is DepthDataset
+    obj = pickle.loads(pickle.dumps(VanillaDataManager[DepthDataset](config)))
+    assert obj.dataset_type is DepthDataset
+    assert isinstance(obj, VanillaDataManager)
+
+    assert VanillaDataManager(config).dataset_type is InputDataset
+    obj = pickle.loads(pickle.dumps(VanillaDataManager(config)))
+    assert obj.dataset_type is InputDataset
+    assert isinstance(obj, VanillaDataManager)
+
+    assert _pickle_enabled_tmp(config).dataset_type is InputDataset
+    obj = pickle.loads(pickle.dumps(_pickle_enabled_tmp(config)))
+    assert obj.dataset_type is InputDataset
+    assert isinstance(obj, _pickle_enabled_tmp)
+
+    assert _pickle_enabled_tmp2(config).dataset_type is DepthDataset
+    obj = pickle.loads(pickle.dumps(_pickle_enabled_tmp2(config)))
+    assert obj.dataset_type is DepthDataset
+    assert isinstance(obj, _pickle_enabled_tmp2)
 
 
 def test_data_manager_type_can_be_serialized(config):