"""Provide the Job class"""
# pyright: reportCallIssue=false
# pyright: reportAttributeAccessIssue=false
# pyright: reportOptionalMemberAccess=false
from __future__ import annotations
import logging
import shlex
from contextlib import suppress
from collections.abc import Iterable
from datetime import datetime
from difflib import unified_diff
from functools import cached_property
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, Mapping
from panpath import PanPath, CloudPath, LocalPath
from diot import OrderedDiot
from xqute import Job as XquteJob
from xqute.path import SpecPath, MountedPath
from ._job_caching import JobCaching
from .defaults import ProcInputType, ProcOutputType
from .exceptions import (
ProcInputTypeError,
ProcOutputNameError,
ProcOutputTypeError,
ProcOutputValueError,
TemplateRenderingError,
)
from .template import Template
from .utils import logger, strsplit, path_is_symlink, path_symlink_to, get_mtime
from .pluginmgr import plugin
if TYPE_CHECKING: # pragma: no cover
from .proc import Proc
def _process_input_file_or_dir(
inkey: str,
intype: str,
inval: Any,
index: int | None = None,
proc_name: str | None = None,
) -> CloudPath | MountedPath:
"""Process the input value for file or dir"""
if inval is None or not isinstance(inval, (str, Path)):
msg = (
f"[{proc_name}] Got <{type(inval).__name__}> instead of "
f"path-like object for input: {inkey + ':' + intype!r}"
)
if index is not None:
msg = f"{msg} at index {index}"
raise ProcInputTypeError(msg)
if isinstance(inval, MountedPath):
return inval
if isinstance(inval, SpecPath):
return inval.mounted
if isinstance(inval, CloudPath): # pragma: no cover
return MountedPath(inval)
if not isinstance(inval, str): # other path-like types, should be all local
return MountedPath(PanPath(inval).expanduser().absolute())
# str
# Let's see if it a path in str format, which is path1:path2
# However, there is also a colon in cloud paths
colon_count = inval.count(":")
if colon_count == 0: # a/b
return MountedPath(PanPath(inval).expanduser().absolute())
if colon_count > 3: # a:b:c:d
msg = (
f"[{proc_name}] Invalid input value: {inkey + ':' + intype!r} "
"(too many ':')"
)
if index is not None:
msg = f"{msg} at index {index}"
raise ProcInputTypeError(msg)
if colon_count == 1: # gs://a/b or a/b:c/d
if isinstance(PanPath(inval), CloudPath): # gs://a/b
return MountedPath(inval)
path1, path2 = inval.split(":")
elif inval.count(":") == 3: # gs://a/b:gs://c/d
p1, p2, path2 = inval.split(":", 2)
path1 = p1 + ":" + p2
else: # gs://a/b:c/d or a/b:gs://c/d
p1, p2, p3 = inval.split(":", 2)
path1, path2 = p1 + ":" + p2, p3
if not isinstance(PanPath(path1), CloudPath):
path1, path2 = p1, p2 + ":" + p3
path1 = PanPath(path1) # type: ignore
path2 = PanPath(path2) # type: ignore
if isinstance(path1, LocalPath):
path1 = path1.expanduser().absolute()
if isinstance(path2, LocalPath):
path2 = path2.expanduser().absolute()
return MountedPath(path2, spec=path1)
class Job(XquteJob, JobCaching):DOCS
"""The job for pipen"""
__slots__ = XquteJob.__slots__ + ("proc", "_output_types", "outdir", "output")
def __init__(
self,
*args: Any,
**kwargs: Any,
) -> None:
super().__init__(*args, **kwargs)
self.proc: Proc | None = None
self._output_types: Dict[str, str] = {}
# Where the real output directory is
self.outdir: SpecPath | None = None
self.output: Mapping[str, Any] | None = None
async def prepare(self, proc: Proc) -> None:DOCS
"""Prepare the job by given process
Primarily prepare the script, and provide cmd to the job for xqute
to wrap and run
Args:
proc: the process object
"""
# Attach the process
self.proc = proc
if proc.export:
# Don't put index if it is a single-job process
if not isinstance(proc._export_dir, SpecPath): # type: ignore
self.outdir = SpecPath(proc._export_dir) # type: ignore
else:
self.outdir = proc._export_dir # type: ignore
# Put job output in a subdirectory with index
# if it is a multi-job process
if not proc.output_flatten:
self.outdir = self.outdir / str(self.index) # type: ignore
if proc.scheduler.fs_shared:
await self.outdir.mounted.a_mkdir(parents=True, exist_ok=True)
else:
# For non-export process, the output directory is the metadir
self.outdir = self.metadir / "output"
# compute the output
await self.prepare_output()
await self.prepare_outdir()
if not proc.script:
self.cmd = ("true",)
else:
try:
script = proc.script.render(self.template_data)
except Exception as exc:
raise TemplateRenderingError(
f"[{self.proc.name}] Failed to render script."
) from exc
script_file_exists = await self.script_file.a_is_file()
old_script = (
await self.script_file.a_read_text()
if script_file_exists
else None
)
if not script_file_exists:
await self.script_file.a_write_text(script)
elif script_file_exists and old_script != script:
self.log("debug", "Job script updated.")
old_mtime = await get_mtime(self.script_file, 0)
diff_file = self.script_file.with_name("job.script.diff")
diff = "\n".join(
unified_diff(
old_script.splitlines(),
script.splitlines(),
fromfile=(
f"old ({datetime.fromtimestamp(old_mtime).isoformat()})"
),
tofile=f"new ({datetime.now().isoformat()})",
lineterm="",
)
)
await diff_file.a_write_text(diff)
await self.script_file.a_write_text(script)
lang = proc.lang or proc.pipeline.config.lang
script_file = self.script_file.mounted
script_file = await script_file.get_fspath()
self.cmd = tuple((*shlex.split(lang), script_file))
# self.cmd = tuple(shlex.split(lang) + [self.script_file.mounted.fspath])
await plugin.hooks.on_job_init(self)
@propertyDOCS
def script_file(self) -> SpecPath:
"""Get the path to script file
Returns:
The path to the script file
"""
return self.metadir / "job.script"
async def prepare_outdir(self) -> SpecPath:DOCS
"""Get the path to the output directory.
When proc.export is True, the output directory is based on the
pipeline.outdir and the process name. Otherwise, it is based on
the metadir.
When the job is running in a detached system (a VM, typically),
this will return the mounted path to the output directory.
To access the real path, use self.outdir
Returns:
The path to the job output directory
"""
# if ret is a dead link
# when switching a proc from end/nonend to nonend/end
# if path_is_symlink(self.outdir) and not self.outdir.exists():
if await path_is_symlink(self.outdir) and ( # type: ignore
# A local deak link
not await self.outdir.a_exists()
# A cloud fake link
or isinstance(getattr(self.outdir, "path", self.outdir), CloudPath)
):
await self.outdir.a_unlink() # pragma: no cover
await self.outdir.a_mkdir(parents=True, exist_ok=True)
# If it is somewhere else, make a symbolic link to the metadir
metaout = self.metadir / "output"
if self.outdir != metaout:
if await path_is_symlink(metaout) or await metaout.a_is_file():
await metaout.a_unlink()
elif await metaout.a_is_dir():
# Remove the directory, it is inconsistent with current setting
await metaout.a_rmtree()
await path_symlink_to(metaout, self.outdir) # type: ignore
return self.outdir # type: ignore
@cached_property
def input(self) -> Mapping[str, Any]:
"""Get the input data for this job
Returns:
A key-value map, where keys are the input keys
"""
import pandas
ret = self.proc.input.data.iloc[self.index, :].to_dict()
# check types
for inkey, intype in self.proc.input.type.items():
if intype == ProcInputType.VAR or ret[inkey] is None:
continue # pragma: no cover, covered actually
if intype in (ProcInputType.FILE, ProcInputType.DIR):
ret[inkey] = _process_input_file_or_dir(
inkey, intype, ret[inkey], None, self.proc.name
)
if intype in (ProcInputType.FILES, ProcInputType.DIRS):
if isinstance(ret[inkey], pandas.DataFrame): # pragma: no cover
# // todo: nested dataframe
ret[inkey] = ret[inkey].iloc[0, 0]
if isinstance(ret[inkey], (str, Path)):
# if a single file, convert to list
ret[inkey] = [ret[inkey]]
if not isinstance(ret[inkey], Iterable):
raise ProcInputTypeError(
f"[{self.proc.name}] Expected an iterable for input: "
f"{inkey + ':' + intype!r}, got {type(ret[inkey])}"
)
for i, file in enumerate(ret[inkey]):
ret[inkey][i] = _process_input_file_or_dir(
inkey, intype, file, i, self.proc.name
)
return ret
async def prepare_output(self) -> None:DOCS
"""Get the output data of the job
Returns:
The key-value map where the keys are the output keys
"""
output_template = self.proc.output
if not output_template:
self.output = {}
return
data = {
"job": dict(
index=self.index,
metadir=self.metadir.mounted,
outdir=self.outdir.mounted,
stdout_file=self.stdout_file.mounted,
stderr_file=self.stderr_file.mounted,
jid_file=self.jid_file.mounted,
),
"in": self.input,
"in_": self.input,
"proc": self.proc,
"envs": self.proc.envs,
}
try:
if isinstance(output_template, Template):
# // TODO: check ',' in output value?
outputs = strsplit(output_template.render(data), ",")
else:
outputs = [oput.render(data) for oput in output_template]
except Exception as exc:
raise TemplateRenderingError(
f"[{self.proc.name}] Failed to render output."
) from exc
self.output = ret = OrderedDiot()
for oput in outputs:
if ":" not in oput:
raise ProcOutputNameError(
f"[{self.proc.name}] No name given in output."
)
if oput.count(":") == 1:
output_name, output_value = oput.split(":")
output_type = ProcOutputType.VAR
else:
output_name, output_type, output_value = oput.split(":", 2)
if output_type not in ProcOutputType.__dict__.values():
raise ProcOutputTypeError(
f"[{self.proc.name}] " f"Unsupported output type: {output_type}"
)
self._output_types[output_name] = output_type
if output_type == ProcOutputType.VAR:
ret[output_name] = output_value
else:
if self.proc.export and self.proc.output_flatten:
self.proc._duplicate_fields_check.setdefault(output_name, set())
if output_value in self.proc._duplicate_fields_check[output_name]:
raise ProcOutputValueError(
f"[{self.proc.name}] The output value {output_value!r} for "
f"{output_name!r} (job #{self.index}) is duplicated with "
"another job, and you have output_flatten=True for this "
"process, which may cause conflicts."
)
self.proc._duplicate_fields_check[output_name].add(output_value)
ov = PanPath(output_value)
if isinstance(ov, CloudPath) or (
isinstance(ov, LocalPath) and ov.is_absolute()
):
raise ProcOutputValueError(
f"[{self.proc.name}] "
f"output path must be a segment: {output_value}"
)
# self.outdir is already awaited
out = self.outdir / output_value # type: ignore
if output_type == ProcOutputType.DIR:
with suppress(Exception): # pragma: no cover
# Likely aiohttp.ClientResponseError
# In case we have many jobs and this is running in parallel
# Parallelly creating the same directory may cause
# a rate limit error in cloud storages
await out.a_mkdir(parents=True, exist_ok=True)
ret[output_name] = out.mounted
@cached_property
def template_data(self) -> Mapping[str, Any]:
"""Get the data for template rendering
Returns:
The data for template rendering
"""
return {
"job": dict(
index=self.index,
metadir=self.metadir.mounted,
outdir=self.outdir.mounted,
script_file=self.script_file.mounted,
stdout_file=self.stdout_file.mounted,
stderr_file=self.stderr_file.mounted,
jid_file=self.jid_file.mounted,
),
"in": self.input,
"in_": self.input,
"out": self.output,
"proc": self.proc,
"envs": self.proc.envs,
}
def log(DOCS
self,
level: int | str,
msg: str,
*args,
limit: int = 3,
limit_indicator: bool = True,
logger: logging.LoggerAdapter = logger,
) -> None:
"""Log message for the jobs
Args:
level: The log level of the record
msg: The message to log
*args: The arguments to format the message
limit: limitation of the log (don't log for all jobs)
limit_indicator: Whether to show an indicator saying the log
has been limited (the level of the indicator will be DEBUG)
logger: The logger used to log
"""
if self.index > limit:
return
if self.index == limit:
if limit_indicator:
msg = f"{msg} (not showing similar logs)"
if self.proc.size == 1:
job_index_indicator = ""
else:
job_index_indicator = "[%s/%s] " % (
str(self.index).zfill(len(str(self.proc.size - 1))),
self.proc.size - 1,
)
self.proc.log(level, job_index_indicator + msg, *args, logger=logger)