Merge pull request #40 from SebastianScherer88/fix-k8s-service-naming…

…-issue-for-ddp-components

Fix k8s service naming issue for ddp components

sdk/bettmensch_ai/components/base_component.py

@@ -30,7 +30,8 @@ class BaseComponent(object):
     name: str = None
     func: Callable = None
     base_name: str = None
-    name: str = None
+    name: str = None  # this attribute will hold the name of the ArgoWorkflow
+    # Task and the Template of this component
     hera_template_kwargs: Dict = {}
     template_inputs: Dict[str, Union[InputParameter, InputArtifact]] = None
     template_outputs: Dict[str, Union[OutputParameter, OutputArtifact]] = None

sdk/bettmensch_ai/components/torch_ddp_component.py

@@ -2,7 +2,6 @@
 import inspect
 import textwrap
 from typing import Callable, Dict, List, Optional, Union
-from uuid import uuid4
 
 from bettmensch_ai.components.base_component import BaseComponent
 from bettmensch_ai.components.base_inline_script_runner import (
@@ -87,7 +86,7 @@ def generate_source(self, instance: Script) -> str:
         # add function definition and decoration with `torch_ddp`
         torch_ddp_decoration = [
             "\nfrom bettmensch_ai.components import torch_ddp\n",
-            "torch_ddp_decorator=torch_ddp()\n"
+            "torch_ddp_decorator=torch_ddp()\n",
             f"""torch_ddp_function=torch_ddp_decorator({
                 instance.source.__name__
             })\n""",
@@ -120,7 +119,6 @@ class TorchDDPComponent(BaseComponent):
     nproc_per_node: int
     service_templates: Dict[str, Callable] = None
     k8s_namespace: str = ARGO_NAMESPACE
-    k8s_service_name: str = ""
 
     # if no resources are specified, set minimal requirements derived from
     # testing the ddp example on K8s
@@ -165,11 +163,11 @@ def build_service_templates(self) -> Dict[str, Callable]:
         return {
             "create": create_torch_ddp_service_template(
                 component_base_name=self.base_name,
-                service_name=self.k8s_service_name,
+                component_task_name=self.name,
             ),
             "delete": delete_torch_ddp_service_template(
                 component_base_name=self.base_name,
-                service_name=self.k8s_service_name,
+                component_task_name=self.name,
             ),
         }
 
@@ -232,8 +230,7 @@ def build_script_decorator_kwargs(self, torch_node_rank: int) -> Dict:
             ),
             Env(
                 name=f"{LaunchConfigSettings.model_config['env_prefix']}rdzv_endpoint_url",  # noqa: E501
-                value=f"{self.k8s_service_name}.{self.k8s_namespace}"
-                + ".svc.cluster.local",
+                value=f"{self.name}-{{{{workflow.uid}}}}.{self.k8s_namespace}.svc.cluster.local",  # noqa: E501
             ),
             Env(
                 name=f"{LaunchConfigSettings.model_config['env_prefix']}rdzv_endpoint_port",  # noqa: E501
@@ -255,7 +252,7 @@ def build_script_decorator_kwargs(self, torch_node_rank: int) -> Dict:
         labels.update(
             {
                 "torch-node": torch_node_rank,
-                "torch-job": self.k8s_service_name,
+                "torch-job": self.name,
             }
         )
         script_decorator_kwargs["labels"] = labels
@@ -273,8 +270,6 @@ def build_hera_task_factory(self) -> List[Callable]:
                 node in the distributed torch run.
         """
 
-        self.k8s_service_name = f"{self.name}-{uuid4()}"
-
         # add torch run environment variables to script kwargs
         task_factory = []
 

sdk/bettmensch_ai/components/torch_utils.py

@@ -140,7 +140,7 @@ def wrapper(*function_args):
 
 def create_torch_ddp_service_template(
     component_base_name: str,
-    service_name: str,
+    component_task_name: str,
     namespace: str = ARGO_NAMESPACE,  # noqa: E501
 ) -> Resource:
     """Utility for a template creating the service resource required for
@@ -152,26 +152,28 @@ def create_torch_ddp_service_template(
         manifest=f"""apiVersion: v1
 kind: Service
 metadata:
-  name: {service_name}
+  name: {component_task_name}-{{{{workflow.uid}}}}
   namespace: {namespace}
   labels:
-    app: {service_name}
+    workflows.argoproj.io/workflow: {{{{workflow.name}}}}
+    torch-job: {component_task_name}
 spec:
   clusterIP: None  # ClusterIP set to None for headless service.
   ports:
   - name: {DDP_PORT_NAME}  # Port for torchrun master<->worker node coms.
     port: {DDP_PORT_NUMBER}
     targetPort: {DDP_PORT_NUMBER}
   selector:
-    torch-job: {service_name}
+    workflows.argoproj.io/workflow: {{{{workflow.name}}}}
+    torch-job: {component_task_name}
     torch-node: '0'  # Selector for pods associated with this service.
 """,
     )
 
 
 def delete_torch_ddp_service_template(
     component_base_name: str,
-    service_name: str,
+    component_task_name: str,
     namespace: str = ARGO_NAMESPACE,  # noqa: E501
 ) -> Resource:
     """Utility for a template deleting the service resource required for
@@ -183,7 +185,7 @@ def delete_torch_ddp_service_template(
         flags=[
             "service",
             "--selector",
-            f"app={service_name}",
+            f"torch-job={component_task_name},workflows.argoproj.io/workflow={{{{workflow.name}}}}",  # noqa: E501
             "-n",
             namespace,
         ],

sdk/test/unit/test_pipeline.py

@@ -74,11 +74,7 @@ def parameter_to_artifact(
     ]
 
     assert wft.templates[3].labels["torch-node"] == "0"
-    assert (
-        wft.templates[3]
-        .labels["torch-job"]
-        .startswith("convert-to-artifact-0-")
-    )
+    assert wft.templates[3].labels["torch-job"] == "convert-to-artifact-0"
     assert wft.templates[4].labels["torch-node"] == "1"
 
     parameter_to_artifact.export(test_output_dir)
@@ -146,7 +142,7 @@ def adding_parameters(
     ]
 
     assert wft.templates[3].labels["torch-node"] == "0"
-    assert wft.templates[3].labels["torch-job"].startswith("a-plus-b-0-")
+    assert wft.templates[3].labels["torch-job"] == "a-plus-b-0"
     assert wft.templates[4].labels["torch-node"] == "1"
 
     adding_parameters.export(test_output_dir)

sdk/test/unit/test_torch_ddp_component.py

@@ -212,7 +212,9 @@ def test_function(
     assert test_component.task_factory is None
 
 
-def test_parameter_torch_component_to_hera(test_mock_pipeline):
+def test_parameter_torch_component_to_hera(
+    test_output_dir, test_mock_pipeline
+):
     """Declaration of Component using InputParameter and OutputParameter"""
 
     # mock active pipeline with 2 inputs
@@ -255,7 +257,7 @@ def test_parameter_torch_component_to_hera(test_mock_pipeline):
     a_plus_b_plus_2.task_factory = a_plus_b_plus_2.build_hera_task_factory()
 
     with WorkflowTemplate(
-        name="test-parameter-component-workflow-template",
+        name="test-parameter-torch-ddp-component-workflow-template",
         entrypoint="test_dag",
         namespace="argo",
         arguments=[
@@ -273,6 +275,8 @@ def test_parameter_torch_component_to_hera(test_mock_pipeline):
             a_plus_b.to_hera()
             a_plus_b_plus_2.to_hera()
 
+    wft.to_file(test_output_dir)
+
     task_names = [task.name for task in wft.templates[4].tasks]
     assert task_names == [
         "a-plus-b-create-torch-ddp-service",
@@ -299,17 +303,18 @@ def test_parameter_torch_component_to_hera(test_mock_pipeline):
     ]
 
     assert wft.templates[5].labels["torch-node"] == "0"
-    assert wft.templates[5].labels["torch-job"].startswith("a-plus-b-0-")
+    assert wft.templates[5].labels["torch-job"] == "a-plus-b-0"
     assert wft.templates[6].labels["torch-node"] == "1"
 
     assert wft.templates[7].labels["torch-node"] == "0"
     assert (
-        wft.templates[7].labels["torch-job"].startswith("a-plus-b-plus-2-0-")
+        wft.templates[7].labels["torch-job"] == "a-plus-b-plus-2-0"
     )  # noqa: E501
     assert wft.templates[8].labels["torch-node"] == "1"
 
 
 def test_artifact_torch_component_to_hera(
+    test_output_dir,
     test_mock_pipeline,
 ):
 
@@ -346,7 +351,7 @@ def test_artifact_torch_component_to_hera(
     show.task_factory = show.build_hera_task_factory()
 
     with WorkflowTemplate(
-        name="test-artifact-component-workflow-template",
+        name="test-artifact-torch-ddp-component-workflow-template",
         entrypoint="test_dag",
         namespace="argo",
         arguments=[
@@ -361,6 +366,8 @@ def test_artifact_torch_component_to_hera(
             convert.to_hera()
             show.to_hera()
 
+    wft.to_file(test_output_dir)
+
     task_names = [task.name for task in wft.templates[4].tasks]
     assert task_names == [
         "convert-parameters-create-torch-ddp-service",
@@ -386,11 +393,10 @@ def test_artifact_torch_component_to_hera(
 
     assert wft.templates[5].labels["torch-node"] == "0"
     assert (
-        wft.templates[5]
-        .labels["torch-job"]
-        .startswith("convert-parameters-0-")  # noqa: E501
+        wft.templates[5].labels["torch-job"]
+        == "convert-parameters-0"  # noqa: E501
     )
     assert wft.templates[6].labels["torch-node"] == "1"
 
     assert wft.templates[7].labels["torch-node"] == "0"
-    assert wft.templates[7].labels["torch-job"].startswith("show-artifacts-0-")
+    assert wft.templates[7].labels["torch-job"] == "show-artifacts-0"