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__
[docs]
@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})"
[docs]
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
[docs]
@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",
},
}
[docs]
@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
[docs]
@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
[docs]
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())
[docs]
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)
[docs]
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)
[docs]
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
[docs]
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})"
[docs]
class FileGetterArgs(Model):
file_stack: FileStack = Reference()
local_resource: Resource
local_dir: Path
[docs]
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())