async def inspect(
root: Root,
bake: str,
attempt: int,
output_graph: Optional[str],
dot: bool,
pdf: bool,
view: bool,
) -> None:
"""Inspect a bake.
Display a list of started/finished tasks of BAKE\\_ID.
"""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
BatchRunner(root.config_dir, root.console, client, storage, root))
if output_graph is not None:
real_output: Optional[LocalPath] = LocalPath(output_graph)
else:
real_output = None
bake_id = await resolve_bake(bake,
project=runner.project_id,
storage=storage)
await runner.inspect(
bake_id,
attempt_no=attempt,
output=real_output,
save_dot=dot,
save_pdf=pdf,
view_pdf=view,
)
async def get_worker_hostnames() -> List[str]:
result = []
async with get() as client:
async for job in client.jobs.list(statuses={JobStatus.RUNNING},
tags={"job:nni_worker"}):
result.append(job.internal_hostname)
return result
async def restart(
root: Root,
bake: str,
attempt: int,
from_failed: bool,
local_executor: bool,
) -> None:
"""Start a batch.
Run BATCH pipeline remotely on the cluster.
"""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
BatchRunner(root.config_dir, root.console, client, storage, root))
bake_id = await resolve_bake(bake,
project=runner.project_id,
storage=storage)
await runner.restart(
bake_id,
attempt_no=attempt,
from_failed=from_failed,
local_executor=local_executor,
)
async def run(
root: Root,
job_id: str,
suffix: Optional[str],
dry_run: bool,
args: Optional[Tuple[str]],
param: List[Tuple[str, str]],
) -> None:
"""Run a job.
RUN job JOB-ID or ATTACH to it if the job is already running
For multi-jobs an explicit job suffix can be used with explicit job arguments.
"""
if args:
root.console.print(
"[yellow]args are deprecated, use --param instead", )
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
LiveRunner(root.config_dir, root.console, client, storage, root,
dry_run))
await runner.run(
job_id,
suffix=suffix,
args=args,
params={key: value
for key, value in param},
)
async def ps(root: Root, ) -> None:
"""List all jobs"""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
LiveRunner(root.config_dir, root.console, client, storage, root))
await runner.ps()
async def mkvolumes(root: Root, ) -> None:
"""Create all remote folders for volumes."""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
LiveRunner(root.config_dir, root.console, client, storage, root))
await runner.mkvolumes()
async def create_config() -> None:
async with get() as client:
config = {
"1": {
"name": f"{client.config.cluster_name} registry",
"url": str(client.config.registry_url),
"user": client.config.username,
"password": await client.config.token()
}
}
with open("db.json", "w") as db_file:
json.dump(config, db_file)
async def delete_project(
root: Root,
project_ids: Sequence[str],
) -> None:
"""Completely remove project with all related entities"""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
for project_id in project_ids:
await storage.project(yaml_id=project_id).delete()
if root.verbosity >= 0:
root.console.print(
f"Project '{project_id}' was successfully removed.")
async def build(root: Root, force_overwrite: bool, image: str) -> None:
"""Build an image.
Assemble the IMAGE remotely and publish it.
"""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
LiveRunner(root.config_dir, root.console, client, storage, root))
if image == "ALL":
await runner.build_all(force_overwrite=force_overwrite)
else:
await runner.build(image, force_overwrite=force_overwrite)
async def logs(
root: Root,
job_id: str,
suffix: Optional[str],
) -> None:
"""Print logs.
Display logs for JOB-ID
"""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
LiveRunner(root.config_dir, root.console, client, storage, root))
await runner.logs(job_id, suffix)
async def dockerhub_auth_secret() -> AsyncIterator[Secret]:
async with neuro_sdk.get() as neuro_client:
secret_name = f"{KANIKO_AUTH_PREFIX}_{uuid.uuid4().hex}"
auth_data = _build_registy_auth(
# Why not v2: https://github.com/GoogleContainerTools/kaniko/pull/1209
registry_uri="https://index.docker.io/v1/",
username=os.environ["DOCKER_CI_USERNAME"],
password=os.environ["DOCKER_CI_TOKEN"],
)
secret = Secret(secret_name, auth_data)
try:
await neuro_client.secrets.add(secret_name, auth_data.encode())
yield secret
finally:
await neuro_client.secrets.rm(secret_name)
async def async_shell_complete(self, root: Root, ctx: click.Context,
param: click.Parameter,
incomplete: str) -> List[CompletionItem]:
job_id = ctx.params[self._job_id_param_name]
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(
ApiStorage(client))
runner = await stack.enter_async_context(
LiveRunner(root.config_dir, root.console, client, storage,
root))
variants = await runner.list_suffixes(job_id)
return [
CompletionItem(suffix) for suffix in variants
if suffix.startswith(incomplete)
]
async def async_shell_complete(self, root: Root, ctx: click.Context,
param: click.Parameter,
incomplete: str) -> List[CompletionItem]:
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(
ApiStorage(client))
runner = await stack.enter_async_context(
LiveRunner(root.config_dir, root.console, client, storage,
root))
variants = list(runner.flow.job_ids)
if self._allow_all:
variants += ["ALL"]
return [
CompletionItem(job_id) for job_id in variants
if job_id.startswith(incomplete)
]
async def async_shell_complete(self, root: Root, ctx: click.Context,
param: click.Parameter,
incomplete: str) -> List[CompletionItem]:
variants = []
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(
ApiStorage(client))
try:
async for project in storage.list_projects():
variants.append(project.yaml_id)
except ValueError:
pass
return [
CompletionItem(yaml_id) for yaml_id in variants
if yaml_id.startswith(incomplete)
]
async def clean(
root: Root,
volume: str,
) -> None:
"""Clean volume.
Clean remote files on VOLUME,
use `clean ALL` for cleaning up all volumes."""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
LiveRunner(root.config_dir, root.console, client, storage, root))
if volume != "ALL":
await runner.clean(volume)
else:
await runner.clean_all()
async def download(
root: Root,
volume: str,
) -> None:
"""Download volume.
Download remote files to local for VOLUME,
use `download ALL` for downloading all volumes."""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
LiveRunner(root.config_dir, root.console, client, storage, root))
if volume != "ALL":
await runner.download(volume)
else:
await runner.download_all()
async def cancel(
root: Root,
bake: str,
attempt: int,
) -> None:
"""Cancel a bake.
Cancel a bake execution by stopping all started tasks.
"""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
BatchRunner(root.config_dir, root.console, client, storage, root))
bake_id = await resolve_bake(bake,
project=runner.project_id,
storage=storage)
await runner.cancel(bake_id, attempt_no=attempt)
async def bakes(
root: Root,
tag: Sequence[str],
since: Optional[str],
until: Optional[str],
recent_first: bool,
) -> None:
"""List existing bakes."""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
BatchRunner(root.config_dir, root.console, client, storage, root))
await runner.list_bakes(
tags=set(tag),
since=_parse_date(since),
until=_parse_date(until),
recent_first=recent_first,
)
async def async_shell_complete(self, root: Root, ctx: click.Context,
param: click.Parameter,
incomplete: str) -> List[CompletionItem]:
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(
ApiStorage(client))
runner = await stack.enter_async_context(
LiveRunner(root.config_dir, root.console, client, storage,
root))
variants = [
volume.id for volume in runner.flow.volumes.values()
if volume.local is not None
]
if self._allow_all:
variants += ["ALL"]
return [
CompletionItem(image) for image in variants
if image.startswith(incomplete)
]
async def kill(
root: Root,
job_id: str,
suffix: Optional[str],
) -> None:
"""Kill a job.
Kill JOB-ID, use `kill ALL` for killing all jobs."""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
LiveRunner(root.config_dir, root.console, client, storage, root))
if job_id != "ALL":
await runner.kill(job_id, suffix)
else:
if suffix is not None:
raise click.BadArgumentUsage(
"Suffix is not supported when killing ALL jobs")
await runner.kill_all()
async def show(
root: Root,
bake: str,
attempt: int,
task_id: str,
raw: bool,
) -> None:
"""Show output of baked task.
Display a logged output of TASK\\_ID from BAKE\\_ID.
"""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
BatchRunner(root.config_dir, root.console, client, storage, root))
bake_id = await resolve_bake(bake,
project=runner.project_id,
storage=storage)
await runner.logs(bake_id, task_id, attempt_no=attempt, raw=raw)
async def async_shell_complete(self, root: Root, ctx: click.Context,
param: click.Parameter,
incomplete: str) -> List[CompletionItem]:
variants = []
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(
ApiStorage(client))
runner: BatchRunner = await stack.enter_async_context(
BatchRunner(root.config_dir, root.console, client, storage,
root))
try:
async for bake in runner.get_bakes():
variants.append(bake.id)
if bake.name is not None:
variants.append(bake.name)
except ValueError:
pass
return [
CompletionItem(bake) for bake in variants
if bake.startswith(incomplete)
]
async def clear_cache(
root: Root,
batch: str,
task_id: Optional[str],
) -> None:
"""Clear cache.
Use `neuro-flow clear-cache <BATCH>` for cleaning up the cache for BATCH;
Use `neuro-flow clear-cache <BATCH> <TASK_ID>` for cleaning up the cache
for TASK_ID in BATCH;
`neuro-flow clear-cache ALL` clears all caches.
"""
async with AsyncExitStack() as stack:
client = await stack.enter_async_context(neuro_sdk.get())
storage: Storage = await stack.enter_async_context(ApiStorage(client))
runner = await stack.enter_async_context(
BatchRunner(root.config_dir, root.console, client, storage, root))
if batch == "ALL":
await runner.clear_cache(None)
else:
await runner.clear_cache(batch, task_id)