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, )