Rework code calculating memory usage - output and Docker container (#…

…1161)

* Methods to calculate memory stats of Container and output run

Install psutil

Functions to calculate memory usage in Docker container and outputs

Removing start_tracking malloc function

* resolve conflicts

* Fix log_files_metric

* Implement the log_file_metrics method

* Reformat with black

* Revert "Fix log_files_metric"

This reverts commit c2993d3.

* Revert "Implement the log_file_metrics method"

This reverts commit e264286.

* Revert "Reformat with black"

This reverts commit 3d5ea63.

* Docker container memory info methods moved to docker.py

Reformatted with black

Rework the peak function method

Docker container memory info methods moved to docker.py

* Resolve conflicts

ersilia/cli/commands/close.py

@@ -4,7 +4,7 @@
 from .. import echo
 from ... import ErsiliaModel
 from ...core.session import Session
-from ...core.tracking import check_file_exists, close_persistent_file
+from ...core.tracking import check_file_exists,  close_persistent_file
 
 
 def close_cmd():
@@ -23,4 +23,4 @@ def close():
 
         # Close our persistent tracking file
         if check_file_exists(model_id):
-            close_persistent_file(mdl.model_id)
+           close_persistent_file(mdl.model_id)

ersilia/cli/commands/serve.py

@@ -1,9 +1,9 @@
 import click
-from . import ersilia_cli
+
 from .. import echo
+from . import ersilia_cli
 from ... import ErsiliaModel
 from ..messages import ModelNotFound
-from ...core.tracking import create_persistent_file
 
 
 def serve_cmd():
@@ -68,7 +68,3 @@ def serve(model, lake, docker, port, track):
         echo("")
         echo(":person_tipping_hand: Information:", fg="blue")
         echo("   - info", fg="blue")
-
-        # Setup persistent tracking
-        if track:
-            create_persistent_file(mdl.model_id)

ersilia/core/model.py

@@ -1,33 +1,32 @@
 import os
+import csv
 import json
-import tempfile
+import time
 import types
-import collections
+import tempfile
 import importlib
+import collections
 import __main__ as main
-import time
-import csv
 
 from .. import logger
+from ..serve.api import Api
+from .session import Session
+from datetime import datetime
 from .base import ErsiliaBase
+from ..lake.base import LakeBase
+from ..utils import tmp_pid_file
 from .modelbase import ModelBase
-from .session import Session
-from .tracking import RunTracker
-from ..serve.autoservice import AutoService
 from ..serve.schema import ApiSchema
-from ..serve.api import Api
-from ..serve.standard_api import StandardCSVRunApi
-from ..io.input import ExampleGenerator, BaseIOGetter
-from ..io.output import TabularOutputStacker
-from ..io.readers.file import FileTyper, TabularFileReader
-from ..utils import tmp_pid_file
 from ..utils.hdf5 import Hdf5DataLoader
 from ..utils.csvfile import CsvDataLoader
 from ..utils.terminal import yes_no_input
-from ..lake.base import LakeBase
-
+from ..serve.autoservice import AutoService
+from ..io.output import TabularOutputStacker
+from ..serve.standard_api import StandardCSVRunApi
+from ..io.input import ExampleGenerator, BaseIOGetter
+from .tracking import RunTracker, create_persistent_file
+from ..io.readers.file import FileTyper, TabularFileReader
 from ..utils.exceptions_utils.api_exceptions import ApiSpecifiedOutputError
-
 from ..default import FETCHED_MODELS_FILENAME, MODEL_SIZE_FILE, CARD_FILE, EOS
 from ..default import DEFAULT_BATCH_SIZE, APIS_LIST_FILE, INFORMATION_FILE
 
@@ -48,7 +47,6 @@ def __init__(
         verbose=None,
         fetch_if_not_available=True,
         preferred_port=None,
-        log_runs=True,
         track_runs=False,
     ):
         ErsiliaBase.__init__(
@@ -126,15 +124,11 @@ def __init__(
         )
         self._set_apis()
         self.session = Session(config_json=self.config_json)
-        if log_runs:
-            self._run_logger = RunTracker(
-                model_id=self.model_id, config_json=self.config_json
-            )
-        else:
-            self._run_logger = None
 
         if track_runs:
-            self._run_tracker = self._run_logger
+            self._run_tracker = RunTracker(
+                model_id=self.model_id, config_json=self.config_json
+            )
         else:
             self._run_tracker = None
 
@@ -430,6 +424,17 @@ def serve(self):
         self.scl = self.autoservice._service_class
         # self.update_model_usage_time(self.model_id) TODO: Check and reactivate
 
+        # Start tracking to get the peak memory, memory usage and cpu time of the Model server(autoservice)
+        if self._run_tracker is not None:
+            create_persistent_file(self.model_id)
+            memory_usage_serve, cpu_time_serve = self._run_tracker.get_memory_info()
+            peak_memory_serve = self._run_tracker.get_peak_memory()
+
+            session = Session(config_json=None)
+            session.update_peak_memory(peak_memory_serve)
+            session.update_total_memory(memory_usage_serve)
+            session.update_cpu_time(cpu_time_serve)
+
     def close(self):
         self.autoservice.close()
         self.session.close()
@@ -440,18 +445,12 @@ def get_apis(self):
     def _run(
         self, input=None, output=None, batch_size=DEFAULT_BATCH_SIZE, track_run=False
     ):
-        # Init some tracking before the run starts
-        if self._run_tracker is not None and track_run:
-            self._run_tracker.start_tracking()
 
         api_name = self.get_apis()[0]
         result = self.api(
             api_name=api_name, input=input, output=output, batch_size=batch_size
         )
-        if self._run_logger is not None:
-            self._run_logger.log(result=result, meta=self._model_info)
-        if self._run_tracker is not None and track_run:
-            self._run_tracker.track(input=input, result=result, meta=self._model_info)
+
         return result
 
     def _standard_run(self, input=None, output=None):
@@ -497,6 +496,12 @@ def run(
             result = self._run(
                 input=input, output=output, batch_size=batch_size, track_run=track_run
             )
+            # Start tracking model run if track flag is used in serve
+            if self._run_tracker is not None and track_run:
+                self._run_tracker.track(
+                    input=input, result=result, meta=self._model_info
+                )
+                self._run_tracker.log(result=result, meta=self._model_info)
             return result
 
     @property

ersilia/core/session.py

@@ -1,12 +1,12 @@
+import os
+import csv
 import json
 import time
 import uuid
-import os
-import csv
 import shutil
 
-from .base import ErsiliaBase
 from ..default import EOS
+from .base import ErsiliaBase
 
 
 class Session(ErsiliaBase):
@@ -69,6 +69,46 @@ def get(self):
             self.logger.debug("No session exists")
             return None
 
+    def update_total_memory(self, additional_memory):
+        """
+        Methods to get the the total memory of processess during serving and running.
+        """
+        data = self.get()
+        if data is None:
+            data = {}
+        current_memory = float(data.get("total memory used by model(MB)", 0))
+        new_memory = current_memory + additional_memory
+        data["total memory used by model(MB)"] = f"{new_memory:.5f}"
+        with open(self.session_file, "w") as f:
+            json.dump(data, f, indent=4)
+
+    def update_cpu_time(self, cpu_time):
+        """
+        Methods to get the total cpu time of processess during serving and running
+        """
+        data = self.get()
+        if data is None:
+            data = {}
+        current_cpu = float(data.get("CPU time used by model(seconds)", 0))
+        new_cpu = current_cpu + cpu_time
+        data["CPU time used by model(seconds)"] = f"{new_cpu}"
+        with open(self.session_file, "w") as f:
+            json.dump(data, f, indent=4)
+
+    def update_peak_memory(self, peak_memory):
+        """
+        Update peak memory usage in session data if the new peak is higher.
+        """
+        data = self.get()
+        if "peak memory used by model(MiB)" in data:
+            stored_peak_memory = float(data["peak memory used by model(MiB)"])
+            if peak_memory > stored_peak_memory:
+                data["peak memory used by model(MiB)"] = f"{peak_memory:.5f}"
+        else:
+            data["peak memory used by model(MiB)"] = f"{peak_memory:.5f}"
+        with open(self.session_file, "w") as f:
+            json.dump(data, f, indent=4)
+
     def close(self):
         self.logger.debug("Closing session {0}".format(self.session_file))
         if os.path.isfile(self.session_file):