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..af8d69f7 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 known 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..710c1f6e 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 known 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..0abd5ffe 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 known 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 known 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/commands/test_backtest.py b/tests/commands/test_backtest.py index b2d5f4d7..200c15bf 100644 --- a/tests/commands/test_backtest.py +++ b/tests/commands/test_backtest.py @@ -57,7 +57,8 @@ def test_backtest_calls_lean_runner_with_correct_algorithm_file() -> None: ENGINE_IMAGE, None, False, - False) + False, + {}) def test_backtest_calls_lean_runner_with_default_output_directory() -> None: @@ -88,7 +89,8 @@ def test_backtest_calls_lean_runner_with_custom_output_directory() -> None: ENGINE_IMAGE, None, False, - False) + False, + {}) def test_backtest_calls_lean_runner_with_release_mode() -> None: @@ -105,7 +107,8 @@ def test_backtest_calls_lean_runner_with_release_mode() -> None: ENGINE_IMAGE, None, True, - False) + False, + {}) def test_backtest_calls_lean_runner_with_detach() -> None: @@ -122,7 +125,8 @@ def test_backtest_calls_lean_runner_with_detach() -> None: ENGINE_IMAGE, None, False, - True) + True, + {}) def test_backtest_aborts_when_project_does_not_exist() -> None: @@ -163,7 +167,8 @@ def test_backtest_forces_update_when_update_option_given() -> None: ENGINE_IMAGE, None, False, - False) + False, + {}) def test_backtest_passes_custom_image_to_lean_runner_when_set_in_config() -> None: @@ -182,7 +187,8 @@ def test_backtest_passes_custom_image_to_lean_runner_when_set_in_config() -> Non DockerImage(name="custom/lean", tag="123"), None, False, - False) + False, + {}) def test_backtest_passes_custom_image_to_lean_runner_when_given_as_option() -> None: @@ -201,7 +207,8 @@ def test_backtest_passes_custom_image_to_lean_runner_when_given_as_option() -> N DockerImage(name="custom/lean", tag="456"), None, False, - False) + False, + {}) @pytest.mark.parametrize("python_venv", ["Custom-venv", @@ -289,7 +296,8 @@ def test_backtest_passes_correct_debugging_method_to_lean_runner(value: str, deb ENGINE_IMAGE, debugging_method, False, - False) + False, + {}) def test_backtest_auto_updates_outdated_python_pycharm_debug_config() -> None: @@ -649,3 +657,23 @@ def test_backtest_adds_python_libraries_path_to_lean_config() -> None: expected_library_path = (Path("/") / library_path.relative_to(lean_cli_root_dir)).as_posix() assert expected_library_path in lean_config.get('python-additional-paths') + + +def test_backtest_calls_lean_runner_with_extra_docker_config() -> None: + create_fake_lean_cli_directory() + + result = CliRunner().invoke(lean, ["backtest", "Python Project", + "--extra-docker-config", + '{"device_requests": [{"count": -1, "capabilities": [["compute"]]}]}']) + + assert result.exit_code == 0 + + container.lean_runner.run_lean.assert_called_once_with(mock.ANY, + "backtesting", + Path("Python Project/main.py").resolve(), + mock.ANY, + ENGINE_IMAGE, + None, + False, + False, + {"device_requests": [{"count": -1, "capabilities": [["compute"]]}]}) diff --git a/tests/commands/test_live.py b/tests/commands/test_live.py index 1e8c1747..7eddebb9 100644 --- a/tests/commands/test_live.py +++ b/tests/commands/test_live.py @@ -84,7 +84,34 @@ def test_live_calls_lean_runner_with_correct_algorithm_file() -> None: ENGINE_IMAGE, None, False, - False) + False, + {}) + + +def test_live_calls_lean_runner_with_extra_docker_config() -> None: + # TODO: currently it is not using the live-paper environment + create_fake_lean_cli_directory() + create_fake_environment("live-paper", True) + + result = CliRunner().invoke(lean, ["live", "Python Project", + "--environment", + "live-paper", + "--extra-docker-config", + '{"device_requests": [{"count": -1, "capabilities": [["compute"]]}]}']) + + traceback.print_exception(*result.exc_info) + + assert result.exit_code == 0 + + container.lean_runner.run_lean.assert_called_once_with(mock.ANY, + "live-paper", + Path("Python Project/main.py").resolve(), + mock.ANY, + ENGINE_IMAGE, + None, + False, + False, + {"device_requests": [{"count": -1, "capabilities": [["compute"]]}]}) def test_live_aborts_when_environment_does_not_exist() -> None: @@ -159,7 +186,8 @@ def test_live_calls_lean_runner_with_release_mode() -> None: ENGINE_IMAGE, None, True, - False) + False, + {}) def test_live_calls_lean_runner_with_detach() -> None: @@ -178,7 +206,8 @@ def test_live_calls_lean_runner_with_detach() -> None: ENGINE_IMAGE, None, False, - True) + True, + {}) def test_live_aborts_when_project_does_not_exist() -> None: @@ -366,7 +395,8 @@ def test_live_calls_lean_runner_with_data_provider(data_provider: str) -> None: ENGINE_IMAGE, None, False, - False) + False, + {}) @pytest.mark.parametrize("brokerage", brokerage_required_options.keys() - ["Paper Trading"]) @@ -468,7 +498,8 @@ def test_live_non_interactive_do_not_store_non_persistent_properties_in_lean_con ENGINE_IMAGE, None, False, - False) + False, + {}) config = container.lean_config_manager.get_lean_config() if brokerage in brokerage_required_options_not_persistently_save_in_lean_config: @@ -509,7 +540,8 @@ def test_live_non_interactive_calls_run_lean_when_all_options_given(brokerage: s ENGINE_IMAGE, None, False, - False) + False, + {}) @pytest.mark.parametrize("brokerage,data_feed1,data_feed2",[(brokerage, *data_feeds) for brokerage, data_feeds in itertools.product(brokerage_required_options.keys(), itertools.combinations(data_feed_required_options.keys(), 2))]) @@ -548,7 +580,8 @@ def test_live_non_interactive_calls_run_lean_when_all_options_given_with_multipl ENGINE_IMAGE, None, False, - False) + False, + {}) @pytest.mark.parametrize("brokerage", brokerage_required_options.keys() - ["Paper Trading"]) @@ -606,7 +639,8 @@ def test_live_non_interactive_falls_back_to_lean_config_for_brokerage_settings(b ENGINE_IMAGE, None, False, - False) + False, + {}) @pytest.mark.parametrize("data_feed", data_feed_required_options.keys()) @@ -653,7 +687,8 @@ def test_live_non_interactive_falls_back_to_lean_config_for_data_feed_settings(d ENGINE_IMAGE, None, False, - False) + False, + {}) @pytest.mark.parametrize("data_feed1,data_feed2", itertools.combinations(data_feed_required_options.keys(), 2)) @@ -702,7 +737,8 @@ def test_live_non_interactive_falls_back_to_lean_config_for_multiple_data_feed_s ENGINE_IMAGE, None, False, - False) + False, + {}) def test_live_forces_update_when_update_option_given() -> None: @@ -721,7 +757,8 @@ def test_live_forces_update_when_update_option_given() -> None: ENGINE_IMAGE, None, False, - False) + False, + {}) def test_live_passes_custom_image_to_lean_runner_when_set_in_config() -> None: @@ -741,7 +778,8 @@ def test_live_passes_custom_image_to_lean_runner_when_set_in_config() -> None: DockerImage(name="custom/lean", tag="123"), None, False, - False) + False, + {}) def test_live_passes_custom_image_to_lean_runner_when_given_as_option() -> None: @@ -762,7 +800,8 @@ def test_live_passes_custom_image_to_lean_runner_when_given_as_option() -> None: DockerImage(name="custom/lean", tag="456"), None, False, - False) + False, + {}) @pytest.mark.parametrize("python_venv", ["Custom-venv", diff --git a/tests/commands/test_optimize.py b/tests/commands/test_optimize.py index 08a25874..62ee0919 100644 --- a/tests/commands/test_optimize.py +++ b/tests/commands/test_optimize.py @@ -753,3 +753,29 @@ def run_image_for_estimate(image: DockerImage, **kwargs) -> bool: args, kwargs = docker_manager.run_image.call_args assert any(command == 'dotnet QuantConnect.Optimizer.Launcher.dll --estimate' for command in kwargs["commands"]) + + +def test_optimize_runs_lean_container_with_extra_docker_config() -> None: + import docker.types + + create_fake_lean_cli_directory() + + docker_manager = mock.MagicMock() + docker_manager.run_image.side_effect = run_image + container.initialize(docker_manager=docker_manager) + container.optimizer_config_manager = _get_optimizer_config_manager_mock() + + Storage(str(Path.cwd() / "Python Project" / "config.json")).set("parameters", {"param1": "1"}) + + result = CliRunner().invoke(lean, ["optimize", "Python Project", + "--extra-docker-config", + '{"device_requests": [{"count": -1, "capabilities": [["compute"]]}]}']) + + assert result.exit_code == 0 + + docker_manager.run_image.assert_called_once() + args, kwargs = docker_manager.run_image.call_args + + assert args[0] == ENGINE_IMAGE + assert "device_requests" in kwargs + assert kwargs["device_requests"] == [docker.types.DeviceRequest(count=-1, capabilities=[["compute"]])] diff --git a/tests/commands/test_research.py b/tests/commands/test_research.py index 09a2b12d..a10541f5 100644 --- a/tests/commands/test_research.py +++ b/tests/commands/test_research.py @@ -228,3 +228,25 @@ def test_research_runs_custom_image_when_given_as_option() -> None: args, kwargs = docker_manager.run_image.call_args assert args[0] == DockerImage(name="custom/research", tag="456") + + +def test_optimize_runs_lean_container_with_extra_docker_config() -> None: + import docker.types + + create_fake_lean_cli_directory() + + docker_manager = mock.MagicMock() + container.initialize(docker_manager) + + result = CliRunner().invoke(lean, ["research", "Python Project", + "--extra-docker-config", + '{"device_requests": [{"count": -1, "capabilities": [["compute"]]}]}']) + + assert result.exit_code == 0 + + docker_manager.run_image.assert_called_once() + args, kwargs = docker_manager.run_image.call_args + + assert args[0] == RESEARCH_IMAGE + assert "device_requests" in kwargs + assert kwargs["device_requests"] == [docker.types.DeviceRequest(count=-1, capabilities=[["compute"]])] diff --git a/tests/components/docker/test_lean_runner.py b/tests/components/docker/test_lean_runner.py index 6d7594b2..59b3ad08 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,51 @@ 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_lean_runner_parses_device_requests_from_extra_docker_configs() -> None: + create_fake_lean_cli_directory() + + run_options = {} + LeanRunner.parse_extra_docker_config(run_options, + {"device_requests": [{"count": -1, "capabilities": [["compute"]]}]}) + + assert "device_requests" in run_options + + device_requests = run_options["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] == "compute") + assert device_request.driver == "" + assert device_request.device_ids == [] + assert device_request.options == {} + + +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": [["compute"]]}]}) + + docker_manager.run_image.assert_called_once() + args, kwargs = docker_manager.run_image.call_args + + assert "device_requests" in kwargs + assert kwargs["device_requests"] == [docker.types.DeviceRequest(count=-1, capabilities=[["compute"]])]