Upload classes.py

classes.py

@@ -196,17 +196,24 @@ class TDocIndexer:
         self.total_count += len([m for m in meeting_folders if m not in ['./', '../']])
         
         # Utiliser ThreadPoolExecutor pour traiter les réunions en parallèle
-        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
-            futures = [executor.submit(self.process_meeting, meeting, wg_url) 
-                      for meeting in meeting_folders if meeting not in ['./', '../']]
-            
-            total = len(futures)
-            done_count = 0
-            yield f"event: get-maximum\ndata: {total}\n\n"
+        executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers)
+        futures = [executor.submit(self.process_meeting, meeting, wg_url)
+                  for meeting in meeting_folders if meeting not in ['./', '../']]
+
+        total = len(futures)
+        done_count = 0
+        yield f"event: get-maximum\ndata: {total}\n\n"
 
+        try:
             for future in concurrent.futures.as_completed(futures):
                 done_count += 1
                 yield f"event: progress\ndata: {done_count}\n\n"
+        except GeneratorExit:
+            for f in futures:
+                f.cancel()
+            executor.shutdown(wait=False, cancel_futures=True)
+            return
+        executor.shutdown(wait=False)
 
     def index_all_tdocs(self):
         """Indexer tous les documents ZIP dans la structure FTP 3GPP avec multithreading"""
@@ -257,17 +264,24 @@ class TDocIndexer:
         self.total_count += len([m for m in meeting_folders if m not in ['./', '../']])
         
         # Utiliser ThreadPoolExecutor pour traiter les réunions en parallèle
-        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
-            futures = [executor.submit(self.process_meeting, meeting, main_url, workshop=True) 
-                      for meeting in meeting_folders if meeting not in ['./', '../']]
-            total = len(futures)
-            done_count = 0
+        executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers)
+        futures = [executor.submit(self.process_meeting, meeting, main_url, workshop=True)
+                  for meeting in meeting_folders if meeting not in ['./', '../']]
+        total = len(futures)
+        done_count = 0
 
-            yield f"event: get-maximum\ndata: {total}\n\n"
+        yield f"event: get-maximum\ndata: {total}\n\n"
 
+        try:
             for future in concurrent.futures.as_completed(futures):
                 done_count += 1
                 yield f"event: progress\ndata: {done_count}\n\n"
+        except GeneratorExit:
+            for f in futures:
+                f.cancel()
+            executor.shutdown(wait=False, cancel_futures=True)
+            return
+        executor.shutdown(wait=False)
         
         docs_count_after = len(self.indexer)
         new_docs_count = docs_count_after - docs_count_before
@@ -287,6 +301,14 @@ class Spec3GPPIndexer:
         self.processed_count = 0
         self.total_count = 0
 
+        try:
+            self.failed_specifications = set(
+                item["spec_id"] for item in load_dataset("OrganizedProgrammers/3GPPFailedSpecs")["train"].to_list()
+            )
+            print(f"Loaded {len(self.failed_specifications)} previously failed specifications")
+        except (EmptyDatasetError, Exception):
+            self.failed_specifications = set()
+
         self.DICT_LOCK = threading.Lock()
         self.DOCUMENT_LOCK = threading.Lock()
         self.STOP_EVENT = threading.Event()
@@ -369,7 +391,7 @@ class Spec3GPPIndexer:
                                     '--outdir', tmpdir,
                                     extracted_path
                                 ]
-                                subprocess.run(cmd, check=True, timeout=60*5, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+                                subprocess.run(cmd, check=True, timeout=60*2, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
                                 txt_file = os.path.splitext(extracted_path)[0] + '.txt'
                                 if os.path.exists(txt_file):
@@ -382,6 +404,10 @@ class Spec3GPPIndexer:
 
         except Exception as e:
             print(f"Error getting text for {specification} v{version_code}: {e}")
+            if isinstance(e, (subprocess.TimeoutExpired, subprocess.CalledProcessError)):
+                with self.DICT_LOCK:
+                    self.failed_specifications.add(specification)
+                print(f"Spec {specification}: marked as failed for future indexation runs")
             return []
 
     def get_spec_content(self, specification, version_code):
@@ -445,6 +471,12 @@ class Spec3GPPIndexer:
                     self.processed_count += 1
                 return
 
+            if doc_id in self.failed_specifications:
+                with self.DICT_LOCK:
+                    self.processed_count += 1
+                    print(f"Spec {doc_id} ({spec.get('title', '')}): skipped (previously failed) - Progress {self.processed_count}/{self.total_count}")
+                return
+
             document = None
             already_indexed = False
             needs_fetch = False
@@ -562,17 +594,24 @@ class Spec3GPPIndexer:
         specifications = self.fetch_spec_table()
         self.total_count = len(specifications)
         print(f"Processing {self.total_count} specs with {self.max_workers} threads...")
-        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
-            futures = [executor.submit(self.process_specification, spec) for spec in specifications]
-            total = len(futures)
-            done_count = 0
-            yield f"event: get-maximum\ndata: {total}\n\n"
+        executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers)
+        futures = [executor.submit(self.process_specification, spec) for spec in specifications]
+        total = len(futures)
+        done_count = 0
+        yield f"event: get-maximum\ndata: {total}\n\n"
 
+        try:
             for future in concurrent.futures.as_completed(futures):
                 done_count += 1
                 yield f"event: progress\ndata: {done_count}\n\n"
                 if self.STOP_EVENT.is_set():
                     break
+        except GeneratorExit:
+            for f in futures:
+                f.cancel()
+            executor.shutdown(wait=False, cancel_futures=True)
+            return
+        executor.shutdown(wait=False)
         print("All specs processed.")
 
     # Sauvegarde (identique au script original)
@@ -591,7 +630,12 @@ class Spec3GPPIndexer:
         print("Pushing ...")
         push_spec_content.push_to_hub("OrganizedProgrammers/3GPPSpecContent", token=os.environ["HF"])
         push_spec_metadata.push_to_hub("OrganizedProgrammers/3GPPSpecMetadata", token=os.environ["HF"])
-        
+
+        if self.failed_specifications:
+            failed_list = [{"spec_id": spec_id} for spec_id in sorted(self.failed_specifications)]
+            Dataset.from_list(failed_list).push_to_hub("OrganizedProgrammers/3GPPFailedSpecs", token=os.environ["HF"])
+            print(f"Saved {len(failed_list)} failed specifications to 3GPPFailedSpecs")
+
         self.spec_contents = load_dataset("OrganizedProgrammers/3GPPSpecContent")["train"].to_list()
         self.documents_by_spec_num = self._make_doc_index(self.spec_contents)
         print("Save finished.")
@@ -830,18 +874,24 @@ class SpecETSIIndexer:
         self.total_count = len(specifications)
         print(f"Traitement de {self.total_count} specs avec {self.max_workers} threads...\n")
 
-        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
-            futures = [executor.submit(self.process_specification, spec) for spec in specifications]
-            total = len(futures)
-            done_count = 0
-            yield f"event: get-maximum\ndata: {total}\n\n"
+        executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers)
+        futures = [executor.submit(self.process_specification, spec) for spec in specifications]
+        total = len(futures)
+        done_count = 0
+        yield f"event: get-maximum\ndata: {total}\n\n"
 
+        try:
             for future in concurrent.futures.as_completed(futures):
                 done_count += 1
                 yield f"event: progress\ndata: {done_count}\n\n"
                 if self.STOP_EVENT.is_set():
                     break
-
+        except GeneratorExit:
+            for f in futures:
+                f.cancel()
+            executor.shutdown(wait=False, cancel_futures=True)
+            return
+        executor.shutdown(wait=False)
         print(f"\nAll {self.processed_count}/{self.total_count} specs processed.")
 
     def save(self):