diff --git a/src/remote_jupyter/kube_exec.py b/src/remote_jupyter/kube_exec.py deleted file mode 100644 index 38213587eca1ff9e1caf215cbd044b0e23b09e8c..0000000000000000000000000000000000000000 --- a/src/remote_jupyter/kube_exec.py +++ /dev/null @@ -1,197 +0,0 @@ -from pathlib import Path -import sys -import time -import uuid - -from .remote import get_client_v1, connect_pod, delete_pod, copy_file, create_kernel, execute_python, list_pods - -from IPython.core.magic import ( - Magics, - cell_magic, - line_magic, - magics_class, - needs_local_scope, -) -from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring - -CURRENT_PATH = Path(__file__).parent.resolve() - -@magics_class -class ExternalKubernetesMagic(Magics): - - def __init__(self, shell): - # You must call the parent constructor - super().__init__(shell) - self.client = None - self.name = None - self.namespace = None - self.session = uuid.uuid1().hex - - def prepare_pod(self, result): - if result['success']: - print(f'Container is running: {self.name}') - print('Transfer utilities ... ', end='') - copy_file( - self.client, - namespace=self.namespace, - pod_name=self.name, - source_file=CURRENT_PATH / 'create_kernel.py', - dest_path='/tmp/' - ) - copy_file( - self.client, - namespace=self.namespace, - pod_name=self.name, - source_file=CURRENT_PATH / 'execute_in_kernel.py', - dest_path='/tmp/' - ) - print('Done.') - # kernel is a list [kernel id, base url, token] - self.kernel = create_kernel(self.client, self.name, self.namespace) - print(self.kernel) - print(f'Created kernel.') - print('Connected.') - else: - # delete container? - print(f'Failed to connect: {result["message"]}', file=sys.stderr) - - @needs_local_scope - @line_magic("kube_exec") - @cell_magic("kube_exec") - @magic_arguments() - @argument( - "line", - default="", - type=str, - nargs='?', - # choices=[ - # '', 'execute', - # 'launch', - # 'connect', - # 'list', - # 'destroy', - # ], - help="Command") - @argument( - "-k", - "--kubeconfig", - default='.kube/config', - type=str, - help="Path to the kubeconfig file, default '.kube/config'", - ) - @argument( - "-ns", - "--namespace", - default='default', - type=str, - help="Kubernetes namespace", - ) - @argument( - "-p", - "--pod-name", - type=str, - help="Name of a running pod to connect to", - ) - @argument( - "-pp", - "--pod-prefix", - type=str, - default='', - help="Pod name prefix for new pod (optional)" - ) - @argument( - "-i", - "--image", - type=str, - help="Pod image (must be running Jupyter on launch)", - ) - @argument( - "-c", - "--cores", - type=int, - default=2, - help="Number of cores in remote executor pod, default 2 cores", - ) - @argument( - "-m", - "--memory", - type=int, - default=2, - help="Memory in remote executor pod, in GiB, default 2 (GiB)", - ) - @argument( - "-g", - "--gpus", - type=int, - default=1, - help="Number of GPUs in remote executor pod, default 1.", - ) - @argument( - "-gp", - "--gpu-product", - type=str, - choices=[ - 'NVIDIA-A100-SXM4-40GB-MIG-1g.5gb', - 'NVIDIA-A100-SXM4-40GB-MIG-3g.20gb', - 'NVIDIA-A100-SXM4-40GB', - 'NVIDIA-A100-SXM4-80GB', - 'NVIDIA-H100-80GB-HBM3', - ], - default='NVIDIA-A100-SXM4-40GB-MIG-1g.5gb', - help="GPU selector for remote executor pod", - ) - def execute(self, line="", cell="", local_ns=None): - args = parse_argstring(self.execute, line) - # print(args) - if args.line == 'launch': - if not Path(args.kubeconfig).is_file(): - print(f'Error: Invalid kube-config file. No configuration found at {args.kubeconfig}', file=sys.stderr) - return - self.name = f'{args.pod_prefix}{uuid.uuid4().hex}' - self.namespace = args.namespace - print(f'Creating pod {self.name}') - result = connect_pod(self.name, args) - self.client = result['client'] - self.prepare_pod(result) - elif args.line == 'connect': - # connect to running container - if not Path(args.kubeconfig).is_file(): - print(f'Error: Invalid kube-config file. No configuration found at {args.kubeconfig}', file=sys.stderr) - return - self.name = args.pod_name - self.namespace = args.namespace - # check the container is running - result = connect_pod(self.name, args, launch=False) - self.client = result['client'] - self.prepare_pod(result) - elif args.line == 'destroy': - # delete the container - if not Path(args.kubeconfig).is_file(): - print(f'Error: Invalid kube-config file. No configuration found at {args.kubeconfig}', file=sys.stderr) - return - print('Destroying container') - delete_pod(self.client, self.name, self.namespace) - self.name = None - # create a new session id - self.session = uuid.uuid1().hex - elif args.line == 'list': - if not Path(args.kubeconfig).is_file(): - print(f'Error: Invalid kube-config file. No configuration found at {args.kubeconfig}', file=sys.stderr) - return - self.client = get_client_v1(args) - return list_pods(self.client, args.namespace) - else: - code = cell or line - if self.client is None or self.name is None or self.namespace is None: - raise Exception('No remote execution container. Please launch a container first.') - result = execute_python( - self.client, - pod_name=self.name, - namespace=self.namespace, - kernel=self.kernel, - session=self.session, - code=code, - ) - return result - - \ No newline at end of file