diff --git a/src/runloop_api_client/resources/blueprints.py b/src/runloop_api_client/resources/blueprints.py index 124f0a139..b471f11f7 100644 --- a/src/runloop_api_client/resources/blueprints.py +++ b/src/runloop_api_client/resources/blueprints.py @@ -20,6 +20,9 @@ async_to_raw_response_wrapper, async_to_streamed_response_wrapper, ) +from .._exceptions import RunloopError +from ..lib.polling import PollingConfig, poll_until +from ..lib.polling_async import async_poll_until from ..pagination import SyncBlueprintsCursorIDPage, AsyncBlueprintsCursorIDPage from .._base_client import AsyncPaginator, make_request_options from ..types.blueprint_view import BlueprintView @@ -154,6 +157,119 @@ def retrieve( cast_to=BlueprintView, ) + def await_build_complete( + self, + id: str, + *, + polling_config: PollingConfig | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + ) -> BlueprintView: + """Wait for a blueprint to finish building. + + Args: + id: The ID of the blueprint to wait for + polling_config: Optional polling configuration + extra_headers: Send extra headers + extra_query: Add additional query parameters to the request + extra_body: Add additional JSON properties to the request + timeout: Override the client-level default timeout for this request, in seconds + + Returns: + The blueprint in built state + + Raises: + PollingTimeout: If polling times out before blueprint is built + RunloopError: If blueprint enters a non-built terminal state + """ + def retrieve_blueprint() -> BlueprintView: + return self.retrieve( + id, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout + ) + + def is_done_building(blueprint: BlueprintView) -> bool: + return blueprint.status not in ["building", "provisioning"] + + blueprint = poll_until(retrieve_blueprint, is_done_building, polling_config) + + if blueprint.status != "build_complete": + raise RunloopError( + f"Blueprint entered non-built terminal state: {blueprint.status}" + ) + + return blueprint + + def create_and_await_build_complete( + self, + *, + name: str, + code_mounts: Optional[Iterable[CodeMountParameters]] | NotGiven = NOT_GIVEN, + dockerfile: Optional[str] | NotGiven = NOT_GIVEN, + file_mounts: Optional[Dict[str, str]] | NotGiven = NOT_GIVEN, + launch_parameters: Optional[LaunchParameters] | NotGiven = NOT_GIVEN, + system_setup_commands: Optional[List[str]] | NotGiven = NOT_GIVEN, + polling_config: PollingConfig | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + idempotency_key: str | None = None, + ) -> BlueprintView: + """Create a new Blueprint and wait for it to finish building. + + Args: + dockerfile: The Dockerfile contents to use for building the Blueprint + file_mounts: Files to mount into the Blueprint + launch_parameters: Launch parameters for Devboxes created from this Blueprint + name: Name for the Blueprint + system_setup_commands: Commands to run during Blueprint build + polling_config: Optional polling configuration + extra_headers: Send extra headers + extra_query: Add additional query parameters to the request + extra_body: Add additional JSON properties to the request + timeout: Override the client-level default timeout for this request, in seconds + idempotency_key: Specify a custom idempotency key for this request + + Returns: + The built blueprint + + Raises: + PollingTimeout: If polling times out before blueprint is built + RunloopError: If blueprint enters a non-built terminal state + """ + blueprint = self.create( + name=name, + dockerfile=dockerfile, + code_mounts=code_mounts, + file_mounts=file_mounts, + launch_parameters=launch_parameters, + system_setup_commands=system_setup_commands, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + idempotency_key=idempotency_key, + ) + + return self.await_build_complete( + blueprint.id, + polling_config=polling_config, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + def list( self, *, @@ -429,6 +545,119 @@ async def retrieve( ), cast_to=BlueprintView, ) + + async def await_build_complete( + self, + id: str, + *, + polling_config: PollingConfig | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + ) -> BlueprintView: + """Wait for a blueprint to finish building. + + Args: + id: The ID of the blueprint to wait for + polling_config: Optional polling configuration + extra_headers: Send extra headers + extra_query: Add additional query parameters to the request + extra_body: Add additional JSON properties to the request + timeout: Override the client-level default timeout for this request, in seconds + + Returns: + The blueprint in built state + + Raises: + PollingTimeout: If polling times out before blueprint is built + RunloopError: If blueprint enters a non-built terminal state + """ + async def retrieve_blueprint() -> BlueprintView: + return await self.retrieve( + id, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout + ) + + def is_done_building(blueprint: BlueprintView) -> bool: + return blueprint.status not in ["building", "provisioning"] + + blueprint = await async_poll_until(retrieve_blueprint, is_done_building, polling_config) + + if blueprint.status != "build_complete": + raise RunloopError( + f"Blueprint entered non-built terminal state: {blueprint.status}" + ) + + return blueprint + + async def create_and_await_build_complete( + self, + *, + name: str, + code_mounts: Optional[Iterable[CodeMountParameters]] | NotGiven = NOT_GIVEN, + dockerfile: Optional[str] | NotGiven = NOT_GIVEN, + file_mounts: Optional[Dict[str, str]] | NotGiven = NOT_GIVEN, + launch_parameters: Optional[LaunchParameters] | NotGiven = NOT_GIVEN, + system_setup_commands: Optional[List[str]] | NotGiven = NOT_GIVEN, + polling_config: PollingConfig | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + idempotency_key: str | None = None, + ) -> BlueprintView: + """Create a new Blueprint and wait for it to finish building. + + Args: + dockerfile: The Dockerfile contents to use for building the Blueprint + file_mounts: Files to mount into the Blueprint + launch_parameters: Launch parameters for Devboxes created from this Blueprint + name: Name for the Blueprint + system_setup_commands: Commands to run during Blueprint build + polling_config: Optional polling configuration + extra_headers: Send extra headers + extra_query: Add additional query parameters to the request + extra_body: Add additional JSON properties to the request + timeout: Override the client-level default timeout for this request, in seconds + idempotency_key: Specify a custom idempotency key for this request + + Returns: + The built blueprint + + Raises: + PollingTimeout: If polling times out before blueprint is built + RunloopError: If blueprint enters a non-built terminal state + """ + blueprint = await self.create( + name=name, + dockerfile=dockerfile, + code_mounts=code_mounts, + file_mounts=file_mounts, + launch_parameters=launch_parameters, + system_setup_commands=system_setup_commands, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + idempotency_key=idempotency_key, + ) + + return await self.await_build_complete( + blueprint.id, + polling_config=polling_config, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) def list( self, diff --git a/src/runloop_api_client/resources/scenarios/scenarios.py b/src/runloop_api_client/resources/scenarios/scenarios.py index 2501f4986..88d6b9563 100644 --- a/src/runloop_api_client/resources/scenarios/scenarios.py +++ b/src/runloop_api_client/resources/scenarios/scenarios.py @@ -25,6 +25,7 @@ maybe_transform, async_maybe_transform, ) +from ...lib.polling import PollingConfig from ..._compat import cached_property from ..._resource import SyncAPIResource, AsyncAPIResource from ..._response import ( @@ -317,6 +318,63 @@ def start_run( cast_to=ScenarioRunView, ) + def start_run_and_await_env_ready( + self, + *, + scenario_id: str, + benchmark_run_id: Optional[str] | NotGiven = NOT_GIVEN, + run_name: Optional[str] | NotGiven = NOT_GIVEN, + polling_config: PollingConfig | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + idempotency_key: str | None = None, + ) -> ScenarioRunView: + """Start a new ScenarioRun and wait for its environment to be ready. + + Args: + scenario_id: ID of the Scenario to run + benchmark_run_id: Benchmark to associate the run + run_name: Display name of the run + polling_config: Optional polling configuration + extra_headers: Send extra headers + extra_query: Add additional query parameters to the request + extra_body: Add additional JSON properties to the request + timeout: Override the client-level default timeout for this request, in seconds + idempotency_key: Specify a custom idempotency key for this request + + Returns: + The scenario run in running state + + Raises: + PollingTimeout: If polling times out before environment is ready + RunloopError: If environment enters a non-running terminal state + """ + run = self.start_run( + scenario_id=scenario_id, + benchmark_run_id=benchmark_run_id, + run_name=run_name, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + idempotency_key=idempotency_key, + ) + + self._client.devboxes.await_running( + run.devbox_id, + polling_config=polling_config, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + + return run + class AsyncScenariosResource(AsyncAPIResource): @cached_property @@ -591,6 +649,40 @@ async def start_run( cast_to=ScenarioRunView, ) + async def start_run_and_await_env_ready( + self, + scenario_id: str, + benchmark_run_id: Optional[str] | NotGiven = NOT_GIVEN, + run_name: Optional[str] | NotGiven = NOT_GIVEN, + polling_config: PollingConfig | None = None, + ) -> ScenarioRunView: + """Start a new ScenarioRun and wait for its environment to be ready. + + Args: + scenario_id: ID of the Scenario to run + benchmark_run_id: Benchmark to associate the run + run_name: Display name of the run + polling_config: Optional polling configuration + + Returns: + The scenario run in running state + + Raises: + PollingTimeout: If polling times out before environment is ready + RunloopError: If environment enters a non-running terminal state + """ + run = await self.start_run( + scenario_id=scenario_id, + benchmark_run_id=benchmark_run_id, + run_name=run_name, + ) + + await self._client.devboxes.await_running( + run.devbox_id, + polling_config=polling_config, + ) + + return run class ScenariosResourceWithRawResponse: def __init__(self, scenarios: ScenariosResource) -> None: