@@ -29,60 +29,166 @@ def __init__(self, mode="server", config_file=None, config_name=None):
29
29
file = sys .stderr )
30
30
sys .exit (1 )
31
31
32
- def _load_env_from_config (self ):
32
+ def _load_env_from_defaults (self ):
33
+ """
34
+ Loads default environment variables from a YAML file based on the mode.
35
+ For each section starting with 'defaults_', if model is in the section's
36
+ 'MODELS' list, loads the environment variables from that section.
37
+ If no section matches, loads nothing.
38
+ If the file does not exist, it returns an empty dictionary.
39
+ """
40
+ defaults_file = ("server_configurations/defaults.yaml"
41
+ if self .mode == "server" else
42
+ "benchmark_configurations/defaults.yaml" )
43
+ try :
44
+ with open (defaults_file ) as f :
45
+ config = yaml .safe_load (f )
46
+ found = False
47
+ for section_name , section in config .items ():
48
+ if section_name .startswith ("defaults_" ) and isinstance (
49
+ section , dict ):
50
+ models = section .get ("MODELS" , [])
51
+ if (isinstance (models , list )
52
+ and self .config_envs .get ("MODEL" ) in models ):
53
+ env_vars = {
54
+ k : v
55
+ for k , v in section .items () if k != "MODELS"
56
+ }
57
+ self .config_envs .update (env_vars )
58
+ print (
59
+ f"[INFO] Loaded default configuration section "
60
+ f"'{ section_name } ' for model "
61
+ f"'{ self .config_envs .get ('MODEL' )} ' from file: "
62
+ f"{ defaults_file } " )
63
+ for key , value in env_vars .items ():
64
+ print (f" { key } : { value } " )
65
+ found = True
66
+ if not found :
67
+ print (f"[WARNING] No defaults section found for model "
68
+ f"'{ self .config_envs .get ('MODEL' )} ' in "
69
+ f"'{ defaults_file } '." )
70
+ except FileNotFoundError :
71
+ print (f"[WARNING] Defaults file '{ defaults_file } ' not found. "
72
+ "No defaults loaded." )
73
+ except Exception as e :
74
+ print (
75
+ f"[ERROR] Failed to load defaults: { e } " ,
76
+ file = sys .stderr ,
77
+ )
78
+ sys .exit (1 )
79
+
80
+ def _load_env_from_config_file (self ):
81
+ """
82
+ Loads a specific configuration section from a YAML file and updates the
83
+ current environment configuration with the values from that section.
84
+ If a key already exists (e.g., from defaults), it will be overwritten
85
+ by the value from the file. Exits the program with an error message if
86
+ the section is missing or invalid, or if the file cannot be read.
87
+
88
+ Raises:
89
+ SystemExit: If the configuration file or section is missing,
90
+ invalid, or cannot be loaded.
91
+ """
33
92
try :
34
93
with open (self .config_file ) as f :
35
94
config = yaml .safe_load (f )
36
95
section = config .get (self .config_name )
37
- if section is None :
38
- print (
39
- f"[ERROR] Section '{ self .config_name } ' not found in "
40
- f"'{ self .config_file } '." ,
41
- file = sys .stderr )
42
- sys .exit (1 )
43
- if not isinstance (section , dict ):
96
+ if section is None or not isinstance (section , dict ):
44
97
print (
45
- f"[ERROR] Section '{ self .config_name } ' is not a "
46
- f"dictionary in '{ self .config_file } '." ,
47
- file = sys .stderr )
98
+ f"[ERROR] Section '{ self .config_name } ' not found or "
99
+ f"is not a dictionary in '{ self .config_file } '." ,
100
+ file = sys .stderr ,
101
+ )
48
102
sys .exit (1 )
49
- self .config_envs = section
50
- print (f"[INFO] Loaded configuration from file: "
51
- f"{ self .config_file } , section: { self .config_name } " )
52
- print ("[INFO] The following parameters and values were loaded "
53
- "from the config file:" )
54
- for key , value in self .config_envs .items ():
103
+ print (f"[INFO] Loaded configuration section "
104
+ f"'{ self .config_name } ' from file: { self .config_file } " )
105
+ for key , value in section .items ():
55
106
print (f" { key } : { value } " )
107
+ return section
56
108
except Exception as e :
57
- print (f"[ERROR] Failed to load config: { e } " , file = sys .stderr )
109
+ print (
110
+ f"[ERROR] Failed to load config: { e } " ,
111
+ file = sys .stderr ,
112
+ )
113
+ sys .exit (1 )
114
+
115
+ def _update_benchmark_envs_from_user_vars (self ):
116
+ """
117
+ Loads a list of variable names from a YAML file and, for each variable
118
+ present in the current environment, updates the internal configuration
119
+ dictionary with the environment value. If the YAML file is missing or
120
+ empty, no variables are updated.
121
+ """
122
+ user_vars_file = "benchmark_configurations/user_vars.yaml"
123
+ try :
124
+ with open (user_vars_file ) as f :
125
+ user_vars = yaml .safe_load (f )
126
+ if user_vars and isinstance (user_vars , dict ):
127
+ variables = user_vars .get ("variables" , [])
128
+ for var in variables :
129
+ if var in os .environ :
130
+ self .config_envs [var ] = os .environ [var ]
131
+ print (f"[INFO] Overwriting { var } with value from "
132
+ f"environment: { self .config_envs [var ]} " )
133
+ else :
134
+ print (f"[WARNING] No user-defined variables found in "
135
+ f"'{ user_vars_file } '." )
136
+ except FileNotFoundError :
137
+ print (
138
+ f"[WARNING] User variables file '{ user_vars_file } ' not found. "
139
+ "No user-defined variables loaded." )
140
+ except Exception as e :
141
+ print (
142
+ f"[ERROR] Failed to load user-defined variables: { e } " ,
143
+ file = sys .stderr ,
144
+ )
58
145
sys .exit (1 )
59
146
60
147
def run (self ):
148
+ model_conf = {}
61
149
if self .config_file and self .config_name :
62
- self ._load_env_from_config ()
150
+ model_conf = self ._load_env_from_config_file ()
151
+ if "MODEL" in model_conf :
152
+ self .config_envs ["MODEL" ] = model_conf ["MODEL" ]
153
+
154
+ env_model = os .environ .get ("MODEL" )
155
+ if env_model :
156
+ self .config_envs ["MODEL" ] = env_model
157
+
158
+ if not self .config_envs .get ("MODEL" ):
159
+ print ("[ERROR] MODEL is not set. Exiting." , file = sys .stderr )
160
+ sys .exit (1 )
161
+
162
+ self ._load_env_from_defaults ()
163
+
164
+ if model_conf :
165
+ self .config_envs .update (model_conf )
63
166
64
167
if self .mode == "server" :
65
168
print ("[INFO] Starting container in server mode." )
66
- # VarsGenerator will read variables from the environment
67
169
for key , value in self .config_envs .items ():
68
170
os .environ [str (key )] = str (value )
69
171
variables = VarsGenerator (
70
172
defaults_path = "server_autoconfig/defaults.yaml" ,
71
173
varlist_conf_path = "server_autoconfig/varlist_conf.yaml" ,
72
- model_def_settings_path = ("server_autoconfig/settings_vllm.csv"
73
- )).calculate_variables ()
174
+ model_def_settings_path = (
175
+ "server_autoconfig/settings_vllm.csv" ),
176
+ ).calculate_variables ()
74
177
ScriptGenerator (
75
178
template_script_path = "templates/template_vllm_server.sh" ,
76
179
output_script_path = "vllm_server.sh" ,
77
180
variables = variables ,
78
- log_dir = "logs" ).create_and_run ()
181
+ log_dir = "logs" ,
182
+ ).create_and_run ()
79
183
elif self .mode == "benchmark" :
80
184
print ("[INFO] Starting container in benchmark mode." )
185
+ self ._update_benchmark_envs_from_user_vars ()
81
186
ScriptGenerator (
82
187
template_script_path = "templates/template_vllm_benchmark.sh" ,
83
188
output_script_path = "vllm_benchmark.sh" ,
84
189
variables = self .config_envs ,
85
- log_dir = "logs" ).create_and_run ()
190
+ log_dir = "logs" ,
191
+ ).create_and_run ()
86
192
elif self .mode == "test" :
87
193
print ("[INFO] Test mode: keeping container active. "
88
194
"Press Ctrl+C to exit." )
@@ -102,17 +208,21 @@ def run(self):
102
208
if __name__ == "__main__" :
103
209
parser = argparse .ArgumentParser (
104
210
description = "EntrypointMain for vllm docker" )
105
- parser .add_argument ("mode" ,
106
- nargs = "?" ,
107
- default = "server" ,
108
- choices = ["server" , "benchmark" , "test" ],
109
- help = "Mode to run: server, benchmark, or test" )
211
+ parser .add_argument (
212
+ "mode" ,
213
+ nargs = "?" ,
214
+ default = "server" ,
215
+ choices = ["server" , "benchmark" , "test" ],
216
+ help = "Mode to run: server, benchmark, or test" ,
217
+ )
110
218
parser .add_argument ("--config-file" , type = str , help = "Path to config file" )
111
219
parser .add_argument ("--config-name" ,
112
220
type = str ,
113
221
help = "Config name in the config file" )
114
222
args = parser .parse_args ()
115
223
116
- EntrypointMain (mode = args .mode ,
117
- config_file = args .config_file ,
118
- config_name = args .config_name ).run ()
224
+ EntrypointMain (
225
+ mode = args .mode ,
226
+ config_file = args .config_file ,
227
+ config_name = args .config_name ,
228
+ ).run ()
0 commit comments