Skip to content

Commit 18088f8

Browse files
abergeronabergeron
committed
more fixes 3
1 parent e34598e commit 18088f8

20 files changed

+1064
-342
lines changed

‎sarc/allocations/allocations.py

+7-3
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@
55

66
import pandas as pd
77
from flatten_dict import flatten
8-
from pydantic import BeforeValidator, ByteSize
9-
from pydantic_mongo import AbstractRepository, ObjectIdField
8+
from pydantic import BeforeValidator, ByteSize, field_serializer
9+
from pydantic_mongo import AbstractRepository, PydanticObjectId
1010

1111
from sarc.config import BaseModel, config, validate_date
1212

@@ -39,7 +39,7 @@ def _convert_date_to_iso(date_value: date) -> datetime:
3939

4040
class Allocation(BaseModel):
4141
# Database ID
42-
id: ObjectIdField = None
42+
id: PydanticObjectId = None
4343

4444
cluster_name: str
4545
resource_name: str
@@ -49,6 +49,10 @@ class Allocation(BaseModel):
4949
end: Annotated[date, BeforeValidator(validate_date)]
5050
resources: AllocationRessources
5151

52+
@field_serializer('start', 'end')
53+
def save_as_datetime(self, value, info):
54+
return datetime(year=value.year, month=value.month, day=value.day)
55+
5256

5357
class AllocationsRepository(AbstractRepository[Allocation]):
5458
class Meta:

‎sarc/client/gpumetrics.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from typing import Dict, List
66

77
from pydantic import field_validator
8-
from pydantic_mongo import AbstractRepository, ObjectIdField
8+
from pydantic_mongo import AbstractRepository, PydanticObjectId
99

1010
from sarc.config import MTL, UTC, BaseModel, config, scraping_mode_required
1111

@@ -16,7 +16,7 @@ class GPUBilling(BaseModel):
1616
"""Holds data for a GPU Billing."""
1717

1818
# # Database ID
19-
id: ObjectIdField = None
19+
id: PydanticObjectId = None
2020

2121
cluster_name: str
2222
since: datetime

‎sarc/client/job.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from typing import Iterable, Optional
66

77
from pydantic import field_validator
8-
from pydantic_mongo import AbstractRepository, ObjectIdField
8+
from pydantic_mongo import AbstractRepository, PydanticObjectId
99

1010
from sarc.traces import trace_decorator
1111

@@ -108,7 +108,7 @@ class SlurmJob(BaseModel):
108108
"""Holds data for a Slurm job."""
109109

110110
# Database ID
111-
id: ObjectIdField = None
111+
id: PydanticObjectId = None
112112

113113
# job identification
114114
cluster_name: str
@@ -400,7 +400,7 @@ class SlurmCLuster(BaseModel):
400400
"""Hold data for a Slurm cluster."""
401401

402402
# Database ID
403-
id: ObjectIdField = None
403+
id: PydanticObjectId = None
404404

405405
cluster_name: str
406406
start_date: Optional[str] = None

‎sarc/client/users/api.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from datetime import date
1010
from typing import Optional
1111

12-
from pydantic_mongo import AbstractRepository, ObjectIdField
12+
from pydantic_mongo import AbstractRepository, PydanticObjectId
1313

1414
from sarc.config import BaseModel, config
1515
from sarc.users.revision import query_latest_records
@@ -22,7 +22,7 @@ class Credentials(BaseModel):
2222

2323

2424
class User(BaseModel):
25-
id: ObjectIdField = None
25+
id: PydanticObjectId = None
2626

2727
name: str
2828

‎sarc/config.py

+6-22
Original file line numberDiff line numberDiff line change
@@ -50,22 +50,6 @@ class BaseModel(_BaseModel):
5050
arbitrary_types_allowed=True,
5151
)
5252

53-
def dict(self, *args, **kwargs) -> dict[str, Any]:
54-
d = super().dict(*args, **kwargs)
55-
for k, v in list(d.items()):
56-
if isinstance(getattr(type(self), k, None), cached_property):
57-
del d[k]
58-
continue
59-
60-
for k, v in d.items():
61-
if isinstance(v, date) and not isinstance(v, datetime):
62-
d[k] = datetime(
63-
year=v.year,
64-
month=v.month,
65-
day=v.day,
66-
)
67-
return d
68-
6953
def replace(self, **replacements):
7054
new_arguments = {**self.dict(), **replacements}
7155
return type(self)(**new_arguments)
@@ -88,8 +72,8 @@ class ClusterConfig(BaseModel):
8872
duc_storage_command: str | None = None
8973
diskusage_report_command: str | None = None
9074
start_date: str = "2022-04-01"
91-
rgu_start_date: str = None
92-
gpu_to_rgu_billing: Path = None
75+
rgu_start_date: str | None = None
76+
gpu_to_rgu_billing: Path | None = None
9377
slurm_conf_host_path: str = "/etc/slurm/slurm.conf"
9478

9579
# Tell if billing (in job's requested|allocated field) is number of GPUs (True) or RGU (False)
@@ -247,8 +231,8 @@ def _absolute_path(value):
247231
class Config(BaseModel):
248232
mongo: MongoConfig
249233
cache: Annotated[Path, AfterValidator(_absolute_path)] = None
250-
loki: LokiConfig = None
251-
tempo: TempoConfig = None
234+
loki: LokiConfig | None = None
235+
tempo: TempoConfig | None = None
252236

253237

254238
class ScraperConfig(BaseModel):
@@ -261,8 +245,8 @@ class ScraperConfig(BaseModel):
261245
sshconfig: Annotated[Path | None, AfterValidator(_absolute_path)] = None
262246
clusters: dict[str, ClusterConfig] = None
263247
logging: LoggingConfig = None
264-
loki: LokiConfig = None
265-
tempo: TempoConfig = None
248+
loki: LokiConfig | None = None
249+
tempo: TempoConfig | None = None
266250

267251
@field_validator("clusters", mode="after")
268252
@classmethod

‎sarc/jobs/node_gpu_mapping.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from typing import Dict, Optional
77

88
from pydantic import field_validator
9-
from pydantic_mongo import AbstractRepository, ObjectIdField
9+
from pydantic_mongo import AbstractRepository, PydanticObjectId
1010

1111
from sarc.config import MTL, UTC, BaseModel, config, scraping_mode_required
1212

@@ -17,7 +17,7 @@ class NodeGPUMapping(BaseModel):
1717
"""Holds data for a mapping <node name> -> <GPU type>."""
1818

1919
# # Database ID
20-
id: ObjectIdField = None
20+
id: PydanticObjectId = None
2121

2222
cluster_name: str
2323
since: datetime

‎sarc/storage/diskusage.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from datetime import date, datetime
44

55
from pydantic import ByteSize
6-
from pydantic_mongo import AbstractRepository, ObjectIdField
6+
from pydantic_mongo import AbstractRepository, PydanticObjectId
77

88
from sarc.config import BaseModel, config
99

@@ -25,7 +25,7 @@ class DiskUsage(BaseModel):
2525
"""
2626

2727
# # Database ID
28-
id: ObjectIdField = None
28+
id: PydanticObjectId = None
2929

3030
cluster_name: str
3131
groups: list[DiskUsageGroup]

‎tests/functional/allocations/test_func_allocations.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ def test_get_allocations(params, data_regression):
2828
data = get_allocations(**params)
2929

3030
data_regression.check(
31-
[allocation.json(exclude={"id": True}) for allocation in data]
31+
[allocation.model_dump(exclude={"id": True}) for allocation in data]
3232
)
3333

3434

‎tests/functional/allocations/test_func_allocations/test_get_allocations_another_name_.yml

+100-27
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,100 @@
1-
- '{"cluster_name": "patate", "resource_name": "patate-gpu", "group_name": "rrg-bonhomme-ad",
2-
"timestamp": "2023-02-01T00:00:00", "start": "2019-04-01", "end": "2020-04-01",
3-
"resources": {"compute": {"gpu_year": 190, "cpu_year": null, "vcpu_year": null,
4-
"vgpu_year": null}, "storage": {"project_size": null, "project_inodes": null, "nearline":
5-
null, "dCache": 0, "object": 0, "cloud_volume": 0, "cloud_shared": 0}}}'
6-
- '{"cluster_name": "patate", "resource_name": "patate-storage", "group_name": "rrg-bonhomme-ad",
7-
"timestamp": "2023-02-01T00:00:00", "start": "2019-04-01", "end": "2020-04-01",
8-
"resources": {"compute": {"gpu_year": null, "cpu_year": null, "vcpu_year": null,
9-
"vgpu_year": null}, "storage": {"project_size": 90000000000000, "project_inodes":
10-
5000000.0, "nearline": 90000000000000, "dCache": 0, "object": 0, "cloud_volume":
11-
0, "cloud_shared": 0}}}'
12-
- '{"cluster_name": "patate", "resource_name": "patate-compute", "group_name": "rrg-bonhomme-ad",
13-
"timestamp": "2023-02-01T00:00:00", "start": "2020-04-01", "end": "2021-04-01",
14-
"resources": {"compute": {"gpu_year": null, "cpu_year": 219, "vcpu_year": null,
15-
"vgpu_year": null}, "storage": {"project_size": null, "project_inodes": null, "nearline":
16-
null, "dCache": 0, "object": 0, "cloud_volume": 0, "cloud_shared": 0}}}'
17-
- '{"cluster_name": "patate", "resource_name": "patate-gpu", "group_name": "rrg-bonhomme-ad",
18-
"timestamp": "2023-02-01T00:00:00", "start": "2020-04-01", "end": "2021-04-01",
19-
"resources": {"compute": {"gpu_year": 200, "cpu_year": null, "vcpu_year": null,
20-
"vgpu_year": null}, "storage": {"project_size": null, "project_inodes": null, "nearline":
21-
null, "dCache": 0, "object": 0, "cloud_volume": 0, "cloud_shared": 0}}}'
22-
- '{"cluster_name": "patate", "resource_name": "patate-storage", "group_name": "rrg-bonhomme-ad",
23-
"timestamp": "2023-02-01T00:00:00", "start": "2020-04-01", "end": "2021-04-01",
24-
"resources": {"compute": {"gpu_year": null, "cpu_year": null, "vcpu_year": null,
25-
"vgpu_year": null}, "storage": {"project_size": 70000000000000, "project_inodes":
26-
5000000.0, "nearline": 80000000000000, "dCache": 0, "object": 0, "cloud_volume":
27-
0, "cloud_shared": 0}}}'
1+
- cluster_name: patate
2+
end: 2020-04-01
3+
group_name: rrg-bonhomme-ad
4+
resource_name: patate-gpu
5+
resources:
6+
compute:
7+
cpu_year: null
8+
gpu_year: 190
9+
vcpu_year: null
10+
vgpu_year: null
11+
storage:
12+
cloud_shared: 0
13+
cloud_volume: 0
14+
dCache: 0
15+
nearline: null
16+
object: 0
17+
project_inodes: null
18+
project_size: null
19+
start: 2019-04-01
20+
timestamp: 2023-02-01 00:00:00
21+
- cluster_name: patate
22+
end: 2020-04-01
23+
group_name: rrg-bonhomme-ad
24+
resource_name: patate-storage
25+
resources:
26+
compute:
27+
cpu_year: null
28+
gpu_year: null
29+
vcpu_year: null
30+
vgpu_year: null
31+
storage:
32+
cloud_shared: 0
33+
cloud_volume: 0
34+
dCache: 0
35+
nearline: 90000000000000
36+
object: 0
37+
project_inodes: 5000000.0
38+
project_size: 90000000000000
39+
start: 2019-04-01
40+
timestamp: 2023-02-01 00:00:00
41+
- cluster_name: patate
42+
end: 2021-04-01
43+
group_name: rrg-bonhomme-ad
44+
resource_name: patate-compute
45+
resources:
46+
compute:
47+
cpu_year: 219
48+
gpu_year: null
49+
vcpu_year: null
50+
vgpu_year: null
51+
storage:
52+
cloud_shared: 0
53+
cloud_volume: 0
54+
dCache: 0
55+
nearline: null
56+
object: 0
57+
project_inodes: null
58+
project_size: null
59+
start: 2020-04-01
60+
timestamp: 2023-02-01 00:00:00
61+
- cluster_name: patate
62+
end: 2021-04-01
63+
group_name: rrg-bonhomme-ad
64+
resource_name: patate-gpu
65+
resources:
66+
compute:
67+
cpu_year: null
68+
gpu_year: 200
69+
vcpu_year: null
70+
vgpu_year: null
71+
storage:
72+
cloud_shared: 0
73+
cloud_volume: 0
74+
dCache: 0
75+
nearline: null
76+
object: 0
77+
project_inodes: null
78+
project_size: null
79+
start: 2020-04-01
80+
timestamp: 2023-02-01 00:00:00
81+
- cluster_name: patate
82+
end: 2021-04-01
83+
group_name: rrg-bonhomme-ad
84+
resource_name: patate-storage
85+
resources:
86+
compute:
87+
cpu_year: null
88+
gpu_year: null
89+
vcpu_year: null
90+
vgpu_year: null
91+
storage:
92+
cloud_shared: 0
93+
cloud_volume: 0
94+
dCache: 0
95+
nearline: 80000000000000
96+
object: 0
97+
project_inodes: 5000000.0
98+
project_size: 70000000000000
99+
start: 2020-04-01
100+
timestamp: 2023-02-01 00:00:00

0 commit comments

Comments
 (0)