class DaskHandler(IProcessingHandler): """This class wraps all Dask related functions.""" def __init__(self, number_of_workers, class_cb: Callable, brain_class, worker_log_level=logging.WARNING): super().__init__(number_of_workers) self._client: Optional[Client] = None self._cluster: Optional[LocalCluster] = None self.class_cb = class_cb self.brain_class = brain_class self.worker_log_level = worker_log_level def init_framework(self): if self._client: raise RuntimeError("Dask client already initialized.") # threads_per_worker must be one, because atari-env is not thread-safe. # And because lower the thread-count from the default, we must increase the number of workers self._cluster = LocalCluster(processes=True, asynchronous=False, threads_per_worker=1, silence_logs=self.worker_log_level, n_workers=self.number_of_workers, memory_pause_fraction=False, lifetime='1 hour', lifetime_stagger='5 minutes', lifetime_restart=True, interface="lo") self._client = Client(self._cluster) self._client.register_worker_plugin(_CreatorPlugin(self.class_cb, self.brain_class), name="creator-plugin") logging.info("Dask dashboard available at port: " + str(self._client.scheduler_info()["services"]["dashboard"])) def map(self, func, *iterable): if not self._client: raise RuntimeError("Dask client not initialized. Call \"init_framework\" before calling \"map\"") return self._client.gather(self._client.map(func, *iterable)) def cleanup_framework(self): self._client.shutdown()
def add_global_preloader( preloader: WorkerPreloader, client: Client = None, ) -> None: with open(PRELOADER_PATH, 'wb') as f: cloudpickle.dump(preloader, f) if client is not None: client.register_worker_plugin(preloader, name="global_preloader") LOCAL_MOCK_WORKER.plugins["global_preloader"] = preloader LOCAL_MOCK_WORKER.plugins["global_preloader"].setup(LOCAL_MOCK_WORKER)
def get_dask_client(n_workers, processes=True, n_threads=1, max_mem_fraction=.9, plugins=None): import psutil from dask.distributed import Client ml = psutil.virtual_memory().total * max(min(max_mem_fraction, 1), 0) ml = str(int(ml // 1e9) // n_workers) client = Client(processes=processes, threads_per_worker=n_threads, n_workers=n_workers, memory_limit=ml + 'GB') if plugins is not None and 'pgb' in plugins: register_codec(codecs.PackGeneticBits) client.register_worker_plugin(codecs.CodecPlugin()) return client
download_if_not_exists( "data/model_kf{0}.h5".format(i), "https://jpata.web.cern.ch/jpata/hepaccelerate/model_kf{0}.h5". format(i), ) print( "Trying to connect to dask cluster, please start it with examples/dask_cluster.sh or examples/dask_cluster_gpu.sh" ) if args.dask_server == "debug": multiprocessing_initializer(args) else: client = Client(args.dask_server) plugin = InitializerPlugin(args) client.register_worker_plugin(plugin) print("Processing all datasets") arglist = [] walltime_t0 = time.time() for dataset, fn_pattern, ismc in datasets: filenames = glob.glob(args.datapath + fn_pattern) if len(filenames) == 0: raise Exception( "Could not find any filenames for dataset={0}: {{datapath}}/{{fn_pattern}}={1}/{2}" .format(dataset, args.datapath, fn_pattern)) ichunk = 0 for fn in filenames: nev = len(uproot.open(fn).get("Events")) # Process in chunks of 500k events to limit peak memory usage
def run(self, current_date: datetime, dry_run: bool = False) -> None: """ Run analysis using mozanalysis for a specific experiment. """ global _dask_cluster logger.info("Analysis.run invoked for experiment %s", self.config.experiment.normandy_slug) self.check_runnable(current_date) assert self.config.experiment.start_date is not None # for mypy self.ensure_enrollments(current_date) # set up dask _dask_cluster = _dask_cluster or LocalCluster( dashboard_address=DASK_DASHBOARD_ADDRESS, processes=True, threads_per_worker=1, n_workers=DASK_N_PROCESSES, ) client = Client(_dask_cluster) results = [] if self.log_config: log_plugin = LogPlugin(self.log_config) client.register_worker_plugin(log_plugin) # add profiling plugins # resource_profiling_plugin = ResourceProfilingPlugin( # scheduler=_dask_cluster.scheduler, # project_id=self.log_config.log_project_id, # dataset_id=self.log_config.log_dataset_id, # table_id=self.log_config.task_profiling_log_table_id, # experiment=self.config.experiment.normandy_slug, # ) # _dask_cluster.scheduler.add_plugin(resource_profiling_plugin) # task_monitoring_plugin = TaskMonitoringPlugin( # scheduler=_dask_cluster.scheduler, # project_id=self.log_config.log_project_id, # dataset_id=self.log_config.log_dataset_id, # table_id=self.log_config.task_monitoring_log_table_id, # experiment=self.config.experiment.normandy_slug, # ) # _dask_cluster.scheduler.add_plugin(task_monitoring_plugin) table_to_dataframe = dask.delayed(self.bigquery.table_to_dataframe) for period in self.config.metrics: segment_results = [] time_limits = self._get_timelimits_if_ready(period, current_date) if time_limits is None: logger.info( "Skipping %s (%s); not ready", self.config.experiment.normandy_slug, period.value, ) continue exp = mozanalysis.experiment.Experiment( experiment_slug=self.config.experiment.normandy_slug, start_date=self.config.experiment.start_date.strftime("%Y-%m-%d"), app_id=self._app_id_to_bigquery_dataset(self.config.experiment.app_id), ) analysis_bases = [] for m in self.config.metrics[period]: for analysis_basis in m.metric.analysis_bases: analysis_bases.append(analysis_basis) analysis_bases = list(set(analysis_bases)) if len(analysis_bases) == 0: continue for analysis_basis in analysis_bases: metrics_table = self.calculate_metrics( exp, time_limits, period, analysis_basis, dry_run ) if dry_run: results.append(metrics_table) else: metrics_dataframe = table_to_dataframe(metrics_table) if dry_run: logger.info( "Not calculating statistics %s (%s); dry run", self.config.experiment.normandy_slug, period.value, ) continue segment_labels = ["all"] + [s.name for s in self.config.experiment.segments] for segment in segment_labels: segment_data = self.subset_to_segment(segment, metrics_dataframe) for m in self.config.metrics[period]: segment_results += self.calculate_statistics( m, segment_data, segment, analysis_basis, ).to_dict()["data"] segment_results += self.counts(segment_data, segment, analysis_basis).to_dict()[ "data" ] results.append( self.save_statistics( period, segment_results, self._table_name(period.value, len(time_limits.analysis_windows)), ) ) result_futures = client.compute(results) client.gather(result_futures) # block until futures have finished
]) sitepackages = os.path.join(installdir, 'lib', 'python' + sys.version[:3], 'site-packages') if sitepackages not in sys.path: sys.path.insert(0, sitepackages) def teardown(self, worker): pass client = Client(os.environ['DASK_SCHEDULER']) # one-time setup if True: client.register_worker_plugin(ConfigureXRootD(), 'user_proxy') client.register_worker_plugin( InstallPackage( 'https://github.com/nsmith-/boostedhiggs/archive/dev.zip'), 'boostedhiggs') # newcoffea = DistributeZipball('/home/ncsmith/coffea/dist/coffea-0.6.23.zip') # client.register_worker_plugin(newcoffea, 'coffeaupdate') # jump_assignment = WorkerJumpAssignment() # def put(dask_scheduler=None): # dask_scheduler.add_plugin(WorkerJumpAssignment(list(dask_scheduler.workers.keys()))) # def get(dask_scheduler=None): # for p in dask_scheduler.plugins: # try: # return p.get_jump_mapping() # except AttributeError:
if args.scheduler_file: # We're wanting to submit workers onto other nodes, and *not* run # them locally because we went through the trouble of specifying # a scheduler file that the scheduler and workers will use to # coordinate with one another. logger.info('Using a remote distributed model') client = Client(scheduler_file=args.scheduler_file) else: logger.info('Using a local distributed model') cluster = LocalCluster(n_workers=args.workers, processes=False, silence_logs=logger.level) logger.info("Cluster: %s", cluster) client = Client(cluster) client.register_worker_plugin(WorkerLoggerPlugin(verbose=args.verbose)) logger.info('Client: %s', client) if args.track_workers_file: track_workers_stream = open(args.track_workers_file, 'w') track_workers_func = log_worker_location(track_workers_stream) if args.track_pop_file is not None: track_pop_stream = open(args.track_pop_file, 'w') track_pop_func = log_pop(args.update_interval, track_pop_stream) final_pop = asynchronous.steady_state( client, births=args.max_births, init_pop_size=5,
def start_dask_cluster(environment=os.path.basename( os.environ['CONDA_PREFIX']), worker_profile='Medium Worker', profile='default', region='us-west-2', endpoint=None, worker_min=2, worker_max=20, adaptive_scaling=True, wait_for_cluster=True, cfile=None, use_existing_cluster=True, propagate_env=False): ''' environment - should match the kernel running, and will be set autmatically worker profile - 'Small Worker', 'Medium Worker', or 'Pangeo Worker' (determines available memory in a worker) profile - 'default' is good, but others can be used region - AWS region endpoint - None by default matches region. Set correct endpoint to s3 buckets worker_min - minumum number of workers (for adaptive scaling) worker_max - maximum number of workers adaptive_scaling - Default True. If False, launches worker_max workers wait_for_cluster - Default True. cfile - None. Finds aws credentials in this file use_existing_cluster - Default True. propagate_env - Default False. Set to True when working with Cloud VRTs ''' if not endpoint: endpoint = f's3.{region}.amazonaws.com' set_credentials2(profile=profile, region=region, endpoint=endpoint, cfile=cfile) try: gateway.list_clusters() except: gateway = Gateway() if gateway.list_clusters(): print('Existing Dask clusters:') j = 0 for c in gateway.list_clusters(): print(f'Cluster Index c_idx: {j} / Name:', c.name, c.status) j += 1 else: print('No Cluster running.') # TODO Check if worker_profile is the same, otherwise start new cluster if gateway.list_clusters() and use_existing_cluster: print('Using existing cluster [0].') cluster = gateway.connect(gateway.list_clusters()[0].name) else: print('Starting new cluster.') cluster = gateway.new_cluster(environment=environment, profile=worker_profile) if adaptive_scaling: print(f'Setting Adaptive Scaling min={worker_min}, max={worker_max}') cluster.adapt(minimum=worker_min, maximum=worker_max) else: print(f'Setting Fixed Scaling workers={worker_max}') cluster.scale(worker_max) try: client = Client(cluster) client.close() print('Reconnect client to clear cache') except: pass client = Client(cluster) print( f'client.dashboard_link (for new browser tab/window or dashboard searchbar in Jupyterhub):\n{client.dashboard_link}' ) if wait_for_cluster: target_workers = worker_min if adaptive_scaling else worker_max live_workers = len(list(cluster.scheduler_info['workers'])) t = 0 interval = 2 print( f'Elapsed time to wait for {target_workers} live workers:\n{live_workers}/{target_workers} workers - {t} seconds', end='') while not live_workers >= target_workers: sleep(interval) t += interval print(f'\r{live_workers}/{target_workers} workers - {t} seconds', end='') live_workers = len(client.scheduler_info()['workers']) print(f'\r{live_workers}/{target_workers} workers - {t} seconds') # We need to propagate credentials to the workers #set_credentials(profile=profile,region=region,endpoint=endpoint) if propagate_env: print('Propagating environment variables to workers') class InitWorker(WorkerPlugin): name = "init_worker" def __init__(self, filepath=None, script=None): self.data = {} if filepath: if isinstance(filepath, str): filepath = [filepath] for file_ in filepath: with open(file_, "rb") as f: filename = os.path.basename(file_) self.data[filename] = f.read() if script: filename = f"{uuid.uuid1()}.py" self.data[filename] = script async def setup(self, worker): responses = await asyncio.gather(*[ worker.upload_file( comm=None, filename=filename, data=data, load=True) for filename, data in self.data.items() ]) assert all( len(data) == r["nbytes"] for r, data in zip(responses, self.data.values())) script = f""" \rimport os \ros.environ["AWS_ACCESS_KEY_ID"] = "{os.getenv("AWS_ACCESS_KEY_ID")}" \ros.environ["AWS_SECRET_ACCESS_KEY"] = "{os.getenv("AWS_SECRET_ACCESS_KEY")}" \ros.environ["AWS_DEFAULT_REGION"] = "{os.getenv("AWS_DEFAULT_REGION")}" \ros.environ["GDAL_DISABLE_READDIR_ON_OPEN"] ="EMPTY_DIR" """ plugin = InitWorker(script=script) client.register_worker_plugin(plugin) return client, cluster