def collect_local(
*,
task_group_activity_id: int,
task_group_id: int,
resource: Resource,
profile: Profile,
wheel_file: FractalUploadedFile | None = None,
) -> None:
"""
Collect a task package.
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: Resource
wheel_file:
"""
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
# Check that the (local) task_group path does exist
if Path(task_group.path).exists():
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
try:
# Create task_group.path folder
Path(task_group.path).mkdir(parents=True)
logger.info(f"Created {task_group.path}")
# Write txt file in task_group.path
txt_path = (
Path(task_group.path) / TASK_GROUP_ID_FILENAME
).as_posix()
logger.info(f"Write txt-file contents into {txt_path}")
with open(txt_path, "w") as f:
f.write(str(task_group_id))
# Write wheel file and set task_group.archive_path
if wheel_file is not None:
archive_path = (
Path(task_group.path) / wheel_file.filename
).as_posix()
logger.info(
f"Write wheel-file contents into {archive_path}"
)
with open(archive_path, "wb") as f:
f.write(wheel_file.contents)
task_group.archive_path = archive_path
task_group = add_commit_refresh(obj=task_group, db=db)
# Prepare replacements for templates
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(
replacements=replacements,
script_dir=(
Path(task_group.path) / SCRIPTS_SUBFOLDER
).as_posix(),
prefix=(
f"{int(time.time())}_{TaskGroupActivityAction.COLLECT}"
),
logger_name=LOGGER_NAME,
)
# 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",
**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",
**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",
**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",
**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
task_group = db.get_one(TaskGroupV2, task_group_id)
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 = pkg_attrs.pop("package_root_parent")
# TODO : 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 = (
Path(package_root_parent) / package_name_underscore
).as_posix()
# Read and validate manifest file
manifest_path = pkg_attrs.pop("manifest_path")
logger.info(f"now loading {manifest_path=}")
with open(manifest_path) as json_data:
pkg_manifest_dict = json.load(json_data)
logger.info(f"loaded {manifest_path=}")
logger.info("now validating manifest content")
pkg_manifest = ManifestV2(**pkg_manifest_dict)
logger.info("validated manifest content")
activity.log = get_current_log(log_file_path)
activity = add_commit_refresh(obj=activity, db=db)
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),
python_bin=Path(python_bin),
)
check_task_files_exist(task_list=task_list)
logger.info("_prepare_tasks_metadata - end")
activity.log = get_current_log(log_file_path)
activity = add_commit_refresh(obj=activity, db=db)
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:
rmtree_nofail(
folder_path=task_group.path,
logger_name=LOGGER_NAME,
)
fail_and_cleanup(
task_group=task_group,
task_group_activity=activity,
logger_name=LOGGER_NAME,
log_file_path=log_file_path,
exception=collection_e,
db=db,
)
return