Amazon SageMaker Unified Studio is in preview release and is subject to change.

Remote Execution APIs

Use the following APIs to start, stop, get, or list executions running on remote compute.

StartExecution

You can start a notebook execution headlessly on a remote compute specified in the StartExecution request.

from sagemaker_ui_helper.sagemaker_ui_helper_api import SageMakerUIHelperAPI
from sagemaker_ui_helper._openapi.models import (
    StartExecutionRequest,
)
from sagemaker_ui_helper import ClientConfig

config = ClientConfig(region="us-west-2")
sagemaker_ui_helper_api = SageMakerUIHelperAPI(config)

result = sagemaker_ui_helper_api.execution_client.start_execution(
    StartExecutionRequest(
        execution_name="my-execution",
        execution_type="NOTEBOOK",
        input_config={"notebook_config": {"input_path": "src/folder2/test.ipynb"}},
        output_config={"notebook_config": {"output_formats": ["NOTEBOOK", "HTML"]}},
        termination_condition={"max_runtime_in_seconds": 9000},
        compute={
            "instance_type": "ml.c5.xlarge",
            "image_details": {
                # provide either ecr_uri or (image_name and image_version)
"image_name": "sagemaker-distribution-2.2",
                "image_version": "2.2",
                "ecr_uri": "123456123456.dkr.ecr.us-west-2.amazonaws.com/sagemaker-distribution-embargoed-prod:3.0.0-reinvent2024-cpu"
            }
        }
    )
)
print(result)
                            

GetExecution

You can retrieve details about an execution running on remote compute using the GetExecution API.

from sagemaker_ui_helper.sagemaker_ui_helper_api import SageMakerUIHelperAPI
from sagemaker_ui_helper._openapi.models import (
    GetExecutionRequest,
)
from sagemaker_ui_helper import ClientConfig

config = ClientConfig(region="us-west-2")
sagemaker_ui_helper_api = SageMakerUIHelperAPI(config)

get_response = sagemaker_ui_helper_api.execution_client.get_execution(GetExecutionRequest(execution_id="asdf-3b998be2-02dd-42af-8802-593d48d04daa"))
print(get_response)
                            

ListExecutions

You can use the ListExecutions API to list all the headless executions that ran on remote compute.

from sagemaker_ui_helper.sagemaker_ui_helper_api import SageMakerUIHelperAPI
from sagemaker_ui_helper._openapi.models import (
    ListExecutionsRequest,
)
from sagemaker_ui_helper import ClientConfig

config = ClientConfig(region="us-west-2")
sagemaker_ui_helper_api = SageMakerUIHelperAPI(config)

list_executions_response = sagemaker_ui_helper_api.execution_client.list_executions(ListExecutionsRequest(status="COMPLETED"))
print(list_executions_response)
                            

StopExecution

You can use the StopExecution API to stop an execution that's running on remote compute.

from sagemaker_ui_helper.sagemaker_ui_helper_api import SageMakerUIHelperAPI
from sagemaker_ui_helper._openapi.models import (
    StopExecutionRequest,
)
from sagemaker_ui_helper import ClientConfig

config = ClientConfig(region="us-west-2")
sagemaker_ui_helper_api = SageMakerUIHelperAPI(config)

stop_response = sagemaker_ui_helper_api.execution_client.stop_execution(StopExecutionRequest(execution_id="asdf-3b998be2-02dd-42af-8802-593d48d04daa"))
print(stop_response)