diff --git a/packages/core/canyonos_core/controller/cloud_provider_logic/EC2/_runtime.py b/packages/core/canyonos_core/controller/cloud_provider_logic/EC2/_runtime.py index b011b682..93580bec 100644 --- a/packages/core/canyonos_core/controller/cloud_provider_logic/EC2/_runtime.py +++ b/packages/core/canyonos_core/controller/cloud_provider_logic/EC2/_runtime.py @@ -30,6 +30,9 @@ from canyonos_core.controller.cloud_provider_logic.shared_utils.llm_proxy_env import ( llm_proxy_docker_env_args, ) +from canyonos_core.controller.cloud_provider_logic.shared_utils.resource_limits import ( + resource_limit_args, +) logger = logging.getLogger(__name__) @@ -223,6 +226,7 @@ def _bootstrap_instance( "-d", "--name", redis_container, + *resource_limit_args(), "-p", f"{redis_port}:6379", "redis:alpine", @@ -286,6 +290,7 @@ def _bootstrap_instance( "--name", container, *port_args, + *resource_limit_args(spec.get("resources")), "-e", f"CANYONOS_REDIS_HOST={redis_host}", "-e", diff --git a/packages/core/canyonos_core/controller/cloud_provider_logic/Local/_runtime.py b/packages/core/canyonos_core/controller/cloud_provider_logic/Local/_runtime.py index 6fe54e4a..1fb8dbde 100644 --- a/packages/core/canyonos_core/controller/cloud_provider_logic/Local/_runtime.py +++ b/packages/core/canyonos_core/controller/cloud_provider_logic/Local/_runtime.py @@ -15,6 +15,9 @@ from canyonos_core.controller.cloud_provider_logic.shared_utils.llm_proxy_env import ( llm_proxy_docker_env_args, ) +from canyonos_core.controller.cloud_provider_logic.shared_utils.resource_limits import ( + resource_limit_args, +) logger = logging.getLogger(__name__) @@ -70,7 +73,6 @@ def provision_instance(spec, replica_index, next_host_port): def bootstrap_instance(provisioned, spec, replica_index, agent_id): agent_name = spec["name"] - resources = spec.get("resources", {}) ctrl_type = spec.get("type", "agent") image = f"canyonos-{agent_name.lower()}" host = provisioned["host"] @@ -146,12 +148,7 @@ def bootstrap_instance(provisioned, spec, replica_index, agent_id): cmd.extend(["-e", f"CANYONOS_DATABASE_URL={db_url}"]) if project_id: cmd.extend(["-e", f"CANYONOS_PROJECT_ID={project_id}"]) - if resources.get("cpu"): - cmd.extend(["--cpus", str(resources["cpu"])]) - if resources.get("memory"): - cmd.extend(["--memory", f"{resources['memory']}m"]) - if resources.get("gpu"): - cmd.extend(["--gpus", str(resources["gpu"])]) + cmd.extend(resource_limit_args(spec.get("resources"))) # User secrets from `env_file`. Explicit -e flags above still win, so a # stray CANYONOS_* line in someone's .env cannot break agent wiring. diff --git a/packages/core/canyonos_core/controller/cloud_provider_logic/shared_utils/resource_limits.py b/packages/core/canyonos_core/controller/cloud_provider_logic/shared_utils/resource_limits.py new file mode 100644 index 00000000..f97383af --- /dev/null +++ b/packages/core/canyonos_core/controller/cloud_provider_logic/shared_utils/resource_limits.py @@ -0,0 +1,24 @@ +"""Default CPU/memory/GPU bounds applied to every container the controller launches.""" + +# Matches the per-agent resources GlobalController already publishes to Redis, +# so an agent that configures nothing is bounded to what it advertises. +DEFAULT_CPU = 1 +DEFAULT_MEMORY_MB = 512 +DEFAULT_GPU = 0 + + +def resource_limit_args(resources=None): + """Turn an agent's `resources:` -- {"cpu": cores, "memory": MB, "gpu": count}, any key absent or None -- into `docker run` --cpus/--memory/--gpus flags.""" + resources = resources or {} + gpu = resources.get("gpu") or DEFAULT_GPU + args = [ + "--cpus", + str(resources.get("cpu") or DEFAULT_CPU), + "--memory", + f"{resources.get('memory') or DEFAULT_MEMORY_MB}m", + ] + # `--gpus 0` is not a no-op: on a host without the NVIDIA runtime Docker + # refuses the request, so zero GPUs means omitting the flag entirely. + if gpu: + args.extend(["--gpus", str(gpu)]) + return args diff --git a/packages/core/canyonos_core/controller/global_controller.py b/packages/core/canyonos_core/controller/global_controller.py index 7e250636..017d588e 100644 --- a/packages/core/canyonos_core/controller/global_controller.py +++ b/packages/core/canyonos_core/controller/global_controller.py @@ -30,6 +30,9 @@ send_runtime_information, send_agent_information, ) +from canyonos_core.controller.cloud_provider_logic.shared_utils.resource_limits import ( + resource_limit_args, +) from canyonos_core.controller.utils.redis_client import RedisClient from canyonos_core.controller.utils.grpc_options import GRPC_CHANNEL_OPTIONS @@ -467,6 +470,7 @@ def _launch_redis_containers(self): "--name", container_name, *network_args, + *resource_limit_args(), "-p", f"{redis_port}:6379", "redis:alpine", diff --git a/packages/core/tests/test_container_resource_limits.py b/packages/core/tests/test_container_resource_limits.py new file mode 100644 index 00000000..89345aca --- /dev/null +++ b/packages/core/tests/test_container_resource_limits.py @@ -0,0 +1,221 @@ +"""CAN-402: every container the controller launches carries --cpus/--memory bounds.""" + +import os +import sys +import tempfile +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from canyonos_core.controller.cloud_provider_logic.EC2 import _runtime as ec2_runtime +from canyonos_core.controller.cloud_provider_logic.Local import ( + _runtime as local_runtime, +) +from canyonos_core.controller.cloud_provider_logic.shared_utils.resource_limits import ( + resource_limit_args, +) +from canyonos_core.controller.global_controller import GlobalController + + +def _limits_of(cmd): + """Return the {flag: value} pairs the given `docker run` argv bounds the container with.""" + return { + flag: cmd[cmd.index(flag) + 1] + for flag in ("--cpus", "--memory", "--gpus") + if flag in cmd + } + + +class ResourceLimitArgsTests(unittest.TestCase): + def test_absent_resources_fall_back_to_the_defaults(self): + self.assertEqual(resource_limit_args(), ["--cpus", "1", "--memory", "512m"]) + self.assertEqual(resource_limit_args({}), ["--cpus", "1", "--memory", "512m"]) + + def test_global_controller_yaml_values_win_over_the_defaults(self): + self.assertEqual( + resource_limit_args({"cpu": 2, "memory": 2048}), + ["--cpus", "2", "--memory", "2048m"], + ) + + def test_a_partial_resources_block_defaults_only_the_missing_half(self): + self.assertEqual( + resource_limit_args({"memory": 2048}), + ["--cpus", "1", "--memory", "2048m"], + ) + + def test_a_gpu_request_is_passed_through(self): + self.assertEqual( + resource_limit_args({"gpu": 2}), + ["--cpus", "1", "--memory", "512m", "--gpus", "2"], + ) + + def test_zero_gpus_omits_the_flag_rather_than_requesting_none(self): + """`--gpus 0` fails outright on a host with no NVIDIA runtime.""" + for resources in ({}, {"gpu": 0}, {"gpu": None}): + with self.subTest(resources=resources): + self.assertNotIn("--gpus", resource_limit_args(resources)) + + +class LocalRuntimeLimitTests(unittest.TestCase): + def setUp(self): + self.original_controller = local_runtime._controller + self.controller = SimpleNamespace( + config={"poll_interval": 5}, + redis=MagicMock(), + _run_cmd=MagicMock(return_value=SimpleNamespace(returncode=0, stdout="")), + ) + local_runtime._controller = self.controller + + def tearDown(self): + local_runtime._controller = self.original_controller + + def _launch(self, spec): + provisioned = local_runtime.provision_instance(spec, 0, lambda host: 8000) + with patch.object(local_runtime, "_port_bound", return_value=False): + local_runtime.bootstrap_instance(provisioned, spec, 0, "agent-id-0") + return self.controller._run_cmd.call_args.args[0] + + def test_an_agent_without_a_resources_block_is_still_bounded(self): + cmd = self._launch({"name": "Alpha", "provider": "local"}) + + self.assertEqual(_limits_of(cmd), {"--cpus": "1", "--memory": "512m"}) + + def test_a_local_gpu_agent_still_gets_its_gpu_flag(self): + cmd = self._launch( + { + "name": "Vllm", + "provider": "local", + "resources": {"cpu": 2, "memory": 2048, "gpu": 1}, + } + ) + + self.assertEqual( + _limits_of(cmd), + {"--cpus": "2", "--memory": "2048m", "--gpus": "1"}, + ) + + def test_an_agent_with_a_resources_block_is_bounded_to_it(self): + cmd = self._launch( + { + "name": "Alpha", + "provider": "local", + "resources": {"cpu": 2, "memory": 2048}, + } + ) + + self.assertEqual(_limits_of(cmd), {"--cpus": "2", "--memory": "2048m"}) + + +class EC2RuntimeLimitTests(unittest.TestCase): + def setUp(self): + self.original_controller = ec2_runtime._controller + key_file = tempfile.NamedTemporaryFile(delete=False) + key_file.close() + self.key_path = key_file.name + os.chmod(self.key_path, 0o600) + self.controller = SimpleNamespace( + config={ + "poll_interval": 5, + "ec2": {"ssh_user": "ubuntu", "ssh_private_key_path": self.key_path}, + }, + redis_containers={}, + node_redis={}, + _run_cmd=MagicMock( + return_value=SimpleNamespace(returncode=0, stdout="", stderr="") + ), + ) + ec2_runtime._controller = self.controller + + def tearDown(self): + os.unlink(self.key_path) + ec2_runtime._controller = self.original_controller + + def _bootstrap(self, spec): + with ( + patch.object(ec2_runtime.time, "sleep"), + patch.object( + ec2_runtime.subprocess, + "run", + return_value=SimpleNamespace(returncode=0, stderr="", stdout=""), + ), + patch.object(ec2_runtime, "RedisClient", return_value=MagicMock()), + patch.object(ec2_runtime, "_wait_for_redis"), + ): + ec2_runtime._bootstrap_instance( + "10.0.0.30", + spec, + 0, + self.controller.config["ec2"], + redis_host="10.0.0.30", + redis_port=6379, + agent_id="agent-id-0", + ) + redis_cmd, agent_cmd = ( + call.args[0] + for call in self.controller._run_cmd.call_args_list + if call.args[0][:2] == ["docker", "run"] + ) + return redis_cmd, agent_cmd + + def test_the_ec2_agent_container_is_bounded_by_its_resources_block(self): + _, agent_cmd = self._bootstrap( + { + "name": "Tagged", + "provider": "EC2", + "resources": {"cpu": 2, "memory": 2048}, + } + ) + + self.assertEqual(_limits_of(agent_cmd), {"--cpus": "2", "--memory": "2048m"}) + + def test_the_ec2_agent_gets_its_gpu_flag_too(self): + _, agent_cmd = self._bootstrap( + {"name": "Vllm", "provider": "EC2", "resources": {"gpu": 1}} + ) + + self.assertEqual( + _limits_of(agent_cmd), + {"--cpus": "1", "--memory": "512m", "--gpus": "1"}, + ) + + def test_the_ec2_agent_and_redis_containers_default_when_unconfigured(self): + redis_cmd, agent_cmd = self._bootstrap({"name": "Tagged", "provider": "EC2"}) + + self.assertEqual(_limits_of(agent_cmd), {"--cpus": "1", "--memory": "512m"}) + self.assertEqual(_limits_of(redis_cmd), {"--cpus": "1", "--memory": "512m"}) + + +class GlobalControllerRedisLimitTests(unittest.TestCase): + def test_a_newly_launched_redis_container_is_bounded(self): + controller = GlobalController.__new__(GlobalController) + controller.controllers = [ + {"name": "Workflow", "replicas": 1, "redis_port": 6379} + ] + controller.redis_containers = {} + controller.node_redis = {} + controller.redis = None + controller._run_cmd = MagicMock( + return_value=SimpleNamespace(returncode=0, stdout="", stderr="") + ) + + with ( + patch( + "canyonos_core.controller.global_controller.RedisClient", + return_value=MagicMock(), + ), + patch("canyonos_core.controller.global_controller._wait_for_redis"), + ): + controller._launch_redis_containers() + + redis_cmd = next( + call.args[0] + for call in controller._run_cmd.call_args_list + if call.args[0][:2] == ["docker", "run"] + ) + self.assertEqual(_limits_of(redis_cmd), {"--cpus": "1", "--memory": "512m"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/packages/core/tests/test_instance_manager_runtime.py b/packages/core/tests/test_instance_manager_runtime.py index 3b967a11..372416e2 100644 --- a/packages/core/tests/test_instance_manager_runtime.py +++ b/packages/core/tests/test_instance_manager_runtime.py @@ -192,6 +192,10 @@ def test_local_instances_keep_default_host_and_increment_host_ports(self): "ANTHROPIC_API_BASE=http://127.0.0.1:8081/anthropic", "-e", "CANYONOS_LLM_STUB_TEXT=", + "--cpus", + "1", + "--memory", + "512m", "canyonos-alpha", ], "localhost",