refactor upload_analysis_to_jumpbox

jerry-ng2 · jerry-ng2 · commit 66e6c9f4e351 · 2026-03-26T14:46:50.000-04:00
diff --git a/skills/root-cause-analysis/scripts/cli.py b/skills/root-cause-analysis/scripts/cli.py
@@ -15,6 +15,7 @@
     from scripts.config import Config
     from scripts.correlator import build_correlation_timeline, fetch_correlated_logs
     from scripts.job_parser import parse_job_log
+    from scripts.jumpbox_io import upload_to_jumpbox
     from scripts.log_fetcher import fetch_job_log
     from scripts.setup import print_checks, run_checks
     from scripts.step4_fetch_github import GitHubClient, Step4Analyzer
@@ -24,6 +25,7 @@
     from .config import Config
     from .correlator import build_correlation_timeline, fetch_correlated_logs
     from .job_parser import parse_job_log
+    from .jumpbox_io import upload_to_jumpbox
     from .log_fetcher import fetch_job_log
     from .setup import print_checks, run_checks
     from .step4_fetch_github import GitHubClient, Step4Analyzer
@@ -47,95 +49,17 @@ def save_step(analysis_dir: Path, step: int, data: dict) -> Path:
 
 
 @trace(name="Upload analysis", span_type=SpanType.CHAIN if SpanType else None)
-def upload_analysis_to_jumpbox(args: argparse.Namespace, config: Config, span=None) -> int:
-    """Upload analysis directory to Jumpbox in /usr/local/mlflow/{job_id}/ with session.json."""
-    analysis_dir = config.analysis_dir / args.job_id
-
-    if not analysis_dir.exists():
-        error_message = f"No analysis found for job {args.job_id}"
-        print(error_message)
-        if span:
-            span.set_outputs({"error": error_message})
-        return 1
-
-    print(f"Uploading analysis for job {args.job_id}...")
-
-    if not config.jumpbox_uri:
-        print("  Skipping upload: JUMPBOX_URI not configured")
-        return 1
-
-    # Parse JUMPBOX_URI format: "user@host -p port"
-    parts = config.jumpbox_uri.split()
-    if len(parts) < 1:
-        print("  Error: Invalid JUMPBOX_URI format")
-        return 1
-
-    ssh_target = parts[0]  # user@host
-    ssh_port = None
-
-    # Extract port if present
-    if "-p" in parts:
-        try:
-            port_idx = parts.index("-p")
-            if port_idx + 1 < len(parts):
-                ssh_port = parts[port_idx + 1]
-        except (ValueError, IndexError):
-            pass
-
-    session_id = os.environ.get("CLAUDE_SESSION_ID", "unknown")
+def cmd_upload(args: argparse.Namespace, config: Config, span=None) -> int:
+    """Upload analysis directory to Jumpbox via jumpbox_io.upload_to_jumpbox."""
     job_id = args.job_id
-    remote_base_dir = f"/usr/local/mlflow/{job_id}"
-
-    # Create session.json file locally (will be overwritten if exists)
-    session_file = analysis_dir / "session.json"
-    try:
-        with open(session_file, "w") as f:
-            json.dump({"session_id": session_id, "job_id": job_id}, f, indent=2)
-    except Exception as e:
-        print(f"  Warning: Could not create session.json: {e}")
-
-    # Build SSH command to create remote directory
-    ssh_cmd = ["ssh"]
-    if ssh_port:
-        ssh_cmd.extend(["-p", ssh_port])
-    ssh_cmd.extend([ssh_target, f"mkdir -p {remote_base_dir}"])
-
-    # Create remote directory structure
-    try:
-        subprocess.run(
-            ssh_cmd,
-            check=True,
-            capture_output=True,
-            timeout=30,
-        )
-    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
-        print(f"  Error creating remote directory: {e}")
-        return 1
+    analysis_dir = config.analysis_dir / job_id
+    session_id = os.environ.get("CLAUDE_SESSION_ID", "unknown")
 
-    # Upload analysis directory contents using rsync
-    try:
-        rsync_cmd = ["rsync"]
-        if ssh_port:
-            rsync_cmd.extend(["-e", f"ssh -p {ssh_port}"])
-        # Add trailing slash to source to copy contents, not the directory itself
-        rsync_cmd.extend(
-            ["-az", "--quiet", f"{str(analysis_dir)}/", f"{ssh_target}:{remote_base_dir}/"]
-        )
+    success = upload_to_jumpbox(job_id, analysis_dir, config.jumpbox_uri, session_id)
 
-        subprocess.run(
-            rsync_cmd,
-            check=True,
-            timeout=60,
-        )
-        print(f"  Uploaded to Jumpbox ({ssh_target}): {remote_base_dir}/")
-        if span:
-            span.set_outputs({"job_id": job_id, "success": True})
-        return 0
-    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
-        print(f"  Error uploading to Jumpbox: {e}")
-        if span:
-            span.set_outputs({"job_id": job_id, "success": False, "error": str(e)})
-        return 1
+    if span:
+        span.set_outputs({"job_id": job_id, "success": success})
+    return 0 if success else 1
 
 
 def get_step_name(step: int) -> str:
@@ -614,7 +538,7 @@ def main() -> int:
         "query": cmd_query,
         "setup": cmd_setup,
         "status": cmd_status,
-        "upload": upload_analysis_to_jumpbox,
+        "upload": cmd_upload,
     }
 
     exit_code = commands[args.command](args, config, span)
diff --git a/skills/root-cause-analysis/scripts/jumpbox_io.py b/skills/root-cause-analysis/scripts/jumpbox_io.py
@@ -0,0 +1,124 @@
+"""
+Jumpbox I/O functions for Root-Cause-Analysis.
+
+Provides functions to parse JUMPBOX_URI and upload analysis directories
+to the jumpbox via rsync.
+"""
+
+import json
+import os
+import subprocess
+from pathlib import Path
+
+
+def parse_jumpbox_uri(jumpbox_uri: str) -> tuple[str, str | None]:
+    """
+    Parse JUMPBOX_URI format: "user@host -p port" or "user@host".
+
+    Returns:
+        Tuple of (ssh_target, ssh_port) where ssh_port may be None
+
+    Raises:
+        ValueError: If format is invalid
+    """
+    if not jumpbox_uri:
+        raise ValueError("JUMPBOX_URI is empty")
+
+    parts = jumpbox_uri.split()
+    if len(parts) < 1:
+        raise ValueError("Invalid JUMPBOX_URI format")
+
+    ssh_target = parts[0]
+    ssh_port = None
+
+    if "-p" in parts:
+        try:
+            port_idx = parts.index("-p")
+            if port_idx + 1 < len(parts):
+                ssh_port = parts[port_idx + 1]
+        except (ValueError, IndexError):
+            pass
+
+    return ssh_target, ssh_port
+
+
+def upload_to_jumpbox(
+    job_id: str,
+    analysis_dir: Path,
+    jumpbox_uri: str | None = None,
+    session_id: str | None = None,
+) -> bool:
+    """
+    Upload analysis directory to jumpbox at /usr/local/mlflow/<job_id>/.
+
+    Optionally writes a session.json before uploading. Uses rsync for
+    efficient directory transfer.
+
+    Args:
+        job_id: Job ID for remote path
+        analysis_dir: Local analysis directory to upload
+        jumpbox_uri: JUMPBOX_URI connection string (defaults to env var)
+        session_id: If provided, writes session.json with this ID before upload
+
+    Returns:
+        True on success, False on failure
+    """
+    if "/" in job_id or "." in job_id:
+        print("  Error: Invalid job_id format. Cannot contain . or /")
+        return False
+
+    if not analysis_dir.exists():
+        print(f"  Error: Analysis directory not found: {analysis_dir}")
+        return False
+
+    if jumpbox_uri is None:
+        jumpbox_uri = os.environ.get("JUMPBOX_URI", "")
+
+    if not jumpbox_uri:
+        print("  JUMPBOX_URI not set. Skipping upload.")
+        return False
+
+    try:
+        ssh_target, ssh_port = parse_jumpbox_uri(jumpbox_uri)
+    except ValueError as e:
+        print(f"  Error: {e}")
+        return False
+
+    remote_dir = f"/usr/local/mlflow/{job_id}"
+
+    if session_id:
+        session_file = analysis_dir / "session.json"
+        try:
+            with open(session_file, "w") as f:
+                json.dump({"session_id": session_id, "job_id": job_id}, f, indent=2)
+        except Exception as e:
+            print(f"  Warning: Could not create session.json: {e}")
+
+    print(f"  Uploading analysis for job {job_id}...")
+    print(f"    Local:  {analysis_dir}/")
+    print(f"    Remote: {ssh_target}:{remote_dir}/")
+
+    ssh_cmd = ["ssh"]
+    if ssh_port:
+        ssh_cmd.extend(["-p", ssh_port])
+    ssh_cmd.extend([ssh_target, f"mkdir -p {remote_dir}"])
+
+    try:
+        subprocess.run(ssh_cmd, check=True, capture_output=True, timeout=30)
+    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
+        print(f"  Error creating remote directory: {e}")
+        return False
+
+    rsync_cmd = ["rsync"]
+    if ssh_port:
+        rsync_cmd.extend(["-e", f"ssh -p {ssh_port}"])
+    rsync_cmd.extend(["-az", "--quiet", f"{analysis_dir}/", f"{ssh_target}:{remote_dir}/"])
+
+    try:
+        subprocess.run(rsync_cmd, check=True, timeout=60)
+    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
+        print(f"  Error uploading to jumpbox: {e}")
+        return False
+
+    print(f"   Uploaded to jumpbox: {ssh_target}:{remote_dir}/")
+    return True
