idds/workflowv2/tree.py

0001 #!/usr/bin/env python
0002 #
0003 # Licensed under the Apache License, Version 2.0 (the "License");
0004 # You may not use this file except in compliance with the License.
0005 # You may obtain a copy of the License at
0006 # http://www.apache.org/licenses/LICENSE-2.0OA
0007 #
0008 # Authors:
0009 # - Wen Guan, <wen.guan@cern.ch>, 2023
0010
0011
0012 """
0013 Construct tree from a workflow
0014 """
0015
0016 from idds.common.utils import setup_logging
0017
0018 from .base import Base
0019
0020
0021 setup_logging(__name__)
0022
0023
0024 class Node(Base):
0025     def __init__(self, name, children=None, parents=None, index=None, level=None, groups=None, disable_grouping=False, **kwargs):
0026         self.__dict__.update(kwargs)
0027
0028         if children is None:
0029             children = []
0030         if parents is None:
0031             parents = []
0032         if groups is None:
0033             groups = []
0034         self.name = name
0035         self.index = index
0036         self.level = level
0037         self.children = children
0038         self.parents = parents
0039         self.groups = []
0040         self.order_groups = []
0041         self.group_id = None
0042         self.order_id = None
0043         self.disable_grouping = disable_grouping
0044
0045         self.setup_logger()
0046
0047     @property
0048     def index(self):
0049         return str(self.__index)
0050
0051     @index.setter
0052     def index(self, value):
0053         if value:
0054             self.__index = str(value)
0055         else:
0056             self.__index = value
0057
0058     @property
0059     def level(self):
0060         return str(self.__level)
0061
0062     @level.setter
0063     def level(self, value):
0064         if value is not None:
0065             if self.__level is None or (self.__level is not None and value > self.__level):
0066                 self.__level = value
0067                 for child in self.children:
0068                     child.level = value + 1
0069         else:
0070             self.__level = value
0071
0072     def add_child(self, obj):
0073         self.children.append(obj)
0074         obj.parents.append(self)
0075
0076     def add_group(self, group):
0077         if group and group not in self.groups:
0078             self.groups.append(group)
0079             for child in self.children:
0080                 if not child.group_id:
0081                     child.add_group(group)
0082             for parent in self.parents:
0083                 if not parent.group_id:
0084                     parent.add_group(group)
0085
0086     def get_potential_group_id(self):
0087         if not self.groups:
0088             return "None"
0089         return "_".join([str(i) for i in sorted(self.groups)])
0090
0091     @property
0092     def group_id(self):
0093         return self.__group_id
0094
0095     @group_id.setter
0096     def group_id(self, group_id):
0097         self.__group_id = group_id
0098         if group_id:
0099             for child in self.children:
0100                 if not child.group_id:
0101                     child.add_group(group_id)
0102             for parent in self.parents:
0103                 if not parent.group_id:
0104                     parent.add_group(group_id)
0105
0106     def add_order_group(self, group):
0107         if group is not None and group not in self.order_groups:
0108             self.order_groups.append(group)
0109             for child in self.children:
0110                 if child.order_id is None:
0111                     child.add_order_group(group)
0112             for parent in self.parents:
0113                 if parent.order_id is None:
0114                     parent.add_order_group(group)
0115
0116     def get_potential_order_id(self):
0117         if not self.order_groups:
0118             return 0
0119         return "_".join([str(i) for i in sorted(self.order_groups)])
0120
0121     @property
0122     def order_id(self):
0123         return self.__order_id
0124
0125     @order_id.setter
0126     def order_id(self, value):
0127         self.__order_id = value
0128         if value is not None:
0129             self.__order_id = int(value)
0130             for child in self.children:
0131                 if child.order_id is None:
0132                     child.add_order_group(value)
0133             for parent in self.parents:
0134                 if parent.order_id is None:
0135                     parent.add_order_group(value)
0136
0137     def get_node_name(self):
0138         return get_node_name(self.index, self.name)
0139
0140     def __repr__(self):
0141         # if self.__level:
0142         #     return "_" * self.__level + "Node(name: %s, level: %s)" % (self.name, self.__level)
0143         # return "Node(name: %s, level: %s, group_id: %s, groups: %s, parents: %s, children: %s)" % (self.name, self.__level, self.group_id, self.groups, self.parents, self.children)
0144         return "Node(name: %s, level: %s)" % (self.name, self.__level)
0145
0146
0147 class WorkNode(Node):
0148     def __init__(self, name, work=None, children=None, parents=None, index=None, level=None, groups=None,
0149                  disable_grouping=False, **kwargs):
0150         super(WorkNode, self).__init__(name=name, work=work, children=children, parents=parents,
0151                                        index=index, level=level, groups=groups,
0152                                        disable_grouping=disable_grouping, **kwargs)
0153
0154     def __repr__(self):
0155         return "WorkNode(name: %s, level: %s)" % (self.name, self.level)
0156
0157
0158 class LabelNode(Node):
0159     def __init__(self, name, jobs=None, children=None, parents=None, index=None, level=None, groups=None,
0160                  disable_grouping=False, **kwargs):
0161         super(LabelNode, self).__init__(name=name, jobs=jobs, children=children, parents=parents,
0162                                         index=index, level=level, groups=groups,
0163                                         disable_grouping=disable_grouping, **kwargs)
0164
0165     def __repr__(self):
0166         return "LabelNode(name: %s, level: %s)" % (self.name, self.level)
0167
0168
0169 class JobNode(Node):
0170     def __init__(self, name, work_node=None, children=None, parents=None, index=None, level=None,
0171                  groups=None, disable_grouping=False, **kwargs):
0172         super(JobNode, self).__init__(name=name, work=work_node, children=children, parents=parents,
0173                                       index=index, level=level, groups=groups,
0174                                       disable_grouping=disable_grouping, **kwargs)
0175
0176     def __repr__(self):
0177         return "JobNode(name: %s, level: %s)" % (self.name, self.level)
0178
0179     def get_node_name(self):
0180         index = self.index
0181         if self.work_node:
0182             index = self.work_node.index
0183         return get_node_name(index, self.name)
0184
0185
0186 def get_node_name(index=None, name=None):
0187     if index:
0188         return "%s:%s" % (index, name)
0189     return name
0190
0191
0192 class Tree(Base):
0193     def __init__(self, name, **kwargs):
0194         self.__dict__.update(kwargs)
0195         self.name = name
0196         self.roots = []
0197
0198         self.setup_logger()