argo_container.hpp 4.92 KB
Newer Older
1 2 3 4 5 6 7 8 9 10
/*******************************************************************************
 * Copyright 2019 UChicago Argonne, LLC.
 * (c.f. AUTHORS, LICENSE)
 *
 * This file is part of the Argo containers project.
 * For more info, see https://xgitlab.cels.anl.gov/argo/containers
 *
 * SPDX-License-Identifier: BSD-3-Clause
*******************************************************************************/

Swann Perarnau's avatar
Swann Perarnau committed
11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172
#ifndef __COMPUTE_CONTAINER_HPP__
#define	__COMPUTE_CONTAINER_HPP__

#include "utils.hpp" 
#include <string>
#include <vector>
#include <algorithm>
#include <sys/types.h>
#include "resource_type.hpp"
using std::string;
using std::vector;

#include "string_parser.hpp"

/*The type below would be an enum in real word; but need it to be a used as a
 * bitmask. It can only take the values defined below though*/
typedef uint64_t Container_commit_type;
#define CCT_NONE 0
#define CCT_NAME 1LLU
#define CCT_CPUS (1LLU<<1)
#define CCT_MEMS (1LLU<<2)
#define CCT_TASKS (1LLU<<3)
#define CCT_MEM_EXCLUSIVE (1LLU<<4)
#define	CCT_CPU_EXCLUSIVE (1LLU<<5)
#define	CCT_LOAD_BALANCING (1LLU<<6)
#define CCT_MEM_MIGRATE (1LLU<<7)

#define CCT_ALL	0xffffffffffffffff

class Argo_container
{
	private:
		bool _created;

		string _parent_path;
		Container_commit_type _commit_type;

		string _older_name;  //The last committed one
		string _name;
		uid_t _owner;
		vector<int> _cpus;
		vector<int> _mems;
		bool _cpu_exclusive;
		bool _mem_exclusive;
		bool _load_balancing;
		bool _mem_migrate;

		vector<int> _tasks;
		vector<int> _added_tasks; //added and not committed

		Argo_container(){}
		Argo_container(Argo_container& original);
		Argo_container& operator = (const Argo_container & orig);

	public:
		Argo_container(const string& parent_path, String_parser& sp);

		/*This one builds a non-functional (empty) Argo_container and then 
		 * reloads it. The state of the container is read from the cgroup fs*/
		Argo_container(const string& parent_path,
				const string& name, uid_t owner); 
		~Argo_container();


		inline const string& get_name() const {return _name;}
		inline const string& get_parent_path() const {return _parent_path;}
		inline const uid_t get_owner() const {return _owner;}
		inline void set_owner(uid_t owner) {_owner = owner;}

		inline bool has_cpu(int cpu) const 
		{return std::find(_cpus.begin(), _cpus.end(), cpu) != _cpus.end();}

		/*If no argument, checks if exclusive_cpus is enabled;
		 * otherwise, checks if a cpu is owned exclusively*/
		inline bool has_exclusive_cpu(int cpu = -1) const
		{return _cpu_exclusive && (cpu == -1 || has_cpu(cpu));}

		void set_cpu_exclusive(bool status);

		inline bool has_mem(int mem) const 
		{return std::find(_mems.begin(), _mems.end(), mem) != _mems.end();}

		/*If no argument, checks if _mem_exclusive is enabled;
		 * otherwise, checks if a memory node is owned exclusively*/
		inline bool has_exclusive_mem(int mem = -1) const
		{return _mem_exclusive && (mem == -1 || has_mem(mem));}

		void set_mem_exclusive(bool status);

		inline bool has_load_balancing() const {return _load_balancing;}
		void set_load_balancing(bool status);
		void set_mem_migrate(bool status);

		void add_cpu(int cpu);
		void add_cpus(const vector<int>& cpus);
		void replace_cpus(const vector<int>& cpus);
		void remove_cpu(int cpu);
		void remove_cpus(const vector<int>& cpus);
		void get_cpus(vector<int>& out) const;

		void add_mem(int mem);
		void add_mems(const vector<int>& mems);
		void replace_mems(const vector<int>& mems);
		void remove_mem(int mem);
		void remove_mems(const vector<int>& mems);
		void get_mems(vector<int>& out) const;

		template <typename T>
		void get_resource(vector<T>& out, Resource_type rtype)
		{
			switch(rtype)
			{
				case RT_CPU:
					get_cpus(out); break;
				case RT_MEM:
					get_mems(out); break;
				default:
					ASSERT_MSG(false, "Unsupported resource");
					THROW("Unsupported resource");
			}
		}

		void add_task(int task);
		void add_tasks(const vector<int>& tasks);
		void remove_task(int task);
		void remove_tasks(const vector<int>& tasks);
		void get_tasks(vector<int>& out) const;
		int get_nb_tasks() const {return (int) _tasks.size();}
		inline bool has_any_task() const {return _tasks.size() != 0;}
		inline bool has_task(int task)
		{return get_object_index<int>(_tasks, task) != -1;}
		inline bool has_all_these_tasks(
				const vector<int>& tasks){return is_superset(_tasks, tasks);}
		inline bool has_any_of_these_tasks(
				const vector<int>& tasks){return has_intersect(_tasks, tasks);}


		void rename(const string& new_name);

		void reload_task_list();

		void reload(bool load_tasks = true);
			
		void display() const;
		void display_resources() const;

		void commit_last_changes();
		void commit_all();

		void create();
		void remove();

		void alter(String_parser& sp);

		bool check_task_ownership(bool reload=true, bool throw_on_failure=false);

		/*It is recommended to invoke this method after having invoked
		 * check_task_ownership*/
		void kill_content(bool reload=true);
};

#endif //__COMPUTE_CONTAINER_HPP__