10
0
mirror of https://github.com/LCPQ/quantum_package synced 2024-12-22 20:35:19 +01:00
quantum_package/ocaml/TaskServer.ml

728 lines
19 KiB
OCaml
Raw Normal View History

2017-08-18 18:28:33 +02:00
open Core
2015-12-03 23:51:10 +01:00
open Qptypes
2016-07-26 18:29:52 +02:00
type pub_state =
| Waiting
| Running of string
| Stopped
let pub_state_of_string = function
| "Waiting" -> Waiting
| "Stopped" -> Stopped
| s -> Running s
let string_of_pub_state = function
| Waiting -> "Waiting"
| Stopped -> "Stopped"
| Running s -> s
2016-02-19 00:20:28 +01:00
type t =
{
queue : Queuing_system.t ;
state : Message.State.t option ;
address_tcp : Address.Tcp.t option ;
address_inproc : Address.Inproc.t option ;
progress_bar : Progress_bar.t option ;
running : bool;
2017-11-26 23:47:11 +01:00
data : (string, string) Hashtbl.t;
2016-02-19 00:20:28 +01:00
}
2015-12-03 23:51:10 +01:00
2016-02-19 00:20:28 +01:00
let debug_env =
match Sys.getenv "QP_TASK_DEBUG" with
| Some x -> x <> ""
| None -> false
2015-12-03 23:51:10 +01:00
2016-02-19 00:20:28 +01:00
let debug str =
if debug_env then
2017-11-24 01:02:53 +01:00
Printf.eprintf "TASK : %s%!" str
2016-02-19 00:20:28 +01:00
let zmq_context =
ZMQ.Context.create ()
2017-04-07 18:58:08 +02:00
let () =
2017-05-18 14:36:06 +02:00
ZMQ.Context.set_io_threads zmq_context 8
2017-04-07 18:58:08 +02:00
2015-12-03 23:51:10 +01:00
2016-10-12 11:26:21 +02:00
let bind_socket ~socket_type ~socket ~port =
2016-03-04 12:11:38 +01:00
let rec loop = function
| 0 -> failwith @@ Printf.sprintf
2016-10-12 11:26:21 +02:00
"Unable to bind the %s socket to port : %d "
socket_type port
2016-03-07 20:13:43 +01:00
| -1 -> ()
2016-03-04 12:11:38 +01:00
| i ->
try
2016-10-12 11:26:21 +02:00
ZMQ.Socket.bind socket @@ Printf.sprintf "tcp://*:%d" port;
2016-03-04 12:11:38 +01:00
loop (-1)
with
2017-08-18 18:28:33 +02:00
| Unix.Unix_error _ -> (Time.pause @@ Time.Span.of_sec 1. ; loop (i-1) )
2016-03-04 12:11:38 +01:00
| other_exception -> raise other_exception
2017-05-04 15:51:33 +02:00
in loop 60
2015-12-03 23:51:10 +01:00
let hostname = lazy (
try
Unix.gethostname ()
with
| _ -> "localhost"
2016-02-19 00:20:28 +01:00
)
2015-12-03 23:51:10 +01:00
let ip_address = lazy (
match Sys.getenv "QP_NIC" with
| None ->
begin
try
Lazy.force hostname
|> Unix.Inet_addr.of_string_or_getbyname
|> Unix.Inet_addr.to_string
with
| Unix.Unix_error _ ->
failwith "Unable to find IP address from host name."
end
| Some interface ->
begin
try
ok_exn Linux_ext.get_ipv4_address_for_interface interface
with
| Unix.Unix_error _ ->
Lazy.force hostname
|> Unix.Inet_addr.of_string_or_getbyname
|> Unix.Inet_addr.to_string
end
)
2016-02-19 00:20:28 +01:00
let reply_ok rep_socket =
2017-02-27 19:24:19 +01:00
Message.Ok_msg.create
2016-02-19 00:20:28 +01:00
|> Message.Ok_msg.to_string
|> ZMQ.Socket.send rep_socket
let reply_wrong_state rep_socket =
Printf.printf "WRONG STATE\n%!";
Message.Error_msg.create "Wrong state"
|> Message.Error_msg.to_string
|> ZMQ.Socket.send rep_socket
2015-12-07 22:03:33 +01:00
let stop ~port =
2016-02-19 00:20:28 +01:00
debug "STOP";
let req_socket =
ZMQ.Socket.create zmq_context ZMQ.Socket.req
and address =
2017-05-04 15:51:33 +02:00
Printf.sprintf "tcp://localhost:%d" port
2016-02-19 00:20:28 +01:00
in
ZMQ.Socket.set_linger_period req_socket 1_000_000;
ZMQ.Socket.connect req_socket address;
2015-12-03 23:51:10 +01:00
2017-02-27 19:24:19 +01:00
Message.Terminate (Message.Terminate_msg.create)
2016-02-19 00:20:28 +01:00
|> Message.to_string
|> ZMQ.Socket.send req_socket ;
2015-12-03 23:51:10 +01:00
2016-02-19 00:20:28 +01:00
let msg =
ZMQ.Socket.recv req_socket
|> Message.of_string
2015-12-03 23:51:10 +01:00
in
2016-02-19 00:20:28 +01:00
let () =
match msg with
| Message.Ok _ -> ()
| _ -> failwith "Problem in termination"
in
ZMQ.Socket.set_linger_period req_socket 1_000;
ZMQ.Socket.close req_socket
2015-12-03 23:51:10 +01:00
2016-07-26 18:29:52 +02:00
let new_job msg program_state rep_socket pair_socket =
2015-12-03 23:51:10 +01:00
2016-02-19 00:20:28 +01:00
let state =
msg.Message.Newjob_msg.state
in
let progress_bar =
Progress_bar.init
~start_value:0.
~end_value:1.
~bar_length:20
~title:(Message.State.to_string state)
in
let result =
{ program_state with
state = Some state ;
progress_bar = Some progress_bar ;
address_tcp = Some msg.Message.Newjob_msg.address_tcp;
address_inproc = Some msg.Message.Newjob_msg.address_inproc;
}
in
reply_ok rep_socket;
2016-08-01 16:05:40 +02:00
string_of_pub_state Waiting
2016-07-26 18:29:52 +02:00
|> ZMQ.Socket.send pair_socket ;
2016-02-19 00:20:28 +01:00
result
2016-08-01 16:05:40 +02:00
let change_pub_state msg program_state rep_socket pair_socket =
let msg =
match msg with
| `Waiting -> Waiting
| `Stopped -> Stopped
| `Running ->
begin
let state =
match program_state.state with
| Some x -> x
| None -> failwith "Trying to change pub state while no job is ready"
in
Running (Message.State.to_string state)
end
in
reply_ok rep_socket;
string_of_pub_state msg
|> ZMQ.Socket.send pair_socket ;
program_state
2016-02-19 00:20:28 +01:00
2016-07-26 18:29:52 +02:00
let end_job msg program_state rep_socket pair_socket =
2016-02-19 00:20:28 +01:00
let failure () =
reply_wrong_state rep_socket;
program_state
and success state =
reply_ok rep_socket;
{ program_state with
state = None ;
2017-05-02 16:18:02 +02:00
progress_bar = Progress_bar.clear ();
2016-02-19 00:20:28 +01:00
}
in
match program_state.state with
| None -> failure ()
| Some state ->
begin
2016-02-19 00:20:28 +01:00
if (msg.Message.Endjob_msg.state = state) then
2016-07-26 18:29:52 +02:00
begin
string_of_pub_state Waiting
|> ZMQ.Socket.send pair_socket ;
success state
end
2016-02-19 00:20:28 +01:00
else
failure ()
end
let connect msg program_state rep_socket =
let state =
match program_state.state with
| Some state -> state
| None -> assert false
in
let push_address =
match msg with
| Message.Connect_msg.Tcp ->
begin
2016-02-19 00:20:28 +01:00
match program_state.address_tcp with
| Some address -> Address.Tcp address
| None -> failwith "Error: No TCP address"
end
2016-02-19 00:20:28 +01:00
| Message.Connect_msg.Inproc ->
begin
2016-02-19 00:20:28 +01:00
match program_state.address_inproc with
| Some address -> Address.Inproc address
| None -> failwith "Error: No inproc address"
end
2016-02-19 00:20:28 +01:00
| Message.Connect_msg.Ipc -> assert false
in
let new_queue, client_id =
Queuing_system.add_client program_state.queue
in
Message.ConnectReply (Message.ConnectReply_msg.create
~state:state ~client_id ~push_address)
|> Message.to_string
|> ZMQ.Socket.send rep_socket ;
{ program_state with
queue = new_queue
}
let disconnect msg program_state rep_socket =
let state, client_id =
msg.Message.Disconnect_msg.state,
msg.Message.Disconnect_msg.client_id
in
let failure () =
reply_wrong_state rep_socket;
program_state
and success () =
let new_program_state =
{ program_state with
queue = Queuing_system.del_client ~client_id program_state.queue
}
in
Message.DisconnectReply (Message.DisconnectReply_msg.create ~state)
|> Message.to_string
|> ZMQ.Socket.send rep_socket ;
new_program_state
in
match program_state.state with
| None -> assert false
| Some state' ->
begin
if (state = state') then
success ()
else
failure ()
end
2016-02-19 00:20:28 +01:00
let del_task msg program_state rep_socket =
2017-02-28 13:28:36 +01:00
let state, task_ids =
2016-02-19 00:20:28 +01:00
msg.Message.DelTask_msg.state,
2017-02-28 13:28:36 +01:00
msg.Message.DelTask_msg.task_ids
2016-02-19 00:20:28 +01:00
in
let failure () =
reply_wrong_state rep_socket;
program_state
and success () =
let new_program_state =
{ program_state with
2017-02-28 13:28:36 +01:00
queue = List.fold ~f:(fun queue task_id -> Queuing_system.del_task ~task_id queue)
~init:program_state.queue task_ids
2016-02-19 00:20:28 +01:00
}
in
let more =
2016-10-13 12:32:22 +02:00
(Queuing_system.number_of_tasks new_program_state.queue > 0)
2016-02-19 00:20:28 +01:00
in
2017-02-28 13:28:36 +01:00
Message.DelTaskReply (Message.DelTaskReply_msg.create ~task_ids ~more)
2016-02-19 00:20:28 +01:00
|> Message.to_string
|> ZMQ.Socket.send ~block:true rep_socket ; (** /!\ Has to be blocking *)
new_program_state
in
match program_state.state with
| None -> assert false
| Some state' ->
begin
if (state = state') then
success ()
else
failure ()
end
let add_task msg program_state rep_socket =
2017-02-28 13:28:36 +01:00
let state, tasks =
2016-02-19 00:20:28 +01:00
msg.Message.AddTask_msg.state,
2017-02-28 13:28:36 +01:00
msg.Message.AddTask_msg.tasks
2016-02-19 00:20:28 +01:00
in
let increment_progress_bar = function
| Some bar -> Some (Progress_bar.increment_end bar)
| None -> None
in
let result =
2017-02-28 14:46:32 +01:00
let new_queue, new_bar =
List.fold ~f:(fun (queue, bar) task ->
Queuing_system.add_task ~task queue,
increment_progress_bar bar)
~init:(program_state.queue, program_state.progress_bar) tasks
in
2017-02-28 13:28:36 +01:00
{ program_state with
2017-02-28 14:46:32 +01:00
queue = new_queue;
progress_bar = new_bar
2017-02-28 13:28:36 +01:00
}
2016-02-19 00:20:28 +01:00
in
reply_ok rep_socket;
result
2016-07-26 18:29:52 +02:00
let get_task msg program_state rep_socket pair_socket =
2016-02-19 00:20:28 +01:00
let state, client_id =
msg.Message.GetTask_msg.state,
msg.Message.GetTask_msg.client_id
in
let failure () =
reply_wrong_state rep_socket;
program_state
and success () =
let new_queue, task_id, task =
Queuing_system.pop_task ~client_id program_state.queue
in
2017-05-03 19:23:12 +02:00
let no_task =
Queuing_system.number_of_queued new_queue = 0
in
if no_task then
2016-07-26 18:29:52 +02:00
string_of_pub_state Waiting
|> ZMQ.Socket.send pair_socket
else
string_of_pub_state (Running (Message.State.to_string state))
|> ZMQ.Socket.send pair_socket;
2016-02-19 00:20:28 +01:00
let new_program_state =
{ program_state with
queue = new_queue
}
in
2016-08-01 17:34:36 +02:00
Message.GetTaskReply (Message.GetTaskReply_msg.create ~task ~task_id)
|> Message.to_string
|> ZMQ.Socket.send rep_socket ;
new_program_state
2016-02-19 00:20:28 +01:00
in
match program_state.state with
| None -> assert false
| Some state' ->
begin
if (state = state') then
success ()
else
failure ()
end
let task_done msg program_state rep_socket =
2017-02-28 13:28:36 +01:00
let state, client_id, task_ids =
2015-12-03 23:51:10 +01:00
msg.Message.TaskDone_msg.state,
msg.Message.TaskDone_msg.client_id,
2017-02-28 13:28:36 +01:00
msg.Message.TaskDone_msg.task_ids
2015-12-03 23:51:10 +01:00
in
2016-02-19 00:20:28 +01:00
let increment_progress_bar = function
| Some bar -> Some (Progress_bar.increment_cur bar)
| None -> None
in
let failure () =
reply_wrong_state rep_socket;
program_state
and success () =
2017-02-28 14:46:32 +01:00
let new_queue, new_bar =
List.fold ~f:(fun (queue, bar) task_id ->
Queuing_system.end_task ~task_id ~client_id queue,
increment_progress_bar bar)
~init:(program_state.queue, program_state.progress_bar) task_ids
in
2016-02-19 00:20:28 +01:00
let result =
{ program_state with
2017-02-28 14:46:32 +01:00
queue = new_queue;
progress_bar = new_bar
2016-02-19 00:20:28 +01:00
}
in
reply_ok rep_socket;
result
in
match program_state.state with
| None -> assert false
| Some state' ->
begin
if (state = state') then
success ()
else
failure ()
end
2017-11-26 23:47:11 +01:00
let put_data msg rest_of_msg program_state rep_socket =
2017-05-03 21:15:54 +02:00
2017-11-26 23:47:11 +01:00
debug (Message.PutData_msg.to_string msg);
let () =
let key, value =
msg.Message.PutData_msg.key,
match rest_of_msg with
| [ x ] -> x
| _ -> failwith "Badly formed put_data message"
in
Hashtbl.set program_state.data ~key ~data:value ;
Message.PutDataReply (Message.PutDataReply_msg.create ())
|> Message.to_string
|> ZMQ.Socket.send rep_socket
2017-05-03 21:15:54 +02:00
in
2017-11-26 23:47:11 +01:00
program_state
2017-05-03 21:15:54 +02:00
2017-11-26 23:47:11 +01:00
let get_data msg program_state rep_socket =
2017-05-03 21:15:54 +02:00
2017-11-26 23:47:11 +01:00
debug (Message.GetData_msg.to_string msg);
let () =
let key =
msg.Message.GetData_msg.key
2017-05-03 21:15:54 +02:00
in
2017-11-26 23:47:11 +01:00
let value =
match Hashtbl.find program_state.data key with
| Some value -> value
| None -> ""
in
Message.GetDataReply (Message.GetDataReply_msg.create ~value)
|> Message.to_string_list
|> ZMQ.Socket.send_all rep_socket
in
program_state
2017-05-03 21:15:54 +02:00
2016-02-19 00:20:28 +01:00
let terminate program_state rep_socket =
reply_ok rep_socket;
{ program_state with
2016-08-01 16:05:40 +02:00
address_tcp = None;
address_inproc = None;
2016-02-19 00:20:28 +01:00
running = false
}
2017-05-10 00:04:34 +02:00
let abort program_state rep_socket =
let queue, client_id =
Queuing_system.add_client program_state.queue
in
let rec aux accu queue = function
| 0 -> (queue, accu)
| rest ->
let new_queue, task_id, _ =
Queuing_system.pop_task ~client_id queue
in
let new_accu =
match task_id with
| Some task_id -> task_id::accu
| None -> accu
in
Queuing_system.number_of_queued new_queue
|> aux new_accu new_queue
in
let queue, tasks =
aux [] queue 1
in
let queue =
List.fold ~f:(fun queue task_id ->
Queuing_system.end_task ~task_id ~client_id queue)
~init:queue tasks
in
let queue =
List.fold ~f:(fun queue task_id -> Queuing_system.del_task ~task_id queue)
~init:queue tasks
in
reply_ok rep_socket;
{ program_state with
queue
}
2016-02-19 00:20:28 +01:00
let error msg program_state rep_socket =
Message.Error (Message.Error_msg.create msg)
|> Message.to_string
|> ZMQ.Socket.send rep_socket ;
program_state
2016-07-26 18:29:52 +02:00
let start_pub_thread ~port =
Thread.create (fun () ->
let timeout =
1000
in
let pair_socket =
ZMQ.Socket.create zmq_context ZMQ.Socket.pair
and address =
"inproc://pair"
in
ZMQ.Socket.connect pair_socket address;
2016-02-19 00:20:28 +01:00
2016-07-26 18:29:52 +02:00
let pub_socket =
ZMQ.Socket.create zmq_context ZMQ.Socket.pub
in
2016-10-12 11:26:21 +02:00
bind_socket ~socket_type:"PUB" ~socket:pub_socket ~port;
2016-07-26 18:29:52 +02:00
let pollitem =
ZMQ.Poll.mask_of
[| (pair_socket, ZMQ.Poll.In) |]
in
let rec run state =
let new_state =
let polling =
ZMQ.Poll.poll ~timeout pollitem
in
if (polling.(0) = Some ZMQ.Poll.In) then
ZMQ.Socket.recv ~block:false pair_socket
|> pub_state_of_string
else
state
in
ZMQ.Socket.send pub_socket @@ string_of_pub_state new_state;
match state with
| Stopped -> ()
| _ -> run new_state
in
run Waiting;
ZMQ.Socket.set_linger_period pair_socket 1000 ;
ZMQ.Socket.close pair_socket;
ZMQ.Socket.set_linger_period pub_socket 1000 ;
ZMQ.Socket.close pub_socket;
)
2016-02-19 00:20:28 +01:00
let run ~port =
2016-07-26 18:29:52 +02:00
(** Bind inproc socket for changing state of pub *)
let pair_socket =
ZMQ.Socket.create zmq_context ZMQ.Socket.pair
and address =
"inproc://pair"
in
2016-10-12 11:26:21 +02:00
ZMQ.Socket.bind pair_socket address;
2016-07-26 18:29:52 +02:00
let pub_thread =
start_pub_thread ~port:(port+1) ()
in
2016-02-19 00:20:28 +01:00
(** Bind REP socket *)
let rep_socket =
ZMQ.Socket.create zmq_context ZMQ.Socket.rep
in
ZMQ.Socket.set_linger_period rep_socket 1_000_000;
2016-10-12 11:26:21 +02:00
bind_socket "REP" rep_socket port;
2016-02-19 00:20:28 +01:00
let initial_program_state =
{ queue = Queuing_system.create () ;
running = true ;
state = None;
address_tcp = None;
address_inproc = None;
progress_bar = None ;
2017-11-26 23:47:11 +01:00
data = Hashtbl.create ~hashable:String.hashable ();
2016-02-19 00:20:28 +01:00
}
in
(** ZMR polling item *)
let pollitem =
ZMQ.Poll.mask_of
[| (rep_socket, ZMQ.Poll.In) |]
in
2016-07-26 18:29:52 +02:00
let address =
Printf.sprintf "tcp://%s:%d" (Lazy.force ip_address) port
in
2016-02-19 00:20:28 +01:00
Printf.printf "Task server running : %s\n%!" address;
(** Main loop *)
let rec main_loop program_state = function
| false -> ()
| true ->
let polling =
ZMQ.Poll.poll ~timeout:1000 pollitem
2015-12-03 23:51:10 +01:00
in
2016-02-19 00:20:28 +01:00
if (polling.(0) <> Some ZMQ.Poll.In) then
main_loop program_state true
else
begin
let program_state =
match program_state.progress_bar with
| None -> program_state
| Some bar ->
if bar.Progress_bar.dirty then
{ program_state with
progress_bar = Some (Progress_bar.display bar)
}
else
program_state
in
(** Extract message *)
let raw_message, rest =
match ZMQ.Socket.recv_all rep_socket with
| x :: rest -> x, rest
| [] -> failwith "Badly formed message"
in
let message =
Message.of_string raw_message
in
(** Debug input *)
2017-05-03 19:23:12 +02:00
let () =
if debug_env then
begin
Printf.sprintf "q:%d r:%d n:%d : %s\n%!"
(Queuing_system.number_of_queued program_state.queue)
(Queuing_system.number_of_running program_state.queue)
(Queuing_system.number_of_tasks program_state.queue)
(Message.to_string message)
|> debug
end
in
2016-02-19 00:20:28 +01:00
let new_program_state =
try
match program_state.state, message with
| _ , Message.Terminate _ -> terminate program_state rep_socket
2017-05-10 00:04:34 +02:00
| _ , Message.Abort _ -> abort program_state rep_socket
2017-11-26 23:47:11 +01:00
| _ , Message.PutData x -> put_data x rest program_state rep_socket
| _ , Message.GetData x -> get_data x program_state rep_socket
2016-07-26 18:29:52 +02:00
| None , Message.Newjob x -> new_job x program_state rep_socket pair_socket
2016-02-19 00:20:28 +01:00
| _ , Message.Newjob _ -> error "A job is already running" program_state rep_socket
2016-07-26 18:29:52 +02:00
| Some _, Message.Endjob x -> end_job x program_state rep_socket pair_socket
2016-08-01 16:05:40 +02:00
| Some _, Message.SetRunning -> change_pub_state `Running program_state rep_socket pair_socket
| _, Message.SetWaiting -> change_pub_state `Waiting program_state rep_socket pair_socket
| _, Message.SetStopped -> change_pub_state `Stopped program_state rep_socket pair_socket
2016-02-19 00:20:28 +01:00
| None , _ -> error "No job is running" program_state rep_socket
| Some _, Message.Connect x -> connect x program_state rep_socket
| Some _, Message.Disconnect x -> disconnect x program_state rep_socket
| Some _, Message.AddTask x -> add_task x program_state rep_socket
| Some _, Message.DelTask x -> del_task x program_state rep_socket
2016-07-26 18:29:52 +02:00
| Some _, Message.GetTask x -> get_task x program_state rep_socket pair_socket
2016-02-19 00:20:28 +01:00
| Some _, Message.TaskDone x -> task_done x program_state rep_socket
| _ , _ ->
error ("Invalid message : "^(Message.to_string message)) program_state rep_socket
with
| Failure f ->
error (f^" : "^raw_message) program_state rep_socket
| Assert_failure (f,i,j) ->
error (Printf.sprintf "%s:%d:%d : %s" f i j raw_message) program_state rep_socket
in
main_loop new_program_state new_program_state.running
end
in main_loop initial_program_state true;
2016-07-26 18:29:52 +02:00
ZMQ.Socket.send pair_socket @@ string_of_pub_state Stopped;
Thread.join pub_thread;
2016-10-12 11:26:21 +02:00
ZMQ.Socket.close rep_socket
2016-07-26 18:29:52 +02:00
2015-12-03 23:51:10 +01:00
2016-02-19 00:20:28 +01:00