#696 fix affinities for argo workflows

Author: filippomc

libraries/cloudharness-common/cloudharness/workflows/operations.py

@@ -178,7 +178,11 @@ def add_on_exit_notify_handler(self, spec):
     def modify_template(self, template):
         """Hook to modify templates (e.g. add volumes)"""
 
-        template["metadata"] = {"labels": {c.key: c.value for c in self.pod_contexts}}
+        if 'metadata' not in template:
+            template['metadata'] = dict()
+        if 'labels' not in template["metadata"]:
+            template["metadata"]["labels"] = dict()
+        template["metadata"]["labels"] = template["metadata"]["labels"] | {c.key: c.value for c in self.pod_contexts}
 
         if self.volumes:
             if 'container' in template:

libraries/cloudharness-common/cloudharness/workflows/tasks.py

@@ -1,7 +1,7 @@
 from . import argo
 
 from cloudharness.utils.env import get_cloudharness_variables, get_image_full_tag
-from .utils import WORKFLOW_NAME_VARIABLE_NAME, PodExecutionContext, affinity_spec, is_accounts_present, volume_mount_template
+from .utils import WORKFLOW_NAME_VARIABLE_NAME, PodExecutionContext, affinity_spec, is_accounts_present, volume_mount_template, volume_requires_affinity
 
 SERVICE_ACCOUNT = 'argo-workflows'
 
@@ -16,9 +16,25 @@ def __init__(self, name, resources={}, volume_mounts=[],  **env_args):
         self.resources = resources
         self.__envs = get_cloudharness_variables()
         self.volume_mounts = volume_mounts
+        self.external_volumes = [
+            v.split(':')[0] for v in self.volume_mounts if volume_requires_affinity(v)]
         for k in env_args:
             self.__envs[k] = str(env_args[k])
 
+    def metadata_spec(self):
+        return {
+            'labels': {
+                'usesvolume': self.external_volumes[0],
+                **{f'usesvolume-{v}': 'true' for v in self.external_volumes}
+            }
+        } if self.external_volumes else {}
+
+    def affinity_spec(self):
+        return affinity_spec(([PodExecutionContext('usesvolume', self.external_volumes[0], True)] if self.external_volumes else []) + [
+            PodExecutionContext(f'usesvolume-{v}', 'true', True)
+            for v in self.external_volumes
+        ] )
+
     @property
     def image_name(self):
         raise NotImplemented
@@ -75,6 +91,7 @@ def __init__(self, name, resources={}, image_pull_policy='IfNotPresent', command
         self.command = command
 
     def spec(self):
+
         spec = {
             'container': {
                 'image': self.image_name,
@@ -84,11 +101,10 @@ def spec(self):
                 'volumeMounts': self.volumes_mounts_spec(),
             },
             'inputs': {},
-            'metadata': {},
+            'metadata': self.metadata_spec(),
             'name': self.name,
             'outputs': {},
-            'affinity': affinity_spec([PodExecutionContext('usesvolume', v.split(':')[0], True) for v in self.volume_mounts if
-                                       ':' in v])
+            'affinity': self.affinity_spec()
 
         }
         if self.command is not None:
@@ -106,12 +122,11 @@ def __init__(self, name, source, **kwargs):
         self.source = source
 
     def spec(self):
+
         return {
             'name': self.name,
-            'affinity': affinity_spec([
-                PodExecutionContext('usesvolume', v.split(':')[0], True)
-                for v in self.volume_mounts if ':' in v
-            ]),
+            'affinity': self.affinity_spec(),
+            'metadata': self.metadata_spec(),
             'script':
                 {
                     'image': self.image_name,

libraries/cloudharness-common/cloudharness/workflows/utils.py

@@ -25,6 +25,8 @@ def get_workflow_name():
     remove = name.split("-")[-1]
     return name[0:-len(remove) - 1]
 
+def volume_requires_affinity(v):
+    return ':' in v and 'rwx' not in v[-4:]
 
 def get_shared_directory():
     return os.getenv(SHARED_DIRECTORY_VARIABLE_NAME)

libraries/cloudharness-common/tests/test_workflow.py

@@ -7,7 +7,7 @@
 
 set_test_environment()
 
-from cloudharness.workflows import operations, tasks
+from cloudharness.workflows import operations, tasks, utils
 from cloudharness import set_debug
 from cloudharness.workflows import argo
 from cloudharness.utils.config import CloudharnessConfig
@@ -23,7 +23,12 @@ def check_wf(wf):
 
     assert wf["kind"] == "Workflow"
     assert "spec" in wf
-    
+
+def test_volume_affinity_check():
+    assert not utils.volume_requires_affinity("a")
+    assert utils.volume_requires_affinity("a:b")
+    assert utils.volume_requires_affinity("a:b:ro")
+    assert not utils.volume_requires_affinity("a:b:rwx")
 
 def test_sync_workflow():
     def f():
@@ -119,9 +124,44 @@ def test_single_task_shared():
     if accounts_offset == 1:
         assert wf['spec']['volumes'][1]['secret']['secretName'] == 'accounts'
     assert len(wf['spec']['templates'][0]['container']['volumeMounts']) == 2 + accounts_offset
+    assert wf['spec']['templates'][0]['metadata']['labels']['usesvolume'] == 'myclaim'
+
+    affinity_glob = \
+        wf['spec']['affinity']['podAffinity']['requiredDuringSchedulingIgnoredDuringExecution'][0]['labelSelector'][
+            'matchExpressions'][0]
+    assert affinity_glob['key'] == 'usesvolume'
+    assert affinity_glob['values'][0] == 'myclaim'
+
     if execute:
         print(op.execute())
 
+def test_pipeline_shared():
+    shared_directory = 'myclaim:/mnt/shared'
+    task_write = operations.CustomTask('download-file', 'workflows-extract-download',
+                                       url='https://raw.githubusercontent.com/openworm/org.geppetto/master/README.md')
+    task_script = tasks.BashTask('print-file', source="ls -la")
+    op = operations.PipelineOperation('test-custom-connected-op-', [task_write, task_script],
+                                        shared_directory=shared_directory, shared_volume_size=100)
+    wf = op.to_workflow()
+
+    accounts_offset = 1 if is_accounts_present() else 0
+    assert len(op.volumes) == 1
+    assert len(wf['spec']['volumes']) == 2 + accounts_offset
+    assert wf['spec']['volumes'][1+accounts_offset]['persistentVolumeClaim']['claimName'] == 'myclaim'
+    if accounts_offset == 1:
+        assert wf['spec']['volumes'][1]['secret']['secretName'] == 'accounts'
+    assert len(wf['spec']['templates'][1]['container']['volumeMounts']) == 2 + accounts_offset
+    assert wf['spec']['templates'][1]['metadata']['labels']['usesvolume'] == 'myclaim'
+    assert wf['spec']['templates'][2]['metadata']['labels']['usesvolume'] == 'myclaim'
+    affinity_glob = \
+        wf['spec']['affinity']['podAffinity']['requiredDuringSchedulingIgnoredDuringExecution'][0]['labelSelector'][
+            'matchExpressions'][0]
+    assert affinity_glob['key'] == 'usesvolume'
+    assert affinity_glob['values'][0] == 'myclaim'
+
+    if execute:
+        print(op.execute())        
+
 def test_single_task_shared_rwx():
     shared_directory = 'myclaim:/mnt/shared:rwx'
     task_write = operations.CustomTask('download-file', 'workflows-extract-download',
@@ -139,9 +179,39 @@ def test_single_task_shared_rwx():
         assert wf['spec']['volumes'][1]['secret']['secretName'] == 'accounts'
     assert len(wf['spec']['templates'][0]['container']['volumeMounts']) == 2 + accounts_offset
 
+
+
+
     assert not 'affinity' in wf['spec'], "Pod affinity should not be added for rwx volumes"
-    
+
 def test_single_task_volume_notshared():
+
+    task_write = operations.CustomTask('download-file', 'workflows-extract-download', volume_mounts=["a:b"],
+                                       url='https://raw.githubusercontent.com/openworm/org.geppetto/master/README.md')
+    op = operations.SingleTaskOperation('test-custom-connected-op-', task_write, shared_volume_size=100)
+    wf = op.to_workflow()
+
+    accounts_offset = 1 if is_accounts_present() else 0
+    assert len(op.volumes) == 0
+    assert len(wf['spec']['volumes']) == 2 + accounts_offset
+    assert wf['spec']['volumes'][1+accounts_offset]['persistentVolumeClaim']['claimName'] == 'a'
+    if accounts_offset == 1:
+        assert wf['spec']['volumes'][1]['secret']['secretName'] == 'accounts'
+    assert len(wf['spec']['templates'][0]['container']['volumeMounts']) == 2 + accounts_offset
+
+    
+    assert 'affinity' not in wf['spec']
+
+    affinity_tpl = \
+        wf['spec']['templates'][0]['affinity']['podAffinity']['requiredDuringSchedulingIgnoredDuringExecution'][0]['labelSelector'][
+            'matchExpressions'][0]
+    assert affinity_tpl['key'] == 'usesvolume'
+    assert affinity_tpl['values'][0] == 'a'
+    assert wf['spec']['templates'][0]['metadata']['labels']['usesvolume'] == 'a'
+    if execute:
+        print(op.execute())
+
+def test_single_task_volumes_notshared():
     shared_directory = 'myclaim:/mnt/shared'
     task_write = operations.CustomTask('download-file', 'workflows-extract-download', volume_mounts=["a:b"],
                                        url='https://raw.githubusercontent.com/openworm/org.geppetto/master/README.md')
@@ -167,7 +237,7 @@ def test_single_task_volume_notshared():
             'matchExpressions'][0]
     assert affinity_tpl['key'] == 'usesvolume'
     assert affinity_tpl['values'][0] == 'a'
-
+    assert wf['spec']['templates'][0]['metadata']['labels']['usesvolume']
     if execute:
         print(op.execute())
 
@@ -202,6 +272,7 @@ def test_single_task_shared_multiple():
         print(op.execute())
 
 
+
 def test_single_task_shared_script():
     shared_directory = 'myclaim:/mnt/shared'
     task_write = tasks.BashTask('download-file', source="ls -la")