v1.7.2: removed debug print, scheduler fix

Author: dholzmueller

README.md

@@ -200,6 +200,9 @@ and https://docs.ray.io/en/latest/cluster/vms/user-guides/community/slurm.html
 
 ## Releases (see git tags)
 
+- v1.7.2: 
+    - Removed debug print in RealMLP.
+    - fixed device memory estimation error in the scheduler when `CUDA_VISIBLE_DEVICES` was used.
 - v1.7.1:
     - LightGBM now processes the `extra_trees`, `max_cat_to_onehot`, and `min_data_per_group` parameters 
       used in the `'tabarena'` search space, which should improve results.

pytabkit/__about__.py

@@ -2,4 +2,4 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "1.7.1"
+__version__ = "1.7.2"

pytabkit/bench/scheduling/execution.py

@@ -13,6 +13,38 @@
 from pytabkit.models.utils import FunctionProcess
 
 
+def get_gpu_rams_gb(use_reserved: bool = True):
+    """
+    Returns:
+      gpu_rams_gb: total GPU memory per visible device (GB)
+      gpu_rams_fixed_gb: this process GPU memory per visible device (GB)
+        - reserved (default): torch caching allocator reserved bytes (often matches "process used" better)
+        - allocated: live tensor bytes only
+    """
+    # do it in torch, it respects CUDA_VISIBLE_DEVICES and doesn't need the pynvml dependency
+    BYTES_TO_GB = 1024.0 ** 3
+    import torch
+
+    gpu_rams_gb = []
+    gpu_rams_fixed_gb = []
+
+    n = torch.cuda.device_count()  # respects CUDA_VISIBLE_DEVICES ("" => 0)
+    for i in range(n):
+        with torch.cuda.device(i):
+            _free_b, total_b = torch.cuda.mem_get_info()
+
+        gpu_rams_gb.append(total_b / BYTES_TO_GB)
+
+        if use_reserved:
+            used_b = torch.cuda.memory_reserved(i)
+        else:
+            used_b = torch.cuda.memory_allocated(i)
+
+        gpu_rams_fixed_gb.append(used_b / BYTES_TO_GB)
+
+    return gpu_rams_gb, gpu_rams_fixed_gb
+
+
 def measure_node_resources(node_id: int) -> Tuple[NodeResources, NodeResources]:
     """
     Function that measures available resources.
@@ -29,21 +61,22 @@ def measure_node_resources(node_id: int) -> Tuple[NodeResources, NodeResources]:
         # init cuda
         # alloc dummy tensors to know how much memory PyTorch uses for its runtime
         dummy_tensors = [torch.ones(1).to(f'cuda:{i}') for i in range(n_gpus)]
-        import pynvml
-        pynvml.nvmlInit()
-
-        gpu_rams_gb = []
-        gpu_rams_fixed_gb = []
-
-        for i in range(n_gpus):
-            # adapted torch.cuda.list_gpu_processes(gpu)
-            h = pynvml.nvmlDeviceGetHandleByIndex(i)
-            info = pynvml.nvmlDeviceGetMemoryInfo(h)
-            total = info.total
-            # print(f'free     : {info.free}')
-            used = info.used
-            gpu_rams_gb.append(total / (1024. ** 3))
-            gpu_rams_fixed_gb.append(used / (1024. ** 3))
+        gpu_rams_gb, gpu_rams_fixed_gb = get_gpu_rams_gb()
+        # import pynvml
+        # pynvml.nvmlInit()
+        #
+        # gpu_rams_gb = []
+        # gpu_rams_fixed_gb = []
+        #
+        # for i in range(n_gpus):
+        #     # adapted torch.cuda.list_gpu_processes(gpu)
+        #     h = pynvml.nvmlDeviceGetHandleByIndex(i)
+        #     info = pynvml.nvmlDeviceGetMemoryInfo(h)
+        #     total = info.total
+        #     # print(f'free     : {info.free}')
+        #     used = info.used
+        #     gpu_rams_gb.append(total / (1024. ** 3))
+        #     gpu_rams_fixed_gb.append(used / (1024. ** 3))
     else:
         gpu_rams_gb = []
         gpu_rams_fixed_gb = []

pytabkit/bench/scheduling/resources.py

@@ -147,6 +147,8 @@ def try_assign(self, required_resources: RequiredResources,
         gpu_rams_gb_all = fr.get_gpu_rams_gb() + rr.gpu_ram_gb
         gpu_availability = np.logical_and(gpu_usages_all <= self.get_gpu_usages() + 1e-8,
                                           gpu_rams_gb_all <= self.get_gpu_rams_gb())
+        # print(f'{fr.get_gpu_rams_gb()=}, {rr.gpu_ram_gb=}')
+        # print(f'{gpu_usages_all=}, {gpu_rams_gb_all=}, {self.get_gpu_usages()=}, {self.get_gpu_rams_gb()=}, {gpu_availability=}')
         available_gpus = np.argwhere(gpu_availability)[:, 0]  # squeeze second dimension
         # sort available gpus by usage
         available_gpu_usages = self.get_gpu_usages()[available_gpus]

pytabkit/bench/scheduling/schedulers.py

@@ -318,7 +318,11 @@ def _submit_more_jobs(self) -> None:
 
             # otherwise, try assigning the job
             for node_idx, r in enumerate(free_resources.resources):
+                # print(f'{fixed_resources.__dict__=}')
+                # print(f'{job_info.required_resources.__dict__=}')
+                # print(f'{r.data=}, {r.get_resource_vector()=}, {node_idx=}')
                 assigned_resources = r.try_assign(job_info.required_resources, fixed_resources)
+                # print(f'{bool(assigned_resources)=}')
                 if assigned_resources is not None:
                     job_info.set_started(assigned_resources)
                     self.job_manager.submit_job(job_info)

pytabkit/models/alg_interfaces/other_interfaces.py

@@ -1250,7 +1250,7 @@ def _create_sklearn_model(self, seed: int, n_threads: int, gpu_devices: List[str
         # print(f'{gpu_devices=}')
         if self.n_classes > 0:
             if self.config.get('use_tabiclex', False):
-                from tabiclex import TabICLClassifier
+                from tabiclv2 import TabICLClassifier
             else:
                 from tabicl import TabICLClassifier
             return TabICLClassifier(random_state=seed,
@@ -1304,9 +1304,9 @@ def get_required_resources(self, ds: DictDataset, n_cv: int, n_refit: int, n_spl
         assert n_splits == 1
         updated_config = utils.join_dicts(dict(n_estimators=100), self.config)
         time_params = {'': 0.5, 'ds_size_gb': 10.0, '1/n_threads*n_samples*n_estimators*n_tree_repeats': 4e-8}
-        ram_params = {'': 0.5, 'ds_size_gb': 3.0, 'n_samples*n_estimators*n_tree_repeats': 3e-9}
+        ram_params = {'': 0.5}
         rc = ResourcePredictor(config=updated_config, time_params=time_params,
-                               cpu_ram_params=ram_params, n_gpus=1, gpu_usage=1.0, gpu_ram_params={'': 10.0})
+                               cpu_ram_params=ram_params, n_gpus=1, gpu_usage=0.999, gpu_ram_params={'': 10.0})
         return rc.get_required_resources(ds)
 
 

pytabkit/models/training/lightning_modules.py

@@ -341,5 +341,5 @@ def on_predict_model_eval(self) -> None:  # redundant with on_predict_start()
 
     def to(self, *args: Any, **kwargs: Any) -> 'TabNNModule':
         super().to(*args, **kwargs)
-        print(f'moving static model to {args} {kwargs}')
+        # print(f'moving static model to {args} {kwargs}')
         self.creator.static_model.to(*args, **kwargs)