django-lambda-tasks 0.6.1__tar.gz → 0.6.2__tar.gz

{django_lambda_tasks-0.6.1 → django_lambda_tasks-0.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: django-lambda-tasks
-Version: 0.6.1
+Version: 0.6.2
 Summary: Run async tasks in a lambda function
 Requires-Python: >=3.10
 Requires-Dist: awslambdaric

{django_lambda_tasks-0.6.1 → django_lambda_tasks-0.6.2}/lambda_tasks/handler.py

@@ -17,6 +17,7 @@ import os
 import resource
 import sys
 import uuid
+from pathlib import Path
 
 import django
 from django.apps import apps
@@ -31,27 +32,43 @@ _cold_start_done: bool = False
 
 _MEMORY_RESERVED_MB = 128
 _MEMORY_MINIMUM_LIMIT_MB = 64
+_CGROUP_MEMORY_MAX_PATH = Path("/sys/fs/cgroup/memory.max")
+
+
+def _get_memory_mb() -> int | None:
+    """Get container memory limit from Lambda env var or cgroup v2."""
+    memory_mb_str = os.environ.get("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")
+    if memory_mb_str is not None:
+        return int(memory_mb_str)
+
+    try:
+        value = _CGROUP_MEMORY_MAX_PATH.read_text().strip()
+        if value == "max":
+            return None
+        return int(value) // (1024 * 1024)
+    except (FileNotFoundError, ValueError):
+        return None
 
 
 def _set_memory_limit() -> None:
-    """Set RLIMIT_DATA from AWS_LAMBDA_FUNCTION_MEMORY_SIZE if available.
+    """Set RLIMIT_AS from the container memory limit.
 
-    When running in Lambda, this causes Python to raise MemoryError on
-    excessive allocation instead of being killed by the OOM killer.
+    On Lambda, reads AWS_LAMBDA_FUNCTION_MEMORY_SIZE. On Fargate/Batch,
+    reads the cgroup v2 memory limit. This causes Python to raise
+    MemoryError on excessive allocation instead of being killed by the
+    OOM killer.
 
     128 MB is reserved for the Python runtime, shared libraries, and OS
-    overhead. The limit is floored at 64 MB so that even small Lambdas
+    overhead. The limit is floored at 64 MB so that even small containers
     get some protection.
     """
-    memory_mb = os.environ.get("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")
+    memory_mb = _get_memory_mb()
 
     if memory_mb is None:
         return
 
     limit_bytes = (
-        max(int(memory_mb) - _MEMORY_RESERVED_MB, _MEMORY_MINIMUM_LIMIT_MB)
-        * 1024
-        * 1024
+        max(memory_mb - _MEMORY_RESERVED_MB, _MEMORY_MINIMUM_LIMIT_MB) * 1024 * 1024
     )
     resource.setrlimit(resource.RLIMIT_AS, (limit_bytes, limit_bytes))
     logger.info(

{django_lambda_tasks-0.6.1 → django_lambda_tasks-0.6.2}/pyproject.toml

@@ -7,7 +7,7 @@ packages = ["lambda_tasks"]
 
 [project]
 name = "django-lambda-tasks"
-version = "0.6.1"
+version = "0.6.2"
 description = "Run async tasks in a lambda function"
 readme = "README.md"
 requires-python = ">=3.10"

{django_lambda_tasks-0.6.1 → django_lambda_tasks-0.6.2}/tests/test_memory_limit.py

@@ -7,11 +7,24 @@ killer.
 """
 
 import resource
-from unittest.mock import patch
 
 import lambda_tasks.handler as handler_module
 
 
+class _FakePath:
+    def __init__(self, *, content: str | None) -> None:
+        self._content = content
+
+    def read_text(self) -> str:
+        if self._content is None:
+            raise FileNotFoundError
+        return self._content
+
+
+def fake_path(content: str | None) -> _FakePath:
+    return _FakePath(content=content)
+
+
 class TestMemoryLimitSetDuringColdStart:
     def test_rlimit_as_set_from_env_var(self, monkeypatch):
         """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE is set, RLIMIT_AS is configured with 128 MB reserved."""
@@ -35,14 +48,65 @@ class TestMemoryLimitSetDuringColdStart:
         expected = (512 - 128) * 1024 * 1024
         assert calls == [(resource.RLIMIT_AS, (expected, expected))]
 
-    def test_rlimit_as_not_set_when_env_var_missing(self, monkeypatch):
-        """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE is not set, RLIMIT_AS is unchanged."""
+    def test_rlimit_as_not_set_when_env_var_missing_and_no_cgroup(self, monkeypatch):
+        """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE is not set and cgroup is unavailable, RLIMIT_AS is unchanged."""
+        monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
+        monkeypatch.delenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", raising=False)
+        monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
+        monkeypatch.setattr(
+            "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
+        )
+        monkeypatch.setattr(handler_module, "_CGROUP_MEMORY_MAX_PATH", fake_path(None))
+
+        calls: list[tuple] = []
+
+        def fake_setrlimit(which: int, limits: tuple[int, int]) -> None:
+            calls.append((which, limits))
+
+        monkeypatch.setattr(resource, "setrlimit", fake_setrlimit)
+
+        handler_module._cold_start_done = False
+        handler_module.handler(event={"Records": []}, context=None)
+
+        assert calls == []
+
+    def test_rlimit_as_set_from_cgroup(self, monkeypatch):
+        """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE is not set, falls back to cgroup v2."""
+        monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
+        monkeypatch.delenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", raising=False)
+        monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
+        monkeypatch.setattr(
+            "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
+        )
+        # 2 GB in bytes
+        monkeypatch.setattr(
+            handler_module,
+            "_CGROUP_MEMORY_MAX_PATH",
+            fake_path(str(2048 * 1024 * 1024)),
+        )
+
+        calls: list[tuple] = []
+
+        def fake_setrlimit(which: int, limits: tuple[int, int]) -> None:
+            calls.append((which, limits))
+
+        monkeypatch.setattr(resource, "setrlimit", fake_setrlimit)
+
+        handler_module._cold_start_done = False
+        handler_module.handler(event={"Records": []}, context=None)
+
+        expected = (2048 - 128) * 1024 * 1024
+        assert calls == [(resource.RLIMIT_AS, (expected, expected))]
+
+    def test_rlimit_as_not_set_when_cgroup_is_max(self, monkeypatch):
+        """When cgroup memory.max is 'max' (unlimited), RLIMIT_AS is unchanged."""
         monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
         monkeypatch.delenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", raising=False)
         monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
         monkeypatch.setattr(
             "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
         )
+        monkeypatch.setattr(handler_module, "_CGROUP_MEMORY_MAX_PATH", fake_path("max"))
 
         calls: list[tuple] = []