feat: materialize distributed inputs in task pool
This commit is contained in:
+21
-1
@@ -93,6 +93,14 @@ def _default_runtime_root(project_root: str) -> str:
|
||||
return os.path.join(normalized_root, "runtime")
|
||||
|
||||
|
||||
def _default_task_pool_root(project_root: str) -> str:
|
||||
normalized_root = os.path.normpath(project_root)
|
||||
drive, _tail = os.path.splitdrive(normalized_root)
|
||||
if drive:
|
||||
return os.path.join(drive + os.sep, "Task_Pool")
|
||||
return os.path.join(normalized_root, "Task_Pool")
|
||||
|
||||
|
||||
def _default_runtime_dir(project_root: str, *parts: str) -> str:
|
||||
return os.path.join(_default_runtime_root(project_root), *parts)
|
||||
|
||||
@@ -187,6 +195,9 @@ class Settings(BaseSettings):
|
||||
DB_SCHEMA_RESET_CONFIRM: bool = False
|
||||
|
||||
UNPACK_SOURCE_DIRS: str = ""
|
||||
TASK_POOL_ROOT: str = ""
|
||||
DINSAR_TASK_POOL_ROOT: str = ""
|
||||
SBAS_TASK_POOL_ROOT: str = ""
|
||||
SOURCE_PRODUCT_DIRS: str = ""
|
||||
SENTINEL1_STORAGE_DIRS: str = ""
|
||||
ORBIT_SOURCE_DIRS: str = ""
|
||||
@@ -441,6 +452,12 @@ class Settings(BaseSettings):
|
||||
"WATER_RESULTS_DIR",
|
||||
os.path.join(backend_dir, "water_results"),
|
||||
)
|
||||
if not self.TASK_POOL_ROOT:
|
||||
object.__setattr__(self, "TASK_POOL_ROOT", _default_task_pool_root(project_root))
|
||||
if not self.DINSAR_TASK_POOL_ROOT:
|
||||
object.__setattr__(self, "DINSAR_TASK_POOL_ROOT", os.path.join(self.TASK_POOL_ROOT, "DInSAR"))
|
||||
if not self.SBAS_TASK_POOL_ROOT:
|
||||
object.__setattr__(self, "SBAS_TASK_POOL_ROOT", os.path.join(self.TASK_POOL_ROOT, "SBAS"))
|
||||
if not self.SAR_ANALYSIS_READY_ROOT:
|
||||
object.__setattr__(
|
||||
self,
|
||||
@@ -780,7 +797,7 @@ class Settings(BaseSettings):
|
||||
object.__setattr__(
|
||||
self,
|
||||
"GAMMA_SBAS_WORK_ROOT",
|
||||
os.path.join(_default_runtime_root(project_root), "sbas_insar_work"),
|
||||
self.SBAS_TASK_POOL_ROOT,
|
||||
)
|
||||
if not self.GAMMA_SBAS_PRODUCT_ROOT:
|
||||
object.__setattr__(
|
||||
@@ -990,6 +1007,9 @@ class Settings(BaseSettings):
|
||||
os.makedirs(settings.RESULT_QUARANTINE_ROOT, exist_ok=True)
|
||||
os.makedirs(settings.SAR_ANALYSIS_READY_ROOT, exist_ok=True)
|
||||
os.makedirs(settings.SAR_ANALYSIS_WORK_ROOT, exist_ok=True)
|
||||
os.makedirs(settings.TASK_POOL_ROOT, exist_ok=True)
|
||||
os.makedirs(settings.DINSAR_TASK_POOL_ROOT, exist_ok=True)
|
||||
os.makedirs(settings.SBAS_TASK_POOL_ROOT, exist_ok=True)
|
||||
for path in split_env_paths(settings.GF3_ARCHIVE_SOURCE_DIRS):
|
||||
os.makedirs(path, exist_ok=True)
|
||||
for path in split_env_paths(settings.GF3_SARSCAPE_NATIVE_DIRS):
|
||||
|
||||
+118
-2
@@ -2,6 +2,7 @@
|
||||
import shutil
|
||||
import asyncio
|
||||
import tempfile
|
||||
import tarfile
|
||||
import zipfile
|
||||
import json
|
||||
import hashlib
|
||||
@@ -43,6 +44,100 @@ def find_dinsar_source_to_copy(path: str) -> str:
|
||||
return path
|
||||
|
||||
|
||||
_DINSAR_ARCHIVE_SUFFIXES = (".tar.gz", ".tgz", ".zip", ".tar")
|
||||
|
||||
|
||||
def _is_supported_archive(path: str) -> bool:
|
||||
lower = str(path or "").lower()
|
||||
return any(lower.endswith(suffix) for suffix in _DINSAR_ARCHIVE_SUFFIXES)
|
||||
|
||||
|
||||
def _safe_archive_member_name(member_name: str, archive_path: str) -> str:
|
||||
name = str(member_name or "").replace("\\", "/").strip("/")
|
||||
if not name or name.startswith("../") or "/../" in f"/{name}/":
|
||||
raise ValueError(f"Unsafe archive member path in {archive_path}: {member_name}")
|
||||
if os.path.isabs(name) or os.path.splitdrive(name)[0]:
|
||||
raise ValueError(f"Unsafe archive member path in {archive_path}: {member_name}")
|
||||
return name
|
||||
|
||||
|
||||
def _extract_archive_to_dir(archive_path: str, dest_dir: str) -> int:
|
||||
if zipfile.is_zipfile(archive_path):
|
||||
extracted = 0
|
||||
with zipfile.ZipFile(archive_path) as zip_obj:
|
||||
for info in zip_obj.infolist():
|
||||
rel_name = _safe_archive_member_name(info.filename, archive_path)
|
||||
dest_path = os.path.abspath(os.path.join(dest_dir, rel_name))
|
||||
if not dest_path.startswith(os.path.abspath(dest_dir) + os.sep):
|
||||
raise ValueError(f"Unsafe ZIP member path: {info.filename}")
|
||||
if info.is_dir():
|
||||
os.makedirs(dest_path, exist_ok=True)
|
||||
continue
|
||||
os.makedirs(os.path.dirname(dest_path), exist_ok=True)
|
||||
with zip_obj.open(info, "r") as source, open(dest_path, "wb") as target:
|
||||
shutil.copyfileobj(source, target, length=1024 * 1024)
|
||||
extracted += 1
|
||||
return extracted
|
||||
|
||||
if tarfile.is_tarfile(archive_path):
|
||||
extracted = 0
|
||||
with tarfile.open(archive_path, "r:*") as tar_obj:
|
||||
for member in tar_obj:
|
||||
rel_name = _safe_archive_member_name(member.name, archive_path)
|
||||
dest_path = os.path.abspath(os.path.join(dest_dir, rel_name))
|
||||
if not dest_path.startswith(os.path.abspath(dest_dir) + os.sep):
|
||||
raise ValueError(f"Unsafe TAR member path: {member.name}")
|
||||
if member.isdir():
|
||||
os.makedirs(dest_path, exist_ok=True)
|
||||
continue
|
||||
if not member.isfile():
|
||||
continue
|
||||
source = tar_obj.extractfile(member)
|
||||
if source is None:
|
||||
continue
|
||||
os.makedirs(os.path.dirname(dest_path), exist_ok=True)
|
||||
with source, open(dest_path, "wb") as target:
|
||||
shutil.copyfileobj(source, target, length=1024 * 1024)
|
||||
extracted += 1
|
||||
return extracted
|
||||
|
||||
raise ValueError(f"Unsupported archive format: {archive_path}")
|
||||
|
||||
|
||||
def _materialize_dinsar_source(source_path: str, dest_dir: str) -> Dict[str, Any]:
|
||||
normalized = os.path.normpath(os.path.abspath(str(source_path or "")))
|
||||
if not os.path.exists(normalized):
|
||||
raise FileNotFoundError(normalized)
|
||||
|
||||
if os.path.isdir(normalized):
|
||||
shutil.copytree(normalized, dest_dir, dirs_exist_ok=True)
|
||||
return {"mode": "copy_directory", "source_path": normalized}
|
||||
|
||||
if os.path.isfile(normalized) and _is_supported_archive(normalized):
|
||||
os.makedirs(dest_dir, exist_ok=True)
|
||||
extracted = _extract_archive_to_dir(normalized, dest_dir)
|
||||
if extracted <= 0:
|
||||
raise OSError(f"Archive extraction produced no files: {normalized}")
|
||||
return {
|
||||
"mode": "extract_archive",
|
||||
"source_path": normalized,
|
||||
"archive_path": normalized,
|
||||
"extracted_files": extracted,
|
||||
}
|
||||
|
||||
if os.path.isfile(normalized):
|
||||
os.makedirs(dest_dir, exist_ok=True)
|
||||
target = os.path.join(dest_dir, os.path.basename(normalized))
|
||||
shutil.copy2(normalized, target)
|
||||
return {
|
||||
"mode": "copy_file",
|
||||
"source_path": normalized,
|
||||
"relative_path": os.path.basename(target),
|
||||
}
|
||||
|
||||
raise FileNotFoundError(normalized)
|
||||
|
||||
|
||||
def _resolve_orbit_dest_path(
|
||||
orbit_dir: str,
|
||||
role: str,
|
||||
@@ -191,6 +286,7 @@ def _build_dinsar_pair_metadata(
|
||||
package_format: str,
|
||||
include_orbit_files: bool,
|
||||
orbit_entries: List[Dict[str, Any]],
|
||||
source_materialization: Optional[Dict[str, Any]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
return {
|
||||
"pair_key": item.get("pair_key"),
|
||||
@@ -214,6 +310,7 @@ def _build_dinsar_pair_metadata(
|
||||
"master_orbit_file_path": item.get("master_orbit_file_path"),
|
||||
"slave_orbit_file_path": item.get("slave_orbit_file_path"),
|
||||
"orbit_files": orbit_entries,
|
||||
"source_materialization": source_materialization or {},
|
||||
"scene_pair_uid": item.get("scene_pair_uid") or item.get("pair_uid"),
|
||||
"pair_uid": item.get("pair_uid") or item.get("scene_pair_uid"),
|
||||
"network_run_id": item.get("network_run_id"),
|
||||
@@ -1094,8 +1191,26 @@ async def run_dinsar_copy_items(
|
||||
failed_count += 1
|
||||
continue
|
||||
|
||||
await asyncio.to_thread(shutil.copytree, master_src_path, master_dir, dirs_exist_ok=True)
|
||||
await asyncio.to_thread(shutil.copytree, slave_src_path, slave_dir, dirs_exist_ok=True)
|
||||
master_materialization = await asyncio.to_thread(
|
||||
_materialize_dinsar_source,
|
||||
master_src_path,
|
||||
master_dir,
|
||||
)
|
||||
slave_materialization = await asyncio.to_thread(
|
||||
_materialize_dinsar_source,
|
||||
slave_src_path,
|
||||
slave_dir,
|
||||
)
|
||||
source_materialization = {
|
||||
"master": {
|
||||
**master_materialization,
|
||||
"target_relative_path": "master",
|
||||
},
|
||||
"slave": {
|
||||
**slave_materialization,
|
||||
"target_relative_path": "slave",
|
||||
},
|
||||
}
|
||||
orbit_entries = await _copy_dinsar_orbit_files(
|
||||
task_id,
|
||||
item,
|
||||
@@ -1112,6 +1227,7 @@ async def run_dinsar_copy_items(
|
||||
"zip" if export_zip else "folder",
|
||||
include_orbit_files,
|
||||
orbit_entries,
|
||||
source_materialization,
|
||||
),
|
||||
)
|
||||
if export_zip and zip_path:
|
||||
|
||||
@@ -458,6 +458,35 @@ def _build_root_specs_from_settings() -> List[RootSpec]:
|
||||
owner_engine="pyint",
|
||||
)
|
||||
)
|
||||
specs.extend(
|
||||
_iter_single_root_specs(
|
||||
env_var="TASK_POOL_ROOT",
|
||||
path=settings.TASK_POOL_ROOT,
|
||||
root_role="task_pool_root",
|
||||
display_name="Task Pool Root",
|
||||
scan_mode="workspace",
|
||||
)
|
||||
)
|
||||
specs.extend(
|
||||
_iter_single_root_specs(
|
||||
env_var="DINSAR_TASK_POOL_ROOT",
|
||||
path=settings.DINSAR_TASK_POOL_ROOT,
|
||||
root_role="task_pool_dinsar",
|
||||
display_name="D-InSAR Task Pool",
|
||||
scan_mode="workspace",
|
||||
owner_engine="dinsar",
|
||||
)
|
||||
)
|
||||
specs.extend(
|
||||
_iter_single_root_specs(
|
||||
env_var="SBAS_TASK_POOL_ROOT",
|
||||
path=settings.SBAS_TASK_POOL_ROOT,
|
||||
root_role="task_pool_sbas",
|
||||
display_name="SBAS Task Pool",
|
||||
scan_mode="workspace",
|
||||
owner_engine="sbas",
|
||||
)
|
||||
)
|
||||
specs.extend(
|
||||
_iter_single_root_specs(
|
||||
env_var="GAMMA_SBAS_WORK_ROOT",
|
||||
|
||||
Reference in New Issue
Block a user