diff --git a/src/cryo_challenge/_map_to_map/gromov_wasserstein/gw_weighted_voxels.py b/src/cryo_challenge/_map_to_map/gromov_wasserstein/gw_weighted_voxels.py index a95fbe7..827a8f8 100644 --- a/src/cryo_challenge/_map_to_map/gromov_wasserstein/gw_weighted_voxels.py +++ b/src/cryo_challenge/_map_to_map/gromov_wasserstein/gw_weighted_voxels.py @@ -6,7 +6,7 @@ from dask import delayed, compute from dask.distributed import Client from dask.diagnostics import ProgressBar -from dask_hpc_runner import SlurmRunner +# from dask_hpc_runner import SlurmRunner from cryo_challenge._preprocessing.fourier_utils import downsample_volume @@ -410,13 +410,14 @@ def main(args): if __name__ == "__main__": args = parse_args() if args.slurm: - job_id = os.environ["SLURM_JOB_ID"] - with SlurmRunner( - scheduler_file=args.scheduler_file, - ) as runner: - # The runner object contains the scheduler address and can be passed directly to a client - with Client(runner) as client: - get_distance_matrix_dask_gw = main(args) + pass + # job_id = os.environ["SLURM_JOB_ID"] + # with SlurmRunner( + # scheduler_file=args.scheduler_file, + # ) as runner: + # # The runner object contains the scheduler address and can be passed directly to a client + # with Client(runner) as client: + # get_distance_matrix_dask_gw = main(args) else: with Client(local_directory=args.local_directory) as client: diff --git a/src/cryo_challenge/_map_to_map/map_to_map_distance.py b/src/cryo_challenge/_map_to_map/map_to_map_distance.py index 329e7c4..6a37e64 100644 --- a/src/cryo_challenge/_map_to_map/map_to_map_distance.py +++ b/src/cryo_challenge/_map_to_map/map_to_map_distance.py @@ -7,7 +7,7 @@ import mrcfile import numpy as np from dask.distributed import Client -from dask_hpc_runner import SlurmRunner +# from dask_hpc_runner import SlurmRunner from .gromov_wasserstein.gw_weighted_voxels import get_distance_matrix_dask_gw @@ -499,25 +499,26 @@ def get_distance_matrix(self, maps1, maps2, global_store_of_running_results): maps2 = maps2.reshape((len(maps2),) + maps1.shape[1:]) if extra_params["slurm"]: - job_id = os.environ["SLURM_JOB_ID"] - scheduler_file = os.path.join( - extra_params["scheduler_file_dir"], f"scheduler-{job_id}.json" - ) - with SlurmRunner( - scheduler_file=scheduler_file, - ) as runner: - # The runner object contains the scheduler address and can be passed directly to a client - with Client(runner) as client: - distance_matrix_dask_gw = get_distance_matrix_dask_gw( - volumes_i=maps1, - volumes_j=maps2, - top_k=extra_params["top_k"], - n_downsample_pix=extra_params["n_downsample_pix"], - exponent=extra_params["exponent"], - cost_scale_factor=extra_params["cost_scale_factor"], - scheduler=extra_params["scheduler"], - element_wise=extra_params["element_wise"], - ) + pass + # job_id = os.environ["SLURM_JOB_ID"] + # scheduler_file = os.path.join( + # extra_params["scheduler_file_dir"], f"scheduler-{job_id}.json" + # ) + # with SlurmRunner( + # scheduler_file=scheduler_file, + # ) as runner: + # # The runner object contains the scheduler address and can be passed directly to a client + # with Client(runner) as client: + # distance_matrix_dask_gw = get_distance_matrix_dask_gw( + # volumes_i=maps1, + # volumes_j=maps2, + # top_k=extra_params["top_k"], + # n_downsample_pix=extra_params["n_downsample_pix"], + # exponent=extra_params["exponent"], + # cost_scale_factor=extra_params["cost_scale_factor"], + # scheduler=extra_params["scheduler"], + # element_wise=extra_params["element_wise"], + # ) else: local_directory = extra_params["local_directory"]