AWS Batch

import time
import datetime
import boto3  # Install it e.g. with `pip install boto3`
from lightly.api import ApiWorkflowClient
from lightly.openapi_generated.swagger_client import DatasetType, DatasourcePurpose

"""
Step 0: Set the Lightly token and the AWS Batch job queue ARN.
"""

# You must change this. Find it under the preferences page in the Lightly Platform: https://app.lightly.ai/preferences
MY_LIGHTLY_TOKEN = "CHANGEME"

# You must change this. Find it under https://console.aws.amazon.com/batch/home#jobDefinitions by clicking on the job queue name.
# It has the format `anr:aws:batch:REGION:AWS_ID:job-queue/JOB_QUEUE_NAME`
MY_AWS_BATCH_JOB_QUEUE_ARN = "CHANGEME"

# Optional to change.
# It is used to make sure that the worker name, worker label, job description name and job name are unique.
# Useful when processing multiple Lightly scheduled runs in parallel.
# Must not contain special characters except for `-` and `_`.
unique_id = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")


# Create a Lightly API client.
client = ApiWorkflowClient(token=MY_LIGHTLY_TOKEN)
# Create a Boto3 client for AWS Batch.
batch_client = boto3.client("batch")


"""
Step 1: Register a Lightly Worker.
A unique Lightly Worker allows scheduling runs specifically for that worker.
This is only needed once for each Lightly Worker you want to run in parallel.
"""

worker_name = f"aws-batch-worker {unique_id}"
worker_labels = [f"{unique_id}"]
print("Worker labels:", worker_labels)
worker_id = client.register_compute_worker(name=worker_name, labels=worker_labels)


"""
Step 2: Create a new dataset on the Lightly Platform and configure the Input and Lightly datasources for the dataset.
This step is tha same as without AWS Batch.
"""
client.create_dataset(
    dataset_name="dataset-name",
    dataset_type=DatasetType.IMAGES,  # must be DatasetType.VIDEOS when working with videos
)
my_dataset_id = client.dataset_id
print(f"Dataset ID: {my_dataset_id}")
# Configure the Input datasource.
client.set_s3_delegated_access_config(
    resource_path="s3://bucket/input/",
    region="eu-central-1",
    role_arn="S3-ROLE-ARN",
    external_id="S3-EXTERNAL-ID",
    purpose=DatasourcePurpose.INPUT
)
# Configure the Lightly datasource.
client.set_s3_delegated_access_config(
    resource_path="s3://bucket/lightly/",
    region="eu-central-1",
    role_arn="S3-ROLE-ARN",
    external_id="S3-EXTERNAL-ID",
    purpose=DatasourcePurpose.LIGHTLY
)


"""
Step 3: Configure and schedule a run.
This differs from the usual way of scheduling a run in two aspects:

i) By setting `worker_config={"shutdown_when_job_finished": True}`, the Lightly Worker docker container will shut down
as soon as it finishes the first scheduled run. This will then cause the AWS batch job to finish and the EC2 instance to shut down.
Otherwise, the job and machine will stay running with the Lightly Worker waiting for the next scheduled run to process.

ii) The `runs_on` argument was set to make the scheduled run be processed by the Lightly Worker created by AWS batch.
Please be aware that a Lightly Worker without any labels might pick up the job instead, see the label matching docs:
https://lightly-docs.readme.io/docs/assign-scheduled-runs-to-specific-workers#label-matching
"""
scheduled_run_id = client.schedule_compute_worker_run(
    worker_config={
        "shutdown_when_job_finished": True,
    },
    selection_config={
        "n_samples": 50,
        "strategies": [
            {"input": {"type": "EMBEDDINGS"}, "strategy": {"type": "DIVERSITY"}}
        ],
    },
    runs_on=worker_labels,
)


"""
Step 5: Set and register the AWS Batch job definition.
This is only needed once unless you change the job definition.
"""
job_definition_name = f"aws-batch-lightly--job-definition--{unique_id}"
job_name = f"aws-batch-lightly--job--{unique_id}"
print(
    f"Registering AWS Batch job definition {job_definition_name} and submitting job {job_name}."
)
job_definition = {
    "jobDefinitionName": job_definition_name,
    "type": "container",
    "containerProperties": {
        "image": "lightly/worker:latest",
        # Resource requirements for a g4dn.2xlarge instance. Change to your needs.
        # Hardware recommendations: https://docs.lightly.ai/docs/hardware-recommendations
        "resourceRequirements": [
            {"type": "MEMORY", "value": "32768"},  # 32 GB
            {"type": "VCPU", "value": "8"},
        ],
        "command": [
            f"token={MY_LIGHTLY_TOKEN}",
            f"worker.worker_id={worker_id}",
        ],
    },
}

# Register the job definition. This is only needed once unless you change the job definition.
response = batch_client.register_job_definition(**job_definition)
if "jobDefinitionName" in response and "jobDefinitionArn" in response:
    print("Job Definition Registered Successfully")
    print("Job Definition Name:", response["jobDefinitionName"])
    print("Job Definition ARN:", response["jobDefinitionArn"])
else:
    raise RuntimeError("Job Definition Registration Failed")


"""
Step 6: Submit a job for the job definition.
This is needed once each time you want to process one or a series of Lightly scheduled runs.
"""
submit_response = batch_client.submit_job(
    jobName=job_name,
    jobQueue=MY_AWS_BATCH_JOB_QUEUE_ARN,
    jobDefinition=job_definition_name,
)
if "jobId" in submit_response and "jobName" in submit_response:
    print("Job Submitted Successfully")
    print("Job Name:", submit_response["jobName"])
    print("Job ID:", submit_response["jobId"])
else:
    raise RuntimeError("Job Submission Failed")
job_id = submit_response["jobId"]
print(f"Job ID: {job_id}")


"""
Step 7: Check the AWS batch and Lightly scheduled run  job status.
This step is optional, but it is useful for monitoring and debugging.

i) At the beginning, the AWS batch job status should be `SUBMITTED` and the Lightly scheduled run state should be `OPEN`.
ii) The AWS batch job status should go into state RUNNABLE within a few seconds.
During state RUNNABLE, a new EC2 instance has to be created and started, which may take a few minutes.
If the job is stuck in `RUNNABLE`, use the AWS batch debugging troubleshooting automation
console.aws.amazon.com/systems-manager/automation/execute/AWSSupport-TroubleshootAWSBatchJob
iii) Once the EC2 instance is running, the AWS batch job status should go into state `STARTING`. 
Now the Lightly Worker docker image is downloaded. This may take another few minutes.
iv) After the Lightly Worker docker image was downloaded, the AWS batch job status should go into state `RUNNING`.
It should pick up the scheduled run and start processing it.

Now the AWS batch job status should stay in state `RUNNING` until the Lightly scheduled went through all its steps.

v) Once the Lightly Worker finished the scheduled run, the AWS batch job status should go into state `SUCCEEDED`.
"""
while True:
    aws_batch_job_status = batch_client.describe_jobs(jobs=[job_id])["jobs"][0][
        "status"
    ]
    ligthly_scheduled_run_info = client.get_compute_worker_run_info(
        scheduled_run_id=scheduled_run_id
    )

    current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(
        f"{current_time} - Current Status of AWS batch job: {aws_batch_job_status} and state of Lightly Worker run: {ligthly_scheduled_run_info.state}"
    )

    if aws_batch_job_status in ["RUNNING"]:
        print("Lightly Worker is running and searches for scheduled jobs.")
        print(
            "Check the Worker status in the Lightly Platform: https://app.lightly.ai/compute/workers"
        )
    elif aws_batch_job_status in ["SUCCEEDED", "FAILED"]:
        print(f"AWS batch job {aws_batch_job_status}. Exiting status check loop.")
        break

    if ligthly_scheduled_run_info.in_end_state():
        print(
            f"Lightly Worker run finished with state {ligthly_scheduled_run_info.state}. Exiting status check loop."
        )
        break

    time.sleep(5)  # Sleep for 5 seconds before next check