Seperate file / s3 lifetimes / writes / reads from http calls

- python-runner no longer simulates a single http api that does
everything (incl s3 interaction)
- Seperates file lifetimes (shared storage) from execution
- Easier to see how to avoid reading into strings (e.g. for piping into
tmp files, follow up)

Earlier on we had a service that managed much of the I/O, metabase now
coordinates it so it makes little sense pretending there is still
an overall http api.

enterprise/backend/src/metabase_enterprise/transforms/execute.clj

@@ -238,38 +238,53 @@
     (driver/insert-from-source! driver (:id db) table-name (mapv :name (:columns table-schema)) data-source)))
 
 (defn- run-python-transform! [{:keys [source] :as transform} db run-id cancel-chan message-log]
-  (with-open [log-thread-ref (open-log-thread! run-id message-log)]
-    (let [driver                           (:engine db)
-          {:keys [source-tables body]}     source
-          {:keys [body status] :as result} (call-python-runner-api! body source-tables run-id cancel-chan)
-          {:keys [events]}   body]
+  (with-open [log-thread-ref     (open-log-thread! run-id message-log)
+              shared-storage-ref (python-runner/open-s3-shared-storage! (:source-tables source))]
+    (let [driver                        (:engine db)
+          server-url                    (transforms.settings/python-execution-server-url)
+          _                             (python-runner/copy-tables-to-s3! {:run-id                  run-id
+                                                                           :shared-storage @shared-storage-ref
+                                                                           :table-name->id (:source-tables source)
+                                                                           :cancel-chan    cancel-chan})
+          _                             (python-runner/open-cancellation-process! server-url run-id cancel-chan) ; inherits lifetime of cancel-chan
+          {:keys [status body] :as response}
+          (python-runner/execute-python-code-http-call!
+           {:server-url     server-url
+            :code           (:body source)
+            :run-id         run-id
+            :table-name->id (:source-tables source)
+            :shared-storage @shared-storage-ref})
+          {:keys [exit_code]} body
+          ;; TODO temporary to keep more code stable while refactoring
+          ;; no need to materialize these early (i.e output we can stream directly into a tmp file or db if small)
+          {:keys [output output-manifest events]} (python-runner/read-output-objects @shared-storage-ref)]
       (.close log-thread-ref)           ; early close to force any writes to flush
       (when (seq events)
         (replace-python-logs! message-log events))
       (if (not= 200 status)
-        (throw (ex-info (debug-info-str body)
+        (throw (ex-info (debug-info-str {:exit-code exit_code :events events}) ;; todo do better here
                         {:status-code     400
                          :api-status-code status
                          :body            body
+                         :events          events
                          :error           (:error body)}))
         (try
-          (let [temp-file (File/createTempFile "transform-output-" ".csv")
-                csv-data  (:output body)
-                metadata  (-> body :metadata json/decode+kw)]
-            (when-not (seq (:fields metadata))
+          (let [temp-file (File/createTempFile "transform-output-" ".csv")]
+            (when-not (seq (:fields output-manifest))
               (throw (ex-info "No fields in metadata"
-                              {:metadata metadata
-                               :raw-body body})))
+                              {:metadata output-manifest
+                               :raw-body body
+                               :events   events})))
             (try
               (with-open [writer (io/writer temp-file)]
-                (.write writer ^String csv-data))
+                (.write writer ^String output))
               (let [file-size (.length temp-file)]
                 (transforms.instrumentation/with-stage-timing [run-id :data-transfer :file-to-dwh]
-                  (transfer-file-to-db driver db transform metadata temp-file))
+                  (transfer-file-to-db driver db transform output-manifest temp-file))
                 (transforms.instrumentation/record-data-transfer! run-id :file-to-dwh file-size nil))
               (finally
                 (.delete temp-file))))
-          result
+          response
           (catch Exception e
             (log/error e "Failed to to create resulting table")
             (throw (ex-info "Failed to create the resulting table" {:error (.getMessage e)}))))))))

enterprise/backend/src/metabase_enterprise/transforms/python_runner.clj

@@ -14,7 +14,8 @@
    [metabase.util.log :as log]
    [toucan2.core :as t2])
   (:import
-   (java.io BufferedWriter File OutputStream OutputStreamWriter)
+   (clojure.lang IDeref)
+   (java.io BufferedWriter Closeable File OutputStream OutputStreamWriter)
    (java.net URI)
    (java.nio.charset StandardCharsets)
    (java.time Duration)
@@ -262,20 +263,20 @@
 (defn- cleanup-s3-objects [^S3Client s3-client bucket-name s3-keys]
   (run! (partial delete-s3-object s3-client bucket-name) s3-keys))
 
-(defn- read-from-s3 [^S3Client s3-client ^String bucket-name ^String key & [fallback-content]]
-  (try
-    (let [^GetObjectRequest request (build-get-object-request bucket-name key)
-          response                  (.getObject s3-client request)]
-      (slurp response))
-    (catch NoSuchKeyException e
-      (if fallback-content
-        fallback-content
-        (throw e)))))
+(defn- read-from-s3
+  ([s3-client bucket-name key] (read-from-s3 s3-client bucket-name key ::throw))
+  ([^S3Client s3-client ^String bucket-name ^String key not-found]
+   (try
+     (let [^GetObjectRequest request (build-get-object-request bucket-name key)
+           response                  (.getObject s3-client request)]
+       (slurp response))
+     (catch NoSuchKeyException e
+       (if (identical? ::throw not-found)
+         (throw e)
+         not-found)))))
 
 (defn get-logs
   "Return the logs of the current running python process"
-  ;; TODO: we should be given an id for the expected job, se we don't return unrelated logs
-  ;;       if the job has already finished, we could fethc the logs from the db instead
   [run-id]
   (let [server-url (transforms.settings/python-execution-server-url)]
     (http/get (str server-url "/logs")
@@ -285,130 +286,121 @@
                :as               :json
                :query-params     {:request_id run-id}})))
 
-(defn execute-python-code
-  "Execute Python code using the Python execution server."
-  [run-id code table-name->id cancel-chan]
-  (let [prefix        (some-> (transforms.settings/python-storage-s-3-prefix) (str "/"))
-        work-dir-name (str prefix "run-" (System/nanoTime) "-" (rand-int 10000))]
+(defn- s3-shared-storage [table-name->id]
+  (let [prefix              (some-> (transforms.settings/python-storage-s-3-prefix) (str "/"))
+        work-dir-name       (str prefix "run-" (System/nanoTime) "-" (rand-int 10000))
+        container-presigner (create-s3-presigner-for-container)
+        bucket-name         (transforms.settings/python-storage-s-3-bucket)
+        loc
+        (fn [method relative-path]
+          (let [path (str work-dir-name "/" relative-path)]
+            {:path   path
+             :method method
+             :url    (case method :put (generate-presigned-put-url container-presigner bucket-name path)
+                           :get (generate-presigned-get-url container-presigner bucket-name path))}))]
+    {:s3-client   (create-s3-client)                        ;; do not like mixing interactive things with descriptions, but its damn convenient to have it here for now
+     :bucket-name bucket-name
+     :objects
+     (-> (into
+          {:output          (loc :put "output.csv")
+           :output-manifest (loc :put "output-manifest.json")
+           :events          (loc :put "events.jsonl")}
+          (for [[table-name id] table-name->id]
+            {[:table id :manifest] (loc :get (str "-table-" (name table-name) "-" id ".manifest.json"))
+             [:table id :data]     (loc :get (str "-table-" (name table-name) "-" id ".jsonl"))})))}))
 
-    (try
-      (let [server-url               (transforms.settings/python-execution-server-url)
-            bucket-name              (transforms.settings/python-storage-s-3-bucket)
-            s3-client                (create-s3-client)
-            container-presigner      (create-s3-presigner-for-container)
+(defn open-s3-shared-storage!
+  "Returns a deref'able shared storage value, (.close) will cleanup any s3 objects named in storage (data files for tables and so on)."
+  ^Closeable [table-name->id]
+  (let [shared-storage (s3-shared-storage table-name->id)]
+    (reify IDeref
+      (deref [_] shared-storage)
+      Closeable
+      (close [_] (cleanup-s3-objects (:s3-client shared-storage) (:bucket-name shared-storage) (map :path (vals (:objects shared-storage))))))))
 
-            ;; Generate S3 keys for output files
-            output-key               (str work-dir-name "/output.csv")
-            output-manifest-key      (str work-dir-name "/output.manifest.json")
-            events-key               (str work-dir-name "/events.jsonl")
-            ;; Generate presigned URLs for writing (using container client)
-            output-url               (generate-presigned-put-url container-presigner bucket-name output-key)
-            output-manifest-url      (generate-presigned-put-url container-presigner bucket-name output-manifest-key)
-            events-url               (generate-presigned-put-url container-presigner bucket-name events-key)
-            ;; Upload input table data (write to disk first, then upload to S3)
-            table-results            (for [[table-name id] table-name->id]
-                                       (let [temp-file       (File/createTempFile
-                                                              (str work-dir-name "-table-" (name table-name) "-" id)
-                                                              ".jsonl")
-                                             manifest-file   (File/createTempFile
-                                                              (str work-dir-name "-table-" (name table-name) "-" id)
-                                                              ".manifest.json")
-                                             s3-key          (str work-dir-name "/" (.getName temp-file))
-                                             manifest-s3-key (str work-dir-name "/" (.getName manifest-file))]
-                                         (try
-                                           ;; Write table data to temporary file and get manifest
-                                           (let [manifest (transforms.instrumentation/with-stage-timing [run-id :data-transfer :dwh-to-file]
-                                                            (write-table-data-to-file! id temp-file cancel-chan))]
-                                             ;; Write manifest to file
-                                             (with-open [writer (io/writer manifest-file)]
-                                               (json/encode-to manifest writer {}))
-                                             (let [file-size (.length temp-file)
-                                                   manifest-size (.length manifest-file)]
-                                               (transforms.instrumentation/record-data-transfer! run-id :dwh-to-file file-size nil)
+(defn copy-tables-to-s3! [{:keys [run-id
+                                  shared-storage
+                                  table-name->id
+                                  cancel-chan]}]
+  (doseq [id (vals table-name->id)
+          :let [{:keys [s3-client bucket-name objects]} shared-storage
+                {data-path :path}     (get objects [:table id :data])
+                {manifest-path :path} (get objects [:table id :manifest])]]
+    (let [temp-file       (File/createTempFile data-path "")
+          manifest-file   (File/createTempFile manifest-path "")]
+      (try
+        ;; Write table data to temporary file and get manifest
+        (let [manifest (transforms.instrumentation/with-stage-timing [run-id :data-transfer :dwh-to-file]
+                         (write-table-data-to-file! id temp-file cancel-chan))]
+          ;; Write manifest to file
+          (with-open [writer (io/writer manifest-file)]
+            (json/encode-to manifest writer {}))
+          (let [file-size (.length temp-file)
+                manifest-size (.length manifest-file)]
+            (transforms.instrumentation/record-data-transfer! run-id :dwh-to-file file-size nil)
 
-                                               ;; Upload both files to S3
-                                               (transforms.instrumentation/with-stage-timing [run-id :data-transfer :file-to-s3]
-                                                 (upload-file-to-s3 s3-client bucket-name s3-key temp-file)
-                                                 (upload-file-to-s3 s3-client bucket-name manifest-s3-key manifest-file))
+            ;; Upload both files to S3
+            (transforms.instrumentation/with-stage-timing [run-id :data-transfer :file-to-s3]
+              (upload-file-to-s3 s3-client bucket-name data-path temp-file)
+              (upload-file-to-s3 s3-client bucket-name manifest-path manifest-file))
 
-                                               (let [data-url     (generate-presigned-get-url container-presigner bucket-name s3-key)
-                                                     manifest-url (generate-presigned-get-url container-presigner bucket-name manifest-s3-key)]
-                                                 (transforms.instrumentation/record-data-transfer! run-id :file-to-s3 (+ file-size manifest-size) nil)
-                                                 {:table-name      (name table-name)
-                                                  :url             data-url
-                                                  :manifest-url    manifest-url
-                                                  :s3-key          s3-key
-                                                  :manifest-s3-key manifest-s3-key})))
-                                           (finally
-                                             ;; Clean up temporary files
-                                             (safe-delete temp-file)
-                                             (safe-delete manifest-file)))))
-            table-name->url          (into {} (map (juxt :table-name :url) table-results))
-            table-name->manifest-url (into {} (map (juxt :table-name :manifest-url) table-results))
-            all-s3-keys              (concat [output-key output-manifest-key events-key]
-                                             (map :s3-key table-results)
-                                             (map :manifest-s3-key table-results))
-            canc                     (a/go (when (a/<! cancel-chan)
-                                             (http/post (str server-url "/cancel")
-                                                        {:content-type :json
-                                                         :body         (json/encode {:request_id run-id})
-                                                         :async?       true}
-                                                        identity identity)))
-            payload                  {:code                code
-                                      :timeout             30
-                                      :request_id          run-id
-                                      :output_url          output-url
-                                      :output_manifest_url output-manifest-url
-                                      :events_url          events-url
-                                      :table_mapping       table-name->url
-                                      :manifest_mapping    table-name->manifest-url}
-            response                 (http/post (str server-url "/execute")
-                                                {:content-type     :json
-                                                 :accept           :json
-                                                 :body             (json/encode payload)
-                                                 :throw-exceptions false
-                                                 :as               :json})
-            _                        (a/close! canc)
+            (transforms.instrumentation/record-data-transfer! run-id :file-to-s3 (+ file-size manifest-size) nil)))
+        (finally
+          ;; Clean up temporary files
+          (safe-delete temp-file)
+          (safe-delete manifest-file))))))
 
-            result                   (:body response)
-            ;; TODO look into why some tests return json and others strings
-            result                   (if (string? result) (json/decode result keyword) result)]
+(defn execute-python-code-http-call! [{:keys [server-url code run-id table-name->id shared-storage]}]
+  (let [{:keys [objects]} shared-storage
+        {:keys [output output-manifest events]} objects
 
-        (try
-          (if (and (= 200 (:status response))
-                   (zero? (:exit_code result)))
-           ;; Success - read the output from S3
-            (let [output-content  (read-from-s3 s3-client bucket-name output-key)
-                  output-manifest (read-from-s3 s3-client bucket-name output-manifest-key "{}")
-                  events-content  (read-from-s3 s3-client bucket-name events-key "[]")]
-              (if (not (str/blank? output-content))
-                {:status 200
-                 :body   {:output output-content
-                          :metadata output-manifest
-                          :events (mapv json/decode+kw (str/split-lines events-content))}}
-                {:status 500
-                 :body   {:error  "Transform did not produce output CSV"
-                          :events (mapv json/decode+kw (str/split-lines events-content))}}))
-           ;; Error from execution server - read events .jsonl (including stderr/stdout) from S3
-            (let [events-content (read-from-s3 s3-client bucket-name events-key "[]")]
-              {:status 500
-               :body
-               {:error       (or (:error result) "Execution failed")
-                :exit-code   (:exit_code result)
-                :status-code (:status response)
-                :timeout     (:timeout result)
-                :events      (mapv json/decode+kw (str/split-lines events-content))}}))
-          (finally
-            ;; Clean up S3 objects
-            (try
-              (cleanup-s3-objects s3-client bucket-name all-s3-keys)
-              (catch Exception _)))))
+        table-name->url          (update-vals table-name->id (comp :url #(get objects [:table % :data])))
+        table-name->manifest-url (update-vals table-name->id (comp :url #(get objects [:table % :manifest])))
 
-      (catch CancellationException _
-        {:status 408
-         :body   {:error "Interrupted"}})
+        payload                  {:code                code
+                                  :timeout             30
+                                  :request_id          run-id
+                                  :output_url          (:url output)
+                                  :output_manifest_url (:url output-manifest)
+                                  :events_url          (:url events)
+                                  :table_mapping       table-name->url
+                                  :manifest_mapping    table-name->manifest-url}
 
-      (catch Exception e
-        (.printStackTrace e)
-        {:status 500
-         :body   {:error (str "Failed to connect to Python execution server: " (.getMessage e))}}))))
+        response                 (http/post (str server-url "/execute")
+                                            {:content-type     :json
+                                             :accept           :json
+                                             :body             (json/encode payload)
+                                             :throw-exceptions false
+                                             :as               :json})]
+    ;; when a 500 is returned we observe a string in the body (despite the python returning json)
+    ;; always try to parse the returned string as json before yielding (could tighten this up at some point)
+    (update response :body (fn [string-if-error]
+                             (if (string? string-if-error)
+                               (try
+                                 (json/decode+kw string-if-error)
+                                 (catch Exception _
+                                   {:error string-if-error}))
+                               string-if-error)))))
+
+(defn open-cancellation-process!
+  "Starts a core.async process that optimistically sends a cancellation request to the python executor if cancel-chan receives a value.
+  Returns a channel that will receive either the async http call j.u.c.FutureTask in the case of cancellation, or nil when the cancel-chan is closed."
+  [server-url run-id cancel-chan]
+  (a/go (when (a/<! cancel-chan)
+          (http/post (str server-url "/cancel")
+                     {:content-type :json
+                      :body         (json/encode {:request_id run-id})
+                      :async?       true}
+                     identity identity))))
+
+;; temporary (just keep the files alive longer and stream them directly (i.e into tmp file for write)
+(defn read-output-objects
+  "Temporary function that strings/jsons stuff in S3 and returns it for compatibility."
+  [{:keys [s3-client bucket-name objects]}]
+  (let [{:keys [output output-manifest events]} objects
+        output-content          (read-from-s3 s3-client bucket-name (:path output) nil)
+        output-manifest-content (read-from-s3 s3-client bucket-name (:path output-manifest) "{}")
+        events-content          (read-from-s3 s3-client bucket-name (:path events))]
+    {:output output-content
+     :output-manifest (json/decode+kw output-manifest-content)
+     :events (mapv json/decode+kw (str/split-lines events-content))}))

enterprise/backend/test/metabase_enterprise/transforms/python_runner_test.clj

@@ -12,6 +12,7 @@
    [metabase-enterprise.transforms.test-util :as transforms.tu]
    [metabase-enterprise.transforms.util :as transforms.util]
    [metabase.driver :as driver]
+   [metabase-enterprise.transforms.settings :as transforms.settings]
    [metabase.driver.sql-jdbc.connection :as sql-jdbc.conn]
    [metabase.sync.core :as sync]
    [metabase.task.core :as task]
@@ -50,12 +51,26 @@
 
 (def ^:private test-id 1)
 
-(defn- reconstruct-logs [{:keys [events] :as body}]
-  (assoc body :stdout (->> events (filter #(= "stdout" (:stream %))) (map :message) (str/join "\n"))
-         :stderr (->> events (filter #(= "stderr" (:stream %))) (map :message) (str/join "\n"))))
-
 (defn- execute [{:keys [code tables]}]
-  (reconstruct-logs (:body (python-runner/execute-python-code test-id code (or tables {}) (a/promise-chan)))))
+  (with-open [shared-storage-ref (python-runner/open-s3-shared-storage! (or tables {}))]
+    (let [server-url     (transforms.settings/python-execution-server-url)
+          cancel-chan    (a/promise-chan)
+          table-name->id (or tables {})
+          _              (python-runner/copy-tables-to-s3! {:run-id         test-id
+                                                            :shared-storage @shared-storage-ref
+                                                            :table-name->id table-name->id
+                                                            :cancel-chan    cancel-chan})
+          response       (python-runner/execute-python-code-http-call! {:server-url     server-url
+                                                                        :code           code
+                                                                        :run-id         test-id
+                                                                        :table-name->id table-name->id
+                                                                        :shared-storage @shared-storage-ref})
+          {:keys [output events]} (python-runner/read-output-objects @shared-storage-ref)]
+      ;; not sure about munging this all together but its what tests expect for now
+      (merge (:body response)
+             {:output output
+              :stdout (->> events (filter #(= "stdout" (:stream %))) (map :message) (str/join "\n"))
+              :stderr (->> events (filter #(= "stderr" (:stream %))) (map :message) (str/join "\n"))}))))
 
 (deftest ^:parallel transform-function-basic-test
   (testing "executes transform function and returns CSV output"
@@ -74,8 +89,7 @@
     (let [result (execute {:code (str "import pandas as pd\n"
                                       "\n"
                                       "# No transform function defined")})]
-      (is (=? {:error     "Execution failed"
-               :exit-code 1
+      (is (=? {:exit_code 1
                :stderr    "ERROR: User script must define a 'transform()' function"
                :stdout    ""}
               result)))))
@@ -84,8 +98,7 @@
   (testing "handles transform function returning non-DataFrame"
     (let [result (execute {:code (str "def transform():\n"
                                       "    return 'not a dataframe'")})]
-      (is (=? {:error     "Execution failed"
-               :exit-code 1
+      (is (=? {:exit_code 1
                :stderr    "ERROR: Transform function must return a pandas DataFrame, got <class 'str'>"
                :stdout    ""}
               result)))))
@@ -103,8 +116,7 @@
                                  "    raise ValueError('Something went wrong')\n"
                                  "ValueError: Something went wrong")
           stderr-pattern    (template->regex expected-template)]
-      (is (=? {:error     "Execution failed"
-               :exit-code 1
+      (is (=? {:exit_code 1
                :stderr    #(re-matches stderr-pattern %)
                :stdout    ""
                :timeout   false}