diff --git a/tesseract_core/sdk/cli.py b/tesseract_core/sdk/cli.py index c5f01110..edc837a2 100755 --- a/tesseract_core/sdk/cli.py +++ b/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: diff --git a/tesseract_core/sdk/docker_client.py b/tesseract_core/sdk/docker_client.py index 2ec0acfc..c1bbd69b 100644 --- a/tesseract_core/sdk/docker_client.py +++ b/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}"']) diff --git a/tesseract_core/sdk/engine.py b/tesseract_core/sdk/engine.py index f6577f01..3df1b256 100644 --- a/tesseract_core/sdk/engine.py +++ b/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) diff --git a/tesseract_core/sdk/tesseract.py b/tesseract_core/sdk/tesseract.py index d6ddfb95..50893345 100644 --- a/tesseract_core/sdk/tesseract.py +++ b/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, diff --git a/tests/endtoend_tests/test_endtoend.py b/tests/endtoend_tests/test_endtoend.py index a0afa118..ce5efd2e 100644 --- a/tests/endtoend_tests/test_endtoend.py +++ b/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 ): diff --git a/tests/sdk_tests/test_engine.py b/tests/sdk_tests/test_engine.py index e87acceb..1de2db64 100644 --- a/tests/sdk_tests/test_engine.py +++ b/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.""" diff --git a/tests/sdk_tests/test_tesseract.py b/tests/sdk_tests/test_tesseract.py index e42f8d23..85caa0b8 100644 --- a/tests/sdk_tests/test_tesseract.py +++ b/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",