File size: 4,004 Bytes
51e84d2 fc1dcbc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | # License: Apache-2.0 License
# Created by: Patrick Lumbantobing, VertoX-AI
# Copyright (c) 2026 VertoX-AI. All rights reserved.
#
# This work is licensed under the Apache-2.0 License.
# To view a copy of this license, visit
# https://www.apache.org/licenses/LICENSE-2.0
"""Manage safe per-run temporary workspaces and optional diagnostic retention."""
from __future__ import annotations
import shutil
import tempfile
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from enum import StrEnum
from pathlib import Path
from speech_text_data_aligners.core.errors import ConfigurationError
class ArtifactPolicy(StrEnum):
"""Select when a completed temporary workspace should be retained."""
NEVER = "never"
ON_FAILURE = "on_failure"
ALWAYS = "always"
@dataclass(slots=True)
class Workspace:
"""Expose a scoped temporary path and its eventual retained copy.
Parameters
----------
path: Unique temporary work directory.
retained_path: Final copied artifact directory, populated during cleanup.
Invariants:
``path`` exists only within the surrounding :func:`workspace` context.
"""
path: Path
retained_path: Path | None = None
failed: bool = False
def mark_failed(self) -> None:
"""Mark a handled failure so ``ON_FAILURE`` policy retains artifacts.
Side Effects:
Changes this workspace's cleanup disposition.
"""
self.failed = True
@contextmanager
def workspace(
*,
run_id: str,
policy: ArtifactPolicy = ArtifactPolicy.NEVER,
temporary_root: Path | None = None,
retention_root: Path | None = None,
) -> Iterator[Workspace]:
"""Create a unique temporary workspace and safely apply retention policy.
Parameters
----------
run_id: Path-component-free identifier used for a retained directory name.
policy: Retain never, only after failure, or always.
temporary_root: Optional parent for temporary workspace creation.
retention_root: Required destination parent when retention can occur.
Yields
------
Mutable workspace status object valid for the context lifetime.
Raises
------
ConfigurationError: If identifiers/roots are unsafe or retention would
overwrite an existing artifact directory.
Side Effects:
Creates a temporary directory, optionally copies it to ``retention_root``, and
always removes the temporary directory at context exit.
"""
if not run_id or Path(run_id).name != run_id or run_id in {".", ".."}:
raise ConfigurationError("workspace run_id must be one safe path component")
if policy is not ArtifactPolicy.NEVER and retention_root is None:
raise ConfigurationError("retention_root is required by artifact policy")
if temporary_root is not None:
temporary_root.mkdir(parents=True, exist_ok=True)
with tempfile.TemporaryDirectory(
prefix=f"speech-aligners-{run_id}-",
dir=temporary_root,
) as temporary_name:
state = Workspace(Path(temporary_name))
raised = False
try:
yield state
except BaseException:
raised = True
state.failed = True
raise
finally:
should_retain = policy is ArtifactPolicy.ALWAYS or (
policy is ArtifactPolicy.ON_FAILURE and (raised or state.failed)
)
if should_retain:
assert retention_root is not None
retention_root.mkdir(parents=True, exist_ok=True)
destination = retention_root / run_id
if destination.exists():
raise ConfigurationError(
f"retained artifact destination already exists: {destination}"
)
shutil.copytree(state.path, destination)
state.retained_path = destination
|