async def test_batch_with_project_mixins(assets: pathlib.Path,
client: Client) -> None:
ws = assets / "with_project_yaml"
config_dir = ConfigDir(
workspace=ws,
config_dir=ws,
)
cl = BatchLocalCL(config_dir, client)
try:
flow = await RunningBatchFlow.create(cl, "batch", "bake-id")
task = await flow.get_task((), "test_mixin", needs={}, state={})
assert task.image == "mixin-image"
assert task.preset == "mixin-preset"
task = await flow.get_task((), "test_mixin_cmd", needs={}, state={})
assert task.image == "mixin-image"
assert task.cmd == "command -o --option arg1 arg2"
task = await flow.get_task((), "test_mixin_bash", needs={}, state={})
assert task.image == "mixin-image"
assert task.cmd == (
"bash -euo pipefail -c 'command -o --option arg1 arg2\n"
"command2 -o --option arg1 arg2\n'")
task = await flow.get_task((), "test_mixin_python", needs={}, state={})
assert task.image == "mixin-image"
assert task.cmd == "python3 -uc 'print(\"hello neuro-flow\")\n'"
finally:
await cl.close()
async def test_project_level_defaults_live(assets: pathlib.Path,
client: Client) -> None:
ws = assets / "with_project_yaml"
config_dir = ConfigDir(
workspace=ws,
config_dir=ws,
)
cl = LiveLocalCL(config_dir, client)
try:
flow = await RunningLiveFlow.create(cl, "live")
job = await flow.get_job("test", {})
assert "tag-a" in job.tags
assert "tag-b" in job.tags
assert job.env["global_a"] == "val-a"
assert job.env["global_b"] == "val-b"
assert job.env["global_b"] == "val-b"
assert job.volumes == [
"storage:common:/mnt/common:rw",
"storage:dir:/var/dir:ro",
]
assert job.workdir == RemotePath("/global/dir")
assert job.life_span == 100800.0
assert job.preset == "cpu-large"
assert job.schedule_timeout == 2157741.0
assert job.image == "image:banana"
finally:
await cl.close()
async def test_early_images(assets: pathlib.Path, client: Client) -> None:
ws = assets / "batch_images"
config_dir = ConfigDir(
workspace=ws,
config_dir=ws,
)
cl = BatchLocalCL(config_dir, client)
try:
flow = await RunningBatchFlow.create(cl, "batch", "bake-id")
assert flow.early_images["image1"].ref == "image:main"
assert flow.early_images["image1"].context == ws / "dir"
assert flow.early_images["image1"].dockerfile == ws / "dir/Dockerfile"
action = await flow.get_action_early("action")
assert action.early_images["image_early"].ref == "image:banana1"
assert action.early_images["image_early"].context == ws / "dir"
assert action.early_images[
"image_early"].dockerfile == ws / "dir/Dockerfile"
assert action.early_images["image_late"].ref == "image:banana2"
assert action.early_images["image_late"].context is None
assert action.early_images["image_late"].dockerfile is None
finally:
await cl.close()
async def test_batch_with_project_globals(assets: pathlib.Path,
client: Client) -> None:
ws = assets / "with_project_yaml"
config_dir = ConfigDir(
workspace=ws,
config_dir=ws,
)
cl = BatchLocalCL(config_dir, client)
try:
flow = await RunningBatchFlow.create(cl, "batch", "bake-id")
task = await flow.get_task((), "task", needs={}, state={})
assert "tag-a" in task.tags
assert "tag-b" in task.tags
assert task.env["global_a"] == "val-a"
assert task.env["global_b"] == "val-b"
assert task.volumes == [
"storage:common:/mnt/common:rw",
"storage:dir:/var/dir:ro",
]
assert task.workdir == RemotePath("/global/dir")
assert task.life_span == 100800.0
assert task.preset == "cpu-large"
assert task.schedule_timeout == 2157741.0
assert task.image == "image:main"
assert not task.strategy.fail_fast
assert task.strategy.max_parallel == 20
assert task.cache.strategy == CacheStrategy.NONE
assert task.cache.life_span == 9000.0
finally:
await cl.close()
async def _factory(subpath: str = "") -> AsyncIterator[ConfigLoader]:
config_dir = ConfigDir(
workspace=assets / subpath,
config_dir=assets / subpath,
)
cl = BatchLocalCL(config_dir, client)
yield cl
await cl.close()
async def batch_config_loader(
loop: None,
assets: pathlib.Path,
client: Client,
) -> AsyncIterator[ConfigLoader]:
config_dir = ConfigDir(
workspace=assets,
config_dir=assets,
)
cl = BatchLocalCL(config_dir, client)
yield cl
await cl.close()
async def test_project_level_mixins_live(assets: pathlib.Path,
client: Client) -> None:
ws = assets / "with_project_yaml"
config_dir = ConfigDir(
workspace=ws,
config_dir=ws,
)
cl = LiveLocalCL(config_dir, client)
try:
flow = await RunningLiveFlow.create(cl, "live")
job = await flow.get_job("test_mixin", {})
assert job.image == "mixin-image"
assert job.preset == "mixin-preset"
finally:
await cl.close()
async def test_stateful_call_with_cache_invalid(
assets: pathlib.Path,
client: Client,
) -> None:
config_dir = ConfigDir(
workspace=assets / "stateful_actions",
config_dir=assets / "stateful_actions",
)
cl = BatchLocalCL(config_dir, client)
with pytest.raises(
EvalError,
match=r"Specifying cache in action call to the action "
r"ws:with-state of kind stateful is not supported.",
):
await RunningBatchFlow.create(cl, "bad-call-with-cache", "bake-id", {})
async def test_batch_task_with_no_image(assets: pathlib.Path,
client: Client) -> None:
ws = assets / "batch_mixins"
config_dir = ConfigDir(
workspace=ws,
config_dir=ws,
)
cl = BatchLocalCL(config_dir, client)
try:
with pytest.raises(
EvalError,
match=r"Image for task test is not specified",
):
await RunningBatchFlow.create(cl, "batch-task-no-image", "bake-id")
finally:
await cl.close()
async def create(path: Path) -> BatchRunner:
config_dir = ConfigDir(
workspace=path,
config_dir=path,
)
nonlocal runner
# BatchRunner should not use client in this case
runner = BatchRunner(
config_dir,
get_console(),
client,
batch_storage,
FakeGlobalOptions(),
run_neuro_cli=mock_neuro_cli_runner.run,
)
await runner.__aenter__()
return runner
async def test_batch_module_call_to_remote_invalid(assets: pathlib.Path,
client: Client) -> None:
ws = assets / "batch_module"
config_dir = ConfigDir(
workspace=ws,
config_dir=ws,
)
cl = BatchLocalCL(config_dir, client)
try:
with pytest.raises(
EvalError,
match=r"Module call to non local action 'gh:username/repo@tag' "
r"is forbidden",
):
await RunningBatchFlow.create(cl, "batch-module-remote-call",
"bake-id")
finally:
await cl.close()
async def test_batch_action_no_access_to_mixin(assets: pathlib.Path,
client: Client) -> None:
ws = assets / "batch_mixins"
config_dir = ConfigDir(
workspace=ws,
config_dir=ws,
)
cl = BatchLocalCL(config_dir, client)
try:
flow = await RunningBatchFlow.create(cl, "batch-action-call",
"bake-id")
with pytest.raises(
EvalError,
match=r"Unknown mixin 'basic'",
):
await flow.get_action("test", needs={})
finally:
await cl.close()
async def test_early_images_include_globals(assets: pathlib.Path,
client: Client) -> None:
ws = assets / "with_project_yaml"
config_dir = ConfigDir(
workspace=ws,
config_dir=ws,
)
cl = BatchLocalCL(config_dir, client)
try:
flow = await RunningBatchFlow.create(cl, "batch", "bake-id")
assert flow.early_images["image_a"].ref == "image:banana"
assert flow.early_images["image_a"].context == ws / "dir"
assert flow.early_images["image_a"].dockerfile == ws / "dir/Dockerfile"
assert flow.early_images["image_b"].ref == "image:main"
assert flow.early_images["image_b"].context == ws / "dir"
assert flow.early_images["image_b"].dockerfile == ws / "dir/Dockerfile"
finally:
await cl.close()
async def test_batch_module_with_mixin(assets: pathlib.Path,
client: Client) -> None:
ws = assets / "batch_mixins"
config_dir = ConfigDir(
workspace=ws,
config_dir=ws,
)
cl = BatchLocalCL(config_dir, client)
try:
flow = await RunningBatchFlow.create(cl, "batch-module-call",
"bake-id")
module_flow = await flow.get_action("test", needs={})
task = await module_flow.get_task(("test", ),
"task_1",
needs={},
state={})
assert task.image == "ubuntu"
assert task.preset == "cpu-micro"
assert task.cmd == "bash -euo pipefail -c 'echo abc'"
finally:
await cl.close()
def _process_args(
self,
ctx: click.Context,
config: Optional[str],
fake_workspace: bool,
verbose: int,
quiet: int,
show_traceback: bool,
) -> None:
if fake_workspace:
config_dir = ConfigDir(
LocalPath("running-with-fake-workspace"),
LocalPath("running-with-fake-workspace"),
)
else:
config_dir = find_workspace(config)
console = Console(
highlight=False,
log_path=False,
log_time_format="[%X %z]",
get_datetime=lambda: datetime.now().astimezone(),
)
verbosity = verbose - quiet
setup_logging(
color=bool(console.color_system),
verbosity=verbosity,
show_traceback=show_traceback,
)
ctx.obj = Root(
config_dir=config_dir,
console=console,
verbosity=verbosity,
show_traceback=show_traceback,
)
