def reactivate_ssh_pixi(
*,
task_group_activity_id: int,
task_group_id: int,
resource: Resource,
profile: Profile,
) -> None:
"""
Reactivate a task group venv.
This function is run as a background task, therefore exceptions must be
handled.
Args:
task_group_id:
task_group_activity_id:
resource:
profile:
"""
LOGGER_NAME = f"{__name__}.ID{task_group_activity_id}"
with TemporaryDirectory() as tmpdir:
log_file_path = get_log_path(Path(tmpdir))
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 source_dir does not
# exist
source_dir = Path(
task_group.path, SOURCE_DIR_NAME
).as_posix()
if fractal_ssh.remote_exists(source_dir):
error_msg = f"{source_dir} 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
pixi_home = resource.tasks_pixi_config["versions"][
task_group.pixi_version
]
pixi_cache_dir = profile.pixi_cache_dir or os.path.join(
pixi_home, "cache"
)
logger.info(f"Setting PIXI_CACHE_DIR to {pixi_cache_dir}")
replacements: set[tuple[str, str]] = {
("__PIXI_HOME__", pixi_home),
("__PIXI_CACHE_DIR__", pixi_cache_dir),
("__PACKAGE_DIR__", task_group.path),
("__TAR_GZ_PATH__", task_group.archive_path),
(
"__IMPORT_PACKAGE_NAME__",
task_group.pkg_name.replace("-", "_"),
),
("__SOURCE_DIR_NAME__", SOURCE_DIR_NAME),
("__FROZEN_OPTION__", "--frozen"),
(
"__TOKIO_WORKER_THREADS__",
str(
resource.tasks_pixi_config[
"TOKIO_WORKER_THREADS"
]
),
),
(
"__PIXI_CONCURRENT_SOLVES__",
str(
resource.tasks_pixi_config[
"PIXI_CONCURRENT_SOLVES"
]
),
),
(
"__PIXI_CONCURRENT_DOWNLOADS__",
str(
resource.tasks_pixi_config[
"PIXI_CONCURRENT_DOWNLOADS"
]
),
),
}
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)
script_dir_remote = Path(
task_group.path, SCRIPTS_SUBFOLDER
).as_posix()
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.REACTIVATE}"
),
logger_name=LOGGER_NAME,
)
# Run script 1 - extract tar.gz into `source_dir`
stdout = _customize_and_run_template(
template_filename="pixi_1_extract.sh",
replacements=replacements,
fractal_ssh=fractal_ssh,
**common_args,
)
logger.debug(f"STDOUT: {stdout}")
activity.log = get_current_log(log_file_path)
activity = add_commit_refresh(obj=activity, db=db)
# Simplify `pyproject.toml`
source_dir = Path(
task_group.path, SOURCE_DIR_NAME
).as_posix()
pyproject_toml_path = Path(source_dir, "pyproject.toml")
edit_pyproject_toml_in_place_ssh(
fractal_ssh=fractal_ssh,
pyproject_toml_path=pyproject_toml_path,
resource=resource,
)
# Write pixi.lock into `source_dir`
pixi_lock_local = Path(tmpdir, "pixi.lock").as_posix()
pixi_lock_remote = Path(
task_group.path, SOURCE_DIR_NAME, "pixi.lock"
).as_posix()
logger.info(
f"Write `env_info` contents into {pixi_lock_local}"
)
with open(pixi_lock_local, "w") as f:
f.write(task_group.env_info)
fractal_ssh.send_file(
local=pixi_lock_local,
remote=pixi_lock_remote,
)
# Prepare scripts 2 and 3
remote_script2_path = _customize_and_send_template(
template_filename="pixi_2_install.sh",
replacements=replacements,
fractal_ssh=fractal_ssh,
**common_args,
)
remote_script3_path = _customize_and_send_template(
template_filename="pixi_3_post_install.sh",
replacements=replacements,
fractal_ssh=fractal_ssh,
**common_args,
)
logger.debug(
"Post-installation script written to "
f"{remote_script3_path=}."
)
logger.debug(
"Installation script written to "
f"{remote_script2_path=}."
)
activity.log = get_current_log(log_file_path)
activity = add_commit_refresh(obj=activity, db=db)
# Run scripts 2 and 3
stdout = run_script_on_remote_slurm(
job_name=LOGGER_NAME,
script_paths=[
remote_script2_path,
remote_script3_path,
],
final_commands=[f"chmod -R 755 {source_dir}"],
slurm_config=resource.tasks_pixi_config["SLURM_CONFIG"],
fractal_ssh=fractal_ssh,
logger_name=LOGGER_NAME,
prefix=common_args["prefix"],
db=db,
activity=activity,
log_file_path=log_file_path,
poll_interval=resource.jobs_poll_interval,
)
activity.log = get_current_log(log_file_path)
activity = add_commit_refresh(obj=activity, db=db)
# Finalize (write metadata to DB)
activity.status = TaskGroupActivityStatus.OK
activity.timestamp_ended = get_timestamp()
activity = add_commit_refresh(obj=activity, db=db)
task_group.active = True
task_group = add_commit_refresh(obj=task_group, db=db)
logger.info("END")
reset_logger_handlers(logger)
except Exception as reactivate_e:
# Delete corrupted source_dir
try:
logger.info(f"Now delete folder {source_dir}")
fractal_ssh.remove_folder(folder=source_dir)
logger.info(f"Deleted folder {source_dir}")
except Exception as rm_e:
logger.error(
"Removing folder failed. "
f"Original error: {str(rm_e)}"
)
fail_and_cleanup(
task_group=task_group,
task_group_activity=activity,
logger_name=LOGGER_NAME,
log_file_path=log_file_path,
exception=reactivate_e,
db=db,
)