def tls_cluster_context( worker_kwargs=None, scheduler_kwargs=None, security=None, **kwargs ): security = security or tls_only_security() worker_kwargs = assoc(worker_kwargs or {}, "security", security) scheduler_kwargs = assoc(scheduler_kwargs or {}, "security", security) with cluster( worker_kwargs=worker_kwargs, scheduler_kwargs=scheduler_kwargs, **kwargs ) as (s, workers): yield s, workers
async def run_traffic_jam(nsends, nbytes): # This test eats `nsends * nbytes` bytes in RAM np = pytest.importorskip("numpy") from distributed.protocol import to_serialize data = bytes(np.random.randint(0, 255, size=(nbytes, )).astype("u1").data) async with EchoServer() as e: comm = await connect(e.address) b = BatchedSend(interval=0.01) b.start(comm) msg = {"x": to_serialize(data)} for i in range(nsends): b.send(assoc(msg, "i", i)) if np.random.random() > 0.5: await asyncio.sleep(0.001) results = [] count = 0 while len(results) < nsends: # If this times out then I think it's a backpressure issue # Somehow we're able to flood the socket so that the receiving end # loses some of our messages L = await asyncio.wait_for(comm.read(), 5) count += 1 results.extend(r["i"] for r in L) assert count == b.batch_count == e.count assert b.message_count == nsends assert results == list(range(nsends)) await comm.close() # external closing await b.close()