src/gha_runner/clouddeployment.py
Rama main · réplica de lectura del repositorio original.
from abc import ABC, abstractmethod
from gha_runner.gh import GitHubInstance, MissingRunnerLabel
from gha_runner.helper.workflow_cmds import warning, error
from dataclasses import dataclass, field
from typing import Type
class CreateCloudInstance(ABC):
"""Abstract base class for starting a cloud instance.
This class defines the interface for starting a cloud instance.
"""
@abstractmethod
def create_instances(self) -> dict[str, str]:
"""Create instances in the cloud provider and return their IDs.
The number of instances to create is defined by the
implementation.
Returns
-------
dict[str, str]
A dictionary of instance IDs and their corresponding github
runner labels.
"""
raise NotImplementedError
@abstractmethod
def wait_until_ready(self, ids: list[str], **kwargs):
"""Wait until instances are in a ready state.
Parameters
----------
ids : list[str]
A list of instance IDs to wait for.
**kwargs : dict, optional
Additional arguments to pass to the waiter.
"""
raise NotImplementedError
@abstractmethod
def set_instance_mapping(self, mapping: dict[str, str]):
"""Set the instance mapping in the environment.
Parameters
----------
mapping : dict[str, str]
A dictionary of instance IDs and their corresponding github
runner labels.
"""
raise NotImplementedError
class StopCloudInstance(ABC):
"""Abstract base class for stopping a cloud instance.
This class defines the interface for stopping a cloud instance.
"""
@abstractmethod
def remove_instances(self, ids: list[str]):
"""Remove instances from the cloud provider.
Parameters
----------
ids : list[str]
A list of instance IDs to remove.
"""
raise NotImplementedError
@abstractmethod
def wait_until_removed(self, ids: list[str], **kwargs):
"""Wait until instances are removed.
Parameters
----------
ids : list[str]
A list of instance IDs to wait for.
**kwargs : dict, optional
Additional arguments to pass to the waiter.
"""
raise NotImplementedError
@abstractmethod
def get_instance_mapping(self) -> dict[str, str]:
"""Get the instance mapping from the environment.
Returns
-------
dict[str, str]
A dictionary of instance IDs and their corresponding github
runner labels.
"""
raise NotImplementedError
@dataclass
class DeployInstance:
"""Class that is used to deploy instances and runners.
Parameters
----------
provider_type : Type[CreateCloudInstance]
The type of cloud provider to use.
cloud_params : dict
The parameters to pass to the cloud provider.
gh : GitHubInstance
The GitHub instance to use.
count : int
The number of instances to create.
timeout : int
The timeout to use when waiting for the runner to come online
Attributes
----------
provider : CreateCloudInstance
The cloud provider instance
provider_type : Type[CreateCloudInstance]
cloud_params : dict
gh : GitHubInstance
count : int
timeout : int
"""
provider_type: Type[CreateCloudInstance]
cloud_params: dict
gh: GitHubInstance
count: int
timeout: int
provider: CreateCloudInstance = field(init=False)
def __post_init__(self):
"""Initialize the cloud provider.
This function is called after the object is created to correctly
init the provider.
"""
# We need to create runner tokens for use by the provider
runner_tokens = self.gh.create_runner_tokens(self.count)
self.cloud_params["gh_runner_tokens"] = runner_tokens
architecture = self.cloud_params.get("arch", "x64")
release = self.gh.get_latest_runner_release(
platform="linux", architecture=architecture
)
self.cloud_params["runner_release"] = release
self.provider = self.provider_type(**self.cloud_params)
def start_runner_instances(self):
"""Start the runner instances.
This function starts the runner instances and waits for them to be
ready.
"""
print("Starting up...")
# Create a GitHub instance
print("Creating GitHub Actions Runner")
mappings = self.provider.create_instances()
instance_ids = list(mappings.keys())
github_labels = list(mappings.values())
# Output the instance mapping and labels so the stop action can use them
self.provider.set_instance_mapping(mappings)
# Wait for the instance to be ready
print("Waiting for instance to be ready...")
self.provider.wait_until_ready(instance_ids)
print("Instance is ready!")
# Confirm the runner is registered with GitHub
for label in github_labels:
print(f"Waiting for {label}...")
self.gh.wait_for_runner(label, self.timeout)
@dataclass
class TeardownInstance:
"""Class that is used to teardown instances and runners.
Parameters
----------
provider_type : Type[StopCloudInstance]
The type of cloud provider to use.
cloud_params : dict
The parameters to pass to the cloud provider.
gh : GitHubInstance
The GitHub instance to use.
Attributes
----------
provider : StopCloudInstance
The cloud provider instance
provider_type : Type[StopCloudInstance]
cloud_params : dict
gh : GitHub
"""
provider_type: Type[StopCloudInstance]
cloud_params: dict
gh: GitHubInstance
provider: StopCloudInstance = field(init=False)
def __post_init__(self):
"""Initialize the cloud provider.
This function is called after the object is created to correctly
stop the provider.
"""
self.provider = self.provider_type(**self.cloud_params)
def stop_runner_instances(self):
"""Stop the runner instances.
This function stops the runner instances and waits for them to be
removed.
"""
print("Shutting down...")
try:
# Get the instance mapping from our input
mappings = self.provider.get_instance_mapping()
except Exception as e:
error(title="Malformed instance mapping", message=e)
exit(1)
# Remove the runners and instances
print("Removing GitHub Actions Runner")
instance_ids = list(mappings.keys())
labels = list(mappings.values())
for label in labels:
try:
print(f"Removing runner {label}")
self.gh.remove_runner(label)
# This occurs when we have a runner that might already be
# shutdown. Since we are mainly using the ephemeral runners,
# we expect this to happen
except MissingRunnerLabel:
print(f"Runner {label} does not exist, skipping...")
continue
# This is more of the case when we have a failure to remove the
# runner. This is not a concern for the user (because we will
# remove the instance anyways), but we should log it for
# debugging purposes.
except Exception as e:
warning(title="Failed to remove runner", message=e)
print("Removing instances...")
self.provider.remove_instances(instance_ids)
print("Waiting for instance to be removed...")
try:
self.provider.wait_until_removed(instance_ids)
except Exception as e:
# Print to stdout
print(
f"Failed to remove instances check your provider console: {e}"
)
# Print to Annotations
error(
title="Failed to remove instances, check your provider console",
message=e,
)
exit(1)
else:
print("Instances removed!")