thestage 0.5.471__py3-none-any.whl → 0.6.1__py3-none-any.whl

thestage/services/clients/thestage_api/api_client.py

@@ -1,720 +1,718 @@
-from typing import Optional, List, Tuple, Dict, Iterator
-
-import httpx
-
-from thestage.helpers.error_handler import error_handler
-from thestage.services.clients.thestage_api.core.api_client_core import TheStageApiClientCore
-from thestage.services.clients.thestage_api.dtos.docker_container_controller.docker_container_list_request import \
-    DockerContainerListRequest
-from thestage.services.clients.thestage_api.dtos.docker_container_controller.docker_container_list_response import \
-    DockerContainerListResponse
-from thestage.services.clients.thestage_api.dtos.entity_filter_request import EntityFilterRequest
-from thestage.services.clients.thestage_api.dtos.inference_controller.deploy_inference_model_to_instance_request import \
-    DeployInferenceModelToInstanceRequest
-from thestage.services.clients.thestage_api.dtos.inference_controller.deploy_inference_model_to_instance_response import \
-    DeployInferenceModelToInstanceResponse
-from thestage.services.clients.thestage_api.dtos.inference_controller.deploy_inference_model_to_sagemaker_request import \
-    DeployInferenceModelToSagemakerRequest
-from thestage.services.clients.thestage_api.dtos.inference_controller.deploy_inference_model_to_sagemaker_response import \
-    DeployInferenceModelToSagemakerResponse
-from thestage.services.clients.thestage_api.dtos.inference_controller.get_inference_simulator_request import \
-    GetInferenceSimulatorRequest
-from thestage.services.clients.thestage_api.dtos.inference_controller.get_inference_simulator_response import \
-    GetInferenceSimulatorResponse
-from thestage.services.clients.thestage_api.dtos.inference_controller.inference_simulator_list_for_project_request import \
-    InferenceSimulatorListForProjectRequest
-from thestage.services.clients.thestage_api.dtos.inference_controller.inference_simulator_list_for_project_response import \
-    InferenceSimulatorListForProjectResponse
-from thestage.services.clients.thestage_api.dtos.inference_controller.inference_simulator_model_list_for_project_request import \
-    InferenceSimulatorModelListForProjectRequest
-from thestage.services.clients.thestage_api.dtos.inference_controller.inference_simulator_model_list_for_project_response import \
-    InferenceSimulatorModelListForProjectResponse
-from thestage.services.clients.thestage_api.dtos.inference_simulator_model_response import \
-    InferenceSimulatorModelStatusMapperResponse
-from thestage.services.clients.thestage_api.dtos.inference_simulator_response import \
-    InferenceSimulatorStatusMapperResponse
-from thestage.services.clients.thestage_api.dtos.logging_controller.log_polling_request import LogPollingRequest
-from thestage.services.clients.thestage_api.dtos.logging_controller.log_polling_response import LogPollingResponse
-from thestage.services.clients.thestage_api.dtos.logging_controller.user_logs_query_request import UserLogsQueryRequest
-from thestage.services.clients.thestage_api.dtos.logging_controller.user_logs_query_response import \
-    UserLogsQueryResponse
-from thestage.services.clients.thestage_api.dtos.paginated_entity_list import PaginatedEntityList
-from thestage.services.clients.thestage_api.dtos.project_controller.project_get_deploy_ssh_key_request import \
-    ProjectGetDeploySshKeyRequest
-from thestage.services.clients.thestage_api.dtos.project_controller.project_get_deploy_ssh_key_response import \
-    ProjectGetDeploySshKeyResponse
-from thestage.services.clients.thestage_api.dtos.project_controller.project_push_inference_simulator_model_request import \
-    ProjectPushInferenceSimulatorModelRequest
-from thestage.services.clients.thestage_api.dtos.project_controller.project_push_inference_simulator_model_response import \
-    ProjectPushInferenceSimulatorModelResponse
-from thestage.services.clients.thestage_api.dtos.project_controller.project_run_task_request import \
-    ProjectRunTaskRequest
-from thestage.services.clients.thestage_api.dtos.project_controller.project_run_task_response import \
-    ProjectRunTaskResponse
-from thestage.services.clients.thestage_api.dtos.project_controller.project_start_inference_simulator_request import \
-    ProjectStartInferenceSimulatorRequest
-from thestage.services.clients.thestage_api.dtos.project_controller.project_start_inference_simulator_response import \
-    ProjectStartInferenceSimulatorResponse
-from thestage.services.clients.thestage_api.dtos.project_response import ProjectDto, ProjectViewResponse
-from thestage.services.clients.thestage_api.dtos.selfhosted_instance_response import SelfHostedInstanceListResponse, \
-    SelfHostedInstanceDto, SelfHostedRentedItemResponse, SelfHostedRentedRentedBusinessStatusMapperResponse
-from thestage.services.clients.thestage_api.dtos.container_param_request import DockerContainerActionRequestDto
-from thestage.services.clients.thestage_api.dtos.container_response import DockerContainerDto, \
-    DockerContainerViewResponse, ContainerBusinessStatusMapperResponse
-from thestage.entities.enums.order_direction_type import OrderDirectionType
-from thestage.services.clients.thestage_api.dtos.ssh_key_controller.add_ssh_key_to_user_request import \
-    AddSshKeyToUserRequest
-from thestage.services.clients.thestage_api.dtos.ssh_key_controller.add_ssh_key_to_user_response import \
-    AddSshKeyToUserResponse
-from thestage.services.clients.thestage_api.dtos.ssh_key_controller.add_ssh_public_key_to_instance_request import \
-    AddSshPublicKeyToInstanceRequest
-from thestage.services.clients.thestage_api.dtos.ssh_key_controller.add_ssh_public_key_to_instance_response import \
-    AddSshPublicKeyToInstanceResponse
-from thestage.services.clients.thestage_api.dtos.ssh_key_controller.is_user_has_public_ssh_key_request import \
-    IsUserHasSshPublicKeyRequest
-from thestage.services.clients.thestage_api.dtos.ssh_key_controller.is_user_has_public_ssh_key_response import \
-    IsUserHasSshPublicKeyResponse
-from thestage.services.clients.thestage_api.dtos.task_controller.task_list_for_project_request import \
-    TaskListForProjectRequest
-from thestage.services.clients.thestage_api.dtos.task_controller.task_list_for_project_response import \
-    TaskListForProjectResponse
-from thestage.services.clients.thestage_api.dtos.task_controller.task_status_localized_map_response import \
-    TaskStatusLocalizedMapResponse
-from thestage.services.clients.thestage_api.dtos.task_controller.task_view_response import TaskViewResponse
-from thestage.services.clients.thestage_api.dtos.instance_rented_response import InstanceRentedListResponse, \
-    InstanceRentedDto, InstanceRentedItemResponse, InstanceRentedBusinessStatusMapperResponse
-
-from thestage.services.clients.thestage_api.dtos.base_response import TheStageBaseResponse
-from thestage.services.config_provider.config_provider import ConfigProvider
-from thestage.services.project.dto.inference_simulator_dto import InferenceSimulatorDto
-from thestage.services.project.dto.inference_simulator_model_dto import InferenceSimulatorModelDto
-from thestage.services.task.dto.task_dto import TaskDto
-
-
-class TheStageApiClient(TheStageApiClientCore):
-    __config_provider: ConfigProvider = None
-
-    def __init__(
-            self,
-            config_provider: ConfigProvider,
-    ):
-        # TODO this is super bullshit, remove all abstract classes
-        super().__init__(url=config_provider.get_config().main.thestage_api_url)
-        self.__config_provider = config_provider
-
-    def get_task(
-            self,
-            task_id: int,
-    ) -> Optional[TaskViewResponse]:
-        data = {
-            "taskId": task_id,
-        }
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/task/view',
-            data=data,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        result = TaskViewResponse.model_validate(response) if response else None
-        return result if result and result.is_success else None
-
-    def get_task_list_for_project(
-            self,
-            project_slug: str,
-            page: int = 1,
-            limit: int = 10,
-    ) -> Optional[PaginatedEntityList[TaskDto]]:
-        request = TaskListForProjectRequest(
-            projectSlug=project_slug,
-            entityFilterRequest=EntityFilterRequest(
-                orderByField="id",
-                orderByDirection=OrderDirectionType.DESC,
-                page=page,
-                limit=limit,
-            ),
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/task/list-for-project',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        result = TaskListForProjectResponse.model_validate(response) if response else None
-        return result.tasks if result and result.is_success else None
-
-    def get_inference_simulator_list_for_project(
-            self,
-            project_slug: str,
-            statuses: List[str] = [],
-            page: int = 1,
-            limit: int = 10,
-    ) -> Optional[PaginatedEntityList[InferenceSimulatorDto]]:
-        request = InferenceSimulatorListForProjectRequest(
-            projectSlug=project_slug,
-            statuses=statuses,
-            entityFilterRequest=EntityFilterRequest(
-                orderByField="id",
-                orderByDirection=OrderDirectionType.DESC,
-                page=page,
-                limit=limit,
-            ),
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/project/inference-simulator/list',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-        result = InferenceSimulatorListForProjectResponse.model_validate(response) if response else None
-        return result.inferenceSimulators if result and result.is_success else None
-
-    def get_inference_simulator_model_list_for_project(
-            self,
-            project_slug: str,
-            statuses: Optional[List[str]] = None,
-            page: int = 1,
-            limit: int = 10,
-    ) -> Optional[PaginatedEntityList[InferenceSimulatorModelDto]]:
-        request = InferenceSimulatorModelListForProjectRequest(
-            projectSlug=project_slug,
-            statuses=statuses,
-            entityFilterRequest=EntityFilterRequest(
-                orderByField="id",
-                orderByDirection=OrderDirectionType.DESC,
-                page=page,
-                limit=limit,
-            ),
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/project/inference-simulator-model/list',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-        result = InferenceSimulatorModelListForProjectResponse.model_validate(response) if response else None
-        return result.inferenceSimulatorModels if result and result.is_success else None
-
-    def get_rented_instance_list(
-            self,
-            statuses: List[str],
-            query: Optional[str] = None,
-            page: int = 1,
-            limit: int = 10,
-    ) -> PaginatedEntityList[InstanceRentedDto]:
-        data = {
-            #"statuses": [item.value for item in statuses],
-            "entityFilterRequest": {
-                "orderByField": "id",
-                "orderByDirection": "DESC",
-                "page": page,
-                "limit": limit
-            },
-            "queryString": query
-        }
-
-        if statuses:
-            data['businessStatuses'] = statuses
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v2/instance-rented/list',
-            data=data,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        result = InstanceRentedListResponse.model_validate(response) if response else None
-        return result.paginated_list if result and result.paginated_list else ([], None)
-
-    def get_rented_business_status_map(self) -> Optional[Dict[str, str]]:
-        response = self._request(
-            method='POST',
-            url='/user-api/v2/instance-rented/business-status-localized-map',
-            data=None,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        data = InstanceRentedBusinessStatusMapperResponse.model_validate(response) if response else None
-
-        return data.instance_rented_business_status_map if data else None
-
-    def get_task_localized_status_map(self) -> Optional[Dict[str, str]]:
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/task/status-localized-mapping',
-            data=None,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        data = TaskStatusLocalizedMapResponse.model_validate(response) if response else None
-
-        return data.taskStatusMap if data else None
-
-    def get_rented_instance(
-            self,
-            instance_slug: Optional[str] = None,
-            instance_id: Optional[int] = None,
-    ) -> Optional[InstanceRentedDto]:
-        if not instance_slug and not instance_id:
-            return None
-
-        data = {
-            "instanceRentedSlug": instance_slug,
-        }
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v2/instance-rented/view-by-slug',
-            data=data,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        return InstanceRentedItemResponse.model_validate(response).instance_rented if response else None
-
-    def get_selfhosted_instance(
-            self,
-            instance_slug: Optional[str] = None,
-            instance_id: Optional[int] = None,
-    ) -> Optional[SelfHostedInstanceDto]:
-        if not instance_slug and not instance_id:
-            return None
-
-        data = {
-            "selfhostedInstanceSlug": instance_slug,
-        }
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v2/self-hosted-instance/view-by-slug',
-            data=data,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        return SelfHostedRentedItemResponse.model_validate(response).selfhosted_instance if response else None
-
-    def get_selfhosted_instance_list(
-            self,
-            statuses: List[str],
-            query: Optional[str] = None,
-            page: int = 1,
-            limit: int = 10,
-    ) -> PaginatedEntityList[SelfHostedInstanceDto]:
-        data = {
-            #"statuses": [item.value for item in statuses] if statuses else [],
-            "entityFilterRequest": {
-                "orderByField": "id",
-                "orderByDirection": "DESC",
-                "page": page,
-                "limit": limit
-            },
-            "queryString": query
-        }
-
-        if statuses:
-            data['businessStatuses'] = statuses
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v2/self-hosted-instance/list',
-            data=data,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        result = SelfHostedInstanceListResponse.model_validate(response) if response else None
-        return result.paginated_list if result and result.paginated_list else ([], None)
-
-    def get_selfhosted_business_status_map(self) -> Optional[Dict[str, str]]:
-        response = self._request(
-            method='POST',
-            url='/user-api/v2/self-hosted-instance/business-status-localized-map',
-            data=None,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        data = SelfHostedRentedRentedBusinessStatusMapperResponse.model_validate(response) if response else None
-        return data.selfhosted_instance_business_status_map if data else None
-
-
-    def cancel_task(
-            self,
-            task_id: int,
-    ) -> Optional[TheStageBaseResponse]:
-        data = {
-            "taskId": task_id,
-        }
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/task/cancel-task',
-            data=data,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        result = TheStageBaseResponse.model_validate(response) if response else None
-        return result if result else None
-
-    def get_container_list(
-            self,
-            project_id: Optional[int] = None,
-            statuses: List[str] = [],
-            page: int = 1,
-            limit: int = 10,
-    ) -> PaginatedEntityList[DockerContainerDto]:
-        request = DockerContainerListRequest(
-            statuses=statuses,
-            projectId=project_id,
-            entityFilterRequest=EntityFilterRequest(
-                orderByField="id",
-                orderByDirection=OrderDirectionType.DESC,
-                page=page,
-                limit=limit,
-            ),
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/docker-container/list',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token
-        )
-
-        result = DockerContainerListResponse.model_validate(response) if response else None
-        # return result.paginatedList.entities, result.paginatedList.pagination_data.total_pages if result and result.is_success else None
-        return result.paginatedList if result and result.is_success else None
-
-    def get_container(
-            self,
-            container_slug: Optional[str] = None,
-            container_id: Optional[int] = None,
-    ) -> Optional[DockerContainerDto]:
-        if not container_slug and not container_id:
-            return None
-
-        if container_id:
-            data = {
-                "dockerContainerId": container_id,
-            }
-        else:
-            data = {
-                "dockerContainerSlug": container_slug,
-            }
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/docker-container/view',
-            data=data,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        return DockerContainerViewResponse.model_validate(response).docker_container if response else None
-
-    def container_action(
-            self,
-            request_param: DockerContainerActionRequestDto,
-    ) -> TheStageBaseResponse:
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/docker-container/action',
-            data=request_param.model_dump(by_alias=True),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        result = TheStageBaseResponse.model_validate(response) if response else None
-        return result
-
-    def get_container_business_status_map(self) -> Optional[Dict[str, str]]:
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/docker-container/status-localized-mapping',
-            data=None,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        data = ContainerBusinessStatusMapperResponse.model_validate(response) if response else None
-        return data.docker_container_status_map if data else None
-
-    def get_project_by_slug(self, slug: str) -> Optional[ProjectDto]:
-        data = {
-            "projectSlug": slug,
-        }
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/project/view-by-slug',
-            data=data,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        result = ProjectViewResponse.model_validate(response) if response else None
-        project = ProjectDto.model_validate(result.project) if result else None
-        return project if result and result.is_success else None
-
-    def get_project_deploy_ssh_key(self, slug: str) -> str:
-        request = ProjectGetDeploySshKeyRequest(
-            projectSlug=slug
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/project/get-deploy-ssh-key',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        result = ProjectGetDeploySshKeyResponse.model_validate(response) if response else None
-        return result.privateKey if result and result.is_success else None
-
-    def execute_project_task(
-            self,
-            project_slug: str,
-            run_command: str,
-            task_title: str,
-            docker_container_slug: Optional[str] = None,
-            commit_hash: Optional[str] = None,
-    ) -> Optional[ProjectRunTaskResponse]:
-        request = ProjectRunTaskRequest(
-            projectSlug=project_slug,
-            dockerContainerSlug=docker_container_slug,
-            commitHash=commit_hash,
-            runCommand=run_command,
-            taskTitle=task_title,
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/project/execute-task',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        return ProjectRunTaskResponse.model_validate(response) if response else None
-
-    async def poll_logs_httpx(self, docker_container_id: Optional[int], last_log_timestamp: str,
-                              last_log_id: str, task_id: Optional[int] = None,
-                              inference_simulator_id: Optional[int] = None) -> Optional[LogPollingResponse]:
-        request_headers = {'Content-Type': 'application/json'}
-        token = self.__config_provider.get_config().main.thestage_auth_token
-        if token: request_headers['Authorization'] = f"Bearer {token}"
-
-        request = LogPollingRequest(
-            inferenceSimulatorId=inference_simulator_id,
-            taskId=task_id,
-            dockerContainerId=docker_container_id,
-            lastLogTimestamp=last_log_timestamp,
-            lastLogId=last_log_id
-        )
-
-        async with httpx.AsyncClient() as client:
-            response = await client.post(
-                url=f"{self._get_host()}/user-api/v1/logging/poll",
-                headers=request_headers,
-                json=request.model_dump(),
-                timeout=3.5
-            )
-
-            return LogPollingResponse.model_validate(response.json()) if response else None
-
-    def add_public_ssh_key_to_user(self, public_key: str, note: str) -> AddSshKeyToUserResponse:
-        request = AddSshKeyToUserRequest(
-            sshKey=public_key,
-            note=note
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/ssh-key/add-public-key-to-user',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        result = AddSshKeyToUserResponse.model_validate(response) if response else None
-        return result
-
-    def is_user_has_ssh_public_key(self, public_key: str) -> IsUserHasSshPublicKeyResponse:
-        request = IsUserHasSshPublicKeyRequest(
-            sshKey=public_key,
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/ssh-key/is-user-has-public-ssh-key',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        result = IsUserHasSshPublicKeyResponse.model_validate(response) if response else None
-        return result
-
-    def add_public_ssh_key_to_instance_rented(self, instance_rented_id: int,
-                                              ssh_key_pair_id: int) -> AddSshPublicKeyToInstanceResponse:
-        request = AddSshPublicKeyToInstanceRequest(
-            instanceRentedId=instance_rented_id,
-            sshPublicKeyId=ssh_key_pair_id,
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/ssh-key/add-public-ssh-key-to-instance-rented',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        result = AddSshPublicKeyToInstanceResponse.model_validate(response) if response else None
-        return result
-
-    def start_project_inference_simulator(
-            self,
-            project_slug: str,
-            slug: str,
-            commit_hash: Optional[str] = None,
-            rented_instance_unique_id: Optional[str] = None,
-            self_hosted_instance_unique_id: Optional[str] = None,
-            inference_dir: Optional[str] = None,
-            is_skip_installation: Optional[bool] = False,
-    ) -> Optional[ProjectStartInferenceSimulatorResponse]:
-        request = ProjectStartInferenceSimulatorRequest(
-            projectSlug=project_slug,
-            commitHash=commit_hash,
-            slug=slug,
-            instanceRentedUId=rented_instance_unique_id,
-            selfhostedInstanceUId=self_hosted_instance_unique_id,
-            inferenceDir=inference_dir,
-            isSkipInstallation=is_skip_installation,
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/project/inference-simulator/create',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        return ProjectStartInferenceSimulatorResponse.model_validate(response) if response else None
-
-    def push_project_inference_simulator_model(
-            self,
-            slug: str,
-    ) -> Optional[ProjectPushInferenceSimulatorModelResponse]:
-        request = ProjectPushInferenceSimulatorModelRequest(
-            slug=slug,
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/project/inference-simulator/push-model',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        return ProjectPushInferenceSimulatorModelResponse.model_validate(response) if response else None
-
-    def get_inference_simulator_business_status_map(self) -> Optional[Dict[str, str]]:
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/inference-simulator/status-localized-mapping',
-            data=None,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        data = InferenceSimulatorStatusMapperResponse.model_validate(response) if response else None
-
-        return data.inference_simulator_status_map if data else None
-
-    def get_inference_simulator_model_business_status_map(self) -> Optional[Dict[str, str]]:
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/inference-simulator-model/status-localized-mapping',
-            data=None,
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        data = InferenceSimulatorModelStatusMapperResponse.model_validate(response) if response else None
-
-        return data.inference_simulator_model_status_map if data else None
-
-    @error_handler()
-    def get_inference_simulator(
-            self,
-            slug: str,
-    ) -> Optional[GetInferenceSimulatorResponse]:
-        request = GetInferenceSimulatorRequest(
-            slug=slug,
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/inference-simulator/get',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-        return GetInferenceSimulatorResponse.model_validate(response) if response else None
-
-    @error_handler()
-    def deploy_inference_model_to_instance(
-            self,
-            unique_id: str,
-            unique_id_with_timestamp: str,
-            rented_instance_unique_id: Optional[str] = None,
-            self_hosted_instance_unique_id: Optional[str] = None,
-
-    ) -> Optional[DeployInferenceModelToInstanceResponse]:
-        request = DeployInferenceModelToInstanceRequest(
-            inferenceSimulatorSlug=unique_id_with_timestamp,
-            modelSlug=unique_id,
-            instanceRentedSlug=rented_instance_unique_id,
-            selfhostedInstanceSlug=self_hosted_instance_unique_id
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/inference-simulator-model/deploy/instance',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-        return DeployInferenceModelToInstanceResponse.model_validate(response) if response else None
-
-    @error_handler()
-    def deploy_inference_model_to_sagemaker(
-            self,
-            unique_id: str,
-            arn: Optional[str] = None,
-    ) -> Optional[DeployInferenceModelToSagemakerResponse]:
-        request = DeployInferenceModelToSagemakerRequest(
-            modelSlug=unique_id,
-            arn=arn,
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/inference-simulator-model/grant-user-arn-access',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-        return DeployInferenceModelToSagemakerResponse.model_validate(response) if response else None
-
-    def query_user_logs(self, limit: int, task_id: Optional[int] = None,
-                        inference_simulator_id: Optional[int] = None,
-                        container_id: Optional[int] = None) -> UserLogsQueryResponse:
-        request = UserLogsQueryRequest(
-            inferenceSimulatorId=inference_simulator_id,
-            taskId=task_id,
-            containerId=container_id,
-            limit=limit,
-            ascendingOrder=False,
-        )
-
-        response = self._request(
-            method='POST',
-            url='/user-api/v1/logging/query-user-logs',
-            data=request.model_dump(),
-            token=self.__config_provider.get_config().main.thestage_auth_token,
-        )
-
-        result = UserLogsQueryResponse.model_validate(response) if response else None
-        return result
+from typing import Optional, List, Tuple, Dict, Iterator
+
+import httpx
+
+from thestage.helpers.error_handler import error_handler
+from thestage.services.clients.thestage_api.core.api_client_core import TheStageApiClientCore
+from thestage.services.clients.thestage_api.dtos.docker_container_controller.docker_container_list_request import \
+    DockerContainerListRequest
+from thestage.services.clients.thestage_api.dtos.docker_container_controller.docker_container_list_response import \
+    DockerContainerListResponse
+from thestage.services.clients.thestage_api.dtos.entity_filter_request import EntityFilterRequest
+from thestage.services.clients.thestage_api.dtos.inference_controller.deploy_inference_model_to_instance_request import \
+    DeployInferenceModelToInstanceRequest
+from thestage.services.clients.thestage_api.dtos.inference_controller.deploy_inference_model_to_instance_response import \
+    DeployInferenceModelToInstanceResponse
+from thestage.services.clients.thestage_api.dtos.inference_controller.deploy_inference_model_to_sagemaker_request import \
+    DeployInferenceModelToSagemakerRequest
+from thestage.services.clients.thestage_api.dtos.inference_controller.deploy_inference_model_to_sagemaker_response import \
+    DeployInferenceModelToSagemakerResponse
+from thestage.services.clients.thestage_api.dtos.inference_controller.get_inference_simulator_request import \
+    GetInferenceSimulatorRequest
+from thestage.services.clients.thestage_api.dtos.inference_controller.get_inference_simulator_response import \
+    GetInferenceSimulatorResponse
+from thestage.services.clients.thestage_api.dtos.inference_controller.inference_simulator_list_for_project_request import \
+    InferenceSimulatorListForProjectRequest
+from thestage.services.clients.thestage_api.dtos.inference_controller.inference_simulator_list_for_project_response import \
+    InferenceSimulatorListForProjectResponse
+from thestage.services.clients.thestage_api.dtos.inference_controller.inference_simulator_model_list_for_project_request import \
+    InferenceSimulatorModelListForProjectRequest
+from thestage.services.clients.thestage_api.dtos.inference_controller.inference_simulator_model_list_for_project_response import \
+    InferenceSimulatorModelListForProjectResponse
+from thestage.services.clients.thestage_api.dtos.inference_simulator_model_response import \
+    InferenceSimulatorModelStatusMapperResponse
+from thestage.services.clients.thestage_api.dtos.inference_simulator_response import \
+    InferenceSimulatorStatusMapperResponse
+from thestage.services.clients.thestage_api.dtos.logging_controller.log_polling_request import LogPollingRequest
+from thestage.services.clients.thestage_api.dtos.logging_controller.log_polling_response import LogPollingResponse
+from thestage.services.clients.thestage_api.dtos.logging_controller.user_logs_query_request import UserLogsQueryRequest
+from thestage.services.clients.thestage_api.dtos.logging_controller.user_logs_query_response import \
+    UserLogsQueryResponse
+from thestage.services.clients.thestage_api.dtos.paginated_entity_list import PaginatedEntityList
+from thestage.services.clients.thestage_api.dtos.project_controller.project_get_deploy_ssh_key_request import \
+    ProjectGetDeploySshKeyRequest
+from thestage.services.clients.thestage_api.dtos.project_controller.project_get_deploy_ssh_key_response import \
+    ProjectGetDeploySshKeyResponse
+from thestage.services.clients.thestage_api.dtos.project_controller.project_push_inference_simulator_model_request import \
+    ProjectPushInferenceSimulatorModelRequest
+from thestage.services.clients.thestage_api.dtos.project_controller.project_push_inference_simulator_model_response import \
+    ProjectPushInferenceSimulatorModelResponse
+from thestage.services.clients.thestage_api.dtos.project_controller.project_run_task_request import \
+    ProjectRunTaskRequest
+from thestage.services.clients.thestage_api.dtos.project_controller.project_run_task_response import \
+    ProjectRunTaskResponse
+from thestage.services.clients.thestage_api.dtos.project_controller.project_start_inference_simulator_request import \
+    ProjectStartInferenceSimulatorRequest
+from thestage.services.clients.thestage_api.dtos.project_controller.project_start_inference_simulator_response import \
+    ProjectStartInferenceSimulatorResponse
+from thestage.services.clients.thestage_api.dtos.project_response import ProjectDto, ProjectViewResponse
+from thestage.services.clients.thestage_api.dtos.selfhosted_instance_response import SelfHostedInstanceListResponse, \
+    SelfHostedInstanceDto, SelfHostedRentedItemResponse, SelfHostedRentedRentedBusinessStatusMapperResponse
+from thestage.services.clients.thestage_api.dtos.container_param_request import DockerContainerActionRequestDto
+from thestage.services.clients.thestage_api.dtos.container_response import DockerContainerDto, \
+    DockerContainerViewResponse, ContainerBusinessStatusMapperResponse
+from thestage.entities.enums.order_direction_type import OrderDirectionType
+from thestage.services.clients.thestage_api.dtos.ssh_key_controller.add_ssh_key_to_user_request import \
+    AddSshKeyToUserRequest
+from thestage.services.clients.thestage_api.dtos.ssh_key_controller.add_ssh_key_to_user_response import \
+    AddSshKeyToUserResponse
+from thestage.services.clients.thestage_api.dtos.ssh_key_controller.add_ssh_public_key_to_instance_request import \
+    AddSshPublicKeyToInstanceRequest
+from thestage.services.clients.thestage_api.dtos.ssh_key_controller.add_ssh_public_key_to_instance_response import \
+    AddSshPublicKeyToInstanceResponse
+from thestage.services.clients.thestage_api.dtos.ssh_key_controller.is_user_has_public_ssh_key_request import \
+    IsUserHasSshPublicKeyRequest
+from thestage.services.clients.thestage_api.dtos.ssh_key_controller.is_user_has_public_ssh_key_response import \
+    IsUserHasSshPublicKeyResponse
+from thestage.services.clients.thestage_api.dtos.task_controller.task_list_for_project_request import \
+    TaskListForProjectRequest
+from thestage.services.clients.thestage_api.dtos.task_controller.task_list_for_project_response import \
+    TaskListForProjectResponse
+from thestage.services.clients.thestage_api.dtos.task_controller.task_status_localized_map_response import \
+    TaskStatusLocalizedMapResponse
+from thestage.services.clients.thestage_api.dtos.task_controller.task_view_response import TaskViewResponse
+from thestage.services.clients.thestage_api.dtos.instance_rented_response import InstanceRentedListResponse, \
+    InstanceRentedDto, InstanceRentedItemResponse, InstanceRentedBusinessStatusMapperResponse
+
+from thestage.services.clients.thestage_api.dtos.base_response import TheStageBaseResponse
+from thestage.services.config_provider.config_provider import ConfigProvider
+from thestage.services.project.dto.inference_simulator_dto import InferenceSimulatorDto
+from thestage.services.project.dto.inference_simulator_model_dto import InferenceSimulatorModelDto
+from thestage.services.task.dto.task_dto import TaskDto
+
+
+class TheStageApiClient(TheStageApiClientCore):
+    __config_provider: ConfigProvider = None
+
+    def __init__(
+            self,
+            config_provider: ConfigProvider,
+    ):
+        # TODO this is super bullshit, remove all abstract classes
+        super().__init__(url=config_provider.get_config().main.thestage_api_url)
+        self.__config_provider = config_provider
+
+    def get_task(
+            self,
+            task_id: int,
+    ) -> Optional[TaskViewResponse]:
+        data = {
+            "taskId": task_id,
+        }
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/task/view',
+            data=data,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = TaskViewResponse.model_validate(response) if response else None
+        return result if result and result.is_success else None
+
+    def get_task_list_for_project(
+            self,
+            project_slug: str,
+            page: int = 1,
+            limit: int = 10,
+    ) -> Optional[PaginatedEntityList[TaskDto]]:
+        request = TaskListForProjectRequest(
+            projectSlug=project_slug,
+            entityFilterRequest=EntityFilterRequest(
+                orderByField="id",
+                orderByDirection=OrderDirectionType.DESC,
+                page=page,
+                limit=limit,
+            ),
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/task/list-for-project',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = TaskListForProjectResponse.model_validate(response) if response else None
+        return result.tasks if result and result.is_success else None
+
+    def get_inference_simulator_list_for_project(
+            self,
+            project_slug: str,
+            statuses: List[str] = [],
+            page: int = 1,
+            limit: int = 10,
+    ) -> Optional[PaginatedEntityList[InferenceSimulatorDto]]:
+        request = InferenceSimulatorListForProjectRequest(
+            projectSlug=project_slug,
+            statuses=statuses,
+            entityFilterRequest=EntityFilterRequest(
+                orderByField="id",
+                orderByDirection=OrderDirectionType.DESC,
+                page=page,
+                limit=limit,
+            ),
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/project/inference-simulator/list',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+        result = InferenceSimulatorListForProjectResponse.model_validate(response) if response else None
+        return result.inferenceSimulators if result and result.is_success else None
+
+    def get_inference_simulator_model_list_for_project(
+            self,
+            project_slug: str,
+            statuses: Optional[List[str]] = None,
+            page: int = 1,
+            limit: int = 10,
+    ) -> Optional[PaginatedEntityList[InferenceSimulatorModelDto]]:
+        request = InferenceSimulatorModelListForProjectRequest(
+            projectSlug=project_slug,
+            statuses=statuses,
+            entityFilterRequest=EntityFilterRequest(
+                orderByField="id",
+                orderByDirection=OrderDirectionType.DESC,
+                page=page,
+                limit=limit,
+            ),
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/project/inference-simulator-model/list',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+        result = InferenceSimulatorModelListForProjectResponse.model_validate(response) if response else None
+        return result.inferenceSimulatorModels if result and result.is_success else None
+
+    def get_rented_instance_list(
+            self,
+            statuses: List[str],
+            query: Optional[str] = None,
+            page: int = 1,
+            limit: int = 10,
+    ) -> PaginatedEntityList[InstanceRentedDto]:
+        data = {
+            #"statuses": [item.value for item in statuses],
+            "entityFilterRequest": {
+                "orderByField": "id",
+                "orderByDirection": "DESC",
+                "page": page,
+                "limit": limit
+            },
+            "queryString": query
+        }
+
+        if statuses:
+            data['businessStatuses'] = statuses
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v2/instance-rented/list',
+            data=data,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = InstanceRentedListResponse.model_validate(response) if response else None
+        return result.paginated_list if result and result.paginated_list else ([], None)
+
+    def get_rented_business_status_map(self) -> Optional[Dict[str, str]]:
+        response = self._request(
+            method='POST',
+            url='/user-api/v2/instance-rented/business-status-localized-map',
+            data=None,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        data = InstanceRentedBusinessStatusMapperResponse.model_validate(response) if response else None
+
+        return data.instance_rented_business_status_map if data else None
+
+    def get_task_localized_status_map(self) -> Optional[Dict[str, str]]:
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/task/status-localized-mapping',
+            data=None,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        data = TaskStatusLocalizedMapResponse.model_validate(response) if response else None
+
+        return data.taskStatusMap if data else None
+
+    def get_rented_instance(
+            self,
+            instance_slug: Optional[str] = None,
+            instance_id: Optional[int] = None,
+    ) -> Optional[InstanceRentedDto]:
+        if not instance_slug and not instance_id:
+            return None
+
+        data = {
+            "instanceRentedSlug": instance_slug,
+        }
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v2/instance-rented/view-by-slug',
+            data=data,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        return InstanceRentedItemResponse.model_validate(response).instance_rented if response else None
+
+    def get_selfhosted_instance(
+            self,
+            instance_slug: Optional[str] = None,
+            instance_id: Optional[int] = None,
+    ) -> Optional[SelfHostedInstanceDto]:
+        if not instance_slug and not instance_id:
+            return None
+
+        data = {
+            "selfhostedInstanceSlug": instance_slug,
+        }
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v2/self-hosted-instance/view-by-slug',
+            data=data,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        return SelfHostedRentedItemResponse.model_validate(response).selfhosted_instance if response else None
+
+    def get_selfhosted_instance_list(
+            self,
+            statuses: List[str],
+            query: Optional[str] = None,
+            page: int = 1,
+            limit: int = 10,
+    ) -> PaginatedEntityList[SelfHostedInstanceDto]:
+        data = {
+            #"statuses": [item.value for item in statuses] if statuses else [],
+            "entityFilterRequest": {
+                "orderByField": "id",
+                "orderByDirection": "DESC",
+                "page": page,
+                "limit": limit
+            },
+            "queryString": query
+        }
+
+        if statuses:
+            data['businessStatuses'] = statuses
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v2/self-hosted-instance/list',
+            data=data,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = SelfHostedInstanceListResponse.model_validate(response) if response else None
+        return result.paginated_list if result and result.paginated_list else ([], None)
+
+    def get_selfhosted_business_status_map(self) -> Optional[Dict[str, str]]:
+        response = self._request(
+            method='POST',
+            url='/user-api/v2/self-hosted-instance/business-status-localized-map',
+            data=None,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        data = SelfHostedRentedRentedBusinessStatusMapperResponse.model_validate(response) if response else None
+        return data.selfhosted_instance_business_status_map if data else None
+
+
+    def cancel_task(
+            self,
+            task_id: int,
+    ) -> Optional[TheStageBaseResponse]:
+        data = {
+            "taskId": task_id,
+        }
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/task/cancel-task',
+            data=data,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = TheStageBaseResponse.model_validate(response) if response else None
+        return result if result else None
+
+    def get_container_list(
+            self,
+            project_id: Optional[int] = None,
+            statuses: List[str] = [],
+            page: int = 1,
+            limit: int = 10,
+    ) -> PaginatedEntityList[DockerContainerDto]:
+        request = DockerContainerListRequest(
+            statuses=statuses,
+            projectId=project_id,
+            entityFilterRequest=EntityFilterRequest(
+                orderByField="id",
+                orderByDirection=OrderDirectionType.DESC,
+                page=page,
+                limit=limit,
+            ),
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/docker-container/list',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token
+        )
+
+        result = DockerContainerListResponse.model_validate(response) if response else None
+        # return result.paginatedList.entities, result.paginatedList.pagination_data.total_pages if result and result.is_success else None
+        return result.paginatedList if result and result.is_success else None
+
+    def get_container(
+            self,
+            container_slug: Optional[str] = None,
+            container_id: Optional[int] = None,
+    ) -> Optional[DockerContainerDto]:
+        if not container_slug and not container_id:
+            return None
+
+        if container_id:
+            data = {
+                "dockerContainerId": container_id,
+            }
+        else:
+            data = {
+                "dockerContainerSlug": container_slug,
+            }
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/docker-container/view',
+            data=data,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        return DockerContainerViewResponse.model_validate(response).docker_container if response else None
+
+    def container_action(
+            self,
+            request_param: DockerContainerActionRequestDto,
+    ) -> TheStageBaseResponse:
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/docker-container/action',
+            data=request_param.model_dump(by_alias=True),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = TheStageBaseResponse.model_validate(response) if response else None
+        return result
+
+    def get_container_business_status_map(self) -> Optional[Dict[str, str]]:
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/docker-container/status-localized-mapping',
+            data=None,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        data = ContainerBusinessStatusMapperResponse.model_validate(response) if response else None
+        return data.docker_container_status_map if data else None
+
+    def get_project_by_slug(self, slug: str) -> Optional[ProjectDto]:
+        data = {
+            "projectSlug": slug,
+        }
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/project/view-by-slug',
+            data=data,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = ProjectViewResponse.model_validate(response) if response else None
+        project = ProjectDto.model_validate(result.project) if result else None
+        return project if result and result.is_success else None
+
+    def get_project_deploy_ssh_key(self, slug: str) -> str:
+        request = ProjectGetDeploySshKeyRequest(
+            projectSlug=slug
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/project/get-deploy-ssh-key',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = ProjectGetDeploySshKeyResponse.model_validate(response) if response else None
+        return result.privateKey if result and result.is_success else None
+
+    def execute_project_task(
+            self,
+            project_slug: str,
+            run_command: str,
+            task_title: str,
+            docker_container_slug: Optional[str] = None,
+            commit_hash: Optional[str] = None,
+    ) -> Optional[ProjectRunTaskResponse]:
+        request = ProjectRunTaskRequest(
+            projectSlug=project_slug,
+            dockerContainerSlug=docker_container_slug,
+            commitHash=commit_hash,
+            runCommand=run_command,
+            taskTitle=task_title,
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/project/execute-task',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        return ProjectRunTaskResponse.model_validate(response) if response else None
+
+    async def poll_logs_httpx(self, docker_container_id: Optional[int], last_log_timestamp: str,
+                              last_log_id: str, task_id: Optional[int] = None,
+                              inference_simulator_id: Optional[int] = None) -> Optional[LogPollingResponse]:
+        request_headers = {'Content-Type': 'application/json'}
+        token = self.__config_provider.get_config().main.thestage_auth_token
+        if token: request_headers['Authorization'] = f"Bearer {token}"
+
+        request = LogPollingRequest(
+            inferenceSimulatorId=inference_simulator_id,
+            taskId=task_id,
+            dockerContainerId=docker_container_id,
+            lastLogTimestamp=last_log_timestamp,
+            lastLogId=last_log_id
+        )
+
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                url=f"{self._get_host()}/user-api/v1/logging/poll",
+                headers=request_headers,
+                json=request.model_dump(),
+                timeout=3.5
+            )
+
+            return LogPollingResponse.model_validate(response.json()) if response else None
+
+    def add_public_ssh_key_to_user(self, public_key: str, note: str) -> AddSshKeyToUserResponse:
+        request = AddSshKeyToUserRequest(
+            sshKey=public_key,
+            note=note
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/ssh-key/add-public-key-to-user',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = AddSshKeyToUserResponse.model_validate(response) if response else None
+        return result
+
+    def is_user_has_ssh_public_key(self, public_key: str) -> IsUserHasSshPublicKeyResponse:
+        request = IsUserHasSshPublicKeyRequest(
+            sshKey=public_key,
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/ssh-key/is-user-has-public-ssh-key',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = IsUserHasSshPublicKeyResponse.model_validate(response) if response else None
+        return result
+
+    def add_public_ssh_key_to_instance_rented(self, instance_rented_id: int,
+                                              ssh_key_pair_id: int) -> AddSshPublicKeyToInstanceResponse:
+        request = AddSshPublicKeyToInstanceRequest(
+            instanceRentedId=instance_rented_id,
+            sshPublicKeyId=ssh_key_pair_id,
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/ssh-key/add-public-ssh-key-to-instance-rented',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = AddSshPublicKeyToInstanceResponse.model_validate(response) if response else None
+        return result
+
+    def start_project_inference_simulator(
+            self,
+            project_slug: str,
+            commit_hash: Optional[str] = None,
+            rented_instance_unique_id: Optional[str] = None,
+            self_hosted_instance_unique_id: Optional[str] = None,
+            inference_dir: Optional[str] = None,
+            is_skip_installation: Optional[bool] = False,
+    ) -> Optional[ProjectStartInferenceSimulatorResponse]:
+        request = ProjectStartInferenceSimulatorRequest(
+            projectSlug=project_slug,
+            commitHash=commit_hash,
+            instanceRentedUId=rented_instance_unique_id,
+            selfhostedInstanceUId=self_hosted_instance_unique_id,
+            inferenceDir=inference_dir,
+            isSkipInstallation=is_skip_installation,
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/project/inference-simulator/create',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        return ProjectStartInferenceSimulatorResponse.model_validate(response) if response else None
+
+    def push_project_inference_simulator_model(
+            self,
+            slug: str,
+    ) -> Optional[ProjectPushInferenceSimulatorModelResponse]:
+        request = ProjectPushInferenceSimulatorModelRequest(
+            slug=slug,
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/project/inference-simulator/push-model',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        return ProjectPushInferenceSimulatorModelResponse.model_validate(response) if response else None
+
+    def get_inference_simulator_business_status_map(self) -> Optional[Dict[str, str]]:
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/inference-simulator/status-localized-mapping',
+            data=None,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        data = InferenceSimulatorStatusMapperResponse.model_validate(response) if response else None
+
+        return data.inference_simulator_status_map if data else None
+
+    def get_inference_simulator_model_business_status_map(self) -> Optional[Dict[str, str]]:
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/inference-simulator-model/status-localized-mapping',
+            data=None,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        data = InferenceSimulatorModelStatusMapperResponse.model_validate(response) if response else None
+
+        return data.inference_simulator_model_status_map if data else None
+
+    @error_handler()
+    def get_inference_simulator(
+            self,
+            slug: str,
+    ) -> Optional[GetInferenceSimulatorResponse]:
+        request = GetInferenceSimulatorRequest(
+            slug=slug,
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/inference-simulator/get',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+        return GetInferenceSimulatorResponse.model_validate(response) if response else None
+
+    @error_handler()
+    def deploy_inference_model_to_instance(
+            self,
+            unique_id: str,
+            unique_id_with_timestamp: str,
+            rented_instance_unique_id: Optional[str] = None,
+            self_hosted_instance_unique_id: Optional[str] = None,
+
+    ) -> Optional[DeployInferenceModelToInstanceResponse]:
+        request = DeployInferenceModelToInstanceRequest(
+            inferenceSimulatorSlug=unique_id_with_timestamp,
+            modelSlug=unique_id,
+            instanceRentedSlug=rented_instance_unique_id,
+            selfhostedInstanceSlug=self_hosted_instance_unique_id
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/inference-simulator-model/deploy/instance',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+        return DeployInferenceModelToInstanceResponse.model_validate(response) if response else None
+
+    @error_handler()
+    def deploy_inference_model_to_sagemaker(
+            self,
+            unique_id: str,
+            arn: Optional[str] = None,
+    ) -> Optional[DeployInferenceModelToSagemakerResponse]:
+        request = DeployInferenceModelToSagemakerRequest(
+            modelSlug=unique_id,
+            arn=arn,
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/inference-simulator-model/grant-user-arn-access',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+        return DeployInferenceModelToSagemakerResponse.model_validate(response) if response else None
+
+    def query_user_logs(self, limit: int, task_id: Optional[int] = None,
+                        inference_simulator_id: Optional[int] = None,
+                        container_id: Optional[int] = None) -> UserLogsQueryResponse:
+        request = UserLogsQueryRequest(
+            inferenceSimulatorId=inference_simulator_id,
+            taskId=task_id,
+            containerId=container_id,
+            limit=limit,
+            ascendingOrder=False,
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v1/logging/query-user-logs',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = UserLogsQueryResponse.model_validate(response) if response else None
+        return result