paasta-tools 1.21.4__py3-none-any.whl → 1.22.0__py3-none-any.whl

paasta_tools/__init__.py

@@ -17,4 +17,4 @@
 # setup phase, the dependencies may not exist on disk yet.
 #
 # Don't bump version manually. See `make release` docs in ./Makefile
-__version__ = "1.21.4"
+__version__ = "1.22.0"

paasta_tools/api/api.py

@@ -189,6 +189,11 @@ def make_app(global_config=None):
         "/v1/service_autoscaler/pause",
         request_method="GET",
     )
+    config.add_route(
+        "service_autoscaler.autoscaling_override.post",
+        "/v1/service_autoscaler/{service}/{instance}/autoscaling_override",
+        request_method="POST",
+    )
     config.add_route(
         "remote_run.start",
         "/v1/remote_run/{service}/{instance}/start",

paasta_tools/api/api_docs/swagger.json

@@ -279,6 +279,79 @@
                 "operationId": "delete_service_autoscaler_pause"
             }
         },
+        "/service_autoscaler/{service}/{instance}/autoscaling_override": {
+            "post": {
+                "operationId": "set_autoscaling_override",
+                "parameters": [
+                    {
+                        "description": "Service name",
+                        "in": "path",
+                        "name": "service",
+                        "required": true,
+                        "type": "string"
+                    },
+                    {
+                        "description": "Instance name",
+                        "in": "path",
+                        "name": "instance",
+                        "required": true,
+                        "type": "string"
+                    },
+                    {
+                        "in": "body",
+                        "name": "json_body",
+                        "required": true,
+                        "schema": {
+                            "$ref": "#/definitions/AutoscalingOverride"
+                        }
+                    }
+                ],
+                "responses": {
+                    "202": {
+                        "description": "Successfully set autoscaling override",
+                        "schema": {
+                            "type": "object",
+                            "properties": {
+                                "service": {
+                                    "type": "string",
+                                    "description": "Service name"
+                                },
+                                "instance": {
+                                    "type": "string",
+                                    "description": "Instance name"
+                                },
+                                "min_instances": {
+                                    "type": "integer",
+                                    "description": "Minimum number of instances to run"
+                                },
+                                "expire_after": {
+                                    "type": "number",
+                                    "format": "float",
+                                    "description": "Unix timestamp after which the override is no longer valid"
+                                },
+                                "status": {
+                                    "type": "string",
+                                    "description": "Status of the operation"
+                                }
+                            }
+                        }
+                    },
+                    "400": {
+                        "description": "Invalid request parameters"
+                    },
+                    "501": {
+                        "description": "Autoscaling not supported for this instance type"
+                    },
+                    "500": {
+                        "description": "Failed to set autoscaling override"
+                    }
+                },
+                "summary": "Set a temporary autoscaling override for a service instance",
+                "tags": [
+                    "autoscaler"
+                ]
+            }
+        },
         "/resources/utilization": {
             "get": {
                 "responses": {
@@ -1627,6 +1700,25 @@
                 }
             }
         },
+        "AutoscalingOverride": {
+            "type": "object",
+            "properties": {
+                "min_instances": {
+                    "type": "integer",
+                    "description": "Minimum number of instances to run",
+                    "minimum": 1
+                },
+                "expire_after": {
+                    "type": "number",
+                    "format": "float",
+                    "description": "Unix timestamp when this override is no longer valid"
+                }
+            },
+            "required": [
+                "min_instances",
+                "expire_after"
+            ]
+        },
         "KubernetesReplicaSet": {
             "type": "object",
             "properties": {

paasta_tools/api/views/autoscaler.py

@@ -15,13 +15,27 @@
 """
 PaaSTA service list (instances) etc.
 """
+import json
+import logging
+from datetime import datetime
+from datetime import timezone
+from typing import Tuple
+
+from kubernetes.client import V1ConfigMap
 from pyramid.response import Response
 from pyramid.view import view_config
 
 from paasta_tools.api import settings
 from paasta_tools.api.views.exception import ApiFailure
 from paasta_tools.cli.utils import get_instance_config
+from paasta_tools.kubernetes_tools import AUTOSCALING_OVERRIDES_CONFIGMAP_NAME
+from paasta_tools.kubernetes_tools import AUTOSCALING_OVERRIDES_CONFIGMAP_NAMESPACE
+from paasta_tools.kubernetes_tools import get_or_create_namespaced_configmap
 from paasta_tools.kubernetes_tools import KubernetesDeploymentConfig
+from paasta_tools.kubernetes_tools import patch_namespaced_configmap
+
+
+log = logging.getLogger(__name__)
 
 
 @view_config(route_name="service.autoscaler.get", request_method="GET", renderer="json")
@@ -98,3 +112,111 @@ def update_autoscaler_count(request):
 
     response_body = {"desired_instances": desired_instances, "status": status}
     return Response(json_body=response_body, status_code=202)
+
+
+def get_or_create_autoscaling_overrides_configmap() -> Tuple[V1ConfigMap, bool]:
+    return get_or_create_namespaced_configmap(
+        AUTOSCALING_OVERRIDES_CONFIGMAP_NAME,
+        namespace=AUTOSCALING_OVERRIDES_CONFIGMAP_NAMESPACE,
+        kube_client=settings.kubernetes_client,
+    )
+
+
+@view_config(
+    route_name="service_autoscaler.autoscaling_override.post",
+    request_method="POST",
+    renderer="json",
+)
+def set_autoscaling_override(request):
+    """Set a temporary autoscaling override for a service/instance.
+
+    This endpoint creates or updates a ConfigMap entry with override information
+    including expiration time. The override will be applied by the autoscaler.
+
+    Required parameters:
+    - service: The service name
+    - instance: The instance name
+    - min_instances: The minimum number of instances to enforce
+    - expires_after: unix timestamp after which the override is no longer valid
+    """
+    service = request.swagger_data.get("service")
+    instance = request.swagger_data.get("instance")
+    cluster = settings.cluster
+    soa_dir = settings.soa_dir
+
+    instance_config = get_instance_config(
+        service, instance, cluster, soa_dir, load_deployments=False
+    )
+    if not isinstance(instance_config, KubernetesDeploymentConfig):
+        error_message = (
+            f"Autoscaling is not supported for {service}.{instance} because instance type is not "
+            f"kubernetes."
+        )
+        raise ApiFailure(error_message, 501)
+
+    json_body = request.swagger_data.get("json_body", {})
+    min_instances_override = json_body.get("min_instances")
+    expire_after = json_body.get("expire_after")
+
+    if not isinstance(min_instances_override, int) or min_instances_override < 1:
+        raise ApiFailure("min_instances must be a positive integer", 400)
+
+    if not expire_after:
+        raise ApiFailure("expire_after is required", 400)
+
+    max_instances = instance_config.get_max_instances()
+    if max_instances is None:
+        raise ApiFailure(f"Autoscaling is not enabled for {service}.{instance}", 400)
+
+    if max_instances < min_instances_override:
+        raise ApiFailure(
+            f"min_instances ({min_instances_override}) cannot be greater than max_instances ({max_instances})",
+            400,
+        )
+
+    configmap, created = get_or_create_autoscaling_overrides_configmap()
+    if created:
+        log.info("Created new autoscaling overrides ConfigMap")
+    # i dunno why this is necessary, but a newly created configmap doesn't have a data field
+    # even when we set it in the create call
+    if not configmap.data:
+        configmap.data = {}
+
+    override_data = {
+        "min_instances": min_instances_override,
+        "created_at": datetime.now(timezone.utc).isoformat(),
+        # NOTE: we may want to also allow setting a max_instances override in the future, but if we do that
+        # we'd probably want to force folks to either set one or both and share the same expiration time
+        "expire_after": expire_after,
+    }
+
+    service_instance = f"{service}.{instance}"
+    existing_overrides = (
+        json.loads(configmap.data[service_instance])
+        if service_instance in configmap.data
+        else {}
+    )
+    merged_overrides = {**existing_overrides, **override_data}
+    serialized_overrides = json.dumps(merged_overrides)
+
+    patch_namespaced_configmap(
+        name=AUTOSCALING_OVERRIDES_CONFIGMAP_NAME,
+        namespace=AUTOSCALING_OVERRIDES_CONFIGMAP_NAMESPACE,
+        # this should only update the single entry for the $service.$instance key
+        # ain't k8s grand?
+        body={"data": {service_instance: serialized_overrides}},
+        kube_client=settings.kubernetes_client,
+    )
+
+    response_body = {
+        "service": service,
+        "instance": instance,
+        "cluster": cluster,
+        "min_instances": min_instances_override,
+        "expire_after": expire_after,
+        "status": "SUCCESS",
+    }
+    # NOTE: this is an HTTP 202 since actually updating the HPA happens asynchronously
+    # through setup_kubernetes_job
+    # XXX: should we try to patch things here as well?
+    return Response(json_body=response_body, status_code=202)

paasta_tools/cli/cmds/autoscale.py

@@ -13,6 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
+import time
+from datetime import datetime
+from typing import Optional
+
+import pytz
+from pytimeparse.timeparse import timeparse
 
 import paasta_tools.paastaapi.models as paastamodels
 from paasta_tools.api import client
@@ -58,6 +64,24 @@ def add_subparser(subparsers):
     autoscale_parser.add_argument(
         "--set", help="Set the number to scale to. Must be an Int.", type=int
     )
+
+    # Temporary override options
+    override_group = autoscale_parser.add_argument_group("Temporary Override Options")
+    override_group.add_argument(
+        "--set-min",
+        help="Set the minimum number of replicas (must be >= 1). Requires --for parameter.",
+        type=lambda x: int(x)
+        if int(x) >= 1
+        else autoscale_parser.error("Minimum instances must be >= 1"),
+        default=None,
+    )
+    override_group.add_argument(
+        "--for",
+        dest="duration",
+        help="Duration for the temporary override (e.g. '3h', '30m'). Required when using --set-min.",
+        default=None,
+    )
+
     autoscale_parser.add_argument(
         "-d",
         "--soa-dir",
@@ -69,9 +93,46 @@ def add_subparser(subparsers):
     autoscale_parser.set_defaults(command=paasta_autoscale)
 
 
+def parse_duration_to_seconds(duration: str) -> Optional[int]:
+    """Parse a duration string like '3h' or '30m' into seconds.
+
+    Args:
+        duration_str: A string representing a duration (e.g., "3h", "30m", "1d")
+
+    Returns:
+        The duration in seconds, or None if parsing failed
+    """
+    if not duration:
+        return None
+
+    seconds = timeparse(duration)
+    return seconds
+
+
 def paasta_autoscale(args):
     log.setLevel(logging.DEBUG)
     service = figure_out_service_name(args)
+
+    if args.set_min is not None and not args.duration:
+        print(
+            PaastaColors.yellow(
+                "WARNING: --set-min requires --for parameter to specify duration - defaulting to 30m"
+            )
+        )
+        args.duration = "30m"
+
+    if args.duration is not None and args.set_min is None:
+        print(PaastaColors.red("Error: --for requires --set-min parameter"))
+        return 1
+
+    if args.set is not None and args.set_min is not None:
+        print(
+            PaastaColors.red(
+                "Error: Cannot use both --set and --set-min at the same time"
+            )
+        )
+        return 1
+
     instance_config = next(
         get_instance_configs_for_service(
             service=service,
@@ -83,7 +144,7 @@ def paasta_autoscale(args):
     )
     if not instance_config:
         print(
-            "Could not find config files for this service instance in soaconfigs. Maybe you mispelled an argument?"
+            "Could not find config files for this service instance in soaconfigs. Maybe you misspelled an argument?"
         )
         return 1
 
@@ -99,12 +160,15 @@ def paasta_autoscale(args):
         return 1
 
     try:
-        if args.set is None:
+        # get current autoscaler count
+        if args.set is None and args.set_min is None:
             log.debug("Getting the current autoscaler count...")
             res, status, _ = api.autoscaler.get_autoscaler_count(
                 service=service, instance=args.instance, _return_http_data_only=False
             )
-        else:
+
+        # set desired instances
+        elif args.set is not None:
             log.debug(f"Setting desired instances to {args.set}.")
             msg = paastamodels.AutoscalerCountMsg(desired_instances=int(args.set))
             res, status, _ = api.autoscaler.update_autoscaler_count(
@@ -121,16 +185,53 @@ def paasta_autoscale(args):
                 instance=args.instance,
                 cluster=args.cluster,
             )
+
+        # set lower bound
+        elif args.set_min is not None:
+            duration_seconds = parse_duration_to_seconds(args.duration)
+            if not duration_seconds:
+                print(
+                    PaastaColors.red(
+                        f"Error: Invalid duration format '{args.duration}'. "
+                        f"Please use a format like '3h' or '30m'."
+                    )
+                )
+                return 1
+            # NOTE: this is explicitly using time.time() since we're doing everything using epoch time
+            # for simplicity
+            expiration_time = time.time() + duration_seconds
+
+            log.debug(
+                f"Setting minimum instances to {args.set_min} for duration {args.duration}."
+            )
+            msg = paastamodels.AutoscalingOverride(
+                min_instances=args.set_min,
+                expire_after=expiration_time,
+            )
+
+            res, status, _ = api.autoscaler.set_autoscaling_override(
+                service=service,
+                instance=args.instance,
+                autoscaling_override=msg,
+                _return_http_data_only=False,
+            )
+            _log_audit(
+                action="manual-autoscale-override",
+                action_details=str(msg),
+                service=service,
+                instance=args.instance,
+                cluster=args.cluster,
+            )
     except api.api_error as exc:
         status = exc.status
 
     if not 200 <= status <= 299:
         print(
             PaastaColors.red(
-                f"ERROR: '{args.instance}' is not configured to autoscale, "
-                f"so paasta autoscale could not scale it up on demand. "
+                f"ERROR: '{args.instance}' is not configured to autoscale OR you set min_instances above the current max_instances, "
+                f"and `paasta autoscale` could not update it. "
                 f"If you want to be able to boost this service, please configure autoscaling for the service "
-                f"in its config file by setting min and max instances. Example: \n"
+                f"in its config file by setting min and max instances appropriately. Example: \n"
                 f"{args.instance}:\n"
                 f"     min_instances: 5\n"
                 f"     max_instances: 50"
@@ -139,5 +240,23 @@ def paasta_autoscale(args):
         return 0
 
     log.debug(f"Res: {res} Http: {status}")
-    print(res.desired_instances)
+    if not args.set_min:
+        print(f"Desired instances: {res.desired_instances}")
+    elif args.set_min:
+        print(
+            f"Temporary override set for {args.service}.{args.instance} with minimum instances: {args.set_min}"
+        )
+        # folks using this might be in different timezones, so let's convert the expiration time to a few common ones
+        # to make it extra clear when the override will expire
+        epoch_time = datetime.fromtimestamp(res.expire_after)
+        eastern_time = epoch_time.astimezone(pytz.timezone("US/Eastern"))
+        pacific_time = epoch_time.astimezone(pytz.timezone("US/Pacific"))
+        london_time = epoch_time.astimezone(pytz.timezone("Europe/London"))
+
+        time_format = "%Y-%m-%d %H:%M:%S %Z%z"
+        print(f"The {args.duration} override will expire at:")
+        print(f"Eastern Time: {eastern_time.strftime(time_format)}")
+        print(f"Pacific Time: {pacific_time.strftime(time_format)}")
+        print(f"London Time:  {london_time.strftime(time_format)}")
+
     return 0

paasta_tools/kubernetes/application/controller_wrappers.py

@@ -18,6 +18,7 @@ from paasta_tools.kubernetes_tools import create_job
 from paasta_tools.kubernetes_tools import create_pod_disruption_budget
 from paasta_tools.kubernetes_tools import create_stateful_set
 from paasta_tools.kubernetes_tools import ensure_service_account
+from paasta_tools.kubernetes_tools import HpaOverride
 from paasta_tools.kubernetes_tools import KubeClient
 from paasta_tools.kubernetes_tools import KubeDeployment
 from paasta_tools.kubernetes_tools import KubernetesDeploymentConfig
@@ -219,6 +220,15 @@ class Application(ABC):
 
 
 class DeploymentWrapper(Application):
+    def __init__(
+        self,
+        item: Union[V1Deployment, V1StatefulSet],
+        logging=logging.getLogger(__name__),
+        hpa_override: Optional[HpaOverride] = None,
+    ) -> None:
+        super().__init__(item, logging)
+        self.hpa_override = hpa_override
+
     def deep_delete(
         self, kube_client: KubeClient, propagation_policy="Foreground"
     ) -> None:
@@ -290,6 +300,9 @@ class DeploymentWrapper(Application):
             cluster=self.soa_config.cluster,
             kube_client=kube_client,
             namespace=self.item.metadata.namespace,
+            min_instances_override=(
+                self.hpa_override["min_instances"] if self.hpa_override else None
+            ),
         )
 
         hpa_exists = self.exists_hpa(kube_client)
@@ -461,11 +474,12 @@ class JobWrapper(Application):
 
 
 def get_application_wrapper(
-    formatted_application: Union[V1Deployment, V1StatefulSet, V1Job]
+    formatted_application: Union[V1Deployment, V1StatefulSet, V1Job],
+    hpa_override: Optional[HpaOverride] = None,
 ) -> Application:
     app: Application
     if isinstance(formatted_application, V1Deployment):
-        app = DeploymentWrapper(formatted_application)
+        app = DeploymentWrapper(formatted_application, hpa_override=hpa_override)
     elif isinstance(formatted_application, V1StatefulSet):
         app = StatefulSetWrapper(formatted_application)
     elif isinstance(formatted_application, V1Job):

paasta_tools/kubernetes_tools.py

@@ -233,6 +233,9 @@ DEFAULT_SIDECAR_REQUEST: KubeContainerResourceRequest = {
 DEFAULT_PROJECTED_SA_EXPIRATION_SECONDS = 3600
 PROJECTED_SA_TOKEN_PATH = "token"
 
+AUTOSCALING_OVERRIDES_CONFIGMAP_NAME = "paasta-autoscaling-overrides"
+AUTOSCALING_OVERRIDES_CONFIGMAP_NAMESPACE = "paasta"
+
 
 # conditions is None when creating a new HPA, but the client raises an error in that case.
 # For detail, https://github.com/kubernetes-client/python/issues/553
@@ -320,6 +323,11 @@ class DatastoreCredentialsConfig(TypedDict, total=False):
     mysql: List[str]
 
 
+class HpaOverride(TypedDict):
+    min_instances: int
+    expire_after: str
+
+
 def _set_disrupted_pods(self: Any, disrupted_pods: Mapping[str, datetime]) -> None:
     """Private function used to patch the setter for V1PodDisruptionBudgetStatus.
     Can be removed once https://github.com/kubernetes-client/python/issues/466 is resolved
@@ -890,6 +898,7 @@ class KubernetesDeploymentConfig(LongRunningServiceConfig):
         cluster: str,
         kube_client: KubeClient,
         namespace: str,
+        min_instances_override: Optional[int] = None,
     ) -> Optional[V2HorizontalPodAutoscaler]:
         # Returns None if an HPA should not be attached based on the config,
         # or the config is invalid.
@@ -904,7 +913,7 @@ class KubernetesDeploymentConfig(LongRunningServiceConfig):
         if autoscaling_params["metrics_providers"][0]["decision_policy"] == "bespoke":
             return None
 
-        min_replicas = self.get_min_instances()
+        min_replicas = min_instances_override or self.get_min_instances()
         max_replicas = self.get_max_instances()
         if min_replicas == 0 or max_replicas == 0:
             log.error(
@@ -4647,6 +4656,73 @@ def get_kubernetes_secret_volumes(
     return secret_volumes
 
 
+def get_namespaced_configmap(
+    name: str, *, namespace: str, kube_client: KubeClient
+) -> Optional[V1ConfigMap]:
+    try:
+        return kube_client.core.read_namespaced_config_map(
+            name=name, namespace=namespace
+        )
+    except ApiException as e:
+        if e.status == 404:
+            return None
+        else:
+            raise
+
+
+def patch_namespaced_configmap(
+    name: str,
+    body: Dict[str, str],
+    *,
+    namespace: str,
+    kube_client: KubeClient,
+) -> V1ConfigMap:
+    """
+    Patches a configmap with the given body. The body should be a dictionary of key-value pairs.
+    """
+    try:
+        return kube_client.core.patch_namespaced_config_map(
+            name=name, namespace=namespace, body=body
+        )
+    except ApiException as e:
+        if e.status == 404:
+            raise ValueError(f"ConfigMap {name} not found in namespace {namespace}")
+        else:
+            raise
+
+
+def get_or_create_namespaced_configmap(
+    configmap: str,
+    *,
+    namespace: str,
+    kube_client: KubeClient,
+) -> Tuple[V1ConfigMap, bool]:
+    """
+    Returns a 2-tuple of (the configmap, a bool representing whether it was just created)
+    """
+    try:
+        return (
+            kube_client.core.read_namespaced_config_map(
+                name=configmap, namespace=namespace
+            ),
+            False,
+        )
+    except ApiException as e:
+        if e.status == 404:
+            configmap = V1ConfigMap(
+                metadata=V1ObjectMeta(name=configmap, namespace=namespace),
+                data={},
+            )
+            return (
+                kube_client.core.create_namespaced_config_map(
+                    namespace=namespace, body=configmap
+                ),
+                True,
+            )
+        else:
+            raise
+
+
 @lru_cache()
 def get_authenticating_services(soa_dir: str = DEFAULT_SOA_DIR) -> Set[str]:
     """Load list of services participating in authenticated traffic"""