truefoundry 0.5.9__py3-none-any.whl → 0.5.10rc2__py3-none-any.whl

truefoundry/cli/display_util.py

@@ -38,7 +38,7 @@ def stringify(x):
         return str(x)
 
 
-def display_time_passed(seconds: int):
+def display_time_passed(seconds: float) -> str:
     m, s = divmod(seconds, 60)
     h, m = divmod(m, 60)
     d, h = divmod(h, 24)

truefoundry/deploy/auto_gen/models.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  application.json
-#   timestamp: 2025-02-19T09:48:05+00:00
+#   timestamp: 2025-02-26T17:51:02+00:00
 
 from __future__ import annotations
 

truefoundry/deploy/io/output_callback.py

@@ -1,23 +1,30 @@
+from typing import Any, List, Optional
+
+
 class OutputCallBack:
-    def print_header(self, line):
+    def print_header(self, line: Any) -> None:
         print(line)
 
-    def _print_separator(self):
+    def _print_separator(self) -> None:
         print("-" * 80)
 
-    def print_line(self, line):
+    def print_line(self, line: str) -> None:
         print(line)
 
-    def print_lines_in_panel(self, lines, header=None):
+    def print_lines_in_panel(
+        self, lines: List[str], header: Optional[str] = None
+    ) -> None:
         self.print_header(header)
         self._print_separator()
         for line in lines:
             self.print_line(line)
         self._print_separator()
 
-    def print_code_in_panel(self, lines, header=None, lang="python"):
+    def print_code_in_panel(
+        self, lines: List[str], header: Optional[str] = None
+    ) -> None:
         self.print_lines_in_panel(lines, header)
 
-    def print(self, line):
+    def print(self, line: Any) -> None:
         # just an alias
         self.print_line(line)

truefoundry/deploy/io/rich_output_callback.py

@@ -1,3 +1,5 @@
+from typing import Any, List, Optional
+
 from rich.console import Console
 from rich.highlighter import ReprHighlighter
 from rich.panel import Panel
@@ -14,14 +16,16 @@ class RichOutputCallBack(OutputCallBack):
     console = Console(soft_wrap=True)
     highlighter = ReprHighlighter()
 
-    def print_header(self, line):
+    def print_header(self, line: Any) -> None:
         self.console.rule(_text(line), style="cyan")
 
-    def print_line(self, line):
+    def print_line(self, line: Any) -> None:
         self.console.print(_text(line))
 
-    def print_lines_in_panel(self, lines, header=None):
+    def print_lines_in_panel(
+        self, lines: List[str], header: Optional[str] = None
+    ) -> None:
         self.console.print(Panel(self.highlighter("\n".join(lines)), title=header))
 
-    def print_code_in_panel(self, lines, header=None, lang="python"):
+    def print_code_in_panel(self, lines: List[str], header: Optional[str] = None):
         self.print_lines_in_panel(lines, header)

truefoundry/deploy/lib/clients/servicefoundry_client.py

@@ -4,7 +4,7 @@ import json
 import os
 import time
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, cast
 from urllib.parse import urljoin
 
 import requests
@@ -32,6 +32,8 @@ from truefoundry.deploy.lib.model.entity import (
     Deployment,
     DockerRegistryCredentials,
     JobRun,
+    LogBody,
+    SocketEvent,
     TriggerJobResult,
     Workspace,
     WorkspaceResources,
@@ -51,7 +53,7 @@ BUILD_LOGS_SUBSCRIBE_MESSAGE = "BUILD_LOGS"
 MAX_RETRIES_WORKFLOW_TRIGGER = 3
 
 
-def _upload_packaged_code(metadata, package_file):
+def _upload_packaged_code(metadata: Dict[str, Any], package_file: str) -> None:
     file_size = os.stat(package_file).st_size
     with open(package_file, "rb") as file_to_upload:
         with tqdm(
@@ -74,7 +76,9 @@ def _upload_packaged_code(metadata, package_file):
 
 
 class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
-    def __init__(self, init_session: bool = True, tfy_host: Optional[str] = None):
+    def __init__(
+        self, init_session: bool = True, tfy_host: Optional[str] = None
+    ) -> None:
         self._session: Optional[Session] = None
         if init_session:
             if tfy_host:
@@ -85,25 +89,19 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             raise Exception("Neither session, not tfy_host provided")
         super().__init__(tfy_host=tfy_host)
 
-    def _get_header(self):
+    def _get_header(self) -> Dict[str, str]:
         if not self._session:
             return {}
         return {"Authorization": f"Bearer {self._session.access_token}"}
 
     @check_min_cli_version
-    def get_id_from_fqn(self, fqn_type: str, fqn: str):
+    def get_id_from_fqn(self, fqn_type: str, fqn: str) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/fqn/{fqn_type}"
         response = session_with_retries().get(
             url, headers=self._get_header(), params={"fqn": fqn}
         )
         return request_handling(response)
 
-    @check_min_cli_version
-    def list_workspace(self):
-        url = f"{self._api_server_url}/{VERSION_PREFIX}/workspace"
-        response = session_with_retries().get(url, headers=self._get_header())
-        return request_handling(response)
-
     @check_min_cli_version
     def list_workspaces(
         self,
@@ -148,33 +146,17 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         return Workspace.parse_obj(response_data)
 
     @check_min_cli_version
-    def remove_workspace(self, workspace_id, force=False) -> Workspace:
+    def remove_workspace(self, workspace_id: str, force: bool = False) -> Workspace:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/workspace/{workspace_id}"
-        force = json.dumps(
+        force_str = json.dumps(
             force
         )  # this dumb conversion is required because `params` just casts as str
         response = session_with_retries().delete(
-            url, headers=self._get_header(), params={"force": force}
+            url, headers=self._get_header(), params={"force": force_str}
         )
-        response = request_handling(response)
+        response = cast(Dict[str, Any], request_handling(response))
         return Workspace.parse_obj(response["workspace"])
 
-    @check_min_cli_version
-    def get_workspace_by_name(self, workspace_name, cluster_id):
-        url = f"{self._api_server_url}/{VERSION_PREFIX}/workspace"
-        response = session_with_retries().get(
-            url,
-            headers=self._get_header(),
-            params={"name": workspace_name, "clusterId": cluster_id},
-        )
-        return request_handling(response)
-
-    @check_min_cli_version
-    def get_workspace(self, workspace_id):
-        url = f"{self._api_server_url}/{VERSION_PREFIX}/workspace/{workspace_id}"
-        response = session_with_retries().get(url, headers=self._get_header())
-        return request_handling(response)
-
     @check_min_cli_version
     def get_workspace_by_fqn(self, workspace_fqn: str) -> List[Workspace]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/workspace"
@@ -187,30 +169,15 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         return parse_obj_as(List[Workspace], response)
 
     @check_min_cli_version
-    def list_deployments(self, workspace_id: str = None):
-        url = f"{self._api_server_url}/{VERSION_PREFIX}/deployment"
-        params = {}
-        if workspace_id:
-            params["workspaceId"] = workspace_id
-        response = session_with_retries().get(
-            url=url, params=params, headers=self._get_header()
-        )
-        return request_handling(response)
-
-    @check_min_cli_version
-    def list_cluster(self):
-        url = f"{self._api_server_url}/{VERSION_PREFIX}/cluster"
-        response = session_with_retries().get(url, headers=self._get_header())
-        return request_handling(response)
-
-    @check_min_cli_version
-    def get_cluster(self, cluster_id):
+    def get_cluster(self, cluster_id: str) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/cluster/{cluster_id}"
         response = session_with_retries().get(url, headers=self._get_header())
-        return request_handling(response)
+        return cast(Dict[str, Any], request_handling(response))
 
     @check_min_cli_version
-    def get_presigned_url(self, space_name, service_name, env):
+    def get_presigned_url(
+        self, space_name: str, service_name: str, env: str
+    ) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/deployment/code-upload-url"
         response = session_with_retries().post(
             url,
@@ -221,7 +188,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             },
             headers=self._get_header(),
         )
-        return request_handling(response)
+        return cast(Dict[str, Any], request_handling(response))
 
     @check_min_cli_version
     def upload_code_package(
@@ -231,7 +198,6 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             space_name=workspace_fqn, service_name=component_name, env="default"
         )
         _upload_packaged_code(metadata=http_response, package_file=package_local_path)
-
         return http_response["uri"]
 
     @check_min_cli_version
@@ -252,28 +218,28 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         response = session_with_retries().post(
             url, json=data, headers=self._get_header()
         )
-        response_data = request_handling(response)
+        response_data = cast(Dict[str, Any], request_handling(response))
         return Deployment.parse_obj(response_data["deployment"])
 
-    def _get_log_print_line(self, log: dict):
-        timestamp = int(log["time"]) / 1e6
+    def _get_log_print_line(self, log_data: LogBody) -> str:
+        timestamp = int(log_data.time) / 1e6
 
         time_obj = datetime.fromtimestamp(timestamp / 1000.0, tz=timezone.utc)
         time_obj.replace(tzinfo=timezone.utc)
         local_time = time_obj.astimezone(tzlocal())
         local_time_str = local_time.isoformat()
-        return f"[{local_time_str}] {log['log'].strip()}"
+        return f"[{local_time_str}] {log_data.log.strip()}"
 
     def _tail_logs(
         self,
         tail_logs_url: str,
-        query_dict: dict,
+        query_dict: Dict[str, Any],
         # NOTE: Rather making this printer callback an argument,
         # we should have global printer callback
         # which will be initialized based on the running env (cli, lib, notebook)
         subscribe_message: str,
         socketio_path: str = "socket.io",
-        callback=None,
+        callback: Optional[OutputCallBack] = None,
     ) -> socketio.Client:
         callback = callback or OutputCallBack()
         sio = socketio.Client(request_timeout=60)
@@ -281,18 +247,20 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         next_log_start_timestamp = query_dict.get("startTs")
 
         @sio.on(subscribe_message)
-        def logs(data):
+        def logs(data: Any) -> None:
             try:
-                _log = json.loads(data)
-                body = _log["body"]
-                callback.print_line(self._get_log_print_line(body))
+                event = SocketEvent.parse_obj(json.loads(data))
+                if not isinstance(event.body, LogBody):
+                    logger.debug(f"Skipped log for {data!r}")
+                    return
+                callback.print_line(self._get_log_print_line(event.body))
                 nonlocal next_log_start_timestamp
-                next_log_start_timestamp = body["time"]
+                next_log_start_timestamp = event.body.time
             except Exception:
-                logger.exception(f"Error while parsing log line, {data!r}")
+                logger.debug(f"Error while parsing log line, {data!r}")
 
         @sio.on("connect")
-        def on_connect():
+        def on_connect() -> None:
             # TODO: We should have have a timeout here. `emit` does
             #   not support timeout. Explore `sio.call`.
             query_dict["startTs"] = next_log_start_timestamp
@@ -301,7 +269,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                 json.dumps(query_dict),
             )
 
-        def sio_disconnect_no_exception():
+        def sio_disconnect_no_exception() -> None:
             try:
                 sio.disconnect()
             except Exception:
@@ -351,11 +319,11 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         end_ts_nano: Optional[int] = None,
         limit: Optional[int] = None,
         num_logs_to_ignore: Optional[int] = None,
-    ) -> List:
+    ) -> List[LogBody]:
         get_logs_query = {"applicationId": application_id}
         if deployment_id:
             get_logs_query["deploymentId"] = deployment_id
-        data = {"getLogsQuery": json.dumps(get_logs_query)}
+        data: Dict[str, Any] = {"getLogsQuery": json.dumps(get_logs_query)}
         if start_ts_nano:
             data["startTs"] = str(start_ts_nano)
         if end_ts_nano:
@@ -371,14 +339,14 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         response = session_with_retries().get(
             url=url, params=data, headers=self._get_header()
         )
-        response_data = request_handling(response)
-        return list(response_data["logs"])
+        response_data = cast(Dict[str, Any], request_handling(response))
+        return parse_obj_as(List[LogBody], response_data["data"])
 
     @check_min_cli_version
     def tail_build_logs(
         self,
         build_response: BuildResponse,
-        callback=None,
+        callback: Optional[OutputCallBack] = None,
     ) -> socketio.Client:
         callback = callback or OutputCallBack()
         tail_logs_obj = json.loads(build_response.tailLogsUrl)
@@ -404,9 +372,9 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         deployment_id: str,
         start_ts: int,
         limit: int,
-        callback=None,
+        callback: Optional[OutputCallBack] = None,
         wait: bool = True,
-    ):
+    ) -> None:
         callback = callback or OutputCallBack()
         self._tail_logs(
             tail_logs_url=urljoin(
@@ -435,8 +403,8 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         start_ts: int,
         limit: int,
         poll_interval_seconds: int,
-        callback=None,
-    ):
+        callback: Optional[OutputCallBack] = None,
+    ) -> None:
         callback = callback or OutputCallBack()
         start_ts_nano = int(start_ts * 1e6)
 
@@ -444,7 +412,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             num_logs_to_ignore = 0
 
             while True:
-                logs_list = self._get_deployment_logs(
+                logs = self._get_deployment_logs(
                     workspace_id=workspace_id,
                     application_id=application_id,
                     deployment_id=deployment_id,
@@ -454,18 +422,18 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                     num_logs_to_ignore=num_logs_to_ignore,
                 )
 
-                if len(logs_list) == 0:
+                if not logs:
                     logger.warning("Did not receive any logs")
                     time.sleep(poll_interval_seconds)
                     continue
 
-                for log in logs_list:
+                for log in logs:
                     callback.print_line(self._get_log_print_line(log))
 
-                last_log_time = logs_list[-1]["time"]
+                last_log_time = logs[-1].time
                 num_logs_to_ignore = 0
-                for log in reversed(logs_list):
-                    if log["time"] != last_log_time:
+                for log in reversed(logs):
+                    if log.time != last_log_time:
                         break
                     num_logs_to_ignore += 1
 
@@ -483,34 +451,37 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         start_ts: Optional[int],
         end_ts: Optional[int],
         limit: Optional[int],
-        callback=None,
-    ):
+        callback: Optional[OutputCallBack] = None,
+    ) -> None:
         callback = callback or OutputCallBack()
-        logs_list = self._get_deployment_logs(
+        logs = self._get_deployment_logs(
             workspace_id=workspace_id,
             application_id=application_id,
             deployment_id=deployment_id,
             job_run_name=job_run_name,
-            start_ts_nano=int(start_ts * 1e6),
-            end_ts_nano=int(end_ts * 1e6),
+            start_ts_nano=int(start_ts * 1e6) if start_ts else None,
+            end_ts_nano=int(end_ts * 1e6) if end_ts else None,
             limit=limit,
         )
-        for log in logs_list:
+        for log in logs:
             callback.print_line(self._get_log_print_line(log))
 
     @check_min_cli_version
     def fetch_build_logs(
         self,
         build_response: BuildResponse,
-        callback=None,
+        callback: Optional[OutputCallBack] = None,
     ) -> None:
         callback = callback or OutputCallBack()
         url = build_response.getLogsUrl
         response = session_with_retries().get(url=url, headers=self._get_header())
-        logs_list = request_handling(response)
-        for log in logs_list["logs"]:
-            # TODO: Have to establish a log line format that includes timestamp, level, message
-            callback.print_line(self._get_log_print_line(log))
+        logs_response = cast(Dict[str, Any], request_handling(response))
+        for _log_body in logs_response["logs"]:
+            try:
+                log_body = LogBody.parse_obj(_log_body)
+                callback.print_line(self._get_log_print_line(log_body))
+            except Exception:
+                logger.debug(f"Failed to parse log body {_log_body}")
 
     @check_min_cli_version
     def get_deployment_info_by_fqn(self, deployment_fqn: str) -> DeploymentFqnResponse:
@@ -533,10 +504,10 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         return ApplicationFqnResponse.parse_obj(response_data)
 
     @check_min_cli_version
-    def remove_application(self, application_id: str):
+    def remove_application(self, application_id: str) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/app/{application_id}"
         response = session_with_retries().delete(url, headers=self._get_header())
-        response = request_handling(response)
+        response = cast(Dict[str, Any], request_handling(response))
         # TODO: Add pydantic here.
         return response
 
@@ -565,14 +536,14 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         response = session_with_retries().get(
             url, headers=self._get_header(), params=params
         )
-        response_data = request_handling(response)
+        response_data = cast(Dict[str, Any], request_handling(response))
         return parse_obj_as(List[JobRun], response_data["data"])
 
     def get_job_run(
         self,
         application_id: str,
         job_run_name: str,
-    ):
+    ) -> JobRun:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/jobs/{application_id}/runs/{job_run_name}"
         response = session_with_retries().get(url, headers=self._get_header())
         response_data = request_handling(response)
@@ -599,13 +570,15 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         response = request_handling(response)
         return TriggerJobResult.parse_obj(response)
 
-    def trigger_workflow(self, application_id: str, inputs: Dict[str, Any]):
+    def trigger_workflow(
+        self, application_id: str, inputs: Dict[str, Any]
+    ) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/workflow/{application_id}/executions"
         body = {"inputs": inputs}
         response = session_with_retries(retries=MAX_RETRIES_WORKFLOW_TRIGGER).post(
             url, json=body, headers=self._get_header()
         )
-        response = request_handling(response)
+        response = cast(Dict[str, Any], request_handling(response))
         return response
 
     @check_min_cli_version
@@ -684,30 +657,30 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         return parse_obj_as(List[Deployment], response)
 
     @check_min_cli_version
-    def apply(self, manifest: Dict[str, Any], dry_run: bool = False):
+    def apply(self, manifest: Dict[str, Any], dry_run: bool = False) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/apply"
         body = {"manifest": manifest, "dryRun": dry_run}
         response = session_with_retries().put(
             url, headers=self._get_header(), json=body
         )
-        response_data = request_handling(response)
+        response_data = cast(Dict[str, Any], request_handling(response))
         return response_data
 
     @check_min_cli_version
-    def delete(self, manifest: Dict[str, Any]):
+    def delete(self, manifest: Dict[str, Any]) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/delete"
         body = {"manifest": manifest}
         response = session_with_retries().post(
             url, headers=self._get_header(), json=body
         )
-        response_data = request_handling(response)
+        response_data = cast(Dict[str, Any], request_handling(response))
         return response_data
 
     def terminate_job_run(
         self,
         deployment_id: str,
         job_run_name: str,
-    ):
+    ) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/jobs/terminate"
         body = {
             "deploymentId": deployment_id,
@@ -720,5 +693,5 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             json=body,
             headers=self._get_header(),
         )
-        response_data = request_handling(response)
+        response_data = cast(Dict[str, Any], request_handling(response))
         return response_data

truefoundry/deploy/lib/dao/apply.py

@@ -6,7 +6,7 @@ import yaml
 from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,
 )
-from truefoundry.deploy.lib.model.entity import ApplyResult, Manifest
+from truefoundry.deploy.lib.model.entity import ApplyResult, ManifestLike
 from truefoundry.pydantic_v1 import ValidationError
 
 
@@ -26,7 +26,7 @@ def _apply_manifest(
         file_metadata += f" from file {filename}"
 
     try:
-        manifest = Manifest.parse_obj(manifest)
+        manifest = ManifestLike.parse_obj(manifest)
     except ValidationError as ex:
         return ApplyResult(
             success=False,

truefoundry/deploy/lib/dao/delete.py

@@ -6,7 +6,7 @@ import yaml
 from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,
 )
-from truefoundry.deploy.lib.model.entity import DeleteResult, Manifest
+from truefoundry.deploy.lib.model.entity import DeleteResult, ManifestLike
 from truefoundry.pydantic_v1 import ValidationError
 
 
@@ -25,7 +25,7 @@ def _delete_manifest(
         file_metadata += f" from file {filename}"
 
     try:
-        manifest = Manifest.parse_obj(manifest)
+        manifest = ManifestLike.parse_obj(manifest)
     except ValidationError as ex:
         return DeleteResult(
             success=False,

truefoundry/deploy/lib/model/entity.py

@@ -71,25 +71,24 @@ class PortMetadata(BaseModel):
 
 
 class DeploymentTransitionStatus(str, Enum):
-    INITIALIZED: str = "INITIALIZED"
-    BUILDING: str = "BUILDING"
-    DEPLOYING: str = "DEPLOYING"
-    BUILD_SUCCESS: str = "BUILD_SUCCESS"
-    DEPLOY_SUCCESS: str = "DEPLOY_SUCCESS"
-    DEPLOY_FAILED: str = "DEPLOY_FAILED"
-    BUILD_FAILED: str = "BUILD_FAILED"
-    CANCELLED: str = "CANCELLED"
-    COMPONENTS_DEPLOYING: str = "COMPONENTS_DEPLOYING"
-    FAILED: str = "FAILED"
-    REUSING_EXISTING_BUILD: str = "REUSING_EXISTING_BUILD"
-    REDEPLOY_STARTED: str = "REDEPLOY_STARTED"
-    PAUSED: str = "PAUSED"
-    ROLLOUT_STARTED: str = "ROLLOUT_STARTED"
-    SET_TRAFFIC: str = "SET_TRAFFIC"
-    DEPLOY_FAILED_WITH_RETRY: str = "DEPLOY_FAILED_WITH_RETRY"
-    WAITING: str = "WAITING"
-
-    _: str = ""
+    INITIALIZED = "INITIALIZED"
+    BUILDING = "BUILDING"
+    DEPLOYING = "DEPLOYING"
+    BUILD_SUCCESS = "BUILD_SUCCESS"
+    DEPLOY_SUCCESS = "DEPLOY_SUCCESS"
+    DEPLOY_FAILED = "DEPLOY_FAILED"
+    BUILD_FAILED = "BUILD_FAILED"
+    CANCELLED = "CANCELLED"
+    COMPONENTS_DEPLOYING = "COMPONENTS_DEPLOYING"
+    FAILED = "FAILED"
+    REUSING_EXISTING_BUILD = "REUSING_EXISTING_BUILD"
+    REDEPLOY_STARTED = "REDEPLOY_STARTED"
+    PAUSED = "PAUSED"
+    ROLLOUT_STARTED = "ROLLOUT_STARTED"
+    SET_TRAFFIC = "SET_TRAFFIC"
+    DEPLOY_FAILED_WITH_RETRY = "DEPLOY_FAILED_WITH_RETRY"
+    WAITING = "WAITING"
+    _ = ""
 
     @classmethod
     def is_failure_state(cls, state: DeploymentTransitionStatus) -> bool:
@@ -112,6 +111,9 @@ class DeploymentManifest(Base):
 
 
 class Deployment(Entity):
+    class Config:
+        extra = Extra.allow
+
     id: str = Field(repr=False)
     fqn: str
     version: int
@@ -123,7 +125,8 @@ class Deployment(Entity):
     metadata: Optional[
         Union[Dict[str, Any], List[Dict[str, Any]]]
     ]  # TODO (chiragjn): revisit the type of this field
-    currentStatus: Optional[DeploymentStatus]  # Server was not returning CurrentStatus
+    # Server was not returning CurrentStatus
+    currentStatus: Optional[DeploymentStatus]
     application_fqn: str
 
     def __init__(self, **kwargs) -> None:
@@ -161,9 +164,6 @@ class Deployment(Entity):
             "updated_at": self.updatedAt,
         }
 
-    class Config:
-        extra = Extra.allow
-
 
 class ApplicationWorkspace(Base):
     name: str
@@ -285,6 +285,19 @@ class DeleteResult(BaseModel):
     message: str
 
 
-class Manifest(Base):
+class ManifestLike(Base):
     type: str
     name: str
+
+
+class LogBody(Base):
+    time: float
+    log: str
+
+
+class SocketEvent(Base):
+    class Config:
+        smart_union = True
+
+    type: str
+    body: Union[LogBody, str, Any]