Skip to content

feat: Model Versioning #284

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Apr 21, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 22 additions & 1 deletion docs/tutorials/create_a_model.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
"\n",
"# Create a new Model Card\n",
"model = Model(\n",
" name=\"MolGPS\", \n",
" name=\"MolGPS\",\n",
" description=\"Graph transformer foundation model for molecular modeling\",\n",
" code_url=\"https://github.com/datamol-io/graphium\"\n",
")"
Expand All @@ -44,6 +44,27 @@
"model.upload_to_hub(owner=\"your-username\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"If you want to upload a new version of your model, you can specify its previous version with the `parent_artifact_id` parameter. Don't forget to add a changelog describing your updates!"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"model.artifact_changelog = \"In this version, I added...\"\n",
"\n",
"model.upload_to_hub(\n",
" owner=\"your-username\",\n",
" parent_artifact_id=\"your-username/tutorial-example\"\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand Down
2 changes: 1 addition & 1 deletion polaris/benchmark/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,7 @@ class BenchmarkSpecification(
readme: str = ""

# Version-related fields
artifact_version: int = 1
artifact_version: int = Field(default=1, frozen=True)
artifact_changelog: str | None = None

@computed_field
Expand Down
2 changes: 1 addition & 1 deletion polaris/dataset/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ class BaseDataset(BaseArtifactModel, abc.ABC):
curation_reference: HttpUrlString | None = None

# Version-related fields
artifact_version: int = 1
artifact_version: int = Field(default=1, frozen=True)
artifact_changelog: str | None = None

# Private attributes
Expand Down
11 changes: 9 additions & 2 deletions polaris/hub/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -942,6 +942,7 @@ def upload_model(
model: Model,
access: AccessType = "private",
owner: HubOwner | str | None = None,
parent_artifact_id: str | None = None,
):
"""Upload a model to the Polaris Hub.

Expand All @@ -959,6 +960,7 @@ def upload_model(
model: The model to upload.
access: Grant public or private access to result
owner: Which Hub user or organization owns the artifact. Takes precedence over `model.owner`.
parent_artifact_id: The `owner/slug` of the parent model, if uploading a new version of a model.
"""
with track_progress(description="Uploading model", total=1) as (progress, task):
# Get the serialized model data-structure
Expand All @@ -967,11 +969,16 @@ def upload_model(

# Make a request to the Hub
url = f"/v2/model/{model.artifact_id}"
response = self._base_request_to_hub(url=url, method="PUT", json={"access": access, **model_json})
response = self._base_request_to_hub(
url=url,
method="PUT",
json={"access": access, "parentArtifactId": parent_artifact_id, **model_json},
)

# NOTE: When we merge in the competition model feature, we will need to update the slug with the inserted model slug to make sure we write to the correct storage location.

# Inform the user about where to find their newly created artifact.
model_url = urljoin(self.settings.hub_url, response.headers.get("Content-Location"))

progress.log(
f"[green]Your model has been successfully uploaded to the Hub. View it here: {model_url}"
)
16 changes: 14 additions & 2 deletions polaris/model/__init__.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from polaris._artifact import BaseArtifactModel
from polaris.utils.types import HttpUrlString
from polaris.utils.types import AccessType, HubOwner
from pydantic import Field


class Model(BaseArtifactModel):
Expand Down Expand Up @@ -30,6 +31,8 @@ class Model(BaseArtifactModel):
readme (str): A detailed README describing the model.
code_url (HttpUrlString | None): Optional URL pointing to the model's code repository.
report_url (HttpUrlString | None): Optional URL linking to a report or publication related to the model.
artifact_version: The version of the model.
artifact_changelog: A description of the changes made in this model version.

Methods:
upload_to_hub(access: AccessType = "private", owner: HubOwner | str | None = None):
Expand All @@ -44,11 +47,20 @@ class Model(BaseArtifactModel):
code_url: HttpUrlString | None = None
report_url: HttpUrlString | None = None

def upload_to_hub(self, access: AccessType = "private", owner: HubOwner | str | None = None):
# Version-related fields
artifact_version: int = Field(default=1, frozen=True)
artifact_changelog: str | None = None

def upload_to_hub(
self,
access: AccessType = "private",
owner: HubOwner | str | None = None,
parent_artifact_id: str | None = None,
):
"""
Uploads the model to the Polaris Hub.
"""
from polaris.hub.client import PolarisHubClient

with PolarisHubClient() as client:
client.upload_model(self, owner=owner, access=access)
client.upload_model(self, owner=owner, access=access, parent_artifact_id=parent_artifact_id)
Loading