class DaskDelegate(Delegate):
type: str = "dask"
def __init__(self, delegate_config: DaskDelegateConfig):
super()
self.delegate_config = delegate_config
self.cache_provider = self.delegate_config.cache_provider
# Attempt to load the global Dask client.
try:
self.client = get_client()
except ValueError as _:
if self.delegate_config.kube_cluster is not None:
self.client = Client(self.delegate_config.kube_cluster)
print(self.delegate_config.kube_cluster)
else:
self.client = Client(f"{self.delegate_config.dask_cluster_address}:{self.delegate_config.dask_cluster_port}")
# Setup functions to be run on the schedule.
def __scheduler_job_exists(dask_scheduler, job_id: str) -> bool:
return job_id in dask_scheduler.tasks
def __scheduler_job_state(dask_scheduler, job_id: str) -> TaskState:
return dask_scheduler.tasks[job_id].state
self.scheduler_job_exists = __scheduler_job_exists
self.scheduler_job_state = __scheduler_job_state
def __job_state(self, job_id: str) -> TaskState:
return self.client.run_on_scheduler(self.scheduler_job_state, job_id=job_id)
def connect(self) -> bool:
# No need to connect.
return True
def test_connection(self) -> bool:
# Shim this out until I figure out a good way to test a Dask and Redis connection.
return True
def create_job(self, job_id: str) -> bool:
# No concept of creating a job.
return True
def start_job(self, job_id: str, work: Callable, *args, **kwargs) -> bool:
if self.job_exists(job_id) or self.job_complete(job_id):
return False
# Parse and replace instances of the internal `result://` proxy protocol.
# In short, this allows for callees to reference an in-progress or remote job without needing direct access.
function_args = [(self.client.get_dataset(arg.replace("result://", "")) if isinstance(arg, str) and arg.startswith("result://") else arg) for arg in args]
# Create a job to run the desired function.
job_future: Future = self.client.submit(work, *function_args, **kwargs, key=job_id, pure=False)
# Start additional cache job which depends on the results of the previous.
cache_future: Future = self.client.submit(self.cache_provider.put, *[job_id, job_future], pure=False)
# Publish the job as a dataset to maintain state across requests.
self.client.publish_dataset(job_future, name=job_id, override=True)
self.client.publish_dataset(cache_future, override=True)
return True
def stop_job(self, job_id: str) -> bool:
if not self.job_exists(job_id):
return False
try:
# Iterate through the dependencies of this job.
dependencies = self.client.run_on_scheduler(lambda dask_scheduler: [(state.key) for state in dask_scheduler.tasks[id].dependencies])
# Filter out any weak depenencies. Strong dependencies are suffixed with "/" and the name of the job.
dependencies = [(dependency) for dependency in dependencies if dependency.replace(id, "").startswith("/")]
futures = [(Future(key)) for key in dependencies]
futures.append(Future(job_id))
except KeyError:
# do nothing if no dependencies
pass
self.client.cancel(Future(job_id))
self.client.unpublish_dataset(job_id)
# Hacky fix -- Simulation processes continue executing EVEN IF the parent task is killed.
def hacky():
os.system("pkill -f 'Simulation.out'")
self.client.run(hacky, nanny=True)
return True
def job_status(self, job_id: str) -> JobStatus:
# If the job is complete (results exist as a dataset or in the vault).
if self.job_complete(job_id):
status = JobStatus()
status.status_id = JobState.DONE
status.status_text = "The job is complete."
status.has_failed = False
status.is_done = True
return status
# If the job doesn't exist.
if not self.job_exists(job_id):
status = JobStatus()
status.status_id = JobState.DOES_NOT_EXIST
status.status_text = f"A job with job_id: '{job_id}' does not exist."
status.has_failed = True
status.is_done = False
return status
status_mapping = {
"released": (JobState.STOPPED, "The job is known but not actively computing or in memory."),
"waiting": (JobState.WAITING, "The job is waiting for dependencies to arrive in memory."),
"no-worker": (JobState.WAITING, "The job is waiting for a worker to become available."),
"processing": (JobState.RUNNING, "The job is running."),
"memory": (JobState.DONE, "The job is done and is being held in memory."),
"erred": (JobState.FAILED, "The job has failed."),
"done": (JobState.DONE, "The job is done and has been cached / stored on disk.")
}
# Grab the task state from the scheduler.
future_status = self.__job_state(job_id)
status = JobStatus()
status.status_id = status_mapping[future_status][0]
status.status_text = status_mapping[future_status][1]
status.is_done = status.status_id is JobState.DONE
status.has_failed = status.status_id is JobState.FAILED
return status
def job_results(self, job_id: str):
# The results of this job may exist on the client dataset.
if job_id in self.client.datasets:
print("[DEBUG] Getting results from dataset.")
return self.client.get_dataset(name=job_id).result()
# If the results are not in the cache, raise an exception.
if not self.cache_provider.exists(job_id):
raise Exception(f"Result with ID '{job_id}' does not exist in the cache.")
return self.cache_provider.get(job_id)
def job_complete(self, job_id: str) -> bool:
# Finished job results must exist within the cache for it to be considered 'done'.
return self.cache_provider.exists(job_id)
def job_exists(self, job_id: str) -> bool:
# Check if the job exists in the scheduler.
return self.client.run_on_scheduler(self.scheduler_job_exists, job_id=job_id)
def get_remote_dependency(self, dependency_id: str):
# Check to see if the job exists as a dataset.
dependency = self.client.get_dataset(name=dependency_id)
if dependency is not None:
return dependency
raise Exception("Something broke, dependency does not exist within distributed memory.")
class DaskExecutor(BaseExecutor):
"""
DaskExecutor submits tasks to a Dask Distributed cluster.
"""
def __init__(self, cluster_address=None):
super().__init__(parallelism=0)
if cluster_address is None:
cluster_address = conf.get('dask', 'cluster_address')
if not cluster_address:
raise ValueError(
'Please provide a Dask cluster address in airflow.cfg')
self.cluster_address = cluster_address
# ssl / tls parameters
self.tls_ca = conf.get('dask', 'tls_ca')
self.tls_key = conf.get('dask', 'tls_key')
self.tls_cert = conf.get('dask', 'tls_cert')
self.client: Optional[Client] = None
self.futures: Optional[Dict[Future, TaskInstanceKeyType]] = None
def start(self) -> None:
if self.tls_ca or self.tls_key or self.tls_cert:
security = Security(
tls_client_key=self.tls_key,
tls_client_cert=self.tls_cert,
tls_ca_file=self.tls_ca,
require_encryption=True,
)
else:
security = None
self.client = Client(self.cluster_address, security=security)
self.futures = {}
def execute_async(self,
key: TaskInstanceKeyType,
command: CommandType,
queue: Optional[str] = None,
executor_config: Optional[Any] = None) -> None:
def airflow_run():
return subprocess.check_call(command, close_fds=True)
if not self.client:
raise AirflowException(NOT_STARTED_MESSAGE)
future = self.client.submit(airflow_run, pure=False)
self.futures[future] = key # type: ignore
def _process_future(self, future: Future) -> None:
if not self.futures:
raise AirflowException(NOT_STARTED_MESSAGE)
if future.done():
key = self.futures[future]
if future.exception():
self.log.error("Failed to execute task: %s",
repr(future.exception()))
self.fail(key)
elif future.cancelled():
self.log.error("Failed to execute task")
self.fail(key)
else:
self.success(key)
self.futures.pop(future)
def sync(self) -> None:
if not self.futures:
raise AirflowException(NOT_STARTED_MESSAGE)
# make a copy so futures can be popped during iteration
for future in self.futures.copy():
self._process_future(future)
def end(self) -> None:
if not self.client:
raise AirflowException(NOT_STARTED_MESSAGE)
if not self.futures:
raise AirflowException(NOT_STARTED_MESSAGE)
self.client.cancel(list(self.futures.keys()))
for future in as_completed(self.futures.copy()):
self._process_future(future)
def terminate(self):
if not self.futures:
raise AirflowException(NOT_STARTED_MESSAGE)
self.client.cancel(self.futures.keys())
self.end()
class Runner:
def __init__(self, input_file):
import yaml
try:
with open(input_file) as f:
self.params = yaml.safe_load(f)
except Exception as exc:
raise exc
self.operations = {}
self.operations['computations'] = [
spatial_mean, temporal_mean, climatology, anomaly
]
self.operations['readwrite'] = [
writefile, openfile, readfile, deletefile
]
self.operations['write'] = [writefile]
self.operations['read'] = [openfile, readfile]
self.client = None
def create_cluster(self, job_scheduler, maxcore, walltime, memory, queue,
wpn):
""" Creates a dask cluster using dask_jobqueue
"""
logger.warning('Creating a dask cluster using dask_jobqueue')
logger.warning(f'Job Scheduler: {job_scheduler}')
logger.warning(f'Memory size for each node: {memory}')
logger.warning(f'Number of cores for each node: {maxcore}')
logger.warning(f'Number of workers for each node: {wpn}')
from dask_jobqueue import PBSCluster, SLURMCluster
job_schedulers = {'pbs': PBSCluster, 'slurm': SLURMCluster}
# Note about OMP_NUM_THREADS=1, --threads 1:
# These two lines are to ensure that each benchmark workers
# only use one threads for benchmark.
# in the job script one sees twice --nthreads,
# but it get overwritten by --nthreads 1
cluster = job_schedulers[job_scheduler](
cores=maxcore,
memory=memory,
processes=wpn,
local_directory='$TMPDIR',
interface='ib0',
queue=queue,
walltime=walltime,
env_extra=['OMP_NUM_THREADS=1'],
extra=['--nthreads 1'],
project='ntdd0004',
)
self.client = Client(cluster)
logger.warning('************************************\n'
'Job script created by dask_jobqueue:\n'
f'{cluster.job_script()}\n'
'***************************************')
logger.warning(
f'Dask cluster dashboard_link: {self.client.cluster.dashboard_link}'
)
def run(self):
logger.warning('Reading configuration YAML config file')
operation_choice = self.params['operation_choice']
machine = self.params['machine']
job_scheduler = self.params['job_scheduler']
queue = self.params['queue']
walltime = self.params['walltime']
maxmemory_per_node = self.params['maxmemory_per_node']
maxcore_per_node = self.params['maxcore_per_node']
chunk_per_worker = self.params['chunk_per_worker']
freq = self.params['freq']
spil = self.params['spil']
output_dir = self.params.get('output_dir', results_dir)
now = datetime.datetime.now()
output_dir = os.path.join(output_dir, f'{machine}/{str(now.date())}')
os.makedirs(output_dir, exist_ok=True)
parameters = self.params['parameters']
num_workers = parameters['number_of_workers_per_nodes']
num_threads = parameters.get('number_of_threads_per_workers', 1)
num_nodes = parameters['number_of_nodes']
chunking_schemes = parameters['chunking_scheme']
io_formats = parameters['io_format']
filesystems = parameters['filesystem']
fixed_totalsize = parameters['fixed_totalsize']
chsz = parameters['chunk_size']
writefile_dir = parameters['writefile_dir']
for wpn in num_workers:
self.create_cluster(
job_scheduler=job_scheduler,
maxcore=maxcore_per_node,
walltime=walltime,
memory=maxmemory_per_node,
queue=queue,
wpn=wpn,
)
for num in num_nodes:
self.client.cluster.scale(num * wpn)
cluster_wait(self.client, num * wpn)
timer = DiagnosticTimer()
# dfs = []
logger.warning(
'#####################################################################\n'
f'Dask cluster:\n'
f'\t{self.client.cluster}\n')
now = datetime.datetime.now()
csv_filename = f"{output_dir}/compute_study_{now.strftime('%Y-%m-%d_%H-%M-%S')}.csv"
for chunk_size in chsz:
for io_format in io_formats:
for filesystem in filesystems:
if filesystem == 's3':
profile = parameters['profile']
bucket = parameters['bucket']
endpoint_url = parameters['endpoint_url']
fs = fsspec.filesystem(
's3',
profile=profile,
anon=False,
client_kwargs={
'endpoint_url': endpoint_url
},
)
root = f'{bucket}/test1'
elif filesystem == 'posix':
fs = LocalFileSystem()
root = writefile_dir
if not os.path.isdir(f'{root}'):
os.makedirs(f'{root}')
for chunking_scheme in chunking_schemes:
logger.warning(
f'Benchmark starting with: \n\tworker_per_node = {wpn},'
f'\n\tnum_nodes = {num}, \n\tchunk_size = {chunk_size},'
f'\n\tchunking_scheme = {chunking_scheme},'
f'\n\tchunk per worker = {chunk_per_worker}'
f'\n\tio_format = {io_format}'
f'\n\tfilesystem = {filesystem}')
ds, chunks = timeseries(
fixed_totalsize=fixed_totalsize,
chunk_per_worker=chunk_per_worker,
chunk_size=chunk_size,
chunking_scheme=chunking_scheme,
io_format=io_format,
num_nodes=num,
freq=freq,
worker_per_node=wpn,
)
# wait(ds)
dataset_size = format_bytes(ds.nbytes)
logger.warning(ds)
logger.warning(
f'Dataset total size: {dataset_size}')
for op in self.operations[operation_choice]:
with timer.time(
'runtime',
operation=op.__name__,
fixed_totalsize=fixed_totalsize,
chunk_size=chunk_size,
chunk_per_worker=chunk_per_worker,
dataset_size=dataset_size,
worker_per_node=wpn,
threads_per_worker=num_threads,
num_nodes=num,
chunking_scheme=chunking_scheme,
io_format=io_format,
filesystem=filesystem,
root=root,
machine=machine,
maxmemory_per_node=
maxmemory_per_node,
maxcore_per_node=maxcore_per_node,
spil=spil,
):
fname = f'{chunk_size}{chunking_scheme}{filesystem}{num}'
if op.__name__ == 'writefile':
print(ds.sst.data.chunksize)
filename = op(
ds, fs, io_format, root, fname)
elif op.__name__ == 'openfile':
ds = op(fs, io_format, root,
chunks, chunk_size)
elif op.__name__ == 'deletefile':
ds = op(fs, io_format, root,
filename)
else:
op(ds)
# kills ds, and every other dependent computation
logger.warning('Computation done')
self.client.cancel(ds)
temp_df = timer.dataframe()
temp_df.to_csv(csv_filename, index=False)
# dfs.append(temp_df)
# now = datetime.datetime.now()
# filename = f"{output_dir}/compute_study_{now.strftime('%Y-%m-%d_%H-%M-%S')}.csv"
# df = pd.concat(dfs)
# df.to_csv(filename, index=False)
logger.warning(
f'Persisted benchmark result file: {csv_filename}')
logger.warning(
'Shutting down the client and cluster before changing number of workers per nodes'
)
self.client.cluster.close()
logger.warning('Cluster shutdown finished')
self.client.close()
logger.warning('Client shutdown finished')
logger.warning('=====> The End <=========')
