From 8d9cb4674c3e4fb3c8abb6bc501c85f8454e24b1 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <hyatkath@amazon.com>
Date: Mon, 30 Oct 2023 12:59:15 -0400
Subject: [PATCH] Add some more logging for docker commands and fix call to
 pull_image

---
 src/device.jl     |  2 +-
 src/local_jobs.jl | 33 +++++++++++++++++++++++++--------
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/src/device.jl b/src/device.jl
index 3ecb4381..0b87687f 100644
--- a/src/device.jl
+++ b/src/device.jl
@@ -18,7 +18,7 @@ end
 for (d, d_arn) in zip((:SV1, :DM1, :TN1), ("sv1", "dm1", "tn1"))
     @eval begin
         struct $d <: AmazonDevice end
-        Base.convert(::Type{String}, d::D) where {D<:AmazonDevice} = "arn:aws:braket:::device/quantum-simulator/amazon/" * $d_arn
+        Base.convert(::Type{String}, d::$d) = "arn:aws:braket:::device/quantum-simulator/amazon/" * $d_arn
     end
 end
 
diff --git a/src/local_jobs.jl b/src/local_jobs.jl
index 556cfd38..7436ec13 100644
--- a/src/local_jobs.jl
+++ b/src/local_jobs.jl
@@ -139,12 +139,17 @@ function run_local_job!(c::LocalJobContainer)
 end
 
 function login_to_ecr(account_id::String, ecr_uri::String, config::AWSConfig)
+    @debug "Attempting to log in to ECR..."
+    @debug "Getting authorization token"
     authorization_data_result = EcR.get_authorization_token(Dict("registryIds"=>[account_id]), aws_config=config)
     isnothing(authorization_data_result) && throw(ErrorException("unable to get permissions to access ECR in order to log in to docker. Please pull down the container before proceeding."))
     raw_token = base64decode(authorization_data_result["authorizationData"][1]["authorizationToken"])
     token = String(raw_token)
     token = replace(token, "AWS:"=>"")
+    @debug "Performing docker login"
     proc_out, proc_err, code = capture_docker_cmd(`docker login -u AWS -p $token $ecr_uri`)
+    @debug "docker login complete"
+    code != 0 && throw(ErrorException("Unable to docker login to ECR with error $proc_err"))
     return
 end
 
@@ -156,6 +161,7 @@ function pull_image(image_uri::String, config::AWSConfig)
     login_to_ecr(account_id, ecr_uri, config)
     @warn "Pulling docker image $image_uri. This may take a while."
     proc_out, proc_err, code = capture_docker_cmd(`docker pull $image_uri`)
+    code != 0 && error(proc_err)
     return
 end
 
@@ -175,15 +181,19 @@ function start_container!(c::LocalJobContainer, force_update::Bool)
     get_image_name(image_uri) = capture_docker_cmd(`docker images -q $image_uri`)[1]
     image_name = get_image_name(image_uri)
     if isempty(image_name) || isnothing(image_name)
-        pull_image(image_uri, c.config)
-        image_name = get_image_name(image_uri)
-        (isempty(image_name) || isnothing(image_name)) && throw(ErrorException("The URI $(c.image_uri) is not available locally and cannot be pulled from Amazon ECR. Please pull down the container before proceeding."))
+        try
+            pull_image(image_uri, c.config)
+            image_name = get_image_name(image_uri)
+            (isempty(image_name) || isnothing(image_name)) && throw(ErrorException("The URI $(c.image_uri) is not available locally and cannot be pulled from Amazon ECR. Please pull down the container before proceeding."))
+        catch ex
+            throw(ErrorException("The URI $(c.image_uri) is not available locally and cannot be pulled from Amazon ECR due to $ex. Please pull down the container before proceeding."))
+        end
     elseif force_update
         try
-            pull_image(image_uri)
+            pull_image(image_uri, c.config)
             image_name = get_image_name(image_uri)
         catch e
-            @warn "Unable to update $(c.image_uri)"
+            @warn "Unable to update $(c.image_uri) with error $e"
         end
     end
     container_name, err, code = capture_docker_cmd(`docker run -d --rm $image_name tail -f /dev/null`)
@@ -194,9 +204,16 @@ end
 
 function stop_container!(c::LocalJobContainer)
     # check that the container is still running
-    c_list = read(`docker container ls -q`, String)
-    stop_flag = occursin(first(c.container_name, 10), c_list)
-    stop_flag && read(Cmd(["docker", "stop", c.container_name]), String)
+    cmd = `docker container ls -q`
+    c_list, c_err, code = capture_docker_cmd(cmd)
+    if code == 0 && occursin(first(c.container_name, 10), c_list)
+        stop_out, stop_err, stop_code = capture_docker_cmd(Cmd(["docker", "stop", c.container_name]))
+        if stop_code != 0
+            error("unable to stop docker container $(c.contianer_name)! Error: $stop_err")
+        end
+    else
+        error("unable to read docker container list! Error: $c_err")
+    end
     return
 end