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

rbx/grading/caching.py

@@ -3,13 +3,18 @@ import io
 import os
 import pathlib
 import shelve
+import shutil
+import tempfile
 from typing import Any, Dict, List, Optional
 
 from pydantic import BaseModel
 
 from rbx import console
+from rbx.grading import grading_context
+from rbx.grading.judge.cacher import FileCacher
 from rbx.grading.judge.digester import digest_cooperatively
-from rbx.grading.judge.storage import Storage, copyfileobj
+from rbx.grading.judge.storage import copyfileobj
+from rbx.grading.profiling import Profiler
 from rbx.grading.steps import DigestHolder, GradingArtifacts, GradingLogsHolder
 
 VERBOSE = False
@@ -89,11 +94,15 @@ def _build_digest_list(artifacts_list: List[GradingArtifacts]) -> List[DigestHol
     return digests
 
 
-def _build_fingerprint_list(artifacts_list: List[GradingArtifacts]) -> List[str]:
+def _build_fingerprint_list(
+    artifacts_list: List[GradingArtifacts], cacher: FileCacher
+) -> List[str]:
     fingerprints = []
     for artifacts in artifacts_list:
         for input in artifacts.inputs:
-            if input.src is None:
+            if input.src is None or not input.hash:
+                continue
+            if cacher.digest_from_symlink(input.src) is not None:
                 continue
             with input.src.open('rb') as f:
                 fingerprints.append(digest_cooperatively(f))
@@ -124,9 +133,10 @@ def _build_logs_list(artifacts_list: List[GradingArtifacts]) -> List[GradingLogs
 
 def _build_cache_fingerprint(
     artifacts_list: List[GradingArtifacts],
+    cacher: FileCacher,
 ) -> CacheFingerprint:
     digests = [digest.value for digest in _build_digest_list(artifacts_list)]
-    fingerprints = _build_fingerprint_list(artifacts_list)
+    fingerprints = _build_fingerprint_list(artifacts_list, cacher)
     output_fingerprints = _build_output_fingerprint_list(artifacts_list)
     logs = _build_logs_list(artifacts_list)
     return CacheFingerprint(
@@ -155,6 +165,7 @@ def _build_cache_input(
     commands: List[str],
     artifact_list: List[GradingArtifacts],
     extra_params: Dict[str, Any],
+    cacher: FileCacher,
 ) -> CacheInput:
     cloned_artifact_list = [
         artifacts.model_copy(deep=True) for artifacts in artifact_list
@@ -164,6 +175,15 @@ def _build_cache_input(
         # part of the cache key.
         artifacts.logs = None
 
+        for input in artifacts.inputs:
+            if input.src is None:
+                continue
+            inferred_digest = cacher.digest_from_symlink(input.src)
+            if inferred_digest is not None:
+                # Consume cache from digest instead of file.
+                input.digest = DigestHolder(value=inferred_digest)
+                input.src = None
+
         for output in artifacts.outputs:
             if output.hash:
                 # Cleanup dest field from hash artifacts
@@ -185,7 +205,7 @@ def _build_cache_key(input: CacheInput) -> str:
         return digest_cooperatively(fobj)
 
 
-def _copy_hashed_files(artifact_list: List[GradingArtifacts], storage: Storage):
+def _copy_hashed_files(artifact_list: List[GradingArtifacts], cacher: FileCacher):
     for artifact in artifact_list:
         for output in artifact.outputs:
             if not output.hash or output.dest is None:
@@ -194,19 +214,27 @@ def _copy_hashed_files(artifact_list: List[GradingArtifacts], storage: Storage):
             if output.optional and output.digest.value is None:
                 continue
             assert output.digest.value is not None
-            with storage.get_file(output.digest.value) as fobj:
-                with output.dest.open('wb') as f:
-                    copyfileobj(fobj, f, maxlen=output.maxlen)
+            if (
+                path_to_symlink := cacher.path_for_symlink(output.digest.value)
+            ) is not None:
+                # Use a symlink to the file in the persistent cache, if available.
+                output.dest.unlink(missing_ok=True)
+                output.dest.symlink_to(path_to_symlink)
+            else:
+                # Otherwise, copy it.
+                with cacher.get_file(output.digest.value) as fobj:
+                    with output.dest.open('wb') as f:
+                        copyfileobj(fobj, f, maxlen=output.maxlen)
             if output.executable:
                 output.dest.chmod(0o755)
 
 
-def is_artifact_ok(artifact: GradingArtifacts, storage: Storage) -> bool:
+def is_artifact_ok(artifact: GradingArtifacts, cacher: FileCacher) -> bool:
     for output in artifact.outputs:
         if output.optional or output.intermediate:
             continue
         if output.digest is not None:
-            if output.digest.value is None or not storage.exists(output.digest.value):
+            if output.digest.value is None or not cacher.exists(output.digest.value):
                 return False
             return True
         assert output.dest is not None
@@ -219,9 +247,9 @@ def is_artifact_ok(artifact: GradingArtifacts, storage: Storage) -> bool:
     return True
 
 
-def are_artifacts_ok(artifacts: List[GradingArtifacts], storage: Storage) -> bool:
+def are_artifacts_ok(artifacts: List[GradingArtifacts], cacher: FileCacher) -> bool:
     for artifact in artifacts:
-        if not is_artifact_ok(artifact, storage):
+        if not is_artifact_ok(artifact, cacher):
             return False
     return True
 
@@ -244,53 +272,70 @@ class DependencyCacheBlock:
         self._key = None
 
     def __enter__(self):
-        input = _build_cache_input(
-            commands=self.commands,
-            artifact_list=self.artifact_list,
-            extra_params=self.extra_params,
-        )
-        if VERBOSE:
-            console.console.log(f'Cache input is: {input}')
-        self._key = _build_cache_key(input)
-        if VERBOSE:
-            console.console.log(f'Cache key is: {self._key}')
-        found = self.cache.find_in_cache(
-            self.commands, self.artifact_list, self.extra_params, key=self._key
-        )
-        return found
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        if exc_type is None:
-            self.cache.store_in_cache(
+        with Profiler('enter_in_cache'):
+            if grading_context.is_no_cache():
+                return False
+            input = _build_cache_input(
+                commands=self.commands,
+                artifact_list=self.artifact_list,
+                extra_params=self.extra_params,
+                cacher=self.cache.cacher,
+            )
+            if VERBOSE:
+                console.console.log(f'Cache input is: {input}')
+            self._key = _build_cache_key(input)
+            if VERBOSE:
+                console.console.log(f'Cache key is: {self._key}')
+            found = self.cache.find_in_cache(
                 self.commands, self.artifact_list, self.extra_params, key=self._key
             )
-        if exc_type is NoCacheException:
-            return True
-        return None
+            return found
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        with Profiler('exit_in_cache'):
+            if grading_context.is_no_cache():
+                return True if exc_type is NoCacheException else None
+            if exc_type is None:
+                self.cache.store_in_cache(
+                    self.commands, self.artifact_list, self.extra_params, key=self._key
+                )
+            if exc_type is NoCacheException:
+                return True
+            return None
 
 
 class DependencyCache:
     root: pathlib.Path
-    storage: Storage
+    cacher: FileCacher
 
-    def __init__(self, root: pathlib.Path, storage: Storage):
+    def __init__(self, root: pathlib.Path, cacher: FileCacher):
         self.root = root
-        self.storage = storage
+        self.cacher = cacher
         self.db = shelve.open(self._cache_name())
+        tmp_dir = pathlib.Path(tempfile.mkdtemp())
+        self.transient_db = shelve.open(tmp_dir / '.cache_db')
         atexit.register(lambda: self.db.close())
+        atexit.register(lambda: self.transient_db.close())
+        atexit.register(lambda: shutil.rmtree(tmp_dir))
 
     def _cache_name(self) -> str:
         return str(self.root / '.cache_db')
 
+    def get_db(self) -> shelve.Shelf:
+        if grading_context.is_transient():
+            return self.transient_db
+        return self.db
+
     def _find_in_cache(self, key: str) -> Optional[CacheFingerprint]:
-        return self.db.get(key)
+        return self.get_db().get(key)
 
     def _store_in_cache(self, key: str, fingerprint: CacheFingerprint):
-        self.db[key] = fingerprint
+        self.get_db()[key] = fingerprint
 
     def _evict_from_cache(self, key: str):
-        if key in self.db:
-            del self.db[key]
+        db = self.get_db()
+        if key in db:
+            del db[key]
 
     def __call__(
         self,
@@ -309,7 +354,10 @@ class DependencyCache:
         key: Optional[str] = None,
     ) -> bool:
         input = _build_cache_input(
-            commands=commands, artifact_list=artifact_list, extra_params=extra_params
+            commands=commands,
+            artifact_list=artifact_list,
+            extra_params=extra_params,
+            cacher=self.cacher,
         )
         key = key or _build_cache_key(input)
 
@@ -317,7 +365,7 @@ class DependencyCache:
         if fingerprint is None:
             return False
 
-        reference_fingerprint = _build_cache_fingerprint(artifact_list)
+        reference_fingerprint = _build_cache_fingerprint(artifact_list, self.cacher)
 
         if not _fingerprints_match(fingerprint, reference_fingerprint):
             self._evict_from_cache(key)
@@ -334,7 +382,7 @@ class DependencyCache:
         for digest, reference_digest in zip(fingerprint.digests, reference_digests):
             reference_digest.value = digest
 
-        if not are_artifacts_ok(artifact_list, self.storage):
+        if not are_artifacts_ok(artifact_list, self.cacher):
             # Rollback digest changes.
             for old_digest_value, reference_digest in zip(
                 old_digest_values, reference_digests
@@ -344,7 +392,7 @@ class DependencyCache:
             return False
 
         # Copy hashed files to file system.
-        _copy_hashed_files(artifact_list, self.storage)
+        _copy_hashed_files(artifact_list, self.cacher)
 
         # Apply logs changes.
         for logs, reference_logs in zip(fingerprint.logs, reference_fingerprint.logs):
@@ -366,11 +414,14 @@ class DependencyCache:
         key: Optional[str] = None,
     ):
         input = _build_cache_input(
-            commands=commands, artifact_list=artifact_list, extra_params=extra_params
+            commands=commands,
+            artifact_list=artifact_list,
+            extra_params=extra_params,
+            cacher=self.cacher,
         )
         key = key or _build_cache_key(input)
 
-        if not are_artifacts_ok(artifact_list, self.storage):
+        if not are_artifacts_ok(artifact_list, self.cacher):
             return
 
-        self._store_in_cache(key, _build_cache_fingerprint(artifact_list))
+        self._store_in_cache(key, _build_cache_fingerprint(artifact_list, self.cacher))

rbx/grading/debug_context.py

@@ -0,0 +1,31 @@
+import contextvars
+import dataclasses
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class DebugContext:
+    enable: bool = False
+
+
+debug_var = contextvars.ContextVar('debug', default=DebugContext())
+
+
+def get_debug_context() -> DebugContext:
+    return debug_var.get()
+
+
+class Debug:
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+        self.token = None
+
+    def __enter__(self):
+        self.token = debug_var.set(
+            dataclasses.replace(debug_var.get(), *self.args, **self.kwargs)
+        )
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.token is not None:
+            debug_var.reset(self.token)

rbx/grading/grading_context.py

@@ -0,0 +1,96 @@
+import contextvars
+from enum import Enum
+from typing import Callable, Optional, Union
+
+Condition = Union[bool, Callable[[], bool]]
+
+
+class ConditionedContext:
+    def __init__(self, when: Condition = True):
+        self.when = when
+
+    def should_enter(self) -> bool:
+        if isinstance(self.when, bool):
+            return self.when
+        return self.when()
+
+
+class CacheLevel(Enum):
+    NO_CACHE = 0
+    CACHE_TRANSIENTLY = 1
+    CACHE_COMPILATION = 2
+    CACHE_ALL = 3
+
+
+cache_level_var = contextvars.ContextVar('cache_level', default=CacheLevel.CACHE_ALL)
+
+
+def is_compilation_only() -> bool:
+    return cache_level_var.get() == CacheLevel.CACHE_COMPILATION
+
+
+def is_transient() -> bool:
+    return cache_level_var.get().value <= CacheLevel.CACHE_TRANSIENTLY.value
+
+
+def is_no_cache() -> bool:
+    return cache_level_var.get().value <= CacheLevel.NO_CACHE.value
+
+
+class cache_level(ConditionedContext):
+    def __init__(self, level: CacheLevel, when: Condition = True):
+        super().__init__(when)
+        self.level = level
+        self.token = None
+
+    def __enter__(self):
+        if self.should_enter():
+            self.token = cache_level_var.set(self.level)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.token is not None:
+            cache_level_var.reset(self.token)
+        return None
+
+
+compression_level_var = contextvars.ContextVar('compression_level', default=5)
+use_compression_var = contextvars.ContextVar('use_compression', default=False)
+
+
+def get_compression_level() -> int:
+    return compression_level_var.get()
+
+
+def should_compress() -> bool:
+    return use_compression_var.get()
+
+
+class compression(ConditionedContext):
+    def __init__(
+        self,
+        level: Optional[int] = None,
+        use_compression: Optional[bool] = None,
+        when: Condition = True,
+    ):
+        super().__init__(when)
+        self.level = level
+        self.use_compression = use_compression
+        self.level_token = None
+        self.use_compression_token = None
+
+    def __enter__(self):
+        if not self.should_enter():
+            return self
+        if self.level is not None:
+            self.level_token = compression_level_var.set(self.level)
+        if self.use_compression is not None:
+            self.use_compression_token = use_compression_var.set(self.use_compression)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.level_token is not None:
+            compression_level_var.reset(self.level_token)
+        if self.use_compression_token is not None:
+            use_compression_var.reset(self.use_compression_token)
+        return None

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: