wallaroo.framework


class Framework(builtins.str, enum.Enum):

An Enum to represent the supported frameworks.

ONNX = <Framework.ONNX: 'onnx'>
TENSORFLOW = <Framework.TENSORFLOW: 'tensorflow'>
PYTHON = <Framework.PYTHON: 'python'>
KERAS = <Framework.KERAS: 'keras'>
SKLEARN = <Framework.SKLEARN: 'sklearn'>
PYTORCH = <Framework.PYTORCH: 'pytorch'>
XGBOOST = <Framework.XGBOOST: 'xgboost'>
HUGGING_FACE_AUTOMATIC_SPEECH_RECOGNITION = <Framework.HUGGING_FACE_AUTOMATIC_SPEECH_RECOGNITION: 'hugging-face-automatic-speech-recognition'>
HUGGING_FACE_FEATURE_EXTRACTION = <Framework.HUGGING_FACE_FEATURE_EXTRACTION: 'hugging-face-feature-extraction'>
HUGGING_FACE_IMAGE_CLASSIFICATION = <Framework.HUGGING_FACE_IMAGE_CLASSIFICATION: 'hugging-face-image-classification'>
HUGGING_FACE_IMAGE_SEGMENTATION = <Framework.HUGGING_FACE_IMAGE_SEGMENTATION: 'hugging-face-image-segmentation'>
HUGGING_FACE_IMAGE_TO_TEXT = <Framework.HUGGING_FACE_IMAGE_TO_TEXT: 'hugging-face-image-to-text'>
HUGGING_FACE_OBJECT_DETECTION = <Framework.HUGGING_FACE_OBJECT_DETECTION: 'hugging-face-object-detection'>
HUGGING_FACE_QUESTION_ANSWERING = <Framework.HUGGING_FACE_QUESTION_ANSWERING: 'hugging-face-question-answering'>
HUGGING_FACE_STABLE_DIFFUSION_TEXT_2_IMG = <Framework.HUGGING_FACE_STABLE_DIFFUSION_TEXT_2_IMG: 'hugging-face-stable-diffusion-text-2-img'>
HUGGING_FACE_SUMMARIZATION = <Framework.HUGGING_FACE_SUMMARIZATION: 'hugging-face-summarization'>
HUGGING_FACE_TEXT_CLASSIFICATION = <Framework.HUGGING_FACE_TEXT_CLASSIFICATION: 'hugging-face-text-classification'>
HUGGING_FACE_TRANSLATION = <Framework.HUGGING_FACE_TRANSLATION: 'hugging-face-translation'>
HUGGING_FACE_ZERO_SHOT_CLASSIFICATION = <Framework.HUGGING_FACE_ZERO_SHOT_CLASSIFICATION: 'hugging-face-zero-shot-classification'>
HUGGING_FACE_ZERO_SHOT_IMAGE_CLASSIFICATION = <Framework.HUGGING_FACE_ZERO_SHOT_IMAGE_CLASSIFICATION: 'hugging-face-zero-shot-image-classification'>
HUGGING_FACE_ZERO_SHOT_OBJECT_DETECTION = <Framework.HUGGING_FACE_ZERO_SHOT_OBJECT_DETECTION: 'hugging-face-zero-shot-object-detection'>
HUGGING_FACE_SENTIMENT_ANALYSIS = <Framework.HUGGING_FACE_SENTIMENT_ANALYSIS: 'hugging-face-sentiment-analysis'>
HUGGING_FACE_TEXT_GENERATION = <Framework.HUGGING_FACE_TEXT_GENERATION: 'hugging-face-text-generation'>
CUSTOM = <Framework.CUSTOM: 'custom'>
VLLM = <Framework.VLLM: 'vllm'>
def get_default_config(self) -> Optional[FrameworkConfig]:

Get the default config for the framework.

def requires_config(self) -> bool:

Check if the framework requires a config.

def validate(self, config: Optional[FrameworkConfig]) -> None:

Validate the config for the framework.

Parameters
  • config: The config to validate.
Raises
  • ModelOptimizationConfigError: If the config is invalid.
@runtime_checkable
class OpenapiFrameworkConfig(typing.Protocol):

Base class for protocol classes.

Protocol classes are defined as::

class Proto(Protocol):
    def meth(self) -> int:
        ...

Such classes are primarily used with static type checkers that recognize structural subtyping (static duck-typing), for example::

class C:
    def meth(self) -> int:
        return 0

def func(x: Proto) -> int:
    return x.meth()

func(C())  # Passes static type check

See PEP 544 for details. Protocol classes decorated with @typing.runtime_checkable act as simple-minded runtime protocols that check only the presence of given attributes, ignoring their type signatures. Protocol classes can be generic, they are defined as::

class GenProto(Protocol[T]):
    def meth(self) -> T:
        ...
OpenapiFrameworkConfig(*args, **kwargs)
def to_dict(self) -> Dict[str, Any]:
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> OpenapiFrameworkConfig:
class FrameworkConfig(pydantic.main.BaseModel, abc.ABC):

A base class for all framework configs.

model_config = {'arbitrary_types_allowed': True, 'extra': 'forbid', 'use_enum_values': True, 'protected_namespaces': ()}
framework: Framework

The framework that this config is for.

def to_dict(self) -> Dict[str, Any]:

Convert the config to a dictionary.

def validate_config(self, framework: Framework) -> None:

Validate the config for the framework.

Parameters
  • framework: The framework to validate the config for.
Raises
  • ModelOptimizationConfigError: If the config is invalid.
model_fields = {}
model_computed_fields = {}
class CustomConfig(FrameworkConfig):

A config for the custom framework.

model_config = {'arbitrary_types_allowed': True, 'extra': 'allow', 'use_enum_values': True, 'protected_namespaces': ()}
model_path: str
framework: Framework

The framework that this config is for.

model_fields = {'model_path': FieldInfo(annotation=str, required=False, default='./model/', description="The relative path of the model artifacts to the provided zip file. Defaults to './model/'.", validate_default=True)}
model_computed_fields = {}
class VLLMConfig(FrameworkConfig):

A config for the VLLM framework.

device_group: Union[List[int], NoneType, wallaroo.wallaroo_ml_ops_api_client.types.Unset]
gpu_memory_utilization: Union[wallaroo.wallaroo_ml_ops_api_client.types.Unset, float]
kv_cache_dtype: Union[wallaroo.wallaroo_ml_ops_api_client.types.Unset, wallaroo.wallaroo_ml_ops_api_client.models.kv_cache_dtype.KvCacheDtype]
max_model_len: Union[NoneType, wallaroo.wallaroo_ml_ops_api_client.types.Unset, int]
max_num_seqs: Union[wallaroo.wallaroo_ml_ops_api_client.types.Unset, int]
max_seq_len_to_capture: Union[wallaroo.wallaroo_ml_ops_api_client.types.Unset, int]
quantization: Union[wallaroo.wallaroo_ml_ops_api_client.types.Unset, wallaroo.wallaroo_ml_ops_api_client.models.quantization.Quantization]
block_size: Union[NoneType, wallaroo.wallaroo_ml_ops_api_client.types.Unset, int]
framework: Framework

The framework that this config is for.

model_config = {'arbitrary_types_allowed': True, 'extra': 'forbid', 'use_enum_values': True, 'protected_namespaces': ()}
model_computed_fields = {}