Support mounting volumes using extra docker configs

QuantConnect · Sep 19, 2023 · db4fcc6 · db4fcc6
1 parent 5b1b415
commit db4fcc6
Show file tree

Hide file tree

Showing 6 changed files with 104 additions and 9 deletions.
diff --git a/lean/components/docker/lean_runner.py b/lean/components/docker/lean_runner.py
@@ -774,6 +774,13 @@ def parse_extra_docker_config(run_options: Dict[str, Any], extra_docker_config:
         from docker.types import DeviceRequest
         # 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"] = [DeviceRequest(**device_request)
-                                              for device_request in extra_docker_config["device_requests"]]
+        if extra_docker_config is not None:
+            if "device_requests" in extra_docker_config:
+                run_options["device_requests"] = [DeviceRequest(**device_request)
+                                                  for device_request in extra_docker_config["device_requests"]]
+
+            if "volumes" in extra_docker_config:
+                volumes = run_options.get("volumes")
+                if not volumes:
+                    volumes = run_options["volumes"] = {}
+                volumes.update(extra_docker_config["volumes"])
diff --git a/tests/commands/test_backtest.py b/tests/commands/test_backtest.py
@@ -664,7 +664,8 @@ def test_backtest_calls_lean_runner_with_extra_docker_config() -> None:
 
     result = CliRunner().invoke(lean, ["backtest", "Python Project",
                                        "--extra-docker-config",
-                                       '{"device_requests": [{"count": -1, "capabilities": [["compute"]]}]}'])
+                                       '{"device_requests": [{"count": -1, "capabilities": [["compute"]]}],'
+                                       '"volumes": {"extra/path": {"bind": "/extra/path", "mode": "rw"}}}'])
 
     assert result.exit_code == 0
 
@@ -676,4 +677,11 @@ def test_backtest_calls_lean_runner_with_extra_docker_config() -> None:
                                                            None,
                                                            False,
                                                            False,
-                                                           {"device_requests": [{"count": -1, "capabilities": [["compute"]]}]})
+                                                           {
+                                                               "device_requests": [
+                                                                   {"count": -1, "capabilities": [["compute"]]}
+                                                               ],
+                                                               "volumes": {
+                                                                   "extra/path": {"bind": "/extra/path", "mode": "rw"}
+                                                               }
+                                                           })
diff --git a/tests/commands/test_live.py b/tests/commands/test_live.py
@@ -97,7 +97,8 @@ def test_live_calls_lean_runner_with_extra_docker_config() -> None:
                                        "--environment",
                                        "live-paper",
                                        "--extra-docker-config",
-                                       '{"device_requests": [{"count": -1, "capabilities": [["compute"]]}]}'])
+                                       '{"device_requests": [{"count": -1, "capabilities": [["compute"]]}],'
+                                       '"volumes": {"extra/path": {"bind": "/extra/path", "mode": "rw"}}}'])
 
     traceback.print_exception(*result.exc_info)
 
@@ -111,7 +112,14 @@ def test_live_calls_lean_runner_with_extra_docker_config() -> None:
                                                            None,
                                                            False,
                                                            False,
-                                                           {"device_requests": [{"count": -1, "capabilities": [["compute"]]}]})
+                                                           {
+                                                               "device_requests": [
+                                                                   {"count": -1, "capabilities": [["compute"]]}
+                                                               ],
+                                                               "volumes": {
+                                                                   "extra/path": {"bind": "/extra/path", "mode": "rw"}
+                                                               }
+                                                           })
 
 
 def test_live_aborts_when_environment_does_not_exist() -> None:

diff --git a/tests/commands/test_optimize.py b/tests/commands/test_optimize.py
@@ -769,13 +769,20 @@ def test_optimize_runs_lean_container_with_extra_docker_config() -> None:
 
     result = CliRunner().invoke(lean, ["optimize", "Python Project",
                                        "--extra-docker-config",
-                                       '{"device_requests": [{"count": -1, "capabilities": [["compute"]]}]}'])
+                                       '{"device_requests": [{"count": -1, "capabilities": [["compute"]]}],'
+                                       '"volumes": {"extra/path": {"bind": "/extra/path", "mode": "rw"}}}'])
 
     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"]])]
+
+    assert "volumes" in kwargs
+    volumes = kwargs["volumes"]
+    assert "extra/path" in volumes
+    assert volumes["extra/path"] == {"bind": "/extra/path", "mode": "rw"}
diff --git a/tests/commands/test_research.py b/tests/commands/test_research.py
@@ -240,13 +240,20 @@ def test_optimize_runs_lean_container_with_extra_docker_config() -> None:
 
     result = CliRunner().invoke(lean, ["research", "Python Project",
                                        "--extra-docker-config",
-                                       '{"device_requests": [{"count": -1, "capabilities": [["compute"]]}]}'])
+                                       '{"device_requests": [{"count": -1, "capabilities": [["compute"]]}],'
+                                       '"volumes": {"extra/path": {"bind": "/extra/path", "mode": "rw"}}}'])
 
     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"]])]
+
+    assert "volumes" in kwargs
+    volumes = kwargs["volumes"]
+    assert "extra/path" in volumes
+    assert volumes["extra/path"] == {"bind": "/extra/path", "mode": "rw"}
diff --git a/tests/components/docker/test_lean_runner.py b/tests/components/docker/test_lean_runner.py
@@ -607,6 +607,36 @@ def test_lean_runner_parses_device_requests_from_extra_docker_configs() -> None:
     assert device_request.options == {}
 
 
+def test_lean_runner_parses_volumes_from_extra_docker_configs() -> None:
+    create_fake_lean_cli_directory()
+
+    run_options = {
+        "volumes": {
+            "source/path": {
+                "bind": "/target/path",
+                "mode": "rw"
+            }
+        }
+    }
+    LeanRunner.parse_extra_docker_config(run_options,
+                                         {"volumes": {"extra/path": {"bind": "/extra/bound/path", "mode": "rw"}}})
+
+    assert "volumes" in run_options
+
+    volumes = run_options["volumes"]
+    assert len(volumes) == 2
+    assert "source/path" in volumes
+    assert "extra/path" in volumes
+
+    existing_volume = volumes["source/path"]
+    assert existing_volume["bind"] == "/target/path"
+    assert existing_volume["mode"] == "rw"
+
+    new_volume = volumes["extra/path"]
+    assert new_volume["bind"] == "/extra/bound/path"
+    assert new_volume["mode"] == "rw"
+
+
 def test_run_lean_passes_device_requests() -> None:
     create_fake_lean_cli_directory()
 
@@ -630,3 +660,31 @@ def test_run_lean_passes_device_requests() -> None:
 
     assert "device_requests" in kwargs
     assert kwargs["device_requests"] == [docker.types.DeviceRequest(count=-1, capabilities=[["compute"]])]
+
+
+def test_run_lean_passes_extra_volumes() -> 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={"volumes": {"extra/path": {"bind": "/extra/bound/path", "mode": "rw"}}})
+
+    docker_manager.run_image.assert_called_once()
+    args, kwargs = docker_manager.run_image.call_args
+
+    assert "volumes" in kwargs
+    volumes = kwargs["volumes"]
+    assert "extra/path" in volumes
+    assert volumes["extra/path"]["bind"] == "/extra/bound/path"
+    assert volumes["extra/path"]["mode"] == "rw"