Add ActorRun model

janbuchar · janbuchar · commit eb4284b5177d · 2024-08-28T17:41:00.000+02:00
diff --git a/src/apify/_actor.py b/src/apify/_actor.py
@@ -19,6 +19,7 @@
 from apify._configuration import Configuration
 from apify._consts import EVENT_LISTENERS_TIMEOUT
 from apify._crypto import decrypt_input_secrets, load_private_key
+from apify._models import ActorRun
 from apify._platform_event_manager import EventManager, LocalEventManager, PlatformEventManager
 from apify._proxy_configuration import ProxyConfiguration
 from apify._utils import get_system_info, is_running_in_ipython
@@ -536,7 +537,7 @@ async def start(
         timeout: timedelta | None = None,
         wait_for_finish: int | None = None,
         webhooks: list[Webhook] | None = None,
-    ) -> dict:
+    ) -> ActorRun:
         """Run an Actor on the Apify platform.
 
         Unlike `Actor.call`, this method just starts the run without waiting for finish.
@@ -565,16 +566,20 @@ async def start(
 
         client = self.new_client(token=token) if token else self._apify_client
 
-        return await client.actor(actor_id).start(
-            run_input=run_input,
-            content_type=content_type,
-            build=build,
-            memory_mbytes=memory_mbytes,
-            timeout_secs=int(timeout.total_seconds()) if timeout is not None else None,
-            wait_for_finish=wait_for_finish,
-            webhooks=[hook.model_dump(by_alias=True, exclude_unset=True, exclude_defaults=True) for hook in webhooks]
-            if webhooks
-            else None,
+        return ActorRun.model_validate(
+            await client.actor(actor_id).start(
+                run_input=run_input,
+                content_type=content_type,
+                build=build,
+                memory_mbytes=memory_mbytes,
+                timeout_secs=int(timeout.total_seconds()) if timeout is not None else None,
+                wait_for_finish=wait_for_finish,
+                webhooks=[
+                    hook.model_dump(by_alias=True, exclude_unset=True, exclude_defaults=True) for hook in webhooks
+                ]
+                if webhooks
+                else None,
+            )
         )
 
     async def abort(
@@ -584,7 +589,7 @@ async def abort(
         token: str | None = None,
         status_message: str | None = None,
         gracefully: bool | None = None,
-    ) -> dict:
+    ) -> ActorRun:
         """Abort given Actor run on the Apify platform using the current user account.
 
         The user account is determined by the `APIFY_TOKEN` environment variable.
@@ -607,7 +612,7 @@ async def abort(
         if status_message:
             await client.run(run_id).update(status_message=status_message)
 
-        return await client.run(run_id).abort(gracefully=gracefully)
+        return ActorRun.model_validate(await client.run(run_id).abort(gracefully=gracefully))
 
     async def call(
         self,
@@ -621,7 +626,7 @@ async def call(
         timeout: timedelta | None = None,
         webhooks: list[Webhook] | None = None,
         wait: timedelta | None = None,
-    ) -> dict | None:
+    ) -> ActorRun | None:
         """Start an Actor on the Apify Platform and wait for it to finish before returning.
 
         It waits indefinitely, unless the wait argument is provided.
@@ -650,16 +655,20 @@ async def call(
 
         client = self.new_client(token=token) if token else self._apify_client
 
-        return await client.actor(actor_id).call(
-            run_input=run_input,
-            content_type=content_type,
-            build=build,
-            memory_mbytes=memory_mbytes,
-            timeout_secs=int(timeout.total_seconds()) if timeout is not None else None,
-            webhooks=[hook.model_dump(by_alias=True, exclude_defaults=True, exclude_unset=True) for hook in webhooks]
-            if webhooks
-            else [],
-            wait_secs=int(wait.total_seconds()) if wait is not None else None,
+        return ActorRun.model_validate(
+            await client.actor(actor_id).call(
+                run_input=run_input,
+                content_type=content_type,
+                build=build,
+                memory_mbytes=memory_mbytes,
+                timeout_secs=int(timeout.total_seconds()) if timeout is not None else None,
+                webhooks=[
+                    hook.model_dump(by_alias=True, exclude_defaults=True, exclude_unset=True) for hook in webhooks
+                ]
+                if webhooks
+                else [],
+                wait_secs=int(wait.total_seconds()) if wait is not None else None,
+            )
         )
 
     async def call_task(
@@ -673,7 +682,7 @@ async def call_task(
         webhooks: list[Webhook] | None = None,
         wait: timedelta | None = None,
         token: str | None = None,
-    ) -> dict | None:
+    ) -> ActorRun | None:
         """Start an Actor task on the Apify Platform and wait for it to finish before returning.
 
         It waits indefinitely, unless the wait argument is provided.
@@ -705,15 +714,19 @@ async def call_task(
 
         client = self.new_client(token=token) if token else self._apify_client
 
-        return await client.task(task_id).call(
-            task_input=task_input,
-            build=build,
-            memory_mbytes=memory_mbytes,
-            timeout_secs=int(timeout.total_seconds()) if timeout is not None else None,
-            webhooks=[hook.model_dump(by_alias=True, exclude_defaults=True, exclude_unset=True) for hook in webhooks]
-            if webhooks
-            else [],
-            wait_secs=int(wait.total_seconds()) if wait is not None else None,
+        return ActorRun.model_validate(
+            await client.task(task_id).call(
+                task_input=task_input,
+                build=build,
+                memory_mbytes=memory_mbytes,
+                timeout_secs=int(timeout.total_seconds()) if timeout is not None else None,
+                webhooks=[
+                    hook.model_dump(by_alias=True, exclude_defaults=True, exclude_unset=True) for hook in webhooks
+                ]
+                if webhooks
+                else [],
+                wait_secs=int(wait.total_seconds()) if wait is not None else None,
+            )
         )
 
     async def metamorph(
diff --git a/src/apify/_models.py b/src/apify/_models.py
@@ -1,11 +1,13 @@
 # ruff: noqa: TCH001 TCH002 TCH003 (Pydantic)
 from __future__ import annotations
 
+from datetime import datetime, timedelta
 from typing import Annotated
 
 from pydantic import BaseModel, BeforeValidator, ConfigDict, Field
 
-from apify_shared.consts import WebhookEventType
+from apify_shared.consts import ActorJobStatus, MetaOrigin, WebhookEventType
+from crawlee._utils.models import timedelta_ms
 from crawlee._utils.urls import validate_http_url
 
 
@@ -25,3 +27,82 @@ class Webhook(BaseModel):
         str | None,
         Field(description='Template for the payload sent by the webook'),
     ] = None
+
+
+class ActorRunMeta(BaseModel):
+    __model_config__ = ConfigDict(populate_by_name=True)
+
+    origin: Annotated[MetaOrigin, Field()]
+
+
+class ActorRunStats(BaseModel):
+    __model_config__ = ConfigDict(populate_by_name=True)
+
+    input_body_len: Annotated[int, Field(alias='inputBodyLen')]
+    restart_count: Annotated[int, Field(alias='restartCount')]
+    resurrect_count: Annotated[int, Field(alias='resurrectCount')]
+    mem_avg_bytes: Annotated[float, Field(alias='memAvgBytes')]
+    mem_max_bytes: Annotated[int, Field(alias='memMaxBytes')]
+    mem_current_bytes: Annotated[int, Field(alias='memCurrentBytes')]
+    cpu_avg_usage: Annotated[float, Field(alias='cpuAvgUsage')]
+    cpu_max_usage: Annotated[float, Field(alias='cpuMaxUsage')]
+    cpu_current_usage: Annotated[float, Field(alias='cpuCurrentUsage')]
+    net_rx_bytes: Annotated[int, Field(alias='netRxBytes')]
+    net_tx_bytes: Annotated[int, Field(alias='netTxBytes')]
+    duration: Annotated[timedelta_ms, Field(alias='durationMillis')]
+    run_time: Annotated[timedelta, Field(alias='runTimeSecs')]
+    metamorph: Annotated[int, Field(alias='metamorph')]
+    compute_units: Annotated[float, Field(alias='computeUnits')]
+
+
+class ActorRunOptions(BaseModel):
+    __model_config__ = ConfigDict(populate_by_name=True)
+    build: str
+    timeout: Annotated[timedelta, Field(alias='timeoutSecs')]
+    memory_mbytes: Annotated[int, Field(alias='memoryMbytes')]
+    disk_mbytes: Annotated[int, Field(alias='diskMbytes')]
+
+
+class ActorRunUsage(BaseModel):
+    __model_config__ = ConfigDict(populate_by_name=True)
+    actor_compute_units: Annotated[int | None, Field(alias='ACTOR_COMPUTE_UNITS')] = None
+    dataset_reads: Annotated[int | None, Field(alias='DATASET_READS')] = None
+    dataset_writes: Annotated[int | None, Field(alias='DATASET_WRITES')] = None
+    key_value_store_reads: Annotated[int | None, Field(alias='KEY_VALUE_STORE_READS')] = None
+    key_value_store_writes: Annotated[int | None, Field(alias='KEY_VALUE_STORE_WRITES')] = None
+    key_value_store_lists: Annotated[int | None, Field(alias='KEY_VALUE_STORE_LISTS')] = None
+    request_queue_reads: Annotated[int | None, Field(alias='REQUEST_QUEUE_READS')] = None
+    request_queue_writes: Annotated[int | None, Field(alias='REQUEST_QUEUE_WRITES')] = None
+    data_transfer_internal_gbytes: Annotated[int | None, Field(alias='DATA_TRANSFER_INTERNAL_GBYTES')] = None
+    data_transfer_external_gbytes: Annotated[int | None, Field(alias='DATA_TRANSFER_EXTERNAL_GBYTES')] = None
+    proxy_residential_transfer_gbytes: Annotated[int | None, Field(alias='PROXY_RESIDENTIAL_TRANSFER_GBYTES')] = None
+    proxy_serps: Annotated[int | None, Field(alias='PROXY_SERPS')] = None
+
+
+class ActorRun(BaseModel):
+    __model_config__ = ConfigDict(populate_by_name=True)
+
+    id: Annotated[str, Field(alias='id')]
+    act_id: Annotated[str, Field(alias='actId')]
+    user_id: Annotated[str, Field(alias='userId')]
+    actor_task_id: Annotated[str | None, Field(alias='actorTaskId')] = None
+    started_at: Annotated[datetime, Field(alias='startedAt')]
+    finished_at: Annotated[datetime | None, Field(alias='finishedAt')] = None
+    status: Annotated[ActorJobStatus, Field(alias='status')]
+    status_message: Annotated[str | None, Field(alias='statusMessage')] = None
+    is_status_message_terminal: Annotated[bool, Field(alias='isStatusMessageTerminal')] = False
+    meta: Annotated[ActorRunMeta, Field(alias='meta')]
+    stats: Annotated[ActorRunStats, Field(alias='stats')]
+    options: Annotated[ActorRunOptions, Field(alias='options')]
+    build_id: Annotated[str, Field(alias='buildId')]
+    exit_code: Annotated[int | None, Field(alias='exitCode')] = None
+    default_key_value_store_id: Annotated[str, Field(alias='defaultKeyValueStoreId')]
+    default_dataset_id: Annotated[str, Field(alias='defaultDatasetId')]
+    default_request_queue_id: Annotated[str, Field(alias='defaultRequestQueueId')]
+    build_number: Annotated[str, Field(alias='buildNumber')]
+    container_url: Annotated[str, Field(alias='containerUrl')]
+    is_container_server_ready: Annotated[bool, Field(alias='isContainerServerReady')] = False
+    git_branch_name: Annotated[str | None, Field(alias='gitBranchName')] = None
+    usage: Annotated[ActorRunUsage | None, Field(alias='usage')] = None
+    usage_total_usd: Annotated[float | None, Field(alias='usageTotalUsd')] = None
+    usage_usd: Annotated[float | None, Field(alias='usageUsd')] = None
