Merge branch 'ErikOsinga:main' into main

Author: ErikOsinga

automation/canfar_polling.py

@@ -33,21 +33,24 @@ async def submit_canfar_job():
       this will depend on the type of CANFAR job that we were running though.
   
 """
-
-
 from __future__ import annotations
+import os
+import pprint
 
+from canfar.sessions import Session
 from dataclasses import dataclass
 
-from prefect import task
+from prefect import task, flow
 from prefect.client.orchestration import get_client
 from prefect.server.schemas.filters import FlowRunFilter, FlowRunFilterState
 from prefect.server.schemas.states import StateType
 from prefect.states import Failed
 
 from print_all_open_sessions import get_open_sessions
+from vos import Client
 
 TAG_PREFIX = "canfar_session:"
+VOS_LOG_FOLDER = "arc:projects/CIRADA/polarimetry/ASKAP/Pipeline_logs/canfar_logs"
 
 
 @dataclass(frozen=True)
@@ -206,3 +209,85 @@ async def reconcile_running_prefect_with_canfar_task(limit: int = 200) -> dict:
         "skipped_untagged": result.skipped_untagged,
     }
 
+@flow(name="Get completed sessions logs and push to Canfar directory")
+def get_completed_session_logs():
+    """
+    Get recently completed CANFAR sessions and push their logs to Canfar directory.
+    This includes Succeeded, Completed, Failed or Error sessions. 
+    The directory is /arc/projects/CIRADA/polarimetry/ASKAP/Pipeline_logs/canfar_logs.
+    """
+    session = Session()
+    statuses = ["Succeeded", "Completed", "Failed", "Error"]
+    # Get all session ids of completed sessions
+    all_session_ids = [
+        s["id"]
+        for status in statuses
+        for s in (session.fetch(kind="headless", status=status) or [])
+    ]
+    for session_id in all_session_ids:
+        # Get logs for each session
+        get_logs(session, session_id)
+
+def push_logs_to_canfar(tmp_file):
+    """
+    Copy Canfar session logs to Canfar directory and remove local file
+    Args:
+        tmp_file: Local log file to copy and delete
+    """
+    client = Client()
+    
+    # Ensure remote directory exists
+    try:
+        client.mkdir(VOS_LOG_FOLDER)
+    except:
+        pass  # Directory might already exist
+    
+    try:
+        remote_file = f"{VOS_LOG_FOLDER}/{tmp_file}"
+        client.copy(tmp_file, remote_file)
+        print(f"Upload Canfar log to {remote_file} completed successfully!\n")
+    except Exception as e:
+        print(f"Failed to upload {tmp_file} to Canfar: {e}\n")
+        return
+    
+    # Clean up local file
+    os.remove(tmp_file)    
+
+def log_exists_in_canfar_dir(session_id):
+    """
+    Check if log already exists in Canfar directory.
+    Args:
+        session_id: ID of the session
+    Returns:
+        True if log exists, False otherwise
+    """
+    client = Client()
+    file_name = f"{VOS_LOG_FOLDER}/{session_id}.log"
+    try:
+        return client.isfile(file_name)
+    except:
+        # Folder doesn't exist yet
+        return False
+        
+
+def get_logs(session, session_id):
+    """
+    Get logs for a Canfar session.
+
+    Args:
+        session: Canfar session object
+        session_id: ID of the session to monitor
+        output_file: File to write logs to
+    """
+    if log_exists_in_canfar_dir(session_id):
+        return
+    
+    tmp_file = f"{session_id}.log"
+    
+    with open(tmp_file, "w") as f:
+        # Fetch logs before exiting
+        logs_dict = session.logs(session_id).get(session_id)
+        f.write(pprint.pformat(logs_dict, depth=4))
+        f.flush
+    push_logs_to_canfar(tmp_file)       
+

automation/prefect/prefect-git-clone.yaml

@@ -62,3 +62,10 @@ deployments:
     work_pool:
       name: local-pool
       work_queue_name: default    
+
+  - name: push-session-logs-to-canfar
+    entrypoint: automation/canfar_polling.py:get_completed_session_logs    
+    description: 'Schedule this to run after each parent flow to push Canfar session logs to Canfar directory.'
+    work_pool:
+      name: local-pool 
+      work_queue_name: default

possum_pipeline_control/check_ingest_3Dpipeline.py

@@ -91,12 +91,14 @@ def launch_ingest(tilenumber, band):
         replicas=1,
     )
 
+    session_id_str = session_id_str = session_id[0] if len(session_id) > 0 else None
+
     print("Check sessions at https://ws-uv.canfar.net/skaha/v1/session")
     print(
-        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id[0]}?view=logs"
+        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id_str}?view=logs"
     )
 
-    return session_id[0]
+    return session_id_str
 
 def update_status(tile_number, band, status, conn):
     """

possum_pipeline_control/check_status_and_launch_3Dpipeline_v2.py

@@ -280,11 +280,14 @@ def launch_download_session(jobname="3dtile-dl"):
         replicas=1,
     )
 
+    session_id_str = session_id_str = session_id[0] if len(session_id) > 0 else None
+
     print("Check sessions at https://ws-uv.canfar.net/skaha/v1/session")
     print(
-        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id[0]}?view=logs"
+        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id_str}?view=logs"
     )
-    return session_id[0]
+
+    return session_id_str
 
 def launch_create_symlinks(jobname="3dsymlinks"):
     """
@@ -317,11 +320,14 @@ def launch_create_symlinks(jobname="3dsymlinks"):
         replicas=1,
     )
 
+    session_id_str = session_id_str = session_id[0] if len(session_id) > 0 else None
+
     print("Check sessions at https://ws-uv.canfar.net/skaha/v1/session")
     print(
-        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id[0]}?view=logs"
+        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id_str}?view=logs"
     )
-    return session_id[0]
+
+    return session_id_str
 
 def needs_prefect_db_backup(
     home_dir: str | Path,

possum_pipeline_control/launch_1Dpipeline_PartialTiles_band1.py

@@ -42,12 +42,14 @@ def launch_session(run_name, field_ID, tilenumbers, SBnumber, image, cores, ram)
         replicas=1,
     )
 
+    session_id_str = session_id_str = session_id[0] if len(session_id) > 0 else None
+
     print("Check sessions at https://ws-uv.canfar.net/skaha/v1/session")
     print(
-        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id[0]}?view=logs"
+        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id_str}?view=logs"
     )
 
-    return session_id[0]
+    return session_id_str
 
 
 @flow(name="launch_1d_partialtiles", log_prints=True)

possum_pipeline_control/launch_1Dpipeline_PartialTiles_band1_pre_or_post.py

@@ -3,9 +3,7 @@
 import os
 from datetime import datetime
 
-# from skaha.session import Session
 from canfar.sessions import Session
-
 from possum_pipeline_control.control_1D_pipeline_PartialTiles import get_open_sessions
 
 """
@@ -81,12 +79,14 @@ def launch_session(
         replicas=1,
     )
 
+    session_id_str = session_id_str = session_id[0] if len(session_id) > 0 else None
+
     print("Check sessions at https://ws-uv.canfar.net/skaha/v1/session")
     print(
-        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id[0]}?view=logs"
+        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id_str}?view=logs"
     )
 
-    return session_id[0] 
+    return session_id_str
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(

possum_pipeline_control/launch_3Dpipeline_band1.py

@@ -28,12 +28,14 @@ def launch_session(run_name, tilenumber, image, cores, ram):
         replicas=1,
     )
 
+    session_id_str = session_id_str = session_id[0] if len(session_id) > 0 else None
+
     print("Check sessions at https://ws-uv.canfar.net/skaha/v1/session")
     print(
-        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id[0]}?view=logs"
+        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id_str}?view=logs"
     )
 
-    return session_id[0]
+    return session_id_str
 
 
 def main_launch3D(tilenumber: str):

possum_pipeline_control/submit_download.py

@@ -54,12 +54,13 @@ def launch_download():
         replicas=1,
     )
 
+    session_id_str = session_id_str = session_id[0] if len(session_id) > 0 else None
+
     print("Check sessions at https://ws-uv.canfar.net/skaha/v1/session")
     print(
-        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id[0]}?view=logs"
+        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id_str}?view=logs"
     )
-
-    return session_id[0]
+    return session_id_str
 
 
 if __name__ == "__main__":

possum_pipeline_control/update_partialtile_google_sheet.py

@@ -25,6 +25,7 @@
 from pathlib import Path
 
 import astropy.table as at
+from canfar.sessions import Session
 import gspread
 import matplotlib.dates as mdates
 import matplotlib.pyplot as plt
@@ -37,6 +38,7 @@
 from possum_pipeline_control import util
 
 
+
 def get_sheet_table(band):
     """
     Connects to the POSSUM Status Monitor Google Sheet and returns a sub-table
@@ -461,9 +463,6 @@ def launch_collate_job():
     """
     Launches the collate job for the 1D pipeline. Once per day
     """
-    # from skaha.session import Session
-    from canfar.sessions import Session
-
     session = Session()
 
 
@@ -495,11 +494,14 @@ def launch_collate_job():
         replicas=1,
     )
 
+    session_id_str = session_id_str = session_id[0] if len(session_id) > 0 else None
+
     print("Check sessions at https://ws-uv.canfar.net/skaha/v1/session")
     print(
-        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id[0]}?view=logs"
+        f"Check logs at https://ws-uv.canfar.net/skaha/v1/session/{session_id_str}?view=logs"
     )
-    return session_id[0]
+
+    return session_id_str
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Update Partial Tile Google Sheet")