Fix gateway server bug and interrupt handling for chunk dispatching

pyproject.toml

@@ -62,6 +62,7 @@ pynacl = { version = "^1.5.0", optional = true }
 pyopenssl = { version = "^22.0.0", optional = true }
 werkzeug = { version = "^2.1.2", optional = true }
 pyarrow = "^10.0.1"
+pytest = "^7.3.2"
 
 [tool.poetry.extras]
 aws = ["boto3"]

scripts/gen_data/gen_many_small.py

@@ -26,4 +26,4 @@ def make_file(data, fname):
 
     files = [f"{outdir}/{i:08d}.bin" for i in range(args.nfiles)]
     data = np.arange(args.size // 4, dtype=np.int32).tobytes()
-    do_parallel(partial(make_file, data), files, desc="Generating files", spinner=True, spinner_persist=True)
+    do_parallel(partial(make_file, data), files, desc="Generating files", spinner=True, spinner_persist=True, n=16)

skyplane/api/dataplane.py

@@ -246,7 +246,7 @@ def copy_gateway_logs(self):
         # copy logs from all gateways in parallel
         do_parallel(self.copy_gateway_log, self.bound_nodes.values(), n=-1)
 
-    def deprovision(self, max_jobs: int = 64, spinner: bool = False):
+    def deprovision(self, max_jobs: int = 64, spinner: bool = True):
         """
         Deprovision the remote gateways
 
@@ -267,6 +267,8 @@ def deprovision(self, max_jobs: int = 64, spinner: bool = False):
                 for task in self.pending_transfers:
                     logger.fs.warning(f"Before deprovisioning, waiting for jobs to finish: {list(task.jobs.keys())}")
                     task.join()
+                    for thread in threading.enumerate():
+                        assert "_run_multipart_chunk_thread" not in thread.name, f"thread {thread.name} is still running"
             except KeyboardInterrupt:
                 logger.warning("Interrupted while waiting for transfers to finish, deprovisioning anyway.")
                 raise

skyplane/api/tracker.py

@@ -1,10 +1,12 @@
 import functools
+import signal
+
 from pprint import pprint
 import json
 import time
 from abc import ABC
 from datetime import datetime
-from threading import Thread
+from threading import Thread, Event
 
 import urllib3
 from typing import TYPE_CHECKING, Dict, List, Optional, Set
@@ -97,6 +99,14 @@ def __init__(self, dataplane, jobs: List["TransferJob"], transfer_config: Transf
         self.jobs = {job.uuid: job for job in jobs}
         self.transfer_config = transfer_config
 
+        # exit handling
+        self.exit_flag = Event()
+
+        def signal_handler(signal, frame):
+            self.exit_flag.set()
+
+        signal.signal(signal.SIGINT, signal_handler)
+
         if hooks is None:
             self.hooks = EmptyTransferHook()
         else:
@@ -138,16 +148,20 @@ def run(self):
         session_start_timestamp_ms = int(time.time() * 1000)
         try:
             # pre-dispatch chunks to begin pre-buffering chunks
-            chunk_streams = {
-                job_uuid: job.dispatch(self.dataplane, transfer_config=self.transfer_config) for job_uuid, job in self.jobs.items()
-            }
+            chunk_streams = {job_uuid: job.dispatch(self.dataplane) for job_uuid, job in self.jobs.items()}
             for job_uuid, job in self.jobs.items():
                 logger.fs.debug(f"[TransferProgressTracker] Dispatching job {job.uuid}")
                 self.job_chunk_requests[job_uuid] = {}
                 self.job_pending_chunk_ids[job_uuid] = {region: set() for region in self.dataplane.topology.dest_region_tags}
                 self.job_complete_chunk_ids[job_uuid] = {region: set() for region in self.dataplane.topology.dest_region_tags}
 
                 for chunk in chunk_streams[job_uuid]:
+                    if self.exit_flag.is_set():
+                        logger.fs.debug(f"[TransferProgressTracker] Exiting due to signal")
+                        self.hooks.on_dispatch_end()
+                        self.hooks.on_transfer_end()
+                        job.stop()  # stop threads in chunk stream
+                        return
                     chunks_dispatched = [chunk]
                     self.job_chunk_requests[job_uuid][chunk.chunk_id] = chunk
                     self.hooks.on_chunk_dispatched(chunks_dispatched)