view gcp_batch_netcat.py @ 3:0ea626b10557 draft

planemo upload for repository https://github.com/afgane/gcp_batch_netcat

import json
import subprocess
import argparse
import uuid
import time
import os

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--nfs_address', required=True)
    parser.add_argument('--output', required=True)
    parser.add_argument('--project', required=True)
    parser.add_argument('--region', required=True)
    parser.add_argument('--port', default='2049')
    parser.add_argument('--service_account_key', required=True)
    args = parser.parse_args()

    # Set up authentication using the service account key
    os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = args.service_account_key

    # Ensure gcloud uses a writable config directory
    os.environ['CLOUDSDK_CONFIG'] = '/tmp/gcloud-config'

    # Create the temp config directory if it doesn't exist
    os.makedirs('/tmp/gcloud-config', exist_ok=True)

    job_name = f'netcat-job-{uuid.uuid4()}'

    job_spec = {
        "taskGroups": [
            {
                "taskSpec": {
                    "runnables": [
                        {
                            "script": {
                                "text": f"/usr/bin/nc -z -v {args.nfs_address} {args.port}"
                            }
                        }
                    ],
                    "computeResource": {
                        "cpuMilli": 1000,
                        "memoryMib": 1024
                    },
                    "environment": {
                        "imageUri": "afgane/gcp-batch-netcat:0.1.0"
                    }
                },
                "taskCount": 1,
                "parallelism": 1
            }
        ],
        "logsPolicy": {
            "destination": "CLOUD_LOGGING"
        }
    }

    job_spec_file = 'job.json'
    with open(job_spec_file, 'w') as f:
        json.dump(job_spec, f)

    command = [
        'gcloud', 'batch', 'jobs', 'submit', job_name,
        '--location', args.region,
        '--project', args.project,
        '--config', job_spec_file,
        '--format=text',
        '--verbosity=debug'
    ]

    # Wait 4 minutes before submitting the job
    time.sleep(240)

    try:
        result = subprocess.run(command, capture_output=True, text=True, check=True)
        with open(args.output, 'w') as f:
            f.write("Job output:\n")
            f.write(result.stdout)
            f.write(result.stderr)
    except subprocess.CalledProcessError as e:
        with open(args.output, 'w') as f:
            f.write("Error submitting job:\n")
            f.write(e.stderr)

if __name__ == '__main__':
    main()