feat: add helpers for blueprint and scenario run creation

Author: evan-runloopai

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,

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: