diff --git a/lean/commands/backtest.py b/lean/commands/backtest.py index 62f232ae..d1c14f18 100644 --- a/lean/commands/backtest.py +++ b/lean/commands/backtest.py @@ -11,7 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. - +import json from pathlib import Path from typing import List, Optional, Tuple from click import command, option, argument, Choice @@ -289,6 +289,10 @@ def _select_organization() -> QCMinimalOrganization: type=(str, str), multiple=True, hidden=True) +@option("--extra-docker-config", + type=str, + default="{}", + hidden=True) @option("--no-update", is_flag=True, default=False, @@ -307,6 +311,7 @@ def backtest(project: Path, backtest_name: str, addon_module: Optional[List[str]], extra_config: Optional[Tuple[str, str]], + extra_docker_config: Optional[str], no_update: bool, **kwargs) -> None: """Backtest a project locally using Docker. @@ -407,4 +412,5 @@ def backtest(project: Path, engine_image, debugging_method, release, - detach) + detach, + json.loads(extra_docker_config)) diff --git a/lean/commands/live/deploy.py b/lean/commands/live/deploy.py index f4ccc878..314ef7ab 100644 --- a/lean/commands/live/deploy.py +++ b/lean/commands/live/deploy.py @@ -11,6 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json from pathlib import Path from typing import Any, Dict, List, Optional, Tuple from click import option, argument, Choice @@ -255,6 +256,10 @@ def _get_default_value(key: str) -> Optional[Any]: type=(str, str), multiple=True, hidden=True) +@option("--extra-docker-config", + type=str, + default="{}", + hidden=True) @option("--no-update", is_flag=True, default=False, @@ -275,6 +280,7 @@ def deploy(project: Path, show_secrets: bool, addon_module: Optional[List[str]], extra_config: Optional[Tuple[str, str]], + extra_docker_config: Optional[str], no_update: bool, **kwargs) -> None: """Start live trading a project locally using Docker. @@ -430,4 +436,4 @@ def deploy(project: Path, raise RuntimeError(f"InteractiveBrokers is currently not supported for ARM hosts") lean_runner = container.lean_runner - lean_runner.run_lean(lean_config, environment_name, algorithm_file, output, engine_image, None, release, detach) + lean_runner.run_lean(lean_config, environment_name, algorithm_file, output, engine_image, None, release, detach, json.loads(extra_docker_config)) diff --git a/lean/commands/optimize.py b/lean/commands/optimize.py index 4478dfd1..5a41d5aa 100644 --- a/lean/commands/optimize.py +++ b/lean/commands/optimize.py @@ -11,6 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json from pathlib import Path from typing import Optional, List, Tuple from datetime import datetime, timedelta @@ -18,6 +19,7 @@ from click import command, argument, option, Choice, IntRange from lean.click import LeanCommand, PathParameter, ensure_options +from lean.components.docker.lean_runner import LeanRunner from lean.constants import DEFAULT_ENGINE_IMAGE from lean.container import container from lean.models.api import QCParameter, QCBacktest @@ -119,6 +121,10 @@ def get_filename_timestamp(path: Path) -> datetime: type=(str, str), multiple=True, hidden=True) +@option("--extra-docker-config", + type=str, + default="{}", + hidden=True) @option("--no-update", is_flag=True, default=False, @@ -139,6 +145,7 @@ def optimize(project: Path, max_concurrent_backtests: Optional[int], addon_module: Optional[List[str]], extra_config: Optional[Tuple[str, str]], + extra_docker_config: Optional[str], no_update: bool) -> None: """Optimize a project's parameters locally using Docker. @@ -308,6 +315,9 @@ def optimize(project: Path, ) container.update_manager.pull_docker_image_if_necessary(engine_image, update, no_update) + # Add know additional run options from the extra docker config + LeanRunner.parse_extra_docker_config(run_options, json.loads(extra_docker_config)) + project_manager.copy_code(algorithm_file.parent, output / "code") success = container.docker_manager.run_image(engine_image, **run_options) diff --git a/lean/commands/research.py b/lean/commands/research.py index ad03b626..d30cbbec 100644 --- a/lean/commands/research.py +++ b/lean/commands/research.py @@ -11,10 +11,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json from pathlib import Path from typing import Optional, Tuple from click import command, argument, option, Choice from lean.click import LeanCommand, PathParameter +from lean.components.docker.lean_runner import LeanRunner from lean.constants import DEFAULT_RESEARCH_IMAGE, LEAN_ROOT_PATH from lean.container import container from lean.models.data_providers import QuantConnectDataProvider, all_data_providers @@ -65,6 +67,10 @@ def _check_docker_output(chunk: str, port: int) -> None: type=(str, str), multiple=True, hidden=True) +@option("--extra-docker-config", + type=str, + default="{}", + hidden=True) @option("--no-update", is_flag=True, default=False, @@ -79,6 +85,7 @@ def research(project: Path, image: Optional[str], update: bool, extra_config: Optional[Tuple[str, str]], + extra_docker_config: Optional[str], no_update: bool, **kwargs) -> None: """Run a Jupyter Lab environment locally using Docker. @@ -116,7 +123,7 @@ def research(project: Path, # Set extra config for key, value in extra_config: lean_config[key] = value - + run_options = lean_runner.get_basic_docker_config(lean_config, algorithm_file, temp_manager.create_temporary_directory(), @@ -160,6 +167,9 @@ def research(project: Path, # Run the script that starts Jupyter Lab when all set up has been done run_options["commands"].append("./start.sh") + # Add know additional run options from the extra docker config + LeanRunner.parse_extra_docker_config(run_options, json.loads(extra_docker_config)) + project_config_manager = container.project_config_manager cli_config_manager = container.cli_config_manager diff --git a/lean/components/docker/lean_runner.py b/lean/components/docker/lean_runner.py index ad18efd4..b309183f 100644 --- a/lean/components/docker/lean_runner.py +++ b/lean/components/docker/lean_runner.py @@ -13,6 +13,8 @@ from pathlib import Path from typing import Any, Dict, Optional, List +import docker.types + from lean.components.cloud.module_manager import ModuleManager from lean.components.config.lean_config_manager import LeanConfigManager from lean.components.config.output_config_manager import OutputConfigManager @@ -70,7 +72,8 @@ def run_lean(self, image: DockerImage, debugging_method: Optional[DebuggingMethod], release: bool, - detach: bool) -> None: + detach: bool, + extra_docker_config: Optional[Dict[str, Any]] = None) -> None: """Runs the LEAN engine locally in Docker. Raises an error if something goes wrong. @@ -83,6 +86,7 @@ def run_lean(self, :param debugging_method: the debugging method if debugging needs to be enabled, None if not :param release: whether C# projects should be compiled in release configuration instead of debug :param detach: whether LEAN should run in a detached container + :param extra_docker_config: additional docker configurations """ project_dir = algorithm_file.parent @@ -95,6 +99,9 @@ def run_lean(self, release, detach) + # Add know additional run options from the extra docker config + self.parse_extra_docker_config(run_options, extra_docker_config) + # Set up PTVSD debugging if debugging_method == DebuggingMethod.PTVSD: run_options["ports"]["5678"] = "5678" @@ -762,3 +769,11 @@ def mount_project_and_library_directories(self, project_dir: Path, run_options: "bind": "/Library", "mode": "rw" } + + @staticmethod + def parse_extra_docker_config(run_options: Dict[str, Any], extra_docker_config: Optional[Dict[str, Any]]) -> None: + # Add know additional run options from the extra docker config. + # For now, only device_requests is supported + if extra_docker_config is not None and "device_requests" in extra_docker_config: + run_options["device_requests"] = [docker.types.DeviceRequest(**device_request) + for device_request in extra_docker_config["device_requests"]] diff --git a/tests/components/docker/test_lean_runner.py b/tests/components/docker/test_lean_runner.py index 6d7594b2..9436779d 100644 --- a/tests/components/docker/test_lean_runner.py +++ b/tests/components/docker/test_lean_runner.py @@ -14,6 +14,7 @@ from pathlib import Path from unittest import mock +import docker.types import pytest from lean.components.config.lean_config_manager import LeanConfigManager @@ -581,3 +582,40 @@ def test_run_lean_compiles_csharp_project_that_is_part_of_a_solution(in_solution assert project_dir_str in kwargs["volumes"] assert kwargs["volumes"][project_dir_str]["bind"] == "/LeanCLI" assert str(root_dir) not in kwargs["volumes"] + + +def test_run_lean_passes_device_requests() -> None: + create_fake_lean_cli_directory() + + docker_manager = mock.Mock() + docker_manager.run_image.return_value = True + + lean_runner = create_lean_runner(docker_manager) + + lean_runner.run_lean({"transaction-log": "transaction-log.log"}, + "backtesting", + Path.cwd() / "Python Project" / "main.py", + Path.cwd() / "output", + ENGINE_IMAGE, + None, + False, + False, + extra_docker_config={"device_requests": [{"count": -1, "capabilities": [["gpu"]]}]}) + + docker_manager.run_image.assert_called_once() + args, kwargs = docker_manager.run_image.call_args + + assert "device_requests" in kwargs + + device_requests = kwargs["device_requests"] + assert len(device_requests) == 1 + + device_request: docker.types.DeviceRequest = device_requests[0] + assert isinstance(device_request, docker.types.DeviceRequest) + assert device_request.count == -1 + assert (len(device_request.capabilities) == 1 and + len(device_request.capabilities[0]) == 1 and + device_request.capabilities[0][0] == "gpu") + assert device_request.driver == "" + assert device_request.device_ids == [] + assert device_request.options == {}