refactor: Linting

bb_tasks/tasks/callable.py

@@ -18,102 +18,105 @@
 #                                Scripted Tasks                                #
 # ---------------------------------------------------------------------------- #
 
+# ---------------------------------- New VM ---------------------------------- #
 @shared_task
-def new_vm_subprocess(instance_id, xml, root_user):
+def new_vm_subprocess(instance_id):
     '''
     Called to start the creation of a VM in the background.
     '''
+    brick = VirtualBrick.objects.get(id=instance_id)
+    host = brick.host
+
     catch_clone_errors.apply_async((instance_id,), countdown=60)
     remove_stale_clone.apply_async((instance_id,), countdown=180)
 
     new_vm_script = [
                         f'{DIR}brick_connect.sh',
-                        'brick_img',
-                        f'{str(Site.objects.get_current().domain)}',
-                        f'{str(instance_id)}',
-                        f'{str(xml)}',
-                        f'{str(root_user)}'
+                        f'{str(host.ssh_username)}', f'{str(host.ssh_port)}',
+                        'brick_img', f'{str(Site.objects.get_current().domain)}',
+                        f'{str(instance_id)}', f'{str(brick.assigned_gpus[0])}'
                     ]
+
     with subprocess.Popen(new_vm_script) as script:
         print(script)
 
 
+# -------------------------------- Shutdown VM ------------------------------- #
 @shared_task
-def pause_vm_subprocess(instance_id, root_user):
+def pause_vm_subprocess(instance_id):
     '''
     Called to power down a VM.
     '''
+    host = VirtualBrick.objects.get(id=instance_id).host
+
     pause_vm_script = [
                         f'{DIR}brick_connect.sh',
-                        'brick_pause',
-                        f'{str(Site.objects.get_current().domain)}',
-                        f'{str(instance_id)}',
-                        f'{str(root_user)}'
+                        f'{str(host.ssh_username)}', f'{str(host.ssh_port)}',
+                        'brick_pause', f'{str(Site.objects.get_current().domain)}',
+                        f'{str(instance_id)}'
                     ]
+
     with subprocess.Popen(pause_vm_script) as script:
         print(script)
 
-    # with subprocess.Popen([f'{DIR}brick_pause.sh', f'{str(instance_id)}']) as script:
-    #     print(script)
-
 
+# --------------------------------- Bootup VM -------------------------------- #
 @shared_task
-def play_vm_subprocess(instance_id, root_user):
+def play_vm_subprocess(instance_id):
     '''
     Resume VM from off state.
     '''
+    host = VirtualBrick.objects.get(id=instance_id).host
+
     play_vm_script = [
                         f'{DIR}brick_connect.sh',
-                        'brick_play',
-                        f'{str(Site.objects.get_current().domain)}',
+                        f'{str(host.ssh_username)}', f'{str(host.ssh_port)}',
+                        'brick_play', f'{str(Site.objects.get_current().domain)}',
                         f'{str(instance_id)}',
-                        f'{str(root_user)}'
                     ]
+
     with subprocess.Popen(play_vm_script) as script:
         print(script)
 
-    # with subprocess.Popen([f'{DIR}brick_play.sh', f'{str(instance_id)}']) as script:
-    #     print(script)
-
 
+# --------------------------------- Reboot VM -------------------------------- #
 @shared_task
-def reboot_vm_subprocess(instance_id, root_user):
+def reboot_vm_subprocess(instance_id):
     '''
     Called to reboot a VM.
     '''
+    host = VirtualBrick.objects.get(id=instance_id).host
+
     reboot_vm_script = [
                         f'{DIR}brick_connect.sh',
-                        'brick_reboot',
-                        f'{str(Site.objects.get_current().domain)}',
+                        f'{str(host.ssh_username)}', f'{str(host.ssh_port)}',
+                        'brick_reboot', f'{str(Site.objects.get_current().domain)}',
                         f'{str(instance_id)}',
-                        f'{str(root_user)}'
                     ]
     with subprocess.Popen(reboot_vm_script) as script:
         print(script)
 
-    # with subprocess.Popen([f'{DIR}brick_reboot.sh', f'{str(instance_id)}']) as script:
-    #     print(script)
-
 
+# --------------------------------- Delete VM -------------------------------- #
 @shared_task
-def destroy_vm_subprocess(instance_id, root_user):
+def destroy_vm_subprocess(instance_id):
     '''
     Called to destroy VM.
     '''
+    host = VirtualBrick.objects.get(id=instance_id).host
+
     destroy_vm_script = [
                         f'{DIR}brick_connect.sh',
-                        'brick_destroy',
-                        f'{str(Site.objects.get_current().domain)}',
+                        f'{str(host.ssh_username)}', f'{str(host.ssh_port)}',
+                        'brick_destroy', f'{str(Site.objects.get_current().domain)}',
                         f'{str(instance_id)}',
-                        f'{str(root_user)}'
                     ]
+
     with subprocess.Popen(destroy_vm_script) as script:
         print(script)
 
-    # with subprocess.Popen([f'{DIR}brick_destroy.sh', f'{str(instance_id)}']) as script:
-    #     print(script)
-
 
+# ---------------------------- Terminate SSH Port ---------------------------- #
 @shared_task
 def close_ssh_port(port_number):
     '''
@@ -143,6 +146,7 @@ def remove_stale_clone(instance_id):
     Last resort to remove clones that did not sucessfully start and never given a port.
     '''
     brick = VirtualBrick.objects.get(id=instance_id)
+
     if brick.ssh_port is None and brick.is_on is False:
         VirtualBrickOwner.objects.filter(virt_brick=instance_id).delete()
         brick.delete()

bb_tasks/tasks/periodic.py

@@ -100,30 +100,28 @@ def reconnect_host(host):
 
     preperation_script = [
                             f'{DIR}brick_connect.sh',
-                            'brick_prep',
-                            f'{str(Site.objects.get_current().domain)}',
-                            f'{host.ssh_username}'
+                            f'{str(host.ssh_username)}', f'{str(host.ssh_port)}',
+                            'brick_prep', f'{str(Site.objects.get_current().domain)}',
                         ]
+
     with Popen(preperation_script, stdout=PIPE) as script:
         try:
             prep_script_result = f"{script.stdout.read().decode('ascii')}"
         except AttributeError:
             prep_script_result = 'No output'
 
-
-
     reconnect_script = "Not GPUs"
     reconnect_script_result = "Not Ran"
 
     for gpu in GPU.objects.filter(host=host):
         reconnect_script = [
                                 f'{DIR}brick_connect.sh',
-                                'brick_reconnect',
-                                f'{str(Site.objects.get_current().domain)}',
+                                f'{str(host.ssh_username)}', f'{str(host.ssh_port)}',
+                                'brick_reconnect', f'{str(Site.objects.get_current().domain)}',
                                 f'{gpu.device}',
                                 f'{gpu.pcie}'
-                                f'{host.ssh_username}'
                             ]
+
         with Popen(reconnect_script) as script:
             try:
                 reconnect_script_result = f"{script.stdout.read().decode('ascii')}"
@@ -153,10 +151,10 @@ def restore_vm_state(instance):
 
     play_vm_script = [
                         f'{DIR}brick_connect.sh',
-                        'brick_play',
-                        f'{str(Site.objects.get_current().domain)}',
+                        f'{str(host.ssh_username)}', f'{str(host.ssh_port)}',
+                        'brick_play', f'{str(Site.objects.get_current().domain)}',
                         f'{str(instance)}'
-                        f'{host.ssh_username}'
                     ]
+
     with Popen(play_vm_script) as script:
         print(script)

bb_vm/bash_scripts/brick_connect.sh

@@ -1,19 +1,34 @@
 #!/bin/bash
 
+# Connects to a remote host through a SSH tunnel and runs the provided script.
+
 # --------------------------------- Arguments -------------------------------- #
-# - (1) URL Endpint
-# - (2) Instance ID that has been provided to refrence this new VM.
-# - (3) XML formatted data to attach GPU to VM
 
-action_script=$1
-url=$2
-instance=$3
-xml_data=$4
-root_user=$5
+# $1: Remote (root) username
+# $2: Host assigned port tunnel
+# $3: The script that will be executed on the remote host
+
+## The Following arguments are passed along to the script
+
+# $4: URL Endpint
+# $5: Instance ID that has been provided to refrence this new VM
+# $6: XML formatted data to attach GPU to VM
+
+host_user=$1
+port=$2
+action=$3
+
+url=$4
+instance=$5
+xml_data=$6
+
 
-if lsof -i tcp:9002; then
+# Check that port is active before trying to connect.
+if lsof -i tcp:"$port"; then
 
-    sshpass -p "Password@1" ssh -i /opt/brickbox/bb_vm/keys/bb_root -o StrictHostKeyChecking=no -p 9002 "$root_user"@localhost 'bash -s' < /opt/brickbox/bb_vm/bash_scripts/"$action_script".sh "$url" "$instance" \""$xml_data"\" 2>> /opt/brickbox/bb_vm/bash_scripts/bash_errors.log
+    ssh -i /opt/brickbox/bb_vm/keys/"$host_user" -o StrictHostKeyChecking=no -p "$port" "$host_user"@localhost \
+    'bash -s' < /opt/brickbox/bb_vm/bash_scripts/"$action".sh "$url" "$instance" \""$xml_data"\" 2>> \
+    /opt/brickbox/bb_vm/bash_scripts/logs/brick_connect_errors.log
 
 else
 

bb_vm/views.py

@@ -49,12 +49,7 @@ def clone_img(request):
             brick_owner = VirtualBrickOwner(owner=profile, virt_brick=instance)
             brick_owner.save()
 
-            new_vm_subprocess.delay(instance.id, designated_gpu_xml, instance.host.ssh_username)
-
-            # subprocess.Popen([
-            #     '/opt/brickbox/bb_vm/bash_scripts/clone_img.sh',
-            #     f'{str(instance.id)}', f'{str(designated_gpu_xml)}'
-            # ])
+            new_vm_subprocess.delay(instance.id)
 
             bricks = VirtualBrickOwner.objects.filter(owner=profile) # All bricks owned.
             response_data = {}
@@ -122,8 +117,7 @@ def brick_pause(request):
     brick.is_on = False
     brick.save()
 
-    pause_vm_subprocess.delay(vm_id, brick.host.ssh_username)
-    # subprocess.Popen(['/opt/brickbox/bb_vm/bash_scripts/brick_pause.sh', f'{str(vm_id)}'])
+    pause_vm_subprocess.delay(vm_id)
 
     return HttpResponse(status=200)
 
@@ -143,8 +137,7 @@ def brick_play(request):
     brick.is_on = True
     brick.save()
 
-    play_vm_subprocess.delay(vm_id, brick.host.ssh_username)
-    # subprocess.Popen(['/opt/brickbox/bb_vm/bash_scripts/brick_play.sh', f'{str(vm_id)}'])
+    play_vm_subprocess.delay(vm_id)
 
     return HttpResponse(status=200)
 
@@ -187,9 +180,7 @@ def brick_destroy(request):
     # brick.ssh_port.delete()
     close_ssh_port.apply_async((brick.ssh_port.port_number,), countdown=43200)
 
-    destroy_vm_subprocess.delay(vm_id, brick.host.ssh_username)
-
-    # subprocess.Popen(['/opt/brickbox/bb_vm/bash_scripts/brick_destroy.sh', f'{str(vm_id)}'])
+    destroy_vm_subprocess.delay(vm_id)
 
     profile = UserProfile.objects.get(user = request.user)
     bricks = VirtualBrickOwner.objects.filter(owner=profile) # All bricks owned.