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

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))

rbx/grading/profiling.py

@@ -0,0 +1,126 @@
+import contextvars
+import functools
+import math
+import threading
+import time
+
+ALL_CONTEXTS_BY_NAME = {}
+_ALL_CONTEXTS_BY_NAME_LOCK = threading.Lock()
+
+
+@functools.cache
+def _get_threadsafe_context(name: str) -> 'Context':
+    with _ALL_CONTEXTS_BY_NAME_LOCK:
+        if name not in ALL_CONTEXTS_BY_NAME:
+            ALL_CONTEXTS_BY_NAME[name] = Context(name)
+        return ALL_CONTEXTS_BY_NAME[name]
+
+
+class Distribution:
+    def __init__(self):
+        self.values = []
+
+    def add(self, value: float):
+        self.values.append(value)
+
+    def mean(self) -> float:
+        return sum(self.values) / len(self.values)
+
+    def median(self) -> float:
+        return sorted(self.values)[len(self.values) // 2]
+
+    def stddev(self) -> float:
+        mean = self.mean()
+        return math.sqrt(sum((x - mean) ** 2 for x in self.values) / len(self.values))
+
+
+class Context:
+    def __init__(self, name: str):
+        self.name = name
+        self._lock = threading.Lock()
+        self.distributions = {}
+        self.counters = {}
+
+    def add_to_distribution(self, name: str, value: float):
+        with self._lock:
+            if name not in self.distributions:
+                self.distributions[name] = Distribution()
+            self.distributions[name].add(value)
+
+    def add_to_counter(self, name: str):
+        with self._lock:
+            if name not in self.counters:
+                self.counters[name] = 0
+            self.counters[name] += 1
+
+    def print_summary(self):
+        with self._lock:
+            print(f'{self.name}:')
+            for name, distribution in sorted(self.distributions.items()):
+                print(f'  ~ {name}: {distribution.mean():.2f}')
+            for name, count in sorted(self.counters.items()):
+                print(f'  + {name}: {count}')
+
+
+profiling_stack_var = contextvars.ContextVar(
+    'profiling_stack', default=[_get_threadsafe_context('root')]
+)
+
+
+def _push_profiling_stack(name: str):
+    return profiling_stack_var.set(
+        profiling_stack_var.get() + [_get_threadsafe_context(name)]
+    )
+
+
+class PushContext:
+    def __init__(self, name: str):
+        self.name = name
+        self.token = None
+
+    def __enter__(self):
+        self.token = _push_profiling_stack(self.name)
+        return profiling_stack_var.get()[-1]
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.token is not None:
+            profiling_stack_var.reset(self.token)
+
+
+def print_summary():
+    print('\n' + ('-') * 3 + '\n')
+    with _ALL_CONTEXTS_BY_NAME_LOCK:
+        for context in ALL_CONTEXTS_BY_NAME.values():
+            context.print_summary()
+
+
+### Public API
+
+
+class Profiler:
+    def __init__(self, name: str, start: bool = False):
+        self.name = name
+        self.start_time = 0
+        if start:
+            self.start()
+
+    def start(self):
+        self.start_time = time.monotonic()
+        return self
+
+    def stop(self):
+        self.end_time = time.monotonic()
+        self.duration = self.end_time - self.start_time
+        for context in profiling_stack_var.get():
+            context.add_to_distribution(self.name, self.duration)
+
+    def __enter__(self):
+        return self.start()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.stop()
+
+
+def add_to_counter(name: str):
+    for context in profiling_stack_var.get():
+        context.add_to_counter(name)

rbx/grading/steps.py

@@ -20,9 +20,10 @@ from rich.text import Text
 from rbx import utils
 from rbx.config import get_bits_stdcpp, get_jngen, get_testlib
 from rbx.console import console
-from rbx.grading import processing_context
+from rbx.grading import grading_context, processing_context
+from rbx.grading.judge.cacher import FileCacher
 from rbx.grading.judge.sandbox import SandboxBase, SandboxParams
-from rbx.grading.judge.storage import Storage, copyfileobj
+from rbx.grading.judge.storage import copyfileobj
 from rbx.grading.limits import Limits
 
 MAX_STDOUT_LEN = 1024 * 1024 * 128  # 128 MB
@@ -139,6 +140,8 @@ class GradingFileInput(BaseModel):
     digest: Optional[DigestHolder] = None
     # Whether the destination file should be marked as an executable.
     executable: bool = False
+    # Whether to track file through its hash (disable for optimization).
+    hash: bool = True
 
 
 class GradingFileOutput(BaseModel):
@@ -161,15 +164,15 @@ class GradingFileOutput(BaseModel):
     # Whether to touch the file before the command runs.
     touch: bool = False
 
-    def get_file(self, storage: Storage) -> Optional[IO[bytes]]:
+    def get_file(self, cacher: FileCacher) -> Optional[IO[bytes]]:
         if self.dest is not None:
             if self.optional and not self.dest.exists():
                 return None
             return self.dest.open('rb')
         if self.digest is not None and self.digest.value is not None:
-            if self.optional and not storage.exists(self.digest.value):
+            if self.optional and not cacher.exists(self.digest.value):
                 return None
-            return storage.get_file(self.digest.value)
+            return cacher.get_file(self.digest.value)
         raise ValueError('No file to get')
 
 
@@ -308,11 +311,14 @@ def _process_input_artifacts(artifacts: GradingArtifacts, sandbox: SandboxBase):
 
 
 def _process_output_artifacts(
-    artifacts: GradingArtifacts, sandbox: SandboxBase
+    artifacts: GradingArtifacts,
+    sandbox: SandboxBase,
 ) -> bool:
     for output_artifact in artifacts.outputs:
         if output_artifact.hash and output_artifact.digest is None:
-            output_artifact.digest = DigestHolder()
+            if not grading_context.is_no_cache():
+                # If cache is enabled, track this file in cache.
+                output_artifact.digest = DigestHolder()
         if not sandbox.file_exists(output_artifact.src):
             if output_artifact.optional:
                 continue
@@ -322,21 +328,43 @@ def _process_output_artifacts(
             return False
 
         if output_artifact.digest is not None:
-            output_artifact.digest.value = sandbox.get_file_to_storage(
-                output_artifact.src, trunc_len=output_artifact.maxlen
-            )
+            # Put it in the cache, possibly compressing it if it's an executable.
+            with grading_context.compression(
+                use_compression=True,
+                when=output_artifact.executable,
+            ):
+                output_artifact.digest.value = sandbox.get_file_to_storage(
+                    output_artifact.src, trunc_len=output_artifact.maxlen
+                )
         if output_artifact.dest is None:
             continue
         dst: pathlib.Path = artifacts.root / output_artifact.dest
         # Ensure dst directory exists.
+
         dst.parent.mkdir(parents=True, exist_ok=True)
-        with dst.open('wb') as f:
-            with sandbox.get_file(output_artifact.src) as sb_f:
-                copyfileobj(
-                    sb_f,
-                    f,
-                    maxlen=output_artifact.maxlen,
+
+        if (
+            output_artifact.digest is not None
+            and output_artifact.digest.value is not None
+            and (
+                path_to_symlink := sandbox.file_cacher.path_for_symlink(
+                    output_artifact.digest.value
                 )
+            )
+            is not None
+        ):
+            # File is in the persistent cache, store a symlink to it.
+            dst.unlink(missing_ok=True)
+            dst.symlink_to(path_to_symlink)
+        else:
+            # File is not in the persistent cache, copy it.
+            with dst.open('wb') as f:
+                with sandbox.get_file(output_artifact.src) as sb_f:
+                    copyfileobj(
+                        sb_f,
+                        f,
+                        maxlen=output_artifact.maxlen,
+                    )
         if output_artifact.executable:
             dst.chmod(0o755)
     return True