blaxel 0.2.32__py3-none-any.whl → 0.2.34__py3-none-any.whl

blaxel/core/sandbox/default/process.py

@@ -184,55 +184,170 @@ class SandboxProcess(SandboxAction):
     ) -> Union[ProcessResponse, ProcessResponseWithLog]:
         """Execute a process in the sandbox."""
         on_log = None
+        on_stdout = None
+        on_stderr = None
+
         if isinstance(process, ProcessRequestWithLog):
             on_log = process.on_log
+            on_stdout = process.on_stdout
+            on_stderr = process.on_stderr
             process = process.to_dict()
 
         if isinstance(process, dict):
             if "on_log" in process:
                 on_log = process["on_log"]
                 del process["on_log"]
+            if "on_stdout" in process:
+                on_stdout = process["on_stdout"]
+                del process["on_stdout"]
+            if "on_stderr" in process:
+                on_stderr = process["on_stderr"]
+                del process["on_stderr"]
             process = ProcessRequest.from_dict(process)
 
         # Store original wait_for_completion setting
         should_wait_for_completion = process.wait_for_completion
 
-        # Always start process without wait_for_completion to avoid server-side blocking
-        if should_wait_for_completion and on_log is not None:
-            process.wait_for_completion = False
-
-        client = self.get_client()
-        response = await client.post("/process", json=process.to_dict())
-        try:
-            content_bytes = await response.aread()
-            self.handle_response_error(response)
-            import json
-
-            response_data = json.loads(content_bytes) if content_bytes else None
-            result = ProcessResponse.from_dict(response_data)
-        finally:
-            await response.aclose()
-
-        # Handle wait_for_completion with parallel log streaming
-        if should_wait_for_completion and on_log is not None:
-            stream_control = self._stream_logs(result.pid, {"on_log": on_log})
+        # When waiting for completion with streaming callbacks, use streaming endpoint
+        if should_wait_for_completion and (on_log or on_stdout or on_stderr):
+            return await self._exec_with_streaming(
+                process, on_log=on_log, on_stdout=on_stdout, on_stderr=on_stderr
+            )
+        else:
+            client = self.get_client()
+            response = await client.post("/process", json=process.to_dict())
             try:
-                # Wait for process completion
-                result = await self.wait(result.pid, interval=500, max_wait=1000 * 60 * 60)
+                content_bytes = await response.aread()
+                self.handle_response_error(response)
+                import json
+
+                response_data = json.loads(content_bytes) if content_bytes else None
+                result = ProcessResponse.from_dict(response_data)
             finally:
-                # Clean up log streaming
-                if stream_control:
-                    stream_control["close"]()
-        else:
-            # For non-blocking execution, set up log streaming immediately if requested
-            if on_log is not None:
-                stream_control = self._stream_logs(result.pid, {"on_log": on_log})
+                await response.aclose()
+
+            if on_log or on_stdout or on_stderr:
+                stream_control = self._stream_logs(
+                    result.pid, {"on_log": on_log, "on_stdout": on_stdout, "on_stderr": on_stderr}
+                )
                 return ProcessResponseWithLog(
                     result,
                     lambda: stream_control["close"]() if stream_control else None,
                 )
 
-        return result
+            return result
+
+    async def _exec_with_streaming(
+        self,
+        process_request: ProcessRequest,
+        on_log: Callable[[str], None] | None = None,
+        on_stdout: Callable[[str], None] | None = None,
+        on_stderr: Callable[[str], None] | None = None,
+    ) -> ProcessResponseWithLog:
+        """Execute a process with streaming response handling for NDJSON."""
+        import json
+
+        headers = (
+            self.sandbox_config.headers
+            if self.sandbox_config.force_url
+            else {**settings.headers, **self.sandbox_config.headers}
+        )
+
+        async with httpx.AsyncClient() as client_instance:
+            async with client_instance.stream(
+                "POST",
+                f"{self.url}/process",
+                headers={
+                    **headers,
+                    "Content-Type": "application/json",
+                    "Accept": "text/event-stream",
+                },
+                json=process_request.to_dict(),
+                timeout=None,
+            ) as response:
+                if response.status_code >= 400:
+                    error_text = await response.aread()
+                    raise Exception(f"Failed to execute process: {error_text}")
+
+                content_type = response.headers.get("Content-Type", "")
+                is_streaming = "application/x-ndjson" in content_type
+
+                # Fallback: server doesn't support streaming, use legacy approach
+                if not is_streaming:
+                    content = await response.aread()
+                    data = json.loads(content)
+                    result = ProcessResponse.from_dict(data)
+
+                    # If process already completed (server waited), emit logs through callbacks
+                    if result.status == "completed" or result.status == "failed":
+                        if result.stdout:
+                            for line in result.stdout.split("\n"):
+                                if line:
+                                    if on_stdout:
+                                        on_stdout(line)
+                        if result.stderr:
+                            for line in result.stderr.split("\n"):
+                                if line:
+                                    if on_stderr:
+                                        on_stderr(line)
+                        if result.logs:
+                            for line in result.logs.split("\n"):
+                                if line:
+                                    if on_log:
+                                        on_log(line)
+
+                    return ProcessResponseWithLog(result, lambda: None)
+
+                # Streaming response handling
+                buffer = ""
+                result = None
+
+                async for chunk in response.aiter_text():
+                    buffer += chunk
+                    lines = buffer.split("\n")
+                    buffer = lines.pop()
+
+                    for line in lines:
+                        if not line.strip():
+                            continue
+                        try:
+                            parsed = json.loads(line)
+                            parsed_type = parsed.get("type", "")
+                            parsed_data = parsed.get("data", "")
+
+                            if parsed_type == "stdout":
+                                if parsed_data:
+                                    if on_stdout:
+                                        on_stdout(parsed_data)
+                                    if on_log:
+                                        on_log(parsed_data)
+                            elif parsed_type == "stderr":
+                                if parsed_data:
+                                    if on_stderr:
+                                        on_stderr(parsed_data)
+                                    if on_log:
+                                        on_log(parsed_data)
+                            elif parsed_type == "result":
+                                try:
+                                    result = ProcessResponse.from_dict(json.loads(parsed_data))
+                                except Exception:
+                                    raise Exception(f"Failed to parse result JSON: {parsed_data}")
+                        except json.JSONDecodeError:
+                            continue
+
+                # Process any remaining buffer
+                if buffer.strip():
+                    if buffer.startswith("result:"):
+                        json_str = buffer[7:]
+                        try:
+                            result = ProcessResponse.from_dict(json.loads(json_str))
+                        except Exception:
+                            raise Exception(f"Failed to parse result JSON: {json_str}")
+
+                if not result:
+                    raise Exception("No result received from streaming response")
+
+                return ProcessResponseWithLog(result, lambda: None)
 
     async def wait(
         self, identifier: str, max_wait: int = 60000, interval: int = 1000

blaxel/core/sandbox/default/sandbox.py

@@ -1,6 +1,6 @@
 import logging
 import uuid
-from typing import Any, Dict, List, Union
+from typing import Any, Callable, Dict, List, Union
 
 from ...client.api.compute.create_sandbox import asyncio as create_sandbox
 from ...client.api.compute.delete_sandbox import asyncio as delete_sandbox
@@ -26,6 +26,24 @@ from .session import SandboxSessions
 logger = logging.getLogger(__name__)
 
 
+class _AsyncDeleteDescriptor:
+    """Descriptor that provides both class-level and instance-level delete functionality."""
+
+    def __init__(self, delete_func: Callable):
+        self._delete_func = delete_func
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            # Called on the class: SandboxInstance.delete("name")
+            return self._delete_func
+        else:
+            # Called on an instance: instance.delete()
+            async def instance_delete() -> Sandbox:
+                return await self._delete_func(instance.metadata.name)
+
+            return instance_delete
+
+
 class SandboxInstance:
     def __init__(
         self,
@@ -107,6 +125,7 @@ class SandboxInstance:
                     or "lifecycle" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                     or "snapshot_enabled"
                     in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
+                    or "labels" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                 )
             )
         ):
@@ -135,7 +154,7 @@ class SandboxInstance:
 
             # Create full Sandbox object
             sandbox = Sandbox(
-                metadata=Metadata(name=name),
+                metadata=Metadata(name=name, labels=config.labels),
                 spec=SandboxSpec(
                     runtime=Runtime(
                         image=image,
@@ -202,14 +221,6 @@ class SandboxInstance:
         response = await list_sandboxes()
         return [cls(sandbox) for sandbox in response]
 
-    @classmethod
-    async def delete(cls, sandbox_name: str) -> Sandbox:
-        response = await delete_sandbox(
-            sandbox_name,
-            client=client,
-        )
-        return response
-
     @classmethod
     async def update_metadata(
         cls, sandbox_name: str, metadata: SandboxUpdateMetadata
@@ -319,3 +330,16 @@ class SandboxInstance:
             headers={"X-Blaxel-Preview-Token": session.token},
             params={"bl_preview_token": session.token},
         )
+
+
+async def _delete_sandbox_by_name(sandbox_name: str) -> Sandbox:
+    """Delete a sandbox by name."""
+    response = await delete_sandbox(
+        sandbox_name,
+        client=client,
+    )
+    return response
+
+
+# Assign the delete descriptor to support both class-level and instance-level calls
+SandboxInstance.delete = _AsyncDeleteDescriptor(_delete_sandbox_by_name)

blaxel/core/sandbox/sync/process.py

@@ -143,43 +143,169 @@ class SyncSandboxProcess(SyncSandboxAction):
         process: Union[ProcessRequest, ProcessRequestWithLog, Dict[str, Any]],
     ) -> Union[ProcessResponse, ProcessResponseWithLog]:
         on_log = None
+        on_stdout = None
+        on_stderr = None
+
         if isinstance(process, ProcessRequestWithLog):
             on_log = process.on_log
+            on_stdout = process.on_stdout
+            on_stderr = process.on_stderr
             process = process.to_dict()
+
         if isinstance(process, dict):
             if "on_log" in process:
                 on_log = process["on_log"]
                 del process["on_log"]
+            if "on_stdout" in process:
+                on_stdout = process["on_stdout"]
+                del process["on_stdout"]
+            if "on_stderr" in process:
+                on_stderr = process["on_stderr"]
+                del process["on_stderr"]
             process = ProcessRequest.from_dict(process)
+
         should_wait_for_completion = process.wait_for_completion
-        if should_wait_for_completion and on_log is not None:
-            process.wait_for_completion = False
-        with self.get_client() as client_instance:
-            response = client_instance.post("/process", json=process.to_dict())
-            response_data = None
-            if response.content:
-                try:
-                    response_data = response.json()
-                except Exception:
-                    self.handle_response_error(response)
-                    raise
-            self.handle_response_error(response)
-            result = ProcessResponse.from_dict(response_data)
-            if should_wait_for_completion and on_log is not None:
-                stream_control = self._stream_logs(result.pid, {"on_log": on_log})
-                try:
-                    result = self.wait(result.pid, interval=500, max_wait=1000 * 60 * 60)
-                finally:
-                    if stream_control:
-                        stream_control["close"]()
-            else:
-                if on_log is not None:
-                    stream_control = self._stream_logs(result.pid, {"on_log": on_log})
+
+        # When waiting for completion with streaming callbacks, use streaming endpoint
+        if should_wait_for_completion and (on_log or on_stdout or on_stderr):
+            return self._exec_with_streaming(
+                process, on_log=on_log, on_stdout=on_stdout, on_stderr=on_stderr
+            )
+        else:
+            with self.get_client() as client_instance:
+                response = client_instance.post("/process", json=process.to_dict())
+                response_data = None
+                if response.content:
+                    try:
+                        response_data = response.json()
+                    except Exception:
+                        self.handle_response_error(response)
+                        raise
+                self.handle_response_error(response)
+                result = ProcessResponse.from_dict(response_data)
+
+                if on_log or on_stdout or on_stderr:
+                    stream_control = self._stream_logs(
+                        result.pid, {"on_log": on_log, "on_stdout": on_stdout, "on_stderr": on_stderr}
+                    )
                     return ProcessResponseWithLog(
                         result,
                         lambda: stream_control["close"]() if stream_control else None,
                     )
-            return result
+
+                return result
+
+    def _exec_with_streaming(
+        self,
+        process_request: ProcessRequest,
+        on_log: Callable[[str], None] | None = None,
+        on_stdout: Callable[[str], None] | None = None,
+        on_stderr: Callable[[str], None] | None = None,
+    ) -> ProcessResponseWithLog:
+        """Execute a process with streaming response handling for NDJSON."""
+        import json
+
+        headers = (
+            self.sandbox_config.headers
+            if self.sandbox_config.force_url
+            else {**settings.headers, **self.sandbox_config.headers}
+        )
+
+        with httpx.Client() as client_instance:
+            with client_instance.stream(
+                "POST",
+                f"{self.url}/process",
+                headers={
+                    **headers,
+                    "Content-Type": "application/json",
+                    "Accept": "text/event-stream",
+                },
+                json=process_request.to_dict(),
+                timeout=None,
+            ) as response:
+                if response.status_code >= 400:
+                    error_text = response.read()
+                    raise Exception(f"Failed to execute process: {error_text}")
+
+                content_type = response.headers.get("Content-Type", "")
+                is_streaming = "application/x-ndjson" in content_type
+
+                # Fallback: server doesn't support streaming, use legacy approach
+                if not is_streaming:
+                    content = response.read()
+                    data = json.loads(content)
+                    result = ProcessResponse.from_dict(data)
+
+                    # If process already completed (server waited), emit logs through callbacks
+                    if result.status == "completed" or result.status == "failed":
+                        if result.stdout:
+                            for line in result.stdout.split("\n"):
+                                if line:
+                                    if on_stdout:
+                                        on_stdout(line)
+                        if result.stderr:
+                            for line in result.stderr.split("\n"):
+                                if line:
+                                    if on_stderr:
+                                        on_stderr(line)
+                        if result.logs:
+                            for line in result.logs.split("\n"):
+                                if line:
+                                    if on_log:
+                                        on_log(line)
+
+                    return ProcessResponseWithLog(result, lambda: None)
+
+                # Streaming response handling
+                buffer = ""
+                result = None
+
+                for chunk in response.iter_text():
+                    buffer += chunk
+                    lines = buffer.split("\n")
+                    buffer = lines.pop()
+
+                    for line in lines:
+                        if not line.strip():
+                            continue
+                        try:
+                            parsed = json.loads(line)
+                            parsed_type = parsed.get("type", "")
+                            parsed_data = parsed.get("data", "")
+
+                            if parsed_type == "stdout":
+                                if parsed_data:
+                                    if on_stdout:
+                                        on_stdout(parsed_data)
+                                    if on_log:
+                                        on_log(parsed_data)
+                            elif parsed_type == "stderr":
+                                if parsed_data:
+                                    if on_stderr:
+                                        on_stderr(parsed_data)
+                                    if on_log:
+                                        on_log(parsed_data)
+                            elif parsed_type == "result":
+                                try:
+                                    result = ProcessResponse.from_dict(json.loads(parsed_data))
+                                except Exception:
+                                    raise Exception(f"Failed to parse result JSON: {parsed_data}")
+                        except json.JSONDecodeError:
+                            continue
+
+                # Process any remaining buffer
+                if buffer.strip():
+                    if buffer.startswith("result:"):
+                        json_str = buffer[7:]
+                        try:
+                            result = ProcessResponse.from_dict(json.loads(json_str))
+                        except Exception:
+                            raise Exception(f"Failed to parse result JSON: {json_str}")
+
+                if not result:
+                    raise Exception("No result received from streaming response")
+
+                return ProcessResponseWithLog(result, lambda: None)
 
     def wait(self, identifier: str, max_wait: int = 60000, interval: int = 1000) -> ProcessResponse:
         start_time = time.monotonic() * 1000

blaxel/core/sandbox/sync/sandbox.py

@@ -1,6 +1,6 @@
 import logging
 import uuid
-from typing import Any, Dict, List, Union
+from typing import Any, Callable, Dict, List, Union
 
 from ...client.api.compute.create_sandbox import sync as create_sandbox
 from ...client.api.compute.delete_sandbox import sync as delete_sandbox
@@ -26,6 +26,24 @@ from .session import SyncSandboxSessions
 logger = logging.getLogger(__name__)
 
 
+class _SyncDeleteDescriptor:
+    """Descriptor that provides both class-level and instance-level delete functionality."""
+
+    def __init__(self, delete_func: Callable):
+        self._delete_func = delete_func
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            # Called on the class: SyncSandboxInstance.delete("name")
+            return self._delete_func
+        else:
+            # Called on an instance: instance.delete()
+            def instance_delete() -> Sandbox:
+                return self._delete_func(instance.metadata.name)
+
+            return instance_delete
+
+
 class SyncSandboxInstance:
     def __init__(
         self,
@@ -102,6 +120,7 @@ class SyncSandboxInstance:
                     or "lifecycle" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                     or "snapshot_enabled"
                     in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
+                    or "labels" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                 )
             )
         ):
@@ -125,7 +144,7 @@ class SyncSandboxInstance:
             region = config.region
             lifecycle = config.lifecycle
             sandbox = Sandbox(
-                metadata=Metadata(name=name),
+                metadata=Metadata(name=name, labels=config.labels),
                 spec=SandboxSpec(
                     runtime=Runtime(
                         image=image,
@@ -184,14 +203,6 @@ class SyncSandboxInstance:
         response = list_sandboxes(client=client)
         return [cls(sandbox) for sandbox in response]
 
-    @classmethod
-    def delete(cls, sandbox_name: str) -> Sandbox:
-        response = delete_sandbox(
-            sandbox_name,
-            client=client,
-        )
-        return response
-
     @classmethod
     def update_metadata(
         cls, sandbox_name: str, metadata: SandboxUpdateMetadata
@@ -261,3 +272,16 @@ class SyncSandboxInstance:
             headers={"X-Blaxel-Preview-Token": session.token},
             params={"bl_preview_token": session.token},
         )
+
+
+def _delete_sandbox_by_name(sandbox_name: str) -> Sandbox:
+    """Delete a sandbox by name."""
+    response = delete_sandbox(
+        sandbox_name,
+        client=client,
+    )
+    return response
+
+
+# Assign the delete descriptor to support both class-level and instance-level calls
+SyncSandboxInstance.delete = _SyncDeleteDescriptor(_delete_sandbox_by_name)

blaxel/core/sandbox/types.py

@@ -155,6 +155,7 @@ class SandboxCreateConfiguration:
         region: str | None = None,
         lifecycle: Union[SandboxLifecycle, Dict[str, Any]] | None = None,
         snapshot_enabled: bool | None = None,
+        labels: Dict[str, str] | None = None,
     ):
         self.name = name
         self.image = image
@@ -167,6 +168,7 @@ class SandboxCreateConfiguration:
         self.region = region
         self.lifecycle = lifecycle
         self.snapshot_enabled = snapshot_enabled
+        self.labels = labels
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "SandboxCreateConfiguration":
@@ -190,6 +192,7 @@ class SandboxCreateConfiguration:
             region=data.get("region"),
             lifecycle=lifecycle,
             snapshot_enabled=data.get("snapshot_enabled"),
+            labels=data.get("labels"),
         )
 
     def _normalize_ports(self) -> List[Port] | None:
@@ -282,6 +285,8 @@ class SandboxCreateConfiguration:
 @_attrs_define
 class ProcessRequestWithLog(ProcessRequest):
     on_log: Callable[[str], None] | None = None
+    on_stdout: Callable[[str], None] | None = None
+    on_stderr: Callable[[str], None] | None = None
 
 
 class ProcessResponseWithLog:

blaxel/core/volume/__init__.py

@@ -1,5 +1,5 @@
 """Volume module for persistent storage management."""
 
-from .volume import VolumeCreateConfiguration, VolumeInstance
+from .volume import SyncVolumeInstance, VolumeCreateConfiguration, VolumeInstance
 
-__all__ = ["VolumeInstance", "VolumeCreateConfiguration"]
+__all__ = ["VolumeInstance", "SyncVolumeInstance", "VolumeCreateConfiguration"]