deploy¶

class Component:
    """Base class for pipeline components, with name, description, and a list of inputs and outputs.

    A new component is created with the given name and an empty input and output list.
    """
    reserved_names = ["timestamp"]

    @dataclass
    class BatchInfo:
        """
        Batch information for the component.

        This attribute specifies whether the component can handle batch input or output data.
        When set to True, the component will receive data in the form of a list of dictionaries instead of a single dictionary.
        It is important to note that the input and output variables on the component should still be defined as if they are single variables.

        If the input of the pipeline is configured for batch processing, it is recommended not to configure timeshifting, as the list will have the same timestamp for all elements, potentially resulting in data loss.
        """
        inputBatch: bool = False
        outputBatch: bool = False

        def dict(self):
            return {
                'inputBatch': 'Yes' if self.inputBatch is True else 'No',
                'outputBatch': 'Yes' if self.outputBatch is True else 'No'
            }

    def __init__(self, name: str, desc: str = ""):
        """
        Creates a new component with the given name and an empty input and output list.

        Args:
            name (str): Name of the component.
            desc (str): Optional description of the component
        """

        if _allowed_characters_in_names.match(name) is None:
            raise AssertionError("Component name contains invalid character. The allowed characters are [-a-zA-Z0-9_].")
        self.name = name
        self.desc = desc
        self.inputs = {}
        self.outputs = {}

        self.batch = self.BatchInfo(False, False)

    def __repr__(self) -> str:
        text = f"[{self.__class__.__name__}] {self.name}\n"
        if self.desc != "":
            text += f"{self.desc}\n"
        if len(self.inputs) > 0:
            text += "\nComponent Inputs:\n"
            for name, input in self.inputs.items():
                text += f"> {name} ({input['type']}){': ' + input['desc'] if input.get('desc') is not None else ''}\n"
        if len(self.outputs) > 0:
            text += "\nComponent Outputs:\n"
            for name, output in self.outputs.items():
                text += f"< {name} ({output['type']}){': ' + output['desc'] if output.get('desc') is not None else ''}\n"
        return text

    def add_input(self, name: str, _type: str, desc: Optional[str] = None):
        """
        Adds a new input to the component with its type.
        Name of the variables cannot be reserved name like 'timestamp'.
        Input variable 'timestamp' is a prebuilt key in the payload and its value contains the timestamp when the payload is created by AI Inference Server.

        Types supported by AI Inference Server version 1.6 are contained in the `type_dictionary`. Newer AI Inference server version may support additional types.
        In case the type is not known by the AI SDK, a warning message will be printed.
        The most frequently used types are
        - String:
            Typically used for data received from Databus
        - Object:
            Object type variables are designed to receive from Vision Connect or transfer images between components
        - Numeric scalar types:
            Typically used for data received from S7 Connector

        The example payload below shows the format of image received from VCA Connector
        ```python
            payload = { "image":
                {
                    "resolutionWidth": image.width,
                    "resolutionHeight": image.height,
                    "mimeType": ["image/raw"],
                    "dataType": "uint8",
                    "channelsPerPixel": 3,
                    "image": _swap_bytes(image.tobytes())
                }
            }
        ```
        Between components the format is the same format as the format of Object as an output.
        ```python
            "processedImage": {
                "metadata": json.dumps( {
                                "resolutionWidth": image.width,
                                "resolutionHeight": image.height
                                }
                            ),
                "bytes": image.tobytes()
            }
        ```

        Args:
            name (str): Name of the new input.
            _type (str): Type of the new input.
            desc (str): Description of the input. (optional)
        """
        if self.inputs is None:
            self.inputs = {}
        if name in self.inputs:
            raise AssertionError(f"Input '{name}' already exists.")
        if name.lower() in self.reserved_names:
            raise AssertionError(f"Input '{name}' is a reserved keyword.")

        flags = type_validation_flags(_type)
        if flags:
            flags_str = ", ".join([f.value for f in flags])
            _logger.warning(f"WARNING! Type `{_type}` for input variable `{name}` has the following flags: {flags_str}. Please check if the target Inference Server supports this type.")

        self.inputs[name] = {
            "type": _type,
        }
        if desc is not None:
            self.inputs[name]['desc'] = desc

    def change_input(self, name: str, _type: str, desc: Optional[str] = None):
        """
        Changes one of the inputs of the component.

        Args:
            name (str): Name of the input to be changed.
            _type (str): New type of the input.
            desc (str): Description of the input. (optional)
        """
        if name not in self.inputs:
            raise AssertionError(f"There is no input with name '{name}'")
        flags = type_validation_flags(_type)
        if flags:
            flags_str = ", ".join([f.value for f in flags])
            _logger.warning(f"WARNING! Type `{_type}` for input variable `{name}` has the following flags: {flags_str}. Please check if the target Inference Server supports this type.")
        self.inputs[name]['type'] = _type
        if desc is not None:
            self.inputs[name]['desc'] = desc

    def delete_input(self, name: str):
        """
        Deletes an input from the component by name.
        Once the package has been created with the given component, it is recommended not to change the component directly.
        Instead, all necessary methods to change it are available through the package to avoid component inconsistencies.
        It is recommended to use `package.delete_input_wire(...)` with default parameter `with_input=True`.

        Args:
            name (str): Name of the input to be deleted.
        """
        if name not in self.inputs:
            raise AssertionError(f"Component '{self.name}' has no input '{name}'")
        self.inputs.pop(name)

    def add_output(self, name: str, _type: str, desc: Optional[str] = None):
        """
        Adds a new output to the component.

        Types supported by AI Inference Server version 1.6 are contained in the `type_dictionary`. Newer AI Inference server version may support additional types.
        In case the type is not known by the AI SDK, a warning message will be printed.
        The most frequently used types are
        - String:
            Typically used for data to be sent to Databus
        - Object:
            Typically used for images to be sent to ZMQ Connector
        - Numeric scalar types:
            Typically used for data sent to S7 Connector

        For outputs of type `Object` the entrypoint must return with a `dictionary` containing two fields, where one field has type `str` and the other field has type `bytes`.
        The example below shows the required format, assuming that 'image' is a PIL Image.
        ```python
            "processedImage": {
                "metadata": json.dumps( {
                                "resolutionWidth": image.width,
                                "resolutionHeight": image.height
                                }
                            ),
                "bytes": image.tobytes()
            }
        ```

        Args:
            name (str): Name of the new output.
            _type (str): Type of the new output.
            desc (str): Description of the output. (optional)
        """
        if self.outputs is None:
            self.outputs = {}
        if name in self.outputs:
            raise AssertionError(f"Output '{name}' already exists")
        if name.lower() in self.reserved_names:
            raise AssertionError(f"Output '{name}' is a reserved keyword.")
        flags = type_validation_flags(_type)
        if flags:
            flags_str = ", ".join([f.value for f in flags])
            _logger.warning(f"WARNING! Type `{_type}` for output variable `{name}` has the following flags: {flags_str}. Please check if the target Inference Server supports this type.")
        self.outputs[name] = {
            "type": _type,
        }
        if desc is not None:
            self.outputs[name]['desc'] = desc

    def change_output(self, name: str, _type: str, desc: Optional[str] = None):
        """
        Changes one of the outputs of the component.

        Args:
            name (str): Name of the output to be changed.
            _type (str): The new type of the output.
            desc (str): Description of the output. (optional)
        """
        if name not in self.outputs:
            raise AssertionError(f"There is no output with name '{name}'")
        flags = type_validation_flags(_type)
        if flags:
            flags_str = ", ".join([f.value for f in flags])
            _logger.warning(f"WARNING! Type `{_type}` for output variable `{name}` has the following flags: {flags_str}. Please check if the target Inference Server supports this type.")
        self.outputs[name]['type'] = _type
        if desc is not None:
            self.outputs[name]['desc'] = desc

    def delete_output(self, name: str):
        """
        Deletes an output from the component by name.
        Once the package has been created with the given component, it is recommended not to change the component directly.
        Instead, all necessary methods to change it are available through the package to avoid component inconsistencies.
        Deleting an output which is represented in any wire will cause package inconsistency.

        Args:
            name (str): Name of the output to be deleted.
        """
        if name not in self.outputs:
            raise AssertionError(f"Component '{self.name}' has no output '{name}'")
        self.outputs.pop(name)

    def _to_dict(self):
        inputs = []
        inputs += [{
            'name': name,
            'type': self.inputs[name]['type'],
        } for name in self.inputs]

        outputs = []
        outputs += [{
            'name': name,
            'type': self.outputs[name]['type'],
            'metric': False,
        } for name in self.outputs]

        return {
            'name': self.name,
            'description': self.desc,
            'batch': self.batch.dict(),
            'inputType': inputs,
            'outputType': outputs,
        }

    def validate(self):
        """
        Empty method for child classess to implement.
        """
        pass

    def save(self, destination, validate):
        """
        Empty method for child classess to implement.
        """
        pass

class GPURuntimeComponent(Component):
    """
    The GPURuntimeComponent is used to define a component that runs on a GPU device.
    The component works only with ONNX models and can be used in an Inference Pipeline.

    Attributes:
        name (str): Component name.
        version (str): Component version.
        desc (str): Component description.

    Methods:
        use_model(self, path: Path | str, max_batch_size: int, optimization: Optional[model_config.TensorRTOptimization] = None, warmup: model_config.Warmup = None): Add an ONNX model file for the component.
        use_config(self, path: Path | str): Use a custom config.pbtxt file instead of the autogenerated one.
        save(self, destination: Path | str, validate = False): Saves the component to a folder structure, so it can be used as part of a pipeline configuration package.
    """

    def __init__(self, name: str = "inference", version: str = "1", desc: str = ""):
        """
        Creates a new, empty GPU Runtime component.

        Args:
            name (str): Component name. (default: inference)
            version (str): Component version. (default: 1)
            desc (str): Component description (optional)
        """
        super().__init__(name=name, desc=desc)
        self.version: str = version
        self.entrypoint: Path | None = None
        self.model_path: Path | None = None
        self.model_version: str = "1"
        self.config: Path | None = None
        self.auto_config = None
        self.warmup: bool = False
        self.optimization: bool = False

    def _to_dict(self):
        return {
            **super()._to_dict(),
            'version': self.version,
            'entrypoint': f"{self.model_version}/{self.entrypoint.name}",
            'hwType': 'GPU',
            'warmup': self.warmup,
            'optimization': self.optimization,
            'runtime': {
                'type': 'gpuruntime',
                'version': '0.1.0',
            }
        }

    def use_model(self, path: Path | str,
                  max_batch_size: int,
                  optimization: model_config.TensorRTOptimization | None = None,
                  warmup: model_config.Warmup | bool | None = None):
        """
        Add the ONNX model file for the component.

        Args:
            path (Path | str): The path to the ONNX model file.
            max_batch_size (int): The maximum batch size for the model.
            optimization (model_config.TensorRTOptimization, optional): The optimization configuration for the model. Defaults to None.
            warmup (model_config.Warmup, optional): The warmup configuration for the model. Defaults to None.

        Raises:
            AssertionError: If the specified model file is not found, has an invalid extension, or if max_batch_size is less than 0.

        """
        path = Path(path)

        if not path.is_file():
            raise AssertionError(f"specified model file not found: '{path}'")

        if path.suffix != ".onnx":
            raise AssertionError(f"model file extension is not '.onnx': '{path}'")

        if max_batch_size < 0:
            raise AssertionError("max_batch_size must be greater or equal to 0")

        self.entrypoint = Path("model.onnx")
        self.model_path = path
        if self.config is not None:
            _logger.warning("Previously added configuration was removed. Component will use the default configuration unless you specify your own.")
        self.config = None

        # Removing old automatic variables
        if self.auto_config is not None:
            for var in self.auto_config.inputs:
                self.delete_input(var["name"])
            for var in self.auto_config.outputs:
                self.delete_output(var["name"])

        model_warmup = warmup
        if isinstance(warmup, bool):
            self.warmup = warmup
            model_warmup = model_config.Warmup.ZERO_DATA if warmup else model_config.Warmup.DISABLED
        elif isinstance(warmup, model_config.Warmup):
            self.warmup = warmup != model_config.Warmup.DISABLED
        else:
            self.warmup = False
            model_warmup = model_config.Warmup.DISABLED

        self.optimization = isinstance(optimization, model_config.TensorRTOptimization)

        self.auto_config = model_config.ModelConfig(onnx_path=path,
                                                    max_batch_size=max_batch_size,
                                                    warmup=model_warmup,
                                                    optimization=optimization)
        for var in self.auto_config.inputs:
            self.add_input(var["name"], var["type"])
        for var in self.auto_config.outputs:
            self.add_output(var["name"], var["type"])

    def use_config(self, path: Path | str):
        """
        Sets the configuration file to be used for inference.
        Intended usage is to use a custom configuration file instead of the autogenerated one.
        This way extra configurations can be added to the component, such as the execution accelerator.

        Args:
            path (Path | str): The path to the configuration file.

        Raises:
            AssertionError: If the specified config file is not found or has an invalid extension.

        """
        path = Path(path)

        if not path.is_file():
            raise AssertionError(f"specified config file not found: '{path}'")

        if path.suffix != ".pbtxt":
            raise AssertionError(f"config file extension is not '.pbtxt': '{path}'")

        config_data = _validate_and_load_gpuruntime_config(path)

        self.warmup = (0 < len(config_data.model_warmup))
        self.optimization = (0 < len(f"{config_data.optimization}"))

        self.config = path

    def save(self, destination: Path | str, validate = False):
        """
        Saves the component to a folder structure, so it can be used as part of a pipeline configuration package.

        The component folder contains the following:

        - An `.onnx` model file
        - A `.pbtxt` configuration file

        Args:
            destination (path-like): Target directory to which the component will be saved.
        """
        if self.entrypoint is None:
            raise AssertionError("An ONNX model file must be specified before the component can be saved.")

        component_dir = Path(destination) / self.name
        component_dir.mkdir(parents = True, exist_ok = True)

        model_dir = component_dir / self.model_version
        model_dir.mkdir(exist_ok = True)

        shutil.copy(str(self.model_path), str(model_dir / "model.onnx"))

        if self.config is None:
            _logger.warning("Configuration was not specified. Model will be saved with default configuration.")
            (component_dir / "config.pbtxt").write_text(f"{self.auto_config}")
        else:
            shutil.copy(self.config, component_dir)

class PipelineData:

    def __init__(self, name: str, version: Optional[str] = None, desc: str = ""):
        """
        A newly initialized `Pipeline` will contain no `Component` or wire, just its name and version will be set.
        The name and version will define together the name of the zip file when the package is saved.

        Args:
            name (str): Name of the package
            desc (str): Package description (optional)
            version (str): Version of the package
        """

        self.name = name
        self.desc = desc

        self.version = version
        self.init_version = version  # initial version; used when version is not set in Pipeline.export() and save()
        self.save_version = version  # contains the version determined at exporting/saving
        self.package_id: Optional[uuid.UUID] = None

        self.author = 'AI SDK'

        self.components = {}
        self.wiring = {}
        self.parameters: list[PipelineParameter] = []

        self.periodicity = None
        self.timeshift_reference = []

        self.inputs: list[PipelineVariable] = []
        self.outputs: list[PipelineVariable] = []
        self.log_level = logging.INFO

        self.cycle_time = None

    # Helper method returns index or None if the pipeline inputs/outputs contain the given variable
    def _get_io_index(self, io_list: list[PipelineVariable], variable: str) -> Optional[int]:
        for i, io in enumerate(io_list):
            if io.variableName == variable:
                return i
        return None

    def _get_io_variable(self, variable_name: str) -> Optional[PipelineVariable]:
        for io in self.inputs + self.outputs:
            if io.variableName == variable_name:
                return io
        return None

    def _get_parameter(self, parameter_name: str) -> Optional[PipelineParameter]:
        for param in self.parameters:
            if param.name == parameter_name:
                return param
        return None

    def _get_parameter_index(self, parameter_name: str) -> Optional[int]:
        for i, param in enumerate(self.parameters):
            if param.name == parameter_name:
                return i
        return None