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:
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])
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
python_load_timeout_secs( self, timeout_secs: int) -> wallaroo.deployment_config.DeploymentConfigBuilder:
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.