simstack 2.0.1__py3-none-any.whl

simstack/models/file_instance.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import logging
+import shutil
+import uuid
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Optional, Union
+
+from odmantic import EmbeddedModel, Field, ObjectId
+from pydantic import model_validator
+
+from simstack.models import simstack_model
+from simstack.models.parameters import Resource
+from simstack.util.file_transfer_client import transfer_resource_name
+from simstack.util.file_hashing import hash_file
+
+logger = logging.getLogger("file_instance")
+
+
+@simstack_model
+class FileInstance(EmbeddedModel):
+    """
+    Represents an embedded model for a file instance.
+
+    The `FileInstance` class is used to encapsulate details about a file,
+    including its path, associated resource, and creation timestamp.
+    It provides a class method for initializing a `FileInstance` object
+    from a local file path, ensuring proper handling of file-related operations.
+
+    Attributes:
+        path (str): Path to the file relative to the host work directory.
+        resource (Resource): Name of the resource associated with the file.
+        created_at (datetime): Timestamp indicating when the file instance was created.
+    """
+
+    id: Optional[str] = Field(default=None, description="File instance id")
+    path: str = Field(
+        description="Path to the file relative to the host work directory"
+    )
+    resource: Resource = Field(description="Resource name")
+    created_at: datetime = Field(description="Creation timestamp")
+    runner_id: Optional[str] = Field(
+        default=None, description="Runner identity that created the instance"
+    )
+    location_type: str = Field(default="local_path", description="File location type")
+    size_bytes: Optional[int] = Field(default=None, description="File size in bytes")
+    checksum_sha256: Optional[str] = Field(default=None, description="SHA256 checksum")
+    last_accessed_at: Optional[datetime] = Field(
+        default=None, description="Last access timestamp"
+    )
+    expires_at: Optional[datetime] = Field(
+        default=None, description="Expiration timestamp"
+    )
+    is_authoritative: bool = Field(
+        default=True, description="Whether this instance is authoritative"
+    )
+    is_cached: bool = Field(
+        default=False,
+        description="Whether this instance is a reusable local cache copy",
+    )
+    status: str = Field(default="available", description="Instance lifecycle status")
+
+    @model_validator(mode="before")
+    def migration(cls, values: Any) -> Any:
+        if isinstance(values.get("resource"), str):
+            values["resource"] = Resource(value=values["resource"])
+        if "path" in values and isinstance(values["path"], Path):
+            values["path"] = str(values["path"])
+        if not values.get("id"):
+            values["id"] = str(uuid.uuid4())
+        if not values.get("location_type"):
+            values["location_type"] = "local_path"
+        if not values.get("status"):
+            values["status"] = "available"
+        if values.get("is_authoritative") is None:
+            values["is_authoritative"] = True
+        if values.get("is_cached") is None:
+            values["is_cached"] = False
+        return values
+
+    @classmethod
+    def from_local_file(
+        cls,
+        path: Union[Path, str],
+        file_stack_id: ObjectId,
+        make_copy: bool = True,
+        tasks_id: str = "",
+    ) -> FileInstance:
+        """
+        Creates a FileInstance object from a local file path.
+
+        This class method is responsible for initializing a `FileInstance` based
+        on a local file's path. It supports options to hash the file, make a
+        user-specific copy, and tracks additional metadata. The method handles
+        local file operations such as copying files to a secure directory when
+        necessary and organizes resources under a configurable working directory.
+
+        :param file_stack_id: the id of the filestack where the file is in
+        :param path: The file path to the local file. Can be either a string
+            or `Path`.
+        :param make_copy: Indicates whether a secure local copy of the file should
+            be made within the application's working directory. Defaults to True.
+        :return: A `FileInstance` object initialized with file details.
+        :rtype: FileInstance
+        :raises ValueError: If there are issues during the creation of the
+            `FileInstance` from the specified local file.
+        """
+        source_path = path if isinstance(path, Path) else Path(path)
+
+        # Prepare the content field if in_memory is True
+        resolved_path = Path(path).resolve()
+        # Find 'simstack' in the path and compute the relative path from its parent
+        from simstack.core.context import context
+
+        workdir = context.config.workdir
+        resolved_workdir = Path(workdir).resolve()
+        if tasks_id == "":
+            logger.debug(f"workdir is {resolved_workdir} path is {resolved_path}")
+        else:
+            logger.debug(
+                f"task_id: {tasks_id} workdir is {resolved_workdir} path is {resolved_path}"
+            )
+        try:
+            relative_path = resolved_path.relative_to(resolved_workdir)
+        except ValueError:
+            if tasks_id == "":
+                logger.error(
+                    f"Path {resolved_path} is not under workdir {resolved_workdir}"
+                )
+            else:
+                logger.error(
+                    f"Path {resolved_path} is not under workdir {resolved_workdir} for task_id: {tasks_id}"
+                )
+            import getpass
+
+            username = getpass.getuser()
+            relative_path = Path(username) / str(file_stack_id)
+            absolute_dir = Path(context.config.workdir) / relative_path
+            absolute_dir.mkdir(parents=True, exist_ok=True)
+            shutil.copy(source_path, absolute_dir)
+
+        file_instance = FileInstance(
+            path=str(relative_path),
+            resource=Resource(value=transfer_resource_name(context.config.resource)),
+            created_at=datetime.now(),
+            size_bytes=source_path.stat().st_size if source_path.exists() else None,
+            checksum_sha256=hash_file(source_path) if source_path.exists() else None,
+        )
+        return file_instance

simstack/models/file_list.py

@@ -0,0 +1,45 @@
+import re
+from typing import List, Optional, Iterator
+
+from odmantic import EmbeddedModel, Field, Model, ObjectId
+
+from simstack.models import simstack_model
+from simstack.util.object_list_mixin import ObjectListMixin
+from simstack.models.files import FileStack
+
+
+@simstack_model
+class FileList(EmbeddedModel, ObjectListMixin[FileStack]):
+    elements: List[ObjectId] = Field(default_factory=list)
+
+    def __init__(self, **data):
+        data, cache = self._normalize_elements_for_init(data)
+        EmbeddedModel.__init__(self, **data)
+        if cache is not None:
+            self._set_cache(cache)
+
+    def __iter__(self) -> Iterator[FileStack]:
+        return ObjectListMixin.__iter__(self)
+
+
+@simstack_model
+class FileListModel(Model, ObjectListMixin[FileStack]):
+    elements: List[ObjectId] = Field(default_factory=list)
+
+    def __init__(self, **data):
+        data, cache = self._normalize_elements_for_init(data)
+        Model.__init__(self, **data)
+        if cache is not None:
+            self._set_cache(cache)
+
+    def __iter__(self) -> Iterator[FileStack]:
+        return ObjectListMixin.__iter__(self)
+
+
+@simstack_model
+class FileListIO(Model):
+    file_list: FileList = Field(default_factory=FileList)
+    task_status: Optional[str] = None
+    error: Optional[str] = None
+    message: Optional[str] = None
+

simstack/models/files.py

@@ -0,0 +1,424 @@
+from __future__ import annotations
+
+import logging
+import os
+import zlib
+from datetime import datetime
+from pathlib import Path
+from typing import List, Optional, Union, Dict, Any
+
+from odmantic import Model, Field, ObjectId, Reference
+
+from simstack.models import simstack_model
+from simstack.models.file_instance import FileInstance
+from simstack.models.parameters import Resource
+from simstack.util.file_transfer_client import (
+    FileTransferClient,
+    first_available_remote_location,
+    path_for_file_instance,
+    transfer_resource_name,
+)
+from simstack.util.file_hashing import hash_file, hash_string
+
+logger = logging.getLogger(__name__)
+
+# MongoDB has a 16MB document size limit
+MONGODB_MAX_DOCUMENT_SIZE = 16 * 1024 * 1024  # 16 MB in bytes
+
+# TODO: its a mystery to me why we should have files with no name or unknown size, both values are actually used in __str__
+
+@simstack_model
+class FileStack(Model):
+    name: Optional[str] = Field(description="Name of the file", default=None)
+    size: Optional[int] = Field(description="Size of the file in bytes", default=None)
+    is_hashable: bool = Field(
+        default=False, description="Whether the file stack is hashable"
+    )
+    hash: Optional[str] = Field(
+        default=None, description="Hash of the file stack content", index=True
+    )
+    in_memory: bool = Field(
+        default=False, description="Whether the file stack is in memory"
+    )
+    content: Optional[bytes] = Field(
+        description="Compressed file content", default=None
+    )
+
+    locations: List[FileInstance] = Field(
+        default_factory=list, description="List of file locations"
+    )
+
+    def __str__(self) -> str:
+        return f"FileStack(name={self.name}, size={self.size}, is_hashable={self.is_hashable}, in_memory={self.in_memory}, locations={self.locations})"
+
+    def __repr__(self) -> str:
+        return f"FileStack(name={self.name}, size={self.size}, is_hashable={self.is_hashable}, in_memory={self.in_memory}, locations={self.locations})"
+
+    async def custom_model_dump(self, **kwargs: Any) -> Dict[str, Any]:
+        dumped_data: Dict[str, Any] = self.model_dump()
+        del dumped_data["content"]  # Exclude content from the dumped data
+        return dumped_data
+
+    @classmethod
+    def ui_base_schema(cls, **kwargs: Any) -> Dict[str, Any]:
+        # TODO get the model programatically
+        return {
+            "ui:field": "FileField",
+            "ui:options": {
+                "model": "simstack.models.files.FileStack",
+            },
+        }
+
+    @classmethod
+    def from_string(cls, data_string: str, file_name: str) -> FileStack:
+        content = zlib.compress(data_string.encode("utf-8"))
+        file_hash = hash_string(data_string)
+        size = len(content)
+
+        if len(content) > MONGODB_MAX_DOCUMENT_SIZE:
+            logger.error(
+                f"Compressed content size {len(content)} bytes exceeds MongoDB limit of {MONGODB_MAX_DOCUMENT_SIZE} bytes for file {file_name}"
+            )
+            raise ValueError(
+                f"Compressed content size {len(content)} bytes exceeds MongoDB document size limit of {MONGODB_MAX_DOCUMENT_SIZE} bytes"
+            )
+
+        file_stack = cls(
+            name=file_name,
+            in_memory=True,
+            content=content,
+            is_hashable=True,
+            size=size,
+            hash=file_hash,
+        )
+        return file_stack
+
+    @classmethod
+    def from_local_file(
+        cls,
+        path: Union[Path, str],
+        is_hashable: bool = True,
+        in_memory: bool = True,
+        secure_source: bool = False,
+        task_id: str = "",
+    ) -> FileStack:
+        """
+        Creates a FileStack object from a local file path.
+
+        :param task_id: task_id of the task that created the file stack, used for logging and tracking
+        :type task_id: str
+        :param secure_source: specifies if the source is secure (already in a directory generated within Simstack II)
+        :type secure_source: bool
+        :param path: The path to the local file. Can be provided as a string or Path object.
+        :type path: Union[Path, str]
+        :param is_hashable: A flag indicating whether the file hash needs to be calculated.
+        :type is_hashable: bool
+        :param in_memory: Whether to store the compressed file content in memory. Defaults to True.
+        :type in_memory: bool
+
+        :return: A FileStack object containing FileInstances for the file.
+        :rtype: FileStack
+        """
+        # TODO make a second version where a unique directory is already there
+        source_path = path if isinstance(path, Path) else Path(path)
+
+        # Check if the source exists
+        if not source_path.exists():
+            logger.error(f"Source file not found: {source_path}")
+            raise FileNotFoundError(f"Source file not found: {source_path}")
+
+        # Check if it's a file (not a directory)
+        if not source_path.is_file():
+            logger.error(f"Source file {path} is not a file")
+            raise ValueError(f"Source file is not a file: {path}")
+
+        # Check read permission using os.access
+        if not os.access(path, os.R_OK):
+            logger.error(f"No read permission for file: {path}")
+            raise ValueError(f"No read permission for file: {path}")
+
+        content = None
+        file_hash = hash_file(source_path) if is_hashable else None
+        name = source_path.name
+        size = source_path.stat().st_size if source_path.exists() else None
+
+        if in_memory:
+            try:
+                # Read the file content
+                with open(source_path, "rb") as f:
+                    file_content = f.read()
+                # Compress the content using zlib
+                content = zlib.compress(file_content)
+                if task_id == "":
+                    logger.debug(
+                        f"Compressed file {source_path} from {len(file_content)} bytes to {len(content)} bytes"
+                    )
+                else:
+                    logger.debug(
+                        f"task_id: {task_id} Compressed file {source_path} from {len(file_content)} bytes to {len(content)} bytes"
+                    )
+                # Check if compressed content exceeds MongoDB document size limit
+                if len(content) > 0.9 * MONGODB_MAX_DOCUMENT_SIZE:
+                    if task_id == "":
+                        logger.error(
+                            f"Compressed content size {len(content)} bytes exceeds MongoDB limit of {MONGODB_MAX_DOCUMENT_SIZE} bytes for file {source_path}"
+                        )
+                        logger.error(
+                            f"Setting in_memory to False for file {source_path} and clear content"
+                        )
+                    else:
+                        logger.error(
+                            f"task_id: {task_id} Compressed content size {len(content)} bytes exceeds MongoDB limit of {MONGODB_MAX_DOCUMENT_SIZE} bytes for file {source_path}"
+                        )
+                        logger.error(
+                            f"task_id: {task_id} Setting in_memory to False for file {source_path} and clear content"
+                        )
+                    in_memory = False
+                    content = None
+            except Exception as e:
+                logger.warning(f"Failed to compress file {source_path}: {e}")
+                raise
+
+        file_stack = cls(
+            name=name,
+            in_memory=in_memory,
+            content=content,
+            is_hashable=is_hashable,
+            size=size,
+            hash=file_hash,
+        )
+
+        if not in_memory:
+            location = FileInstance.from_local_file(
+                path=path,
+                file_stack_id=file_stack.id,
+                make_copy=not secure_source and not in_memory,
+            )
+            file_stack.locations.append(location)
+
+        return file_stack
+
+    def complex_hash(self) -> str:
+        if self.is_hashable:
+            if self.hash:
+                return self.hash
+            else:
+                raise ValueError("FileStack is hashable but hash is not set.")
+            # elif self.in_memory and self.content:
+            #    # If the content is in memory, hash the compressed content
+            #    return complex_hash_function(zlib.decompress(self.content))
+            # else:
+            #    temp_dir = Path(tempfile.mkdtemp())
+            #    local_file = self.get(None, local_dir=temp_dir)
+            #    return complex_hash_function(local_file.read_bytes())
+        else:
+            logger.warning(
+                f"FileStack {self.id} is not hashable, returning unique hash."
+            )
+            return str(ObjectId())
+
+    def append(self, file_instance: FileInstance) -> None:
+        """
+        Appends a FileInstance to the file stack.
+
+        :param file_instance: The FileInstance to append.
+        :type file_instance: FileInstance
+        """
+        self.locations.append(file_instance)
+
+    def get(self, local_dir: Path | None = None) -> Path:
+        """
+        Copies the file stack to a local directory. This is the version to be used in applications
+
+        :param local_dir: The local directory to copy the file stack to.
+        :type local_dir: Path
+        """
+        from simstack.core.context import context
+
+        return self.get_raw(context.config.resource, local_dir)
+
+    def get_raw(self, local_resource: Resource, local_dir: Path | None = None) -> Path:
+        """
+        Copies the file stack to a local directory, assumes no context.
+
+        :param local_resource: the local resource to copy the file stack to. Defaults to the current resource.
+        :param local_dir: The local directory to copy the file stack to.
+        :type local_dir: Path
+        """
+
+        # select the best instance
+        # first search for an instance with "in_memory" set to True
+
+        if local_dir is None:
+            local_dir = Path.cwd()
+
+        if self.in_memory:
+            local_dir.mkdir(parents=True, exist_ok=True)
+            try:
+                if self.content is None:
+                    raise ValueError("FileStack is in-memory but content is missing.")
+                file_name = self.name or "file"
+                # Decompress the content
+                decompressed_content = zlib.decompress(self.content)
+                # Write the decompressed content to the local directory
+                target_path = local_dir / file_name
+                if target_path.exists(): # the user must make sure that the file does not exist, otherwise the file will be overwritten
+                    raise FileExistsError(f"File {target_path} already exists.")
+                with open(target_path, "wb") as f:
+                    f.write(decompressed_content)
+                return target_path
+            except Exception as e:
+                logger.error(f"Failed to decompress and write file {self.name}: {e}")
+                raise ValueError(
+                    f"Failed to decompress and write file {self.name}: {e}"
+                )
+
+        local_resource_name = transfer_resource_name(local_resource)
+        local_transfer_resource = Resource(value=local_resource_name)
+
+        # If in-memory instance not found or decompression failed, try finding instance on same resource
+        same_resource_instance = next(
+            (
+                f
+                for f in self.locations
+                if transfer_resource_name(f.resource) == local_resource_name
+                and getattr(f, "status", "available") == "available"
+            ),
+            None,
+        )
+
+        if same_resource_instance is not None:
+            same_resource_instance.last_accessed_at = datetime.now()
+            # Return the absolute path by joining with the resource's workdir if it's relative
+            path = Path(same_resource_instance.path)
+            if not path.is_absolute():
+                from simstack.core.context import context
+
+                return Path(context.config.workdir) / path
+            logger.info(f"Using existing instance {path} for {self.name}")
+            return path
+
+        remote_instance = first_available_remote_location(
+            self.locations, local_transfer_resource
+        )
+        if remote_instance is None:
+            logger.error("No suitable file instance found for copying.")
+            raise ValueError(
+                f"FileStack {self.id} {self.name} is unavailable: no accessible file instance for resource {local_transfer_resource}."
+            )
+
+        local_dir.mkdir(parents=True, exist_ok=True)
+        return self._get_via_server_transfer(
+            remote_instance, local_transfer_resource, local_dir
+        )
+
+    def _get_via_server_transfer(
+        self,
+        remote_instance: FileInstance,
+        local_resource: Resource,
+        local_dir: Path,
+    ) -> Path:
+        from simstack.core.context import context
+
+        client = FileTransferClient.from_context(required=True)
+        assert client is not None
+
+        transfer = client.create_transfer(
+            file_stack_id=str(self.id),
+            source_file_instance_id=getattr(remote_instance, "id", None),
+            source_resource_name=transfer_resource_name(remote_instance.resource),
+            target_resource_name=transfer_resource_name(local_resource),
+        )
+        transfer_id = str(transfer["transfer_id"])
+        logger.info(
+            "Created FileStack transfer %s for file_stack=%s from resource=%s to resource=%s",
+            transfer_id,
+            self.id,
+            transfer_resource_name(remote_instance.resource),
+            transfer_resource_name(local_resource),
+        )
+
+        try:
+            client.wait_until_uploaded(transfer_id)
+        except Exception as exc:
+            try:
+                client.fail_transfer(
+                    transfer_id,
+                    error_message=str(exc),
+                    error_code="SOURCE_RUNNER_UNAVAILABLE",
+                )
+            except Exception:
+                logger.warning(
+                    "Failed to report failed FileStack transfer %s", transfer_id
+                )
+            raise
+        target_path = local_dir / (self.name or Path(str(remote_instance.path)).name)
+        if target_path.exists():
+            raise FileExistsError(f"File {target_path} already exists.")
+        downloaded = client.download_file(transfer_id, target_path)
+        instance_path = path_for_file_instance(downloaded.path, context.config.workdir)
+        completed = client.complete_transfer(
+            transfer_id=transfer_id,
+            target_resource_name=transfer_resource_name(local_resource),
+            target_path=instance_path,
+            size_bytes=downloaded.size_bytes,
+            checksum_sha256=downloaded.checksum_sha256,
+        )
+
+        file_instance_id = completed.get("file_instance_id")
+        existing = next(
+            (
+                location
+                for location in self.locations
+                if getattr(location, "id", None) == file_instance_id
+            ),
+            None,
+        )
+        if existing is None:
+            self.locations.append(
+                FileInstance(
+                    id=str(file_instance_id) if file_instance_id else None,
+                    path=instance_path,
+                    resource=local_resource,
+                    created_at=datetime.now(),
+                    size_bytes=downloaded.size_bytes,
+                    checksum_sha256=downloaded.checksum_sha256,
+                    location_type="local_path",
+                    is_cached=True,
+                    status="available",
+                )
+            )
+
+        return downloaded.path
+
+    def str(self) -> str:
+        return f"FileStack(name={self.name}, size={self.size}, is_hashable={self.is_hashable}, in_memory={self.in_memory}, locations={self.locations})"
+
+
+class FileGetterArgs(Model):
+    file_stack: FileStack = Reference()
+    local_resource: Resource
+    local_dir: Path
+
+
+async def main() -> None:
+    from simstack.core.context import context
+
+    await context.initialize()
+
+    # write a file test.txt
+    with open("test.txt", "w") as f:
+        f.write("Hello World")
+    file_stack = FileStack.from_local_file("test.txt", is_hashable=True, in_memory=True)
+    print(file_stack)
+    await context.db.save(file_stack)
+    local_dir = Path(context.config.workdir) / "samira" / str(file_stack.id)
+    retrieved = file_stack.get(local_dir=local_dir)
+    print("Retrieved file path:", retrieved)
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    logging.basicConfig(level=logging.DEBUG)
+    asyncio.run(main())

simstack/models/fire_and_forget_result.py

@@ -0,0 +1,19 @@
+from typing import Dict, Any
+from odmantic import Model, Field
+from simstack.models.simstack_model import simstack_model
+
+@simstack_model
+class FireAndForgetResult(Model):
+    """
+    Model representing the result of a fire-and-forget task.
+
+    Attributes:
+        call_path: The full call path of the node.
+        models: A dictionary mapping argument/result names to their values.
+        success: Whether the task was successful.
+        next_step: Whether this result triggers a next step.
+    """
+    call_path: str = Field(description="The full call path of the node")
+    models: Dict[str, Any] = Field(description="A dictionary mapping argument/result names to their values")
+    success: bool = Field(description="Whether the task was successful")
+    next_step: bool = Field(default=False, description="Whether this result triggers a next step")