11import json
22import logging
33import os
4+ import shutil
45import subprocess
56import sys
67from pathlib import Path
8+ from typing import Dict
79
810logging .basicConfig (level = logging .INFO )
911logger = logging .getLogger ("osparc-python-main" )
1012
1113
12- INPUT_1 = Path (os .environ ["INPUT_1" ])
14+ ENVIRONS = ["INPUT_FOLDER" , "OUTPUT_FOLDER" ]
15+ try :
16+ INPUT_FOLDER , OUTPUT_FOLDER = [Path (os .environ [v ]) for v in ENVIRONS ]
17+ except KeyError :
18+ raise ValueError ("Required env vars {ENVIRONS} were not set" )
19+
20+ # NOTE: sync with schema in metadata!!
21+ NUM_INPUTS = 5
22+ NUM_OUTPUTS = 4
23+ OUTPUT_SUBFOLDER_ENV_TEMPLATE = "OUTPUT_{}"
24+ OUTPUT_SUBFOLDER_TEMPLATE = "output_{}"
25+ OUTPUT_FILE_TEMPLATE = "output_{}.zip"
1326
1427
1528def _find_user_code_entrypoint (code_dir : Path ) -> Path :
@@ -45,44 +58,59 @@ def _ensure_pip_requirements(code_dir: Path) -> Path:
4558 f"pipreqs --savepath={ requirements } { code_dir } .split (),
4659 shell = False ,
4760 check = True ,
48- cwd = INPUT_1 ,
61+ cwd = INPUT_FOLDER ,
4962 )
5063
5164 # TODO log subprocess.run
5265
5366 else :
5467 requirements = requirements [0 ]
55- logger .info ("Found: %s" , requirements )
68+ logger .info (f "Found: { requirements } "
5669 return requirements
5770
5871
59- def _show_io_environments () -> None :
60- for io_type in ["input" , "output" ]:
61- logger .info (
62- "%s ENVs available: %s" ,
63- io_type .capitalize (),
64- json .dumps (
65- list (
66- filter (
67- lambda x , io_type = io_type : f"{ io_type .upper ()} in x ,
68- os .environ ,
69- )
70- ),
71- indent = 2 ,
72- ),
73- )
72+ # TODO: Next version of integration will take care of this and maybe the ENVs as well
73+ def _ensure_output_subfolders_exist () -> Dict [str , str ]:
74+ output_envs = {}
75+ for n in range (1 , NUM_OUTPUTS + 1 ):
76+ output_sub_folder_env = f"OUTPUT_{ n }
77+ output_sub_folder = OUTPUT_FOLDER / OUTPUT_SUBFOLDER_TEMPLATE .format (n )
78+ # NOTE: exist_ok for forward compatibility in case they are already created
79+ output_sub_folder .mkdir (parents = True , exist_ok = True )
80+ output_envs [output_sub_folder_env ] = f"{ output_sub_folder }
81+ logger .info (
82+ "Output ENVs available: %s" ,
83+ json .dumps (output_envs , indent = 2 ),
84+ )
85+ return output_envs
86+
87+
88+ def _ensure_input_environment () -> Dict [str , str ]:
89+ input_envs = {
90+ f"INPUT_{ n } : os .environ [f"INPUT_{ n } ] for n in range (1 , NUM_INPUTS + 1 )
91+ }
92+ logger .info (
93+ "Input ENVs available: %s" ,
94+ json .dumps (input_envs , indent = 2 ),
95+ )
96+ return input_envs
7497
7598
7699def setup ():
77- _show_io_environments ()
100+ input_envs = _ensure_input_environment ()
101+ output_envs = _ensure_output_subfolders_exist ()
78102 logger .info ("Available data:" )
79103 os .system ("ls -tlah" )
80104
81- user_code_entrypoint = _find_user_code_entrypoint (INPUT_1 )
82- requirements_txt = _ensure_pip_requirements (INPUT_1 )
105+ user_code_entrypoint = _find_user_code_entrypoint (INPUT_FOLDER )
106+ requirements_txt = _ensure_pip_requirements (INPUT_FOLDER )
83107
84108 logger .info ("Preparing launch script ..." )
85109 venv_dir = Path .home () / ".venv"
110+ bash_input_env_export = [f"export { env } { path } for env , path in input_envs .items ()]
111+ bash_output_env_export = [
112+ f"export { env } { path } for env , path in output_envs .items ()
113+ ]
86114 script = [
87115 "#!/bin/sh" ,
88116 "set -o errexit" ,
@@ -92,16 +120,30 @@ def setup():
92120 f'python3 -m venv --system-site-packages --symlinks --upgrade "{ venv_dir } ,
93121 f'"{ venv_dir } ,
94122 f'"{ venv_dir } { requirements_txt } ,
123+ "\n " .join (bash_input_env_export ),
124+ "\n " .join (bash_output_env_export ),
95125 f'echo "Executing code { user_code_entrypoint .name } ,
96126 f'"{ venv_dir } { user_code_entrypoint } ,
97127 'echo "DONE ..."' ,
98128 ]
99129 main_script_path = Path ("main.sh" )
100- main_script_path .write_text ("\n " .join (script ), encoding = "utf-8" )
130+ main_script_path .write_text ("\n " .join (script ))
101131
102132
103133def teardown ():
104- logger .info ("Completed" )
134+ logger .info ("Zipping output..." )
135+ for n in range (1 , NUM_OUTPUTS + 1 ):
136+ output_path = OUTPUT_FOLDER / f"output_{ n }
137+ archive_file_path = OUTPUT_FOLDER / OUTPUT_FILE_TEMPLATE .format (n )
138+ logger .info ("Zipping %s into %s..." , output_path , archive_file_path )
139+ shutil .make_archive (
140+ f"{ (archive_file_path .parent / archive_file_path .stem )} ,
141+ format = "zip" ,
142+ root_dir = output_path ,
143+ logger = logger ,
144+ )
145+ logger .info ("Zipping %s into %s done" , output_path , archive_file_path )
146+ logger .info ("Zipping done." )
105147
106148
107149if __name__ == "__main__" :
0 commit comments