Extend job report with an optional name property

alembic/versions/20250318_160027_afe52ed2f369_.py

@@ -0,0 +1,31 @@
+"""empty message
+
+Revision ID: afe52ed2f369
+Revises: a4cab64aa586
+Create Date: 2025-03-18 16:00:27.684926
+
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "afe52ed2f369"
+down_revision: str | None = "a4cab64aa586"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("job", sa.Column("name", sa.String(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("job", "name")
+    # ### end Alembic commands ###

app/db/model.py

@@ -5,7 +5,16 @@
 from typing import Any, ClassVar
 from uuid import UUID
 
-from sqlalchemy import DateTime, ForeignKey, Identity, Index, MetaData, SmallInteger, text, true
+from sqlalchemy import (
+    DateTime,
+    ForeignKey,
+    Identity,
+    Index,
+    MetaData,
+    SmallInteger,
+    text,
+    true,
+)
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 
@@ -59,6 +68,7 @@ class Job(Base):
     group_id: Mapped[UUID | None]
     vlab_id: Mapped[UUID] = mapped_column(ForeignKey("account.id"), index=True)
     proj_id: Mapped[UUID] = mapped_column(ForeignKey("account.id"), index=True)
+    name: Mapped[str | None]
     user_id: Mapped[UUID | None]
     service_type: Mapped[ServiceType]
     service_subtype: Mapped[ServiceSubtype]

app/repository/report.py

@@ -54,6 +54,7 @@ async def get_job_reports(
                 Job.service_subtype.label("subtype"),
                 Job.user_id,
                 Job.group_id,
+                Job.name,
                 Job.reserved_at,
                 Job.started_at,
                 Job.finished_at,

app/schema/api.py

@@ -107,6 +107,7 @@ class BaseMakeReservationIn(BaseModel):
     proj_id: UUID
     group_id: UUID | None = None
     user_id: UUID
+    name: str | None = None
     type: ServiceType
     subtype: ServiceSubtype
 
@@ -269,6 +270,7 @@ class OneshotReportOut(BaseModel, from_attributes=True):
 
     job_id: UUID
     user_id: UUID
+    name: str | None = None
     group_id: UUID | None = None
     vlab_id: UUID | None = None
     proj_id: UUID | None = None
@@ -287,6 +289,7 @@ class LongrunReportOut(BaseModel, from_attributes=True):
 
     job_id: UUID
     user_id: UUID
+    name: str | None = None
     group_id: UUID | None = None
     vlab_id: UUID | None = None
     proj_id: UUID | None = None

app/schema/queue.py

@@ -26,6 +26,7 @@ class OneshotEvent(BaseModel):
     subtype: ServiceSubtype
     proj_id: UUID
     job_id: UUID
+    name: str | None = None
     count: Annotated[int, Field(ge=0)]
     timestamp: RecentTimeStamp
 
@@ -37,6 +38,7 @@ class LongrunEvent(BaseModel):
     subtype: ServiceSubtype
     proj_id: UUID
     job_id: UUID
+    name: str | None = None
     status: LongrunStatus
     instances: Annotated[int | None, Field(ge=0)] = None
     instance_type: str | None = None

app/service/reservation.py

@@ -59,6 +59,7 @@ async def _make_reservation(
         vlab_id=accounts.vlab.id,
         proj_id=accounts.proj.id,
         user_id=reservation_request.user_id,
+        name=reservation_request.name,
         service_type=reservation_request.type,
         service_subtype=reservation_request.subtype,
         reserved_at=reserving_at,

app/task/queue_consumer/longrun.py

@@ -37,12 +37,17 @@ async def _handle_running(repos: RepositoryGroup, event: LongrunEvent, accounts:
 
 
 async def _handle_finished(repos: RepositoryGroup, event: LongrunEvent, accounts: Accounts) -> Job:
+    optional_kwargs = {}
+    if event.name is not None:
+        optional_kwargs["name"] = event.name
+
     return await repos.job.update_job(
         job_id=event.job_id,
         vlab_id=accounts.vlab.id,
         proj_id=accounts.proj.id,
         last_alive_at=event.timestamp,
         finished_at=event.timestamp,
+        **optional_kwargs,
     )
 
 

app/task/queue_consumer/oneshot.py

@@ -52,6 +52,7 @@ async def _consume(self, msg: dict[str, Any], db: AsyncSession) -> UUID:
             started_at=event.timestamp,
             last_alive_at=event.timestamp,
             finished_at=event.timestamp,
+            name=event.name or job.name,
             usage_params={
                 "count": event.count,
             },

tests/api/test_report.py

@@ -38,6 +38,7 @@ async def test_get_report(api_client, url, extra_expected_data):
                 **extra_expected_data,
                 "job_id": str(UUIDS.JOB[0]),
                 "user_id": str(UUIDS.USER[0]),
+                "name": "test job 0",
                 "group_id": str(UUIDS.GROUP[0]),
                 "type": "oneshot",
                 "subtype": "ml-llm",

tests/api/test_reservation.py

@@ -46,6 +46,28 @@ async def test_make_oneshot_reservation(api_client, mock_uuid):
     assert mock_uuid.call_count == 0
 
 
+@pytest.mark.usefixtures("_db_account", "_db_price")
+async def test_make_oneshot_reservation_with_job_name(api_client, mock_uuid):
+    # make a valid reservation
+    request_payload = {
+        "proj_id": PROJ_ID,
+        "user_id": USER_ID,
+        "name": "test job",
+        "group_id": GROUP_ID,
+        "type": ServiceType.ONESHOT,
+        "subtype": ServiceSubtype.ML_LLM,
+        "count": "1000000",
+    }
+    response = await api_client.post("/reservation/oneshot", json=request_payload)
+
+    assert response.json()["data"] == {
+        "job_id": "00000000-0000-0000-0000-000000000001",
+        "amount": "10.00000",
+    }
+    assert response.status_code == 201
+    assert mock_uuid.call_count == 1
+
+
 @pytest.mark.usefixtures("_db_account", "_db_price")
 async def test_make_longrun_reservation(api_client, mock_uuid):
     # make a valid reservation
@@ -89,3 +111,27 @@ async def test_make_longrun_reservation(api_client, mock_uuid):
     }
     assert response.status_code == 402
     assert mock_uuid.call_count == 0
+
+
+@pytest.mark.usefixtures("_db_account", "_db_price")
+async def test_make_longrun_reservation_with_job_name(api_client, mock_uuid):
+    # make a valid reservation
+    request_payload = {
+        "proj_id": PROJ_ID,
+        "user_id": USER_ID,
+        "name": "test job",
+        "group_id": GROUP_ID,
+        "type": ServiceType.LONGRUN,
+        "subtype": ServiceSubtype.SINGLE_CELL_SIM,
+        "instances": "2",
+        "instance_type": "",
+        "duration": "3600",
+    }
+    response = await api_client.post("/reservation/longrun", json=request_payload)
+
+    assert response.json()["data"] == {
+        "job_id": "00000000-0000-0000-0000-000000000001",
+        "amount": "73.50",
+    }
+    assert response.status_code == 201
+    assert mock_uuid.call_count == 1

tests/api/test_usage.py

@@ -34,13 +34,51 @@ async def test_add_oneshot_usage(api_client, mock_sqs_manager):
     assert mock_sqs_manager.client.send_message.await_count == 1
 
 
-async def test_add_longrun_usage(api_client, mock_sqs_manager):
+async def test_add_oneshot_usage_with_job_name(api_client, mock_sqs_manager):
+    timestamp = since_unix_epoch()
+    request_payload = {
+        "job_id": "00000000-0000-0000-0000-000000000001",
+        "proj_id": PROJ_ID,
+        "type": ServiceType.ONESHOT,
+        "subtype": ServiceSubtype.ML_LLM,
+        "name": "test_job",
+        "count": "1000000",
+        "timestamp": timestamp,
+    }
+    response = await api_client.post("/usage/oneshot", json=request_payload)
+
+    assert response.json()["data"] is None
+    assert response.status_code == 201
+    assert mock_sqs_manager.client.send_message.await_count == 1
+
+
+async def test_add_longrun_usage_with_started_status(api_client, mock_sqs_manager):
+    timestamp = since_unix_epoch()
+    request_payload = {
+        "job_id": "00000000-0000-0000-0000-000000000001",
+        "proj_id": PROJ_ID,
+        "type": ServiceType.LONGRUN,
+        "subtype": ServiceSubtype.SINGLE_CELL_SIM,
+        "instances": "2",
+        "instance_type": "",
+        "status": "started",
+        "timestamp": timestamp,
+    }
+    response = await api_client.post("/usage/longrun", json=request_payload)
+
+    assert response.json()["data"] is None
+    assert response.status_code == 201
+    assert mock_sqs_manager.client.send_message.await_count == 1
+
+
+async def test_add_longrun_usage_with_job_name(api_client, mock_sqs_manager):
     timestamp = since_unix_epoch()
     request_payload = {
         "job_id": "00000000-0000-0000-0000-000000000001",
         "proj_id": PROJ_ID,
         "type": ServiceType.LONGRUN,
         "subtype": ServiceSubtype.SINGLE_CELL_SIM,
+        "name": "test_job",
         "instances": "2",
         "instance_type": "",
         "status": "started",

tests/conftest.py

@@ -223,6 +223,7 @@ async def _db_job(db, _db_account):
                 "vlab_id": UUIDS.VLAB[0],
                 "proj_id": UUIDS.PROJ[0],
                 "user_id": USER_ID,
+                "name": "test job 0",
                 "service_type": ServiceType.ONESHOT,
                 "service_subtype": ServiceSubtype.ML_LLM,
                 "reserved_at": dt,

tests/repository/test_report.py

@@ -17,6 +17,7 @@
         "type": ServiceType.ONESHOT,
         "subtype": ServiceSubtype.ML_LLM,
         "user_id": UUID(USER_ID),
+        "name": "test job 0",
         "reserved_at": ANY,
         "started_at": ANY,
         "finished_at": ANY,
@@ -35,6 +36,7 @@
         "type": ServiceType.LONGRUN,
         "subtype": ServiceSubtype.SINGLE_CELL_SIM,
         "user_id": UUID(USER_ID),
+        "name": None,
         "reserved_at": ANY,
         "started_at": ANY,
         "finished_at": ANY,
@@ -53,6 +55,7 @@
         "type": ServiceType.STORAGE,
         "subtype": ServiceSubtype.STORAGE,
         "user_id": None,
+        "name": None,
         "reserved_at": None,
         "started_at": ANY,
         "finished_at": ANY,