def collect_ssh(
*,
task_group_id: int,
task_group_activity_id: int,
resource: Resource,
profile: Profile,
wheel_file: FractalUploadedFile | None = None,
) -> None:
"""
Collect a task package over SSH
This function runs as a background task, therefore exceptions must be
handled.
NOTE: since this function is sync, it runs within a thread - due to
starlette/fastapi handling of background tasks (see
https://github.com/encode/starlette/blob/master/starlette/background.py).
Args:
task_group_id:
task_group_activity_id:
resource:
profile:
wheel_file:
"""
LOGGER_NAME = f"{__name__}.ID{task_group_activity_id}"
# Work within a temporary folder, where also logs will be placed
with TemporaryDirectory() as tmpdir:
log_file_path = Path(tmpdir) / "log"
logger = set_logger(
logger_name=LOGGER_NAME,
log_file_path=log_file_path,
)
logger.info("START")
with next(get_sync_db()) as db:
try:
task_group, activity = get_activity_and_task_group(
task_group_activity_id=task_group_activity_id,
task_group_id=task_group_id,
db=db,
logger_name=LOGGER_NAME,
)
except NoResultFound:
return
with SingleUseFractalSSH(
ssh_config=SSHConfig(
host=resource.host,
user=profile.username,
key_path=profile.ssh_key_path,
),
logger_name=LOGGER_NAME,
) as fractal_ssh:
try:
# Check SSH connection
ssh_ok = check_ssh_or_fail_and_cleanup(
fractal_ssh=fractal_ssh,
task_group=task_group,
task_group_activity=activity,
logger_name=LOGGER_NAME,
log_file_path=log_file_path,
db=db,
)
if not ssh_ok:
return
# Check that the (remote) task_group path does not exist
# NOTE: this is not part of the try/except below, in order
# to avoid removing the existing folder (as part of the
# exception-handling).
if fractal_ssh.remote_exists(task_group.path):
error_msg = f"{task_group.path} already exists."
logger.error(error_msg)
fail_and_cleanup(
task_group=task_group,
task_group_activity=activity,
logger_name=LOGGER_NAME,
log_file_path=log_file_path,
exception=FileExistsError(error_msg),
db=db,
)
return
# Create remote `task_group.path` and `script_dir_remote`
# folders
script_dir_remote = (
Path(task_group.path) / SCRIPTS_SUBFOLDER
).as_posix()
fractal_ssh.mkdir(folder=task_group.path, parents=True)
fractal_ssh.mkdir(folder=script_dir_remote, parents=True)
fractal_ssh.write_remote_file(
path=f"{task_group.path}/{TASK_GROUP_ID_FILENAME}",
content=str(task_group_id),
)
# Write wheel file locally and send it to remote path,
# and set task_group.archive_path
if wheel_file is not None:
wheel_filename = wheel_file.filename
archive_path = (
Path(task_group.path) / wheel_filename
).as_posix()
tmp_archive_path = (
Path(tmpdir) / wheel_filename
).as_posix()
logger.info(f"Write wheel file into {tmp_archive_path}")
with open(tmp_archive_path, "wb") as f:
f.write(wheel_file.contents)
fractal_ssh.send_file(
local=tmp_archive_path,
remote=archive_path,
)
task_group.archive_path = archive_path
task_group = add_commit_refresh(obj=task_group, db=db)
python_bin = get_python_interpreter(
python_version=task_group.python_version,
resource=resource,
)
replacements = get_collection_replacements(
task_group=task_group,
python_bin=python_bin,
resource=resource,
)
# Prepare common arguments for _customize_and_run_template
common_args = dict(
script_dir_local=Path(
tmpdir, SCRIPTS_SUBFOLDER
).as_posix(),
script_dir_remote=script_dir_remote,
prefix=(
f"{int(time.time())}_"
f"{TaskGroupActivityAction.COLLECT}"
),
logger_name=LOGGER_NAME,
)
logger.info("installing - START")
# Set status to ONGOING and refresh logs
activity.status = TaskGroupActivityStatus.ONGOING
activity.log = get_current_log(log_file_path)
activity = add_commit_refresh(obj=activity, db=db)
# Run script 1
stdout = _customize_and_run_template(
template_filename="1_create_venv.sh",
fractal_ssh=fractal_ssh,
replacements=replacements,
**common_args,
)
activity.log = get_current_log(log_file_path)
activity = add_commit_refresh(obj=activity, db=db)
# Run script 2
stdout = _customize_and_run_template(
template_filename="2_pip_install.sh",
fractal_ssh=fractal_ssh,
replacements=replacements,
**common_args,
)
activity.log = get_current_log(log_file_path)
activity = add_commit_refresh(obj=activity, db=db)
# Run script 3
pip_freeze_stdout = _customize_and_run_template(
template_filename="3_pip_freeze.sh",
fractal_ssh=fractal_ssh,
replacements=replacements,
**common_args,
)
activity.log = get_current_log(log_file_path)
activity = add_commit_refresh(obj=activity, db=db)
# Run script 4
stdout = _customize_and_run_template(
template_filename="4_pip_show.sh",
fractal_ssh=fractal_ssh,
replacements=replacements,
**common_args,
)
activity.log = get_current_log(log_file_path)
activity = add_commit_refresh(obj=activity, db=db)
pkg_attrs = parse_script_pip_show_stdout(stdout)
for key, value in pkg_attrs.items():
logger.debug(f"parsed from pip-show: {key}={value}")
# Check package_name match between pip show and task-group
package_name_pip_show = pkg_attrs.get("package_name")
package_name_task_group = task_group.pkg_name
compare_package_names(
pkg_name_pip_show=package_name_pip_show,
pkg_name_task_group=package_name_task_group,
logger_name=LOGGER_NAME,
)
# Extract/drop parsed attributes
package_name = package_name_task_group
python_bin = pkg_attrs.pop("python_bin")
package_root_parent_remote = pkg_attrs.pop(
"package_root_parent"
)
manifest_path_remote = pkg_attrs.pop("manifest_path")
# TODO SSH: Use more robust logic to determine
# `package_root`. Examples: use `importlib.util.find_spec`
# or parse the output of `pip show --files {package_name}`.
package_name_underscore = package_name.replace("-", "_")
package_root_remote = (
Path(package_root_parent_remote)
/ package_name_underscore
).as_posix()
# Read and validate remote manifest file
pkg_manifest_dict = fractal_ssh.read_remote_json_file(
manifest_path_remote
)
logger.info(f"Loaded {manifest_path_remote=}")
pkg_manifest = ManifestV2(**pkg_manifest_dict)
logger.info("Manifest is a valid ManifestV2")
logger.info("_prepare_tasks_metadata - start")
task_list = prepare_tasks_metadata(
package_manifest=pkg_manifest,
package_version=task_group.version,
package_root=Path(package_root_remote),
python_bin=Path(python_bin),
)
logger.info("_prepare_tasks_metadata - end")
logger.info("create_db_tasks_and_update_task_group - start")
create_db_tasks_and_update_task_group_sync(
task_list=task_list,
task_group_id=task_group.id,
db=db,
)
logger.info("create_db_tasks_and_update_task_group - end")
# Update task_group data
logger.info("Add env_info to TaskGroupV2 - start")
task_group.env_info = pip_freeze_stdout
task_group = add_commit_refresh(obj=task_group, db=db)
logger.info("Add env_info to TaskGroupV2 - end")
# Finalize (write metadata to DB)
logger.info("finalising - START")
activity.status = TaskGroupActivityStatus.OK
activity.timestamp_ended = get_timestamp()
activity = add_commit_refresh(obj=activity, db=db)
logger.info("finalising - END")
logger.info("END")
reset_logger_handlers(logger)
except Exception as collection_e:
# Delete corrupted package dir
try:
logger.info(
f"Now delete remote folder {task_group.path}"
)
fractal_ssh.remove_folder(folder=task_group.path)
logger.info(f"Deleted remoted folder {task_group.path}")
except Exception as e_rm:
logger.error(
"Removing folder failed. "
f"Original error: {str(e_rm)}"
)
fail_and_cleanup(
task_group=task_group,
task_group_activity=activity,
log_file_path=log_file_path,
logger_name=LOGGER_NAME,
exception=collection_e,
db=db,
)