triton-model-analyzer 1.48.0__py3-none-any.whl

model_analyzer/monitor/dcgm/DcgmHandle.py

@@ -0,0 +1,141 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import model_analyzer.monitor.dcgm.pydcgm as pydcgm
+import model_analyzer.monitor.dcgm.dcgm_structs as dcgm_structs
+import model_analyzer.monitor.dcgm.dcgm_agent as dcgm_agent
+
+
+class DcgmHandle:
+    '''
+    Class to encapsulate a handle to DCGM and global methods to control + query the host engine
+    '''
+
+    def __init__(self,
+                 handle=None,
+                 ipAddress=None,
+                 opMode=dcgm_structs.DCGM_OPERATION_MODE_AUTO,
+                 persistAfterDisconnect=False,
+                 unixSocketPath=None,
+                 timeoutMs=0):
+        '''
+        Constructor
+
+        handle is an existing handle from dcgmInit(). Pass None if you want this object to handle DCGM initialization for you
+        ipAddress is the host to connect to. None = start embedded host engine
+        opMode is a dcgm_structs.DCGM_OPERATION_MODE_* constant for how the host engine should run (embedded mode only)
+        persistAfterDisconnect (TCP-IP connections only) is whether the host engine should persist all of our watches
+                               after we disconnect. 1=persist our watches. 0=clean up after our connection
+        unixSocketPath is a path to a path on the local filesystem that is a unix socket that the host engine is listening on.
+                       This option is mutually exclusive with ipAddress
+        timeoutMs is how long to wait for TCP/IP or Unix domain connections to establish in ms. 0=Default timeout (5000ms)
+        '''
+        self._handleCreated = False
+        self._persistAfterDisconnect = persistAfterDisconnect
+
+        if handle is not None:
+            self.handle = handle
+            return
+
+        self._ipAddress = ipAddress
+
+        #Can't provide both unix socket and ip address
+        if ipAddress is not None and unixSocketPath is not None:
+            raise dcgm_structs.dcgmExceptionClass(dcgm_structs.DCGM_ST_BADPARAM)
+
+        #Initialize the DCGM client library
+        dcgm_structs._dcgmInit()
+        dcgm_agent.dcgmInit(
+        )  #Not harmful to call this multiple times in a process
+
+        #If neither ipAddress nor unixSocketPath are present, start an embedded host engine
+        if ipAddress is None and unixSocketPath is None:
+            self.handle = dcgm_agent.dcgmStartEmbedded(opMode)
+            self.isEmbedded = True
+            self._handleCreated = True
+            return
+
+        #Set up connection parameters. We're connecting to something
+        connectParams = dcgm_structs.c_dcgmConnectV2Params_v2()
+        connectParams.version = dcgm_structs.c_dcgmConnectV2Params_version
+        connectParams.timeoutMs = timeoutMs
+        if self._persistAfterDisconnect:
+            connectParams.persistAfterDisconnect = 1
+        else:
+            connectParams.persistAfterDisconnect = 0
+
+        if ipAddress is not None:
+            connectToAddress = ipAddress
+            connectParams.addressIsUnixSocket = 0
+        else:
+            connectToAddress = unixSocketPath
+            connectParams.addressIsUnixSocket = 1
+
+        self.handle = dcgm_agent.dcgmConnect_v2(connectToAddress, connectParams)
+        self.isEmbedded = False
+        self._handleCreated = True
+
+    def __del__(self):
+        '''
+        Destructor
+        '''
+        if self._handleCreated:
+            self.Shutdown()
+
+    def GetSystem(self):
+        '''
+        Get a DcgmSystem instance for this handle
+        '''
+        return pydcgm.DcgmSystem(self)
+
+    def __StopDcgm__(self):
+        '''
+        Shuts down either the hostengine or the embedded server
+        '''
+        if self.isEmbedded:
+            dcgm_agent.dcgmStopEmbedded(self.handle)
+        else:
+            dcgm_agent.dcgmDisconnect(self.handle)
+
+    def Shutdown(self):
+        '''
+        Shutdown DCGM hostengine
+        '''
+        if not self._handleCreated:
+            return
+
+        try:
+            self.__StopDcgm__()
+        except AttributeError as e:
+            # Due to multi-threading, sometimes this is called after the modules have been unloaded, making
+            # dcgm_agent effectively NoneType and resulting in this error being thrown.
+            pass
+
+        self._handleCreated = False
+        self.handle = None
+
+    @staticmethod
+    def Unload():
+        '''
+        Unload DCGM, removing any memory it is pointing at. Use this if you really
+        want DCGM gone from your process. Shutdown() only closes the connection/embedded host engine
+        that was create in __init__().
+        '''
+        dcgm_agent.dcgmShutdown()
+
+    def GetIpAddress(self):
+        '''
+        Returns the IP address associated with this handle. None=embedded connection
+        '''
+        return self._ipAddress

model_analyzer/monitor/dcgm/DcgmJsonReader.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from model_analyzer.monitor.dcgm.DcgmReader import DcgmReader
+from json import dumps as toJson
+from os import environ
+from socket import socket, AF_INET, SOCK_DGRAM
+from time import sleep
+import model_analyzer.monitor.dcgm.dcgm_fields as dcgm_fields
+import logging
+
+
+class DcgmJsonReader(DcgmReader):
+
+    ###########################################################################
+    def ConvertFieldIdToTag(self, fieldId):
+        return self.m_fieldIdToInfo[fieldId].tag
+
+    ###########################################################################
+    def PrepareJson(self, gpuId, obj):
+        '''
+        Receive an object with measurements turn it into an equivalent JSON. We
+        add the GPU UUID first.
+        '''
+        uuid = self.m_gpuIdToUUId[gpuId]
+        # This mutates the original object, but it shouldn't be a problem here
+        obj['gpu_uuid'] = uuid
+        return toJson(obj)
+
+    ###########################################################################
+    def CustomDataHandler(self, fvs):
+        for gpuId in list(fvs.keys()):
+            # We don't need the keys because each value has a `fieldId`
+            # So just get the values
+            gpuData = list(fvs[gpuId].values())
+
+            # Get the values from FV (which is a list of values)
+            valuesListOfLists = [datum.values for datum in gpuData]
+
+            # We only want the last measurement
+            lastValueList = [l[-1] for l in valuesListOfLists]
+
+            # Turn FV into a conventional Python Object which can be converted to JSON
+            outObject = {
+                self.ConvertFieldIdToTag(i.fieldId): i.value
+                for i in lastValueList
+            }
+            outJson = self.PrepareJson(gpuId, outObject)
+
+            self.CustomJsonHandler(outJson)
+
+    ###########################################################################
+    def CustomJsonHandler(self, outJson):
+        '''
+        This method should be overriden by subclasses to handle the JSON objects
+        received.
+        '''
+        logging.warning('CustomJsonHandler has not been overriden')
+        logging.info(outJson)