<gh_stars>100-1000
import os
import tempfile

import pytest

from hypertunity.domain import Domain, Sample
from hypertunity.optimisation import base

from ..jobs import Job, SlurmJob
from ..scheduler import Scheduler
from . import script


@pytest.fixture(scope="module")
def shared_slurm_tmp_dir():
    return "/tmp"


def square(sample: Sample) -> base.EvaluationScore:
    return base.EvaluationScore(sample["x"]**2)


def run_jobs(jobs):
    with Scheduler(n_parallel=2) as scheduler:
        scheduler.dispatch(jobs)
        results = scheduler.collect(n_results=len(jobs), timeout=60.0)
    assert len(results) == len(jobs)
    assert all([r.id == j.id for r, j in zip(results, jobs)])
    return results


@pytest.mark.timeout(10.0)
def test_local_from_script_and_function():
    domain = Domain({
        "x": {0, 1, 2, 3},
        "y": [-1., 1.],
        "z": {"123", "abc"}
    }, seed=7)
    jobs = [Job(task="hypertunity/scheduling/tests/script.py::main",
                args=(*domain.sample().as_namedtuple(),)) for _ in range(10)]
    results = run_jobs(jobs)
    assert all([r.data == script.main(*j.args) for r, j in zip(results, jobs)])


@pytest.mark.timeout(10.0)
def test_local_from_script_and_cmdline_args():
    domain = Domain({
        "x": {0, 1, 2, 3},
        "y": [-1., 1.],
        "z": {"123", "abc"}
    }, seed=7)
    jobs = [Job(task="hypertunity/scheduling/tests/script.py",
                args=(*domain.sample().as_namedtuple(),),
                meta={"binary": "python"}) for _ in range(10)]
    results = run_jobs(jobs)
    assert all([r.data == script.main(*j.args) for r, j in zip(results, jobs)])


@pytest.mark.timeout(10.0)
def test_local_from_script_and_cmdline_named_args():
    domain = Domain({
        "--x": {0, 1, 2, 3},
        "--y": [-1., 1.],
        "--z": {"acb123", "abc"}
    }, seed=7)
    jobs = [Job(task="hypertunity/scheduling/tests/script.py",
                args=domain.sample().as_dict(),
                meta={"binary": "python"}) for _ in range(10)]
    results = run_jobs(jobs)
    assert all([
        r.data == script.main(**{k.lstrip("-"): v for k, v in j.args.items()})
        for r, j in zip(results, jobs)
    ])


@pytest.mark.timeout(10.0)
def test_local_from_fn():
    domain = Domain({"x": [0., 1.]}, seed=7)
    jobs = [Job(task=square, args=(domain.sample(),)) for _ in range(10)]
    results = run_jobs(jobs)
    assert all([r.data.value == square(*j.args).value
                for r, j in zip(results, jobs)])


@pytest.mark.slurm
@pytest.mark.timeout(60.0)
def test_slurm_from_script(shared_slurm_tmp_dir):
    domain = Domain({
        "x": {0, 1, 2, 3},
        "y": [-1., 1.],
        "z": {"123", "abc"}
    }, seed=7)
    jobs, dirs = [], []
    n_jobs = 4
    for i in range(n_jobs):
        sample = domain.sample()
        dirs.append(tempfile.TemporaryDirectory(dir=shared_slurm_tmp_dir))
        jobs.append(SlurmJob(
            task="hypertunity/scheduling/tests/script.py",
            args=(*sample.as_namedtuple(),),
            output_file=f"{os.path.join(dirs[-1].name, 'results.pkl')}",
            meta={"binary": "python", "resources": {"cpu": 1}}
        ))
    results = run_jobs(jobs)
    assert all([r.data == script.main(*j.args) for r, j in zip(results, jobs)])
    # clean-up the temporary dirs
    for d in dirs:
        d.cleanup()

import os
import tempfile

import pytest

from hypertunity.domain import Domain, Sample
from hypertunity.optimisation import base

from ..jobs import Job, SlurmJob
from ..scheduler import Scheduler
from . import script


@pytest.fixture(scope="module")
def shared_slurm_tmp_dir():
    return "/tmp"


def square(sample: Sample) -> base.EvaluationScore:
    return base.EvaluationScore(sample["x"]**2)


def run_jobs(jobs):
    with Scheduler(n_parallel=2) as scheduler:
        scheduler.dispatch(jobs)
        results = scheduler.collect(n_results=len(jobs), timeout=60.0)
    assert len(results) == len(jobs)
    assert all([r.id == j.id for r, j in zip(results, jobs)])
    return results


@pytest.mark.timeout(10.0)
def test_local_from_script_and_function():
    domain = Domain({
        "x": {0, 1, 2, 3},
        "y": [-1., 1.],
        "z": {"123", "abc"}
    }, seed=7)
    jobs = [Job(task="hypertunity/scheduling/tests/script.py::main",
                args=(*domain.sample().as_namedtuple(),)) for _ in range(10)]
    results = run_jobs(jobs)
    assert all([r.data == script.main(*j.args) for r, j in zip(results, jobs)])


@pytest.mark.timeout(10.0)
def test_local_from_script_and_cmdline_args():
    domain = Domain({
        "x": {0, 1, 2, 3},
        "y": [-1., 1.],
        "z": {"123", "abc"}
    }, seed=7)
    jobs = [Job(task="hypertunity/scheduling/tests/script.py",
                args=(*domain.sample().as_namedtuple(),),
                meta={"binary": "python"}) for _ in range(10)]
    results = run_jobs(jobs)
    assert all([r.data == script.main(*j.args) for r, j in zip(results, jobs)])


@pytest.mark.timeout(10.0)
def test_local_from_script_and_cmdline_named_args():
    domain = Domain({
        "--x": {0, 1, 2, 3},
        "--y": [-1., 1.],
        "--z": {"acb123", "abc"}
    }, seed=7)
    jobs = [Job(task="hypertunity/scheduling/tests/script.py",
                args=domain.sample().as_dict(),
                meta={"binary": "python"}) for _ in range(10)]
    results = run_jobs(jobs)
    assert all([
        r.data == script.main(**{k.lstrip("-"): v for k, v in j.args.items()})
        for r, j in zip(results, jobs)
    ])


@pytest.mark.timeout(10.0)
def test_local_from_fn():
    domain = Domain({"x": [0., 1.]}, seed=7)
    jobs = [Job(task=square, args=(domain.sample(),)) for _ in range(10)]
    results = run_jobs(jobs)
    assert all([r.data.value == square(*j.args).value
                for r, j in zip(results, jobs)])


@pytest.mark.slurm
@pytest.mark.timeout(60.0)
def test_slurm_from_script(shared_slurm_tmp_dir):
    domain = Domain({
        "x": {0, 1, 2, 3},
        "y": [-1., 1.],
        "z": {"123", "abc"}
    }, seed=7)
    jobs, dirs = [], []
    n_jobs = 4
    for i in range(n_jobs):
        sample = domain.sample()
        dirs.append(tempfile.TemporaryDirectory(dir=shared_slurm_tmp_dir))
        jobs.append(SlurmJob(
            task="hypertunity/scheduling/tests/script.py",
            args=(*sample.as_namedtuple(),),
            output_file=f"{os.path.join(dirs[-1].name, 'results.pkl')}",
            meta={"binary": "python", "resources": {"cpu": 1}}
        ))
    results = run_jobs(jobs)
    assert all([r.data == script.main(*j.args) for r, j in zip(results, jobs)])
    # clean-up the temporary dirs
    for d in dirs:
        d.cleanup()
# clean-up the temporary dirs