diff --git a/docs/setup.md b/docs/setup.md index fb3ecb0..0dbea2e 100644 --- a/docs/setup.md +++ b/docs/setup.md @@ -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`. @@ -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. @@ -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. diff --git a/open-webui/functions/run_code.py b/open-webui/functions/run_code.py index d7fe2a8..24ad0ec 100644 --- a/open-webui/functions/run_code.py +++ b/open-webui/functions/run_code.py @@ -110,6 +110,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.", @@ -271,6 +275,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(): @@ -306,6 +311,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, ) @@ -1252,10 +1258,11 @@ class _Switcheroo: _CGROUP_SUPERVISOR_NAME = "supervisor" _CGROUP_LEAF = "leaf" - def __init__(self, libc, log_path, max_sandbox_ram_bytes): + def __init__(self, libc, log_path, max_sandbox_ram_bytes, do_resource_limiting): self._libc = libc self._log_path = log_path self._max_sandbox_ram_bytes = max_sandbox_ram_bytes + self._do_resource_limiting = do_resource_limiting self._my_euid = None self._my_egid = None self._checkpoint = None @@ -1266,7 +1273,7 @@ def __init__(self, libc, log_path, max_sandbox_ram_bytes): self._initial_cgroup_name = None self._codeeval_cgroup_name = None self._moved = False - self._operations = ( + self._operations = [ # Save EUID and EGID before we move to a new user namespace. ("save_euid", self._save_euid), ("save_egid", self._save_egid), @@ -1275,124 +1282,144 @@ def __init__(self, libc, log_path, max_sandbox_ram_bytes): ("write_uid_map", self._write_uid_map), ("write_setgroups", self._write_setgroups), ("write_gid_map", self._write_gid_map), - # cgroupfs's view does not take into account cgroup namespaces. - # Weird, right? - # This means `/proc/PID/cgroup` will show the namespaced view of - # the cgroup that the PID is in, but `/sys/fs/cgroup` will still - # contain the whole system cgroup hierarchy regardless of namespace. - # Instead, namespaces act as "boundary box" around process movement - # requests when writing to cgroup.procs or creating new cgroups. - # So our first order of business here is to find out which cgroup we - # are running in. We do this by scanning the whole cgroupfs hierarchy - # and looking for our PID. This will populate - # `self._initial_cgroup_name`. - ("find_self_in_cgroup_hierarchy", self._find_self_in_cgroup_hierarchy), - # The cgroup nesting rules are complicated, but the short of it is: - # A cgroup can either **contain processes** OR **have limits**. - # Also, cgroups that contain processes must be leaf nodes. - # Also, cgroups that enforce limits must have their parent cgroup - # also have the same limit "controller" be active. - # So we will have two types of cgroups: - # - Leaf nodes with no controllers - # - Non-leaf nodes with controllers - # So initially, all the processes in the container's initial - # namespace need to be moved out to a new leaf node, - # otherwise we cannot turn on controllers on the initial - # cgroup. - # So we will set up the following hierarchy: - # /sys/fs/cgroup/$INITIAL: - # The cgroup where the container's processes were running - # the first time we run any Sandbox in the container. - # It may initially have no controllers enabled, but we will - # turn them on later. - # /sys/fs/cgroup/$INITIAL/leaf: - # The cgroup where the container's processes are moved to - # from the $INITIAL cgroup upon first run of any Sandbox in - # this container. When this code runs again, processes that - # are already in `$INITIAL/leaf` are not moved. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM: - # A per-Sandbox cgroup that never contains any processes. - # It will have controllers enabled on it but will never have - # specific limits enforced. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox: - # A per-Sandbox cgroup that never contains any processes. - # It will have controllers enabled on it and will enforce - # resource limits for the processes running in its /leaf. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox/leaf: - # A per-Sandbox cgroup that is running `runsc` (gVisor). - # It has no controllers enabled on it, but resources are - # being enforced by virtue of being a child of - # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor: - # A per-Sandbox cgroup that never contains any processes. - # It will have controllers enabled on it and will enforce - # resource limits for the processes running in its /leaf. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor/leaf: - # A per-Sandbox cgroup that is running a Python interpreter - # that manages the lifetime of the `runsc` process. - # It will run `Sandbox.maybe_main`. - # It has no controllers enabled on it, but resources are - # being enforced by virtue of being a child of - # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. - # - # This particular step creates the `$INITIAL/leaf` cgroup. - # If already created, it does nothing. - ("create_initial_leaf_cgroup", self._create_initial_leaf_cgroup), - # Move all processes in `$INITIAL` to `$INITIAL/leaf`. - ( - "move_initial_cgroup_processes_to_initial_leaf_cgroup", - self._move_initial_cgroup_processes_to_initial_leaf_cgroup, - ), - # Read the cgroup controllers enabled in `$INITIAL`. This acts - # as a bounding set on the ones we can enable in any child of it. - ("read_cgroup_controllers", self._read_cgroup_controllers), - # Cleanup old `$INITIAL/codeeval_*` cgroups that may be lying - # around from past runs. - ("cleanup_old_cgroups", self._cleanup_old_cgroups), - # Create a new `$INITIAL/codeeval_$NUM` cgroup. - ("create_codeeval_cgroup", self._create_codeeval_cgroup), - # Create a new `$INITIAL/codeeval_$NUM/sandbox` cgroup. - ("create_sandbox_cgroup", self._create_sandbox_cgroup), - # Create a new `$INITIAL/codeeval_$NUM/sandbox/leaf` cgroup. - ("create_sandbox_leaf_cgroup", self._create_sandbox_leaf_cgroup), - # Create a new `$INITIAL/codeeval_$NUM/supervisor` cgroup. - ("create_supervisor_cgroup", self._create_supervisor_cgroup), - # Create a new `$INITIAL/codeeval_$NUM/supervisor/leaf` cgroup. - ("create_supervisor_leaf_cgroup", self._create_supervisor_leaf_cgroup), - # Add controllers to `$INITIAL`. - ( - "add_cgroup_controllers_to_root", - self._add_cgroup_controllers_to_root, - ), - # Add controllers to `$INITIAL/codeeval_$NUM`. - ( - "add_cgroup_controllers_to_codeeval", - self._add_cgroup_controllers_to_codeeval, - ), - # Add controllers to `$INITIAL/codeeval_$NUM/sandbox`. - ( - "add_cgroup_controllers_to_sandbox", - self._add_cgroup_controllers_to_sandbox, - ), - # Set resource limits on `$INITIAL/codeeval_$NUM`. - ("set_sandbox_cgroup_limits", self._set_sandbox_cgroup_limits), - # Add controllers to `$INITIAL/codeeval_$NUM/supervisor`. - ( - "add_cgroup_controllers_to_supervisor", - self._add_cgroup_controllers_to_supervisor, - ), - # Set resource limits on `$INITIAL/codeeval_$NUM/supervisor`. - ("set_supervisor_cgroup_limits", self._set_supervisor_cgroup_limits), - # Move current process to - # `$INITIAL/codeeval_$NUM/supervisor/leaf`. - ( - "move_process_to_supervisor_leaf", - self._move_process_to_supervisor_leaf, - ), - # Double-check that we have moved to - # `$INITIAL/codeeval_$NUM/supervisor/leaf`. - ("sanity_check_own_cgroup", self._sanity_check_own_cgroup), - ) + ] + if do_resource_limiting: + self._operations.extend( + ( + # cgroupfs's view does not take into account cgroup namespaces. + # Weird, right? + # This means `/proc/PID/cgroup` will show the namespaced view of + # the cgroup that the PID is in, but `/sys/fs/cgroup` will still + # contain the whole system cgroup hierarchy regardless of namespace. + # Instead, namespaces act as "boundary box" around process movement + # requests when writing to cgroup.procs or creating new cgroups. + # So our first order of business here is to find out which cgroup we + # are running in. We do this by scanning the whole cgroupfs hierarchy + # and looking for our PID. This will populate + # `self._initial_cgroup_name`. + ( + "find_self_in_cgroup_hierarchy", + self._find_self_in_cgroup_hierarchy, + ), + # The cgroup nesting rules are complicated, but the short of it is: + # A cgroup can either **contain processes** OR **have limits**. + # Also, cgroups that contain processes must be leaf nodes. + # Also, cgroups that enforce limits must have their parent cgroup + # also have the same limit "controller" be active. + # So we will have two types of cgroups: + # - Leaf nodes with no controllers + # - Non-leaf nodes with controllers + # So initially, all the processes in the container's initial + # namespace need to be moved out to a new leaf node, + # otherwise we cannot turn on controllers on the initial + # cgroup. + # So we will set up the following hierarchy: + # /sys/fs/cgroup/$INITIAL: + # The cgroup where the container's processes were running + # the first time we run any Sandbox in the container. + # It may initially have no controllers enabled, but we will + # turn them on later. + # /sys/fs/cgroup/$INITIAL/leaf: + # The cgroup where the container's processes are moved to + # from the $INITIAL cgroup upon first run of any Sandbox in + # this container. When this code runs again, processes that + # are already in `$INITIAL/leaf` are not moved. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM: + # A per-Sandbox cgroup that never contains any processes. + # It will have controllers enabled on it but will never have + # specific limits enforced. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox: + # A per-Sandbox cgroup that never contains any processes. + # It will have controllers enabled on it and will enforce + # resource limits for the processes running in its /leaf. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox/leaf: + # A per-Sandbox cgroup that is running `runsc` (gVisor). + # It has no controllers enabled on it, but resources are + # being enforced by virtue of being a child of + # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor: + # A per-Sandbox cgroup that never contains any processes. + # It will have controllers enabled on it and will enforce + # resource limits for the processes running in its /leaf. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor/leaf: + # A per-Sandbox cgroup that is running a Python interpreter + # that manages the lifetime of the `runsc` process. + # It will run `Sandbox.maybe_main`. + # It has no controllers enabled on it, but resources are + # being enforced by virtue of being a child of + # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. + # + # This particular step creates the `$INITIAL/leaf` cgroup. + # If already created, it does nothing. + ( + "create_initial_leaf_cgroup", + self._create_initial_leaf_cgroup, + ), + # Move all processes in `$INITIAL` to `$INITIAL/leaf`. + ( + "move_initial_cgroup_processes_to_initial_leaf_cgroup", + self._move_initial_cgroup_processes_to_initial_leaf_cgroup, + ), + # Read the cgroup controllers enabled in `$INITIAL`. This acts + # as a bounding set on the ones we can enable in any child of it. + ("read_cgroup_controllers", self._read_cgroup_controllers), + # Cleanup old `$INITIAL/codeeval_*` cgroups that may be lying + # around from past runs. + ("cleanup_old_cgroups", self._cleanup_old_cgroups), + # Create a new `$INITIAL/codeeval_$NUM` cgroup. + ("create_codeeval_cgroup", self._create_codeeval_cgroup), + # Create a new `$INITIAL/codeeval_$NUM/sandbox` cgroup. + ("create_sandbox_cgroup", self._create_sandbox_cgroup), + # Create a new `$INITIAL/codeeval_$NUM/sandbox/leaf` cgroup. + ( + "create_sandbox_leaf_cgroup", + self._create_sandbox_leaf_cgroup, + ), + # Create a new `$INITIAL/codeeval_$NUM/supervisor` cgroup. + ("create_supervisor_cgroup", self._create_supervisor_cgroup), + # Create a new `$INITIAL/codeeval_$NUM/supervisor/leaf` cgroup. + ( + "create_supervisor_leaf_cgroup", + self._create_supervisor_leaf_cgroup, + ), + # Add controllers to `$INITIAL`. + ( + "add_cgroup_controllers_to_root", + self._add_cgroup_controllers_to_root, + ), + # Add controllers to `$INITIAL/codeeval_$NUM`. + ( + "add_cgroup_controllers_to_codeeval", + self._add_cgroup_controllers_to_codeeval, + ), + # Add controllers to `$INITIAL/codeeval_$NUM/sandbox`. + ( + "add_cgroup_controllers_to_sandbox", + self._add_cgroup_controllers_to_sandbox, + ), + # Set resource limits on `$INITIAL/codeeval_$NUM`. + ("set_sandbox_cgroup_limits", self._set_sandbox_cgroup_limits), + # Add controllers to `$INITIAL/codeeval_$NUM/supervisor`. + ( + "add_cgroup_controllers_to_supervisor", + self._add_cgroup_controllers_to_supervisor, + ), + # Set resource limits on `$INITIAL/codeeval_$NUM/supervisor`. + ( + "set_supervisor_cgroup_limits", + self._set_supervisor_cgroup_limits, + ), + # Move current process to + # `$INITIAL/codeeval_$NUM/supervisor/leaf`. + ( + "move_process_to_supervisor_leaf", + self._move_process_to_supervisor_leaf, + ), + # Double-check that we have moved to + # `$INITIAL/codeeval_$NUM/supervisor/leaf`. + ("sanity_check_own_cgroup", self._sanity_check_own_cgroup), + ) + ) def _status(self): """ @@ -1406,62 +1433,65 @@ def _status(self): if self._checkpoint == self._operations[-1][0]: main_status = "OK" my_pid = os.getpid() - status_line = f"{main_status} (euid={self._my_euid} egid={self._my_egid} pid={my_pid} initial_cgroup_name={self._initial_cgroup_name} codeeval_cgroup_name={self._codeeval_cgroup_name} controllers={self._cgroup_controllers})" - cgroupfs_data = [] - for cgroup_components in ( - (self._initial_cgroup_name,), - (self._initial_cgroup_name, self._CGROUP_LEAF), - (self._initial_cgroup_name, self._codeeval_cgroup_name), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_LEAF, - ), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_SUPERVISOR_NAME, - ), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_SUPERVISOR_NAME, - self._CGROUP_LEAF, - ), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_SANDBOX_NAME, - ), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_SANDBOX_NAME, - self._CGROUP_LEAF, - ), - ): - if any(c is None for c in cgroup_components): - continue - file_data = [] - for filename in ("procs", "controllers", "subtree_control"): - data = None - try: - with self._open( - self._cgroup_path( - *(cgroup_components + (f"cgroup.{filename}",)) - ), - "rb", - ) as f: - data = f.read().decode("ascii").replace("\n", " ") - except Exception as e: - data = f"[fail: {e}]" - file_data.append(f"{filename}: {data}") - cgroup_components_joined = os.sep.join(cgroup_components) - file_data_joined = ", ".join(file_data) - cgroupfs_data.append(f"{cgroup_components_joined}=[{file_data_joined}]") - if len(cgroupfs_data) > 0: - cgroupfs_data_joined = " ".join(cgroupfs_data) - status_line += f" {cgroupfs_data_joined}" + status_line = f"{main_status} (euid={self._my_euid} egid={self._my_egid} pid={my_pid} do_resource_limiting={self._do_resource_limiting} initial_cgroup_name={self._initial_cgroup_name} codeeval_cgroup_name={self._codeeval_cgroup_name} controllers={self._cgroup_controllers})" + if self._do_resource_limiting: + cgroupfs_data = [] + for cgroup_components in ( + (self._initial_cgroup_name,), + (self._initial_cgroup_name, self._CGROUP_LEAF), + (self._initial_cgroup_name, self._codeeval_cgroup_name), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_LEAF, + ), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_SUPERVISOR_NAME, + ), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_SUPERVISOR_NAME, + self._CGROUP_LEAF, + ), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_SANDBOX_NAME, + ), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_SANDBOX_NAME, + self._CGROUP_LEAF, + ), + ): + if any(c is None for c in cgroup_components): + continue + file_data = [] + for filename in ("procs", "controllers", "subtree_control"): + data = None + try: + with self._open( + self._cgroup_path( + *(cgroup_components + (f"cgroup.{filename}",)) + ), + "rb", + ) as f: + data = f.read().decode("ascii").replace("\n", " ") + except Exception as e: + data = f"[fail: {e}]" + file_data.append(f"{filename}: {data}") + cgroup_components_joined = os.sep.join(cgroup_components) + file_data_joined = ", ".join(file_data) + cgroupfs_data.append( + f"{cgroup_components_joined}=[{file_data_joined}]" + ) + if len(cgroupfs_data) > 0: + cgroupfs_data_joined = " ".join(cgroupfs_data) + status_line += f" {cgroupfs_data_joined}" return status_line def _cgroup_path(self, *components): @@ -1926,6 +1956,8 @@ def move_process_to_sandbox_leaf_cgroup_lambda(self): :return: A function to move the current process to the sandbox cgroup. :raises SandboxException: If not queried after we have already chosen a new cgroup name. """ + if not self._do_resource_limiting: + return lambda: None if self._codeeval_cgroup_name is None: raise Sandbox.SandboxException( "Tried to move process to sandbox leaf cgroup before we know it" @@ -1976,13 +2008,15 @@ def _move(cgroup_path): def monitor_cgroup_resources(self): """ - Spawns a background thread that monitors resources. + Spawns a background thread that monitors resources, if limiting is enabled. cgroups should be taking care of this, but some systems do not enforce this. So this does the same in userspace. Better than nothing. - :return: A function to cancel the monitor thread. + :return: A function to cancel the monitor thread, if resource limiting is enabled. """ + if not self._do_resource_limiting: + return lambda: None self_memory_path = self._cgroup_path( self._initial_cgroup_name, self._codeeval_cgroup_name, @@ -2211,20 +2245,6 @@ def check_cgroups(cls): "cgroupfs does not have the 'memory' controller enabled, necessary to enforce memory limits; please enable it" ) - @classmethod - def cgroups_available(cls) -> bool: - """ - Returns whether cgroupfs is mounted and usable for resource limit enforcement. - - :return: Whether cgroupfs is mounted and usable for resource limit enforcement. - """ - try: - cls.check_cgroups() - except Exception: - return False - else: - return True - @classmethod def check_procfs(cls): """ @@ -2368,12 +2388,18 @@ def install_runsc(cls): os.rename(download_path, cls.AUTO_INSTALLATION_PATH) @classmethod - def check_setup(cls, language: str, auto_install_allowed: bool): + def check_setup( + cls, + language: str, + auto_install_allowed: bool, + require_resource_limiting: bool, + ): """ Verifies that the environment is compatible with running sandboxes. :param language: The programming language to run. :param auto_install_allowed: Whether auto-installation of `runsc` is allowed. + :param require_resource_limiting: Check that the host supports resource limiting via cgroups. :return: Nothing. :raises ValueError: If provided an invalid language name. @@ -2392,7 +2418,8 @@ def check_setup(cls, language: str, auto_install_allowed: bool): ) cls.check_platform() cls.check_unshare() - cls.check_cgroups() + if require_resource_limiting: + cls.check_cgroups() cls.check_procfs() if not auto_install_allowed and cls.get_runsc_path() is None: raise cls.GVisorNotInstalledException( @@ -2462,8 +2489,9 @@ def __init__( debug: bool, networking_allowed: bool, max_runtime_seconds: int, - max_ram_bytes: typing.Optional[int], - persistent_home_dir: typing.Optional[str], + max_ram_bytes: typing.Optional[int] = None, + require_resource_limiting: bool = False, + persistent_home_dir: typing.Optional[str] = None, ): """ Constructor. @@ -2475,6 +2503,7 @@ def __init__( :param networking_allowed: Whether the code should be given access to the network. :param max_runtime_seconds: How long the code should be allowed to run, in seconds. :param max_ram_bytes: How many bytes of RAM the interpreter should be allowed to use, or `None` for no limit. + :param require_resource_limiting: If true, refuse to launch a sandbox if the host doesn't support resource limiting via cgroups. :param persistent_home_dir: Optional directory which will be mapped read-write to this real host directory. """ self._init( @@ -2486,6 +2515,7 @@ def __init__( "networking_allowed": networking_allowed, "max_runtime_seconds": max_runtime_seconds, "max_ram_bytes": max_ram_bytes, + "require_resource_limiting": require_resource_limiting, "persistent_home_dir": persistent_home_dir, } ) @@ -2504,13 +2534,12 @@ def _init(self, settings): self._networking_allowed = self._settings["networking_allowed"] self._max_runtime_seconds = self._settings["max_runtime_seconds"] self._max_ram_bytes = self._settings["max_ram_bytes"] + self._require_resource_limiting = self._settings[ + "require_resource_limiting" + ] or all((self._max_ram_bytes is None,)) self._persistent_home_dir = self._settings["persistent_home_dir"] self._sandboxed_command = None - self._switcheroo = self._Switcheroo( - libc=self._libc(), - log_path=os.path.join(self._logs_path, "switcheroo.txt"), - max_sandbox_ram_bytes=self._max_ram_bytes, - ) + self._switcheroo = None def _setup_sandbox(self): """ @@ -2537,7 +2566,18 @@ def _setup_sandbox(self): f"Persistent home directory {self._persistent_home_dir} does not exist" ) oci_config["root"]["path"] = rootfs_path - + do_resource_limiting = True + if not self._require_resource_limiting: + try: + self.check_cgroups() + except self.EnvironmentNeedsSetupException: + do_resource_limiting = False + self._switcheroo = self._Switcheroo( + libc=self._libc(), + log_path=os.path.join(self._logs_path, "switcheroo.txt"), + max_sandbox_ram_bytes=self._max_ram_bytes, + do_resource_limiting=do_resource_limiting, + ) try: self._switcheroo.do() except Exception as e: @@ -3531,6 +3571,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: @@ -3551,7 +3597,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 = { @@ -3562,7 +3609,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()) diff --git a/open-webui/tools/run_code.py b/open-webui/tools/run_code.py index cac9350..3257f27 100644 --- a/open-webui/tools/run_code.py +++ b/open-webui/tools/run_code.py @@ -79,6 +79,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.", @@ -225,6 +229,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(): @@ -258,6 +263,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, ) @@ -700,10 +706,11 @@ class _Switcheroo: _CGROUP_SUPERVISOR_NAME = "supervisor" _CGROUP_LEAF = "leaf" - def __init__(self, libc, log_path, max_sandbox_ram_bytes): + def __init__(self, libc, log_path, max_sandbox_ram_bytes, do_resource_limiting): self._libc = libc self._log_path = log_path self._max_sandbox_ram_bytes = max_sandbox_ram_bytes + self._do_resource_limiting = do_resource_limiting self._my_euid = None self._my_egid = None self._checkpoint = None @@ -714,7 +721,7 @@ def __init__(self, libc, log_path, max_sandbox_ram_bytes): self._initial_cgroup_name = None self._codeeval_cgroup_name = None self._moved = False - self._operations = ( + self._operations = [ # Save EUID and EGID before we move to a new user namespace. ("save_euid", self._save_euid), ("save_egid", self._save_egid), @@ -723,124 +730,144 @@ def __init__(self, libc, log_path, max_sandbox_ram_bytes): ("write_uid_map", self._write_uid_map), ("write_setgroups", self._write_setgroups), ("write_gid_map", self._write_gid_map), - # cgroupfs's view does not take into account cgroup namespaces. - # Weird, right? - # This means `/proc/PID/cgroup` will show the namespaced view of - # the cgroup that the PID is in, but `/sys/fs/cgroup` will still - # contain the whole system cgroup hierarchy regardless of namespace. - # Instead, namespaces act as "boundary box" around process movement - # requests when writing to cgroup.procs or creating new cgroups. - # So our first order of business here is to find out which cgroup we - # are running in. We do this by scanning the whole cgroupfs hierarchy - # and looking for our PID. This will populate - # `self._initial_cgroup_name`. - ("find_self_in_cgroup_hierarchy", self._find_self_in_cgroup_hierarchy), - # The cgroup nesting rules are complicated, but the short of it is: - # A cgroup can either **contain processes** OR **have limits**. - # Also, cgroups that contain processes must be leaf nodes. - # Also, cgroups that enforce limits must have their parent cgroup - # also have the same limit "controller" be active. - # So we will have two types of cgroups: - # - Leaf nodes with no controllers - # - Non-leaf nodes with controllers - # So initially, all the processes in the container's initial - # namespace need to be moved out to a new leaf node, - # otherwise we cannot turn on controllers on the initial - # cgroup. - # So we will set up the following hierarchy: - # /sys/fs/cgroup/$INITIAL: - # The cgroup where the container's processes were running - # the first time we run any Sandbox in the container. - # It may initially have no controllers enabled, but we will - # turn them on later. - # /sys/fs/cgroup/$INITIAL/leaf: - # The cgroup where the container's processes are moved to - # from the $INITIAL cgroup upon first run of any Sandbox in - # this container. When this code runs again, processes that - # are already in `$INITIAL/leaf` are not moved. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM: - # A per-Sandbox cgroup that never contains any processes. - # It will have controllers enabled on it but will never have - # specific limits enforced. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox: - # A per-Sandbox cgroup that never contains any processes. - # It will have controllers enabled on it and will enforce - # resource limits for the processes running in its /leaf. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox/leaf: - # A per-Sandbox cgroup that is running `runsc` (gVisor). - # It has no controllers enabled on it, but resources are - # being enforced by virtue of being a child of - # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor: - # A per-Sandbox cgroup that never contains any processes. - # It will have controllers enabled on it and will enforce - # resource limits for the processes running in its /leaf. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor/leaf: - # A per-Sandbox cgroup that is running a Python interpreter - # that manages the lifetime of the `runsc` process. - # It will run `Sandbox.maybe_main`. - # It has no controllers enabled on it, but resources are - # being enforced by virtue of being a child of - # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. - # - # This particular step creates the `$INITIAL/leaf` cgroup. - # If already created, it does nothing. - ("create_initial_leaf_cgroup", self._create_initial_leaf_cgroup), - # Move all processes in `$INITIAL` to `$INITIAL/leaf`. - ( - "move_initial_cgroup_processes_to_initial_leaf_cgroup", - self._move_initial_cgroup_processes_to_initial_leaf_cgroup, - ), - # Read the cgroup controllers enabled in `$INITIAL`. This acts - # as a bounding set on the ones we can enable in any child of it. - ("read_cgroup_controllers", self._read_cgroup_controllers), - # Cleanup old `$INITIAL/codeeval_*` cgroups that may be lying - # around from past runs. - ("cleanup_old_cgroups", self._cleanup_old_cgroups), - # Create a new `$INITIAL/codeeval_$NUM` cgroup. - ("create_codeeval_cgroup", self._create_codeeval_cgroup), - # Create a new `$INITIAL/codeeval_$NUM/sandbox` cgroup. - ("create_sandbox_cgroup", self._create_sandbox_cgroup), - # Create a new `$INITIAL/codeeval_$NUM/sandbox/leaf` cgroup. - ("create_sandbox_leaf_cgroup", self._create_sandbox_leaf_cgroup), - # Create a new `$INITIAL/codeeval_$NUM/supervisor` cgroup. - ("create_supervisor_cgroup", self._create_supervisor_cgroup), - # Create a new `$INITIAL/codeeval_$NUM/supervisor/leaf` cgroup. - ("create_supervisor_leaf_cgroup", self._create_supervisor_leaf_cgroup), - # Add controllers to `$INITIAL`. - ( - "add_cgroup_controllers_to_root", - self._add_cgroup_controllers_to_root, - ), - # Add controllers to `$INITIAL/codeeval_$NUM`. - ( - "add_cgroup_controllers_to_codeeval", - self._add_cgroup_controllers_to_codeeval, - ), - # Add controllers to `$INITIAL/codeeval_$NUM/sandbox`. - ( - "add_cgroup_controllers_to_sandbox", - self._add_cgroup_controllers_to_sandbox, - ), - # Set resource limits on `$INITIAL/codeeval_$NUM`. - ("set_sandbox_cgroup_limits", self._set_sandbox_cgroup_limits), - # Add controllers to `$INITIAL/codeeval_$NUM/supervisor`. - ( - "add_cgroup_controllers_to_supervisor", - self._add_cgroup_controllers_to_supervisor, - ), - # Set resource limits on `$INITIAL/codeeval_$NUM/supervisor`. - ("set_supervisor_cgroup_limits", self._set_supervisor_cgroup_limits), - # Move current process to - # `$INITIAL/codeeval_$NUM/supervisor/leaf`. - ( - "move_process_to_supervisor_leaf", - self._move_process_to_supervisor_leaf, - ), - # Double-check that we have moved to - # `$INITIAL/codeeval_$NUM/supervisor/leaf`. - ("sanity_check_own_cgroup", self._sanity_check_own_cgroup), - ) + ] + if do_resource_limiting: + self._operations.extend( + ( + # cgroupfs's view does not take into account cgroup namespaces. + # Weird, right? + # This means `/proc/PID/cgroup` will show the namespaced view of + # the cgroup that the PID is in, but `/sys/fs/cgroup` will still + # contain the whole system cgroup hierarchy regardless of namespace. + # Instead, namespaces act as "boundary box" around process movement + # requests when writing to cgroup.procs or creating new cgroups. + # So our first order of business here is to find out which cgroup we + # are running in. We do this by scanning the whole cgroupfs hierarchy + # and looking for our PID. This will populate + # `self._initial_cgroup_name`. + ( + "find_self_in_cgroup_hierarchy", + self._find_self_in_cgroup_hierarchy, + ), + # The cgroup nesting rules are complicated, but the short of it is: + # A cgroup can either **contain processes** OR **have limits**. + # Also, cgroups that contain processes must be leaf nodes. + # Also, cgroups that enforce limits must have their parent cgroup + # also have the same limit "controller" be active. + # So we will have two types of cgroups: + # - Leaf nodes with no controllers + # - Non-leaf nodes with controllers + # So initially, all the processes in the container's initial + # namespace need to be moved out to a new leaf node, + # otherwise we cannot turn on controllers on the initial + # cgroup. + # So we will set up the following hierarchy: + # /sys/fs/cgroup/$INITIAL: + # The cgroup where the container's processes were running + # the first time we run any Sandbox in the container. + # It may initially have no controllers enabled, but we will + # turn them on later. + # /sys/fs/cgroup/$INITIAL/leaf: + # The cgroup where the container's processes are moved to + # from the $INITIAL cgroup upon first run of any Sandbox in + # this container. When this code runs again, processes that + # are already in `$INITIAL/leaf` are not moved. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM: + # A per-Sandbox cgroup that never contains any processes. + # It will have controllers enabled on it but will never have + # specific limits enforced. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox: + # A per-Sandbox cgroup that never contains any processes. + # It will have controllers enabled on it and will enforce + # resource limits for the processes running in its /leaf. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox/leaf: + # A per-Sandbox cgroup that is running `runsc` (gVisor). + # It has no controllers enabled on it, but resources are + # being enforced by virtue of being a child of + # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor: + # A per-Sandbox cgroup that never contains any processes. + # It will have controllers enabled on it and will enforce + # resource limits for the processes running in its /leaf. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor/leaf: + # A per-Sandbox cgroup that is running a Python interpreter + # that manages the lifetime of the `runsc` process. + # It will run `Sandbox.maybe_main`. + # It has no controllers enabled on it, but resources are + # being enforced by virtue of being a child of + # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. + # + # This particular step creates the `$INITIAL/leaf` cgroup. + # If already created, it does nothing. + ( + "create_initial_leaf_cgroup", + self._create_initial_leaf_cgroup, + ), + # Move all processes in `$INITIAL` to `$INITIAL/leaf`. + ( + "move_initial_cgroup_processes_to_initial_leaf_cgroup", + self._move_initial_cgroup_processes_to_initial_leaf_cgroup, + ), + # Read the cgroup controllers enabled in `$INITIAL`. This acts + # as a bounding set on the ones we can enable in any child of it. + ("read_cgroup_controllers", self._read_cgroup_controllers), + # Cleanup old `$INITIAL/codeeval_*` cgroups that may be lying + # around from past runs. + ("cleanup_old_cgroups", self._cleanup_old_cgroups), + # Create a new `$INITIAL/codeeval_$NUM` cgroup. + ("create_codeeval_cgroup", self._create_codeeval_cgroup), + # Create a new `$INITIAL/codeeval_$NUM/sandbox` cgroup. + ("create_sandbox_cgroup", self._create_sandbox_cgroup), + # Create a new `$INITIAL/codeeval_$NUM/sandbox/leaf` cgroup. + ( + "create_sandbox_leaf_cgroup", + self._create_sandbox_leaf_cgroup, + ), + # Create a new `$INITIAL/codeeval_$NUM/supervisor` cgroup. + ("create_supervisor_cgroup", self._create_supervisor_cgroup), + # Create a new `$INITIAL/codeeval_$NUM/supervisor/leaf` cgroup. + ( + "create_supervisor_leaf_cgroup", + self._create_supervisor_leaf_cgroup, + ), + # Add controllers to `$INITIAL`. + ( + "add_cgroup_controllers_to_root", + self._add_cgroup_controllers_to_root, + ), + # Add controllers to `$INITIAL/codeeval_$NUM`. + ( + "add_cgroup_controllers_to_codeeval", + self._add_cgroup_controllers_to_codeeval, + ), + # Add controllers to `$INITIAL/codeeval_$NUM/sandbox`. + ( + "add_cgroup_controllers_to_sandbox", + self._add_cgroup_controllers_to_sandbox, + ), + # Set resource limits on `$INITIAL/codeeval_$NUM`. + ("set_sandbox_cgroup_limits", self._set_sandbox_cgroup_limits), + # Add controllers to `$INITIAL/codeeval_$NUM/supervisor`. + ( + "add_cgroup_controllers_to_supervisor", + self._add_cgroup_controllers_to_supervisor, + ), + # Set resource limits on `$INITIAL/codeeval_$NUM/supervisor`. + ( + "set_supervisor_cgroup_limits", + self._set_supervisor_cgroup_limits, + ), + # Move current process to + # `$INITIAL/codeeval_$NUM/supervisor/leaf`. + ( + "move_process_to_supervisor_leaf", + self._move_process_to_supervisor_leaf, + ), + # Double-check that we have moved to + # `$INITIAL/codeeval_$NUM/supervisor/leaf`. + ("sanity_check_own_cgroup", self._sanity_check_own_cgroup), + ) + ) def _status(self): """ @@ -854,62 +881,65 @@ def _status(self): if self._checkpoint == self._operations[-1][0]: main_status = "OK" my_pid = os.getpid() - status_line = f"{main_status} (euid={self._my_euid} egid={self._my_egid} pid={my_pid} initial_cgroup_name={self._initial_cgroup_name} codeeval_cgroup_name={self._codeeval_cgroup_name} controllers={self._cgroup_controllers})" - cgroupfs_data = [] - for cgroup_components in ( - (self._initial_cgroup_name,), - (self._initial_cgroup_name, self._CGROUP_LEAF), - (self._initial_cgroup_name, self._codeeval_cgroup_name), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_LEAF, - ), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_SUPERVISOR_NAME, - ), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_SUPERVISOR_NAME, - self._CGROUP_LEAF, - ), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_SANDBOX_NAME, - ), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_SANDBOX_NAME, - self._CGROUP_LEAF, - ), - ): - if any(c is None for c in cgroup_components): - continue - file_data = [] - for filename in ("procs", "controllers", "subtree_control"): - data = None - try: - with self._open( - self._cgroup_path( - *(cgroup_components + (f"cgroup.{filename}",)) - ), - "rb", - ) as f: - data = f.read().decode("ascii").replace("\n", " ") - except Exception as e: - data = f"[fail: {e}]" - file_data.append(f"{filename}: {data}") - cgroup_components_joined = os.sep.join(cgroup_components) - file_data_joined = ", ".join(file_data) - cgroupfs_data.append(f"{cgroup_components_joined}=[{file_data_joined}]") - if len(cgroupfs_data) > 0: - cgroupfs_data_joined = " ".join(cgroupfs_data) - status_line += f" {cgroupfs_data_joined}" + status_line = f"{main_status} (euid={self._my_euid} egid={self._my_egid} pid={my_pid} do_resource_limiting={self._do_resource_limiting} initial_cgroup_name={self._initial_cgroup_name} codeeval_cgroup_name={self._codeeval_cgroup_name} controllers={self._cgroup_controllers})" + if self._do_resource_limiting: + cgroupfs_data = [] + for cgroup_components in ( + (self._initial_cgroup_name,), + (self._initial_cgroup_name, self._CGROUP_LEAF), + (self._initial_cgroup_name, self._codeeval_cgroup_name), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_LEAF, + ), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_SUPERVISOR_NAME, + ), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_SUPERVISOR_NAME, + self._CGROUP_LEAF, + ), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_SANDBOX_NAME, + ), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_SANDBOX_NAME, + self._CGROUP_LEAF, + ), + ): + if any(c is None for c in cgroup_components): + continue + file_data = [] + for filename in ("procs", "controllers", "subtree_control"): + data = None + try: + with self._open( + self._cgroup_path( + *(cgroup_components + (f"cgroup.{filename}",)) + ), + "rb", + ) as f: + data = f.read().decode("ascii").replace("\n", " ") + except Exception as e: + data = f"[fail: {e}]" + file_data.append(f"{filename}: {data}") + cgroup_components_joined = os.sep.join(cgroup_components) + file_data_joined = ", ".join(file_data) + cgroupfs_data.append( + f"{cgroup_components_joined}=[{file_data_joined}]" + ) + if len(cgroupfs_data) > 0: + cgroupfs_data_joined = " ".join(cgroupfs_data) + status_line += f" {cgroupfs_data_joined}" return status_line def _cgroup_path(self, *components): @@ -1374,6 +1404,8 @@ def move_process_to_sandbox_leaf_cgroup_lambda(self): :return: A function to move the current process to the sandbox cgroup. :raises SandboxException: If not queried after we have already chosen a new cgroup name. """ + if not self._do_resource_limiting: + return lambda: None if self._codeeval_cgroup_name is None: raise Sandbox.SandboxException( "Tried to move process to sandbox leaf cgroup before we know it" @@ -1424,13 +1456,15 @@ def _move(cgroup_path): def monitor_cgroup_resources(self): """ - Spawns a background thread that monitors resources. + Spawns a background thread that monitors resources, if limiting is enabled. cgroups should be taking care of this, but some systems do not enforce this. So this does the same in userspace. Better than nothing. - :return: A function to cancel the monitor thread. + :return: A function to cancel the monitor thread, if resource limiting is enabled. """ + if not self._do_resource_limiting: + return lambda: None self_memory_path = self._cgroup_path( self._initial_cgroup_name, self._codeeval_cgroup_name, @@ -1659,20 +1693,6 @@ def check_cgroups(cls): "cgroupfs does not have the 'memory' controller enabled, necessary to enforce memory limits; please enable it" ) - @classmethod - def cgroups_available(cls) -> bool: - """ - Returns whether cgroupfs is mounted and usable for resource limit enforcement. - - :return: Whether cgroupfs is mounted and usable for resource limit enforcement. - """ - try: - cls.check_cgroups() - except Exception: - return False - else: - return True - @classmethod def check_procfs(cls): """ @@ -1816,12 +1836,18 @@ def install_runsc(cls): os.rename(download_path, cls.AUTO_INSTALLATION_PATH) @classmethod - def check_setup(cls, language: str, auto_install_allowed: bool): + def check_setup( + cls, + language: str, + auto_install_allowed: bool, + require_resource_limiting: bool, + ): """ Verifies that the environment is compatible with running sandboxes. :param language: The programming language to run. :param auto_install_allowed: Whether auto-installation of `runsc` is allowed. + :param require_resource_limiting: Check that the host supports resource limiting via cgroups. :return: Nothing. :raises ValueError: If provided an invalid language name. @@ -1840,7 +1866,8 @@ def check_setup(cls, language: str, auto_install_allowed: bool): ) cls.check_platform() cls.check_unshare() - cls.check_cgroups() + if require_resource_limiting: + cls.check_cgroups() cls.check_procfs() if not auto_install_allowed and cls.get_runsc_path() is None: raise cls.GVisorNotInstalledException( @@ -1910,8 +1937,9 @@ def __init__( debug: bool, networking_allowed: bool, max_runtime_seconds: int, - max_ram_bytes: typing.Optional[int], - persistent_home_dir: typing.Optional[str], + max_ram_bytes: typing.Optional[int] = None, + require_resource_limiting: bool = False, + persistent_home_dir: typing.Optional[str] = None, ): """ Constructor. @@ -1923,6 +1951,7 @@ def __init__( :param networking_allowed: Whether the code should be given access to the network. :param max_runtime_seconds: How long the code should be allowed to run, in seconds. :param max_ram_bytes: How many bytes of RAM the interpreter should be allowed to use, or `None` for no limit. + :param require_resource_limiting: If true, refuse to launch a sandbox if the host doesn't support resource limiting via cgroups. :param persistent_home_dir: Optional directory which will be mapped read-write to this real host directory. """ self._init( @@ -1934,6 +1963,7 @@ def __init__( "networking_allowed": networking_allowed, "max_runtime_seconds": max_runtime_seconds, "max_ram_bytes": max_ram_bytes, + "require_resource_limiting": require_resource_limiting, "persistent_home_dir": persistent_home_dir, } ) @@ -1952,13 +1982,12 @@ def _init(self, settings): self._networking_allowed = self._settings["networking_allowed"] self._max_runtime_seconds = self._settings["max_runtime_seconds"] self._max_ram_bytes = self._settings["max_ram_bytes"] + self._require_resource_limiting = self._settings[ + "require_resource_limiting" + ] or all((self._max_ram_bytes is None,)) self._persistent_home_dir = self._settings["persistent_home_dir"] self._sandboxed_command = None - self._switcheroo = self._Switcheroo( - libc=self._libc(), - log_path=os.path.join(self._logs_path, "switcheroo.txt"), - max_sandbox_ram_bytes=self._max_ram_bytes, - ) + self._switcheroo = None def _setup_sandbox(self): """ @@ -1985,7 +2014,18 @@ def _setup_sandbox(self): f"Persistent home directory {self._persistent_home_dir} does not exist" ) oci_config["root"]["path"] = rootfs_path - + do_resource_limiting = True + if not self._require_resource_limiting: + try: + self.check_cgroups() + except self.EnvironmentNeedsSetupException: + do_resource_limiting = False + self._switcheroo = self._Switcheroo( + libc=self._libc(), + log_path=os.path.join(self._logs_path, "switcheroo.txt"), + max_sandbox_ram_bytes=self._max_ram_bytes, + do_resource_limiting=do_resource_limiting, + ) try: self._switcheroo.do() except Exception as e: @@ -2725,6 +2765,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: @@ -2745,17 +2791,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()) diff --git a/src/openwebui/functions/run_code.py b/src/openwebui/functions/run_code.py index d8cb9e2..7b2851a 100644 --- a/src/openwebui/functions/run_code.py +++ b/src/openwebui/functions/run_code.py @@ -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.", @@ -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(): @@ -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, ) @@ -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: @@ -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 = { @@ -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()) diff --git a/src/openwebui/tools/run_code.py b/src/openwebui/tools/run_code.py index a9cbf22..37e6fd4 100644 --- a/src/openwebui/tools/run_code.py +++ b/src/openwebui/tools/run_code.py @@ -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.", @@ -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(): @@ -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, ) @@ -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: @@ -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()) diff --git a/src/safecode/sandbox.py b/src/safecode/sandbox.py index 4ad5420..3b4c25d 100644 --- a/src/safecode/sandbox.py +++ b/src/safecode/sandbox.py @@ -265,10 +265,11 @@ class _Switcheroo: _CGROUP_SUPERVISOR_NAME = "supervisor" _CGROUP_LEAF = "leaf" - def __init__(self, libc, log_path, max_sandbox_ram_bytes): + def __init__(self, libc, log_path, max_sandbox_ram_bytes, do_resource_limiting): self._libc = libc self._log_path = log_path self._max_sandbox_ram_bytes = max_sandbox_ram_bytes + self._do_resource_limiting = do_resource_limiting self._my_euid = None self._my_egid = None self._checkpoint = None @@ -279,7 +280,7 @@ def __init__(self, libc, log_path, max_sandbox_ram_bytes): self._initial_cgroup_name = None self._codeeval_cgroup_name = None self._moved = False - self._operations = ( + self._operations = [ # Save EUID and EGID before we move to a new user namespace. ("save_euid", self._save_euid), ("save_egid", self._save_egid), @@ -288,124 +289,144 @@ def __init__(self, libc, log_path, max_sandbox_ram_bytes): ("write_uid_map", self._write_uid_map), ("write_setgroups", self._write_setgroups), ("write_gid_map", self._write_gid_map), - # cgroupfs's view does not take into account cgroup namespaces. - # Weird, right? - # This means `/proc/PID/cgroup` will show the namespaced view of - # the cgroup that the PID is in, but `/sys/fs/cgroup` will still - # contain the whole system cgroup hierarchy regardless of namespace. - # Instead, namespaces act as "boundary box" around process movement - # requests when writing to cgroup.procs or creating new cgroups. - # So our first order of business here is to find out which cgroup we - # are running in. We do this by scanning the whole cgroupfs hierarchy - # and looking for our PID. This will populate - # `self._initial_cgroup_name`. - ("find_self_in_cgroup_hierarchy", self._find_self_in_cgroup_hierarchy), - # The cgroup nesting rules are complicated, but the short of it is: - # A cgroup can either **contain processes** OR **have limits**. - # Also, cgroups that contain processes must be leaf nodes. - # Also, cgroups that enforce limits must have their parent cgroup - # also have the same limit "controller" be active. - # So we will have two types of cgroups: - # - Leaf nodes with no controllers - # - Non-leaf nodes with controllers - # So initially, all the processes in the container's initial - # namespace need to be moved out to a new leaf node, - # otherwise we cannot turn on controllers on the initial - # cgroup. - # So we will set up the following hierarchy: - # /sys/fs/cgroup/$INITIAL: - # The cgroup where the container's processes were running - # the first time we run any Sandbox in the container. - # It may initially have no controllers enabled, but we will - # turn them on later. - # /sys/fs/cgroup/$INITIAL/leaf: - # The cgroup where the container's processes are moved to - # from the $INITIAL cgroup upon first run of any Sandbox in - # this container. When this code runs again, processes that - # are already in `$INITIAL/leaf` are not moved. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM: - # A per-Sandbox cgroup that never contains any processes. - # It will have controllers enabled on it but will never have - # specific limits enforced. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox: - # A per-Sandbox cgroup that never contains any processes. - # It will have controllers enabled on it and will enforce - # resource limits for the processes running in its /leaf. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox/leaf: - # A per-Sandbox cgroup that is running `runsc` (gVisor). - # It has no controllers enabled on it, but resources are - # being enforced by virtue of being a child of - # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor: - # A per-Sandbox cgroup that never contains any processes. - # It will have controllers enabled on it and will enforce - # resource limits for the processes running in its /leaf. - # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor/leaf: - # A per-Sandbox cgroup that is running a Python interpreter - # that manages the lifetime of the `runsc` process. - # It will run `Sandbox.maybe_main`. - # It has no controllers enabled on it, but resources are - # being enforced by virtue of being a child of - # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. - # - # This particular step creates the `$INITIAL/leaf` cgroup. - # If already created, it does nothing. - ("create_initial_leaf_cgroup", self._create_initial_leaf_cgroup), - # Move all processes in `$INITIAL` to `$INITIAL/leaf`. - ( - "move_initial_cgroup_processes_to_initial_leaf_cgroup", - self._move_initial_cgroup_processes_to_initial_leaf_cgroup, - ), - # Read the cgroup controllers enabled in `$INITIAL`. This acts - # as a bounding set on the ones we can enable in any child of it. - ("read_cgroup_controllers", self._read_cgroup_controllers), - # Cleanup old `$INITIAL/codeeval_*` cgroups that may be lying - # around from past runs. - ("cleanup_old_cgroups", self._cleanup_old_cgroups), - # Create a new `$INITIAL/codeeval_$NUM` cgroup. - ("create_codeeval_cgroup", self._create_codeeval_cgroup), - # Create a new `$INITIAL/codeeval_$NUM/sandbox` cgroup. - ("create_sandbox_cgroup", self._create_sandbox_cgroup), - # Create a new `$INITIAL/codeeval_$NUM/sandbox/leaf` cgroup. - ("create_sandbox_leaf_cgroup", self._create_sandbox_leaf_cgroup), - # Create a new `$INITIAL/codeeval_$NUM/supervisor` cgroup. - ("create_supervisor_cgroup", self._create_supervisor_cgroup), - # Create a new `$INITIAL/codeeval_$NUM/supervisor/leaf` cgroup. - ("create_supervisor_leaf_cgroup", self._create_supervisor_leaf_cgroup), - # Add controllers to `$INITIAL`. - ( - "add_cgroup_controllers_to_root", - self._add_cgroup_controllers_to_root, - ), - # Add controllers to `$INITIAL/codeeval_$NUM`. - ( - "add_cgroup_controllers_to_codeeval", - self._add_cgroup_controllers_to_codeeval, - ), - # Add controllers to `$INITIAL/codeeval_$NUM/sandbox`. - ( - "add_cgroup_controllers_to_sandbox", - self._add_cgroup_controllers_to_sandbox, - ), - # Set resource limits on `$INITIAL/codeeval_$NUM`. - ("set_sandbox_cgroup_limits", self._set_sandbox_cgroup_limits), - # Add controllers to `$INITIAL/codeeval_$NUM/supervisor`. - ( - "add_cgroup_controllers_to_supervisor", - self._add_cgroup_controllers_to_supervisor, - ), - # Set resource limits on `$INITIAL/codeeval_$NUM/supervisor`. - ("set_supervisor_cgroup_limits", self._set_supervisor_cgroup_limits), - # Move current process to - # `$INITIAL/codeeval_$NUM/supervisor/leaf`. - ( - "move_process_to_supervisor_leaf", - self._move_process_to_supervisor_leaf, - ), - # Double-check that we have moved to - # `$INITIAL/codeeval_$NUM/supervisor/leaf`. - ("sanity_check_own_cgroup", self._sanity_check_own_cgroup), - ) + ] + if do_resource_limiting: + self._operations.extend( + ( + # cgroupfs's view does not take into account cgroup namespaces. + # Weird, right? + # This means `/proc/PID/cgroup` will show the namespaced view of + # the cgroup that the PID is in, but `/sys/fs/cgroup` will still + # contain the whole system cgroup hierarchy regardless of namespace. + # Instead, namespaces act as "boundary box" around process movement + # requests when writing to cgroup.procs or creating new cgroups. + # So our first order of business here is to find out which cgroup we + # are running in. We do this by scanning the whole cgroupfs hierarchy + # and looking for our PID. This will populate + # `self._initial_cgroup_name`. + ( + "find_self_in_cgroup_hierarchy", + self._find_self_in_cgroup_hierarchy, + ), + # The cgroup nesting rules are complicated, but the short of it is: + # A cgroup can either **contain processes** OR **have limits**. + # Also, cgroups that contain processes must be leaf nodes. + # Also, cgroups that enforce limits must have their parent cgroup + # also have the same limit "controller" be active. + # So we will have two types of cgroups: + # - Leaf nodes with no controllers + # - Non-leaf nodes with controllers + # So initially, all the processes in the container's initial + # namespace need to be moved out to a new leaf node, + # otherwise we cannot turn on controllers on the initial + # cgroup. + # So we will set up the following hierarchy: + # /sys/fs/cgroup/$INITIAL: + # The cgroup where the container's processes were running + # the first time we run any Sandbox in the container. + # It may initially have no controllers enabled, but we will + # turn them on later. + # /sys/fs/cgroup/$INITIAL/leaf: + # The cgroup where the container's processes are moved to + # from the $INITIAL cgroup upon first run of any Sandbox in + # this container. When this code runs again, processes that + # are already in `$INITIAL/leaf` are not moved. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM: + # A per-Sandbox cgroup that never contains any processes. + # It will have controllers enabled on it but will never have + # specific limits enforced. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox: + # A per-Sandbox cgroup that never contains any processes. + # It will have controllers enabled on it and will enforce + # resource limits for the processes running in its /leaf. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox/leaf: + # A per-Sandbox cgroup that is running `runsc` (gVisor). + # It has no controllers enabled on it, but resources are + # being enforced by virtue of being a child of + # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor: + # A per-Sandbox cgroup that never contains any processes. + # It will have controllers enabled on it and will enforce + # resource limits for the processes running in its /leaf. + # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor/leaf: + # A per-Sandbox cgroup that is running a Python interpreter + # that manages the lifetime of the `runsc` process. + # It will run `Sandbox.maybe_main`. + # It has no controllers enabled on it, but resources are + # being enforced by virtue of being a child of + # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. + # + # This particular step creates the `$INITIAL/leaf` cgroup. + # If already created, it does nothing. + ( + "create_initial_leaf_cgroup", + self._create_initial_leaf_cgroup, + ), + # Move all processes in `$INITIAL` to `$INITIAL/leaf`. + ( + "move_initial_cgroup_processes_to_initial_leaf_cgroup", + self._move_initial_cgroup_processes_to_initial_leaf_cgroup, + ), + # Read the cgroup controllers enabled in `$INITIAL`. This acts + # as a bounding set on the ones we can enable in any child of it. + ("read_cgroup_controllers", self._read_cgroup_controllers), + # Cleanup old `$INITIAL/codeeval_*` cgroups that may be lying + # around from past runs. + ("cleanup_old_cgroups", self._cleanup_old_cgroups), + # Create a new `$INITIAL/codeeval_$NUM` cgroup. + ("create_codeeval_cgroup", self._create_codeeval_cgroup), + # Create a new `$INITIAL/codeeval_$NUM/sandbox` cgroup. + ("create_sandbox_cgroup", self._create_sandbox_cgroup), + # Create a new `$INITIAL/codeeval_$NUM/sandbox/leaf` cgroup. + ( + "create_sandbox_leaf_cgroup", + self._create_sandbox_leaf_cgroup, + ), + # Create a new `$INITIAL/codeeval_$NUM/supervisor` cgroup. + ("create_supervisor_cgroup", self._create_supervisor_cgroup), + # Create a new `$INITIAL/codeeval_$NUM/supervisor/leaf` cgroup. + ( + "create_supervisor_leaf_cgroup", + self._create_supervisor_leaf_cgroup, + ), + # Add controllers to `$INITIAL`. + ( + "add_cgroup_controllers_to_root", + self._add_cgroup_controllers_to_root, + ), + # Add controllers to `$INITIAL/codeeval_$NUM`. + ( + "add_cgroup_controllers_to_codeeval", + self._add_cgroup_controllers_to_codeeval, + ), + # Add controllers to `$INITIAL/codeeval_$NUM/sandbox`. + ( + "add_cgroup_controllers_to_sandbox", + self._add_cgroup_controllers_to_sandbox, + ), + # Set resource limits on `$INITIAL/codeeval_$NUM`. + ("set_sandbox_cgroup_limits", self._set_sandbox_cgroup_limits), + # Add controllers to `$INITIAL/codeeval_$NUM/supervisor`. + ( + "add_cgroup_controllers_to_supervisor", + self._add_cgroup_controllers_to_supervisor, + ), + # Set resource limits on `$INITIAL/codeeval_$NUM/supervisor`. + ( + "set_supervisor_cgroup_limits", + self._set_supervisor_cgroup_limits, + ), + # Move current process to + # `$INITIAL/codeeval_$NUM/supervisor/leaf`. + ( + "move_process_to_supervisor_leaf", + self._move_process_to_supervisor_leaf, + ), + # Double-check that we have moved to + # `$INITIAL/codeeval_$NUM/supervisor/leaf`. + ("sanity_check_own_cgroup", self._sanity_check_own_cgroup), + ) + ) def _status(self): """ @@ -419,62 +440,65 @@ def _status(self): if self._checkpoint == self._operations[-1][0]: main_status = "OK" my_pid = os.getpid() - status_line = f"{main_status} (euid={self._my_euid} egid={self._my_egid} pid={my_pid} initial_cgroup_name={self._initial_cgroup_name} codeeval_cgroup_name={self._codeeval_cgroup_name} controllers={self._cgroup_controllers})" - cgroupfs_data = [] - for cgroup_components in ( - (self._initial_cgroup_name,), - (self._initial_cgroup_name, self._CGROUP_LEAF), - (self._initial_cgroup_name, self._codeeval_cgroup_name), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_LEAF, - ), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_SUPERVISOR_NAME, - ), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_SUPERVISOR_NAME, - self._CGROUP_LEAF, - ), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_SANDBOX_NAME, - ), - ( - self._initial_cgroup_name, - self._codeeval_cgroup_name, - self._CGROUP_SANDBOX_NAME, - self._CGROUP_LEAF, - ), - ): - if any(c is None for c in cgroup_components): - continue - file_data = [] - for filename in ("procs", "controllers", "subtree_control"): - data = None - try: - with self._open( - self._cgroup_path( - *(cgroup_components + (f"cgroup.{filename}",)) - ), - "rb", - ) as f: - data = f.read().decode("ascii").replace("\n", " ") - except Exception as e: - data = f"[fail: {e}]" - file_data.append(f"{filename}: {data}") - cgroup_components_joined = os.sep.join(cgroup_components) - file_data_joined = ", ".join(file_data) - cgroupfs_data.append(f"{cgroup_components_joined}=[{file_data_joined}]") - if len(cgroupfs_data) > 0: - cgroupfs_data_joined = " ".join(cgroupfs_data) - status_line += f" {cgroupfs_data_joined}" + status_line = f"{main_status} (euid={self._my_euid} egid={self._my_egid} pid={my_pid} do_resource_limiting={self._do_resource_limiting} initial_cgroup_name={self._initial_cgroup_name} codeeval_cgroup_name={self._codeeval_cgroup_name} controllers={self._cgroup_controllers})" + if self._do_resource_limiting: + cgroupfs_data = [] + for cgroup_components in ( + (self._initial_cgroup_name,), + (self._initial_cgroup_name, self._CGROUP_LEAF), + (self._initial_cgroup_name, self._codeeval_cgroup_name), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_LEAF, + ), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_SUPERVISOR_NAME, + ), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_SUPERVISOR_NAME, + self._CGROUP_LEAF, + ), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_SANDBOX_NAME, + ), + ( + self._initial_cgroup_name, + self._codeeval_cgroup_name, + self._CGROUP_SANDBOX_NAME, + self._CGROUP_LEAF, + ), + ): + if any(c is None for c in cgroup_components): + continue + file_data = [] + for filename in ("procs", "controllers", "subtree_control"): + data = None + try: + with self._open( + self._cgroup_path( + *(cgroup_components + (f"cgroup.{filename}",)) + ), + "rb", + ) as f: + data = f.read().decode("ascii").replace("\n", " ") + except Exception as e: + data = f"[fail: {e}]" + file_data.append(f"{filename}: {data}") + cgroup_components_joined = os.sep.join(cgroup_components) + file_data_joined = ", ".join(file_data) + cgroupfs_data.append( + f"{cgroup_components_joined}=[{file_data_joined}]" + ) + if len(cgroupfs_data) > 0: + cgroupfs_data_joined = " ".join(cgroupfs_data) + status_line += f" {cgroupfs_data_joined}" return status_line def _cgroup_path(self, *components): @@ -939,6 +963,8 @@ def move_process_to_sandbox_leaf_cgroup_lambda(self): :return: A function to move the current process to the sandbox cgroup. :raises SandboxException: If not queried after we have already chosen a new cgroup name. """ + if not self._do_resource_limiting: + return lambda: None if self._codeeval_cgroup_name is None: raise Sandbox.SandboxException( "Tried to move process to sandbox leaf cgroup before we know it" @@ -989,13 +1015,15 @@ def _move(cgroup_path): def monitor_cgroup_resources(self): """ - Spawns a background thread that monitors resources. + Spawns a background thread that monitors resources, if limiting is enabled. cgroups should be taking care of this, but some systems do not enforce this. So this does the same in userspace. Better than nothing. - :return: A function to cancel the monitor thread. + :return: A function to cancel the monitor thread, if resource limiting is enabled. """ + if not self._do_resource_limiting: + return lambda: None self_memory_path = self._cgroup_path( self._initial_cgroup_name, self._codeeval_cgroup_name, @@ -1224,20 +1252,6 @@ def check_cgroups(cls): "cgroupfs does not have the 'memory' controller enabled, necessary to enforce memory limits; please enable it" ) - @classmethod - def cgroups_available(cls) -> bool: - """ - Returns whether cgroupfs is mounted and usable for resource limit enforcement. - - :return: Whether cgroupfs is mounted and usable for resource limit enforcement. - """ - try: - cls.check_cgroups() - except Exception: - return False - else: - return True - @classmethod def check_procfs(cls): """ @@ -1381,12 +1395,18 @@ def install_runsc(cls): os.rename(download_path, cls.AUTO_INSTALLATION_PATH) @classmethod - def check_setup(cls, language: str, auto_install_allowed: bool): + def check_setup( + cls, + language: str, + auto_install_allowed: bool, + require_resource_limiting: bool, + ): """ Verifies that the environment is compatible with running sandboxes. :param language: The programming language to run. :param auto_install_allowed: Whether auto-installation of `runsc` is allowed. + :param require_resource_limiting: Check that the host supports resource limiting via cgroups. :return: Nothing. :raises ValueError: If provided an invalid language name. @@ -1405,7 +1425,8 @@ def check_setup(cls, language: str, auto_install_allowed: bool): ) cls.check_platform() cls.check_unshare() - cls.check_cgroups() + if require_resource_limiting: + cls.check_cgroups() cls.check_procfs() if not auto_install_allowed and cls.get_runsc_path() is None: raise cls.GVisorNotInstalledException( @@ -1475,8 +1496,9 @@ def __init__( debug: bool, networking_allowed: bool, max_runtime_seconds: int, - max_ram_bytes: typing.Optional[int], - persistent_home_dir: typing.Optional[str], + max_ram_bytes: typing.Optional[int] = None, + require_resource_limiting: bool = False, + persistent_home_dir: typing.Optional[str] = None, ): """ Constructor. @@ -1488,6 +1510,7 @@ def __init__( :param networking_allowed: Whether the code should be given access to the network. :param max_runtime_seconds: How long the code should be allowed to run, in seconds. :param max_ram_bytes: How many bytes of RAM the interpreter should be allowed to use, or `None` for no limit. + :param require_resource_limiting: If true, refuse to launch a sandbox if the host doesn't support resource limiting via cgroups. :param persistent_home_dir: Optional directory which will be mapped read-write to this real host directory. """ self._init( @@ -1499,6 +1522,7 @@ def __init__( "networking_allowed": networking_allowed, "max_runtime_seconds": max_runtime_seconds, "max_ram_bytes": max_ram_bytes, + "require_resource_limiting": require_resource_limiting, "persistent_home_dir": persistent_home_dir, } ) @@ -1517,13 +1541,12 @@ def _init(self, settings): self._networking_allowed = self._settings["networking_allowed"] self._max_runtime_seconds = self._settings["max_runtime_seconds"] self._max_ram_bytes = self._settings["max_ram_bytes"] + self._require_resource_limiting = self._settings[ + "require_resource_limiting" + ] or all((self._max_ram_bytes is None,)) self._persistent_home_dir = self._settings["persistent_home_dir"] self._sandboxed_command = None - self._switcheroo = self._Switcheroo( - libc=self._libc(), - log_path=os.path.join(self._logs_path, "switcheroo.txt"), - max_sandbox_ram_bytes=self._max_ram_bytes, - ) + self._switcheroo = None def _setup_sandbox(self): """ @@ -1550,7 +1573,18 @@ def _setup_sandbox(self): f"Persistent home directory {self._persistent_home_dir} does not exist" ) oci_config["root"]["path"] = rootfs_path - + do_resource_limiting = True + if not self._require_resource_limiting: + try: + self.check_cgroups() + except self.EnvironmentNeedsSetupException: + do_resource_limiting = False + self._switcheroo = self._Switcheroo( + libc=self._libc(), + log_path=os.path.join(self._logs_path, "switcheroo.txt"), + max_sandbox_ram_bytes=self._max_ram_bytes, + do_resource_limiting=do_resource_limiting, + ) try: self._switcheroo.do() except Exception as e: diff --git a/tests/openwebui/functions/run_code_test.sh b/tests/openwebui/functions/run_code_test.sh index 0ab6386..57a57b4 100755 --- a/tests/openwebui/functions/run_code_test.sh +++ b/tests/openwebui/functions/run_code_test.sh @@ -4,6 +4,7 @@ set -euo pipefail REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +echo 'Running function self-tests...' >&2 docker run --rm \ --security-opt=seccomp=unconfined \ --security-opt=apparmor=unconfined \ @@ -13,3 +14,31 @@ docker run --rm \ --mount=type=bind,source="$REPO_DIR",target=/test \ ghcr.io/open-webui/open-webui:main \ python3 /test/open-webui/functions/run_code.py --self_test "$@" + +echo 'Checking cgroupfs presence enforcement enabled (this should fail)...' >&2 +docker run --rm \ + --security-opt=seccomp=unconfined \ + --security-opt=apparmor=unconfined \ + --security-opt=label=type:container_engine_t \ + --mount=type=tmpfs,target=/sys/fs/cgroup,readonly=true \ + --mount=type=bind,source=/proc,target=/proc2,readonly=false,bind-recursive=disabled \ + --mount=type=bind,source="$REPO_DIR",target=/test \ + --env=CODE_EVAL_VALVE_OVERRIDE_MAX_RAM_MEGABYTES=32 \ + --env=CODE_EVAL_VALVE_OVERRIDE_REQUIRE_RESOURCE_LIMITING=true \ + ghcr.io/open-webui/open-webui:main \ + python3 /test/open-webui/functions/run_code.py \ + --use_sample_code --want_status=SANDBOX_ERROR "$@" + +echo 'Checking cgroupfs presence enforcement disabled (this should succeed)...' >&2 +docker run --rm \ + --security-opt=seccomp=unconfined \ + --security-opt=apparmor=unconfined \ + --security-opt=label=type:container_engine_t \ + --mount=type=tmpfs,target=/sys/fs/cgroup,readonly=true \ + --mount=type=bind,source=/proc,target=/proc2,readonly=false,bind-recursive=disabled \ + --mount=type=bind,source="$REPO_DIR",target=/test \ + --env=CODE_EVAL_VALVE_OVERRIDE_MAX_RAM_MEGABYTES=32 \ + --env=CODE_EVAL_VALVE_OVERRIDE_REQUIRE_RESOURCE_LIMITING=false \ + ghcr.io/open-webui/open-webui:main \ + python3 /test/open-webui/functions/run_code.py \ + --use_sample_code --want_status=OK "$@" diff --git a/tests/openwebui/tools/run_code_test.sh b/tests/openwebui/tools/run_code_test.sh index 33ba4e4..9f1b9fe 100755 --- a/tests/openwebui/tools/run_code_test.sh +++ b/tests/openwebui/tools/run_code_test.sh @@ -4,6 +4,7 @@ set -euo pipefail REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +echo 'Running tool self-tests...' >&2 docker run --rm \ --security-opt=seccomp=unconfined \ --security-opt=apparmor=unconfined \ @@ -13,3 +14,31 @@ docker run --rm \ --mount=type=bind,source="$REPO_DIR",target=/test \ ghcr.io/open-webui/open-webui:main \ python3 /test/open-webui/tools/run_code.py --self_test "$@" + +echo 'Checking cgroupfs presence enforcement enabled (this should fail)...' >&2 +docker run --rm \ + --security-opt=seccomp=unconfined \ + --security-opt=apparmor=unconfined \ + --security-opt=label=type:container_engine_t \ + --mount=type=tmpfs,target=/sys/fs/cgroup,readonly=true \ + --mount=type=bind,source=/proc,target=/proc2,readonly=false,bind-recursive=disabled \ + --mount=type=bind,source="$REPO_DIR",target=/test \ + --env=CODE_EVAL_VALVE_OVERRIDE_MAX_RAM_MEGABYTES=32 \ + --env=CODE_EVAL_VALVE_OVERRIDE_REQUIRE_RESOURCE_LIMITING=true \ + ghcr.io/open-webui/open-webui:main \ + python3 /test/open-webui/tools/run_code.py \ + --use_sample_code --want_status=SANDBOX_ERROR "$@" + +echo 'Checking cgroupfs presence enforcement disabled (this should succeed)...' >&2 +docker run --rm \ + --security-opt=seccomp=unconfined \ + --security-opt=apparmor=unconfined \ + --security-opt=label=type:container_engine_t \ + --mount=type=tmpfs,target=/sys/fs/cgroup,readonly=true \ + --mount=type=bind,source=/proc,target=/proc2,readonly=false,bind-recursive=disabled \ + --mount=type=bind,source="$REPO_DIR",target=/test \ + --env=CODE_EVAL_VALVE_OVERRIDE_MAX_RAM_MEGABYTES=32 \ + --env=CODE_EVAL_VALVE_OVERRIDE_REQUIRE_RESOURCE_LIMITING=false \ + ghcr.io/open-webui/open-webui:main \ + python3 /test/open-webui/tools/run_code.py \ + --use_sample_code --want_status=OK "$@"