package azurebatch import ( "context" "encoding/json" "fmt" "io" "io/ioutil" "log" "net/http" "os" "strings" "time" "github.com/Azure/go-autorest/autorest" "github.com/Azure/go-autorest/autorest/azure" "github.com/Azure/azure-sdk-for-go/services/batch/2017-09-01.6.0/batch" "github.com/Azure/go-autorest/autorest/to" "github.com/lawrencegripper/pod2docker" "github.com/virtual-kubelet/virtual-kubelet/manager" azureCreds "github.com/virtual-kubelet/virtual-kubelet/providers/azure" "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" "k8s.io/client-go/tools/remotecommand" ) const ( podJSONKey string = "virtualkubelet_pod" ) // Provider the base struct for the Azure Batch provider type Provider struct { batchConfig *Config ctx context.Context cancelCtx context.CancelFunc fileClient *batch.FileClient resourceManager *manager.ResourceManager listTasks func() (*[]batch.CloudTask, error) addTask func(batch.TaskAddParameter) (autorest.Response, error) getTask func(taskID string) (batch.CloudTask, error) deleteTask func(taskID string) (autorest.Response, error) getFileFromTask func(taskID, path string) (result batch.ReadCloser, err error) nodeName string operatingSystem string cpu string memory string pods string internalIP string daemonEndpointPort int32 } // Config - Basic azure config used to interact with ARM resources. type Config struct { ClientID string ClientSecret string SubscriptionID string TenantID string ResourceGroup string PoolID string JobID string AccountName string AccountLocation string } // NewBatchProvider Creates a batch provider func NewBatchProvider(configString string, rm *manager.ResourceManager, nodeName, operatingSystem string, internalIP string, daemonEndpointPort int32) (*Provider, error) { fmt.Println("Starting create provider") config := &Config{} if azureCredsFilepath := os.Getenv("AZURE_CREDENTIALS_LOCATION"); azureCredsFilepath != "" { creds, err := azureCreds.NewAcsCredential(azureCredsFilepath) if err != nil { return nil, err } config.ClientID = creds.ClientID config.ClientSecret = creds.ClientSecret config.SubscriptionID = creds.SubscriptionID config.TenantID = creds.TenantID } err := getAzureConfigFromEnv(config) if err != nil { log.Println("Failed to get auth information") } return NewBatchProviderFromConfig(config, rm, nodeName, operatingSystem, internalIP, daemonEndpointPort) } // NewBatchProviderFromConfig Creates a batch provider func NewBatchProviderFromConfig(config *Config, rm *manager.ResourceManager, nodeName, operatingSystem string, internalIP string, daemonEndpointPort int32) (*Provider, error) { p := Provider{} p.batchConfig = config // Set sane defaults for Capacity in case config is not supplied p.cpu = "20" p.memory = "100Gi" p.pods = "20" p.resourceManager = rm p.operatingSystem = operatingSystem p.nodeName = nodeName p.internalIP = internalIP p.daemonEndpointPort = daemonEndpointPort p.ctx, p.cancelCtx = context.WithCancel(context.Background()) auth := getAzureADAuthorizer(config, azure.PublicCloud.BatchManagementEndpoint) batchBaseURL := getBatchBaseURL(config.AccountName, config.AccountLocation) _, err := getPool(p.ctx, batchBaseURL, config.PoolID, auth) if err != nil { log.Panicf("Error retreiving Azure Batch pool: %v", err) } _, err = createOrGetJob(p.ctx, batchBaseURL, config.JobID, config.PoolID, auth) if err != nil { log.Panicf("Error retreiving/creating Azure Batch job: %v", err) } taskClient := batch.NewTaskClientWithBaseURI(batchBaseURL) taskClient.Authorizer = auth p.listTasks = func() (*[]batch.CloudTask, error) { res, err := taskClient.List(p.ctx, config.JobID, "", "", "", nil, nil, nil, nil, nil) if err != nil { return &[]batch.CloudTask{}, err } currentTasks := res.Values() for res.NotDone() { err = res.Next() if err != nil { return &[]batch.CloudTask{}, err } pageTasks := res.Values() if pageTasks != nil || len(pageTasks) != 0 { currentTasks = append(currentTasks, pageTasks...) } } return ¤tTasks, nil } p.addTask = func(task batch.TaskAddParameter) (autorest.Response, error) { return taskClient.Add(p.ctx, config.JobID, task, nil, nil, nil, nil) } p.getTask = func(taskID string) (batch.CloudTask, error) { return taskClient.Get(p.ctx, config.JobID, taskID, "", "", nil, nil, nil, nil, "", "", nil, nil) } p.deleteTask = func(taskID string) (autorest.Response, error) { return taskClient.Delete(p.ctx, config.JobID, taskID, nil, nil, nil, nil, "", "", nil, nil) } p.getFileFromTask = func(taskID, path string) (batch.ReadCloser, error) { return p.fileClient.GetFromTask(p.ctx, config.JobID, taskID, path, nil, nil, nil, nil, "", nil, nil) } fileClient := batch.NewFileClientWithBaseURI(batchBaseURL) fileClient.Authorizer = auth p.fileClient = &fileClient return &p, nil } // CreatePod accepts a Pod definition func (p *Provider) CreatePod(ctx context.Context, pod *v1.Pod) error { log.Println("Creating pod...") podCommand, err := pod2docker.GetBashCommand(pod2docker.PodComponents{ InitContainers: pod.Spec.InitContainers, Containers: pod.Spec.Containers, PodName: pod.Name, Volumes: pod.Spec.Volumes, }) if err != nil { return err } bytes, err := json.Marshal(pod) if err != nil { panic(err) } task := batch.TaskAddParameter{ DisplayName: to.StringPtr(string(pod.UID)), ID: to.StringPtr(getTaskIDForPod(pod.Namespace, pod.Name)), CommandLine: to.StringPtr(fmt.Sprintf(`/bin/bash -c "%s"`, podCommand)), UserIdentity: &batch.UserIdentity{ AutoUser: &batch.AutoUserSpecification{ ElevationLevel: batch.Admin, Scope: batch.Pool, }, }, EnvironmentSettings: &[]batch.EnvironmentSetting{ { Name: to.StringPtr(podJSONKey), Value: to.StringPtr(string(bytes)), }, }, } _, err = p.addTask(task) if err != nil { return err } return nil } // GetPodStatus retrieves the status of a given pod by name. func (p *Provider) GetPodStatus(ctx context.Context, namespace, name string) (*v1.PodStatus, error) { log.Println("Getting pod status ....") pod, err := p.GetPod(ctx, namespace, name) if err != nil { return nil, err } if pod == nil { return nil, nil } return &pod.Status, nil } // UpdatePod accepts a Pod definition func (p *Provider) UpdatePod(ctx context.Context, pod *v1.Pod) error { log.Println("Pod Update called: No-op as not implemented") return nil } // DeletePod accepts a Pod definition func (p *Provider) DeletePod(ctx context.Context, pod *v1.Pod) error { taskID := getTaskIDForPod(pod.Namespace, pod.Name) task, err := p.deleteTask(taskID) if err != nil { log.Println(task) log.Println(err) return wrapError(err) } log.Printf(fmt.Sprintf("Deleting task: %v", taskID)) return nil } // GetPod returns a pod by name func (p *Provider) GetPod(ctx context.Context, namespace, name string) (*v1.Pod, error) { log.Println("Getting Pod ...") task, err := p.getTask(getTaskIDForPod(namespace, name)) if err != nil { if task.Response.StatusCode == http.StatusNotFound { return nil, nil } log.Println(err) return nil, err } pod, err := getPodFromTask(&task) if err != nil { panic(err) } status, _ := convertTaskToPodStatus(&task) pod.Status = *status return pod, nil } // GetContainerLogs returns the logs of a container running in a pod by name. func (p *Provider) GetContainerLogs(ctx context.Context, namespace, podName, containerName string, tail int) (string, error) { log.Println("Getting pod logs ....") taskID := getTaskIDForPod(namespace, podName) logFileLocation := fmt.Sprintf("wd/%s.log", containerName) containerLogReader, err := p.getFileFromTask(taskID, logFileLocation) if containerLogReader.Response.Response != nil && containerLogReader.StatusCode == http.StatusNotFound { stdoutReader, err := p.getFileFromTask(taskID, "stdout.txt") if err != nil { return "", err } stderrReader, err := p.getFileFromTask(taskID, "stderr.txt") if err != nil { return "", err } var builder strings.Builder builderPtr := &builder mustWriteString(builderPtr, "Container still starting....\n") mustWriteString(builderPtr, "Showing startup logs from Azure Batch node instead:\n") mustWriteString(builderPtr, "----- STDOUT -----\n") stdoutBytes, _ := ioutil.ReadAll(*stdoutReader.Value) mustWrite(builderPtr, stdoutBytes) mustWriteString(builderPtr, "\n") mustWriteString(builderPtr, "----- STDERR -----\n") stderrBytes, _ := ioutil.ReadAll(*stderrReader.Value) mustWrite(builderPtr, stderrBytes) mustWriteString(builderPtr, "\n") return builder.String(), nil } if err != nil { return "", err } result, err := formatLogJSON(containerLogReader) if err != nil { return "", fmt.Errorf("Container log formating failed err: %v", err) } return result, nil } // Get full pod name as defined in the provider context // TODO: Implementation func (p *Provider) GetPodFullName(namespace string, pod string) string { return "" } // ExecInContainer executes a command in a container in the pod, copying data // between in/out/err and the container's stdin/stdout/stderr. // TODO: Implementation func (p *Provider) ExecInContainer(name string, uid types.UID, container string, cmd []string, in io.Reader, out, err io.WriteCloser, tty bool, resize <-chan remotecommand.TerminalSize, timeout time.Duration) error { log.Printf("receive ExecInContainer %q\n", container) return nil } // GetPods retrieves a list of all pods scheduled to run. func (p *Provider) GetPods(ctx context.Context) ([]*v1.Pod, error) { log.Println("Getting pods...") tasksPtr, err := p.listTasks() if err != nil { panic(err) } if tasksPtr == nil { return []*v1.Pod{}, nil } tasks := *tasksPtr pods := make([]*v1.Pod, len(tasks), len(tasks)) for i, t := range tasks { pod, err := getPodFromTask(&t) if err != nil { panic(err) } pods[i] = pod } return pods, nil } // Capacity returns a resource list containing the capacity limits func (p *Provider) Capacity(ctx context.Context) v1.ResourceList { return v1.ResourceList{ "cpu": resource.MustParse(p.cpu), "memory": resource.MustParse(p.memory), "pods": resource.MustParse(p.pods), "nvidia.com/gpu": resource.MustParse("1"), } } // NodeConditions returns a list of conditions (Ready, OutOfDisk, etc), for updates to the node status // within Kubernetes. func (p *Provider) NodeConditions(ctx context.Context) []v1.NodeCondition { return []v1.NodeCondition{ { Type: "Ready", Status: v1.ConditionTrue, LastHeartbeatTime: metav1.Now(), LastTransitionTime: metav1.Now(), Reason: "KubeletReady", Message: "kubelet is ready.", }, { Type: "OutOfDisk", Status: v1.ConditionFalse, LastHeartbeatTime: metav1.Now(), LastTransitionTime: metav1.Now(), Reason: "KubeletHasSufficientDisk", Message: "kubelet has sufficient disk space available", }, { Type: "MemoryPressure", Status: v1.ConditionFalse, LastHeartbeatTime: metav1.Now(), LastTransitionTime: metav1.Now(), Reason: "KubeletHasSufficientMemory", Message: "kubelet has sufficient memory available", }, { Type: "DiskPressure", Status: v1.ConditionFalse, LastHeartbeatTime: metav1.Now(), LastTransitionTime: metav1.Now(), Reason: "KubeletHasNoDiskPressure", Message: "kubelet has no disk pressure", }, { Type: "NetworkUnavailable", Status: v1.ConditionFalse, LastHeartbeatTime: metav1.Now(), LastTransitionTime: metav1.Now(), Reason: "RouteCreated", Message: "RouteController created a route", }, } } // NodeAddresses returns a list of addresses for the node status // within Kubernetes. func (p *Provider) NodeAddresses(ctx context.Context) []v1.NodeAddress { // TODO: Make these dynamic and augment with custom ACI specific conditions of interest return []v1.NodeAddress{ { Type: "InternalIP", Address: p.internalIP, }, } } // NodeDaemonEndpoints returns NodeDaemonEndpoints for the node status // within Kubernetes. func (p *Provider) NodeDaemonEndpoints(ctx context.Context) *v1.NodeDaemonEndpoints { return &v1.NodeDaemonEndpoints{ KubeletEndpoint: v1.DaemonEndpoint{ Port: p.daemonEndpointPort, }, } } // OperatingSystem returns the operating system for this provider. func (p *Provider) OperatingSystem() string { return p.operatingSystem }