feat: Expose docker memory limit to CLI and Python API

tesseract_core/sdk/cli.py

@@ -558,6 +558,17 @@ def serve(
             ),
         ),
     ] = None,
+    memory: Annotated[
+        str | None,
+        typer.Option(
+            "--memory",
+            "-m",
+            help=(
+                "Memory limit for the container (e.g., '512m', '2g'). "
+                "Minimum allowed value is 6m (6 megabytes)."
+            ),
+        ),
+    ] = None,
     input_path: Annotated[
         str | None,
         typer.Option(
@@ -618,6 +629,7 @@ def serve(
             debug=debug,
             num_workers=num_workers,
             user=user,
+            memory=memory,
             input_path=input_path,
             output_path=output_path,
             output_format=_enum_to_val(output_format),
@@ -968,6 +980,17 @@ def run_container(
             ),
         ),
     ] = None,
+    memory: Annotated[
+        str | None,
+        typer.Option(
+            "--memory",
+            "-m",
+            help=(
+                "Memory limit for the container (e.g., '512m', '2g'). "
+                "Minimum allowed value is 6m (6 megabytes)."
+            ),
+        ),
+    ] = None,
     invoke_help: Annotated[
         bool,
         typer.Option(
@@ -1042,6 +1065,7 @@ def run_container(
             environment=parsed_environment,
             network=network,
             user=user,
+            memory=memory,
         )
 
     except ImageNotFound as e:

tesseract_core/sdk/docker_client.py

@@ -572,6 +572,7 @@ def run(
         stdout: bool = True,
         stderr: bool = False,
         user: str | None = None,
+        memory: str | None = None,
         extra_args: list_[str] | None = None,  # noqa: UP006
     ) -> Container | tuple[bytes, bytes] | bytes:
         """Run a command in a container from an image.
@@ -595,6 +596,7 @@ def run(
             stdout: If True, return stdout.
             stderr: If True, return stderr.
             environment: Environment variables to set in the container.
+            memory: Memory limit for the container (e.g., "512m", "2g"). Minimum allowed is 6m.
             extra_args: Additional arguments to pass to the `docker run` CLI command.
 
         Returns:
@@ -625,6 +627,9 @@ def run(
         if user:
             optional_args.extend(["-u", user])
 
+        if memory:
+            optional_args.extend(["--memory", memory])
+
         if device_requests:
             gpus_str = ",".join(device_requests)
             optional_args.extend(["--gpus", f'"device={gpus_str}"'])

tesseract_core/sdk/engine.py

@@ -487,6 +487,7 @@ def serve(
     debug: bool = False,
     num_workers: int = 1,
     user: str | None = None,
+    memory: str | None = None,
     input_path: str | Path | None = None,
     output_path: str | Path | None = None,
     output_format: Literal["json", "json+base64", "json+binref"] | None = None,
@@ -510,6 +511,7 @@ def serve(
         num_workers: number of workers to use for serving the Tesseracts.
         user: user to run the Tesseracts as, e.g. '1000' or '1000:1000' (uid:gid).
               Defaults to the current user.
+        memory: Memory limit for the container (e.g., "512m", "2g"). Minimum allowed is 6m.
         input_path: Input path to read input files from, such as local directory or S3 URI.
         output_path: Output path to write output files to, such as local directory or S3 URI.
         output_format: Output format to use for the results.
@@ -603,6 +605,7 @@ def serve(
         detach=True,
         volumes=parsed_volumes,
         user=user,
+        memory=memory,
         environment=environment,
         extra_args=extra_args,
     )
@@ -771,6 +774,7 @@ def run_tesseract(
     environment: dict[str, str] | None = None,
     network: str | None = None,
     user: str | None = None,
+    memory: str | None = None,
     input_path: str | Path | None = None,
     output_path: str | Path | None = None,
     output_format: Literal["json", "json+base64", "json+binref"] | None = None,
@@ -791,6 +795,7 @@ def run_tesseract(
         network: name of the Docker network to connect the container to.
         user: user to run the Tesseract as, e.g. '1000' or '1000:1000' (uid:gid).
             Defaults to the current user.
+        memory: Memory limit for the container (e.g., "512m", "2g"). Minimum allowed is 6m.
         input_path: Input path to read input files from, such as local directory or S3 URI.
         output_path: Output path to write output files to, such as local directory or S3 URI.
         output_format: Format of the output.
@@ -871,6 +876,7 @@ def run_tesseract(
         remove=True,
         stderr=True,
         user=user,
+        memory=memory,
         extra_args=extra_args,
     )
     assert isinstance(result, tuple)

tesseract_core/sdk/tesseract.py

@@ -83,6 +83,7 @@ def from_image(
         gpus: list[str] | None = None,
         num_workers: int = 1,
         user: str | None = None,
+        memory: str | None = None,
         input_path: str | Path | None = None,
         output_path: str | Path | None = None,
         output_format: Literal["json", "json+base64"] = "json+base64",
@@ -112,6 +113,7 @@ def from_image(
             num_workers: number of workers to use for serving the Tesseracts.
             user: user to run the Tesseracts as, e.g. '1000' or '1000:1000' (uid:gid).
                 Defaults to the current user.
+            memory: Memory limit for the container (e.g., "512m", "2g"). Minimum allowed is 6m.
             input_path: Input path to read input files from, such as local directory or S3 URI.
             output_path: Output path to write output files to, such as local directory or S3 URI.
             output_format: Format to use for the output data (json+binref not yet supported).
@@ -140,6 +142,7 @@ def from_image(
             network=network,
             network_alias=network_alias,
             user=user,
+            memory=memory,
             input_path=input_path,
             output_path=output_path,
             output_format=output_format,

tests/endtoend_tests/test_endtoend.py

@@ -202,6 +202,27 @@ def test_tesseract_run_stdout(cli_runner, built_image_name):
             raise
 
 
+def test_run_with_memory(cli_runner, built_image_name):
+    """Ensure we can run a Tesseract command with memory limits."""
+    run_res = cli_runner.invoke(
+        app,
+        [
+            "run",
+            built_image_name,
+            "health",
+            "--memory",
+            "512m",
+        ],
+        catch_exceptions=False,
+    )
+    assert run_res.exit_code == 0, run_res.stderr
+    assert run_res.stdout
+
+    # Verify the command executed successfully
+    result = json.loads(run_res.stdout)
+    assert result["status"] == "ok"
+
+
 @pytest.mark.parametrize("user", [None, "root", "1000:1000"])
 def test_run_as_user(cli_runner, docker_client, built_image_name, user, docker_cleanup):
     """Ensure we can run a basic Tesseract image as any user."""
@@ -233,6 +254,39 @@ def test_run_as_user(cli_runner, docker_client, built_image_name, user, docker_c
     assert output.decode("utf-8").strip() == str(expected_user)
 
 
+@pytest.mark.parametrize("memory", ["512m", "1g", "256m"])
+def test_serve_with_memory(
+    cli_runner, docker_client, built_image_name, memory, docker_cleanup
+):
+    """Ensure we can serve a Tesseract with memory limits."""
+    run_res = cli_runner.invoke(
+        app,
+        [
+            "serve",
+            built_image_name,
+            "--memory",
+            memory,
+        ],
+        catch_exceptions=False,
+    )
+    assert run_res.exit_code == 0, run_res.stderr
+
+    serve_meta = json.loads(run_res.stdout)
+    container = docker_client.containers.get(serve_meta["container_name"])
+    docker_cleanup["containers"].append(container)
+
+    # Verify memory limit was set on container
+    container_inspect = docker_client.containers.get(container.id)
+    memory_limit = container_inspect.attrs["HostConfig"]["Memory"]
+
+    # Convert memory string to bytes for comparison
+    memory_value = int(memory[:-1])
+    memory_unit = memory[-1].lower()
+    expected_bytes = memory_value * (1024**2 if memory_unit == "m" else 1024**3)
+
+    assert memory_limit == expected_bytes
+
+
 def test_tesseract_serve_pipeline(
     cli_runner, docker_client, built_image_name, docker_cleanup
 ):

tests/sdk_tests/test_engine.py

@@ -153,6 +153,19 @@ def test_run_gpu(mocked_docker):
     assert res["device_requests"] == ["all"]
 
 
+def test_run_memory(mocked_docker):
+    """Test running a tesseract with memory limit."""
+    res_out, _ = engine.run_tesseract(
+        "foobar",
+        "apply",
+        ['{"inputs": {"a": [1, 2, 3], "b": [4, 5, 6]}}'],
+        memory="512m",
+    )
+
+    res = json.loads(res_out)
+    assert res["memory"] == "512m"
+
+
 def test_run_tesseract_file_input(mocked_docker, tmpdir):
     """Test running a tesseract with file input / output."""
     outdir = Path(tmpdir) / "output"
@@ -274,6 +287,33 @@ def test_serve_tesseracts(mocked_docker):
     # Teardown valid
     engine.teardown(json.loads(container_name_multi_tesseract)["name"])
 
+    # Serve with memory
+    container_name_with_memory, _ = engine.serve("vectoradd", memory="512m")
+    assert container_name_with_memory
+
+    # Teardown valid
+    engine.teardown(json.loads(container_name_with_memory)["name"])
+
+
+def test_serve_memory(mocked_docker):
+    """Test serving a tesseract with memory limit."""
+    res, _ = engine.serve(
+        "foobar",
+        memory="512m",
+    )
+
+    res = json.loads(res)
+    assert res["memory"] == "512m"
+
+    # Test with different memory values
+    res, _ = engine.serve(
+        "foobar",
+        memory="2g",
+    )
+
+    res = json.loads(res)
+    assert res["memory"] == "2g"
+
 
 def test_serve_tesseract_volumes(mocked_docker, tmpdir):
     """Test running a tesseract with volumes."""

tests/sdk_tests/test_tesseract.py

@@ -116,6 +116,7 @@ def test_serve_lifecycle(mock_serving, mock_clients):
         network_alias=None,
         host_ip="127.0.0.1",
         user=None,
+        memory=None,
         input_path=None,
         output_path=None,
         output_format="json+base64",