removed outdated prefect class

als-computing · Jan 18, 2024 · 6016ef1 · 6016ef1
1 parent f5afdfa
commit 6016ef1
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/orchestration/prefect.py b/orchestration/prefect.py
@@ -3,9 +3,9 @@
 import logging
 
 from prefect import get_run_logger, task
-from prefect.client import get_client
-from prefect.orion.utilities.schemas import DateTimeTZ
-from prefect.orion.schemas.states import Scheduled
+from prefect import get_client
+
+from prefect.states import Scheduled
 
 
 logger = logging.getLogger("orchestration.prefect")
@@ -23,7 +23,7 @@ async def schedule(
         assert (
             deployment
         ), f"No deployment found in config for deploymnent_name {deploymnent_name}"
-        date_time_tz = DateTimeTZ.now() + duration_from_now
+        date_time_tz = datetime.datetime.now() + duration_from_now
         await client.create_flow_run_from_deployment(
             deployment.id,
             state=Scheduled(scheduled_time=date_time_tz),