feat: UN-2166 Passed optional container name from backend to runner (#1158)

* UN-2166 Passed optional container name from backend to runner

* Removed tool-sandbox from lockfile automation script - since core is a private lib that's added, the lockfile can't be resolved and pushed

* minor: Addressed review comments for renaming a variable

Author: chandrasekharan-zipstack

backend/workflow_manager/workflow_v2/execution.py

@@ -231,9 +231,7 @@ def build(self) -> None:
             )
             raise WorkflowExecutionError(self.compilation_result["problems"][0])
 
-    def execute(
-        self, file_execution_id: str, file_name: str, single_step: bool = False
-    ) -> None:
+    def execute(self, file_execution_id: str, single_step: bool = False) -> None:
         execution_type = ExecutionType.COMPLETE
         if single_step:
             execution_type = ExecutionType.STEP
@@ -256,7 +254,6 @@ def execute(
         try:
             self.execute_workflow(
                 file_execution_id=file_execution_id,
-                file_name=file_name,
                 execution_type=execution_type,
             )
             end_time = time.time()
@@ -356,7 +353,13 @@ def execute_input_file(
         )
         workflow_file_execution.update_status(ExecutionStatus.EXECUTING)
 
-        self.execute(file_execution_id, file_name, single_step)
+        logger.info(
+            f"Running execution: '{self.execution_id}',  "
+            f"file_execution_id: '{file_execution_id}', "
+            f"file '{file_name}'"
+        )
+
+        self.execute(file_execution_id, single_step)
         self.publish_log(f"Tool executed successfully for '{file_name}'")
         self._handle_execution_type(execution_type)
 

docker/scripts/pdm-lock-gen/pdm-lock.sh

@@ -79,7 +79,6 @@ directories=(
     "platform-service"
     "x2text-service"
     "unstract/connectors"
-    "unstract/tool-sandbox"
 )
 
 # If directories are passed as arguments, override the default

runner/src/unstract/runner/clients/docker.py

@@ -158,17 +158,24 @@ def get_image(self) -> str:
     def get_container_run_config(
         self,
         command: list[str],
-        run_id: str,
+        file_execution_id: str,
+        container_name: Optional[str] = None,
         envs: Optional[dict[str, Any]] = None,
         auto_remove: bool = False,
     ) -> dict[str, Any]:
         if envs is None:
             envs = {}
         mounts = []
+
+        if not container_name:
+            container_name = UnstractUtils.build_tool_container_name(
+                tool_image=self.image_name,
+                tool_version=self.image_tag,
+                file_execution_id=file_execution_id,
+            )
+
         return {
-            "name": UnstractUtils.build_tool_container_name(
-                tool_image=self.image_name, tool_version=self.image_tag, run_id=run_id
-            ),
+            "name": container_name,
             "image": self.get_image(),
             "command": command,
             "detach": True,

runner/src/unstract/runner/clients/interface.py

@@ -62,7 +62,8 @@ def get_image(self) -> str:
     def get_container_run_config(
         self,
         command: list[str],
-        run_id: str,
+        file_execution_id: str,
+        container_name: Optional[str] = None,
         envs: Optional[dict[str, Any]] = None,
         auto_remove: bool = False,
     ) -> dict[str, Any]:

runner/src/unstract/runner/clients/test_docker.py

@@ -109,19 +109,21 @@ def test_get_image(docker_client, mocker):
 def test_get_container_run_config(docker_client, mocker):
     """Test the get_container_run_config method."""
     command = ["echo", "hello"]
-    run_id = "run123"
+    file_execution_id = "run123"
 
     mocker.patch.object(docker_client, "_Client__image_exists", return_value=True)
     mocker_normalize = mocker.patch(
         "unstract.core.utilities.UnstractUtils.build_tool_container_name",
         return_value="test-image",
     )
     config = docker_client.get_container_run_config(
-        command, run_id, envs={"KEY": "VALUE"}, auto_remove=True
+        command, file_execution_id, envs={"KEY": "VALUE"}, auto_remove=True
     )
 
     mocker_normalize.assert_called_once_with(
-        tool_image="test-image", tool_version="latest", run_id=run_id
+        tool_image="test-image",
+        tool_version="latest",
+        file_execution_id=file_execution_id,
     )
     assert config["name"] == "test-image"
     assert config["image"] == "test-image:latest"
@@ -134,17 +136,21 @@ def test_get_container_run_config_without_mount(docker_client, mocker):
     """Test the get_container_run_config method."""
     os.environ[Env.EXECUTION_DATA_DIR] = "/source"
     command = ["echo", "hello"]
-    run_id = "run123"
+    file_execution_id = "run123"
 
     mocker.patch.object(docker_client, "_Client__image_exists", return_value=True)
     mocker_normalize = mocker.patch(
         "unstract.core.utilities.UnstractUtils.build_tool_container_name",
         return_value="test-image",
     )
-    config = docker_client.get_container_run_config(command, run_id, auto_remove=True)
+    config = docker_client.get_container_run_config(
+        command, file_execution_id, auto_remove=True
+    )
 
     mocker_normalize.assert_called_once_with(
-        tool_image="test-image", tool_version="latest", run_id=run_id
+        tool_image="test-image",
+        tool_version="latest",
+        file_execution_id=file_execution_id,
     )
     assert config["name"] == "test-image"
     assert config["image"] == "test-image:latest"

runner/src/unstract/runner/main.py

@@ -28,17 +28,19 @@ def run_container() -> Optional[Any]:
     organization_id = data["organization_id"]
     workflow_id = data["workflow_id"]
     execution_id = data["execution_id"]
-    run_id = data["run_id"]
+    file_execution_id = data["file_execution_id"]
+    container_name = data["container_name"]
     settings = data["settings"]
     envs = data["envs"]
     messaging_channel = data["messaging_channel"]
 
     runner = UnstractRunner(image_name, image_tag, app)
     result = runner.run_container(
+        container_name=container_name,
         organization_id=organization_id,
         workflow_id=workflow_id,
         execution_id=execution_id,
-        run_id=run_id,
+        file_execution_id=file_execution_id,
         settings=settings,
         envs=envs,
         messaging_channel=messaging_channel,

runner/src/unstract/runner/runner.py

@@ -39,7 +39,7 @@ def stream_logs(
         tool_instance_id: str,
         execution_id: str,
         organization_id: str,
-        run_id: str,
+        file_execution_id: str,
         channel: Optional[str] = None,
     ) -> None:
         for line in container.logs(follow=True):
@@ -51,7 +51,7 @@ def stream_logs(
                 channel=channel,
                 execution_id=execution_id,
                 organization_id=organization_id,
-                run_id=run_id,
+                file_execution_id=file_execution_id,
             )
 
     def get_valid_log_message(self, log_message: str) -> Optional[dict[str, Any]]:
@@ -76,7 +76,7 @@ def process_log_message(
         tool_instance_id: str,
         execution_id: str,
         organization_id: str,
-        run_id: str,
+        file_execution_id: str,
         channel: Optional[str] = None,
     ) -> Optional[dict[str, Any]]:
         log_dict = self.get_valid_log_message(log_message)
@@ -99,7 +99,7 @@ def process_log_message(
             log_dict[LogFieldName.EXECUTION_ID] = execution_id
             log_dict[LogFieldName.ORGANIZATION_ID] = organization_id
             log_dict[LogFieldName.TIMESTAMP] = datetime.now(timezone.utc).timestamp()
-            log_dict[LogFieldName.FILE_EXECUTION_ID] = run_id
+            log_dict[LogFieldName.FILE_EXECUTION_ID] = file_execution_id
 
             # Publish to channel of socket io
             LogPublisher.publish(channel, log_dict)
@@ -139,7 +139,7 @@ def run_command(self, command: str) -> Optional[Any]:
         """
         command = command.upper()
         container_config = self.client.get_container_run_config(
-            command=["--command", command], run_id="", auto_remove=True
+            command=["--command", command], file_execution_id="", auto_remove=True
         )
         container = None
 
@@ -163,10 +163,11 @@ def run_container(
         organization_id: str,
         workflow_id: str,
         execution_id: str,
-        run_id: str,
+        file_execution_id: str,
         settings: dict[str, Any],
         envs: dict[str, Any],
         messaging_channel: Optional[str] = None,
+        container_name: Optional[str] = None,
     ) -> Optional[Any]:
         """RUN container With RUN Command.
 
@@ -200,7 +201,8 @@ def run_container(
                 "--log-level",
                 "DEBUG",
             ],
-            run_id=run_id,
+            file_execution_id=file_execution_id,
+            container_name=container_name,
             envs=envs,
         )
         # Add labels to container for logging with Loki.
@@ -217,7 +219,7 @@ def run_container(
         try:
             self.logger.info(
                 f"Execution ID: {execution_id}, running docker "
-                f"container: {container_config.get('name')}"
+                f"container: {container_name}"
             )
             container: ContainerInterface = self.client.run_container(container_config)
             tool_instance_id = str(settings.get(ToolKey.TOOL_INSTANCE_ID))
@@ -228,7 +230,7 @@ def run_container(
                 channel=messaging_channel,
                 execution_id=execution_id,
                 organization_id=organization_id,
-                run_id=run_id,
+                file_execution_id=file_execution_id,
             )
         except ToolRunException as te:
             self.logger.error(

unstract/core/src/unstract/core/utilities.py

@@ -30,12 +30,21 @@ def get_env(env_key: str, default: Optional[str] = None, raise_err=False) -> str
 
     @staticmethod
     def build_tool_container_name(
-        tool_image: str, tool_version: str, run_id: str
+        tool_image: str,
+        tool_version: str,
+        file_execution_id: str,
+        retry_count: Optional[int] = None,
     ) -> str:
         tool_name = tool_image.split("/")[-1]
-        # TODO: Add execution attempt to better track instead of uuid
-        short_uuid = uuid.uuid4().hex[:6]  # To avoid duplicate name collision
-        container_name = f"{tool_name}-{tool_version}-{short_uuid}-{run_id}"
+        # To avoid duplicate name collision
+        if retry_count:
+            unique_suffix = retry_count
+        else:
+            unique_suffix = uuid.uuid4().hex[:6]
+
+        container_name = (
+            f"{tool_name}-{tool_version}-{unique_suffix}-{file_execution_id}"
+        )
 
         # To support limits of container clients like K8s
         if len(container_name) > 63:

unstract/tool-sandbox/pyproject.toml

@@ -10,7 +10,18 @@ authors = [
     {name = "Zipstack Inc.", email = "devsupport@zipstack.com"},
 ]
 dependencies = [
-    "requests==2.31.0"
+    "requests==2.31.0",
+    # ! IMPORTANT!
+    # Local dependencies usually need to be added as:
+    # https://pdm-project.org/latest/usage/dependency/#local-dependencies
+    #
+    # However, local dependencies which are not direct depedency of main project
+    # appear as absolute paths in pdm.lock of main project, making it impossible
+    # to check in the lock file.
+    #
+    # Hence instead, add the dependencies without version constraints where the
+    # assumption is they are added as direct dependencies in main project itself.
+    "unstract-core",
 ]
 requires-python = ">=3.9"
 readme = "README.md"

unstract/tool-sandbox/src/unstract/tool_sandbox/helper.py

@@ -6,6 +6,8 @@
 import requests
 from unstract.tool_sandbox.constants import UnstractRunner
 
+from unstract.core.utilities import UnstractUtils
+
 logger = logging.getLogger(__name__)
 
 
@@ -69,10 +71,11 @@ def make_get_request(
 
     def call_tool_handler(
         self,
-        run_id: str,
+        file_execution_id: str,
         image_name: str,
         image_tag: str,
         settings: dict[str, Any],
+        retry_count: Optional[int] = None,
     ) -> Optional[dict[str, Any]]:
         """Calling unstract runner to run the required tool.
 
@@ -87,10 +90,7 @@ def call_tool_handler(
         """
         url = f"{self.base_url}{UnstractRunner.RUN_API_ENDPOINT}"
         data = self.create_tool_request_data(
-            run_id,
-            image_name,
-            image_tag,
-            settings,
+            file_execution_id, image_name, image_tag, settings, retry_count
         )
 
         response = requests.post(url, json=data)
@@ -110,18 +110,26 @@ def call_tool_handler(
 
     def create_tool_request_data(
         self,
-        run_id: str,
+        file_execution_id: str,
         image_name: str,
         image_tag: str,
         settings: dict[str, Any],
+        retry_count: Optional[int] = None,
     ) -> dict[str, Any]:
+        container_name = UnstractUtils.build_tool_container_name(
+            tool_image=image_name,
+            tool_version=image_tag,
+            file_execution_id=file_execution_id,
+            retry_count=retry_count,
+        )
         data = {
             "image_name": image_name,
             "image_tag": image_tag,
             "organization_id": self.organization_id,
             "workflow_id": self.workflow_id,
             "execution_id": self.execution_id,
-            "run_id": run_id,
+            "file_execution_id": file_execution_id,
+            "container_name": container_name,
             "settings": settings,
             "envs": self.envs,
             "messaging_channel": self.messaging_channel,