simstack 2.0.1__py3-none-any.whl

simstack/core/submit_node.py

@@ -0,0 +1,217 @@
+import os
+import re
+import stat
+import subprocess
+from pathlib import Path
+
+from simstack.core.context import context
+from simstack.core.definitions import TaskStatus
+from simstack.models import NodeRegistry
+from simstack.util.submit_to_watchdog import submit_to_watchdog
+import logging
+
+logger = logging.getLogger("submit_node")
+
+
+def make_executable(file_path: str | os.PathLike[str]) -> None:
+    # Get current permissions
+    current_permissions = os.stat(file_path).st_mode
+
+    # Add executable bit for user, group and others
+    executable_mode = current_permissions | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
+
+    # Apply new permissions
+    os.chmod(file_path, executable_mode)
+
+
+async def submit_node(registry_entry: NodeRegistry) -> bool:
+    """Submit a node to the SLURM queue"""
+    task_id = registry_entry.id
+    try:
+        logger.info(f"Submitting task_id: {task_id} to SLURM queue")
+        # Implement SLURM submission logic here
+        base_path = context.config.project_root
+
+        python_path = ":".join(str(path) for path in context.config.python_paths)
+        work_dir = context.config.workdir / registry_entry.name / str(registry_entry.id)
+        job_name = registry_entry.name + "." + str(registry_entry.id)
+
+        logger.info(f"task_id: {task_id} workdir {work_dir} python path {python_path}")
+        # write a slurm script that starts a python script "run_node.py" with the node id as an argument
+
+        original_slurm_parameters = registry_entry.parameters.slurm_parameters
+        if original_slurm_parameters is None:
+            logger.error(f"Task task_id: {task_id} has no slurm parameters -- failing")
+            registry_entry.status = TaskStatus.FAILED
+            await context.db.save(registry_entry)
+            return False
+        slurm_parameters = original_slurm_parameters.model_copy(deep=True)
+        slurm_parameters.output = f"{work_dir}/%j.out"
+        slurm_parameters.error = f"{work_dir}/%j.err"
+        slurm_parameters.job_name = f"{job_name}"
+        slurm_parameters.chdir = str(work_dir)
+
+        slurm_parameters.startup_commands.append("source ~/.bashrc")
+        slurm_parameters.startup_commands.append(f"{context.config.environment_start}")
+        slurm_parameters.startup_commands.append(
+            f"export PYTHONPATH={python_path}:$PYTHONPATH"
+        )
+
+        sync_data_start = """# 1. Define the rsync function
+        sync_data() {
+            echo "=== Time limit approaching! Starting rsync at $(date) ==="
+
+            # Run rsync safely 
+        """
+        sync_data_end = """
+            echo "=== Rsync completed at $(date) ==="
+            exit 0
+        }        
+        """
+
+        if context.resource_config is not None:
+            program_config = context.resource_config.get_program(registry_entry.name)
+        else:
+            program_config = {}
+
+        logger.info(f"task_id: {task_id} program_config {program_config}")
+        if program_config.get("use_tmp", False):
+            tmp_dir = context.resource_config.tmp_dir(registry_entry.id)
+
+            if not tmp_dir.exists() or not tmp_dir.is_dir():
+                logger.error(
+                    f"Task task_id: {task_id} tmp_dir {tmp_dir} does not exist or is not a directory -- failing")
+                registry_entry.status = TaskStatus.FAILED
+                await context.db.save(registry_entry)
+                return False
+
+            full_sync_data = sync_data_start + f'cp -a {tmp_dir}/. {work_dir}/recovery\n' + sync_data_end
+
+            logger.info(f"task_id: {task_id} full_sync_data {full_sync_data}")
+            slurm_parameters.startup_commands.append(full_sync_data)
+            slurm_parameters.startup_commands.append("trap 'sync_data' SIGUSR1")
+
+        if context.config.docker:
+            external_work_dir = (
+                context.config.external_workdir
+                / registry_entry.name
+                / str(registry_entry.id)
+            )
+
+            watcher_file = (
+                context.config.project_root
+                / "src"
+                / "simstack"
+                / "util"
+                / "queue_watcher.py"
+            )
+            with open(watcher_file, "r") as f:
+                watcher_content = f.read()
+            slurm_parameters.startup_commands.append(
+                f"cat > watcher.py <<'EOF'\n{watcher_content}\nEOF"
+            )
+            slurm_parameters.startup_commands.append("python watcher.py &")
+            slurm_parameters.startup_commands.append("WATCHER_PID=$!")
+
+            slurm_parameters.startup_commands.append("set -euo pipefail")
+            slurm_parameters.startup_commands.append(
+                "cleanup() { "
+                'if [ -n "${WATCHER_PID:-}" ] && kill -0 "$WATCHER_PID" 2>/dev/null; then '
+                '  kill "$WATCHER_PID" 2>/dev/null || true; '
+                "  sleep 2; "
+                '  kill -9 "$WATCHER_PID" 2>/dev/null || true; '
+                "fi; "
+                "}"
+            )
+
+            slurm_parameters.startup_commands.append("trap cleanup EXIT INT TERM")
+            toml_path = context.config.project_root / "simstack_docker.toml"
+            if not toml_path.exists():
+                logger.error(
+                    f"Task task_id: {task_id} has no simstack.toml file -- failing"
+                )
+                registry_entry.status = TaskStatus.FAILED
+                await context.db.save(registry_entry)
+                return False
+
+            docker_start = "udocker run "
+            docker_start += "-e GIT_TOKEN=XXX"
+            docker_start += f"-e NODE_ID={registry_entry.id} "
+            docker_start += f"-e RESOURCE={str(context.config.resource)} "
+            docker_start += f"-v {external_work_dir}:/home/appuser/simstack "
+            docker_start += (
+                f"-v {context.config.external_source_dir}:/app/simstack-model "
+            )
+            docker_start += f"-v {toml_path}:/app/simstack-model/simstack.toml "
+            docker_start += " simstack-runner"
+            slurm_parameters.startup_commands.append(docker_start)
+        else:
+            slurm_parameters.startup_commands.append(
+                f"uv run --directory {base_path} run_node --node-id {registry_entry.id} --resource {str(context.config.resource)} &"
+            )
+            slurm_parameters.startup_commands.append("wait")
+        slurm_parameters.signal = "B:SIGUSR1@60"
+
+
+        slurm_script = slurm_parameters.to_sbatch_header()
+
+        # write the script to a file in the work_dir
+        os.makedirs(work_dir, exist_ok=True)
+        logger.info(f"task_id: {task_id} workdir {work_dir} python path {python_path}")
+        script_path = os.path.join(work_dir, "slurm_script.sh")
+        with open(script_path, "w") as f:
+            f.write(slurm_script)
+
+        make_executable(script_path)
+        # submit the script to the slurm queue
+        if context.config.docker:
+            job_id = "slurm_" + str(registry_entry.id)
+            queue_dir = context.config.workdir / "queue"
+            result = submit_to_watchdog(
+                f"/usr/bin/sbatch {os.path.join(external_work_dir, 'slurm_script.sh')}",
+                job_id,
+                queue_dir,
+            )
+        else:
+            result = subprocess.run(
+                f"/usr/bin/sbatch {os.path.join(work_dir, 'slurm_script.sh')}",
+                shell=True,
+                capture_output=True,
+                text=True,
+                timeout=30,  # Add timeout to prevent hanging
+            )
+
+        logger.info(
+            f"submitting job task_id: {task_id} returns: {result.returncode} {result.stdout} {result.stderr}"
+        )
+        if result.returncode == 0:
+            # Extract job ID using a regex pattern
+            match = re.search(r"Submitted batch job (\d+)", result.stdout)
+            if match:
+                job_id = match.group(1)
+                logger.info(
+                    f"task_id: {task_id} job successfully submitted with job_id: {job_id}"
+                )
+                registry_entry.job_id = job_id
+            else:
+                logger.warning(
+                    f"task_id: {task_id} job submitted but could not extract job_id from output: {result.stdout}"
+                )
+        else:
+            logger.error(
+                f"error submitting job for task_id: {task_id} return code: {result.returncode} stdout: {result.stdout}"
+            )
+            logger.error(
+                f"submitting job for task_id: {task_id} stderr: {result.stderr}"
+            )
+            registry_entry.status = TaskStatus.FAILED
+            await context.db.save(registry_entry)
+            return False
+        registry_entry.status = TaskStatus.SLURM_QUEUED
+        await context.db.save(registry_entry)
+        return True
+    except Exception as e:
+        logger.exception(f"fatal error in submitting task_id: {task_id} {str(e)}")
+        registry_entry.status = TaskStatus.FAILED
+        await context.db.save(registry_entry)
+        return False

simstack/core/task_id.py

@@ -0,0 +1,22 @@
+from contextvars import ContextVar
+from typing import Optional
+
+from bson.objectid import ObjectId
+
+# this variable is meant to give access to the task id to user functions that do not have access to it (e.g. for logging)
+_task_id: ContextVar[Optional[ObjectId]] = ContextVar("task_id", default=None)
+
+
+def get_task_id() -> Optional[ObjectId]:
+    """Get the current task ID."""
+    return _task_id.get()
+
+
+def set_task_id(task_id: Optional[ObjectId]) -> None:
+    """Set the current task ID."""
+    _task_id.set(task_id)
+
+
+def clear_task_id() -> None:
+    """Clear the current task ID."""
+    _task_id.set(None)

simstack/methods/archive_files.py

@@ -0,0 +1,292 @@
+"""Disabled archive prototype retained for a future, separately reviewed change.
+
+The public entrypoints intentionally fail closed. Permanent storage ownership and
+safe source cleanup must be designed and covered by production-path tests before
+this module can be enabled.
+"""
+
+import inspect
+import shutil
+from datetime import datetime
+from pathlib import Path
+from typing import List, Optional
+
+from odmantic import Model, Field
+
+from simstack.core.context import context
+from simstack.core.node import node
+from simstack.models import FileStack, FileList, simstack_model
+from simstack.models.base_lists import BooleanDataList
+from simstack.models.file_instance import FileInstance
+
+
+ARCHIVE_DISABLED_MESSAGE = (
+    "File archiving is disabled for this migration until permanent-resource "
+    "semantics and non-destructive cleanup are implemented and tested."
+)
+
+
+class ArchiveFeatureDisabledError(RuntimeError):
+    """Raised when code tries to use the excluded archive prototype."""
+
+
+def _raise_archive_disabled() -> None:
+    raise ArchiveFeatureDisabledError(ARCHIVE_DISABLED_MESSAGE)
+
+
+def archive_one_file(file_stack: FileStack,**kwargs):
+    _raise_archive_disabled()
+    node_runner = kwargs["node_runner"]
+    node_runner.info(f"Archiving file: {file_stack.name}")
+
+    config = context.resource_config.get_program("archive_one_file")
+    if config.get("archive_path", None) is None:
+        raise ValueError("archive_path not set in config")
+    archive_path = Path(config.get("archive_path"))
+    full_path = archive_path / str(file_stack.id)
+    full_path.mkdir(parents=True, exist_ok=True)
+    if full_path.exists() and full_path.is_dir():
+        try:
+            archive_location = file_stack.get(local_dir=full_path)
+            archive_file_instance = FileInstance.from_local_file(archive_location)
+            node_runner.info(f"Archived file to {archive_file_instance.path}")
+            file_stack.locations.append(archive_file_instance)
+            return True
+        except Exception as e:
+            node_runner.error(f"Error archiving file: {str(e)}")
+            return False
+    else:
+        raise ValueError(f"Archive path {archive_path} does not exist or is not a directory")
+
+@node
+async def archive_file(file_stack: FileStack,**kwargs):
+    _raise_archive_disabled()
+    return archive_one_file(file_stack, **kwargs)
+
+@node
+def archive_files(file_list: FileList,**kwargs):
+    _raise_archive_disabled()
+    archive_results = BooleanDataList(field_name="archive_result")
+    for file_stack in file_list:
+        archive_results.append(archive_one_file(file_stack, **kwargs))
+    return archive_results
+
+
+@simstack_model
+class ArchiveConfig(Model):
+    call_paths: List[str] = Field(default_factory=list)
+    archive_resource: str = Field(default="archive_resource", description="Name of the resource where the files will be archived")
+    in_memory: bool = Field(default=False, description="Whether to archive files in memory")
+    use_time_window: bool = Field(default=False, description="Whether to archive files within a time window")
+    start_date: Optional[datetime] = Field(default=None, description="Start of the period to archive files for")
+    end_date: Optional[datetime] = Field(default=None, description="End end of the period to archive files for")
+    include_patterns: List[str] = Field(default_factory=list, description="List of file patterns to include")
+    exclude_patterns: List[str] = Field(default_factory=list, description="List of file patterns to exclude")
+    min_size: int = Field(default=0, description="Minimum size of files to archive")
+    filter_by_resource: bool = Field(default=False, description="Whether to only archive files that have a location on the current resource")
+
+    model_config = {"collection": "archive_config"}
+
+
+@node
+async def archive_node(archive_config: ArchiveConfig, **kwargs):
+    """
+    Archives FileStacks matching the criteria in archive_config and deletes local instances upon success.
+
+    Args:
+        archive_config (ArchiveConfig): Configuration for the archive process.
+        **kwargs: Additional arguments.
+    """
+    _raise_archive_disabled()
+
+    with open("TEST_EXECUTION.log", "a") as f:
+        f.write(f"archive_node CALLED at {datetime.now()}\n")
+    
+    node_runner = kwargs["node_runner"]
+    db = context.db
+    my_resource = context.config.resource
+    
+    # Discover FileStacks
+    all_file_stacks = []
+    import sys
+    sys.stderr.write(f"DEBUG: Discovery starting for resource {my_resource}\n")
+    sys.stderr.flush()
+    if archive_config.call_paths:
+        from simstack.models import NodeRegistry, FileListModel
+
+        # Find matching NodeRegistry entries
+        node_query = {
+            "call_path": {"$in": archive_config.call_paths},
+            "parameters.resource.value": my_resource
+        }
+        matching_nodes = await db.find(NodeRegistry, node_query)
+        
+        # We also need mappings to know which collection to look into
+        # context.model_mappings has by_mapping and by_name
+        filestack_mapping = context.model_mappings.get_by_name("FileStack")
+        filelist_mapping = context.model_mappings.get_by_name("FileList")
+        filelistmodel_mapping = context.model_mappings.get_by_name("FileListModel")
+        
+        fs_mapping_str = filestack_mapping.mapping if filestack_mapping else ""
+        fl_mapping_str = filelist_mapping.mapping if filelist_mapping else ""
+        flm_mapping_str = filelistmodel_mapping.mapping if filelistmodel_mapping else ""
+
+        found_fs_ids = set()
+        found_fl_ids = set()
+        found_flm_ids = set()
+
+        for node in matching_nodes:
+            # Check info_files
+            for fs in node.info_files:
+                found_fs_ids.add(fs.id)
+            
+            # Check input_references and results_references
+            for ref in node.input_references + node.results_references:
+                if ref.variable_mapping == fs_mapping_str:
+                    found_fs_ids.add(ref.reference)
+                elif ref.variable_mapping == fl_mapping_str:
+                    found_fl_ids.add(ref.reference)
+                elif ref.variable_mapping == flm_mapping_str:
+                    found_flm_ids.add(ref.reference)
+
+        # Retrieve FileStacks from found IDs
+        if found_fs_ids:
+            file_stacks = await db.find(FileStack, {"_id": {"$in": list(found_fs_ids)}})
+            all_file_stacks.extend(file_stacks)
+        
+        # Retrieve FileLists and their FileStacks
+        if found_flm_ids:
+            file_list_models = await db.find(FileListModel, {"_id": {"$in": list(found_flm_ids)}})
+            for flm in file_list_models:
+                # FileListModel uses ObjectListMixin[FileStack], elements are ObjectIds
+                # We need to fetch the actual FileStack objects.
+                if flm.elements:
+                    stacks = await db.find(FileStack, {"_id": {"$in": flm.elements}})
+                    all_file_stacks.extend(stacks)
+
+        # Deduplicate all_file_stacks by ID
+        unique_stacks = {}
+        for fs in all_file_stacks:
+            unique_stacks[fs.id] = fs
+        all_file_stacks = list(unique_stacks.values())
+
+    else:
+        # Existing logic: Query FileStack collection directly
+        query = {}
+        if archive_config.min_size > 0:
+            query["size"] = {"$gte": archive_config.min_size}
+
+        if archive_config.filter_by_resource:
+            query["locations.resource.value"] = my_resource
+
+        all_file_stacks = await db.find(FileStack, query)
+    
+    print(f"DEBUG: Found {len(all_file_stacks)} FileStacks")
+    filtered_file_stacks = []
+
+    for fs in all_file_stacks:
+        # Filter by date if specified
+        # Since FileStack doesn't have a date, we check its locations.
+        if archive_config.start_date or archive_config.end_date:
+            valid_date = False
+            for loc in fs.locations:
+                if loc.created_at:
+                    if archive_config.start_date and loc.created_at < archive_config.start_date:
+                        continue
+                    if archive_config.end_date and loc.created_at > archive_config.end_date:
+                        continue
+                    valid_date = True
+                    break
+            if not valid_date:
+                continue
+
+        # Filter by patterns
+        if archive_config.include_patterns:
+            import fnmatch
+            if not any(fnmatch.fnmatch(fs.name, pat) for pat in archive_config.include_patterns):
+                continue
+
+        if archive_config.exclude_patterns:
+            import fnmatch
+            if any(fnmatch.fnmatch(fs.name, pat) for pat in archive_config.exclude_patterns):
+                continue
+
+        filtered_file_stacks.append(fs)
+
+    if not filtered_file_stacks:
+        node_runner.info("No files found matching the archival criteria.")
+        return BooleanDataList(field_name="archive_results")
+
+    # 2. Add them to FileList
+    archive_file_list = FileList(field_name="archive_files")
+    for fs in filtered_file_stacks:
+        archive_file_list.append(fs)
+
+    print(f"DEBUG: archive_file_list has {len(archive_file_list)} files")
+    # 3. Call archive_files on the filelist
+    # Note: archive_files is a @node, we can call it directly.
+    # It returns a BooleanDataList
+  
+    # To handle mocks and different return types, we ensure we get a BooleanDataList
+    # When calling a @node decorated function, it returns a SimstackResult or Model in sync mode,
+    # or the result of the function if it was run locally.
+    # If we are in the same process and calling it, it might be running via NodeRunner.
+    raw_results = await archive_files(archive_file_list, **kwargs)
+    print(f"DEBUG: raw_results type={type(raw_results)}")
+    if isinstance(raw_results, BooleanDataList):
+        archive_results = raw_results
+    else:
+         # If mocked or returned as list/iterable
+         archive_results = BooleanDataList(field_name="archive_results")
+         from simstack.models import BooleanData
+         if raw_results is not None:
+             # Handle possible async result if not awaited properly by decorator
+             if inspect.isawaitable(raw_results):
+                 raw_results = await raw_results
+             
+             for r in raw_results:
+                 if isinstance(r, bool):
+                     archive_results.append(BooleanData(value=r))
+                 elif isinstance(r, BooleanData):
+                     archive_results.append(r)
+                 else:
+                     # try to save it if it's a model
+                     archive_results.append(r)
+
+    # 4. Delete local FileInstances for successfull archival
+    # archive_results is a BooleanDataList, elements are BooleanData (presumably)
+    # Wait, archive_files returns BooleanDataList. Let's check BooleanDataList implementation.
+    
+    # Actually archive_files returns BooleanDataList where each element corresponds to a FileStack in archive_file_list
+    for i, success_data in enumerate(archive_results):
+        if success_data.value:
+            fs = archive_file_list[i]
+            # Delete local FileInstances
+            # We need to know which ones are "local".
+            # Usually those are on the current resource.
+            current_resource = context.config.resource
+            
+            remaining_locations = []
+            for loc in fs.locations:
+                # If it's on the current resource and is a local path, delete it
+                if loc.resource.value == current_resource and loc.location_type == "local_path":
+                    try:
+                        local_path = Path(context.config.workdir) / loc.path
+                        if local_path.exists():
+                            if local_path.is_file():
+                                local_path.unlink()
+                            elif local_path.is_dir():
+                                shutil.rmtree(local_path)
+                            node_runner.info(f"Deleted local instance of {fs.name} at {local_path}")
+                        else:
+                            node_runner.warning(f"Local path {local_path} for {fs.name} does not exist")
+                    except Exception as e:
+                        node_runner.error(f"Failed to delete local instance of {fs.name}: {e}")
+                        remaining_locations.append(loc)
+                else:
+                    remaining_locations.append(loc)
+            
+            fs.locations = remaining_locations
+            await db.save(fs)
+
+    return archive_results

simstack/methods/fire_and_forget_runner.py

@@ -0,0 +1,113 @@
+import asyncio
+import inspect
+from typing import List, Callable, Any, Optional
+
+from odmantic import Model
+
+from simstack.core.context import context
+from simstack.core.definitions import TaskStatus
+from simstack.core.simstack_result import SimstackResult
+from simstack.core.node_runner import NodeRunner
+from simstack.core.process_results import process_result_helper
+from simstack.models import FireAndForgetResult
+
+class FireAndForgetRunner(NodeRunner):
+    def __init__(self, node: Callable[..., Any], max_concurrency: Optional[int] = None, **kwargs):
+        # We need to handle the case where "node_runner" might be in kwargs, just like MassRunner does.
+        # However, NodeRunner.__init__ expects (name, task_id, logger).
+        # MassRunner does: super().__init__(kwargs["node_runner"].name, kwargs["node_runner"].task_id, kwargs["node_runner"].logger )
+        
+        node_runner = kwargs.get("node_runner")
+        if node_runner:
+            super().__init__(node_runner.name, node_runner.task_id, node_runner.logger)
+            self.call_path = getattr(node_runner, "call_path", "NA")
+        else:
+            # Fallback if not provided, though it's expected in this context
+            super().__init__(kwargs.get("name", "FireAndForgetRunner"), kwargs.get("task_id", "NA"), kwargs.get("logger"))
+            self.call_path = "NA"
+
+        self._node = node
+        self._max_concurrency = max_concurrency
+        self._semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None
+        self._tasks = []
+        self._kwargs = kwargs.copy()
+
+    async def _run_node(self, args: List[Model]):
+        sig = inspect.signature(self._node)
+        param_names = list(sig.parameters.keys())
+
+        node_name = getattr(self._node, "__name__", "unknown_node")
+        full_call_path = f"{self.call_path}/{node_name}"
+
+        models_dict = {}
+        for i, m in enumerate(args):
+            if i < len(param_names):
+                models_dict[f"arg_{param_names[i]}"] = m
+            else:
+                models_dict[f"arg_{i}"] = m
+
+        success = False
+
+        try:
+            if self._semaphore:
+                async with self._semaphore:
+                    if asyncio.iscoroutinefunction(self._node):
+                        result = await self._node(*args, **self._kwargs)
+                    else:
+                        result = self._node(*args, **self._kwargs)
+            else:
+                if asyncio.iscoroutinefunction(self._node):
+                    result = await self._node(*args, **self._kwargs)
+                else:
+                    result = self._node(*args, **self._kwargs)
+            
+            if result is None:
+                success = False
+            elif isinstance(result, bool):
+                success = result
+            elif isinstance(result, (SimstackResult, Model)):
+                references, models = await process_result_helper(result, self.task_id)
+                if isinstance(result, SimstackResult):
+                    success = (result.status == TaskStatus.COMPLETED)
+                else:
+                    success = True
+                
+                if references:
+                    for reference, model in zip(references, models):
+                        models_dict[f"result_{reference.variable_name}"] = model
+            elif isinstance(result, (list, tuple)) and all(isinstance(m, Model) for m in result):
+                success = True
+                for i, m in enumerate(result):
+                    models_dict[f"result_{i}"] = m
+            else:
+                # Other types of results
+                success = True
+
+        except Exception as e:
+            self.error(f"Error running node in FireAndForgetRunner: {e}")
+            success = False
+            models_dict["error"] = str(e)
+
+        # Immediately write to DB
+        ff_result = FireAndForgetResult(
+            call_path=full_call_path,
+            models=models_dict,
+            success=success
+        )
+        
+        db = context.db
+        await db.save(ff_result)
+        return ff_result
+
+    def create_tasks(self, *args: Model):
+        task = asyncio.create_task(self._run_node(list(args)))
+        self._tasks.append(task)
+        return task
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self._tasks:
+            await asyncio.gather(*self._tasks, return_exceptions=True)
+        return self.succeed()