async def test_create_cluster():
skip_without_credentials()
async with GCPCluster(asynchronous=True,
env_vars={"FOO": "bar"},
security=True) as cluster:
assert cluster.status == Status.running
cluster.scale(2)
await cluster
assert len(cluster.workers) == 2
async with Client(cluster, asynchronous=True) as client:
def inc(x):
return x + 1
def check_env():
import os
return os.environ["FOO"]
assert await client.submit(inc, 10).result() == 11
assert await client.submit(check_env).result() == "bar"
async def test_get_cloud_init():
skip_without_credentials()
cloud_init = GCPCluster.get_cloud_init(docker_args="--privileged")
assert "dask-scheduler" in cloud_init
assert "# Bootstrap" in cloud_init
assert " --privileged " in cloud_init
async def test_create_cluster_sync():
skip_without_credentials()
cluster = GCPCluster(n_workers=1)
client = Client(cluster)
def inc(x):
return x + 1
assert client.submit(inc, 10).result() == 11
def create(n_workers=None, env_var_file=None, name=None, **kwargs):
env_vars = {}
if env_var_file:
with open(env_var_file, "r") as f:
env_vars = json.load(f)
global cluster
cluster = GCPCluster(name=name,
n_workers=n_workers,
env_vars=env_vars,
**kwargs)
print("Cluster created")
def test_create_rapids_cluster_sync():
skip_without_credentials()
cluster = GCPCluster(
source_image="projects/nv-ai-infra/global/images/packer-1607527229",
network="dask-gcp-network-test",
zone="us-east1-c",
machine_type="n1-standard-1",
filesystem_size=50,
ngpus=2,
gpu_type="nvidia-tesla-t4",
docker_image="rapidsai/rapidsai:cuda11.0-runtime-ubuntu18.04-py3.8",
worker_class="dask_cuda.CUDAWorker",
worker_options={"rmm_pool_size": "15GB"},
asynchronous=False,
bootstrap=False,
)
cluster.scale(1)
client = Client(cluster) # noqa
client.wait_for_workers(2)
def gpu_mem():
from pynvml.smi import nvidia_smi
nvsmi = nvidia_smi.getInstance()
return nvsmi.DeviceQuery("memory.free, memory.total")
results = client.run(gpu_mem)
for w, res in results.items():
assert "total" in res["gpu"][0]["fb_memory_usage"].keys()
print(res)
cluster.close()
async def test_create_rapids_cluster():
skip_without_credentials()
async with GCPCluster(
source_image=
"projects/nv-ai-infra/global/images/ngc-docker-11-20200916",
zone="us-east1-c",
machine_type="n1-standard-1",
filesystem_size=50,
ngpus=2,
gpu_type="nvidia-tesla-t4",
docker_image="rapidsai/rapidsai:cuda11.0-runtime-ubuntu18.04-py3.8",
worker_class="dask_cuda.CUDAWorker",
worker_options={"rmm_pool_size": "15GB"},
asynchronous=True,
auto_shutdown=True,
bootstrap=False,
) as cluster:
assert cluster.status == Status.running
cluster.scale(1)
await cluster
assert len(cluster.workers) == 1
client = Client(cluster, asynchronous=True) # noqa
await client
def gpu_mem():
from pynvml.smi import nvidia_smi
nvsmi = nvidia_smi.getInstance()
return nvsmi.DeviceQuery("memory.free, memory.total")
results = await client.run(gpu_mem)
for w, res in results.items():
assert "total" in res["gpu"][0]["fb_memory_usage"].keys()
print(res)
async def test_init():
skip_without_credentials()
cluster = GCPCluster(asynchronous=True)
assert cluster.status == Status.created
async def test_get_cloud_init():
skip_without_credentials()
cloud_init = GCPCluster.get_cloud_init()
assert "dask-scheduler" in cloud_init
assert "# Bootstrap" in cloud_init
