make dataset location an env variable

h2oai · Sep 21, 2022 · 0eac3fe · 0eac3fe
1 parent 3289617
commit 0eac3fe
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 3 deletions.
diff --git a/sales-forecasting/app.toml b/sales-forecasting/app.toml
@@ -1,6 +1,6 @@
 [App]
 Name = "ai.h2o.wave.walmart-forecasting"
-Version = "0.2.1"
+Version = "0.2.2"
 Title = "Sales Forecasting"
 Description = "Demo to explore demand and forecasts for supply chain optimization."
 LongDescription = "about.md"
@@ -19,3 +19,8 @@ SecretKey = "accessKey"
 Name = "AWS_SECRET_ACCESS_KEY"
 Secret = "h2o-s3-readonly"
 SecretKey = "secretKey"
+
+[[Env]]
+Name = "WALMART_SALES_FORECASTING_DATA_DIR"
+Secret = "h2o-s3-readonly"
+SecretKey = "walmartSalesForecastingDataDir"
diff --git a/sales-forecasting/wave_forecast.py b/sales-forecasting/wave_forecast.py
@@ -10,8 +10,21 @@
 
 
 # Inputs for the app, Should be read from a config file
-walmart_train_s3 = "s3://h2o-benchmark/walmart-sales-forecasting/walmart_train.csv"
-walmart_predictions_s3 = "s3://h2o-benchmark/walmart-sales-forecasting/walmart_test_preds.csv"
+
+# Valid values for WALMART_SALES_FORECASTING_DATA_DIR
+# "s3://h2o-benchmark/walmart-sales-forecasting"
+# "s3://h2o-benchmark"
+# "s3://h2o-benchmark/walmart-sales-forecasting/"
+# "s3://h2o-benchmark/"
+walmart_sales_forecasting_data_env = os.getenv("WALMART_SALES_FORECASTING_DATA_DIR")
+
+if not walmart_sales_forecasting_data_env:
+    sys.exit(1)
+
+walmart_sales_forecasting_data_dir = walmart_sales_forecasting_data_env.rstrip("/")
+walmart_train_s3 = walmart_sales_forecasting_data_dir + "/walmart_train.csv"
+walmart_predictions_s3 = walmart_sales_forecasting_data_dir + "/walmart_test_preds.csv"
+
 walmart_train = './walmart_train.csv'
 walmart_predictions = './walmart_test_preds.csv'