@@ -75,9 +75,10 @@ def update_nodes(item, appwrapper_name, cpu, memory, gpu, workers):
75
75
update_affinity (spec , appwrapper_name )
76
76
update_resources (spec , cpu , memory , gpu )
77
77
78
- def write_user_appwrapper (user_yaml , appwrapper_name ):
79
- with open (f' { appwrapper_name } .yaml' ,'w' ) as outfile :
78
+ def write_user_appwrapper (user_yaml , output_file_name ):
79
+ with open (output_file_name ,'w' ) as outfile :
80
80
yaml .dump (user_yaml , outfile , default_flow_style = False )
81
+ print (f"Written to: { output_file_name } " )
81
82
82
83
def generate_appwrapper (cpu , memory , gpu , workers , template ):
83
84
user_yaml = read_template (template )
@@ -87,7 +88,10 @@ def generate_appwrapper(cpu, memory, gpu, workers, template):
87
88
update_names (user_yaml , item , appwrapper_name , cluster_name )
88
89
update_custompodresources (item , cpu , memory , gpu , workers )
89
90
update_nodes (item , appwrapper_name , cpu , memory , gpu , workers )
90
- write_user_appwrapper (user_yaml , appwrapper_name )
91
+ outfile = appwrapper_name + ".yaml"
92
+ write_user_appwrapper (user_yaml , outfile )
93
+ return outfile
94
+
91
95
92
96
def main ():
93
97
parser = argparse .ArgumentParser (description = 'Generate user AppWrapper' )
@@ -104,7 +108,8 @@ def main():
104
108
workers = args .workers
105
109
template = args .template
106
110
107
- generate_appwrapper (cpu , memory , gpu , workers , template )
111
+ outfile = generate_appwrapper (cpu , memory , gpu , workers , template )
112
+ return outfile
108
113
109
114
if __name__ == "__main__" :
110
115
main ()
0 commit comments