feat: Support for handling Benchmarks V2 (#234)

* wip * Make `BenchmarkSpecification` into an abstract class, and have V1 and V2 implementations. Add Pydantic models for the V2 split, and V2 validations. * Post-rebase cleanup * Further split V1 and V2 into two classes. Tighten model validation so everything stays coherent. Serialize the artifact version when dumping a model. Various fixes. * Add upload in client. Rework the path handling in the storage session. * Fix some errors with get functions * Fix splitfiles upload and failing tests * Add some test coverage * Review feedback --------- Co-authored-by: Honore Hounwanou <[email protected]>
polaris-hub · Dec 17, 2024 · abb9b29 · abb9b29
1 parent ff0bb3b
commit abb9b29
Show file tree

Hide file tree

Showing 23 changed files with 1,057 additions and 520 deletions.
diff --git a/env.yml b/env.yml
@@ -14,7 +14,7 @@ dependencies:
   - yaspin
   - typing-extensions >=4.12.0
   - boto3 >=1.35.0
-
+  - pyroaring
 
   # Hub client
   - authlib
@@ -25,7 +25,7 @@ dependencies:
   - filelock
 
   # Scientific
-  - numpy < 2  # We need to pin numpy to avoid issues with fastpdb/biotite.
+  - numpy < 2 # We need to pin numpy to avoid issues with fastpdb/biotite.
   - pandas
   - scipy
   - scikit-learn

diff --git a/polaris/_artifact.py b/polaris/_artifact.py
@@ -1,5 +1,5 @@
 import json
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 import fsspec
 from loguru import logger
@@ -41,6 +41,7 @@ class BaseArtifactModel(BaseModel):
         polaris_version: The version of the Polaris library that was used to create the artifact.
     """
 
+    _version: ClassVar[Literal[1]] = 1
     _artifact_type: ClassVar[str]
 
     model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True, arbitrary_types_allowed=True)
@@ -72,6 +73,11 @@ def urn(self) -> ArtifactUrn | None:
             return self.urn_for(self.owner, self.slug)
         return None
 
+    @computed_field
+    @property
+    def version(self) -> int:
+        return self._version
+
     @field_validator("polaris_version")
     @classmethod
     def _validate_version(cls, value: str) -> str:

diff --git a/polaris/benchmark/__init__.py b/polaris/benchmark/__init__.py
@@ -1,11 +1,9 @@
-from polaris.benchmark._base import BenchmarkSpecification
-from polaris.benchmark._definitions import (
-    MultiTaskBenchmarkSpecification,
-    SingleTaskBenchmarkSpecification,
-)
+from polaris.benchmark._base import BenchmarkSpecification, BenchmarkV1Specification
+from polaris.benchmark._definitions import MultiTaskBenchmarkSpecification, SingleTaskBenchmarkSpecification
 
 __all__ = [
     "BenchmarkSpecification",
+    "BenchmarkV1Specification",
     "SingleTaskBenchmarkSpecification",
     "MultiTaskBenchmarkSpecification",
 ]