Skip to content

Commit

Permalink
fix for beam shell with pod specs
Browse files Browse the repository at this point in the history
  • Loading branch information
Luke Lombardi committed Mar 4, 2025
1 parent 78fe5b6 commit 8d4f79c
Show file tree
Hide file tree
Showing 4 changed files with 8 additions and 4 deletions.
1 change: 1 addition & 0 deletions sdk/src/beta9/abstractions/pod.py
Original file line number Diff line number Diff line change
Expand Up @@ -299,6 +299,7 @@ def cleanup_deployment_artifacts(self):

@with_grpc_error_handling
def shell(self, url_type: str = ""):
self.authorized = True
stub_type = SHELL_STUB_TYPE

if not self.prepare_runtime(stub_type=stub_type, force_create_stub=True):
Expand Down
5 changes: 4 additions & 1 deletion sdk/src/beta9/abstractions/shell.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,9 @@ def wait_for_ok(sock: socket.socket, max_retries: int = 5, delay: float = 0.25):
raise ConnectionError(f"Failed to setup socket after {max_retries} retries")


EXIT_STATUS_CTRL_C = 130


@dataclass
class SSHShell:
"""Interactive ssh shell that can be used as a context manager - for use with 'shell' command"""
Expand Down Expand Up @@ -136,7 +139,7 @@ def start(self):

# Check the exit status after the shell session ends
exit_status = self.channel.recv_exit_status()
if exit_status != 0:
if exit_status != 0 and exit_status != EXIT_STATUS_CTRL_C:
terminal.warn("Lost connection to shell, attempting to reconnect in 5 seconds...")
time.sleep(5)

Expand Down
4 changes: 2 additions & 2 deletions sdk/src/beta9/abstractions/volume.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,9 +85,9 @@ class CloudBucketConfig:
read_only (bool):
Whether the volume is read-only.
access_key (str):
The name of the beam secret containing the S3 access key for the external provider.
The name of the secret containing the S3 access key for the external provider.
secret_key (str):
The name of the beam secret containing the S3 secret key for the external provider.
The name of the secret containing the S3 secret key for the external provider.
endpoint (Optional[str]):
The S3 endpoint for the external provider.
region (Optional[str]):
Expand Down
2 changes: 1 addition & 1 deletion sdk/src/beta9/terminal.py
Original file line number Diff line number Diff line change
Expand Up @@ -197,7 +197,7 @@ def reset_terminal() -> None:
def progress_description(name: str, max_width: Optional[int] = None):
max_desc_width = max_width or len(name)
if len(name) > max_desc_width:
text = f"...{name[-(max_desc_width - 3):]}"
text = f"...{name[-(max_desc_width - 3) :]}"
else:
text = name.ljust(max_desc_width)

Expand Down

0 comments on commit 8d4f79c

Please sign in to comment.