From b9a187b332e9b388d23aeec6e52df3266f26827d Mon Sep 17 00:00:00 2001
From: Evan Krall <krall@yelp.com>
Date: Wed, 3 Apr 2024 16:34:26 -0700
Subject: [PATCH] Add --memory, --cpu, --cpu-limit, and --disk options to
 paasta local-run to allow overriding of settings from yelpsoa-configs.

---
 paasta_tools/cli/cmds/local_run.py | 55 ++++++++++++++++++++++++++++++
 tests/cli/test_cmds_local_run.py   |  2 ++
 2 files changed, 57 insertions(+)

diff --git a/paasta_tools/cli/cmds/local_run.py b/paasta_tools/cli/cmds/local_run.py
index ef478cb772..7f321384e0 100755
--- a/paasta_tools/cli/cmds/local_run.py
+++ b/paasta_tools/cli/cmds/local_run.py
@@ -66,6 +66,7 @@
 from paasta_tools.utils import NoDockerImageError
 from paasta_tools.utils import PaastaColors
 from paasta_tools.utils import PaastaNotConfiguredError
+from paasta_tools.utils import PATH_TO_SYSTEM_PAASTA_CONFIG_DIR
 from paasta_tools.utils import SystemPaastaConfig
 from paasta_tools.utils import Timeout
 from paasta_tools.utils import TimeoutError
@@ -386,6 +387,51 @@ def add_subparser(subparsers):
         required=False,
         default=None,
     )
+    list_parser.add_argument(
+        "--memory",
+        "--mem",
+        help=(
+            "Specify a memory limit (in megabytes) for the container."
+            " If not specified, the memory limit is pulled from yelpsoa-configs for your --instance/--cluster."
+        ),
+        required=False,
+        default=None,
+        type=float,
+    )
+    list_parser.add_argument(
+        "--cpus",
+        help=(
+            "Specify a CPU limit (in cores) for the container."
+            " If not specified, the cpu limit is pulled from yelpsoa-configs for your --instance/--cluster."
+            " cpu_burst_add/--cpu-burst-add is added to this limit, to match the behavior of real clusters."
+        ),
+        required=False,
+        default=None,
+        type=float,
+    )
+    list_parser.add_argument(
+        "--cpu-burst-add",
+        help=(
+            "Override cpu_burst_add for the container."
+            " If not specified, cpu_burst_add is pulled from yelpsoa-configs for your --instance/--cluster."
+            " This value is added to your cpus/--cpus limit to determine the CPU limit for the container."
+        ),
+        required=False,
+        default=None,
+        type=float,
+    )
+    list_parser.add_argument(
+        "--disk",
+        help=(
+            "Override disk limit for the container."
+            " If not specified, the disk limit is pulled from yelpsoa-configs for your --instance/--cluster."
+            f" (Disk limits are only enforced if enforce_disk_quota is set in {PATH_TO_SYSTEM_PAASTA_CONFIG_DIR})"
+        ),
+        required=False,
+        default=None,
+        type=float,
+    )
+
     list_parser.add_argument(
         "-i",
         "--instance",
@@ -1174,6 +1220,15 @@ def configure_and_run_docker_container(
         )
         return 1
 
+    if args.memory:
+        instance_config.config_dict["mem"] = args.memory
+    if args.cpus:
+        instance_config.config_dict["cpus"] = args.cpus
+    if args.cpu_burst_add:
+        instance_config.config_dict["cpu_burst_add"] = args.cpu_burst_add
+    if args.disk:
+        instance_config.config_dict["disk"] = args.disk
+
     if docker_sha is not None:
         instance_config.branch_dict = {
             "git_sha": docker_sha,
diff --git a/tests/cli/test_cmds_local_run.py b/tests/cli/test_cmds_local_run.py
index cce8a48ed8..72fcc566e4 100644
--- a/tests/cli/test_cmds_local_run.py
+++ b/tests/cli/test_cmds_local_run.py
@@ -534,6 +534,7 @@ def test_configure_and_run_pulls_image_when_asked(
     fake_instance_config.get_docker_registry.return_value = "fake_registry"
     fake_instance_config.get_docker_image.return_value = "fake_image"
     fake_instance_config.get_docker_url.return_value = "fake_registry/fake_image"
+    fake_instance_config.config_dict = {}
     mock_get_instance_config.return_value = fake_instance_config
     fake_service = "fake_service"
     args = mock.MagicMock()
@@ -633,6 +634,7 @@ def test_configure_and_run_docker_container_defaults_to_interactive_instance(
         args.use_okta_role = False
 
         mock_config = mock.create_autospec(AdhocJobConfig)
+        mock_config.config_dict = {}
         mock_get_default_interactive_config.return_value = mock_config
         return_code = configure_and_run_docker_container(
             docker_client=mock_docker_client,