Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Properly handle dask change #3361

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 1 addition & 3 deletions ci/build_cpp.sh
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,6 @@ rapids-print-env

rapids-logger "Begin cpp build"

LIBRAFT_CHANNEL=$(rapids-get-artifact ci/raft/pull-request/1333/c60c0cb/raft_conda_cpp_cuda11_$(arch).tar.gz)

rapids-mamba-retry mambabuild --channel "${LIBRAFT_CHANNEL}" conda/recipes/libcugraph
rapids-mamba-retry mambabuild conda/recipes/libcugraph

rapids-upload-conda-to-s3 cpp
13 changes: 0 additions & 13 deletions ci/build_python.sh
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,6 @@ export CMAKE_GENERATOR=Ninja
rapids-print-env

CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
PY_VER=${RAPIDS_PY_VERSION//./}
LIBRAFT_CHANNEL=$(rapids-get-artifact ci/raft/pull-request/1333/c60c0cb/raft_conda_cpp_cuda11_$(arch).tar.gz)
RAFT_CHANNEL=$(rapids-get-artifact ci/raft/pull-request/1333/c60c0cb/raft_conda_python_cuda11_${PY_VER}_$(arch).tar.gz)

rapids-logger "Begin py build"

Expand All @@ -21,39 +18,29 @@ rapids-logger "Begin py build"
rapids-mamba-retry mambabuild \
--no-test \
--channel "${CPP_CHANNEL}" \
--channel "${LIBRAFT_CHANNEL}" \
--channel "${RAFT_CHANNEL}" \
conda/recipes/pylibcugraph

rapids-mamba-retry mambabuild \
--no-test \
--channel "${CPP_CHANNEL}" \
--channel "${LIBRAFT_CHANNEL}" \
--channel "${RAFT_CHANNEL}" \
--channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
conda/recipes/cugraph

rapids-mamba-retry mambabuild \
--no-test \
--channel "${CPP_CHANNEL}" \
--channel "${LIBRAFT_CHANNEL}" \
--channel "${RAFT_CHANNEL}" \
--channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
conda/recipes/cugraph-service

rapids-mamba-retry mambabuild \
--no-test \
--channel "${CPP_CHANNEL}" \
--channel "${LIBRAFT_CHANNEL}" \
--channel "${RAFT_CHANNEL}" \
--channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
conda/recipes/cugraph-pyg

rapids-mamba-retry mambabuild \
--no-test \
--channel "${CPP_CHANNEL}" \
--channel "${LIBRAFT_CHANNEL}" \
--channel "${RAFT_CHANNEL}" \
--channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
--channel dglteam \
--channel pytorch \
Expand Down
2 changes: 0 additions & 2 deletions ci/test_cpp.sh
Original file line number Diff line number Diff line change
Expand Up @@ -19,15 +19,13 @@ conda activate test
set -u

CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
LIBRAFT_CHANNEL=$(rapids-get-artifact ci/raft/pull-request/1333/c60c0cb/raft_conda_cpp_cuda11_$(arch).tar.gz)
RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}/
mkdir -p "${RAPIDS_TESTS_DIR}"

rapids-print-env

rapids-mamba-retry install \
--channel "${CPP_CHANNEL}" \
--channel "${LIBRAFT_CHANNEL}" \
libcugraph libcugraph_etl libcugraph-tests

rapids-logger "Check GPU usage"
Expand Down
5 changes: 0 additions & 5 deletions ci/test_python.sh
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,6 @@ set -u
rapids-logger "Downloading artifacts from previous jobs"
CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
PY_VER=${RAPIDS_PY_VERSION//./}
LIBRAFT_CHANNEL=$(rapids-get-artifact ci/raft/pull-request/1333/c60c0cb/raft_conda_cpp_cuda11_$(arch).tar.gz)
RAFT_CHANNEL=$(rapids-get-artifact ci/raft/pull-request/1333/c60c0cb/raft_conda_python_cuda11_${PY_VER}_$(arch).tar.gz)

RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}
RAPIDS_COVERAGE_DIR=${RAPIDS_COVERAGE_DIR:-"${PWD}/coverage-results"}
Expand All @@ -34,8 +31,6 @@ rapids-print-env
rapids-mamba-retry install \
--channel "${CPP_CHANNEL}" \
--channel "${PYTHON_CHANNEL}" \
--channel "${LIBRAFT_CHANNEL}" \
--channel "${RAFT_CHANNEL}" \
libcugraph \
pylibcugraph \
cugraph \
Expand Down
10 changes: 5 additions & 5 deletions python/cugraph/cugraph/dask/comms/comms.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# Copyright (c) 2018-2022, NVIDIA CORPORATION.
# Copyright (c) 2018-2023, NVIDIA CORPORATION.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
Expand Down Expand Up @@ -28,7 +28,7 @@
raise
from pylibraft.common.handle import Handle
from cugraph.dask.comms.comms_wrapper import init_subcomms as c_init_subcomms
from dask.distributed import default_client
from dask.distributed import default_client, get_worker
from cugraph.dask.common import read_utils
import math

Expand Down Expand Up @@ -234,15 +234,15 @@ def get_handle(sID):
"""
Returns the handle from within the worker using the sessionstate.
"""
sessionstate = get_raft_comm_state(sID)
sessionstate = get_raft_comm_state(sID, get_worker())
return sessionstate["handle"]


def get_worker_id(sID):
"""
Returns the worker's sessionId from within the worker.
"""
sessionstate = get_raft_comm_state(sID)
sessionstate = get_raft_comm_state(sID, get_worker())
return sessionstate["wid"]


Expand All @@ -257,5 +257,5 @@ def get_n_workers(sID=None):
if sID is None:
return read_utils.get_n_workers()
else:
sessionstate = get_raft_comm_state(sID)
sessionstate = get_raft_comm_state(sID, get_worker())
return sessionstate["nworkers"]