debug loggers for checking data and network memory usage

Author: ravinkohli

autoPyTorch/pipeline/components/setup/early_preprocessor/EarlyPreprocessing.py

@@ -1,4 +1,7 @@
 from typing import Any, Dict, Optional, Union
+import logging.handlers
+import time
+import psutil
 
 from ConfigSpace.configuration_space import ConfigurationSpace
 
@@ -12,6 +15,7 @@
 from autoPyTorch.pipeline.components.setup.base_setup import autoPyTorchSetupComponent
 from autoPyTorch.pipeline.components.setup.early_preprocessor.utils import get_preprocess_transforms, preprocess
 from autoPyTorch.utils.common import FitRequirement
+from autoPyTorch.utils.logging_ import get_named_client_logger
 
 
 class EarlyPreprocessing(autoPyTorchSetupComponent):
@@ -25,7 +29,12 @@ def __init__(self, random_state: Optional[np.random.RandomState] = None) -> None
 
     def fit(self, X: Dict[str, Any], y: Any = None) -> "EarlyPreprocessing":
         self.check_requirements(X, y)
-
+        self.logger = get_named_client_logger(
+            name=f"{X['num_run']}_{self.__class__.__name__}_{time.time()}",
+            # Log to a user provided port else to the default logging port
+            port=X['logger_port'
+                   ] if 'logger_port' in X else logging.handlers.DEFAULT_TCP_LOGGING_PORT,
+        )
         return self
 
     def transform(self, X: Dict[str, Any]) -> Dict[str, Any]:
@@ -37,7 +46,9 @@ def transform(self, X: Dict[str, Any]) -> Dict[str, Any]:
             # Incorporate the transform to the dataset
             X_train = X['backend'].load_datamanager().train_tensors[0]
 
+        self.logger.debug(f"Available virtual memory: {psutil.virtual_memory().available/1024/1024}, total virtual memroy: {psutil.virtual_memory().total/1024/1024}")
         X['X_train'] = preprocess(dataset=X_train, transforms=transforms)
+        self.logger.debug(f"After preprocessing Available virtual memory: {psutil.virtual_memory().available/1024/1024}, total virtual memroy: {psutil.virtual_memory().total/1024/1024}")
 
         # We need to also save the preprocess transforms for inference
         X.update({

autoPyTorch/pipeline/components/setup/network_backbone/base_network_backbone.py

@@ -1,5 +1,8 @@
 from abc import abstractmethod
 from typing import Any, Dict, Iterable, Optional, Tuple
+import logging.handlers
+import time
+import psutil
 
 import numpy as np
 
@@ -16,6 +19,8 @@
 )
 from autoPyTorch.pipeline.components.setup.network_backbone.utils import get_output_shape
 from autoPyTorch.utils.common import FitRequirement
+from autoPyTorch.utils.logging_ import get_named_client_logger
+
 
 
 class NetworkBackboneComponent(autoPyTorchComponent):
@@ -48,11 +53,17 @@ def fit(self, X: Dict[str, Any], y: Any = None) -> BaseEstimator:
             Self
         """
         self.check_requirements(X, y)
-
+        self.logger = get_named_client_logger(
+            name=f"{X['num_run']}_{self.__class__.__name__}_{time.time()}",
+            # Log to a user provided port else to the default logging port
+            port=X['logger_port'
+                   ] if 'logger_port' in X else logging.handlers.DEFAULT_TCP_LOGGING_PORT,
+        )
         input_shape = X['shape_after_preprocessing']
 
         input_shape = get_output_shape(X['network_embedding'], input_shape=input_shape)
         self.input_shape = input_shape
+        self.logger.debug(f"Before building backbone Available virtual memory: {psutil.virtual_memory().available/1024/1024}, total virtual memroy: {psutil.virtual_memory().total/1024/1024}")
 
         self.backbone = self.build_backbone(
             input_shape=input_shape,

autoPyTorch/pipeline/components/setup/network_embedding/base_network_embedding.py

@@ -1,4 +1,8 @@
-from typing import Any, Dict, List, Optional, Tuple
+import copy
+from typing import Any, Dict, Optional, Tuple
+import logging.handlers
+import time
+import psutil
 
 import numpy as np
 
@@ -7,6 +11,8 @@
 from torch import nn
 
 from autoPyTorch.pipeline.components.setup.base_setup import autoPyTorchSetupComponent
+from autoPyTorch.utils.logging_ import get_named_client_logger
+
 from autoPyTorch.utils.common import FitRequirement
 
 
@@ -20,12 +26,20 @@ def __init__(self, random_state: Optional[np.random.RandomState] = None):
         self.embedding: Optional[nn.Module] = None
 
     def fit(self, X: Dict[str, Any], y: Any = None) -> BaseEstimator:
+        self.logger = get_named_client_logger(
+            name=f"{X['num_run']}_{self.__class__.__name__}_{time.time()}",
+            # Log to a user provided port else to the default logging port
+            port=X['logger_port'
+                   ] if 'logger_port' in X else logging.handlers.DEFAULT_TCP_LOGGING_PORT,
+        )
+        self.logger.debug(f"Before getting info for embedding Available virtual memory: {psutil.virtual_memory().available/1024/1024}, total virtual memroy: {psutil.virtual_memory().total/1024/1024}")
 
         num_features_excl_embed, num_categories_per_col = self._get_required_info_from_data(X)
-
+        self.logger.debug(f"Before building embedding Available virtual memory: {psutil.virtual_memory().available/1024/1024}, total virtual memroy: {psutil.virtual_memory().total/1024/1024}")
         self.embedding = self.build_embedding(
             num_categories_per_col=num_categories_per_col,
             num_features_excl_embed=num_features_excl_embed)
+        self.logger.debug(f"After building embedding Available virtual memory: {psutil.virtual_memory().available/1024/1024}, total virtual memroy: {psutil.virtual_memory().total/1024/1024}")
         return self
 
     def transform(self, X: Dict[str, Any]) -> Dict[str, Any]: