diff --git a/src/aignostics.py b/src/aignostics.py index d74c56317..4f03ea5c8 100644 --- a/src/aignostics.py +++ b/src/aignostics.py @@ -29,25 +29,12 @@ boot(SENTRY_INTEGRATIONS) -EXEC_SCRIPT_FLAG = "--exec-script" -MIN_ARGS_FOR_SCRIPT = 3 # program name, flag, and script content MODULE_FLAG = "--run-module" MIN_ARGS_FOR_MODULE = 3 # program name, flag, and module name DEBUG_FLAG = "--debug" -if len(sys.argv) > 1 and sys.argv[1] == EXEC_SCRIPT_FLAG: - if len(sys.argv) >= MIN_ARGS_FOR_SCRIPT: - script_content = sys.argv[2] - try: - exec(script_content) # noqa: S102 - except Exception: - logger.exception("Failed to execute script") - sys.exit(1) - else: - logger.error("No script content provided") - sys.exit(1) -elif len(sys.argv) > 1 and sys.argv[1] == MODULE_FLAG: +if len(sys.argv) > 1 and sys.argv[1] == MODULE_FLAG: if pyi_splash and pyi_splash.is_alive(): pyi_splash.close() diff --git a/src/aignostics/dataset/_download_worker.py b/src/aignostics/dataset/_download_worker.py new file mode 100644 index 000000000..9e75b9f5c --- /dev/null +++ b/src/aignostics/dataset/_download_worker.py @@ -0,0 +1,42 @@ +"""IDC download worker — runs as a subprocess, invoked via --run-module or -m.""" + +import json +import sys +from pathlib import Path +from typing import cast + +MIN_ARGS = 2 # program name + config file path + + +def main() -> None: + """Read a JSON config file and run IDC download_from_selection.""" + if len(sys.argv) < MIN_ARGS: + print("Usage: _download_worker ", file=sys.stderr) + sys.exit(1) + + config_path = Path(sys.argv[1]) + if config_path.suffix != ".json" or not config_path.is_file() or not config_path.is_absolute(): + print(f"Invalid config file path: {config_path}", file=sys.stderr) + sys.exit(1) + + config: dict[str, object] = json.loads(config_path.read_text()) + + from aignostics.third_party.idc_index import IDCClient # noqa: PLC0415 + + client = IDCClient.client() + client.fetch_index("sm_instance_index") + kwarg_name = str(config["kwarg_name"]) + matched_ids = cast("list[str]", config["matched_ids"]) + client.download_from_selection( # type: ignore[no-untyped-call] # pyright: ignore[reportArgumentType] + **{kwarg_name: matched_ids}, # pyright: ignore[reportArgumentType] + downloadDir=str(config["download_dir"]), + dirTemplate=str(config["dir_template"]), + quiet=False, + show_progress_bar=True, + use_s5cmd_sync=True, + dry_run=bool(config["dry_run"]), + ) + + +if __name__ == "__main__": + main() diff --git a/src/aignostics/dataset/_service.py b/src/aignostics/dataset/_service.py index aa3d66ed7..aff763d23 100644 --- a/src/aignostics/dataset/_service.py +++ b/src/aignostics/dataset/_service.py @@ -1,9 +1,11 @@ """Service of dataset module.""" import atexit +import json import re import subprocess import sys +import tempfile import threading import time from collections.abc import Callable @@ -218,53 +220,37 @@ def check_and_download(column_name: str, item_ids: list[str], target_directory: logger.debug("Identified matching {}: {}", column_name, matched_ids) queue.put_nowait(0.04) - # Properly handle Windows paths - convert to raw string format - safe_target_dir = str(target_directory).replace("\\", "\\\\") - - # Create command for the subprocess - script_content = f""" -import sys -from aignostics.third_party.idc_index import IDCClient - -client = IDCClient.client() -client.fetch_index("sm_instance_index") -client.download_from_selection( - {kwarg_name}={matched_ids!r}, - downloadDir="{safe_target_dir}", - dirTemplate="{target_layout}", - quiet=False, - show_progress_bar=True, - use_s5cmd_sync=True, - dry_run={dry_run!r} -) -""" - - # Run the download in a subprocess + # Write download parameters to a temp file so the worker subprocess + # receives structured data instead of executable code. + config = { + "kwarg_name": kwarg_name, + "matched_ids": matched_ids, + "download_dir": str(target_directory), + "dir_template": target_layout, + "dry_run": dry_run, + } + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tf: + json.dump(config, tf) + config_file = Path(tf.name) + + worker = "aignostics.dataset._download_worker" if getattr(sys, "frozen", False): - # When running under PyInstaller, sys.executable points to the PyInstaller executable. - # We use a special flag to execute the script without launching the GUI. - # See src/aignostics.py - logger.trace("Running under PyInstaller - using --exec-script flag") - process = subprocess.Popen( # noqa: S603 - [sys.executable, "--exec-script", script_content], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - bufsize=1, - creationflags=SUBPROCESS_CREATION_FLAGS, - ) + # sys.executable is the PyInstaller binary; use --run-module to + # invoke the worker without launching the GUI. See src/aignostics.py + logger.trace("Running under PyInstaller - using --run-module flag") + cmd = [sys.executable, "--run-module", worker, str(config_file)] else: - logger.trace( - "Starting download subprocess with executable '{}' and script:\n{}", sys.executable, script_content - ) - process = subprocess.Popen( # noqa: S603 - [sys.executable, "-c", script_content], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - bufsize=1, - creationflags=SUBPROCESS_CREATION_FLAGS, - ) + logger.trace("Starting download subprocess with executable '{}'", sys.executable) + cmd = [sys.executable, "-m", worker, str(config_file)] + + process = subprocess.Popen( # noqa: S603 + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + creationflags=SUBPROCESS_CREATION_FLAGS, + ) # Register process for cleanup _active_processes.append(process) @@ -295,9 +281,9 @@ def check_and_download(column_name: str, item_ids: list[str], target_directory: queue.put_nowait(1.0) return True finally: - # Clean up process reference if process in _active_processes: _active_processes.remove(process) + config_file.unlink(missing_ok=True) matches_found = 0 matches_found += check_and_download("collection_id", item_ids, target_directory, "collection_id")