ozzyable commited on
Commit
e89d17f
·
1 Parent(s): a32f421

Added cancelling and runsync

Browse files
Files changed (1) hide show
  1. app.py +33 -12
app.py CHANGED
@@ -12,6 +12,12 @@ model_selection = st.radio("Select model to use:", ["Tinyllama", "Gemma"])
12
  new_transactions_input = st.text_area("Enter your transactions (comma-separated)", key="input_area")
13
  submit_button = st.button("Submit New Transactions", type="primary")
14
 
 
 
 
 
 
 
15
  if submit_button:
16
  # Split transactions and strip whitespace
17
  new_transactions = [i.strip() for i in new_transactions_input.split(',') if i.strip()]
@@ -24,10 +30,9 @@ if model_selection == "Gemma":
24
  elif model_selection == "Tinyllama":
25
  base_url = "https://api.runpod.ai/v2/0wnm75vx5o77s1/"
26
 
27
-
28
  # When submit button is clicked and there are transactions to process
29
  if submit_button and new_transactions:
30
- url = base_url + "run"
31
 
32
  # Retrieve API key from Streamlit secrets
33
  api_key = st.secrets["api_key"]
@@ -47,23 +52,24 @@ if submit_button and new_transactions:
47
  # Show a spinner while waiting for the response
48
  with st.spinner("Processing..."):
49
  try:
50
-
51
  # Send POST request to start processing
52
  response = requests.post(url, headers=headers, data=json_data)
53
  response.raise_for_status() # Raise an error for bad status codes
54
 
55
  # Parse response to get job ID
56
  result = response.json()
57
- job_id = result['id']
58
- st.write(f"New Job ID: {job_id}")
59
 
60
- # Keep checking status until it's no longer 'IN_QUEUE'
61
- status_url = f"{base_url}status/{job_id}"
62
- status = "IN_QUEUE"
63
- while status != "COMPLETED":
64
  status_response = requests.get(status_url, headers=headers)
65
  status_data = status_response.json()
66
- status = status_data.get('status', '')
 
 
67
  time.sleep(2) # Adjust interval as needed
68
 
69
  # Once status changes, display final status
@@ -71,8 +77,23 @@ if submit_button and new_transactions:
71
 
72
  except requests.exceptions.RequestException as e:
73
  st.error(f"An error occurred: {e}")
74
- finally:
75
- st.spinner(False) # Turn off the spinner after response is received
76
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
77
  if st.button("Reset", key="reset_button"):
 
 
78
  st.text_area("Enter your transactions (comma-separated)", value="", key="reset_area")
 
12
  new_transactions_input = st.text_area("Enter your transactions (comma-separated)", key="input_area")
13
  submit_button = st.button("Submit New Transactions", type="primary")
14
 
15
+ # Session state to keep track of job ID and status
16
+ if "job_id" not in st.session_state:
17
+ st.session_state.job_id = None
18
+ if "job_status" not in st.session_state:
19
+ st.session_state.job_status = None
20
+
21
  if submit_button:
22
  # Split transactions and strip whitespace
23
  new_transactions = [i.strip() for i in new_transactions_input.split(',') if i.strip()]
 
30
  elif model_selection == "Tinyllama":
31
  base_url = "https://api.runpod.ai/v2/0wnm75vx5o77s1/"
32
 
 
33
  # When submit button is clicked and there are transactions to process
34
  if submit_button and new_transactions:
35
+ url = base_url + "runsync"
36
 
37
  # Retrieve API key from Streamlit secrets
38
  api_key = st.secrets["api_key"]
 
52
  # Show a spinner while waiting for the response
53
  with st.spinner("Processing..."):
54
  try:
 
55
  # Send POST request to start processing
56
  response = requests.post(url, headers=headers, data=json_data)
57
  response.raise_for_status() # Raise an error for bad status codes
58
 
59
  # Parse response to get job ID
60
  result = response.json()
61
+ st.session_state.job_id = result['id']
62
+ st.write(f"New Job ID: {st.session_state.job_id}")
63
 
64
+ # Keep checking status until it's no longer 'IN_QUEUE' or cancelled
65
+ status_url = f"{base_url}status/{st.session_state.job_id}"
66
+ st.session_state.job_status = "IN_QUEUE"
67
+ while st.session_state.job_status == "IN_QUEUE":
68
  status_response = requests.get(status_url, headers=headers)
69
  status_data = status_response.json()
70
+ st.session_state.job_status = status_data.get('status', '')
71
+ if st.session_state.job_status != "IN_QUEUE":
72
+ break
73
  time.sleep(2) # Adjust interval as needed
74
 
75
  # Once status changes, display final status
 
77
 
78
  except requests.exceptions.RequestException as e:
79
  st.error(f"An error occurred: {e}")
 
 
80
 
81
+ # Cancel button
82
+ if st.session_state.job_id and st.session_state.job_status == "IN_QUEUE":
83
+ cancel_button = st.button("Cancel Request")
84
+ if cancel_button:
85
+ cancel_url = f"{base_url}cancel/{st.session_state.job_id}"
86
+ try:
87
+ cancel_response = requests.post(cancel_url, headers=headers)
88
+ cancel_response.raise_for_status()
89
+ cancel_result = cancel_response.json()
90
+ st.session_state.job_status = "CANCELLED"
91
+ st.write(f"Job {st.session_state.job_id} has been cancelled.")
92
+ except requests.exceptions.RequestException as e:
93
+ st.error(f"An error occurred while cancelling: {e}")
94
+
95
+ # Reset button
96
  if st.button("Reset", key="reset_button"):
97
+ st.session_state.job_id = None
98
+ st.session_state.job_status = None
99
  st.text_area("Enter your transactions (comma-separated)", value="", key="reset_area")