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

model_analyzer/state/analyzer_state_manager.py

@@ -0,0 +1,215 @@
+#!/usr/bin/env python3
+
+# Copyright 2021-2023, 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 glob
+import json
+import logging
+import os
+import signal
+import sys
+import traceback
+
+from model_analyzer.constants import LOGGER_NAME, MAX_NUMBER_OF_INTERRUPTS
+from model_analyzer.model_analyzer_exceptions import TritonModelAnalyzerException
+from model_analyzer.state.analyzer_state import AnalyzerState
+
+logger = logging.getLogger(LOGGER_NAME)
+
+
+class AnalyzerStateManager:
+    """
+    Maintains the state of the Model Analyzer
+    """
+
+    def __init__(self, config, server):
+        """
+        Parameters
+        ----------
+        config: ConfigCommand
+            The analyzer's config
+        server : TritonServer
+            Handle for tritonserver instance
+        """
+
+        self._config = config
+        self._server = server
+        self._exiting = 0
+        self._checkpoint_dir = config.checkpoint_directory
+        self._state_changed = False
+
+        if os.path.exists(self._checkpoint_dir):
+            self._checkpoint_index = self._latest_checkpoint() + 1
+        else:
+            os.makedirs(self._checkpoint_dir)
+            self._checkpoint_index = 0
+        signal.signal(signal.SIGINT, self.interrupt_handler)
+
+        self._current_state = AnalyzerState()
+        self._starting_fresh_run = True
+
+    def starting_fresh_run(self):
+        """
+        Returns
+        -------
+        True if starting a fresh run
+        False if checkpoint found and loaded
+        """
+
+        return self._starting_fresh_run
+
+    def exiting(self):
+        """
+        Returns
+        -------
+        True if interrupt handler ran
+        even once, False otherwise
+        """
+
+        return self._exiting > 0
+
+    def get_state_variable(self, name):
+        """
+        Get a named variable from
+        the current AnalyzerState
+
+        Parameters
+        ----------
+        name : str
+            The name of the variable
+        """
+        return self._current_state.get(name)
+
+    def set_state_variable(self, name, value):
+        """
+        Set a named variable from
+        the current AnalyzerState
+
+        Parameters
+        ----------
+        name: str
+            The name of the variable
+        value: Any
+            the value to set for that variable
+        """
+
+        self._state_changed = True
+        self._current_state.set(name, value)
+
+    def load_checkpoint(self, checkpoint_required):
+        """
+        Load the state of the Model Analyzer from
+        most recent checkpoint file, also
+        set whether we are starting a fresh run
+
+        Parameters
+        ----------
+        checkpoint_required : bool
+            If true, an existing checkpoint is required to run MA
+        """
+
+        latest_checkpoint_file = os.path.join(
+            self._checkpoint_dir, f"{self._latest_checkpoint()}.ckpt"
+        )
+        if os.path.exists(latest_checkpoint_file):
+            logger.info(f"Loaded checkpoint from file {latest_checkpoint_file}")
+            with open(latest_checkpoint_file, "r") as f:
+                try:
+                    self._current_state = AnalyzerState.from_dict(json.load(f))
+                except EOFError:
+                    raise TritonModelAnalyzerException(
+                        f"Checkpoint file {latest_checkpoint_file} is"
+                        " empty or corrupted. Remove it from checkpoint"
+                        " directory."
+                    )
+            self._starting_fresh_run = False
+        else:
+            if checkpoint_required:
+                raise TritonModelAnalyzerException(f"No checkpoint file found")
+            else:
+                logger.info("No checkpoint file found, starting a fresh run.")
+
+    def default_encode(self, obj):
+        if isinstance(obj, bytes):
+            return obj.decode("utf-8")
+        elif hasattr(obj, "to_dict"):
+            return obj.to_dict()
+        else:
+            return obj.__dict__
+
+    def save_checkpoint(self):
+        """
+        Saves the state of the model analyzer to disk
+        if there has been a change since the last checkpoint
+
+        Parameters
+        ----------
+        state: AnalyzerState
+            The state object to be saved
+        """
+
+        ckpt_filename = os.path.join(
+            self._checkpoint_dir, f"{self._checkpoint_index}.ckpt"
+        )
+        if self._state_changed:
+            with open(ckpt_filename, "w") as f:
+                json.dump(self._current_state, f, default=self.default_encode)
+            logger.info(f"Saved checkpoint to {ckpt_filename}")
+
+            self._state_changed = False
+        else:
+            logger.info(f"No changes made to analyzer data, no checkpoint saved.")
+
+    def interrupt_handler(self, signal, frame):
+        """
+        A signal handler to properly
+        shutdown the model analyzer on
+        interrupt
+        """
+
+        self._exiting += 1
+        if logger.getEffectiveLevel() <= logging.DEBUG:
+            traceback.print_stack(limit=15)
+        logger.info(
+            f"Received SIGINT {self._exiting}/{MAX_NUMBER_OF_INTERRUPTS}. "
+            "Will attempt to exit after current measurement."
+        )
+        if self._exiting >= MAX_NUMBER_OF_INTERRUPTS:
+            logger.info(
+                f"Received SIGINT maximum number of times. Saving state and exiting immediately. "
+                "perf_analyzer may still be running"
+            )
+            self.save_checkpoint()
+
+            # Exit server
+            if self._server:
+                self._server.stop()
+            sys.exit(1)
+
+    def _latest_checkpoint(self):
+        """
+        Get the highest index checkpoint file in the
+        checkpoint directory, return its index.
+        """
+
+        checkpoint_files = glob.glob(os.path.join(self._checkpoint_dir, "*.ckpt"))
+        if not checkpoint_files:
+            return -1
+        try:
+            return max(
+                [int(os.path.split(f)[1].split(".")[0]) for f in checkpoint_files]
+            )
+        except Exception as e:
+            raise TritonModelAnalyzerException(e)

model_analyzer/triton/__init__.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+
+# Copyright 2020-2023, 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.

model_analyzer/triton/client/__init__.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+
+# Copyright 2020-2023, 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.

model_analyzer/triton/client/client.py

@@ -0,0 +1,234 @@
+#!/usr/bin/env python3
+
+# Copyright 2020-2023, 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 logging
+import time
+from subprocess import DEVNULL
+
+from model_analyzer.constants import LOGGER_NAME
+from model_analyzer.model_analyzer_exceptions import TritonModelAnalyzerException
+
+logger = logging.getLogger(LOGGER_NAME)
+
+
+class TritonClient:
+    """
+    Defines the interface for the objects created by
+    TritonClientFactory
+    """
+
+    def wait_for_server_ready(
+        self,
+        num_retries,
+        sleep_time=1,
+        log_file=None,
+    ):
+        """
+        Parameters
+        ----------
+        num_retries : int
+            number of times to send a ready status
+            request to the server before raising
+            an exception
+        sleep_time: int
+            amount of time in seconds to sleep between retries
+        log_file: TextIOWrapper
+            file that contains the server's output log
+        Raises
+        ------
+        TritonModelAnalyzerException
+            If server readiness could not be
+            determined in given num_retries
+        """
+
+        retries = num_retries
+        while retries > 0:
+            try:
+                if self._client.is_server_ready():
+                    time.sleep(sleep_time)
+                    return
+                else:
+                    self._check_for_triton_log_errors(log_file)
+                    time.sleep(sleep_time)
+                    retries -= 1
+            except Exception as e:
+                # Log connection failures with more detail for debugging
+                if retries == num_retries or retries % 10 == 0:
+                    logger.debug(
+                        f"Failed to connect to Triton server (attempt {num_retries - retries + 1}/{num_retries}): {e}"
+                    )
+                self._check_for_triton_log_errors(log_file)
+                time.sleep(sleep_time)
+                retries -= 1
+                if retries == 0:
+                    raise TritonModelAnalyzerException(e)
+        raise TritonModelAnalyzerException(
+            "Could not determine server readiness. " "Number of retries exceeded."
+        )
+
+    def load_model(self, model_name, variant_name="", config_str=None):
+        """
+        Request the inference server to load
+        a particular model in explicit model
+        control mode.
+
+        Parameters
+        ----------
+        model_name : str
+            Name of the model
+
+        variant_name: str
+            Name of the model variant
+
+        config_str: str
+            Optional config string used to load the model
+
+        Returns
+        ------
+        int or None
+            Returns -1 if the failed.
+        """
+
+        variant_name = variant_name if variant_name else model_name
+
+        try:
+            self._client.load_model(model_name, config=config_str)
+            logger.debug(f"Model {variant_name} loaded")
+            return None
+        except Exception as e:
+            logger.info(f"Model {variant_name} load failed: {e}")
+            if "polling is enabled" in e.message():
+                raise TritonModelAnalyzerException(
+                    "The remote Tritonserver needs to be launched in EXPLICIT mode"
+                )
+            return -1
+
+    def unload_model(self, model_name):
+        """
+        Request the inference server to unload
+        a particular model in explicit model
+        control mode.
+
+        Parameters
+        ----------
+        model_name : str
+            name of the model to load from repository
+
+        Raises
+        ------
+        TritonModelAnalyzerException
+            If server throws Exception
+
+        Returns
+        ------
+        int or None
+            Returns -1 if the failed.
+        """
+
+        try:
+            self._client.unload_model(model_name)
+            logger.debug(f"Model {model_name} unloaded")
+            return None
+        except Exception as e:
+            logger.info(f"Model {model_name} unload failed: {e}")
+            return -1
+
+    def wait_for_model_ready(self, model_name, num_retries, sleep_time=1):
+        """
+        Returns when model is ready.
+
+        Parameters
+        ----------
+        model_name : str
+            name of the model to load from repository
+        num_retries : int
+            number of times to send a ready status
+            request to the server before raising
+            an exception
+
+        Raises
+        ------
+        TritonModelAnalyzerException
+            If could not determine model readiness
+            in given num_retries
+
+        Returns
+        ------
+        int or None
+            Returns -1 if the failed.
+        """
+
+        retries = num_retries
+        error = None
+        while retries > 0:
+            try:
+                if self._client.is_model_ready(model_name):
+                    return None
+                else:
+                    time.sleep(sleep_time)
+                    retries -= 1
+            except Exception as e:
+                error = e
+                time.sleep(sleep_time)
+                retries -= 1
+
+        logger.info(f"Model readiness failed for model {model_name}. Error {error}")
+        return -1
+
+    def get_model_config(self, model_name, num_retries):
+        """
+        Model name to get the config for.
+
+        Parameters
+        ----------
+        model_name : str
+            Name of the model to find the config.
+
+        num_retries : int
+            Number of times to wait for the model load
+
+        Returns
+        -------
+        dict or None
+            A dictionary containing the model config.
+        """
+
+        self.wait_for_model_ready(model_name, num_retries)
+        model_config_dict = self._client.get_model_config(model_name)
+        return model_config_dict
+
+    def is_server_ready(self):
+        """
+        Returns true if the server is ready. Else False
+        """
+        return self._client.is_server_ready()
+
+    def _check_for_triton_log_errors(self, log_file):
+        if not log_file or log_file == DEVNULL:
+            return
+
+        log_file.seek(0)
+        log_output = log_file.read()
+
+        if not type(log_output) == str:
+            log_output = log_output.decode("utf-8")
+
+        if log_output:
+            if "Unexpected argument:" in log_output:
+                error_start = log_output.find("Unexpected argument:")
+                raise TritonModelAnalyzerException(
+                    f"Error: TritonServer did not launch successfully\n\n{log_output[error_start:]}"
+                )

model_analyzer/triton/client/client_factory.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+
+# Copyright 2020-2023, 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 .grpc_client import TritonGRPCClient
+from .http_client import TritonHTTPClient
+
+
+class TritonClientFactory:
+    """
+    Base client creator class that declares
+    a factory method
+    """
+
+    @staticmethod
+    def create_grpc_client(server_url, ssl_options={}):
+        """
+        Parameters
+        ----------
+        server_url : str
+            The url for Triton server's GRPC endpoint
+        ssl_options : dict
+            Dictionary of SSL options for gRPC python client
+
+        Returns
+        -------
+        TritonGRPCClient
+        """
+        return TritonGRPCClient(server_url=server_url, ssl_options=ssl_options)
+
+    @staticmethod
+    def create_http_client(server_url, ssl_options={}):
+        """
+        Parameters
+        ----------
+        server_url : str
+            The url for Triton server's HTTP endpoint
+        ssl_options : dict
+            Dictionary of SSL options for HTTP python client
+
+        Returns
+        -------
+        TritonHTTPClient
+        """
+        return TritonHTTPClient(server_url=server_url, ssl_options=ssl_options)

model_analyzer/triton/client/grpc_client.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+
+# Copyright 2020-2023, 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 tritonclient.grpc as grpcclient
+
+from .client import TritonClient
+
+
+class TritonGRPCClient(TritonClient):
+    """
+    Concrete implementation of TritonClient
+    for GRPC
+    """
+
+    def __init__(self, server_url, ssl_options={}):
+        """
+        Parameters
+        ----------
+        server_url : str
+            The url for Triton server's GRPC endpoint
+        ssl_options : dict
+            Dictionary of SSL options for gRPC python client
+        """
+
+        ssl = False
+        root_certificates = None
+        private_key = None
+        certificate_chain = None
+
+        if "ssl-grpc-use-ssl" in ssl_options:
+            ssl = ssl_options["ssl-grpc-use-ssl"].lower() == "true"
+        if "ssl-grpc-root-certifications-file" in ssl_options:
+            root_certificates = ssl_options["ssl-grpc-root-certifications-file"]
+        if "ssl-grpc-private-key-file" in ssl_options:
+            private_key = ssl_options["ssl-grpc-private-key-file"]
+        if "ssl-grpc-certificate-chain-file" in ssl_options:
+            certificate_chain = ssl_options["ssl-grpc-certificate-chain-file"]
+
+        # Fix for gRPC 1.60.0+: Force IPv4 resolution for localhost connections
+        # gRPC 1.60.0+ prefers IPv6, causing "localhost" to resolve to [::1]
+        # On systems where IPv6 is not properly configured, this causes connection failures
+        # Force IPv4 by using 127.0.0.1, which is more reliable across environments
+        channel_args = None
+        if "localhost" in server_url:
+            server_url = server_url.replace("localhost", "127.0.0.1")
+            # For SSL connections, override target name to match certificate
+            if ssl:
+                channel_args = [("grpc.ssl_target_name_override", "localhost")]
+
+        self._client = grpcclient.InferenceServerClient(
+            url=server_url,
+            ssl=ssl,
+            root_certificates=root_certificates,
+            private_key=private_key,
+            certificate_chain=certificate_chain,
+            channel_args=channel_args,
+        )
+
+    def get_model_config(self, model_name, num_retries):
+        """
+        Model name to get the config for.
+
+        Parameters
+        ----------
+        model_name : str
+            Name of the model to find the config.
+
+        num_retries : int
+            Number of times to wait for the model load
+
+        Returns
+        -------
+        dict
+            A dictionary containing the model config.
+        """
+
+        self.wait_for_model_ready(model_name, num_retries)
+        model_config_dict = self._client.get_model_config(model_name, as_json=True)
+        return model_config_dict["config"]
+
+    def get_model_repository_index(self):
+        """
+        Returns the JSON dict holding the model repository index.
+        """
+        return self._client.get_model_repository_index(as_json=True)["models"]
+
+    def is_model_ready(self, model_name: str) -> bool:
+        """
+        Returns true if the model is loaded on the server
+        """
+        return self._client.is_model_ready(model_name)