refactor: reuse manage_docker_image

QuantConnect · Sep 25, 2024 · 746c23b · 746c23b
1 parent 4789169
commit 746c23b
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 61 deletions.
diff --git a/lean/commands/backtest.py b/lean/commands/backtest.py
@@ -16,7 +16,7 @@
 from click import command, option, argument, Choice
 
 from lean.click import LeanCommand, PathParameter
-from lean.constants import DEFAULT_ENGINE_IMAGE, LEAN_ROOT_PATH, CONTAINER_LABEL_LEAN_VERSION_NAME
+from lean.constants import DEFAULT_ENGINE_IMAGE, LEAN_ROOT_PATH
 from lean.container import container, Logger
 from lean.models.utils import DebuggingMethod
 from lean.models.cli import cli_data_downloaders, cli_addon_modules
@@ -359,16 +359,8 @@ def backtest(project: Path,
     organization_id = container.organization_manager.try_get_working_organization_id()
     paths_to_mount = None
 
-    cli_config_manager = container.cli_config_manager
-    project_config_manager = container.project_config_manager
-
-    project_config = project_config_manager.get_project_config(algorithm_file.parent)
-    engine_image = cli_config_manager.get_engine_image(image or project_config.get("engine-image", None))
-
-    container.update_manager.pull_docker_image_if_necessary(engine_image, update, no_update)
-
-    container_module_version = container.docker_manager.get_image_label(engine_image,
-                                                                        CONTAINER_LABEL_LEAN_VERSION_NAME, None)
+    engine_image, container_module_version, project_config = container.manage_docker_image(image, update, no_update,
+                                                                                           algorithm_file)
 
     if data_provider_historical is not None:
         data_provider = non_interactive_config_build_for_name(lean_config, data_provider_historical,
@@ -379,9 +371,6 @@ def backtest(project: Path,
 
     lean_config_manager.configure_data_purchase_limit(lean_config, data_purchase_limit)
 
-    if str(engine_image) != DEFAULT_ENGINE_IMAGE:
-        logger.warn(f'A custom engine image: "{engine_image}" is being used!')
-
     if not output.exists():
         output.mkdir(parents=True)
 

diff --git a/lean/commands/data/download.py b/lean/commands/data/download.py
@@ -18,7 +18,7 @@
 from click import command, option, confirm, pass_context, Context, Choice, prompt
 from lean.click import LeanCommand, ensure_options
 from lean.components.util.json_modules_handler import config_build_for_name
-from lean.constants import DEFAULT_ENGINE_IMAGE, CONTAINER_LABEL_LEAN_VERSION_NAME
+from lean.constants import DEFAULT_ENGINE_IMAGE
 from lean.container import container
 from lean.models.api import QCDataInformation, QCDataVendor, QCFullOrganization, QCDatasetDelivery, QCResolution, QCSecurityType, QCDataType
 from lean.models.click_options import get_configs_for_options, options_from_json
@@ -675,16 +675,7 @@ def download(ctx: Context,
         logger = container.logger
         lean_config = container.lean_config_manager.get_complete_lean_config(None, None, None)
 
-        engine_image = container.cli_config_manager.get_engine_image(image)
-
-        if str(engine_image) != DEFAULT_ENGINE_IMAGE:
-            # Custom engine image should not be updated.
-            logger.warn(f'A custom engine image: "{engine_image}" is being used!')
-
-        container.update_manager.pull_docker_image_if_necessary(engine_image, update, no_update)
-
-        container_module_version = container.docker_manager.get_image_label(engine_image,
-                                                                            CONTAINER_LABEL_LEAN_VERSION_NAME, None)
+        engine_image, container_module_version = container.manage_docker_image(image, update, no_update)
 
         data_downloader_provider = config_build_for_name(lean_config, data_downloader_provider.get_name(),
                                                          cli_data_downloaders, kwargs, logger, interactive=True)

diff --git a/lean/commands/data/generate.py b/lean/commands/data/generate.py
@@ -208,9 +208,7 @@ def generate(start: datetime,
         }
     }
 
-    engine_image = container.cli_config_manager.get_engine_image(image)
-
-    container.update_manager.pull_docker_image_if_necessary(engine_image, update)
+    engine_image, container_module_version = container.manage_docker_image(image, update, no_update=False)
 
     success = container.docker_manager.run_image(engine_image, **run_options)
     if not success:

diff --git a/lean/commands/live/deploy.py b/lean/commands/live/deploy.py
@@ -16,7 +16,7 @@
 from click import option, argument, Choice
 from lean.click import LeanCommand, PathParameter
 from lean.components.util.name_rename import rename_internal_config_to_user_friendly_format
-from lean.constants import DEFAULT_ENGINE_IMAGE, CONTAINER_LABEL_LEAN_VERSION_NAME
+from lean.constants import DEFAULT_ENGINE_IMAGE
 from lean.container import container
 from lean.models.cli import (cli_brokerages, cli_data_queue_handlers, cli_data_downloaders,
                              cli_addon_modules, cli_history_provider)
@@ -271,16 +271,8 @@ def deploy(project: Path,
                                                                  kwargs, logger, interactive=True,
                                                                  environment_name=environment_name))
 
-    project_config_manager = container.project_config_manager
-    cli_config_manager = container.cli_config_manager
-
-    project_config = project_config_manager.get_project_config(algorithm_file.parent)
-    engine_image = cli_config_manager.get_engine_image(image or project_config.get("engine-image", None))
-
-    container.update_manager.pull_docker_image_if_necessary(engine_image, update, no_update)
-
-    container_module_version = container.docker_manager.get_image_label(engine_image,
-                                                                        CONTAINER_LABEL_LEAN_VERSION_NAME, None)
+    engine_image, container_module_version, project_config = container.manage_docker_image(image, update, no_update,
+                                                                                           algorithm_file)
 
     organization_id = container.organization_manager.try_get_working_organization_id()
     paths_to_mount = {}
@@ -298,8 +290,8 @@ def deploy(project: Path,
     if python_venv is not None and python_venv != "":
         lean_config["python-venv"] = f'{"/" if python_venv[0] != "/" else ""}{python_venv}'
 
-    cash_balance_option, holdings_option, last_cash, last_holdings = get_last_portfolio_cash_holdings(container.api_client, brokerage_instance,
-                                                                                                      project_config.get("cloud-id", None), project)
+    cash_balance_option, holdings_option, last_cash, last_holdings = get_last_portfolio_cash_holdings(
+        container.api_client, brokerage_instance, project_config.get("cloud-id", None), project)
 
     # We cannot create the output directory before calling get_last_portfolio_holdings, since then the most recently
     # deployment would be always the local one (it has the current time in its name), and we would never be able to
@@ -336,9 +328,6 @@ def deploy(project: Path,
             "AveragePrice": holding["averagePrice"]
         } for holding in live_holdings]
 
-    if str(engine_image) != DEFAULT_ENGINE_IMAGE:
-        logger.warn(f'A custom engine image: "{engine_image}" is being used!')
-
     # Set extra config
     given_algorithm_id = None
     for key, value in extra_config:

diff --git a/lean/commands/optimize.py b/lean/commands/optimize.py
@@ -19,7 +19,7 @@
 
 from lean.click import LeanCommand, PathParameter, ensure_options
 from lean.components.docker.lean_runner import LeanRunner
-from lean.constants import DEFAULT_ENGINE_IMAGE, CONTAINER_LABEL_LEAN_VERSION_NAME
+from lean.constants import DEFAULT_ENGINE_IMAGE
 from lean.container import container
 from lean.models.api import QCParameter, QCBacktest
 from lean.models.click_options import options_from_json, get_configs_for_options
@@ -227,6 +227,8 @@ def optimize(project: Path,
     if optimizer_config is not None and strategy is not None:
         raise RuntimeError("--optimizer-config and --strategy are mutually exclusive")
 
+    engine_image, container_module_version, project_config = container.manage_docker_image(image, update, no_update,
+                                                                                           algorithm_file)
     if optimizer_config is not None:
         config = loads(optimizer_config.read_text(encoding="utf-8"))
 
@@ -242,8 +244,6 @@ def optimize(project: Path,
         optimization_parameters = optimizer_config_manager.parse_parameters(parameter)
         optimization_constraints = optimizer_config_manager.parse_constraints(constraint)
     else:
-        project_config_manager = container.project_config_manager
-        project_config = project_config_manager.get_project_config(algorithm_file.parent)
         project_parameters = [QCParameter(key=k, value=v) for k, v in project_config.get("parameters", {}).items()]
 
         if len(project_parameters) == 0:
@@ -286,17 +286,8 @@ def optimize(project: Path,
     with config_path.open("w+", encoding="utf-8") as file:
         file.write(dumps(config, indent=4) + "\n")
 
-    project_config_manager = container.project_config_manager
-    cli_config_manager = container.cli_config_manager
-
-    project_config = project_config_manager.get_project_config(algorithm_file.parent)
-    engine_image = cli_config_manager.get_engine_image(image or project_config.get("engine-image", None))
-
     logger = container.logger
 
-    if str(engine_image) != DEFAULT_ENGINE_IMAGE:
-        logger.warn(f'A custom engine image: "{engine_image}" is being used!')
-
     lean_config_manager = container.lean_config_manager
     lean_config = lean_config_manager.get_complete_lean_config(environment_name, algorithm_file, None)
 
@@ -307,9 +298,6 @@ def optimize(project: Path,
 
     paths_to_mount = None
 
-    container_module_version = container.docker_manager.get_image_label(engine_image,
-                                                                        CONTAINER_LABEL_LEAN_VERSION_NAME, None)
-
     if data_provider_historical is not None:
         data_provider = non_interactive_config_build_for_name(lean_config, data_provider_historical,
                                                               cli_data_downloaders, kwargs, logger, environment_name)
@@ -342,8 +330,6 @@ def optimize(project: Path,
     build_and_configure_modules(addon_module, cli_addon_modules, organization_id, lean_config,
                                 kwargs, logger, environment_name, container_module_version)
 
-    container.update_manager.pull_docker_image_if_necessary(engine_image, update, no_update)
-
     run_options = lean_runner.get_basic_docker_config(lean_config, algorithm_file, output, None, release, should_detach,
                                                       engine_image, paths_to_mount)