Skip to content
11 changes: 6 additions & 5 deletions alchemiscale/compute/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@
"""

import json
from datetime import datetime, timedelta
import datetime
from datetime import timedelta
import random

from fastapi import FastAPI, APIRouter, Body, Depends, Request
Expand Down Expand Up @@ -103,7 +104,7 @@ def register_computeservice(
compute_service_id,
n4js: Neo4jStore = Depends(get_n4js_depends),
):
now = datetime.utcnow()
now = datetime.datetime.now(tz=datetime.UTC)
csreg = ComputeServiceRegistration(
identifier=ComputeServiceID(compute_service_id),
registered=now,
Expand Down Expand Up @@ -134,7 +135,7 @@ def heartbeat_computeservice(
n4js: Neo4jStore = Depends(get_n4js_depends),
settings: ComputeAPISettings = Depends(get_base_api_settings),
):
now = datetime.utcnow()
now = datetime.datetime.now(tz=datetime.UTC)

# expire any stale registrations, along with their claims
expire_delta = timedelta(
Expand Down Expand Up @@ -215,7 +216,7 @@ def claim_tasks(

"""
# check if the compute service can claim tasks
now = datetime.now()
now = datetime.datetime.now(tz=datetime.UTC)
if not n4js.compute_service_can_claim(
compute_service_id,
now - timedelta(seconds=settings.ALCHEMISCALE_COMPUTE_API_FORGIVE_TIME_SECONDS),
Expand Down Expand Up @@ -384,7 +385,7 @@ async def set_task_result(
n4js.set_task_error(tasks=[task_sk])

# report that the compute service experienced a failure
now = datetime.utcnow()
now = datetime.datetime.now(tz=datetime.UTC)
n4js.log_failure_compute_service(compute_service_id, now)
n4js.resolve_task_restarts(task_scoped_keys=[task_sk])

Expand Down
5 changes: 3 additions & 2 deletions alchemiscale/security/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@
"""

import secrets
from datetime import datetime, timedelta
import datetime
from datetime import timedelta

import bcrypt
from fastapi import HTTPException, status
Expand Down Expand Up @@ -112,7 +113,7 @@ def create_access_token(
) -> str:
to_encode = data.copy()

expire = datetime.utcnow() + timedelta(seconds=expires_seconds)
expire = datetime.datetime.now(tz=datetime.UTC) + timedelta(seconds=expires_seconds)
to_encode.update({"exp": expire})

encoded_jwt = jwt.encode(to_encode, secret_key, algorithm=jwt_algorithm)
Expand Down
4 changes: 2 additions & 2 deletions alchemiscale/security/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

"""

from datetime import datetime
import datetime

from pydantic import BaseModel, field_validator

Expand All @@ -23,7 +23,7 @@ class TokenData(BaseModel):

class CredentialedEntity(BaseModel):
hashed_key: str
expires: datetime | None = None
expires: datetime.datetime | None = None


class ScopedIdentity(BaseModel):
Expand Down
16 changes: 9 additions & 7 deletions alchemiscale/storage/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

from abc import abstractmethod
from copy import copy
from datetime import datetime
import datetime
from enum import Enum
from uuid import uuid4
import hashlib
Expand Down Expand Up @@ -39,7 +39,7 @@ def __str__(self):

@classmethod
def from_now(cls, identifier: ComputeServiceID):
now = datetime.utcnow()
now = datetime.datetime.now(tz=datetime.UTC)
return cls(
identifier=identifier, registered=now, heartbeat=now, failure_times=[]
)
Expand Down Expand Up @@ -96,7 +96,7 @@ class Task(GufeTokenizable):
status: TaskStatusEnum
priority: int
claim: str | None
datetime_created: datetime | None
datetime_created: datetime.datetime | None
creator: str | None
extends: str | None

Expand All @@ -105,7 +105,7 @@ def __init__(
*,
status: str | TaskStatusEnum = TaskStatusEnum.waiting,
priority: int = 10,
datetime_created: datetime | None = None,
datetime_created: datetime.datetime | None = None,
creator: str | None = None,
extends: str | None = None,
claim: str | None = None,
Expand All @@ -118,7 +118,9 @@ def __init__(
self.priority = priority

self.datetime_created = (
datetime_created if datetime_created is not None else datetime.utcnow()
datetime_created
if datetime_created is not None
else datetime.datetime.now(tz=datetime.UTC)
)

self.creator = creator
Expand Down Expand Up @@ -422,7 +424,7 @@ def __init__(
obj_key: GufeKey,
scope: Scope,
ok: bool,
datetime_created: datetime | None = None,
datetime_created: datetime.datetime | None = None,
creator: str | None = None,
):
self.location = location
Expand Down Expand Up @@ -450,7 +452,7 @@ def _to_dict(self):
def _from_dict(cls, d):
d_ = copy(d)
d_["datetime_created"] = (
datetime.fromisoformat(d["datetime_created"])
datetime.datetime.fromisoformat(d["datetime_created"])
if d.get("received") is not None
else None
)
Expand Down
4 changes: 2 additions & 2 deletions alchemiscale/storage/objectstore.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
"""

import os
from datetime import datetime
import datetime
from boto3.session import Session
from functools import lru_cache

Expand Down Expand Up @@ -237,7 +237,7 @@ def push_protocoldagresult(
obj_key=protocoldagresult_gufekey,
scope=transformation.scope,
ok=ok,
datetime_created=datetime.utcnow(),
datetime_created=datetime.datetime.now(tz=datetime.UTC),
creator=creator,
)

Expand Down
22 changes: 11 additions & 11 deletions alchemiscale/storage/statestore.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

import abc
import bisect
from datetime import datetime
import datetime
from contextlib import contextmanager
import json
import re
Expand Down Expand Up @@ -115,7 +115,7 @@ def _select_tasks_from_taskpool(taskpool: list[tuple[str, float]], count) -> lis

// create CLAIMS relationship with given compute service
MATCH (csreg:ComputeServiceRegistration {{identifier: $compute_service_id}})
CREATE (t)<-[cl:CLAIMS {{claimed: localdatetime($datetimestr)}}]-(csreg)
CREATE (t)<-[cl:CLAIMS {{claimed: datetime($datetimestr)}}]-(csreg)

SET t.status = '{TaskStatusEnum.running.value}'

Expand Down Expand Up @@ -1255,25 +1255,25 @@ def deregister_computeservice(self, compute_service_id: ComputeServiceID):
return ComputeServiceID(identifier)

def heartbeat_computeservice(
self, compute_service_id: ComputeServiceID, heartbeat: datetime
self, compute_service_id: ComputeServiceID, heartbeat: datetime.datetime
):
"""Update the heartbeat for the given ComputeServiceID."""

q = f"""
MATCH (n:ComputeServiceRegistration {{identifier: $compute_service_id}})
SET n.heartbeat = localdatetime('{heartbeat.isoformat()}')
SET n.heartbeat = datetime('{heartbeat.isoformat()}')

"""
with self.transaction() as tx:
tx.run(q, compute_service_id=str(compute_service_id))

return compute_service_id

def expire_registrations(self, expire_time: datetime):
def expire_registrations(self, expire_time: datetime.datetime):
"""Remove all registrations with last heartbeat prior to the given `expire_time`."""
q = f"""
MATCH (n:ComputeServiceRegistration)
WHERE n.heartbeat < localdatetime('{expire_time.isoformat()}')
WHERE n.heartbeat < datetime('{expire_time.isoformat()}')

WITH n

Expand All @@ -1298,7 +1298,7 @@ def expire_registrations(self, expire_time: datetime):
def log_failure_compute_service(
self,
compute_service_id: ComputeServiceID,
failure_time: datetime,
failure_time: datetime.datetime,
) -> ComputeServiceID:
"""Add a reported compute service failure to the database.

Expand All @@ -1311,7 +1311,7 @@ def log_failure_compute_service(
"""
q = """
MATCH (n:ComputeServiceRegistration {identifier: $compute_service_id})
SET n.failure_times = [localdatetime($failure_time)] + n.failure_times
SET n.failure_times = [datetime($failure_time)] + n.failure_times
"""

with self.transaction() as tx:
Expand All @@ -1326,7 +1326,7 @@ def log_failure_compute_service(
def compute_service_can_claim(
self,
compute_service_id: ComputeServiceID,
forgive_time: datetime,
forgive_time: datetime.datetime,
max_failures: int,
) -> bool:
"""Check if a compute service is able to claim a ``Task``.
Expand All @@ -1345,7 +1345,7 @@ def compute_service_can_claim(
# get the number of failures that occured after `forgive_time`
query = """
MATCH (cs:ComputeServiceRegistration {identifier: $compute_service_id})
SET cs.failure_times = [entry IN cs.failure_times WHERE entry > localdatetime($forgive_time)]
SET cs.failure_times = [entry IN cs.failure_times WHERE entry > datetime($forgive_time)]
RETURN size(cs.failure_times) as n_failures
"""
results = self.execute_query(
Expand Down Expand Up @@ -2074,7 +2074,7 @@ def task_count(task_dict: dict):
tx.run(
CLAIM_QUERY,
tasks_list=[str(task) for task in tasks if task is not None],
datetimestr=str(datetime.utcnow().isoformat()),
datetimestr=str(datetime.datetime.now(tz=datetime.UTC).isoformat()),
compute_service_id=str(compute_service_id),
)

Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import pytest
import json
import os
from datetime import datetime
import datetime
from time import sleep

from gufe.tokenization import JSON_HANDLER
Expand Down Expand Up @@ -378,7 +378,7 @@ def test_set_task_result_legacy(
obj_key=protocoldagresult.key,
scope=tf_sk.scope,
ok=protocoldagresult.ok(),
datetime_created=datetime.utcnow(),
datetime_created=datetime.datetime.now(tz=datetime.UTC),
creator=None,
)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ def test_cycle_max_failures(self, n4js_preloaded, s3os_server_fresh, service):
# create blocking failures
query = """
MATCH (cs:ComputeServiceRegistration {identifier: $compute_service_id})
SET cs.failure_times = [localdatetime()] + cs.failure_times
SET cs.failure_times = [datetime()] + cs.failure_times
"""

for _ in range(4):
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import pytest
from datetime import datetime
import datetime
from time import sleep
import os
from pathlib import Path
Expand Down Expand Up @@ -2212,7 +2212,7 @@ def _push_result_legacy(task_scoped_key, protocoldagresult, n4js, s3os_server):
obj_key=protocoldagresult.key,
scope=transformation_sk.scope,
ok=ok,
datetime_created=datetime.utcnow(),
datetime_created=datetime.datetime.now(tz=datetime.UTC),
creator=None,
)
n4js.set_task_result(
Expand Down
15 changes: 8 additions & 7 deletions alchemiscale/tests/integration/storage/test_statestore.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from datetime import datetime, timedelta
import datetime
from datetime import timedelta
import random
from pathlib import Path
from functools import reduce
Expand Down Expand Up @@ -573,7 +574,7 @@ def test_get_transformation_failures(
### compute

def test_register_computeservice(self, n4js, compute_service_id):
now = datetime.utcnow()
now = datetime.datetime.now(tz=datetime.UTC)
registration = ComputeServiceRegistration(
identifier=compute_service_id,
registered=now,
Expand All @@ -599,7 +600,7 @@ def test_register_computeservice(self, n4js, compute_service_id):
assert int(csreg["heartbeat"].to_native().timestamp()) == int(now.timestamp())

def test_deregister_computeservice(self, n4js, compute_service_id):
now = datetime.utcnow()
now = datetime.datetime.now(tz=datetime.UTC)
registration = ComputeServiceRegistration(
identifier=compute_service_id,
registered=now,
Expand All @@ -622,7 +623,7 @@ def test_deregister_computeservice(self, n4js, compute_service_id):
assert not csreg.records

def test_heartbeat_computeservice(self, n4js, compute_service_id):
now = datetime.utcnow()
now = datetime.datetime.now(tz=datetime.UTC)
registration = ComputeServiceRegistration(
identifier=compute_service_id,
registered=now,
Expand Down Expand Up @@ -651,7 +652,7 @@ def test_heartbeat_computeservice(self, n4js, compute_service_id):
)

def test_expire_registrations(self, n4js, compute_service_id):
now = datetime.utcnow()
now = datetime.datetime.now(tz=datetime.UTC)
yesterday = now - timedelta(days=1)
an_hour_ago = now - timedelta(hours=1)
registration = ComputeServiceRegistration(
Expand All @@ -678,7 +679,7 @@ def test_expire_registrations(self, n4js, compute_service_id):
assert compute_service_id in identities

def test_log_failure_computeservice(self, n4js, compute_service_id):
now = datetime.utcnow()
now = datetime.datetime.now(tz=datetime.UTC)
registration = ComputeServiceRegistration(
identifier=compute_service_id,
registered=now,
Expand Down Expand Up @@ -706,7 +707,7 @@ def test_log_failure_computeservice(self, n4js, compute_service_id):
assert 6 == results.records[0]["n_failures"]

def test_compute_service_can_claim(self, n4js, compute_service_id):
now = datetime.utcnow()
now = datetime.datetime.now(tz=datetime.UTC)
registration = ComputeServiceRegistration(
identifier=compute_service_id,
registered=now,
Expand Down
6 changes: 3 additions & 3 deletions alchemiscale/tests/integration/storage/utils.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from datetime import datetime
import datetime

from gufe.protocols import ProtocolUnitFailure

Expand Down Expand Up @@ -58,7 +58,7 @@ def complete_tasks(
for task in tasks:
ok_pdrr = ProtocolDAGResultRef(
ok=True,
datetime_created=datetime.utcnow(),
datetime_created=datetime.datetime.now(tz=datetime.UTC),
obj_key=task.gufe_key,
scope=task.scope,
)
Expand All @@ -78,7 +78,7 @@ def fail_task(

not_ok_pdrr = ProtocolDAGResultRef(
ok=False,
datetime_created=datetime.utcnow(),
datetime_created=datetime.datetime.now(tz=datetime.UTC),
obj_key=task.gufe_key,
scope=task.scope,
)
Expand Down
Loading