Fix: not pulling docker image (#503)

Romazes · web-flow · commit 2456258699f3 · 2024-09-25T20:30:35.000-03:00
* fix: change position of pull docker image

* feat: generic manage_docker_image

* refactor: reuse manage_docker_image

* feat: add research in manage_docker_image
refactor: use project_directory instead of AlgorithmFile Path

* fix: missed project_config in download/generate
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,14 +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_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.parent)
 
     if data_provider_historical is not None:
         data_provider = non_interactive_config_build_for_name(lean_config, data_provider_historical,
@@ -377,11 +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!')
-
-    container.update_manager.pull_docker_image_if_necessary(engine_image, update, no_update)
-
     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, project_config = 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, project_config = 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,14 +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_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.parent)
 
     organization_id = container.organization_manager.try_get_working_organization_id()
     paths_to_mount = {}
@@ -291,15 +285,13 @@ def deploy(project: Path,
         raise MoreInfoError(f"The '{environment_name}' is not a live trading environment (live-mode is set to false)",
                             "https://www.lean.io/docs/v2/lean-cli/live-trading/brokerages/quantconnect-paper-trading")
 
-    container.update_manager.pull_docker_image_if_necessary(engine_image, update, no_update)
-
     _start_iqconnect_if_necessary(lean_config, environment_name)
 
     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.parent)
     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)
 
diff --git a/lean/commands/report.py b/lean/commands/report.py
@@ -258,17 +258,8 @@ def report(backtest_results: Optional[Path],
                                            type="bind",
                                            read_only=True))
 
-    cli_config_manager = container.cli_config_manager
-    engine_image_override = image
-
-    if engine_image_override is None and project_directory is not None:
-        project_config_manager = container.project_config_manager
-        project_config = project_config_manager.get_project_config(project_directory)
-        engine_image_override = project_config.get("engine-image", None)
-
-    engine_image = cli_config_manager.get_engine_image(engine_image_override)
-
-    container.update_manager.pull_docker_image_if_necessary(engine_image, update)
+    engine_image, container_module_version, project_config = container.manage_docker_image(image, update, False,
+                                                                                           project_directory)
 
     success = container.docker_manager.run_image(engine_image, **run_options)
     if not success:
diff --git a/lean/commands/research.py b/lean/commands/research.py
@@ -16,7 +16,7 @@
 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, CONTAINER_LABEL_LEAN_VERSION_NAME
+from lean.constants import DEFAULT_RESEARCH_IMAGE, LEAN_ROOT_PATH
 from lean.container import container
 from lean.models.cli import cli_data_downloaders
 from lean.components.util.name_extraction import convert_to_class_name
@@ -118,19 +118,9 @@ def research(project: Path,
     if download_data:
         data_provider_historical = "QuantConnect"
 
-    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)
-    research_image = cli_config_manager.get_research_image(image or project_config.get("research-image", None))
-
-    container.update_manager.pull_docker_image_if_necessary(research_image, update, no_update)
-
-    container_module_version = container.docker_manager.get_image_label(research_image,
-                                                                        CONTAINER_LABEL_LEAN_VERSION_NAME, None)
-
-    if str(research_image) != DEFAULT_RESEARCH_IMAGE:
-        logger.warn(f'A custom research image: "{research_image}" is being used!')
+    research_image, container_module_version, project_config = container.manage_docker_image(image, update, no_update,
+                                                                                             algorithm_file.parent,
+                                                                                             False)
 
     paths_to_mount = None
 
diff --git a/lean/container.py b/lean/container.py
@@ -11,7 +11,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Union, Any
+from pathlib import Path
+from typing import Union, Any, Optional, Tuple
 
 from lean.components.api.api_client import APIClient
 from lean.components.config.lean_config_manager import LeanConfigManager
@@ -41,12 +42,16 @@
 from lean.components.util.temp_manager import TempManager
 from lean.components.util.update_manager import UpdateManager
 from lean.components.util.xml_manager import XMLManager
-from lean.constants import CACHE_PATH, CREDENTIALS_CONFIG_PATH, GENERAL_CONFIG_PATH
+from lean.constants import CACHE_PATH, CREDENTIALS_CONFIG_PATH, GENERAL_CONFIG_PATH, DEFAULT_RESEARCH_IMAGE
+from lean.constants import DEFAULT_ENGINE_IMAGE, CONTAINER_LABEL_LEAN_VERSION_NAME
+from lean.models.docker import DockerImage
 
 
 class Container:
 
     def __init__(self):
+        self.project_config_manager = None
+        self.cli_config_manager = None
         self.initialize()
 
     def initialize(self,
@@ -161,6 +166,47 @@ def initialize(self,
 
         self.update_manager = UpdateManager(self.logger, self.http_client, self.cache_storage, self.docker_manager)
 
+    def manage_docker_image(self, image: Optional[str], update: bool, no_update: bool,
+                            project_directory: Path = None,
+                            is_engine_image: bool = True) -> Tuple[DockerImage, str, Optional[Storage]]:
+        """
+        Manages the Docker image for the LEAN engine by:
+        1. Retrieving the engine image from the provided image or project config.
+        2. Pulling the image if necessary based on the update flags.
+        3. Logging a warning if a custom image is used.
+
+        :param project_directory: Path to the project directory, used to get the project configuration.
+        :param image: Optional custom Docker image. Defaults to the project configuration if not provided.
+        :param update: Whether to update the Docker image.
+        :param no_update: Whether to skip updating the Docker image.
+        :param is_engine_image: True to manage the 'engine-image', False to manage the 'research-image'.
+        :return: A tuple containing the engine image, its version label, and the project configuration.
+        """
+
+        project_config = None
+        image_project_config = None
+        image_type_name = "engine-image" if is_engine_image else "research-image"
+        if project_directory:
+            project_config = self.project_config_manager.get_project_config(project_directory)
+            image_project_config = project_config.get(image_type_name, None)
+
+        if is_engine_image:
+            engine_image = self.cli_config_manager.get_engine_image(image or image_project_config)
+        else:
+            engine_image = self.cli_config_manager.get_research_image(image or image_project_config)
+
+        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
+        )
+
+        default_image_name = DEFAULT_ENGINE_IMAGE if is_engine_image else DEFAULT_RESEARCH_IMAGE
+        if str(engine_image) != default_image_name:
+            self.logger.warn(f'A custom {image_type_name} image: "{engine_image}" is being used!')
+
+        return engine_image, container_module_version, project_config
+
 
 container = Container()
 container.data_downloader.update_database_files()

Original file line number	Diff line number	Diff line change
`@@ -208,9 +208,7 @@ def generate(start: datetime,`
`208`	`208`	`}`
`209`	`209`	`}`
`210`	`210`
`211`		`- engine_image = container.cli_config_manager.get_engine_image(image)`
`212`		`-`
`213`		`- container.update_manager.pull_docker_image_if_necessary(engine_image, update)`
	`211`	`+ engine_image, container_module_version, project_config = container.manage_docker_image(image, update, no_update=False)`
`214`	`212`
`215`	`213`	`success = container.docker_manager.run_image(engine_image, **run_options)`
`216`	`214`	`if not success:`