require "ood_core/refinements/hash_extensions" require "json" class OodCore::Job::Adapters::Kubernetes::Batch require_relative "helper" require_relative "k8s_job_info" using OodCore::Refinements::HashExtensions class Error < StandardError; end class NotFoundError < StandardError; end attr_reader :config_file, :bin, :cluster, :mounts attr_reader :all_namespaces, :using_context, :helper attr_reader :username_prefix, :namespace_prefix def initialize(options = {}) options = options.to_h.symbolize_keys @config_file = options.fetch(:config_file, default_config_file) @bin = options.fetch(:bin, '/usr/bin/kubectl') @cluster = options.fetch(:cluster, 'open-ondemand') @mounts = options.fetch(:mounts, []).map { |m| m.to_h.symbolize_keys } @all_namespaces = options.fetch(:all_namespaces, false) @username_prefix = options.fetch(:username_prefix, '') @namespace_prefix = options.fetch(:namespace_prefix, '') @using_context = false @helper = OodCore::Job::Adapters::Kubernetes::Helper.new begin make_kubectl_config(options) rescue # FIXME could use a log here # means you couldn't 'kubectl set config' end end def resource_file(resource_type = 'pod') File.dirname(__FILE__) + "/templates/#{resource_type}.yml.erb" end def submit(script, after: [], afterok: [], afternotok: [], afterany: []) raise ArgumentError, 'Must specify the script' if script.nil? resource_yml, id = generate_id_yml(script) if !script.workdir.nil? && Dir.exist?(script.workdir) File.open(File.join(script.workdir, 'pod.yml'), 'w') { |f| f.write resource_yml } end call("#{formatted_ns_cmd} create -f -", stdin: resource_yml) id end def generate_id(name) # 2_821_109_907_456 = 36**8 name.downcase.tr(' ', '-') + '-' + rand(2_821_109_907_456).to_s(36) end def info_all(attrs: nil) cmd = if all_namespaces "#{base_cmd} get pods -o json --all-namespaces" else "#{namespaced_cmd} get pods -o json" end output = call(cmd) all_pods_to_info(output) end def info_where_owner(owner, attrs: nil) owner = Array.wrap(owner).map(&:to_s) # must at least have job_owner to filter by job_owner attrs = Array.wrap(attrs) | [:job_owner] unless attrs.nil? info_all(attrs: attrs).select { |info| owner.include? info.job_owner } end def info_all_each(attrs: nil) return to_enum(:info_all_each, attrs: attrs) unless block_given? info_all(attrs: attrs).each do |job| yield job end end def info_where_owner_each(owner, attrs: nil) return to_enum(:info_where_owner_each, owner, attrs: attrs) unless block_given? info_where_owner(owner, attrs: attrs).each do |job| yield job end end def info(id) pod_json = safe_call('get', 'pod', id) return OodCore::Job::Info.new({ id: id, status: 'completed' }) if pod_json.empty? service_json = safe_call('get', 'service', service_name(id)) secret_json = safe_call('get', 'secret', secret_name(id)) helper.info_from_json(pod_json: pod_json, service_json: service_json, secret_json: secret_json) end def status(id) info(id).status end def delete(id) safe_call("delete", "pod", id) safe_call("delete", "service", service_name(id)) safe_call("delete", "secret", secret_name(id)) safe_call("delete", "configmap", configmap_name(id)) end private def safe_call(verb, resource, id) begin case verb.to_s when "get" call_json_output('get', resource, id) when "delete" call("#{namespaced_cmd} delete #{resource} #{id}") end rescue NotFoundError {} end end # helper to help format multi-line yaml data from the submit.yml into # mutli-line yaml in the pod.yml.erb def config_data_lines(data) output = [] first = true data.to_s.each_line do |line| output.append(first ? line : line.prepend(" ")) first = false end output end def username @username ||= Etc.getlogin end def k8s_username "#{username_prefix}#{username}" end def user @user ||= Etc.getpwnam(username) end def home_dir user.dir end def run_as_user user.uid end def run_as_group user.gid end def fs_group run_as_group end def group Etc.getgrgid(run_as_group).name end def default_env { USER: username, UID: run_as_user, HOME: home_dir, GROUP: group, GID: run_as_group, KUBECONFIG: '/dev/null', } end # helper to template resource yml you're going to submit and # create an id. def generate_id_yml(script) native_data = script.native container = helper.container_from_native(native_data[:container], default_env) id = generate_id(container.name) configmap = helper.configmap_from_native(native_data, id, script.content) init_containers = helper.init_ctrs_from_native(native_data[:init_containers], container.env) spec = OodCore::Job::Adapters::Kubernetes::Resources::PodSpec.new(container, init_containers: init_containers) all_mounts = native_data[:mounts].nil? ? mounts : mounts + native_data[:mounts] node_selector = native_data[:node_selector].nil? ? {} : native_data[:node_selector] gpu_type = native_data[:gpu_type].nil? ? "nvidia.com/gpu" : native_data[:gpu_type] template = ERB.new(File.read(resource_file), nil, '-') [template.result(binding), id] end # helper to call kubectl and get json data back. # verb, resrouce and id are the kubernetes parlance terms. # example: 'kubectl get pod my-pod-id' is verb=get, resource=pod # and id=my-pod-id def call_json_output(verb, resource, id, stdin: nil) cmd = "#{formatted_ns_cmd} #{verb} #{resource} #{id}" data = call(cmd, stdin: stdin) data = data.empty? ? '{}' : data json_data = JSON.parse(data, symbolize_names: true) json_data end def service_name(id) helper.service_name(id) end def secret_name(id) helper.secret_name(id) end def configmap_name(id) helper.configmap_name(id) end def namespace "#{namespace_prefix}#{username}" end def context cluster end def default_config_file (ENV['KUBECONFIG'] || "#{Dir.home}/.kube/config") end def default_auth { type: 'managaged' }.symbolize_keys end def default_server { endpoint: 'https://localhost:8080', cert_authority_file: nil }.symbolize_keys end def formatted_ns_cmd "#{namespaced_cmd} -o json" end def namespaced_cmd "#{base_cmd} --namespace=#{namespace}" end def base_cmd base = "#{bin} --kubeconfig=#{config_file}" base << " --context=#{context}" if using_context base end def all_pods_to_info(data) json_data = JSON.parse(data, symbolize_names: true) pods = json_data.dig(:items) info_array = [] pods.each do |pod| info = pod_info_from_json(pod) info_array.push(info) unless info.nil? end info_array rescue JSON::ParserError # 'no resources in ' throws parse error [] end def pod_info_from_json(pod) hash = helper.pod_info_from_json(pod) K8sJobInfo.new(hash) rescue Helper::K8sDataError # FIXME: silently eating error, could probably use a logger nil end def make_kubectl_config(config) set_cluster(config.fetch(:server, default_server).to_h.symbolize_keys) configure_auth(config.fetch(:auth, default_auth).to_h.symbolize_keys) end def configure_auth(auth) type = auth.fetch(:type) return if managed?(type) case type when 'gke' set_gke_config(auth) when 'oidc' set_context end end def use_context @using_context = true end def managed?(type) if type.nil? true # maybe should be false? else type.to_s == 'managed' end end def set_gke_config(auth) cred_file = auth.fetch(:svc_acct_file) cmd = "gcloud auth activate-service-account --key-file=#{cred_file}" call(cmd) set_gke_credentials(auth) end def set_gke_credentials(auth) zone = auth.fetch(:zone, nil) region = auth.fetch(:region, nil) locale = '' locale = "--zone=#{zone}" unless zone.nil? locale = "--region=#{region}" unless region.nil? # gke cluster name can probably can differ from what ood calls the cluster cmd = "gcloud container clusters get-credentials #{locale} #{cluster}" env = { 'KUBECONFIG' => config_file } call(cmd, env) end def set_context cmd = "#{base_cmd} config set-context #{cluster}" cmd << " --cluster=#{cluster} --namespace=#{namespace}" cmd << " --user=#{k8s_username}" call(cmd) use_context end def set_cluster(config) server = config.fetch(:endpoint) cert = config.fetch(:cert_authority_file, nil) cmd = "#{base_cmd} config set-cluster #{cluster}" cmd << " --server=#{server}" cmd << " --certificate-authority=#{cert}" unless cert.nil? call(cmd) end def call(cmd = '', env: {}, stdin: nil) o, e, s = Open3.capture3(env, cmd, stdin_data: stdin.to_s) s.success? ? o : interpret_and_raise(e) end def interpret_and_raise(stderr) raise NotFoundError, stderr if /^Error from server \(NotFound\):/.match(stderr) raise(Error, stderr) end end