wallaroo.deployment_config

  1import json
  2import os
  3from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
  4
  5if TYPE_CHECKING:
  6    from .model import Model
  7
  8
  9class DeploymentConfig(Dict):
 10    def guarantee_workspace_id(self, workspace_id: Optional[int]) -> "DeploymentConfig":
 11        if workspace_id is None:
 12            return self
 13        if self.get("workspace_id", None) is None:
 14            self["workspace_id"] = workspace_id
 15        return self
 16
 17
 18class DeploymentConfigBuilder(object):
 19    def __init__(self, workspace_id: Optional[int] = None) -> None:
 20        self._config: Dict[str, Any] = {
 21            "engine": {},
 22            "enginelb": {},
 23            "engineAux": {"images": {}},
 24            **({"workspace_id": workspace_id} if workspace_id is not None else {}),
 25        }
 26
 27        env_config = os.environ.get("DEPLOYMENT_CONFIG", None)
 28
 29        if env_config:
 30            conf = json.loads(env_config)
 31            setters: Dict[str, Callable[[Any], Any]] = {
 32                "image": self.image,
 33                "replicas": self.replica_count,
 34                "autoscale": self._autoscale,
 35                "cpus": self.cpus,
 36                "memory": self.memory,
 37                "lb_cpus": self.lb_cpus,
 38                "lb_memory": self.lb_memory,
 39            }
 40            for key, set in setters.items():
 41                if key in conf:
 42                    set(conf[key])
 43
 44    def image(self, image: str) -> "DeploymentConfigBuilder":
 45        self._config["engine"]["image"] = image
 46        return self
 47
 48    def replica_count(self, count: int) -> "DeploymentConfigBuilder":
 49        if (
 50            "autoscale" in self._config["engine"]
 51            and self._config["engine"]["autoscale"]["replica_max"] < count
 52        ):
 53            raise RuntimeError(
 54                "Replica count must be less than or equal to replica max. Use replica_autoscale_min_max to adjust this."
 55            )
 56        self._config["engine"]["replicas"] = count
 57        return self
 58
 59    def _autoscale(self, autoscale: Dict[str, Any]):
 60        self._config["engine"]["autoscale"] = autoscale
 61        return self
 62
 63    def replica_autoscale_min_max(self, maximum: int, minimum: int = 0):
 64        """Configures the minimum and maximum for autoscaling"""
 65        if minimum > maximum:
 66            raise RuntimeError("Minimum must be less than or equal to maximum")
 67        if minimum < 0:
 68            raise RuntimeError("Minimum must be at least 0")
 69        if "autoscale" not in self._config["engine"]:
 70            self._config["engine"]["autoscale"] = {}
 71        if (
 72            "replicas" in self._config["engine"]
 73            and self._config["engine"]["replicas"] > maximum
 74        ):
 75            raise RuntimeError(
 76                "Maximum must be greater than or equal to number of replicas"
 77            )
 78        self._config["engine"]["autoscale"]["replica_min"] = minimum
 79        self._config["engine"]["autoscale"]["replica_max"] = maximum
 80        self._config["engine"]["autoscale"]["type"] = "cpu"
 81        return self
 82
 83    def autoscale_cpu_utilization(self, cpu_utilization_percentage: int):
 84        """Sets the average CPU metric to scale on in a percentage"""
 85        if "autoscale" not in self._config["engine"]:
 86            print(
 87                "Warn: min and max not set for autoscaling. These must be set to enable autoscaling"
 88            )
 89            self._config["engine"]["autoscale"] = {}
 90        self._config["engine"]["autoscale"][
 91            "cpu_utilization"
 92        ] = cpu_utilization_percentage
 93        return self
 94
 95    def disable_autoscale(self):
 96        """Disables autoscaling in the deployment configuration"""
 97        if "autoscale" in ["engine"]:
 98            del self._config["engine"]["autoscale"]
 99        return self
100
101    def _add_resource(
102        self,
103        component_stanza: Dict[str, Any],
104        resource_name: str,
105        value: Union[int, str],
106    ) -> "DeploymentConfigBuilder":
107        if "resources" not in component_stanza:
108            component_stanza["resources"] = {"limits": {}, "requests": {}}
109        component_stanza["resources"]["limits"][resource_name] = value
110        component_stanza["resources"]["requests"][resource_name] = value
111        return self
112
113    def cpus(self, core_count: int) -> "DeploymentConfigBuilder":
114        self._config["engine"]["cpu"] = core_count
115        return self._add_resource(self._config["engine"], "cpu", core_count)
116
117    def memory(self, memory_spec: str) -> "DeploymentConfigBuilder":
118        return self._add_resource(self._config["engine"], "memory", memory_spec)
119
120    def lb_cpus(self, core_count: int) -> "DeploymentConfigBuilder":
121        return self._add_resource(self._config["enginelb"], "cpu", core_count)
122
123    def lb_memory(self, memory_spec: int) -> "DeploymentConfigBuilder":
124        return self._add_resource(self._config["enginelb"], "memory", memory_spec)
125
126    def python_load_timeout_secs(self, timeout_secs: int) -> "DeploymentConfigBuilder":
127        if "python" not in self._config["engine"]:
128            self._config["engine"]["python"] = {}
129        self._config["engine"]["python"]["load_timeout_millis"] = timeout_secs * 1000
130        return self
131
132    def _guarantee_sidekick_stanza(self, model: "Model") -> Dict[str, Any]:
133        pass
134
135        model_uid = model.uid()
136        if model_uid not in self._config["engineAux"]["images"]:
137            self._config["engineAux"]["images"][model_uid] = {}
138        return self._config["engineAux"]["images"][model_uid]
139
140    def sidekick_cpus(
141        self, model: "Model", core_count: int
142    ) -> "DeploymentConfigBuilder":
143        """Sets the number of CPUs to be used for the model's sidekick container. Only affects
144        image-based models (e.g. MLFlow models) in a deployment.
145
146        :param Model model: The sidekick model to configure.
147        :param int core_count: Number of CPU cores to use in this sidekick.
148        :return: This DeploymentConfigBuilder instance for chaining."""
149
150        return self._add_resource(
151            self._guarantee_sidekick_stanza(model), "cpu", core_count
152        )
153
154    def sidekick_memory(
155        self, model: "Model", memory_spec: str
156    ) -> "DeploymentConfigBuilder":
157        """Sets the memory to be used for the model's sidekick container. Only affects
158        image-based models (e.g. MLFlow models) in a deployment.
159
160        :param Model model: The sidekick model to configure.
161        :param str memory_spec: Specification of amount of memory (e.g., "2Gi", "500Mi") to use in
162        this sidekick.
163        :return: This DeploymentConfigBuilder instance for chaining."""
164
165        return self._add_resource(
166            self._guarantee_sidekick_stanza(model), "memory", memory_spec
167        )
168
169    def sidekick_env(
170        self, model: "Model", environment: Dict[str, str]
171    ) -> "DeploymentConfigBuilder":
172        """Sets the environment variables to be set for the model's sidekick container. Only affects
173        image-based models (e.g. MLFlow models) in a deployment.
174
175        :param Model model: The sidekick model to configure.
176        :param Dict[str, str] environment: Dictionary of environment variables names and their
177        corresponding values to be set in the sidekick container.
178        :return: This DeploymentConfigBuilder instance for chaining."""
179
180        stanza = self._guarantee_sidekick_stanza(model)
181        stanza["env"] = []
182        for name, value in environment.items():
183            stanza["env"].append({"name": name, "value": value})
184
185        return self
186
187    def build(self) -> DeploymentConfig:
188        return DeploymentConfig(self._config)
class DeploymentConfig(typing.Dict):
10class DeploymentConfig(Dict):
11    def guarantee_workspace_id(self, workspace_id: Optional[int]) -> "DeploymentConfig":
12        if workspace_id is None:
13            return self
14        if self.get("workspace_id", None) is None:
15            self["workspace_id"] = workspace_id
16        return self
def guarantee_workspace_id( self, workspace_id: Optional[int]) -> wallaroo.deployment_config.DeploymentConfig:
11    def guarantee_workspace_id(self, workspace_id: Optional[int]) -> "DeploymentConfig":
12        if workspace_id is None:
13            return self
14        if self.get("workspace_id", None) is None:
15            self["workspace_id"] = workspace_id
16        return self
Inherited Members
builtins.dict
get
setdefault
pop
popitem
keys
items
values
update
fromkeys
clear
copy
class DeploymentConfigBuilder:
 19class DeploymentConfigBuilder(object):
 20    def __init__(self, workspace_id: Optional[int] = None) -> None:
 21        self._config: Dict[str, Any] = {
 22            "engine": {},
 23            "enginelb": {},
 24            "engineAux": {"images": {}},
 25            **({"workspace_id": workspace_id} if workspace_id is not None else {}),
 26        }
 27
 28        env_config = os.environ.get("DEPLOYMENT_CONFIG", None)
 29
 30        if env_config:
 31            conf = json.loads(env_config)
 32            setters: Dict[str, Callable[[Any], Any]] = {
 33                "image": self.image,
 34                "replicas": self.replica_count,
 35                "autoscale": self._autoscale,
 36                "cpus": self.cpus,
 37                "memory": self.memory,
 38                "lb_cpus": self.lb_cpus,
 39                "lb_memory": self.lb_memory,
 40            }
 41            for key, set in setters.items():
 42                if key in conf:
 43                    set(conf[key])
 44
 45    def image(self, image: str) -> "DeploymentConfigBuilder":
 46        self._config["engine"]["image"] = image
 47        return self
 48
 49    def replica_count(self, count: int) -> "DeploymentConfigBuilder":
 50        if (
 51            "autoscale" in self._config["engine"]
 52            and self._config["engine"]["autoscale"]["replica_max"] < count
 53        ):
 54            raise RuntimeError(
 55                "Replica count must be less than or equal to replica max. Use replica_autoscale_min_max to adjust this."
 56            )
 57        self._config["engine"]["replicas"] = count
 58        return self
 59
 60    def _autoscale(self, autoscale: Dict[str, Any]):
 61        self._config["engine"]["autoscale"] = autoscale
 62        return self
 63
 64    def replica_autoscale_min_max(self, maximum: int, minimum: int = 0):
 65        """Configures the minimum and maximum for autoscaling"""
 66        if minimum > maximum:
 67            raise RuntimeError("Minimum must be less than or equal to maximum")
 68        if minimum < 0:
 69            raise RuntimeError("Minimum must be at least 0")
 70        if "autoscale" not in self._config["engine"]:
 71            self._config["engine"]["autoscale"] = {}
 72        if (
 73            "replicas" in self._config["engine"]
 74            and self._config["engine"]["replicas"] > maximum
 75        ):
 76            raise RuntimeError(
 77                "Maximum must be greater than or equal to number of replicas"
 78            )
 79        self._config["engine"]["autoscale"]["replica_min"] = minimum
 80        self._config["engine"]["autoscale"]["replica_max"] = maximum
 81        self._config["engine"]["autoscale"]["type"] = "cpu"
 82        return self
 83
 84    def autoscale_cpu_utilization(self, cpu_utilization_percentage: int):
 85        """Sets the average CPU metric to scale on in a percentage"""
 86        if "autoscale" not in self._config["engine"]:
 87            print(
 88                "Warn: min and max not set for autoscaling. These must be set to enable autoscaling"
 89            )
 90            self._config["engine"]["autoscale"] = {}
 91        self._config["engine"]["autoscale"][
 92            "cpu_utilization"
 93        ] = cpu_utilization_percentage
 94        return self
 95
 96    def disable_autoscale(self):
 97        """Disables autoscaling in the deployment configuration"""
 98        if "autoscale" in ["engine"]:
 99            del self._config["engine"]["autoscale"]
100        return self
101
102    def _add_resource(
103        self,
104        component_stanza: Dict[str, Any],
105        resource_name: str,
106        value: Union[int, str],
107    ) -> "DeploymentConfigBuilder":
108        if "resources" not in component_stanza:
109            component_stanza["resources"] = {"limits": {}, "requests": {}}
110        component_stanza["resources"]["limits"][resource_name] = value
111        component_stanza["resources"]["requests"][resource_name] = value
112        return self
113
114    def cpus(self, core_count: int) -> "DeploymentConfigBuilder":
115        self._config["engine"]["cpu"] = core_count
116        return self._add_resource(self._config["engine"], "cpu", core_count)
117
118    def memory(self, memory_spec: str) -> "DeploymentConfigBuilder":
119        return self._add_resource(self._config["engine"], "memory", memory_spec)
120
121    def lb_cpus(self, core_count: int) -> "DeploymentConfigBuilder":
122        return self._add_resource(self._config["enginelb"], "cpu", core_count)
123
124    def lb_memory(self, memory_spec: int) -> "DeploymentConfigBuilder":
125        return self._add_resource(self._config["enginelb"], "memory", memory_spec)
126
127    def python_load_timeout_secs(self, timeout_secs: int) -> "DeploymentConfigBuilder":
128        if "python" not in self._config["engine"]:
129            self._config["engine"]["python"] = {}
130        self._config["engine"]["python"]["load_timeout_millis"] = timeout_secs * 1000
131        return self
132
133    def _guarantee_sidekick_stanza(self, model: "Model") -> Dict[str, Any]:
134        pass
135
136        model_uid = model.uid()
137        if model_uid not in self._config["engineAux"]["images"]:
138            self._config["engineAux"]["images"][model_uid] = {}
139        return self._config["engineAux"]["images"][model_uid]
140
141    def sidekick_cpus(
142        self, model: "Model", core_count: int
143    ) -> "DeploymentConfigBuilder":
144        """Sets the number of CPUs to be used for the model's sidekick container. Only affects
145        image-based models (e.g. MLFlow models) in a deployment.
146
147        :param Model model: The sidekick model to configure.
148        :param int core_count: Number of CPU cores to use in this sidekick.
149        :return: This DeploymentConfigBuilder instance for chaining."""
150
151        return self._add_resource(
152            self._guarantee_sidekick_stanza(model), "cpu", core_count
153        )
154
155    def sidekick_memory(
156        self, model: "Model", memory_spec: str
157    ) -> "DeploymentConfigBuilder":
158        """Sets the memory to be used for the model's sidekick container. Only affects
159        image-based models (e.g. MLFlow models) in a deployment.
160
161        :param Model model: The sidekick model to configure.
162        :param str memory_spec: Specification of amount of memory (e.g., "2Gi", "500Mi") to use in
163        this sidekick.
164        :return: This DeploymentConfigBuilder instance for chaining."""
165
166        return self._add_resource(
167            self._guarantee_sidekick_stanza(model), "memory", memory_spec
168        )
169
170    def sidekick_env(
171        self, model: "Model", environment: Dict[str, str]
172    ) -> "DeploymentConfigBuilder":
173        """Sets the environment variables to be set for the model's sidekick container. Only affects
174        image-based models (e.g. MLFlow models) in a deployment.
175
176        :param Model model: The sidekick model to configure.
177        :param Dict[str, str] environment: Dictionary of environment variables names and their
178        corresponding values to be set in the sidekick container.
179        :return: This DeploymentConfigBuilder instance for chaining."""
180
181        stanza = self._guarantee_sidekick_stanza(model)
182        stanza["env"] = []
183        for name, value in environment.items():
184            stanza["env"].append({"name": name, "value": value})
185
186        return self
187
188    def build(self) -> DeploymentConfig:
189        return DeploymentConfig(self._config)
DeploymentConfigBuilder(workspace_id: Optional[int] = None)
20    def __init__(self, workspace_id: Optional[int] = None) -> None:
21        self._config: Dict[str, Any] = {
22            "engine": {},
23            "enginelb": {},
24            "engineAux": {"images": {}},
25            **({"workspace_id": workspace_id} if workspace_id is not None else {}),
26        }
27
28        env_config = os.environ.get("DEPLOYMENT_CONFIG", None)
29
30        if env_config:
31            conf = json.loads(env_config)
32            setters: Dict[str, Callable[[Any], Any]] = {
33                "image": self.image,
34                "replicas": self.replica_count,
35                "autoscale": self._autoscale,
36                "cpus": self.cpus,
37                "memory": self.memory,
38                "lb_cpus": self.lb_cpus,
39                "lb_memory": self.lb_memory,
40            }
41            for key, set in setters.items():
42                if key in conf:
43                    set(conf[key])
def image(self, image: str) -> wallaroo.deployment_config.DeploymentConfigBuilder:
45    def image(self, image: str) -> "DeploymentConfigBuilder":
46        self._config["engine"]["image"] = image
47        return self
def replica_count(self, count: int) -> wallaroo.deployment_config.DeploymentConfigBuilder:
49    def replica_count(self, count: int) -> "DeploymentConfigBuilder":
50        if (
51            "autoscale" in self._config["engine"]
52            and self._config["engine"]["autoscale"]["replica_max"] < count
53        ):
54            raise RuntimeError(
55                "Replica count must be less than or equal to replica max. Use replica_autoscale_min_max to adjust this."
56            )
57        self._config["engine"]["replicas"] = count
58        return self
def replica_autoscale_min_max(self, maximum: int, minimum: int = 0):
64    def replica_autoscale_min_max(self, maximum: int, minimum: int = 0):
65        """Configures the minimum and maximum for autoscaling"""
66        if minimum > maximum:
67            raise RuntimeError("Minimum must be less than or equal to maximum")
68        if minimum < 0:
69            raise RuntimeError("Minimum must be at least 0")
70        if "autoscale" not in self._config["engine"]:
71            self._config["engine"]["autoscale"] = {}
72        if (
73            "replicas" in self._config["engine"]
74            and self._config["engine"]["replicas"] > maximum
75        ):
76            raise RuntimeError(
77                "Maximum must be greater than or equal to number of replicas"
78            )
79        self._config["engine"]["autoscale"]["replica_min"] = minimum
80        self._config["engine"]["autoscale"]["replica_max"] = maximum
81        self._config["engine"]["autoscale"]["type"] = "cpu"
82        return self

Configures the minimum and maximum for autoscaling

def autoscale_cpu_utilization(self, cpu_utilization_percentage: int):
84    def autoscale_cpu_utilization(self, cpu_utilization_percentage: int):
85        """Sets the average CPU metric to scale on in a percentage"""
86        if "autoscale" not in self._config["engine"]:
87            print(
88                "Warn: min and max not set for autoscaling. These must be set to enable autoscaling"
89            )
90            self._config["engine"]["autoscale"] = {}
91        self._config["engine"]["autoscale"][
92            "cpu_utilization"
93        ] = cpu_utilization_percentage
94        return self

Sets the average CPU metric to scale on in a percentage

def disable_autoscale(self):
 96    def disable_autoscale(self):
 97        """Disables autoscaling in the deployment configuration"""
 98        if "autoscale" in ["engine"]:
 99            del self._config["engine"]["autoscale"]
100        return self

Disables autoscaling in the deployment configuration

def cpus( self, core_count: int) -> wallaroo.deployment_config.DeploymentConfigBuilder:
114    def cpus(self, core_count: int) -> "DeploymentConfigBuilder":
115        self._config["engine"]["cpu"] = core_count
116        return self._add_resource(self._config["engine"], "cpu", core_count)
def memory( self, memory_spec: str) -> wallaroo.deployment_config.DeploymentConfigBuilder:
118    def memory(self, memory_spec: str) -> "DeploymentConfigBuilder":
119        return self._add_resource(self._config["engine"], "memory", memory_spec)
def lb_cpus( self, core_count: int) -> wallaroo.deployment_config.DeploymentConfigBuilder:
121    def lb_cpus(self, core_count: int) -> "DeploymentConfigBuilder":
122        return self._add_resource(self._config["enginelb"], "cpu", core_count)
def lb_memory( self, memory_spec: int) -> wallaroo.deployment_config.DeploymentConfigBuilder:
124    def lb_memory(self, memory_spec: int) -> "DeploymentConfigBuilder":
125        return self._add_resource(self._config["enginelb"], "memory", memory_spec)
def python_load_timeout_secs( self, timeout_secs: int) -> wallaroo.deployment_config.DeploymentConfigBuilder:
127    def python_load_timeout_secs(self, timeout_secs: int) -> "DeploymentConfigBuilder":
128        if "python" not in self._config["engine"]:
129            self._config["engine"]["python"] = {}
130        self._config["engine"]["python"]["load_timeout_millis"] = timeout_secs * 1000
131        return self
def sidekick_cpus( self, model: wallaroo.model.Model, core_count: int) -> wallaroo.deployment_config.DeploymentConfigBuilder:
141    def sidekick_cpus(
142        self, model: "Model", core_count: int
143    ) -> "DeploymentConfigBuilder":
144        """Sets the number of CPUs to be used for the model's sidekick container. Only affects
145        image-based models (e.g. MLFlow models) in a deployment.
146
147        :param Model model: The sidekick model to configure.
148        :param int core_count: Number of CPU cores to use in this sidekick.
149        :return: This DeploymentConfigBuilder instance for chaining."""
150
151        return self._add_resource(
152            self._guarantee_sidekick_stanza(model), "cpu", core_count
153        )

Sets the number of CPUs to be used for the model's sidekick container. Only affects image-based models (e.g. MLFlow models) in a deployment.

Parameters
  • Model model: The sidekick model to configure.
  • int core_count: Number of CPU cores to use in this sidekick.
Returns

This DeploymentConfigBuilder instance for chaining.

def sidekick_memory( self, model: wallaroo.model.Model, memory_spec: str) -> wallaroo.deployment_config.DeploymentConfigBuilder:
155    def sidekick_memory(
156        self, model: "Model", memory_spec: str
157    ) -> "DeploymentConfigBuilder":
158        """Sets the memory to be used for the model's sidekick container. Only affects
159        image-based models (e.g. MLFlow models) in a deployment.
160
161        :param Model model: The sidekick model to configure.
162        :param str memory_spec: Specification of amount of memory (e.g., "2Gi", "500Mi") to use in
163        this sidekick.
164        :return: This DeploymentConfigBuilder instance for chaining."""
165
166        return self._add_resource(
167            self._guarantee_sidekick_stanza(model), "memory", memory_spec
168        )

Sets the memory to be used for the model's sidekick container. Only affects image-based models (e.g. MLFlow models) in a deployment.

Parameters
  • Model model: The sidekick model to configure.
  • str memory_spec: Specification of amount of memory (e.g., "2Gi", "500Mi") to use in this sidekick.
Returns

This DeploymentConfigBuilder instance for chaining.

def sidekick_env( self, model: wallaroo.model.Model, environment: Dict[str, str]) -> wallaroo.deployment_config.DeploymentConfigBuilder:
170    def sidekick_env(
171        self, model: "Model", environment: Dict[str, str]
172    ) -> "DeploymentConfigBuilder":
173        """Sets the environment variables to be set for the model's sidekick container. Only affects
174        image-based models (e.g. MLFlow models) in a deployment.
175
176        :param Model model: The sidekick model to configure.
177        :param Dict[str, str] environment: Dictionary of environment variables names and their
178        corresponding values to be set in the sidekick container.
179        :return: This DeploymentConfigBuilder instance for chaining."""
180
181        stanza = self._guarantee_sidekick_stanza(model)
182        stanza["env"] = []
183        for name, value in environment.items():
184            stanza["env"].append({"name": name, "value": value})
185
186        return self

Sets the environment variables to be set for the model's sidekick container. Only affects image-based models (e.g. MLFlow models) in a deployment.

Parameters
  • Model model: The sidekick model to configure.
  • Dict[str, str] environment: Dictionary of environment variables names and their corresponding values to be set in the sidekick container.
Returns

This DeploymentConfigBuilder instance for chaining.

def build(self) -> wallaroo.deployment_config.DeploymentConfig:
188    def build(self) -> DeploymentConfig:
189        return DeploymentConfig(self._config)