zerohertzLib.mlops.server ¶

def _get_inputs(
    self, requests: list[Any]
) -> tuple[dict[str, NDArray[DTypeLike]], list[int]]:
    batch_index = [0]
    _inputs = defaultdict(list)
    for request in requests:
        for index, cfg_input in enumerate(self.cfg["input"]):
            value = pb_utils.get_input_tensor_by_name(
                request, cfg_input["name"]
            ).as_numpy()
            if index == 0 and 0 < self.max_batch_size:
                batch_index.append(batch_index[-1] + value.shape[0])
            _inputs[cfg_input["name"]].append(value)
    inputs = {}
    for key, value in _inputs.items():
        inputs[key] = np.concatenate(value, axis=0)
    logger.debug(
        "inputs: "
        + ", ".join([f"{key}={value.shape}" for key, value in inputs.items()])
    )
    return inputs, batch_index

@abstractmethod
def _inference(self, **inputs: NDArray[DTypeLike]) -> Any | tuple[Any]:
    """
    Model 추론을 수행하는 private method (상속을 통한 재정의 필수)

    Args:
        inputs: Model 추론 시 사용될 입력 (`config.pbtxt` 의 입력에 따라 입력 결정)

    Returns:
        Model의 추론 결과
    """
    pass

def _set_outputs(self, outputs: tuple[Any], batch_index: list[int]) -> list[Any]:
    responses = []
    if 0 < self.max_batch_size:
        for index in range(len(batch_index) - 1):
            batch_tensors = defaultdict(list)
            for batch in range(batch_index[index], batch_index[index + 1]):
                for cfg_output, value in zip(self.cfg["output"], outputs):
                    _value = value[batch]
                    if cfg_output["name"] == "batch_index":
                        _value -= batch_index[index]
                    batch_tensors[cfg_output["name"]].append(_value)
            output_tensors = []
            for cfg_output in self.cfg["output"]:
                value = np.concatenate(batch_tensors[cfg_output["name"]], axis=0)
                output_tensors.append(
                    pb_utils.Tensor(
                        cfg_output["name"],
                        value.astype(
                            pb_utils.triton_string_to_numpy(cfg_output["data_type"])
                        ),
                    )
                )
            responses.append(
                pb_utils.InferenceResponse(output_tensors=output_tensors)
            )
            logger.debug(
                f"outputs ({batch_index[index]} ~ {batch_index[index + 1] - 1}): "
                + ", ".join(
                    [
                        f"{key}="
                        + ", ".join([f"{_value.shape}" for _value in value])
                        for key, value in batch_tensors.items()
                    ]
                )
            )
        return responses
    output_tensors = []
    for cfg_output, value in zip(self.cfg["output"], outputs):
        output_tensors.append(
            pb_utils.Tensor(
                cfg_output["name"],
                value.astype(
                    pb_utils.triton_string_to_numpy(cfg_output["data_type"])
                ),
            )
        )
    responses.append(pb_utils.InferenceResponse(output_tensors=output_tensors))
    logger.debug(
        "outputs: "
        + ", ".join(
            [
                f"""{key["name"]}={value.shape}"""
                for key, value in zip(self.cfg["output"], outputs)
            ]
        )
    )
    return responses

def execute(self, requests: list[Any]) -> list[Any]:
    """Triton Inference Server 호출 시 수행되는 method

    Args:
        requests: Client에서 전송된 model inputs

    Returns:
        Client에 응답할 model의 추론 결과
    """
    logger.info("Called")
    try:
        inputs, batch_index = self._get_inputs(requests=requests)
        logger.info("Inference start")
        start = time.time()
        outputs = self._inference(**inputs)
        end = time.time()
        logger.info(f"Inference completed ({end - start:.2f}s)")
        if not isinstance(outputs, tuple):
            outputs = tuple([outputs])
        responses = self._set_outputs(outputs=outputs, batch_index=batch_index)
    except Exception as exc:
        logger.exception(exc)
        responses = [
            pb_utils.InferenceResponse(
                output_tensors=[], error=pb_utils.TritonError(exc)
            )
            for _ in requests
        ]
    return responses

def finalize(self) -> None:
    """Triton Inference Server 종료 시 수행되는 method"""
    logger.info("Finalize")

def initialize(self, args: dict[str, str]) -> None:
    """Triton Inference Server 시작 시 수행되는 method

    Args:
        args: `config.pbtxt` 에 포함된 model의 정보
    """
    self.cfg = json.loads(args["model_config"])
    logger.info(f"Initialize: {json.dumps(self.cfg, indent=4)}")
    self.device = "cpu"
    device = args.get("model_instance_device_id", None)
    if device is not None:
        self.device = f"cuda:{device}"
    self.max_batch_size = self.cfg.get("max_batch_size", 0)