Merge pull request #22 from MoseleyBioinformaticsLab/text

Modifies text format to reflect JSON
MoseleyBioinformaticsLab · Mar 31, 2024 · 614e390 · 614e390
2 parents e3c3402 + 3b396aa
commit 614e390
Show file tree

Hide file tree

Showing 34 changed files with 879 additions and 494 deletions.
diff --git a/src/gpu_tracker/tracker.py b/src/gpu_tracker/tracker.py
@@ -21,25 +21,36 @@ class RSSValues:
 
 @dclass.dataclass
 class MaxRAM:
+    unit: str
+    system_capacity: float
+    system: float = 0.
     main: RSSValues = dclass.field(default_factory=RSSValues)
     descendents: RSSValues = dclass.field(default_factory=RSSValues)
     combined: RSSValues = dclass.field(default_factory=RSSValues)
-    system: float = 0.
 
 
 @dclass.dataclass
 class MaxGPURAM:
+    unit: str
     main: float = 0.
     descendents: float = 0.
     combined: float = 0.
 
 
+@dclass.dataclass
+class ComputeTime:
+    unit: str
+    time: float = 0.
+
+
 class Tracker:
     """
     Runs a thread in the background that tracks the compute time, maximum RAM, and maximum GPU RAM usage within a context manager or explicit ``start()`` and ``stop()`` methods.
-    Calculated quantities are scaled depending on the unit chosen for them (e.g. megabytes vs. gigabytes, hours vs. days, etc.).
+    Calculated quantities are scaled depending on the units chosen for them (e.g. megabytes vs. gigabytes, hours vs. days, etc.).
 
-    :ivar dict measurements: The measured values of the computational-resource usage i.e. maximum RAM, maximum GPU RAM, and compute time.
+    :ivar MaxRAM max_ram: Description of the maximum RAM usage of the process, any descendents it may have, and the operating system overall.
+    :ivar MaxGPURAM max_gpu_ram: Description of the maximum GPU RAM usage of the process and any descendents it may have.
+    :ivar ComputeTime compute_time: Description of the real compute time i.e. the duration of tracking.
     """
     NO_PROCESS_WARNING = 'Attempted to obtain RAM information of a process that no longer exists.'
 
@@ -61,9 +72,6 @@ def __init__(
         Tracker._validate_mem_unit(gpu_ram_unit)
         Tracker._validate_unit(time_unit, valid_units={'seconds', 'minutes', 'hours', 'days'}, unit_type='time')
         self.sleep_time = sleep_time
-        self.ram_unit = ram_unit
-        self.gpu_ram_unit = gpu_ram_unit
-        self.time_unit = time_unit
         self._ram_coefficient: float = {
             'bytes': 1.0,
             'kilobytes': 1 / 1e3,
@@ -92,10 +100,9 @@ def __init__(
         self.process_id = process_id if process_id is not None else os.getpid()
         self._main_process = psutil.Process(self.process_id)
         self._is_linux = platform.system().lower() == 'linux'
-        self.system_ram_capacity = psutil.virtual_memory().total * self._ram_coefficient
-        self.max_ram = MaxRAM()
-        self.max_gpu_ram = MaxGPURAM()
-        self.compute_time = 0.
+        self.max_ram = MaxRAM(unit=ram_unit, system_capacity=psutil.virtual_memory().total * self._ram_coefficient)
+        self.max_gpu_ram = MaxGPURAM(unit=gpu_ram_unit)
+        self.compute_time = ComputeTime(unit=time_unit)
 
     @staticmethod
     def _validate_mem_unit(unit: str):
@@ -178,7 +185,7 @@ def _profile(self):
                     process_ids.add(self.process_id)
                     self._update_gpu_ram(attr='combined', process_ids=process_ids, nvidia_smi_output=nvidia_smi_output)
                 # Update compute time
-                self.compute_time = (time.time() - start_time) * self._time_coefficient
+                self.compute_time.time = (time.time() - start_time) * self._time_coefficient
                 _testable_sleep(self.sleep_time)
             except psutil.NoSuchProcess:
                 log.warning('Failed to track a process that does not exist. '
@@ -225,28 +232,38 @@ def __str__(self) -> str:
         """
         Constructs a string representation of the computational-resource-usage measurements and their units.
         """
-        max_ram, max_gpu_ram, compute_time = (
-            f'{measurement:.3f} {unit}' if measurement is not None else 'null' for measurement, unit in (
-                (self.measurements['max_ram']['combined']['total_rss'], self.ram_unit),
-                (self.measurements['max_gpu_ram']['combined'], self.gpu_ram_unit),
-                (self.measurements['compute_time'], self.time_unit)))
-        return \
-            f'Max RAM (combined total RSS): {max_ram}\n' \
-            f'Max GPU RAM (combined): {max_gpu_ram}\n' \
-            f'Compute time: {compute_time}'
+        tracker_json = self.to_json()
+        Tracker._format_float(dictionary=tracker_json)
+        indent = 3
+        text = json.dumps(tracker_json, indent=indent)
+        # Un-indent the lines to the left after removing curley braces.
+        text = '\n'.join(line[indent:] for line in text.split('\n') if line.strip() not in {'{', '}', '},'})
+        text = text.replace(': {', ':').replace('{', '').replace('}', '').replace('_', ' ').replace('"', '').replace(',', '')
+        return text.replace('max', 'Max').replace('ram', 'RAM').replace('unit', 'Unit').replace('system', 'System').replace(
+            'compute', 'Compute').replace('time: ', 'Time: ').replace('rss', 'RSS').replace('total', 'Total').replace(
+            'private', 'Private').replace('shared', 'Shared').replace('main', 'Main').replace('descendents', 'Descendents').replace(
+            'combined', 'Combined').replace('gpu', 'GPU')
 
-    def __repr__(self) -> str:
-        return str(self)  # pragma: no cover
+    @staticmethod
+    def _format_float(dictionary: dict):
+        """
+        Recursively formats the floating points in a dictionary to 3 decimal places.
+        :param dictionary: The dictionary to format.
+        """
+        for key, value in dictionary.items():
+            if type(value) == float:
+                dictionary[key] = float(f'{value:.3f}')
+            elif type(value) == dict:
+                Tracker._format_float(value)
 
     def to_json(self):
+        """
+        Constructs a dictionary of the computational-resource-usage measurements and their units.
+        """
         return {
-            'system_ram_capacity': self.system_ram_capacity,
             'max_ram': dclass.asdict(self.max_ram),
-            'ram_unit': self.ram_unit,
             'max_gpu_ram': dclass.asdict(self.max_gpu_ram),
-            'gpu_ram_unit': self.gpu_ram_unit,
-            'compute_time': self.compute_time,
-            'time_unit': self.time_unit
+            'compute_time': dclass.asdict(self.compute_time),
         }
 
 

diff --git a/tests/data/False-Linux-bytes-megabytes-seconds.json b/tests/data/False-Linux-bytes-megabytes-seconds.json
@@ -1,30 +1,32 @@
 {
- "system_ram_capacity": 67000000000.0,
- "max_ram": {
-  "main": {
-   "total_rss": 16754.0,
-   "private_rss": 6718.0,
-   "shared_rss": 10036.0
+  "max_ram": {
+    "unit": "bytes",
+    "system_capacity": 67000000000.0,
+    "system": 31000000000.0,
+    "main": {
+      "total_rss": 16754.0,
+      "private_rss": 6718.0,
+      "shared_rss": 10036.0
+    },
+    "descendents": {
+      "total_rss": 30208.0,
+      "private_rss": 13818.0,
+      "shared_rss": 16390.0
+    },
+    "combined": {
+      "total_rss": 42553.0,
+      "private_rss": 20536.0,
+      "shared_rss": 22017.0
+    }
   },
-  "descendents": {
-   "total_rss": 30208.0,
-   "private_rss": 13818.0,
-   "shared_rss": 16390.0
+  "max_gpu_ram": {
+    "unit": "megabytes",
+    "main": 1600.0,
+    "descendents": 4300.0,
+    "combined": 5800.0
   },
-  "combined": {
-   "total_rss": 42553.0,
-   "private_rss": 20536.0,
-   "shared_rss": 22017.0
-  },
-  "system": 31000000000.0
- },
- "ram_unit": "bytes",
- "max_gpu_ram": {
-  "main": 1600.0,
-  "descendents": 4300.0,
-  "combined": 5800.0
- },
- "gpu_ram_unit": "megabytes",
- "compute_time": 300.0,
- "time_unit": "seconds"
+  "compute_time": {
+    "unit": "seconds",
+    "time": 300.0
+  }
 }
diff --git a/tests/data/False-Linux-bytes-megabytes-seconds.txt b/tests/data/False-Linux-bytes-megabytes-seconds.txt
@@ -1,3 +1,24 @@
-Max RAM (combined total RSS): 42553.000 bytes
-Max GPU RAM (combined): 5800.000 megabytes
-Compute time: 300.000 seconds
+Max RAM:
+   Unit: bytes
+   System capacity: 67000000000.0
+   System: 31000000000.0
+   Main:
+      Total RSS: 16754.0
+      Private RSS: 6718.0
+      Shared RSS: 10036.0
+   Descendents:
+      Total RSS: 30208.0
+      Private RSS: 13818.0
+      Shared RSS: 16390.0
+   Combined:
+      Total RSS: 42553.0
+      Private RSS: 20536.0
+      Shared RSS: 22017.0
+Max GPU RAM:
+   Unit: megabytes
+   Main: 1600.0
+   Descendents: 4300.0
+   Combined: 5800.0
+Compute time:
+   Unit: seconds
+   Time: 300.0
diff --git a/tests/data/False-Linux-kilobytes-bytes-days.json b/tests/data/False-Linux-kilobytes-bytes-days.json
@@ -1,30 +1,32 @@
 {
- "system_ram_capacity": 67000000.0,
- "max_ram": {
-  "main": {
-   "total_rss": 16.753999999999998,
-   "private_rss": 6.718,
-   "shared_rss": 10.036
+  "max_ram": {
+    "unit": "kilobytes",
+    "system_capacity": 67000000.0,
+    "system": 31000000.0,
+    "main": {
+      "total_rss": 16.753999999999998,
+      "private_rss": 6.718,
+      "shared_rss": 10.036
+    },
+    "descendents": {
+      "total_rss": 30.208,
+      "private_rss": 13.818,
+      "shared_rss": 16.39
+    },
+    "combined": {
+      "total_rss": 42.553,
+      "private_rss": 20.536,
+      "shared_rss": 22.017
+    }
   },
-  "descendents": {
-   "total_rss": 30.208,
-   "private_rss": 13.818,
-   "shared_rss": 16.39
+  "max_gpu_ram": {
+    "unit": "bytes",
+    "main": 1600000000.0,
+    "descendents": 4300000000.0,
+    "combined": 5800000000.0
   },
-  "combined": {
-   "total_rss": 42.553,
-   "private_rss": 20.536,
-   "shared_rss": 22.017
-  },
-  "system": 31000000.0
- },
- "ram_unit": "kilobytes",
- "max_gpu_ram": {
-  "main": 1600000000.0,
-  "descendents": 4300000000.0,
-  "combined": 5800000000.0
- },
- "gpu_ram_unit": "bytes",
- "compute_time": 0.003472222222222222,
- "time_unit": "days"
+  "compute_time": {
+    "unit": "days",
+    "time": 0.003472222222222222
+  }
 }
diff --git a/tests/data/False-Linux-kilobytes-bytes-days.txt b/tests/data/False-Linux-kilobytes-bytes-days.txt
@@ -1,3 +1,24 @@
-Max RAM (combined total RSS): 42.553 kilobytes
-Max GPU RAM (combined): 5800000000.000 bytes
-Compute time: 0.003 days
+Max RAM:
+   Unit: kilobytes
+   System capacity: 67000000.0
+   System: 31000000.0
+   Main:
+      Total RSS: 16.754
+      Private RSS: 6.718
+      Shared RSS: 10.036
+   Descendents:
+      Total RSS: 30.208
+      Private RSS: 13.818
+      Shared RSS: 16.39
+   Combined:
+      Total RSS: 42.553
+      Private RSS: 20.536
+      Shared RSS: 22.017
+Max GPU RAM:
+   Unit: bytes
+   Main: 1600000000.0
+   Descendents: 4300000000.0
+   Combined: 5800000000.0
+Compute time:
+   Unit: days
+   Time: 0.003
diff --git a/tests/data/False-Linux-kilobytes-gigabytes-minutes.json b/tests/data/False-Linux-kilobytes-gigabytes-minutes.json
@@ -1,30 +1,32 @@
 {
- "system_ram_capacity": 67000000.0,
- "max_ram": {
-  "main": {
-   "total_rss": 16.753999999999998,
-   "private_rss": 6.718,
-   "shared_rss": 10.036
+  "max_ram": {
+    "unit": "kilobytes",
+    "system_capacity": 67000000.0,
+    "system": 31000000.0,
+    "main": {
+      "total_rss": 16.753999999999998,
+      "private_rss": 6.718,
+      "shared_rss": 10.036
+    },
+    "descendents": {
+      "total_rss": 30.208,
+      "private_rss": 13.818,
+      "shared_rss": 16.39
+    },
+    "combined": {
+      "total_rss": 42.553,
+      "private_rss": 20.536,
+      "shared_rss": 22.017
+    }
   },
-  "descendents": {
-   "total_rss": 30.208,
-   "private_rss": 13.818,
-   "shared_rss": 16.39
+  "max_gpu_ram": {
+    "unit": "gigabytes",
+    "main": 1.6,
+    "descendents": 4.3,
+    "combined": 5.8
   },
-  "combined": {
-   "total_rss": 42.553,
-   "private_rss": 20.536,
-   "shared_rss": 22.017
-  },
-  "system": 31000000.0
- },
- "ram_unit": "kilobytes",
- "max_gpu_ram": {
-  "main": 1.6,
-  "descendents": 4.3,
-  "combined": 5.8
- },
- "gpu_ram_unit": "gigabytes",
- "compute_time": 5.0,
- "time_unit": "minutes"
+  "compute_time": {
+    "unit": "minutes",
+    "time": 5.0
+  }
 }
diff --git a/tests/data/False-Linux-kilobytes-gigabytes-minutes.txt b/tests/data/False-Linux-kilobytes-gigabytes-minutes.txt
@@ -1,3 +1,24 @@
-Max RAM (combined total RSS): 42.553 kilobytes
-Max GPU RAM (combined): 5.800 gigabytes
-Compute time: 5.000 minutes
+Max RAM:
+   Unit: kilobytes
+   System capacity: 67000000.0
+   System: 31000000.0
+   Main:
+      Total RSS: 16.754
+      Private RSS: 6.718
+      Shared RSS: 10.036
+   Descendents:
+      Total RSS: 30.208
+      Private RSS: 13.818
+      Shared RSS: 16.39
+   Combined:
+      Total RSS: 42.553
+      Private RSS: 20.536
+      Shared RSS: 22.017
+Max GPU RAM:
+   Unit: gigabytes
+   Main: 1.6
+   Descendents: 4.3
+   Combined: 5.8
+Compute time:
+   Unit: minutes
+   Time: 5.0