diff --git a/docs/source/how-to/configure-workflows/systems.md b/docs/source/how-to/configure-workflows/systems.md index 161560bc84..d4954cf799 100644 --- a/docs/source/how-to/configure-workflows/systems.md +++ b/docs/source/how-to/configure-workflows/systems.md @@ -108,6 +108,8 @@ The docker system is configured with the following attributes: ```{Note} * The docker container must have `olive-ai` installed. +* Input resources referenced in the workflow config (input models, data files, user scripts) are mounted into the container read-only. Only write destinations are mounted read-write: `engine.output_dir`, the engine cache directory, `output_dir`/`logging_dir` in pass training arguments and metric `sample_log_dir`. An input located inside one of these directories is also bound read-only at its nested location, so it cannot be modified through the writable mount. +* When `hf_token` is set to `true`, the full HuggingFace token from the host (`HF_TOKEN` environment variable or `~/.huggingface/token`) is passed to the container as the `HF_TOKEN` environment variable, so any code in the image built from `build_context_path` can read it. Only enable it for images you trust, and prefer a fine-grained, read-only token for Docker runs. ``` ### Prerequisites diff --git a/docs/source/reference/options.md b/docs/source/reference/options.md index a55b9cc67d..05f7b7f449 100644 --- a/docs/source/reference/options.md +++ b/docs/source/reference/options.md @@ -122,7 +122,7 @@ information of the system contains following items: - `config: [Dict]` The system config dictionary that contains the system specific information. The fields can be provided directly under the parent dictionary. - `accelerators: [List[str]]` The accelerators that will be used for this workflow. - - `hf_token: [bool]` Whether to use a Huggingface token to access Huggingface resources. If it is set to `True`, For local system, Docker system, and PythonEnvironment system, Olive will retrieve the token from the `HF_TOKEN` environment variable or from the token file located at `~/.huggingface/token`. If set to `False`, no token will be utilized during this workflow run. The default value is `False`. + - `hf_token: [bool]` Whether to use a Huggingface token to access Huggingface resources. If it is set to `True`, For local system, Docker system, and PythonEnvironment system, Olive will retrieve the token from the `HF_TOKEN` environment variable or from the token file located at `~/.huggingface/token`. If set to `False`, no token will be utilized during this workflow run. The default value is `False`. For the Docker system, the token is exposed to everything running inside the container, see [Docker System](../how-to/configure-workflows/systems.md#docker-system). Please refer to [How To Configure System](../how-to/configure-workflows/systems.md) for the more information of the system config dictionary. diff --git a/olive/systems/docker/docker_system.py b/olive/systems/docker/docker_system.py index 8371cccd44..85602b60d2 100644 --- a/olive/systems/docker/docker_system.py +++ b/olive/systems/docker/docker_system.py @@ -30,6 +30,18 @@ logger = logging.getLogger(__name__) +# Resources the workflow writes to, matched by config key prefix or by the name of the config field holding them. +# Every other resource is an input and is mounted read-only. +WRITABLE_RESOURCE_KEY_PREFIXES = (("engine", "cache_config"),) +WRITABLE_RESOURCE_KEY_NAMES = frozenset( + { + "output_dir", # engine output, also Trainer output dir in pass training args + "cache_dir", # engine cache + "logging_dir", # Trainer logging dir in pass training args + "sample_log_dir", # evaluator metric sample logs + } +) + class DockerSystem(OliveSystem): system_type = SystemType.Docker @@ -176,20 +188,47 @@ def _mount_workflow_config( # Find and mount all resources all_resources = find_all_resources(config_copy, ignore_keys=ignore_keys) + writable_mounts = [] + input_mounts = [] for resource_key, resource_path in all_resources.items(): # Create container path - container_path = str(self.work_dir / Path(resource_path.get_path()).name) + host_path = resource_path.get_path() + container_path = str(self.work_dir / Path(host_path).name) - # Add volume mount - mount_string = f"{resource_path.get_path()}:{container_path}" - volumes.append(mount_string) + # Add volume mount, inputs are read-only so the container cannot modify user supplied resources + if self._is_writable_resource(resource_key): + volumes.append(f"{host_path}:{container_path}") + writable_mounts.append((Path(host_path).resolve(), container_path)) + else: + volumes.append(f"{host_path}:{container_path}:ro") + input_mounts.append(host_path) # Update config with container path set_nested_dict_value(config_copy, resource_key, container_path) + # An input located inside a writable directory (e.g. under the default output dir, the working directory) + # would still be writable through that directory's mount. Bind it read-only at its nested location too so + # the read-only mount shadows the writable view. + for host_path in input_mounts: + resolved_host_path = Path(host_path).resolve() + for writable_host_path, writable_container_path in writable_mounts: + if resolved_host_path != writable_host_path and resolved_host_path.is_relative_to(writable_host_path): + nested_container_path = Path(writable_container_path) / resolved_host_path.relative_to( + writable_host_path + ) + volumes.append(f"{host_path}:{nested_container_path}:ro") + logger.debug("Container config: %s", config_copy) return config_copy + @staticmethod + def _is_writable_resource(resource_key: tuple) -> bool: + """Check whether the resource at the given config key is written to by the workflow.""" + # match the field name anywhere in the key since list values (e.g. cache_dir: [...]) end with an index + if any(key in WRITABLE_RESOURCE_KEY_NAMES for key in resource_key): + return True + return any(resource_key[: len(prefix)] == prefix for prefix in WRITABLE_RESOURCE_KEY_PREFIXES) + def _run_container( self, command: str, @@ -289,7 +328,7 @@ def _create_config_mount(self, tmp_dir: str, config: dict) -> tuple[str, str]: # Create mount string container_path = str(self.work_dir / "config.json") - mount_string = f"{config_file_path}:{container_path}" + mount_string = f"{config_file_path}:{container_path}:ro" return container_path, mount_string def _create_runner_script_mount(self) -> tuple[str, str]: @@ -297,7 +336,7 @@ def _create_runner_script_mount(self) -> tuple[str, str]: runner_script_name = "workflow_runner.py" container_path = str(self.work_dir / runner_script_name) host_path = Path(__file__).resolve().parent / runner_script_name - mount_string = f"{host_path}:{container_path}" + mount_string = f"{host_path}:{container_path}:ro" return container_path, mount_string @staticmethod diff --git a/test/systems/docker/test_docker_system.py b/test/systems/docker/test_docker_system.py index 5430b68587..ffa9986f6e 100644 --- a/test/systems/docker/test_docker_system.py +++ b/test/systems/docker/test_docker_system.py @@ -206,6 +206,102 @@ def test_run_workflow_container_error(self, mock_find_resources, mock_tempdir, m # Verify container was still removed mock_container.remove.assert_called_once() + @patch("olive.systems.docker.docker_system.docker.from_env") + def test_mount_workflow_config_mounts_inputs_read_only(self, mock_from_env, tmp_path): + mock_from_env.return_value = MagicMock() + docker_config = self.get_default_docker_config() + docker_system = DockerSystem( + image_name=docker_config.image_name, + build_context_path=docker_config.build_context_path, + dockerfile=docker_config.dockerfile, + work_dir=docker_config.work_dir, + ) + output_dir = tmp_path / "output" + cache_dir = tmp_path / "cache" + output_dir.mkdir() + cache_dir.mkdir() + config = { + "input_model": {"type": "ONNXModel", "config": {"model_path": str(ONNX_MODEL_PATH)}}, + "engine": {"output_dir": str(output_dir), "cache_dir": str(cache_dir)}, + } + + volumes = [] + container_config = docker_system._mount_workflow_config(config, volumes) + + def container_path(host_path): + return str(docker_system.work_dir / host_path.name) + + assert f"{ONNX_MODEL_PATH}:{container_path(ONNX_MODEL_PATH)}:ro" in volumes + assert f"{output_dir}:{container_path(output_dir)}" in volumes + assert f"{cache_dir}:{container_path(cache_dir)}" in volumes + assert len(volumes) == 3 + assert container_config["input_model"]["config"]["model_path"] == container_path(ONNX_MODEL_PATH) + + @patch("olive.systems.docker.docker_system.docker.from_env") + def test_mount_workflow_config_shadows_input_inside_writable_dir(self, mock_from_env, tmp_path): + mock_from_env.return_value = MagicMock() + docker_config = self.get_default_docker_config() + docker_system = DockerSystem( + image_name=docker_config.image_name, + build_context_path=docker_config.build_context_path, + dockerfile=docker_config.dockerfile, + work_dir=docker_config.work_dir, + ) + # the default output dir is the working directory, which usually contains the input model + output_dir = tmp_path / "project" + model_path = output_dir / "models" / "model.onnx" + model_path.parent.mkdir(parents=True) + model_path.write_bytes(b"dummy") + config = { + "input_model": {"type": "ONNXModel", "config": {"model_path": str(model_path)}}, + "engine": {"output_dir": str(output_dir)}, + } + + volumes = [] + docker_system._mount_workflow_config(config, volumes) + + output_container_path = docker_system.work_dir / output_dir.name + assert f"{output_dir}:{output_container_path}" in volumes + assert f"{model_path}:{docker_system.work_dir / model_path.name}:ro" in volumes + assert f"{model_path}:{output_container_path / 'models' / 'model.onnx'}:ro" in volumes + assert len(volumes) == 3 + + @pytest.mark.parametrize( + ("resource_key", "expected"), + [ + (("engine", "output_dir"), True), + (("engine", "cache_dir"), True), + (("engine", "cache_dir", 0), True), + (("engine", "cache_config", "cache_dir"), True), + (("input_model", "config", "model_path"), False), + (("data_configs", 0, "load_dataset_config", "data_dir"), False), + (("engine",), False), + (("passes", "lora", "training_args", "output_dir"), True), + (("passes", "lora", "training_args", "logging_dir"), True), + (("engine", "evaluator", "metrics", 0, "sample_log_dir"), True), + (("passes", "lora", "train_data_config", "load_dataset_config", "data_dir"), False), + ], + ) + def test_is_writable_resource(self, resource_key, expected): + assert DockerSystem._is_writable_resource(resource_key) == expected + + @patch("olive.systems.docker.docker_system.docker.from_env") + def test_config_and_runner_mounts_are_read_only(self, mock_from_env, tmp_path): + mock_from_env.return_value = MagicMock() + docker_config = self.get_default_docker_config() + docker_system = DockerSystem( + image_name=docker_config.image_name, + build_context_path=docker_config.build_context_path, + dockerfile=docker_config.dockerfile, + work_dir=docker_config.work_dir, + ) + + _, config_mount = docker_system._create_config_mount(str(tmp_path), {}) + _, runner_mount = docker_system._create_runner_script_mount() + + assert config_mount.endswith(":ro") + assert runner_mount.endswith(":ro") + @patch("olive.systems.docker.docker_system.docker.from_env") def test_remove_image(self, mock_from_env): mock_docker_client = MagicMock()