thestage 0.5.37__tar.gz → 0.5.38__tar.gz

{thestage-0.5.37 → thestage-0.5.38}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: thestage
-Version: 0.5.37
+Version: 0.5.38
 Summary: 
 Author: TheStage AI team
 Author-email: hello@thestage.ai
@@ -11,6 +11,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: aioconsole (>=0.8.0,<0.9.0)
+Requires-Dist: boto3 (>=1.35.80,<2.0.0)
 Requires-Dist: gitpython (>=3.1.40,<4.0.0)
 Requires-Dist: httpx (>=0.27.2,<0.28.0)
 Requires-Dist: paramiko (>=3.4.0,<4.0.0)

{thestage-0.5.37 → thestage-0.5.38}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "thestage"
-version = "0.5.37"
+version = "0.5.38"
 description = ""
 authors = ["TheStage AI team <hello@thestage.ai>"]
 readme = "README.md"
@@ -33,6 +33,7 @@ paramiko = "^3.4.0"
 thestage-core = "0.0.16"
 aioconsole = "^0.8.0"
 httpx = "^0.27.2"
+boto3 = "^1.35.80"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.4.3"

thestage-0.5.38/thestage/.env

@@ -0,0 +1,6 @@
+THESTAGE_CONFIG_DIR=.thestage
+THESTAGE_CONFIG_FILE=config.json
+THESTAGE_CLI_ENV=DEV
+THESTAGE_API_URL=https://backend-staging.thestage.ai
+THESTAGE_API_URL=https://backend.thestage.ai
+LOG_FILE=thestage.log

{thestage-0.5.37 → thestage-0.5.38}/thestage/__init__.py

@@ -1,3 +1,3 @@
 from . import *
 __app_name__ = "thestage"
-__version__ = "0.5.37"
+__version__ = "0.5.38"

{thestage-0.5.37 → thestage-0.5.38}/thestage/controllers/project_controller.py

@@ -1,3 +1,4 @@
+import time
 from typing import Optional, List
 
 import re
@@ -723,4 +724,138 @@ def inference_simulator_logs(
             logging_service.print_last_inference_simulator_logs(config=config, inference_simulator_id=inference_simulator_id, logs_number=logs_number)
 
     app_logger.info(f'Inference simulator logs - end')
-    raise typer.Exit(0)
+    raise typer.Exit(0)
+
+
+@inference_simulator_model_app.command("deploy-instance", no_args_is_help=True, help=__("Deploy an inference simulator model to an instance"))
+def deploy_inference_simulator_model_to_instance(
+        unique_id: Optional[str] = typer.Argument(help=__("Inference simulator model unique ID"), ),
+        rented_instance_unique_id: Optional[str] = typer.Option(
+            None,
+            '--rented-instance-unique-id',
+            '-ruid',
+            help=__("The rented instance unique ID on which the inference simulator model will be deployed"),
+            is_eager=False,
+        ),
+        self_hosted_instance_unique_id: Optional[str] = typer.Option(
+            None,
+            '--self-hosted-instance-unique-id',
+            '-suid',
+            help=__("The self-hosted instance unique ID on which the inference simulator model will be deployed"),
+            is_eager=False,
+        ),
+        working_directory: Optional[str] = typer.Option(
+            None,
+            "--working-directory",
+            "-wd",
+            help=__("Full path to working directory. By default, the current directory is used"),
+            show_default=False,
+            is_eager=False,
+        ),
+        enable_log_stream: Optional[bool] = typer.Option(
+            True,
+            " /--no-logs",
+            " /-nl",
+            help=__("Disable real-time log streaming"),
+            is_eager=False,
+        ),
+):
+    """
+        Deploy an inference simulator model to an instance
+    """
+
+    if unique_id and not re.match(r"^[a-zA-Z0-9-]+$", unique_id):
+        raise typer.BadParameter(__("Invalid UID format. The UID can only contain letters, numbers, and hyphens."))
+
+    unique_id_with_timestamp = f"{unique_id}-{int(time.time())}"
+
+    app_logger.info(f'Deploying an inference simulator model')
+
+    service_factory = validate_config_and_get_service_factory(working_directory=working_directory)
+    config = service_factory.get_config_provider().get_full_config()
+
+    project_service = service_factory.get_project_service()
+
+    project_service.project_deploy_inference_simulator_model_to_instance(
+        config=config,
+        unique_id=unique_id,
+        unique_id_with_timestamp=unique_id_with_timestamp,
+        rented_instance_unique_id=rented_instance_unique_id,
+        self_hosted_instance_unique_id=self_hosted_instance_unique_id,
+    )
+
+    if enable_log_stream:
+        logging_service: LoggingService = service_factory.get_logging_service()
+
+        logging_service.stream_inference_simulator_logs_with_controls(
+            config=config,
+            slug=unique_id_with_timestamp
+        )
+    raise typer.Exit(0)
+
+
+@inference_simulator_model_app.command("deploy-sagemaker", no_args_is_help=True, help=__("Deploy an inference simulator model to SageMaker"))
+def deploy_inference_simulator_model_to_sagemaker(
+        unique_id: Optional[str] = typer.Argument(help=__("Inference simulator model unique ID"), ),
+        arn: Optional[str] = typer.Option(
+            None,
+            '--amazon-resource-name',
+            '-arn',
+            help=__("The Amazon Resource Name of the IAM Role to use, e.g., arn:aws:iam::{aws_account_id}:role/{role}"),
+            is_eager=False,
+        ),
+        working_directory: Optional[str] = typer.Option(
+            None,
+            "--working-directory",
+            "-wd",
+            help=__("Full path to working directory. By default, the current directory is used"),
+            show_default=False,
+            is_eager=False,
+        ),
+        instance_type: Optional[str] = typer.Option(
+            None,
+            '--instance-type',
+            '-it',
+            help=__("Instance type on which the inference simulator model will be deployed"),
+            is_eager=False,
+        ),
+        initial_variant_weight: Optional[float] = typer.Option(
+            None,
+            "--initial-variant-weight",
+            "-ivw",
+            help=__("Initial Variant Weight. By default 1.0"),
+            show_default=False,
+            is_eager=False,
+        ),
+        initial_instance_count: Optional[int] = typer.Option(
+            None,
+            "--initial-instance-count",
+            "-iic",
+            help=__("Initial Instance Count"),
+            show_default=False,
+            is_eager=False,
+        ),
+
+):
+    """
+        Deploy an inference simulator model to SageMaker
+    """
+
+    if unique_id and not re.match(r"^[a-zA-Z0-9-]+$", unique_id):
+        raise typer.BadParameter(__("Invalid UID format. The UID can only contain letters, numbers, and hyphens."))
+
+    app_logger.info(f'Deploying an inference simulator model')
+
+    service_factory = validate_config_and_get_service_factory(working_directory=working_directory)
+    config = service_factory.get_config_provider().get_full_config()
+
+    project_service = service_factory.get_project_service()
+
+    project_service.project_deploy_inference_simulator_model_to_sagemaker(
+        config=config,
+        unique_id=unique_id,
+        arn=arn,
+        instance_type=instance_type,
+        initial_variant_weight=initial_variant_weight,
+        initial_instance_count=initial_instance_count,
+    )

{thestage-0.5.37 → thestage-0.5.38}/thestage/entities/project_inference_simulator_model.py

@@ -11,6 +11,7 @@ class ProjectInferenceSimulatorModelEntity(BaseModel):
     )
 
     id: Optional[int] = Field(None, alias='ID')
+    slug: Optional[str] = Field(None, alias='UNIQUE_ID')
     status: Optional[str] = Field(None, alias='STATUS')
     commit_hash: Optional[str] = Field(None, alias='COMMIT_HASH')
     environment_metadata: Optional[str] = Field(None, alias='ENVIRONMENT_METADATA')

{thestage-0.5.37 → thestage-0.5.38}/thestage/services/clients/thestage_api/api_client.py

@@ -10,6 +10,14 @@ from thestage.services.clients.thestage_api.dtos.docker_container_controller.doc
 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 \
@@ -736,6 +744,53 @@ class TheStageApiClient(TheStageApiClientCore):
         return GetInferenceSimulatorResponse.model_validate(response) if response else None
 
 
+    @error_handler()
+    def deploy_inference_model_to_instance(
+            self,
+            token: str,
+            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=token,
+        )
+        return DeployInferenceModelToInstanceResponse.model_validate(response) if response else None
+
+
+    @error_handler()
+    def deploy_inference_model_to_sagemaker(
+            self,
+            token: str,
+            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=token,
+        )
+        return DeployInferenceModelToSagemakerResponse.model_validate(response) if response else None
+
+
     def query_user_logs(self, token: str, limit: int, task_id: Optional[int] = None, inference_simulator_id: Optional[int] = None) -> UserLogsQueryResponse:
         request = UserLogsQueryRequest(
             inferenceSimulatorId=inference_simulator_id,

thestage-0.5.38/thestage/services/clients/thestage_api/dtos/inference_controller/deploy_inference_model_to_instance_request.py

@@ -0,0 +1,14 @@
+from typing import Optional, List
+
+from pydantic import Field, ConfigDict, BaseModel
+
+from thestage.services.clients.thestage_api.dtos.entity_filter_request import EntityFilterRequest
+
+
+class DeployInferenceModelToInstanceRequest(BaseModel):
+    model_config = ConfigDict(use_enum_values=True)
+
+    modelSlug: Optional[str] = Field(None, alias='modelSlug')
+    instanceRentedSlug: Optional[str] = Field(None, alias='instanceRentedSlug')
+    selfhostedInstanceSlug: Optional[str] = Field(None, alias='selfhostedInstanceSlug')
+    inferenceSimulatorSlug: Optional[str] = Field(None, alias='inferenceSimulatorSlug')

thestage-0.5.38/thestage/services/clients/thestage_api/dtos/inference_controller/deploy_inference_model_to_instance_response.py

@@ -0,0 +1,13 @@
+from typing import Optional
+
+from pydantic import Field, ConfigDict
+
+from thestage.services.clients.thestage_api.dtos.base_response import TheStageBaseResponse
+from thestage.services.project.dto.inference_simulator_model_dto import InferenceSimulatorModelDto
+
+
+class DeployInferenceModelToInstanceResponse(TheStageBaseResponse):
+    model_config = ConfigDict(use_enum_values=True)
+    inferenceSimulatorModel: Optional[InferenceSimulatorModelDto] = Field(None, alias='inferenceSimulatorModel')
+
+

thestage-0.5.38/thestage/services/clients/thestage_api/dtos/inference_controller/deploy_inference_model_to_sagemaker_request.py

@@ -0,0 +1,12 @@
+from typing import Optional, List
+
+from pydantic import Field, ConfigDict, BaseModel
+
+from thestage.services.clients.thestage_api.dtos.entity_filter_request import EntityFilterRequest
+
+
+class DeployInferenceModelToSagemakerRequest(BaseModel):
+    model_config = ConfigDict(use_enum_values=True)
+
+    modelSlug: Optional[str] = Field(None, alias='modelSlug')
+    arn: Optional[str] = Field(None, alias='arn')

thestage-0.5.38/thestage/services/clients/thestage_api/dtos/inference_controller/deploy_inference_model_to_sagemaker_response.py

@@ -0,0 +1,12 @@
+from pydantic import Field, ConfigDict
+
+from thestage.services.clients.thestage_api.dtos.base_response import TheStageBaseResponse
+
+
+class DeployInferenceModelToSagemakerResponse(TheStageBaseResponse):
+    model_config = ConfigDict(use_enum_values=True)
+    modelId: str = Field(None, alias='modelId')
+    ecrImageUrl: str = Field(None, alias='ecrImageUrl')
+    s3ArtifactsUrl: str = Field(None, alias='s3ArtifactsUrl')
+
+

{thestage-0.5.37 → thestage-0.5.38}/thestage/services/project/dto/inference_simulator_model_dto.py

@@ -5,6 +5,7 @@ class InferenceSimulatorModelDto(BaseModel):
     model_config = ConfigDict(use_enum_values=True)
 
     id: Optional[int] = Field(None, alias='id')
+    slug: Optional[str] = Field(None, alias='slug')
     client_id: Optional[int] = Field(None, alias='clientId')
     instance_rented_id: Optional[int] = Field(None, alias='instanceRentedId')
     selfhosted_instance_id: Optional[int] = Field(None, alias='selfhostedInstanceId')

{thestage-0.5.37 → thestage-0.5.38}/thestage/services/project/mapper/project_inference_simulator_model_mapper.py

@@ -12,6 +12,7 @@ class ProjectInferenceSimulatorModelMapper(AbstractMapper):
 
         return ProjectInferenceSimulatorModelEntity(
             id=item.id,
+            slug=item.slug,
             status=item.status or '',
             commit_hash=item.commit_hash or '',
             environment_metadata=item.environment_metadata or '',

{thestage-0.5.37 → thestage-0.5.38}/thestage/services/project/project_service.py

@@ -1,8 +1,12 @@
 import os
+import time
+from datetime import datetime
 from pathlib import Path
 from typing import Optional, List
 
 import json
+
+import boto3
 import click
 import typer
 from git import Commit
@@ -18,6 +22,10 @@ from thestage.i18n.translation import __
 from thestage.services.clients.git.git_client import GitLocalClient
 from thestage.services.clients.thestage_api.dtos.container_response import DockerContainerDto
 from thestage.services.clients.thestage_api.dtos.enums.container_status import DockerContainerStatus
+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_response import \
+    DeployInferenceModelToSagemakerResponse
 from thestage.services.clients.thestage_api.dtos.inference_controller.get_inference_simulator_response import \
     GetInferenceSimulatorResponse
 from thestage.services.clients.thestage_api.dtos.paginated_entity_list import PaginatedEntityList
@@ -902,3 +910,170 @@ class ProjectService(AbstractService):
 
         return project_config
 
+    @error_handler()
+    def project_deploy_inference_simulator_model_to_instance(
+            self,
+            config: ConfigEntity,
+            unique_id: Optional[str] = None,
+            unique_id_with_timestamp: Optional[str] = None,
+            rented_instance_unique_id: Optional[str] = None,
+            self_hosted_instance_unique_id: Optional[str] = None,
+    ) -> None:
+        project_config: ProjectConfig = self.__get_fixed_project_config(config=config)
+        if not project_config:
+            typer.echo(
+                __("No project found at the path: %path%. Please initialize or clone a project first. Or provide path to project using --working-directory option.",
+                   {"path": config.runtime.working_directory}))
+            raise typer.Exit(1)
+
+        if rented_instance_unique_id and self_hosted_instance_unique_id:
+            typer.echo(__("Error: Cannot provide both rented and self-hosted instance unique IDs."))
+            raise typer.Exit(1)
+
+        if not rented_instance_unique_id and not self_hosted_instance_unique_id:
+            typer.echo(__("Error: Either a rented instance ID or a self-hosted instance unique ID must be provided."))
+            raise typer.Exit(1)
+
+        project_config: ProjectConfig = self.__config_provider.read_project_config()
+        if not project_config:
+            typer.echo(__("No project found at the path: %path%. Please initialize or clone a project first.",
+                          {"path": config.runtime.working_directory}))
+            raise typer.Exit(1)
+
+        typer.echo(__("Creating inference simulator with unique ID: %unique_id_with_timestamp%", {"unique_id_with_timestamp": unique_id_with_timestamp}))
+        deploy_model_to_instance_response: DeployInferenceModelToInstanceResponse = self.__thestage_api_client.deploy_inference_model_to_instance(
+            token=config.main.thestage_auth_token,
+            unique_id=unique_id,
+            unique_id_with_timestamp=unique_id_with_timestamp,
+            rented_instance_unique_id=rented_instance_unique_id,
+            self_hosted_instance_unique_id=self_hosted_instance_unique_id
+        )
+        if deploy_model_to_instance_response:
+            if deploy_model_to_instance_response.message:
+                typer.echo(deploy_model_to_instance_response.message)
+            if deploy_model_to_instance_response.is_success:
+                typer.echo("Inference simulator has been scheduled to run successfully.")
+            else:
+                typer.echo(__(
+                    'Inference simulator failed to run with an error: %server_massage%',
+                    {'server_massage': deploy_model_to_instance_response.message or ""}
+                ))
+                raise typer.Exit(1)
+        else:
+            typer.echo(__("Inference simulator failed to run with an error"))
+            raise typer.Exit(1)
+
+
+    @error_handler()
+    def project_deploy_inference_simulator_model_to_sagemaker(
+            self,
+            config: ConfigEntity,
+            unique_id: Optional[str] = None,
+            arn: Optional[str] = None,
+            instance_type: Optional[str] = None,
+            initial_variant_weight: Optional[float] = 1.0,
+            initial_instance_count: Optional[int] = None,
+    ) -> None:
+        project_config: ProjectConfig = self.__get_fixed_project_config(config=config)
+        if not project_config:
+            typer.echo(
+                __("No project found at the path: %path%. Please initialize or clone a project first. Or provide path to project using --working-directory option.",
+                   {"path": config.runtime.working_directory}))
+            raise typer.Exit(1)
+
+        if not instance_type:
+            typer.echo(__("Error: Instance type is required."))
+            raise typer.Exit(1)
+
+        if not initial_instance_count:
+            typer.echo(__("Error: Initial instance count is required."))
+            raise typer.Exit(1)
+
+        if not arn:
+            typer.echo(__("Error: ARN is required."))
+            raise typer.Exit(1)
+
+        project_config: ProjectConfig = self.__config_provider.read_project_config()
+        if not project_config:
+            typer.echo(__("No project found at the path: %path%. Please initialize or clone a project first.",
+                          {"path": config.runtime.working_directory}))
+            raise typer.Exit(1)
+
+        deploy_model_to_sagemaker_response: DeployInferenceModelToSagemakerResponse = self.__thestage_api_client.deploy_inference_model_to_sagemaker(
+            token=config.main.thestage_auth_token,
+            unique_id=unique_id,
+            arn=arn,
+        )
+
+        if not deploy_model_to_sagemaker_response.is_success:
+            typer.echo(__(
+                'Failed to prepare model for deployment with an error: %server_massage%',
+                {'server_massage': deploy_model_to_sagemaker_response.message or ""}
+            ))
+            raise typer.Exit(1)
+
+        model_id = deploy_model_to_sagemaker_response.modelId
+        image_uri = deploy_model_to_sagemaker_response.ecrImageUrl
+        model_uri = deploy_model_to_sagemaker_response.s3ArtifactsUrl
+        region = "us-east-1"
+        sm_client = boto3.client('sagemaker', region_name=region)
+
+        try:
+            container = {
+                "Image": image_uri,
+                "ModelDataUrl": model_uri,
+                "Environment": {
+                    "SAGEMAKER_TRITON_DEFAULT_MODEL_NAME": model_id,
+                    "THESTAGE_API_URL": config.main.thestage_api_url,
+                    "THESTAGE_AUTH_TOKEN": config.main.thestage_auth_token
+                },
+            }
+
+            sm_model_name = f"{unique_id}-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+            create_model_response = sm_client.create_model(
+                ModelName=sm_model_name,
+                ExecutionRoleArn=arn,
+                PrimaryContainer=container,
+            )
+            typer.echo(f"Model created successfully. Model ARN: {create_model_response['ModelArn']}")
+
+            endpoint_config_name = f"{unique_id}-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+            create_endpoint_config_response = sm_client.create_endpoint_config(
+                EndpointConfigName=endpoint_config_name,
+                ProductionVariants=[
+                    {
+                        "InstanceType": instance_type,
+                        "InitialVariantWeight": initial_variant_weight,
+                        "InitialInstanceCount": initial_instance_count,
+                        "ModelName": sm_model_name,
+                        "VariantName": "AllTraffic",
+                    }
+                ],
+            )
+            typer.echo(
+                f"Endpoint configuration created successfully. Endpoint Config ARN: {create_endpoint_config_response['EndpointConfigArn']}")
+
+            endpoint_name = f"{unique_id}-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+            create_endpoint_response = sm_client.create_endpoint(
+                EndpointName=endpoint_name,
+                EndpointConfigName=endpoint_config_name,
+            )
+            typer.echo(f"Endpoint created successfully. Endpoint ARN: {create_endpoint_response['EndpointArn']}")
+
+            typer.echo("Waiting for the endpoint to become active...")
+            while True:
+                resp = sm_client.describe_endpoint(EndpointName=endpoint_name)
+                status = resp["EndpointStatus"]
+                typer.echo(f"Status: {status}")
+                if status == "InService":
+                    break
+                elif status == "Failed":
+                    typer.echo(f"Endpoint creation failed. Reason: {resp.get('FailureReason', 'Unknown')}")
+                    raise typer.Exit(1)
+                time.sleep(60)
+
+            typer.echo(f"Endpoint is ready. ARN: {resp['EndpointArn']} Status: {status}")
+
+        except Exception as e:
+            typer.echo(__("Failed to deploy the inference simulator model to SageMaker: %error%", {"error": str(e)}))
+            raise typer.Exit(1)

thestage-0.5.37/thestage/.env

@@ -1,5 +0,0 @@
-THESTAGE_CONFIG_DIR=.thestage
-THESTAGE_CONFIG_FILE=config.json
-THESTAGE_API_URL=https://backend-staging3.thestage.ai
-THESTAGE_API_URL=http://localhost:8100
-LOG_FILE=thestage.log