def test_same_input(rsds_env):
url = rsds_env.start([1])
client = Client(url)
f1 = client.submit(comp_fn1, 10)
f2 = client.submit(comp_fn2, f1, f1)
r2 = client.gather(f2)
assert r2 == 0
def test_more_clients(rsds_env):
url = rsds_env.start([1])
client1 = Client(url)
client2 = Client(url)
f1 = client1.submit(comp_fn1, 10)
f2 = client2.submit(comp_fn1, 20)
r2 = client2.gather(f2)
r1 = client1.gather(f1)
assert r1 == 100
assert r2 == 200
class ClusterTest:
Cluster = None
kwargs: dict[str, Any] = {}
def setUp(self):
self.cluster = self.Cluster(2, scheduler_port=0, **self.kwargs)
self.client = Client(self.cluster.scheduler_address)
def tearDown(self):
self.client.close()
self.cluster.close()
@pytest.mark.xfail()
def test_cores(self):
self.client.scheduler_info()
assert len(self.client.nthreads()) == 2
def test_submit(self):
future = self.client.submit(lambda x: x + 1, 1)
assert future.result() == 2
def test_context_manager(self):
with self.Cluster(**self.kwargs) as c:
with Client(c) as e:
assert e.nthreads()
def test_no_workers(self):
with self.Cluster(0, scheduler_port=0, **self.kwargs):
pass
def user_scheduler(
scheduler: Client,
scenario: Scenario,
backend: IUserManagerBackend,
context: dict,
):
"""Schedules users inside user manager on events from scenario.
Args:
scheduler (Client): Dask client to start users
scenario (Scenario): User Scenario
user_manager_id (str): ID of this user manager
backend (IUserManagerBackend): Backend implementation for user manager to use
context (dict): Test context
"""
while True:
for user_id in backend.get_new_users():
fut = scheduler.submit(
user_runner,
scenario=scenario,
user_id=user_id,
backend=backend.get_user_backend(user_id),
context=context,
pure=False,
)
# NOTE: may be better waiting for all futures to finish
fire_and_forget(fut)
backend.send_user_results()
time.sleep(1)
def test_gather_already_finished(rsds_env):
url = rsds_env.start([1])
client = Client(url)
x = client.submit(inc, 10)
assert not x.done()
assert isinstance(x, Future)
assert x.client is client
result = client.gather(x)
assert result == 11
assert x.done()
y = client.submit(inc, 20)
z = client.submit(add, x, y)
result = client.gather(z)
assert result == 11 + 21
def test_dask_cv_single(self):
test_cluster = LocalCluster(1)
test_client = Client(test_cluster)
iris = load_iris()
reg = tree.DecisionTreeClassifier()
cv_score = test_client.submit(cross_val_score, reg, iris.data,
iris.target)
self.assertGreater(cv_score.result().mean(), 0)
test_cluster.scale_up(4)
_cv_results = {
'reg_%i':
test_client.submit(cross_val_score,
tree.DecisionTreeClassifier(min_samples_leaf=i),
iris.data, iris.target)
for i in range(5)
}
cv_results = test_client.gather(list(_cv_results.values()))
for cv_result in cv_results:
self.assertGreaterEqual(cv_result.mean(), 0)
def test_scatter(rsds_env):
url = rsds_env.start([1])
client = Client(url)
client.wait_for_workers(1)
metadata = client.scheduler_info()
worker = list(metadata["workers"].keys())[0]
futures = client.scatter(range(10), workers=[worker])
fut = client.submit(comp_fn3, futures)
assert client.gather(fut) == list(range(1, 11))
def test_recompute_existing(rsds_env):
url = rsds_env.start([1])
client = Client(url)
# assert delayed_fn1(10).compute() == 100
# assert delayed_fn1(10).compute() == 100
f1 = client.submit(comp_fn1, 10)
f2 = client.submit(comp_fn1, 10)
r1, r2 = client.gather([f1, f2])
assert r1 == 100
assert r2 == 100
def test_submit_gather(rsds_env):
url = rsds_env.start([1])
client = Client(url)
f1 = client.submit(comp_fn1, 10)
f2 = client.submit(comp_fn1, 13)
r1, r2 = client.gather([f1, f2])
assert r1 == 100
assert r2 == 130
f1, f2 = client.map(comp_fn1, [20, 30])
r1, r2 = client.gather([f1, f2])
assert r1 == 200
assert r2 == 300
r1 = delayed_fn2(delayed_fn1(2), delayed_fn1(3)).compute()
assert r1 == 50
def run_user(
self,
scenario_name: str,
user_manager_id: str,
backend_address: str,
encoded_context: str,
):
"""Startup function when user is started. Runs scenario user loop.
Args:
scenario_name (str): Name of scenario being run
user_manager_id (str): Unique ID of user manager assigned by scenario
backend_address (str): Address of backend client to receive work and save results
encoded_context (str): Context from test containing previous results
"""
scenario = self.scenarios[scenario_name]
context = decode_context(encoded_context)
client = Client()
# Create buffer actor
buffer_fut = client.submit(
UserBufferActor,
user_manager_id=user_manager_id,
backend_address=backend_address,
actor=True,
)
buffer = buffer_fut.result()
fire_and_forget(buffer_fut)
backend = UserManagerBackend(
user_manager_id=user_manager_id,
buffer=buffer,
address=backend_address,
)
atexit.register(lambda: backend.send_user_results())
user_scheduler(
client,
scenario,
backend,
context,
)
