async def start_scheduler(gateway, security, exit_on_failure=True):
loop = IOLoop.current()
plugin = GatewaySchedulerPlugin(gateway, loop)
services = {("gateway", 0): (GatewaySchedulerService, {"plugin": plugin})}
dashboard = False
with ignoring(ImportError):
from distributed.dashboard.scheduler import BokehScheduler
services[("dashboard", 0)] = (BokehScheduler, {})
dashboard = True
scheduler = Scheduler(loop=loop, services=services, security=security)
scheduler.add_plugin(plugin)
await scheduler
host = urlparse(scheduler.address).hostname
gateway_port = scheduler.services["gateway"].port
api_address = "http://%s:%d" % (host, gateway_port)
if dashboard:
dashboard_port = scheduler.services["dashboard"].port
dashboard_address = "http://%s:%d" % (host, dashboard_port)
else:
dashboard_address = ""
try:
await gateway.send_addresses(scheduler.address, dashboard_address,
api_address)
except Exception as exc:
logger.error("Failed to send addresses to gateway", exc_info=exc)
if exit_on_failure:
sys.exit(1)
return scheduler
async def _inject_viztracer(dask_scheduler: distributed.Scheduler):
plugin = VizTracerScheduler(
# Type "str | Unknown" cannot be assigned to type "Literal['json', 'json.gz', 'html']"
format=format, # type: ignore
tracer_entries=tracer_entries,
verbose=verbose,
max_stack_depth=max_stack_depth,
include_files=include_files,
exclude_files=exclude_files,
ignore_c_function=ignore_c_function,
ignore_frozen=ignore_frozen,
log_func_retval=log_func_retval,
log_func_args=log_func_args,
log_print=log_print,
log_gc=log_gc,
log_sparse=log_sparse,
trace_sparse=trace_sparse,
log_async=log_async,
vdb=vdb,
file_info=file_info,
register_global=register_global,
plugins=plugins,
)
dask_scheduler.add_plugin(plugin)
return await plugin.start(dask_scheduler)
def main(host, port, bokeh_port, show, _bokeh, bokeh_whitelist, bokeh_prefix,
use_xheaders, pid_file, scheduler_file, interface, local_directory,
preload, preload_argv, tls_ca_file, tls_cert, tls_key):
logger = SchedulerLogger.getLogger()
enable_proctitle_on_current()
enable_proctitle_on_children()
log_metrics = EdasEnv.getBool("log.metrics", False)
logger.info(f"Log Metrics: {log_metrics}")
plugins = [EDASSchedulerPlugin(logger)] if log_metrics else []
sec = Security(
tls_ca_file=tls_ca_file,
tls_scheduler_cert=tls_cert,
tls_scheduler_key=tls_key,
)
if not host and (tls_ca_file or tls_cert or tls_key):
host = 'tls://'
if pid_file:
with open(pid_file, 'w') as f:
f.write(str(os.getpid()))
def del_pid_file():
if os.path.exists(pid_file):
os.remove(pid_file)
atexit.register(del_pid_file)
local_directory_created = False
if local_directory:
if not os.path.exists(local_directory):
os.mkdir(local_directory)
local_directory_created = True
else:
local_directory = tempfile.mkdtemp(prefix='scheduler-')
local_directory_created = True
if local_directory not in sys.path:
sys.path.insert(0, local_directory)
if sys.platform.startswith('linux'):
import resource # module fails importing on Windows
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
limit = max(soft, hard // 2)
resource.setrlimit(resource.RLIMIT_NOFILE, (limit, hard))
if interface:
if host:
raise ValueError("Can not specify both interface and host")
else:
host = get_ip_interface(interface)
addr = uri_from_host_port(host, port, 8786)
loop = IOLoop.current()
logger.info('-' * 47)
services = {}
if _bokeh:
try:
from distributed.bokeh.scheduler import BokehScheduler
services[('bokeh', bokeh_port)] = (BokehScheduler, {
'prefix': bokeh_prefix
})
except ImportError as error:
if str(error).startswith('No module named'):
logger.info(
'Web dashboard not loaded. Unable to import bokeh')
else:
logger.info('Unable to import bokeh: %s' % str(error))
scheduler = Scheduler(loop=loop,
services=services,
scheduler_file=scheduler_file,
security=sec)
for plugin in plugins:
logger.info(f"@SP: Adding scheduler plugin: {plugin}")
scheduler.add_plugin(plugin)
scheduler.start(addr)
comm = Comm(scheduler)
comm.start()
if not preload:
preload = dask.config.get('distributed.scheduler.preload', {})
if not preload_argv:
preload_argv = dask.config.get('distributed.scheduler.preload-argv',
{})
preload_modules(preload,
parameter=scheduler,
file_dir=local_directory,
argv=preload_argv)
logger.info('Local Directory: %26s', local_directory)
logger.info('-' * 47)
install_signal_handlers(loop)
def shutdown_scheduler():
comm.terminate()
scheduler.stop()
if local_directory_created:
shutil.rmtree(local_directory)
logger.info("End scheduler at %r", addr)
def close_loop():
loop.stop()
loop.close()
shutdown_scheduler()
atexit.register(close_loop)
try:
loop.start()
loop.close()
finally:
shutdown_scheduler()
