osg_submit.py

#!/usr/bin/env python3
"""
osg_submit.py

OSG submission pipeline:
  1. Check condor job capacity (running+idle vs max_submitted_jobs)
  2. Fetch first non-submitted job from DB          (-b to target a specific ID)
  3. Parse scard; mark job as Processing in DB
  4. Create staging directory ~/osgOutput/<username>/job_<id>/
  5. For type-2: validate lund location and write lund_files via pelican object ls
  6. Generate and stage clas12.condor, nodescript.sh, and functions.sh
  7. Submit to OSG
"""

import argparse
import os
import shutil
import subprocess
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from SConfiguration import SConfiguration
from statuses import (
    FAILED_TO_READ_DIRECTORY,
    NOTSUBMITTED,
    PROCESSING,
    SUBMITTED,
    SCRIPTS_GENERATED,
)

DEFAULT_MAX_SUBMITTED_JOBS = 80000
DEFAULT_OWNER = "gemc"

_REPO_ROOT   = os.path.dirname(os.path.abspath(__file__))
_OUTPUT_ROOT = os.path.join(os.path.expanduser("~"), "osgOutput")


def _job_dir(username, user_submission_id):
    """Return the staging directory path for one submission."""
    return os.path.join(_OUTPUT_ROOT, username, "job_{}".format(user_submission_id))


def _print_test_warning(lines):
    border = "=" * 64
    print(border)
    for line in lines:
        print("  {}".format(line))
    print(border)


def _stage_run_list_files(scard, job_dir):
    """Stage run-by-run inputs into job_dir for a run_list submission.

    Writes run_list.txt (one run per line) and copies runs.json and
    select_run.py so the worker node can perform the weighted draw. Returns the
    list of staged filenames to add to transfer_input_files, or [] when the
    submission has no run_list.
    """
    run_list = getattr(scard, 'run_list', None)
    if not run_list:
        return []

    runs = [r.strip() for r in run_list.split(',') if r.strip()]

    with open(os.path.join(job_dir, "run_list.txt"), 'w') as f:
        f.write("\n".join(runs) + "\n")
    shutil.copy2(
        os.path.join(_REPO_ROOT, "web_portal", "data", "runs.json"),
        os.path.join(job_dir, "runs.json"),
    )
    shutil.copy2(
        os.path.join(_REPO_ROOT, "generators", "bash", "select_run.py"),
        os.path.join(job_dir, "select_run.py"),
    )
    return ["runs.json", "run_list.txt", "select_run.py"]


def _mark_failed_to_read_directory(db_name, user_submission_id, lund_location, reason):
    """Mark a submission as a terminal Lund-directory read failure."""
    from db_io.database import Database

    with Database(database_name=db_name) as db:
        db.mark_submission_terminal(user_submission_id, FAILED_TO_READ_DIRECTORY)
    print("Status → '{}'; priority set to 0.".format(FAILED_TO_READ_DIRECTORY))
    print("Skipping submission for Lund directory {!r}: {}".format(lund_location, reason))


def build_parser():
    # type: () -> argparse.ArgumentParser
    parser = argparse.ArgumentParser(
        description="Submit the next pending OSG job from the database."
    )
    parser.add_argument(
        "--max-submitted-jobs",
        type=int,
        default=DEFAULT_MAX_SUBMITTED_JOBS,
        help="Maximum running+idle jobs allowed before submission is blocked (default: {}).".format(
            DEFAULT_MAX_SUBMITTED_JOBS
        ),
    )
    parser.add_argument(
        "-b", "--user-submission-id",
        type=int,
        default=None,
        metavar="ID",
        help="Process a specific UserSubmissionID instead of the next pending job.",
    )
    parser.add_argument(
        "--test",
        action="store_true",
        default=False,
        help="Test mode: skip htcondor2 capacity check and use a pelican mockup "
             "for type-2 lund files. Without this flag, missing dependencies fail.",
    )
    parser.add_argument(
        "--target-site",
        default=None,
        metavar="SITE",
        help="Pin all jobs to a single GLIDEIN_Site (e.g. CNAF). "
             "Adds (GLIDEIN_Site == \"SITE\") to Requirements.",
    )
    parser.add_argument(
        "--print-condor-card",
        action="store_true",
        default=False,
        help="Print the generated HTCondor submit file to stdout.",
    )
    parser.add_argument(
        "--print-nodescript",
        action="store_true",
        default=False,
        help="Print the generated nodescript.sh to stdout.",
    )
    parser.add_argument(
        "--devel",
        action="store_true",
        default=False,
        help="Use the CLAS12TEST database and the devel singularity image instead of production.",
    )
    return parser


def main(argv=None):
    # type: (list) -> int
    if argv is None:
        argv = sys.argv[1:]

    args = build_parser().parse_args(argv)

    print("\nStep 1: Capacity check")
    # requires htcondor2, only available on the submit node.
    try:
        from condor_io.htcondor_utils import is_under_job_limit
        if not is_under_job_limit(DEFAULT_OWNER, max_jobs=args.max_submitted_jobs):
            print("Job limit reached: '{}' has >= {} running+idle jobs. Skipping.".format(
                DEFAULT_OWNER, args.max_submitted_jobs
            ))
            # Reaching the job limit is a normal back-pressure condition, not an
            # error: there is simply no capacity to submit right now. Return 0 so
            # the caller (cron/wrapper) does not treat the skip as a failure.
            return 0
        print("Capacity check passed: under {} jobs for owner '{}'.".format(
            args.max_submitted_jobs, DEFAULT_OWNER
        ))
    except ImportError:
        if not args.test:
            print("htcondor2 not available. Use --test to skip the capacity check.")
            return 1
        _print_test_warning([
            "TEST MODE: htcondor2 not found.",
            "Skipping job capacity check.",
        ])

    print("\nStep 2: Fetch job from the database")
    try:
        from db_io.database import Database, print_job, current_timestamp
    except ImportError:
        print("pymysql not available — cannot query database.")
        return 1

    db_name = "CLAS12TEST" if args.devel else "CLAS12OCR"
    with Database(database_name=db_name) as db:
        row = db.return_unsubmitted_job(args.user_submission_id)

    if row is None:
        if args.user_submission_id is not None:
            print("No submission found with ID {}.".format(args.user_submission_id))
        else:
            print("No unsubmitted jobs found in the database.")
        return 0

    label = "Targeting specific submission:" if args.user_submission_id else "Next unsubmitted job:"
    print()
    print(label)
    print_job(row)

    print("\nStep 3: Parse submission card")
    # build SConfiguration from scard; backfill username from the DB row.
    scard = SConfiguration.from_string(row['scard'])
    scard.username = row['user']
    user_submission_id = row['user_submission_id']
    marked_processing = False
    condor_submit_succeeded = False
    terminal_status_set = False

    print("\nStep 3a: Mark job as Processing")
    # mark the job as Processing so it is not picked up by a concurrent run.
    if not args.test:
        with Database(database_name=db_name) as db:
            claimed = db.mark_submission_processing(user_submission_id, PROCESSING)
        if claimed != 1:
            print("Submission {} is no longer pending.".format(user_submission_id))
            return 0
        marked_processing = True
        print("Status → '{}'.".format(PROCESSING))
    else:
        print("TEST MODE: skipping status update to '{}'.".format(PROCESSING))

    try:
        print("\nStep 3b: Create staging directory")
        # create staging directory ~/osgOutput/<username>/job_<id>/
        username = scard.username or "unknown"
        job_dir = _job_dir(username, user_submission_id)
        os.makedirs(os.path.join(job_dir, "log"), exist_ok=True)
        print("Staging directory: {}".format(job_dir))

        if scard.type == '2':
            print("\nStep 4: Stage lund files")
        else:
            print("\nStep 4: Stage lund files — skipped (type-1 submission)")

        # for type-2 (lund) submissions, validate the lund location and write lund_files.
        if scard.type == '2':
            lund_location = scard.generator or ""
            if not lund_location.startswith('/volatile/clas12/'):
                reason = "location must start with /volatile/clas12/"
                if not args.test:
                    _mark_failed_to_read_directory(
                        db_name,
                        user_submission_id,
                        lund_location,
                        reason,
                    )
                    terminal_status_set = True
                else:
                    print("TEST MODE: would mark '{}' for {}.".format(
                        FAILED_TO_READ_DIRECTORY,
                        reason,
                    ))
                return 0
            from generators.lund_helper import write_lund_files, LUND_FILES
            try:
                lund_count = write_lund_files(
                    lund_location,
                    output_file=os.path.join(job_dir, LUND_FILES),
                    test=args.test,
                )
            except (subprocess.CalledProcessError, ValueError) as exc:
                if not args.test:
                    _mark_failed_to_read_directory(
                        db_name,
                        user_submission_id,
                        lund_location,
                        exc,
                    )
                    terminal_status_set = True
                else:
                    print("TEST MODE: would mark '{}' after Lund lookup failure: {}".format(
                        FAILED_TO_READ_DIRECTORY,
                        exc,
                    ))
                return 0
            if lund_count == 0:
                reason = "no lund files found"
                if not args.test:
                    _mark_failed_to_read_directory(
                        db_name,
                        user_submission_id,
                        lund_location,
                        reason,
                    )
                    terminal_status_set = True
                else:
                    print("TEST MODE: would mark '{}' because {} at {!r}.".format(
                        FAILED_TO_READ_DIRECTORY,
                        reason,
                        lund_location,
                    ))
                return 0

        # For run_list submissions, stage runs.json, run_list.txt, and select_run.py
        # so the worker node can pick a run weighted by luminosity at runtime.
        extra_input_files = _stage_run_list_files(scard, job_dir)
        if extra_input_files:
            print("Staged run-by-run inputs: {}".format(", ".join(extra_input_files)))

        print("\nStep 5: Generate HTCondor submit file")
        from generators.condor.generate_condor_card import generate_condor_card
        condor_card = generate_condor_card(scard, user_submission_id=user_submission_id,
                                           extra_input_files=extra_input_files,
                                           target_site=args.target_site, devel=args.devel)
        if args.print_condor_card:
            print()
            print(condor_card)

        condor_path = os.path.join(job_dir, "clas12.condor")
        with open(condor_path, 'w') as f:
            f.write(condor_card)
        print("Wrote: {}".format(condor_path))

        print("\nStep 6: Generate nodescript.sh and stage scripts")
        from generators.bash.generate_nodescript import generate_nodescript
        nodescript_path = generate_nodescript(
            scard,
            user_submission_id=user_submission_id,
            test=args.test,
            output_file=os.path.join(job_dir, "nodescript.sh"),
        )
        if args.print_nodescript:
            with open(nodescript_path) as f:
                print()
                print(f.read())

        # Stage functions.sh alongside nodescript.sh (referenced as functions.sh in transfer_input_files).
        shutil.copy2(
            os.path.join(_REPO_ROOT, "generators", "bash", "functions.sh"),
            os.path.join(job_dir, "functions.sh"),
        )

        print("Staged scripts to {}.".format(job_dir))

        # Step 6a: save generated scripts to the DB for later inspection.
        if not args.test:
            with open(nodescript_path) as f:
                nodescript_text = f.read()
            with Database(database_name=db_name) as db:
                db.execute(
                    "UPDATE submissions SET runscript_text = %s, clas12_condor_text = %s"
                    " WHERE user_submission_id = %s",
                    [nodescript_text, condor_card, user_submission_id],
                )
            print("Saved runscript_text and clas12_condor_text to DB.")
        else:
            print("TEST MODE: skipping DB script storage.")

        print("\nStep 7: Submit to HTCondor")
        if args.test:
            print("TEST MODE: skipping condor_submit.")
            return 0

        import re
        result = subprocess.run(
            ["condor_submit", "clas12.condor"],
            cwd=job_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
        )
        if result.returncode != 0:
            error_path = os.path.join(job_dir, "htcondor_submission_error.txt")
            with open(error_path, 'w') as f:
                f.write(result.stdout)
            print("condor_submit failed (exit {}). Error written to {}.".format(
                result.returncode, error_path
            ))
            return 1

        condor_submit_succeeded = True
        print(result.stdout)

        cluster_id = None
        m = re.search(r"submitted to cluster\s+(\d+)", result.stdout, re.IGNORECASE)
        if m:
            cluster_id = int(m.group(1))
            print("HTCondor cluster ID: {}".format(cluster_id))
        else:
            print("Warning: could not parse cluster ID from condor_submit output.")

        with Database(database_name=db_name) as db:
            db.mark_submission_submitted(
                user_submission_id,
                SUBMITTED,
                current_timestamp(),
                cluster_id,
            )
        print("Status → '{}', server_time and pool_node recorded.".format(SUBMITTED))

        return 0
    finally:
        if marked_processing and not condor_submit_succeeded and not terminal_status_set:
            with Database(database_name=db_name) as db:
                db.restore_pending_submission(user_submission_id, PROCESSING)
            print("Status → '{}' because submission did not reach HTCondor.".format(
                NOTSUBMITTED
            ))


if __name__ == "__main__":
    raise SystemExit(main())