# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
import json
import re
import unicodedata
from hastegeo.core.artifact_storage.unified_artifact_storage import (
UnifiedArtifactStorage,
)
from hastegeo.core.config import Config
from hastegeo.core.models.projects import Model, ModelArtifacts, ZipJob
from hastegeo.core.runners.unified_runner import UnifiedRunner
from hastegeo.core.utils.logs import Logger
from hastegeo.core.utils.metadata import MetadataUtils
from hastegeo.core.utils.queues import AzureQueueHandler
BATCH_JOB_WORKDIR = "AZ_BATCH_TASK_WORKING_DIR"
ZIP_PREFIX = "zip"
ZIP_IN_PROGRESS_MESSAGE = "Zipping in progress"
_SLUG_INVALID = re.compile(r"[^A-Za-z0-9._-]+")
def _slugify_model_name(name: str) -> str:
"""Convert a free-form Model.name into a value safe for blob paths.
Normalizes unicode (NFKD decomposition + ASCII-only encoding) so that
accented characters map to their ASCII base letters (e.g. é → e),
collapses any run of characters outside [A-Za-z0-9._-] into a single
'-', strips leading/trailing '-._' so paths don't start with a dot
or hyphen, and falls back to 'model' when the result is empty.
"""
if not name:
return "model"
name = unicodedata.normalize("NFKD", name)
name = name.encode("ascii", "ignore").decode("ascii")
slug = _SLUG_INVALID.sub("-", name)
slug = slug.strip("-._")
return slug or "model"
[docs]class ArtifactProcessor:
[docs] def __init__(
self,
partition_key: str = None,
config: Config = None,
model: Model = None,
model_artifacts: ModelArtifacts = None,
):
self.config = config or Config()
self.storage = UnifiedArtifactStorage(
storage_type=self.config.artifact_storage_type,
partition_key=partition_key,
**self.config.artifact_storage_config,
)
self.logger = Logger.get_logger(__name__)
self.queue_client = AzureQueueHandler(
self.config.queue_config["queue_connection_string"],
self.config.queue_config["zip_queue_name"],
self.config.queue_config["queue_account_url"],
)
self.model_data = model
self.runner = UnifiedRunner(
runner_type=self.config.runner_type,
config=self.config,
pool_id=self.config.get_azure_batch_config()["imageprep_pool_id"],
candidate_pool_ids=self.config.get_azure_batch_config()[
"imageryprep_pool_ids"
],
)
self.model_artifacts = model_artifacts
if self.model_data is not None:
safe_name = _slugify_model_name(self.model_data.name)
self.zip_name = self.config.get_artifact_types().MODEL_ARTIFACTS_ZIP.value.substitute(
modelName=safe_name
)
self.training_zip_name = self.config.get_artifact_types().TRAINING_ARTIFACTS_ZIP.value.substitute(
modelName=safe_name
)
self.inference_zip_name = self.config.get_artifact_types().INFERENCE_ARTIFACTS_ZIP.value.substitute(
modelName=safe_name
)
[docs] def get_download_url(
self,
identifier=None,
artifact_path=None,
extra_partition_keys=None,
):
"""
Get the download URL for the artifact.
"""
return self.storage.get_download_url(
identifier=identifier,
artifact_path=artifact_path,
extra_partition_keys=extra_partition_keys,
)
[docs] def send_to_zip_queue(self):
"""
Put a message to the queue.
"""
self.model_artifacts.zipStatus = (
self.config.get_status_types().PENDING.value
)
self.model_artifacts.zipStatusMessage = (
MetadataUtils.append_status_message("", "Queued for zipping")
)
self.model_artifacts.zipUrl = None
self.model_artifacts.currentZipJobUid = None
# Earlier zip jobs keep their status and dates but not their logs.
# Each copy is a finished run's history that nothing reads, and the
# whole record is queued now and on every poll: keeping one per run
# would let repeated zips outgrow the 64 KiB queue message limit.
for previous_job in self.model_artifacts.zipJobs:
previous_job.logs = ""
# Setting visibility timeout to 0 to make sure the message is processed immediately
self.queue_client.put_message(
self._queue_payload(), visibility_timeout=0
)
return self.model_artifacts
[docs] def process_zip(self):
self.logger.info(
f"{self.__class__.__name__}.process: Processing artifacts for "
f"model {self.model_artifacts.modelId} with status "
f"{self.model_artifacts.zipStatus}"
)
if (
self.model_artifacts.zipStatus
== self.config.get_status_types().PENDING.value
):
self.logger.info(
f"Zipping artifacts for model {self.model_artifacts.modelId}"
)
self._update_zip_progress("Submitting zip task")
self.model_artifacts = self.submit_zip_job()
elif (
self.model_artifacts.zipStatus
== self.config.get_status_types().IN_PROGRESS.value
):
for idx, zip_job in enumerate(self.model_artifacts.zipJobs):
if zip_job.taskId == self.model_artifacts.currentZipJobUid:
break
task_status = self.runner.get_task_status(
job_id=self.model_artifacts.zipJobs[idx].jobId,
task_id=self.model_artifacts.zipJobs[idx].taskId,
)
self.logger.info(
f"Task status of zip job for model {self.model_artifacts.modelId} is {task_status}"
)
if task_status == self.config.get_status_types().COMPLETED.value:
self.model_artifacts.zipStatus = task_status
self.model_artifacts.zipJobs[idx].status = task_status
self.model_artifacts.zipJobs[
idx
].completedDate = MetadataUtils.get_timestamp()
# Read the zip manifest to get individual zip URLs and sizes
zip_task_id = self.model_artifacts.zipJobs[idx].taskId
zip_prefix = self.model_artifacts.zipJobs[idx].dstZipPath
try:
manifest_data = self._read_zip_manifest(zip_prefix)
except Exception as e:
self.logger.warning(
f"Could not read zip manifest: {e}; "
"falling back to combined zip URL"
)
manifest_data = {}
if "training_zip" in manifest_data:
self.model_artifacts.trainingZipUrl = (
self.storage.get_download_url(
identifier=manifest_data["training_zip"][
"filename"
],
extra_partition_keys=zip_task_id,
)
)
self.model_artifacts.trainingZipSize = manifest_data[
"training_zip"
]["size_bytes"]
if "inference_zip" in manifest_data:
self.model_artifacts.inferenceZipUrl = (
self.storage.get_download_url(
identifier=manifest_data["inference_zip"][
"filename"
],
extra_partition_keys=zip_task_id,
)
)
self.model_artifacts.inferenceZipSize = manifest_data[
"inference_zip"
]["size_bytes"]
# Keep legacy zipUrl pointing at the training zip for
# backwards compatibility with older UI versions.
self.model_artifacts.zipUrl = (
self.model_artifacts.trainingZipUrl
or self.model_artifacts.inferenceZipUrl
)
self._update_zip_progress(
"Zipping artifacts completed successfully"
)
self.model_artifacts.zipJobs[
idx
].logs = self.model_artifacts.zipStatusMessage
# Cleanup the task on the runner
self.runner.cleanup_task(
job_id=self.model_artifacts.zipJobs[idx].jobId,
task_id=self.model_artifacts.zipJobs[idx].taskId,
)
elif task_status == self.config.get_status_types().FAILED.value:
self.model_artifacts.zipStatus = task_status
self.model_artifacts.zipJobs[idx].status = task_status
self.model_artifacts.zipJobs[
idx
].completedDate = MetadataUtils.get_timestamp()
self._update_zip_progress("Zip job failed")
self.model_artifacts.zipJobs[
idx
].logs = self.model_artifacts.zipStatusMessage
# Cleanup the task on the runner
self.runner.cleanup_task(
job_id=self.model_artifacts.zipJobs[idx].jobId,
task_id=self.model_artifacts.zipJobs[idx].taskId,
)
else:
self.model_artifacts.zipStatus = task_status
self.model_artifacts.zipJobs[idx].status = task_status
# Replace rather than append: the whole record is re-queued
# on every poll and a queue message cannot exceed 64 KiB.
self._update_zip_progress(
ZIP_IN_PROGRESS_MESSAGE,
replace_prefix=ZIP_IN_PROGRESS_MESSAGE,
)
self.model_artifacts.zipJobs[
idx
].logs = self.model_artifacts.zipStatusMessage
self.queue_client.put_message(self._queue_payload())
else:
self.model_artifacts.zipStatus = (
self.config.get_status_types().FAILED.value
)
self.logger.info(
f"Model {self.model_artifacts.modelId} is not ready for zipping"
)
return self.model_artifacts
[docs] def fetch_artifact(
self,
identifier: str = None,
extra_partition_keys: list | str = None,
src_path: str = None,
dst_path: str = None,
) -> str:
"""
Download the artifact from the storage.
"""
return self.storage.fetch_artifact(
identifier=identifier,
extra_partition_keys=extra_partition_keys,
src_path=src_path,
dst_path=dst_path,
)
[docs] def store_artifact(
self,
artifact_name: str,
data: str = None,
src_path: str = None,
namespace: str | list = None,
*,
overwrite: bool = True,
) -> str:
"""
Store an artifact in artifact storage.
"""
return self.storage.store_artifact(
artifact_name=artifact_name,
data=data,
src_path=src_path,
namespace=namespace,
**({"overwrite": False} if not overwrite else {}),
)
[docs] def prepare_zip_job(self):
zip_input_files = {}
if self.model_data.trainingOutputPath:
zip_input_files["training"] = {
"storage_container_url": f"{self.storage.get_base_url()}",
"blob_prefix": self.model_data.trainingOutputPath,
"file_path": "inputs/",
}
if self.model_data.inferenceOutputPath:
zip_input_files["inference"] = {
"storage_container_url": f"{self.storage.get_base_url()}",
"blob_prefix": self.model_data.inferenceOutputPath,
"file_path": "inputs/",
}
return zip_input_files
[docs] def submit_zip_job(self):
try:
self.logger.info(
f"Adding task for artifact zipping {self.model_artifacts.modelId}"
)
zip_input_files = self.prepare_zip_job()
command = '"python -m hastegeo.workflows.zip_artifacts"'
job_id = self.config.get_azure_batch_config()[
"artifact_batch_job_id"
]
# Trim job_id to 64 characters to comply with Azure Batch limits
job_id = job_id[:64]
task_id = f"{ZIP_PREFIX}-{MetadataUtils.generate_id()}"
zip_output_prefix = f"{MetadataUtils.hash_string(self.model_artifacts.projectId)}/{task_id}"
job_id, task_id = self.runner.add_task(
job_id=job_id,
task_id=task_id,
output_prefix=zip_output_prefix,
resource_files_for_upload=zip_input_files,
file_pattern=f"${BATCH_JOB_WORKDIR}/outputs/*.*",
command=command,
image_name=self.config.get_azure_batch_config()[
"imageprep_docker_image"
],
env_vars={
"INPUT_DIR": f"inputs/{MetadataUtils.hash_string(self.model_artifacts.projectId)}",
"OUTPUT_TRAINING_ZIP_NAME": f"{self.training_zip_name}",
"OUTPUT_INFERENCE_ZIP_NAME": f"{self.inference_zip_name}",
},
)
self.logger.info(
f"Completed add task {task_id} to job id {job_id} for zipping artifacts of model {self.model_artifacts.modelId}"
)
srcArtifactPaths = []
if self.model_data.trainingOutputPath:
srcArtifactPaths.append(self.model_data.trainingOutputPath)
if self.model_data.inferenceOutputPath:
srcArtifactPaths.append(self.model_data.inferenceOutputPath)
self.model_artifacts.zipJobs.append(
ZipJob(
projectId=self.model_artifacts.projectId,
imageLayerId=self.model_artifacts.imageLayerId,
modelId=self.model_artifacts.modelId,
jobId=job_id,
taskId=task_id,
status=self.config.get_status_types().IN_PROGRESS.value,
srcArtifactPaths=srcArtifactPaths,
dstZipPath=zip_output_prefix,
creationDate=MetadataUtils.get_timestamp(),
)
)
self.model_artifacts.currentZipJobUid = task_id
self.model_artifacts.zipStatus = (
self.config.get_status_types().IN_PROGRESS.value
)
self._update_zip_progress(
f"Zipping submitted with task id {task_id}"
)
self.queue_client.put_message(self._queue_payload())
self.logger.info(
f"InProgress message to queue sent for model {self.model_artifacts.modelId}"
)
except Exception as e:
self.logger.error(
f"Error processing model {self.model_artifacts.modelId}: {e}",
stack_info=True,
)
self.model_artifacts.zipStatus = (
self.config.get_status_types().FAILED.value
)
return self.model_artifacts
def _read_zip_manifest(self, zip_prefix: str) -> dict:
"""Read the zip_manifest.json produced by zip_artifacts.py."""
blob_path = f"{zip_prefix}/zip_manifest.json"
blob_client = (
self.storage.artifact_storage.container_client.get_blob_client(
blob_path
)
)
data = blob_client.download_blob().readall()
return json.loads(data)
def _queue_payload(self) -> str:
"""Serialize the artifacts record for the zip queue within its limit.
Finished zip jobs' logs are left out of the queued copy, since only
the current job's logs are read. The status history and the current
job's copy of it share the room the rest of the record leaves.
"""
record = self.model_artifacts.dict()
histories = [(record, "zipStatusMessage")]
for job in record["zipJobs"]:
if job["taskId"] == record["currentZipJobUid"]:
histories.append((job, "logs"))
else:
job["logs"] = ""
return MetadataUtils.fit_queue_payload(record, histories)
def _update_zip_progress(
self, message: str, timestamp: str = None, replace_prefix: str = None
):
if replace_prefix:
status_message = MetadataUtils.upsert_status_message(
self.model_artifacts.zipStatusMessage,
message,
replace_prefix,
timestamp=timestamp,
)
else:
status_message = MetadataUtils.append_status_message(
self.model_artifacts.zipStatusMessage,
message,
timestamp=timestamp,
)
self.model_artifacts.zipStatusMessage = (
MetadataUtils.trim_status_message(status_message)
)