def test_multiple_environments_free_assignment():
# Dummy environments with no components for testing.
environments = [
TaskEnvironment(placement=[cpu(0)],
components=[DummyComponent("foo")]),
TaskEnvironment(placement=[cpu(1)], components=[DummyComponent("bar")])
]
with Parla(environments):
for _ in repetitions():
task_results = []
@spawn(vcus=1)
def task():
sleep(0.1)
task_results.append(thread_locals.value)
@spawn(vcus=1)
def task():
sleep(0.1)
task_results.append(thread_locals.value)
@spawn(vcus=1)
def task():
sleep(0.1)
task_results.append(thread_locals.value)
sleep_until(lambda: len(task_results) == 3)
assert set(task_results) == {"foo", "bar"}
def test_multiple_environments_tagged():
# Dummy environments with no components for testing.
environments = [
TaskEnvironment(placement=[cpu(0)],
components=[DummyComponent("foo")],
tags=(threading, )),
TaskEnvironment(placement=[cpu(1)],
components=[DummyComponent("bar")],
tags=(logging, ))
]
with Parla(environments):
for _ in repetitions():
task_results = []
@spawn(tags=(threading, ))
def task():
task_results.append(thread_locals.value)
sleep_until(lambda: len(task_results) == 1)
assert task_results == ["foo"]
task_results = []
@spawn(tags=(logging, ))
def task():
task_results.append(thread_locals.value)
sleep_until(lambda: len(task_results) == 1)
assert task_results == ["bar"]
def runtime_sched():
try:
from parla.cuda import gpu
except:
pass
with Parla() as p:
yield get_scheduler_context().scheduler
def test_parla_ctx():
# Dummy environments with no components for testing.
environments = [TaskEnvironment(placement=[d], components=[]) for d in cpu.devices]
with Parla(environments):
task_results = []
@spawn()
def task():
task_results.append(1)
sleep_until(lambda: len(task_results) == 1)
assert task_results == [1]
def test_placement_multi():
# Dummy environments with no components for testing.
environments = [TaskEnvironment(placement=d, components=[]) for d in combinations(cpu.devices, 2)]
with Parla(environments):
devices = [frozenset((cpu(0), cpu(1))), frozenset((cpu(1), cpu(2))), frozenset((cpu(4), cpu(3)))]
for rep in repetitions():
task_results = []
for (i, dev) in enumerate(devices):
@spawn(placement=dev, ndevices=2)
def task():
task_results.append(frozenset(get_current_devices()))
sleep_until(lambda: len(task_results) == i+1)
assert task_results == devices
def test_dummy_environment_component():
environments = [
TaskEnvironment(placement=[cpu(0)],
components=[DummyComponent("test")])
]
with Parla(environments):
task_results = []
@spawn()
def task():
assert get_current_devices() == [cpu(0)]
task_results.append(thread_locals.value)
sleep_until(lambda: len(task_results) == 1)
assert task_results == ["test"]
def test_multiple_environments_fixed_assignment():
# Dummy environments with no components for testing.
environments = [
TaskEnvironment(placement=[cpu(0)],
components=[DummyComponent("foo")]),
TaskEnvironment(placement=[cpu(1)], components=[DummyComponent("bar")])
]
with Parla(environments):
task_results = []
@spawn(placement=cpu(0))
def task():
task_results.append(thread_locals.value)
@spawn(placement=cpu(1))
def task():
task_results.append(thread_locals.value)
sleep_until(lambda: len(task_results) == 2)
assert set(task_results) == {"foo", "bar"}
def test_scheduler_hook():
task_results = []
from parla.task_runtime import Scheduler
class MyScheduler(Scheduler):
def _assignment_policy(self, task: Task):
sleep(0.05)
task_results.append((task.name, tuple(d.name for d in task.dependees)))
return super(MyScheduler, self)._assignment_policy(task)
with Parla(scheduler_class=MyScheduler):
@spawn()
def task():
sleep(0.05)
task_results.append(1)
@spawn(None, [task])
def task2():
sleep(0.05)
task_results.append(2)
sleep_until(lambda: len(task_results) == 4)
assert task_results == [("task", ("task2",)), 1, ("task2", ()), 2]
def test_multiple_environments_less_good_fit():
# Dummy environments with no components for testing.
environments = [
TaskEnvironment(placement=[cpu(0), cpu(1)],
components=[DummyComponent("foo")]),
TaskEnvironment(placement=[cpu(2), cpu(3), cpu(4)],
components=[DummyComponent("bar")])
]
with Parla(environments):
for _ in repetitions():
task_results = []
# The first two will fit in the the first environment using 0.5 of the environment.
# The next two will spill into the less good (0.33) fit of the second environment.
@spawn(placement=[cpu(1), cpu(2)], vcus=1)
def task():
sleep(0.1)
task_results.append(thread_locals.value)
@spawn(placement=[cpu(1), cpu(2)], vcus=1)
def task():
sleep(0.1)
task_results.append(thread_locals.value)
@spawn(placement=[cpu(1), cpu(2)], vcus=1)
def task():
sleep(0.1)
task_results.append(thread_locals.value)
@spawn(placement=[cpu(1), cpu(2)], vcus=1)
def task():
sleep(0.1)
task_results.append(thread_locals.value)
sleep_until(lambda: len(task_results) == 4)
task_results.sort()
assert task_results == ["bar", "bar", "foo", "foo"]
np.random.seed(10)
# Construct input data
a = np.random.rand(n, n)
a = a @ a.T
# Copy and layout input
a1 = a.copy()
ap = a1.reshape(n // block_size, block_size, n // block_size,
block_size).swapaxes(1, 2)
start = time.perf_counter()
# Call Parla Cholesky result and wait for completion
await cholesky_blocked_inplace(ap)
end = time.perf_counter()
print(end - start, "seconds")
print("Truth", linalg.cholesky(a).T)
# Check result
computed_L = np.tril(a1)
print("Soln", computed_L)
error = np.max(np.absolute(a - computed_L @ computed_L.T))
print("Error", error)
assert (error < 1E-8)
if __name__ == '__main__':
with Parla():
main()
