rbx.cp 0.6.1__py3-none-any.whl → 0.8.0__py3-none-any.whl

rbx/grading/judge/cacher.py

@@ -7,8 +7,11 @@ import pathlib
 import shutil
 import tempfile
 import typing
-from typing import IO, List, Optional
+from typing import IO, Dict, List, Optional, Type
 
+from pydantic import BaseModel
+
+from rbx.grading import grading_context
 from rbx.grading.judge import digester, storage
 
 logger = logging.getLogger(__name__)
@@ -60,6 +63,7 @@ class FileCacher:
         self.backend = backend
         self.shared = shared
         self.folder = folder
+        self.existing = set()
 
         # First we create the config directories.
         if folder:
@@ -146,6 +150,11 @@ class FileCacher:
 
         logger.debug('File %s not in cache, downloading ' 'from database.', digest)
 
+        if (symlink := self.backend.path_for_symlink(digest)) is not None:
+            cache_file_path.unlink(missing_ok=True)
+            cache_file_path.symlink_to(symlink)
+            return cache_file_path.open('rb') if not cache_only else None
+
         ftmp_handle, temp_file_path = tempfile.mkstemp(dir=self.temp_dir, text=False)
         temp_file_path = pathlib.Path(temp_file_path)
         with open(ftmp_handle, 'wb') as ftmp, self.backend.get_file(digest) as fobj:
@@ -168,6 +177,22 @@ class FileCacher:
         if not cache_only:
             return fd
 
+    def exists(self, digest: str, cache_only: bool = False) -> bool:
+        """Check if a file exists in the cacher.
+
+        cache_only (bool): don't check the backend.
+
+        """
+        cache_file_path = self.file_dir / digest
+        if cache_file_path.exists() or digest in self.existing:
+            return True
+        if cache_only:
+            return False
+        exists = self.backend.exists(digest)
+        if exists:
+            self.existing.add(digest)
+        return exists
+
     def cache_file(self, digest: str):
         """Load a file into the cache.
 
@@ -219,9 +244,18 @@ class FileCacher:
         if digest == storage.TOMBSTONE:
             raise TombstoneError()
 
+        if grading_context.is_transient():
+            return None
+
         logger.debug('Getting symlink file path %s.', digest)
         return self.backend.path_for_symlink(digest)
 
+    def digest_from_symlink(self, link: pathlib.Path) -> Optional[str]:
+        if grading_context.is_transient():
+            return None
+
+        return self.backend.filename_from_symlink(link)
+
     def get_file_content(self, digest: str) -> bytes:
         """Retrieve a file from the storage.
 
@@ -280,7 +314,9 @@ class FileCacher:
             with dst_path.open('wb') as dst:
                 storage.copyfileobj(src, dst, self.CHUNK_SIZE)
 
-    def put_file_from_fobj(self, src: IO[bytes], desc: str = '') -> str:
+    def put_file_from_fobj(
+        self, src: IO[bytes], metadata: Optional[Dict[str, BaseModel]] = None
+    ) -> str:
         """Store a file in the storage.
 
         If it's already (for some reason...) in the cache send that
@@ -292,7 +328,7 @@ class FileCacher:
 
         src (fileobj): a readable binary file-like object from which
             to read the contents of the file.
-        desc (unicode): the (optional) description to associate to the
+        metadata (Dict[str, BaseModel]): the (optional) metadata to associate to the
             file.
 
         return (unicode): the digest of the stored file.
@@ -334,36 +370,45 @@ class FileCacher:
             # We read from the temporary file before moving it to
             # cache_file_path because the latter might be deleted before
             # we get a chance to open it.
-            with open(dst.name, 'rb') as src:
-                pending_file = self.backend.create_file(digest)
-                if pending_file is not None:
-                    storage.copyfileobj(src, pending_file.fd, self.CHUNK_SIZE)
-                    self.backend.commit_file(pending_file, desc)
+            #
+            # Only store file when not in transient mode.
+            if not grading_context.is_transient():
+                with open(dst.name, 'rb') as src:
+                    pending_file = self.backend.create_file(digest)
+                    if pending_file is not None:
+                        storage.copyfileobj(src, pending_file.fd, self.CHUNK_SIZE)
+                        self.backend.commit_file(pending_file, metadata)
 
             os.rename(dst.name, cache_file_path)
 
         return digest
 
-    def put_file_content(self, content: bytes, desc: str = '') -> str:
+    def put_file_content(
+        self, content: bytes, metadata: Optional[Dict[str, BaseModel]] = None
+    ) -> str:
         """Store a file in the storage.
 
         See `put_file_from_fobj'. This method will read the content of
         the file from the given binary string.
 
         content (bytes): the content of the file to store.
-        desc (unicode): the (optional) description to associate to the
+        metadata (Dict[str, BaseModel]): the (optional) metadata to associate to the
             file.
 
         return (unicode): the digest of the stored file.
 
         """
         with io.BytesIO(content) as src:
-            return self.put_file_from_fobj(src, desc)
+            return self.put_file_from_fobj(src, metadata)
 
-    def put_file_text(self, text: str, desc: str = '') -> str:
-        return self.put_file_content(text.encode('utf-8'), desc)
+    def put_file_text(
+        self, text: str, metadata: Optional[Dict[str, BaseModel]] = None
+    ) -> str:
+        return self.put_file_content(text.encode('utf-8'), metadata)
 
-    def put_file_from_path(self, src_path: pathlib.Path, desc: str = '') -> str:
+    def put_file_from_path(
+        self, src_path: pathlib.Path, metadata: Optional[Dict[str, BaseModel]] = None
+    ) -> str:
         """Store a file in the storage.
 
         See `put_file_from_fobj'. This method will read the content of
@@ -371,28 +416,53 @@ class FileCacher:
 
         src_path (Path): an accessible location on the file-system
             from which to read the contents of the file.
-        desc (unicode): the (optional) description to associate to the
+        metadata (Dict[str, BaseModel]): the (optional) metadata to associate to the
             file.
 
         return (unicode): the digest of the stored file.
 
         """
         with src_path.open('rb') as src:
-            return self.put_file_from_fobj(src, desc)
+            return self.put_file_from_fobj(src, metadata)
+
+    def set_metadata(self, digest: str, key: str, value: Optional[BaseModel]):
+        """Set the description of a file given its digest.
+
+        digest (unicode): the digest of the file to add the description.
+        key (str): the key of the metadata to add.
+        value (BaseModel): the value of the metadata to add.
+        """
+        if grading_context.is_transient():
+            return
+        self.backend.set_metadata(digest, key, value)
 
-    def describe(self, digest: str) -> str:
+    def get_metadata(
+        self, digest: str, key: str, model_cls: Type[storage.BaseModelT]
+    ) -> Optional[storage.BaseModelT]:
         """Return the description of a file given its digest.
 
         digest (unicode): the digest of the file to describe.
-
-        return (unicode): the description of the file.
+        key (str): the key of the metadata to get.
+        model_cls (Type[storage.BaseModelT]): the model class of the metadata.
+        return (BaseModel): the metadata of the file.
 
         raise (KeyError): if the file cannot be found.
 
         """
         if digest == storage.TOMBSTONE:
             raise TombstoneError()
-        return self.backend.describe(digest)
+        return typing.cast(
+            Optional[storage.BaseModelT],
+            self.backend.get_metadata(digest, key, model_cls),
+        )
+
+    def list_metadata(self, filename: str) -> List[str]:
+        """List the metadata of a file given its filename.
+
+        filename (str): the filename of the file to list the metadata.
+        return (List[str]): the list of metadata keys.
+        """
+        return self.backend.list_metadata(filename)
 
     def get_size(self, digest: str) -> int:
         """Return the size of a file given its digest.
@@ -431,6 +501,7 @@ class FileCacher:
             return
         cache_file_path: pathlib.Path = self.file_dir / digest
         cache_file_path.unlink(missing_ok=True)
+        self.existing.discard(digest)
 
     def purge_cache(self):
         """Empty the local cache.
@@ -442,6 +513,7 @@ class FileCacher:
         self.file_dir.mkdir(parents=True, exist_ok=True)
         if self.folder is not None:
             self.folder.mkdir(parents=True, exist_ok=True)
+        self.existing.clear()
 
     def destroy_cache(self):
         """Completely remove and destroy the cache.
@@ -456,7 +528,7 @@ class FileCacher:
             raise Exception('You may not destroy a shared cache.')
         shutil.rmtree(str(self.file_dir))
 
-    def list(self) -> List[storage.FileWithDescription]:
+    def list(self) -> List[storage.FileWithMetadata]:
         """List the files available in the storage.
 
         return ([(unicode, unicode)]): a list of pairs, each

rbx/grading/judge/sandbox.py

@@ -647,12 +647,15 @@ class SandboxBase(abc.ABC):
         return self.get_file_to_bytes(path, maxlen).decode('utf-8')
 
     def get_file_to_storage(
-        self, path: pathlib.Path, description: str = '', trunc_len: Optional[int] = None
+        self,
+        path: pathlib.Path,
+        metadata: Optional[Dict[str, pydantic.BaseModel]] = None,
+        trunc_len: Optional[int] = None,
     ) -> str:
         """Put a sandbox file in FS and return its digest.
 
         path (Path): relative path of the file inside the sandbox.
-        description (str): the description for FS.
+        metadata (Dict[str, pydantic.BaseModel]): the metadata for FS.
         trunc_len (int|None): if None, does nothing; otherwise, before
             returning truncate it at the specified length.
 
@@ -660,7 +663,7 @@ class SandboxBase(abc.ABC):
 
         """
         with self.get_file(path, trunc_len=trunc_len) as file_:
-            return self.file_cacher.put_file_from_fobj(file_, description)
+            return self.file_cacher.put_file_from_fobj(file_, metadata)
 
     def stat_file(self, path: pathlib.Path) -> os.stat_result:
         """Return the stats of a file in the sandbox.

rbx/grading/judge/sandboxes/timeit.py

@@ -100,9 +100,9 @@ def create_tee(files, mode, buffer_size=4096, prefix=''):
                         tee.file.write(tee.prefix)
                     tee.file.write(bytes)
                     tee.file.flush()
-                    new = bytes == b'\n'
                     # TODO maybe add in fsync() here if the fileno() method
                     # exists on file
+                new = bytes == b'\n'
         except Exception:
             pass
         finally:

rbx/grading/judge/storage.py

@@ -3,13 +3,21 @@ import io
 import logging
 import pathlib
 import tempfile
+import typing
 from abc import ABC, abstractmethod
-from typing import IO, AnyStr, List, Optional
+from typing import IO, AnyStr, Dict, List, Optional, Type, TypeVar
+
+import lz4.frame
+from pydantic import BaseModel
+
+from rbx.grading import grading_context
 
 logger = logging.getLogger(__name__)
 
 TOMBSTONE = 'x'
 
+BaseModelT = TypeVar('BaseModelT', bound=BaseModel)
+
 
 def copyfileobj(
     source_fobj: IO[AnyStr],
@@ -43,16 +51,24 @@ def copyfileobj(
             maxlen -= written
 
 
+COMPRESSION_LEVEL = 5
+
+
+class CompressionMetadata(BaseModel):
+    compression_level: int
+
+
 @dataclasses.dataclass
 class PendingFile:
     fd: IO[bytes]
     filename: str
+    metadata: Dict[str, Optional[BaseModel]] = dataclasses.field(default_factory=dict)
 
 
 @dataclasses.dataclass
-class FileWithDescription:
+class FileWithMetadata:
     filename: str
-    description: str
+    metadata: List[str]
 
 
 class Storage(ABC):
@@ -81,13 +97,15 @@ class Storage(ABC):
         pass
 
     @abstractmethod
-    def commit_file(self, file: PendingFile, desc: str = '') -> bool:
+    def commit_file(
+        self, file: PendingFile, metadata: Optional[Dict[str, BaseModel]] = None
+    ) -> bool:
         """Commit a file created by create_file() to be stored.
         Given a file object returned by create_file(), this function populates
         the database to record that this file now legitimately exists and can
         be used.
-        fobj (fileobj): the object returned by create_file()
         file (PendingFile): the file to commit.
+        metadata (Dict[str, BaseModel]): the metadata of the file.
         return (bool): True if the file was committed successfully, False if
             there was already a file with the same filename in the database. This
             shouldn't make any difference to the caller, except for testing
@@ -96,19 +114,40 @@ class Storage(ABC):
         pass
 
     @abstractmethod
-    def exists(self, filename: str) -> bool:
-        """Check if a file exists in the storage."""
+    def set_metadata(self, filename: str, key: str, value: Optional[BaseModel]):
+        """Set the metadata of a file given its filename.
+        filename (unicode): the filename of the file to set the metadata.
+        key (unicode): the key of the metadata to set.
+        value (BaseModel): the value of the metadata to set.
+        """
         pass
 
     @abstractmethod
-    def describe(self, filename: str) -> str:
-        """Return the description of a file given its filename.
-        filename (unicode): the filename of the file to describe.
-        return (unicode): the description of the file.
+    def get_metadata(
+        self, filename: str, key: str, model_cls: Type[BaseModel]
+    ) -> Optional[BaseModel]:
+        """Get the metadata of a file given its filename and key.
+        filename (unicode): the filename of the file to get the metadata.
+        key (unicode): the key of the metadata to get.
+        model_cls (Type[BaseModel]): the model class of the metadata.
+        return (BaseModel): the value of the metadata.
         raise (KeyError): if the file cannot be found.
         """
         pass
 
+    @abstractmethod
+    def list_metadata(self, filename: str) -> List[str]:
+        """List the metadata of a file given its filename.
+        filename (unicode): the filename of the file to list the metadata.
+        return (List[str]): the list of metadata keys.
+        """
+        pass
+
+    @abstractmethod
+    def exists(self, filename: str) -> bool:
+        """Check if a file exists in the storage."""
+        pass
+
     @abstractmethod
     def get_size(self, filename: str) -> int:
         """Return the size of a file given its filename.
@@ -127,7 +166,7 @@ class Storage(ABC):
         pass
 
     @abstractmethod
-    def list(self) -> List[FileWithDescription]:
+    def list(self) -> List[FileWithMetadata]:
         """List the files available in the storage.
         return ([(unicode, unicode)]): a list of pairs, each
             representing a file in the form (filename, description).
@@ -138,6 +177,10 @@ class Storage(ABC):
     def path_for_symlink(self, filename: str) -> Optional[pathlib.Path]:
         pass
 
+    @abstractmethod
+    def filename_from_symlink(self, link: pathlib.Path) -> Optional[str]:
+        pass
+
 
 class NullStorage(Storage):
     """This backend is always empty, it just drops each file that
@@ -153,41 +196,54 @@ class NullStorage(Storage):
     def create_file(self, digest: str) -> Optional[PendingFile]:
         return None
 
-    def commit_file(self, file: PendingFile, desc: str = '') -> bool:
+    def commit_file(
+        self, file: PendingFile, metadata: Optional[Dict[str, BaseModel]] = None
+    ) -> bool:
         return False
 
-    def exists(self, filename: str) -> bool:
-        return False
+    def set_metadata(self, filename: str, key: str, value: Optional[BaseModel]):
+        pass
 
-    def describe(self, digest: str) -> str:
+    def get_metadata(
+        self, filename: str, key: str, model_cls: Type[BaseModel]
+    ) -> Optional[BaseModel]:
         raise KeyError('File not found.')
 
+    def list_metadata(self, filename: str) -> List[str]:
+        return []
+
+    def exists(self, filename: str) -> bool:
+        return False
+
     def get_size(self, digest: str) -> int:
         raise KeyError('File not found.')
 
     def delete(self, digest: str):
         pass
 
-    def list(self) -> List[FileWithDescription]:
+    def list(self) -> List[FileWithMetadata]:
         return list()
 
     def path_for_symlink(self, digest: str) -> Optional[pathlib.Path]:
         return None
 
+    def filename_from_symlink(self, link: pathlib.Path) -> Optional[str]:
+        return None
+
 
 class FilesystemStorage(Storage):
     """This class implements a backend for FileCacher that keeps all
     the files in a file system directory, named after their filename.
     """
 
-    def __init__(self, path: pathlib.Path):
+    def __init__(self, path: pathlib.Path, compress: bool = False):
         """Initialize the backend.
         path (string): the base path for the storage.
         """
         self.path = path
-
+        self.compress = compress
         # Create the directory if it doesn't exist
-        path.mkdir(parents=True, exist_ok=True)
+        (path / '.metadata').mkdir(parents=True, exist_ok=True)
 
     def get_file(self, filename: str) -> IO[bytes]:
         """See FileCacherBackend.get_file()."""
@@ -196,6 +252,18 @@ class FilesystemStorage(Storage):
         if not file_path.is_file():
             raise KeyError('File not found.')
 
+        compression_metadata = self.get_metadata(
+            filename, 'compression', CompressionMetadata
+        )
+        if compression_metadata is not None:
+            return typing.cast(
+                IO[bytes],
+                lz4.frame.open(
+                    file_path,
+                    mode='rb',
+                    compression_level=compression_metadata.compression_level,
+                ),
+            )
         return file_path.open('rb')
 
     def create_file(self, filename: str) -> Optional[PendingFile]:
@@ -211,13 +279,39 @@ class FilesystemStorage(Storage):
         temp_file = tempfile.NamedTemporaryFile(
             'wb', delete=False, prefix='.tmp.', suffix=filename, dir=self.path
         )
-        return PendingFile(fd=temp_file, filename=filename)
-
-    def commit_file(self, file: PendingFile, desc: str = '') -> bool:
+        metadata: Dict[str, Optional[BaseModel]] = {'compression': None}
+        if self.compress or grading_context.should_compress():
+            fd_name = temp_file.name
+            level = grading_context.get_compression_level()
+            temp_file = typing.cast(
+                IO[bytes],
+                lz4.frame.open(
+                    temp_file,
+                    mode='wb',
+                    compression_level=level,
+                ),
+            )
+            temp_file.name = fd_name  # type: ignore
+            metadata['compression'] = CompressionMetadata(compression_level=level)
+
+        return PendingFile(fd=temp_file, filename=filename, metadata=metadata)
+
+    def commit_file(
+        self, file: PendingFile, metadata: Optional[Dict[str, BaseModel]] = None
+    ) -> bool:
         """See FileCacherBackend.commit_file()."""
         file.fd.close()
 
         file_path: pathlib.Path = self.path / file.filename
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        for key, value in file.metadata.items():
+            self._set_metadata(file.filename, key, value)
+
+        if metadata is not None:
+            for key, value in metadata.items():
+                self._set_metadata(file.filename, key, value)
+
         # Move it into place in the cache. Skip if it already exists, and
         # delete the temporary file instead.
         if not file_path.is_file():
@@ -231,21 +325,43 @@ class FilesystemStorage(Storage):
             pathlib.PosixPath(file.fd.name).unlink()
             return False
 
+    def _get_metadata_path(self, filename: str, key: str) -> pathlib.Path:
+        return self.path / '.metadata' / f'{filename}__{key}.json'
+
+    def _set_metadata(self, filename: str, key: str, value: Optional[BaseModel]):
+        if value is None:
+            self._get_metadata_path(filename, key).unlink(missing_ok=True)
+        else:
+            metadata_path = self._get_metadata_path(filename, key)
+            metadata_path.parent.mkdir(parents=True, exist_ok=True)
+            metadata_path.write_text(value.model_dump_json())
+
+    def set_metadata(self, filename: str, key: str, value: Optional[BaseModel]):
+        if not self.exists(filename):
+            raise KeyError('File not found.')
+
+        self._set_metadata(filename, key, value)
+
+    def get_metadata(
+        self, filename: str, key: str, model_cls: Type[BaseModelT]
+    ) -> Optional[BaseModelT]:
+        path = self._get_metadata_path(filename, key)
+        if not path.is_file():
+            return None
+        return model_cls.model_validate_json(path.read_text())
+
+    def list_metadata(self, filename: str) -> List[str]:
+        return [
+            path.stem.split('__')[1]
+            for path in (self.path / '.metadata').glob(f'{filename}__*.json')
+        ]
+
     def exists(self, filename: str) -> bool:
         """See FileCacherBackend.exists()."""
         file_path: pathlib.Path = self.path / filename
 
         return file_path.is_file()
 
-    def describe(self, filename: str) -> str:
-        """See FileCacherBackend.describe()."""
-        file_path: pathlib.Path = self.path / filename
-
-        if not file_path.is_file():
-            raise KeyError('File not found.')
-
-        return ''
-
     def get_size(self, filename: str) -> int:
         """See FileCacherBackend.get_size()."""
         file_path: pathlib.Path = self.path / filename
@@ -260,15 +376,19 @@ class FilesystemStorage(Storage):
         file_path: pathlib.Path = self.path / filename
 
         file_path.unlink(missing_ok=True)
+        for key in self.list_metadata(filename):
+            self._get_metadata_path(filename, key).unlink(missing_ok=True)
 
-    def list(self) -> List[FileWithDescription]:
+    def list(self) -> List[FileWithMetadata]:
         """See FileCacherBackend.list()."""
         res = []
         for path in self.path.glob('*'):
             if path.is_file():
+                filename = str(path.relative_to(self.path))
                 res.append(
-                    FileWithDescription(
-                        filename=str(path.relative_to(self.path)), description=''
+                    FileWithMetadata(
+                        filename=filename,
+                        metadata=self.list_metadata(filename),
                     )
                 )
         return res
@@ -277,4 +397,18 @@ class FilesystemStorage(Storage):
         file_path = self.path / filename
         if not file_path.is_file():
             raise KeyError('File not found.')
+
+        compression_metadata = self.get_metadata(
+            filename, 'compression', CompressionMetadata
+        )
+        if compression_metadata is not None:
+            return None
         return file_path
+
+    def filename_from_symlink(self, link: pathlib.Path) -> Optional[str]:
+        if not link.is_symlink():
+            return None
+        filename = link.readlink().resolve()
+        if not filename.is_file():
+            return None
+        return str(filename.relative_to(self.path))