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()