Skip to content

Commit

Permalink
Make resource limiting (and dependency on cgroups v2) optional.
Browse files Browse the repository at this point in the history
This adds a new `REQUIRE_RESOURCE_LIMITING` valve which is enabled by
default, but may be turned off to remove the dependency on cgroups v2
for code evaluation sandbox resource limiting. This is unsafe to do in
multi-user setups but may be OK for trusted single-user setups.

Fixes #14.
  • Loading branch information
EtiennePerot committed Oct 3, 2024
1 parent 5be7e04 commit 8585d81
Show file tree
Hide file tree
Showing 8 changed files with 886 additions and 620 deletions.
6 changes: 6 additions & 0 deletions docs/setup.md
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ This is adequate for single-user setups not exposed to the outside Internet, whi
* On **Docker**: Add `--mount=type=bind,source=/sys/fs/cgroup,target=/sys/fs/cgroup,readonly=false` to `docker run`.
* On **Kubernetes**: Add a [`hostPath` volume](https://kubernetes.io/docs/concepts/storage/volumes/#hostpath) with `path` set to `/sys/fs/cgroup`, then mount it in your container's `volumeMounts` with options `mountPath` set to `/sys/fs/cgroup` and `readOnly` set to `false`.
* **Why**: This is needed so that gVisor can create child [cgroups](https://en.wikipedia.org/wiki/Cgroups), necessary to enforce per-sandbox resource usage limits.
* If you wish to disable resource limiting on code evaluation sandboxes, you can skip this setting and not mount `cgroupfs` at all in the container. Note that this means code evaluation sandboxes will be able to take as much CPU and memory as they want.
* **Mount `procfs` at `/proc2`**:
* On **Docker**: Add `--mount=type=bind,source=/proc,target=/proc2,readonly=false,bind-recursive=disabled` to `docker run`.
* On **Kubernetes**: Add a [`hostPath` volume](https://kubernetes.io/docs/concepts/storage/volumes/#hostpath) with `path` set to `/proc`, then mount it in your container's `volumeMounts` with options `mountPath` set to `/proc2` and `readOnly` set to `false`.
Expand Down Expand Up @@ -179,6 +180,10 @@ While this document will not elaborate on how, it should be fairly obvious how o
* Useful for multi-user setups to avoid denial-of-service, and to avoid running LLM-generated code that contains infinite loops forever.
* **Max RAM**: The maximum amount of memory the sandboxed code will be allowed to use, in megabytes.
* Useful for multi-user setups to avoid denial-of-service.
* **Resource limiting enforcement**: Whether to enforce that code evaluation sandboxes are resource-limited.
* This is enabled by default, and requires cgroups v2 to be present on your system and mounted in the Open WebUI container.
* If you do not mind your code evaluation sandboxes being able to use as much CPU and memory as they want, you may disable this setting (set it to `false`).
* On systems that only have cgroups v1 and not cgroups v2, such as WSL and some old Linux distributions, you may need to disable this.
* **Auto install**: Whether to automatically download and install gVisor if not present in the container.
* If not installed, gVisor will be automatically installed in `/tmp`.
* You can set the HTTPS proxy used for this download using the `HTTPS_PROXY` environment variable.
Expand Down Expand Up @@ -281,6 +286,7 @@ Once you have done this, consider setting the following environment variable:
* `CODE_EVAL_VALVE_OVERRIDE_MAX_FILES_PER_EXECUTION`: The maximum number of newly-created files to retain in each sandbox execution. **This should be non-zero**.
* `CODE_EVAL_VALVE_OVERRIDE_MAX_FILES_PER_USER`: The maximum number of files that can be stored long-term for a single user. **This should be non-zero**.
* `CODE_EVAL_VALVE_OVERRIDE_MAX_MEGABYTES_PER_USER`: The maximum amount of total long-term file storage (in megabytes) that each user may use. **This should be non-zero**.
* `CODE_EVAL_VALVE_OVERRIDE_REQUIRE_RESOURCE_LIMITING`: Whether to require that code evaluation sandboxes are resource-limited. **This should be set to `true`**.
* `CODE_EVAL_VALVE_OVERRIDE_WEB_ACCESSIBLE_DIRECTORY_PATH`: The directory where user files are stored. **This should be overridden** to prevent it from being modified by users to reveal or overwrite sensitive files in the Open WebUI installation.
* `CODE_EVAL_VALVE_OVERRIDE_WEB_ACCESSIBLE_DIRECTORY_URL`: The HTTP URL of the directory specified by `CODE_EVAL_VALVE_OVERRIDE_WEB_ACCESSIBLE_DIRECTORY_PATH`. This can start with a `/` to make it domain-relative. **This should be overridden** to prevent users from modifying it in such a way that it tricks other users into clicking unrelated links.
* `CODE_EVAL_VALVE_OVERRIDE_NETWORKING_ALLOWED`: **This should be set to `false`** if running on a LAN with sensitive services that sandboxes could reach out to. Firewall rules are not yet supported, so this setting is currently all-or-nothing.
Expand Down
469 changes: 264 additions & 205 deletions open-webui/functions/run_code.py

Large diffs are not rendered by default.

471 changes: 265 additions & 206 deletions open-webui/tools/run_code.py

Large diffs are not rendered by default.

31 changes: 28 additions & 3 deletions src/openwebui/functions/run_code.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,10 @@ class Valves(pydantic.BaseModel):
default=256,
description=f"Maximum total size of files to keep around for a given user; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}MAX_MEGABYTES_PER_USER.",
)
REQUIRE_RESOURCE_LIMITING: bool = pydantic.Field(
default=True,
description=f"Whether to enforce resource limiting, which requires cgroups v2 to be available; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}REQUIRE_RESOURCE_LIMITING.",
)
WEB_ACCESSIBLE_DIRECTORY_PATH: str = pydantic.Field(
default="$DATA_DIR/cache/functions/run_code",
description=f"Path of the directory to write files that should be accessible for user download in. If it begins by '$DATA_DIR', this will be replaced with the DATA_DIR environment variable. The whole field may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}WEB_ACCESSIBLE_DIRECTORY_PATH.",
Expand Down Expand Up @@ -230,6 +234,7 @@ async def _fail(error_message, status="SANDBOX_ERROR"):
Sandbox.check_setup(
language=language,
auto_install_allowed=self.valves.AUTO_INSTALL,
require_resource_limiting=self.valves.REQUIRE_RESOURCE_LIMITING,
)

if self.valves.AUTO_INSTALL and Sandbox.runsc_needs_installation():
Expand Down Expand Up @@ -265,6 +270,7 @@ async def _fail(error_message, status="SANDBOX_ERROR"):
networking_allowed=valves.NETWORKING_ALLOWED,
max_runtime_seconds=valves.MAX_RUNTIME_SECONDS,
max_ram_bytes=max_ram_bytes,
require_resource_limiting=valves.REQUIRE_RESOURCE_LIMITING,
persistent_home_dir=sandbox_storage_path,
)

Expand Down Expand Up @@ -1337,6 +1343,12 @@ def _print_output(obj):
parser.add_argument(
"--debug", action="store_true", default=False, help="Enable debug mode."
)
parser.add_argument(
"--want_status",
type=str,
default="",
help="If set, verify that the code evaluation status matches this or exit with error code.",
)
args = parser.parse_args()

if args.debug:
Expand All @@ -1357,7 +1369,8 @@ def _print_output(obj):

async def _local_run():
def _dummy_emitter(event):
print(f"Event: {event}", file=sys.stderr)
if not args.want_status:
print(f"Event: {event}", file=sys.stderr)

action = Action()
body = {
Expand All @@ -1368,7 +1381,19 @@ def _dummy_emitter(event):
},
],
}
output = await action.action(body=body, __event_emitter__=_dummy_emitter)
print(output)
output_str = await action.action(body=body, __event_emitter__=_dummy_emitter)
if args.want_status:
output = json.loads(output_str)
got_status = output["status"]
if got_status != args.want_status:
raise RuntimeError(
f"Code evaluation status is {got_status} but expected {args.want_status}"
)
print(
f"\u2705 Code evaluation status is {got_status} as expected.",
file=sys.stderr,
)
else:
print(output_str)

asyncio.run(_local_run())
33 changes: 29 additions & 4 deletions src/openwebui/tools/run_code.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,10 @@ class Valves(pydantic.BaseModel):
default=128,
description=f"Maximum number of megabytes that the interpreter has when running. Must run as root with host cgroups writable (`--mount=type=bind,source=/sys/fs/cgroup,target=/sys/fs/cgroup,readonly=false`) for this to work. Set to 0 to disable memory limits. May be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}MAX_RAM_MEGABYTES",
)
REQUIRE_RESOURCE_LIMITING: bool = pydantic.Field(
default=True,
description=f"Whether to enforce resource limiting, which requires cgroups v2 to be available; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}REQUIRE_RESOURCE_LIMITING.",
)
AUTO_INSTALL: bool = pydantic.Field(
default=True,
description=f"Whether to automatically install gVisor if not installed on the system; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}AUTO_INSTALL. Use the 'HTTPS_PROXY' environment variable to control the proxy used for download.",
Expand Down Expand Up @@ -178,6 +182,7 @@ async def _fail(error_message, status="SANDBOX_ERROR"):
Sandbox.check_setup(
language=language,
auto_install_allowed=valves.AUTO_INSTALL,
require_resource_limiting=valves.REQUIRE_RESOURCE_LIMITING,
)

if valves.AUTO_INSTALL and Sandbox.runsc_needs_installation():
Expand Down Expand Up @@ -211,6 +216,7 @@ async def _fail(error_message, status="SANDBOX_ERROR"):
networking_allowed=valves.NETWORKING_ALLOWED,
max_runtime_seconds=valves.MAX_RUNTIME_SECONDS,
max_ram_bytes=max_ram_bytes,
require_resource_limiting=valves.REQUIRE_RESOURCE_LIMITING,
persistent_home_dir=None,
)

Expand Down Expand Up @@ -525,6 +531,12 @@ def _print_output(obj):
parser.add_argument(
"--debug", action="store_true", default=False, help="Enable debug mode."
)
parser.add_argument(
"--want_status",
type=str,
default="",
help="If set, verify that the code evaluation status matches this or exit with error code.",
)
args = parser.parse_args()

if args.debug:
Expand All @@ -545,17 +557,30 @@ def _print_output(obj):

async def _local_run():
def _dummy_emitter(event):
print(f"Event: {event}", file=sys.stderr)
if not args.want_status:
print(f"Event: {event}", file=sys.stderr)

tools = Tools()
if args.language == "bash":
output = await tools.run_bash_command(
output_str = await tools.run_bash_command(
bash_command=code, __event_emitter__=_dummy_emitter
)
else:
output = await tools.run_python_code(
output_str = await tools.run_python_code(
python_code=code, __event_emitter__=_dummy_emitter
)
print(output)
if args.want_status:
output = json.loads(output_str)
got_status = output["status"]
if got_status != args.want_status:
raise RuntimeError(
f"Code evaluation status is {got_status} but expected {args.want_status}"
)
print(
f"\u2705 Code evaluation status is {got_status} as expected.",
file=sys.stderr,
)
else:
print(output_str)

asyncio.run(_local_run())
Loading

0 comments on commit 8585d81

Please sign in to comment.