def test_function_name():
def foo():
pass
job = Job()
assert job.function(foo).name == "foo/0"
assert job.function(foo).name == "foo/1"
def test_cluster_add_worker():
with LocalCluster() as cluster:
cluster.start_worker()
client = cluster.client()
job = Job()
job.program(["echo", "hello"], stdout="out.txt")
job_id = client.submit(job)
client.wait_for_jobs([job_id])
check_file_contents("out.txt", "hello\n")
def prepare_job_client(hq_env: HqEnv,
with_worker=True,
**job_args) -> Tuple[Job, Client]:
hq_env.start_server()
if with_worker:
hq_env.start_worker()
client = Client(hq_env.server_dir)
return (Job(**job_args), client)
def test_wait_for_job(hq_env: HqEnv):
(job, client) = prepare_job_client(hq_env)
job.program(
args=bash("exit 1"),
)
job_id = client.submit(job)
with pytest.raises(FailedJobsException):
client.wait_for_jobs([job_id])
job = Job()
job.program(
args=bash("echo Test1 > output"),
)
job_id = client.submit(job)
client.wait_for_jobs([job_id])
check_file_contents("output", "Test1\n")
def test_visualization():
def fn():
pass
job = Job()
a = job.function(fn, name="a")
b1 = job.function(fn, name="b1", deps=[a])
b2 = job.function(fn, name="b2", deps=[a])
c1 = job.function(fn, name="c1", deps=[b1])
c2 = job.function(fn, name="c2", deps=[b2])
job.function(fn, name="d", deps=[c1, c2])
with NamedTemporaryFile() as f:
visualize_job(job, f.name)
check_file_contents(
f.name,
"""digraph job {
a;
b1;
a -> b1;
b2;
a -> b2;
c1;
b1 -> c1;
c2;
b2 -> c2;
d;
c1 -> d;
c2 -> d;
}
""",
)
def test_submit_python_prologue(hq_env: HqEnv):
init_sh = Path("init.sh")
init_sh.write_text("""export ABC=xyz""")
hq_env.start_server()
hq_env.start_worker()
client = Client(
hq_env.server_dir, python_env=PythonEnv(prologue=f"source {init_sh.resolve()}")
)
def body():
print(os.environ.get("ABC"))
job = Job()
job.function(body, stdout="out")
job_id = client.submit(job)
wait_for_job_state(hq_env, job_id, "FINISHED")
check_file_contents("out", "xyz\n")
def test_cluster_create_client():
with LocalCluster() as cluster:
client = cluster.client()
job = Job()
job.program(["hostname"])
client.submit(job)