From 0514147f2b405e12c94ccc17d3e95400bf421c4f Mon Sep 17 00:00:00 2001 From: Justin Lin Date: 2024年1月12日 00:21:17 -0500 Subject: [PATCH 1/4] Adding send and download to runpod pod CLI --- runpod/cli/groups/pod/commands.py | 63 +++++++++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/runpod/cli/groups/pod/commands.py b/runpod/cli/groups/pod/commands.py index de47b75e..74709465 100644 --- a/runpod/cli/groups/pod/commands.py +++ b/runpod/cli/groups/pod/commands.py @@ -2,6 +2,7 @@ RunPod | CLI | Pod | Commands ''' import click +import os from prettytable import PrettyTable from runpod import get_pods, create_pod @@ -58,3 +59,65 @@ def connect_to_pod(pod_id): click.echo(f'Connecting to pod {pod_id}...') ssh = ssh_cmd.SSHConnection(pod_id) ssh.launch_terminal() + +@pod_cli.command("send") +@click.argument("pod_id", required=True) +@click.argument( + "local_path", + required=True, + type=click.Path(exists=True, file_okay=True, dir_okay=False), +) +@click.argument("remote_path", required=True) +def send_file(pod_id, local_path, remote_path): + """ + Send a local file to a specified pod. + ... + """ + try: + absolute_local_path = os.path.abspath(local_path) + + if not os.path.isfile(absolute_local_path): + raise ValueError(f"The local path '{absolute_local_path}' is not a file.") + + # Assuming the remote path is relative to the user's home directory + remote_directory = os.path.dirname(remote_path) + if remote_directory.startswith("."): + remote_directory = remote_directory[1:] # Remove './' if present + remote_directory_display = f"{remote_directory}" if remote_directory else "~" + + click.echo( + f"Sending file from {absolute_local_path} to pod {pod_id}:{remote_path}..." + ) + with ssh_cmd.SSHConnection(pod_id) as ssh: + ssh.put_file(absolute_local_path, remote_path) + click.echo( + f"File sent successfully to {remote_directory_display} on pod {pod_id}." + ) + click.echo(f"To access the file, use: cd {remote_directory_display}. Type pwd to make sure you get put in the right directory.") + + except Exception as e: + click.echo(f"Failed to send file: {e}", err=True) + + +@pod_cli.command("download") +@click.argument("pod_id", required=True) +@click.argument("remote_path", required=True) +@click.argument("local_path", required=True) +def download_file(pod_id, remote_path, local_path): + """ + Download a file from a specified pod to local machine. + ... + """ + try: + absolute_local_path = os.path.abspath(local_path) + + click.echo( + f"Downloading file from pod {pod_id}:{remote_path} to {absolute_local_path}..." + ) + with ssh_cmd.SSHConnection(pod_id) as ssh: + ssh.get_file(remote_path, absolute_local_path) + click.echo(f"File downloaded successfully to {absolute_local_path}.") + + except Exception as e: + click.echo(f"Failed to download file: {e}", err=True) + click.echo(f"Ensure that the remote path exists on pod {pod_id}. \nAnd that the arguments are correct. \nFor example: runpod pod download 1234 /home/REMOTE_POD_PATH/file.txt /home/LOCAL_PATH/file.txt") \ No newline at end of file From 54b342535616beab8564abcfb0f49faa6b95bee2 Mon Sep 17 00:00:00 2001 From: Justin Lin Date: 2024年1月12日 00:23:08 -0500 Subject: [PATCH 2/4] Adjusting terminal messaging --- runpod/cli/groups/pod/commands.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runpod/cli/groups/pod/commands.py b/runpod/cli/groups/pod/commands.py index 74709465..c331fa9f 100644 --- a/runpod/cli/groups/pod/commands.py +++ b/runpod/cli/groups/pod/commands.py @@ -93,7 +93,7 @@ def send_file(pod_id, local_path, remote_path): click.echo( f"File sent successfully to {remote_directory_display} on pod {pod_id}." ) - click.echo(f"To access the file, use: cd {remote_directory_display}. Type pwd to make sure you get put in the right directory.") + click.echo(f"To access the file, type in the terminal on your pod: cd {remote_directory_display}. Type pwd to make sure you get put in the right directory.") except Exception as e: click.echo(f"Failed to send file: {e}", err=True) From 1265c80f96399e1e2a7d703664dc007a649ec023 Mon Sep 17 00:00:00 2001 From: Justin Lin Date: 2024年1月12日 01:06:55 -0500 Subject: [PATCH 3/4] Adding callback for progress indicator --- runpod/cli/groups/pod/commands.py | 18 ++++++++++++------ runpod/cli/utils/ssh_cmd.py | 12 ++++++------ 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/runpod/cli/groups/pod/commands.py b/runpod/cli/groups/pod/commands.py index c331fa9f..d307bcfa 100644 --- a/runpod/cli/groups/pod/commands.py +++ b/runpod/cli/groups/pod/commands.py @@ -4,11 +4,16 @@ import click import os from prettytable import PrettyTable - +import sys from runpod import get_pods, create_pod from ...utils import ssh_cmd +def sftp_progress_callback(transferred, total): + progress_percentage = (transferred / total) * 100 + sys.stdout.write(f'\rTransferring... {progress_percentage:.2f}% \n') + sys.stdout.flush() + @click.group('pod', help='Manage and interact with pods.') def pod_cli(): '''A collection of CLI functions for Pod.''' @@ -83,20 +88,21 @@ def send_file(pod_id, local_path, remote_path): remote_directory = os.path.dirname(remote_path) if remote_directory.startswith("."): remote_directory = remote_directory[1:] # Remove './' if present - remote_directory_display = f"{remote_directory}" if remote_directory else "~" + remote_directory_display = f"{remote_directory}" if remote_directory else "/" click.echo( f"Sending file from {absolute_local_path} to pod {pod_id}:{remote_path}..." ) with ssh_cmd.SSHConnection(pod_id) as ssh: - ssh.put_file(absolute_local_path, remote_path) + ssh.put_file(absolute_local_path, remote_path, callback=sftp_progress_callback) click.echo( - f"File sent successfully to {remote_directory_display} on pod {pod_id}." + f"File sent successfully to directory {remote_directory_display} on pod {pod_id}." ) - click.echo(f"To access the file, type in the terminal on your pod: cd {remote_directory_display}. Type pwd to make sure you get put in the right directory.") + click.echo(f"To access the file, type in the terminal on your pod: cd {remote_directory_display}. If you ls, the file should be there. Type pwd to make sure you get put in the right directory.") except Exception as e: click.echo(f"Failed to send file: {e}", err=True) + click.echo(f"Common reason for failure: a directory in '{remote_directory_display}' does not exist on pod {pod_id}.") @pod_cli.command("download") @@ -115,7 +121,7 @@ def download_file(pod_id, remote_path, local_path): f"Downloading file from pod {pod_id}:{remote_path} to {absolute_local_path}..." ) with ssh_cmd.SSHConnection(pod_id) as ssh: - ssh.get_file(remote_path, absolute_local_path) + ssh.get_file(remote_path, absolute_local_path, sftp_progress_callback) click.echo(f"File downloaded successfully to {absolute_local_path}.") except Exception as e: diff --git a/runpod/cli/utils/ssh_cmd.py b/runpod/cli/utils/ssh_cmd.py index e23ca6ec..f59053d3 100644 --- a/runpod/cli/utils/ssh_cmd.py +++ b/runpod/cli/utils/ssh_cmd.py @@ -89,15 +89,15 @@ def handle_stream(stream, color, prefix): stdout_thread.join() stderr_thread.join() - def put_file(self, local_path, remote_path): - ''' Copy local file to remote machine over SSH. ''' + def put_file(self, local_path, remote_path, callback=None): + ''' Copy local file to remote machine over SSH with optional progress callback. ''' with self.ssh.open_sftp() as sftp: - sftp.put(local_path, remote_path) + sftp.put(local_path, remote_path, callback=callback) - def get_file(self, remote_path, local_path): - ''' Fetch a remote file to local machine over SSH. ''' + def get_file(self, remote_path, local_path, callback=None): + ''' Fetch a remote file to local machine over SSH with optional progress callback. ''' with self.ssh.open_sftp() as sftp: - sftp.get(remote_path, local_path) + sftp.get(remote_path, local_path, callback=callback) def launch_terminal(self): ''' Launch an interactive terminal over SSH. ''' From 08c0d8a05710c51b0876a9c46ff0ae9939741940 Mon Sep 17 00:00:00 2001 From: Justin Lin Date: 2024年1月31日 14:30:28 -0500 Subject: [PATCH 4/4] Fixing pylint errors --- runpod/cli/groups/pod/commands.py | 113 +++++++++++++++++++----------- 1 file changed, 73 insertions(+), 40 deletions(-) diff --git a/runpod/cli/groups/pod/commands.py b/runpod/cli/groups/pod/commands.py index d307bcfa..dc181615 100644 --- a/runpod/cli/groups/pod/commands.py +++ b/runpod/cli/groups/pod/commands.py @@ -1,70 +1,89 @@ -''' +""" RunPod | CLI | Pod | Commands -''' -import click +""" import os -from prettytable import PrettyTable import sys +import click +from prettytable import PrettyTable from runpod import get_pods, create_pod from ...utils import ssh_cmd +# pylint: disable=line-too-long def sftp_progress_callback(transferred, total): + """ + Callback function for SFTP transfers to display progress. + """ progress_percentage = (transferred / total) * 100 - sys.stdout.write(f'\rTransferring... {progress_percentage:.2f}% \n') + sys.stdout.write(f"\rTransferring... {progress_percentage:.2f}% \n") sys.stdout.flush() -@click.group('pod', help='Manage and interact with pods.') + +@click.group("pod", help="Manage and interact with pods.") def pod_cli(): - '''A collection of CLI functions for Pod.''' + """A collection of CLI functions for Pod.""" -@pod_cli.command('list') + +@pod_cli.command("list") def list_pods(): - ''' + """ Lists the pods for the current user. - ''' - table = PrettyTable(['ID', 'Name', 'Status', 'Image']) + """ + table = PrettyTable(["ID", "Name", "Status", "Image"]) for pod in get_pods(): - table.add_row((pod['id'], pod['name'], pod['desiredStatus'], pod['imageName'])) + table.add_row((pod["id"], pod["name"], pod["desiredStatus"], pod["imageName"])) click.echo(table) -@pod_cli.command('create') -@click.argument('name', required=False) -@click.option('--image', default=None, help='The image to use for the pod.') -@click.option('--gpu-type', default=None, help='The GPU type to use for the pod.') -@click.option('--gpu-count', default=1, help='The number of GPUs to use for the pod.') -@click.option('--support-public-ip', default=True, help='Whether or not to support a public IP.') -def create_new_pod(name, image, gpu_type, gpu_count, support_public_ip): # pylint: disable=too-many-arguments - ''' + +@pod_cli.command("create") +@click.argument("name", required=False) +@click.option("--image", default=None, help="The image to use for the pod.") +@click.option("--gpu-type", default=None, help="The GPU type to use for the pod.") +@click.option("--gpu-count", default=1, help="The number of GPUs to use for the pod.") +@click.option( + "--support-public-ip", default=True, help="Whether or not to support a public IP." +) +def create_new_pod( + name, image, gpu_type, gpu_count, support_public_ip +): # pylint: disable=too-many-arguments + """ Creates a pod. - ''' + """ if not name: - name = click.prompt('Enter pod name', default='RunPod-CLI-Pod') + name = click.prompt("Enter pod name", default="RunPod-CLI-Pod") - quick_launch = click.confirm('Would you like to launch default pod?', abort=True) + quick_launch = click.confirm("Would you like to launch default pod?", abort=True) if quick_launch: - image = 'runpod/base:0.0.0' - gpu_type = 'NVIDIA GeForce RTX 3090' - ports ='22/tcp' + image = "runpod/base:0.0.0" + gpu_type = "NVIDIA GeForce RTX 3090" + ports = "22/tcp" - click.echo('Launching default pod...') + click.echo("Launching default pod...") - new_pod = create_pod(name, image, gpu_type, - gpu_count=gpu_count, support_public_ip=support_public_ip, ports=ports) + new_pod = create_pod( + name, + image, + gpu_type, + gpu_count=gpu_count, + support_public_ip=support_public_ip, + ports=ports, + ) click.echo(f'Pod {new_pod["id"]} has been created.') -@pod_cli.command('connect') -@click.argument('pod_id') + +@pod_cli.command("connect") +@click.argument("pod_id") def connect_to_pod(pod_id): - ''' + """ Connects to a pod. - ''' - click.echo(f'Connecting to pod {pod_id}...') + """ + click.echo(f"Connecting to pod {pod_id}...") ssh = ssh_cmd.SSHConnection(pod_id) ssh.launch_terminal() + @pod_cli.command("send") @click.argument("pod_id", required=True) @click.argument( @@ -94,15 +113,21 @@ def send_file(pod_id, local_path, remote_path): f"Sending file from {absolute_local_path} to pod {pod_id}:{remote_path}..." ) with ssh_cmd.SSHConnection(pod_id) as ssh: - ssh.put_file(absolute_local_path, remote_path, callback=sftp_progress_callback) + ssh.put_file( + absolute_local_path, remote_path, callback=sftp_progress_callback + ) click.echo( f"File sent successfully to directory {remote_directory_display} on pod {pod_id}." ) - click.echo(f"To access the file, type in the terminal on your pod: cd {remote_directory_display}. If you ls, the file should be there. Type pwd to make sure you get put in the right directory.") + click.echo( + f"To access the file, type in the terminal on your pod: cd {remote_directory_display}. If you ls, the file should be there. Type pwd to make sure you get put in the right directory." + ) - except Exception as e: + except FileNotFoundError as e: click.echo(f"Failed to send file: {e}", err=True) - click.echo(f"Common reason for failure: a directory in '{remote_directory_display}' does not exist on pod {pod_id}.") + click.echo( + f"Common reason for failure: a directory in '{remote_directory_display}' does not exist on pod {pod_id}." + ) @pod_cli.command("download") @@ -124,6 +149,14 @@ def download_file(pod_id, remote_path, local_path): ssh.get_file(remote_path, absolute_local_path, sftp_progress_callback) click.echo(f"File downloaded successfully to {absolute_local_path}.") - except Exception as e: + except FileNotFoundError as e: click.echo(f"Failed to download file: {e}", err=True) - click.echo(f"Ensure that the remote path exists on pod {pod_id}. \nAnd that the arguments are correct. \nFor example: runpod pod download 1234 /home/REMOTE_POD_PATH/file.txt /home/LOCAL_PATH/file.txt") \ No newline at end of file + click.echo( + f"Ensure that the remote path exists on pod {pod_id} and args are correct." + ) + + except PermissionError as e: + click.echo(f"Failed to download file: {e}", err=True) + click.echo( + f"Ensure that you have the necessary permissions to access the file on pod {pod_id}." + )

AltStyle によって変換されたページ (->オリジナル) /