Skip to content

Commit

Permalink
chore: update python api to new protos
Browse files Browse the repository at this point in the history
move python argument in dataclass to be in the same order than in the protos
  • Loading branch information
ngruelaneo committed May 2, 2024
1 parent 4615295 commit 2eb5370
Show file tree
Hide file tree
Showing 4 changed files with 104 additions and 67 deletions.
82 changes: 45 additions & 37 deletions packages/python/src/armonik/common/enumwrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,17 @@
from enum import IntEnum

from ..protogen.common.task_status_pb2 import (
TASK_STATUS_CANCELLED,
TASK_STATUS_CANCELLING,
TASK_STATUS_COMPLETED,
TASK_STATUS_UNSPECIFIED,
TASK_STATUS_CREATING,
TASK_STATUS_SUBMITTED,
TASK_STATUS_DISPATCHED,
TASK_STATUS_COMPLETED,
TASK_STATUS_ERROR,
TASK_STATUS_PROCESSED,
TASK_STATUS_PROCESSING,
TASK_STATUS_SUBMITTED,
TASK_STATUS_TIMEOUT,
TASK_STATUS_UNSPECIFIED,
TASK_STATUS_CANCELLING,
TASK_STATUS_CANCELLED,
TASK_STATUS_PROCESSING,
TASK_STATUS_PROCESSED,
TASK_STATUS_RETRIED,
)
from ..protogen.common.events_common_pb2 import (
Expand All @@ -27,8 +27,12 @@
SessionStatus as RawSessionStatus,
_SESSIONSTATUS,
SESSION_STATUS_UNSPECIFIED,
SESSION_STATUS_CANCELLED,
SESSION_STATUS_RUNNING,
SESSION_STATUS_CANCELLED,
SESSION_STATUS_PAUSED,
SESSION_STATUS_CLOSED,
SESSION_STATUS_PURGED,
SESSION_STATUS_DELETED,
)
from ..protogen.common.result_status_pb2 import (
ResultStatus as RawResultStatus,
Expand All @@ -52,30 +56,32 @@
# This file is necessary because the grpc types aren't considered proper types


class HealthCheckStatus:
NOT_SERVING = HealthCheckReply.NOT_SERVING
SERVING = HealthCheckReply.SERVING
UNKNOWN = HealthCheckReply.UNKNOWN


class TaskStatus(IntEnum):
CANCELLED = TASK_STATUS_CANCELLED
CANCELLING = TASK_STATUS_CANCELLING
COMPLETED = TASK_STATUS_COMPLETED
UNSPECIFIED = TASK_STATUS_UNSPECIFIED
CREATING = TASK_STATUS_CREATING
SUBMITTED = TASK_STATUS_SUBMITTED
DISPATCHED = TASK_STATUS_DISPATCHED
COMPLETED = TASK_STATUS_COMPLETED
ERROR = TASK_STATUS_ERROR
PROCESSED = TASK_STATUS_PROCESSED
TIMEOUT = TASK_STATUS_TIMEOUT
CANCELLING = TASK_STATUS_CANCELLING
CANCELLED = TASK_STATUS_CANCELLED
PROCESSING = TASK_STATUS_PROCESSING
SUBMITTED = TASK_STATUS_SUBMITTED
PROCESSED = TASK_STATUS_PROCESSED
RETRIED = TASK_STATUS_RETRIED
TIMEOUT = TASK_STATUS_TIMEOUT
UNSPECIFIED = TASK_STATUS_UNSPECIFIED


class Direction:
ASC = SORT_DIRECTION_ASC
DESC = SORT_DIRECTION_DESC
class EventTypes:
UNSPECIFIED = EVENTS_ENUM_UNSPECIFIED
NEW_TASK = EVENTS_ENUM_NEW_TASK
TASK_STATUS_UPDATE = EVENTS_ENUM_TASK_STATUS_UPDATE
NEW_RESULT = EVENTS_ENUM_NEW_RESULT
RESULT_STATUS_UPDATE = EVENTS_ENUM_RESULT_STATUS_UPDATE
RESULT_OWNER_UPDATE = EVENTS_ENUM_RESULT_OWNER_UPDATE

@classmethod
def from_string(cls, name: str):
return getattr(cls, name.upper())


class SessionStatus:
Expand All @@ -86,6 +92,10 @@ def name_from_value(status: RawSessionStatus) -> str:
UNSPECIFIED = SESSION_STATUS_UNSPECIFIED
RUNNING = SESSION_STATUS_RUNNING
CANCELLED = SESSION_STATUS_CANCELLED
PAUSED = SESSION_STATUS_PAUSED
CLOSED = SESSION_STATUS_CLOSED
PURGED = SESSION_STATUS_PURGED
DELETED = SESSION_STATUS_DELETED


class ResultStatus:
Expand All @@ -101,21 +111,19 @@ def name_from_value(status: RawResultStatus) -> str:
NOTFOUND = RESULT_STATUS_NOTFOUND


class EventTypes:
UNSPECIFIED = EVENTS_ENUM_UNSPECIFIED
NEW_TASK = EVENTS_ENUM_NEW_TASK
TASK_STATUS_UPDATE = EVENTS_ENUM_TASK_STATUS_UPDATE
NEW_RESULT = EVENTS_ENUM_NEW_RESULT
RESULT_STATUS_UPDATE = EVENTS_ENUM_RESULT_STATUS_UPDATE
RESULT_OWNER_UPDATE = EVENTS_ENUM_RESULT_OWNER_UPDATE

@classmethod
def from_string(cls, name: str):
return getattr(cls, name.upper())


class ServiceHealthCheckStatus:
UNSPECIFIED = HEALTH_STATUS_ENUM_UNSPECIFIED
HEALTHY = HEALTH_STATUS_ENUM_HEALTHY
DEGRADED = HEALTH_STATUS_ENUM_DEGRADED
UNHEALTHY = HEALTH_STATUS_ENUM_UNHEALTHY


class HealthCheckStatus:
UNKNOWN = HealthCheckReply.UNKNOWN
SERVING = HealthCheckReply.SERVING
NOT_SERVING = HealthCheckReply.NOT_SERVING


class Direction:
ASC = SORT_DIRECTION_ASC
DESC = SORT_DIRECTION_DESC
54 changes: 27 additions & 27 deletions packages/python/src/armonik/common/filter.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,41 +9,41 @@
from google.protobuf.message import Message

from ..protogen.common.filters_common_pb2 import (
FILTER_STRING_OPERATOR_EQUAL,
FILTER_STRING_OPERATOR_NOT_EQUAL,
FILTER_STRING_OPERATOR_CONTAINS,
FILTER_STRING_OPERATOR_NOT_CONTAINS,
FILTER_STRING_OPERATOR_STARTS_WITH,
FILTER_STRING_OPERATOR_ENDS_WITH,
FilterString,
FILTER_NUMBER_OPERATOR_EQUAL,
FILTER_NUMBER_OPERATOR_NOT_EQUAL,
FILTER_NUMBER_OPERATOR_LESS_THAN,
FILTER_NUMBER_OPERATOR_LESS_THAN_OR_EQUAL,
FILTER_NUMBER_OPERATOR_GREATER_THAN_OR_EQUAL,
FILTER_NUMBER_OPERATOR_GREATER_THAN,
FilterNumber,
FILTER_DATE_OPERATOR_EQUAL,
FILTER_DATE_OPERATOR_NOT_EQUAL,
FILTER_DATE_OPERATOR_BEFORE,
FILTER_DATE_OPERATOR_BEFORE_OR_EQUAL,
FILTER_DATE_OPERATOR_AFTER,
FILTER_DATE_OPERATOR_AFTER_OR_EQUAL,
FilterDate,
FILTER_ARRAY_OPERATOR_CONTAINS,
FILTER_ARRAY_OPERATOR_NOT_CONTAINS,
FilterArray,
FILTER_STATUS_OPERATOR_EQUAL,
FILTER_STATUS_OPERATOR_NOT_EQUAL,
FILTER_BOOLEAN_OPERATOR_IS,
FILTER_DATE_OPERATOR_AFTER,
FILTER_DATE_OPERATOR_AFTER_OR_EQUAL,
FILTER_DATE_OPERATOR_BEFORE,
FILTER_DATE_OPERATOR_BEFORE_OR_EQUAL,
FILTER_DATE_OPERATOR_EQUAL,
FILTER_DATE_OPERATOR_NOT_EQUAL,
FilterBoolean,
FILTER_DURATION_OPERATOR_EQUAL,
FILTER_DURATION_OPERATOR_LONGER_THAN,
FILTER_DURATION_OPERATOR_LONGER_THAN_OR_EQUAL,
FILTER_DURATION_OPERATOR_NOT_EQUAL,
FILTER_DURATION_OPERATOR_SHORTER_THAN,
FILTER_DURATION_OPERATOR_SHORTER_THAN_OR_EQUAL,
FILTER_NUMBER_OPERATOR_EQUAL,
FILTER_NUMBER_OPERATOR_GREATER_THAN,
FILTER_NUMBER_OPERATOR_GREATER_THAN_OR_EQUAL,
FILTER_NUMBER_OPERATOR_LESS_THAN,
FILTER_NUMBER_OPERATOR_LESS_THAN_OR_EQUAL,
FILTER_NUMBER_OPERATOR_NOT_EQUAL,
FILTER_STATUS_OPERATOR_EQUAL,
FILTER_STATUS_OPERATOR_NOT_EQUAL,
FILTER_STRING_OPERATOR_CONTAINS,
FILTER_STRING_OPERATOR_ENDS_WITH,
FILTER_STRING_OPERATOR_EQUAL,
FILTER_STRING_OPERATOR_NOT_CONTAINS,
FILTER_STRING_OPERATOR_NOT_EQUAL,
FILTER_STRING_OPERATOR_STARTS_WITH,
FilterArray,
FilterBoolean,
FilterDate,
FILTER_DURATION_OPERATOR_LONGER_THAN,
FILTER_DURATION_OPERATOR_LONGER_THAN_OR_EQUAL,
FilterDuration,
FilterNumber,
FilterString,
)


Expand Down
30 changes: 28 additions & 2 deletions packages/python/src/armonik/common/objects.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
from datetime import datetime, timedelta
from typing import Dict, List, Optional

from deprecation import deprecated

from ..protogen.common.objects_pb2 import (
Empty,
)
Expand Down Expand Up @@ -99,13 +101,16 @@ class Task:
id: Optional[str] = None
session_id: Optional[str] = None
owner_pod_id: Optional[str] = None

initial_task_id: Optional[str] = None
parent_task_ids: List[str] = field(default_factory=list)
data_dependencies: List[str] = field(default_factory=list)
expected_output_ids: List[str] = field(default_factory=list)
retry_of_ids: List[str] = field(default_factory=list)

status: RawTaskStatus = TaskStatus.UNSPECIFIED
payload_id: Optional[str] = None
status_message: Optional[str] = None

options: Optional[TaskOptions] = None
created_at: Optional[datetime] = None
submitted_at: Optional[datetime] = None
Expand All @@ -116,7 +121,13 @@ class Task:
processed_at: Optional[datetime] = None
ended_at: Optional[datetime] = None
pod_ttl: Optional[datetime] = None

creation_to_end_duration: Optional[timedelta] = timedelta(0)
processing_to_end_duration: Optional[timedelta] = timedelta(0)
received_to_end_duration: Optional[timedelta] = timedelta(0)

output: Optional[Output] = None

pod_hostname: Optional[str] = None

def refresh(self, task_client) -> None:
Expand All @@ -128,12 +139,16 @@ def refresh(self, task_client) -> None:
result: "Task" = task_client.get_task(self.id)
self.session_id = result.session_id
self.owner_pod_id = result.owner_pod_id

self.initial_task_id = result.initial_task_id
self.parent_task_ids = result.parent_task_ids
self.data_dependencies = result.data_dependencies
self.expected_output_ids = result.expected_output_ids
self.retry_of_ids = result.retry_of_ids

self.status = result.status
self.status_message = result.status_message

self.options = result.options
self.created_at = result.created_at
self.submitted_at = result.submitted_at
Expand All @@ -144,7 +159,13 @@ def refresh(self, task_client) -> None:
self.processed_at = result.processed_at
self.ended_at = result.ended_at
self.pod_ttl = result.pod_ttl

self.creation_to_end_duration = result.creation_to_end_duration
self.processing_to_end_duration = result.processing_to_end_duration
self.received_to_end_duration = result.received_to_end_duration

self.output = result.output

self.pod_hostname = result.pod_hostname
self.is_init = True

Expand All @@ -154,6 +175,7 @@ def from_message(cls, task_raw: TaskDetailed) -> "Task":
id=task_raw.id,
session_id=task_raw.session_id,
owner_pod_id=task_raw.owner_pod_id,
initial_task_id=task_raw.initial_task_id,
parent_task_ids=list(task_raw.parent_task_ids),
data_dependencies=list(task_raw.data_dependencies),
expected_output_ids=list(task_raw.expected_output_ids),
Expand All @@ -170,12 +192,16 @@ def from_message(cls, task_raw: TaskDetailed) -> "Task":
processed_at=timestamp_to_datetime(task_raw.processed_at),
ended_at=timestamp_to_datetime(task_raw.ended_at),
pod_ttl=timestamp_to_datetime(task_raw.pod_ttl),
creation_to_end_duration=duration_to_timedelta(task_raw.creation_to_end_duration),
processing_to_end_duration=duration_to_timedelta(task_raw.processing_to_end_duration),
received_to_end_duration=duration_to_timedelta(task_raw.received_to_end_duration),
output=Output(error=(task_raw.output.error if not task_raw.output.success else None)),
pod_hostname=task_raw.pod_hostname,
)


@dataclass()
@deprecated(deprecated_in="3.14.0", details="Use sessions, task and results client instead")
@dataclass
class ResultAvailability:
errors: List[str] = field(default_factory=list)

Expand Down
5 changes: 4 additions & 1 deletion packages/python/tests/test_tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,12 @@ class TestArmoniKTasks:
id="task-id",
session_id="session-id",
owner_pod_id="",
initial_task_id="",
parent_task_ids=[],
data_dependencies=[],
expected_output_ids=[],
retry_of_ids=[],
status=4,
payload_id=None,
status_message="",
options=TaskOptions(
max_duration=datetime.timedelta(seconds=1),
Expand All @@ -38,6 +38,9 @@ class TestArmoniKTasks:
processed_at=None,
ended_at=None,
pod_ttl=None,
creation_to_end_duration=timedelta(0),
processing_to_end_duration=timedelta(0),
received_to_end_duration=timedelta(0),
output=Output(error=""),
pod_hostname="",
)
Expand Down

0 comments on commit 2eb5370

Please sign in to comment.