@@ -127,7 +127,7 @@ def create_dependency_watcher(
127
127
) -> tuple [Watcher , CircusSocket , str ]:
128
128
from bentoml .serving import create_watcher
129
129
130
- num_workers , worker_envs = scheduler .get_worker_env (svc )
130
+ num_workers , resource_envs = scheduler .get_resource_envs (svc )
131
131
uri , socket = _get_server_socket (svc , uds_path , port_stack , backlog )
132
132
args = [
133
133
"-m" ,
@@ -141,8 +141,8 @@ def create_dependency_watcher(
141
141
"$(CIRCUS.WID)" ,
142
142
]
143
143
144
- if worker_envs :
145
- args .extend (["--worker-env" , json .dumps (worker_envs )])
144
+ if resource_envs :
145
+ args .extend (["--worker-env" , json .dumps (resource_envs )])
146
146
147
147
watcher = create_watcher (
148
148
name = f"service_{ svc .name } " ,
@@ -171,7 +171,7 @@ def create_dynamo_watcher(
171
171
uri , socket = _get_server_socket (svc , uds_path , port_stack , backlog )
172
172
173
173
# Get worker configuration
174
- num_workers , worker_envs = scheduler .get_worker_env (svc )
174
+ num_workers , resource_envs = scheduler .get_resource_envs (svc )
175
175
176
176
# Create Dynamo-specific worker args
177
177
args = [
@@ -184,8 +184,8 @@ def create_dynamo_watcher(
184
184
"$(CIRCUS.WID)" ,
185
185
]
186
186
187
- if worker_envs :
188
- args .extend (["--worker-env" , json .dumps (worker_envs )])
187
+ if resource_envs :
188
+ args .extend (["--worker-env" , json .dumps (resource_envs )])
189
189
190
190
# Update env to include ServiceConfig and service-specific environment variables
191
191
worker_env = env .copy () if env else {}
@@ -315,7 +315,7 @@ def serve_http(
315
315
316
316
if service_name and service_name != svc .name :
317
317
svc = svc .find_dependent_by_name (service_name )
318
- num_workers , worker_envs = allocator .get_worker_env (svc )
318
+ num_workers , resource_envs = allocator .get_resource_envs (svc )
319
319
server_on_deployment (svc )
320
320
uds_path = tempfile .mkdtemp (prefix = "bentoml-uds-" )
321
321
try :
@@ -419,8 +419,8 @@ def serve_http(
419
419
* timeouts_args ,
420
420
* timeout_args ,
421
421
]
422
- if worker_envs :
423
- server_args .extend (["--worker-env" , json .dumps (worker_envs )])
422
+ if resource_envs :
423
+ server_args .extend (["--worker-env" , json .dumps (resource_envs )])
424
424
if development_mode :
425
425
server_args .append ("--development-mode" )
426
426
@@ -438,13 +438,37 @@ def serve_http(
438
438
"--worker-id" ,
439
439
"$(CIRCUS.WID)" ,
440
440
]
441
+ # resource_envs is the resource allocation (ie CUDA_VISIBLE_DEVICES) for each worker created by the allocator
442
+ # these resource_envs are passed to each individual worker's environment which is set in serve_dynamo
443
+ if resource_envs :
444
+ args .extend (["--worker-env" , json .dumps (resource_envs )])
445
+ # env is the base bentoml environment variables. We make a copy and update it to add any service configurations and additional env vars
446
+ worker_env = env .copy () if env else {}
447
+
448
+ # Pass through the main service config
449
+ if "DYNAMO_SERVICE_CONFIG" in os .environ :
450
+ worker_env ["DYNAMO_SERVICE_CONFIG" ] = os .environ [
451
+ "DYNAMO_SERVICE_CONFIG"
452
+ ]
453
+
454
+ # Get service-specific environment variables from DYNAMO_SERVICE_ENVS
455
+ if "DYNAMO_SERVICE_ENVS" in os .environ :
456
+ try :
457
+ service_envs = json .loads (os .environ ["DYNAMO_SERVICE_ENVS" ])
458
+ if svc .name in service_envs :
459
+ service_args = service_envs [svc .name ].get ("ServiceArgs" , {})
460
+ if "envs" in service_args :
461
+ worker_env .update (service_args ["envs" ])
462
+ except json .JSONDecodeError as e :
463
+ logger .warning (f"Failed to parse DYNAMO_SERVICE_ENVS: { e } " )
464
+
441
465
watcher = create_watcher (
442
466
name = f"dynamo_service_{ svc .name } " ,
443
467
args = args ,
444
468
numprocesses = num_workers ,
445
469
working_dir = str (bento_path .absolute ()),
446
470
close_child_stdin = not development_mode ,
447
- env = env , # Dependency map will be injected by serve_http
471
+ env = worker_env , # Dependency map will be injected by serve_http
448
472
)
449
473
watchers .append (watcher )
450
474
print (f"dynamo_service_{ svc .name } entrypoint created" )
@@ -495,15 +519,15 @@ def serve_http(
495
519
arbiter .start (
496
520
cb = lambda _ : logger .info ( # type: ignore
497
521
(
498
- "Starting Dynamo Service %s (%s/%s) listening on %s://%s:%d ( Press CTRL+C to quit)"
522
+ "Starting Dynamo Service %s (Press CTRL+C to quit)"
499
523
if (
500
524
hasattr (svc , "is_dynamo_component" )
501
525
and svc .is_dynamo_component ()
502
526
)
503
527
else "Starting %s (Press CTRL+C to quit)"
504
528
),
505
529
* (
506
- (svc .name , * svc . dynamo_address (), scheme , log_host , port )
530
+ (svc .name ,)
507
531
if (
508
532
hasattr (svc , "is_dynamo_component" )
509
533
and svc .is_dynamo_component ()
0 commit comments